wehub-resource-sync
504 lines
24 KiB
C++
504 lines
24 KiB
C++
/*!
|
|
* Copyright (c) 2023-2025 by Contributors
|
|
* \file serve/engine_actions/new_request_prefill.cc
|
|
*/
|
|
|
|
#include <tvm/support/cuda/nvtx.h>
|
|
|
|
#include "../sampler/sampler.h"
|
|
#include "batch_prefill_base.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace serve {
|
|
|
|
using tvm::support::NVTXScopedRange;
|
|
|
|
/*!
|
|
* \brief The action that prefills requests in the `waiting_queue` of
|
|
* the engine state.
|
|
* Aside from that, this action sends the computed KV data to remote
|
|
* instances after computing the KV data.
|
|
*/
|
|
class DisaggRemoteSendActionObj : public BatchPrefillBaseActionObj {
|
|
public:
|
|
explicit DisaggRemoteSendActionObj(Array<Model> models,
|
|
std::vector<ModelWorkspace> model_workspaces,
|
|
EngineConfig engine_config,
|
|
std::vector<tvm::ffi::json::Object> model_configs,
|
|
Optional<EventTraceRecorder> trace_recorder,
|
|
FRequestStreamCallback request_stream_callback, Device device)
|
|
: BatchPrefillBaseActionObj(std::move(models), std::move(engine_config),
|
|
std::move(model_configs), std::move(trace_recorder)),
|
|
model_workspaces_(std::move(model_workspaces)),
|
|
request_stream_callback_(std::move(request_stream_callback)),
|
|
device_(device) {
|
|
if (device.device_type == DLDeviceType::kDLCUDA ||
|
|
device.device_type == DLDeviceType::kDLROCM) {
|
|
// The compute stream is the default stream.
|
|
compute_stream_ = DeviceAPI::Get(device)->GetCurrentStream(device);
|
|
}
|
|
}
|
|
|
|
// Mimicked from NewRequestPrefillActionObj::Step
|
|
Array<Request> Step(EngineState estate) final {
|
|
// - Find the requests in `waiting_queue` that can prefill in this step.
|
|
std::vector<PrefillInput> prefill_inputs;
|
|
{
|
|
NVTXScopedRange nvtx_scope("DisaggRemoteSend getting requests");
|
|
prefill_inputs = GetRequestStateEntriesToPrefill(estate);
|
|
if (prefill_inputs.empty()) {
|
|
return {};
|
|
}
|
|
}
|
|
|
|
int num_rsentries = prefill_inputs.size();
|
|
{
|
|
NVTXScopedRange nvtx_scope("DisaggRemoteSend matching prefix");
|
|
for (int i = 0; i < num_rsentries; ++i) {
|
|
MatchPrefixCache(estate, &prefill_inputs[i]);
|
|
}
|
|
}
|
|
|
|
auto tstart = std::chrono::high_resolution_clock::now();
|
|
|
|
// - Update status of request states from pending to alive.
|
|
Array<String> request_ids;
|
|
std::vector<RequestState> rstates_of_entries;
|
|
std::vector<RequestStateStatus> status_before_prefill;
|
|
UpdateRequestToAlive(prefill_inputs, estate, &request_ids, &rstates_of_entries,
|
|
&status_before_prefill);
|
|
|
|
// - Get embedding and run prefill for each model.
|
|
// NOTE: we don't keep the logits as we don't run sampling in this action by design.
|
|
std::vector<int> prefill_lengths;
|
|
prefill_lengths.resize(/*size=*/num_rsentries, /*value=*/-1);
|
|
for (int model_id = 0; model_id < static_cast<int>(models_.size()); ++model_id) {
|
|
std::vector<int64_t> request_internal_ids;
|
|
request_internal_ids.reserve(num_rsentries);
|
|
ObjectRef embeddings = model_workspaces_[model_id].embeddings;
|
|
int cum_prefill_length = 0;
|
|
bool single_input =
|
|
num_rsentries == 1 && prefill_inputs[0].rsentry->mstates[model_id]->inputs.size() == 1;
|
|
std::vector<int64_t> cached_token_data;
|
|
for (int i = 0; i < num_rsentries; ++i) {
|
|
const RequestStateEntry& rsentry = prefill_inputs[i].rsentry;
|
|
RequestModelState mstate = rsentry->mstates[model_id];
|
|
auto [input_data, input_length] =
|
|
ChunkPrefillInputData(mstate, prefill_inputs[i].max_prefill_length);
|
|
if (prefill_lengths[i] == -1) {
|
|
prefill_lengths[i] = input_length;
|
|
} else {
|
|
TVM_FFI_ICHECK_EQ(prefill_lengths[i], input_length);
|
|
}
|
|
mstate->num_prefilled_tokens += input_length;
|
|
|
|
TVM_FFI_ICHECK(mstate->draft_output_tokens.empty());
|
|
TVM_FFI_ICHECK(mstate->draft_token_slots.empty());
|
|
if (status_before_prefill[i] == RequestStateStatus::kPending &&
|
|
!estate->prefix_cache->HasSequence(mstate->internal_id)) {
|
|
// Add the sequence to the model.
|
|
// If the sequence is already in prefix cache, it has also been added/forked in the
|
|
// KVCache.
|
|
TVM_FFI_ICHECK_EQ(rsentry->parent_idx, -1);
|
|
models_[model_id]->AddNewSequence(mstate->internal_id);
|
|
// Enable sliding window for the sequence if it is not a parent.
|
|
if (rsentry->child_indices.empty()) {
|
|
models_[model_id]->EnableSlidingWindowForSeq(mstate->internal_id);
|
|
}
|
|
DisaggConfig disagg_config = mstate->request->generation_cfg->debug_config.disagg_config;
|
|
TVM_FFI_ICHECK(disagg_config.dst_group_offset.has_value());
|
|
models_[model_id]->DisaggMarkKVSend(
|
|
mstate->internal_id, disagg_config.kv_window_begin.value_or(0),
|
|
disagg_config.kv_append_metadata[model_id], disagg_config.dst_group_offset.value());
|
|
}
|
|
request_internal_ids.push_back(mstate->internal_id);
|
|
RECORD_EVENT(trace_recorder_, rsentry->request->id, "start embedding");
|
|
for (int j = 0; j < static_cast<int>(input_data.size()); ++j) {
|
|
if (!model_id && !prefill_inputs[i].is_decode) {
|
|
mstate->prefilled_inputs.push_back(input_data[j]);
|
|
}
|
|
if (const auto* token_data = input_data[j].as<TokenDataNode>()) {
|
|
cached_token_data.insert(cached_token_data.end(), token_data->token_ids.begin(),
|
|
token_data->token_ids.end());
|
|
} else {
|
|
if (!cached_token_data.empty()) {
|
|
embeddings = TokenData(cached_token_data)
|
|
->GetEmbedding(models_[model_id],
|
|
/*dst=*/!single_input ? &embeddings : nullptr,
|
|
/*offset=*/cum_prefill_length);
|
|
cum_prefill_length += cached_token_data.size();
|
|
cached_token_data.clear();
|
|
}
|
|
embeddings = input_data[j]->GetEmbedding(models_[model_id],
|
|
/*dst=*/!single_input ? &embeddings : nullptr,
|
|
/*offset=*/cum_prefill_length);
|
|
cum_prefill_length += input_data[j]->GetLength();
|
|
}
|
|
}
|
|
RECORD_EVENT(trace_recorder_, rsentry->request->id, "finish embedding");
|
|
}
|
|
if (!cached_token_data.empty()) {
|
|
embeddings = TokenData(cached_token_data)
|
|
->GetEmbedding(models_[model_id],
|
|
/*dst=*/!single_input ? &embeddings : nullptr,
|
|
/*offset=*/cum_prefill_length);
|
|
cum_prefill_length += cached_token_data.size();
|
|
cached_token_data.clear();
|
|
}
|
|
|
|
RECORD_EVENT(trace_recorder_, request_ids, "start prefill");
|
|
Tensor logits =
|
|
models_[model_id]->BatchPrefill(embeddings, request_internal_ids, prefill_lengths);
|
|
RECORD_EVENT(trace_recorder_, request_ids, "finish prefill");
|
|
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
|
|
TVM_FFI_ICHECK_EQ(logits->shape[0], 1);
|
|
TVM_FFI_ICHECK_EQ(logits->shape[1], num_rsentries);
|
|
}
|
|
|
|
// - 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();
|
|
|
|
// - We run synchronize to make sure that the prefill is finished.
|
|
// We need explicit synchronization because we don't do sampling in this action.
|
|
DeviceAPI::Get(device_)->StreamSync(device_, compute_stream_);
|
|
|
|
auto tend = std::chrono::high_resolution_clock::now();
|
|
estate->metrics.engine_prefill_time_sum += static_cast<double>((tend - tstart).count()) / 1e9;
|
|
|
|
std::vector<Request> processed_requests =
|
|
RemoveProcessedRequests(prefill_inputs, estate, rstates_of_entries);
|
|
estate->running_rsentries_changed = true;
|
|
return processed_requests;
|
|
}
|
|
|
|
private:
|
|
// Mimicked from BatchPrefillBaseActionObj::GetRequestStateEntriesToPrefill
|
|
std::vector<PrefillInput> 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 {};
|
|
}
|
|
}
|
|
|
|
// Explicitly filter the waiting queue to only keep the requests
|
|
// with disaggregation request kind "kRemoteSend".
|
|
std::vector<Request> waiting_queue;
|
|
waiting_queue.reserve(estate->waiting_queue.size());
|
|
for (Request request : estate->waiting_queue) {
|
|
if (request->generation_cfg->debug_config.disagg_config.kind ==
|
|
DisaggRequestKind::kRemoteSend) {
|
|
waiting_queue.push_back(request);
|
|
}
|
|
}
|
|
if (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_running_rsentries = 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;
|
|
int total_required_pages = 0;
|
|
int num_available_pages;
|
|
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 : waiting_queue) {
|
|
NVTXScopedRange nvtx_scope("Process request " + request->id);
|
|
RequestState rstate = estate->GetRequestState(request);
|
|
TVM_FFI_ICHECK_EQ(rstate->entries.size(), 1) << "n > 1 is not supported.";
|
|
const RequestStateEntry& rsentry = rstate->entries[0];
|
|
TVM_FFI_ICHECK(!rsentry->mstates[i]->inputs.empty())
|
|
<< "The request entry must have pending inputs.";
|
|
|
|
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.
|
|
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.
|
|
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());
|
|
{
|
|
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;
|
|
}
|
|
|
|
// Copied from NewRequestPrefillActionObj::MatchPrefixCache
|
|
/*!
|
|
* \brief Match the request state entry with prefix cache, to skip prefilling common prefix
|
|
* tokens. If the request state entry is not added to KVCache yet, this method will add/fork the
|
|
* request in the KVCache, depending on the matching result from prefix cache.
|
|
* \param estate The engine state.
|
|
* \param[in, out] input The prefill input to be matched and updated.
|
|
* \return The matched length in prefix cache.
|
|
*/
|
|
int MatchPrefixCache(EngineState estate, PrefillInput* input) final {
|
|
RequestStateEntry rsentry = input->rsentry;
|
|
if (estate->prefix_cache->Mode() == PrefixCacheMode::kDisable) {
|
|
return 0;
|
|
}
|
|
if (rsentry->parent_idx == -1 && rsentry->status == RequestStateStatus::kPending &&
|
|
!estate->prefix_cache->HasSequence(rsentry->mstates[0]->internal_id)) {
|
|
std::vector<int32_t> tokens = GetConcatPrefillInputData(rsentry->mstates[0]);
|
|
if (tokens.empty()) {
|
|
// If the RequestStateEntry is of empty input data, or not fully tokenized, do nothing
|
|
// and return.
|
|
return 0;
|
|
}
|
|
PrefixCacheMatchedResult result = estate->prefix_cache->InsertSequence(
|
|
rsentry->mstates[0]->internal_id, tokens, models_[0]->GetSlidingWindowSize(),
|
|
models_[0]->GetAttentionSinkSize());
|
|
|
|
if (result.prefilled_offset == 0) {
|
|
// Add new sequence
|
|
TVM_FFI_ICHECK_EQ(result.forked_seq_id, -1);
|
|
TVM_FFI_ICHECK_EQ(result.reused_seq_id, -1);
|
|
TVM_FFI_ICHECK_EQ(result.reused_seq_pop_last_tokens, 0);
|
|
for (int model_id = 0; model_id < static_cast<int>(models_.size()); ++model_id) {
|
|
Model model = models_[model_id];
|
|
RequestModelState mstate = rsentry->mstates[model_id];
|
|
model->AddNewSequence(rsentry->mstates[0]->internal_id);
|
|
// Enable sliding window for the sequence if it is not a parent.
|
|
if (rsentry->child_indices.empty()) {
|
|
model->EnableSlidingWindowForSeq(rsentry->mstates[0]->internal_id);
|
|
}
|
|
DisaggConfig disagg_config = mstate->request->generation_cfg->debug_config.disagg_config;
|
|
models_[model_id]->DisaggMarkKVSend(
|
|
mstate->internal_id, disagg_config.kv_window_begin.value_or(0),
|
|
disagg_config.kv_append_metadata[model_id], disagg_config.dst_group_offset.value());
|
|
}
|
|
} else {
|
|
if (result.forked_seq_id != -1) {
|
|
TVM_FFI_ICHECK_EQ(result.reused_seq_id, -1);
|
|
TVM_FFI_ICHECK_EQ(result.reused_seq_pop_last_tokens, 0);
|
|
// Fork from active sequence
|
|
for (int model_id = 0; model_id < static_cast<int>(models_.size()); ++model_id) {
|
|
Model model = models_[model_id];
|
|
RequestModelState mstate = rsentry->mstates[model_id];
|
|
model->ForkSequence(result.forked_seq_id, rsentry->mstates[0]->internal_id,
|
|
result.prefilled_offset);
|
|
// Enable sliding window for the sequence if it is not a parent.
|
|
if (rsentry->child_indices.empty()) {
|
|
model->EnableSlidingWindowForSeq(rsentry->mstates[0]->internal_id);
|
|
}
|
|
DisaggConfig disagg_config =
|
|
mstate->request->generation_cfg->debug_config.disagg_config;
|
|
models_[model_id]->DisaggMarkKVSend(
|
|
mstate->internal_id, disagg_config.kv_window_begin.value_or(0),
|
|
disagg_config.kv_append_metadata[model_id], disagg_config.dst_group_offset.value());
|
|
}
|
|
} else {
|
|
// Reuse recycling sequence
|
|
TVM_FFI_ICHECK_EQ(result.forked_seq_id, -1);
|
|
estate->id_manager.RecycleId(rsentry->mstates[0]->internal_id);
|
|
for (int i = 0; i < rsentry->mstates.size(); ++i) {
|
|
rsentry->mstates[i]->internal_id = result.reused_seq_id;
|
|
}
|
|
if (result.reused_seq_pop_last_tokens > 0) {
|
|
for (Model model : models_) {
|
|
model->PopNFromKVCache(rsentry->mstates[0]->internal_id,
|
|
result.reused_seq_pop_last_tokens);
|
|
}
|
|
}
|
|
for (int model_id = 0; model_id < static_cast<int>(models_.size()); ++model_id) {
|
|
RequestModelState mstate = rsentry->mstates[model_id];
|
|
DisaggConfig disagg_config =
|
|
mstate->request->generation_cfg->debug_config.disagg_config;
|
|
models_[model_id]->DisaggMarkKVSend(
|
|
mstate->internal_id, disagg_config.kv_window_begin.value_or(0),
|
|
disagg_config.kv_append_metadata[model_id], disagg_config.dst_group_offset.value());
|
|
}
|
|
}
|
|
}
|
|
// Pop matched prefix
|
|
if (result.prefilled_offset) {
|
|
for (int i = 0; i < rsentry->mstates.size(); ++i) {
|
|
PopPrefillInputData(rsentry->mstates[i], result.prefilled_offset);
|
|
}
|
|
}
|
|
// Update max prefill length
|
|
input->max_prefill_length =
|
|
std::min(input->max_prefill_length, rsentry->mstates[0]->GetInputLength());
|
|
return result.prefilled_offset;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
/*! \brief Workspace of each model. */
|
|
std::vector<ModelWorkspace> model_workspaces_;
|
|
/*! \brief The stream callback function to passes back the sampled results after prefill. */
|
|
FRequestStreamCallback request_stream_callback_;
|
|
/*! \brief The device which we run synchronization for after prefill. */
|
|
Device device_;
|
|
/*! \brief The compute stream to run synchronization for. */
|
|
TVMStreamHandle compute_stream_ = nullptr;
|
|
};
|
|
|
|
EngineAction EngineAction::DisaggRemoteSend(
|
|
Array<Model> models, std::vector<ModelWorkspace> model_workspaces, EngineConfig engine_config,
|
|
std::vector<tvm::ffi::json::Object> model_configs, Optional<EventTraceRecorder> trace_recorder,
|
|
FRequestStreamCallback request_stream_callback, Device device) {
|
|
return EngineAction(tvm::ffi::make_object<DisaggRemoteSendActionObj>(
|
|
std::move(models), std::move(model_workspaces), std::move(engine_config),
|
|
std::move(model_configs), std::move(trace_recorder), std::move(request_stream_callback),
|
|
device));
|
|
}
|
|
|
|
} // namespace serve
|
|
} // namespace llm
|
|
} // namespace mlc
|