LlamaBatch.cc 60.3 KB
Newer Older
Li Zhang's avatar
Li Zhang committed
1
2
// Copyright (c) OpenMMLab. All rights reserved.

lvhan028's avatar
lvhan028 committed
3
4
#include "src/turbomind/models/llama/LlamaBatch.h"
#include "src/turbomind/kernels/decoding_kernels.h"
Li Zhang's avatar
Li Zhang committed
5
#include "src/turbomind/kernels/sampling_topk_kernels.h"
Chen Xin's avatar
Chen Xin committed
6
#include "src/turbomind/macro.h"
lvhan028's avatar
lvhan028 committed
7
8
9
#include "src/turbomind/models/llama/LlamaNcclGuard.h"
#include "src/turbomind/models/llama/LlamaV2.h"
#include "src/turbomind/models/llama/Request.h"
Li Zhang's avatar
Li Zhang committed
10
#include "src/turbomind/models/llama/SequenceManager.h"
11
#include "src/turbomind/models/llama/copy.h"
Li Zhang's avatar
Li Zhang committed
12
#include "src/turbomind/models/llama/llama_kernels.h"
lvhan028's avatar
lvhan028 committed
13
14
#include "src/turbomind/models/llama/llama_utils.h"
#include "src/turbomind/utils/Tensor.h"
Li Zhang's avatar
Li Zhang committed
15
16
17
#include "src/turbomind/utils/cuda_utils.h"
#include "src/turbomind/utils/debug_utils.h"
#include "src/turbomind/utils/gemm_test/gemm_func.h"
lvhan028's avatar
lvhan028 committed
18
#include "src/turbomind/utils/logger.h"
Li Zhang's avatar
Li Zhang committed
19
20
#include <algorithm>
#include <cmath>
Li Zhang's avatar
Li Zhang committed
21
#include <cstddef>
Li Zhang's avatar
Li Zhang committed
22
#include <cstdint>
23
#include <functional>
Li Zhang's avatar
Li Zhang committed
24
#include <iomanip>
Li Zhang's avatar
Li Zhang committed
25
#include <iterator>
Li Zhang's avatar
Li Zhang committed
26
27
#include <mutex>
#include <numeric>
Li Zhang's avatar
Li Zhang committed
28
29
#include <sstream>
#include <unordered_map>
Li Zhang's avatar
Li Zhang committed
30
#include <utility>
Li Zhang's avatar
Li Zhang committed
31

