wehub/mlc-ai--mlc-llm
1
0
Fork 0
Code Issues Pull Requests Actions 9 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
mlc-ai--mlc-llm/cpp/serve/model.cc
T
wehub-resource-sync 770d92cb1f
Lint / lint (push) Waiting to run
Details
Windows CI / Windows (push) Waiting to run
Details
Build Docs / Deploy Docs (push) Waiting to run
Details
chore: import upstream snapshot with attribution
2026-07-13 13:23:58 +08:00

1294 lines
56 KiB
C++
Raw Permalink Blame History

/*!
* Copyright (c) 2023-2025 by Contributors
* \file serve/model.cc
* \brief The implementation of runtime module of LLM functions (prefill/decode/etc.)
*/
#include "model.h"
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/memory/memory_manager.h>
#include <tvm/support/cuda/nvtx.h>
#include <fstream>
#include <unordered_set>
#include "../support/json_parser.h"
#include "../support/vlm_utils.h"
#include "config.h"
#include "logit_processor.h"
namespace mlc {
namespace llm {
namespace serve {
using tvm::support::NVTXScopedRange;
/*********************** Model Implementation ***********************/
TVM_FFI_STATIC_INIT_BLOCK() { ModelObj::RegisterReflection(); }
class ModelImpl;
Model Model::Create(String reload_lib_path, String model_path,
const tvm::ffi::json::Object& model_config, DLDevice device,
const Optional<Session>& session, int num_shards, int num_stages,
bool trace_enabled) {
return Model(tvm::ffi::make_object<ModelImpl>(reload_lib_path, model_path, model_config, device,
session, num_shards, num_stages, trace_enabled));
}
Result<tvm::ffi::json::Object> Model::LoadModelConfig(const String& model_path) {
using TResult = Result<tvm::ffi::json::Object>;
std::ifstream config_istream((model_path + "/mlc-chat-config.json").c_str());
std::ostringstream config_ostream;
TVM_FFI_ICHECK(config_istream);
config_ostream << config_istream.rdbuf();
std::string config_str = config_ostream.str();
tvm::ffi::String err;
auto config_json = tvm::ffi::json::Parse(config_str, &err);
if (!err.empty()) {
return TResult::Error(std::string(err));
}
auto opt = config_json.try_cast<tvm::ffi::json::Object>();
if (!opt.has_value()) {
return TResult::Error("Expected JSON object in model config");
}
return TResult::Ok(*opt);
}
class ModelImpl : public ModelObj {
public:
/*!
* \brief Constructor of ModelImpl.
* \sa Model::Create
*/
explicit ModelImpl(String reload_lib_path, String model_path, tvm::ffi::json::Object model_config,
DLDevice device, const Optional<Session>& session, int num_shards,
int num_stages, bool trace_enabled)
: model_(model_path), device_(device), trace_enabled_(trace_enabled) {
// Step 1. Process model config json string.
LoadModelConfigJSON(model_config);
// Step 2. Initialize vm, we use the packed function mechanism
// so there is no explicit abi dependency on these extra
// classes other than basic tvm runtime.
this->ft_.Init(reload_lib_path, device_, model_config, session, num_shards, num_stages);
this->num_shards_ = ft_.model_metadata_.tensor_parallel_shards;
this->num_stages_ = ft_.model_metadata_.pipeline_parallel_stages;
this->seqlen_padding_factor_ = ft_.model_metadata_.seqlen_padding_factor;
// Step 3. Reset
this->Reset();
// Step 4. Set model type
this->kind = GetMetadata().kv_state_kind;
}
/*********************** Model Computation ***********************/
ObjectRef TokenEmbed(Shape token_ids, ObjectRef* dst, int offset) final {
NVTXScopedRange nvtx_scope("TokenEmbed");
int num_tokens = token_ids.size();
if (seqlen_padding_factor_ > 1) {
num_tokens = (offset + num_tokens + seqlen_padding_factor_ - 1) / seqlen_padding_factor_ *
seqlen_padding_factor_;
}
// Copy input token ids to device.
DLDataType dtype(DLDataType{kDLInt, 32, 1});
Tensor token_ids_nd;
{
NVTXScopedRange nvtx_scope("Allocate token_ids at offset");
token_ids_nd = token_ids_storage_->AllocTensor(offset * 4, {num_tokens}, dtype);
int* p_token_ids = static_cast<int*>(token_ids_nd->data) + (token_ids_nd->byte_offset) / 4;
for (int i = 0; i < static_cast<int>(token_ids.size()); ++i) {
p_token_ids[i] = token_ids[i];
}
for (int i = static_cast<int>(token_ids.size()); i < num_tokens; ++i) {
p_token_ids[i] = 0;
}
}
TVM_FFI_ICHECK_EQ(token_ids_nd->ndim, 1);
TVM_FFI_ICHECK_EQ(token_ids_nd->shape[0], num_tokens);
TVM_FFI_ICHECK_NE(prefill_chunk_size_, -1);
ObjectRef token_ids_dref_or_nd;
{
NVTXScopedRange nvtx_scope("Copy to worker 0");
token_ids_dref_or_nd = ft_.CopyToWorker0(token_ids_nd, "token_ids", {prefill_chunk_size_});
}
ObjectRef embeddings = ft_.embed_func_(token_ids_dref_or_nd, params_).cast<ObjectRef>();
if (dst != nullptr) {
TVM_FFI_ICHECK(dst->defined());
ft_.nd_copy_embedding_to_offset_func_(embeddings, *dst, offset);
return *dst;
} else {
TVM_FFI_ICHECK_EQ(offset, 0);
return embeddings;
}
}
ObjectRef ImageEmbed(const Tensor& image, ObjectRef* dst, int offset) final {
NVTXScopedRange nvtx_scope("ImageEmbed");
TVM_FFI_ICHECK(ft_.image_embed_func_.defined())
<< "`image_embed` function is not found in the model. ";
int tmp_h = 0, tmp_w = 0;
CalculateResizeShape(image, this->model_type_, &tmp_h, &tmp_w);
Shape resize_h = {tmp_h};
Shape resize_w = {tmp_w};
CalculateCropShape(image, this->model_type_, &tmp_h, &tmp_w);
Shape crop_h = {tmp_h};
Shape crop_w = {tmp_w};
auto image_dref_or_nd = ft_.CopyToWorker0(image, "image", image.Shape());
ObjectRef embeddings =
ft_.image_embed_func_(image_dref_or_nd, resize_h, resize_w, crop_h, crop_w, params_)
.cast<ObjectRef>();
if (dst != nullptr) {
TVM_FFI_ICHECK(dst->defined());
ft_.nd_copy_embedding_to_offset_func_(embeddings, *dst, offset);
return *dst;
} else {
TVM_FFI_ICHECK_EQ(offset, 0);
return embeddings;
}
}
bool CanGetLogits() final {
return ft_.get_logits_func_.defined() && ft_.batch_get_logits_func_.defined();
}
Tensor GetLogits(const ObjectRef& hidden_states) final {
NVTXScopedRange nvtx_scope("GetLogits");
TVM_FFI_ICHECK(ft_.get_logits_func_.defined())
<< "`get_logits` function is not found in the model.";
ObjectRef hidden_states_dref_or_nd{nullptr};
if (!ft_.use_disco && hidden_states->IsInstance<DRefObj>()) {
hidden_states_dref_or_nd =
hidden_states.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<ObjectRef>();
} else {
hidden_states_dref_or_nd = hidden_states;
}
ObjectRef ret = ft_.get_logits_func_(hidden_states_dref_or_nd, params_).cast<ObjectRef>();
if (trace_enabled_) {
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
}
Tensor logits{nullptr};
if (ft_.use_disco) {
logits = ret.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<Tensor>();
} else {
logits = ret.as_or_throw<Tensor>();
}
// logits: (b * s, v)
return logits;
}
Array<Tensor> GetMultiStepLogits(const ObjectRef& hidden_states) final {
NVTXScopedRange nvtx_scope("GetMultiStepLogits");
TVM_FFI_ICHECK(ft_.get_logits_func_.defined())
<< "`get_logits` function is not found in the model.";
ObjectRef hidden_states_dref_or_nd{nullptr};
ObjectRef ret = ft_.get_logits_func_(hidden_states, params_).cast<ObjectRef>();
Array<Tensor> logits{nullptr};
if (ft_.use_disco) {
logits = ret.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<Array<Tensor>>();
} else {
logits = ret.as_or_throw<Array<Tensor>>();
}
return logits;
}
ObjectRef FuseEmbedHidden(const ObjectRef& embeddings, const ObjectRef& previous_hidden_states,
int batch_size, int seq_len) final {
NVTXScopedRange nvtx_scope("FuseEmbedHidden");
ObjectRef embeddings_dref_or_nd{nullptr};
if (!embeddings->IsInstance<DRefObj>()) {
// embeddings: (n, h)
Tensor embeddings_nd = embeddings.as_or_throw<Tensor>();
TVM_FFI_ICHECK_NE(hidden_size_, -1);
TVM_FFI_ICHECK_EQ(embeddings_nd->ndim, 2);
TVM_FFI_ICHECK_GE(embeddings_nd->shape[0], batch_size * seq_len);
TVM_FFI_ICHECK_EQ(embeddings_nd->shape[1], hidden_size_);
embeddings_dref_or_nd =
embeddings_nd.CreateView({batch_size * seq_len, hidden_size_}, embeddings_nd->dtype);
} else {
Shape embedding_shape{batch_size * seq_len, hidden_size_};
embeddings_dref_or_nd = ft_.nd_view_func_(embeddings, embedding_shape).cast<ObjectRef>();
}
ObjectRef previous_hidden_states_dref_or_nd{nullptr};
if (!ft_.use_disco && previous_hidden_states->IsInstance<DRefObj>()) {
previous_hidden_states_dref_or_nd =
previous_hidden_states.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<ObjectRef>();
} else {
previous_hidden_states_dref_or_nd = previous_hidden_states;
}
ObjectRef fused = ft_.fuse_embed_hidden_func_(embeddings_dref_or_nd,
previous_hidden_states_dref_or_nd, params_)
.cast<ObjectRef>();
if (trace_enabled_) {
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
}
Shape out_shape{batch_size, seq_len, hidden_size_};
if (ft_.use_disco) {
return ft_.nd_view_func_(fused, out_shape).cast<ObjectRef>();
} else {
Tensor fused_nd = fused.as_or_throw<Tensor>();
TVM_FFI_ICHECK_EQ(fused_nd->ndim, 2);
TVM_FFI_ICHECK_EQ(fused_nd->shape[0], batch_size * seq_len);
return fused_nd.CreateView(out_shape, fused_nd->dtype);
}
}
Tensor BatchPrefill(const ObjectRef& embeddings, const std::vector<int64_t>& seq_ids,
const std::vector<int>& lengths) final {
TVM_FFI_ICHECK(!seq_ids.empty());
TVM_FFI_ICHECK_EQ(seq_ids.size(), lengths.size());
int num_sequences = seq_ids.size();
int total_length = 0;
int* p_logit_pos = static_cast<int*>(logit_pos_arr_->data);
for (int i = 0; i < num_sequences; ++i) {
total_length += lengths[i];
p_logit_pos[i] = total_length - 1;
}
bool padded = total_length % seqlen_padding_factor_ != 0;
if (padded) {
total_length = (total_length + seqlen_padding_factor_ - 1) / seqlen_padding_factor_ *
seqlen_padding_factor_;
}
NVTXScopedRange nvtx_scope("BatchPrefill num_seq=" + std::to_string(num_sequences) +
" total_len=" + std::to_string(total_length));
Tensor logit_pos_nd = logit_pos_arr_.CreateView({num_sequences}, DLDataType{kDLInt, 32, 1});
TVM_FFI_ICHECK(ft_.prefill_func_.defined())
<< "`prefill_with_embed` function is not found in the model. Please make sure the model is "
"compiled with flag `--sep-embed` and `--enable-batching`";
TVM_FFI_ICHECK(ft_.kv_cache_begin_forward_func_.defined());
TVM_FFI_ICHECK(ft_.kv_cache_end_forward_func_.defined());
TVM_FFI_ICHECK(kv_cache_.defined()) << "KV cache has not been initialized.";
// Begin forward with the sequence ids and new lengths.
Shape seq_ids_tuple(seq_ids);
Shape lengths_tuple(lengths.begin(), lengths.end());
ft_.kv_cache_begin_forward_func_(kv_cache_, seq_ids_tuple, lengths_tuple);
if (kind == KVStateKind::kHybrid) {
TVM_FFI_ICHECK(rnn_state_.defined()) << "RNN state has not been initialized.";
ft_.kv_cache_begin_forward_func_(rnn_state_, seq_ids_tuple, lengths_tuple);
}
ObjectRef embeddings_dref_or_nd;
if (!embeddings->IsInstance<DRefObj>()) {
// embeddings: (1, n, h)
Tensor embeddings_nd = embeddings.as_or_throw<Tensor>();
TVM_FFI_ICHECK_NE(hidden_size_, -1);
TVM_FFI_ICHECK_EQ(embeddings_nd->ndim, 2);
TVM_FFI_ICHECK_GE(embeddings_nd->shape[0], total_length);
TVM_FFI_ICHECK_EQ(embeddings_nd->shape[1], hidden_size_);
TVM_FFI_ICHECK_EQ(embeddings_nd->device.device_type, device_.device_type);
TVM_FFI_ICHECK_EQ(embeddings_nd->device.device_id, device_.device_id);
embeddings_dref_or_nd =
embeddings_nd.CreateView({1, total_length, hidden_size_}, embeddings_nd->dtype);
} else {
Shape embedding_shape{1, total_length, hidden_size_};
embeddings_dref_or_nd = ft_.nd_view_func_(embeddings, embedding_shape).cast<ObjectRef>();
}
TVM_FFI_ICHECK_NE(max_num_sequence_, -1);
ObjectRef logit_pos_dref_or_nd =
ft_.CopyToWorker0(logit_pos_nd, "logit_pos", {max_num_sequence_});
Function single_batch_prefill_func = ft_.single_batch_prefill_func_;
Function prefill_func = ft_.prefill_func_;
if (ft_.single_batch_extend_func_.defined()) {
TVM_FFI_ICHECK(ft_.extend_func_.defined())
<< "`batch_extend` function is not found in the model.";
bool has_existing_sequence = false;
for (int64_t seq_id : seq_ids) {
if (prefilled_seq_ids_.count(seq_id)) {
has_existing_sequence = true;
break;
}
}
if (has_existing_sequence) {
single_batch_prefill_func = ft_.single_batch_extend_func_;
prefill_func = ft_.extend_func_;
}
for (int64_t seq_id : seq_ids) {
prefilled_seq_ids_.insert(seq_id);
}
}
// args: embeddings, logit_pos, kv_cache, [rnn_state,] params
ObjectRef ret;
if (kind == KVStateKind::kHybrid) {
// Hybrid always uses batch_prefill (single_batch prefill has tensor-based GDN args).
ret =
prefill_func(embeddings_dref_or_nd, logit_pos_dref_or_nd, kv_cache_, rnn_state_, params_)
.cast<ObjectRef>();
} else if (seq_ids.size() == 1 && !padded) {
ret = single_batch_prefill_func(embeddings_dref_or_nd, kv_cache_, params_).cast<ObjectRef>();
} else {
ret = prefill_func(embeddings_dref_or_nd, logit_pos_dref_or_nd, kv_cache_, params_)
.cast<ObjectRef>();
}
Tensor logits;
if (ft_.use_disco) {
ret = ft_.tuple_getitem_func_(ret, 0).cast<ObjectRef>();
if (num_stages_ > 1) {
// Send the result from the last worker group to worker 0.
Shape shape{1, num_sequences, vocab_size_};
DLDataType dtype = DLDataType{kDLFloat, 32, 1};
ret = ft_.last_group_send_to_worker_0_(ret, disco_logits_arr_, shape, dtype)
.cast<ObjectRef>();
}
logits = ret.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<Tensor>();
} else {
// The function returns a (logits, kv_cache) tuple; extract element 0 as a
// generic Any before casting to Tensor. Casting the whole tuple to
// Array<Tensor> now fails strict element-type checking on the kv_cache.
logits = ret.as_or_throw<Array<Any>>()[0].cast<Tensor>();
}
if (trace_enabled_) {
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
}
ft_.kv_cache_end_forward_func_(kv_cache_);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_end_forward_func_(rnn_state_);
}
// logits: (1, num_sequences, v)
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
TVM_FFI_ICHECK_EQ(logits->shape[0], 1);
TVM_FFI_ICHECK_EQ(logits->shape[1], num_sequences);
return logits;
}
ObjectRef BatchPrefillToLastHidden(const ObjectRef& embedding_or_hidden_states,
const std::vector<int64_t>& seq_ids,
const std::vector<int>& lengths) final {
NVTXScopedRange nvtx_scope("BatchPrefillToLastHidden");
TVM_FFI_ICHECK(!seq_ids.empty());
TVM_FFI_ICHECK_EQ(seq_ids.size(), lengths.size());
int num_sequences = seq_ids.size();
int total_length = 0;
for (int i = 0; i < num_sequences; ++i) {
total_length += lengths[i];
}
ObjectRef embedding_or_hidden_states_dref_or_nd{nullptr};
Shape hidden_states_shape{1, total_length, hidden_size_};
if (!ft_.use_disco) {
Tensor embedding_or_hidden_states_nd = embedding_or_hidden_states.as_or_throw<Tensor>();
embedding_or_hidden_states_dref_or_nd = embedding_or_hidden_states_nd.CreateView(
hidden_states_shape, embedding_or_hidden_states_nd->dtype);
} else {
embedding_or_hidden_states_dref_or_nd =
ft_.nd_view_func_(embedding_or_hidden_states, hidden_states_shape).cast<ObjectRef>();
}
TVM_FFI_ICHECK(ft_.prefill_to_last_hidden_func_.defined())
<< "`prefill_to_last_hidden_states` function is not found in the model.";
TVM_FFI_ICHECK(ft_.kv_cache_begin_forward_func_.defined());
TVM_FFI_ICHECK(ft_.kv_cache_end_forward_func_.defined());
TVM_FFI_ICHECK(kv_cache_.defined()) << "KV cache has not been initialized.";
// Begin forward with the sequence ids and new lengths.
Shape seq_ids_tuple(seq_ids);
Shape lengths_tuple(lengths.begin(), lengths.end());
ft_.kv_cache_begin_forward_func_(kv_cache_, seq_ids_tuple, lengths_tuple);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_begin_forward_func_(rnn_state_, seq_ids_tuple, lengths_tuple);
}
// args: embeddings, logit_pos, kv_cache, params
ObjectRef result{nullptr};
if (seq_ids.size() == 1) {
TVM_FFI_ICHECK(ft_.single_batch_prefill_to_last_hidden_func_.defined())
<< "`single_batch_prefill_to_last_hidden_states` function is not found in the model.";
if (kind == KVStateKind::kHybrid) {
result = ft_.single_batch_prefill_to_last_hidden_func_(
embedding_or_hidden_states_dref_or_nd, kv_cache_, rnn_state_, params_)
.cast<ObjectRef>();
} else {
result = ft_.single_batch_prefill_to_last_hidden_func_(
embedding_or_hidden_states_dref_or_nd, kv_cache_, params_)
.cast<ObjectRef>();
}
} else {
if (kind == KVStateKind::kHybrid) {
result = ft_.prefill_to_last_hidden_func_(embedding_or_hidden_states_dref_or_nd, kv_cache_,
rnn_state_, params_)
.cast<ObjectRef>();
} else {
result = ft_.prefill_to_last_hidden_func_(embedding_or_hidden_states_dref_or_nd, kv_cache_,
params_)
.cast<ObjectRef>();
}
}
ObjectRef hidden_states = ft_.tuple_getitem_func_(result, 0).cast<ObjectRef>();
if (trace_enabled_) {
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
}
ft_.kv_cache_end_forward_func_(kv_cache_);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_end_forward_func_(rnn_state_);
}
Shape out_shape{total_length, hidden_size_};
if (ft_.use_disco) {
return ft_.nd_view_func_(hidden_states, out_shape).cast<ObjectRef>();
} else {
Tensor hidden_states_nd = hidden_states.as_or_throw<Tensor>();
TVM_FFI_ICHECK_EQ(hidden_states_nd->ndim, 3);
TVM_FFI_ICHECK_EQ(hidden_states_nd->shape[0], 1);
TVM_FFI_ICHECK_EQ(hidden_states_nd->shape[1], total_length);
TVM_FFI_ICHECK_EQ(hidden_states_nd->shape[2], hidden_size_);
return hidden_states_nd.CreateView(out_shape, hidden_states_nd->dtype);
}
}
Tensor BatchDecode(const ObjectRef& embeddings, const std::vector<int64_t>& seq_ids) final {
NVTXScopedRange nvtx_scope("BatchDecode num_seqs=" + std::to_string(seq_ids.size()));
int num_sequence = seq_ids.size();
TVM_FFI_ICHECK(ft_.decode_func_.defined())
<< "`decode_with_embed` function is not found in the model. Please make sure the model is "
"compiled with flag `--sep-embed` and `--enable-batching`";
TVM_FFI_ICHECK(ft_.kv_cache_begin_forward_func_.defined());
TVM_FFI_ICHECK(ft_.kv_cache_end_forward_func_.defined());
TVM_FFI_ICHECK(kv_cache_.defined()) << "KV cache has not been initialized.";
// Reserve in KV cache for the lengths of the input.
// Begin forward with the sequence ids and new lengths.
Shape seq_ids_tuple(seq_ids);
Shape lengths_tuple(std::vector<int64_t>(/*n=*/seq_ids.size(), /*v=*/1));
ft_.kv_cache_begin_forward_func_(kv_cache_, seq_ids_tuple, lengths_tuple);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_begin_forward_func_(rnn_state_, seq_ids_tuple, lengths_tuple);
}
ObjectRef embeddings_dref_or_nd;
if (!embeddings->IsInstance<DRefObj>()) {
// embeddings: (1, b, h)
Tensor embeddings_nd = embeddings.as_or_throw<Tensor>();
TVM_FFI_ICHECK_NE(hidden_size_, -1);
TVM_FFI_ICHECK_EQ(embeddings_nd->ndim, 2);
TVM_FFI_ICHECK_GE(embeddings_nd->shape[0], num_sequence);
TVM_FFI_ICHECK_EQ(embeddings_nd->shape[1], hidden_size_);
TVM_FFI_ICHECK_EQ(embeddings_nd->device.device_type, device_.device_type);
TVM_FFI_ICHECK_EQ(embeddings_nd->device.device_id, device_.device_id);
embeddings_dref_or_nd =
embeddings_nd.CreateView({num_sequence, 1, hidden_size_}, embeddings_nd->dtype);
} else {
Shape embedding_shape{num_sequence, 1, hidden_size_};
embeddings_dref_or_nd = ft_.nd_view_func_(embeddings, embedding_shape).cast<ObjectRef>();
}
// args: embeddings, kv_cache, [rnn_state,] params
ObjectRef ret;
if (kind == KVStateKind::kHybrid) {
// Hybrid always uses batch_decode (single_batch decode has tensor-based GDN args).
ret =
ft_.decode_func_(embeddings_dref_or_nd, kv_cache_, rnn_state_, params_).cast<ObjectRef>();
} else if (seq_ids.size() == 1) {
ret = ft_.single_batch_decode_func_(embeddings_dref_or_nd, kv_cache_, params_)
.cast<ObjectRef>();
} else {
ret = ft_.decode_func_(embeddings_dref_or_nd, kv_cache_, params_).cast<ObjectRef>();
}
Tensor logits;
if (ft_.use_disco) {
ret = ft_.tuple_getitem_func_(ret, 0).cast<ObjectRef>();
if (num_stages_ > 1) {
// Send the result from the last worker group to worker 0.
Shape shape{num_sequence, 1, vocab_size_};
DLDataType dtype = DLDataType{kDLFloat, 32, 1};
ret = ft_.last_group_send_to_worker_0_(ret, disco_logits_arr_, shape, dtype)
.cast<ObjectRef>();
}
logits = ret.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<Tensor>();
} else {
// The function returns a (logits, kv_cache) tuple; extract element 0 as a
// generic Any before casting to Tensor. Casting the whole tuple to
// Array<Tensor> now fails strict element-type checking on the kv_cache.
logits = ret.as_or_throw<Array<Any>>()[0].cast<Tensor>();
}
if (trace_enabled_) {
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
}
ft_.kv_cache_end_forward_func_(kv_cache_);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_end_forward_func_(rnn_state_);
}
// logits: (b, 1, v)
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
TVM_FFI_ICHECK_EQ(logits->shape[0], num_sequence);
TVM_FFI_ICHECK_EQ(logits->shape[1], 1);
return logits;
}
Tensor BatchTreeDecode(const ObjectRef& embeddings, const std::vector<int64_t>& seq_ids,
const std::vector<int>& lengths,
const std::vector<int64_t>& token_tree_parent_ptr) {
// This is similar to BatchDecode, except that it takes 'length', so that each sequence can have
// multiple leaf nodes for decoding.
NVTXScopedRange nvtx_scope("BatchTreeDecode num_seqs=" + std::to_string(seq_ids.size()));
int num_sequence = seq_ids.size();
int total_length = 0;
for (int i = 0; i < num_sequence; ++i) {
total_length += lengths[i];
}
TVM_FFI_ICHECK_EQ(total_length, token_tree_parent_ptr.size());
TVM_FFI_ICHECK(ft_.decode_func_.defined())
<< "`tree_decode_with_embed` function is not found in the model. Please make sure the "
"model "
"is compiled with flag `--sep-embed` and `--enable-batching`";
TVM_FFI_ICHECK(ft_.kv_cache_begin_forward_func_.defined());
TVM_FFI_ICHECK(ft_.kv_cache_end_forward_func_.defined());
TVM_FFI_ICHECK(kv_cache_.defined()) << "KV cache has not been initialized.";
// Reserve in KV cache for the lengths of the input.
// Begin forward with the sequence ids and new lengths.
Shape seq_ids_tuple(seq_ids);
Shape lengths_tuple(lengths.begin(), lengths.end());
Shape token_tree_parent_ptr_tuple(token_tree_parent_ptr);
ft_.kv_cache_begin_forward_func_(kv_cache_, seq_ids_tuple, lengths_tuple,
token_tree_parent_ptr_tuple);
ObjectRef embeddings_dref_or_nd;
if (!embeddings->IsInstance<DRefObj>()) {
// embeddings: (1, n, h)
Tensor embeddings_nd = embeddings.as_or_throw<Tensor>();
TVM_FFI_ICHECK_NE(hidden_size_, -1);
TVM_FFI_ICHECK_EQ(embeddings_nd->ndim, 2);
TVM_FFI_ICHECK_GE(embeddings_nd->shape[0], total_length);
TVM_FFI_ICHECK_EQ(embeddings_nd->shape[1], hidden_size_);
TVM_FFI_ICHECK_EQ(embeddings_nd->device.device_type, device_.device_type);
TVM_FFI_ICHECK_EQ(embeddings_nd->device.device_id, device_.device_id);
embeddings_dref_or_nd =
embeddings_nd.CreateView({total_length, 1, hidden_size_}, embeddings_nd->dtype);
} else {
Shape embedding_shape{total_length, 1, hidden_size_};
embeddings_dref_or_nd = ft_.nd_view_func_(embeddings, embedding_shape).cast<ObjectRef>();
}
// same as BatchDecode
ObjectRef ret;
if (0 && seq_ids.size() == 1) {
ret = ft_.single_batch_decode_func_(embeddings_dref_or_nd, kv_cache_, params_)
.cast<ObjectRef>();
} else {
ret = ft_.decode_func_(embeddings_dref_or_nd, kv_cache_, params_).cast<ObjectRef>();
}
Tensor logits;
if (ft_.use_disco) {
Array<ObjectRef> result =
ret.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<Array<ObjectRef>>();
logits = result[0].as_or_throw<Tensor>();
} else {
// The function returns a (logits, kv_cache) tuple; extract element 0 as a
// generic Any before casting to Tensor. Casting the whole tuple to
// Array<Tensor> now fails strict element-type checking on the kv_cache.
logits = ret.as_or_throw<Array<Any>>()[0].cast<Tensor>();
}
if (trace_enabled_) {
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
}
ft_.kv_cache_end_forward_func_(kv_cache_);
// logits: (b, 1, v)
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
TVM_FFI_ICHECK_EQ(logits->shape[0], total_length);
TVM_FFI_ICHECK_EQ(logits->shape[1], 1);
return logits;
}
ObjectRef BatchDecodeToLastHidden(const ObjectRef& hidden_states_dref_or_nd,
const std::vector<int64_t>& seq_ids) final {
NVTXScopedRange nvtx_scope("BatchDecodeToLastHidden num_seqs=" +
std::to_string(seq_ids.size()));
int num_sequence = seq_ids.size();
TVM_FFI_ICHECK(ft_.decode_to_last_hidden_func_.defined())
<< "`batch_decode_to_last_hidden_states` function is not found in the model.";
TVM_FFI_ICHECK(ft_.kv_cache_begin_forward_func_.defined());
TVM_FFI_ICHECK(ft_.kv_cache_end_forward_func_.defined());
TVM_FFI_ICHECK(kv_cache_.defined()) << "KV cache has not been initialized.";
// Reserve in KV cache for the lengths of the input.
// Begin forward with the sequence ids and new lengths.
Shape seq_ids_tuple(seq_ids);
Shape lengths_tuple(std::vector<int64_t>(/*n=*/seq_ids.size(), /*v=*/1));
ft_.kv_cache_begin_forward_func_(kv_cache_, seq_ids_tuple, lengths_tuple);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_begin_forward_func_(rnn_state_, seq_ids_tuple, lengths_tuple);
}
// args: embeddings, kv_cache, params
ObjectRef result{nullptr};
if (seq_ids.size() == 1) {
TVM_FFI_ICHECK(ft_.single_batch_decode_to_last_hidden_func_.defined())
<< "`decode_to_last_hidden_states` function is not found in the model.";
if (kind == KVStateKind::kHybrid) {
result = ft_.single_batch_decode_to_last_hidden_func_(hidden_states_dref_or_nd, kv_cache_,
rnn_state_, params_)
.cast<ObjectRef>();
} else {
result = ft_.single_batch_decode_to_last_hidden_func_(hidden_states_dref_or_nd, kv_cache_,
params_)
.cast<ObjectRef>();
}
} else {
if (kind == KVStateKind::kHybrid) {
result = ft_.decode_to_last_hidden_func_(hidden_states_dref_or_nd, kv_cache_, rnn_state_,
params_)
.cast<ObjectRef>();
} else {
result = ft_.decode_to_last_hidden_func_(hidden_states_dref_or_nd, kv_cache_, params_)
.cast<ObjectRef>();
}
}
ft_.kv_cache_end_forward_func_(kv_cache_);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_end_forward_func_(rnn_state_);
}
ObjectRef hidden_states = ft_.tuple_getitem_func_(result, 0).cast<ObjectRef>();
if (trace_enabled_) {
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
}
// hidden_states: (b, 1, v) to (b, v)
Shape out_shape{num_sequence, hidden_size_};
if (ft_.use_disco) {
return ft_.nd_view_func_(hidden_states, out_shape).cast<ObjectRef>();
} else {
Tensor hidden_states_nd = hidden_states.as_or_throw<Tensor>();
TVM_FFI_ICHECK_EQ(hidden_states_nd->ndim, 3);
TVM_FFI_ICHECK_EQ(hidden_states_nd->shape[0], num_sequence);
TVM_FFI_ICHECK_EQ(hidden_states_nd->shape[1], 1);
TVM_FFI_ICHECK_EQ(hidden_states_nd->shape[2], hidden_size_);
return hidden_states_nd.CreateView(out_shape, hidden_states_nd->dtype);
}
}
Tensor BatchVerify(const ObjectRef& embeddings, const std::vector<int64_t>& seq_ids,
const std::vector<int>& lengths,
const std::vector<int64_t>& token_tree_parent_ptr) final {
TVM_FFI_ICHECK(!seq_ids.empty());
TVM_FFI_ICHECK_EQ(seq_ids.size(), lengths.size());
int num_sequences = seq_ids.size();
int total_length = 0;
for (int i = 0; i < num_sequences; ++i) {
total_length += lengths[i];
}
TVM_FFI_ICHECK_EQ(total_length, token_tree_parent_ptr.size());
NVTXScopedRange nvtx_scope("BatchVerify num_tokens=" + std::to_string(total_length));
TVM_FFI_ICHECK(ft_.verify_func_.defined())
<< "`verify_with_embed` function is not found in the model. Please make sure the model is "
"compiled with flag `--sep-embed` and `--enable-batching`";
TVM_FFI_ICHECK(ft_.kv_cache_begin_forward_func_.defined());
TVM_FFI_ICHECK(ft_.kv_cache_end_forward_func_.defined());
TVM_FFI_ICHECK(kv_cache_.defined()) << "KV cache has not been initialized.";
// Begin forward with the sequence ids and new lengths.
Shape seq_ids_tuple(seq_ids);
Shape lengths_tuple(lengths.begin(), lengths.end());
Shape token_tree_parent_ptr_tuple(token_tree_parent_ptr);
ft_.kv_cache_begin_forward_func_(kv_cache_, seq_ids_tuple, lengths_tuple,
token_tree_parent_ptr_tuple);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_begin_forward_func_(rnn_state_, seq_ids_tuple, lengths_tuple,
token_tree_parent_ptr_tuple);
}
ObjectRef embeddings_dref_or_nd;
if (!embeddings->IsInstance<DRefObj>()) {
// embeddings: (1, n, h)
Tensor embeddings_nd = embeddings.as_or_throw<Tensor>();
TVM_FFI_ICHECK_NE(hidden_size_, -1);
TVM_FFI_ICHECK_EQ(embeddings_nd->ndim, 2);
TVM_FFI_ICHECK_GE(embeddings_nd->shape[0], total_length);
TVM_FFI_ICHECK_EQ(embeddings_nd->shape[1], hidden_size_);
TVM_FFI_ICHECK_EQ(embeddings_nd->device.device_type, device_.device_type);
TVM_FFI_ICHECK_EQ(embeddings_nd->device.device_id, device_.device_id);
embeddings_dref_or_nd =
embeddings_nd.CreateView({1, total_length, hidden_size_}, embeddings_nd->dtype);
} else {
Shape embedding_shape{1, total_length, hidden_size_};
embeddings_dref_or_nd = ft_.nd_view_func_(embeddings, embedding_shape).cast<ObjectRef>();
}
// args: embeddings, kv_cache, [rnn_state,] params
ObjectRef ret;
if (kind == KVStateKind::kHybrid) {
ret =
ft_.verify_func_(embeddings_dref_or_nd, kv_cache_, rnn_state_, params_).cast<ObjectRef>();
} else {
ret = ft_.verify_func_(embeddings_dref_or_nd, kv_cache_, params_).cast<ObjectRef>();
}
Tensor logits;
if (ft_.use_disco) {
ret = ft_.tuple_getitem_func_(ret, 0).cast<ObjectRef>();
if (num_stages_ > 1) {
// Send the result from the last worker group to worker 0.
Shape shape{1, total_length, vocab_size_};
DLDataType dtype = DLDataType{kDLFloat, 32, 1};
ret = ft_.last_group_send_to_worker_0_(ret, disco_logits_arr_, shape, dtype)
.cast<ObjectRef>();
}
logits = ret.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<Tensor>();
} else {
// The function returns a (logits, kv_cache) tuple; extract element 0 as a
// generic Any before casting to Tensor. Casting the whole tuple to
// Array<Tensor> now fails strict element-type checking on the kv_cache.
logits = ret.as_or_throw<Array<Any>>()[0].cast<Tensor>();
}
if (trace_enabled_) {
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
}
ft_.kv_cache_end_forward_func_(kv_cache_);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_end_forward_func_(rnn_state_);
}
// logits: (1, total_length, v)
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
TVM_FFI_ICHECK_EQ(logits->shape[0], 1);
TVM_FFI_ICHECK_EQ(logits->shape[1], total_length);
return logits;
}
ObjectRef BatchVerifyToLastHidden(const ObjectRef& embeddings,
const std::vector<int64_t>& seq_ids,
const std::vector<int>& lengths,
const std::vector<int64_t>& token_tree_parent_ptr) final {
TVM_FFI_ICHECK(!seq_ids.empty());
TVM_FFI_ICHECK_EQ(seq_ids.size(), lengths.size());
int num_sequences = seq_ids.size();
int total_length = 0;
for (int i = 0; i < num_sequences; ++i) {
total_length += lengths[i];
}
TVM_FFI_ICHECK_EQ(total_length, token_tree_parent_ptr.size());
NVTXScopedRange nvtx_scope("BatchVerifyToLastHidden num_tokens=" +
std::to_string(total_length));
TVM_FFI_ICHECK(ft_.verify_to_last_hidden_func_.defined())
<< "`batch_verify_to_last_hidden_states` function is not found in the model.";
TVM_FFI_ICHECK(ft_.kv_cache_begin_forward_func_.defined());
TVM_FFI_ICHECK(ft_.kv_cache_end_forward_func_.defined());
TVM_FFI_ICHECK(kv_cache_.defined()) << "KV cache has not been initialized.";
ObjectRef embeddings_dref_or_nd;
if (!embeddings->IsInstance<DRefObj>()) {
// embeddings: (1, n, h)
Tensor embeddings_nd = embeddings.as_or_throw<Tensor>();
TVM_FFI_ICHECK_NE(hidden_size_, -1);
TVM_FFI_ICHECK_EQ(embeddings_nd->ndim, 2);
TVM_FFI_ICHECK_GE(embeddings_nd->shape[0], total_length);
TVM_FFI_ICHECK_EQ(embeddings_nd->shape[1], hidden_size_);
TVM_FFI_ICHECK_EQ(embeddings_nd->device.device_type, device_.device_type);
TVM_FFI_ICHECK_EQ(embeddings_nd->device.device_id, device_.device_id);
embeddings_dref_or_nd =
embeddings_nd.CreateView({1, total_length, hidden_size_}, embeddings_nd->dtype);
} else {
Shape embedding_shape{1, total_length, hidden_size_};
embeddings_dref_or_nd = ft_.nd_view_func_(embeddings, embedding_shape).cast<ObjectRef>();
}
// Begin forward with the sequence ids and new lengths.
Shape seq_ids_tuple(seq_ids);
Shape lengths_tuple(lengths.begin(), lengths.end());
Shape token_tree_parent_ptr_tuple(token_tree_parent_ptr);
ft_.kv_cache_begin_forward_func_(kv_cache_, seq_ids_tuple, lengths_tuple,
token_tree_parent_ptr_tuple);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_begin_forward_func_(rnn_state_, seq_ids_tuple, lengths_tuple,
token_tree_parent_ptr_tuple);
}
// args: embeddings, logit_pos, kv_cache, params
ObjectRef result;
if (kind == KVStateKind::kHybrid) {
result =
ft_.verify_to_last_hidden_func_(embeddings_dref_or_nd, kv_cache_, rnn_state_, params_)
.cast<ObjectRef>();
} else {
result = ft_.verify_to_last_hidden_func_(embeddings_dref_or_nd, kv_cache_, params_)
.cast<ObjectRef>();
}
ft_.kv_cache_end_forward_func_(kv_cache_);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_end_forward_func_(rnn_state_);
}
ObjectRef hidden_states = ft_.tuple_getitem_func_(result, 0).cast<ObjectRef>();
if (trace_enabled_) {
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
}
Shape out_shape{total_length, hidden_size_};
if (!ft_.use_disco) {
Tensor hidden_states_nd = hidden_states.as_or_throw<Tensor>();
TVM_FFI_ICHECK_EQ(hidden_states_nd->ndim, 3);
TVM_FFI_ICHECK_EQ(hidden_states_nd->shape[0], 1);
TVM_FFI_ICHECK_EQ(hidden_states_nd->shape[1], total_length);
TVM_FFI_ICHECK_EQ(hidden_states_nd->shape[2], hidden_size_);
return hidden_states_nd.CreateView(out_shape, hidden_states_nd->dtype);
} else {
return ft_.nd_view_func_(hidden_states, out_shape).cast<ObjectRef>();
}
}
/*********************** KV Cache Management ***********************/
void CreateKVCache(int page_size, int max_num_sequence, int64_t max_total_sequence_length,
int64_t prefill_chunk_size, int max_history_size,
int prefix_cache_max_num_recycling_seqs = 0) final {
KVStateKind kv_state_kind = GetMetadata().kv_state_kind;
if (kv_state_kind == KVStateKind::kKVCache) {
Shape max_num_sequence_tuple{max_num_sequence};
Shape max_total_sequence_length_tuple{max_total_sequence_length};
Shape prefill_chunk_size_tuple{prefill_chunk_size};
Shape page_size_tuple{page_size};
Shape support_sliding_window{sliding_window_size_ != -1};
kv_cache_ = ft_.create_kv_cache_func_(max_num_sequence_tuple, max_total_sequence_length_tuple,
prefill_chunk_size_tuple, page_size_tuple,
support_sliding_window)
.cast<ObjectRef>();
local_kv_cache_ = ft_.use_disco
? kv_cache_.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<ObjectRef>()
: kv_cache_;
} else if (kv_state_kind == KVStateKind::kRNNState) {
// RNN state needs extra slots for prefix cache recycling sequences.
Shape max_num_sequence_tuple{max_num_sequence + prefix_cache_max_num_recycling_seqs};
Shape max_history_size_tuple = {std::max(max_history_size, 1)};
kv_cache_ = ft_.create_kv_cache_func_(max_num_sequence_tuple, max_history_size_tuple)
.cast<ObjectRef>();
local_kv_cache_ = ft_.use_disco
? kv_cache_.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<ObjectRef>()
: kv_cache_;
} else if (kv_state_kind == KVStateKind::kHybrid) {
// Hybrid: create both PagedKVCache (for attention layers) and RNNState (for GDN layers).
Shape max_num_sequence_tuple{max_num_sequence};
Shape max_total_sequence_length_tuple{max_total_sequence_length};
Shape prefill_chunk_size_tuple{prefill_chunk_size};
Shape page_size_tuple{page_size};
Shape support_sliding_window{sliding_window_size_ != -1};
kv_cache_ = ft_.create_kv_cache_func_(max_num_sequence_tuple, max_total_sequence_length_tuple,
prefill_chunk_size_tuple, page_size_tuple,
support_sliding_window)
.cast<ObjectRef>();
local_kv_cache_ = ft_.use_disco
? kv_cache_.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<ObjectRef>()
: kv_cache_;
// Create RNN state for recurrent layers.
// RNN state needs extra slots for prefix cache recycling sequences that
// coexist with active sequences (e.g., during ForkSequence).
Shape rnn_max_batch{max_num_sequence + prefix_cache_max_num_recycling_seqs};
Shape rnn_max_history{std::max(max_history_size, 1)};
rnn_state_ = ft_.create_rnn_state_func_(rnn_max_batch, rnn_max_history).cast<ObjectRef>();
local_rnn_state_ =
ft_.use_disco ? rnn_state_.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<ObjectRef>()
: rnn_state_;
} else if (kv_state_kind == KVStateKind::kNone) {
// Do nothing
} else {
LOG(FATAL) << "Unknown kv_state_kind: " << static_cast<int>(kv_state_kind);
}
}
void AddNewSequence(int64_t seq_id) final {
if (this->kind == KVStateKind::kNone) {
return;
}
ft_.kv_cache_add_sequence_func_(kv_cache_, seq_id);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_add_sequence_func_(rnn_state_, seq_id);
}
}
void ForkSequence(int64_t parent_seq_id, int64_t child_seq_id, int64_t fork_pos) final {
if (this->kind == KVStateKind::kNone) {
return;
}
ft_.kv_cache_fork_sequence_func_(kv_cache_, parent_seq_id, child_seq_id, fork_pos);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_fork_sequence_func_(rnn_state_, parent_seq_id, child_seq_id, fork_pos);
}
prefilled_seq_ids_.insert(child_seq_id);
}
void RemoveSequence(int64_t seq_id) final {
if (this->kind == KVStateKind::kNone) {
return;
}
prefilled_seq_ids_.erase(seq_id);
ft_.kv_cache_remove_sequence_func_(kv_cache_, seq_id);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_remove_sequence_func_(rnn_state_, seq_id);
}
}
void PopNFromKVCache(int64_t seq_id, int num_tokens) final {
if (this->kind == KVStateKind::kNone) {
return;
}
ft_.kv_cache_popn_func_(kv_cache_, seq_id, num_tokens);
if (kind == KVStateKind::kHybrid) {
ft_.kv_cache_popn_func_(rnn_state_, seq_id, num_tokens);
}
}
void CommitAcceptedTokenTreeNodesToKVCache(
const std::vector<int64_t>& seq_ids,
const std::vector<int64_t>& accepted_leaf_indices) final {
Shape seq_ids_tuple(seq_ids);
Shape accepted_leaf_indices_tuple(accepted_leaf_indices);
ft_.kv_cache_commit_accepted_token_tree_nodes_func_(kv_cache_, seq_ids_tuple,
accepted_leaf_indices_tuple);
}
void EnableSlidingWindowForSeq(int64_t seq_id) final {
if (this->kind == KVStateKind::kNone) {
return;
}
if (sliding_window_size_ != -1) {
ft_.kv_cache_enable_sliding_window_for_seq_(kv_cache_, seq_id, sliding_window_size_,
attention_sink_size_);
}
}
Shape DisaggPrepareKVRecv(int64_t seq_id, int length) final {
NVTXScopedRange nvtx_scope("DisaggPrepareKVRecv length=" + std::to_string(length));
TVM_FFI_ICHECK(ft_.kv_cache_disagg_prepare_recv_func_.defined());
TVM_FFI_ICHECK(kv_cache_.defined()) << "KV cache has not been initialized.";
// Run KV receive preparation.
ObjectRef ret;
ret = ft_.kv_cache_disagg_prepare_recv_func_(kv_cache_, seq_id, length).cast<ObjectRef>();
Shape compressed_kv_append_metadata;
if (ft_.use_disco) {
compressed_kv_append_metadata = ret.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<Shape>();
} else {
compressed_kv_append_metadata = ret.as_or_throw<Shape>();
}
return compressed_kv_append_metadata;
}
void DisaggMarkKVSend(int64_t seq_id, int begin_pos, Shape compressed_kv_append_metadata,
int dst_group_offset) final {
NVTXScopedRange nvtx_scope("DisaggMarkKVSend seq_id=" + std::to_string(seq_id) +
" begin_pos=" + std::to_string(begin_pos));
TVM_FFI_ICHECK(ft_.kv_cache_disagg_mark_send_func_.defined());
TVM_FFI_ICHECK(kv_cache_.defined()) << "KV cache has not been initialized.";
// Run KV send preparation.
ft_.kv_cache_disagg_mark_send_func_(kv_cache_, seq_id, begin_pos, compressed_kv_append_metadata,
dst_group_offset);
}
/************** Raw Info Query **************/
ModelMetadata GetMetadata() const final { return ft_.model_metadata_; }
int GetNumAvailablePages() const final {
if (this->kind == KVStateKind::kRNNState || this->kind == KVStateKind::kNone) {
// RNNState does not introduce new page at runtime
return std::numeric_limits<int>::max();
} else {
// kKVCache and kHybrid both use PagedKVCache for capacity.
return ft_.kv_cache_get_num_available_pages_func_(local_kv_cache_).cast<int>();
}
}
int GetCurrentTotalSequenceLength() const final {
if (this->kind == KVStateKind::kRNNState || this->kind == KVStateKind::kNone) {
// RNNState does not have a total sequence length limit
return 0;
} else {
return ft_.kv_cache_get_total_sequence_length_func_(local_kv_cache_).cast<int>();
}
}
/*********************** Utilities ***********************/
void LoadParams() final { this->params_ = ft_.LoadParams(model_, device_); }
void SetMaxNumSequence(int max_num_sequence) final {
this->max_num_sequence_ = max_num_sequence;
this->logit_pos_arr_ = Tensor::Empty({max_num_sequence}, DLDataType{kDLInt, 32, 1},
Device{DLDeviceType::kDLCPU, 0});
}
void SetPrefillChunkSize(int prefill_chunk_size) final {
this->prefill_chunk_size_ = prefill_chunk_size;
Device preferred_host_device = GetPreferredHostDevice(device_);
memory::Allocator* allocator = memory::MemoryManager::GetOrCreateAllocator(
preferred_host_device, memory::AllocatorType::kNaive);
TVM_FFI_ICHECK_NOTNULL(allocator);
token_ids_storage_ = memory::Storage(
allocator->Alloc(preferred_host_device, {prefill_chunk_size_}, DLDataType{kDLInt, 32, 1}),
allocator);
if (this->num_stages_ > 1) {
// Create a remote Tensor for logits when pipeline parallelism is enabled.
disco_logits_arr_ =
ft_.Empty({prefill_chunk_size_, vocab_size_}, DLDataType{kDLFloat, 32, 1}, device_,
/*worker0_only=*/true);
}
}
LogitProcessor CreateLogitProcessor(int max_num_token,
Optional<EventTraceRecorder> trace_recorder) final {
return LogitProcessor(max_num_token, vocab_size_, &this->ft_, device_,
std::move(trace_recorder));
}
Sampler CreateSampler(int max_num_sample, int num_models,
Optional<EventTraceRecorder> trace_recorder) final {
if (Sampler::SupportGPUSampler(device_)) {
return Sampler::CreateGPUSampler(max_num_sample, vocab_size_, &this->ft_, device_,
std::move(trace_recorder));
} else {
return Sampler::CreateCPUSampler(std::move(trace_recorder));
}
}
int EstimateHostCPURequirement() const final {
TVM_FFI_ICHECK_NE(num_shards_, -1) << "The model has not been initialized";
return num_shards_ > 1 ? num_shards_ : 0;
}
int GetSlidingWindowSize() const final { return sliding_window_size_; }
int GetAttentionSinkSize() const final { return attention_sink_size_; }
ObjectRef AllocEmbeddingTensor() final {
if (!ft_.alloc_embedding_tensor_func_.defined()) {
return ObjectRef{nullptr};
}
// Allocate the embedding tensor.
ObjectRef embedding = ft_.alloc_embedding_tensor_func_().cast<ObjectRef>();
// Get the shape of the embedding tensor for hidden size.
Shape embedding_shape;
if (ft_.use_disco) {
TVM_FFI_ICHECK(embedding->IsInstance<DRefObj>());
ObjectRef shape_ref = ft_.nd_get_shape_func_(embedding).cast<ObjectRef>();
embedding_shape = shape_ref.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<Shape>();
} else {
Tensor embedding_nd = embedding.as_or_throw<Tensor>();
embedding_shape = embedding_nd.Shape();
}
TVM_FFI_ICHECK_NE(prefill_chunk_size_, -1);
TVM_FFI_ICHECK_EQ(embedding_shape.size(), 2);
TVM_FFI_ICHECK_GE(embedding_shape[0], prefill_chunk_size_);
this->hidden_size_ = embedding_shape[1];
return embedding;
}
ObjectRef AllocHiddenStatesTensor() final {
if (!ft_.alloc_embedding_tensor_func_.defined()) {
return ObjectRef{nullptr};
}
// Allocate the hidden_states tensor.
// Use the same function as embeddings.
ObjectRef hidden_states = ft_.alloc_embedding_tensor_func_().cast<ObjectRef>();
Tensor hidden_states_nd{nullptr};
// Get the shape of the hidden_states tensor for hidden size.
if (ft_.use_disco) {
TVM_FFI_ICHECK(hidden_states->IsInstance<DRefObj>());
hidden_states_nd = hidden_states.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<Tensor>();
} else {
hidden_states_nd = hidden_states.as_or_throw<Tensor>();
}
Shape hidden_states_shape = hidden_states_nd.Shape();
TVM_FFI_ICHECK_NE(prefill_chunk_size_, -1);
TVM_FFI_ICHECK_EQ(hidden_states_shape.size(), 2);
TVM_FFI_ICHECK_GE(hidden_states_shape[0], prefill_chunk_size_);
this->hidden_size_ = hidden_states_shape[1];
this->hidden_states_dtype_ = hidden_states_nd->dtype;
return hidden_states;
}
void Reset() final {
// Reset the KV cache.
if (kv_cache_.defined()) {
ft_.reset_kv_cache_func_(kv_cache_);
}
if (rnn_state_.defined()) {
ft_.reset_kv_cache_func_(rnn_state_);
}
}
/********************** Utilities for speculative decoding **********************/
DraftTokenWorkspaceManager CreateDraftTokenWorkspaceManager(int max_num_tokens) {
return DraftTokenWorkspaceManager(max_num_tokens, vocab_size_, hidden_size_,
hidden_states_dtype_, device_, ft_);
}
ObjectRef GatherHiddenStates(const ObjectRef& input, const std::vector<int>& indices,
ObjectRef* dst) final {
ObjectRef dst_view{nullptr};
Shape out_shape{static_cast<int64_t>(indices.size()), hidden_size_};
if ((*dst)->IsInstance<DRefObj>()) {
dst_view = ft_.nd_view_func_(*dst, out_shape).cast<ObjectRef>();
} else {
Tensor dst_nd = (*dst).as_or_throw<Tensor>();
dst_view = dst_nd.CreateView(out_shape, hidden_states_dtype_);
}
Tensor indices_nd = logit_pos_arr_.CreateView({static_cast<int64_t>(indices.size())},
DLDataType{kDLInt, 32, 1});
indices_nd.CopyFromBytes(indices.data(), indices.size() * sizeof(int));
TVM_FFI_ICHECK_NE(max_num_sequence_, -1);
ObjectRef indices_device = ft_.CopyToWorker0(indices_nd, "logit_pos", {max_num_sequence_});
ft_.gather_hidden_states_func_(input, indices_device, dst_view);
return dst_view;
}
void ScatterHiddenStates(const ObjectRef& input, const std::vector<int>& indices,
ObjectRef* dst) final {
Tensor indices_nd = logit_pos_arr_.CreateView({static_cast<int64_t>(indices.size())},
DLDataType{kDLInt, 32, 1});
indices_nd.CopyFromBytes(indices.data(), indices.size() * sizeof(int));
TVM_FFI_ICHECK_NE(max_num_sequence_, -1);
ObjectRef indices_device = ft_.CopyToWorker0(indices_nd, "logit_pos", {max_num_sequence_});
ft_.scatter_hidden_states_func_(input, indices_device, *dst);
}
Tensor GatherDraftProbs(const Tensor& input, const std::vector<int>& indices, Tensor* dst) final {
Tensor dst_view = dst->CreateView({static_cast<int64_t>(indices.size()), vocab_size_},
DLDataType{kDLFloat, 32, 1});
Tensor indices_nd = logit_pos_arr_.CreateView({static_cast<int64_t>(indices.size())},
DLDataType{kDLInt, 32, 1});
indices_nd.CopyFromBytes(indices.data(), indices.size() * sizeof(int));
TVM_FFI_ICHECK_NE(max_num_sequence_, -1);
ObjectRef indices_device =
ft_.CopyToWorker0(indices_nd, "logit_pos_local", {max_num_sequence_}, /*local_only=*/true);
ft_.gather_probs_func_(input, indices_device, dst_view);
return dst_view;
}
void ScatterDraftProbs(const Tensor& input, const std::vector<int>& indices, Tensor* dst) final {
Tensor indices_nd = logit_pos_arr_.CreateView({static_cast<int64_t>(indices.size())},
DLDataType{kDLInt, 32, 1});
indices_nd.CopyFromBytes(indices.data(), indices.size() * sizeof(int));
TVM_FFI_ICHECK_NE(max_num_sequence_, -1);
ObjectRef indices_device =
ft_.CopyToWorker0(indices_nd, "logit_pos_local", {max_num_sequence_}, /*local_only=*/true);
ft_.scatter_probs_func_(input, indices_device, *dst);
}
Array<Tensor> GetMedusaLogits(const ObjectRef& hidden_states) {
ObjectRef result = ft_.get_logits_func_(hidden_states).cast<ObjectRef>();
Array<Tensor> logits{nullptr};
if (ft_.use_disco) {
logits = result.as_or_throw<DRef>()->DebugGetFromRemote(0).cast<Array<Tensor>>();
} else {
logits = result.as_or_throw<Array<Tensor>>();
}
return logits;
}
/************** Debug/Profile **************/
void DebugCallFuncOnAllAllWorker(const String& func_name, Optional<String> func_args) final {
ft_.DebugCallFuncOnAllAllWorker(func_name, func_args);
}
private:
/*! \brief Load model configuration from JSON. */
void LoadModelConfigJSON(const tvm::ffi::json::Object& config) {
this->sliding_window_size_ =
json::LookupOrDefault<int64_t>(config, "sliding_window_size", this->sliding_window_size_);
TVM_FFI_ICHECK(sliding_window_size_ == -1 || sliding_window_size_ > 0)
<< "Sliding window should be either -1 (which means disabled) of positive";
this->attention_sink_size_ =
json::LookupOrDefault<int64_t>(config, "attention_sink_size", this->attention_sink_size_);
this->attention_sink_size_ = std::max(this->attention_sink_size_, 0);
this->vocab_size_ = json::Lookup<int64_t>(config, "vocab_size");
this->model_type_ = json::Lookup<std::string>(config, "model_type");
}
//----------------------------
// Model configurations
//----------------------------
std::string model_;
int sliding_window_size_ = -1;
int attention_sink_size_ = 0;
int num_shards_ = -1;
int num_stages_ = -1;
int max_num_sequence_ = -1;
int prefill_chunk_size_ = -1;
int hidden_size_ = -1;
DLDataType hidden_states_dtype_;
int vocab_size_ = -1;
int image_embed_size_ = -1;
int seqlen_padding_factor_ = 1;
std::string model_type_;
//----------------------------
// TVM related states
//----------------------------
// Packed function table
FunctionTable ft_;
// Paged KV cache.
// - We use `kv_cache_` for general KV cache operations.
// When tensor parallelism is enabled, `kv_cache_` is a DRef object.
// - For efficient KV cache raw info query, we use `local_kv_cache`
// as a local **reference** of `kv_cache_`. It is a pure mirror of `kv_cache_`
// except that it is always a local object.
ObjectRef kv_cache_{nullptr};
ObjectRef local_kv_cache_{nullptr};
// RNN state for hybrid models (GatedDeltaNet recurrent layers).
ObjectRef rnn_state_{nullptr};
ObjectRef local_rnn_state_{nullptr};
// Runtime device
Device device_;
// Model parameters
ObjectRef params_;
// Shared Tensor
memory::Storage token_ids_storage_{nullptr};
Tensor logit_pos_arr_{nullptr};
ObjectRef disco_logits_arr_{nullptr};
// A boolean indicating if tracing is enabled.
bool trace_enabled_;
// An enum indicating whether it's RNN-based.
KVStateKind kind;
// A set of sequence IDs that have been prefilled.
std::unordered_set<int64_t> prefilled_seq_ids_;
};
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
"mlc.copy_embedding_to_offset", [](Tensor embedding, Tensor dst, int offset) {
// embedding: (m, hidden_size)
// dst: (prefill_chunk_size, hidden_size)
TVM_FFI_ICHECK_EQ(embedding->ndim, 2);
TVM_FFI_ICHECK_EQ(dst->ndim, 2);
TVM_FFI_ICHECK_LE(embedding->shape[0] + offset, dst->shape[0]);
TVM_FFI_ICHECK_EQ(embedding->shape[1], dst->shape[1]);
const DLTensor& copy_src = *(embedding.operator->());
const DLTensor* p_copy_dst = dst.operator->();
DLTensor copy_dst = *p_copy_dst;
copy_dst.shape = embedding->shape;
copy_dst.byte_offset = offset * embedding->shape[1] *
((embedding->dtype.bits * embedding->dtype.lanes + 7) / 8);
Tensor::CopyFromTo(&copy_src, &copy_dst);
});
}
} // namespace serve
} // namespace llm
} // namespace mlc
Reference in New Issue View Git Blame Copy Permalink