wehub-resource-sync
241 lines
9.0 KiB
C++
241 lines
9.0 KiB
C++
/*!
|
|
* Copyright (c) 2023-2025 by Contributors
|
|
* \file serve/engine_actions/batch_verify.cc
|
|
*/
|
|
|
|
#include <tvm/support/cuda/nvtx.h>
|
|
|
|
#include <cmath>
|
|
#include <exception>
|
|
|
|
#include "../config.h"
|
|
#include "../model.h"
|
|
#include "../sampler/sampler.h"
|
|
#include "action.h"
|
|
#include "action_commons.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace serve {
|
|
|
|
using tvm::support::NVTXScopedRange;
|
|
|
|
/*!
|
|
* \brief The action that runs verification 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 BatchJumpForwardActionObj : public EngineActionObj {
|
|
public:
|
|
explicit BatchJumpForwardActionObj(Array<Model> models, Tokenizer tokenizer,
|
|
Optional<EventTraceRecorder> trace_recorder)
|
|
: models_(std::move(models)),
|
|
tokenizer_(tokenizer),
|
|
trace_recorder_(std::move(trace_recorder)) {}
|
|
|
|
Array<Request> Step(EngineState estate) final {
|
|
// - Do not run decode when there are multiple models or no running requests.
|
|
if (models_.size() > 1 || estate->running_queue.empty()) {
|
|
return {};
|
|
}
|
|
|
|
// Preempt request state entries when jump-forward decoding cannot apply.
|
|
std::vector<RequestStateEntry> running_rsentries;
|
|
{
|
|
NVTXScopedRange nvtx_scope("BatchJumpForward getting requests");
|
|
running_rsentries = estate->GetRunningRequestStateEntries();
|
|
while (!CheckMemForJumpForward(running_rsentries.size())) {
|
|
if (estate->prefix_cache->TryFreeMemory()) continue;
|
|
RequestStateEntry preempted =
|
|
PreemptLastRunningRequestStateEntry(estate, models_, std::nullopt, trace_recorder_);
|
|
if (preempted.same_as(running_rsentries.back())) {
|
|
running_rsentries.pop_back();
|
|
}
|
|
}
|
|
}
|
|
|
|
if (running_rsentries.empty()) {
|
|
return {};
|
|
}
|
|
|
|
auto tstart = std::chrono::high_resolution_clock::now();
|
|
|
|
for (auto rsentry : running_rsentries) {
|
|
if (!CanJumpForward(rsentry)) {
|
|
continue;
|
|
}
|
|
|
|
auto mstate = rsentry->mstates[0];
|
|
auto jump_forward_str = mstate->grammar_matcher->FindJumpForwardString();
|
|
|
|
if (jump_forward_str.empty()) {
|
|
continue;
|
|
}
|
|
|
|
auto [rollback_cnt, new_tokens, new_string] =
|
|
RetokenizeWithNewString(mstate, jump_forward_str, MAX_ROLLBACK_TOKENS_);
|
|
|
|
HandleRollback(rsentry, mstate, rollback_cnt, new_tokens, new_string);
|
|
|
|
// Commit new tokens (kv cache is handled in the next decode)
|
|
for (auto token_id : new_tokens) {
|
|
mstate->CommitToken({{token_id, 1.0}, {}});
|
|
}
|
|
|
|
mstate->require_retokenization_in_next_decode = true;
|
|
|
|
// Update metrics
|
|
rsentry->rstate->metrics.jump_forward_tokens +=
|
|
std::max(static_cast<int>(new_tokens.size()) - rollback_cnt, 0);
|
|
|
|
rsentry->rstate->metrics.completion_tokens +=
|
|
static_cast<int>(new_tokens.size()) - rollback_cnt;
|
|
}
|
|
|
|
auto tend = std::chrono::high_resolution_clock::now();
|
|
estate->metrics.engine_jump_forward_time_sum +=
|
|
static_cast<double>((tend - tstart).count()) / 1e9;
|
|
|
|
return {};
|
|
}
|
|
|
|
private:
|
|
/*! \brief Check if jump-forward decoding can be executed without exceeding the memory limit. */
|
|
bool CheckMemForJumpForward(int num_rsentries) {
|
|
static constexpr int MAX_AVG_JUMPFORWARD_PAGES_PER_REQUEST = 10;
|
|
int num_available_pages = models_[0]->GetNumAvailablePages();
|
|
return num_rsentries * MAX_AVG_JUMPFORWARD_PAGES_PER_REQUEST <= num_available_pages;
|
|
}
|
|
|
|
/*! \brief Check if the jump-forward can be executed. When logprobs is requested, or the
|
|
* grammar state matcher is not defined, jump-forward is not executed. */
|
|
bool CanJumpForward(const RequestStateEntry& rsentry) {
|
|
if (rsentry->request->generation_cfg->debug_config.grammar_execution_mode !=
|
|
GrammarExecutionMode::kJumpForward) {
|
|
return false;
|
|
}
|
|
if (rsentry->request->generation_cfg->logprobs) {
|
|
return false;
|
|
}
|
|
if (!rsentry->mstates[0]->grammar_matcher) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
/*!
|
|
* \brief Retokenize the input string with a new string.
|
|
* \param mstate The model state.
|
|
* \param new_string The new string to append.
|
|
* \param max_rollback_tokens The maximum number of tokens to rollback.
|
|
* \return The number of tokens to rollback, the new tokens and a delta string of output (equal to
|
|
* new_string if no cutoff happens; shorter than new_string if cutoff happens).
|
|
*/
|
|
std::tuple<int, std::vector<int32_t>, std::string> RetokenizeWithNewString(
|
|
RequestModelState mstate, const std::string& new_string, int max_rollback_tokens) {
|
|
// Step 1. Get past tokens
|
|
// past_tokens = mstate[-max_rollback_tokens:]
|
|
// past_string = detokenize(past_tokens)
|
|
const auto& token_table = tokenizer_->PostProcessedTokenTable();
|
|
std::vector<int32_t> past_tokens;
|
|
std::string past_string;
|
|
auto past_begin_it = mstate->committed_tokens.size() >= max_rollback_tokens
|
|
? mstate->committed_tokens.end() - max_rollback_tokens
|
|
: mstate->committed_tokens.begin();
|
|
for (auto it = past_begin_it; it != mstate->committed_tokens.end(); ++it) {
|
|
past_tokens.push_back(it->GetTokenId());
|
|
past_string += token_table[it->GetTokenId()];
|
|
}
|
|
|
|
// Step 2. Retokenize
|
|
// Compare tokenize(past_string + new_string) and past_tokens
|
|
auto new_tokens = tokenizer_->EncodeNoPrependSpace(past_string + new_string);
|
|
auto delta_string = new_string;
|
|
|
|
// Pop last token if it is a prefix of another token. That's because such tokens will often
|
|
// be rolled back in the next decode, which disturbs the distribution, so we will avoid
|
|
// generating them.
|
|
if (tokenizer_->GetPrefixTokenMask()[new_tokens.back()]) {
|
|
auto last_token = token_table[new_tokens.back()];
|
|
if (last_token.length() >= new_string.length()) {
|
|
return {0, {}, ""};
|
|
}
|
|
|
|
delta_string = delta_string.substr(0, delta_string.length() - last_token.length());
|
|
new_tokens.pop_back();
|
|
}
|
|
|
|
int first_differ_idx = past_tokens.size();
|
|
for (int i = 0; i < static_cast<int>(past_tokens.size()); ++i) {
|
|
if (i == static_cast<int>(new_tokens.size()) || past_tokens[i] != new_tokens[i]) {
|
|
first_differ_idx = i;
|
|
break;
|
|
}
|
|
}
|
|
|
|
return {past_tokens.size() - first_differ_idx,
|
|
std::vector<int32_t>(new_tokens.begin() + first_differ_idx, new_tokens.end()),
|
|
delta_string};
|
|
}
|
|
|
|
/*!
|
|
* \brief Handle rollback for the stream output, the model state and the kv cache.
|
|
* \param rsentry The request state entry.
|
|
* \param mstate The model state.
|
|
* \param rollback_cnt The number of tokens to rollback.
|
|
* \param new_tokens The new tokens. Useful for the stream output.
|
|
* \param new_string The delta string of output. Useful for the stream output.
|
|
*/
|
|
void HandleRollback(const RequestStateEntry& rsentry, RequestModelState mstate, int rollback_cnt,
|
|
const std::vector<int32_t>& new_tokens, const std::string& new_string) {
|
|
// 1. Handle rollback for the stream output
|
|
if (rollback_cnt >
|
|
static_cast<int>(mstate->committed_tokens.size()) - rsentry->next_callback_token_pos) {
|
|
const auto& token_table = tokenizer_->PostProcessedTokenTable();
|
|
for (auto i = rsentry->next_callback_token_pos; i < mstate->committed_tokens.size(); ++i) {
|
|
auto token_id = mstate->committed_tokens[i].GetTokenId();
|
|
rsentry->extra_prefix_string += token_table[token_id];
|
|
}
|
|
rsentry->extra_prefix_string += new_string;
|
|
rsentry->next_callback_token_pos = static_cast<int>(mstate->committed_tokens.size()) -
|
|
rollback_cnt + static_cast<int>(new_tokens.size());
|
|
}
|
|
|
|
// 2. Handle rollback for the model state
|
|
if (rollback_cnt > 0) {
|
|
mstate->RollbackTokens(rollback_cnt);
|
|
}
|
|
|
|
// 3. Handle rollback for the kv cache
|
|
if (rollback_cnt > mstate->num_tokens_for_next_decode) {
|
|
models_[0]->PopNFromKVCache(mstate->internal_id,
|
|
rollback_cnt - mstate->num_tokens_for_next_decode);
|
|
mstate->num_tokens_for_next_decode = 0;
|
|
} else {
|
|
mstate->num_tokens_for_next_decode -= rollback_cnt;
|
|
}
|
|
}
|
|
|
|
/*!
|
|
* \brief The model to run jump-forward decoding. When there are multiple
|
|
* models, the `Step` function of the created action will not take effect.
|
|
*/
|
|
Array<Model> models_;
|
|
/*! \brief Tokenizer for retokenization. */
|
|
Tokenizer tokenizer_;
|
|
/*! \brief Event trace recorder. */
|
|
Optional<EventTraceRecorder> trace_recorder_;
|
|
/*! \brief The maximum number of tokens to rollback. */
|
|
const int MAX_ROLLBACK_TOKENS_ = 10;
|
|
};
|
|
|
|
EngineAction EngineAction::BatchJumpForward(Array<Model> models, Tokenizer tokenizer,
|
|
Optional<EventTraceRecorder> trace_recorder) {
|
|
return EngineAction(tvm::ffi::make_object<BatchJumpForwardActionObj>(
|
|
std::move(models), std::move(tokenizer), std::move(trace_recorder)));
|
|
}
|
|
|
|
} // namespace serve
|
|
} // namespace llm
|
|
} // namespace mlc
|