wehub-resource-sync
535 lines
24 KiB
C++
535 lines
24 KiB
C++
/*!
|
|
* Copyright (c) 2023-2025 by Contributors
|
|
* \file serve/engine_actions/batch_prefill_base.h
|
|
*/
|
|
|
|
#include "batch_prefill_base.h"
|
|
|
|
#include <tvm/support/cuda/nvtx.h>
|
|
|
|
#include <numeric>
|
|
|
|
#include "../../support/json_parser.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace serve {
|
|
|
|
using tvm::support::NVTXScopedRange;
|
|
|
|
bool HasPrefillSpace(int num_required_pages, bool sliding_window_enabled, int new_batch_size,
|
|
int num_available_pages, int current_total_seq_len, int total_input_length,
|
|
int max_total_sequence_length) {
|
|
return num_required_pages + (!sliding_window_enabled ? new_batch_size : 0) <=
|
|
num_available_pages &&
|
|
(sliding_window_enabled ||
|
|
current_total_seq_len + total_input_length + 8 * new_batch_size <=
|
|
max_total_sequence_length);
|
|
}
|
|
|
|
BatchPrefillBaseActionObj::BatchPrefillBaseActionObj(
|
|
Array<Model> models, EngineConfig engine_config,
|
|
std::vector<tvm::ffi::json::Object> model_configs, Optional<EventTraceRecorder> trace_recorder)
|
|
: models_(std::move(models)),
|
|
engine_config_(std::move(engine_config)),
|
|
trace_recorder_(std::move(trace_recorder)) {
|
|
TVM_FFI_ICHECK_EQ(models_.size(), model_configs.size());
|
|
sliding_window_sizes_.reserve(models_.size());
|
|
for (const tvm::ffi::json::Object& model_config : model_configs) {
|
|
// "-1" means the sliding window is disabled.
|
|
sliding_window_sizes_.push_back(
|
|
json::LookupOrDefault<int64_t>(model_config, "sliding_window_size", -1));
|
|
}
|
|
kv_state_kind_ = models_[0]->GetMetadata().kv_state_kind;
|
|
}
|
|
|
|
/*!
|
|
* \brief Find one or multiple request state entries to run prefill.
|
|
* \param estate The engine state.
|
|
* \return The request entries to prefill, together with their input lengths.
|
|
*/
|
|
std::vector<BatchPrefillBaseActionObj::PrefillInput>
|
|
BatchPrefillBaseActionObj::GetRequestStateEntriesToPrefill(EngineState estate) {
|
|
// Preempt request state entries when decode cannot apply.
|
|
const std::vector<RequestStateEntry>* running_rsentries;
|
|
{
|
|
NVTXScopedRange nvtx_scope("BatchDecode getting requests");
|
|
running_rsentries = &estate->GetRunningRequestStateEntries();
|
|
if (!(running_rsentries->size() <= models_[0]->GetNumAvailablePages())) {
|
|
// Even the decode cannot be performed.
|
|
// As a result, directly return without doing prefill.
|
|
return {};
|
|
}
|
|
}
|
|
|
|
if (estate->waiting_queue.empty()) {
|
|
// No request to prefill.
|
|
return {};
|
|
}
|
|
|
|
std::vector<std::vector<PrefillInput>> prefill_inputs_for_all_models;
|
|
prefill_inputs_for_all_models.reserve(models_.size());
|
|
|
|
int num_decode_inputs = static_cast<int>(running_rsentries->size());
|
|
|
|
// We first collect the inputs that can be prefilled for each model.
|
|
// Then we make a reduction to return the maximum common inputs.
|
|
for (int i = 0; i < static_cast<int>(models_.size()); ++i) {
|
|
std::vector<PrefillInput> prefill_inputs;
|
|
// - Try to prefill pending requests.
|
|
int total_input_length = 0;
|
|
for (const RequestStateEntry& rsentry : *running_rsentries) {
|
|
total_input_length += rsentry->mstates[i]->num_tokens_for_next_decode;
|
|
}
|
|
int total_required_pages = num_decode_inputs;
|
|
int num_available_pages;
|
|
int num_running_rsentries = num_decode_inputs;
|
|
int current_total_seq_len;
|
|
{
|
|
NVTXScopedRange nvtx_scope("KV cache GetNumAvailablePages");
|
|
num_available_pages = models_[i]->GetNumAvailablePages();
|
|
}
|
|
{
|
|
NVTXScopedRange nvtx_scope("KV cache GetCurrentTotalSequenceLength");
|
|
current_total_seq_len = models_[i]->GetCurrentTotalSequenceLength();
|
|
}
|
|
|
|
int num_prefill_rsentries = 0;
|
|
for (const Request& request : estate->waiting_queue) {
|
|
NVTXScopedRange nvtx_scope("Process request " + request->id);
|
|
if (request->generation_cfg->debug_config.disagg_config.kind != DisaggRequestKind::kNone) {
|
|
continue;
|
|
}
|
|
RequestState rstate = estate->GetRequestState(request);
|
|
bool prefill_stops = false;
|
|
for (const RequestStateEntry& rsentry : rstate->entries) {
|
|
// A request state entry can be prefilled only when:
|
|
// - it has inputs, and
|
|
// - it has no parent or its parent is alive and has no remaining input.
|
|
if (rsentry->mstates[i]->inputs.empty() ||
|
|
(rsentry->parent_idx != -1 &&
|
|
(rstate->entries[rsentry->parent_idx]->status == RequestStateStatus::kPending ||
|
|
!rstate->entries[rsentry->parent_idx]->mstates[i]->inputs.empty()))) {
|
|
continue;
|
|
}
|
|
|
|
int input_length = rsentry->mstates[i]->GetInputLength();
|
|
int num_require_pages = (input_length + engine_config_->kv_cache_page_size - 1) /
|
|
engine_config_->kv_cache_page_size;
|
|
bool sliding_window_enabled = sliding_window_sizes_[i] != -1;
|
|
int num_required_pages_under_sliding_window = std::numeric_limits<int>::max();
|
|
if (sliding_window_enabled) {
|
|
// Sliding window for model i is enabled.
|
|
int max_single_request_page_requirement =
|
|
1 + (sliding_window_sizes_[i] + engine_config_->kv_cache_page_size - 1) /
|
|
engine_config_->kv_cache_page_size;
|
|
int num_total_prefilled_tokens = rsentry->mstates[i]->num_prefilled_tokens;
|
|
int parent_ptr = rsentry->parent_idx;
|
|
while (parent_ptr != -1) {
|
|
num_total_prefilled_tokens +=
|
|
rstate->entries[parent_ptr]->mstates[i]->num_prefilled_tokens;
|
|
parent_ptr = rstate->entries[parent_ptr]->parent_idx;
|
|
}
|
|
|
|
int num_pages_in_use = (std::min(num_total_prefilled_tokens, sliding_window_sizes_[i]) +
|
|
engine_config_->kv_cache_page_size - 1) /
|
|
engine_config_->kv_cache_page_size;
|
|
num_required_pages_under_sliding_window =
|
|
max_single_request_page_requirement - num_pages_in_use;
|
|
num_require_pages = std::min(num_require_pages, num_required_pages_under_sliding_window);
|
|
TVM_FFI_ICHECK_GE(num_require_pages, 0);
|
|
}
|
|
|
|
total_input_length += input_length;
|
|
total_required_pages += num_require_pages;
|
|
// - Attempt 1. Check if the entire request state entry can fit for prefill.
|
|
bool can_prefill = false;
|
|
{
|
|
NVTXScopedRange nvtx_scope("Attempt 1");
|
|
for (int num_child_to_activate = rsentry->child_indices.size();
|
|
num_child_to_activate >= 0; --num_child_to_activate) {
|
|
while (!HasPrefillSpace(total_required_pages, sliding_window_enabled,
|
|
(num_running_rsentries + num_prefill_rsentries),
|
|
num_available_pages, current_total_seq_len, total_input_length,
|
|
engine_config_->max_total_sequence_length)) {
|
|
if (!estate->prefix_cache->TryFreeMemory()) break;
|
|
// Update number of available pages after memory free.
|
|
num_available_pages = models_[i]->GetNumAvailablePages();
|
|
current_total_seq_len = models_[i]->GetCurrentTotalSequenceLength();
|
|
}
|
|
if (CanPrefill(estate, num_prefill_rsentries + 1 + num_child_to_activate,
|
|
total_input_length, total_required_pages, num_available_pages,
|
|
current_total_seq_len, num_running_rsentries, kv_state_kind_,
|
|
sliding_window_enabled)) {
|
|
prefill_inputs.push_back(
|
|
{rsentry, input_length, num_child_to_activate, /*is_decode=*/false});
|
|
num_prefill_rsentries += 1 + num_child_to_activate;
|
|
can_prefill = true;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
if (can_prefill) {
|
|
continue;
|
|
}
|
|
total_input_length -= input_length;
|
|
total_required_pages -= num_require_pages;
|
|
|
|
// - Attempt 2. Check if the request state entry can partially fit by input chunking.
|
|
TVM_FFI_ICHECK_LE(total_input_length, engine_config_->prefill_chunk_size);
|
|
if (engine_config_->prefill_chunk_size - total_input_length >= input_length ||
|
|
engine_config_->prefill_chunk_size == total_input_length) {
|
|
// 1. If the input length can fit the remaining prefill chunk size,
|
|
// it means the failure of attempt 1 is not because of the input
|
|
// length being too long, and thus chunking does not help.
|
|
// 2. If the total input length already reaches the prefill chunk size,
|
|
// the current request state entry will not be able to be processed.
|
|
// So we can safely return in either case.
|
|
prefill_stops = true;
|
|
break;
|
|
}
|
|
input_length = engine_config_->prefill_chunk_size - total_input_length;
|
|
num_require_pages = (input_length + engine_config_->kv_cache_page_size - 1) /
|
|
engine_config_->kv_cache_page_size;
|
|
if (sliding_window_enabled) {
|
|
// Sliding window for model i is enabled.
|
|
num_require_pages = std::min(num_require_pages, num_required_pages_under_sliding_window);
|
|
TVM_FFI_ICHECK_GE(num_require_pages, 0);
|
|
}
|
|
|
|
{
|
|
NVTXScopedRange nvtx_scope("Attempt 2");
|
|
total_input_length += input_length;
|
|
total_required_pages += num_require_pages;
|
|
if (CanPrefill(estate, num_prefill_rsentries + 1, total_input_length,
|
|
total_required_pages, num_available_pages, current_total_seq_len,
|
|
num_running_rsentries, kv_state_kind_, sliding_window_enabled)) {
|
|
prefill_inputs.push_back({rsentry, input_length, 0, /*is_decode=*/false});
|
|
}
|
|
}
|
|
|
|
// - Prefill stops here.
|
|
prefill_stops = true;
|
|
break;
|
|
}
|
|
if (prefill_stops) {
|
|
break;
|
|
}
|
|
}
|
|
prefill_inputs_for_all_models.push_back(prefill_inputs);
|
|
}
|
|
|
|
// Reduce over the prefill inputs of all models.
|
|
TVM_FFI_ICHECK(!prefill_inputs_for_all_models.empty());
|
|
int num_prefill_inputs = prefill_inputs_for_all_models[0].size();
|
|
for (int i = 1; i < static_cast<int>(prefill_inputs_for_all_models.size()); ++i) {
|
|
num_prefill_inputs =
|
|
std::min(num_prefill_inputs, static_cast<int>(prefill_inputs_for_all_models[i].size()));
|
|
}
|
|
|
|
if (num_prefill_inputs == 0) {
|
|
return {};
|
|
}
|
|
|
|
// Add the decode requests to the prefill inputs if prefill mode is hybrid.
|
|
std::vector<PrefillInput> prefill_inputs(prefill_inputs_for_all_models[0].begin(),
|
|
prefill_inputs_for_all_models[0].end());
|
|
if (engine_config_->prefill_mode == PrefillMode::kHybrid) {
|
|
prefill_inputs.reserve(num_decode_inputs + num_prefill_inputs);
|
|
for (const RequestStateEntry& rsentry : *running_rsentries) {
|
|
prefill_inputs.push_back(
|
|
{rsentry, rsentry->mstates[0]->num_tokens_for_next_decode, 0, /*is_decode=*/true});
|
|
}
|
|
}
|
|
{
|
|
NVTXScopedRange nvtx_scope("reduction");
|
|
for (int i = 1; i < static_cast<int>(prefill_inputs_for_all_models.size()); ++i) {
|
|
// Prefill input lengths except the last one are supposed to be the same for all models.
|
|
for (int j = 0; j < num_prefill_inputs - 1; ++j) {
|
|
TVM_FFI_ICHECK(
|
|
prefill_inputs_for_all_models[i][j].rsentry.same_as(prefill_inputs[j].rsentry));
|
|
TVM_FFI_ICHECK_EQ(prefill_inputs_for_all_models[i][j].max_prefill_length,
|
|
prefill_inputs[j].max_prefill_length);
|
|
prefill_inputs[j].num_child_to_activate =
|
|
std::min(prefill_inputs[j].num_child_to_activate,
|
|
prefill_inputs_for_all_models[i][j].num_child_to_activate);
|
|
}
|
|
// The input length of the last input is the minimum among all models.
|
|
TVM_FFI_ICHECK(prefill_inputs_for_all_models[i][num_prefill_inputs - 1].rsentry.same_as(
|
|
prefill_inputs[num_prefill_inputs - 1].rsentry));
|
|
prefill_inputs[num_prefill_inputs - 1].max_prefill_length =
|
|
std::min(prefill_inputs[num_prefill_inputs - 1].max_prefill_length,
|
|
prefill_inputs_for_all_models[i][num_prefill_inputs - 1].max_prefill_length);
|
|
prefill_inputs[num_prefill_inputs - 1].num_child_to_activate =
|
|
std::min(prefill_inputs[num_prefill_inputs - 1].num_child_to_activate,
|
|
prefill_inputs_for_all_models[i][num_prefill_inputs - 1].num_child_to_activate);
|
|
}
|
|
}
|
|
|
|
return prefill_inputs;
|
|
}
|
|
|
|
bool BatchPrefillBaseActionObj::CanPrefill(EngineState estate, int num_prefill_rsentries,
|
|
int total_input_length, int num_required_pages,
|
|
int num_available_pages, int current_total_seq_len,
|
|
int num_running_rsentries, KVStateKind kv_state_kind,
|
|
bool sliding_window_enabled) {
|
|
TVM_FFI_ICHECK_LE(num_running_rsentries, engine_config_->max_num_sequence);
|
|
|
|
// For pure RNN State, it can prefill as long as it can be instantiated.
|
|
// Hybrid uses KVCache for capacity (PagedKVCache is the constraining factor).
|
|
if (kv_state_kind == KVStateKind::kRNNState || kv_state_kind == KVStateKind::kNone) {
|
|
return true;
|
|
}
|
|
|
|
// No exceeding of the maximum allowed requests that can
|
|
// run simultaneously.
|
|
int spec_factor = engine_config_->speculative_mode != SpeculativeMode::kDisable
|
|
? (estate->spec_draft_length + 1)
|
|
: 1;
|
|
if ((num_running_rsentries + num_prefill_rsentries) * spec_factor >
|
|
std::min(static_cast<int64_t>(engine_config_->max_num_sequence),
|
|
engine_config_->prefill_chunk_size)) {
|
|
return false;
|
|
}
|
|
|
|
// NOTE: The conditions are heuristic and can be revised.
|
|
// Cond 1: total input length <= prefill chunk size.
|
|
// Cond 2: at least one decode can be performed after prefill.
|
|
// Cond 3: number of total tokens after 8 times of decode does not
|
|
// exceed the limit, where 8 is a watermark number can
|
|
// be configured and adjusted in the future.
|
|
return total_input_length <= engine_config_->prefill_chunk_size &&
|
|
HasPrefillSpace(num_required_pages, sliding_window_enabled,
|
|
(num_running_rsentries + num_prefill_rsentries), num_available_pages,
|
|
current_total_seq_len, total_input_length,
|
|
engine_config_->max_total_sequence_length);
|
|
}
|
|
|
|
/*!
|
|
* \brief Chunk the input of the given RequestModelState for prefill
|
|
* with regard to the provided maximum allowed prefill length.
|
|
* Return the list of input for prefill and the total prefill length.
|
|
* The `inputs` field of the given `mstate` will be mutated to exclude
|
|
* the returned input.
|
|
* \param mstate The RequestModelState whose input data is to be chunked.
|
|
* \param max_prefill_length The maximum allowed prefill length for the mstate.
|
|
* \return The list of input for prefill and the total prefill length.
|
|
*/
|
|
std::pair<Array<Data>, int> BatchPrefillBaseActionObj::ChunkPrefillInputData(
|
|
const RequestModelState& mstate, int max_prefill_length) {
|
|
if (mstate->inputs.empty()) {
|
|
// If the request is a hybrid decode request
|
|
TVM_FFI_ICHECK(mstate->num_tokens_for_next_decode > 0);
|
|
int num_tokens = mstate->num_tokens_for_next_decode;
|
|
mstate->num_tokens_for_next_decode = 0;
|
|
std::vector<int32_t> decode_tokens;
|
|
decode_tokens.reserve(num_tokens);
|
|
for (auto begin = mstate->committed_tokens.end() - num_tokens;
|
|
begin != mstate->committed_tokens.end(); ++begin) {
|
|
decode_tokens.push_back(begin->GetTokenId());
|
|
}
|
|
return {{TokenData(decode_tokens)}, num_tokens};
|
|
}
|
|
TVM_FFI_ICHECK(!mstate->inputs.empty());
|
|
std::vector<Data> inputs;
|
|
int cum_input_length = 0;
|
|
inputs.reserve(mstate->inputs.size());
|
|
for (int i = 0; i < static_cast<int>(mstate->inputs.size()); ++i) {
|
|
inputs.push_back(mstate->inputs[i]);
|
|
int input_length = mstate->inputs[i]->GetLength();
|
|
cum_input_length += input_length;
|
|
// Case 0. the cumulative input length does not reach the maximum prefill length.
|
|
if (cum_input_length < max_prefill_length) {
|
|
continue;
|
|
}
|
|
|
|
// Case 1. the cumulative input length equals the maximum prefill length.
|
|
if (cum_input_length == max_prefill_length) {
|
|
if (i == static_cast<int>(mstate->inputs.size()) - 1) {
|
|
// - If `i` is the last input, we just copy and reset `mstate->inputs`.
|
|
mstate->inputs.clear();
|
|
} else {
|
|
// - Otherwise, set the new input array.
|
|
mstate->inputs = Array<Data>{mstate->inputs.begin() + i + 1, mstate->inputs.end()};
|
|
}
|
|
return {inputs, cum_input_length};
|
|
}
|
|
|
|
// Case 2. cum_input_length > max_prefill_length
|
|
// The input `i` itself needs chunking if it is TokenData,
|
|
// or otherwise it cannot be chunked.
|
|
Data input = mstate->inputs[i];
|
|
inputs.pop_back();
|
|
cum_input_length -= input_length;
|
|
const auto* token_input = input.as<TokenDataNode>();
|
|
if (token_input == nullptr) {
|
|
// Cannot chunk the input.
|
|
if (i != 0) {
|
|
mstate->inputs = Array<Data>{mstate->inputs.begin() + i, mstate->inputs.end()};
|
|
}
|
|
return {inputs, cum_input_length};
|
|
}
|
|
|
|
// Split the token data into two parts.
|
|
// Return the first part for prefill, and keep the second part.
|
|
int chunked_input_length = max_prefill_length - cum_input_length;
|
|
TVM_FFI_ICHECK_GT(input_length, chunked_input_length);
|
|
TokenData chunked_input(Shape{token_input->token_ids.begin(),
|
|
token_input->token_ids.begin() + chunked_input_length});
|
|
TokenData remaining_input(
|
|
Shape{token_input->token_ids.begin() + chunked_input_length, token_input->token_ids.end()});
|
|
inputs.push_back(chunked_input);
|
|
cum_input_length += chunked_input_length;
|
|
std::vector<Data> remaining_inputs{mstate->inputs.begin() + i + 1, mstate->inputs.end()};
|
|
remaining_inputs.insert(remaining_inputs.begin(), remaining_input);
|
|
mstate->inputs = remaining_inputs;
|
|
return {inputs, cum_input_length};
|
|
}
|
|
|
|
TVM_FFI_ICHECK(false) << "Cannot reach here";
|
|
}
|
|
|
|
void BatchPrefillBaseActionObj::UpdateRequestToAlive(
|
|
const std::vector<BatchPrefillBaseActionObj::PrefillInput>& prefill_inputs,
|
|
const EngineState& estate, Array<String>* request_ids,
|
|
std::vector<RequestState>* rstates_of_entries,
|
|
std::vector<RequestStateStatus>* status_before_prefill) {
|
|
int num_rsentries = prefill_inputs.size();
|
|
request_ids->reserve(num_rsentries);
|
|
rstates_of_entries->reserve(num_rsentries);
|
|
status_before_prefill->reserve(num_rsentries);
|
|
for (const PrefillInput& prefill_input : prefill_inputs) {
|
|
const RequestStateEntry& rsentry = prefill_input.rsentry;
|
|
const Request& request = rsentry->request;
|
|
RequestState request_rstate = estate->GetRequestState(request);
|
|
request_ids->push_back(request->id);
|
|
status_before_prefill->push_back(rsentry->status);
|
|
rsentry->status = RequestStateStatus::kAlive;
|
|
|
|
if (status_before_prefill->back() == RequestStateStatus::kPending) {
|
|
// - Add the request to running queue if the request state
|
|
// status was pending and all its request states were pending.
|
|
bool alive_state_existed = false;
|
|
for (const RequestStateEntry& rsentry_ : request_rstate->entries) {
|
|
if (rsentry_->status == RequestStateStatus::kAlive && !rsentry_.same_as(rsentry)) {
|
|
alive_state_existed = true;
|
|
}
|
|
}
|
|
if (!alive_state_existed) {
|
|
estate->running_queue.push_back(request);
|
|
}
|
|
}
|
|
rstates_of_entries->push_back(std::move(request_rstate));
|
|
}
|
|
}
|
|
|
|
std::vector<Request> BatchPrefillBaseActionObj::RemoveProcessedRequests(
|
|
const std::vector<BatchPrefillBaseActionObj::PrefillInput>& prefill_inputs,
|
|
const EngineState& estate, const std::vector<RequestState>& rstates_of_entries) {
|
|
// - Remove the request from waiting queue if all its request states
|
|
// are now alive and have no remaining chunked inputs.
|
|
std::vector<Request> processed_requests;
|
|
int num_rsentries = prefill_inputs.size();
|
|
processed_requests.reserve(num_rsentries);
|
|
std::unordered_set<const RequestNode*> dedup_map;
|
|
for (int i = 0; i < num_rsentries; ++i) {
|
|
const RequestStateEntry& rsentry = prefill_inputs[i].rsentry;
|
|
if (dedup_map.find(rsentry->request.operator->()) != dedup_map.end()) {
|
|
continue;
|
|
}
|
|
dedup_map.insert(rsentry->request.operator->());
|
|
processed_requests.push_back(rsentry->request);
|
|
|
|
bool pending_state_exists = false;
|
|
for (const RequestStateEntry& rsentry_ : rstates_of_entries[i]->entries) {
|
|
if (rsentry_->status == RequestStateStatus::kPending ||
|
|
!rsentry_->mstates[0]->inputs.empty()) {
|
|
pending_state_exists = true;
|
|
break;
|
|
}
|
|
}
|
|
if (!pending_state_exists &&
|
|
std::find(estate->waiting_queue.begin(), estate->waiting_queue.end(), rsentry->request) !=
|
|
estate->waiting_queue.end()) {
|
|
auto it =
|
|
std::find(estate->waiting_queue.begin(), estate->waiting_queue.end(), rsentry->request);
|
|
if (it != estate->waiting_queue.end()) {
|
|
estate->waiting_queue.erase(it);
|
|
}
|
|
}
|
|
}
|
|
return processed_requests;
|
|
}
|
|
|
|
void BatchPrefillBaseActionObj::UpdateRequestStateEntriesWithSampleResults(
|
|
const std::vector<RequestStateEntry>& rsentries_for_sample,
|
|
const std::vector<bool>& rsentry_activated, const std::vector<SampleResult>& sample_results) {
|
|
auto tnow = std::chrono::high_resolution_clock::now();
|
|
for (int i = 0; i < static_cast<int>(rsentries_for_sample.size()); ++i) {
|
|
// If the request is a hybrid decode request
|
|
if (rsentries_for_sample[i]->status == RequestStateStatus::kAlive &&
|
|
rsentries_for_sample[i]->child_indices.empty() &&
|
|
rsentries_for_sample[i]->mstates[0]->inputs.empty()) {
|
|
for (const RequestModelState& mstate : rsentries_for_sample[i]->mstates) {
|
|
TVM_FFI_ICHECK(!mstate->require_retokenization_in_next_decode);
|
|
mstate->CommitToken(sample_results[i]);
|
|
// live update the output metrics
|
|
rsentries_for_sample[i]->rstate->metrics.completion_tokens += 1;
|
|
rsentries_for_sample[i]->rstate->metrics.prefill_end_time_point = tnow;
|
|
}
|
|
continue;
|
|
}
|
|
|
|
// Update all model states of the request state entry.
|
|
for (const RequestModelState& mstate : rsentries_for_sample[i]->mstates) {
|
|
mstate->CommitToken(sample_results[i]);
|
|
if (!rsentry_activated[i]) {
|
|
// When the child rsentry is not activated,
|
|
// add the sampled token as an input of the mstate for prefill.
|
|
mstate->inputs.push_back(TokenData(std::vector<int64_t>{sample_results[i].GetTokenId()}));
|
|
}
|
|
}
|
|
// prefill has finished
|
|
if (rsentries_for_sample[i]->mstates[0]->committed_tokens.size() == 1) {
|
|
TVM_FFI_ICHECK(rsentries_for_sample[i]->rstate != nullptr);
|
|
rsentries_for_sample[i]->rstate->metrics.prefill_end_time_point = tnow;
|
|
}
|
|
}
|
|
}
|
|
|
|
std::vector<int32_t> BatchPrefillBaseActionObj::GetConcatPrefillInputData(
|
|
const RequestModelState& mstate) {
|
|
std::vector<int32_t> tokens;
|
|
for (Data data : mstate->inputs) {
|
|
if (const TokenDataNode* token_data = data.as<TokenDataNode>()) {
|
|
tokens.reserve(tokens.size() + token_data->GetLength());
|
|
tokens.insert(tokens.end(), token_data->token_ids.begin(), token_data->token_ids.end());
|
|
} else {
|
|
return {};
|
|
}
|
|
}
|
|
return tokens;
|
|
}
|
|
|
|
void BatchPrefillBaseActionObj::PopPrefillInputData(const RequestModelState& mstate,
|
|
size_t num_tokens) {
|
|
while (mstate->inputs[0]->GetLength() <= num_tokens) {
|
|
num_tokens -= mstate->inputs[0]->GetLength();
|
|
mstate->inputs.erase(mstate->inputs.begin());
|
|
}
|
|
if (num_tokens) {
|
|
const TokenDataNode* token_data = mstate->inputs[0].as<TokenDataNode>();
|
|
std::vector<int32_t> tokens;
|
|
tokens.reserve(token_data->GetLength() - num_tokens);
|
|
tokens.insert(tokens.begin(), token_data->token_ids.begin() + num_tokens,
|
|
token_data->token_ids.end());
|
|
mstate->inputs.erase(mstate->inputs.begin());
|
|
mstate->inputs.insert(mstate->inputs.begin(), TokenData(tokens));
|
|
}
|
|
}
|
|
|
|
} // namespace serve
|
|
} // namespace llm
|
|
} // namespace mlc
|