lvhan028's avatar
lvhan028 committed
32
namespace turbomind {
Li Zhang's avatar
Li Zhang committed
33

34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
void PrintDecodeTokens(
    const int* token_ids, int max_seq_len, int batch_sizse, cudaStream_t stream, const std::string& msg)
{
    // tokens in [S, B] layout
    std::vector<int> tokens(max_seq_len * batch_sizse);
    check_cuda_error(cudaMemcpyAsync(tokens.data(), token_ids, sizeof(int) * tokens.size(), cudaMemcpyDefault, stream));
    check_cuda_error(cudaStreamSynchronize(stream));

    printf("[%s] ", msg.c_str());
    for (int j = 0; j < max_seq_len; ++j) {
        printf("%5d ", j);
    }
    printf("\n");
    for (int i = 0; i < batch_sizse; ++i) {
        printf("[%s] ", msg.c_str());
        for (int j = 0; j < max_seq_len; ++j) {
            // std::cout << sb_tokens[j * batch_size + i] << " ";
            printf("%5d ", tokens[j * batch_sizse + i]);
        }
        printf("\n");
    }
}
Li Zhang's avatar
Li Zhang committed
56
57
58
59
60
61
62
void ClearState(BatchState& s)
{
    std::fill_n(s.requests.begin(), s.size, nullptr);
    std::fill_n(s.sequences.begin(), s.size, nullptr);
    s.size = s.active_size = 0;
}

Li Zhang's avatar
Li Zhang committed
63
template<typename T>
Li Zhang's avatar
Li Zhang committed
64
void LlamaBatch<T>::RejectInvalidRequests(Requests& stop_reqs, Requests& infer_reqs)
Li Zhang's avatar
Li Zhang committed
65
{
AllentDan's avatar
AllentDan committed
66
    std::unordered_map<uint64_t, int> occurrence;
Li Zhang's avatar
Li Zhang committed
67

Li Zhang's avatar
Li Zhang committed
68
    auto count_occurrence = [&occurrence](const Requests& rs) {
Li Zhang's avatar
Li Zhang committed
69
        for (const auto& r : rs) {
AllentDan's avatar
AllentDan committed
70
            ++occurrence[r->id];
Li Zhang's avatar
Li Zhang committed
71
72
73
        }
    };

Li Zhang's avatar
Li Zhang committed
74
75
76
    auto reject = [](const char* type, std::shared_ptr<Request>& req, int ec) {
        TM_LOG_WARNING(
            "[RejectInvalidRequests] Skipping invalid %s request for id %ld, code = %d", type, (long)req->id, ec);
Li Zhang's avatar
Li Zhang committed
77
78
79
80
        req->signal.set_value(ec);
        req.reset();
    };

Li Zhang's avatar
Li Zhang committed
81
    auto handle_conflict_or_invalid = [this, &occurrence, &reject](Requests& rs, const char* type) {
Li Zhang's avatar
Li Zhang committed
82
83
84
85
        for (auto& r : rs) {
            if (r) {
                int ec = 0;

Li Zhang's avatar
Li Zhang committed
86
87
88
89
90
                const int  input_length = r->inputs[rank_].getVal<int>("input_lengths", 0);
                const auto get_offset   = [&](int token_count) {
                    return std::max(0, std::min(token_count, r->inputs[rank_].getVal<int>("step", token_count)));
                };

AllentDan's avatar
AllentDan committed
91
                if (occurrence[r->id] != 1) {
Li Zhang's avatar
Li Zhang committed
92
93
94
95
96
                    ec = Request::kConflict;
                }
                else if (r->start_flag && r->stop_flag) {
                    ec = Request::kInvalid;
                }
Li Zhang's avatar
Li Zhang committed
97
98
99
100
101
102
103
104
105
106
                else if (input_length > session_len_) {
                    ec = Request::kTooLong;
                }
                else if (!r->start_flag) {
                    if (auto seq = sequence_manager_->Get(r->id); seq == nullptr) {
                        ec = Request::kInvalid;
                    }
                    else if (get_offset(seq->tokens.size()) + input_length > session_len_) {
                        ec = Request::kTooLong;
                    }
Li Zhang's avatar
Li Zhang committed
107
108
109
                }

                if (ec) {
Li Zhang's avatar
Li Zhang committed
110
                    reject(type, r, ec);
Li Zhang's avatar
Li Zhang committed
111
112
113
114
115
                }
            }
        }
    };

Li Zhang's avatar
Li Zhang committed
116
    auto drop_invalid = [](Requests& rs) {
Li Zhang's avatar
Li Zhang committed
117
118
119
120
121
122
123
124
125
        int count = 0;
        for (int i = 0; i < rs.size(); ++i) {
            if (rs[i]) {
                rs[count++] = std::move(rs[i]);
            }
        }
        rs.resize(count);
    };

AllentDan's avatar
AllentDan committed
126
127
    count_occurrence(stop_reqs);
    count_occurrence(infer_reqs);
Li Zhang's avatar
Li Zhang committed
128
129
130
131
132
133
134
135

    if (!stop_reqs.empty()) {
        handle_conflict_or_invalid(stop_reqs, "stop");

        // invalidate stop-only requests for inactive sequences
        for (auto& r : stop_reqs) {
            if (r && r->end_flag == false) {
                int ec = Request::kInactive;
Li Zhang's avatar
Li Zhang committed
136
137
                for (int i = 0; i < state_->size; ++i) {
                    if (state_->requests[i] && state_->requests[i]->id == r->id) {
Li Zhang's avatar
Li Zhang committed
138
139
140
141
142
                        ec = 0;
                        break;
                    }
                }
                if (ec) {
Li Zhang's avatar
Li Zhang committed
143
                    reject("stop", r, ec);
Li Zhang's avatar
Li Zhang committed
144
145
146
147
148
149
150
151
152
153
154
155
156
                }
            }
        }

        drop_invalid(stop_reqs);
    }

    if (!infer_reqs.empty()) {
        handle_conflict_or_invalid(infer_reqs, "infer");

        // invalidate requests for busy sequences
        for (auto& r : infer_reqs) {
            if (r) {
Li Zhang's avatar
Li Zhang committed
157
158
159
                for (int i = 0; i < state_->size; ++i) {
                    if (state_->requests[i] && state_->requests[i]->id == r->id) {
                        reject("infer", r, Request::kBusy);
Li Zhang's avatar
Li Zhang committed
160
161
162
163
164
165
166
167
168
169
170
                        break;
                    }
                }
            }
        }

        drop_invalid(infer_reqs);
    }
}

template<typename T>
Li Zhang's avatar
Li Zhang committed
171
auto LlamaBatch<T>::ProcessStopRequests(const Requests& requests) -> std::vector<Signal>
Li Zhang's avatar
Li Zhang committed
172
{
Li Zhang's avatar
Li Zhang committed
173
    NvtxScope           scope("stop_request");
Li Zhang's avatar
Li Zhang committed
174
    std::vector<Signal> signals;
Li Zhang's avatar
Li Zhang committed
175
    int                 count = 0;
Li Zhang's avatar
Li Zhang committed
176
177
178
    for (const auto& r : requests) {
        int ec = Request::kFail;
        // find matching active sequence
Li Zhang's avatar
Li Zhang committed
179
        for (int i = 0; i < state_->size; ++i) {
Li Zhang's avatar
Li Zhang committed
180
            // stop & optionally erase active sequence
Li Zhang's avatar
Li Zhang committed
181
            if (state_->requests[i] && state_->requests[i]->id == r->id) {
Li Zhang's avatar
Li Zhang committed
182
                ec = 0;
Li Zhang's avatar
Li Zhang committed
183
184
                signals.push_back(Interrupt(i, true, r->end_flag));
                ++count;
Li Zhang's avatar
Li Zhang committed
185
186
187
                break;
            }
        }
Li Zhang's avatar
Li Zhang committed
188
        // mismatch, try erase inactive sequence, in this case there is no active request to interrupt
Li Zhang's avatar
Li Zhang committed
189
        if (ec && r->end_flag) {
Li Zhang's avatar
Li Zhang committed
190
191
192
            if (sequence_manager_->Erase(r->id)) {
                ec = 0;
            }
Li Zhang's avatar
Li Zhang committed
193
        }
Li Zhang's avatar
Li Zhang committed
194
        signals.push_back([=] {
Li Zhang's avatar
Li Zhang committed
195
            if (rank_ == 0) {
Li Zhang's avatar
Li Zhang committed
196
                r->signal.set_value(ec);
Li Zhang's avatar
Li Zhang committed
197
            }
Li Zhang's avatar
Li Zhang committed
198
199
200
201
        });
    }
    if (count) {
        check_cuda_error(cudaStreamSynchronize(stream_));
Li Zhang's avatar
Li Zhang committed
202
203
204
    }
    return signals;
}
akhoroshev's avatar
akhoroshev committed
205

Li Zhang's avatar
Li Zhang committed
206
207
208
template<typename T>
void LlamaBatch<T>::ProcessInferRequests(const Requests& requests)
{
Li Zhang's avatar
Li Zhang committed
209
210
    NvtxScope scope("infer_request");
    auto&     state = *incoming_;
Li Zhang's avatar
Li Zhang committed
211
212
213
214

    FT_CHECK(state.size == 0);
    FT_CHECK(state.active_size == 0);

Li Zhang's avatar
Li Zhang committed
215
    std::vector<int> existing_idx;
Li Zhang's avatar
Li Zhang committed
216

Li Zhang's avatar
Li Zhang committed
217
218
219
    int idx = 0;
    for (const auto& r : requests) {
        FT_CHECK(!state.requests[idx]);
Li Zhang's avatar
Li Zhang committed
220

Li Zhang's avatar
Li Zhang committed
221
222
223
        if (rank_ == 0) {
            TM_LOG_WARNING("[ProcessInferRequests] Request for %ld received.", (long)r->id);
        }
Li Zhang's avatar
Li Zhang committed
224

Li Zhang's avatar
Li Zhang committed
225
        state.requests[idx] = r;
Li Zhang's avatar
Li Zhang committed
226
227

        // get sequence for the request
Li Zhang's avatar
Li Zhang committed
228
229
        state.sequences[idx] = r->start_flag ? sequence_manager_->Create(r->id) : sequence_manager_->Get(r->id);
        FT_CHECK(state.sequences[idx]);
Li Zhang's avatar
Li Zhang committed
230

Li Zhang's avatar
Li Zhang committed
231
        auto& seq = *state.sequences[idx];
Li Zhang's avatar
Li Zhang committed
232
233
234
235
236
237
238
239
240
241

        if (int step = r->inputs[rank_].getVal<int>("step", -1); step >= 0) {
            if (step <= seq.tokens.size()) {
                seq.tokens.resize(step);
                seq.cache_len = std::min(seq.cache_len, step);
            }
            else if (rank_ == 0) {
                TM_LOG_WARNING(
                    "[ProcessInferRequests] Skipping invalid step (%d) setting for ID %ld", step, (long)seq.id);
            }
Li Zhang's avatar
Li Zhang committed
242
        }
Li Zhang's avatar
Li Zhang committed
243
244
245
246
247

        const int  input_length = r->inputs[rank_].getVal<int>("input_lengths");
        const int* input_ids    = r->inputs[rank_].getPtr<int>("input_ids");

        // `output_ids` contains all token ids of the sequences
Li Zhang's avatar
Li Zhang committed
248
        const auto output_ids_base = state.output_ids + session_len_ * idx;
Li Zhang's avatar
Li Zhang committed
249
250
251
252
253
254
255
256
257
258
259
260
261
        auto       output_ids      = output_ids_base;

        // copy history tokens
        if (!seq.tokens.empty()) {
            output_ids = Copy(seq.tokens.data(), seq.tokens.size(), output_ids);
        }

        // copy input tokens
        if (input_length) {
            output_ids = Copy(input_ids, input_length, output_ids);
        }

        // total context length (history + input)
Li Zhang's avatar
Li Zhang committed
262
263
        state.h_context_length[idx] = output_ids - output_ids_base;
        state.h_finished[idx]       = false;
Li Zhang's avatar
Li Zhang committed
264

Li Zhang's avatar
Li Zhang committed
265
266
        const int request_output_len = state.requests[idx]->inputs[rank_].getVal<int>("request_output_len");
        state.seq_len_limit[idx]     = state.h_context_length[idx] + request_output_len;
Li Zhang's avatar
Li Zhang committed
267
268
        // `length_criterion` sets finish flag when step >= seq_limit_len, however when step == seq_limit_len
        // the actual sequence length is seq_limit_len + 1, hence seq_limit_len must truncated to session_len - 1
Li Zhang's avatar
Li Zhang committed
269
270
        if (state.seq_len_limit[idx] >= session_len_) {
            state.seq_len_limit[idx] = session_len_ - 1;
Li Zhang's avatar
Li Zhang committed
271
            if (rank_ == 0) {
Li Zhang's avatar
Li Zhang committed
272
                const int trunc_output_len = state.seq_len_limit[idx] - state.h_context_length[idx];
Li Zhang's avatar
Li Zhang committed
273
274
275
                TM_LOG_WARNING(
                    "[ProcessInferRequests] [%ld] total sequence length (%d + %d) exceeds `session_len` (%d), `request_output_len` is truncated to %d",
                    (long)seq.id,
Li Zhang's avatar
Li Zhang committed
276
                    state.h_context_length[idx],
Li Zhang's avatar
Li Zhang committed
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
                    request_output_len,
                    (int)session_len_,
                    trunc_output_len);
            }
        }

        // compute rope scaling factor
        if (r->start_flag) {
            seq.rope_theta      = model_->attn_params_.rotary_embedding_base;
            auto scaling_factor = 1.f;
            if (r->inputs[rank_].isExist("rope_scaling_factor")) {  // runtime scaling factor
                scaling_factor = r->inputs[rank_].getVal<float>("rope_scaling_factor");
            }
            else if (model_->attn_params_.rope_scaling_factor >= 1.f) {  // infer by `seq_len_limit`
                scaling_factor   = model_->attn_params_.rope_scaling_factor;
Li Zhang's avatar
Li Zhang committed
292
                auto max_seq_len = state.seq_len_limit[idx];
Li Zhang's avatar
Li Zhang committed
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
                auto max_pos_emb = model_->attn_params_.max_position_embeddings;
                if (max_seq_len > max_pos_emb) {
                    scaling_factor = scaling_factor * max_seq_len / max_pos_emb - (scaling_factor - 1);
                    // scaling_factor = std::max(exp2f(ceilf(log2f((float)max_seq_len / max_pos_emb) + 1.f))
                    // - 1.f, 1.f);
                }
            }
            if (scaling_factor != 1.f) {
                float rope_dim = model_->attn_params_.rotary_embedding_dim;
                seq.rope_theta *= powf(scaling_factor, rope_dim / (rope_dim - 2.f));
                TM_LOG_INFO("[ProcessInferRequests] %ld rope_scaling_factor: %f, rope_theta = %f",
                            (long)seq.id,
                            scaling_factor,
                            seq.rope_theta);
            }
        }
Li Zhang's avatar
Li Zhang committed
309
        state.h_rope_theta[idx] = seq.rope_theta;
Li Zhang's avatar
Li Zhang committed
310

Li Zhang's avatar
Li Zhang committed
311
312
313
314
315
316
317
318
        if (r->start_flag) {
            // prepare to initialize random state for new sequence
            h_random_seed_[idx] = r->inputs[rank_].getVal<unsigned long long>("random_seed", 0);
        }
        else {
            // Recover device states if not a new sequence
            h_curand_state_[existing_idx.size()] = *(curandState_t*)seq.random_state.data();
            existing_idx.push_back(idx);
Li Zhang's avatar
Li Zhang committed
319
320
        }

Li Zhang's avatar
Li Zhang committed
321
        // ! SHARED STATE IS MODIFIED, BARRIER SYNCHRONIZATION REQUIRED
Li Zhang's avatar
Li Zhang committed
322
        // assign priority based on arrival time
Li Zhang's avatar
Li Zhang committed
323
        if (rank_ == 0) {
324
            r->unique_id = request_count_++;
Li Zhang's avatar
Li Zhang committed
325
        }
Li Zhang's avatar
Li Zhang committed
326
327

        // increment pointer
Li Zhang's avatar
Li Zhang committed
328
        idx++;
Li Zhang's avatar
Li Zhang committed
329
    }
Li Zhang's avatar
Li Zhang committed
330

Li Zhang's avatar
Li Zhang committed
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
    state.size = idx;

    // when there are new sequences
    if (state.size != existing_idx.size()) {
        // copy random seeds to device
        Copy(h_random_seed_, state.size, d_random_seed_);
        // initialize random states
        invokeCurandBatchInitialize(state.curand_state, state.size, d_random_seed_, stream_);
        sync_check_cuda_error();
    }

    if (!existing_idx.empty()) {
        // copy existing curand states to device
        Copy(h_curand_state_, existing_idx.size(), d_curand_state_);
        // insert the states to their correct positions in the batch
        IndexedCopy({}, existing_idx, std::tuple{d_curand_state_, state.curand_state, 1});
    }
Li Zhang's avatar
Li Zhang committed
348
349
350
}

template<typename T>
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
void LlamaBatch<T>::AdjustMaxInputCount(GenerationState&                    g,
                                        const std::vector<const Sequence*>& sequences,
                                        const std::vector<int>&             context_length)
{
    int input_count = 0;
    for (int i = 0; i < sequences.size(); ++i) {
        input_count += context_length[i] - sequences[i]->cache_len;
    }
    const int batch_size = sequences.size();
    input_count -= batch_size;

    // min tokens per iter for satisfying max prefill iters constraint
    input_count = (input_count + max_prefill_iters_ - 1) / max_prefill_iters_;

    if (g.min_input_count.empty()) {
        g.min_input_count.resize(max_prefill_iters_);
    }
    g.min_input_count.pop_front();
    g.min_input_count.push_back(input_count);
    /// TODO: sub-optimal when there are inactive sequences due to memory constraint
    for (auto& x : g.min_input_count) {
        x = std::max(x, input_count);
    }

    input_count = std::max(g.min_input_count.front() + batch_size, num_tokens_per_iter_);
    input_count = std::min(input_count, max_context_token_num_);
    // update max input count
    g.max_input_count1 = input_count;
    g.max_input_count2 = std::min(input_count + extra_tokens_per_iter_, max_context_token_num_);
}

template<typename T>
void LlamaBatch<T>::Initialize(GenerationState& g)
Li Zhang's avatar
Li Zhang committed
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
{
    NvtxScope                                scope("initialize");
    std::vector<const Sequence*>             sequences;
    std::vector<Sequence::Status>            status;
    std::vector<uint64_t>                    priorities;
    std::vector<int>                         context_lengths;
    std::vector<std::pair<BatchState*, int>> coords;

    // count the holes introduced by finished requests in from previous iteration or stop requests from
    // current iteration
    int holes{};
    int active_holes{};
    for (int i = 0; i < state_->size; ++i) {
        if (!state_->requests[i]) {
            ++holes;
            if (i < state_->active_size) {
                ++active_holes;
            }
        }
    }

    auto process = [&](BatchState* state) {
        for (int i = 0; i < state->size; ++i) {
            if (auto& r = state->requests[i]) {
                sequences.push_back(state->sequences[i]);
                status.push_back(state->sequences[i]->status);
410
                priorities.push_back(r->unique_id);
Li Zhang's avatar
Li Zhang committed
411
412
413
414
415
416
417
418
419
                context_lengths.push_back(state->h_context_length[i]);
                coords.emplace_back(state, i);
            }
        }
    };

    process(state_);
    process(incoming_);

420
421
422
423
424
425
426
427
    auto adjust = [this, &g](const Sequences&        sequences,
                             const std::vector<int>& context_length) -> std::pair<int, int> {
        AdjustMaxInputCount(g, sequences, context_length);
        return {g.max_input_count1, g.max_input_count2};
    };

    // TM_LOG_INFO("max_input_count %d", max_input_count);
    auto outcome = sequence_manager_->Materialize(sequences, context_lengths, priorities, step_length_, adjust);
Li Zhang's avatar
Li Zhang committed
428
429
430
431
432
433
434
435
436
437
438
439
440

    if (outcome.allocation || outcome.swap_in || outcome.swap_out) {
        dbg(outcome);
    }

    bool exchange = outcome.swap_in + outcome.swap_out > 0;

    std::vector<int> idxs(sequences.size());
    std::iota(idxs.begin(), idxs.end(), 0);

    if (exchange || holes || incoming_->size) {
        // put active ones first
        auto active_end = std::stable_partition(idxs.begin(), idxs.end(), [&](int idx) {
441
            return sequences[idx]->status == Sequence::kActive;  // current status
Li Zhang's avatar
Li Zhang committed
442
443
444
        });

        // all blocks are not enough to hold a single sequence
Li Zhang's avatar
Li Zhang committed
445
446
447
        if (!sequences.empty()) {
            FT_CHECK_WITH_INFO(active_end != idxs.begin(), "No enough blocks.");
        }
Li Zhang's avatar
Li Zhang committed
448

449
450
451
        // move the partial seq to the back
        auto partial_beg = std::stable_partition(idxs.begin(), active_end, [&](int i) {
            return sequences[i]->cache_len + sequences[i]->input_length == context_lengths[i];
Li Zhang's avatar
Li Zhang committed
452
        });
453
        FT_CHECK(active_end - partial_beg <= 1);
Li Zhang's avatar
Li Zhang committed
454

455
456
457
458
459
460
461
462
463
        auto swapin_beg = std::stable_partition(idxs.begin(), partial_beg, [&](int i) {
            return status[i] == Sequence::kActive;  // past status
        });

        // sort swap-ins according to input length
        if (swapin_beg != partial_beg) {
            std::stable_sort(swapin_beg, partial_beg, [&](int i, int j) {
                return sequences[i]->input_length < sequences[j]->input_length;
            });
Li Zhang's avatar
Li Zhang committed
464
465
466
467
        }

        // Copy sequence states to back buffer
        FT_CHECK(back_->size == 0 && back_->active_size == 0);
Li Zhang's avatar
Li Zhang committed
468
        std::vector<std::tuple<BatchState*, BatchState*, int, int>> cpys;
Li Zhang's avatar
Li Zhang committed
469
470
471
472
473
        for (const auto& i : idxs) {
            auto& s = *sequences[i];
            if (s.status == Sequence::kActive) {
                ++back_->active_size;
            }
Li Zhang's avatar
Li Zhang committed
474
            cpys.emplace_back(coords[i].first, back_, coords[i].second, back_->size++);
Li Zhang's avatar
Li Zhang committed
475
        }
Li Zhang's avatar
Li Zhang committed
476
        CopyState(cpys);
Li Zhang's avatar
Li Zhang committed
477
478
479
480
481
482
483
        // Swap the buffers
        std::swap(state_, back_);

        ClearState(*back_);
        ClearState(*incoming_);
    }

Li Zhang's avatar
Li Zhang committed
484
485
    FT_CHECK(state_->size <= max_batch_size_);

Li Zhang's avatar
Li Zhang committed
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
    /// Update block ptrs when there were
    //  1. swap-in or swap-out
    //  2. holes in the active buffer
    //  3. new allocations (for existing active sequences)
    if (exchange || active_holes || outcome.allocation) {
        // Prepare intermediate buffers
        h_cu_block_counts_[0] = 0;

        auto k_ptrs = h_k_block_ptrs_;
        auto v_ptrs = h_v_block_ptrs_;

        const int batch_size = state_->active_size;

        for (int i = 0; i < batch_size; ++i) {
            const auto& seq = *state_->sequences[i];

            // cumulative num of blocks
            h_cu_block_counts_[i + 1] = h_cu_block_counts_[i] + seq.blocks.size();

Li Zhang's avatar
Li Zhang committed
505
506
507
            FT_CHECK_WITH_INFO(h_cu_block_counts_[i + 1] <= sequence_manager_->max_block_count(),
                               std::to_string(h_cu_block_counts_[i + 1]));

Li Zhang's avatar
Li Zhang committed
508
509
            k_ptrs = std::transform(seq.blocks.cbegin(), seq.blocks.cend(), k_ptrs, [&](int block_id) {
                return reinterpret_cast<uintptr_t>(sequence_manager_->GetKeyPtr(block_id));
Li Zhang's avatar
Li Zhang committed
510
            });
Li Zhang's avatar
Li Zhang committed
511
512
            v_ptrs = std::transform(seq.blocks.cbegin(), seq.blocks.cend(), v_ptrs, [&](int block_id) {
                return reinterpret_cast<uintptr_t>(sequence_manager_->GetValPtr(block_id));
Li Zhang's avatar
Li Zhang committed
513
514
515
516
517
518
519
520
521
522
            });
        }

        static_assert(sizeof(uintptr_t) == sizeof(void*));

        Copy(h_cu_block_counts_, batch_size + 1, cu_block_counts_);
        Copy(h_k_block_ptrs_, h_cu_block_counts_[batch_size], k_block_ptrs_);
        Copy(h_v_block_ptrs_, h_cu_block_counts_[batch_size], v_block_ptrs_);
    }

523
524
525
526
527
    const int batch_size = state_->active_size;

    // check if the last sequence is partial
    int partial     = 0;
    int partial_len = -1;
Li Zhang's avatar
Li Zhang committed
528
    if (state_->active_size) {
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
        const int i = state_->active_size - 1;
        partial = state_->sequences[i]->cache_len + state_->sequences[i]->input_length != state_->h_context_length[i];
        if (partial) {
            // backup full context length of partial
            partial_len = state_->h_context_length[i];
            // replace with partial context length
            state_->h_context_length[i] = state_->sequences[i]->cache_len + state_->sequences[i]->input_length;
        }
    }

    const int max_context_len = *std::max_element(state_->h_context_length, state_->h_context_length + batch_size);

    std::vector<uint64_t> unique_ids(batch_size);
    for (int i = 0; i < batch_size; ++i) {
        unique_ids[i] = state_->requests[i]->unique_id;
    }

    // Real-time context length that will change during generation
    Copy(state_->h_context_length, batch_size, context_length_buf_);
    Copy(state_->h_finished, batch_size, finished_buf_);
    Copy(state_->h_rope_theta, batch_size, rope_theta_);

    // used for dispatching split-k decoding kernels
    const int sum_seq_len =
        std::accumulate(state_->h_context_length, state_->h_context_length + batch_size, -batch_size);
    const int max_seq_len = *std::max_element(state_->h_context_length, state_->h_context_length + batch_size) - 1;

    // TM_LOG_INFO(
    //     "[init] batch_size = %d, max_ctx_len = %d, partial = %d", (int)batch_size, (int)max_context_len, partial);

    bool skip_init_sampling = std::equal(g.unique_ids.begin(),  //
                                         g.unique_ids.end() - g.partial,
                                         unique_ids.begin(),
                                         unique_ids.end() - partial);

    g.sum_seq_len            = sum_seq_len;
    g.max_seq_len            = max_seq_len;
    g.partial                = partial;
    g.partial_context_legnth = partial_len;
    g.unique_ids             = std::move(unique_ids);
    g.finished_count         = 0;

    if (!skip_init_sampling) {
        g.max_init_ctx_len = max_context_len;
        g.step             = max_context_len;
        InitializeSampling(g);
    }
Li Zhang's avatar
Li Zhang committed
576
577
578
}

template<typename T>
Li Zhang's avatar
Li Zhang committed
579
void LlamaBatch<T>::CopyState(const std::vector<std::tuple<BatchState*, BatchState*, int, int>>& desc)
Li Zhang's avatar
Li Zhang committed
580
{
Li Zhang's avatar
Li Zhang committed
581
582
583
584
    if (desc.empty()) {
        return;
    }

Li Zhang's avatar
Li Zhang committed
585
586
    std::vector<int> idxs(desc.size());
    std::iota(idxs.begin(), idxs.end(), 0);
Li Zhang's avatar
Li Zhang committed
587

Li Zhang's avatar
Li Zhang committed
588
    std::sort(idxs.begin(), idxs.end(), [&](int i, int j) { return desc[i] < desc[j]; });
Li Zhang's avatar
Li Zhang committed
589

Li Zhang's avatar
Li Zhang committed
590
591
592
    auto get_signature = [&](int i) -> std::pair<BatchState*, BatchState*> {
        return std::make_pair(std::get<0>(desc[idxs[i]]), std::get<1>(desc[idxs[i]]));
    };
Li Zhang's avatar
Li Zhang committed
593

Li Zhang's avatar
Li Zhang committed
594
595
596
597
598
599
600
601
602
603
    std::vector<int> offsets;
    auto             current = get_signature(0);
    offsets.push_back(0);
    for (int i = 0; i < idxs.size(); ++i) {
        if (auto signature = get_signature(i); signature != current) {
            current = signature;
            offsets.push_back(i);
        }
    }
    offsets.push_back(idxs.size());
Li Zhang's avatar
Li Zhang committed
604

Li Zhang's avatar
Li Zhang committed
605
606
607
    for (int bi = 1; bi < offsets.size(); ++bi) {
        int beg = offsets[bi - 1];
        int end = offsets[bi];
Li Zhang's avatar
Li Zhang committed
608

Li Zhang's avatar
Li Zhang committed
609
610
611
        if (beg == end) {
            continue;
        }
Li Zhang's avatar
Li Zhang committed
612

Li Zhang's avatar
Li Zhang committed
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
        auto [s, d] = get_signature(beg);

        std::vector<int> s_idx;
        std::vector<int> d_idx;
        for (int i = beg; i < end; ++i) {
            s_idx.push_back(std::get<2>(desc[idxs[i]]));
            d_idx.push_back(std::get<3>(desc[idxs[i]]));
        }

        IndexedCopy(s_idx,
                    d_idx,
                    std::tuple{s->output_ids, d->output_ids, session_len_},
                    std::tuple{s->curand_state, d->curand_state, 1});
    }

    for (const auto& [s, d, si, di] : desc) {
        d->h_context_length[di] = s->h_context_length[si];
        d->h_finished[di]       = s->h_finished[si];
        d->h_rope_theta[di]     = s->h_rope_theta[si];
        d->seq_len_limit[di]    = s->seq_len_limit[si];
        d->sequences[di]        = s->sequences[si];
        d->requests[di]         = s->requests[si];
    }
Li Zhang's avatar
Li Zhang committed
636
637
638
639
}

template<typename T>
void LlamaBatch<T>::AllocateBuffer(size_t batch_size, size_t session_len)
Li Zhang's avatar
Li Zhang committed
640
{
lvhan028's avatar
lvhan028 committed
641
    TM_LOG_DEBUG(__PRETTY_FUNCTION__);
Li Zhang's avatar
Li Zhang committed
642
643
    const size_t batchxbeam = batch_size;

Li Zhang's avatar
Li Zhang committed
644
645
646
647
648
649
    const size_t hidden_units      = model_->hidden_units_;
    const size_t vocab_size        = model_->vocab_size_padded_;
    const size_t head_dim          = model_->size_per_head_;
    const size_t local_kv_head_num = model_->local_kv_head_num_;
    // +1 padding, BlockIterator does not use predicate
    const size_t max_block_count = sequence_manager_->max_block_count() + 1;
Li Zhang's avatar
Li Zhang committed
650
651
652

    context_decoder_input_buf_ =
        (T*)allocator_->reMalloc(context_decoder_input_buf_, sizeof(T) * max_context_token_num_ * hidden_units, false);
653
654
    context_decoder_output_buf_ =
        (T*)allocator_->reMalloc(context_decoder_output_buf_, sizeof(T) * max_context_token_num_ * hidden_units, false);
Li Zhang's avatar
Li Zhang committed
655
656
657
    context_decoder_ids_buf_ =
        (int*)allocator_->reMalloc(context_decoder_ids_buf_, sizeof(int) * max_context_token_num_, false);

Li Zhang's avatar
Li Zhang committed
658
659
660
661
662
663
664
665
    tmp_k_cache_buf_ = (T*)allocator_->reMalloc(
        tmp_k_cache_buf_, sizeof(T) * max_context_token_num_ * local_kv_head_num * head_dim, false);
    tmp_v_cache_buf_ = (T*)allocator_->reMalloc(
        tmp_v_cache_buf_, sizeof(T) * max_context_token_num_ * local_kv_head_num * head_dim, false);

    tmp_k_ptrs_ = (void**)allocator_->reMalloc(tmp_k_ptrs_, sizeof(void*) * batch_size, false);
    tmp_v_ptrs_ = (void**)allocator_->reMalloc(tmp_v_ptrs_, sizeof(void*) * batch_size, false);

Li Zhang's avatar
Li Zhang committed
666
667
668
    decoder_input_buf_  = (T*)allocator_->reMalloc(decoder_input_buf_, sizeof(T) * batchxbeam * hidden_units, false);
    decoder_output_buf_ = (T*)allocator_->reMalloc(decoder_output_buf_, sizeof(T) * batchxbeam * hidden_units, false);

669
670
671
672
    input_ids_buf_       = (int*)allocator_->reMalloc(input_ids_buf_, sizeof(int) * batchxbeam * session_len, true);
    input_length_buf_    = (int*)allocator_->reMalloc(input_length_buf_, sizeof(int) * batchxbeam);
    context_length_buf_  = (int*)allocator_->reMalloc(context_length_buf_, sizeof(int) * batchxbeam);
    init_context_length_ = (int*)allocator_->reMalloc(init_context_length_, sizeof(int) * batchxbeam);
Li Zhang's avatar
Li Zhang committed
673

Li Zhang's avatar
Li Zhang committed
674
    sequence_lengths_ = (int*)allocator_->reMalloc(sequence_lengths_, sizeof(int) * batchxbeam, false);
Li Zhang's avatar
Li Zhang committed
675

Li Zhang's avatar
Li Zhang committed
676
677
678
    cu_block_counts_ = (int*)allocator_->reMalloc(cu_block_counts_, sizeof(int) * (batch_size + 1));
    k_block_ptrs_    = (uintptr_t*)allocator_->reMalloc(k_block_ptrs_, sizeof(uintptr_t) * max_block_count);
    v_block_ptrs_    = (uintptr_t*)allocator_->reMalloc(v_block_ptrs_, sizeof(uintptr_t) * max_block_count);
Li Zhang's avatar
Li Zhang committed
679
680
681
682
683
684
685
686
687

    logits_buf_       = (float*)allocator_->reMalloc(logits_buf_, sizeof(float) * batchxbeam * vocab_size, false);
    local_logits_buf_ = (float*)allocator_->reMalloc(local_logits_buf_, sizeof(float) * batchxbeam * vocab_size, false);

    token_ids_buf_ = (int*)allocator_->reMalloc(token_ids_buf_, sizeof(int) * batchxbeam * session_len * 2, true);

    finished_buf_  = (bool*)allocator_->reMalloc(finished_buf_, sizeof(bool) * batchxbeam, false);
    seq_limit_len_ = (uint32_t*)allocator_->reMalloc(seq_limit_len_, sizeof(uint32_t) * batch_size, false);

Li Zhang's avatar
Li Zhang committed
688
689
    rope_theta_ = (float*)allocator_->reMalloc(rope_theta_, sizeof(float) * batch_size, false);

Li Zhang's avatar
Li Zhang committed
690
691
692
693
    is_allocate_buffer_ = true;
}

template<typename T>
Li Zhang's avatar
Li Zhang committed
694
void LlamaBatch<T>::AllocatePersistantBuffer(size_t max_batch_size)
Li Zhang's avatar
Li Zhang committed
695
{
Li Zhang's avatar
Li Zhang committed
696
697
698
699
700
701
    d_stop_words_ = (int*)allocator_->reMalloc(d_stop_words_, sizeof(int) * max_batch_size * kMaxStopBadWordsLen, true);
    d_bad_words_  = (int*)allocator_->reMalloc(d_bad_words_, sizeof(int) * max_batch_size * kMaxStopBadWordsLen, true);
    h_stop_words_ =
        (int*)allocator_->reMalloc(h_stop_words_, sizeof(int) * max_batch_size * kMaxStopBadWordsLen, true, true);
    h_bad_words_ =
        (int*)allocator_->reMalloc(h_bad_words_, sizeof(int) * max_batch_size * kMaxStopBadWordsLen, true, true);
Li Zhang's avatar
Li Zhang committed
702
703
704
705
706
707
708

    h_runtime_top_k_ = (int*)allocator_->reMalloc(h_runtime_top_k_, sizeof(int) * max_batch_size, true, true);
    h_runtime_top_p_ = (float*)allocator_->reMalloc(h_runtime_top_p_, sizeof(float) * max_batch_size, true, true);
    h_temperature_   = (float*)allocator_->reMalloc(h_temperature_, sizeof(float) * max_batch_size, true, true);
    h_repetition_penalty_ =
        (float*)allocator_->reMalloc(h_repetition_penalty_, sizeof(float) * max_batch_size, true, true);

Li Zhang's avatar
Li Zhang committed
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
    h_random_seed_ = (unsigned long long*)allocator_->reMalloc(
        h_random_seed_, sizeof(unsigned long long) * max_batch_size, true, true);
    d_random_seed_ = (unsigned long long*)allocator_->reMalloc(
        d_random_seed_, sizeof(unsigned long long) * max_batch_size, true, false);

    h_curand_state_ =
        (curandState_t*)allocator_->reMalloc(h_curand_state_, sizeof(curandState_t) * max_batch_size, true, true);
    d_curand_state_ =
        (curandState_t*)allocator_->reMalloc(d_curand_state_, sizeof(curandState_t) * max_batch_size, true, false);

    d_end_ids_buf_ = (int*)allocator_->reMalloc(d_end_ids_buf_, sizeof(int) * max_batch_size, false);
    h_end_ids_buf_ = (int*)allocator_->reMalloc(h_end_ids_buf_, sizeof(int) * max_batch_size, false, true);

    sampling_params_ = {
        {"stop_words_list", (std::byte*)h_stop_words_, (std::byte*)d_stop_words_},
        {"bad_words_list", (std::byte*)h_bad_words_, (std::byte*)d_bad_words_},
        {"runtime_top_k", (std::byte*)h_runtime_top_k_, nullptr},
        {"runtime_top_p", (std::byte*)h_runtime_top_p_, nullptr},
        {"temperature", (std::byte*)h_temperature_, nullptr},
        {"repetition_penalty", (std::byte*)h_repetition_penalty_, nullptr},
    };
Li Zhang's avatar
Li Zhang committed
730

Li Zhang's avatar
Li Zhang committed
731
732
    for (auto& s : states_) {
        s.output_ids = (int*)allocator_->reMalloc(s.output_ids, sizeof(int) * max_batch_size * session_len_, true);
Li Zhang's avatar
Li Zhang committed
733
734
        s.curand_state =
            (curandState_t*)allocator_->reMalloc(s.curand_state, sizeof(curandState_t) * max_batch_size, true);
Li Zhang's avatar
Li Zhang committed
735
736
737
    }

    const size_t max_block_count = sequence_manager_->max_block_count();
Li Zhang's avatar
Li Zhang committed
738
739

    {
Li Zhang's avatar
Li Zhang committed
740
        NcclGuard barrier(model_->tensor_para_, stream_, true);
Li Zhang's avatar
Li Zhang committed
741
742
743
744
        h_input_ids_buf_ =
            (int*)allocator_->reMalloc(h_input_ids_buf_, sizeof(int) * max_batch_size * session_len_, false, true);
        h_input_length_buf_ =
            (int*)allocator_->reMalloc(h_input_length_buf_, sizeof(int) * max_batch_size, false, true);
Li Zhang's avatar
Li Zhang committed
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762

        h_tmp_k_ptrs_ = (void**)allocator_->reMalloc(h_tmp_k_ptrs_, sizeof(void*) * max_batch_size, false, true);
        h_tmp_v_ptrs_ = (void**)allocator_->reMalloc(h_tmp_v_ptrs_, sizeof(void*) * max_batch_size, false, true);

        h_cu_block_counts_ =
            (int*)allocator_->reMalloc(h_cu_block_counts_, sizeof(int) * (max_batch_size + 1), false, true);
        h_k_block_ptrs_ =
            (uintptr_t*)allocator_->reMalloc(h_k_block_ptrs_, sizeof(uintptr_t) * max_block_count, false, true);
        h_v_block_ptrs_ =
            (uintptr_t*)allocator_->reMalloc(h_v_block_ptrs_, sizeof(uintptr_t) * max_block_count, false, true);

        for (auto& s : states_) {
            s.h_context_length =
                (int*)allocator_->reMalloc(s.h_context_length, sizeof(int) * max_batch_size, false, true);
            s.h_finished   = (bool*)allocator_->reMalloc(s.h_finished, sizeof(bool) * max_batch_size * 2, false, true);
            s.h_rope_theta = (float*)allocator_->reMalloc(s.h_rope_theta, sizeof(float) * max_batch_size, false, true);
        }

Li Zhang's avatar
Li Zhang committed
763
764
        h_seq_limit_len_ =
            (uint32_t*)allocator_->reMalloc(h_seq_limit_len_, sizeof(uint32_t) * max_batch_size, false, true);
Li Zhang's avatar
Li Zhang committed
765

Li Zhang's avatar
Li Zhang committed
766
767
        h_output_ids_ =
            (int*)allocator_->reMalloc(h_output_ids_, sizeof(int) * max_batch_size * session_len_, false, true);
Li Zhang's avatar
Li Zhang committed
768
769
770
771
772
773
    }

    is_allocate_persistant_buffer_ = true;
}

template<typename T>
Li Zhang's avatar
Li Zhang committed
774
void LlamaBatch<T>::FreeBuffer()
Li Zhang's avatar
Li Zhang committed
775
{
lvhan028's avatar
lvhan028 committed
776
    TM_LOG_DEBUG(__PRETTY_FUNCTION__);
Li Zhang's avatar
Li Zhang committed
777
778
    if (is_allocate_buffer_) {
        allocator_->free((void**)&context_decoder_input_buf_);
779
        allocator_->free((void**)&context_decoder_output_buf_);
Li Zhang's avatar
Li Zhang committed
780
781
        allocator_->free((void**)&context_decoder_ids_buf_);

Li Zhang's avatar
Li Zhang committed
782
783
784
785
786
        allocator_->free((void**)&tmp_k_cache_buf_);
        allocator_->free((void**)&tmp_v_cache_buf_);
        allocator_->free((void**)&tmp_k_ptrs_);
        allocator_->free((void**)&tmp_v_ptrs_);

Li Zhang's avatar
Li Zhang committed
787
788
789
790
791
792
        allocator_->free((void**)&decoder_input_buf_);
        allocator_->free((void**)&decoder_output_buf_);

        allocator_->free((void**)&input_ids_buf_);
        allocator_->free((void**)&input_length_buf_);
        allocator_->free((void**)&context_length_buf_);
793
        allocator_->free((void**)&init_context_length_);
Li Zhang's avatar
Li Zhang committed
794
795
796

        allocator_->free((void**)&sequence_lengths_);

Li Zhang's avatar
Li Zhang committed
797
798
799
        allocator_->free((void**)&cu_block_counts_);
        allocator_->free((void**)&k_block_ptrs_);
        allocator_->free((void**)&v_block_ptrs_);
Li Zhang's avatar
Li Zhang committed
800
801
802
803

        allocator_->free((void**)&logits_buf_);
        allocator_->free((void**)&local_logits_buf_);

804
805
806
807
808
809
810
        if (local_context_logits_buf_) {
            allocator_->free((void**)&local_context_logits_buf_);
        }
        if (context_logits_buf_) {
            allocator_->free((void**)&context_logits_buf_);
        }

Li Zhang's avatar
Li Zhang committed
811
812
        allocator_->free((void**)&token_ids_buf_);

Li Zhang's avatar
Li Zhang committed
813
814
815
        allocator_->free((void**)&d_end_ids_buf_);
        allocator_->free((void**)&h_end_ids_buf_, true);

Li Zhang's avatar
Li Zhang committed
816
817
818
        allocator_->free((void**)&finished_buf_);
        allocator_->free((void**)&seq_limit_len_);

Li Zhang's avatar
Li Zhang committed
819
820
        allocator_->free((void**)&rope_theta_);

Li Zhang's avatar
Li Zhang committed
821
822
823
824
        is_allocate_buffer_ = false;
    }

    if (is_allocate_persistant_buffer_) {
Li Zhang's avatar
Li Zhang committed
825
826
827
828
829
830
831
832
833
834

        allocator_->free((void**)&d_stop_words_);
        allocator_->free((void**)&h_stop_words_, true);
        allocator_->free((void**)&d_bad_words_);
        allocator_->free((void**)&h_bad_words_, true);
        allocator_->free((void**)&d_random_seed_);
        allocator_->free((void**)&h_random_seed_, true);
        allocator_->free((void**)&d_curand_state_);
        allocator_->free((void**)&h_curand_state_, true);

Li Zhang's avatar
Li Zhang committed
835
836
837
838
839
        for (auto& s : states_) {
            allocator_->free((void**)&s.h_context_length, true);
            allocator_->free((void**)&s.h_finished, true);
            allocator_->free((void**)&s.h_rope_theta, true);
            allocator_->free((void**)&s.output_ids);
Li Zhang's avatar
Li Zhang committed
840
            allocator_->free((void**)&s.curand_state);
Li Zhang's avatar
Li Zhang committed
841
842
843
844
845
846
        }
        allocator_->free((void**)&h_tmp_k_ptrs_, true);
        allocator_->free((void**)&h_tmp_v_ptrs_, true);
        allocator_->free((void**)&h_cu_block_counts_, true);
        allocator_->free((void**)&h_k_block_ptrs_, true);
        allocator_->free((void**)&h_v_block_ptrs_, true);
Li Zhang's avatar
Li Zhang committed
847
848
849
850
        allocator_->free((void**)&h_input_ids_buf_, true);
        allocator_->free((void**)&h_input_length_buf_, true);
        allocator_->free((void**)&h_seq_limit_len_, true);

Li Zhang's avatar
Li Zhang committed
851
852
        allocator_->free((void**)&h_output_ids_, true);

Li Zhang's avatar
Li Zhang committed
853
854
855
856
857
        is_allocate_persistant_buffer_ = false;
    }
}

template<typename T>
858
859
860
861
LlamaBatch<T>::LlamaBatch(const EngineParams& params, int cache_block_seq_len, int quant_policy, LlamaV2<T>* model):
    max_batch_size_(params.max_batch_size),
    max_context_token_num_(params.max_context_token_num),
    session_len_(params.session_len),
Li Zhang's avatar
Li Zhang committed
862
863
    rank_(model->tensor_para_.rank_),
    debug_(model->debug_),
864
    step_length_(params.step_length),
Li Zhang's avatar
Li Zhang committed
865
    model_(model),
866
867
868
869
    data_type_(getTensorType<T>()),
    num_tokens_per_iter_(params.num_tokens_per_iter),
    extra_tokens_per_iter_(params.extra_tokens_per_iter),
    max_prefill_iters_(params.max_prefill_iters)
Li Zhang's avatar
Li Zhang committed
870
{
Li Zhang's avatar
Li Zhang committed
871
872
873
874
    stream_         = model_->stream_;
    allocator_      = model_->allocator_;
    cublas_wrapper_ = model_->cublas_wrapper_;

875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
    const size_t elem_bits = (quant_policy & QuantPolicy::kCacheKVInt8) ? 8 : sizeof(T) * 8;

    sequence_manager_.reset(new SequenceManager{model_->num_layer_,
                                                model_->local_kv_head_num_,
                                                model_->size_per_head_,
                                                (size_t)cache_block_seq_len,
                                                params.cache_max_block_count,
                                                params.cache_chunk_size,
                                                elem_bits,
                                                model->tensor_para_.rank_,
                                                allocator_});

    const size_t max_session_len = sequence_manager_->max_block_count() * cache_block_seq_len;
    if (max_session_len < session_len_) {
        if (rank_ == 0) {
            TM_LOG_WARNING("No enough blocks for `session_len` (%d), `session_len` truncated to %d.",
                           session_len_,
                           max_session_len);
        }
        session_len_ = max_session_len;
    }

Li Zhang's avatar
Li Zhang committed
897
    for (auto& s : states_) {
898
899
900
        s.requests.resize(max_batch_size_);
        s.sequences.resize(max_batch_size_);
        s.seq_len_limit.resize(max_batch_size_);
Li Zhang's avatar
Li Zhang committed
901
    }
Li Zhang's avatar
Li Zhang committed
902

Li Zhang's avatar
Li Zhang committed
903
904
905
    state_    = &states_[0];
    back_     = &states_[1];
    incoming_ = &states_[2];
Li Zhang's avatar
Li Zhang committed
906

907
908
    AllocateBuffer(max_batch_size_, session_len_);
    AllocatePersistantBuffer(max_batch_size_);
Li Zhang's avatar
Li Zhang committed
909
910
911
}

template<typename T>
912
void LlamaBatch<T>::InitializeSampling(const GenerationState& g)
Li Zhang's avatar
Li Zhang committed
913
{
Li Zhang's avatar
Li Zhang committed
914
    NvtxScope _("InitSampling");
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
    const int batch_size = state_->active_size - g.partial;
    if (batch_size == 0) {
        return;
    }

    // Context length at initialization, will stay constant until re-initialziation
    Copy(context_length_buf_, batch_size, init_context_length_);

    Copy(context_length_buf_, batch_size, sequence_lengths_);
    // `sequence_lengths_` will be increased by dynamic decode
    // note that in decoder and in output "sequence length" has different semantic
    // - in decoder it means length of sequence that has kv cache already computed
    // - in output it means length of all tokens (the last generated token does not have k/v cache computed yet)
    invokePlusScalar(sequence_lengths_, -1, batch_size, stream_);
    sync_check_cuda_error();

    Clear(token_ids_buf_, batch_size * session_len_);
    invokeTransposeAxis01(token_ids_buf_, state_->output_ids, batch_size, session_len_, 1, stream_);
    sync_check_cuda_error();

    // token_ids_buf_[s, b]
    // ABCDe            ABCDe     e
    // ABCDEFGHIJk      ABCDEFGHIJk
    // ABCDEFGHi    ->  ABCDEFGHi i
    // ABCDEFGh         ABCDEFGh  h
    // ABCd             ABCd      d
    invokePadLastTokenIds(token_ids_buf_, init_context_length_, g.max_init_ctx_len, batch_size, stream_);
    sync_check_cuda_error();

    // seq_limit_len_, will be compared to `step` instead of `sequence_length`, so padding len should be accounted for
    for (int i = 0; i < batch_size; ++i) {
        h_seq_limit_len_[i] = state_->seq_len_limit[i] + (g.max_init_ctx_len - state_->h_context_length[i]);
    }
    Copy(h_seq_limit_len_, batch_size, seq_limit_len_);

Li Zhang's avatar
Li Zhang committed
950
    TensorMap inputs;
Li Zhang's avatar
Li Zhang committed
951
    for (const auto& [name, h_ptr, d_ptr] : sampling_params_) {
Li Zhang's avatar
Li Zhang committed
952
        // find an exemplar that matches the param name
Li Zhang's avatar
Li Zhang committed
953
        const Tensor* ptr{};
Li Zhang's avatar
Li Zhang committed
954
        for (int i = 0; i < batch_size; ++i) {
Li Zhang's avatar
Li Zhang committed
955
956
            if (state_->requests[i]->inputs[rank_].isExist(name)) {
                ptr = &state_->requests[i]->inputs[rank_].at(name);
Li Zhang's avatar
Li Zhang committed
957
958
959
                break;
            }
        }
Li Zhang's avatar
Li Zhang committed
960
        // fill the batch of the param
Li Zhang's avatar
Li Zhang committed
961
962
963
964
        if (ptr) {
            const auto& ref   = *ptr;
            auto        shape = ref.shape;
            FT_CHECK(shape[0] == 1);
Li Zhang's avatar
Li Zhang committed
965
            shape[0]                = batch_size;
Li Zhang's avatar
Li Zhang committed
966
            const int size_in_bytes = ref.sizeBytes();
Li Zhang's avatar
Li Zhang committed
967
            memset(h_ptr, 0, size_in_bytes * batch_size);
Li Zhang's avatar
Li Zhang committed
968
            for (int i = 0; i < batch_size; ++i) {
969
                FT_CHECK(state_->requests[i] != nullptr);
Li Zhang's avatar
Li Zhang committed
970
971
                if (state_->requests[i]->inputs[rank_].isExist(name)) {
                    Tensor& src = state_->requests[i]->inputs[rank_].at(name);
Li Zhang's avatar
Li Zhang committed
972
                    FT_CHECK(ref.shape == src.shape);
Li Zhang's avatar
Li Zhang committed
973
                    std::copy_n(src.getPtr<std::byte>(), size_in_bytes, h_ptr + size_in_bytes * i);
Li Zhang's avatar
Li Zhang committed
974
975
                }
            }
Li Zhang's avatar
Li Zhang committed
976
977
978
979
            if (d_ptr) {
                Copy(h_ptr, batch_size * size_in_bytes, d_ptr);
            }
            inputs.insert({name, {d_ptr ? MEMORY_GPU : MEMORY_CPU, ref.type, shape, d_ptr ? d_ptr : h_ptr}});
Li Zhang's avatar
Li Zhang committed
980
            if (debug_ && rank_ == 0) {
Li Zhang's avatar
Li Zhang committed
981
                TM_LOG_INFO("[initializeSampling] %s", format({name, inputs.at(name)}).c_str());
Li Zhang's avatar
Li Zhang committed
982
983
984
985
            }
        }
    }

Li Zhang's avatar
Li Zhang committed
986
987
988
989
990
    // init for eos
    std::fill_n(h_end_ids_buf_, batch_size, model_->end_id_);
    Copy(h_end_ids_buf_, batch_size, d_end_ids_buf_);
    inputs.insert({"end_id", {MEMORY_GPU, TYPE_INT32, {(size_t)batch_size}, d_end_ids_buf_}});

Li Zhang's avatar
Li Zhang committed
991
992
    inputs_ = std::move(inputs);

Li Zhang's avatar
Li Zhang committed
993
    model_->dynamic_decode_layer_->setup(batch_size, 1, &inputs_);
Li Zhang's avatar
Li Zhang committed
994
995
}

996
template<typename T>
Li Zhang's avatar
Li Zhang committed
997
void LlamaBatch<T>::OutputContextLogits(T*                      context_decoder_output,
998
999
1000
1001
1002
1003
1004
1005
                                        const std::vector<int>& indices,
                                        const std::vector<int>& lengths)
{
    std::vector<float*> output_logits;
    int                 num_token = 0;
    {
        bool is_return_logits = false;
        for (int k = 0; k < indices.size(); ++k) {
Li Zhang's avatar
Li Zhang committed
1006
            auto& request = state_->requests[indices[k]];
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
            output_logits.push_back(request->outputs[rank_].getPtr<float>("logits", nullptr));
            num_token += lengths[k];
            if (output_logits.back()) {
                is_return_logits = true;
            }
        }
        if (!is_return_logits) {
            return;
        }
    }

    if (context_logits_buf_ == nullptr) {
Li Zhang's avatar
Li Zhang committed
1019
        NcclGuard guard(model_->tensor_para_, stream_, true);
Chen Xin's avatar
Chen Xin committed
1020
        context_logits_buf_ =
Li Zhang's avatar
Li Zhang committed
1021
1022
            (float*)allocator_->malloc(sizeof(float) * model_->vocab_size_padded_ * max_context_token_num_);
        const auto tp = model_->tensor_para_.world_size_;
1023
        if (tp > 1) {
Li Zhang's avatar
Li Zhang committed
1024
1025
            FT_CHECK(model_->vocab_size_padded_ % tp == 0);
            const auto local_vocab_size = model_->vocab_size_padded_ / tp;
1026
1027
1028
1029
1030
            local_context_logits_buf_ =
                (float*)allocator_->malloc(sizeof(float) * local_vocab_size * max_context_token_num_);
        }
    }

Li Zhang's avatar
Li Zhang committed
1031
    model_->postDecodeEmbedding(context_logits_buf_, local_context_logits_buf_, context_decoder_output, num_token);
1032
1033
1034
1035
1036

    auto logits = context_logits_buf_;

    for (int k = 0; k < indices.size(); ++k) {
        if (output_logits[k]) {
Li Zhang's avatar
Li Zhang committed
1037
            Copy(logits, model_->vocab_size_ * lengths[k], output_logits[k]);
1038
        }
Li Zhang's avatar
Li Zhang committed
1039
        logits += model_->vocab_size_padded_ * lengths[k];
1040
1041
1042
    }
}

Li Zhang's avatar
Li Zhang committed
1043
template<typename T>
1044
auto LlamaBatch<T>::Finish(GenerationState& g) -> std::vector<Signal>
Li Zhang's avatar
Li Zhang committed
1045
{
Li Zhang's avatar
Li Zhang committed
1046
1047
    NvtxScope scope("Finish");
    const int batch_size = state_->active_size;
Li Zhang's avatar
Li Zhang committed
1048

1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
    if (batch_size - g.partial) {
        FT_CHECK(g.step >= 0);

        // [s,b] -> [b,s] and skip padding in [context_len, max_context_len)
        invokeGatherOutput(state_->output_ids,
                           token_ids_buf_,
                           init_context_length_,
                           g.max_init_ctx_len,
                           g.step,
                           session_len_,
                           batch_size - g.partial,
                           stream_);
        sync_check_cuda_error();
    }
Li Zhang's avatar
Li Zhang committed
1063

Li Zhang's avatar
Li Zhang committed
1064
1065
    Copy(state_->output_ids, batch_size * session_len_, h_output_ids_);
    Copy(finished_buf_, batch_size, state_->h_finished);
Li Zhang's avatar
Li Zhang committed
1066
1067
    Copy(sequence_lengths_, batch_size, state_->h_context_length);

Li Zhang's avatar
Li Zhang committed
1068
    check_cuda_error(cudaStreamSynchronize(stream_));
Li Zhang's avatar
Li Zhang committed
1069

1070
1071
    // invariant: context_length = sequence_length + 1, so that h_context_length include all (including the one just
    // generated) tokens
Li Zhang's avatar
Li Zhang committed
1072
1073
    for (int i = 0; i < batch_size; ++i) {
        ++state_->h_context_length[i];
Li Zhang's avatar
Li Zhang committed
1074
    }
Li Zhang's avatar
Li Zhang committed
1075

Li Zhang's avatar
Li Zhang committed
1076
1077
    {  // set output tokens ids and sequence length
        int* output_ptr = h_output_ids_;
1078
        for (int i = 0; i < batch_size - g.partial; ++i) {
Li Zhang's avatar
Li Zhang committed
1079
            if (state_->requests[i] && (state_->requests[i]->stream_cb || state_->h_finished[i])) {
1080
1081
1082
1083
1084
1085
                auto      output_ids = state_->requests[i]->outputs[rank_].getPtr<int>("output_ids");
                auto      output_len = state_->requests[i]->outputs[rank_].getPtr<int>("sequence_length");
                const int count      = state_->h_context_length[i];
                // TODO: sync history output tokens at when receiving the request and copy the last token here
                std::copy(output_ptr, output_ptr + count, output_ids);
                *output_len = count;
Li Zhang's avatar
Li Zhang committed
1086
            }
Li Zhang's avatar
Li Zhang committed
1087
            output_ptr += session_len_;
Li Zhang's avatar
Li Zhang committed
1088
        }
Chen Xin's avatar
Chen Xin committed
1089
    }
Li Zhang's avatar
Li Zhang committed
1090
1091

    if (debug_ && rank_ == 0) {
Li Zhang's avatar
Li Zhang committed
1092
        for (int i = 0; i < batch_size; ++i) {
1093
1094
1095
1096
1097
1098
1099
1100
1101
            // ss << (i ? ", " : "") << "(" << state_->h_context_length[i] << "," << state_->h_finished[i] << ")";
            std::vector<int> tokens(state_->h_context_length[i]);
            Copy(state_->output_ids + i * session_len_, tokens.size(), tokens.data());
            cudaStreamSynchronize(stream_);
            std::stringstream ss;
            for (const auto& t : tokens) {
                ss << " " << t;
            }
            TM_LOG_INFO("[Finish] slot %d, tokens [%s]", i, ss.str().c_str());
Li Zhang's avatar
Li Zhang committed
1102
1103
1104
        }
    }

Li Zhang's avatar
Li Zhang committed
1105
1106
    std::vector<Signal> signals;
    {
Li Zhang's avatar
Li Zhang committed
1107
        NvtxScope _("stream_and_completion_signal");
1108
        for (int i = 0; i < batch_size - g.partial; ++i) {
Li Zhang's avatar
Li Zhang committed
1109
1110
1111
1112
            if (state_->requests[i]) {
                if (state_->h_finished[i]) {
                    // Interrupt finished sequences and move the request handle into the signal closure
                    signals.push_back(Interrupt(i));
1113
                    ++g.finished_count;
Li Zhang's avatar
Li Zhang committed
1114
1115
1116
1117
1118
1119
1120
1121
1122
                }
                else if (state_->requests[i]->stream_cb) {
                    // Create signals by copying the request handles for non-finished streaming requests
                    signals.push_back([this, r = state_->requests[i]] {
                        if (rank_ == 0) {
                            r->stream_cb(&r->outputs[rank_].get());
                        }
                    });
                }
Li Zhang's avatar
Li Zhang committed
1123
1124
            }
        }
1125
        if (g.finished_count) {
Li Zhang's avatar
Li Zhang committed
1126
1127
1128
            // synchronize for interrupted sequences
            check_cuda_error(cudaStreamSynchronize(stream_));
        }
Li Zhang's avatar
Li Zhang committed
1129
    }
1130
1131
1132
1133
1134
1135
1136

    if (g.partial) {
        const int i = batch_size - 1;
        // recover full context length of partial
        state_->h_context_length[i] = g.partial_context_legnth;
    }

Li Zhang's avatar
Li Zhang committed
1137
    return signals;
Li Zhang's avatar
Li Zhang committed
1138
1139
1140
}

template<typename T>
Li Zhang's avatar
Li Zhang committed
1141
auto LlamaBatch<T>::Interrupt(int index, bool force_stop, bool force_end) -> Signal
Li Zhang's avatar
Li Zhang committed
1142
1143
{
    if (rank_ == 0) {
Li Zhang's avatar
Li Zhang committed
1144
        TM_LOG_INFO("[Interrupt] slot = %d, id = %lu", index, (long)state_->requests[index]->id);
Li Zhang's avatar
Li Zhang committed
1145
1146
1147
    }

    if (debug_ && rank_ == 0) {
Li Zhang's avatar
Li Zhang committed
1148
1149
        std::vector<int> tokens(state_->h_context_length[index]);
        Copy(state_->output_ids + index * session_len_, tokens.size(), tokens.data());
Li Zhang's avatar
Li Zhang committed
1150
1151
1152
1153
1154
        cudaStreamSynchronize(stream_);
        std::stringstream ss;
        for (const auto& t : tokens) {
            ss << " " << t;
        }
Li Zhang's avatar
Li Zhang committed
1155
        TM_LOG_INFO("[Interrupt] slot %d, tokens [%s]", index, ss.str().c_str());
Li Zhang's avatar
Li Zhang committed
1156
1157
    }

Li Zhang's avatar
Li Zhang committed
1158
1159
1160
    if (state_->requests[index]->end_flag || force_end) {
        // Sequence is ending this round or a stop request is issued to end it
        FT_CHECK(sequence_manager_->Erase(state_->requests[index]->id));
Li Zhang's avatar
Li Zhang committed
1161
1162
    }
    else {
1163
        const int output_len = state_->h_context_length[index];
Li Zhang's avatar
Li Zhang committed
1164
        auto&     seq        = *state_->sequences[index];
Li Zhang's avatar
Li Zhang committed
1165

Li Zhang's avatar
Li Zhang committed
1166
        // Update token IDs
Li Zhang's avatar
Li Zhang committed
1167
1168
        seq.tokens.resize(output_len);
        const auto output_ids_data = state_->requests[index]->outputs[rank_].at("output_ids").getPtr<int>();
Li Zhang's avatar
Li Zhang committed
1169
        std::copy_n(output_ids_data, output_len, seq.tokens.data());
Li Zhang's avatar
Li Zhang committed
1170

Li Zhang's avatar
Li Zhang committed
1171
1172
1173
1174
        // Save random state in host memory
        seq.random_state.resize(sizeof(curandState_t));
        // This async copy must be synchronized by the caller
        Copy(state_->curand_state + index, 1, (curandState_t*)seq.random_state.data());
Li Zhang's avatar
Li Zhang committed
1175

Li Zhang's avatar
Li Zhang committed
1176
        // Set unlock flag for corresponding blocks, will be unlocked in the next `Materialize()`
Li Zhang's avatar
Li Zhang committed
1177
1178
1179
1180
        sequence_manager_->UpdateAndSetUnlock(seq);
    }

    state_->sequences[index] = nullptr;
Li Zhang's avatar
Li Zhang committed
1181
1182
1183
1184
1185
1186
1187

    // move the request handle into the signal
    return [this, r = std::move(state_->requests[index])] {
        if (rank_ == 0) {
            r->signal.set_value(0);
        }
    };
Li Zhang's avatar
Li Zhang committed
1188
1189
1190
1191
1192
}

template<typename T>
void LlamaBatch<T>::InternalThreadEntry(int device_id)
{
Li Zhang's avatar
Li Zhang committed
1193
    // TM_LOG_INFO("[InternalThreadEntry] %d", (int)rank_);
Li Zhang's avatar
Li Zhang committed
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
    check_cuda_error(cudaSetDevice(device_id));

    auto& shared_state = model_->shared_state_;

    auto& request_queue  = shared_state->request_queue;
    auto& infer_requests = shared_state->infer_requests;
    auto& stop_requests  = shared_state->stop_requests;

    GenerationState g{};

Li Zhang's avatar
Li Zhang committed
1204
1205
1206
    constexpr int request_interval = 1;
    long          request_counter  = 0;

Li Zhang's avatar
Li Zhang committed
1207
1208
    while (1) {
        if (rank_ == 0) {
1209
            const int  free_slot_count = max_batch_size_ - state_->size + g.finished_count;
Li Zhang's avatar
Li Zhang committed
1210
            const bool is_empty        = (free_slot_count == max_batch_size_);
Li Zhang's avatar
Li Zhang committed
1211
1212
1213
1214
1215
1216
1217
1218
            stop_requests.clear();
            infer_requests.clear();
            if (is_empty || request_counter % request_interval == 0) {
                // Block if batch is empty
                request_queue.dequeue(stop_requests, infer_requests, free_slot_count, is_empty, shared_state->abort);
                if (!shared_state->abort) {
                    RejectInvalidRequests(stop_requests, infer_requests);
                }
Li Zhang's avatar
Li Zhang committed
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
            }
        }

        NvtxScope scope("mainloop");

        // wait while rank-0 is dequeueing
        shared_state->barrier->wait();

        if (shared_state->abort) {
            TM_LOG_INFO("[InternalThreadEntry] stop requested.");
            return;
        }

        auto signals = ProcessStopRequests(stop_requests);

Li Zhang's avatar
Li Zhang committed
1234
        // Shared `priority` field will be assigned by rank-0
Li Zhang's avatar
Li Zhang committed
1235
1236
        ProcessInferRequests(infer_requests);

Li Zhang's avatar
Li Zhang committed
1237
        // Wait while shared `requests` is being used
Li Zhang's avatar
Li Zhang committed
1238
1239
        shared_state->barrier->wait();

Li Zhang's avatar
Li Zhang committed
1240
1241
        SendSignals(std::move(signals));

1242
        Initialize(g);
Li Zhang's avatar
Li Zhang committed
1243

1244
        FT_CHECK(step_length_ == 1);
Li Zhang's avatar
Li Zhang committed
1245

1246
        if (state_->active_size) {
Li Zhang's avatar
Li Zhang committed
1247
            for (int i = 0; i < step_length_; ++i) {
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
                //
                auto cont = Forward(g, i);
                //
                if (auto signals = Finish(g); !signals.empty()) {
                    if (g.finished_count) {
                        // Finished requests and corresponding output tensors will be released when notified
                        // wait for all ranks to ensure no rank (except for output thread) will access related
                        // resources
                        shared_state->barrier->wait();
                    }
                    SendSignals(std::move(signals));
Li Zhang's avatar
Li Zhang committed
1259
                }
1260
1261
                if (!cont) {  // early exit
                    break;
Li Zhang's avatar
Li Zhang committed
1262
1263
                }
            }
Li Zhang's avatar
Li Zhang committed
1264
        }
Li Zhang's avatar
Li Zhang committed
1265
1266

        ++request_counter;
Li Zhang's avatar
Li Zhang committed
1267
1268
    }

Li Zhang's avatar
Li Zhang committed
1269
1270
1271
1272
    FT_CHECK(0);
}

template<typename T>
Li Zhang's avatar
Li Zhang committed
1273
void LlamaBatch<T>::SendSignals(std::vector<Signal> signals)
Li Zhang's avatar
Li Zhang committed
1274
{
Li Zhang's avatar
Li Zhang committed
1275
1276
1277
1278
1279
1280
1281
1282
    if (rank_ != 0 || signals.empty()) {
        return;
    }
    {
        std::lock_guard lock{output_mutex_};
        output_signals_.insert(output_signals_.end(),  //
                               std::move_iterator{signals.begin()},
                               std::move_iterator{signals.end()});
Li Zhang's avatar
Li Zhang committed
1283
    }
Li Zhang's avatar
Li Zhang committed
1284
    output_cv_.notify_one();
Li Zhang's avatar
Li Zhang committed
1285
1286
1287
1288
1289
}

template<typename T>
void LlamaBatch<T>::Start()
{
Li Zhang's avatar
Li Zhang committed
1290
    TM_LOG_INFO("LlamaBatch<T>::Start()");
Li Zhang's avatar
Li Zhang committed
1291
1292
1293
    int device_id = -1;
    check_cuda_error(cudaGetDevice(&device_id));
    internal_thread_ = std::thread(&LlamaBatch::InternalThreadEntry, this, device_id);
Li Zhang's avatar
Li Zhang committed
1294
    if (rank_ == 0) {
Li Zhang's avatar
Li Zhang committed
1295
        output_thread_ = std::thread(&LlamaBatch::OutputThreadEntry, this);
Li Zhang's avatar
Li Zhang committed
1296
    }
Li Zhang's avatar
Li Zhang committed
1297
}
Li Zhang's avatar
Li Zhang committed
1298

Li Zhang's avatar
Li Zhang committed
1299
1300
1301
1302
template<typename T>
void LlamaBatch<T>::OutputThreadEntry()
{
    while (true) {
Li Zhang's avatar
Li Zhang committed
1303
        std::vector<Signal> signals;
Li Zhang's avatar
Li Zhang committed
1304
        {
Li Zhang's avatar
Li Zhang committed
1305
            // Wait for signals to come
Li Zhang's avatar
Li Zhang committed
1306
            std::unique_lock lock(output_mutex_);
Li Zhang's avatar
Li Zhang committed
1307
            output_cv_.wait(lock, [&] { return !output_signals_.empty() || output_stop_token_; });
Li Zhang's avatar
Li Zhang committed
1308
1309
1310
1311
            if (output_stop_token_) {
                TM_LOG_INFO("[OutputThreadEntry] stop requested.");
                return;
            }
Li Zhang's avatar
Li Zhang committed
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
            signals = std::move(output_signals_);
        }
        if (rank_ == 0 && model_->ffi_lock_) {
            model_->ffi_lock_(1);
        }
        // invoke stream cbs & signals
        for (const auto& s : signals) {
            s();
        }
        if (rank_ == 0 && model_->ffi_lock_) {
            model_->ffi_lock_(0);
Li Zhang's avatar
Li Zhang committed
1323
1324
        }
    }
Li Zhang's avatar
Li Zhang committed
1325
1326
}

1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
template<typename T>
bool LlamaBatch<T>::Forward(GenerationState& g, int iter)
{
    NvtxScope _("Forward");

    FT_CHECK(max_context_token_num_ >= max_batch_size_);

    const int active_size = state_->active_size;

    constexpr int kLogInterval = 10;
    if (rank_ == 0 && (g.step - 1) % kLogInterval == 0) {
        TM_LOG_INFO("------------------------- step = %d -------------------------", g.step - 1);
    }

    int               pf_offset = -1;
    std::vector<int*> input_d_ptrs(active_size);

    if (iter == 0) {  // The first iter may have pre-fill tokens
        for (int i = 0; i < active_size; ++i) {
            const auto& seq = *state_->sequences[i];
            // const int   missing    = state_->h_context_length[i] - seq.cache_len;
            FT_CHECK(seq.input_length >= 1);
            h_input_length_buf_[i] = seq.input_length;
            input_d_ptrs[i]        = state_->output_ids + i * session_len_ + seq.cache_len;
            if (seq.input_length > 1 && pf_offset < 0) {
                pf_offset = i;
            }
        }
        if (pf_offset < 0) {
            pf_offset = active_size;
        }
    }
    else {
        for (int i = 0; i < active_size; ++i) {
            h_input_length_buf_[i] = 1;
            input_d_ptrs[i]        = state_->output_ids + i * session_len_ + state_->h_context_length[i] - 1;
        }
        pf_offset = active_size;
    }

    // These buffers are only accessed when there are prefill workloads
    if (pf_offset != active_size) {
        Copy(state_->h_context_length, active_size, context_length_buf_);
        Copy(h_input_length_buf_, active_size, input_length_buf_);
    }

    // Find mini-batch offsets: input length > 1 ? prefill() : decode()
    // Constraints on mini-batches
    // - `context_decoder_input` and `context_decoder_output` can hold `max_context_token_num_` tokens w/o padding
    // - prefill() use `tmp_k_cache_buf_` and `tmp_k_cache_buf_`, they can hold `max_context_token_num_` tokens
    //     but each sequence is padded to the maximum context length in the batch
    std::vector<int> offsets{0};
    std::vector<int> max_context_cnts;
    // initialize first mini-batch with decode tokens
    int accum_size        = pf_offset;
    int accum_token_count = pf_offset;
    int max_context_count = 0;
    for (int i = pf_offset; i < active_size; ++i) {
        FT_CHECK(iter == 0);
        int size          = accum_size + 1;
        int input_count   = accum_token_count + h_input_length_buf_[i];
        int context_count = std::max(max_context_count, state_->h_context_length[i]);
        // correct pre-fill batch size for the first batch
        int pf_size = offsets.size() == 1 ? size - pf_offset : size;
        // we have `cu_seqlens` on q so no padding for input is needed
        // prefill kernels are expecting uniform k/v cache length -> `max_context_count * size <=
        // max_context_token_num_`
        if (input_count <= max_context_token_num_ && context_count * pf_size <= max_context_token_num_) {
            accum_size        = size;
            accum_token_count = input_count;
            max_context_count = context_count;
        }
        else {
            offsets.push_back(i);
            max_context_cnts.push_back(max_context_count);
            accum_size        = 1;
            accum_token_count = h_input_length_buf_[i];
            max_context_count = state_->h_context_length[i];
        }
    }
    offsets.push_back(active_size);
    max_context_cnts.push_back(max_context_count);

    // forward on mini-batches
    for (int p = 0; p < (int)offsets.size() - 1; ++p) {
        int  first           = offsets[p];
        int  last            = offsets[p + 1];
        int  mini_batch_size = last - first;
        T*   k_ptr           = tmp_k_cache_buf_;
        T*   v_ptr           = tmp_v_cache_buf_;
        int  max_input_len{};
        auto input_ids = context_decoder_ids_buf_;
        //
        std::vector<int> decode_indices{};
        std::vector<int> decode_lengths{};

        BatchedCopy batched_copy;
        for (int i = first; i < last; ++i) {
            input_ids = batched_copy.Add(input_d_ptrs[i], h_input_length_buf_[i], input_ids);
            dbg(i, h_input_length_buf_[i]);
            // allocate tmp k/v buffer for pre-fill sequences
            if (i < pf_offset) {
                h_tmp_k_ptrs_[i] = h_tmp_v_ptrs_[i] = nullptr;
            }
            else {
                h_tmp_k_ptrs_[i] = k_ptr;
                h_tmp_v_ptrs_[i] = v_ptr;
                k_ptr += model_->local_kv_head_num_ * max_context_cnts[p] * model_->size_per_head_;
                v_ptr += model_->local_kv_head_num_ * max_context_cnts[p] * model_->size_per_head_;
            }
            decode_indices.push_back(i);
            decode_lengths.push_back(h_input_length_buf_[i]);
            max_input_len = std::max(max_input_len, h_input_length_buf_[i]);
        }
        int token_count = input_ids - context_decoder_ids_buf_;

        batched_copy.Submit(stream_);

        Copy(h_tmp_k_ptrs_ + first, mini_batch_size, tmp_k_ptrs_ + first);
        Copy(h_tmp_v_ptrs_ + first, mini_batch_size, tmp_v_ptrs_ + first);

        const int dc_batch_size = p ? 0 : pf_offset;
        const int pf_batch_size = mini_batch_size - dc_batch_size;

        if (rank_ == 0) {
            if (pf_batch_size) {
                TM_LOG_INFO("[Forward] [%d, %d), dc_bsz = %d, pf_bsz = %d, n_tok = %d, max_q = %d, max_k = %d",
                            first,
                            last,
                            dc_batch_size,
                            pf_batch_size,
                            token_count,
                            max_input_len,
                            max_context_cnts[p]);
            }
        }

        model_->forwardUnified(decoder_output_buf_ + first * model_->hidden_units_,
                               context_decoder_output_buf_,  // temp
                               context_decoder_input_buf_,   // temp
                               (void**)k_block_ptrs_,
                               (void**)v_block_ptrs_,
                               context_decoder_ids_buf_,  // temp
                               cu_block_counts_ + first,
                               rope_theta_ + first,
                               finished_buf_ + first,
                               input_length_buf_ + first,
                               context_length_buf_ + first,
                               (T**)tmp_k_ptrs_ + first,
                               (T**)tmp_v_ptrs_ + first,
                               token_count,
                               dc_batch_size,
                               g.step,
                               g.sum_seq_len,
                               g.max_seq_len,
                               pf_batch_size,
                               max_input_len,
                               max_context_cnts[p],
                               max_context_cnts[p]);

        if (iter == 0) {
            // compute logits of inputs if requested
            OutputContextLogits(context_decoder_output_buf_, decode_indices, decode_lengths);
        }
    }

    std::fill(h_input_length_buf_, h_input_length_buf_ + active_size, 0);

    // `SequenceManager` needs real-time value of cache length
    for (int i = 0; i < active_size; ++i) {
        if (state_->requests[i]) {
            FT_CHECK(state_->sequences[i]);
            state_->sequences[i]->cache_len += state_->sequences[i]->input_length;
        }
    }

    bool should_stop{};

    if (active_size > g.partial) {
        model_->postDecodeEmbedding(logits_buf_, local_logits_buf_, decoder_output_buf_, active_size - g.partial);

        FT_CHECK(g.step >= 0);

        // TM_LOG_INFO("dyn decode bsz %d, partial %d", active_size, g.partial);

        // stop-words & bad-words require the matched tokens to be contiguous, so item size > 1 is
        // not supported yet.
        model_->dynamicDecode(token_ids_buf_,
                              finished_buf_,
                              sequence_lengths_,
                              &should_stop,
                              state_->curand_state,
                              &inputs_,
                              &outputs_,
                              logits_buf_,
                              seq_limit_len_,
                              init_context_length_,
                              d_end_ids_buf_,
                              g.step,
                              0,
                              g.max_init_ctx_len,
                              session_len_ * 2,
                              active_size - g.partial);
    }

    if (debug_ && rank_ == 0) {
        std::vector<int> curr(active_size);
        Copy(token_ids_buf_ + g.step * active_size, active_size, curr.data());
        cudaStreamSynchronize(stream_);
        std::stringstream scurr;
        for (int k = 0; k < curr.size(); ++k) {
            scurr << std::setw(6) << curr[k];
        }
        TM_LOG_INFO("[Forward] step = %d, [%s]", g.step - 1, scurr.str().c_str());
    }

    // check_cuda_error(cudaStreamSynchronize(stream_));

    ////////////////////////////////////////////////
    /// ! increase the counters
    g.step += 1;
    g.max_seq_len += 1;
    g.sum_seq_len += state_->active_size;

    // PrintDecodeTokens(token_ids_buf_, g.step, active_size, stream_, "Forward");

    return !should_stop;
}

Li Zhang's avatar
Li Zhang committed
1556
1557
1558
template class LlamaBatch<half>;
template class LlamaBatch<float>;

lvhan028's avatar
lvhan028 committed
1559
}  // namespace turbomind