wehub-resource-sync
408 lines
19 KiB
C++
408 lines
19 KiB
C++
/*!
|
|
* Copyright (c) 2023-2025 by Contributors
|
|
* \file serve/engine_actions/batch_draft.cc
|
|
*/
|
|
|
|
#include <numeric>
|
|
|
|
#include "../config.h"
|
|
#include "../model.h"
|
|
#include "../sampler/sampler.h"
|
|
#include "action.h"
|
|
#include "action_commons.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace serve {
|
|
|
|
/*!
|
|
* \brief The action that runs draft proposal for requests in the
|
|
* `running_queue` of engine state. Preempt low-priority requests
|
|
* accordingly when it is impossible to decode all the running requests.
|
|
*/
|
|
class BatchDraftActionObj : public EngineActionObj {
|
|
public:
|
|
explicit BatchDraftActionObj(Array<Model> models, LogitProcessor logit_processor, Sampler sampler,
|
|
std::vector<ModelWorkspace> model_workspaces,
|
|
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
|
EngineConfig engine_config,
|
|
Optional<EventTraceRecorder> trace_recorder)
|
|
: models_(std::move(models)),
|
|
logit_processor_(std::move(logit_processor)),
|
|
sampler_(std::move(sampler)),
|
|
model_workspaces_(std::move(model_workspaces)),
|
|
draft_token_workspace_manager_(std::move(draft_token_workspace_manager)),
|
|
engine_config_(std::move(engine_config)),
|
|
trace_recorder_(std::move(trace_recorder)) {}
|
|
|
|
Array<Request> Step(EngineState estate) final {
|
|
// - Only run spec decode when there are two models (llm+ssm) and >=1 running requests.
|
|
if (models_.size() != 2 || estate->running_queue.empty()) {
|
|
return {};
|
|
}
|
|
|
|
// Preempt request state entries when decode cannot apply.
|
|
std::vector<RequestStateEntry> running_rsentries = estate->GetRunningRequestStateEntries();
|
|
while (!CanDecode(running_rsentries.size())) {
|
|
if (estate->prefix_cache->TryFreeMemory()) continue;
|
|
RequestStateEntry preempted = PreemptLastRunningRequestStateEntry(
|
|
estate, models_, draft_token_workspace_manager_, trace_recorder_);
|
|
if (preempted.same_as(running_rsentries.back())) {
|
|
running_rsentries.pop_back();
|
|
}
|
|
}
|
|
while (running_rsentries.size() * (engine_config_->spec_draft_length + 1) >
|
|
std::min(static_cast<int64_t>(engine_config_->max_num_sequence),
|
|
engine_config_->prefill_chunk_size)) {
|
|
running_rsentries.pop_back();
|
|
}
|
|
|
|
auto tstart = std::chrono::high_resolution_clock::now();
|
|
|
|
int num_rsentries = running_rsentries.size();
|
|
TVM_FFI_ICHECK_GT(num_rsentries, 0)
|
|
<< "There should be at least one request state entry that can run decode. "
|
|
"Possible failure reason: none of the prefill phase of the running requests is finished";
|
|
TVM_FFI_ICHECK_LE(num_rsentries, engine_config_->max_num_sequence)
|
|
<< "The number of running requests exceeds the max number of sequence in EngineConfig. "
|
|
"Possible failure reason: the prefill action allows new sequence in regardless of the "
|
|
"max num sequence.";
|
|
Array<String> request_ids;
|
|
std::vector<int64_t> request_internal_ids;
|
|
Array<String> request_ids_per_leaf_node;
|
|
Array<GenerationConfig> generation_cfg;
|
|
Array<GenerationConfig> generation_cfg_for_logitproc;
|
|
std::vector<RandomGenerator*> rngs;
|
|
std::vector<std::vector<int>> draft_token_indices;
|
|
// Number of input tokens for each request. Each request can have multiple leaf tokens for the
|
|
// next forward when multiple tokens are drafted.
|
|
std::vector<int> cum_num_tokens;
|
|
std::vector<int64_t> token_tree_parent_ptr;
|
|
request_ids.reserve(num_rsentries);
|
|
request_internal_ids.reserve(num_rsentries);
|
|
generation_cfg.reserve(num_rsentries);
|
|
generation_cfg_for_logitproc.reserve(num_rsentries);
|
|
draft_token_indices.reserve(num_rsentries);
|
|
cum_num_tokens.reserve(num_rsentries + 1);
|
|
for (const RequestStateEntry& rsentry : running_rsentries) {
|
|
request_ids.push_back(rsentry->request->id);
|
|
request_internal_ids.push_back(rsentry->mstates[0]->internal_id);
|
|
}
|
|
|
|
TVM_FFI_ICHECK_GT(estate->spec_draft_length, 0)
|
|
<< "The speculative decoding draft length must be positive.";
|
|
// The first model doesn't get involved in draft proposal.
|
|
for (int model_id = 1; model_id < static_cast<int>(models_.size()); ++model_id) {
|
|
// Collect
|
|
// - the last committed token,
|
|
// - the request model state of each request,
|
|
// - the number of tokens for each request to send into the model (it may
|
|
// be more than one if the draft model is lagging behind the main model, when
|
|
// the engine switches from normal batch decode mode to speculative decoding mode).
|
|
std::vector<int> input_tokens;
|
|
Array<RequestModelState> mstates;
|
|
std::vector<int> input_lengths;
|
|
input_tokens.reserve(num_rsentries);
|
|
mstates.reserve(num_rsentries);
|
|
input_lengths.reserve(num_rsentries);
|
|
for (const RequestStateEntry& rsentry : running_rsentries) {
|
|
mstates.push_back(rsentry->mstates[model_id]);
|
|
}
|
|
// "Draft length" rounds of draft proposal.
|
|
for (int draft_id = 0; draft_id < estate->spec_draft_length; ++draft_id) {
|
|
auto tdraft_start = std::chrono::high_resolution_clock::now();
|
|
// prepare new input tokens
|
|
input_tokens.clear();
|
|
input_lengths.clear();
|
|
token_tree_parent_ptr.clear();
|
|
generation_cfg.clear();
|
|
generation_cfg_for_logitproc.clear();
|
|
rngs.clear();
|
|
cum_num_tokens.clear();
|
|
cum_num_tokens.push_back(0);
|
|
request_ids_per_leaf_node.clear();
|
|
std::vector<int> draft_token_parent_idx;
|
|
draft_token_indices.clear();
|
|
|
|
if (draft_id == 0) {
|
|
// Compute the total length that needs to be processed by the draft model,
|
|
// including the lagging-behind part of hte draft model.
|
|
// When the total length to be processed is larger than the prefill chunk
|
|
// size, we must do the prefill with multiple rounds by chunk.
|
|
int total_length = 0;
|
|
for (int i = 0; i < num_rsentries; ++i) {
|
|
TVM_FFI_ICHECK_LE(mstates[i]->committed_tokens.size(),
|
|
running_rsentries[i]->mstates[0]->committed_tokens.size());
|
|
total_length += running_rsentries[i]->mstates[0]->committed_tokens.size() -
|
|
mstates[i]->committed_tokens.size() + 1;
|
|
}
|
|
if (total_length > engine_config_->prefill_chunk_size) {
|
|
PrefillLaggedTokensByChunk(mstates, running_rsentries, models_[model_id],
|
|
total_length - engine_config_->prefill_chunk_size);
|
|
}
|
|
}
|
|
|
|
for (int i = 0; i < num_rsentries; ++i) {
|
|
int num_leaf_nodes = 0;
|
|
// Starting from last committed tokens
|
|
if (draft_id == 0) {
|
|
TVM_FFI_ICHECK_LE(mstates[i]->committed_tokens.size(),
|
|
running_rsentries[i]->mstates[0]->committed_tokens.size());
|
|
TVM_FFI_ICHECK_EQ(mstates[i]->num_tokens_for_next_decode, 1);
|
|
input_tokens.push_back(mstates[i]->committed_tokens.back().GetTokenId());
|
|
input_lengths.push_back(running_rsentries[i]->mstates[0]->committed_tokens.size() -
|
|
mstates[i]->committed_tokens.size() + 1);
|
|
for (size_t j = mstates[i]->committed_tokens.size();
|
|
j < running_rsentries[i]->mstates[0]->committed_tokens.size(); ++j) {
|
|
// This draft model is lagging behind the main model.
|
|
// It may happen when the engine just switches from the normal batch decode
|
|
// mode to the speculative decoding mode.
|
|
// In this case, we need to prefill the misaligned tokens into the draft model.
|
|
mstates[i]->CommitToken(running_rsentries[i]->mstates[0]->committed_tokens[j]);
|
|
input_tokens.push_back(
|
|
running_rsentries[i]->mstates[0]->committed_tokens[j].GetTokenId());
|
|
}
|
|
mstates[i]->num_tokens_for_next_decode = 0;
|
|
draft_token_indices.emplace_back(std::vector<int>{-1});
|
|
rngs.push_back(&running_rsentries[i]->rng);
|
|
draft_token_parent_idx.push_back(-1);
|
|
request_ids_per_leaf_node.push_back(request_ids[i]);
|
|
num_leaf_nodes = 1;
|
|
cum_num_tokens.push_back(cum_num_tokens.back() + 1);
|
|
} else {
|
|
TVM_FFI_ICHECK_EQ(mstates[i]->committed_tokens.size(),
|
|
running_rsentries[i]->mstates[0]->committed_tokens.size());
|
|
TVM_FFI_ICHECK(!mstates[i]->draft_output_tokens.empty());
|
|
draft_token_indices.emplace_back(std::vector<int>{});
|
|
// Get all leaf nodes
|
|
for (int j = 0; j < static_cast<int>(mstates[i]->draft_output_tokens.size()); ++j) {
|
|
if (mstates[i]->draft_token_first_child_idx[j] == -1) {
|
|
int64_t parent_idx = mstates[i]->draft_token_parent_idx[j];
|
|
token_tree_parent_ptr.push_back(parent_idx);
|
|
input_tokens.push_back(mstates[i]->draft_output_tokens[j].GetTokenId());
|
|
draft_token_indices.back().push_back(j);
|
|
rngs.push_back(&running_rsentries[i]->rng);
|
|
num_leaf_nodes++;
|
|
request_ids_per_leaf_node.push_back(request_ids[i]);
|
|
draft_token_parent_idx.push_back(j);
|
|
}
|
|
}
|
|
input_lengths.push_back(num_leaf_nodes);
|
|
cum_num_tokens.push_back(cum_num_tokens.back() + input_lengths.back());
|
|
}
|
|
GenerationConfig generation_cfg_for_draft = [&]() {
|
|
if (engine_config_->spec_tree_width == 1) {
|
|
return mstates[i]->request->generation_cfg;
|
|
}
|
|
auto spec_generation_cfg = tvm::ffi::make_object<GenerationConfigNode>(
|
|
*(mstates[i]->request->generation_cfg.get()));
|
|
spec_generation_cfg->top_logprobs = engine_config_->spec_tree_width;
|
|
spec_generation_cfg->logprobs = true;
|
|
spec_generation_cfg->temperature = 1.0;
|
|
return GenerationConfig(spec_generation_cfg);
|
|
}();
|
|
for (int j = 0; j < num_leaf_nodes; ++j) {
|
|
generation_cfg.push_back(generation_cfg_for_draft);
|
|
}
|
|
generation_cfg_for_logitproc.push_back(generation_cfg_for_draft);
|
|
}
|
|
|
|
// - Compute embeddings.
|
|
RECORD_EVENT(trace_recorder_, request_ids, "start proposal embedding");
|
|
TVM_FFI_ICHECK_LE(input_tokens.size(), engine_config_->prefill_chunk_size);
|
|
ObjectRef embeddings =
|
|
models_[model_id]->TokenEmbed({Shape{input_tokens.begin(), input_tokens.end()}});
|
|
RECORD_EVENT(trace_recorder_, request_ids, "finish proposal embedding");
|
|
|
|
// - Invoke model decode.
|
|
RECORD_EVENT(trace_recorder_, request_ids, "start proposal decode");
|
|
Tensor logits{nullptr};
|
|
|
|
if (input_tokens.size() == num_rsentries) {
|
|
// Each request entry only has one token to feed into the draft model.
|
|
logits = models_[model_id]->BatchDecode(embeddings, request_internal_ids);
|
|
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
|
|
TVM_FFI_ICHECK_EQ(logits->shape[0], num_rsentries);
|
|
TVM_FFI_ICHECK_EQ(logits->shape[1], 1);
|
|
} else if (draft_id == 0) {
|
|
// There exists some request entry which has more than one token to feed.
|
|
// It may happen when the engine just switches from the normal batch decode
|
|
// mode to the speculative decoding mode.
|
|
logits = models_[model_id]->BatchPrefill(embeddings, request_internal_ids, input_lengths);
|
|
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
|
|
TVM_FFI_ICHECK_EQ(logits->shape[0], 1);
|
|
TVM_FFI_ICHECK_EQ(logits->shape[1], num_rsentries);
|
|
} else {
|
|
TVM_FFI_ICHECK_GT(engine_config_->spec_tree_width, 1);
|
|
logits = models_[model_id]->BatchTreeDecode(embeddings, request_internal_ids,
|
|
input_lengths, token_tree_parent_ptr);
|
|
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
|
|
TVM_FFI_ICHECK_EQ(logits->shape[0], cum_num_tokens.back());
|
|
TVM_FFI_ICHECK_EQ(logits->shape[1], 1);
|
|
}
|
|
TVM_FFI_ICHECK_EQ(input_lengths.size(), num_rsentries);
|
|
RECORD_EVENT(trace_recorder_, request_ids, "finish proposal decode");
|
|
|
|
// - Update logits.
|
|
logits = logits.CreateView({cum_num_tokens.back(), logits->shape[2]}, logits->dtype);
|
|
|
|
logit_processor_->InplaceUpdateLogits(logits, generation_cfg_for_logitproc, mstates,
|
|
request_ids, &cum_num_tokens, &mstates,
|
|
&draft_token_indices);
|
|
|
|
// - Compute probability distributions.
|
|
Tensor probs_on_device = logit_processor_->ComputeProbsFromLogits(
|
|
logits, generation_cfg_for_logitproc, request_ids, &cum_num_tokens);
|
|
|
|
// - Commit the prefix cache changes from previous round of action.
|
|
// Note: we commit prefix cache changes here to overlap this commit with the GPU execution.
|
|
estate->prefix_cache->CommitSequenceExtention();
|
|
|
|
// - Sample tokens.
|
|
// Fill range [0, num_rsentries) into `sample_indices`.
|
|
std::vector<int> sample_indices(cum_num_tokens.back());
|
|
std::iota(sample_indices.begin(), sample_indices.end(), 0);
|
|
std::vector<Tensor> prob_dist;
|
|
Tensor renormalized_probs = sampler_->BatchRenormalizeProbsByTopP(
|
|
probs_on_device, sample_indices, request_ids_per_leaf_node, generation_cfg);
|
|
std::vector<SampleResult> sample_results = sampler_->BatchSampleTokensWithProbAfterTopP(
|
|
renormalized_probs, sample_indices, request_ids_per_leaf_node, generation_cfg, rngs);
|
|
TVM_FFI_ICHECK_EQ(sample_results.size(), cum_num_tokens.back());
|
|
|
|
// - Add draft token to the state.
|
|
draft_token_workspace_manager_->AllocSlots(cum_num_tokens.back(), &draft_token_slots_);
|
|
models_[model_id]->ScatterDraftProbs(probs_on_device, draft_token_slots_,
|
|
&model_workspaces_[0].draft_probs_storage);
|
|
for (int i = 0; i < num_rsentries; ++i) {
|
|
for (int j = cum_num_tokens[i]; j < cum_num_tokens[i + 1]; ++j) {
|
|
int parent_idx = draft_token_parent_idx[j];
|
|
if (engine_config_->spec_tree_width == 1) {
|
|
mstates[i]->AddDraftToken(sample_results[j], draft_token_slots_[j], parent_idx);
|
|
continue;
|
|
}
|
|
for (int k = 0; k < sample_results[j].top_prob_tokens.size(); ++k) {
|
|
SampleResult top_k_token{sample_results[j].top_prob_tokens[k]};
|
|
mstates[i]->AddDraftToken(top_k_token, draft_token_slots_[j], parent_idx);
|
|
}
|
|
}
|
|
}
|
|
|
|
auto tdraft_end = std::chrono::high_resolution_clock::now();
|
|
estate->metrics.UpdateDraftTimeByBatchSize(
|
|
num_rsentries, static_cast<double>((tdraft_end - tdraft_start).count()) / 1e9);
|
|
}
|
|
}
|
|
|
|
auto tend = std::chrono::high_resolution_clock::now();
|
|
estate->metrics.engine_decode_time_sum += static_cast<double>((tend - tstart).count()) / 1e9;
|
|
|
|
return {};
|
|
}
|
|
|
|
private:
|
|
/*! \brief Check if the input requests can be decoded under conditions. */
|
|
bool CanDecode(int num_rsentries) {
|
|
// The first model is not involved in draft proposal.
|
|
for (int model_id = 1; model_id < static_cast<int>(models_.size()); ++model_id) {
|
|
// Check if the model has enough available pages.
|
|
int num_available_pages = models_[model_id]->GetNumAvailablePages();
|
|
if (num_rsentries > num_available_pages) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
void PrefillLaggedTokensByChunk(const Array<RequestModelState>& mstates,
|
|
const std::vector<RequestStateEntry>& running_rsentries,
|
|
Model model, int remaining_prefill_length) {
|
|
int num_rsentries = mstates.size();
|
|
std::vector<int> input_tokens;
|
|
std::vector<int64_t> request_internal_ids;
|
|
std::vector<int> lengths;
|
|
input_tokens.reserve(engine_config_->prefill_chunk_size);
|
|
request_internal_ids.reserve(num_rsentries);
|
|
lengths.reserve(num_rsentries);
|
|
|
|
auto f_run_prefill = [&model, &input_tokens, &request_internal_ids, &lengths]() {
|
|
ObjectRef embeddings = model->TokenEmbed({Shape{input_tokens.begin(), input_tokens.end()}});
|
|
model->BatchPrefill(embeddings, request_internal_ids, lengths);
|
|
};
|
|
|
|
for (int i = 0; i < num_rsentries; ++i) {
|
|
int prefill_length =
|
|
std::min({static_cast<int>(running_rsentries[i]->mstates[0]->committed_tokens.size() -
|
|
mstates[i]->committed_tokens.size()),
|
|
static_cast<int>(engine_config_->prefill_chunk_size - input_tokens.size()),
|
|
remaining_prefill_length});
|
|
if (prefill_length == 0) {
|
|
// This rsentry is done.
|
|
continue;
|
|
}
|
|
|
|
TVM_FFI_ICHECK(!mstates[i]->committed_tokens.empty());
|
|
for (size_t j = mstates[i]->committed_tokens.size();
|
|
j < running_rsentries[i]->mstates[0]->committed_tokens.size(); ++j) {
|
|
// Commit the lagging-behind tokens to the draft model.
|
|
mstates[i]->CommitToken(running_rsentries[i]->mstates[0]->committed_tokens[j - 1]);
|
|
input_tokens.push_back(
|
|
running_rsentries[i]->mstates[0]->committed_tokens[j - 1].GetTokenId());
|
|
}
|
|
lengths.push_back(prefill_length);
|
|
request_internal_ids.push_back(running_rsentries[i]->mstates[0]->internal_id);
|
|
mstates[i]->num_tokens_for_next_decode = 1;
|
|
remaining_prefill_length -= prefill_length;
|
|
if (remaining_prefill_length == 0) {
|
|
// All rsentries are done.
|
|
break;
|
|
}
|
|
|
|
if (input_tokens.size() == engine_config_->prefill_chunk_size) {
|
|
// Run prefill if the pending part total length reaches the prefill chunk size.
|
|
f_run_prefill();
|
|
input_tokens.clear();
|
|
request_internal_ids.clear();
|
|
lengths.clear();
|
|
--i;
|
|
continue;
|
|
}
|
|
}
|
|
|
|
if (!input_tokens.empty()) {
|
|
f_run_prefill();
|
|
}
|
|
}
|
|
|
|
/*! \brief The model to run draft generation in speculative decoding. */
|
|
Array<Model> models_;
|
|
/*! \brief The logit processor. */
|
|
LogitProcessor logit_processor_;
|
|
/*! \brief The sampler to sample new tokens. */
|
|
Sampler sampler_;
|
|
/*! \brief The model workspaces. */
|
|
std::vector<ModelWorkspace> model_workspaces_;
|
|
/*! \brief The draft token workspace manager. */
|
|
DraftTokenWorkspaceManager draft_token_workspace_manager_;
|
|
/*! \brief The engine config. */
|
|
EngineConfig engine_config_;
|
|
/*! \brief Event trace recorder. */
|
|
Optional<EventTraceRecorder> trace_recorder_;
|
|
/*! \brief Temporary buffer to store the slots of the current draft tokens */
|
|
std::vector<int> draft_token_slots_;
|
|
};
|
|
|
|
EngineAction EngineAction::BatchDraft(Array<Model> models, LogitProcessor logit_processor,
|
|
Sampler sampler, std::vector<ModelWorkspace> model_workspaces,
|
|
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
|
EngineConfig engine_config,
|
|
Optional<EventTraceRecorder> trace_recorder) {
|
|
return EngineAction(tvm::ffi::make_object<BatchDraftActionObj>(
|
|
std::move(models), std::move(logit_processor), std::move(sampler),
|
|
std::move(model_workspaces), std::move(draft_token_workspace_manager),
|
|
std::move(engine_config), std::move(trace_recorder)));
|
|
}
|
|
|
|
} // namespace serve
|
|
} // namespace llm
|
|
} // namespace mlc
|