chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
+16
@@ -0,0 +1,16 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file base.h
|
||||
*/
|
||||
|
||||
#ifndef MLC_LLM_DLL
|
||||
#ifdef _WIN32
|
||||
#ifdef MLC_LLM_EXPORTS
|
||||
#define MLC_LLM_DLL __declspec(dllexport)
|
||||
#else
|
||||
#define MLC_LLM_DLL __declspec(dllimport)
|
||||
#endif
|
||||
#else
|
||||
#define MLC_LLM_DLL __attribute__((visibility("default")))
|
||||
#endif
|
||||
#endif
|
||||
@@ -0,0 +1,618 @@
|
||||
#include "conv_template.h"
|
||||
|
||||
#include <tvm/ffi/function.h>
|
||||
|
||||
#include "../support/json_parser.h"
|
||||
#include "image_utils.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace json_ffi {
|
||||
|
||||
using namespace mlc::llm;
|
||||
|
||||
/****************** Model vision config ******************/
|
||||
|
||||
ModelVisionConfig ModelVisionConfig::FromJSON(const tvm::ffi::json::Object& json_obj) {
|
||||
ModelVisionConfig config;
|
||||
|
||||
Result<int64_t> hidden_size_res = json::LookupWithResultReturn<int64_t>(json_obj, "hidden_size");
|
||||
if (hidden_size_res.IsOk()) {
|
||||
config.hidden_size = static_cast<int>(hidden_size_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> image_size_res = json::LookupWithResultReturn<int64_t>(json_obj, "image_size");
|
||||
if (image_size_res.IsOk()) {
|
||||
config.image_size = static_cast<int>(image_size_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> intermediate_size_res =
|
||||
json::LookupWithResultReturn<int64_t>(json_obj, "intermediate_size");
|
||||
if (intermediate_size_res.IsOk()) {
|
||||
config.intermediate_size = static_cast<int>(intermediate_size_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> num_attention_heads_res =
|
||||
json::LookupWithResultReturn<int64_t>(json_obj, "num_attention_heads");
|
||||
if (num_attention_heads_res.IsOk()) {
|
||||
config.num_attention_heads = static_cast<int>(num_attention_heads_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> num_hidden_layers_res =
|
||||
json::LookupWithResultReturn<int64_t>(json_obj, "num_hidden_layers");
|
||||
if (num_hidden_layers_res.IsOk()) {
|
||||
config.num_hidden_layers = static_cast<int>(num_hidden_layers_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> patch_size_res = json::LookupWithResultReturn<int64_t>(json_obj, "patch_size");
|
||||
if (patch_size_res.IsOk()) {
|
||||
config.patch_size = static_cast<int>(patch_size_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> projection_dim_res =
|
||||
json::LookupWithResultReturn<int64_t>(json_obj, "projection_dim");
|
||||
if (projection_dim_res.IsOk()) {
|
||||
config.projection_dim = static_cast<int>(projection_dim_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> vocab_size_res = json::LookupWithResultReturn<int64_t>(json_obj, "vocab_size");
|
||||
if (vocab_size_res.IsOk()) {
|
||||
config.vocab_size = static_cast<int>(vocab_size_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<std::string> dtype_res = json::LookupWithResultReturn<std::string>(json_obj, "dtype");
|
||||
if (dtype_res.IsOk()) {
|
||||
config.dtype = dtype_res.Unwrap();
|
||||
}
|
||||
|
||||
Result<int64_t> num_channels_res =
|
||||
json::LookupWithResultReturn<int64_t>(json_obj, "num_channels");
|
||||
if (num_channels_res.IsOk()) {
|
||||
config.num_channels = static_cast<int>(num_channels_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<double> layer_norm_eps_res =
|
||||
json::LookupWithResultReturn<double>(json_obj, "layer_norm_eps");
|
||||
if (layer_norm_eps_res.IsOk()) {
|
||||
config.layer_norm_eps = layer_norm_eps_res.Unwrap();
|
||||
}
|
||||
|
||||
return config;
|
||||
}
|
||||
|
||||
/****************** Model config ******************/
|
||||
|
||||
ModelConfig ModelConfig::FromJSON(const tvm::ffi::json::Object& json_obj) {
|
||||
ModelConfig config;
|
||||
|
||||
Result<int64_t> vocab_size_res = json::LookupWithResultReturn<int64_t>(json_obj, "vocab_size");
|
||||
if (vocab_size_res.IsOk()) {
|
||||
config.vocab_size = static_cast<int>(vocab_size_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> context_window_size_res =
|
||||
json::LookupWithResultReturn<int64_t>(json_obj, "context_window_size");
|
||||
if (context_window_size_res.IsOk()) {
|
||||
config.context_window_size = static_cast<int>(context_window_size_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> sliding_window_size_res =
|
||||
json::LookupWithResultReturn<int64_t>(json_obj, "sliding_window_size");
|
||||
if (sliding_window_size_res.IsOk()) {
|
||||
config.sliding_window_size = static_cast<int>(sliding_window_size_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> prefill_chunk_size_res =
|
||||
json::LookupWithResultReturn<int64_t>(json_obj, "prefill_chunk_size");
|
||||
if (prefill_chunk_size_res.IsOk()) {
|
||||
config.prefill_chunk_size = static_cast<int>(prefill_chunk_size_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> tensor_parallel_shards_res =
|
||||
json::LookupWithResultReturn<int64_t>(json_obj, "tensor_parallel_shards");
|
||||
if (tensor_parallel_shards_res.IsOk()) {
|
||||
config.tensor_parallel_shards = static_cast<int>(tensor_parallel_shards_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> pipeline_parallel_stages_res =
|
||||
json::LookupWithResultReturn<int64_t>(json_obj, "pipeline_parallel_stages");
|
||||
if (pipeline_parallel_stages_res.IsOk()) {
|
||||
config.pipeline_parallel_stages = static_cast<int>(pipeline_parallel_stages_res.Unwrap());
|
||||
}
|
||||
|
||||
Result<int64_t> max_batch_size_res =
|
||||
json::LookupWithResultReturn<int64_t>(json_obj, "max_batch_size");
|
||||
if (max_batch_size_res.IsOk()) {
|
||||
config.max_batch_size = static_cast<int>(max_batch_size_res.Unwrap());
|
||||
}
|
||||
|
||||
if (json_obj.count("vision_config")) {
|
||||
const tvm::ffi::json::Object& vision_config_obj =
|
||||
json_obj.at("vision_config").cast<tvm::ffi::json::Object>();
|
||||
config.vision_config = ModelVisionConfig::FromJSON(vision_config_obj);
|
||||
}
|
||||
|
||||
return config;
|
||||
}
|
||||
|
||||
/****************** Conversation template ******************/
|
||||
|
||||
std::unordered_map<MessagePlaceholders, std::string> PLACEHOLDERS = {
|
||||
{MessagePlaceholders::SYSTEM, "{system_message}"},
|
||||
{MessagePlaceholders::USER, "{user_message}"},
|
||||
{MessagePlaceholders::ASSISTANT, "{assistant_message}"},
|
||||
{MessagePlaceholders::TOOL, "{tool_message}"},
|
||||
{MessagePlaceholders::FUNCTION, "{function_string}"}};
|
||||
|
||||
MessagePlaceholders MessagePlaceholderFromString(const std::string& role) {
|
||||
static const std::unordered_map<std::string, MessagePlaceholders> enum_map = {
|
||||
{"system", MessagePlaceholders::SYSTEM}, {"user", MessagePlaceholders::USER},
|
||||
{"assistant", MessagePlaceholders::ASSISTANT}, {"tool", MessagePlaceholders::TOOL},
|
||||
{"function", MessagePlaceholders::FUNCTION},
|
||||
};
|
||||
|
||||
return enum_map.at(role);
|
||||
}
|
||||
|
||||
Conversation::Conversation()
|
||||
: role_templates({{"user", PLACEHOLDERS[MessagePlaceholders::USER]},
|
||||
{"assistant", PLACEHOLDERS[MessagePlaceholders::ASSISTANT]},
|
||||
{"tool", PLACEHOLDERS[MessagePlaceholders::TOOL]}}) {}
|
||||
|
||||
std::string Conversation::GetSystemText(const std::string& system_msg) const {
|
||||
std::string system_text = this->system_template;
|
||||
static std::string system_placeholder = PLACEHOLDERS[MessagePlaceholders::SYSTEM];
|
||||
size_t pos = system_text.find(system_placeholder);
|
||||
if (pos != std::string::npos) {
|
||||
system_text.replace(pos, system_placeholder.length(), system_msg);
|
||||
}
|
||||
return system_text;
|
||||
}
|
||||
|
||||
std::string Conversation::GetRoleText(const std::string& role, const std::string& content,
|
||||
const std::optional<std::string>& fn_call_string) const {
|
||||
std::string role_text = this->role_templates.at(role);
|
||||
std::string placeholder = PLACEHOLDERS[MessagePlaceholderFromString(role)];
|
||||
size_t pos = role_text.find(placeholder);
|
||||
if (pos != std::string::npos) {
|
||||
role_text.replace(pos, placeholder.length(), content);
|
||||
}
|
||||
if (fn_call_string) {
|
||||
// replace placeholder[FUNCTION] with function_string
|
||||
// this assumes function calling is used for a single request scenario only
|
||||
pos = role_text.find(PLACEHOLDERS[MessagePlaceholders::FUNCTION]);
|
||||
if (pos != std::string::npos) {
|
||||
role_text.replace(pos, PLACEHOLDERS[MessagePlaceholders::FUNCTION].length(),
|
||||
fn_call_string.value());
|
||||
}
|
||||
}
|
||||
return role_text;
|
||||
}
|
||||
|
||||
/// Try to detect if function calling is needed, if so, return the function calling string
|
||||
Result<std::optional<std::string>> TryGetFunctionCallingString(
|
||||
const Conversation& conv, const ChatCompletionRequest& request) {
|
||||
using TResult = Result<std::optional<std::string>>;
|
||||
if (!request.tools.has_value() ||
|
||||
(request.tool_choice.has_value() && request.tool_choice.value() == "none")) {
|
||||
return TResult::Ok(std::nullopt);
|
||||
}
|
||||
std::vector<ChatTool> tools_ = request.tools.value();
|
||||
std::string tool_choice_ = request.tool_choice.value();
|
||||
|
||||
// TODO: support with tool choice as dict
|
||||
for (const auto& tool : tools_) {
|
||||
if (tool.function.name == tool_choice_) {
|
||||
tvm::ffi::json::Value function_str(tool.function.AsJSON());
|
||||
return TResult::Ok(tvm::ffi::json::Stringify(function_str));
|
||||
}
|
||||
}
|
||||
|
||||
if (tool_choice_ != "auto") {
|
||||
return TResult::Error("Invalid tool_choice value in the request: " + tool_choice_);
|
||||
}
|
||||
|
||||
tvm::ffi::json::Array function_list;
|
||||
for (const auto& tool : tools_) {
|
||||
function_list.push_back(tool.function.AsJSON());
|
||||
}
|
||||
|
||||
tvm::ffi::json::Value function_list_json(function_list);
|
||||
return TResult::Ok(tvm::ffi::json::Stringify(function_list_json));
|
||||
};
|
||||
|
||||
Result<std::vector<Data>> CreatePrompt(const Conversation& conv,
|
||||
const ChatCompletionRequest& request,
|
||||
const ModelConfig& config, DLDevice device) {
|
||||
using TResult = Result<std::vector<Data>>;
|
||||
|
||||
Result<std::optional<std::string>> fn_call_str_tmp = TryGetFunctionCallingString(conv, request);
|
||||
if (fn_call_str_tmp.IsErr()) {
|
||||
return TResult::Error(fn_call_str_tmp.UnwrapErr());
|
||||
}
|
||||
std::optional<std::string> fn_call_string = fn_call_str_tmp.Unwrap();
|
||||
|
||||
// Handle system message
|
||||
// concz
|
||||
bool has_custom_system = false;
|
||||
std::string custom_system_inputs;
|
||||
|
||||
auto f_populate_system_message = [&](const std::vector<ChatCompletionMessage>& msg_vec) {
|
||||
for (ChatCompletionMessage msg : msg_vec) {
|
||||
if (msg.role == "system") {
|
||||
TVM_FFI_ICHECK(msg.content.IsText()) << "System message must be text";
|
||||
custom_system_inputs += msg.content.Text();
|
||||
has_custom_system = true;
|
||||
}
|
||||
}
|
||||
};
|
||||
// go through messages in template and passed in.
|
||||
f_populate_system_message(conv.messages);
|
||||
f_populate_system_message(request.messages);
|
||||
|
||||
// pending text records the text to be put into data
|
||||
// we lazily accumulate the pending text
|
||||
// to reduce amount of segments in the Data vector
|
||||
std::string pending_text =
|
||||
conv.GetSystemText(has_custom_system ? custom_system_inputs : conv.system_message);
|
||||
|
||||
// Get the message strings
|
||||
std::vector<Data> message_list;
|
||||
size_t non_system_msg_count = 0;
|
||||
|
||||
// returns error if error happens
|
||||
auto f_process_messages =
|
||||
[&](const std::vector<ChatCompletionMessage>& msg_vec) -> std::optional<TResult> {
|
||||
for (size_t i = 0; i < msg_vec.size(); ++i) {
|
||||
const ChatCompletionMessage& msg = msg_vec[i];
|
||||
// skip system message as it is already processed
|
||||
if (msg.role == "system") continue;
|
||||
|
||||
auto role_it = conv.roles.find(msg.role);
|
||||
if (role_it == conv.roles.end()) {
|
||||
return TResult::Error("Role \"" + msg.role + "\" is not supported");
|
||||
}
|
||||
const std::string& role_name = role_it->second;
|
||||
// skip when content is empty
|
||||
if (msg.content.IsNull()) {
|
||||
pending_text += role_name + conv.role_empty_sep;
|
||||
continue;
|
||||
}
|
||||
++non_system_msg_count;
|
||||
// assistant uses conv.seps[1] if there are two seps
|
||||
int sep_offset = msg.role == "assistant" ? 1 : 0;
|
||||
const std::string& seperator = conv.seps[sep_offset % conv.seps.size()];
|
||||
// setup role prefix
|
||||
std::string role_prefix = "";
|
||||
// Do not append role prefix if this is the first message and there is already a system
|
||||
// message
|
||||
if (conv.add_role_after_system_message || pending_text.empty() || non_system_msg_count != 1) {
|
||||
role_prefix = role_name + conv.role_content_sep;
|
||||
}
|
||||
pending_text += role_prefix;
|
||||
|
||||
if (msg.content.IsParts()) {
|
||||
for (const auto& item : msg.content.Parts()) {
|
||||
auto it_type = item.find("type");
|
||||
if (it_type == item.end()) {
|
||||
return TResult::Error("The content of a message does not have \"type\" field");
|
||||
}
|
||||
if (it_type->second == "text") {
|
||||
auto it_text = item.find("text");
|
||||
if (it_text == item.end()) {
|
||||
return TResult::Error(
|
||||
"The text type content of a message does not have \"text\" field");
|
||||
}
|
||||
// replace placeholder[ROLE] with input message from role
|
||||
pending_text += conv.GetRoleText(msg.role, it_text->second, fn_call_string);
|
||||
} else if (it_type->second == "image_url") {
|
||||
if (item.find("image_url") == item.end()) {
|
||||
return TResult::Error("Content should have an image_url field");
|
||||
}
|
||||
std::string image_url =
|
||||
item.at("image_url"); // TODO(mlc-team): According to OpenAI API reference this
|
||||
// should be a map, with a "url" key containing the URL, but
|
||||
// we are just assuming this as the URL for now
|
||||
std::string base64_image = image_url.substr(image_url.find(",") + 1);
|
||||
Result<Tensor> image_data_res = LoadImageFromBase64(base64_image);
|
||||
if (image_data_res.IsErr()) {
|
||||
return TResult::Error(image_data_res.UnwrapErr());
|
||||
}
|
||||
if (!config.vision_config.has_value()) {
|
||||
return TResult::Error("Vision config is required for image input");
|
||||
}
|
||||
int image_size = config.vision_config.value().image_size;
|
||||
int patch_size = config.vision_config.value().patch_size;
|
||||
|
||||
int embed_size = (image_size * image_size) / (patch_size * patch_size);
|
||||
|
||||
Tensor image_data = image_data_res.Unwrap();
|
||||
std::vector<int64_t> new_shape = {1, image_size, image_size, 3};
|
||||
Tensor image_tensor = image_data.CreateView(new_shape, image_data.DataType());
|
||||
// TODO: Not sure if commenting will affect other functions. But
|
||||
// python part will do clip preprocessing. auto image_tensor =
|
||||
// ClipPreprocessor(image_data_res.Unwrap(), image_size, device);
|
||||
// lazily commit text data
|
||||
if (pending_text.length() != 0) {
|
||||
message_list.push_back(TextData(pending_text));
|
||||
pending_text = "";
|
||||
}
|
||||
message_list.push_back(ImageData(image_tensor, embed_size));
|
||||
} else {
|
||||
return TResult::Error("Unsupported content type: " + it_type->second);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
TVM_FFI_ICHECK(msg.content.IsText());
|
||||
pending_text += conv.GetRoleText(msg.role, msg.content.Text(), fn_call_string);
|
||||
}
|
||||
pending_text += seperator;
|
||||
}
|
||||
return std::nullopt;
|
||||
};
|
||||
|
||||
// Optionally strip `<think>...</think>` blocks from historical assistant
|
||||
// messages (those before the last user message), mirroring Qwen3's HF chat
|
||||
// template. See mlc-ai/mlc-llm#3482.
|
||||
const std::vector<ChatCompletionMessage>* conv_messages_ptr = &conv.messages;
|
||||
const std::vector<ChatCompletionMessage>* request_messages_ptr = &request.messages;
|
||||
std::vector<ChatCompletionMessage> stripped_conv_messages;
|
||||
std::vector<ChatCompletionMessage> stripped_request_messages;
|
||||
if (conv.strip_reasoning_in_history) {
|
||||
const size_t conv_size = conv.messages.size();
|
||||
int64_t last_user_idx = -1;
|
||||
for (size_t i = 0; i < conv_size; ++i) {
|
||||
if (conv.messages[i].role == "user") last_user_idx = static_cast<int64_t>(i);
|
||||
}
|
||||
for (size_t i = 0; i < request.messages.size(); ++i) {
|
||||
if (request.messages[i].role == "user") last_user_idx = static_cast<int64_t>(conv_size + i);
|
||||
}
|
||||
const std::string kCloseTag = "</think>";
|
||||
auto strip_range = [&](const std::vector<ChatCompletionMessage>& in, size_t offset,
|
||||
std::vector<ChatCompletionMessage>& out) {
|
||||
out.reserve(in.size());
|
||||
for (size_t i = 0; i < in.size(); ++i) {
|
||||
ChatCompletionMessage msg = in[i];
|
||||
const int64_t global_idx = static_cast<int64_t>(offset + i);
|
||||
if (msg.role == "assistant" && global_idx < last_user_idx && msg.content.IsText()) {
|
||||
const std::string& text = msg.content.Text();
|
||||
const size_t close_pos = text.rfind(kCloseTag);
|
||||
if (close_pos != std::string::npos) {
|
||||
size_t start = close_pos + kCloseTag.size();
|
||||
while (start < text.size() && text[start] == '\n') ++start;
|
||||
msg.content = ChatCompletionMessageContent(text.substr(start));
|
||||
}
|
||||
}
|
||||
out.push_back(std::move(msg));
|
||||
}
|
||||
};
|
||||
strip_range(conv.messages, 0, stripped_conv_messages);
|
||||
strip_range(request.messages, conv_size, stripped_request_messages);
|
||||
conv_messages_ptr = &stripped_conv_messages;
|
||||
request_messages_ptr = &stripped_request_messages;
|
||||
}
|
||||
|
||||
if (auto err = f_process_messages(*conv_messages_ptr)) {
|
||||
return err.value();
|
||||
}
|
||||
if (auto err = f_process_messages(*request_messages_ptr)) {
|
||||
return err.value();
|
||||
}
|
||||
// append last assistant begin message
|
||||
ChatCompletionMessage last_assistant_begin;
|
||||
last_assistant_begin.role = "assistant";
|
||||
last_assistant_begin.content = std::nullopt;
|
||||
if (auto err = f_process_messages({last_assistant_begin})) {
|
||||
return err.value();
|
||||
}
|
||||
if (pending_text.length() != 0) {
|
||||
message_list.push_back(TextData(pending_text));
|
||||
}
|
||||
// Handle system_prefix_token_ids
|
||||
if (conv.system_prefix_token_ids.has_value()) {
|
||||
message_list.insert(message_list.begin(), TokenData(conv.system_prefix_token_ids.value()));
|
||||
}
|
||||
return TResult::Ok(message_list);
|
||||
}
|
||||
|
||||
Result<Conversation> Conversation::FromJSON(const tvm::ffi::json::Object& json_obj) {
|
||||
using TResult = Result<Conversation>;
|
||||
Conversation conv;
|
||||
|
||||
Result<std::optional<std::string>> name_res =
|
||||
json::LookupOptionalWithResultReturn<std::string>(json_obj, "name");
|
||||
if (name_res.IsErr()) {
|
||||
return TResult::Error(name_res.UnwrapErr());
|
||||
}
|
||||
conv.name = name_res.Unwrap();
|
||||
|
||||
Result<std::string> system_template_res =
|
||||
json::LookupWithResultReturn<std::string>(json_obj, "system_template");
|
||||
if (system_template_res.IsErr()) {
|
||||
return TResult::Error(system_template_res.UnwrapErr());
|
||||
}
|
||||
conv.system_template = system_template_res.Unwrap();
|
||||
|
||||
Result<std::string> system_message_res =
|
||||
json::LookupWithResultReturn<std::string>(json_obj, "system_message");
|
||||
if (system_message_res.IsErr()) {
|
||||
return TResult::Error(system_message_res.UnwrapErr());
|
||||
}
|
||||
conv.system_message = system_message_res.Unwrap();
|
||||
|
||||
Result<std::optional<tvm::ffi::json::Array>> system_prefix_token_ids_arr_res =
|
||||
json::LookupOptionalWithResultReturn<tvm::ffi::json::Array>(json_obj,
|
||||
"system_prefix_token_ids");
|
||||
if (system_prefix_token_ids_arr_res.IsErr()) {
|
||||
return TResult::Error(system_prefix_token_ids_arr_res.UnwrapErr());
|
||||
}
|
||||
std::optional<tvm::ffi::json::Array> system_prefix_token_ids_arr =
|
||||
system_prefix_token_ids_arr_res.Unwrap();
|
||||
if (system_prefix_token_ids_arr.has_value()) {
|
||||
std::vector<int> system_prefix_token_ids;
|
||||
system_prefix_token_ids.reserve(system_prefix_token_ids_arr.value().size());
|
||||
for (const auto& token_id : system_prefix_token_ids_arr.value()) {
|
||||
if (!token_id.try_cast<int64_t>().has_value()) {
|
||||
return TResult::Error("A system prefix token id is not integer.");
|
||||
}
|
||||
system_prefix_token_ids.push_back(static_cast<int>(token_id.cast<int64_t>()));
|
||||
}
|
||||
conv.system_prefix_token_ids = std::move(system_prefix_token_ids);
|
||||
}
|
||||
|
||||
Result<bool> add_role_after_system_message_res =
|
||||
json::LookupWithResultReturn<bool>(json_obj, "add_role_after_system_message");
|
||||
if (add_role_after_system_message_res.IsErr()) {
|
||||
return TResult::Error(add_role_after_system_message_res.UnwrapErr());
|
||||
}
|
||||
conv.add_role_after_system_message = add_role_after_system_message_res.Unwrap();
|
||||
|
||||
Result<tvm::ffi::json::Object> roles_object_res =
|
||||
json::LookupWithResultReturn<tvm::ffi::json::Object>(json_obj, "roles");
|
||||
if (roles_object_res.IsErr()) {
|
||||
return TResult::Error(roles_object_res.UnwrapErr());
|
||||
}
|
||||
for (const auto& role : roles_object_res.Unwrap()) {
|
||||
if (!role.second.try_cast<std::string>().has_value()) {
|
||||
return TResult::Error("A role value in the conversation template is not a string.");
|
||||
}
|
||||
conv.roles[role.first.cast<tvm::ffi::String>()] = role.second.cast<std::string>();
|
||||
}
|
||||
|
||||
Result<std::optional<tvm::ffi::json::Object>> role_templates_object_res =
|
||||
json::LookupOptionalWithResultReturn<tvm::ffi::json::Object>(json_obj, "role_templates");
|
||||
if (role_templates_object_res.IsErr()) {
|
||||
return TResult::Error(role_templates_object_res.UnwrapErr());
|
||||
}
|
||||
std::optional<tvm::ffi::json::Object> role_templates_object = role_templates_object_res.Unwrap();
|
||||
if (role_templates_object.has_value()) {
|
||||
for (const auto& [role, msg] : role_templates_object.value()) {
|
||||
if (!msg.try_cast<std::string>().has_value()) {
|
||||
return TResult::Error("A value in \"role_templates\" is not a string.");
|
||||
}
|
||||
conv.role_templates[role.cast<tvm::ffi::String>()] = msg.cast<std::string>();
|
||||
}
|
||||
}
|
||||
|
||||
Result<tvm::ffi::json::Array> messages_arr_res =
|
||||
json::LookupWithResultReturn<tvm::ffi::json::Array>(json_obj, "messages");
|
||||
if (messages_arr_res.IsErr()) {
|
||||
return TResult::Error(messages_arr_res.UnwrapErr());
|
||||
}
|
||||
for (const auto& message : messages_arr_res.Unwrap()) {
|
||||
if (!message.try_cast<tvm::ffi::json::Array>().has_value() ||
|
||||
message.cast<tvm::ffi::json::Array>().size() != 2) {
|
||||
return TResult::Error(
|
||||
"A message in the conversation template is not an array of [role, content].");
|
||||
}
|
||||
tvm::ffi::json::Array message_arr = message.cast<tvm::ffi::json::Array>();
|
||||
if (!message_arr[0].try_cast<std::string>().has_value()) {
|
||||
return TResult::Error("The role of a message in the conversation template is not a string.");
|
||||
}
|
||||
std::string role = message_arr[0].cast<std::string>();
|
||||
// content can be a string or an array of objects
|
||||
if (message_arr[1].try_cast<std::string>().has_value()) {
|
||||
ChatCompletionMessage msg;
|
||||
msg.role = role;
|
||||
msg.content = message_arr[1].cast<std::string>();
|
||||
conv.messages.push_back(msg);
|
||||
continue;
|
||||
} else if (message_arr[1].try_cast<tvm::ffi::json::Array>().has_value()) {
|
||||
tvm::ffi::json::Array content_arr = message_arr[1].cast<tvm::ffi::json::Array>();
|
||||
std::vector<std::unordered_map<std::string, std::string>> content;
|
||||
content.reserve(content_arr.size());
|
||||
for (const auto& item : content_arr) {
|
||||
if (!item.try_cast<tvm::ffi::json::Object>().has_value()) {
|
||||
return TResult::Error("The content of conversation template message is not an object");
|
||||
}
|
||||
std::unordered_map<std::string, std::string> item_map;
|
||||
for (const auto& [key, value] : item.cast<tvm::ffi::json::Object>()) {
|
||||
item_map[key.cast<tvm::ffi::String>()] = tvm::ffi::json::Stringify(value);
|
||||
}
|
||||
content.push_back(std::move(item_map));
|
||||
}
|
||||
ChatCompletionMessage msg;
|
||||
msg.role = role;
|
||||
msg.content = content;
|
||||
conv.messages.push_back(msg);
|
||||
continue;
|
||||
} else {
|
||||
return TResult::Error(
|
||||
"The content of a message in the conversation template is not a string or an array.");
|
||||
}
|
||||
}
|
||||
|
||||
Result<tvm::ffi::json::Array> seps_arr_res =
|
||||
json::LookupWithResultReturn<tvm::ffi::json::Array>(json_obj, "seps");
|
||||
if (seps_arr_res.IsErr()) {
|
||||
return TResult::Error(seps_arr_res.UnwrapErr());
|
||||
}
|
||||
std::vector<std::string> seps;
|
||||
for (const auto& sep : seps_arr_res.Unwrap()) {
|
||||
if (!sep.try_cast<std::string>().has_value()) {
|
||||
return TResult::Error("A separator (\"seps\") of the conversation template is not a string");
|
||||
}
|
||||
conv.seps.push_back(sep.cast<std::string>());
|
||||
}
|
||||
|
||||
Result<std::string> role_content_sep_res =
|
||||
json::LookupWithResultReturn<std::string>(json_obj, "role_content_sep");
|
||||
if (role_content_sep_res.IsErr()) {
|
||||
return TResult::Error(role_content_sep_res.UnwrapErr());
|
||||
}
|
||||
conv.role_content_sep = role_content_sep_res.Unwrap();
|
||||
|
||||
Result<std::string> role_empty_sep_res =
|
||||
json::LookupWithResultReturn<std::string>(json_obj, "role_empty_sep");
|
||||
if (role_empty_sep_res.IsErr()) {
|
||||
return TResult::Error(role_empty_sep_res.UnwrapErr());
|
||||
}
|
||||
conv.role_empty_sep = role_empty_sep_res.Unwrap();
|
||||
|
||||
Result<tvm::ffi::json::Array> stop_str_arr_res =
|
||||
json::LookupWithResultReturn<tvm::ffi::json::Array>(json_obj, "stop_str");
|
||||
if (stop_str_arr_res.IsErr()) {
|
||||
return TResult::Error(stop_str_arr_res.UnwrapErr());
|
||||
}
|
||||
for (const auto& stop : stop_str_arr_res.Unwrap()) {
|
||||
if (!stop.try_cast<std::string>().has_value()) {
|
||||
return TResult::Error(
|
||||
"A stop string (\"stop_str\") of the conversation template is not a string.");
|
||||
}
|
||||
conv.stop_str.push_back(stop.cast<std::string>());
|
||||
}
|
||||
|
||||
Result<tvm::ffi::json::Array> stop_token_ids_arr_res =
|
||||
json::LookupWithResultReturn<tvm::ffi::json::Array>(json_obj, "stop_token_ids");
|
||||
if (stop_token_ids_arr_res.IsErr()) {
|
||||
return TResult::Error(stop_token_ids_arr_res.UnwrapErr());
|
||||
}
|
||||
for (const auto& stop : stop_token_ids_arr_res.Unwrap()) {
|
||||
if (!stop.try_cast<int64_t>().has_value()) {
|
||||
return TResult::Error(
|
||||
"A stop token id (\"stop_token_ids\") of the conversation template is not an integer.");
|
||||
}
|
||||
conv.stop_token_ids.push_back(static_cast<int>(stop.cast<int64_t>()));
|
||||
}
|
||||
|
||||
Result<std::optional<bool>> strip_reasoning_res =
|
||||
json::LookupOptionalWithResultReturn<bool>(json_obj, "strip_reasoning_in_history");
|
||||
if (strip_reasoning_res.IsErr()) {
|
||||
return TResult::Error(strip_reasoning_res.UnwrapErr());
|
||||
}
|
||||
conv.strip_reasoning_in_history = strip_reasoning_res.Unwrap().value_or(false);
|
||||
|
||||
return TResult::Ok(conv);
|
||||
}
|
||||
|
||||
Result<Conversation> Conversation::FromJSON(const std::string& json_str) {
|
||||
Result<tvm::ffi::json::Object> json_obj = json::ParseToJSONObjectWithResultReturn(json_str);
|
||||
if (json_obj.IsErr()) {
|
||||
return Result<Conversation>::Error(json_obj.UnwrapErr());
|
||||
}
|
||||
return Conversation::FromJSON(json_obj.Unwrap());
|
||||
}
|
||||
|
||||
} // namespace json_ffi
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,161 @@
|
||||
#ifndef MLC_LLM_JSON_FFI_CONV_TEMPLATE_H
|
||||
#define MLC_LLM_JSON_FFI_CONV_TEMPLATE_H
|
||||
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
|
||||
#include <iostream>
|
||||
#include <map>
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <typeinfo>
|
||||
#include <variant>
|
||||
#include <vector>
|
||||
|
||||
#include "../serve/data.h"
|
||||
#include "../support/result.h"
|
||||
#include "openai_api_protocol.h"
|
||||
|
||||
using namespace mlc::llm::serve;
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace json_ffi {
|
||||
|
||||
/****************** Model vision config ******************/
|
||||
|
||||
/*! \brief Defines the Vision config of the model (if present) */
|
||||
class ModelVisionConfig {
|
||||
public:
|
||||
int hidden_size;
|
||||
int image_size;
|
||||
int intermediate_size;
|
||||
int num_attention_heads;
|
||||
int num_hidden_layers;
|
||||
int patch_size;
|
||||
int projection_dim;
|
||||
int vocab_size;
|
||||
std::string dtype;
|
||||
int num_channels;
|
||||
double layer_norm_eps;
|
||||
|
||||
static ModelVisionConfig FromJSON(const tvm::ffi::json::Object& json_obj);
|
||||
};
|
||||
|
||||
/****************** Model config ******************/
|
||||
|
||||
/*! \brief Defines the config of the model.
|
||||
Populated from "model_config" field in mlc-chat-config.json */
|
||||
class ModelConfig {
|
||||
public:
|
||||
int vocab_size;
|
||||
int context_window_size;
|
||||
int sliding_window_size;
|
||||
int prefill_chunk_size;
|
||||
int tensor_parallel_shards;
|
||||
int pipeline_parallel_stages;
|
||||
int max_batch_size;
|
||||
std::optional<ModelVisionConfig> vision_config = std::nullopt;
|
||||
|
||||
static ModelConfig FromJSON(const tvm::ffi::json::Object& json_obj);
|
||||
};
|
||||
|
||||
/****************** Conversation template ******************/
|
||||
|
||||
enum class MessagePlaceholders { SYSTEM, USER, ASSISTANT, TOOL, FUNCTION };
|
||||
|
||||
MessagePlaceholders MessagePlaceholderFromString(const std::string& role);
|
||||
|
||||
/**
|
||||
* @brief A struct that specifies the convention template of conversation
|
||||
* and contains the conversation history.
|
||||
*/
|
||||
struct Conversation {
|
||||
// Optional name of the template.
|
||||
std::optional<std::string> name = std::nullopt;
|
||||
|
||||
// The system prompt template, it optionally contains the system
|
||||
// message placeholder, and the placeholder will be replaced with
|
||||
// the system message below.
|
||||
std::string system_template;
|
||||
|
||||
// The content of the system prompt (without the template format).
|
||||
std::string system_message;
|
||||
|
||||
// The system token ids to be prepended at the beginning of tokenized
|
||||
// generated prompt.
|
||||
std::optional<std::vector<int>> system_prefix_token_ids = std::nullopt;
|
||||
|
||||
// Whether or not to append user role and separator after the system message.
|
||||
// This is mainly for [INST] [/INST] style prompt format
|
||||
bool add_role_after_system_message = true;
|
||||
|
||||
// The conversation roles
|
||||
std::unordered_map<std::string, std::string> roles;
|
||||
|
||||
// The roles prompt template, it optionally contains the defaults
|
||||
// message placeholders and will be replaced by actual content
|
||||
std::unordered_map<std::string, std::string> role_templates;
|
||||
|
||||
// The conversation history messages.
|
||||
// Each message is a pair of strings, denoting "(role, content)".
|
||||
// The content can be None.
|
||||
std::vector<ChatCompletionMessage> messages;
|
||||
|
||||
// The separators between messages when concatenating into a single prompt.
|
||||
// List size should be either 1 or 2.
|
||||
// - When size is 1, the separator will be used between adjacent messages.
|
||||
// - When size is 2, seps[0] is used after user message, and
|
||||
// seps[1] is used after assistant message.
|
||||
std::vector<std::string> seps;
|
||||
|
||||
// The separator between the role and the content in a message.
|
||||
std::string role_content_sep;
|
||||
|
||||
// The separator between the role and empty contents.
|
||||
std::string role_empty_sep;
|
||||
|
||||
// The stop criteria
|
||||
std::vector<std::string> stop_str;
|
||||
std::vector<int> stop_token_ids;
|
||||
|
||||
// When true, strip `<think>...</think>` blocks (and any trailing whitespace)
|
||||
// from historical assistant messages before rendering the prompt, mirroring
|
||||
// Qwen3's official HF chat template. Only assistant messages that appear
|
||||
// before the last user message are affected.
|
||||
bool strip_reasoning_in_history = false;
|
||||
|
||||
Conversation();
|
||||
|
||||
/*!
|
||||
* \brief Get the system text(with the prompt template) given the system prompt message
|
||||
* \param system_msg The system prompt message.
|
||||
* \return The created system text.
|
||||
*/
|
||||
std::string GetSystemText(const std::string& system_msg) const;
|
||||
|
||||
/*!
|
||||
* \brief replace the content from role by the correct role text in template
|
||||
* \param role The input role
|
||||
* \param content The input content from the role
|
||||
* \param fn_call_str The function calling string if any.
|
||||
* \return The created text.
|
||||
*/
|
||||
std::string GetRoleText(const std::string& role, const std::string& content,
|
||||
const std::optional<std::string>& fn_call_str) const;
|
||||
|
||||
/*! \brief Create a Conversation instance from the given JSON object. */
|
||||
static Result<Conversation> FromJSON(const tvm::ffi::json::Object& json);
|
||||
/*! \brief Parse and create a Conversation instance from the given JSON string. */
|
||||
static Result<Conversation> FromJSON(const std::string& json_str);
|
||||
};
|
||||
|
||||
/*! \brief Create the list of prompts from the messages based on the conversation template. */
|
||||
Result<std::vector<Data>> CreatePrompt(const Conversation& conv,
|
||||
const ChatCompletionRequest& request,
|
||||
const ModelConfig& config, DLDevice device);
|
||||
|
||||
} // namespace json_ffi
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_JSON_FFI_CONV_TEMPLATE_H
|
||||
@@ -0,0 +1,160 @@
|
||||
#include "image_utils.h"
|
||||
|
||||
#include <tvm/support/io.h>
|
||||
|
||||
#include "../../3rdparty/tvm/src/support/base64.h"
|
||||
#define STB_IMAGE_IMPLEMENTATION
|
||||
#include "stb_image.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace json_ffi {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
|
||||
class MemoryBufferStream : public tvm::support::Stream {
|
||||
public:
|
||||
using Stream::Read;
|
||||
using Stream::Write;
|
||||
|
||||
MemoryBufferStream(const char* data, size_t size) : data_(data), size_(size), pos_(0) {}
|
||||
|
||||
size_t Read(void* ptr, size_t size) override {
|
||||
size_t remaining = size_ - pos_;
|
||||
if (size > remaining) {
|
||||
size = remaining;
|
||||
}
|
||||
if (size == 0) {
|
||||
return 0;
|
||||
}
|
||||
std::memcpy(ptr, data_ + pos_, size);
|
||||
pos_ += size;
|
||||
return size;
|
||||
}
|
||||
|
||||
size_t Write(const void* ptr, size_t size) override {
|
||||
TVM_FFI_THROW(InternalError) << "MemoryBufferStream does not support write";
|
||||
return 0;
|
||||
}
|
||||
|
||||
private:
|
||||
const char* data_;
|
||||
size_t size_;
|
||||
size_t pos_;
|
||||
};
|
||||
|
||||
size_t Base64DecodedSize(const std::string& base64_str) {
|
||||
size_t len = base64_str.size();
|
||||
size_t padding = 0;
|
||||
if (base64_str[len - 1] == '=') {
|
||||
padding++;
|
||||
}
|
||||
if (base64_str[len - 2] == '=') {
|
||||
padding++;
|
||||
}
|
||||
return 3 * len / 4 - padding;
|
||||
}
|
||||
|
||||
Result<Tensor> LoadImageFromBase64(const std::string& base64_str) {
|
||||
using TResult = Result<Tensor>;
|
||||
MemoryBufferStream stream(base64_str.c_str(), base64_str.size());
|
||||
tvm::support::Base64InStream base64_stream(&stream);
|
||||
size_t decoded_size = Base64DecodedSize(base64_str);
|
||||
std::vector<unsigned char> decoded(decoded_size);
|
||||
base64_stream.InitPosition();
|
||||
base64_stream.Read((void*)decoded.data(), decoded_size);
|
||||
int width, height, num_channels;
|
||||
unsigned char* image_data =
|
||||
stbi_load_from_memory(decoded.data(), decoded_size, &width, &height, &num_channels, 3);
|
||||
if (!image_data) {
|
||||
return TResult::Error(stbi_failure_reason());
|
||||
}
|
||||
auto image_tensor = Tensor::Empty({height, width, 3}, {kDLUInt, 8, 1}, {kDLCPU, 0});
|
||||
image_tensor.CopyFromBytes((void*)image_data, width * height * 3);
|
||||
stbi_image_free(image_data);
|
||||
return TResult::Ok(image_tensor);
|
||||
}
|
||||
|
||||
Tensor ClipPreprocessor(Tensor image_data, int target_size, DLDevice device) {
|
||||
int height = image_data->shape[0];
|
||||
int width = image_data->shape[1];
|
||||
// Resize
|
||||
const int short_side = width < height ? width : height;
|
||||
const int long_side = width > height ? width : height;
|
||||
const int new_short_side = target_size;
|
||||
const int new_long_side = (int)(new_short_side * (long_side / (float)short_side));
|
||||
const int new_width = width < height ? new_short_side : new_long_side;
|
||||
const int new_height = width > height ? new_short_side : new_long_side;
|
||||
|
||||
std::vector<float> processed_image_data(new_width * new_height * 3);
|
||||
|
||||
// Bilinear Interpolation
|
||||
for (int y = 0; y < new_height; y++) {
|
||||
for (int x = 0; x < new_width; x++) {
|
||||
const float x_ratio = float(width - 1) / new_width;
|
||||
const float y_ratio = float(height - 1) / new_height;
|
||||
const int x1 = int(x_ratio * x);
|
||||
const int y1 = int(y_ratio * y);
|
||||
const int x2 = x1 + 1;
|
||||
const int y2 = y1 + 1;
|
||||
const float x_diff = x_ratio * x - x1;
|
||||
const float y_diff = y_ratio * y - y1;
|
||||
for (int c = 0; c < 3; c++) {
|
||||
const uint8_t top_left = ((uint8_t*)image_data->data)[(y1 * width + x1) * 3 + c];
|
||||
const uint8_t top_right = ((uint8_t*)image_data->data)[(y1 * width + x2) * 3 + c];
|
||||
const uint8_t bottom_left = ((uint8_t*)image_data->data)[(y2 * width + x1) * 3 + c];
|
||||
const uint8_t bottom_right = ((uint8_t*)image_data->data)[(y2 * width + x2) * 3 + c];
|
||||
processed_image_data[(y * new_width + x) * 3 + c] =
|
||||
(float)(int(top_left * (1 - x_diff) * (1 - y_diff) + top_right * x_diff * (1 - y_diff) +
|
||||
bottom_left * y_diff * (1 - x_diff) + bottom_right * x_diff * y_diff));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Center crop
|
||||
const int crop_x = (new_width - target_size) / 2;
|
||||
const int crop_y = (new_height - target_size) / 2;
|
||||
std::vector<float> cropped_image_data(target_size * target_size * 3);
|
||||
for (int y = 0; y < target_size; y++) {
|
||||
for (int x = 0; x < target_size; x++) {
|
||||
for (int c = 0; c < 3; c++) {
|
||||
cropped_image_data[(y * target_size + x) * 3 + c] =
|
||||
processed_image_data[((y + crop_y) * new_width + x + crop_x) * 3 + c];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rescale
|
||||
for (int i = 0; i < target_size * target_size * 3; i++) {
|
||||
cropped_image_data[i] = cropped_image_data[i] / 255.0f;
|
||||
}
|
||||
|
||||
// Normalize
|
||||
const float IMAGE_MEAN[] = {0.48145466f, 0.4578275f, 0.40821073f};
|
||||
const float IMAGE_STD[] = {0.26862954f, 0.26130258f, 0.27577711f};
|
||||
for (int i = 0; i < target_size * target_size * 3; i++) {
|
||||
const int c = i % 3;
|
||||
cropped_image_data[i] = (cropped_image_data[i] - IMAGE_MEAN[c]) / IMAGE_STD[c];
|
||||
}
|
||||
|
||||
std::vector<float> image_data_channel_first(target_size * target_size * 3);
|
||||
for (int y = 0; y < target_size; y++) {
|
||||
for (int x = 0; x < target_size; x++) {
|
||||
for (int c = 0; c < 3; c++) {
|
||||
image_data_channel_first[c * target_size * target_size + y * target_size + x] =
|
||||
cropped_image_data[(y * target_size + x) * 3 + c];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Create Tensor
|
||||
auto image_tensor = Tensor::Empty({1, 3, target_size, target_size}, {kDLFloat, 32, 1}, device);
|
||||
image_tensor.CopyFromBytes((void*)image_data_channel_first.data(),
|
||||
target_size * target_size * 3 * sizeof(float));
|
||||
|
||||
return image_tensor;
|
||||
}
|
||||
|
||||
} // namespace json_ffi
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,31 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file json_ffi/image_utils.h
|
||||
* \brief The header of Image utils for JSON FFI Engine in MLC LLM.
|
||||
*/
|
||||
#ifndef MLC_LLM_JSON_FFI_IMAGE_UTILS_H_
|
||||
#define MLC_LLM_JSON_FFI_IMAGE_UTILS_H_
|
||||
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
#include <optional>
|
||||
#include <string>
|
||||
|
||||
#include "../support/result.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace json_ffi {
|
||||
|
||||
/*! \brief Load a base64 encoded image string into a CPU Tensor of shape {height, width, 3} */
|
||||
Result<tvm::runtime::Tensor> LoadImageFromBase64(const std::string& base64_str);
|
||||
|
||||
/*! \brief Preprocess the CPU image for CLIP encoder and return an Tensor on the given device */
|
||||
tvm::runtime::Tensor ClipPreprocessor(tvm::runtime::Tensor image_data, int target_size,
|
||||
DLDevice device);
|
||||
|
||||
} // namespace json_ffi
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_JSON_FFI_IMAGE_UTILS_H_
|
||||
@@ -0,0 +1,310 @@
|
||||
#include "json_ffi_engine.h"
|
||||
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
#include <tvm/ffi/extra/module.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
|
||||
#include "../serve/model.h"
|
||||
#include "../support/json_parser.h"
|
||||
#include "../support/module_vtable.h"
|
||||
#include "../support/result.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace json_ffi {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::Shape;
|
||||
|
||||
JSONFFIEngine::JSONFFIEngine() { engine_ = serve::ThreadedEngine::Create(); }
|
||||
|
||||
bool JSONFFIEngine::ChatCompletion(std::string request_json_str, std::string request_id) {
|
||||
bool success = this->AddRequest(request_json_str, request_id);
|
||||
if (!success) {
|
||||
this->StreamBackError(request_id);
|
||||
}
|
||||
return success;
|
||||
}
|
||||
|
||||
void JSONFFIEngine::StreamBackError(std::string request_id) {
|
||||
ChatCompletionMessage delta;
|
||||
delta.content = this->err_;
|
||||
delta.role = "assistant";
|
||||
|
||||
ChatCompletionStreamResponseChoice choice;
|
||||
choice.finish_reason = FinishReason::error;
|
||||
choice.index = 0;
|
||||
choice.delta = delta;
|
||||
|
||||
ChatCompletionStreamResponse response;
|
||||
response.id = request_id;
|
||||
response.choices = std::vector<ChatCompletionStreamResponseChoice>{choice};
|
||||
response.model = "json_ffi"; // TODO: Return model name from engine (or from args)
|
||||
response.system_fingerprint = "";
|
||||
|
||||
tvm::ffi::json::Array response_arr;
|
||||
response_arr.push_back(response.AsJSON());
|
||||
|
||||
// now stream back the final usage block, which is required.
|
||||
// NOTE: always stream back final usage block as it is an
|
||||
// invariant of the system
|
||||
response.choices.clear();
|
||||
tvm::ffi::json::Object dummy_usage;
|
||||
dummy_usage.Set("prompt_tokens", static_cast<int64_t>(0));
|
||||
dummy_usage.Set("completion_tokens", static_cast<int64_t>(0));
|
||||
dummy_usage.Set("total_tokens", static_cast<int64_t>(0));
|
||||
response.usage = tvm::ffi::json::Value(dummy_usage);
|
||||
response_arr.push_back(response.AsJSON());
|
||||
|
||||
std::string stream_back_json = tvm::ffi::json::Stringify(response_arr);
|
||||
this->request_stream_callback_(stream_back_json);
|
||||
}
|
||||
|
||||
bool JSONFFIEngine::AddRequest(std::string request_json_str, std::string request_id) {
|
||||
Result<ChatCompletionRequest> request_res = ChatCompletionRequest::FromJSON(request_json_str);
|
||||
if (request_res.IsErr()) {
|
||||
err_ = request_res.UnwrapErr();
|
||||
return false;
|
||||
}
|
||||
ChatCompletionRequest request = request_res.Unwrap();
|
||||
Array<Data> inputs;
|
||||
Array<String> stop_strs;
|
||||
bool is_special_request =
|
||||
(request.debug_config.has_value() &&
|
||||
request.debug_config.value().special_request != SpecialRequestKind::kNone);
|
||||
// special request does not have to go through prompt construction
|
||||
if (!is_special_request) {
|
||||
// get prompt: note, assistant was appended in the end.
|
||||
Result<std::vector<Data>> inputs_obj =
|
||||
CreatePrompt(this->conv_template_, request, this->model_config_, this->device_);
|
||||
if (inputs_obj.IsErr()) {
|
||||
err_ = inputs_obj.UnwrapErr();
|
||||
return false;
|
||||
}
|
||||
inputs = inputs_obj.Unwrap();
|
||||
|
||||
stop_strs.reserve(this->conv_template_.stop_str.size());
|
||||
for (const std::string& stop_str : this->conv_template_.stop_str) {
|
||||
stop_strs.push_back(stop_str);
|
||||
}
|
||||
if (request.stop.has_value()) {
|
||||
stop_strs.reserve(stop_strs.size() + request.stop.value().size());
|
||||
for (const std::string& stop_str : request.stop.value()) {
|
||||
stop_strs.push_back(stop_str);
|
||||
}
|
||||
}
|
||||
}
|
||||
// create a generation config from request
|
||||
const auto& default_gen_cfg = default_generation_config_;
|
||||
auto gen_cfg = tvm::ffi::make_object<GenerationConfigNode>();
|
||||
gen_cfg->n = request.n;
|
||||
gen_cfg->temperature = request.temperature.value_or(default_gen_cfg->temperature);
|
||||
gen_cfg->top_p = request.top_p.value_or(default_gen_cfg->top_p);
|
||||
gen_cfg->frequency_penalty =
|
||||
request.frequency_penalty.value_or(default_gen_cfg->frequency_penalty);
|
||||
gen_cfg->presence_penalty = request.presence_penalty.value_or(default_gen_cfg->presence_penalty);
|
||||
gen_cfg->logprobs = request.logprobs;
|
||||
gen_cfg->top_logprobs = request.top_logprobs;
|
||||
gen_cfg->logit_bias = request.logit_bias.value_or(default_gen_cfg->logit_bias);
|
||||
gen_cfg->seed = request.seed.value_or(std::random_device{}());
|
||||
gen_cfg->max_tokens = request.max_tokens.value_or(default_gen_cfg->max_tokens);
|
||||
gen_cfg->stop_strs = std::move(stop_strs);
|
||||
gen_cfg->stop_token_ids = conv_template_.stop_token_ids;
|
||||
gen_cfg->response_format = request.response_format.value_or(ResponseFormat());
|
||||
gen_cfg->debug_config = request.debug_config.value_or(DebugConfig());
|
||||
|
||||
Result<GenerationConfig> res_gen_config = GenerationConfig::Validate(GenerationConfig(gen_cfg));
|
||||
if (res_gen_config.IsErr()) {
|
||||
err_ = res_gen_config.UnwrapErr();
|
||||
return false;
|
||||
}
|
||||
|
||||
Request engine_request(request_id, inputs, res_gen_config.Unwrap());
|
||||
|
||||
// setup request state
|
||||
RequestState rstate;
|
||||
rstate.model = request.model.value_or("");
|
||||
rstate.streamer.reserve(gen_cfg->n);
|
||||
for (int i = 0; i < gen_cfg->n; ++i) {
|
||||
rstate.streamer.push_back(TextStreamer(tokenizer_));
|
||||
}
|
||||
request_map_[request_id] = std::move(rstate);
|
||||
|
||||
this->engine_->AddRequest(engine_request);
|
||||
return true;
|
||||
}
|
||||
|
||||
bool JSONFFIEngine::Abort(std::string request_id) {
|
||||
this->engine_->AbortRequest(request_id);
|
||||
auto it = request_map_.find(request_id);
|
||||
if (it != request_map_.end()) {
|
||||
request_map_.erase(it);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
std::string JSONFFIEngine::GetLastError() { return err_; }
|
||||
|
||||
void JSONFFIEngine::ExitBackgroundLoop() { this->engine_->ExitBackgroundLoop(); }
|
||||
|
||||
JSONFFIEngine::~JSONFFIEngine() { this->ExitBackgroundLoop(); }
|
||||
|
||||
class JSONFFIEngineImpl : public JSONFFIEngine, public ffi::ModuleObj {
|
||||
public:
|
||||
TVM_MODULE_VTABLE_BEGIN("mlc.json_ffi");
|
||||
TVM_MODULE_VTABLE_ENTRY("init_background_engine", &JSONFFIEngineImpl::InitBackgroundEngine);
|
||||
TVM_MODULE_VTABLE_ENTRY("reload", &JSONFFIEngineImpl::Reload);
|
||||
TVM_MODULE_VTABLE_ENTRY("unload", &JSONFFIEngineImpl::Unload);
|
||||
TVM_MODULE_VTABLE_ENTRY("reset", &JSONFFIEngineImpl::Reset);
|
||||
TVM_MODULE_VTABLE_ENTRY("chat_completion", &JSONFFIEngineImpl::ChatCompletion);
|
||||
TVM_MODULE_VTABLE_ENTRY("abort", &JSONFFIEngineImpl::Abort);
|
||||
TVM_MODULE_VTABLE_ENTRY("get_last_error", &JSONFFIEngineImpl::GetLastError);
|
||||
TVM_MODULE_VTABLE_ENTRY("run_background_loop", &JSONFFIEngineImpl::RunBackgroundLoop);
|
||||
TVM_MODULE_VTABLE_ENTRY("run_background_stream_back_loop",
|
||||
&JSONFFIEngineImpl::RunBackgroundStreamBackLoop);
|
||||
TVM_MODULE_VTABLE_ENTRY("exit_background_loop", &JSONFFIEngineImpl::ExitBackgroundLoop);
|
||||
TVM_MODULE_VTABLE_END();
|
||||
|
||||
void InitBackgroundEngine(int device_type, int device_id,
|
||||
Optional<Function> request_stream_callback) {
|
||||
DLDevice device{static_cast<DLDeviceType>(device_type), device_id};
|
||||
this->device_ = device;
|
||||
TVM_FFI_ICHECK(request_stream_callback.has_value())
|
||||
<< "JSONFFIEngine requires request stream callback function, but it is not given.";
|
||||
this->request_stream_callback_ = request_stream_callback.value();
|
||||
|
||||
auto frequest_stream_callback_wrapper = [this](ffi::PackedArgs args, ffi::Any* ret) {
|
||||
TVM_FFI_ICHECK_EQ(args.size(), 1);
|
||||
Array<RequestStreamOutput> delta_outputs = args[0].cast<Array<RequestStreamOutput>>();
|
||||
std::string responses = this->GetResponseFromStreamOutput(delta_outputs);
|
||||
this->request_stream_callback_(responses);
|
||||
};
|
||||
|
||||
request_stream_callback = Function(frequest_stream_callback_wrapper);
|
||||
this->engine_->InitThreadedEngine(device, std::move(request_stream_callback), std::nullopt);
|
||||
}
|
||||
|
||||
void Reload(String engine_config_json_str) {
|
||||
this->engine_->Reload(engine_config_json_str);
|
||||
this->default_generation_config_ = this->engine_->GetDefaultGenerationConfig();
|
||||
auto engine_config = this->engine_->GetCompleteEngineConfig();
|
||||
|
||||
// Load conversation template.
|
||||
Result<tvm::ffi::json::Object> model_config_json =
|
||||
serve::Model::LoadModelConfig(engine_config->model);
|
||||
TVM_FFI_ICHECK(model_config_json.IsOk()) << model_config_json.UnwrapErr();
|
||||
const tvm::ffi::json::Object& model_config_json_unwrapped = model_config_json.Unwrap();
|
||||
Result<Conversation> conv_template = Conversation::FromJSON(
|
||||
json::Lookup<tvm::ffi::json::Object>(model_config_json_unwrapped, "conv_template"));
|
||||
TVM_FFI_ICHECK(!conv_template.IsErr())
|
||||
<< "Invalid conversation template JSON: " << conv_template.UnwrapErr();
|
||||
this->conv_template_ = conv_template.Unwrap();
|
||||
this->model_config_ = ModelConfig::FromJSON(
|
||||
json::Lookup<tvm::ffi::json::Object>(model_config_json_unwrapped, "model_config"));
|
||||
this->tokenizer_ = Tokenizer::FromPath(engine_config->model);
|
||||
}
|
||||
|
||||
void Unload() { this->engine_->Unload(); }
|
||||
|
||||
void Reset() { this->engine_->Reset(); }
|
||||
|
||||
void RunBackgroundLoop() { this->engine_->RunBackgroundLoop(); }
|
||||
|
||||
void RunBackgroundStreamBackLoop() { this->engine_->RunBackgroundStreamBackLoop(); }
|
||||
|
||||
String GetResponseFromStreamOutput(Array<RequestStreamOutput> delta_outputs) {
|
||||
tvm::ffi::json::Array json_response_arr;
|
||||
for (const auto& delta_output : delta_outputs) {
|
||||
std::string request_id = delta_output->request_id;
|
||||
auto request_state_it = request_map_.find(request_id);
|
||||
if (request_state_it == request_map_.end()) continue;
|
||||
RequestState& rstate = request_state_it->second;
|
||||
|
||||
// build the final usage messages
|
||||
// invariant, we can always let other messages to come first
|
||||
// then the final usage messages, as final usage is always last
|
||||
if (delta_output->request_final_usage_json_str.has_value()) {
|
||||
ChatCompletionStreamResponse response;
|
||||
response.id = request_id;
|
||||
response.model = rstate.model;
|
||||
response.system_fingerprint = "";
|
||||
std::string usage_json_str = delta_output->request_final_usage_json_str.value();
|
||||
tvm::ffi::String parse_err;
|
||||
auto usage_json = tvm::ffi::json::Parse(usage_json_str, &parse_err);
|
||||
if (!parse_err.empty()) {
|
||||
err_ = parse_err;
|
||||
} else {
|
||||
response.usage = usage_json;
|
||||
}
|
||||
json_response_arr.push_back(response.AsJSON());
|
||||
request_map_.erase(request_state_it);
|
||||
continue;
|
||||
}
|
||||
TVM_FFI_ICHECK_NE(delta_output->group_finish_reason.size(), 0);
|
||||
TVM_FFI_ICHECK_EQ(delta_output->group_delta_token_ids.size(),
|
||||
delta_output->group_finish_reason.size());
|
||||
TVM_FFI_ICHECK_EQ(delta_output->group_delta_token_ids.size(), rstate.streamer.size());
|
||||
|
||||
ChatCompletionStreamResponse response;
|
||||
response.id = request_id;
|
||||
response.model = rstate.model;
|
||||
response.system_fingerprint = "";
|
||||
|
||||
for (size_t i = 0; i < delta_output->group_finish_reason.size(); ++i) {
|
||||
// choice
|
||||
ChatCompletionStreamResponseChoice choice;
|
||||
Optional<String> finish_reason = delta_output->group_finish_reason[i];
|
||||
if (finish_reason.has_value()) {
|
||||
if (finish_reason.value() == "stop") {
|
||||
choice.finish_reason = FinishReason::stop;
|
||||
} else if (finish_reason.value() == "length") {
|
||||
choice.finish_reason = FinishReason::length;
|
||||
} else if (finish_reason.value() == "tool_calls") {
|
||||
choice.finish_reason = FinishReason::tool_calls;
|
||||
} else if (finish_reason.value() == "error") {
|
||||
choice.finish_reason = FinishReason::error;
|
||||
}
|
||||
} else {
|
||||
choice.finish_reason = std::nullopt;
|
||||
}
|
||||
choice.index = static_cast<int>(i);
|
||||
ChatCompletionMessage delta;
|
||||
// Size of delta_output->group_delta_token_ids Array should be 1
|
||||
const Shape& delta_token_ids = delta_output->group_delta_token_ids[i];
|
||||
std::vector<int32_t> delta_token_ids_vec(delta_token_ids.begin(), delta_token_ids.end());
|
||||
std::string content = rstate.streamer[i]->Put(delta_token_ids_vec);
|
||||
if (finish_reason.has_value()) {
|
||||
content += rstate.streamer[i]->Finish();
|
||||
}
|
||||
if (!content.empty()) {
|
||||
delta.content = content;
|
||||
}
|
||||
delta.role = "assistant";
|
||||
choice.delta = delta;
|
||||
if (!choice.delta.content.IsNull() || choice.finish_reason.has_value()) {
|
||||
response.choices.push_back(choice);
|
||||
}
|
||||
}
|
||||
// if it is not the usage block, choices cannot be empty
|
||||
if (!response.choices.empty()) {
|
||||
json_response_arr.push_back(response.AsJSON());
|
||||
}
|
||||
}
|
||||
return tvm::ffi::json::Stringify(json_response_arr);
|
||||
}
|
||||
};
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("mlc.json_ffi.CreateJSONFFIEngine",
|
||||
[]() { return ffi::Module(tvm::ffi::make_object<JSONFFIEngineImpl>()); });
|
||||
}
|
||||
|
||||
} // namespace json_ffi
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,74 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file json_ffi/json_ffi_engine.h
|
||||
* \brief The header of JSON FFI engine in MLC LLM.
|
||||
*/
|
||||
#ifndef MLC_LLM_JSON_FFI_JSON_FFI_ENGINE_H_
|
||||
#define MLC_LLM_JSON_FFI_JSON_FFI_ENGINE_H_
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "../serve/threaded_engine.h"
|
||||
#include "../tokenizers/streamer.h"
|
||||
#include "conv_template.h"
|
||||
#include "openai_api_protocol.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace json_ffi {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using namespace mlc::llm::serve;
|
||||
|
||||
/*!
|
||||
* \brief // Todo: document this class, fields and member functions
|
||||
*/
|
||||
class JSONFFIEngine {
|
||||
public:
|
||||
JSONFFIEngine();
|
||||
|
||||
~JSONFFIEngine();
|
||||
|
||||
bool ChatCompletion(std::string request_json_str, std::string request_id);
|
||||
|
||||
bool AddRequest(std::string request_json_str, std::string request_id);
|
||||
|
||||
void StreamBackError(std::string request_id);
|
||||
|
||||
bool Abort(std::string request_id);
|
||||
|
||||
std::string GetLastError();
|
||||
|
||||
void ExitBackgroundLoop();
|
||||
|
||||
protected:
|
||||
/*! \brief local request state entry, one per reply stream. */
|
||||
struct RequestState {
|
||||
/*! \brief model to fill in reply. */
|
||||
std::string model;
|
||||
/*! \brief text streamer for each stream */
|
||||
std::vector<TextStreamer> streamer;
|
||||
};
|
||||
|
||||
std::unique_ptr<ThreadedEngine> engine_;
|
||||
std::string err_;
|
||||
Function request_stream_callback_;
|
||||
// tokenizer
|
||||
Tokenizer tokenizer_;
|
||||
// conversation template
|
||||
Conversation conv_template_;
|
||||
// generation config
|
||||
GenerationConfig default_generation_config_;
|
||||
// model config
|
||||
ModelConfig model_config_;
|
||||
// local device
|
||||
DLDevice device_;
|
||||
// request state map
|
||||
std::unordered_map<String, RequestState> request_map_;
|
||||
};
|
||||
|
||||
} // namespace json_ffi
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_JSON_FFI_JSON_FFI_ENGINE_H_
|
||||
@@ -0,0 +1,543 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file json_ffi/openai_api_protocol.cc
|
||||
* \brief The implementation of OpenAI API Protocol in MLC LLM.
|
||||
*/
|
||||
#include "openai_api_protocol.h"
|
||||
|
||||
#include "../support/json_parser.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace json_ffi {
|
||||
|
||||
Result<ChatFunction> ChatFunction::FromJSON(const tvm::ffi::json::Object& json_obj) {
|
||||
using TResult = Result<ChatFunction>;
|
||||
ChatFunction chat_func;
|
||||
|
||||
// description
|
||||
Result<std::optional<std::string>> description_res =
|
||||
json::LookupOptionalWithResultReturn<std::string>(json_obj, "description");
|
||||
if (description_res.IsErr()) {
|
||||
return TResult::Error(description_res.UnwrapErr());
|
||||
}
|
||||
chat_func.description = description_res.Unwrap();
|
||||
|
||||
// name
|
||||
Result<std::string> name_res = json::LookupWithResultReturn<std::string>(json_obj, "name");
|
||||
if (name_res.IsErr()) {
|
||||
return TResult::Error(name_res.UnwrapErr());
|
||||
}
|
||||
chat_func.name = name_res.Unwrap();
|
||||
|
||||
// parameters
|
||||
Result<tvm::ffi::json::Object> parameters_obj_res =
|
||||
json::LookupWithResultReturn<tvm::ffi::json::Object>(json_obj, "parameters");
|
||||
if (parameters_obj_res.IsErr()) {
|
||||
return TResult::Error(parameters_obj_res.UnwrapErr());
|
||||
}
|
||||
tvm::ffi::json::Object parameters_obj = parameters_obj_res.Unwrap();
|
||||
chat_func.parameters.reserve(parameters_obj.size());
|
||||
for (const auto& [key, value] : parameters_obj) {
|
||||
chat_func.parameters[key.cast<tvm::ffi::String>()] = tvm::ffi::json::Stringify(value);
|
||||
}
|
||||
|
||||
return TResult::Ok(chat_func);
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object ChatFunction::AsJSON() const {
|
||||
tvm::ffi::json::Object obj;
|
||||
if (this->description.has_value()) {
|
||||
obj.Set("description", this->description.value());
|
||||
}
|
||||
obj.Set("name", this->name);
|
||||
tvm::ffi::json::Object parameters_obj;
|
||||
for (const auto& pair : this->parameters) {
|
||||
parameters_obj.Set(pair.first, pair.second);
|
||||
}
|
||||
obj.Set("parameters", parameters_obj);
|
||||
return obj;
|
||||
}
|
||||
|
||||
Result<ChatTool> ChatTool::FromJSON(const tvm::ffi::json::Object& json_obj) {
|
||||
using TResult = Result<ChatTool>;
|
||||
ChatTool chatTool;
|
||||
|
||||
// function
|
||||
Result<tvm::ffi::json::Object> function_obj_res =
|
||||
json::LookupWithResultReturn<tvm::ffi::json::Object>(json_obj, "function");
|
||||
if (function_obj_res.IsErr()) {
|
||||
return TResult::Error(function_obj_res.UnwrapErr());
|
||||
}
|
||||
Result<ChatFunction> function = ChatFunction::FromJSON(function_obj_res.Unwrap());
|
||||
if (function.IsErr()) {
|
||||
return TResult::Error(function.UnwrapErr());
|
||||
}
|
||||
chatTool.function = function.Unwrap();
|
||||
|
||||
return TResult::Ok(chatTool);
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object ChatTool::AsJSON() const {
|
||||
tvm::ffi::json::Object obj;
|
||||
obj.Set("type", "function");
|
||||
obj.Set("function", this->function.AsJSON());
|
||||
return obj;
|
||||
}
|
||||
|
||||
Result<ChatFunctionCall> ChatFunctionCall::FromJSON(const tvm::ffi::json::Object& json_obj) {
|
||||
using TResult = Result<ChatFunctionCall>;
|
||||
ChatFunctionCall chat_func_call;
|
||||
|
||||
// name
|
||||
Result<std::string> name_res = json::LookupWithResultReturn<std::string>(json_obj, "name");
|
||||
if (name_res.IsErr()) {
|
||||
return TResult::Error(name_res.UnwrapErr());
|
||||
}
|
||||
chat_func_call.name = name_res.Unwrap();
|
||||
|
||||
// arguments
|
||||
Result<std::optional<tvm::ffi::json::Object>> arguments_obj_res =
|
||||
json::LookupOptionalWithResultReturn<tvm::ffi::json::Object>(json_obj, "arguments");
|
||||
if (arguments_obj_res.IsErr()) {
|
||||
return TResult::Error(arguments_obj_res.UnwrapErr());
|
||||
}
|
||||
std::optional<tvm::ffi::json::Object> arguments_obj = arguments_obj_res.Unwrap();
|
||||
if (arguments_obj.has_value()) {
|
||||
std::unordered_map<std::string, std::string> arguments;
|
||||
arguments.reserve(arguments_obj.value().size());
|
||||
for (const auto& [key, value] : arguments_obj.value()) {
|
||||
arguments[key.cast<tvm::ffi::String>()] = tvm::ffi::json::Stringify(value);
|
||||
}
|
||||
chat_func_call.arguments = std::move(arguments);
|
||||
}
|
||||
|
||||
return TResult::Ok(chat_func_call);
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object ChatFunctionCall::AsJSON() const {
|
||||
tvm::ffi::json::Object obj;
|
||||
tvm::ffi::json::Object arguments_obj;
|
||||
if (this->arguments.has_value()) {
|
||||
for (const auto& pair : this->arguments.value()) {
|
||||
arguments_obj.Set(pair.first, pair.second);
|
||||
}
|
||||
obj.Set("arguments", arguments_obj);
|
||||
}
|
||||
|
||||
obj.Set("name", this->name);
|
||||
return obj;
|
||||
}
|
||||
|
||||
Result<ChatToolCall> ChatToolCall::FromJSON(const tvm::ffi::json::Object& json_obj) {
|
||||
using TResult = Result<ChatToolCall>;
|
||||
ChatToolCall chat_tool_call;
|
||||
|
||||
// function
|
||||
Result<tvm::ffi::json::Object> function_obj_res =
|
||||
json::LookupWithResultReturn<tvm::ffi::json::Object>(json_obj, "function");
|
||||
if (function_obj_res.IsErr()) {
|
||||
return TResult::Error(function_obj_res.UnwrapErr());
|
||||
}
|
||||
Result<ChatFunctionCall> function_res = ChatFunctionCall::FromJSON(function_obj_res.Unwrap());
|
||||
if (function_res.IsErr()) {
|
||||
return TResult::Error(function_res.UnwrapErr());
|
||||
}
|
||||
chat_tool_call.function = function_res.Unwrap();
|
||||
|
||||
// overwrite default id
|
||||
Result<std::optional<std::string>> id_res =
|
||||
json::LookupOptionalWithResultReturn<std::string>(json_obj, "id");
|
||||
if (id_res.IsErr()) {
|
||||
return TResult::Error(id_res.UnwrapErr());
|
||||
}
|
||||
std::optional<std::string> id = id_res.UnwrapErr();
|
||||
if (id.has_value()) {
|
||||
chat_tool_call.id = id.value();
|
||||
}
|
||||
|
||||
return TResult::Ok(chat_tool_call);
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object ChatToolCall::AsJSON() const {
|
||||
tvm::ffi::json::Object obj;
|
||||
obj.Set("id", this->id);
|
||||
obj.Set("function", this->function.AsJSON());
|
||||
obj.Set("type", "function");
|
||||
return obj;
|
||||
}
|
||||
|
||||
Result<ChatCompletionMessage> ChatCompletionMessage::FromJSON(
|
||||
const tvm::ffi::json::Object& json_obj) {
|
||||
using TResult = Result<ChatCompletionMessage>;
|
||||
ChatCompletionMessage message;
|
||||
ChatCompletionMessageContent content;
|
||||
|
||||
// content
|
||||
if (json_obj.count("content") == 0) {
|
||||
return TResult::Error("ValueError: key \"content\" not found in the chat completion.");
|
||||
}
|
||||
tvm::ffi::json::Value content_val = json_obj.at("content");
|
||||
if (content_val.try_cast<std::string>().has_value()) {
|
||||
content = content_val.cast<std::string>();
|
||||
} else if (content_val == nullptr) {
|
||||
// skip
|
||||
} else {
|
||||
// most complicated case
|
||||
std::vector<std::unordered_map<std::string, std::string>> parts;
|
||||
Result<tvm::ffi::json::Array> content_arr_res =
|
||||
json::LookupWithResultReturn<tvm::ffi::json::Array>(json_obj, "content");
|
||||
if (content_arr_res.IsErr()) {
|
||||
return TResult::Error(content_arr_res.UnwrapErr());
|
||||
}
|
||||
tvm::ffi::json::Array content_arr = content_arr_res.Unwrap();
|
||||
for (const auto& item : content_arr) {
|
||||
if (!item.try_cast<tvm::ffi::json::Object>().has_value()) {
|
||||
return TResult::Error("The content of chat completion message is not an object");
|
||||
}
|
||||
tvm::ffi::json::Object item_obj = item.cast<tvm::ffi::json::Object>();
|
||||
std::unordered_map<std::string, std::string> item_map;
|
||||
for (const auto& [key, value] : item_obj) {
|
||||
item_map[key.cast<tvm::ffi::String>()] = tvm::ffi::json::Stringify(value);
|
||||
}
|
||||
parts.push_back(std::move(item_map));
|
||||
}
|
||||
content = parts;
|
||||
}
|
||||
message.content = content;
|
||||
|
||||
// role
|
||||
Result<std::string> role_str_res = json::LookupWithResultReturn<std::string>(json_obj, "role");
|
||||
if (role_str_res.IsErr()) {
|
||||
return TResult::Error(role_str_res.UnwrapErr());
|
||||
}
|
||||
std::string role_str = role_str_res.Unwrap();
|
||||
if (role_str == "system" || role_str == "user" || role_str == "assistant" || role_str == "tool") {
|
||||
message.role = role_str;
|
||||
} else {
|
||||
return TResult::Error("Invalid role in chat completion message: " + role_str);
|
||||
}
|
||||
|
||||
// name
|
||||
Result<std::optional<std::string>> name_res =
|
||||
json::LookupOptionalWithResultReturn<std::string>(json_obj, "name");
|
||||
if (name_res.IsErr()) {
|
||||
return TResult::Error(name_res.UnwrapErr());
|
||||
}
|
||||
message.name = name_res.Unwrap();
|
||||
|
||||
// tool calls
|
||||
Result<std::optional<tvm::ffi::json::Array>> tool_calls_arr_res =
|
||||
json::LookupOptionalWithResultReturn<tvm::ffi::json::Array>(json_obj, "tool_calls");
|
||||
if (tool_calls_arr_res.IsErr()) {
|
||||
return TResult::Error(tool_calls_arr_res.UnwrapErr());
|
||||
}
|
||||
std::optional<tvm::ffi::json::Array> tool_calls_arr = tool_calls_arr_res.Unwrap();
|
||||
if (tool_calls_arr.has_value()) {
|
||||
std::vector<ChatToolCall> tool_calls;
|
||||
tool_calls.reserve(tool_calls_arr.value().size());
|
||||
for (const auto& item : tool_calls_arr.value()) {
|
||||
if (!item.try_cast<tvm::ffi::json::Object>().has_value()) {
|
||||
return TResult::Error("A tool call item in the chat completion message is not an object");
|
||||
}
|
||||
Result<ChatToolCall> tool_call = ChatToolCall::FromJSON(item.cast<tvm::ffi::json::Object>());
|
||||
if (tool_call.IsErr()) {
|
||||
return TResult::Error(tool_call.UnwrapErr());
|
||||
}
|
||||
tool_calls.push_back(tool_call.Unwrap());
|
||||
}
|
||||
message.tool_calls = tool_calls;
|
||||
}
|
||||
|
||||
// tool call id
|
||||
Result<std::optional<std::string>> tool_call_id_res =
|
||||
json::LookupOptionalWithResultReturn<std::string>(json_obj, "tool_call_id");
|
||||
if (tool_call_id_res.IsErr()) {
|
||||
return TResult::Error(tool_call_id_res.UnwrapErr());
|
||||
}
|
||||
message.tool_call_id = tool_call_id_res.Unwrap();
|
||||
|
||||
return TResult::Ok(message);
|
||||
}
|
||||
|
||||
Result<ChatCompletionRequest> ChatCompletionRequest::FromJSON(const std::string& json_str) {
|
||||
using TResult = Result<ChatCompletionRequest>;
|
||||
Result<tvm::ffi::json::Object> json_obj_res = json::ParseToJSONObjectWithResultReturn(json_str);
|
||||
if (json_obj_res.IsErr()) {
|
||||
return TResult::Error(json_obj_res.UnwrapErr());
|
||||
}
|
||||
tvm::ffi::json::Object json_obj = json_obj_res.Unwrap();
|
||||
ChatCompletionRequest request;
|
||||
|
||||
// messages
|
||||
Result<tvm::ffi::json::Array> messages_arr_res =
|
||||
json::LookupWithResultReturn<tvm::ffi::json::Array>(json_obj, "messages");
|
||||
if (messages_arr_res.IsErr()) {
|
||||
return TResult::Error(messages_arr_res.UnwrapErr());
|
||||
}
|
||||
std::vector<ChatCompletionMessage> messages;
|
||||
tvm::ffi::json::Array messages_arr = messages_arr_res.Unwrap();
|
||||
for (const auto& item : messages_arr) {
|
||||
if (!item.try_cast<tvm::ffi::json::Object>().has_value()) {
|
||||
return TResult::Error("A message in chat completion request is not object");
|
||||
}
|
||||
tvm::ffi::json::Object item_obj = item.cast<tvm::ffi::json::Object>();
|
||||
Result<ChatCompletionMessage> message = ChatCompletionMessage::FromJSON(item_obj);
|
||||
if (message.IsErr()) {
|
||||
return TResult::Error(message.UnwrapErr());
|
||||
}
|
||||
messages.push_back(message.Unwrap());
|
||||
}
|
||||
request.messages = messages;
|
||||
|
||||
// model
|
||||
Result<std::optional<std::string>> model_res =
|
||||
json::LookupOptionalWithResultReturn<std::string>(json_obj, "model");
|
||||
if (model_res.IsErr()) {
|
||||
return TResult::Error(model_res.UnwrapErr());
|
||||
}
|
||||
request.model = model_res.Unwrap();
|
||||
|
||||
// temperature
|
||||
Result<std::optional<double>> temperature_res =
|
||||
json::LookupOptionalWithResultReturn<double>(json_obj, "temperature");
|
||||
if (temperature_res.IsErr()) {
|
||||
return TResult::Error(temperature_res.UnwrapErr());
|
||||
}
|
||||
request.temperature = temperature_res.Unwrap();
|
||||
// top_p
|
||||
Result<std::optional<double>> top_p_res =
|
||||
json::LookupOptionalWithResultReturn<double>(json_obj, "top_p");
|
||||
if (top_p_res.IsErr()) {
|
||||
return TResult::Error(top_p_res.UnwrapErr());
|
||||
}
|
||||
request.top_p = top_p_res.Unwrap();
|
||||
// max_tokens
|
||||
Result<std::optional<int64_t>> max_tokens_res =
|
||||
json::LookupOptionalWithResultReturn<int64_t>(json_obj, "max_tokens");
|
||||
if (max_tokens_res.IsErr()) {
|
||||
return TResult::Error(max_tokens_res.UnwrapErr());
|
||||
}
|
||||
request.max_tokens = max_tokens_res.Unwrap();
|
||||
// n
|
||||
Result<int64_t> n_res = json::LookupOrDefaultWithResultReturn<int64_t>(json_obj, "n", 1);
|
||||
if (n_res.IsErr()) {
|
||||
return TResult::Error(n_res.UnwrapErr());
|
||||
}
|
||||
request.n = n_res.Unwrap();
|
||||
// frequency_penalty
|
||||
Result<std::optional<double>> frequency_penalty_res =
|
||||
json::LookupOptionalWithResultReturn<double>(json_obj, "frequency_penalty");
|
||||
if (frequency_penalty_res.IsErr()) {
|
||||
return TResult::Error(frequency_penalty_res.UnwrapErr());
|
||||
}
|
||||
request.frequency_penalty = frequency_penalty_res.Unwrap();
|
||||
// presence_penalty
|
||||
Result<std::optional<double>> presence_penalty_res =
|
||||
json::LookupOptionalWithResultReturn<double>(json_obj, "presence_penalty");
|
||||
if (presence_penalty_res.IsErr()) {
|
||||
return TResult::Error(presence_penalty_res.UnwrapErr());
|
||||
}
|
||||
request.presence_penalty = presence_penalty_res.Unwrap();
|
||||
// seed
|
||||
Result<std::optional<int64_t>> seed_res =
|
||||
json::LookupOptionalWithResultReturn<int64_t>(json_obj, "seed");
|
||||
if (seed_res.IsErr()) {
|
||||
return TResult::Error(seed_res.UnwrapErr());
|
||||
}
|
||||
request.seed = seed_res.Unwrap();
|
||||
|
||||
// stop strings
|
||||
Result<std::optional<tvm::ffi::json::Array>> stop_strs_res =
|
||||
json::LookupOptionalWithResultReturn<tvm::ffi::json::Array>(json_obj, "stop");
|
||||
if (stop_strs_res.IsErr()) {
|
||||
return TResult::Error(stop_strs_res.UnwrapErr());
|
||||
}
|
||||
std::optional<tvm::ffi::json::Array> stop_strs = stop_strs_res.Unwrap();
|
||||
if (stop_strs.has_value()) {
|
||||
std::vector<std::string> stop;
|
||||
for (const auto& stop_str_value : stop_strs.value()) {
|
||||
if (!stop_str_value.try_cast<std::string>().has_value()) {
|
||||
return TResult::Error("One given value in field \"stop\" is not a string.");
|
||||
}
|
||||
stop.push_back(stop_str_value.cast<std::string>());
|
||||
}
|
||||
request.stop = std::move(stop);
|
||||
}
|
||||
|
||||
// tool_choice
|
||||
Result<std::string> tool_choice_res =
|
||||
json::LookupOrDefaultWithResultReturn<std::string>(json_obj, "tool_choice", "auto");
|
||||
if (tool_choice_res.IsErr()) {
|
||||
return TResult::Error(tool_choice_res.UnwrapErr());
|
||||
}
|
||||
request.tool_choice = tool_choice_res.Unwrap();
|
||||
|
||||
// tools
|
||||
Result<std::optional<tvm::ffi::json::Array>> tools_arr_res =
|
||||
json::LookupOptionalWithResultReturn<tvm::ffi::json::Array>(json_obj, "tools");
|
||||
if (tool_choice_res.IsErr()) {
|
||||
return TResult::Error(tool_choice_res.UnwrapErr());
|
||||
}
|
||||
std::optional<tvm::ffi::json::Array> tools_arr = tools_arr_res.Unwrap();
|
||||
if (tools_arr.has_value()) {
|
||||
std::vector<ChatTool> tools;
|
||||
tools.reserve(tools_arr.value().size());
|
||||
for (const auto& item : tools_arr.value()) {
|
||||
if (!item.try_cast<tvm::ffi::json::Object>().has_value()) {
|
||||
return TResult::Error("A tool of the chat completion request is not an object");
|
||||
}
|
||||
Result<ChatTool> tool = ChatTool::FromJSON(item.cast<tvm::ffi::json::Object>());
|
||||
if (tool.IsErr()) {
|
||||
return TResult::Error(tool.UnwrapErr());
|
||||
}
|
||||
tools.push_back(tool.Unwrap());
|
||||
}
|
||||
request.tools = tools;
|
||||
}
|
||||
|
||||
// response format
|
||||
std::optional<tvm::ffi::json::Object> response_format_obj =
|
||||
json::LookupOptional<tvm::ffi::json::Object>(json_obj, "response_format");
|
||||
if (response_format_obj.has_value()) {
|
||||
Result<ResponseFormat> response_format_res =
|
||||
ResponseFormat::FromJSON(response_format_obj.value());
|
||||
if (response_format_res.IsErr()) {
|
||||
return TResult::Error(response_format_res.UnwrapErr());
|
||||
}
|
||||
request.response_format = response_format_res.Unwrap();
|
||||
}
|
||||
|
||||
// debug_config
|
||||
Result<std::optional<tvm::ffi::json::Object>> debug_config_opt_res =
|
||||
json::LookupOptionalWithResultReturn<tvm::ffi::json::Object>(json_obj, "debug_config");
|
||||
if (debug_config_opt_res.IsErr()) {
|
||||
return TResult::Error(debug_config_opt_res.UnwrapErr());
|
||||
}
|
||||
auto debug_config_opt = debug_config_opt_res.Unwrap();
|
||||
if (debug_config_opt.has_value()) {
|
||||
Result<DebugConfig> debug_config_res = DebugConfig::FromJSON(debug_config_opt.value());
|
||||
if (debug_config_res.IsErr()) {
|
||||
return TResult::Error(debug_config_res.UnwrapErr());
|
||||
}
|
||||
request.debug_config = debug_config_res.Unwrap();
|
||||
}
|
||||
|
||||
// TODO: Other parameters
|
||||
return TResult::Ok(request);
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object ChatCompletionMessage::AsJSON() const {
|
||||
tvm::ffi::json::Object obj;
|
||||
|
||||
if (this->content.IsText()) {
|
||||
obj.Set("content", this->content.Text());
|
||||
} else if (this->content.IsParts()) {
|
||||
tvm::ffi::json::Array content_arr;
|
||||
for (const auto& item : this->content.Parts()) {
|
||||
tvm::ffi::json::Object item_obj;
|
||||
for (const auto& pair : item) {
|
||||
item_obj.Set(pair.first, pair.second);
|
||||
}
|
||||
content_arr.push_back(item_obj);
|
||||
}
|
||||
obj.Set("content", content_arr);
|
||||
}
|
||||
|
||||
obj.Set("role", this->role);
|
||||
|
||||
if (this->name.has_value()) {
|
||||
obj.Set("name", this->name.value());
|
||||
}
|
||||
if (this->tool_call_id.has_value()) {
|
||||
obj.Set("tool_call_id", this->tool_call_id.value());
|
||||
}
|
||||
if (this->tool_calls.has_value()) {
|
||||
tvm::ffi::json::Array tool_calls_arr;
|
||||
for (const auto& tool_call : this->tool_calls.value()) {
|
||||
tool_calls_arr.push_back(tool_call.AsJSON());
|
||||
}
|
||||
obj.Set("tool_calls", tool_calls_arr);
|
||||
}
|
||||
return obj;
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object ChatCompletionResponseChoice::AsJSON() const {
|
||||
tvm::ffi::json::Object obj;
|
||||
if (!this->finish_reason.has_value()) {
|
||||
obj.Set("finish_reason", nullptr);
|
||||
} else {
|
||||
if (this->finish_reason == FinishReason::stop) {
|
||||
obj.Set("finish_reason", "stop");
|
||||
} else if (this->finish_reason == FinishReason::length) {
|
||||
obj.Set("finish_reason", "length");
|
||||
} else if (this->finish_reason == FinishReason::tool_calls) {
|
||||
obj.Set("finish_reason", "tool_calls");
|
||||
} else if (this->finish_reason == FinishReason::error) {
|
||||
obj.Set("finish_reason", "error");
|
||||
}
|
||||
}
|
||||
obj.Set("index", static_cast<int64_t>(this->index));
|
||||
obj.Set("message", this->message.AsJSON());
|
||||
return obj;
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object ChatCompletionStreamResponseChoice::AsJSON() const {
|
||||
tvm::ffi::json::Object obj;
|
||||
if (!this->finish_reason.has_value()) {
|
||||
obj.Set("finish_reason", nullptr);
|
||||
} else {
|
||||
if (this->finish_reason.value() == FinishReason::stop) {
|
||||
obj.Set("finish_reason", "stop");
|
||||
} else if (this->finish_reason.value() == FinishReason::length) {
|
||||
obj.Set("finish_reason", "length");
|
||||
} else if (this->finish_reason.value() == FinishReason::tool_calls) {
|
||||
obj.Set("finish_reason", "tool_calls");
|
||||
} else if (this->finish_reason.value() == FinishReason::error) {
|
||||
obj.Set("finish_reason", "error");
|
||||
}
|
||||
}
|
||||
|
||||
obj.Set("index", static_cast<int64_t>(this->index));
|
||||
obj.Set("delta", this->delta.AsJSON());
|
||||
return obj;
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object ChatCompletionResponse::AsJSON() const {
|
||||
tvm::ffi::json::Object obj;
|
||||
obj.Set("id", this->id);
|
||||
tvm::ffi::json::Array choices_arr;
|
||||
for (const auto& choice : this->choices) {
|
||||
choices_arr.push_back(choice.AsJSON());
|
||||
}
|
||||
obj.Set("choices", choices_arr);
|
||||
obj.Set("created", static_cast<int64_t>(this->created));
|
||||
obj.Set("model", this->model);
|
||||
obj.Set("system_fingerprint", this->system_fingerprint);
|
||||
obj.Set("object", this->object);
|
||||
return obj;
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object ChatCompletionStreamResponse::AsJSON() const {
|
||||
tvm::ffi::json::Object obj;
|
||||
obj.Set("id", this->id);
|
||||
|
||||
tvm::ffi::json::Array choices_arr;
|
||||
for (const auto& choice : this->choices) {
|
||||
choices_arr.push_back(choice.AsJSON());
|
||||
}
|
||||
obj.Set("choices", choices_arr);
|
||||
|
||||
obj.Set("created", static_cast<int64_t>(this->created));
|
||||
obj.Set("model", this->model);
|
||||
obj.Set("system_fingerprint", this->system_fingerprint);
|
||||
obj.Set("object", this->object);
|
||||
if (usage.has_value()) {
|
||||
obj.Set("usage", usage.value());
|
||||
}
|
||||
return obj;
|
||||
}
|
||||
|
||||
} // namespace json_ffi
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,208 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file json_ffi/openai_api_protocol.h
|
||||
* \brief The header of OpenAI API Protocol in MLC LLM.
|
||||
*/
|
||||
#ifndef MLC_LLM_JSON_FFI_OPENAI_API_PROTOCOL_H
|
||||
#define MLC_LLM_JSON_FFI_OPENAI_API_PROTOCOL_H
|
||||
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
|
||||
#include <ctime>
|
||||
#include <optional>
|
||||
#include <random>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "../serve/config.h"
|
||||
#include "../support/result.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace json_ffi {
|
||||
|
||||
using serve::DebugConfig;
|
||||
using serve::ResponseFormat;
|
||||
|
||||
enum class Type { text, json_object, function };
|
||||
enum class FinishReason { stop, length, tool_calls, error };
|
||||
|
||||
inline std::string GenerateUUID(size_t length) {
|
||||
auto randchar = []() -> char {
|
||||
const char charset[] =
|
||||
"0123456789"
|
||||
"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
|
||||
"abcdefghijklmnopqrstuvwxyz";
|
||||
const size_t max_index = (sizeof(charset) - 1);
|
||||
return charset[rand() % max_index];
|
||||
};
|
||||
std::string str(length, 0);
|
||||
std::generate_n(str.begin(), length, randchar);
|
||||
return str;
|
||||
}
|
||||
|
||||
class ChatFunction {
|
||||
public:
|
||||
std::optional<std::string> description = std::nullopt;
|
||||
std::string name;
|
||||
// Todo: change to std::vector<std::pair<std::string, std::string>>?
|
||||
std::unordered_map<std::string, std::string>
|
||||
parameters; // Assuming parameters are string key-value pairs
|
||||
|
||||
static Result<ChatFunction> FromJSON(const tvm::ffi::json::Object& json);
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
class ChatTool {
|
||||
public:
|
||||
Type type = Type::function;
|
||||
ChatFunction function;
|
||||
|
||||
static Result<ChatTool> FromJSON(const tvm::ffi::json::Object& json);
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
class ChatFunctionCall {
|
||||
public:
|
||||
std::string name;
|
||||
std::optional<std::unordered_map<std::string, std::string>> arguments =
|
||||
std::nullopt; // Assuming arguments are string key-value pairs
|
||||
|
||||
static Result<ChatFunctionCall> FromJSON(const tvm::ffi::json::Object& json);
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
class ChatToolCall {
|
||||
public:
|
||||
std::string id = "call_" + GenerateUUID(8);
|
||||
Type type = Type::function;
|
||||
ChatFunctionCall function;
|
||||
|
||||
static Result<ChatToolCall> FromJSON(const tvm::ffi::json::Object& json);
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
class ChatCompletionMessageContent {
|
||||
public:
|
||||
ChatCompletionMessageContent() = default;
|
||||
|
||||
ChatCompletionMessageContent(std::nullopt_t) {} // NOLINT(*)
|
||||
|
||||
ChatCompletionMessageContent(std::string text) : text_(text) {} // NOLINT(*)
|
||||
|
||||
ChatCompletionMessageContent(
|
||||
std::vector<std::unordered_map<std::string, std::string>> parts) // NOLINT(*)
|
||||
: parts_(parts) {}
|
||||
|
||||
bool IsNull() const { return !IsText() && !IsParts(); }
|
||||
|
||||
bool IsText() const { return text_.operator bool(); }
|
||||
|
||||
bool IsParts() const { return parts_.operator bool(); }
|
||||
|
||||
const std::string& Text() const { return text_.value(); }
|
||||
|
||||
const std::vector<std::unordered_map<std::string, std::string>>& Parts() const {
|
||||
return parts_.value();
|
||||
}
|
||||
|
||||
private:
|
||||
/*! \brief used to store text content */
|
||||
std::optional<std::string> text_;
|
||||
std::optional<std::vector<std::unordered_map<std::string, std::string>>> parts_;
|
||||
};
|
||||
|
||||
class ChatCompletionMessage {
|
||||
public:
|
||||
ChatCompletionMessageContent content =
|
||||
std::nullopt; // Assuming content is a list of string key-value pairs
|
||||
std::string role;
|
||||
std::optional<std::string> name = std::nullopt;
|
||||
std::optional<std::vector<ChatToolCall>> tool_calls = std::nullopt;
|
||||
std::optional<std::string> tool_call_id = std::nullopt;
|
||||
|
||||
static Result<ChatCompletionMessage> FromJSON(const tvm::ffi::json::Object& json);
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
class ChatCompletionRequest {
|
||||
public:
|
||||
std::vector<ChatCompletionMessage> messages;
|
||||
std::optional<std::string> model = std::nullopt;
|
||||
std::optional<double> frequency_penalty = std::nullopt;
|
||||
std::optional<double> presence_penalty = std::nullopt;
|
||||
bool logprobs = false;
|
||||
int top_logprobs = 0;
|
||||
std::optional<std::vector<std::pair<int, float>>> logit_bias = std::nullopt;
|
||||
std::optional<int> max_tokens = std::nullopt;
|
||||
int n = 1;
|
||||
std::optional<int> seed = std::nullopt;
|
||||
std::optional<std::vector<std::string>> stop = std::nullopt;
|
||||
bool stream = false;
|
||||
std::optional<double> temperature = std::nullopt;
|
||||
std::optional<double> top_p = std::nullopt;
|
||||
std::optional<std::vector<ChatTool>> tools = std::nullopt;
|
||||
std::optional<std::string> tool_choice = std::nullopt;
|
||||
std::optional<std::string> user = std::nullopt;
|
||||
bool ignore_eos = false;
|
||||
std::optional<ResponseFormat> response_format = std::nullopt;
|
||||
std::optional<DebugConfig> debug_config = std::nullopt;
|
||||
|
||||
/*! \brief Parse and create a ChatCompletionRequest instance from the given JSON string. */
|
||||
static Result<ChatCompletionRequest> FromJSON(const std::string& json_str);
|
||||
|
||||
// TODO: check_penalty_range, check_logit_bias, check_logprobs
|
||||
};
|
||||
|
||||
class ChatCompletionResponseChoice {
|
||||
public:
|
||||
std::optional<FinishReason> finish_reason;
|
||||
int index = 0;
|
||||
ChatCompletionMessage message;
|
||||
// TODO: logprobs
|
||||
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
class ChatCompletionStreamResponseChoice {
|
||||
public:
|
||||
std::optional<FinishReason> finish_reason;
|
||||
int index = 0;
|
||||
ChatCompletionMessage delta;
|
||||
// TODO: logprobs
|
||||
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
class ChatCompletionResponse {
|
||||
public:
|
||||
std::string id;
|
||||
std::vector<ChatCompletionResponseChoice> choices;
|
||||
int created = static_cast<int>(std::time(nullptr));
|
||||
std::string model;
|
||||
std::string system_fingerprint;
|
||||
std::string object = "chat.completion";
|
||||
// TODO: usage_info
|
||||
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
class ChatCompletionStreamResponse {
|
||||
public:
|
||||
std::string id;
|
||||
std::vector<ChatCompletionStreamResponseChoice> choices;
|
||||
int created = static_cast<int>(std::time(nullptr));
|
||||
std::string model;
|
||||
std::string system_fingerprint;
|
||||
std::string object = "chat.completion.chunk";
|
||||
std::optional<tvm::ffi::json::Value> usage;
|
||||
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
} // namespace json_ffi
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_JSON_FFI_OPENAI_API_PROTOCOL_H
|
||||
@@ -0,0 +1,155 @@
|
||||
#include "./model.h"
|
||||
|
||||
#include <unordered_map>
|
||||
|
||||
#include "../support/json_parser.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Function;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::Optional;
|
||||
|
||||
ModelMetadata::Param::Preproc ModelMetadata::Param::Preproc::FromJSON(
|
||||
const tvm::ffi::json::Object& js, const tvm::ffi::json::Object& model_config) {
|
||||
Preproc preproc;
|
||||
TVM_FFI_ICHECK_GE(js.size(), 3) << "ValueError: Invalid preprocessing info in JSON";
|
||||
preproc.func_name = json::Lookup<std::string>(js, "func_name");
|
||||
json::SymShapeTuple sym_out_shape = json::Lookup<json::SymShapeTuple>(js, "out_shape");
|
||||
preproc.out_shape = sym_out_shape.ToStatic(model_config);
|
||||
json::SymShapeTuple sym_in_shape =
|
||||
json::LookupOrDefault<json::SymShapeTuple>(js, "in_shape", sym_out_shape);
|
||||
preproc.in_shape = sym_in_shape.ToStatic(model_config);
|
||||
preproc.out_dtype = json::Lookup<DLDataType>(js, "out_dtype");
|
||||
return preproc;
|
||||
}
|
||||
|
||||
ModelMetadata::Param ModelMetadata::Param::FromJSON(const tvm::ffi::json::Object& param,
|
||||
const tvm::ffi::json::Object& model_config) {
|
||||
Param result;
|
||||
result.name = json::Lookup<std::string>(param, "name");
|
||||
result.dtype = json::Lookup<DLDataType>(param, "dtype");
|
||||
// A shape being `-1` means that it is dynamic
|
||||
json::SymShapeTuple sym_shape = json::Lookup<json::SymShapeTuple>(param, "shape");
|
||||
result.shape = sym_shape.ToStatic(model_config);
|
||||
// - "preproc"
|
||||
tvm::ffi::json::Array preprocs = json::Lookup<tvm::ffi::json::Array>(param, "preprocs");
|
||||
result.preprocs.reserve(preprocs.size());
|
||||
for (int i = 0; i < preprocs.size(); i++) {
|
||||
result.preprocs.emplace_back(ModelMetadata::Param::Preproc::FromJSON(
|
||||
json::Lookup<tvm::ffi::json::Object>(preprocs, i), model_config));
|
||||
}
|
||||
// - "pipeline_stages"
|
||||
int pipeline_parallel_stages =
|
||||
json::LookupOrDefault<int64_t>(model_config, "pipeline_parallel_stages", 1);
|
||||
std::optional<tvm::ffi::json::Array> opt_pipeline_stages =
|
||||
json::LookupOptional<tvm::ffi::json::Array>(param, "pipeline_stages");
|
||||
if (pipeline_parallel_stages > 1) {
|
||||
TVM_FFI_ICHECK(opt_pipeline_stages.has_value())
|
||||
<< "The pipeline stage is undefined for parameter \"" << result.name
|
||||
<< "\" when the number of pipeline parallel stages is " << pipeline_parallel_stages;
|
||||
}
|
||||
if (opt_pipeline_stages.has_value()) {
|
||||
result.pipeline_stages.reserve(opt_pipeline_stages.value().size());
|
||||
for (const tvm::ffi::json::Value& v : opt_pipeline_stages.value()) {
|
||||
auto int_opt = v.try_cast<int64_t>();
|
||||
TVM_FFI_ICHECK(int_opt.has_value()) << "Pipeline stage is not a integer.";
|
||||
result.pipeline_stages.push_back(*int_opt);
|
||||
}
|
||||
} else {
|
||||
result.pipeline_stages = {0};
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
ModelMetadata::KVCacheMetadata ModelMetadata::KVCacheMetadata::FromJSON(
|
||||
const tvm::ffi::json::Object& json) {
|
||||
KVCacheMetadata kv_cache_metadata;
|
||||
kv_cache_metadata.num_hidden_layers = json::Lookup<int64_t>(json, "num_hidden_layers");
|
||||
kv_cache_metadata.head_dim = json::Lookup<int64_t>(json, "head_dim");
|
||||
kv_cache_metadata.num_attention_heads = json::Lookup<int64_t>(json, "num_attention_heads");
|
||||
kv_cache_metadata.num_key_value_heads = json::Lookup<int64_t>(json, "num_key_value_heads");
|
||||
return kv_cache_metadata;
|
||||
}
|
||||
|
||||
ModelMetadata ModelMetadata::FromJSON(const tvm::ffi::json::Object& metadata,
|
||||
const tvm::ffi::json::Object& model_config) {
|
||||
ModelMetadata result;
|
||||
result.model_type = json::Lookup<std::string>(metadata, "model_type");
|
||||
result.quantization = json::Lookup<std::string>(metadata, "quantization");
|
||||
result.context_window_size = json::Lookup<int64_t>(metadata, "context_window_size");
|
||||
result.prefill_chunk_size = json::Lookup<int64_t>(metadata, "prefill_chunk_size");
|
||||
result.max_batch_size = json::Lookup<int64_t>(metadata, "max_batch_size");
|
||||
if (metadata.count("sliding_window_size"))
|
||||
result.sliding_window_size = json::Lookup<int64_t>(metadata, "sliding_window_size");
|
||||
if (metadata.count("sliding_window")) // to be removed after SLM migration
|
||||
result.sliding_window_size = json::Lookup<int64_t>(metadata, "sliding_window");
|
||||
if (metadata.count("attention_sink_size")) // remove after sink is decoupled from model lib
|
||||
result.attention_sink_size = json::Lookup<int64_t>(metadata, "attention_sink_size");
|
||||
result.seqlen_padding_factor =
|
||||
json::LookupOrDefault<int64_t>(metadata, "seqlen_padding_factor", 1);
|
||||
result.tensor_parallel_shards = json::Lookup<int64_t>(metadata, "tensor_parallel_shards");
|
||||
result.pipeline_parallel_stages =
|
||||
json::LookupOrDefault<int64_t>(metadata, "pipeline_parallel_stages", 1);
|
||||
result.disaggregation = json::LookupOrDefault<bool>(metadata, "disaggregation", false);
|
||||
result.model_task = json::LookupOrDefault<std::string>(metadata, "model_task", "chat");
|
||||
if (metadata.count("embedding_metadata")) {
|
||||
tvm::ffi::json::Object emb =
|
||||
json::Lookup<tvm::ffi::json::Object>(metadata, "embedding_metadata");
|
||||
result.embedding_model_type = json::LookupOrDefault<std::string>(emb, "model_type", "");
|
||||
result.embedding_pooling_strategy =
|
||||
json::LookupOrDefault<std::string>(emb, "pooling_strategy", "");
|
||||
result.embedding_normalize = json::LookupOrDefault<bool>(emb, "normalize", false);
|
||||
}
|
||||
result.kv_state_kind = KVStateKindFromString(
|
||||
json::LookupOrDefault<std::string>(metadata, "kv_state_kind", "kv_cache"));
|
||||
if (result.kv_state_kind != KVStateKind::kNone &&
|
||||
result.kv_state_kind != KVStateKind::kRNNState) {
|
||||
result.kv_cache_metadata =
|
||||
KVCacheMetadata::FromJSON(json::Lookup<tvm::ffi::json::Object>(metadata, "kv_cache"));
|
||||
} else {
|
||||
result.kv_cache_metadata = {/*num_hidden_layers=*/0,
|
||||
/*head_dim=*/0,
|
||||
/*num_attention_heads=*/0,
|
||||
/*num_key_value_heads=*/0};
|
||||
}
|
||||
{
|
||||
std::vector<ModelMetadata::Param>& params = result.params;
|
||||
tvm::ffi::json::Array json_params = json::Lookup<tvm::ffi::json::Array>(metadata, "params");
|
||||
params.reserve(json_params.size());
|
||||
for (int i = 0, n = json_params.size(); i < n; ++i) {
|
||||
params.emplace_back(ModelMetadata::Param::FromJSON(
|
||||
json::Lookup<tvm::ffi::json::Object>(json_params, i), model_config));
|
||||
}
|
||||
}
|
||||
{
|
||||
std::unordered_map<std::string, int64_t>& memory_usage = result.memory_usage;
|
||||
tvm::ffi::json::Object json_memory_usage =
|
||||
json::Lookup<tvm::ffi::json::Object>(metadata, "memory_usage");
|
||||
memory_usage.reserve(json_memory_usage.size());
|
||||
for (const auto& [key, val] : json_memory_usage) {
|
||||
std::string func_name = key.cast<tvm::ffi::String>();
|
||||
memory_usage[func_name] = json::Lookup<int64_t>(json_memory_usage, func_name);
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
ModelMetadata ModelMetadata::FromModule(Module module, const tvm::ffi::json::Object& model_config) {
|
||||
std::string json_str = "";
|
||||
Optional<Function> pf = module->GetFunction("_metadata");
|
||||
TVM_FFI_ICHECK(pf.has_value()) << "ValueError: _metadata function not found in module";
|
||||
json_str = pf.value()().cast<String>();
|
||||
tvm::ffi::json::Object json = json::ParseToJSONObject(json_str);
|
||||
try {
|
||||
return ModelMetadata::FromJSON(json, model_config);
|
||||
} catch (const std::exception& e) {
|
||||
LOG(WARNING) << "Failed to parse metadata:\n" << json_str << "\nerror: " << e.what();
|
||||
throw e;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,115 @@
|
||||
/*!
|
||||
* \file model.h
|
||||
* \brief Metadata stored in model lib
|
||||
*/
|
||||
#ifndef MLC_LLM_CPP_MODEL_METADATA_H_
|
||||
#define MLC_LLM_CPP_MODEL_METADATA_H_
|
||||
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/dtype.h>
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
#include <tvm/ffi/extra/module.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
#include <tvm/runtime/logging.h>
|
||||
|
||||
#include <unordered_map>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
using tvm::ffi::Module;
|
||||
using tvm::ffi::Shape;
|
||||
using tvm::ffi::String;
|
||||
|
||||
/*! \brief The kind of cache. */
|
||||
enum class KVStateKind : int {
|
||||
kKVCache = 0,
|
||||
kRNNState = 1,
|
||||
kNone = 2,
|
||||
kHybrid = 3,
|
||||
};
|
||||
|
||||
inline std::string KVStateKindToString(KVStateKind kv_state_kind) {
|
||||
if (kv_state_kind == KVStateKind::kKVCache) {
|
||||
return "kv_cache";
|
||||
} else if (kv_state_kind == KVStateKind::kRNNState) {
|
||||
return "rnn_state";
|
||||
} else if (kv_state_kind == KVStateKind::kNone) {
|
||||
return "none";
|
||||
} else if (kv_state_kind == KVStateKind::kHybrid) {
|
||||
return "hybrid";
|
||||
} else {
|
||||
LOG(FATAL) << "Invalid kv state kind: " << static_cast<int>(kv_state_kind);
|
||||
}
|
||||
}
|
||||
|
||||
inline KVStateKind KVStateKindFromString(const std::string& kv_state_kind) {
|
||||
if (kv_state_kind == "kv_cache") {
|
||||
return KVStateKind::kKVCache;
|
||||
} else if (kv_state_kind == "rnn_state") {
|
||||
return KVStateKind::kRNNState;
|
||||
} else if (kv_state_kind == "none") {
|
||||
return KVStateKind::kNone;
|
||||
} else if (kv_state_kind == "hybrid") {
|
||||
return KVStateKind::kHybrid;
|
||||
} else {
|
||||
LOG(FATAL) << "Invalid kv state kind string: " << kv_state_kind;
|
||||
}
|
||||
}
|
||||
struct ModelMetadata {
|
||||
struct Param {
|
||||
struct Preproc {
|
||||
String func_name;
|
||||
Shape in_shape;
|
||||
Shape out_shape;
|
||||
DLDataType out_dtype;
|
||||
static Preproc FromJSON(const tvm::ffi::json::Object& js,
|
||||
const tvm::ffi::json::Object& model_config);
|
||||
};
|
||||
|
||||
String name;
|
||||
Shape shape;
|
||||
DLDataType dtype;
|
||||
std::vector<Preproc> preprocs;
|
||||
std::vector<int> pipeline_stages;
|
||||
static Param FromJSON(const tvm::ffi::json::Object& param_obj,
|
||||
const tvm::ffi::json::Object& model_config);
|
||||
};
|
||||
|
||||
struct KVCacheMetadata {
|
||||
int64_t num_hidden_layers;
|
||||
int64_t num_attention_heads;
|
||||
int64_t num_key_value_heads;
|
||||
int64_t head_dim;
|
||||
static KVCacheMetadata FromJSON(const tvm::ffi::json::Object& json);
|
||||
};
|
||||
|
||||
std::string model_type;
|
||||
std::string quantization;
|
||||
int64_t context_window_size;
|
||||
int64_t prefill_chunk_size;
|
||||
int64_t max_batch_size;
|
||||
int64_t sliding_window_size;
|
||||
int64_t tensor_parallel_shards;
|
||||
int64_t pipeline_parallel_stages;
|
||||
bool disaggregation;
|
||||
int64_t attention_sink_size;
|
||||
int64_t seqlen_padding_factor;
|
||||
std::vector<Param> params;
|
||||
std::unordered_map<std::string, int64_t> memory_usage;
|
||||
KVStateKind kv_state_kind;
|
||||
KVCacheMetadata kv_cache_metadata;
|
||||
std::string model_task;
|
||||
std::string embedding_model_type;
|
||||
std::string embedding_pooling_strategy;
|
||||
bool embedding_normalize = false;
|
||||
|
||||
static ModelMetadata FromJSON(const tvm::ffi::json::Object& json_str,
|
||||
const tvm::ffi::json::Object& model_config);
|
||||
static ModelMetadata FromModule(Module module, const tvm::ffi::json::Object& model_config);
|
||||
};
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_CPP_MODEL_METADATA_H_
|
||||
@@ -0,0 +1,100 @@
|
||||
/*!
|
||||
* \file builtin.cc
|
||||
* \brief Multi-GPU builtin functions in MLC LLM.
|
||||
*/
|
||||
#ifndef MLC_SINGLE_GPU_ONLY
|
||||
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/optional.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/runtime/disco/builtin.h>
|
||||
#include <tvm/runtime/disco/disco_worker.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
#include <tvm/runtime/vm/vm.h>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace multi_gpu {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Array;
|
||||
using tvm::ffi::ObjectRef;
|
||||
using tvm::ffi::Optional;
|
||||
using tvm::ffi::Shape;
|
||||
|
||||
ObjectRef DispatchFunctionByGroup(tvm::ffi::AnyView vm_arg,
|
||||
Array<Array<ObjectRef>> funcs_and_args) {
|
||||
using namespace vm;
|
||||
VirtualMachine* vm = VirtualMachine::GetContextPtr(vm_arg);
|
||||
DiscoWorker* worker = DiscoWorker::ThreadLocal();
|
||||
int world_size = worker->num_workers;
|
||||
int group_size = worker->num_workers / worker->num_groups;
|
||||
int num_group = world_size / group_size;
|
||||
TVM_FFI_ICHECK_EQ(funcs_and_args.size(), num_group)
|
||||
<< "Number of groups mismatches. There are " << num_group
|
||||
<< " groups while the function/arg array has " << funcs_and_args.size() << " elements.";
|
||||
|
||||
int group_id = worker->worker_id / group_size;
|
||||
TVM_FFI_ICHECK(!funcs_and_args[group_id].empty())
|
||||
<< "No function is provided for group " << group_id;
|
||||
VMClosure func = funcs_and_args[group_id][0].as_or_throw<VMClosure>();
|
||||
|
||||
int num_args = static_cast<int>(funcs_and_args[group_id].size()) - 1;
|
||||
std::vector<tvm::ffi::AnyView> packed_args(num_args);
|
||||
for (int i = 0; i < num_args; ++i) {
|
||||
// NOTE: Need explicily define `arg` so that the argument does not
|
||||
// have type code kTVMObjectRValueRefArg.
|
||||
packed_args[i] = funcs_and_args[group_id][1 + i];
|
||||
}
|
||||
|
||||
tvm::ffi::Any rv;
|
||||
vm->InvokeClosurePacked(funcs_and_args[group_id][0].as_or_throw<VMClosure>(),
|
||||
tvm::ffi::PackedArgs(packed_args.data(), packed_args.size()), &rv);
|
||||
return rv.cast<ObjectRef>();
|
||||
}
|
||||
|
||||
ObjectRef SendFromLastGroupToWorker0(Tensor send, Optional<Tensor> recv, Shape shape,
|
||||
DLDataType dtype) {
|
||||
DiscoWorker* worker = DiscoWorker::ThreadLocal();
|
||||
int worker_id = worker->worker_id;
|
||||
int world_size = worker->num_workers;
|
||||
int group_size = worker->num_workers / worker->num_groups;
|
||||
TVM_FFI_ICHECK_NE(world_size, group_size) << "Cannot perform when there is only one group.";
|
||||
int sender_id = world_size - group_size;
|
||||
if (worker_id == 0) {
|
||||
TVM_FFI_ICHECK(recv.has_value()) << "The receive Tensor is undefined for worker 0.";
|
||||
Tensor recv_arr = recv.value().CreateView(shape, dtype);
|
||||
RecvFromWorker(recv_arr, sender_id);
|
||||
return recv_arr;
|
||||
} else if (worker_id == sender_id) {
|
||||
TVM_FFI_ICHECK_EQ(send->dtype, dtype)
|
||||
<< "The src Tensor has mismatched dtype than the expected dtype.";
|
||||
TVM_FFI_ICHECK_EQ(send->ndim, shape.size())
|
||||
<< "The src Tensor has mismatched shape than the expected shape.";
|
||||
for (int i = 0; i < send->ndim; ++i) {
|
||||
TVM_FFI_ICHECK_EQ(send->shape[i], shape[i])
|
||||
<< "The src Tensor has mismatched shape than the expected shape.";
|
||||
}
|
||||
SendToWorker(send, /*receiver_id=*/0);
|
||||
return recv.value_or(Tensor(nullptr));
|
||||
}
|
||||
|
||||
// We only process for worker 0 and the first worker of the last group.
|
||||
// For other workers, we return the input object.
|
||||
return recv.value_or(Tensor(nullptr));
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef()
|
||||
.def("mlc.multi_gpu.DispatchFunctionByGroup", DispatchFunctionByGroup)
|
||||
.def("mlc.multi_gpu.SendFromLastGroupToWorker0", SendFromLastGroupToWorker0);
|
||||
}
|
||||
|
||||
} // namespace multi_gpu
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_SINGLE_GPU_ONLY
|
||||
@@ -0,0 +1,327 @@
|
||||
/*!
|
||||
* \file multi_gpu_loader.cc
|
||||
* \brief Implementation of a multi-GPU loader with loading-time sharding.
|
||||
*/
|
||||
#ifndef MLC_SINGLE_GPU_ONLY
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/optional.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/runtime/device_api.h>
|
||||
#include <tvm/runtime/disco/builtin.h>
|
||||
#include <tvm/runtime/disco/disco_worker.h>
|
||||
#include <tvm/runtime/vm/tensor_cache_support.h>
|
||||
|
||||
#include <chrono>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <functional>
|
||||
#include <numeric>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "../metadata/model.h"
|
||||
#include "../support/progress_bar.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace multi_gpu {
|
||||
|
||||
using tvm::Device;
|
||||
using tvm::runtime::vm::TensorCacheMetadata;
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Array;
|
||||
using tvm::ffi::Function;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::Optional;
|
||||
using tvm::ffi::Shape;
|
||||
using tvm::ffi::TypedFunction;
|
||||
using DurationType = std::chrono::microseconds;
|
||||
|
||||
class RangeTimer {
|
||||
public:
|
||||
explicit RangeTimer(DurationType* result)
|
||||
: start(std::chrono::high_resolution_clock::now()), result(result) {}
|
||||
|
||||
~RangeTimer() {
|
||||
std::chrono::time_point<std::chrono::high_resolution_clock> end =
|
||||
std::chrono::high_resolution_clock::now(); //
|
||||
auto duration = end - start;
|
||||
(*result) += std::chrono::duration_cast<DurationType>(end - start);
|
||||
}
|
||||
|
||||
private:
|
||||
std::chrono::time_point<std::chrono::high_resolution_clock> start;
|
||||
DurationType* result;
|
||||
};
|
||||
|
||||
class PreprocessorPool {
|
||||
public:
|
||||
explicit PreprocessorPool(const ModelMetadata& model_metadata, Module vm_module) {
|
||||
for (const ModelMetadata::Param& param : model_metadata.params) {
|
||||
for (const ModelMetadata::Param::Preproc& preproc : param.preprocs) {
|
||||
const std::string& func_name = preproc.func_name;
|
||||
if (Function f = vm_module.defined()
|
||||
? vm_module->GetFunction(func_name, true).value_or(Function(nullptr))
|
||||
: nullptr;
|
||||
f != nullptr) {
|
||||
preproc_funcs[func_name] = f;
|
||||
} else if (const auto f = Function::GetGlobal(func_name); f.has_value()) {
|
||||
preproc_funcs[func_name] = *f;
|
||||
} else {
|
||||
LOG(FATAL) << "ValueError: Undefined function: " << func_name;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Tensor Apply(Tensor param, const ModelMetadata::Param& param_info) const {
|
||||
for (const ModelMetadata::Param::Preproc& preproc : param_info.preprocs) {
|
||||
const std::string& func_name = preproc.func_name;
|
||||
Tensor param_in = param;
|
||||
param = Tensor::Empty(preproc.out_shape, preproc.out_dtype, param->device);
|
||||
TVM_FFI_ICHECK(preproc_funcs.count(func_name));
|
||||
DLTensor dl_param_in = *param_in.operator->();
|
||||
DLTensor dl_param = *param.operator->();
|
||||
preproc_funcs.at(func_name)(&dl_param_in, &dl_param);
|
||||
}
|
||||
return param;
|
||||
}
|
||||
|
||||
private:
|
||||
std::unordered_map<std::string, TypedFunction<void(DLTensor*, DLTensor*)>> preproc_funcs;
|
||||
};
|
||||
|
||||
struct ParamInfo {
|
||||
const TensorCacheMetadata::FileRecord* file;
|
||||
const TensorCacheMetadata::FileRecord::ParamRecord* param;
|
||||
};
|
||||
|
||||
Tensor RecvFromGlobalWorker0(Device device, const ModelMetadata::Param& param_info) {
|
||||
Shape shape = param_info.preprocs.empty() ? param_info.shape : param_info.preprocs[0].in_shape;
|
||||
Tensor result = Tensor::Empty(shape, param_info.dtype, device);
|
||||
RecvFromWorker0(result);
|
||||
return result;
|
||||
}
|
||||
|
||||
Tensor BroadcastOrShardAndScatter(Tensor param, const ModelMetadata::Param& param_info,
|
||||
int num_shards, const PreprocessorPool& preprocs) {
|
||||
bool needs_sharding = !param_info.preprocs.empty();
|
||||
if (!needs_sharding) {
|
||||
BroadcastFromWorker0(param, /*in_group=*/true, param);
|
||||
return param;
|
||||
}
|
||||
Device device = param->device;
|
||||
Shape shape = param_info.preprocs.back().out_shape;
|
||||
DLDataType dtype = param_info.preprocs.back().out_dtype;
|
||||
TVM_FFI_ICHECK(shape.size() >= 1 && shape[0] == num_shards)
|
||||
<< "ValueError: The first dimension of the output shape must be equal to the "
|
||||
<< "number of shards, but got: " << shape << " and num_shards = " << num_shards;
|
||||
param = preprocs.Apply(param, param_info);
|
||||
Tensor result = Tensor::Empty(Shape(shape.begin() + 1, shape.end()), dtype, device);
|
||||
ScatterFromWorker0(param, /*in_group=*/true, result);
|
||||
return result;
|
||||
}
|
||||
|
||||
Tensor ReceiveBroadcastedOrSharded(Device device, const ModelMetadata::Param& param_info,
|
||||
int num_shards) {
|
||||
bool needs_sharding = !param_info.preprocs.empty();
|
||||
Tensor result;
|
||||
if (needs_sharding) {
|
||||
Shape shape = param_info.preprocs.back().out_shape;
|
||||
DLDataType dtype = param_info.preprocs.back().out_dtype;
|
||||
result = Tensor::Empty(Shape(shape.begin() + 1, shape.end()), dtype, device);
|
||||
ScatterFromWorker0(std::nullopt, /*in_group=*/true, result);
|
||||
} else {
|
||||
result = Tensor::Empty(param_info.shape, param_info.dtype, device);
|
||||
BroadcastFromWorker0(result, /*in_group=*/true, result);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
std::string FormatDuration(DurationType duration) {
|
||||
std::ostringstream os;
|
||||
auto float_seconds = std::chrono::duration_cast<std::chrono::duration<float>>(duration).count();
|
||||
os << std::fixed << std::setprecision(3) << float_seconds << " s";
|
||||
return os.str();
|
||||
}
|
||||
|
||||
Array<Optional<Tensor>> LoadMultiGPU(const std::string& model_path, Module vm_module,
|
||||
const std::string& model_config_str) {
|
||||
DiscoWorker* worker = DiscoWorker::ThreadLocal();
|
||||
Device device = worker->default_device;
|
||||
int worker_id = worker->worker_id;
|
||||
int group_size = worker->num_workers / worker->num_groups;
|
||||
int num_shards = group_size;
|
||||
int group_id = worker_id / group_size;
|
||||
LOG(INFO) << "[Worker #" << worker_id << "] Loading model to device: " << device;
|
||||
// Step 0. Initialize metadata and paths
|
||||
TensorCacheMetadata tensor_cache_metadata = TensorCacheMetadata::Load(model_path);
|
||||
tvm::ffi::json::Value model_config = tvm::ffi::json::Parse(model_config_str);
|
||||
ModelMetadata model_metadata =
|
||||
ModelMetadata::FromModule(vm_module, model_config.cast<tvm::ffi::json::Object>());
|
||||
TVM_FFI_ICHECK_EQ(model_metadata.tensor_parallel_shards, num_shards)
|
||||
<< "ValueError: The model is compiled using `--tensor-parallel-shards="
|
||||
<< model_metadata.tensor_parallel_shards
|
||||
<< "`, but mlc-chat-config.json is configured to use " << num_shards << " GPUs. "
|
||||
<< "Please set \"tensor_parallel_shards\" in mlc-chat-config.json to "
|
||||
<< model_metadata.tensor_parallel_shards;
|
||||
// Step 1. Extract auxiliary information
|
||||
PreprocessorPool preprocs(model_metadata, vm_module);
|
||||
std::unordered_map<std::string, ModelMetadata::Param> param_name2info;
|
||||
for (const ModelMetadata::Param& param : model_metadata.params) {
|
||||
param_name2info[param.name] = param;
|
||||
}
|
||||
// Step 2. Load, preprocess and shard all the parameters
|
||||
std::unordered_map<std::string, Tensor> sharded_params;
|
||||
if (worker_id == 0) {
|
||||
DurationType time_loading(0);
|
||||
DurationType time_preproc(0);
|
||||
ProgressBar progress_bar(model_metadata.params.size());
|
||||
LOG(INFO) << "Loading parameters...";
|
||||
for (const TensorCacheMetadata::FileRecord& record : tensor_cache_metadata.records) {
|
||||
Array<Tensor> loaded_params;
|
||||
{
|
||||
RangeTimer _(&time_loading);
|
||||
std::string raw_data_buffer;
|
||||
loaded_params = record.Load(device, model_path, &raw_data_buffer);
|
||||
DeviceAPI::Get(device)->StreamSync(device, nullptr);
|
||||
}
|
||||
// For each parameter in the shard file, preprocess and shard it
|
||||
for (size_t i = 0; i < record.records.size(); ++i, progress_bar.Progress()) {
|
||||
RangeTimer _(&time_preproc);
|
||||
const std::string& param_name = record.records[i].name;
|
||||
const ModelMetadata::Param& param_info = param_name2info.at(param_name);
|
||||
for (int group_id : param_info.pipeline_stages) {
|
||||
if (group_id == 0) {
|
||||
// Broadcast or shard-scatter this parameter to all workers in worker group 0.
|
||||
sharded_params[param_name] =
|
||||
BroadcastOrShardAndScatter(loaded_params[i], param_info, num_shards, preprocs);
|
||||
} else {
|
||||
// Send this parameter to the first worker of the worker group of "group_id",
|
||||
// and let that first worker to process this parameter.
|
||||
SendToWorker(loaded_params[i], /*receiver_id=*/group_id * group_size);
|
||||
}
|
||||
}
|
||||
DeviceAPI::Get(device)->StreamSync(device, nullptr);
|
||||
}
|
||||
}
|
||||
LOG(INFO) << "Loading done. Time used:" << std::fixed << std::setprecision(3) //
|
||||
<< " Loading " << FormatDuration(time_loading) << " Preprocessing "
|
||||
<< FormatDuration(time_preproc) << ".";
|
||||
} else {
|
||||
for (const TensorCacheMetadata::FileRecord& record : tensor_cache_metadata.records) {
|
||||
for (size_t i = 0; i < record.records.size(); ++i) {
|
||||
const std::string& param_name = record.records[i].name;
|
||||
const ModelMetadata::Param& param_info = param_name2info.at(param_name);
|
||||
if (std::find(param_info.pipeline_stages.begin(), param_info.pipeline_stages.end(),
|
||||
group_id) == param_info.pipeline_stages.end()) {
|
||||
// This worker group doesn't need to hold a copy of this parameter.
|
||||
continue;
|
||||
}
|
||||
|
||||
if (worker_id % group_size == 0) {
|
||||
// The worker is the first worker of its worker group (while not the first worker group).
|
||||
// Receive the full parameter from the global worker 0.
|
||||
Tensor full_param = RecvFromGlobalWorker0(device, param_info);
|
||||
// Broadcast or shard-scatter this parameter to all workers in its worker group.
|
||||
sharded_params[param_name] =
|
||||
BroadcastOrShardAndScatter(full_param, param_info, num_shards, preprocs);
|
||||
} else {
|
||||
// The worker is not the first worker of its worker group.
|
||||
// Receive from the first worker in the its worker group.
|
||||
sharded_params[param_name] = ReceiveBroadcastedOrSharded(device, param_info, num_shards);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Step 3. Reorder the sharded parameters according to the order in model_metadata
|
||||
Array<Optional<Tensor>> shards;
|
||||
shards.reserve(model_metadata.params.size());
|
||||
for (const ModelMetadata::Param& param : model_metadata.params) {
|
||||
const auto& it = sharded_params.find(param.name);
|
||||
shards.push_back(it == sharded_params.end() ? Optional<Tensor>() : it->second);
|
||||
}
|
||||
return shards;
|
||||
}
|
||||
|
||||
Array<Optional<Tensor>> LoadMultiGPUPresharded(const std::string& model_path, Module vm_module,
|
||||
const std::string& model_config_str) {
|
||||
DiscoWorker* worker = DiscoWorker::ThreadLocal();
|
||||
Device device = worker->default_device;
|
||||
int worker_id = worker->worker_id;
|
||||
int group_size = worker->num_workers / worker->num_groups;
|
||||
int num_shards = group_size;
|
||||
int group_id = worker_id / group_size;
|
||||
int local_worker_id = worker_id % group_size;
|
||||
LOG(INFO) << "[Worker #" << worker_id << "] Loading model to device: " << device;
|
||||
// Step 0. Initialize metadata and paths
|
||||
TensorCacheMetadata tensor_cache_metadata = TensorCacheMetadata::Load(model_path);
|
||||
tvm::ffi::json::Value model_config = tvm::ffi::json::Parse(model_config_str);
|
||||
ModelMetadata model_metadata =
|
||||
ModelMetadata::FromModule(vm_module, model_config.cast<tvm::ffi::json::Object>());
|
||||
|
||||
std::unordered_map<std::string, ParamInfo> param_info_map;
|
||||
for (const TensorCacheMetadata::FileRecord& file_record : tensor_cache_metadata.records) {
|
||||
for (const TensorCacheMetadata::FileRecord::ParamRecord& param_record : file_record.records) {
|
||||
const std::string& param_name = param_record.name;
|
||||
param_info_map[param_name] = ParamInfo{&file_record, ¶m_record};
|
||||
}
|
||||
}
|
||||
|
||||
Array<Optional<Tensor>> params;
|
||||
const TensorCacheMetadata::FileRecord* current_file_;
|
||||
std::string current_file_stream_;
|
||||
params.reserve(model_metadata.params.size());
|
||||
DurationType time_loading(0);
|
||||
for (const ModelMetadata::Param& param : model_metadata.params) {
|
||||
RangeTimer _(&time_loading);
|
||||
if (std::find(param.pipeline_stages.begin(), param.pipeline_stages.end(), group_id) ==
|
||||
param.pipeline_stages.end()) {
|
||||
// This worker group doesn't need to hold a copy of this parameter.
|
||||
params.push_back(Optional<Tensor>());
|
||||
continue;
|
||||
}
|
||||
bool needs_sharding = !param.preprocs.empty();
|
||||
std::string param_name =
|
||||
needs_sharding ? static_cast<const std::stringstream&>(
|
||||
std::stringstream() << param.name << "_shard-" << local_worker_id)
|
||||
.str()
|
||||
: std::string(param.name);
|
||||
auto it = param_info_map.find(param_name);
|
||||
TVM_FFI_ICHECK(it != param_info_map.end())
|
||||
<< "ValueError: Cannot find parameter: " << param_name;
|
||||
const ParamInfo& param_info = (*it).second;
|
||||
const TensorCacheMetadata::FileRecord::ParamRecord* param_record = param_info.param;
|
||||
const TensorCacheMetadata::FileRecord* file_record = param_info.file;
|
||||
|
||||
if (file_record != current_file_) {
|
||||
current_file_ = file_record;
|
||||
file_record->Load(device, model_path, ¤t_file_stream_);
|
||||
}
|
||||
|
||||
params.push_back(param_record->Load(device, ¤t_file_stream_));
|
||||
}
|
||||
SyncWorker();
|
||||
if (worker_id == 0) {
|
||||
LOG(INFO) << "Loading done. Time used: " << FormatDuration(time_loading) << ".";
|
||||
}
|
||||
return params;
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef()
|
||||
.def("mlc.multi_gpu.LoadMultiGPU", LoadMultiGPU)
|
||||
.def("mlc.multi_gpu.LoadMultiGPUPresharded", LoadMultiGPUPresharded);
|
||||
}
|
||||
|
||||
} // namespace multi_gpu
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_SINGLE_GPU_ONLY
|
||||
+1082
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,470 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/config.h
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_CONFIG_H_
|
||||
#define MLC_LLM_SERVE_CONFIG_H_
|
||||
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
#include <tvm/runtime/device_api.h>
|
||||
|
||||
#include <optional>
|
||||
|
||||
#include "../metadata/model.h"
|
||||
#include "../support/result.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm;
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Array;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectPtr;
|
||||
using tvm::ffi::ObjectRef;
|
||||
using tvm::ffi::Optional;
|
||||
using tvm::ffi::Shape;
|
||||
using tvm::ffi::String;
|
||||
|
||||
/****************** GenerationConfig ******************/
|
||||
|
||||
/*! \brief The response format of a request. */
|
||||
struct ResponseFormat {
|
||||
String type = "text";
|
||||
Optional<String> schema = std::nullopt;
|
||||
/*!
|
||||
* \brief Create debug config from JSON.
|
||||
* \param config_json The json string for generation config
|
||||
* \returns The converted result.
|
||||
*/
|
||||
static Result<ResponseFormat> FromJSON(const tvm::ffi::json::Object& config_json);
|
||||
|
||||
/**
|
||||
* \return serialized json value of the config.
|
||||
*/
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
enum class SpecialRequestKind : int {
|
||||
kNone = 0,
|
||||
kQueryEngineMetrics = 1,
|
||||
};
|
||||
|
||||
enum class DisaggRequestKind : int {
|
||||
kNone = 0,
|
||||
kPrepareReceive = 1,
|
||||
kRemoteSend = 2,
|
||||
kStartGeneration = 3,
|
||||
};
|
||||
|
||||
/*! \brief Controls the behavior of inference with grammar constraint. */
|
||||
enum class GrammarExecutionMode : int {
|
||||
/*! \brief If grammar is provided for a request, use the grammar to constrain the output token. */
|
||||
kConstraint = 0,
|
||||
/*! \brief If grammar is provided for a request, not only constrain the output, but also use the
|
||||
* jump-forward decoding to predict the next tokens. This is the default option. */
|
||||
kJumpForward = 1,
|
||||
};
|
||||
|
||||
/*! \brief The config for disaggregation requests. */
|
||||
class DisaggConfig {
|
||||
public:
|
||||
DisaggRequestKind kind = DisaggRequestKind::kNone;
|
||||
std::vector<Shape> kv_append_metadata;
|
||||
// "kv_window_begin" and "kv_window_end" denote the KV interval of interests.
|
||||
// "kv_window_end" supports Python style negative indexing.
|
||||
// The concrete meaning varies for different special request kind:
|
||||
// - For "prepare_receive", the begin is always 0, and "[0:end]" denotes
|
||||
// the KV range to prefill on a prefill instance.
|
||||
// - For "remote_send", "[begin:end]" means the KV range to compute prefill
|
||||
// and send to the decode instance.
|
||||
// - For "start_generation", the end is always nullopt, and "[begin:]" denotes
|
||||
// the KV range to prefill locally on the decode instance.
|
||||
std::optional<int> kv_window_begin = std::nullopt;
|
||||
std::optional<int> kv_window_end = std::nullopt;
|
||||
std::optional<int> dst_group_offset = std::nullopt;
|
||||
|
||||
static Result<DisaggConfig> FromJSON(const tvm::ffi::json::Object& config_json);
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
/*! \brief The debug configuration of a request. */
|
||||
class DebugConfig {
|
||||
public:
|
||||
bool ignore_eos = false;
|
||||
bool pinned_system_prompt = false;
|
||||
SpecialRequestKind special_request = SpecialRequestKind::kNone;
|
||||
/*! \brief The grammar execution mode. */
|
||||
GrammarExecutionMode grammar_execution_mode = GrammarExecutionMode::kJumpForward;
|
||||
DisaggConfig disagg_config;
|
||||
|
||||
/*!
|
||||
* \brief Create debug config from JSON.
|
||||
* \param config_json The json string for generation config
|
||||
* \returns The converted result.
|
||||
*/
|
||||
static Result<DebugConfig> FromJSON(const tvm::ffi::json::Object& config_json);
|
||||
|
||||
/**
|
||||
* \return serialized json value of the config.
|
||||
*/
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
/*! \brief The generation configuration of a request. */
|
||||
class GenerationConfigNode : public Object {
|
||||
public:
|
||||
int n = 1;
|
||||
double temperature = 1.0;
|
||||
double top_p = 1.0;
|
||||
double frequency_penalty = 0.0;
|
||||
double presence_penalty = 0.0;
|
||||
double repetition_penalty = 1.0;
|
||||
bool logprobs = false;
|
||||
int top_logprobs = 0;
|
||||
std::vector<std::pair<int, float>> logit_bias;
|
||||
int seed;
|
||||
// -1 means infinite
|
||||
int max_tokens = -1;
|
||||
Array<String> stop_strs;
|
||||
std::vector<int> stop_token_ids;
|
||||
|
||||
ResponseFormat response_format;
|
||||
DebugConfig debug_config;
|
||||
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<GenerationConfigNode>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.GenerationConfig", GenerationConfigNode, Object);
|
||||
};
|
||||
|
||||
class GenerationConfig : public ObjectRef {
|
||||
public:
|
||||
/*!
|
||||
* \brief Run validation of generation config and ensure values are in bound.
|
||||
* \return The validtaed Generation config or error.
|
||||
*/
|
||||
static Result<GenerationConfig> Validate(GenerationConfig cfg);
|
||||
|
||||
/*!
|
||||
* \brief Create generation config from JSON.
|
||||
* \param config_json The json string for generation config
|
||||
* \param default_config The default config
|
||||
*/
|
||||
static Result<GenerationConfig> FromJSON(const tvm::ffi::json::Object& config_json,
|
||||
const GenerationConfig& default_config);
|
||||
|
||||
/*! \brief Get the default generation config from the model config. */
|
||||
static GenerationConfig GetDefaultFromModelConfig(const tvm::ffi::json::Object& json);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GenerationConfig, ObjectRef, GenerationConfigNode);
|
||||
};
|
||||
|
||||
/****************** Engine config ******************/
|
||||
|
||||
/*!
|
||||
* \brief The engine mode in MLC LLM.
|
||||
* We provide three preset modes: "local", "interactive" and "server".
|
||||
* The default mode is "local".
|
||||
* The choice of mode decides the values of "max_batch_size", "max_total_sequence_length"
|
||||
* and "prefill_chunk_size" when they are not explicitly specified.
|
||||
* 1. Mode "local" refers to the local server deployment which has low
|
||||
* request concurrency. So the max batch size will be set to 4, and max
|
||||
* total sequence length and prefill chunk size are set to the context
|
||||
* window size (or sliding window size) of the model.
|
||||
* 2. Mode "interactive" refers to the interactive use of server, which
|
||||
* has at most 1 concurrent request. So the max batch size will be set to 1,
|
||||
* and max total sequence length and prefill chunk size are set to the context
|
||||
* window size (or sliding window size) of the model.
|
||||
* 3. Mode "server" refers to the large server use case which may handle
|
||||
* many concurrent request and want to use GPU memory as much as possible.
|
||||
* In this mode, we will automatically infer the largest possible max batch
|
||||
* size and max total sequence length.
|
||||
*/
|
||||
enum class EngineMode : int {
|
||||
kLocal = 0,
|
||||
kInteractive = 1,
|
||||
kServer = 2,
|
||||
};
|
||||
|
||||
/*! \brief The prefix cache mode. */
|
||||
enum class PrefixCacheMode : int {
|
||||
/*! \brief Disable prefix cache. */
|
||||
kDisable = 0,
|
||||
/*! \brief The paged radix tree based prefix cache mode. */
|
||||
kRadix = 1,
|
||||
};
|
||||
|
||||
/*! \brief The speculative mode. */
|
||||
enum class SpeculativeMode : int {
|
||||
/*! \brief Disable speculative decoding. */
|
||||
kDisable = 0,
|
||||
/*! \brief The normal speculative decoding (small draft) mode. */
|
||||
kSmallDraft = 1,
|
||||
/*! \brief The eagle-style speculative decoding. */
|
||||
kEagle = 2,
|
||||
/*! \brief The Medusa-style speculative decoding. */
|
||||
kMedusa = 3,
|
||||
};
|
||||
|
||||
/*! \brief The prefill mode. */
|
||||
enum class PrefillMode : int {
|
||||
/*! \brief Only chunked prefill is enabled. */
|
||||
kChunked = 0,
|
||||
/*!
|
||||
* \brief The hybrid prefill or split-fuse prefill is enabled, some decode steps will be fused
|
||||
* to prefill
|
||||
*/
|
||||
kHybrid = 1,
|
||||
};
|
||||
|
||||
class InferrableEngineConfig;
|
||||
|
||||
/*! \brief The configuration of engine execution config. */
|
||||
class EngineConfigNode : public Object {
|
||||
public:
|
||||
/*************** Models ***************/
|
||||
|
||||
/*! \brief The path to the model directory. */
|
||||
String model;
|
||||
/*! \brief The path or identifier to the model library. */
|
||||
String model_lib;
|
||||
/*! \brief The path to the additional models' directories. */
|
||||
Array<String> additional_models;
|
||||
/*! \brief The path to the additional models' libraries. */
|
||||
Array<String> additional_model_libs;
|
||||
|
||||
/*************** KV cache config and engine capacities ***************/
|
||||
|
||||
/*!
|
||||
* \brief The engine mode in MLC LLM.
|
||||
* \sa EngineMode
|
||||
*/
|
||||
EngineMode mode = EngineMode::kLocal;
|
||||
/*!
|
||||
* \brief A number in (0, 1) denoting the fraction of GPU memory used by the server in total.
|
||||
* It is used to infer to maximum possible KV cache capacity.
|
||||
* When it is unspecified, it defaults to 0.85.
|
||||
* Under mode "local" or "interactive", the actual memory usage may be
|
||||
* significantly smaller than this number. Under mode "server", the actual
|
||||
* memory usage may be slightly larger than this number.
|
||||
*/
|
||||
float gpu_memory_utilization = 0.85;
|
||||
/*! \brief The number of consecutive tokens handled in each page in paged KV cache. */
|
||||
int kv_cache_page_size = 16;
|
||||
/*!
|
||||
* \brief The maximum number of sequences that are allowed to be
|
||||
* processed by the KV cache at any time.
|
||||
*/
|
||||
int max_num_sequence = 4;
|
||||
/*! \brief The maximum length allowed for a single sequence in the engine. */
|
||||
int64_t max_total_sequence_length = 4096;
|
||||
/*!
|
||||
* \brief The maximum total number of tokens whose KV data are allowed
|
||||
* to exist in the KV cache at any time.
|
||||
*/
|
||||
int64_t max_single_sequence_length = 4096;
|
||||
/*! \brief The maximum total sequence length in a prefill. */
|
||||
int64_t prefill_chunk_size = 1024;
|
||||
/*! \brief The maximum history size for RNN state. KV cache does not need this. */
|
||||
int max_history_size = 0;
|
||||
|
||||
/*************** Prefix cache ***************/
|
||||
|
||||
/*! \brief The prefix cache mode. */
|
||||
PrefixCacheMode prefix_cache_mode = PrefixCacheMode::kRadix;
|
||||
/*! \brief The maximum number of recycling sequences in prefix cache, default as max_num_sequence.
|
||||
* And set 0 to disable prefix cache, set -1 to have infinite capacity prefix cache. */
|
||||
int prefix_cache_max_num_recycling_seqs = -1;
|
||||
|
||||
/*************** Speculative decoding ***************/
|
||||
|
||||
/*! \brief The speculative mode. */
|
||||
SpeculativeMode speculative_mode = SpeculativeMode::kDisable;
|
||||
/*!
|
||||
* \brief The number of tokens to generate in speculative proposal (draft).
|
||||
* Being 0 means to enable adaptive speculative mode, where the draft length
|
||||
* will be automatically adjusted based on engine state.
|
||||
*/
|
||||
int spec_draft_length = 0;
|
||||
/*! \brief The number of tokens to generate in speculative tree decoding */
|
||||
int spec_tree_width = 1;
|
||||
|
||||
/*************** Prefill mode ***************/
|
||||
|
||||
/*! \brief The prefill mode. */
|
||||
PrefillMode prefill_mode = PrefillMode::kHybrid;
|
||||
|
||||
/*************** Debug ***************/
|
||||
bool verbose = false;
|
||||
|
||||
String AsJSONString() const;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<EngineConfigNode>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.EngineConfig", EngineConfigNode, Object);
|
||||
};
|
||||
|
||||
class EngineConfig : public ObjectRef {
|
||||
public:
|
||||
/*! \brief Create EngineConfig from JSON object and inferred config. */
|
||||
static EngineConfig FromJSONAndInferredConfig(const tvm::ffi::json::Object& json,
|
||||
const InferrableEngineConfig& inferred_config);
|
||||
|
||||
/*!
|
||||
* \brief Get all the models and model libs from the JSON string for engine initialization.
|
||||
* \return The parsed models/model libs from config or error message.
|
||||
*/
|
||||
static Result<std::vector<std::pair<std::string, std::string>>>
|
||||
GetModelsAndModelLibsFromJSONString(const std::string& json_str);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(EngineConfig, ObjectRef, EngineConfigNode);
|
||||
};
|
||||
|
||||
/*! \brief A subset of engine config that is inferrable. */
|
||||
struct InferrableEngineConfig {
|
||||
std::optional<int64_t> max_num_sequence;
|
||||
std::optional<int64_t> max_total_sequence_length;
|
||||
std::optional<int64_t> max_single_sequence_length;
|
||||
std::optional<int64_t> prefill_chunk_size;
|
||||
std::optional<int64_t> max_history_size;
|
||||
|
||||
/*! \brief Infer the config for KV cache from a given initial config. */
|
||||
static Result<InferrableEngineConfig> InferForKVCache(
|
||||
EngineMode mode, Device device, double gpu_memory_utilization,
|
||||
const std::vector<tvm::ffi::json::Object>& model_configs,
|
||||
const std::vector<ModelMetadata>& model_metadata, InferrableEngineConfig init_config,
|
||||
bool verbose);
|
||||
/*! \brief Infer the config for RNN state from a given initial config. */
|
||||
static Result<InferrableEngineConfig> InferForRNNState(
|
||||
EngineMode mode, Device device, double gpu_memory_utilization,
|
||||
const std::vector<tvm::ffi::json::Object>& model_configs,
|
||||
const std::vector<ModelMetadata>& model_metadata, InferrableEngineConfig init_config,
|
||||
bool verbose);
|
||||
};
|
||||
|
||||
/****************** Config utils ******************/
|
||||
|
||||
/*! \brief Check if the models use KV cache or RNN state. */
|
||||
Result<bool> ModelsUseKVCache(const std::vector<tvm::ffi::json::Object>& model_configs);
|
||||
|
||||
inline std::string EngineModeToString(EngineMode mode) {
|
||||
if (mode == EngineMode::kLocal) {
|
||||
return "local";
|
||||
} else if (mode == EngineMode::kInteractive) {
|
||||
return "interactive";
|
||||
} else if (mode == EngineMode::kServer) {
|
||||
return "server";
|
||||
} else {
|
||||
LOG(FATAL) << "Invalid engine mode: " << static_cast<int>(mode);
|
||||
throw;
|
||||
}
|
||||
}
|
||||
|
||||
inline EngineMode EngineModeFromString(const std::string& mode) {
|
||||
if (mode == "local") {
|
||||
return EngineMode::kLocal;
|
||||
} else if (mode == "interactive") {
|
||||
return EngineMode::kInteractive;
|
||||
} else if (mode == "server") {
|
||||
return EngineMode::kServer;
|
||||
} else {
|
||||
LOG(FATAL) << "Invalid engine mode string: " << mode;
|
||||
throw;
|
||||
}
|
||||
}
|
||||
|
||||
inline std::string PrefixCacheModeToString(PrefixCacheMode prefix_cache_mode) {
|
||||
if (prefix_cache_mode == PrefixCacheMode::kDisable) {
|
||||
return "disable";
|
||||
} else if (prefix_cache_mode == PrefixCacheMode::kRadix) {
|
||||
return "radix";
|
||||
} else {
|
||||
LOG(FATAL) << "Invalid prefix cache mode: " << static_cast<int>(prefix_cache_mode);
|
||||
}
|
||||
}
|
||||
|
||||
inline PrefixCacheMode PrefixCacheModeFromString(const std::string& prefix_cache_mode) {
|
||||
if (prefix_cache_mode == "disable") {
|
||||
return PrefixCacheMode::kDisable;
|
||||
} else if (prefix_cache_mode == "radix") {
|
||||
return PrefixCacheMode::kRadix;
|
||||
} else {
|
||||
LOG(FATAL) << "Invalid prefix cache mode string: " << prefix_cache_mode;
|
||||
throw;
|
||||
}
|
||||
}
|
||||
|
||||
inline std::string SpeculativeModeToString(SpeculativeMode speculative_mode) {
|
||||
if (speculative_mode == SpeculativeMode::kDisable) {
|
||||
return "disable";
|
||||
} else if (speculative_mode == SpeculativeMode::kSmallDraft) {
|
||||
return "small_draft";
|
||||
} else if (speculative_mode == SpeculativeMode::kEagle) {
|
||||
return "eagle";
|
||||
} else if (speculative_mode == SpeculativeMode::kMedusa) {
|
||||
return "medusa";
|
||||
} else {
|
||||
LOG(FATAL) << "Invalid speculative mode: " << static_cast<int>(speculative_mode);
|
||||
}
|
||||
}
|
||||
|
||||
inline SpeculativeMode SpeculativeModeFromString(const std::string& speculative_mode) {
|
||||
if (speculative_mode == "disable") {
|
||||
return SpeculativeMode::kDisable;
|
||||
} else if (speculative_mode == "small_draft") {
|
||||
return SpeculativeMode::kSmallDraft;
|
||||
} else if (speculative_mode == "eagle") {
|
||||
return SpeculativeMode::kEagle;
|
||||
} else if (speculative_mode == "medusa") {
|
||||
return SpeculativeMode::kMedusa;
|
||||
} else {
|
||||
LOG(FATAL) << "Invalid speculative mode string: " << speculative_mode;
|
||||
throw;
|
||||
}
|
||||
}
|
||||
|
||||
inline std::string PrefillModeToString(PrefillMode prefill_mode) {
|
||||
if (prefill_mode == PrefillMode::kChunked) {
|
||||
return "chunked";
|
||||
} else if (prefill_mode == PrefillMode::kHybrid) {
|
||||
return "hybrid";
|
||||
} else {
|
||||
LOG(FATAL) << "Invalid prefill mode: " << static_cast<int>(prefill_mode);
|
||||
}
|
||||
}
|
||||
|
||||
inline PrefillMode PrefillModeFromString(const std::string& prefill_mode) {
|
||||
if (prefill_mode == "chunked") {
|
||||
return PrefillMode::kChunked;
|
||||
} else if (prefill_mode == "hybrid") {
|
||||
return PrefillMode::kHybrid;
|
||||
} else {
|
||||
LOG(FATAL) << "Invalid prefill mode string: " << prefill_mode;
|
||||
throw;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_CONFIG_H_
|
||||
@@ -0,0 +1,266 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/data.cc
|
||||
*/
|
||||
#include "data.h"
|
||||
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
|
||||
#include "model.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
DataNode::RegisterReflection();
|
||||
TextDataNode::RegisterReflection();
|
||||
TokenDataNode::RegisterReflection();
|
||||
ImageDataNode::RegisterReflection();
|
||||
RequestStreamOutputObj::RegisterReflection();
|
||||
}
|
||||
|
||||
/****************** Data ******************/
|
||||
|
||||
std::pair<Array<Data>, Array<Data>> SplitData(const Array<Data>& original_data, int total_length,
|
||||
int split_pos) {
|
||||
TVM_FFI_ICHECK_GE(split_pos, 0);
|
||||
TVM_FFI_ICHECK_GE(total_length, split_pos)
|
||||
<< "Cannot truncate when the current length is already less than the target length";
|
||||
std::vector<Data> lhs(original_data.begin(), original_data.end());
|
||||
std::vector<Data> rhs;
|
||||
while (total_length > split_pos) {
|
||||
TVM_FFI_ICHECK(!lhs.empty());
|
||||
Data last_data = lhs.back();
|
||||
int last_data_length = last_data->GetLength();
|
||||
TVM_FFI_ICHECK_GE(total_length - last_data_length, 0);
|
||||
if (total_length - last_data_length >= split_pos) {
|
||||
// Pop the entire last data.
|
||||
rhs.push_back(lhs.back());
|
||||
lhs.pop_back();
|
||||
total_length -= last_data_length;
|
||||
continue;
|
||||
}
|
||||
// Partially truncate the last data.
|
||||
const auto* token_data = last_data.as<TokenDataNode>();
|
||||
TVM_FFI_ICHECK(token_data != nullptr) << "Only TokenData supports partial truncation.";
|
||||
int length_to_truncate = total_length - split_pos;
|
||||
TVM_FFI_ICHECK_GT(length_to_truncate, 0);
|
||||
TVM_FFI_ICHECK_LT(length_to_truncate, last_data_length);
|
||||
TokenData lhs_token_data(
|
||||
Shape{token_data->token_ids.begin(), token_data->token_ids.end() - length_to_truncate});
|
||||
TokenData rhs_token_data(
|
||||
Shape{token_data->token_ids.end() - length_to_truncate, token_data->token_ids.end()});
|
||||
TVM_FFI_ICHECK_EQ(total_length - last_data_length + lhs_token_data->GetLength(), split_pos);
|
||||
lhs.pop_back();
|
||||
lhs.push_back(lhs_token_data);
|
||||
rhs.push_back(rhs_token_data);
|
||||
std::reverse(rhs.begin(), rhs.end());
|
||||
total_length = split_pos;
|
||||
}
|
||||
return {lhs, rhs};
|
||||
}
|
||||
|
||||
/****************** TextData ******************/
|
||||
|
||||
TextData::TextData(String text) {
|
||||
ObjectPtr<TextDataNode> n = tvm::ffi::make_object<TextDataNode>();
|
||||
n->text = std::move(text);
|
||||
data_ = std::move(n);
|
||||
}
|
||||
|
||||
int TextDataNode::GetLength() const {
|
||||
LOG(FATAL) << "\"GetLength\" for TextData is not supported. "
|
||||
"Please tokenize the text and construct a TokenData object.";
|
||||
}
|
||||
|
||||
ObjectRef TextDataNode::GetEmbedding(Model model, ObjectRef* dst, int offset) const {
|
||||
LOG(FATAL) << "\"GetEmbedding\" for TextData is not supported. "
|
||||
"Please tokenize the text and construct a TokenData object.";
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef()
|
||||
.def("mlc.serve.TextData", [](String text) { return TextData(std::move(text)); })
|
||||
.def("mlc.serve.TextDataGetTextString", [](TextData data) { return data->text; });
|
||||
}
|
||||
|
||||
/****************** TokenData ******************/
|
||||
|
||||
TokenData::TokenData(Shape token_ids) {
|
||||
ObjectPtr<TokenDataNode> n = tvm::ffi::make_object<TokenDataNode>();
|
||||
n->token_ids = std::move(token_ids);
|
||||
data_ = std::move(n);
|
||||
}
|
||||
|
||||
TokenData::TokenData(std::vector<int32_t> token_ids) {
|
||||
ObjectPtr<TokenDataNode> n = tvm::ffi::make_object<TokenDataNode>();
|
||||
n->token_ids = Shape(token_ids.begin(), token_ids.end());
|
||||
data_ = std::move(n);
|
||||
}
|
||||
|
||||
int TokenDataNode::GetLength() const { return token_ids.size(); }
|
||||
|
||||
ObjectRef TokenDataNode::GetEmbedding(Model model, ObjectRef* dst, int offset) const {
|
||||
return model->TokenEmbed(token_ids, dst, offset);
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef()
|
||||
.def_packed("mlc.serve.TokenData",
|
||||
[](ffi::PackedArgs args, ffi::Any* rv) {
|
||||
std::vector<int32_t> token_ids;
|
||||
token_ids.reserve(args.size());
|
||||
for (int i = 0; i < args.size(); i++) {
|
||||
token_ids.push_back(args[i].cast<int32_t>());
|
||||
}
|
||||
*rv = TokenData(std::move(token_ids));
|
||||
})
|
||||
.def("mlc.serve.TokenDataGetTokenIds", [](TokenData data) { return data->token_ids; });
|
||||
}
|
||||
|
||||
/****************** ImageData ******************/
|
||||
|
||||
ImageData::ImageData(Tensor image, int embed_size) {
|
||||
ObjectPtr<ImageDataNode> n = tvm::ffi::make_object<ImageDataNode>();
|
||||
n->image = std::move(image);
|
||||
n->embed_size = embed_size;
|
||||
data_ = std::move(n);
|
||||
}
|
||||
|
||||
int ImageDataNode::GetLength() const { return embed_size; }
|
||||
|
||||
ObjectRef ImageDataNode::GetEmbedding(Model model, ObjectRef* dst, int offset) const {
|
||||
return model->ImageEmbed(image, dst, offset);
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef()
|
||||
.def("mlc.serve.ImageData",
|
||||
[](Tensor image, int embed_size) { return ImageData(std::move(image), embed_size); })
|
||||
.def("mlc.serve.ImageDataGetImage", [](ImageData data) { return data->image; });
|
||||
}
|
||||
|
||||
/****************** SampleResult ******************/
|
||||
|
||||
/*! \brief Convert a single token with probability to JSON string. */
|
||||
inline void TokenToLogProbJSON(const Tokenizer& tokenizer, const TokenProbPair& token_prob,
|
||||
std::ostringstream* os) {
|
||||
const std::string& token = tokenizer->PostProcessedTokenTable()[token_prob.first];
|
||||
|
||||
(*os) << "\"token\": \"";
|
||||
for (char ch : token) {
|
||||
if (ch >= 33 && ch <= 126) {
|
||||
// The character is in ASCII visible range.
|
||||
// Handle escape characters in JSON.
|
||||
if (ch == '"') {
|
||||
(*os) << "\\\"";
|
||||
} else if (ch == '\\') {
|
||||
(*os) << "\\\\";
|
||||
} else {
|
||||
(*os) << ch;
|
||||
}
|
||||
}
|
||||
}
|
||||
(*os) << "\", ";
|
||||
(*os) << "\"logprob\": " << std::log(std::max(token_prob.second, 1e-10f)) << ", ";
|
||||
(*os) << "\"bytes\": [";
|
||||
int token_len = token.size();
|
||||
for (int pos = 0; pos < token_len; ++pos) {
|
||||
(*os) << static_cast<int>(static_cast<unsigned char>(token[pos]));
|
||||
if (pos != token_len - 1) {
|
||||
(*os) << ", ";
|
||||
}
|
||||
}
|
||||
(*os) << "]";
|
||||
}
|
||||
|
||||
int32_t SampleResult::GetTokenId() const { return this->sampled_token_id.first; }
|
||||
|
||||
std::string SampleResult::GetLogProbJSON(const Tokenizer& tokenizer, bool logprob) const {
|
||||
TVM_FFI_ICHECK(top_prob_tokens.empty() || logprob);
|
||||
if (!logprob) {
|
||||
// Logprob is not needed.
|
||||
return "";
|
||||
}
|
||||
|
||||
std::ostringstream os;
|
||||
os << "{";
|
||||
// - Convert the sampled token to JSON.
|
||||
TokenToLogProbJSON(tokenizer, sampled_token_id, &os);
|
||||
// - Convert the tokens with top probabilities.
|
||||
os << ", \"top_logprobs\": [";
|
||||
int num_top = top_prob_tokens.size();
|
||||
for (int i = 0; i < num_top; ++i) {
|
||||
os << "{";
|
||||
TokenToLogProbJSON(tokenizer, top_prob_tokens[i], &os);
|
||||
os << "}";
|
||||
if (i != num_top - 1) {
|
||||
os << ", ";
|
||||
}
|
||||
}
|
||||
os << "]}";
|
||||
return os.str();
|
||||
}
|
||||
|
||||
/****************** RequestStreamOutput ******************/
|
||||
|
||||
RequestStreamOutput::RequestStreamOutput(
|
||||
String request_id, std::vector<std::vector<int64_t>> group_delta_token_ids,
|
||||
std::optional<std::vector<std::vector<String>>> group_delta_logprob_json_strs,
|
||||
std::vector<Optional<String>> group_finish_reason,
|
||||
std::vector<String> group_extra_prefix_string) {
|
||||
ObjectPtr<RequestStreamOutputObj> n = tvm::ffi::make_object<RequestStreamOutputObj>();
|
||||
n->request_id = std::move(request_id);
|
||||
n->group_delta_token_ids = std::move(group_delta_token_ids);
|
||||
n->group_delta_logprob_json_strs = std::move(group_delta_logprob_json_strs);
|
||||
n->group_finish_reason = std::move(group_finish_reason);
|
||||
n->group_extra_prefix_string = std::move(group_extra_prefix_string);
|
||||
data_ = std::move(n);
|
||||
}
|
||||
|
||||
RequestStreamOutput RequestStreamOutput::Usage(String request_id,
|
||||
String request_final_usage_json_str) {
|
||||
ObjectPtr<RequestStreamOutputObj> n = tvm::ffi::make_object<RequestStreamOutputObj>();
|
||||
n->request_id = std::move(request_id);
|
||||
n->request_final_usage_json_str = std::move(request_final_usage_json_str);
|
||||
return RequestStreamOutput(n);
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("mlc.serve.RequestStreamOutputUnpack", [](RequestStreamOutput output) {
|
||||
TVM_FFI_ICHECK(!output->unpacked)
|
||||
<< "One RequestStreamOutput can be unpacked for at most once.";
|
||||
std::vector<Shape> group_delta_token_ids;
|
||||
std::vector<Array<String>> group_delta_logprob_json_strs;
|
||||
group_delta_token_ids.reserve(output->group_delta_token_ids.size());
|
||||
if (output->group_delta_logprob_json_strs.has_value()) {
|
||||
group_delta_logprob_json_strs.reserve(output->group_delta_token_ids.size());
|
||||
}
|
||||
for (int i = 0; i < static_cast<int>(output->group_delta_token_ids.size()); ++i) {
|
||||
group_delta_token_ids.push_back(output->group_delta_token_ids[i]);
|
||||
if (output->group_delta_logprob_json_strs.has_value()) {
|
||||
group_delta_logprob_json_strs.push_back(output->group_delta_logprob_json_strs.value()[i]);
|
||||
}
|
||||
}
|
||||
Array<Any> ret = {output->request_id,
|
||||
Array<Shape>(std::move(group_delta_token_ids)),
|
||||
output->group_delta_logprob_json_strs.has_value()
|
||||
? Array<Array<String>>(std::move(group_delta_logprob_json_strs))
|
||||
: Optional<Array<Array<String>>>(),
|
||||
Array<Optional<String>>(output->group_finish_reason),
|
||||
output->request_final_usage_json_str,
|
||||
Array<String>(output->group_extra_prefix_string)};
|
||||
output->unpacked = true;
|
||||
return ret;
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,257 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/data.h
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_DATA_H_
|
||||
#define MLC_LLM_SERVE_DATA_H_
|
||||
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
#include <tvm/ffi/optional.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
#include <atomic>
|
||||
#include <optional>
|
||||
|
||||
#include "../tokenizers/tokenizers.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectPtr;
|
||||
using tvm::ffi::ObjectRef;
|
||||
using tvm::ffi::Optional;
|
||||
using tvm::ffi::Shape;
|
||||
|
||||
class Model;
|
||||
|
||||
/****************** DataNode ******************/
|
||||
|
||||
/*! \brief The base class of multi-modality data (text, tokens, embedding, etc). */
|
||||
class DataNode : public Object {
|
||||
public:
|
||||
/*! \brief Get the length (equivalent number of tokens) of the data. */
|
||||
virtual int GetLength() const = 0;
|
||||
|
||||
/*!
|
||||
* \brief Compute the embedding of this data with regard to the input model.
|
||||
* When the input destination pointer is not nullptr, it in-place writes the
|
||||
* embedding into the input destination array at the given offset.
|
||||
* Otherwise, the embeddings will be directly returned back.
|
||||
* \param model The model to take embeddings from.
|
||||
* \param dst The destination array of the embedding lookup.
|
||||
* \param offset The token offset where the computed embeddings will be written
|
||||
* into the destination array.
|
||||
* \return The updated destination embedding array or the computed embeddings.
|
||||
* \note When `dst` is nullptr, we require `offset` to be 0.
|
||||
*/
|
||||
virtual ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const = 0;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<DataNode>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const uint32_t _type_child_slots = 3;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.Data", DataNode, Object);
|
||||
};
|
||||
|
||||
class Data : public ObjectRef {
|
||||
public:
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Data, ObjectRef, DataNode);
|
||||
};
|
||||
|
||||
/*! \brief Split the given data array into two arrays at the "split_pos" position. */
|
||||
std::pair<Array<Data>, Array<Data>> SplitData(const Array<Data>& original_data, int total_length,
|
||||
int split_pos);
|
||||
|
||||
/****************** TextDataNode ******************/
|
||||
|
||||
/*! \brief The class of text data, containing a text string. */
|
||||
class TextDataNode : public DataNode {
|
||||
public:
|
||||
/*! \brief The text string. */
|
||||
tvm::ffi::String text;
|
||||
|
||||
int GetLength() const final;
|
||||
ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<TextDataNode>();
|
||||
}
|
||||
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.TextData", TextDataNode, DataNode);
|
||||
};
|
||||
|
||||
class TextData : public Data {
|
||||
public:
|
||||
explicit TextData(String text);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TextData, Data, TextDataNode);
|
||||
};
|
||||
|
||||
/****************** TokenDataNode ******************/
|
||||
|
||||
/*! \brief The class of token data, containing a list of token ids. */
|
||||
class TokenDataNode : public DataNode {
|
||||
public:
|
||||
/*! \brief The token ids. */
|
||||
Shape token_ids;
|
||||
|
||||
int GetLength() const final;
|
||||
ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<TokenDataNode>();
|
||||
}
|
||||
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.TokenData", TokenDataNode, DataNode);
|
||||
};
|
||||
|
||||
class TokenData : public Data {
|
||||
public:
|
||||
explicit TokenData(Shape token_ids);
|
||||
|
||||
explicit TokenData(std::vector<int32_t> token_ids);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TokenData, Data, TokenDataNode);
|
||||
};
|
||||
|
||||
/****************** ImageDataNode ******************/
|
||||
|
||||
/*! \brief The class of image data, containing a 3D array of pixel values. */
|
||||
class ImageDataNode : public DataNode {
|
||||
public:
|
||||
/*! \brief The pixel values. */
|
||||
Tensor image;
|
||||
int embed_size;
|
||||
|
||||
int GetLength() const final;
|
||||
ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<ImageDataNode>();
|
||||
}
|
||||
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.ImageData", ImageDataNode, DataNode);
|
||||
};
|
||||
|
||||
class ImageData : public Data {
|
||||
public:
|
||||
explicit ImageData(Tensor image, int embed_size);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ImageData, Data, ImageDataNode);
|
||||
};
|
||||
|
||||
/****************** SampleResult ******************/
|
||||
|
||||
// The pair of a token id and its probability in sampling.
|
||||
using TokenProbPair = std::pair<int32_t, float>;
|
||||
|
||||
/*!
|
||||
* \brief The class of sampler's sampling result.
|
||||
* It's not a TVM object since it will not be used directly on Python side.
|
||||
*/
|
||||
struct SampleResult {
|
||||
/*! \brief The token id and probability of the sampled token. */
|
||||
TokenProbPair sampled_token_id;
|
||||
/*! \brief The token id and probability of the tokens with top probabilities. */
|
||||
std::vector<TokenProbPair> top_prob_tokens;
|
||||
|
||||
/*! \brief Get the sampled token id. */
|
||||
int32_t GetTokenId() const;
|
||||
|
||||
/*!
|
||||
* \brief Get the logprob JSON string of this token with regard
|
||||
* to OpenAI API at https://platform.openai.com/docs/api-reference/chat/object.
|
||||
* \param tokenizer The tokenizer for token table lookup.
|
||||
* \param logprob A boolean indicating if need to return log probability.
|
||||
* \return A JSON string that conforms to the logprob spec in OpenAI API.
|
||||
*/
|
||||
std::string GetLogProbJSON(const Tokenizer& tokenizer, bool logprob) const;
|
||||
};
|
||||
|
||||
/****************** RequestStreamOutput ******************/
|
||||
|
||||
/*!
|
||||
* \brief The generated delta request output that is streamed back
|
||||
* through callback stream function.
|
||||
*
|
||||
* \note: This output object corresponds to parallel generated outputs when n != 1.
|
||||
*
|
||||
* For example, if n=2, then group_delta_token_ids[0] matches to the output stream 0
|
||||
* and group_delta_token_ids[1] matches to the output stream 1
|
||||
*/
|
||||
class RequestStreamOutputObj : public Object {
|
||||
public:
|
||||
/*! \brief The id of the request that the function is invoked for. */
|
||||
String request_id;
|
||||
/*!
|
||||
* \brief The new generated token ids since the last callback invocation
|
||||
* for the input request.
|
||||
*/
|
||||
std::vector<std::vector<int64_t>> group_delta_token_ids;
|
||||
/*! \brief The logprobs JSON strings of the new generated tokens since last invocation. */
|
||||
std::optional<std::vector<std::vector<String>>> group_delta_logprob_json_strs;
|
||||
/*!
|
||||
* \brief The finish reason of the request when it is finished,
|
||||
* of None if the request has not finished yet.
|
||||
*/
|
||||
std::vector<Optional<String>> group_finish_reason;
|
||||
/*!
|
||||
* \brief The usage field of the response, this is global to all streams.
|
||||
*/
|
||||
Optional<String> request_final_usage_json_str;
|
||||
|
||||
/*!
|
||||
* \brief The extra prefix string of all requests.
|
||||
*/
|
||||
std::vector<String> group_extra_prefix_string;
|
||||
|
||||
std::atomic<bool> unpacked = false;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<RequestStreamOutputObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.RequestStreamOutput", RequestStreamOutputObj, Object);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Managed reference to RequestStreamOutputObj.
|
||||
* \sa RequestStreamOutputObj
|
||||
*/
|
||||
class RequestStreamOutput : public ObjectRef {
|
||||
public:
|
||||
explicit RequestStreamOutput(
|
||||
String request_id, std::vector<std::vector<int64_t>> group_delta_token_ids,
|
||||
std::optional<std::vector<std::vector<String>>> group_delta_logprob_json_strs,
|
||||
std::vector<Optional<String>> group_finish_reason,
|
||||
std::vector<String> group_extra_prefix_string);
|
||||
|
||||
static RequestStreamOutput Usage(String request_id, String request_final_usage_json_str);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(RequestStreamOutput, ObjectRef,
|
||||
RequestStreamOutputObj);
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_DATA_H_
|
||||
@@ -0,0 +1,77 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/draft_token_workspace_manager.cc
|
||||
*/
|
||||
|
||||
#include "draft_token_workspace_manager.h"
|
||||
|
||||
#include "model.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() { DraftTokenWorkspaceManagerObj::RegisterReflection(); }
|
||||
|
||||
DraftTokenWorkspaceManagerObj::DraftTokenWorkspaceManagerObj(int max_num_tokens, int vocab_size,
|
||||
int hidden_size,
|
||||
DLDataType hidden_states_dtype,
|
||||
DLDevice device,
|
||||
const FunctionTable& ft)
|
||||
: max_num_tokens_(max_num_tokens),
|
||||
vocab_size_(vocab_size),
|
||||
hidden_size_(hidden_size),
|
||||
hidden_states_dtype_(hidden_states_dtype),
|
||||
device_(device),
|
||||
ft_(ft) {
|
||||
free_slots_.resize(max_num_tokens);
|
||||
std::iota(free_slots_.begin(), free_slots_.end(), 0);
|
||||
}
|
||||
|
||||
void DraftTokenWorkspaceManagerObj::AllocSlots(int num_slots, std::vector<int>* result) {
|
||||
TVM_FFI_ICHECK_LE(num_slots, free_slots_.size());
|
||||
result->assign(free_slots_.rbegin(), free_slots_.rbegin() + num_slots);
|
||||
free_slots_.resize(free_slots_.size() - num_slots);
|
||||
for (int slot : (*result)) {
|
||||
ref_count_[slot] = 1;
|
||||
}
|
||||
}
|
||||
|
||||
void DraftTokenWorkspaceManagerObj::AllocSlots(int num_slots,
|
||||
const std::vector<int>& initial_ref_count,
|
||||
std::vector<int>* result) {
|
||||
TVM_FFI_ICHECK_LE(num_slots, free_slots_.size());
|
||||
TVM_FFI_ICHECK_EQ(num_slots, initial_ref_count.size());
|
||||
result->assign(free_slots_.rbegin(), free_slots_.rbegin() + num_slots);
|
||||
free_slots_.resize(free_slots_.size() - num_slots);
|
||||
for (int i = 0; i < num_slots; ++i) {
|
||||
int slot = (*result)[i];
|
||||
TVM_FFI_ICHECK(initial_ref_count[i] > 0);
|
||||
ref_count_[slot] = initial_ref_count[i];
|
||||
}
|
||||
}
|
||||
|
||||
void DraftTokenWorkspaceManagerObj::FreeSlots(const std::vector<int>& slots) {
|
||||
for (int slot : slots) {
|
||||
if (--ref_count_.at(slot) == 0) {
|
||||
free_slots_.push_back(slot);
|
||||
ref_count_.erase(slot);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void DraftTokenWorkspaceManagerObj::AllocWorkspace(ModelWorkspace* workspace,
|
||||
bool require_hidden_states) {
|
||||
workspace->draft_probs =
|
||||
Tensor::Empty({max_num_tokens_, vocab_size_}, DLDataType{kDLFloat, 32, 1}, device_);
|
||||
workspace->draft_probs_storage =
|
||||
Tensor::Empty({max_num_tokens_, vocab_size_}, DLDataType{kDLFloat, 32, 1}, device_);
|
||||
if (require_hidden_states) {
|
||||
workspace->draft_hidden_states_storage = ft_.Empty(
|
||||
{max_num_tokens_, hidden_size_}, hidden_states_dtype_, device_, /*worker0_only=*/false);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,117 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/draft_token_workspace_manager.h
|
||||
*/
|
||||
|
||||
#ifndef MLC_LLM_SERVE_DRAFT_TOKEN_WORKSPACE_MANAGER_H_
|
||||
#define MLC_LLM_SERVE_DRAFT_TOKEN_WORKSPACE_MANAGER_H_
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/runtime/device_api.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <optional>
|
||||
#include <vector>
|
||||
|
||||
#include "data.h"
|
||||
#include "function_table.h"
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::Device;
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
|
||||
struct ModelWorkspace;
|
||||
|
||||
/*!
|
||||
* \brief Managing the workspace for draft token generation.
|
||||
*
|
||||
* The workspace is used to store the associated states for each draft token, including the
|
||||
* probability distribution of the draft token, the hidden states, etc. The workspace manager
|
||||
* maintains a pool of slots for the draft tokens to store the states.
|
||||
*/
|
||||
class DraftTokenWorkspaceManagerObj : public Object {
|
||||
public:
|
||||
/*!
|
||||
* \brief Constructor
|
||||
* \param max_num_tokens The maximum number of draft tokens that can be stored in the workspace.
|
||||
* \param vocab_size The size of the vocabulary.
|
||||
* \param hidden_size The size of the hidden states.
|
||||
* \param hidden_states_dtype The data type of the hidden states.
|
||||
* \param device The device running the model.
|
||||
* \param ft The function table.
|
||||
*/
|
||||
DraftTokenWorkspaceManagerObj(int max_num_tokens, int vocab_size, int hidden_size,
|
||||
DLDataType hidden_states_dtype, DLDevice device,
|
||||
const FunctionTable& ft);
|
||||
|
||||
/*!
|
||||
* \brief Allocate the workspace for draft tokens and update `ModelWorkspace` data structure.
|
||||
* \param workspace The object to stored the allocated draft token workspace.
|
||||
* \param require_hidden_states Whether to allocate workspace for the hidden states.
|
||||
*/
|
||||
void AllocWorkspace(ModelWorkspace* workspace, bool require_hidden_states);
|
||||
|
||||
/*!
|
||||
* \brief Allocate slots for the draft tokens.
|
||||
* \param num_slots The number of slots to allocate.
|
||||
* \param result The vector to store the allocated slots.
|
||||
*/
|
||||
void AllocSlots(int num_slots, std::vector<int>* result);
|
||||
|
||||
/*!
|
||||
* \brief Allocate slots for the draft tokens.
|
||||
* \param num_slots The number of slots to allocate.
|
||||
* \param initial_ref_count The initial reference count for each slot.
|
||||
* \param result The vector to store the allocated slots.
|
||||
*/
|
||||
void AllocSlots(int num_slots, const std::vector<int>& initial_ref_count,
|
||||
std::vector<int>* result);
|
||||
|
||||
/*!
|
||||
* \brief Free the slots.
|
||||
* \param slots The slots to free.
|
||||
*/
|
||||
void FreeSlots(const std::vector<int>& slots);
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<DraftTokenWorkspaceManagerObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.DraftTokenWorkspaceManager",
|
||||
DraftTokenWorkspaceManagerObj, Object);
|
||||
|
||||
private:
|
||||
std::vector<int> free_slots_;
|
||||
int max_num_tokens_;
|
||||
int vocab_size_;
|
||||
int hidden_size_;
|
||||
DLDataType hidden_states_dtype_;
|
||||
DLDevice device_;
|
||||
const FunctionTable& ft_;
|
||||
std::unordered_map<int, int> ref_count_;
|
||||
};
|
||||
|
||||
class DraftTokenWorkspaceManager : public ObjectRef {
|
||||
public:
|
||||
DraftTokenWorkspaceManager(int max_num_tokens, int vocab_size, int hidden_size,
|
||||
DLDataType hidden_states_dtype, DLDevice device,
|
||||
const FunctionTable& ft) {
|
||||
data_ = tvm::ffi::make_object<DraftTokenWorkspaceManagerObj>(
|
||||
max_num_tokens, vocab_size, hidden_size, hidden_states_dtype, device, ft);
|
||||
}
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DraftTokenWorkspaceManager, ObjectRef,
|
||||
DraftTokenWorkspaceManagerObj);
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_DRAFT_TOKEN_WORKSPACE_MANAGER_H_
|
||||
+1102
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,126 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine.h
|
||||
* \brief The header of serving engine in MLC LLM.
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_ENGINE_H_
|
||||
#define MLC_LLM_SERVE_ENGINE_H_
|
||||
|
||||
#include "data.h"
|
||||
#include "engine_state.h"
|
||||
#include "event_trace_recorder.h"
|
||||
#include "request.h"
|
||||
#include "request_state.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
|
||||
class Engine;
|
||||
|
||||
/*!
|
||||
* \brief The output of engine creation, including the created engine and
|
||||
* the default generation config for requests.
|
||||
*/
|
||||
struct EngineCreationOutput {
|
||||
std::unique_ptr<Engine> reloaded_engine;
|
||||
EngineConfig completed_engine_config;
|
||||
GenerationConfig default_generation_cfg;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief The engine interface for request serving in MLC LLM.
|
||||
* The engine can run one or multiple LLM models internally for
|
||||
* text generation. Usually, when there are multiple models,
|
||||
* speculative inference will be activated, where the first model
|
||||
* (index 0) is the main "large model" that has better generation
|
||||
* quality, and all other models are "small" models that used for
|
||||
* speculation.
|
||||
* The engine receives requests from the "AddRequest" method. For
|
||||
* an given request, the engine will keep generating new tokens for
|
||||
* the request until finish (under certain criterion). After finish,
|
||||
* the engine will return the generation result through the callback
|
||||
* function provided by the request.
|
||||
* \note For now only one model run in the engine is supported.
|
||||
* Multiple model support such as speculative inference will
|
||||
* be followed soon in the future.
|
||||
*
|
||||
* The public interface of Engine has the following three categories:
|
||||
* - engine management,
|
||||
* - high-level request management,
|
||||
* - engine "step" action.
|
||||
*/
|
||||
class Engine {
|
||||
public:
|
||||
/********************** Engine Management **********************/
|
||||
virtual ~Engine() = default;
|
||||
|
||||
/*!
|
||||
* \brief Create an engine in unique pointer.
|
||||
* \param engine_config_json_str The serialized JSON string of the engine config.
|
||||
* \param device The device where the run models.
|
||||
* \param request_stream_callback The request stream callback function to.
|
||||
* \param trace_recorder Event trace recorder for requests.
|
||||
* \return The created Engine in pointer, and the default generation config.
|
||||
*/
|
||||
static Result<EngineCreationOutput> Create(const std::string& engine_config_json_str,
|
||||
Device device,
|
||||
FRequestStreamCallback request_stream_callback,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
|
||||
/*! \brief Reset the engine, clean up all running data and metrics. */
|
||||
virtual void Reset() = 0;
|
||||
|
||||
/*! \brief Check if the engine has no request to process. */
|
||||
virtual bool Empty() = 0;
|
||||
|
||||
/*! \brief Get the request stream callback function of the engine. */
|
||||
virtual FRequestStreamCallback GetRequestStreamCallback() = 0;
|
||||
|
||||
/*! \brief Set the request stream callback function of the engine. */
|
||||
virtual void SetRequestStreamCallback(FRequestStreamCallback request_stream_callback) = 0;
|
||||
|
||||
/***************** High-level Request Management *****************/
|
||||
|
||||
/*! \brief Add a new request to the engine. */
|
||||
virtual void AddRequest(Request request) = 0;
|
||||
|
||||
/*! \brief Abort the input request (specified by id string) from engine. */
|
||||
virtual void AbortRequest(const String& request_id) = 0;
|
||||
|
||||
/*! \brief Abort all requests from the engine. */
|
||||
virtual void AbortAllRequests() = 0;
|
||||
|
||||
/*********************** Engine Action ***********************/
|
||||
|
||||
/*!
|
||||
* \brief The main function that the engine takes a step of action.
|
||||
* At each step, the engine may decide to
|
||||
* - run prefill for one (or more) requests,
|
||||
* - run one-step decode for the all existing requests
|
||||
* ...
|
||||
* In the end of certain actions (e.g., decode), the engine will
|
||||
* check if any request has finished, and will return the
|
||||
* generation results for those finished requests.
|
||||
*/
|
||||
virtual void Step() = 0;
|
||||
|
||||
/************** Debug/Profile **************/
|
||||
|
||||
/*! \brief Internal engine metrics. */
|
||||
virtual String JSONMetrics() = 0;
|
||||
|
||||
/*! \brief Call the given global function on all workers. Only for debug purpose. */
|
||||
virtual void DebugCallFuncOnAllAllWorker(const String& func_name, Optional<String> func_args) = 0;
|
||||
};
|
||||
|
||||
void AbortRequestImpl(EngineState estate, const Array<Model>& models, const String& request_id,
|
||||
String finish_reason = "abort");
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_ENGINE_H_
|
||||
@@ -0,0 +1,16 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/action.cc
|
||||
*/
|
||||
|
||||
#include "action.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() { EngineActionObj::RegisterReflection(); }
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,264 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/action.h
|
||||
* \brief The abstraction of actions (e.g., prefill/decode) that an
|
||||
* Engine can take at each time step.
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_ENGINE_ACTIONS_ACTION_H_
|
||||
#define MLC_LLM_SERVE_ENGINE_ACTIONS_ACTION_H_
|
||||
|
||||
#include "../config.h"
|
||||
#include "../draft_token_workspace_manager.h"
|
||||
#include "../engine.h"
|
||||
#include "../engine_state.h"
|
||||
#include "../event_trace_recorder.h"
|
||||
#include "../model.h"
|
||||
#include "../sampler/sampler.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
|
||||
/*!
|
||||
* \brief The abstraction of actions that an Engine can take at each time step.
|
||||
* The only core interface of an action is the `Step` function.
|
||||
* At high level, the Step function takes the current engine state
|
||||
* as input, invokes model functions (such as batched-prefill or
|
||||
* batched-decode), run sampler to sample new tokens, and update
|
||||
* the engine state.
|
||||
*/
|
||||
class EngineActionObj : public Object {
|
||||
public:
|
||||
/*!
|
||||
* \brief The behavior of the engine action in a single step.
|
||||
* \param estate The engine state to be analyzed and updated.
|
||||
* \return The processed requests in this step.
|
||||
*/
|
||||
virtual Array<Request> Step(EngineState estate) = 0;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<EngineActionObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.EngineAction", EngineActionObj, Object);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Managed reference of EngineActionObj.
|
||||
* It declares the full list of supported actions.
|
||||
* \sa EngineActionObj
|
||||
*/
|
||||
class EngineAction : public ObjectRef {
|
||||
public:
|
||||
/*!
|
||||
* \brief Create the action that prefills requests in the `waiting_queue`
|
||||
* of the engine state.
|
||||
* \param models The models to run prefill in.
|
||||
* \param logit_processor The logit processor.
|
||||
* \param sampler The sampler to sample new tokens.
|
||||
* \param model_workspaces The workspace of each model.
|
||||
* \param engine_config The engine config.
|
||||
* \param model_configs The config of each model.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* \return The created action object.
|
||||
*/
|
||||
static EngineAction NewRequestPrefill(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
/*!
|
||||
* \brief Create the action that prefills requests in the `waiting_queue`
|
||||
* of the engine state.
|
||||
* \param models The models to run prefill in.
|
||||
* \param logit_processor The logit processor.
|
||||
* \param sampler The sampler to sample new tokens.
|
||||
* \param model_workspaces The workspace of each model.
|
||||
* \param draft_token_workspace_manager The draft token workspace manager.
|
||||
* \param engine_config The engine config.
|
||||
* \param model_configs The config of each model.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* \return The created action object.
|
||||
*/
|
||||
static EngineAction EagleNewRequestPrefill(
|
||||
Array<Model> models, LogitProcessor logit_processor, Sampler sampler,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager, EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
/*!
|
||||
* \brief Create the action that runs one-step decode for requests in the
|
||||
* `running_queue` of engine state. Preempt low-priority requests
|
||||
* accordingly when it is impossible to decode all the running requests.
|
||||
* \note The BatchDecode action **does not** take effect for speculative
|
||||
* decoding scenarios where there are multiple models. For speculative
|
||||
* decoding in the future, we will use other specific actions.
|
||||
* \param models The model to run decode in. When there are multiple
|
||||
* models, the `Step` function of the created action will not take effect.
|
||||
* \param tokenizer The tokenizer of the engine.
|
||||
* \param sampler The sampler to sample new tokens.
|
||||
* \param engine_config The engine config.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* \return The created action object.
|
||||
*/
|
||||
static EngineAction BatchDecode(Array<Model> models, Tokenizer tokenizer,
|
||||
LogitProcessor logit_processor, Sampler sampler,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
|
||||
/*!
|
||||
* \brief Create the action that runs one-step speculative draft proposal for
|
||||
* requests in the `running_queue` of engine state. Preempt low-priority requests
|
||||
* accordingly when it is impossible to decode all the running requests.
|
||||
* \param models The model to run decode in. When there are multiple
|
||||
* models, the `Step` function of the created action will not take effect.
|
||||
* \param sampler The sampler to sample new tokens.
|
||||
* \param model_workspaces The workspace of each model.
|
||||
* \param draft_token_workspace_manager The draft token workspace manager.
|
||||
* \param engine_config The engine config.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* \return The created action object.
|
||||
*/
|
||||
static EngineAction BatchDraft(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler, std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
|
||||
/*!
|
||||
* \brief Create the action that runs one-step speculative draft proposal for
|
||||
* requests in the `running_queue` of engine state. Preempt low-priority requests
|
||||
* accordingly when it is impossible to decode all the running requests.
|
||||
* \param models The model to run decode in. When there are multiple
|
||||
* models, the `Step` function of the created action will not take effect.
|
||||
* \param sampler The sampler to sample new tokens.
|
||||
* \param model_workspaces The workspace of each model.
|
||||
* \param draft_token_workspace_manager The draft token workspace manager.
|
||||
* \param engine_config The engine config.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* \return The created action object.
|
||||
*/
|
||||
static EngineAction EagleBatchDraft(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler, std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
|
||||
/*!
|
||||
* \brief Create the action that runs one-step speculative verification for requests in the
|
||||
* `running_queue` of engine state. Preempt low-priority requests
|
||||
* accordingly when it is impossible to decode all the running requests.
|
||||
* \param models The model to run decode in. When there are multiple
|
||||
* models, the `Step` function of the created action will not take effect.
|
||||
* \param model_workspaces The workspace of each model.
|
||||
* \param draft_token_workspace_manager The draft token workspace manager.
|
||||
* \param sampler The sampler to sample new tokens.
|
||||
* \param engine_config The engine config.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* \return The created action object.
|
||||
*/
|
||||
static EngineAction BatchVerify(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler, std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
|
||||
/*!
|
||||
* \brief Create the action that runs one-step speculative verification for requests in the
|
||||
* `running_queue` of engine state. Preempt low-priority requests
|
||||
* accordingly when it is impossible to decode all the running requests.
|
||||
* \param models The model to run decode in. When there are multiple
|
||||
* models, the `Step` function of the created action will not take effect.
|
||||
* \param sampler The sampler to sample new tokens.
|
||||
* \param model_workspaces The workspace of each model.
|
||||
* \param draft_token_workspace_manager The draft token workspace manager.
|
||||
* \param engine_config The engine config.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* \return The created action object.
|
||||
*/
|
||||
static EngineAction EagleBatchVerify(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
/*!
|
||||
* \brief Create the action that executes the jump-forward decoding to predict the next tokens
|
||||
* according to the grammar constraint. Does nothing for the requests without grammar. The
|
||||
* predicted tokens will be fed to the next BatchDecode action. Retokenization may happen when
|
||||
* the predicted string breaks the tokenization boundary.
|
||||
* \param models The model to run decode in. When there are multiple
|
||||
* models, the `Step` function of the created action will not take effect.
|
||||
* \param tokenizer The tokenizer of the engine.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* \return The created action object.
|
||||
*/
|
||||
static EngineAction BatchJumpForward(Array<Model> models, Tokenizer tokenizer,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
|
||||
/*!
|
||||
* \brief Create the action that first makes a decision on whether to run speculative
|
||||
* decoding or normal mode batch decode, and then runs the selected actions.
|
||||
* \param spec_decode_actions The actions for speculative decoding.
|
||||
* \param batch_decode_actions The actions for normal mode batch decoding.
|
||||
* \param engine_config The engine config.
|
||||
* \return The created action object
|
||||
*/
|
||||
static EngineAction AutoSpecDecode(std::vector<EngineAction> spec_decode_actions,
|
||||
std::vector<EngineAction> batch_decode_actions,
|
||||
EngineConfig engine_config);
|
||||
|
||||
/*!
|
||||
* \brief Create the action that runs the disaggregation preparation for prefill.
|
||||
* \param models The underlying models whose KV cache are to be updated.
|
||||
* \param engine_config The engine config.
|
||||
* \param model_configs The config of each model.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* \param request_stream_callback The stream callback function to pass the prefill
|
||||
* preparation result back, including the KV cache append metadata and the prefix
|
||||
* matched length in the prefix cache.
|
||||
* \return The created action object.
|
||||
*/
|
||||
static EngineAction DisaggPrepareReceive(Array<Model> models, EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
Optional<EventTraceRecorder> trace_recorder,
|
||||
FRequestStreamCallback request_stream_callback);
|
||||
|
||||
/*!
|
||||
* \brief Create the action that runs the prefill and sends KV data to remote instance.
|
||||
* \param models The underlying models whose KV cache are to be updated.
|
||||
* \param model_workspaces The workspace of each model.
|
||||
* \param engine_config The engine config.
|
||||
* \param model_configs The config of each model.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* \param request_stream_callback The stream callback function to pass the prefill
|
||||
* preparation result back, including the KV cache append metadata and the prefix
|
||||
* matched length in the prefix cache.
|
||||
* \param device The device of the model for synchronization.
|
||||
* \return The created action object.
|
||||
*/
|
||||
static 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);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(EngineAction, ObjectRef, EngineActionObj);
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_ENGINE_ACTIONS_ACTION_H_
|
||||
@@ -0,0 +1,452 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/action_commons.cc
|
||||
*/
|
||||
|
||||
#include "action_commons.h"
|
||||
|
||||
#include <tvm/support/cuda/nvtx.h>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::support::NVTXScopedRange;
|
||||
|
||||
Array<EngineAction> CreateEngineActions(Array<Model> models, EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
LogitProcessor logit_processor, Sampler sampler,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
Tokenizer tokenizer,
|
||||
Optional<EventTraceRecorder> trace_recorder,
|
||||
FRequestStreamCallback request_stream_callback,
|
||||
Device device) {
|
||||
Array<EngineAction> actions;
|
||||
ModelMetadata model_metadata = models[0]->GetMetadata();
|
||||
if (engine_config->speculative_mode != SpeculativeMode::kDisable) {
|
||||
// Speculative decoding is only possible for more than one model.
|
||||
TVM_FFI_ICHECK_GT(models.size(), 1U);
|
||||
if (engine_config->speculative_mode == SpeculativeMode::kEagle) {
|
||||
TVM_FFI_ICHECK_GT(engine_config->spec_draft_length, 0)
|
||||
<< "The automatic spec decoding does not support Eagle mode as of now.";
|
||||
actions = {EngineAction::EagleNewRequestPrefill(models, //
|
||||
logit_processor, //
|
||||
sampler, //
|
||||
model_workspaces, //
|
||||
draft_token_workspace_manager, //
|
||||
engine_config, //
|
||||
model_configs, //
|
||||
trace_recorder),
|
||||
EngineAction::EagleBatchDraft(models, logit_processor, sampler, model_workspaces,
|
||||
draft_token_workspace_manager, engine_config,
|
||||
trace_recorder),
|
||||
EngineAction::EagleBatchVerify(models, logit_processor, sampler, model_workspaces,
|
||||
draft_token_workspace_manager, engine_config,
|
||||
trace_recorder)};
|
||||
} else if (engine_config->speculative_mode == SpeculativeMode::kMedusa) {
|
||||
TVM_FFI_ICHECK_GT(engine_config->spec_draft_length, 0)
|
||||
<< "The automatic spec decoding does not support Eagle mode as of now.";
|
||||
actions = {EngineAction::EagleNewRequestPrefill(models, //
|
||||
logit_processor, //
|
||||
sampler, //
|
||||
model_workspaces, //
|
||||
draft_token_workspace_manager, //
|
||||
engine_config, //
|
||||
model_configs, //
|
||||
trace_recorder),
|
||||
EngineAction::EagleBatchVerify(models, logit_processor, sampler, model_workspaces,
|
||||
draft_token_workspace_manager, engine_config,
|
||||
trace_recorder)};
|
||||
} else if (engine_config->spec_draft_length > 0) {
|
||||
// The "small draft" mode speculative decoding.
|
||||
// If "engine_config->spec_draft_length" > 0, it means the draft length is
|
||||
// configured to be a fixed value.
|
||||
actions = {
|
||||
EngineAction::NewRequestPrefill(models, //
|
||||
logit_processor, //
|
||||
sampler, //
|
||||
model_workspaces, //
|
||||
engine_config, //
|
||||
model_configs, //
|
||||
trace_recorder),
|
||||
EngineAction::BatchDraft(models, logit_processor, sampler, model_workspaces,
|
||||
draft_token_workspace_manager, engine_config, trace_recorder),
|
||||
EngineAction::BatchVerify(models, logit_processor, sampler, model_workspaces,
|
||||
draft_token_workspace_manager, engine_config, trace_recorder)};
|
||||
} else {
|
||||
// The "small draft" mode speculative decoding.
|
||||
// "engine_config->spec_draft_length" being 0 means we want to enable
|
||||
// automatic speculative decoding, which decides the spec decoding draft length
|
||||
// automatically.
|
||||
actions = {EngineAction::NewRequestPrefill(models, //
|
||||
logit_processor, //
|
||||
sampler, //
|
||||
model_workspaces, //
|
||||
engine_config, //
|
||||
model_configs, //
|
||||
trace_recorder),
|
||||
EngineAction::AutoSpecDecode(
|
||||
/*spec_decode_actions=*/{EngineAction::BatchDraft(
|
||||
models, logit_processor, sampler,
|
||||
model_workspaces, draft_token_workspace_manager,
|
||||
engine_config, trace_recorder),
|
||||
EngineAction::BatchVerify(
|
||||
models, logit_processor, sampler,
|
||||
model_workspaces, draft_token_workspace_manager,
|
||||
engine_config, trace_recorder)},
|
||||
/*batch_decode_actions=*/
|
||||
{EngineAction::BatchDecode(models, tokenizer, logit_processor, sampler,
|
||||
engine_config, trace_recorder)},
|
||||
engine_config)};
|
||||
}
|
||||
} else if (model_metadata.disaggregation) {
|
||||
actions = {EngineAction::NewRequestPrefill(models, //
|
||||
logit_processor, //
|
||||
sampler, //
|
||||
model_workspaces, //
|
||||
engine_config, //
|
||||
model_configs, //
|
||||
trace_recorder),
|
||||
EngineAction::BatchDecode(models, tokenizer, logit_processor, sampler, engine_config,
|
||||
trace_recorder)};
|
||||
} else {
|
||||
// The normal mode.
|
||||
actions = {EngineAction::NewRequestPrefill(models, //
|
||||
logit_processor, //
|
||||
sampler, //
|
||||
model_workspaces, //
|
||||
engine_config, //
|
||||
model_configs, //
|
||||
trace_recorder),
|
||||
EngineAction::BatchJumpForward(models, tokenizer, trace_recorder),
|
||||
EngineAction::BatchDecode(models, tokenizer, logit_processor, sampler, engine_config,
|
||||
trace_recorder)};
|
||||
}
|
||||
|
||||
if (model_metadata.disaggregation) {
|
||||
// Insert the disaggregation actions.
|
||||
Array<EngineAction> disaggregation_actions = {
|
||||
EngineAction::DisaggPrepareReceive(models, engine_config, model_configs, trace_recorder,
|
||||
request_stream_callback),
|
||||
EngineAction::DisaggRemoteSend(models, model_workspaces, engine_config, model_configs,
|
||||
trace_recorder, request_stream_callback, device)};
|
||||
actions.insert(actions.begin(), disaggregation_actions.begin(), disaggregation_actions.end());
|
||||
}
|
||||
return actions;
|
||||
}
|
||||
|
||||
void RemoveRequestFromModel(EngineState estate, int64_t req_internal_id,
|
||||
const Array<Model>& models) {
|
||||
// Remove the request from all models (usually the KV cache).
|
||||
for (Model model : models) {
|
||||
model->RemoveSequence(req_internal_id);
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Remove the given request state entry.
|
||||
* \param estate The engine state to update after removal.
|
||||
* \param models The models to remove the given request from.
|
||||
* \param rsentry The request state entry to remove.
|
||||
*/
|
||||
void RemoveRequestStateEntry(EngineState estate, const Array<Model>& models,
|
||||
RequestStateEntry rsentry,
|
||||
Optional<DraftTokenWorkspaceManager> draft_token_workspace_manager) {
|
||||
if (draft_token_workspace_manager.has_value()) {
|
||||
std::vector<int> draft_token_slots;
|
||||
for (const RequestModelState& mstate : rsentry->mstates) {
|
||||
mstate->RemoveAllDraftTokens(&draft_token_slots);
|
||||
draft_token_workspace_manager.value()->FreeSlots(draft_token_slots);
|
||||
}
|
||||
}
|
||||
if (estate->prefix_cache->HasSequence(rsentry->mstates[0]->internal_id)) {
|
||||
// If the sequence is stored in prefix cache, call prefix cache to remove.
|
||||
if (!(rsentry->request->generation_cfg->debug_config.pinned_system_prompt)) {
|
||||
// If the request is not pinned, recycle the request.
|
||||
estate->prefix_cache->RecycleSequence(rsentry->mstates[0]->internal_id, /*lazy=*/true);
|
||||
}
|
||||
// If the request is pinned, do nothing over the prefix cache and KVCache.
|
||||
} else {
|
||||
// If the sequence is not stored in prefix cache, remove it directly.
|
||||
RemoveRequestFromModel(estate, rsentry->mstates[0]->internal_id, models);
|
||||
estate->id_manager.RecycleId(rsentry->mstates[0]->internal_id);
|
||||
}
|
||||
}
|
||||
|
||||
void ProcessFinishedRequestStateEntries(
|
||||
const std::vector<RequestStateEntry>& finished_rsentries, EngineState estate,
|
||||
const Array<Model>& models, int max_single_sequence_length,
|
||||
Optional<DraftTokenWorkspaceManager> draft_token_workspace_manager,
|
||||
Array<RequestStreamOutput>* callback_delta_outputs) {
|
||||
NVTXScopedRange nvtx_scope("Process finished requests");
|
||||
// - Remove the finished request state entries.
|
||||
for (const RequestStateEntry& rsentry : finished_rsentries) {
|
||||
// The finished entry must be a leaf.
|
||||
TVM_FFI_ICHECK(rsentry->child_indices.empty());
|
||||
// Mark the status of this entry as finished.
|
||||
rsentry->status = RequestStateStatus::kFinished;
|
||||
// Remove the request state entry from all the models.
|
||||
RemoveRequestStateEntry(estate, models, rsentry, draft_token_workspace_manager);
|
||||
|
||||
RequestState rstate = estate->GetRequestState(rsentry->request);
|
||||
int parent_idx = rsentry->parent_idx;
|
||||
while (parent_idx != -1) {
|
||||
bool all_children_finished = true;
|
||||
for (int child_idx : rstate->entries[parent_idx]->child_indices) {
|
||||
if (rstate->entries[child_idx]->status != RequestStateStatus::kFinished) {
|
||||
all_children_finished = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!all_children_finished) {
|
||||
break;
|
||||
}
|
||||
|
||||
// All the children of the parent request state entry have finished.
|
||||
// So we mark the parent entry as finished.
|
||||
rstate->entries[parent_idx]->status = RequestStateStatus::kFinished;
|
||||
// Remove the request state entry from all the models.
|
||||
|
||||
RemoveRequestStateEntry(estate, models, rstate->entries[parent_idx],
|
||||
draft_token_workspace_manager);
|
||||
// Climb up to the parent.
|
||||
parent_idx = rstate->entries[parent_idx]->parent_idx;
|
||||
}
|
||||
|
||||
if (parent_idx == -1) {
|
||||
// Remove from running queue and engine state.
|
||||
auto it =
|
||||
std::find(estate->running_queue.begin(), estate->running_queue.end(), rsentry->request);
|
||||
TVM_FFI_ICHECK(it != estate->running_queue.end());
|
||||
estate->running_queue.erase(it);
|
||||
estate->request_states.erase(rsentry->request->id);
|
||||
|
||||
// Update engine metrics.
|
||||
const RequestStateEntry& root_rsentry = rstate->entries[0];
|
||||
auto trequest_finish = std::chrono::high_resolution_clock::now();
|
||||
|
||||
rstate->metrics.finish_time_point = trequest_finish;
|
||||
estate->metrics.RequestFinishUpdate(rstate->metrics);
|
||||
|
||||
// always stream back usage in backend
|
||||
callback_delta_outputs->push_back(RequestStreamOutput::Usage(
|
||||
root_rsentry->request->id, rstate->metrics.AsUsageJSONStr(true)));
|
||||
}
|
||||
estate->running_rsentries_changed = true;
|
||||
}
|
||||
}
|
||||
|
||||
void ActionStepPostProcess(Array<Request> requests, EngineState estate, const Array<Model>& models,
|
||||
const Tokenizer& tokenizer,
|
||||
FRequestStreamCallback request_stream_callback,
|
||||
int64_t max_single_sequence_length,
|
||||
Optional<DraftTokenWorkspaceManager> draft_token_workspace_manager,
|
||||
Optional<EventTraceRecorder> trace_recorder) {
|
||||
NVTXScopedRange nvtx_scope("EngineAction postproc");
|
||||
int num_requests = requests.size();
|
||||
estate->postproc_workspace.finished_rsentries.clear();
|
||||
estate->postproc_workspace.callback_delta_outputs.clear();
|
||||
estate->postproc_workspace.finished_rsentries.reserve(num_requests);
|
||||
estate->postproc_workspace.callback_delta_outputs.reserve(num_requests * 2);
|
||||
|
||||
// - Collect new generated tokens and finish reasons for requests.
|
||||
for (int r = 0; r < num_requests; ++r) {
|
||||
Request request = requests[r];
|
||||
int n = request->generation_cfg->n;
|
||||
RequestState rstate = estate->GetRequestState(requests[r]);
|
||||
|
||||
bool invoke_callback = false;
|
||||
RequestStreamOutput stream_output = rstate->postproc_states.GetStreamOutput();
|
||||
for (int i = 0; i < n; ++i) {
|
||||
const RequestStateEntry& rsentry = n == 1 ? rstate->entries[0] : rstate->entries[i + 1];
|
||||
rsentry->GetDeltaRequestReturn(tokenizer, max_single_sequence_length, &stream_output, i);
|
||||
if (stream_output->group_finish_reason[i].has_value()) {
|
||||
invoke_callback = true;
|
||||
estate->postproc_workspace.finished_rsentries.push_back(rsentry);
|
||||
}
|
||||
if (!stream_output->group_delta_token_ids[i].empty() ||
|
||||
!stream_output->group_extra_prefix_string[i].empty()) {
|
||||
invoke_callback = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (invoke_callback) {
|
||||
stream_output->unpacked = false;
|
||||
estate->postproc_workspace.callback_delta_outputs.push_back(std::move(stream_output));
|
||||
}
|
||||
|
||||
// Update prefix cache and metrics.
|
||||
for (const RequestStateEntry& rsentry : rstate->entries) {
|
||||
std::vector<int32_t>& token_ids = rsentry->token_ids_for_prefix_cache_update;
|
||||
token_ids.clear();
|
||||
if (!rsentry->mstates[0]->prefilled_inputs.empty()) {
|
||||
// Notify the prefix cache of the newly prefilled data.
|
||||
for (const Data& data : rsentry->mstates[0]->prefilled_inputs) {
|
||||
const TokenDataNode* token_data = data.as<TokenDataNode>();
|
||||
if (token_data == nullptr) continue;
|
||||
token_ids.insert(token_ids.end(), token_data->token_ids->data,
|
||||
token_data->token_ids->data + token_data->token_ids.size());
|
||||
// note that we are counting prefill tokens across all branches
|
||||
rstate->metrics.prefill_tokens += data->GetLength();
|
||||
}
|
||||
rsentry->mstates[0]->prefilled_inputs.clear();
|
||||
}
|
||||
int64_t num_committed_tokens = rsentry->mstates[0]->committed_tokens.size();
|
||||
if (rsentry->mstates[0]->cached_committed_tokens < num_committed_tokens - 1) {
|
||||
// Notify the prefix cache of the newly decoded data, except the last token as it is not
|
||||
// in KVCache yet.
|
||||
for (int64_t& i = rsentry->mstates[0]->cached_committed_tokens;
|
||||
i < num_committed_tokens - 1; ++i) {
|
||||
token_ids.push_back(rsentry->mstates[0]->committed_tokens[i].sampled_token_id.first);
|
||||
}
|
||||
}
|
||||
if (!token_ids.empty()) {
|
||||
estate->prefix_cache->ExtendSequence(rsentry->mstates[0]->internal_id, token_ids);
|
||||
}
|
||||
}
|
||||
|
||||
// - For all disaggregation requests with "remote_send",
|
||||
// if it does not appear in the waiting queue, it means the prefill has been finished.
|
||||
// In this case, we mark the request as finished.
|
||||
if (request->generation_cfg->debug_config.disagg_config.kind ==
|
||||
DisaggRequestKind::kRemoteSend) {
|
||||
auto it = std::find(estate->waiting_queue.begin(), estate->waiting_queue.end(), request);
|
||||
if (it == estate->waiting_queue.end()) {
|
||||
TVM_FFI_ICHECK_EQ(rstate->entries.size(), 1);
|
||||
estate->postproc_workspace.finished_rsentries.push_back(rstate->entries[0]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ProcessFinishedRequestStateEntries(estate->postproc_workspace.finished_rsentries, estate, models,
|
||||
max_single_sequence_length, draft_token_workspace_manager,
|
||||
&estate->postproc_workspace.callback_delta_outputs);
|
||||
|
||||
if (!estate->postproc_workspace.callback_delta_outputs.empty()) {
|
||||
NVTXScopedRange nvtx_scope("Call request stream callback");
|
||||
// - Invoke the stream callback function once for all collected requests.
|
||||
request_stream_callback(estate->postproc_workspace.callback_delta_outputs);
|
||||
}
|
||||
} // namespace serve
|
||||
|
||||
RequestStateEntry PreemptLastRunningRequestStateEntry(
|
||||
EngineState estate, const Array<Model>& models,
|
||||
Optional<DraftTokenWorkspaceManager> draft_token_workspace_manager,
|
||||
Optional<EventTraceRecorder> trace_recorder) {
|
||||
TVM_FFI_ICHECK(!estate->running_queue.empty());
|
||||
Request request = estate->running_queue.back();
|
||||
|
||||
// Find the last alive request state entry, which is what we want to preempt.
|
||||
RequestState rstate = estate->GetRequestState(request);
|
||||
int preempt_rstate_idx = -1;
|
||||
for (int i = static_cast<int>(rstate->entries.size()) - 1; i >= 0; --i) {
|
||||
if (rstate->entries[i]->status == RequestStateStatus::kAlive) {
|
||||
preempt_rstate_idx = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
TVM_FFI_ICHECK_NE(preempt_rstate_idx, -1);
|
||||
RequestStateEntry rsentry = rstate->entries[preempt_rstate_idx];
|
||||
if (estate->disaggregation) {
|
||||
AbortRequestImpl(estate, models, request->id, "preempt");
|
||||
return rsentry;
|
||||
}
|
||||
// When the request state entry still has pending inputs,
|
||||
// it means the request is still in the waiting queue.
|
||||
bool partially_alive = !rsentry->mstates[0]->inputs.empty();
|
||||
|
||||
// Remove from models.
|
||||
// - Clear model speculation draft.
|
||||
// - Update `inputs` for future prefill.
|
||||
RECORD_EVENT(trace_recorder, rsentry->request->id, "preempt");
|
||||
rsentry->status = RequestStateStatus::kPending;
|
||||
std::vector<int> draft_token_slots;
|
||||
for (RequestModelState mstate : rsentry->mstates) {
|
||||
if (draft_token_workspace_manager.has_value()) {
|
||||
mstate->RemoveAllDraftTokens(&draft_token_slots);
|
||||
draft_token_workspace_manager.value()->FreeSlots(draft_token_slots);
|
||||
}
|
||||
|
||||
// If the commited tokens of the current model lags behind the
|
||||
// committed tokens of the main model (models[0]), we commit those
|
||||
// new tokens to this model.
|
||||
for (size_t i = mstate->committed_tokens.size();
|
||||
i < rsentry->mstates[0]->committed_tokens.size(); ++i) {
|
||||
mstate->CommitToken(rsentry->mstates[0]->committed_tokens[i]);
|
||||
}
|
||||
|
||||
std::vector<int32_t> committed_token_ids;
|
||||
committed_token_ids.reserve(mstate->committed_tokens.size());
|
||||
for (const SampleResult& committed_token : mstate->committed_tokens) {
|
||||
committed_token_ids.push_back(committed_token.GetTokenId());
|
||||
}
|
||||
mstate->num_prefilled_tokens = 0;
|
||||
|
||||
Array<Data> inputs;
|
||||
if (rsentry->parent_idx == -1) {
|
||||
inputs = request->inputs;
|
||||
if (const auto* token_input = inputs.back().as<TokenDataNode>()) {
|
||||
// Merge the TokenData so that a single time TokenEmbed is needed.
|
||||
std::vector<int> token_ids{token_input->token_ids->data,
|
||||
token_input->token_ids->data + token_input->token_ids.size()};
|
||||
token_ids.insert(token_ids.end(), committed_token_ids.begin(), committed_token_ids.end());
|
||||
inputs.Set(static_cast<int64_t>(inputs.size()) - 1, TokenData(token_ids));
|
||||
} else if (!committed_token_ids.empty()) {
|
||||
inputs.push_back(TokenData(committed_token_ids));
|
||||
}
|
||||
} else if (!committed_token_ids.empty()) {
|
||||
inputs.push_back(TokenData(committed_token_ids));
|
||||
}
|
||||
mstate->inputs = std::move(inputs);
|
||||
mstate->prefilled_inputs.clear();
|
||||
mstate->cached_committed_tokens = 0;
|
||||
mstate->num_tokens_for_next_decode = 0;
|
||||
}
|
||||
if (estate->prefix_cache->HasSequence(rsentry->mstates[0]->internal_id)) {
|
||||
estate->prefix_cache->RecycleSequence(rsentry->mstates[0]->internal_id, /*lazy=*/false);
|
||||
} else {
|
||||
RemoveRequestFromModel(estate, rsentry->mstates[0]->internal_id, models);
|
||||
}
|
||||
// Since the sequence has been removed from model, assign a new sequence ID.
|
||||
int64_t new_seq_id = estate->id_manager.GetNewId();
|
||||
for (RequestModelState mstate : rsentry->mstates) {
|
||||
mstate->internal_id = new_seq_id;
|
||||
}
|
||||
|
||||
if (preempt_rstate_idx == 0) {
|
||||
// Remove from running queue.
|
||||
estate->running_queue.erase(estate->running_queue.end() - 1);
|
||||
}
|
||||
if (!partially_alive && preempt_rstate_idx == static_cast<int>(rstate->entries.size()) - 1) {
|
||||
// Add to the front of waiting queue.
|
||||
estate->waiting_queue.insert(estate->waiting_queue.begin(), request);
|
||||
}
|
||||
estate->running_rsentries_changed = true;
|
||||
return rsentry;
|
||||
}
|
||||
|
||||
std::pair<Tensor, std::vector<SampleResult>> ApplyLogitProcessorAndSample(
|
||||
const LogitProcessor& logit_processor, const Sampler& sampler, const Tensor& logits,
|
||||
const Array<GenerationConfig>& generation_cfg, const Array<String>& request_ids,
|
||||
const Array<RequestModelState>& mstates, const std::vector<RandomGenerator*>& rngs,
|
||||
const std::vector<int>& sample_indices, const Array<GenerationConfig>& child_generation_cfg,
|
||||
const Array<String>& child_request_ids, const std::vector<int>& child_sample_indices) {
|
||||
// - Update logits.
|
||||
logit_processor->InplaceUpdateLogits(logits, generation_cfg, mstates, request_ids);
|
||||
|
||||
// - Compute probability distributions.
|
||||
Tensor probs_on_device =
|
||||
logit_processor->ComputeProbsFromLogits(logits, generation_cfg, request_ids);
|
||||
|
||||
// - Sample tokens.
|
||||
Tensor renormalized_probs = sampler->BatchRenormalizeProbsByTopP(probs_on_device, sample_indices,
|
||||
request_ids, generation_cfg);
|
||||
std::vector<SampleResult> sample_results = sampler->BatchSampleTokensWithProbAfterTopP(
|
||||
renormalized_probs, child_sample_indices, child_request_ids, child_generation_cfg, rngs);
|
||||
return {std::move(probs_on_device), std::move(sample_results)};
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,119 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/action_commons.h
|
||||
* \brief Common functions that may be used in multiple EngineActions.
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_ENGINE_ACTIONS_ACTION_COMMONS_H_
|
||||
#define MLC_LLM_SERVE_ENGINE_ACTIONS_ACTION_COMMONS_H_
|
||||
|
||||
#include <tvm/ffi/container/array.h>
|
||||
|
||||
#include "../../tokenizers/tokenizers.h"
|
||||
#include "../draft_token_workspace_manager.h"
|
||||
#include "../engine.h"
|
||||
#include "../engine_state.h"
|
||||
#include "../event_trace_recorder.h"
|
||||
#include "../model.h"
|
||||
#include "action.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
|
||||
/*! \brief Create the engine actions based on engine config. */
|
||||
Array<EngineAction> CreateEngineActions(Array<Model> models, EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
LogitProcessor logit_processor, Sampler sampler,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
Tokenizer tokenizer,
|
||||
Optional<EventTraceRecorder> trace_recorder,
|
||||
FRequestStreamCallback request_stream_callback,
|
||||
Device device);
|
||||
|
||||
/*!
|
||||
* \brief Remove the given request from models.
|
||||
* \param estate The engine state to update after removal.
|
||||
* \param req_internal_id The internal id of the request to remove.
|
||||
* \param models The models to remove the given request from.
|
||||
*/
|
||||
void RemoveRequestFromModel(EngineState estate, int64_t req_internal_id,
|
||||
const Array<Model>& models);
|
||||
|
||||
/*!
|
||||
* \brief The request post-processing after an engine action step.
|
||||
* It includes
|
||||
* - invoke the request function callback to return new generated tokens,
|
||||
* - update the engine state for finished requests.
|
||||
* \note This function may remove requests from the `running_queue`.
|
||||
* \param requests The requests to process.
|
||||
* \param estate The engine state.
|
||||
* \param models The models to remove the finished from.
|
||||
* \param tokenizer The tokenizer for logprob process.
|
||||
* \param request_stream_callback The request stream callback function.
|
||||
* \param max_single_sequence_length The max single sequence length to help decide
|
||||
* \param draft_token_workspace_manager The draft token workspace manager.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* if a request is finished.
|
||||
*/
|
||||
void ActionStepPostProcess(Array<Request> requests, EngineState estate, const Array<Model>& models,
|
||||
const Tokenizer& tokenizer,
|
||||
FRequestStreamCallback request_stream_callback,
|
||||
int64_t max_single_sequence_length,
|
||||
Optional<DraftTokenWorkspaceManager> draft_token_workspace_manager,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
|
||||
/*!
|
||||
* \brief Preempt the last running request state entry from `running_queue`.
|
||||
* If all entries of the selected request have been preempted,
|
||||
* remove it from running request.
|
||||
* If it is not in the waiting request queue, add it to the waiting queue.
|
||||
* \param estate The engine state to update due to preemption.
|
||||
* \param models The models to remove preempted requests from.
|
||||
* \param draft_token_workspace_manager The draft token workspace manager for requests. Must be
|
||||
* provided if speculative decoding is enabled.
|
||||
* \param trace_recorder The event trace recorder for requests.
|
||||
* \return The preempted request state.
|
||||
*/
|
||||
RequestStateEntry PreemptLastRunningRequestStateEntry(
|
||||
EngineState estate, const Array<Model>& models,
|
||||
Optional<DraftTokenWorkspaceManager> draft_token_workspace_manager,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
|
||||
/*!
|
||||
* \brief Apply the logit processor to the logits and sample one token for each request.
|
||||
*
|
||||
* Both the parent request configurations and the child request configurations need to be provided.
|
||||
* The parent request configurations are used to process the logits, normalize the probabilities.
|
||||
* The child request configurations are used to sample the tokens.
|
||||
*
|
||||
* When the request doesn't have children, the parent and child configurations are the same.
|
||||
*
|
||||
* \param logit_processor The logit processor to apply.
|
||||
* \param sampler The sampler to sample tokens.
|
||||
* \param logits The logits to process.
|
||||
* \param generation_cfg The generation configurations of the requests.
|
||||
* \param request_ids The request ids.
|
||||
* \param mstates The model states of the requests.
|
||||
* \param rngs The random generators of the requests.
|
||||
* \param sample_indices The indices of the requests to sample.
|
||||
* \param child_generation_cfg The generation configurations of the child requests.
|
||||
* \param child_request_ids The request ids of the child requests.
|
||||
* \param child_sample_indices The indices of the child requests to sample.
|
||||
* \return The processed logits and the sampled results.
|
||||
*/
|
||||
std::pair<Tensor, std::vector<SampleResult>> ApplyLogitProcessorAndSample(
|
||||
const LogitProcessor& logit_processor, const Sampler& sampler, const Tensor& logits,
|
||||
const Array<GenerationConfig>& generation_cfg, const Array<String>& request_ids,
|
||||
const Array<RequestModelState>& mstates, const std::vector<RandomGenerator*>& rngs,
|
||||
const std::vector<int>& sample_indices, const Array<GenerationConfig>& child_generation_cfg,
|
||||
const Array<String>& child_request_ids, const std::vector<int>& child_sample_indices);
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_ENGINE_ACTIONS_ACTION_COMMONS_H_
|
||||
@@ -0,0 +1,88 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/auto_spec_decode.cc
|
||||
*/
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "../config.h"
|
||||
#include "action.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
/*!
|
||||
* \brief The action that first makes a decision on whether to run speculative
|
||||
* decoding or normal mode batch decode, and then runs the selected actions.
|
||||
*/
|
||||
class AutoSpecDecodeActionObj : public EngineActionObj {
|
||||
public:
|
||||
explicit AutoSpecDecodeActionObj(Array<EngineAction> spec_decode_actions,
|
||||
Array<EngineAction> batch_decode_actions,
|
||||
EngineConfig engine_config)
|
||||
: spec_decode_actions_(std::move(spec_decode_actions)),
|
||||
batch_decode_actions_(std::move(batch_decode_actions)),
|
||||
engine_config_(std::move(engine_config)) {}
|
||||
|
||||
Array<Request> Step(EngineState estate) final {
|
||||
int num_running_rsentries = estate->GetRunningRequestStateEntries().size();
|
||||
if (num_running_rsentries == 0) {
|
||||
return {};
|
||||
}
|
||||
|
||||
// Calculate the draft length to use for the next round decode.
|
||||
estate->spec_draft_length = CalculateDraftLength(estate, num_running_rsentries);
|
||||
TVM_FFI_ICHECK_GE(estate->spec_draft_length, 0);
|
||||
Array<Request> processed_requests;
|
||||
// Use speculative decoding when the computed draft length is positive.
|
||||
// Otherwise use normal mode batch decode.
|
||||
Array<EngineAction> actions =
|
||||
estate->spec_draft_length > 0 ? spec_decode_actions_ : batch_decode_actions_;
|
||||
for (EngineAction action : actions) {
|
||||
processed_requests = action->Step(estate);
|
||||
}
|
||||
|
||||
// Reset the draft length.
|
||||
estate->spec_draft_length = 0;
|
||||
return processed_requests;
|
||||
}
|
||||
|
||||
private:
|
||||
int CalculateDraftLength(EngineState estate, int num_running_rsentries) {
|
||||
// Right now we use the fixed table to select the draft length (only based on
|
||||
// the batch size). We will follow up to adopt powerful draft length selection.
|
||||
int draft_length = 0;
|
||||
if (num_running_rsentries < 10) {
|
||||
draft_length = 4;
|
||||
} else if (num_running_rsentries < 20) {
|
||||
draft_length = 3;
|
||||
} else if (num_running_rsentries < 30) {
|
||||
draft_length = 2;
|
||||
} else {
|
||||
draft_length = 0;
|
||||
}
|
||||
|
||||
int effective_batch_size = num_running_rsentries * (draft_length + 1);
|
||||
return effective_batch_size > engine_config_->max_num_sequence ? 0 : draft_length;
|
||||
}
|
||||
|
||||
/*! \brief The speculative decode actions. */
|
||||
Array<EngineAction> spec_decode_actions_;
|
||||
/*! \brief The normal mode decode actions. */
|
||||
Array<EngineAction> batch_decode_actions_;
|
||||
/*! \brief The engine config. */
|
||||
EngineConfig engine_config_;
|
||||
};
|
||||
|
||||
EngineAction EngineAction::AutoSpecDecode(std::vector<EngineAction> spec_decode_actions_,
|
||||
std::vector<EngineAction> batch_decode_actions_,
|
||||
EngineConfig engine_config) {
|
||||
return EngineAction(tvm::ffi::make_object<AutoSpecDecodeActionObj>(
|
||||
Array<EngineAction>(spec_decode_actions_), Array<EngineAction>(batch_decode_actions_),
|
||||
std::move(engine_config)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,329 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/batch_decode.cc
|
||||
*/
|
||||
|
||||
#include <tvm/support/cuda/nvtx.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "../../support/random.h"
|
||||
#include "../config.h"
|
||||
#include "../model.h"
|
||||
#include "../sampler/sampler.h"
|
||||
#include "action.h"
|
||||
#include "action_commons.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::support::NVTXScopedRange;
|
||||
|
||||
/*!
|
||||
* \brief The action that runs one-step decode for requests in the
|
||||
* `running_queue` of engine state. Preempt low-priority requests
|
||||
* accordingly when it is impossible to decode all the running requests.
|
||||
* \note The BatchDecode action **does not** take effect for speculative
|
||||
* decoding scenarios where there are multiple models. For speculative
|
||||
* decoding in the future, we will use other specific actions.
|
||||
*/
|
||||
class BatchDecodeActionObj : public EngineActionObj {
|
||||
public:
|
||||
explicit BatchDecodeActionObj(Array<Model> models, Tokenizer tokenizer,
|
||||
LogitProcessor logit_processor, Sampler sampler,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder)
|
||||
: models_(std::move(models)),
|
||||
tokenizer_(std::move(tokenizer)),
|
||||
logit_processor_(std::move(logit_processor)),
|
||||
sampler_(std::move(sampler)),
|
||||
engine_config_(std::move(engine_config)),
|
||||
trace_recorder_(std::move(trace_recorder)) {}
|
||||
|
||||
Array<Request> Step(EngineState estate) final {
|
||||
// - Do not run decode when there is no running request.
|
||||
if (estate->running_queue.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
// Preempt request state entries when decode cannot apply.
|
||||
std::vector<RequestStateEntry> running_rsentries;
|
||||
{
|
||||
NVTXScopedRange nvtx_scope("BatchDecode getting requests");
|
||||
running_rsentries = estate->GetRunningRequestStateEntries();
|
||||
while (!CanDecode(running_rsentries.size())) {
|
||||
if (estate->prefix_cache->TryFreeMemory()) continue;
|
||||
RequestStateEntry preempted =
|
||||
PreemptLastRunningRequestStateEntry(estate, models_, std::nullopt, trace_recorder_);
|
||||
if (preempted.same_as(running_rsentries.back())) {
|
||||
running_rsentries.pop_back();
|
||||
}
|
||||
}
|
||||
while (running_rsentries.size() >
|
||||
std::min(static_cast<int64_t>(engine_config_->max_num_sequence),
|
||||
engine_config_->prefill_chunk_size)) {
|
||||
running_rsentries.pop_back();
|
||||
}
|
||||
}
|
||||
|
||||
auto tstart = std::chrono::high_resolution_clock::now();
|
||||
|
||||
// NOTE: Right now we only support decode all the running request states at a time.
|
||||
int num_rsentries = running_rsentries.size();
|
||||
TVM_FFI_ICHECK_GT(num_rsentries, 0)
|
||||
<< "There should be at least one request state entry that can run decode. "
|
||||
"Possible failure reason: none of the prefill phase of the running requests is finished";
|
||||
TVM_FFI_ICHECK_LE(num_rsentries, engine_config_->max_num_sequence)
|
||||
<< "The number of running requests exceeds the max number of sequence in EngineConfig. "
|
||||
"Possible failure reason: the prefill action allows new sequence in regardless of the "
|
||||
"max num sequence.";
|
||||
// Collect
|
||||
// - the last committed token,
|
||||
// - the request id,
|
||||
// - the generation config,
|
||||
// - the random number generator,
|
||||
// of each request state entry.
|
||||
std::vector<int> input_tokens;
|
||||
std::vector<int> lengths;
|
||||
Array<String> request_ids;
|
||||
std::vector<int64_t> request_internal_ids;
|
||||
Array<RequestModelState> mstates;
|
||||
Array<GenerationConfig> generation_cfg;
|
||||
std::vector<RandomGenerator*> rngs;
|
||||
|
||||
input_tokens.reserve(num_rsentries);
|
||||
request_ids.reserve(num_rsentries);
|
||||
request_internal_ids.reserve(num_rsentries);
|
||||
mstates.reserve(num_rsentries);
|
||||
generation_cfg.reserve(num_rsentries);
|
||||
rngs.reserve(num_rsentries);
|
||||
|
||||
{
|
||||
NVTXScopedRange nvtx_scope("BatchDecode setting batch info");
|
||||
for (const RequestStateEntry& rsentry : running_rsentries) {
|
||||
auto mstate = rsentry->mstates[0];
|
||||
TVM_FFI_ICHECK(mstate->num_tokens_for_next_decode > 0 &&
|
||||
mstate->num_tokens_for_next_decode <=
|
||||
static_cast<int>(mstate->committed_tokens.size()));
|
||||
|
||||
for (auto begin = mstate->committed_tokens.end() - mstate->num_tokens_for_next_decode;
|
||||
begin != mstate->committed_tokens.end(); ++begin) {
|
||||
input_tokens.push_back(begin->GetTokenId());
|
||||
}
|
||||
|
||||
lengths.push_back(mstate->num_tokens_for_next_decode);
|
||||
mstate->num_tokens_for_next_decode = 0;
|
||||
|
||||
request_ids.push_back(rsentry->request->id);
|
||||
request_internal_ids.push_back(mstate->internal_id);
|
||||
mstates.push_back(mstate);
|
||||
generation_cfg.push_back(rsentry->request->generation_cfg);
|
||||
rngs.push_back(&rsentry->rng);
|
||||
}
|
||||
}
|
||||
|
||||
// - Compute embeddings.
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start embedding");
|
||||
ObjectRef embeddings =
|
||||
models_[0]->TokenEmbed({Shape(input_tokens.begin(), input_tokens.end())});
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish embedding");
|
||||
|
||||
// - Invoke model decode.
|
||||
// If every request only requires to process one token, batch decode kernel is called.
|
||||
// Otherwise, batch prefill kernel is called.
|
||||
bool is_every_request_single_token =
|
||||
std::all_of(lengths.begin(), lengths.end(), [](int len) { return len == 1; });
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start decode");
|
||||
Tensor logits;
|
||||
if (is_every_request_single_token) {
|
||||
logits = models_[0]->BatchDecode(embeddings, request_internal_ids);
|
||||
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[0], num_rsentries);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[1], 1);
|
||||
} else {
|
||||
logits = models_[0]->BatchPrefill(embeddings, request_internal_ids, lengths);
|
||||
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[0], 1);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[1], num_rsentries);
|
||||
}
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish decode");
|
||||
|
||||
// - Update logits.
|
||||
logits = logits.CreateView({num_rsentries, logits->shape[2]}, logits->dtype);
|
||||
logit_processor_->InplaceUpdateLogits(logits, generation_cfg, mstates, request_ids);
|
||||
|
||||
// - Compute probability distributions.
|
||||
Tensor probs_on_device =
|
||||
logit_processor_->ComputeProbsFromLogits(logits, generation_cfg, request_ids);
|
||||
|
||||
// - Commit the prefix cache changes from previous round of action.
|
||||
// Note: we commit prefix cache changes here to overlap this commit with the GPU execution.
|
||||
estate->prefix_cache->CommitSequenceExtention();
|
||||
|
||||
// - Sample tokens.
|
||||
// Fill range [0, num_rsentries) into `sample_indices`.
|
||||
std::vector<int> sample_indices(num_rsentries);
|
||||
std::iota(sample_indices.begin(), sample_indices.end(), 0);
|
||||
Tensor renormalized_probs = sampler_->BatchRenormalizeProbsByTopP(
|
||||
probs_on_device, sample_indices, request_ids, generation_cfg);
|
||||
std::vector<SampleResult> sample_results = sampler_->BatchSampleTokensWithProbAfterTopP(
|
||||
renormalized_probs, sample_indices, request_ids, generation_cfg, rngs);
|
||||
TVM_FFI_ICHECK_EQ(sample_results.size(), num_rsentries);
|
||||
|
||||
// - Update the committed tokens of states.
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
auto mstate = mstates[i];
|
||||
|
||||
if (!mstate->require_retokenization_in_next_decode) {
|
||||
mstates[i]->CommitToken(sample_results[i]);
|
||||
// live update the output metrics
|
||||
running_rsentries[i]->rstate->metrics.completion_tokens += 1;
|
||||
} else {
|
||||
// Retokenize and commit tokens.
|
||||
CommitTokenMayRetokenize(running_rsentries[i], mstate, sample_results[i]);
|
||||
mstate->require_retokenization_in_next_decode = false;
|
||||
}
|
||||
|
||||
running_rsentries[i]->rstate->metrics.decode_tokens += lengths[i];
|
||||
}
|
||||
|
||||
double elapsed_time;
|
||||
{
|
||||
NVTXScopedRange nvtx_scope("BatchDecode get time");
|
||||
auto tend = std::chrono::high_resolution_clock::now();
|
||||
elapsed_time = static_cast<double>((tend - tstart).count()) / 1e9;
|
||||
}
|
||||
estate->metrics.engine_decode_time_sum += elapsed_time;
|
||||
estate->metrics.UpdateDecodeTimeByBatchSize(num_rsentries, elapsed_time);
|
||||
|
||||
return estate->running_queue;
|
||||
}
|
||||
|
||||
private:
|
||||
/*! \brief Check if the input request state entries can be decoded under conditions. */
|
||||
bool CanDecode(int num_rsentries) {
|
||||
int num_available_pages = models_[0]->GetNumAvailablePages();
|
||||
return num_rsentries <= num_available_pages;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Retokenize the past tokens with a new token.
|
||||
* \param mstate The model state.
|
||||
* \param token_id The new token id.
|
||||
* \param max_rollback_tokens The maximum number of tokens to rollback.
|
||||
* \return The number of tokens to rollback and the new tokens.
|
||||
*/
|
||||
std::pair<int, std::vector<int32_t>> RetokenizeWithNewToken(RequestModelState mstate,
|
||||
int32_t token_id,
|
||||
int max_rollback_tokens) {
|
||||
// Step 1. Get past tokens
|
||||
// past_tokens = mstate[-max_rollback_tokens:]
|
||||
// past_string = detokenize(past_tokens)
|
||||
const auto& token_table = tokenizer_->PostProcessedTokenTable();
|
||||
std::vector<int32_t> past_tokens;
|
||||
std::string past_string;
|
||||
auto past_begin_it = mstate->committed_tokens.size() >= max_rollback_tokens
|
||||
? mstate->committed_tokens.end() - max_rollback_tokens
|
||||
: mstate->committed_tokens.begin();
|
||||
for (auto it = past_begin_it; it != mstate->committed_tokens.end(); ++it) {
|
||||
past_tokens.push_back(it->GetTokenId());
|
||||
past_string += token_table[it->GetTokenId()];
|
||||
}
|
||||
|
||||
// Step 2. Retokenize
|
||||
// Compare tokenize(past_string + new_string) and past_tokens
|
||||
auto new_tokens = tokenizer_->EncodeNoPrependSpace(past_string + token_table[token_id]);
|
||||
|
||||
int first_differ_idx = past_tokens.size();
|
||||
for (int i = 0; i < static_cast<int>(past_tokens.size()); ++i) {
|
||||
if (i == static_cast<int>(new_tokens.size()) || past_tokens[i] != new_tokens[i]) {
|
||||
first_differ_idx = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return {past_tokens.size() - first_differ_idx,
|
||||
std::vector<int32_t>(new_tokens.begin() + first_differ_idx, new_tokens.end())};
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Commit the token and may retokenize the past tokens.
|
||||
* \param rsentry The request state entry.
|
||||
* \param mstate The model state.
|
||||
* \param sample_result The sampled token.
|
||||
*/
|
||||
void CommitTokenMayRetokenize(RequestStateEntry rsentry, RequestModelState mstate,
|
||||
const SampleResult& sample_result) {
|
||||
auto generation_cfg = rsentry->request->generation_cfg;
|
||||
// 1. If EOS token is generated, jump commit it
|
||||
if (!generation_cfg->debug_config.ignore_eos &&
|
||||
std::any_of(generation_cfg->stop_token_ids.begin(), generation_cfg->stop_token_ids.end(),
|
||||
[&](int32_t token) { return token == sample_result.GetTokenId(); })) {
|
||||
mstate->CommitToken(sample_result);
|
||||
rsentry->rstate->metrics.completion_tokens += 1;
|
||||
return;
|
||||
}
|
||||
|
||||
// 2. Check retokenization
|
||||
const auto& committed_tokens = mstate->committed_tokens;
|
||||
auto [rollback_cnt, new_tokens] =
|
||||
RetokenizeWithNewToken(mstate, sample_result.GetTokenId(), MAX_ROLLBACK_TOKENS_);
|
||||
|
||||
// 3. Handle output when retokenization happens
|
||||
if (rollback_cnt >
|
||||
static_cast<int>(committed_tokens.size()) - rsentry->next_callback_token_pos) {
|
||||
const auto& token_table = tokenizer_->PostProcessedTokenTable();
|
||||
for (auto i = rsentry->next_callback_token_pos; i < committed_tokens.size(); ++i) {
|
||||
auto token_id = committed_tokens[i].GetTokenId();
|
||||
rsentry->extra_prefix_string += token_table[token_id];
|
||||
}
|
||||
rsentry->extra_prefix_string += token_table[sample_result.GetTokenId()];
|
||||
rsentry->next_callback_token_pos = static_cast<int>(committed_tokens.size()) - rollback_cnt +
|
||||
static_cast<int>(new_tokens.size());
|
||||
}
|
||||
|
||||
if (rollback_cnt > 0) {
|
||||
mstate->RollbackTokens(rollback_cnt);
|
||||
models_[0]->PopNFromKVCache(mstate->internal_id, rollback_cnt);
|
||||
}
|
||||
|
||||
for (auto token_id : new_tokens) {
|
||||
mstate->CommitToken({{token_id, 1.0}, {}});
|
||||
}
|
||||
|
||||
rsentry->rstate->metrics.completion_tokens +=
|
||||
static_cast<int>(new_tokens.size()) - rollback_cnt;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief The model to run decode in. When there are multiple
|
||||
* models, the `Step` function of the created action will not take effect.
|
||||
*/
|
||||
Array<Model> models_;
|
||||
/*! \brief The tokenizer of the engine. */
|
||||
Tokenizer tokenizer_;
|
||||
/*! \brief The logit processor. */
|
||||
LogitProcessor logit_processor_;
|
||||
/*! \brief The sampler to sample new tokens. */
|
||||
Sampler sampler_;
|
||||
/*! \brief The engine config. */
|
||||
EngineConfig engine_config_;
|
||||
/*! \brief Event trace recorder. */
|
||||
Optional<EventTraceRecorder> trace_recorder_;
|
||||
/*! \brief The maximum number of tokens to retokenize and may be rolled back. */
|
||||
const int MAX_ROLLBACK_TOKENS_ = 10;
|
||||
};
|
||||
|
||||
EngineAction EngineAction::BatchDecode(Array<Model> models, Tokenizer tokenizer,
|
||||
LogitProcessor logit_processor, Sampler sampler,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder) {
|
||||
return EngineAction(tvm::ffi::make_object<BatchDecodeActionObj>(
|
||||
std::move(models), std::move(tokenizer), std::move(logit_processor), std::move(sampler),
|
||||
std::move(engine_config), std::move(trace_recorder)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,407 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/batch_draft.cc
|
||||
*/
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "../config.h"
|
||||
#include "../model.h"
|
||||
#include "../sampler/sampler.h"
|
||||
#include "action.h"
|
||||
#include "action_commons.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
/*!
|
||||
* \brief The action that runs draft proposal for requests in the
|
||||
* `running_queue` of engine state. Preempt low-priority requests
|
||||
* accordingly when it is impossible to decode all the running requests.
|
||||
*/
|
||||
class BatchDraftActionObj : public EngineActionObj {
|
||||
public:
|
||||
explicit BatchDraftActionObj(Array<Model> models, LogitProcessor logit_processor, Sampler sampler,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder)
|
||||
: models_(std::move(models)),
|
||||
logit_processor_(std::move(logit_processor)),
|
||||
sampler_(std::move(sampler)),
|
||||
model_workspaces_(std::move(model_workspaces)),
|
||||
draft_token_workspace_manager_(std::move(draft_token_workspace_manager)),
|
||||
engine_config_(std::move(engine_config)),
|
||||
trace_recorder_(std::move(trace_recorder)) {}
|
||||
|
||||
Array<Request> Step(EngineState estate) final {
|
||||
// - Only run spec decode when there are two models (llm+ssm) and >=1 running requests.
|
||||
if (models_.size() != 2 || estate->running_queue.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
// Preempt request state entries when decode cannot apply.
|
||||
std::vector<RequestStateEntry> running_rsentries = estate->GetRunningRequestStateEntries();
|
||||
while (!CanDecode(running_rsentries.size())) {
|
||||
if (estate->prefix_cache->TryFreeMemory()) continue;
|
||||
RequestStateEntry preempted = PreemptLastRunningRequestStateEntry(
|
||||
estate, models_, draft_token_workspace_manager_, trace_recorder_);
|
||||
if (preempted.same_as(running_rsentries.back())) {
|
||||
running_rsentries.pop_back();
|
||||
}
|
||||
}
|
||||
while (running_rsentries.size() * (engine_config_->spec_draft_length + 1) >
|
||||
std::min(static_cast<int64_t>(engine_config_->max_num_sequence),
|
||||
engine_config_->prefill_chunk_size)) {
|
||||
running_rsentries.pop_back();
|
||||
}
|
||||
|
||||
auto tstart = std::chrono::high_resolution_clock::now();
|
||||
|
||||
int num_rsentries = running_rsentries.size();
|
||||
TVM_FFI_ICHECK_GT(num_rsentries, 0)
|
||||
<< "There should be at least one request state entry that can run decode. "
|
||||
"Possible failure reason: none of the prefill phase of the running requests is finished";
|
||||
TVM_FFI_ICHECK_LE(num_rsentries, engine_config_->max_num_sequence)
|
||||
<< "The number of running requests exceeds the max number of sequence in EngineConfig. "
|
||||
"Possible failure reason: the prefill action allows new sequence in regardless of the "
|
||||
"max num sequence.";
|
||||
Array<String> request_ids;
|
||||
std::vector<int64_t> request_internal_ids;
|
||||
Array<String> request_ids_per_leaf_node;
|
||||
Array<GenerationConfig> generation_cfg;
|
||||
Array<GenerationConfig> generation_cfg_for_logitproc;
|
||||
std::vector<RandomGenerator*> rngs;
|
||||
std::vector<std::vector<int>> draft_token_indices;
|
||||
// Number of input tokens for each request. Each request can have multiple leaf tokens for the
|
||||
// next forward when multiple tokens are drafted.
|
||||
std::vector<int> cum_num_tokens;
|
||||
std::vector<int64_t> token_tree_parent_ptr;
|
||||
request_ids.reserve(num_rsentries);
|
||||
request_internal_ids.reserve(num_rsentries);
|
||||
generation_cfg.reserve(num_rsentries);
|
||||
generation_cfg_for_logitproc.reserve(num_rsentries);
|
||||
draft_token_indices.reserve(num_rsentries);
|
||||
cum_num_tokens.reserve(num_rsentries + 1);
|
||||
for (const RequestStateEntry& rsentry : running_rsentries) {
|
||||
request_ids.push_back(rsentry->request->id);
|
||||
request_internal_ids.push_back(rsentry->mstates[0]->internal_id);
|
||||
}
|
||||
|
||||
TVM_FFI_ICHECK_GT(estate->spec_draft_length, 0)
|
||||
<< "The speculative decoding draft length must be positive.";
|
||||
// The first model doesn't get involved in draft proposal.
|
||||
for (int model_id = 1; model_id < static_cast<int>(models_.size()); ++model_id) {
|
||||
// Collect
|
||||
// - the last committed token,
|
||||
// - the request model state of each request,
|
||||
// - the number of tokens for each request to send into the model (it may
|
||||
// be more than one if the draft model is lagging behind the main model, when
|
||||
// the engine switches from normal batch decode mode to speculative decoding mode).
|
||||
std::vector<int> input_tokens;
|
||||
Array<RequestModelState> mstates;
|
||||
std::vector<int> input_lengths;
|
||||
input_tokens.reserve(num_rsentries);
|
||||
mstates.reserve(num_rsentries);
|
||||
input_lengths.reserve(num_rsentries);
|
||||
for (const RequestStateEntry& rsentry : running_rsentries) {
|
||||
mstates.push_back(rsentry->mstates[model_id]);
|
||||
}
|
||||
// "Draft length" rounds of draft proposal.
|
||||
for (int draft_id = 0; draft_id < estate->spec_draft_length; ++draft_id) {
|
||||
auto tdraft_start = std::chrono::high_resolution_clock::now();
|
||||
// prepare new input tokens
|
||||
input_tokens.clear();
|
||||
input_lengths.clear();
|
||||
token_tree_parent_ptr.clear();
|
||||
generation_cfg.clear();
|
||||
generation_cfg_for_logitproc.clear();
|
||||
rngs.clear();
|
||||
cum_num_tokens.clear();
|
||||
cum_num_tokens.push_back(0);
|
||||
request_ids_per_leaf_node.clear();
|
||||
std::vector<int> draft_token_parent_idx;
|
||||
draft_token_indices.clear();
|
||||
|
||||
if (draft_id == 0) {
|
||||
// Compute the total length that needs to be processed by the draft model,
|
||||
// including the lagging-behind part of hte draft model.
|
||||
// When the total length to be processed is larger than the prefill chunk
|
||||
// size, we must do the prefill with multiple rounds by chunk.
|
||||
int total_length = 0;
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
TVM_FFI_ICHECK_LE(mstates[i]->committed_tokens.size(),
|
||||
running_rsentries[i]->mstates[0]->committed_tokens.size());
|
||||
total_length += running_rsentries[i]->mstates[0]->committed_tokens.size() -
|
||||
mstates[i]->committed_tokens.size() + 1;
|
||||
}
|
||||
if (total_length > engine_config_->prefill_chunk_size) {
|
||||
PrefillLaggedTokensByChunk(mstates, running_rsentries, models_[model_id],
|
||||
total_length - engine_config_->prefill_chunk_size);
|
||||
}
|
||||
}
|
||||
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
int num_leaf_nodes = 0;
|
||||
// Starting from last committed tokens
|
||||
if (draft_id == 0) {
|
||||
TVM_FFI_ICHECK_LE(mstates[i]->committed_tokens.size(),
|
||||
running_rsentries[i]->mstates[0]->committed_tokens.size());
|
||||
TVM_FFI_ICHECK_EQ(mstates[i]->num_tokens_for_next_decode, 1);
|
||||
input_tokens.push_back(mstates[i]->committed_tokens.back().GetTokenId());
|
||||
input_lengths.push_back(running_rsentries[i]->mstates[0]->committed_tokens.size() -
|
||||
mstates[i]->committed_tokens.size() + 1);
|
||||
for (size_t j = mstates[i]->committed_tokens.size();
|
||||
j < running_rsentries[i]->mstates[0]->committed_tokens.size(); ++j) {
|
||||
// This draft model is lagging behind the main model.
|
||||
// It may happen when the engine just switches from the normal batch decode
|
||||
// mode to the speculative decoding mode.
|
||||
// In this case, we need to prefill the misaligned tokens into the draft model.
|
||||
mstates[i]->CommitToken(running_rsentries[i]->mstates[0]->committed_tokens[j]);
|
||||
input_tokens.push_back(
|
||||
running_rsentries[i]->mstates[0]->committed_tokens[j].GetTokenId());
|
||||
}
|
||||
mstates[i]->num_tokens_for_next_decode = 0;
|
||||
draft_token_indices.emplace_back(std::vector<int>{-1});
|
||||
rngs.push_back(&running_rsentries[i]->rng);
|
||||
draft_token_parent_idx.push_back(-1);
|
||||
request_ids_per_leaf_node.push_back(request_ids[i]);
|
||||
num_leaf_nodes = 1;
|
||||
cum_num_tokens.push_back(cum_num_tokens.back() + 1);
|
||||
} else {
|
||||
TVM_FFI_ICHECK_EQ(mstates[i]->committed_tokens.size(),
|
||||
running_rsentries[i]->mstates[0]->committed_tokens.size());
|
||||
TVM_FFI_ICHECK(!mstates[i]->draft_output_tokens.empty());
|
||||
draft_token_indices.emplace_back(std::vector<int>{});
|
||||
// Get all leaf nodes
|
||||
for (int j = 0; j < static_cast<int>(mstates[i]->draft_output_tokens.size()); ++j) {
|
||||
if (mstates[i]->draft_token_first_child_idx[j] == -1) {
|
||||
int64_t parent_idx = mstates[i]->draft_token_parent_idx[j];
|
||||
token_tree_parent_ptr.push_back(parent_idx);
|
||||
input_tokens.push_back(mstates[i]->draft_output_tokens[j].GetTokenId());
|
||||
draft_token_indices.back().push_back(j);
|
||||
rngs.push_back(&running_rsentries[i]->rng);
|
||||
num_leaf_nodes++;
|
||||
request_ids_per_leaf_node.push_back(request_ids[i]);
|
||||
draft_token_parent_idx.push_back(j);
|
||||
}
|
||||
}
|
||||
input_lengths.push_back(num_leaf_nodes);
|
||||
cum_num_tokens.push_back(cum_num_tokens.back() + input_lengths.back());
|
||||
}
|
||||
GenerationConfig generation_cfg_for_draft = [&]() {
|
||||
if (engine_config_->spec_tree_width == 1) {
|
||||
return mstates[i]->request->generation_cfg;
|
||||
}
|
||||
auto spec_generation_cfg = tvm::ffi::make_object<GenerationConfigNode>(
|
||||
*(mstates[i]->request->generation_cfg.get()));
|
||||
spec_generation_cfg->top_logprobs = engine_config_->spec_tree_width;
|
||||
spec_generation_cfg->logprobs = true;
|
||||
spec_generation_cfg->temperature = 1.0;
|
||||
return GenerationConfig(spec_generation_cfg);
|
||||
}();
|
||||
for (int j = 0; j < num_leaf_nodes; ++j) {
|
||||
generation_cfg.push_back(generation_cfg_for_draft);
|
||||
}
|
||||
generation_cfg_for_logitproc.push_back(generation_cfg_for_draft);
|
||||
}
|
||||
|
||||
// - Compute embeddings.
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start proposal embedding");
|
||||
TVM_FFI_ICHECK_LE(input_tokens.size(), engine_config_->prefill_chunk_size);
|
||||
ObjectRef embeddings =
|
||||
models_[model_id]->TokenEmbed({Shape{input_tokens.begin(), input_tokens.end()}});
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish proposal embedding");
|
||||
|
||||
// - Invoke model decode.
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start proposal decode");
|
||||
Tensor logits{nullptr};
|
||||
|
||||
if (input_tokens.size() == num_rsentries) {
|
||||
// Each request entry only has one token to feed into the draft model.
|
||||
logits = models_[model_id]->BatchDecode(embeddings, request_internal_ids);
|
||||
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[0], num_rsentries);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[1], 1);
|
||||
} else if (draft_id == 0) {
|
||||
// There exists some request entry which has more than one token to feed.
|
||||
// It may happen when the engine just switches from the normal batch decode
|
||||
// mode to the speculative decoding mode.
|
||||
logits = models_[model_id]->BatchPrefill(embeddings, request_internal_ids, input_lengths);
|
||||
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[0], 1);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[1], num_rsentries);
|
||||
} else {
|
||||
TVM_FFI_ICHECK_GT(engine_config_->spec_tree_width, 1);
|
||||
logits = models_[model_id]->BatchTreeDecode(embeddings, request_internal_ids,
|
||||
input_lengths, token_tree_parent_ptr);
|
||||
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[0], cum_num_tokens.back());
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[1], 1);
|
||||
}
|
||||
TVM_FFI_ICHECK_EQ(input_lengths.size(), num_rsentries);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish proposal decode");
|
||||
|
||||
// - Update logits.
|
||||
logits = logits.CreateView({cum_num_tokens.back(), logits->shape[2]}, logits->dtype);
|
||||
|
||||
logit_processor_->InplaceUpdateLogits(logits, generation_cfg_for_logitproc, mstates,
|
||||
request_ids, &cum_num_tokens, &mstates,
|
||||
&draft_token_indices);
|
||||
|
||||
// - Compute probability distributions.
|
||||
Tensor probs_on_device = logit_processor_->ComputeProbsFromLogits(
|
||||
logits, generation_cfg_for_logitproc, request_ids, &cum_num_tokens);
|
||||
|
||||
// - Commit the prefix cache changes from previous round of action.
|
||||
// Note: we commit prefix cache changes here to overlap this commit with the GPU execution.
|
||||
estate->prefix_cache->CommitSequenceExtention();
|
||||
|
||||
// - Sample tokens.
|
||||
// Fill range [0, num_rsentries) into `sample_indices`.
|
||||
std::vector<int> sample_indices(cum_num_tokens.back());
|
||||
std::iota(sample_indices.begin(), sample_indices.end(), 0);
|
||||
std::vector<Tensor> prob_dist;
|
||||
Tensor renormalized_probs = sampler_->BatchRenormalizeProbsByTopP(
|
||||
probs_on_device, sample_indices, request_ids_per_leaf_node, generation_cfg);
|
||||
std::vector<SampleResult> sample_results = sampler_->BatchSampleTokensWithProbAfterTopP(
|
||||
renormalized_probs, sample_indices, request_ids_per_leaf_node, generation_cfg, rngs);
|
||||
TVM_FFI_ICHECK_EQ(sample_results.size(), cum_num_tokens.back());
|
||||
|
||||
// - Add draft token to the state.
|
||||
draft_token_workspace_manager_->AllocSlots(cum_num_tokens.back(), &draft_token_slots_);
|
||||
models_[model_id]->ScatterDraftProbs(probs_on_device, draft_token_slots_,
|
||||
&model_workspaces_[0].draft_probs_storage);
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
for (int j = cum_num_tokens[i]; j < cum_num_tokens[i + 1]; ++j) {
|
||||
int parent_idx = draft_token_parent_idx[j];
|
||||
if (engine_config_->spec_tree_width == 1) {
|
||||
mstates[i]->AddDraftToken(sample_results[j], draft_token_slots_[j], parent_idx);
|
||||
continue;
|
||||
}
|
||||
for (int k = 0; k < sample_results[j].top_prob_tokens.size(); ++k) {
|
||||
SampleResult top_k_token{sample_results[j].top_prob_tokens[k]};
|
||||
mstates[i]->AddDraftToken(top_k_token, draft_token_slots_[j], parent_idx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
auto tdraft_end = std::chrono::high_resolution_clock::now();
|
||||
estate->metrics.UpdateDraftTimeByBatchSize(
|
||||
num_rsentries, static_cast<double>((tdraft_end - tdraft_start).count()) / 1e9);
|
||||
}
|
||||
}
|
||||
|
||||
auto tend = std::chrono::high_resolution_clock::now();
|
||||
estate->metrics.engine_decode_time_sum += static_cast<double>((tend - tstart).count()) / 1e9;
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
private:
|
||||
/*! \brief Check if the input requests can be decoded under conditions. */
|
||||
bool CanDecode(int num_rsentries) {
|
||||
// The first model is not involved in draft proposal.
|
||||
for (int model_id = 1; model_id < static_cast<int>(models_.size()); ++model_id) {
|
||||
// Check if the model has enough available pages.
|
||||
int num_available_pages = models_[model_id]->GetNumAvailablePages();
|
||||
if (num_rsentries > num_available_pages) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
void PrefillLaggedTokensByChunk(const Array<RequestModelState>& mstates,
|
||||
const std::vector<RequestStateEntry>& running_rsentries,
|
||||
Model model, int remaining_prefill_length) {
|
||||
int num_rsentries = mstates.size();
|
||||
std::vector<int> input_tokens;
|
||||
std::vector<int64_t> request_internal_ids;
|
||||
std::vector<int> lengths;
|
||||
input_tokens.reserve(engine_config_->prefill_chunk_size);
|
||||
request_internal_ids.reserve(num_rsentries);
|
||||
lengths.reserve(num_rsentries);
|
||||
|
||||
auto f_run_prefill = [&model, &input_tokens, &request_internal_ids, &lengths]() {
|
||||
ObjectRef embeddings = model->TokenEmbed({Shape{input_tokens.begin(), input_tokens.end()}});
|
||||
model->BatchPrefill(embeddings, request_internal_ids, lengths);
|
||||
};
|
||||
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
int prefill_length =
|
||||
std::min({static_cast<int>(running_rsentries[i]->mstates[0]->committed_tokens.size() -
|
||||
mstates[i]->committed_tokens.size()),
|
||||
static_cast<int>(engine_config_->prefill_chunk_size - input_tokens.size()),
|
||||
remaining_prefill_length});
|
||||
if (prefill_length == 0) {
|
||||
// This rsentry is done.
|
||||
continue;
|
||||
}
|
||||
|
||||
TVM_FFI_ICHECK(!mstates[i]->committed_tokens.empty());
|
||||
for (size_t j = mstates[i]->committed_tokens.size();
|
||||
j < running_rsentries[i]->mstates[0]->committed_tokens.size(); ++j) {
|
||||
// Commit the lagging-behind tokens to the draft model.
|
||||
mstates[i]->CommitToken(running_rsentries[i]->mstates[0]->committed_tokens[j - 1]);
|
||||
input_tokens.push_back(
|
||||
running_rsentries[i]->mstates[0]->committed_tokens[j - 1].GetTokenId());
|
||||
}
|
||||
lengths.push_back(prefill_length);
|
||||
request_internal_ids.push_back(running_rsentries[i]->mstates[0]->internal_id);
|
||||
mstates[i]->num_tokens_for_next_decode = 1;
|
||||
remaining_prefill_length -= prefill_length;
|
||||
if (remaining_prefill_length == 0) {
|
||||
// All rsentries are done.
|
||||
break;
|
||||
}
|
||||
|
||||
if (input_tokens.size() == engine_config_->prefill_chunk_size) {
|
||||
// Run prefill if the pending part total length reaches the prefill chunk size.
|
||||
f_run_prefill();
|
||||
input_tokens.clear();
|
||||
request_internal_ids.clear();
|
||||
lengths.clear();
|
||||
--i;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
if (!input_tokens.empty()) {
|
||||
f_run_prefill();
|
||||
}
|
||||
}
|
||||
|
||||
/*! \brief The model to run draft generation in speculative decoding. */
|
||||
Array<Model> models_;
|
||||
/*! \brief The logit processor. */
|
||||
LogitProcessor logit_processor_;
|
||||
/*! \brief The sampler to sample new tokens. */
|
||||
Sampler sampler_;
|
||||
/*! \brief The model workspaces. */
|
||||
std::vector<ModelWorkspace> model_workspaces_;
|
||||
/*! \brief The draft token workspace manager. */
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager_;
|
||||
/*! \brief The engine config. */
|
||||
EngineConfig engine_config_;
|
||||
/*! \brief Event trace recorder. */
|
||||
Optional<EventTraceRecorder> trace_recorder_;
|
||||
/*! \brief Temporary buffer to store the slots of the current draft tokens */
|
||||
std::vector<int> draft_token_slots_;
|
||||
};
|
||||
|
||||
EngineAction EngineAction::BatchDraft(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler, std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder) {
|
||||
return EngineAction(tvm::ffi::make_object<BatchDraftActionObj>(
|
||||
std::move(models), std::move(logit_processor), std::move(sampler),
|
||||
std::move(model_workspaces), std::move(draft_token_workspace_manager),
|
||||
std::move(engine_config), std::move(trace_recorder)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,240 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/batch_verify.cc
|
||||
*/
|
||||
|
||||
#include <tvm/support/cuda/nvtx.h>
|
||||
|
||||
#include <cmath>
|
||||
#include <exception>
|
||||
|
||||
#include "../config.h"
|
||||
#include "../model.h"
|
||||
#include "../sampler/sampler.h"
|
||||
#include "action.h"
|
||||
#include "action_commons.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::support::NVTXScopedRange;
|
||||
|
||||
/*!
|
||||
* \brief The action that runs verification for requests in the
|
||||
* `running_queue` of engine state. Preempt low-priority requests
|
||||
* accordingly when it is impossible to decode all the running requests.
|
||||
*/
|
||||
class BatchJumpForwardActionObj : public EngineActionObj {
|
||||
public:
|
||||
explicit BatchJumpForwardActionObj(Array<Model> models, Tokenizer tokenizer,
|
||||
Optional<EventTraceRecorder> trace_recorder)
|
||||
: models_(std::move(models)),
|
||||
tokenizer_(tokenizer),
|
||||
trace_recorder_(std::move(trace_recorder)) {}
|
||||
|
||||
Array<Request> Step(EngineState estate) final {
|
||||
// - Do not run decode when there are multiple models or no running requests.
|
||||
if (models_.size() > 1 || estate->running_queue.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
// Preempt request state entries when jump-forward decoding cannot apply.
|
||||
std::vector<RequestStateEntry> running_rsentries;
|
||||
{
|
||||
NVTXScopedRange nvtx_scope("BatchJumpForward getting requests");
|
||||
running_rsentries = estate->GetRunningRequestStateEntries();
|
||||
while (!CheckMemForJumpForward(running_rsentries.size())) {
|
||||
if (estate->prefix_cache->TryFreeMemory()) continue;
|
||||
RequestStateEntry preempted =
|
||||
PreemptLastRunningRequestStateEntry(estate, models_, std::nullopt, trace_recorder_);
|
||||
if (preempted.same_as(running_rsentries.back())) {
|
||||
running_rsentries.pop_back();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (running_rsentries.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
auto tstart = std::chrono::high_resolution_clock::now();
|
||||
|
||||
for (auto rsentry : running_rsentries) {
|
||||
if (!CanJumpForward(rsentry)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
auto mstate = rsentry->mstates[0];
|
||||
auto jump_forward_str = mstate->grammar_matcher->FindJumpForwardString();
|
||||
|
||||
if (jump_forward_str.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
auto [rollback_cnt, new_tokens, new_string] =
|
||||
RetokenizeWithNewString(mstate, jump_forward_str, MAX_ROLLBACK_TOKENS_);
|
||||
|
||||
HandleRollback(rsentry, mstate, rollback_cnt, new_tokens, new_string);
|
||||
|
||||
// Commit new tokens (kv cache is handled in the next decode)
|
||||
for (auto token_id : new_tokens) {
|
||||
mstate->CommitToken({{token_id, 1.0}, {}});
|
||||
}
|
||||
|
||||
mstate->require_retokenization_in_next_decode = true;
|
||||
|
||||
// Update metrics
|
||||
rsentry->rstate->metrics.jump_forward_tokens +=
|
||||
std::max(static_cast<int>(new_tokens.size()) - rollback_cnt, 0);
|
||||
|
||||
rsentry->rstate->metrics.completion_tokens +=
|
||||
static_cast<int>(new_tokens.size()) - rollback_cnt;
|
||||
}
|
||||
|
||||
auto tend = std::chrono::high_resolution_clock::now();
|
||||
estate->metrics.engine_jump_forward_time_sum +=
|
||||
static_cast<double>((tend - tstart).count()) / 1e9;
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
private:
|
||||
/*! \brief Check if jump-forward decoding can be executed without exceeding the memory limit. */
|
||||
bool CheckMemForJumpForward(int num_rsentries) {
|
||||
static constexpr int MAX_AVG_JUMPFORWARD_PAGES_PER_REQUEST = 10;
|
||||
int num_available_pages = models_[0]->GetNumAvailablePages();
|
||||
return num_rsentries * MAX_AVG_JUMPFORWARD_PAGES_PER_REQUEST <= num_available_pages;
|
||||
}
|
||||
|
||||
/*! \brief Check if the jump-forward can be executed. When logprobs is requested, or the
|
||||
* grammar state matcher is not defined, jump-forward is not executed. */
|
||||
bool CanJumpForward(const RequestStateEntry& rsentry) {
|
||||
if (rsentry->request->generation_cfg->debug_config.grammar_execution_mode !=
|
||||
GrammarExecutionMode::kJumpForward) {
|
||||
return false;
|
||||
}
|
||||
if (rsentry->request->generation_cfg->logprobs) {
|
||||
return false;
|
||||
}
|
||||
if (!rsentry->mstates[0]->grammar_matcher) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Retokenize the input string with a new string.
|
||||
* \param mstate The model state.
|
||||
* \param new_string The new string to append.
|
||||
* \param max_rollback_tokens The maximum number of tokens to rollback.
|
||||
* \return The number of tokens to rollback, the new tokens and a delta string of output (equal to
|
||||
* new_string if no cutoff happens; shorter than new_string if cutoff happens).
|
||||
*/
|
||||
std::tuple<int, std::vector<int32_t>, std::string> RetokenizeWithNewString(
|
||||
RequestModelState mstate, const std::string& new_string, int max_rollback_tokens) {
|
||||
// Step 1. Get past tokens
|
||||
// past_tokens = mstate[-max_rollback_tokens:]
|
||||
// past_string = detokenize(past_tokens)
|
||||
const auto& token_table = tokenizer_->PostProcessedTokenTable();
|
||||
std::vector<int32_t> past_tokens;
|
||||
std::string past_string;
|
||||
auto past_begin_it = mstate->committed_tokens.size() >= max_rollback_tokens
|
||||
? mstate->committed_tokens.end() - max_rollback_tokens
|
||||
: mstate->committed_tokens.begin();
|
||||
for (auto it = past_begin_it; it != mstate->committed_tokens.end(); ++it) {
|
||||
past_tokens.push_back(it->GetTokenId());
|
||||
past_string += token_table[it->GetTokenId()];
|
||||
}
|
||||
|
||||
// Step 2. Retokenize
|
||||
// Compare tokenize(past_string + new_string) and past_tokens
|
||||
auto new_tokens = tokenizer_->EncodeNoPrependSpace(past_string + new_string);
|
||||
auto delta_string = new_string;
|
||||
|
||||
// Pop last token if it is a prefix of another token. That's because such tokens will often
|
||||
// be rolled back in the next decode, which disturbs the distribution, so we will avoid
|
||||
// generating them.
|
||||
if (tokenizer_->GetPrefixTokenMask()[new_tokens.back()]) {
|
||||
auto last_token = token_table[new_tokens.back()];
|
||||
if (last_token.length() >= new_string.length()) {
|
||||
return {0, {}, ""};
|
||||
}
|
||||
|
||||
delta_string = delta_string.substr(0, delta_string.length() - last_token.length());
|
||||
new_tokens.pop_back();
|
||||
}
|
||||
|
||||
int first_differ_idx = past_tokens.size();
|
||||
for (int i = 0; i < static_cast<int>(past_tokens.size()); ++i) {
|
||||
if (i == static_cast<int>(new_tokens.size()) || past_tokens[i] != new_tokens[i]) {
|
||||
first_differ_idx = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return {past_tokens.size() - first_differ_idx,
|
||||
std::vector<int32_t>(new_tokens.begin() + first_differ_idx, new_tokens.end()),
|
||||
delta_string};
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Handle rollback for the stream output, the model state and the kv cache.
|
||||
* \param rsentry The request state entry.
|
||||
* \param mstate The model state.
|
||||
* \param rollback_cnt The number of tokens to rollback.
|
||||
* \param new_tokens The new tokens. Useful for the stream output.
|
||||
* \param new_string The delta string of output. Useful for the stream output.
|
||||
*/
|
||||
void HandleRollback(const RequestStateEntry& rsentry, RequestModelState mstate, int rollback_cnt,
|
||||
const std::vector<int32_t>& new_tokens, const std::string& new_string) {
|
||||
// 1. Handle rollback for the stream output
|
||||
if (rollback_cnt >
|
||||
static_cast<int>(mstate->committed_tokens.size()) - rsentry->next_callback_token_pos) {
|
||||
const auto& token_table = tokenizer_->PostProcessedTokenTable();
|
||||
for (auto i = rsentry->next_callback_token_pos; i < mstate->committed_tokens.size(); ++i) {
|
||||
auto token_id = mstate->committed_tokens[i].GetTokenId();
|
||||
rsentry->extra_prefix_string += token_table[token_id];
|
||||
}
|
||||
rsentry->extra_prefix_string += new_string;
|
||||
rsentry->next_callback_token_pos = static_cast<int>(mstate->committed_tokens.size()) -
|
||||
rollback_cnt + static_cast<int>(new_tokens.size());
|
||||
}
|
||||
|
||||
// 2. Handle rollback for the model state
|
||||
if (rollback_cnt > 0) {
|
||||
mstate->RollbackTokens(rollback_cnt);
|
||||
}
|
||||
|
||||
// 3. Handle rollback for the kv cache
|
||||
if (rollback_cnt > mstate->num_tokens_for_next_decode) {
|
||||
models_[0]->PopNFromKVCache(mstate->internal_id,
|
||||
rollback_cnt - mstate->num_tokens_for_next_decode);
|
||||
mstate->num_tokens_for_next_decode = 0;
|
||||
} else {
|
||||
mstate->num_tokens_for_next_decode -= rollback_cnt;
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief The model to run jump-forward decoding. When there are multiple
|
||||
* models, the `Step` function of the created action will not take effect.
|
||||
*/
|
||||
Array<Model> models_;
|
||||
/*! \brief Tokenizer for retokenization. */
|
||||
Tokenizer tokenizer_;
|
||||
/*! \brief Event trace recorder. */
|
||||
Optional<EventTraceRecorder> trace_recorder_;
|
||||
/*! \brief The maximum number of tokens to rollback. */
|
||||
const int MAX_ROLLBACK_TOKENS_ = 10;
|
||||
};
|
||||
|
||||
EngineAction EngineAction::BatchJumpForward(Array<Model> models, Tokenizer tokenizer,
|
||||
Optional<EventTraceRecorder> trace_recorder) {
|
||||
return EngineAction(tvm::ffi::make_object<BatchJumpForwardActionObj>(
|
||||
std::move(models), std::move(tokenizer), std::move(trace_recorder)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,534 @@
|
||||
/*!
|
||||
* 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
|
||||
@@ -0,0 +1,144 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/batch_prefill_base.h
|
||||
*/
|
||||
|
||||
#include "../config.h"
|
||||
#include "../model.h"
|
||||
#include "action.h"
|
||||
#include "action_commons.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
/*!
|
||||
* \brief The base action of that prefills requests in the `waiting_queue` of
|
||||
* the engine state.
|
||||
*/
|
||||
class BatchPrefillBaseActionObj : public EngineActionObj {
|
||||
protected:
|
||||
/*! \brief The class of request state entry and its maximum allowed length for prefill. */
|
||||
struct PrefillInput {
|
||||
RequestStateEntry rsentry;
|
||||
int max_prefill_length = 0;
|
||||
int num_child_to_activate = 0;
|
||||
bool is_decode = false;
|
||||
};
|
||||
|
||||
BatchPrefillBaseActionObj(Array<Model> models, EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
|
||||
/*!
|
||||
* \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<PrefillInput> GetRequestStateEntriesToPrefill(EngineState estate);
|
||||
|
||||
/*! \brief Check if the input requests can be prefilled under conditions. */
|
||||
bool 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);
|
||||
|
||||
/*!
|
||||
* \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> ChunkPrefillInputData(const RequestModelState& mstate,
|
||||
int max_prefill_length);
|
||||
|
||||
/*!
|
||||
* \brief Update status of request states from pending to alive and collect request state entries
|
||||
* from the prefill input.
|
||||
* \param prefill_inputs The prefill input.
|
||||
* \param estate The engine state.
|
||||
* \param[out] request_ids The array to store the request ids of the request state entries.
|
||||
* \param[out] rstates_of_entries The vector to store the request state entries.
|
||||
* \param[out] status_before_prefill The vector to store the status of the request state entries
|
||||
* before prefill.
|
||||
*/
|
||||
void UpdateRequestToAlive(const std::vector<PrefillInput>& prefill_inputs,
|
||||
const EngineState& estate, Array<String>* request_ids,
|
||||
std::vector<RequestState>* rstates_of_entries,
|
||||
std::vector<RequestStateStatus>* status_before_prefill);
|
||||
|
||||
/*!
|
||||
* \brief Remove the request from waiting queue if all its request states are now alive and have
|
||||
* no remaining chunked inputs.
|
||||
* \param prefill_inputs The prefill input.
|
||||
* \param estate The engine state.
|
||||
* \param rstates_of_entries The request state entries for each prefill input.
|
||||
* \return The processed requests.
|
||||
*/
|
||||
std::vector<Request> RemoveProcessedRequests(const std::vector<PrefillInput>& prefill_inputs,
|
||||
const EngineState& estate,
|
||||
const std::vector<RequestState>& rstates_of_entries);
|
||||
|
||||
/*!
|
||||
* \brief Update the committed tokens of states. If a request is first-time prefilled, set the
|
||||
* prefill finish time.
|
||||
* \param rsentries_for_sample The request state entries for sample.
|
||||
* \param rsentry_activated The activation status of the request state entries.
|
||||
* \param sample_results The sample results.
|
||||
*/
|
||||
void UpdateRequestStateEntriesWithSampleResults(
|
||||
const std::vector<RequestStateEntry>& rsentries_for_sample,
|
||||
const std::vector<bool>& rsentry_activated, const std::vector<SampleResult>& sample_results);
|
||||
|
||||
/*!
|
||||
* \brief Get the concatenated Shape of RequestModelState input data, return empty Shape if
|
||||
* there is untokenized data.
|
||||
* \param mstate The RequestModelState whose input data is to be concatenated.
|
||||
* \return The concatenate Shape.
|
||||
*/
|
||||
std::vector<int32_t> GetConcatPrefillInputData(const RequestModelState& mstate);
|
||||
|
||||
/*!
|
||||
* \brief Pop the prefix tokens of the RequestModelState input data array.
|
||||
* \param mstate The RequestModelState to be popped.
|
||||
* \param num_tokens The number of prefix tokens to be popped.
|
||||
*/
|
||||
void PopPrefillInputData(const RequestModelState& mstate, size_t num_tokens);
|
||||
|
||||
/*!
|
||||
* \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.
|
||||
*/
|
||||
virtual int MatchPrefixCache(EngineState estate, PrefillInput* input) = 0;
|
||||
|
||||
/*! \brief The models to run prefill in. */
|
||||
Array<Model> models_;
|
||||
/*! \brief The engine config. */
|
||||
EngineConfig engine_config_;
|
||||
/*! \brief The KV state kind. */
|
||||
KVStateKind kv_state_kind_;
|
||||
/*! \brief The sliding window size of each model. */
|
||||
std::vector<int> sliding_window_sizes_;
|
||||
/*! \brief Event trace recorder. */
|
||||
Optional<EventTraceRecorder> trace_recorder_;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief A utility function to check whether there is enough spare space in
|
||||
* KV cache for the number of required pages and total input length.
|
||||
*/
|
||||
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);
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,381 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/batch_verify.cc
|
||||
*/
|
||||
|
||||
#include <cmath>
|
||||
#include <exception>
|
||||
#include <numeric>
|
||||
|
||||
#include "../../support/random.h"
|
||||
#include "../config.h"
|
||||
#include "../model.h"
|
||||
#include "../sampler/sampler.h"
|
||||
#include "action.h"
|
||||
#include "action_commons.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
/*!
|
||||
* \brief The action that runs verification for requests in the
|
||||
* `running_queue` of engine state. Preempt low-priority requests
|
||||
* accordingly when it is impossible to decode all the running requests.
|
||||
*/
|
||||
class BatchVerifyActionObj : public EngineActionObj {
|
||||
public:
|
||||
explicit BatchVerifyActionObj(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler, std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder)
|
||||
: models_(std::move(models)),
|
||||
logit_processor_(std::move(logit_processor)),
|
||||
sampler_(std::move(sampler)),
|
||||
model_workspaces_(std::move(model_workspaces)),
|
||||
draft_token_workspace_manager_(std::move(draft_token_workspace_manager)),
|
||||
engine_config_(std::move(engine_config)),
|
||||
trace_recorder_(std::move(trace_recorder)),
|
||||
rng_(RandomGenerator::GetInstance()) {}
|
||||
|
||||
Array<Request> Step(EngineState estate) final {
|
||||
// - Only run spec decode when there are two models (llm+ssm) and >=1 running requests.
|
||||
if (models_.size() != 2 || estate->running_queue.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
const auto& [rsentries, verify_lengths, total_verify_length] = GetDraftsToVerify(estate);
|
||||
TVM_FFI_ICHECK_EQ(rsentries.size(), verify_lengths.size());
|
||||
if (rsentries.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
auto tstart = std::chrono::high_resolution_clock::now();
|
||||
int num_rsentries = rsentries.size();
|
||||
Array<String> request_ids =
|
||||
rsentries.Map([](const RequestStateEntry& rstate) { return rstate->request->id; });
|
||||
|
||||
// - Get embedding and run verify.
|
||||
std::vector<int64_t> request_internal_ids;
|
||||
std::vector<int32_t> all_tokens_to_verify;
|
||||
Array<RequestModelState> verify_request_mstates;
|
||||
Array<RequestModelState> draft_request_mstates;
|
||||
Array<GenerationConfig> generation_cfg;
|
||||
Array<GenerationConfig> generation_cfg_for_top_p_norm;
|
||||
std::vector<RandomGenerator*> rngs;
|
||||
std::vector<std::vector<SampleResult>> draft_output_tokens;
|
||||
std::vector<int64_t> token_tree_parent_ptr;
|
||||
std::vector<std::vector<int>> draft_token_indices;
|
||||
token_tree_parent_ptr.reserve(total_verify_length);
|
||||
request_internal_ids.reserve(num_rsentries);
|
||||
all_tokens_to_verify.reserve(total_verify_length);
|
||||
draft_token_indices.reserve(num_rsentries);
|
||||
verify_request_mstates.reserve(num_rsentries);
|
||||
draft_request_mstates.reserve(num_rsentries);
|
||||
rngs.reserve(num_rsentries);
|
||||
generation_cfg.reserve(num_rsentries);
|
||||
generation_cfg_for_top_p_norm.reserve(total_verify_length);
|
||||
draft_output_tokens.reserve(num_rsentries);
|
||||
draft_token_slots_.clear();
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
RequestModelState verify_mstate = rsentries[i]->mstates[verify_model_id_];
|
||||
RequestModelState draft_mstate = rsentries[i]->mstates[draft_model_id_];
|
||||
request_internal_ids.push_back(verify_mstate->internal_id);
|
||||
TVM_FFI_ICHECK(!verify_lengths.empty());
|
||||
TVM_FFI_ICHECK_EQ(verify_lengths[i], draft_mstate->draft_output_tokens.size() + 1);
|
||||
TVM_FFI_ICHECK_EQ(verify_lengths[i], draft_mstate->draft_token_slots.size() + 1);
|
||||
// the last committed token + all the draft tokens.
|
||||
draft_token_slots_.push_back(0); // placeholder for the last committed token
|
||||
all_tokens_to_verify.push_back(draft_mstate->committed_tokens.back().GetTokenId());
|
||||
token_tree_parent_ptr.push_back(-1);
|
||||
generation_cfg_for_top_p_norm.push_back(rsentries[i]->request->generation_cfg);
|
||||
std::vector<int> cur_draft_token_indices;
|
||||
cur_draft_token_indices.resize(draft_mstate->draft_output_tokens.size() + 1);
|
||||
std::iota(cur_draft_token_indices.begin(), cur_draft_token_indices.end(), -1);
|
||||
for (int j = 0; j < static_cast<int>(draft_mstate->draft_output_tokens.size()); ++j) {
|
||||
all_tokens_to_verify.push_back(draft_mstate->draft_output_tokens[j].GetTokenId());
|
||||
draft_token_slots_.push_back(draft_mstate->draft_token_slots[j]);
|
||||
token_tree_parent_ptr.push_back(draft_mstate->draft_token_parent_idx[j] + 1);
|
||||
generation_cfg_for_top_p_norm.push_back(rsentries[i]->request->generation_cfg);
|
||||
}
|
||||
draft_token_indices.emplace_back(std::move(cur_draft_token_indices));
|
||||
verify_request_mstates.push_back(verify_mstate);
|
||||
draft_request_mstates.push_back(draft_mstate);
|
||||
generation_cfg.push_back(rsentries[i]->request->generation_cfg);
|
||||
rngs.push_back(&rsentries[i]->rng);
|
||||
draft_output_tokens.push_back(draft_mstate->draft_output_tokens);
|
||||
}
|
||||
Tensor draft_probs_on_device = models_[draft_model_id_]->GatherDraftProbs(
|
||||
model_workspaces_[verify_model_id_].draft_probs_storage, draft_token_slots_,
|
||||
&model_workspaces_[verify_model_id_].draft_probs);
|
||||
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start verify embedding");
|
||||
ObjectRef embeddings = models_[verify_model_id_]->TokenEmbed(
|
||||
{Shape{all_tokens_to_verify.begin(), all_tokens_to_verify.end()}});
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish verify embedding");
|
||||
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start verify");
|
||||
Tensor logits = models_[verify_model_id_]->BatchVerify(embeddings, request_internal_ids,
|
||||
verify_lengths, token_tree_parent_ptr);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish verify");
|
||||
TVM_FFI_ICHECK_EQ(logits->ndim, 3);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[0], 1);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[1], total_verify_length);
|
||||
|
||||
// - Update logits.
|
||||
std::vector<int> cum_verify_lengths = {0};
|
||||
cum_verify_lengths.reserve(num_rsentries + 1);
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
cum_verify_lengths.push_back(cum_verify_lengths.back() + verify_lengths[i]);
|
||||
}
|
||||
logits = logits.CreateView({total_verify_length, logits->shape[2]}, logits->dtype);
|
||||
logit_processor_->InplaceUpdateLogits(logits, generation_cfg, verify_request_mstates,
|
||||
request_ids, &cum_verify_lengths, &draft_request_mstates,
|
||||
&draft_token_indices);
|
||||
|
||||
// - Compute probability distributions.
|
||||
Tensor probs_on_device = logit_processor_->ComputeProbsFromLogits(
|
||||
logits, generation_cfg, request_ids, &cum_verify_lengths);
|
||||
|
||||
// - 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();
|
||||
|
||||
// Fill range [0, total_verify_length) into `sample_indices`.
|
||||
std::vector<int> sample_indices(total_verify_length);
|
||||
std::iota(sample_indices.begin(), sample_indices.end(), 0);
|
||||
Tensor renormalized_probs = sampler_->BatchRenormalizeProbsByTopP(
|
||||
probs_on_device, sample_indices, request_ids, generation_cfg_for_top_p_norm);
|
||||
auto [sample_results_arr, last_accepted_tree_node_verify_model] =
|
||||
sampler_->BatchVerifyDraftTokensWithProbAfterTopP(
|
||||
renormalized_probs, request_ids, cum_verify_lengths, generation_cfg, rngs,
|
||||
draft_output_tokens, token_tree_parent_ptr, draft_probs_on_device);
|
||||
TVM_FFI_ICHECK_EQ(sample_results_arr.size(), num_rsentries);
|
||||
|
||||
// We collect the requests whose drafts are fully accepted.
|
||||
// When a request's draft is fully accepted, there is an extra token proposed
|
||||
// by the draft model but not added into the draft model's KV cache.
|
||||
// In this case, an additional batch decode step is needed for these requests.
|
||||
std::vector<int64_t> fully_accepted_rsentries;
|
||||
std::vector<int64_t> verify_model_seq_internal_ids;
|
||||
std::vector<int64_t> draft_model_seq_internal_ids;
|
||||
fully_accepted_rsentries.reserve(num_rsentries);
|
||||
verify_model_seq_internal_ids.reserve(num_rsentries);
|
||||
draft_model_seq_internal_ids.reserve(num_rsentries);
|
||||
|
||||
// The index of the last accepted tree node in the draft model. This is different from the
|
||||
// last accepted tree node in the verify model because the first round of draft does not
|
||||
// use tree attention.
|
||||
std::vector<int64_t> last_accepted_tree_node_draft_model;
|
||||
last_accepted_tree_node_draft_model.reserve(num_rsentries);
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
const std::vector<SampleResult>& sample_results = sample_results_arr[i];
|
||||
int accept_length = sample_results.size();
|
||||
for (SampleResult sample_result : sample_results) {
|
||||
rsentries[i]->mstates[verify_model_id_]->CommitToken(sample_result);
|
||||
rsentries[i]->mstates[draft_model_id_]->CommitToken(sample_result);
|
||||
}
|
||||
// Metrics update
|
||||
// live update the output metrics
|
||||
rsentries[i]->rstate->metrics.completion_tokens += accept_length;
|
||||
estate->metrics.spec_decode.Update(cum_verify_lengths[i + 1] - cum_verify_lengths[i],
|
||||
accept_length);
|
||||
if (engine_config_->spec_tree_width == 1) {
|
||||
// The roll back is needed for the chain draft case.
|
||||
int rollback_length =
|
||||
std::max(cum_verify_lengths[i + 1] - cum_verify_lengths[i] - accept_length, 0);
|
||||
if (rollback_length > 0) {
|
||||
// The last accepted token is not yet added into the draft model.
|
||||
// Therefore, the rollback length for the draft model is one less.
|
||||
models_[draft_model_id_]->PopNFromKVCache(
|
||||
rsentries[i]->mstates[draft_model_id_]->internal_id, rollback_length - 1);
|
||||
}
|
||||
}
|
||||
// Commit accepted tokens to the "verify_model", rollback kv cache
|
||||
// in the "draft_model".
|
||||
// NOTE: when number of small models is more than 1 (in the future),
|
||||
// it is possible to re-compute prefill for the small models.
|
||||
verify_model_seq_internal_ids.push_back(rsentries[i]->mstates[verify_model_id_]->internal_id);
|
||||
draft_model_seq_internal_ids.push_back(rsentries[i]->mstates[draft_model_id_]->internal_id);
|
||||
int last_accepted = last_accepted_tree_node_verify_model[i] -
|
||||
1; // minus one to get the index in the draft tokens
|
||||
if (last_accepted >= 0 &&
|
||||
rsentries[i]->mstates[draft_model_id_]->draft_token_first_child_idx[last_accepted] ==
|
||||
-1) { // minus one to get the index in the draft tokens
|
||||
// is leaf node, fully accepted
|
||||
last_accepted_tree_node_draft_model.push_back(
|
||||
rsentries[i]->mstates[draft_model_id_]->draft_token_parent_idx[last_accepted]);
|
||||
fully_accepted_rsentries.push_back(i);
|
||||
} else {
|
||||
last_accepted_tree_node_draft_model.push_back(last_accepted);
|
||||
}
|
||||
}
|
||||
models_[verify_model_id_]->CommitAcceptedTokenTreeNodesToKVCache(
|
||||
verify_model_seq_internal_ids,
|
||||
std::vector<int64_t>{last_accepted_tree_node_verify_model.begin(),
|
||||
last_accepted_tree_node_verify_model.end()});
|
||||
if (engine_config_->spec_tree_width > 1) {
|
||||
models_[draft_model_id_]->CommitAcceptedTokenTreeNodesToKVCache(
|
||||
draft_model_seq_internal_ids, last_accepted_tree_node_draft_model);
|
||||
}
|
||||
|
||||
if (!fully_accepted_rsentries.empty()) {
|
||||
// - Run a step of batch decode for requests whose drafts are fully accepted.
|
||||
// When a request's draft is fully accepted, there is an extra token proposed
|
||||
// by the draft model but not added into the draft model's KV cache.
|
||||
// In this case, an additional batch decode step is needed for these requests.
|
||||
std::vector<int> input_tokens;
|
||||
std::vector<int64_t> fully_accepted_request_internal_ids;
|
||||
input_tokens.reserve(fully_accepted_rsentries.size());
|
||||
fully_accepted_request_internal_ids.reserve(fully_accepted_rsentries.size());
|
||||
for (int rsentry_id : fully_accepted_rsentries) {
|
||||
int num_committed_tokens =
|
||||
rsentries[rsentry_id]->mstates[verify_model_id_]->committed_tokens.size();
|
||||
// When a request's draft is fully accepted, an additional new token is sampled.
|
||||
// So the token needed to fill in the draft model is the committed_token[-2].
|
||||
TVM_FFI_ICHECK_GE(num_committed_tokens, 2);
|
||||
input_tokens.push_back(rsentries[rsentry_id]
|
||||
->mstates[verify_model_id_]
|
||||
->committed_tokens[num_committed_tokens - 2]
|
||||
.GetTokenId());
|
||||
fully_accepted_request_internal_ids.push_back(
|
||||
rsentries[rsentry_id]->mstates[draft_model_id_]->internal_id);
|
||||
}
|
||||
// - Compute embeddings.
|
||||
ObjectRef embeddings =
|
||||
models_[draft_model_id_]->TokenEmbed({Shape{input_tokens.begin(), input_tokens.end()}});
|
||||
// - Invoke model decode.
|
||||
Tensor logits =
|
||||
models_[draft_model_id_]->BatchDecode(embeddings, fully_accepted_request_internal_ids);
|
||||
// - We explicitly synchronize to avoid the input tokens getting overriden in the
|
||||
// next runs of BatchDecode.
|
||||
// This is because we do not do sample for this round of batch decode.
|
||||
DeviceAPI::Get(logits->device)->StreamSync(logits->device, nullptr);
|
||||
}
|
||||
|
||||
// clear the draft model state entries
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
rsentries[i]->mstates[draft_model_id_]->RemoveAllDraftTokens(&draft_token_slots_);
|
||||
draft_token_workspace_manager_->FreeSlots(draft_token_slots_);
|
||||
// reset num_tokens_for_next_decode to 1
|
||||
rsentries[i]->mstates[verify_model_id_]->num_tokens_for_next_decode = 1;
|
||||
rsentries[i]->mstates[draft_model_id_]->num_tokens_for_next_decode = 1;
|
||||
}
|
||||
|
||||
auto tend = std::chrono::high_resolution_clock::now();
|
||||
double elapsed_time = static_cast<double>((tend - tstart).count()) / 1e9;
|
||||
estate->metrics.engine_decode_time_sum += elapsed_time;
|
||||
estate->metrics.UpdateVerifyTimeByBatchSize(total_verify_length, elapsed_time);
|
||||
|
||||
return estate->running_queue;
|
||||
}
|
||||
|
||||
private:
|
||||
struct DraftRequestStateEntries {
|
||||
/*! \brief The request state entries to verify. */
|
||||
Array<RequestStateEntry> draft_rsentries;
|
||||
/*! \brief The length to verify for each request state. */
|
||||
std::vector<int> verify_lengths;
|
||||
/*! \brief The total draft length. */
|
||||
int total_verify_length;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Decide whether to run verify for the draft of each request.
|
||||
* \param estate The engine state.
|
||||
* \return The drafts to verify, together with their respective
|
||||
* state and input length.
|
||||
*/
|
||||
DraftRequestStateEntries GetDraftsToVerify(EngineState estate) {
|
||||
std::vector<int> verify_lengths;
|
||||
int total_verify_length = 0;
|
||||
int total_required_pages = 0;
|
||||
int num_available_pages = models_[verify_model_id_]->GetNumAvailablePages();
|
||||
|
||||
// Preempt the request state entries that cannot fit the large model for verification.
|
||||
std::vector<RequestStateEntry> init_running_rsentries = estate->GetRunningRequestStateEntries();
|
||||
std::vector<int> num_page_requirement;
|
||||
num_page_requirement.reserve(init_running_rsentries.size());
|
||||
std::vector<RequestStateEntry> running_rsentries;
|
||||
running_rsentries.reserve(init_running_rsentries.size());
|
||||
for (const RequestStateEntry& rsentry : init_running_rsentries) {
|
||||
int draft_length = rsentry->mstates[draft_model_id_]->draft_output_tokens.size();
|
||||
if (draft_length == 0) {
|
||||
continue;
|
||||
}
|
||||
running_rsentries.push_back(rsentry);
|
||||
int num_require_pages = (draft_length + engine_config_->kv_cache_page_size - 1) /
|
||||
engine_config_->kv_cache_page_size;
|
||||
verify_lengths.push_back(draft_length + 1);
|
||||
num_page_requirement.push_back(num_require_pages);
|
||||
total_verify_length += draft_length + 1;
|
||||
total_required_pages += num_require_pages;
|
||||
}
|
||||
while (!CanVerify(total_required_pages)) {
|
||||
if (estate->prefix_cache->TryFreeMemory()) continue;
|
||||
RequestStateEntry preempted = PreemptLastRunningRequestStateEntry(
|
||||
estate, models_, draft_token_workspace_manager_, trace_recorder_);
|
||||
if (preempted.same_as(running_rsentries.back())) {
|
||||
total_verify_length -= verify_lengths.back();
|
||||
total_required_pages -= num_page_requirement.back();
|
||||
verify_lengths.pop_back();
|
||||
num_page_requirement.pop_back();
|
||||
running_rsentries.pop_back();
|
||||
}
|
||||
}
|
||||
TVM_FFI_ICHECK_LE(total_verify_length,
|
||||
std::min(static_cast<int64_t>(engine_config_->max_num_sequence),
|
||||
engine_config_->prefill_chunk_size))
|
||||
<< total_verify_length << " " << engine_config_->max_num_sequence;
|
||||
|
||||
return {running_rsentries, verify_lengths, total_verify_length};
|
||||
}
|
||||
|
||||
bool CanVerify(int num_required_pages) {
|
||||
int num_available_pages = models_[0]->GetNumAvailablePages();
|
||||
return num_required_pages <= num_available_pages;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief The model to run decode in. When there are multiple
|
||||
* models, the `Step` function of the created action will not take effect.
|
||||
*/
|
||||
Array<Model> models_;
|
||||
/*! \brief The logit processor. */
|
||||
LogitProcessor logit_processor_;
|
||||
/*! \brief The sampler to sample new tokens. */
|
||||
Sampler sampler_;
|
||||
/*! \brief The model workspaces. */
|
||||
std::vector<ModelWorkspace> model_workspaces_;
|
||||
/*! \brief The draft token workspace manager. */
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager_;
|
||||
/*! \brief The engine config. */
|
||||
EngineConfig engine_config_;
|
||||
/*! \brief Event trace recorder. */
|
||||
Optional<EventTraceRecorder> trace_recorder_;
|
||||
/*! \brief Random number generator. */
|
||||
RandomGenerator& rng_;
|
||||
/*! \brief The ids of verify/draft models. */
|
||||
const int verify_model_id_ = 0;
|
||||
const int draft_model_id_ = 1;
|
||||
const float eps_ = 1e-5;
|
||||
/*! \brief Temporary buffer to store the slots of the current draft tokens */
|
||||
std::vector<int> draft_token_slots_;
|
||||
};
|
||||
|
||||
EngineAction EngineAction::BatchVerify(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder) {
|
||||
return EngineAction(tvm::ffi::make_object<BatchVerifyActionObj>(
|
||||
std::move(models), std::move(logit_processor), std::move(sampler),
|
||||
std::move(model_workspaces), std::move(draft_token_workspace_manager),
|
||||
std::move(engine_config), std::move(trace_recorder)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,447 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/new_request_prefill.cc
|
||||
*/
|
||||
|
||||
#include <tvm/support/cuda/nvtx.h>
|
||||
|
||||
#include <optional>
|
||||
|
||||
#include "../../support/utils.h"
|
||||
#include "../sampler/sampler.h"
|
||||
#include "batch_prefill_base.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::support::NVTXScopedRange;
|
||||
|
||||
/*!
|
||||
* \brief The action that runs prefill preparation in disaggregation system.
|
||||
* It picks a new request, reserve its KV data locations, and returns the
|
||||
* KV data locations and the matched prefix length in prefix cache.
|
||||
*/
|
||||
class DisaggPrepareReceiveActionObj : public BatchPrefillBaseActionObj {
|
||||
public:
|
||||
explicit DisaggPrepareReceiveActionObj(Array<Model> models, EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
Optional<EventTraceRecorder> trace_recorder,
|
||||
FRequestStreamCallback request_stream_callback)
|
||||
: BatchPrefillBaseActionObj(std::move(models), std::move(engine_config),
|
||||
std::move(model_configs), std::move(trace_recorder)),
|
||||
request_stream_callback_(std::move(request_stream_callback)) {
|
||||
TVM_FFI_ICHECK(kv_state_kind_ == KVStateKind::kKVCache)
|
||||
<< "Only PagedKVCache supports prefill preparation and KV migration";
|
||||
}
|
||||
|
||||
Array<Request> Step(EngineState estate) final {
|
||||
std::vector<Request> processed_requests;
|
||||
|
||||
// - Find the requests in `waiting_queue` that can prefill in this step.
|
||||
std::optional<PrefillInput> prefill_input_opt;
|
||||
while (true) {
|
||||
prefill_input_opt = GetRequestStateEntriesToPrefill(estate);
|
||||
if (!prefill_input_opt.has_value()) {
|
||||
break;
|
||||
}
|
||||
PrefillInput prefill_input = prefill_input_opt.value();
|
||||
int prefix_matched_length = 0;
|
||||
Request request = prefill_input.rsentry->request;
|
||||
processed_requests.push_back(request);
|
||||
int total_input_length = 0;
|
||||
for (const Data& data : request->inputs) {
|
||||
total_input_length += data->GetLength();
|
||||
}
|
||||
|
||||
{
|
||||
NVTXScopedRange nvtx_scope("DisaggPrepareReceive matching prefix");
|
||||
prefix_matched_length = MatchPrefixCache(estate, &prefill_input);
|
||||
}
|
||||
|
||||
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_input}, estate, &request_ids, &rstates_of_entries,
|
||||
&status_before_prefill);
|
||||
// "UpdateRequestToAlive" may add the request to the engine's running request queue.
|
||||
// We erase it since it's pending for the prefill instance to send the KV data over.
|
||||
if (!estate->running_queue.empty() && estate->running_queue.back().same_as(request)) {
|
||||
estate->running_queue.pop_back();
|
||||
}
|
||||
|
||||
// - Add the sequence to each model.
|
||||
int prefill_length = -1;
|
||||
Tensor logits_for_sample{nullptr};
|
||||
std::vector<Shape> kv_append_metadata;
|
||||
kv_append_metadata.reserve(models_.size());
|
||||
for (int model_id = 0; model_id < static_cast<int>(models_.size()); ++model_id) {
|
||||
const RequestStateEntry& rsentry = prefill_input.rsentry;
|
||||
RequestModelState mstate = rsentry->mstates[model_id];
|
||||
Array<Data> input_data = mstate->inputs;
|
||||
mstate->inputs.clear();
|
||||
int input_length = prefill_input.max_prefill_length;
|
||||
if (prefill_length == -1) {
|
||||
prefill_length = input_length;
|
||||
} else {
|
||||
TVM_FFI_ICHECK_EQ(prefill_length, 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[0] == RequestStateStatus::kPending &&
|
||||
!estate->prefix_cache->HasSequence(mstate->internal_id)) {
|
||||
// Add the sequence to the model, or fork the sequence from its parent.
|
||||
// If the sequence is already in prefix cache, it has also been added/forked in the
|
||||
// KVCache.
|
||||
if (rsentry->parent_idx == -1) {
|
||||
models_[model_id]->AddNewSequence(mstate->internal_id);
|
||||
} else {
|
||||
models_[model_id]->ForkSequence(
|
||||
rstates_of_entries[0]->entries[rsentry->parent_idx]->mstates[model_id]->internal_id,
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
// Record the of the prefilled inputs for prefix cache update.
|
||||
for (int j = 0; j < static_cast<int>(input_data.size()); ++j) {
|
||||
if (!model_id && !prefill_input.is_decode) {
|
||||
mstate->prefilled_inputs.push_back(input_data[j]);
|
||||
}
|
||||
}
|
||||
|
||||
int64_t request_internal_id = mstate->internal_id;
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start prefill");
|
||||
Shape compressed_kv_append_metadata = {0};
|
||||
if (prefill_length > 0) {
|
||||
compressed_kv_append_metadata =
|
||||
models_[model_id]->DisaggPrepareKVRecv(request_internal_id, prefill_length);
|
||||
}
|
||||
kv_append_metadata.push_back(compressed_kv_append_metadata);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish prefill");
|
||||
}
|
||||
|
||||
// - 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();
|
||||
|
||||
auto tend = std::chrono::high_resolution_clock::now();
|
||||
|
||||
// - Remove the request from the waiting queue.
|
||||
auto it_request =
|
||||
std::find(estate->waiting_queue.begin(), estate->waiting_queue.end(), request);
|
||||
TVM_FFI_ICHECK(it_request != estate->waiting_queue.end());
|
||||
estate->waiting_queue.erase(it_request);
|
||||
|
||||
{
|
||||
NVTXScopedRange nvtx_scope("Call request stream callback");
|
||||
tvm::ffi::json::Object response_body;
|
||||
response_body.Set("prompt_length", static_cast<int64_t>(total_input_length));
|
||||
response_body.Set("prefix_matched_length", static_cast<int64_t>(prefix_matched_length));
|
||||
// We further flatten the metadata array of all models into a single array.
|
||||
tvm::ffi::json::Array kv_append_metadata_arr;
|
||||
for (const Shape& compressed_kv_append_metadata : kv_append_metadata) {
|
||||
for (int64_t value : compressed_kv_append_metadata) {
|
||||
kv_append_metadata_arr.push_back(value);
|
||||
}
|
||||
TVM_FFI_ICHECK(!compressed_kv_append_metadata.empty());
|
||||
int num_segments = compressed_kv_append_metadata[0];
|
||||
TVM_FFI_ICHECK_EQ(compressed_kv_append_metadata.size(), num_segments * 2 + 1);
|
||||
int transmission_length = 0;
|
||||
for (int i = 0; i < num_segments; ++i) {
|
||||
transmission_length += compressed_kv_append_metadata[i * 2 + 2];
|
||||
}
|
||||
TVM_FFI_ICHECK_EQ(transmission_length, prefill_length);
|
||||
}
|
||||
|
||||
response_body.Set(
|
||||
"kv_append_metadata",
|
||||
Base64Encode(std::string(tvm::ffi::json::Stringify(kv_append_metadata_arr))));
|
||||
|
||||
tvm::ffi::json::Object usage;
|
||||
usage.Set("prompt_tokens", static_cast<int64_t>(0));
|
||||
usage.Set("completion_tokens", static_cast<int64_t>(0));
|
||||
usage.Set("total_tokens", static_cast<int64_t>(0));
|
||||
usage.Set("extra", response_body);
|
||||
RequestStreamOutput stream_output =
|
||||
RequestStreamOutput::Usage(request->id, std::string(tvm::ffi::json::Stringify(usage)));
|
||||
// - Invoke the stream callback function once for all collected requests.
|
||||
request_stream_callback_(Array<RequestStreamOutput>{stream_output});
|
||||
}
|
||||
}
|
||||
|
||||
for (const Request& request : processed_requests) {
|
||||
TVM_FFI_ICHECK(std::find(estate->running_queue.begin(), estate->running_queue.end(),
|
||||
request) == estate->running_queue.end());
|
||||
}
|
||||
return {processed_requests};
|
||||
}
|
||||
|
||||
private:
|
||||
// Mimicked from BatchPrefillBaseActionObj::GetRequestStateEntriesToPrefill
|
||||
std::optional<PrefillInput> GetRequestStateEntriesToPrefill(EngineState estate) {
|
||||
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 {};
|
||||
}
|
||||
}
|
||||
int num_running_rsentries = static_cast<int>(running_rsentries->size());
|
||||
|
||||
Request request{nullptr};
|
||||
for (const Request& request_candidate : estate->waiting_queue) {
|
||||
if (request_candidate->generation_cfg->debug_config.disagg_config.kind ==
|
||||
DisaggRequestKind::kPrepareReceive) {
|
||||
request = request_candidate;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!request.defined()) {
|
||||
// No request to prepare for prefill.
|
||||
return {};
|
||||
}
|
||||
TVM_FFI_ICHECK_EQ(
|
||||
request->generation_cfg->debug_config.disagg_config.kv_window_begin.value_or(0), 0);
|
||||
|
||||
std::vector<PrefillInput> prefill_input_for_all_models;
|
||||
prefill_input_for_all_models.reserve(models_.size());
|
||||
|
||||
// We first collect the inputs that can be prefilled for each model.
|
||||
// The inputs for each model are expected to be exactly the same.
|
||||
for (int i = 0; i < static_cast<int>(models_.size()); ++i) {
|
||||
NVTXScopedRange nvtx_scope("Process request " + request->id);
|
||||
|
||||
PrefillInput prefill_input;
|
||||
// - Try to prefill pending requests.
|
||||
int num_available_pages = models_[i]->GetNumAvailablePages();
|
||||
int current_total_seq_len = models_[i]->GetCurrentTotalSequenceLength();
|
||||
|
||||
RequestState rstate = estate->GetRequestState(request);
|
||||
bool prefill_stops = false;
|
||||
for (int j = 1; j < static_cast<int>(rstate->entries.size()); ++j) {
|
||||
TVM_FFI_ICHECK(rstate->entries[j]->mstates[i]->inputs.empty())
|
||||
<< "Re-prefill of preempted requests is not supported by prefill preparation.";
|
||||
}
|
||||
const RequestStateEntry& rsentry = rstate->entries[0];
|
||||
TVM_FFI_ICHECK(!rsentry->mstates[i]->inputs.empty())
|
||||
<< "The request entry must have pending inputs.";
|
||||
|
||||
// Todo: handle the case that input length is 1.
|
||||
|
||||
int input_length = rsentry->mstates[i]->GetInputLength();
|
||||
// Update the input length with the requested KV window, where "[begin:end]"
|
||||
// means the KV range to prefill on a prefill instance.
|
||||
int kv_window_begin =
|
||||
request->generation_cfg->debug_config.disagg_config.kv_window_begin.value_or(0);
|
||||
int kv_window_end =
|
||||
request->generation_cfg->debug_config.disagg_config.kv_window_end.value_or(input_length);
|
||||
TVM_FFI_ICHECK_EQ(kv_window_begin, 0);
|
||||
if (kv_window_end < 0) {
|
||||
kv_window_end = input_length + kv_window_end;
|
||||
}
|
||||
TVM_FFI_ICHECK_GE(kv_window_end, 0);
|
||||
TVM_FFI_ICHECK_LT(kv_window_end, input_length)
|
||||
<< "Prefill the full input on the remote machine is not supported.";
|
||||
int orig_input_length = input_length;
|
||||
input_length = kv_window_end;
|
||||
|
||||
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 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);
|
||||
}
|
||||
|
||||
// Check if the entire request state entry can fit for prefill.
|
||||
bool can_prefill = false;
|
||||
{
|
||||
NVTXScopedRange nvtx_scope("Attempt");
|
||||
for (int num_child_to_activate = rsentry->child_indices.size(); num_child_to_activate >= 0;
|
||||
--num_child_to_activate) {
|
||||
while (!HasPrefillSpace(num_require_pages, sliding_window_enabled, num_running_rsentries,
|
||||
num_available_pages, current_total_seq_len, 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, 1 + num_child_to_activate, input_length, num_require_pages,
|
||||
num_available_pages, current_total_seq_len, num_running_rsentries,
|
||||
kv_state_kind_, sliding_window_enabled)) {
|
||||
prefill_input = {rsentry, input_length, num_child_to_activate, /*is_decode=*/false};
|
||||
can_prefill = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!can_prefill) {
|
||||
return std::nullopt;
|
||||
}
|
||||
rsentry->mstates[i]->inputs =
|
||||
SplitData(rsentry->mstates[i]->inputs, orig_input_length, kv_window_end).first;
|
||||
prefill_input_for_all_models.push_back(prefill_input);
|
||||
}
|
||||
|
||||
// Prefill inputs of all models should be the same.
|
||||
TVM_FFI_ICHECK(!prefill_input_for_all_models.empty());
|
||||
PrefillInput prefill_input = prefill_input_for_all_models[0];
|
||||
{
|
||||
NVTXScopedRange nvtx_scope("reduction");
|
||||
for (int i = 1; i < static_cast<int>(prefill_input_for_all_models.size()); ++i) {
|
||||
TVM_FFI_ICHECK(prefill_input_for_all_models[i].rsentry.same_as(prefill_input.rsentry));
|
||||
TVM_FFI_ICHECK_EQ(prefill_input_for_all_models[i].max_prefill_length,
|
||||
prefill_input.max_prefill_length);
|
||||
TVM_FFI_ICHECK_EQ(prefill_input_for_all_models[i].num_child_to_activate,
|
||||
prefill_input.num_child_to_activate);
|
||||
}
|
||||
}
|
||||
|
||||
return prefill_input;
|
||||
}
|
||||
|
||||
// Mimicked from BatchPrefillBaseActionObj::CanPrefill
|
||||
bool 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) {
|
||||
// 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: at least one decode can be performed after prefill.
|
||||
// Cond 2: number of total tokens after "x" times of decode does not
|
||||
// exceed the limit, where "x" is a watermark number can
|
||||
// be configured and adjusted in the future.
|
||||
if (num_required_pages + 400 > num_available_pages) {
|
||||
return false;
|
||||
}
|
||||
return 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);
|
||||
}
|
||||
|
||||
// Mimicked from NewRequestPrefillActionObj::MatchPrefixCache
|
||||
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 (Model model : models_) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
} 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 (Model model : models_) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
} 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// 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 The stream callback function to passes back the KV cache metadata
|
||||
* and prefix matched length in prefix cache.
|
||||
*/
|
||||
FRequestStreamCallback request_stream_callback_;
|
||||
};
|
||||
|
||||
EngineAction EngineAction::DisaggPrepareReceive(Array<Model> models, EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
Optional<EventTraceRecorder> trace_recorder,
|
||||
FRequestStreamCallback request_stream_callback) {
|
||||
return EngineAction(tvm::ffi::make_object<DisaggPrepareReceiveActionObj>(
|
||||
std::move(models), std::move(engine_config), std::move(model_configs),
|
||||
std::move(trace_recorder), std::move(request_stream_callback)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,503 @@
|
||||
/*!
|
||||
* 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
|
||||
@@ -0,0 +1,234 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/eagle_batch_draft.cc
|
||||
*/
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "../config.h"
|
||||
#include "../model.h"
|
||||
#include "../sampler/sampler.h"
|
||||
#include "action.h"
|
||||
#include "action_commons.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
/*!
|
||||
* \brief The action that runs draft proposal for requests in the
|
||||
* `running_queue` of engine state. Preempt low-priority requests
|
||||
* accordingly when it is impossible to decode all the running requests.
|
||||
*/
|
||||
class EagleBatchDraftActionObj : public EngineActionObj {
|
||||
public:
|
||||
explicit EagleBatchDraftActionObj(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler, std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder)
|
||||
: models_(std::move(models)),
|
||||
logit_processor_(std::move(logit_processor)),
|
||||
sampler_(std::move(sampler)),
|
||||
model_workspaces_(std::move(model_workspaces)),
|
||||
draft_token_workspace_manager_(std::move(draft_token_workspace_manager)),
|
||||
engine_config_(std::move(engine_config)),
|
||||
trace_recorder_(std::move(trace_recorder)) {}
|
||||
|
||||
Array<Request> Step(EngineState estate) final {
|
||||
// - Only run spec decode when there are two models (llm+ssm) and >=1 running requests.
|
||||
if (models_.size() != 2 || estate->running_queue.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
// Preempt request state entries when decode cannot apply.
|
||||
std::vector<RequestStateEntry> running_rsentries = estate->GetRunningRequestStateEntries();
|
||||
while (!CanDecode(running_rsentries.size())) {
|
||||
if (estate->prefix_cache->TryFreeMemory()) continue;
|
||||
RequestStateEntry preempted = PreemptLastRunningRequestStateEntry(
|
||||
estate, models_, draft_token_workspace_manager_, trace_recorder_);
|
||||
if (preempted.same_as(running_rsentries.back())) {
|
||||
running_rsentries.pop_back();
|
||||
}
|
||||
}
|
||||
|
||||
auto tstart = std::chrono::high_resolution_clock::now();
|
||||
|
||||
int num_rsentries = running_rsentries.size();
|
||||
TVM_FFI_ICHECK_GT(num_rsentries, 0)
|
||||
<< "There should be at least one request state entry that can run decode. "
|
||||
"Possible failure reason: none of the prefill phase of the running requests is finished";
|
||||
TVM_FFI_ICHECK_LE(num_rsentries, engine_config_->max_num_sequence)
|
||||
<< "The number of running requests exceeds the max number of sequence in EngineConfig. "
|
||||
"Possible failure reason: the prefill action allows new sequence in regardless of the "
|
||||
"max num sequence.";
|
||||
|
||||
Array<String> request_ids;
|
||||
std::vector<int64_t> request_internal_ids;
|
||||
Array<GenerationConfig> generation_cfg;
|
||||
std::vector<RandomGenerator*> rngs;
|
||||
std::vector<std::vector<int>> draft_token_indices;
|
||||
request_ids.reserve(num_rsentries);
|
||||
request_internal_ids.reserve(num_rsentries);
|
||||
generation_cfg.reserve(num_rsentries);
|
||||
draft_token_indices.reserve(num_rsentries);
|
||||
for (const RequestStateEntry& rsentry : running_rsentries) {
|
||||
request_ids.push_back(rsentry->request->id);
|
||||
request_internal_ids.push_back(rsentry->mstates[0]->internal_id);
|
||||
generation_cfg.push_back(rsentry->request->generation_cfg);
|
||||
rngs.push_back(&rsentry->rng);
|
||||
}
|
||||
|
||||
TVM_FFI_ICHECK_GT(estate->spec_draft_length, 0)
|
||||
<< "The speculative decoding draft length must be positive.";
|
||||
// The first model doesn't get involved in draft proposal.
|
||||
for (int model_id = 1; model_id < static_cast<int>(models_.size()); ++model_id) {
|
||||
// Collect
|
||||
// - the last committed token,
|
||||
// - the request model state
|
||||
// of each request.
|
||||
std::vector<int> input_tokens;
|
||||
Array<RequestModelState> mstates;
|
||||
input_tokens.reserve(num_rsentries);
|
||||
mstates.reserve(num_rsentries);
|
||||
for (const RequestStateEntry& rsentry : running_rsentries) {
|
||||
mstates.push_back(rsentry->mstates[model_id]);
|
||||
}
|
||||
// draft_length_ rounds of draft proposal.
|
||||
ObjectRef hidden_states = model_workspaces_[model_id].hidden_states;
|
||||
// Concat last hidden_states
|
||||
draft_token_slots_.clear();
|
||||
if (estate->spec_draft_length > 1) {
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
draft_token_slots_.push_back(mstates[i]->draft_token_slots.back());
|
||||
}
|
||||
hidden_states = models_[model_id]->GatherHiddenStates(
|
||||
model_workspaces_[0].draft_hidden_states_storage, draft_token_slots_, &hidden_states);
|
||||
}
|
||||
// The first draft token has been generated in prefill/verify stage
|
||||
for (int draft_id = 1; draft_id < estate->spec_draft_length; ++draft_id) {
|
||||
draft_token_indices.clear();
|
||||
auto tdraft_start = std::chrono::high_resolution_clock::now();
|
||||
// prepare new input tokens
|
||||
input_tokens.clear();
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
TVM_FFI_ICHECK(!mstates[i]->draft_output_tokens.empty());
|
||||
input_tokens.push_back(mstates[i]->draft_output_tokens.back().GetTokenId());
|
||||
draft_token_indices.emplace_back(
|
||||
std::vector<int>{static_cast<int>(mstates[i]->draft_output_tokens.size() - 1)});
|
||||
}
|
||||
|
||||
// - Compute embeddings.
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start proposal embedding");
|
||||
ObjectRef embeddings =
|
||||
models_[model_id]->TokenEmbed({Shape{input_tokens.begin(), input_tokens.end()}});
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish proposal embedding");
|
||||
|
||||
// - Invoke model decode.
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start proposal decode");
|
||||
ObjectRef fused_embedding_hidden_states = models_[model_id]->FuseEmbedHidden(
|
||||
embeddings, hidden_states, /*batch_size*/ num_rsentries, /*seq_len*/ 1);
|
||||
hidden_states = models_[model_id]->BatchDecodeToLastHidden(fused_embedding_hidden_states,
|
||||
request_internal_ids);
|
||||
Tensor logits;
|
||||
if (models_[model_id]->CanGetLogits()) {
|
||||
logits = models_[model_id]->GetLogits(hidden_states);
|
||||
} else {
|
||||
// - Use base model's head.
|
||||
logits = models_[0]->GetLogits(hidden_states);
|
||||
}
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish proposal decode");
|
||||
TVM_FFI_ICHECK_EQ(logits->ndim, 2);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[0], num_rsentries);
|
||||
|
||||
// - Update logits.
|
||||
logit_processor_->InplaceUpdateLogits(logits, generation_cfg, mstates, request_ids, nullptr,
|
||||
&mstates, &draft_token_indices);
|
||||
|
||||
// - Compute probability distributions.
|
||||
Tensor probs_on_device =
|
||||
logit_processor_->ComputeProbsFromLogits(logits, generation_cfg, request_ids);
|
||||
|
||||
// - Commit the prefix cache changes from previous round of action.
|
||||
// Note: we commit prefix cache changes here to overlap this commit with the GPU execution.
|
||||
estate->prefix_cache->CommitSequenceExtention();
|
||||
|
||||
// - Sample tokens.
|
||||
// Fill range [0, num_rsentries) into `sample_indices`.
|
||||
std::vector<int> sample_indices(num_rsentries);
|
||||
std::iota(sample_indices.begin(), sample_indices.end(), 0);
|
||||
Tensor renormalized_probs = sampler_->BatchRenormalizeProbsByTopP(
|
||||
probs_on_device, sample_indices, request_ids, generation_cfg);
|
||||
std::vector<SampleResult> sample_results = sampler_->BatchSampleTokensWithProbAfterTopP(
|
||||
renormalized_probs, sample_indices, request_ids, generation_cfg, rngs);
|
||||
TVM_FFI_ICHECK_EQ(sample_results.size(), num_rsentries);
|
||||
|
||||
// - Add draft token to the state.
|
||||
draft_token_workspace_manager_->AllocSlots(num_rsentries, &draft_token_slots_);
|
||||
models_[model_id]->ScatterDraftProbs(probs_on_device, draft_token_slots_,
|
||||
&model_workspaces_[0].draft_probs_storage);
|
||||
// No need to save hidden states as they are not used by subsequent engine actions
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
int64_t parent_idx = static_cast<int64_t>(mstates[i]->draft_output_tokens.size()) - 1;
|
||||
mstates[i]->AddDraftToken(sample_results[i], draft_token_slots_[i], parent_idx);
|
||||
}
|
||||
|
||||
auto tdraft_end = std::chrono::high_resolution_clock::now();
|
||||
estate->metrics.UpdateDraftTimeByBatchSize(
|
||||
num_rsentries, static_cast<double>((tdraft_end - tdraft_start).count()) / 1e9);
|
||||
}
|
||||
}
|
||||
|
||||
auto tend = std::chrono::high_resolution_clock::now();
|
||||
estate->metrics.engine_decode_time_sum += static_cast<double>((tend - tstart).count()) / 1e9;
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
private:
|
||||
/*! \brief Check if the input requests can be decoded under conditions. */
|
||||
bool CanDecode(int num_rsentries) {
|
||||
// The first model is not involved in draft proposal.
|
||||
for (int model_id = 1; model_id < static_cast<int>(models_.size()); ++model_id) {
|
||||
// Check if the model has enough available pages.
|
||||
int num_available_pages = models_[model_id]->GetNumAvailablePages();
|
||||
if (num_rsentries > num_available_pages) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/*! \brief The model to run draft generation in speculative decoding. */
|
||||
Array<Model> models_;
|
||||
/*! \brief The logit processor. */
|
||||
LogitProcessor logit_processor_;
|
||||
/*! \brief The sampler to sample new tokens. */
|
||||
Sampler sampler_;
|
||||
/*! \brief Workspace of each model. */
|
||||
std::vector<ModelWorkspace> model_workspaces_;
|
||||
/*! \brief The draft token workspace manager. */
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager_;
|
||||
/*! \brief The engine config. */
|
||||
EngineConfig engine_config_;
|
||||
/*! \brief Event trace recorder. */
|
||||
Optional<EventTraceRecorder> trace_recorder_;
|
||||
/*! \brief Temporary buffer to store the slots of the current draft tokens */
|
||||
std::vector<int> draft_token_slots_;
|
||||
};
|
||||
|
||||
EngineAction EngineAction::EagleBatchDraft(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder) {
|
||||
return EngineAction(tvm::ffi::make_object<EagleBatchDraftActionObj>(
|
||||
std::move(models), std::move(logit_processor), std::move(sampler),
|
||||
std::move(model_workspaces), std::move(draft_token_workspace_manager),
|
||||
std::move(engine_config), std::move(trace_recorder)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,460 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/eagle_batch_verify.cc
|
||||
*/
|
||||
|
||||
#include <cmath>
|
||||
#include <exception>
|
||||
#include <numeric>
|
||||
|
||||
#include "../../support/random.h"
|
||||
#include "../config.h"
|
||||
#include "../model.h"
|
||||
#include "../sampler/sampler.h"
|
||||
#include "action.h"
|
||||
#include "action_commons.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
/*!
|
||||
* \brief The action that runs verification for requests in the
|
||||
* `running_queue` of engine state. Preempt low-priority requests
|
||||
* accordingly when it is impossible to decode all the running requests.
|
||||
*/
|
||||
class EagleBatchVerifyActionObj : public EngineActionObj {
|
||||
public:
|
||||
explicit EagleBatchVerifyActionObj(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler, std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder)
|
||||
: models_(std::move(models)),
|
||||
logit_processor_(std::move(logit_processor)),
|
||||
sampler_(std::move(sampler)),
|
||||
model_workspaces_(std::move(model_workspaces)),
|
||||
draft_token_workspace_manager_(std::move(draft_token_workspace_manager)),
|
||||
engine_config_(std::move(engine_config)),
|
||||
trace_recorder_(std::move(trace_recorder)),
|
||||
rng_(RandomGenerator::GetInstance()) {}
|
||||
|
||||
Array<Request> Step(EngineState estate) final {
|
||||
// - Only run spec decode when there are two models (llm+ssm) and >=1 running requests.
|
||||
if (models_.size() != 2 || estate->running_queue.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
const auto& [rsentries, draft_lengths, total_draft_length] = GetDraftsToVerify(estate);
|
||||
TVM_FFI_ICHECK_EQ(rsentries.size(), draft_lengths.size());
|
||||
if (rsentries.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
auto tstart = std::chrono::high_resolution_clock::now();
|
||||
int num_rsentries = rsentries.size();
|
||||
Array<String> request_ids =
|
||||
rsentries.Map([](const RequestStateEntry& rstate) { return rstate->request->id; });
|
||||
|
||||
// - Get embedding and run verify.
|
||||
std::vector<int64_t> request_internal_ids;
|
||||
std::vector<int32_t> all_tokens_to_verify;
|
||||
Array<RequestModelState> verify_request_mstates;
|
||||
Array<RequestModelState> draft_request_mstates;
|
||||
Array<GenerationConfig> generation_cfg;
|
||||
std::vector<RandomGenerator*> rngs;
|
||||
std::vector<std::vector<SampleResult>> draft_output_tokens;
|
||||
std::vector<std::vector<int>> draft_token_indices;
|
||||
std::vector<int64_t> token_tree_parent_ptr;
|
||||
request_internal_ids.reserve(num_rsentries);
|
||||
all_tokens_to_verify.reserve(total_draft_length);
|
||||
token_tree_parent_ptr.reserve(total_draft_length);
|
||||
verify_request_mstates.reserve(num_rsentries);
|
||||
draft_request_mstates.reserve(num_rsentries);
|
||||
rngs.reserve(num_rsentries);
|
||||
generation_cfg.reserve(num_rsentries);
|
||||
draft_output_tokens.reserve(num_rsentries);
|
||||
draft_token_indices.reserve(num_rsentries);
|
||||
draft_token_slots_.clear();
|
||||
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
RequestModelState verify_mstate = rsentries[i]->mstates[verify_model_id_];
|
||||
RequestModelState draft_mstate = rsentries[i]->mstates[draft_model_id_];
|
||||
request_internal_ids.push_back(verify_mstate->internal_id);
|
||||
TVM_FFI_ICHECK(!draft_lengths.empty());
|
||||
TVM_FFI_ICHECK_EQ(draft_lengths[i], draft_mstate->draft_output_tokens.size());
|
||||
TVM_FFI_ICHECK_EQ(draft_lengths[i], draft_mstate->draft_token_slots.size());
|
||||
// the last committed token + all the draft tokens but the last one.
|
||||
all_tokens_to_verify.push_back(draft_mstate->committed_tokens.back().GetTokenId());
|
||||
draft_token_slots_.push_back(0); // placeholder for the last committed token
|
||||
token_tree_parent_ptr.push_back(-1);
|
||||
|
||||
for (int j = 0; j < static_cast<int>(draft_mstate->draft_output_tokens.size()); ++j) {
|
||||
all_tokens_to_verify.push_back(draft_mstate->draft_output_tokens[j].GetTokenId());
|
||||
draft_token_slots_.push_back(draft_mstate->draft_token_slots[j]);
|
||||
token_tree_parent_ptr.push_back(draft_mstate->draft_token_parent_idx[j] + 1);
|
||||
}
|
||||
std::vector<int> cur_draft_token_indices(draft_mstate->draft_output_tokens.size() + 1);
|
||||
std::iota(cur_draft_token_indices.begin(), cur_draft_token_indices.end(), -1);
|
||||
draft_token_indices.emplace_back(std::move(cur_draft_token_indices));
|
||||
verify_request_mstates.push_back(verify_mstate);
|
||||
draft_request_mstates.push_back(draft_mstate);
|
||||
generation_cfg.push_back(rsentries[i]->request->generation_cfg);
|
||||
rngs.push_back(&rsentries[i]->rng);
|
||||
draft_output_tokens.push_back(draft_mstate->draft_output_tokens);
|
||||
}
|
||||
|
||||
Tensor draft_probs_on_device = models_[draft_model_id_]->GatherDraftProbs(
|
||||
model_workspaces_[verify_model_id_].draft_probs_storage, draft_token_slots_,
|
||||
&model_workspaces_[verify_model_id_].draft_probs);
|
||||
|
||||
std::vector<int> cum_verify_lengths = {0};
|
||||
cum_verify_lengths.reserve(num_rsentries + 1);
|
||||
std::vector<int> verify_lengths;
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
// Add one committed token.
|
||||
verify_lengths.push_back(draft_lengths[i] + 1);
|
||||
cum_verify_lengths.push_back(cum_verify_lengths.back() + verify_lengths.back());
|
||||
}
|
||||
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start verify embedding");
|
||||
ObjectRef embeddings = models_[verify_model_id_]->TokenEmbed(
|
||||
{Shape{all_tokens_to_verify.begin(), all_tokens_to_verify.end()}});
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish verify embedding");
|
||||
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start verify");
|
||||
ObjectRef hidden_states = models_[verify_model_id_]->BatchVerifyToLastHidden(
|
||||
embeddings, request_internal_ids, verify_lengths, token_tree_parent_ptr);
|
||||
Tensor logits = models_[verify_model_id_]->GetLogits(hidden_states);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish verify");
|
||||
TVM_FFI_ICHECK_EQ(logits->ndim, 2);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[0], cum_verify_lengths.back());
|
||||
|
||||
// - Update logits.
|
||||
logit_processor_->InplaceUpdateLogits(logits, generation_cfg, verify_request_mstates,
|
||||
request_ids, &cum_verify_lengths, &draft_request_mstates,
|
||||
&draft_token_indices);
|
||||
|
||||
// - Compute probability distributions.
|
||||
Tensor probs_on_device = logit_processor_->ComputeProbsFromLogits(
|
||||
logits, generation_cfg, request_ids, &cum_verify_lengths);
|
||||
|
||||
// - 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();
|
||||
|
||||
std::vector<int> sample_indices(num_rsentries);
|
||||
std::iota(sample_indices.begin(), sample_indices.end(), 0);
|
||||
Tensor renormalized_probs = sampler_->BatchRenormalizeProbsByTopP(
|
||||
probs_on_device, sample_indices, request_ids, generation_cfg);
|
||||
auto [sample_results_arr, _] = sampler_->BatchVerifyDraftTokensWithProbAfterTopP(
|
||||
renormalized_probs, request_ids, cum_verify_lengths, generation_cfg, rngs,
|
||||
draft_output_tokens, token_tree_parent_ptr, draft_probs_on_device);
|
||||
TVM_FFI_ICHECK_EQ(sample_results_arr.size(), num_rsentries);
|
||||
|
||||
// We collect the requests whose drafts are fully accepted.
|
||||
// When a request's draft is fully accepted, there is an extra token proposed
|
||||
// by the draft model but not added into the draft model's KV cache.
|
||||
// In this case, an additional batch decode step is needed for these requests.
|
||||
std::vector<int64_t> fully_accepted_rsentries;
|
||||
std::vector<int64_t> verify_model_seq_internal_ids;
|
||||
std::vector<int64_t> accepted_token_tree_leaf_nodes;
|
||||
fully_accepted_rsentries.reserve(num_rsentries);
|
||||
verify_model_seq_internal_ids.reserve(num_rsentries);
|
||||
accepted_token_tree_leaf_nodes.reserve(num_rsentries);
|
||||
|
||||
std::vector<int> last_accepted_hidden_positions;
|
||||
last_accepted_hidden_positions.reserve(num_rsentries);
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
const std::vector<SampleResult>& sample_results = sample_results_arr[i];
|
||||
int accept_length = sample_results.size();
|
||||
TVM_FFI_ICHECK_GE(accept_length, 1);
|
||||
for (SampleResult sample_result : sample_results) {
|
||||
rsentries[i]->mstates[verify_model_id_]->CommitToken(sample_result);
|
||||
rsentries[i]->mstates[draft_model_id_]->CommitToken(sample_result);
|
||||
}
|
||||
// Metrics update
|
||||
// live update the output metrics
|
||||
rsentries[i]->rstate->metrics.completion_tokens += accept_length;
|
||||
rsentries[i]->rstate->metrics.decode_tokens += accept_length;
|
||||
estate->metrics.spec_decode.Update(cum_verify_lengths[i + 1] - cum_verify_lengths[i],
|
||||
accept_length);
|
||||
// - Minus one because the last draft token has no kv cache entry
|
||||
// - Take max with 0 in case of all accepted.
|
||||
int rollback_length =
|
||||
std::max(cum_verify_lengths[i + 1] - cum_verify_lengths[i] - accept_length, 0);
|
||||
|
||||
// Commit accepted tokens to the "verify_model", rollback kv cache
|
||||
// in the "draft_model".
|
||||
// NOTE: when number of small models is more than 1 (in the future),
|
||||
// it is possible to re-compute prefill for the small models.
|
||||
verify_model_seq_internal_ids.push_back(rsentries[i]->mstates[verify_model_id_]->internal_id);
|
||||
accepted_token_tree_leaf_nodes.push_back(accept_length - 1);
|
||||
if (rollback_length > 0) {
|
||||
// Draft model rollback minus one because verify uses one more token.
|
||||
models_[draft_model_id_]->PopNFromKVCache(
|
||||
rsentries[i]->mstates[draft_model_id_]->internal_id, rollback_length - 1);
|
||||
} else {
|
||||
fully_accepted_rsentries.push_back(i);
|
||||
}
|
||||
// clear the draft model state entries
|
||||
rsentries[i]->mstates[draft_model_id_]->RemoveAllDraftTokens(&draft_token_slots_);
|
||||
draft_token_workspace_manager_->FreeSlots(draft_token_slots_);
|
||||
// - Slice and save hidden_states_for_sample
|
||||
last_accepted_hidden_positions.push_back(cum_verify_lengths[i] + accept_length - 1);
|
||||
}
|
||||
models_[verify_model_id_]->CommitAcceptedTokenTreeNodesToKVCache(
|
||||
verify_model_seq_internal_ids, accepted_token_tree_leaf_nodes);
|
||||
if (!fully_accepted_rsentries.empty() &&
|
||||
engine_config_->speculative_mode == SpeculativeMode::kEagle) {
|
||||
// - Run a step of batch decode for requests whose drafts are fully accepted.
|
||||
// When a request's draft is fully accepted, there is an extra token proposed
|
||||
// by the draft model but not added into the draft model's KV cache.
|
||||
// In this case, an additional batch decode step is needed for these requests.
|
||||
std::vector<int> input_tokens;
|
||||
std::vector<int64_t> fully_accepted_request_internal_ids;
|
||||
input_tokens.reserve(fully_accepted_rsentries.size());
|
||||
fully_accepted_request_internal_ids.reserve(fully_accepted_rsentries.size());
|
||||
|
||||
std::vector<int> hidden_states_positions_for_fully_accepted;
|
||||
hidden_states_positions_for_fully_accepted.reserve(fully_accepted_rsentries.size());
|
||||
|
||||
for (int rsentry_id : fully_accepted_rsentries) {
|
||||
int num_committed_tokens =
|
||||
rsentries[rsentry_id]->mstates[verify_model_id_]->committed_tokens.size();
|
||||
// When a request's draft is fully accepted, an additional new token is sampled.
|
||||
// So the token needed to fill in the draft model is the committed_token[-2].
|
||||
TVM_FFI_ICHECK_GE(num_committed_tokens, 2);
|
||||
input_tokens.push_back(rsentries[rsentry_id]
|
||||
->mstates[verify_model_id_]
|
||||
->committed_tokens[num_committed_tokens - 2]
|
||||
.GetTokenId());
|
||||
|
||||
// Taking the hidden states of the token before the last token
|
||||
hidden_states_positions_for_fully_accepted.push_back(
|
||||
last_accepted_hidden_positions[rsentry_id] - 1);
|
||||
fully_accepted_request_internal_ids.push_back(
|
||||
rsentries[rsentry_id]->mstates[draft_model_id_]->internal_id);
|
||||
}
|
||||
|
||||
// - Compute embeddings.
|
||||
ObjectRef embeddings =
|
||||
models_[draft_model_id_]->TokenEmbed({Shape{input_tokens.begin(), input_tokens.end()}});
|
||||
// - Gather hidden states
|
||||
ObjectRef hidden_states_for_fully_accepted = models_[draft_model_id_]->GatherHiddenStates(
|
||||
hidden_states, hidden_states_positions_for_fully_accepted,
|
||||
&model_workspaces_[draft_model_id_].hidden_states);
|
||||
// - Invoke model decode.
|
||||
ObjectRef fused_embedding_hidden_states = models_[draft_model_id_]->FuseEmbedHidden(
|
||||
embeddings, hidden_states_for_fully_accepted,
|
||||
/*batch_size*/ fully_accepted_rsentries.size(), /*seq_len*/ 1);
|
||||
hidden_states_for_fully_accepted = models_[draft_model_id_]->BatchDecodeToLastHidden(
|
||||
fused_embedding_hidden_states, fully_accepted_request_internal_ids);
|
||||
// - We explicitly synchronize to avoid the input tokens getting overriden in the
|
||||
// next runs of BatchDecode.
|
||||
// This is because we do not do sample for this round of batch decode.
|
||||
if (hidden_states_for_fully_accepted->IsInstance<DRefObj>()) {
|
||||
(hidden_states_for_fully_accepted.as_or_throw<DRef>()->session)
|
||||
.as_or_throw<Session>()
|
||||
->SyncWorker(0);
|
||||
} else {
|
||||
Tensor hidden_states_for_fully_accepted_nd =
|
||||
hidden_states_for_fully_accepted.as_or_throw<Tensor>();
|
||||
DeviceAPI::Get(hidden_states_for_fully_accepted_nd->device)
|
||||
->StreamSync(hidden_states_for_fully_accepted_nd->device, nullptr);
|
||||
}
|
||||
}
|
||||
{
|
||||
// One step draft for the following steps
|
||||
|
||||
// Gather hidden states for the last accepted tokens.
|
||||
// Use the function and the workspace of the verify model because the information about the
|
||||
// hidden states is not available in the draft model for medusa.
|
||||
hidden_states = models_[0]->GatherHiddenStates(hidden_states, last_accepted_hidden_positions,
|
||||
&model_workspaces_[0].hidden_states);
|
||||
|
||||
std::vector<int> input_tokens;
|
||||
Array<RequestModelState> mstates;
|
||||
input_tokens.reserve(num_rsentries);
|
||||
mstates.reserve(num_rsentries);
|
||||
for (const RequestStateEntry& rsentry : rsentries) {
|
||||
mstates.push_back(rsentry->mstates[draft_model_id_]);
|
||||
}
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
TVM_FFI_ICHECK(!mstates[i]->committed_tokens.empty());
|
||||
input_tokens.push_back(mstates[i]->committed_tokens.back().GetTokenId());
|
||||
}
|
||||
|
||||
Array<Tensor> multi_step_logits{nullptr}; // for medusa output
|
||||
if (engine_config_->speculative_mode == SpeculativeMode::kEagle) {
|
||||
// - Compute embeddings.
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start proposal embedding");
|
||||
embeddings =
|
||||
models_[draft_model_id_]->TokenEmbed({Shape{input_tokens.begin(), input_tokens.end()}});
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish proposal embedding");
|
||||
|
||||
// - Invoke model decode.
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start proposal decode");
|
||||
ObjectRef fused_embedding_hidden_states = models_[draft_model_id_]->FuseEmbedHidden(
|
||||
embeddings, hidden_states, /*batch_size*/ num_rsentries, /*seq_len*/ 1);
|
||||
hidden_states = models_[draft_model_id_]->BatchDecodeToLastHidden(
|
||||
fused_embedding_hidden_states, request_internal_ids);
|
||||
|
||||
int lm_head_model_id = models_[draft_model_id_]->CanGetLogits() ? draft_model_id_ : 0;
|
||||
logits = models_[lm_head_model_id]->GetLogits(hidden_states);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish proposal decode");
|
||||
TVM_FFI_ICHECK_EQ(logits->ndim, 2);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[0], num_rsentries);
|
||||
} else if (engine_config_->speculative_mode == SpeculativeMode::kMedusa) {
|
||||
multi_step_logits = models_[draft_model_id_]->GetMultiStepLogits(hidden_states);
|
||||
}
|
||||
|
||||
// Fill range [0, num_rsentries) into `sample_indices`.
|
||||
std::vector<int> sample_indices(num_rsentries);
|
||||
std::iota(sample_indices.begin(), sample_indices.end(), 0);
|
||||
|
||||
if (engine_config_->speculative_mode == SpeculativeMode::kEagle) {
|
||||
const auto& [renormalized_probs, sample_results] = ApplyLogitProcessorAndSample(
|
||||
logit_processor_, sampler_, logits, generation_cfg, request_ids, mstates, rngs,
|
||||
sample_indices, generation_cfg, request_ids, sample_indices);
|
||||
UpdateRequestStatesWithDraftProposals(mstates, sample_results, draft_model_id_,
|
||||
renormalized_probs, hidden_states, estate);
|
||||
} else if (engine_config_->speculative_mode == SpeculativeMode::kMedusa) {
|
||||
TVM_FFI_ICHECK_NE(estate->spec_draft_length, 0);
|
||||
for (int draft_id = 0; draft_id < estate->spec_draft_length; draft_id++) {
|
||||
const auto& [renormalized_probs, sample_results] = ApplyLogitProcessorAndSample(
|
||||
logit_processor_, sampler_, multi_step_logits[draft_id], generation_cfg, request_ids,
|
||||
mstates, rngs, sample_indices, generation_cfg, request_ids, sample_indices);
|
||||
UpdateRequestStatesWithDraftProposals(mstates, sample_results, draft_model_id_,
|
||||
renormalized_probs, hidden_states, estate);
|
||||
}
|
||||
}
|
||||
}
|
||||
// reset num_tokens_for_next_decode
|
||||
for (const RequestStateEntry& rsentry : rsentries) {
|
||||
rsentry->mstates[verify_model_id_]->num_tokens_for_next_decode = 0;
|
||||
rsentry->mstates[draft_model_id_]->num_tokens_for_next_decode = 0;
|
||||
}
|
||||
auto tend = std::chrono::high_resolution_clock::now();
|
||||
double elapsed_time = static_cast<double>((tend - tstart).count()) / 1e9;
|
||||
estate->metrics.engine_decode_time_sum += elapsed_time;
|
||||
estate->metrics.UpdateVerifyTimeByBatchSize(cum_verify_lengths.back(), elapsed_time);
|
||||
|
||||
return estate->running_queue;
|
||||
}
|
||||
|
||||
private:
|
||||
struct DraftRequestStateEntries {
|
||||
/*! \brief The request state entries to verify. */
|
||||
Array<RequestStateEntry> draft_rsentries;
|
||||
/*! \brief The draft length of each request state. */
|
||||
std::vector<int> draft_lengths;
|
||||
/*! \brief The total draft length. */
|
||||
int total_draft_length;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Decide whether to run verify for the draft of each request.
|
||||
* \param estate The engine state.
|
||||
* \return The drafts to verify, together with their respective
|
||||
* state and input length.
|
||||
*/
|
||||
DraftRequestStateEntries GetDraftsToVerify(EngineState estate) {
|
||||
std::vector<int> draft_lengths;
|
||||
int total_draft_length = 0;
|
||||
int total_required_pages = 0;
|
||||
int num_available_pages = models_[verify_model_id_]->GetNumAvailablePages();
|
||||
|
||||
// Preempt the request state entries that cannot fit the large model for verification.
|
||||
std::vector<RequestStateEntry> running_rsentries = estate->GetRunningRequestStateEntries();
|
||||
std::vector<int> num_page_requirement;
|
||||
num_page_requirement.reserve(running_rsentries.size());
|
||||
for (const RequestStateEntry& rsentry : running_rsentries) {
|
||||
int draft_length = rsentry->mstates[draft_model_id_]->draft_output_tokens.size();
|
||||
int num_require_pages = (draft_length + engine_config_->kv_cache_page_size - 1) /
|
||||
engine_config_->kv_cache_page_size;
|
||||
draft_lengths.push_back(draft_length);
|
||||
num_page_requirement.push_back(num_require_pages);
|
||||
total_draft_length += draft_length;
|
||||
total_required_pages += num_require_pages;
|
||||
}
|
||||
while (!CanVerify(total_required_pages)) {
|
||||
if (estate->prefix_cache->TryFreeMemory()) continue;
|
||||
RequestStateEntry preempted = PreemptLastRunningRequestStateEntry(
|
||||
estate, models_, draft_token_workspace_manager_, trace_recorder_);
|
||||
if (preempted.same_as(running_rsentries.back())) {
|
||||
total_draft_length -= draft_lengths.back();
|
||||
total_required_pages -= num_page_requirement.back();
|
||||
draft_lengths.pop_back();
|
||||
num_page_requirement.pop_back();
|
||||
running_rsentries.pop_back();
|
||||
}
|
||||
}
|
||||
|
||||
return {running_rsentries, draft_lengths, total_draft_length};
|
||||
}
|
||||
|
||||
bool CanVerify(int num_required_pages) {
|
||||
int num_available_pages = models_[0]->GetNumAvailablePages();
|
||||
return num_required_pages <= num_available_pages;
|
||||
}
|
||||
|
||||
void UpdateRequestStatesWithDraftProposals(const Array<RequestModelState>& mstates,
|
||||
const std::vector<SampleResult>& sample_results,
|
||||
int model_id, const Tensor& renormalized_probs,
|
||||
const ObjectRef& hidden_states_for_sample,
|
||||
EngineState estate) {
|
||||
draft_token_workspace_manager_->AllocSlots(mstates.size(), &draft_token_slots_);
|
||||
models_[0]->ScatterDraftProbs(renormalized_probs, draft_token_slots_,
|
||||
&model_workspaces_[0].draft_probs_storage);
|
||||
if (engine_config_->speculative_mode == SpeculativeMode::kEagle &&
|
||||
estate->spec_draft_length > 1) {
|
||||
models_[0]->ScatterHiddenStates(hidden_states_for_sample, draft_token_slots_,
|
||||
&model_workspaces_[0].draft_hidden_states_storage);
|
||||
}
|
||||
for (int i = 0; i < static_cast<int>(mstates.size()); ++i) {
|
||||
int64_t parent_idx = static_cast<int64_t>(mstates[i]->draft_output_tokens.size()) - 1;
|
||||
mstates[i]->AddDraftToken(sample_results[i], draft_token_slots_[i], parent_idx);
|
||||
}
|
||||
}
|
||||
/*!
|
||||
* \brief The model to run decode in. When there are multiple
|
||||
* models, the `Step` function of the created action will not take effect.
|
||||
*/
|
||||
Array<Model> models_;
|
||||
/*! \brief The logit processor. */
|
||||
LogitProcessor logit_processor_;
|
||||
/*! \brief The sampler to sample new tokens. */
|
||||
Sampler sampler_;
|
||||
/*! \brief Workspace of each model. */
|
||||
std::vector<ModelWorkspace> model_workspaces_;
|
||||
/*! \brief The draft token workspace manager. */
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager_;
|
||||
/*! \brief The engine config. */
|
||||
EngineConfig engine_config_;
|
||||
/*! \brief Event trace recorder. */
|
||||
Optional<EventTraceRecorder> trace_recorder_;
|
||||
/*! \brief Random number generator. */
|
||||
RandomGenerator& rng_;
|
||||
/*! \brief The ids of verify/draft models. */
|
||||
const int verify_model_id_ = 0;
|
||||
const int draft_model_id_ = 1;
|
||||
const float eps_ = 1e-5;
|
||||
/*! \brief Temporary buffer to store the slots of the current draft tokens */
|
||||
std::vector<int> draft_token_slots_;
|
||||
};
|
||||
|
||||
EngineAction EngineAction::EagleBatchVerify(
|
||||
Array<Model> models, LogitProcessor logit_processor, Sampler sampler,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager, EngineConfig engine_config,
|
||||
Optional<EventTraceRecorder> trace_recorder) {
|
||||
return EngineAction(tvm::ffi::make_object<EagleBatchVerifyActionObj>(
|
||||
std::move(models), std::move(logit_processor), std::move(sampler),
|
||||
std::move(model_workspaces), std::move(draft_token_workspace_manager),
|
||||
std::move(engine_config), std::move(trace_recorder)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,502 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_actions/eagle_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.
|
||||
*/
|
||||
class EagleNewRequestPrefillActionObj : public BatchPrefillBaseActionObj {
|
||||
public:
|
||||
explicit EagleNewRequestPrefillActionObj(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager,
|
||||
EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
Optional<EventTraceRecorder> trace_recorder)
|
||||
: BatchPrefillBaseActionObj(std::move(models), std::move(engine_config),
|
||||
std::move(model_configs), std::move(trace_recorder)),
|
||||
logit_processor_(std::move(logit_processor)),
|
||||
sampler_(std::move(sampler)),
|
||||
model_workspaces_(std::move(model_workspaces)),
|
||||
draft_token_workspace_manager_(std::move(draft_token_workspace_manager)) {}
|
||||
|
||||
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("NewRequestPrefill getting requests");
|
||||
prefill_inputs = GetRequestStateEntriesToPrefill(estate);
|
||||
if (prefill_inputs.empty()) {
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
int num_rsentries = prefill_inputs.size();
|
||||
{
|
||||
NVTXScopedRange nvtx_scope("NewRequestPrefill 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.
|
||||
std::vector<int> prefill_lengths;
|
||||
prefill_lengths.resize(/*size=*/num_rsentries, /*value=*/-1);
|
||||
ObjectRef hidden_states_for_input{nullptr};
|
||||
ObjectRef hidden_states_for_sample{nullptr};
|
||||
Tensor logits_for_sample{nullptr};
|
||||
// A map used to record the entry and child_idx pair needed to fork sequence.
|
||||
// The base model (id 0) should record all the pairs and all the small models
|
||||
// fork sequences according to this map.
|
||||
std::unordered_map<int, std::unordered_set<int>> fork_rsentry_child_map;
|
||||
std::vector<bool> extra_prefill_tokens;
|
||||
prefill_lengths.resize(/*size=*/num_rsentries, /*value=*/false);
|
||||
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;
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
const RequestStateEntry& rsentry = prefill_inputs[i].rsentry;
|
||||
RequestModelState mstate = rsentry->mstates[model_id];
|
||||
TVM_FFI_ICHECK(mstate->draft_output_tokens.empty());
|
||||
TVM_FFI_ICHECK(mstate->draft_token_slots.empty());
|
||||
if (status_before_prefill[i] == RequestStateStatus::kPending) {
|
||||
if (!estate->prefix_cache->HasSequence(mstate->internal_id)) {
|
||||
// Add the sequence to the model, or fork the sequence from its parent.
|
||||
// If the sequence is already in prefix cache, it has also been added/forked in the
|
||||
// KVCache.
|
||||
if (rsentry->parent_idx == -1) {
|
||||
models_[model_id]->AddNewSequence(mstate->internal_id);
|
||||
} else {
|
||||
models_[model_id]->ForkSequence(rstates_of_entries[i]
|
||||
->entries[rsentry->parent_idx]
|
||||
->mstates[model_id]
|
||||
->internal_id,
|
||||
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);
|
||||
}
|
||||
// Shift the input tokens by 1 for eagle models.
|
||||
if (model_id == 0) {
|
||||
for (int j = 1; j < static_cast<int>(models_.size()); ++j) {
|
||||
TVM_FFI_ICHECK(rsentry->mstates[j]->inputs.size());
|
||||
TokenData token_data = rsentry->mstates[j]->inputs[0].as_or_throw<TokenData>();
|
||||
rsentry->mstates[j]->inputs.Set(0, TokenData(Shape(token_data->token_ids.begin() + 1,
|
||||
token_data->token_ids.end())));
|
||||
}
|
||||
}
|
||||
}
|
||||
request_internal_ids.push_back(mstate->internal_id);
|
||||
|
||||
if (engine_config_->speculative_mode == SpeculativeMode::kMedusa && model_id > 0) {
|
||||
// Embedding is only needed for the base model in Medusa.
|
||||
continue;
|
||||
}
|
||||
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;
|
||||
|
||||
RECORD_EVENT(trace_recorder_, prefill_inputs[i].rsentry->request->id, "start embedding");
|
||||
// Speculative models shift left the input tokens by 1 when base model has committed tokens.
|
||||
// Note: for n > 1 cases Eagle doesn't work because parent entry doesn't shift input tokens.
|
||||
for (int j = 0; j < static_cast<int>(input_data.size()); ++j) {
|
||||
if (model_id == 0) {
|
||||
mstate->prefilled_inputs.push_back(input_data[j]);
|
||||
}
|
||||
embeddings = input_data[j]->GetEmbedding(
|
||||
models_[model_id],
|
||||
/*dst=*/!single_input ? &model_workspaces_[model_id].embeddings : nullptr,
|
||||
/*offset=*/cum_prefill_length);
|
||||
cum_prefill_length += input_data[j]->GetLength();
|
||||
}
|
||||
RECORD_EVENT(trace_recorder_, rsentry->request->id, "finish embedding");
|
||||
}
|
||||
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start prefill");
|
||||
|
||||
Array<Tensor> multi_step_logits{nullptr};
|
||||
|
||||
if (model_id == 0 || engine_config_->speculative_mode == SpeculativeMode::kEagle) {
|
||||
ObjectRef embedding_or_hidden_states{nullptr};
|
||||
if (model_id == 0) {
|
||||
embedding_or_hidden_states = embeddings;
|
||||
} else {
|
||||
embedding_or_hidden_states =
|
||||
models_[model_id]->FuseEmbedHidden(embeddings, hidden_states_for_input,
|
||||
/*batch_size*/ 1, /*seq_len*/ cum_prefill_length);
|
||||
}
|
||||
// hidden_states: (b * s, h)
|
||||
ObjectRef hidden_states = models_[model_id]->BatchPrefillToLastHidden(
|
||||
embedding_or_hidden_states, request_internal_ids, prefill_lengths);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish prefill");
|
||||
|
||||
if (model_id == 0) {
|
||||
// We only need to sample for model 0 in prefill.
|
||||
hidden_states_for_input = hidden_states;
|
||||
|
||||
// - 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();
|
||||
}
|
||||
|
||||
// Whether to use base model to get logits.
|
||||
int sample_model_id = !models_[model_id]->CanGetLogits() ? 0 : model_id;
|
||||
|
||||
std::vector<int> logit_positions;
|
||||
{
|
||||
// Prepare the logit positions
|
||||
logit_positions.reserve(prefill_lengths.size());
|
||||
int total_len = 0;
|
||||
for (int i = 0; i < prefill_lengths.size(); ++i) {
|
||||
total_len += prefill_lengths[i];
|
||||
logit_positions.push_back(total_len - 1);
|
||||
}
|
||||
}
|
||||
// hidden_states_for_sample: (b * s, h)
|
||||
hidden_states_for_sample = models_[sample_model_id]->GatherHiddenStates(
|
||||
hidden_states, logit_positions, &model_workspaces_[model_id].hidden_states);
|
||||
// logits_for_sample: (b * s, v)
|
||||
logits_for_sample = models_[sample_model_id]->GetLogits(hidden_states_for_sample);
|
||||
} else if (engine_config_->speculative_mode == SpeculativeMode::kMedusa) {
|
||||
// Note: spec_draft_length in engine config has to be match the model config in Medusa.
|
||||
multi_step_logits = models_[model_id]->GetMultiStepLogits(hidden_states_for_sample);
|
||||
} else {
|
||||
LOG(FATAL) << "unreachable";
|
||||
}
|
||||
|
||||
Array<String> child_request_ids;
|
||||
// - Prepare the configurations for the sampler.
|
||||
// For prefill_inputs which have children, sample
|
||||
// one token for each rstate that is depending.
|
||||
// Otherwise, sample a token for the current rstate.
|
||||
std::vector<int> child_sample_indices;
|
||||
std::vector<RequestStateEntry> rsentries_for_sample;
|
||||
std::vector<RandomGenerator*> rngs;
|
||||
std::vector<bool> rsentry_activated;
|
||||
Array<GenerationConfig> child_generation_cfg;
|
||||
child_sample_indices.reserve(num_rsentries);
|
||||
child_generation_cfg.reserve(num_rsentries);
|
||||
child_request_ids.reserve(num_rsentries);
|
||||
rsentries_for_sample.reserve(num_rsentries);
|
||||
rngs.reserve(num_rsentries);
|
||||
rsentry_activated.reserve(num_rsentries);
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
const RequestStateEntry& rsentry = prefill_inputs[i].rsentry;
|
||||
// No sample for rsentries with remaining inputs.
|
||||
if (!rsentry->mstates[0]->inputs.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
int remaining_num_child_to_activate = prefill_inputs[i].num_child_to_activate;
|
||||
for (int child_idx : rsentry->child_indices) {
|
||||
// Only use base model to judge if we need to add child entries.
|
||||
if ((rstates_of_entries[i]->entries[child_idx]->status == RequestStateStatus::kPending &&
|
||||
rstates_of_entries[i]
|
||||
->entries[child_idx]
|
||||
->mstates[0]
|
||||
->committed_tokens.empty() ||
|
||||
fork_rsentry_child_map[i].count(child_idx))) {
|
||||
// If rstates_of_entries[i]->entries[child_idx] has no committed token,
|
||||
// the prefill of the current rsentry will unblock
|
||||
// rstates_of_entries[i]->entries[child_idx],
|
||||
// and thus we want to sample a token for rstates_of_entries[i]->entries[child_idx].
|
||||
fork_rsentry_child_map[i].insert(child_idx);
|
||||
child_sample_indices.push_back(i);
|
||||
rsentries_for_sample.push_back(rstates_of_entries[i]->entries[child_idx]);
|
||||
child_request_ids.push_back(rsentry->request->id);
|
||||
child_generation_cfg.push_back(rsentry->request->generation_cfg);
|
||||
rngs.push_back(&rstates_of_entries[i]->entries[child_idx]->rng);
|
||||
|
||||
// We only fork the first `num_child_to_activate` children.
|
||||
// The children not being forked will be forked via later prefills.
|
||||
// Usually `num_child_to_activate` is the same as the number of children.
|
||||
// But it can be fewer subject to the KV cache max num sequence limit.
|
||||
if (remaining_num_child_to_activate == 0) {
|
||||
rsentry_activated.push_back(false);
|
||||
continue;
|
||||
}
|
||||
rsentry_activated.push_back(true);
|
||||
--remaining_num_child_to_activate;
|
||||
if (model_id == 0) {
|
||||
TVM_FFI_ICHECK(rstates_of_entries[i]->entries[child_idx]->status ==
|
||||
RequestStateStatus::kPending);
|
||||
rstates_of_entries[i]->entries[child_idx]->status = RequestStateStatus::kAlive;
|
||||
}
|
||||
int64_t child_internal_id =
|
||||
rstates_of_entries[i]->entries[child_idx]->mstates[model_id]->internal_id;
|
||||
models_[model_id]->ForkSequence(rsentry->mstates[model_id]->internal_id,
|
||||
child_internal_id);
|
||||
// Enable sliding window for the child sequence if the child is not a parent.
|
||||
if (rstates_of_entries[i]->entries[child_idx]->child_indices.empty()) {
|
||||
models_[model_id]->EnableSlidingWindowForSeq(child_internal_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (rsentry->child_indices.empty()) {
|
||||
// If rsentry has no child, we sample a token for itself.
|
||||
child_sample_indices.push_back(i);
|
||||
rsentries_for_sample.push_back(rsentry);
|
||||
child_request_ids.push_back(rsentry->request->id);
|
||||
child_generation_cfg.push_back(rsentry->request->generation_cfg);
|
||||
rngs.push_back(&rsentry->rng);
|
||||
rsentry_activated.push_back(true);
|
||||
}
|
||||
}
|
||||
|
||||
// - Prepare input for logit processor.
|
||||
TVM_FFI_ICHECK(logits_for_sample.defined());
|
||||
Array<GenerationConfig> generation_cfg;
|
||||
Array<RequestModelState> mstates_for_logitproc;
|
||||
std::vector<int> sample_indices(num_rsentries);
|
||||
generation_cfg.reserve(num_rsentries);
|
||||
mstates_for_logitproc.reserve(num_rsentries);
|
||||
std::iota(sample_indices.begin(), sample_indices.end(), 0);
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
generation_cfg.push_back(prefill_inputs[i].rsentry->request->generation_cfg);
|
||||
mstates_for_logitproc.push_back(prefill_inputs[i].rsentry->mstates[model_id]);
|
||||
}
|
||||
if (model_id == 0 || engine_config_->speculative_mode == SpeculativeMode::kEagle) {
|
||||
const auto& [renormalized_probs, sample_results] = ApplyLogitProcessorAndSample(
|
||||
logit_processor_, sampler_, logits_for_sample, generation_cfg, request_ids,
|
||||
mstates_for_logitproc, rngs, sample_indices, child_generation_cfg, child_request_ids,
|
||||
child_sample_indices);
|
||||
if (model_id == 0) {
|
||||
UpdateRequestStateEntriesWithSampleResults(rsentries_for_sample, rsentry_activated,
|
||||
sample_results);
|
||||
// Add the sampled token as an input of the eagle models.
|
||||
if (engine_config_->speculative_mode == SpeculativeMode::kEagle) {
|
||||
for (int i = 0; i < static_cast<int>(rsentries_for_sample.size()); ++i) {
|
||||
for (int mid = 1; mid < static_cast<int>(models_.size()); ++mid) {
|
||||
TokenData token_data =
|
||||
rsentries_for_sample[i]->mstates[mid]->inputs.back().as_or_throw<TokenData>();
|
||||
std::vector<int32_t> token_ids = {token_data->token_ids.begin(),
|
||||
token_data->token_ids.end()};
|
||||
token_ids.push_back(sample_results[i].GetTokenId());
|
||||
int ninputs =
|
||||
static_cast<int>(rsentries_for_sample[i]->mstates[mid]->inputs.size());
|
||||
rsentries_for_sample[i]->mstates[mid]->inputs.Set(
|
||||
ninputs - 1, TokenData(Shape(token_ids.begin(), token_ids.end())));
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// - Slice and save hidden_states_for_sample
|
||||
UpdateRequestStatesWithDraftProposals(rsentries_for_sample, sample_results, model_id,
|
||||
renormalized_probs, hidden_states_for_sample,
|
||||
estate, child_sample_indices);
|
||||
}
|
||||
} else if (engine_config_->speculative_mode == SpeculativeMode::kMedusa) {
|
||||
TVM_FFI_ICHECK_NE(estate->spec_draft_length, 0);
|
||||
for (int draft_id = 0; draft_id < estate->spec_draft_length; ++draft_id) {
|
||||
const auto& [renormalized_probs, sample_results] = ApplyLogitProcessorAndSample(
|
||||
logit_processor_, sampler_, multi_step_logits[draft_id], generation_cfg, request_ids,
|
||||
mstates_for_logitproc, rngs, sample_indices, child_generation_cfg, child_request_ids,
|
||||
child_sample_indices);
|
||||
|
||||
UpdateRequestStatesWithDraftProposals(
|
||||
rsentries_for_sample, sample_results, model_id, renormalized_probs,
|
||||
/*hidden_states=*/ObjectRef{nullptr}, estate, child_sample_indices);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
void UpdateRequestStatesWithDraftProposals(
|
||||
const std::vector<RequestStateEntry>& rsentries_for_sample,
|
||||
const std::vector<SampleResult>& sample_results, int model_id,
|
||||
const Tensor& renormalized_probs, const ObjectRef& hidden_states_for_sample,
|
||||
EngineState estate, const std::vector<int>& sample_indices) {
|
||||
std::vector<int> reuse_count(renormalized_probs->shape[0], 0);
|
||||
for (int i = 0; i < static_cast<int>(sample_indices.size()); ++i) {
|
||||
// The same probability may be sampled multiple times.
|
||||
reuse_count[sample_indices[i]]++;
|
||||
}
|
||||
draft_token_workspace_manager_->AllocSlots(renormalized_probs->shape[0], reuse_count,
|
||||
&draft_token_slots_);
|
||||
|
||||
models_[0]->ScatterDraftProbs(renormalized_probs, draft_token_slots_,
|
||||
&model_workspaces_[0].draft_probs_storage);
|
||||
if (engine_config_->speculative_mode == SpeculativeMode::kEagle &&
|
||||
estate->spec_draft_length > 1) {
|
||||
models_[0]->ScatterHiddenStates(hidden_states_for_sample, draft_token_slots_,
|
||||
&model_workspaces_[0].draft_hidden_states_storage);
|
||||
}
|
||||
for (int i = 0; i < static_cast<int>(rsentries_for_sample.size()); ++i) {
|
||||
int parent_idx =
|
||||
rsentries_for_sample[i]->mstates[model_id]->draft_output_tokens.empty()
|
||||
? -1
|
||||
: rsentries_for_sample[i]->mstates[model_id]->draft_output_tokens.size() - 1;
|
||||
rsentries_for_sample[i]->mstates[model_id]->AddDraftToken(
|
||||
sample_results[i], draft_token_slots_[sample_indices[i]], parent_idx);
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
/*! \brief The logit processor. */
|
||||
LogitProcessor logit_processor_;
|
||||
/*! \brief The sampler to sample new tokens. */
|
||||
Sampler sampler_;
|
||||
/*! \brief Workspace of each model. */
|
||||
std::vector<ModelWorkspace> model_workspaces_;
|
||||
/*! \brief The draft token workspace manager. */
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager_;
|
||||
/*! \brief Temporary buffer to store the slots of the current draft tokens */
|
||||
std::vector<int> draft_token_slots_;
|
||||
|
||||
/*!
|
||||
* \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.
|
||||
*/
|
||||
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.
|
||||
// Note: Almost same as without eagle speculative decoding. But in prefill step, the
|
||||
// prefill embedding input in draft model will be shifted one token, compared to the base
|
||||
// model. Just the new sequence without prefix cache. Here we merely add the new sequence
|
||||
// in advance of prefill step.
|
||||
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 i = 0; i < models_.size(); ++i) {
|
||||
models_[i]->AddNewSequence(rsentry->mstates[0]->internal_id);
|
||||
// Enable sliding window for the sequence if it is not a parent.
|
||||
if (rsentry->child_indices.empty()) {
|
||||
models_[i]->EnableSlidingWindowForSeq(rsentry->mstates[0]->internal_id);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (result.forked_seq_id != -1) {
|
||||
// Fork from active sequence
|
||||
// Note: Due to the shifted KVCache between base model and draft model, we do a trick
|
||||
// over forking sequence:
|
||||
// For example. we have a sequence of [0, 1, 2] in base model KVCache, and the
|
||||
// corresponding sequence of [1, 2, 3] in draft model KVCache, where token [3] was
|
||||
// sampled from base model, but not appended in base model KVCache. Then we get a new
|
||||
// sequence [0, 1, 4] to prefill. Although the new sequence matches first two tokens
|
||||
// with the sequence [0, 1, 2], we have to fork from the first token 0, not the second
|
||||
// token 1. Because if we fork from the second token, we will prefill like: Base model:
|
||||
// [0, 1] + prefill([4]) => [5] Draft model: [1] + prefill([4, 5]) The lengths to
|
||||
// prefill is different between base model and draft model, which is illegal. So we roll
|
||||
// back one token in prefix cache to fork from the first token. Then the prefill will be
|
||||
// like: Base model: [0] + prefill([1, 4]) => [5] Draft model: [1] + prefill([4, 5]) And
|
||||
// we shift the input prefill data as other new sequence, to avoid double prefilling
|
||||
// token 1, and make the prefill length aligned between base model and draft model.
|
||||
TVM_FFI_ICHECK_EQ(result.reused_seq_id, -1);
|
||||
TVM_FFI_ICHECK_EQ(result.reused_seq_pop_last_tokens, 0);
|
||||
estate->prefix_cache->RollBackSequence(rsentry->mstates[0]->internal_id, 1);
|
||||
for (int i = 0; i < models_.size(); ++i) {
|
||||
models_[i]->ForkSequence(result.forked_seq_id, rsentry->mstates[0]->internal_id,
|
||||
result.prefilled_offset - 1);
|
||||
// Enable sliding window for the sequence if it is not a parent.
|
||||
if (rsentry->child_indices.empty()) {
|
||||
models_[i]->EnableSlidingWindowForSeq(rsentry->mstates[0]->internal_id);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Reuse recycling sequence
|
||||
// Note: The processing for reusing recycling sequence is like forking sequence. And we
|
||||
// also roll back one token due to the reason mentioned above.
|
||||
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;
|
||||
}
|
||||
estate->prefix_cache->RollBackSequence(rsentry->mstates[0]->internal_id, 1);
|
||||
for (int i = 0; i < models_.size(); ++i) {
|
||||
models_[i]->PopNFromKVCache(rsentry->mstates[0]->internal_id,
|
||||
result.reused_seq_pop_last_tokens + 1);
|
||||
}
|
||||
result.prefilled_offset -= 1;
|
||||
}
|
||||
}
|
||||
// Pop matched prefix
|
||||
if (result.prefilled_offset > 0) {
|
||||
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 - 1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
};
|
||||
|
||||
EngineAction EngineAction::EagleNewRequestPrefill(
|
||||
Array<Model> models, LogitProcessor logit_processor, Sampler sampler,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
DraftTokenWorkspaceManager draft_token_workspace_manager, EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
Optional<EventTraceRecorder> trace_recorder) {
|
||||
return EngineAction(tvm::ffi::make_object<EagleNewRequestPrefillActionObj>(
|
||||
std::move(models), std::move(logit_processor), std::move(sampler),
|
||||
std::move(model_workspaces), std::move(draft_token_workspace_manager),
|
||||
std::move(engine_config), std::move(model_configs), std::move(trace_recorder)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,370 @@
|
||||
/*!
|
||||
* 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.
|
||||
*/
|
||||
class NewRequestPrefillActionObj : public BatchPrefillBaseActionObj {
|
||||
public:
|
||||
explicit NewRequestPrefillActionObj(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler, std::vector<ModelWorkspace> model_workspaces,
|
||||
EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
Optional<EventTraceRecorder> trace_recorder)
|
||||
: BatchPrefillBaseActionObj(std::move(models), std::move(engine_config),
|
||||
std::move(model_configs), std::move(trace_recorder)),
|
||||
logit_processor_(std::move(logit_processor)),
|
||||
sampler_(std::move(sampler)),
|
||||
model_workspaces_(std::move(model_workspaces)) {}
|
||||
|
||||
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("NewRequestPrefill getting requests");
|
||||
prefill_inputs = GetRequestStateEntriesToPrefill(estate);
|
||||
if (prefill_inputs.empty()) {
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
int num_rsentries = prefill_inputs.size();
|
||||
{
|
||||
NVTXScopedRange nvtx_scope("NewRequestPrefill 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.
|
||||
std::vector<int> prefill_lengths;
|
||||
prefill_lengths.resize(/*size=*/num_rsentries, /*value=*/-1);
|
||||
Tensor logits_for_sample{nullptr};
|
||||
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, or fork the sequence from its parent.
|
||||
// If the sequence is already in prefix cache, it has also been added/forked in the
|
||||
// KVCache.
|
||||
if (rsentry->parent_idx == -1) {
|
||||
models_[model_id]->AddNewSequence(mstate->internal_id);
|
||||
} else {
|
||||
models_[model_id]->ForkSequence(
|
||||
rstates_of_entries[i]->entries[rsentry->parent_idx]->mstates[model_id]->internal_id,
|
||||
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);
|
||||
}
|
||||
}
|
||||
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);
|
||||
|
||||
if (model_id == 0) {
|
||||
// We only need to sample for model 0 in prefill.
|
||||
logits_for_sample = logits;
|
||||
}
|
||||
}
|
||||
|
||||
// - Update logits.
|
||||
TVM_FFI_ICHECK(logits_for_sample.defined());
|
||||
Array<GenerationConfig> generation_cfg;
|
||||
Array<RequestModelState> mstates_for_logitproc;
|
||||
generation_cfg.reserve(num_rsentries);
|
||||
mstates_for_logitproc.reserve(num_rsentries);
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
generation_cfg.push_back(prefill_inputs[i].rsentry->request->generation_cfg);
|
||||
mstates_for_logitproc.push_back(prefill_inputs[i].rsentry->mstates[0]);
|
||||
}
|
||||
logits_for_sample = logits_for_sample.CreateView({num_rsentries, logits_for_sample->shape[2]},
|
||||
logits_for_sample->dtype);
|
||||
logit_processor_->InplaceUpdateLogits(logits_for_sample, generation_cfg, mstates_for_logitproc,
|
||||
request_ids);
|
||||
|
||||
// - Compute probability distributions.
|
||||
Tensor probs_on_device =
|
||||
logit_processor_->ComputeProbsFromLogits(logits_for_sample, generation_cfg, request_ids);
|
||||
|
||||
// - Commit the prefix cache changes from previous round of action.
|
||||
// Note: we commit prefix cache changes here to overlap this commit with the GPU execution.
|
||||
estate->prefix_cache->CommitSequenceExtention();
|
||||
|
||||
// - Sample tokens.
|
||||
// For rsentries which have children, sample
|
||||
// one token for each rstate that is depending.
|
||||
// Otherwise, sample a token for the current rstate.
|
||||
std::vector<int> sample_indices;
|
||||
std::vector<RequestStateEntry> rsentries_for_sample;
|
||||
std::vector<RandomGenerator*> rngs;
|
||||
std::vector<bool> rsentry_activated;
|
||||
sample_indices.reserve(num_rsentries);
|
||||
rsentries_for_sample.reserve(num_rsentries);
|
||||
rngs.reserve(num_rsentries);
|
||||
rsentry_activated.reserve(num_rsentries);
|
||||
request_ids.clear();
|
||||
generation_cfg.clear();
|
||||
for (int i = 0; i < num_rsentries; ++i) {
|
||||
const RequestStateEntry& rsentry = prefill_inputs[i].rsentry;
|
||||
// No sample for rsentries with remaining inputs.
|
||||
if (!rsentry->mstates[0]->inputs.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
int remaining_num_child_to_activate = prefill_inputs[i].num_child_to_activate;
|
||||
for (int child_idx : rsentry->child_indices) {
|
||||
// If rstates_of_entries[i]->entries[child_idx] has no committed token,
|
||||
// the prefill of the current rsentry will unblock
|
||||
// rstates_of_entries[i]->entries[child_idx],
|
||||
// and thus we want to sample a token for rstates_of_entries[i]->entries[child_idx].
|
||||
if (rstates_of_entries[i]->entries[child_idx]->status != RequestStateStatus::kPending ||
|
||||
!rstates_of_entries[i]->entries[child_idx]->mstates[0]->committed_tokens.empty()) {
|
||||
continue;
|
||||
}
|
||||
sample_indices.push_back(i);
|
||||
rsentries_for_sample.push_back(rstates_of_entries[i]->entries[child_idx]);
|
||||
request_ids.push_back(rsentry->request->id);
|
||||
generation_cfg.push_back(rsentry->request->generation_cfg);
|
||||
rngs.push_back(&rstates_of_entries[i]->entries[child_idx]->rng);
|
||||
|
||||
TVM_FFI_ICHECK(rstates_of_entries[i]->entries[child_idx]->status ==
|
||||
RequestStateStatus::kPending);
|
||||
// We only fork the first `num_child_to_activate` children.
|
||||
// The children not being forked will be forked via later prefills.
|
||||
// Usually `num_child_to_activate` is the same as the number of children.
|
||||
// But it can be fewer subject to the KV cache max num sequence limit.
|
||||
if (remaining_num_child_to_activate == 0) {
|
||||
rsentry_activated.push_back(false);
|
||||
continue;
|
||||
}
|
||||
rsentry_activated.push_back(true);
|
||||
--remaining_num_child_to_activate;
|
||||
rstates_of_entries[i]->entries[child_idx]->status = RequestStateStatus::kAlive;
|
||||
for (int model_id = 0; model_id < static_cast<int>(models_.size()); ++model_id) {
|
||||
int64_t child_internal_id =
|
||||
rstates_of_entries[i]->entries[child_idx]->mstates[model_id]->internal_id;
|
||||
models_[model_id]->ForkSequence(rsentry->mstates[model_id]->internal_id,
|
||||
child_internal_id);
|
||||
// Enable sliding window for the child sequence if the child is not a parent.
|
||||
if (rstates_of_entries[i]->entries[child_idx]->child_indices.empty()) {
|
||||
models_[model_id]->EnableSlidingWindowForSeq(child_internal_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (rsentry->child_indices.empty()) {
|
||||
// If rsentry has no child, we sample a token for itself.
|
||||
sample_indices.push_back(i);
|
||||
rsentries_for_sample.push_back(rsentry);
|
||||
request_ids.push_back(rsentry->request->id);
|
||||
generation_cfg.push_back(rsentry->request->generation_cfg);
|
||||
rngs.push_back(&rsentry->rng);
|
||||
rsentry_activated.push_back(true);
|
||||
}
|
||||
}
|
||||
Tensor renormalized_probs = sampler_->BatchRenormalizeProbsByTopP(
|
||||
probs_on_device, sample_indices, request_ids, generation_cfg);
|
||||
std::vector<SampleResult> sample_results = sampler_->BatchSampleTokensWithProbAfterTopP(
|
||||
renormalized_probs, sample_indices, request_ids, generation_cfg, rngs);
|
||||
TVM_FFI_ICHECK_EQ(sample_results.size(), rsentries_for_sample.size());
|
||||
|
||||
// - Update the committed tokens of states.
|
||||
// - If a request is first-time prefilled, set the prefill finish time.
|
||||
UpdateRequestStateEntriesWithSampleResults(rsentries_for_sample, rsentry_activated,
|
||||
sample_results);
|
||||
|
||||
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:
|
||||
/*! \brief The logit processor. */
|
||||
LogitProcessor logit_processor_;
|
||||
/*! \brief The sampler to sample new tokens. */
|
||||
Sampler sampler_;
|
||||
/*! \brief Workspace of each model. */
|
||||
std::vector<ModelWorkspace> model_workspaces_;
|
||||
|
||||
/*!
|
||||
* \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 (Model model : models_) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
} 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 (Model model : models_) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
} 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// 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;
|
||||
}
|
||||
}; // namespace serve
|
||||
|
||||
EngineAction EngineAction::NewRequestPrefill(Array<Model> models, LogitProcessor logit_processor,
|
||||
Sampler sampler,
|
||||
std::vector<ModelWorkspace> model_workspaces,
|
||||
EngineConfig engine_config,
|
||||
std::vector<tvm::ffi::json::Object> model_configs,
|
||||
Optional<EventTraceRecorder> trace_recorder) {
|
||||
return EngineAction(tvm::ffi::make_object<NewRequestPrefillActionObj>(
|
||||
std::move(models), std::move(logit_processor), std::move(sampler),
|
||||
std::move(model_workspaces), std::move(engine_config), std::move(model_configs),
|
||||
std::move(trace_recorder)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,54 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_state.cc
|
||||
*/
|
||||
#include "engine_state.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() { EngineStateObj::RegisterReflection(); }
|
||||
|
||||
EngineState::EngineState() { data_ = tvm::ffi::make_object<EngineStateObj>(); }
|
||||
|
||||
void EngineStateObj::Reset() {
|
||||
running_queue.clear();
|
||||
waiting_queue.clear();
|
||||
request_states.clear();
|
||||
id_manager.Reset();
|
||||
metrics.Reset();
|
||||
if (prefix_cache.defined()) {
|
||||
prefix_cache->Reset();
|
||||
}
|
||||
running_rsentries_changed = true;
|
||||
postproc_workspace = ActionPostProcessWorkspace();
|
||||
}
|
||||
|
||||
RequestState EngineStateObj::GetRequestState(Request request) {
|
||||
TVM_FFI_ICHECK(request->rstate != nullptr) << "The state of the request has not been defined.";
|
||||
return GetRef<RequestState>(static_cast<RequestStateNode*>(request->rstate));
|
||||
}
|
||||
|
||||
const std::vector<RequestStateEntry>& EngineStateObj::GetRunningRequestStateEntries() {
|
||||
if (running_rsentries_changed) {
|
||||
cached_running_rsentries_.clear();
|
||||
for (const Request& request : running_queue) {
|
||||
for (const RequestStateEntry& rsentry : GetRequestState(request)->entries) {
|
||||
// One request entry is considered as running for decode if it is a leaf and has
|
||||
// finished all input prefill.
|
||||
if (rsentry->status == RequestStateStatus::kAlive && rsentry->child_indices.empty() &&
|
||||
rsentry->mstates[0]->inputs.empty()) {
|
||||
cached_running_rsentries_.push_back(rsentry);
|
||||
}
|
||||
}
|
||||
}
|
||||
running_rsentries_changed = false;
|
||||
}
|
||||
return cached_running_rsentries_;
|
||||
//
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,133 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/engine_state.h
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_ENGINE_STATE_H_
|
||||
#define MLC_LLM_SERVE_ENGINE_STATE_H_
|
||||
|
||||
#include <tvm/ffi/cast.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
|
||||
#include "config.h"
|
||||
#include "metrics.h"
|
||||
#include "prefix_cache.h"
|
||||
#include "request.h"
|
||||
#include "request_state.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::GetRef;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
|
||||
typedef TypedFunction<void(Array<RequestStreamOutput>)> FRequestStreamCallback;
|
||||
|
||||
/*! \brief The manager of internal id for requests in engine. */
|
||||
struct EngineInternalIDManager {
|
||||
std::vector<int64_t> available_ids;
|
||||
int64_t id_cnt = 0;
|
||||
|
||||
/*! \brief Return an unused id. */
|
||||
int64_t GetNewId() {
|
||||
if (!available_ids.empty()) {
|
||||
int64_t id = available_ids.back();
|
||||
available_ids.pop_back();
|
||||
return id;
|
||||
} else {
|
||||
return id_cnt++;
|
||||
}
|
||||
}
|
||||
|
||||
/*! \brief Recycle an id. */
|
||||
void RecycleId(int64_t id) { available_ids.push_back(id); }
|
||||
|
||||
/*! \brief Reset the manager. */
|
||||
void Reset() {
|
||||
available_ids.clear();
|
||||
id_cnt = 0;
|
||||
}
|
||||
};
|
||||
|
||||
/*! \brief The data structures used in the action post-process. */
|
||||
struct ActionPostProcessWorkspace {
|
||||
std::vector<RequestStateEntry> finished_rsentries;
|
||||
Array<RequestStreamOutput> callback_delta_outputs;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief The state of the running engine.
|
||||
* It contains the requests and their states submitted to the Engine.
|
||||
*/
|
||||
class EngineStateObj : public Object {
|
||||
public:
|
||||
/*! \brief The requests being processed. */
|
||||
std::vector<Request> running_queue;
|
||||
/*! \brief The requests that have not started for process yet. */
|
||||
std::vector<Request> waiting_queue;
|
||||
/*! \brief The states of all requests. */
|
||||
std::unordered_map<String, RequestState> request_states;
|
||||
/*! \brief The internal id manager. */
|
||||
EngineInternalIDManager id_manager;
|
||||
/*! \brief Runtime metrics. */
|
||||
EngineMetrics metrics;
|
||||
/*! \brief The prefix cache. */
|
||||
PrefixCache prefix_cache{nullptr};
|
||||
/*! \brief A boolean flag denoting whether the running request state entry list has changed. */
|
||||
bool running_rsentries_changed = true;
|
||||
/*!
|
||||
* \brief The current engine speculative decoding draft length.
|
||||
* The length may change across time under the auto speculative decoding mode.
|
||||
* Value 0 means undefined. It must have a positive value for speculative decoding to
|
||||
* properly work.
|
||||
*/
|
||||
int spec_draft_length = 0;
|
||||
/*! \brief A boolean flag denoting whether the engine is in disaggregation mode. */
|
||||
bool disaggregation = false;
|
||||
// Request stream callback function
|
||||
FRequestStreamCallback request_stream_callback_;
|
||||
/*!
|
||||
* \brief The post-process data structures.
|
||||
* We make it a workspace to avoid repetitive memory allocation/free in the action post process.
|
||||
*/
|
||||
ActionPostProcessWorkspace postproc_workspace;
|
||||
|
||||
/*! \brief Reset the engine state and clear the metrics. */
|
||||
void Reset();
|
||||
/*! \brief Get the request state of the given request. */
|
||||
RequestState GetRequestState(Request request);
|
||||
/*! \brief Return the running request state entries*/
|
||||
const std::vector<RequestStateEntry>& GetRunningRequestStateEntries();
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<EngineStateObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.EngineState", EngineStateObj, Object);
|
||||
|
||||
private:
|
||||
std::vector<RequestStateEntry> cached_running_rsentries_;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Managed reference of EngineStateObj.
|
||||
* \sa EngineStateObj
|
||||
*/
|
||||
class EngineState : public ObjectRef {
|
||||
public:
|
||||
explicit EngineState();
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(EngineState, ObjectRef, EngineStateObj);
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_ENGINE_STATE_H_
|
||||
@@ -0,0 +1,163 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/event_trace_recorder.cc
|
||||
*/
|
||||
#include "event_trace_recorder.h"
|
||||
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <chrono>
|
||||
#include <mutex>
|
||||
#include <unordered_map>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::ffi::String;
|
||||
|
||||
namespace detail {
|
||||
|
||||
struct PairHash {
|
||||
template <class T1, class T2>
|
||||
std::size_t operator()(const std::pair<T1, T2>& p) const {
|
||||
auto h1 = std::hash<T1>{}(p.first);
|
||||
auto h2 = std::hash<T2>{}(p.second);
|
||||
return h1 ^ h2;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace detail
|
||||
|
||||
/*! \brief The implementation of event trace recorder. */
|
||||
class EventTraceRecorderImpl : public EventTraceRecorderObj {
|
||||
public:
|
||||
void AddEvent(const String& request_id, const std::string& event) final {
|
||||
double event_time = std::chrono::duration_cast<std::chrono::duration<double>>(
|
||||
std::chrono::system_clock::now().time_since_epoch())
|
||||
.count();
|
||||
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
AddEventInternal(request_id, event, event_time);
|
||||
}
|
||||
}
|
||||
|
||||
void AddEvent(const Array<String>& request_ids, const std::string& event) final {
|
||||
double event_time = std::chrono::duration_cast<std::chrono::duration<double>>(
|
||||
std::chrono::system_clock::now().time_since_epoch())
|
||||
.count(); // in seconds
|
||||
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
for (const String& request_id : request_ids) {
|
||||
AddEventInternal(request_id, event, event_time);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::string DumpJSON() final {
|
||||
std::unordered_map<std::string, std::vector<std::pair<std::string, double>>> local_events;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
local_events = events_;
|
||||
}
|
||||
|
||||
auto fcmp_events = [](const std::pair<int64_t, tvm::ffi::json::Value>& lhs,
|
||||
const std::pair<int64_t, tvm::ffi::json::Value>& rhs) {
|
||||
return lhs.first < rhs.first;
|
||||
};
|
||||
|
||||
tvm::ffi::json::Array event_array;
|
||||
for (const std::string& request_id : request_id_in_order_) {
|
||||
std::vector<std::pair<std::string, double>> event_pairs = local_events.at(request_id);
|
||||
std::vector<std::pair<int64_t, tvm::ffi::json::Value>> events_to_sort;
|
||||
events_to_sort.reserve(event_pairs.size());
|
||||
for (int i = 0; i < static_cast<int>(event_pairs.size()); ++i) {
|
||||
std::string event = event_pairs[i].first;
|
||||
double event_time = event_pairs[i].second;
|
||||
std::string name;
|
||||
std::string phase;
|
||||
if (event.compare(0, 6, "start ") == 0) {
|
||||
// Duration begin.
|
||||
name = event.substr(6);
|
||||
phase = "B";
|
||||
} else if (event.compare(0, 7, "finish ") == 0) {
|
||||
// Duration end.
|
||||
name = event.substr(7);
|
||||
phase = "E";
|
||||
} else {
|
||||
// Instant event.
|
||||
name = event;
|
||||
phase = "i";
|
||||
}
|
||||
int64_t event_time_in_us = static_cast<int64_t>(event_time * 1e6);
|
||||
|
||||
tvm::ffi::json::Object event_json;
|
||||
event_json.Set("name", name);
|
||||
event_json.Set("ph", phase);
|
||||
event_json.Set("ts", event_time_in_us);
|
||||
event_json.Set("pid", static_cast<int64_t>(1));
|
||||
event_json.Set("tid", request_id);
|
||||
|
||||
events_to_sort.push_back({event_time_in_us, event_json});
|
||||
}
|
||||
std::sort(events_to_sort.begin(), events_to_sort.end(), fcmp_events);
|
||||
for (auto [timestamp, event] : events_to_sort) {
|
||||
event_array.push_back(std::move(event));
|
||||
}
|
||||
}
|
||||
return tvm::ffi::json::Stringify(event_array);
|
||||
}
|
||||
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.EventTraceRecorder", EventTraceRecorderImpl,
|
||||
EventTraceRecorderObj);
|
||||
|
||||
private:
|
||||
/*! \brief The internal impl of AddEvent, taking the event time as input. */
|
||||
void AddEventInternal(const std::string& request_id, const std::string& event,
|
||||
double event_time) {
|
||||
if (std::find(request_id_in_order_.begin(), request_id_in_order_.end(), request_id) ==
|
||||
request_id_in_order_.end()) {
|
||||
request_id_in_order_.push_back(request_id);
|
||||
}
|
||||
int event_cnt = event_counter_[{request_id, event}]++;
|
||||
events_[request_id].push_back({event + " (" + std::to_string(event_cnt) + ")", event_time});
|
||||
}
|
||||
|
||||
/*! \brief The mutex ensuring only one thread can access critical regions. */
|
||||
std::mutex mutex_;
|
||||
|
||||
/************** Critical Regions **************/
|
||||
/*! \brief The request ids in time order. Each id only appears once. */
|
||||
std::vector<std::string> request_id_in_order_;
|
||||
/*! \brief The number of a certain event for a request. */
|
||||
std::unordered_map<std::pair<std::string, std::string>, int, detail::PairHash> event_counter_;
|
||||
/*! \brief The event list of each request together with the timestamps. */
|
||||
std::unordered_map<std::string, std::vector<std::pair<std::string, double>>> events_;
|
||||
};
|
||||
|
||||
EventTraceRecorder EventTraceRecorder::Create() {
|
||||
return EventTraceRecorder(tvm::ffi::make_object<EventTraceRecorderImpl>());
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
EventTraceRecorderImpl::RegisterReflection();
|
||||
refl::GlobalDef()
|
||||
.def("mlc.serve.EventTraceRecorder", []() { return EventTraceRecorder::Create(); })
|
||||
.def("mlc.serve.EventTraceRecorderAddEvent",
|
||||
[](const EventTraceRecorder& trace_recorder, const String& request_id,
|
||||
const std::string& event) { trace_recorder->AddEvent(request_id, event); })
|
||||
.def_method("mlc.serve.EventTraceRecorderDumpJSON", &EventTraceRecorderObj::DumpJSON);
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,81 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/event_trace_recorder.h
|
||||
* \brief The event trace recorder for requests in MLC LLM.
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_EVENT_TRACE_RECORDER_H_
|
||||
#define MLC_LLM_SERVE_EVENT_TRACE_RECORDER_H_
|
||||
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
|
||||
#include <string>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::ffi::Array;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
using tvm::ffi::String;
|
||||
|
||||
/*! \brief The event trace recorder for requests. */
|
||||
class EventTraceRecorderObj : public Object {
|
||||
public:
|
||||
/*!
|
||||
* \brief Record a event for the input request in the trace recorder.
|
||||
* \param request_id The subject request of the event.
|
||||
* \param event The event in a string name.
|
||||
* It can have one of the following patterns:
|
||||
* - "start xxx", which marks the start of event "xxx",
|
||||
* - "finish xxx", which marks the finish of event "xxx",
|
||||
* - "yyy", which marks the instant event "yyy".
|
||||
* The "starts" and "finishes" will be automatically paired in the trace recorder.
|
||||
*/
|
||||
virtual void AddEvent(const String& request_id, const std::string& event) = 0;
|
||||
|
||||
/*! \brief Record a event for the list of input requests. */
|
||||
virtual void AddEvent(const Array<String>& request_ids, const std::string& event) = 0;
|
||||
|
||||
/*! \brief Dump the logged events in Chrome Trace Event Format in JSON string. */
|
||||
virtual std::string DumpJSON() = 0;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<EventTraceRecorderObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.EventTraceRecorder", EventTraceRecorderObj, Object);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Managed reference to EventTraceRecorderObj.
|
||||
* \sa EventTraceRecorderObj
|
||||
*/
|
||||
class EventTraceRecorder : public ObjectRef {
|
||||
public:
|
||||
/*! \brief Create an event trace recorder. */
|
||||
static EventTraceRecorder Create();
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(EventTraceRecorder, ObjectRef, EventTraceRecorderObj);
|
||||
};
|
||||
|
||||
/****************** Helper macro ******************/
|
||||
|
||||
/*! \brief Record a event for the input request or list or requests. */
|
||||
#define RECORD_EVENT(trace_recorder, request_ids, event) \
|
||||
if (trace_recorder.has_value()) { \
|
||||
trace_recorder.value()->AddEvent(request_ids, event); \
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_EVENT_TRACE_RECORDER_H_
|
||||
@@ -0,0 +1,373 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/function_table.cc
|
||||
* \brief The implementation of function table in serving for distributed inference.
|
||||
*/
|
||||
|
||||
#include "function_table.h"
|
||||
|
||||
#include <tvm/ffi/extra/module.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/runtime/disco/session.h>
|
||||
#include <tvm/runtime/memory/memory_manager.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
#include <cstdlib>
|
||||
#include <filesystem>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "../support/load_bytes_from_file.h"
|
||||
#include "../support/utils.h"
|
||||
#include "sampler/sampler.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
Optional<Shape> GetDiscoWorkerCPUBinding(int num_workers) {
|
||||
const char* raw_cpu_binding = std::getenv("MLC_DISCO_WORKER_CPU_BINDING");
|
||||
if (raw_cpu_binding == nullptr) {
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
std::string cpu_binding_str(raw_cpu_binding);
|
||||
std::vector<std::string> cpu_ids_str = Split(cpu_binding_str, ',');
|
||||
std::vector<int64_t> cpu_ids;
|
||||
for (const std::string& cpu_id_str : cpu_ids_str) {
|
||||
try {
|
||||
cpu_ids.push_back(std::stol(cpu_id_str));
|
||||
} catch (std::invalid_argument const& ex) {
|
||||
LOG(FATAL) << "Invalid MLC_DISCO_WORKER_CPU_BINDING \"" << cpu_binding_str << "\"";
|
||||
}
|
||||
}
|
||||
if (static_cast<int>(cpu_ids.size()) < num_workers) {
|
||||
LOG(FATAL) << "Insufficient number of specified CPU workers in MLC_DISCO_WORKER_CPU_BINDING, "
|
||||
"expecting at least "
|
||||
<< num_workers << "CPU ids but only " << cpu_ids.size() << " are given.";
|
||||
}
|
||||
|
||||
return Shape{cpu_ids};
|
||||
}
|
||||
|
||||
Function FunctionTable::SessionFuncAsPackedFunc(Session sess, DRef sess_func, String name) {
|
||||
return Function([sess, func = std::move(sess_func), name = std::move(name)](
|
||||
ffi::PackedArgs args, ffi::Any* rv) -> void {
|
||||
std::vector<AnyView> packed_args(args.size() + 3);
|
||||
packed_args[0] = static_cast<int>(DiscoAction::kCallPacked);
|
||||
packed_args[1] = 0;
|
||||
packed_args[2] = func;
|
||||
for (int i = 0; i < args.size(); ++i) {
|
||||
packed_args[i + 3] = args[i];
|
||||
}
|
||||
*rv = sess->CallWithPacked(tvm::ffi::PackedArgs(packed_args.data(), packed_args.size()));
|
||||
});
|
||||
}
|
||||
|
||||
void FunctionTable::Init(String reload_lib_path, Device device, tvm::ffi::json::Object model_config,
|
||||
Optional<Session> session, int num_shards, int num_stages) {
|
||||
local_gpu_device = device;
|
||||
this->model_config = model_config;
|
||||
this->cached_buffers = Map<String, ObjectRef>();
|
||||
|
||||
int num_workers = num_shards * num_stages;
|
||||
if (num_workers > 1) {
|
||||
TVM_FFI_ICHECK(session.has_value());
|
||||
this->sess = session.value();
|
||||
this->use_disco = true;
|
||||
this->disco_mod = sess->CallPacked(sess->GetGlobalFunc("runtime.disco.load_vm_module"),
|
||||
reload_lib_path, Optional<Device>(std::nullopt));
|
||||
this->mod_get_func = [this, fmodule_get_function = sess->GetGlobalFunc(
|
||||
"ffi.ModuleGetFunction")](const std::string& name) -> Function {
|
||||
DRef func = sess->CallPacked(fmodule_get_function, this->disco_mod, name, true);
|
||||
bool exists = (func->DebugGetFromRemote(0).as<Function>()) != nullptr;
|
||||
if (!exists) {
|
||||
return Function(nullptr);
|
||||
}
|
||||
return SessionFuncAsPackedFunc(sess, func, name);
|
||||
};
|
||||
if (num_stages == 1) {
|
||||
if (Optional<Shape> cpu_ids = GetDiscoWorkerCPUBinding(/*num_workers=*/num_shards)) {
|
||||
Shape cpu_ids_value = cpu_ids.value();
|
||||
sess->CallPacked(sess->GetGlobalFunc("runtime.disco.bind_worker_to_cpu_core"),
|
||||
cpu_ids_value);
|
||||
}
|
||||
}
|
||||
this->get_global_func = [this](const std::string& name) -> Function {
|
||||
return SessionFuncAsPackedFunc(sess, sess->GetGlobalFunc(name), name);
|
||||
};
|
||||
this->model_metadata_ = ModelMetadata::FromModule(
|
||||
this->disco_mod.value()->DebugGetFromRemote(0).cast<Module>(), std::move(model_config));
|
||||
this->_InitFunctions();
|
||||
} else {
|
||||
TVM_FFI_ICHECK(!session.has_value());
|
||||
Optional<Module> executable = std::nullopt;
|
||||
Optional<Function> fload_exec;
|
||||
if (StartsWith(reload_lib_path, "system://")) {
|
||||
static Function f_load_system_lib = Function::GetGlobalRequired("ffi.SystemLib");
|
||||
std::string system_lib_prefix = std::string(reload_lib_path).substr(9);
|
||||
std::replace(system_lib_prefix.begin(), system_lib_prefix.end(), /*old=*/'-', /*new=*/'_');
|
||||
executable = f_load_system_lib(system_lib_prefix + "_").cast<Module>();
|
||||
fload_exec = executable.value()->GetFunction("vm_load_executable");
|
||||
TVM_FFI_ICHECK(fload_exec.has_value())
|
||||
<< "Cannot find system lib with " << system_lib_prefix
|
||||
<< ", please make sure you set model_lib field consistently with the compilation ";
|
||||
} else {
|
||||
executable = tvm::ffi::Module::LoadFromFile(reload_lib_path);
|
||||
fload_exec = executable.value()->GetFunction("vm_load_executable");
|
||||
/* precompile opencl kernel programs */
|
||||
if (device.device_type == kDLOpenCL) {
|
||||
auto f_get = executable.value()->GetFunction("opencl.GetPreCompiledPrograms", true);
|
||||
TVM_FFI_ICHECK(f_get.has_value()) << "Cannot find opencl.GetPreCompiledPrograms";
|
||||
tvm::ffi::String bytes = f_get.value()().cast<String>();
|
||||
auto f_set = executable.value()->GetFunction("opencl.SetPreCompiledPrograms", true);
|
||||
TVM_FFI_ICHECK(f_set.has_value()) << "Cannot find opencl.SetPreCompiledPrograms";
|
||||
f_set.value()(tvm::ffi::String(bytes));
|
||||
}
|
||||
TVM_FFI_ICHECK(fload_exec.has_value()) << "TVM runtime cannot find vm_load_executable";
|
||||
}
|
||||
this->use_disco = false;
|
||||
this->local_vm = fload_exec.value()().cast<Module>();
|
||||
this->local_vm.value()
|
||||
->GetFunction("vm_initialization")
|
||||
.value()(static_cast<int>(device.device_type), device.device_id,
|
||||
static_cast<int>(tvm::runtime::memory::AllocatorType::kPooled),
|
||||
static_cast<int>(kDLCPU), 0,
|
||||
static_cast<int>(tvm::runtime::memory::AllocatorType::kPooled));
|
||||
this->mod_get_func = [this](const std::string& name) -> Function {
|
||||
return this->local_vm.value()->GetFunction(name, true).value_or(Function(nullptr));
|
||||
};
|
||||
this->get_global_func = [](const std::string& name) -> Function {
|
||||
return Function::GetGlobalRequired(name);
|
||||
};
|
||||
this->model_metadata_ =
|
||||
ModelMetadata::FromModule(this->local_vm.value(), std::move(model_config));
|
||||
this->_InitFunctions();
|
||||
}
|
||||
TVM_FFI_ICHECK_EQ(this->model_metadata_.tensor_parallel_shards, num_shards);
|
||||
TVM_FFI_ICHECK_EQ(this->model_metadata_.pipeline_parallel_stages, num_stages);
|
||||
// Invoke the CUDA graph allocation init function if it is defined.
|
||||
if (cuda_graph_alloc_init_func_.defined()) {
|
||||
this->cuda_graph_alloc_init_func_();
|
||||
}
|
||||
}
|
||||
|
||||
ObjectRef FunctionTable::LoadParams(const std::string& model_path, Device device) {
|
||||
if (this->use_disco) {
|
||||
Optional<DRef> params = std::nullopt;
|
||||
if (this->model_metadata_.params.empty()) {
|
||||
std::filesystem::path fs_model_path = model_path;
|
||||
std::string metadata_path = (fs_model_path / "tensor-cache.json").string();
|
||||
std::string tensor_cache_metadata = LoadBytesFromFile(metadata_path);
|
||||
Function loader_create = this->get_global_func("runtime.disco.ShardLoader");
|
||||
|
||||
auto load_all_func_name = "runtime.disco.ShardLoaderLoadAll";
|
||||
Function loader_load_all = this->get_global_func(load_all_func_name);
|
||||
TVM_FFI_ICHECK(loader_create != nullptr);
|
||||
TVM_FFI_ICHECK(loader_load_all != nullptr);
|
||||
DRef loader =
|
||||
loader_create(metadata_path, tensor_cache_metadata, "", this->disco_mod).cast<DRef>();
|
||||
params = loader_load_all(loader).cast<DRef>();
|
||||
} else {
|
||||
auto load_func_name = getenv("MLC_INTERNAL_PRESHARD_NUM") == nullptr
|
||||
? "mlc.multi_gpu.LoadMultiGPU"
|
||||
: "mlc.multi_gpu.LoadMultiGPUPresharded";
|
||||
Function loader = this->get_global_func(load_func_name);
|
||||
params = loader(model_path, this->disco_mod, tvm::ffi::json::Stringify(this->model_config))
|
||||
.cast<DRef>();
|
||||
}
|
||||
return params.value();
|
||||
} else {
|
||||
static Function fload_cache = Function::GetGlobalRequired("vm.builtin.tensor_cache.load");
|
||||
fload_cache(model_path, static_cast<int32_t>(device.device_type), device.device_id);
|
||||
Array<Tensor> params;
|
||||
if (this->model_metadata_.params.empty()) {
|
||||
constexpr const char* name_loader = "vm.builtin.param_array_from_cache";
|
||||
static Function fload_params = Function::GetGlobalRequired(name_loader);
|
||||
params = fload_params("param", -1).cast<Array<Tensor>>();
|
||||
} else {
|
||||
constexpr const char* name_loader = "vm.builtin.param_array_from_cache_by_name";
|
||||
static Function fload_params = Function::GetGlobalRequired(name_loader);
|
||||
Array<String> param_names;
|
||||
param_names.reserve(this->model_metadata_.params.size());
|
||||
for (const auto& param : this->model_metadata_.params) {
|
||||
param_names.push_back(param.name);
|
||||
}
|
||||
params = fload_params(param_names).cast<Array<Tensor>>();
|
||||
}
|
||||
// after we get params, it is safe to simply clear the cached version
|
||||
// as these params are referenced by params_
|
||||
static Function fclear_tensor_cache =
|
||||
Function::GetGlobalRequired("vm.builtin.tensor_cache.clear");
|
||||
fclear_tensor_cache();
|
||||
return params;
|
||||
}
|
||||
}
|
||||
|
||||
void FunctionTable::_InitFunctions() {
|
||||
this->embed_func_ = mod_get_func("embed");
|
||||
this->image_embed_func_ = mod_get_func("image_embed");
|
||||
this->single_batch_prefill_func_ = mod_get_func("prefill");
|
||||
this->single_batch_decode_func_ = mod_get_func("decode");
|
||||
this->single_batch_extend_func_ = mod_get_func("extend");
|
||||
this->prefill_func_ = mod_get_func("batch_prefill");
|
||||
this->decode_func_ = mod_get_func("batch_decode");
|
||||
this->extend_func_ = mod_get_func("batch_extend");
|
||||
this->verify_func_ = mod_get_func("batch_verify");
|
||||
this->single_batch_prefill_to_last_hidden_func_ = mod_get_func("prefill_to_last_hidden_states");
|
||||
this->single_batch_decode_to_last_hidden_func_ = mod_get_func("decode_to_last_hidden_states");
|
||||
this->prefill_to_last_hidden_func_ = mod_get_func("batch_prefill_to_last_hidden_states");
|
||||
this->decode_to_last_hidden_func_ = mod_get_func("batch_decode_to_last_hidden_states");
|
||||
this->verify_to_last_hidden_func_ = mod_get_func("batch_verify_to_last_hidden_states");
|
||||
this->fuse_embed_hidden_func_ = mod_get_func("fuse_embed_hidden_states");
|
||||
Module mod = this->use_disco ? this->disco_mod.value()->DebugGetFromRemote(0).cast<Module>()
|
||||
: this->local_vm.value();
|
||||
this->get_logits_func_ = mod_get_func("get_logits");
|
||||
this->batch_get_logits_func_ = mod_get_func("batch_get_logits");
|
||||
this->batch_select_last_hidden_func_ = mod_get_func("batch_select_last_hidden_states");
|
||||
this->softmax_func_ =
|
||||
mod->GetFunction("softmax_with_temperature", true).value_or(Function(nullptr));
|
||||
this->apply_logit_bias_func_ =
|
||||
mod->GetFunction("apply_logit_bias_inplace", true).value_or(Function(nullptr));
|
||||
this->apply_penalty_func_ =
|
||||
mod->GetFunction("apply_penalty_inplace", true).value_or(Function(nullptr));
|
||||
this->apply_bitmask_func_ =
|
||||
mod->GetFunction("apply_bitmask_inplace", true).value_or(Function(nullptr));
|
||||
this->alloc_embedding_tensor_func_ = mod_get_func("alloc_embedding_tensor");
|
||||
this->cuda_graph_alloc_init_func_ = mod_get_func("cuda_graph_alloc_init");
|
||||
this->create_kv_cache_func_ = mod_get_func("create_flashinfer_paged_kv_cache");
|
||||
if (this->model_metadata_.sliding_window_size != -1 || !this->create_kv_cache_func_.defined()) {
|
||||
Function f_create_rnn_state = mod_get_func("create_rnn_state");
|
||||
if (this->model_metadata_.kv_state_kind == KVStateKind::kHybrid) {
|
||||
// Hybrid models need both KV cache and RNN state.
|
||||
this->create_kv_cache_func_ = mod_get_func("create_tir_paged_kv_cache");
|
||||
this->create_rnn_state_func_ = f_create_rnn_state;
|
||||
} else if (f_create_rnn_state.defined()) {
|
||||
this->create_kv_cache_func_ = f_create_rnn_state;
|
||||
} else {
|
||||
this->create_kv_cache_func_ = mod_get_func("create_tir_paged_kv_cache");
|
||||
}
|
||||
}
|
||||
this->reset_kv_cache_func_ = get_global_func("vm.builtin.kv_state_clear");
|
||||
this->kv_cache_add_sequence_func_ = get_global_func("vm.builtin.kv_state_add_sequence");
|
||||
this->kv_cache_fork_sequence_func_ = get_global_func("vm.builtin.kv_state_fork_sequence");
|
||||
this->kv_cache_enable_sliding_window_for_seq_ =
|
||||
get_global_func("vm.builtin.attention_kv_cache_enable_sliding_window_for_seq");
|
||||
this->kv_cache_remove_sequence_func_ = get_global_func("vm.builtin.kv_state_remove_sequence");
|
||||
this->kv_cache_begin_forward_func_ = get_global_func("vm.builtin.kv_state_begin_forward");
|
||||
this->kv_cache_end_forward_func_ = get_global_func("vm.builtin.kv_state_end_forward");
|
||||
this->kv_cache_disagg_prepare_recv_func_ =
|
||||
get_global_func("vm.builtin.kv_cache_disagg_prepare_recv");
|
||||
this->kv_cache_disagg_mark_send_func_ = get_global_func("vm.builtin.kv_cache_disagg_mark_send");
|
||||
this->kv_cache_popn_func_ = get_global_func("vm.builtin.kv_state_popn");
|
||||
this->kv_cache_commit_accepted_token_tree_nodes_func_ =
|
||||
get_global_func("vm.builtin.attention_kv_cache_commit_accepted_token_tree_nodes");
|
||||
this->kv_cache_get_num_available_pages_func_ =
|
||||
Function::GetGlobalRequired("vm.builtin.attention_kv_cache_get_num_available_pages");
|
||||
this->kv_cache_get_total_sequence_length_func_ =
|
||||
Function::GetGlobalRequired("vm.builtin.attention_kv_cache_get_total_sequence_length");
|
||||
if (Sampler::SupportGPUSampler(local_gpu_device)) {
|
||||
gpu_multinomial_from_uniform_func_ =
|
||||
mod->GetFunction("multinomial_from_uniform", true).value_or(Function(nullptr));
|
||||
gpu_argsort_probs_func_ = mod->GetFunction("argsort_probs", true).value_or(Function(nullptr));
|
||||
gpu_sample_with_top_p_func_ =
|
||||
mod->GetFunction("sample_with_top_p", true).value_or(Function(nullptr));
|
||||
gpu_sampler_take_probs_func_ =
|
||||
mod->GetFunction("sampler_take_probs", true).value_or(Function(nullptr));
|
||||
gpu_verify_draft_tokens_func_ =
|
||||
mod->GetFunction("sampler_verify_draft_tokens", true).value_or(Function(nullptr));
|
||||
gpu_renormalize_by_top_p_func_ =
|
||||
mod->GetFunction("renormalize_by_top_p", true).value_or(Function(nullptr));
|
||||
}
|
||||
this->nd_view_func_ = get_global_func("vm.builtin.reshape");
|
||||
this->nd_get_shape_func_ = get_global_func("vm.builtin.shape_of");
|
||||
this->nd_copy_embedding_to_offset_func_ = get_global_func("mlc.copy_embedding_to_offset");
|
||||
support_backtracking_kv_ = true;
|
||||
this->tuple_getitem_func_ = get_global_func("vm.builtin.tuple_getitem");
|
||||
if (use_disco) {
|
||||
this->last_group_send_to_worker_0_ =
|
||||
get_global_func("mlc.multi_gpu.SendFromLastGroupToWorker0");
|
||||
}
|
||||
|
||||
this->gather_probs_func_ = mod->GetFunction("gather_probs", true).value_or(Function(nullptr));
|
||||
this->scatter_probs_func_ = mod->GetFunction("scatter_probs", true).value_or(Function(nullptr));
|
||||
this->gather_hidden_states_func_ = mod_get_func("gather_hidden_states");
|
||||
this->scatter_hidden_states_func_ = mod_get_func("scatter_hidden_states");
|
||||
}
|
||||
|
||||
ObjectRef FunctionTable::Empty(Shape shape, DLDataType dtype, Device device,
|
||||
bool worker0_only) const {
|
||||
if (this->use_disco) {
|
||||
DRef empty_func = sess->GetGlobalFunc("runtime.disco.empty");
|
||||
return sess->CallPacked(empty_func, shape, dtype, Optional<Device>(std::nullopt), worker0_only,
|
||||
/*in_group=*/false);
|
||||
} else {
|
||||
return Tensor::Empty(shape, dtype, device);
|
||||
}
|
||||
}
|
||||
|
||||
ObjectRef FunctionTable::CopyToWorker0(const Tensor& host_array, String buffer_cache_key,
|
||||
Shape max_reserved_shape, bool local_only) {
|
||||
Map<String, ObjectRef> cached_buffers = this->cached_buffers.value();
|
||||
if (this->use_disco && !local_only) {
|
||||
Device null_device{DLDeviceType(0), 0};
|
||||
Optional<DRef> buffer = std::nullopt;
|
||||
auto it = cached_buffers.find(buffer_cache_key);
|
||||
if (it != cached_buffers.end()) {
|
||||
buffer = (*it).second.as_or_throw<DRef>();
|
||||
} else {
|
||||
buffer = this->Empty(max_reserved_shape, host_array.DataType(), null_device,
|
||||
/*worker0_only=*/false)
|
||||
.as_or_throw<DRef>();
|
||||
cached_buffers.Set(buffer_cache_key, buffer.value());
|
||||
}
|
||||
Shape real_shape = host_array.Shape();
|
||||
DRef buffer_view = nd_view_func_(buffer.value(), real_shape).cast<DRef>();
|
||||
sess->CopyToWorker0(host_array, buffer_view);
|
||||
return buffer_view;
|
||||
} else {
|
||||
auto it = cached_buffers.find(buffer_cache_key);
|
||||
Tensor buffer{nullptr};
|
||||
if (it != cached_buffers.end()) {
|
||||
buffer = (*it).second.as_or_throw<Tensor>();
|
||||
if (buffer_cache_key == "image") {
|
||||
if (tvm::ffi::GetDataSize(*buffer.operator->()) <
|
||||
tvm::ffi::GetDataSize(*host_array.operator->())) {
|
||||
buffer = Tensor::Empty(max_reserved_shape, host_array->dtype, local_gpu_device);
|
||||
cached_buffers.Set(buffer_cache_key, buffer);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
buffer = Tensor::Empty(max_reserved_shape, host_array->dtype, local_gpu_device);
|
||||
cached_buffers.Set(buffer_cache_key, buffer);
|
||||
}
|
||||
buffer = buffer.CreateView(host_array.Shape(), host_array->dtype);
|
||||
DLTensor copy_dst = *(buffer.operator->());
|
||||
Tensor::CopyFromTo(host_array.operator->(), ©_dst);
|
||||
return buffer;
|
||||
}
|
||||
}
|
||||
|
||||
void FunctionTable::DebugCallFuncOnAllAllWorker(const String& func_name,
|
||||
Optional<String> func_args) const {
|
||||
if (func_args) {
|
||||
std::string args = func_args.value();
|
||||
if (this->use_disco) {
|
||||
sess->CallPacked(sess->GetGlobalFunc(func_name), args);
|
||||
} else {
|
||||
static Function func = Function::GetGlobalRequired(func_name);
|
||||
func(args);
|
||||
}
|
||||
} else {
|
||||
if (this->use_disco) {
|
||||
sess->CallPacked(sess->GetGlobalFunc(func_name));
|
||||
} else {
|
||||
static Function func = Function::GetGlobalRequired(func_name);
|
||||
func();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,152 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/function_table.h
|
||||
* \brief The header for function table in serving for distributed inference.
|
||||
*/
|
||||
|
||||
#ifndef MLC_LLM_SERVE_FUNCTION_TABLE_H_
|
||||
#define MLC_LLM_SERVE_FUNCTION_TABLE_H_
|
||||
|
||||
#include <tvm/ffi/container/map.h>
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
#include <tvm/ffi/extra/module.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/optional.h>
|
||||
#include <tvm/runtime/disco/session.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "../metadata/model.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::Device;
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Function;
|
||||
using tvm::ffi::Map;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
using tvm::ffi::Optional;
|
||||
using tvm::ffi::Shape;
|
||||
using tvm::ffi::TypedFunction;
|
||||
|
||||
//--------------------------------------------------------
|
||||
// The function table under batching settings.
|
||||
// The implementation is mostly the same as the one for
|
||||
// single-sequence distributed inference in llm_chat.cc.
|
||||
// The only difference is that the function table for
|
||||
// batching uses a different set of packed functions.
|
||||
//
|
||||
// Here we choose to have the duplicate code instead of
|
||||
// reusing the existing function table. This is mainly
|
||||
// for the independent development of batching/serving
|
||||
// and make the codebase manageable.
|
||||
// We will eventually merge two implementation into one
|
||||
// after the batching development becomes stable.
|
||||
//--------------------------------------------------------
|
||||
struct FunctionTable {
|
||||
static Function SessionFuncAsPackedFunc(Session sess, DRef sess_func, String name);
|
||||
|
||||
void Init(String reload_lib_path, Device device, tvm::ffi::json::Object model_config,
|
||||
Optional<Session> session, int num_shards, int num_stages);
|
||||
|
||||
ObjectRef LoadParams(const std::string& model_path, Device device);
|
||||
|
||||
void _InitFunctions();
|
||||
|
||||
ObjectRef Empty(Shape shape, DLDataType dtype, Device device, bool worker0_only) const;
|
||||
|
||||
/*!
|
||||
* \brief Copy a host array to the worker or local gpu.
|
||||
* \param host_array The host array to be copied.
|
||||
* \param buffer_cache_key The key to the buffer cache.
|
||||
* \param max_reserved_shape The maximum shape to be reserved in the buffer cache.
|
||||
* \param local_only Whether to copy the array to the local gpu only. If true, the use_disco
|
||||
* flag will be ignored. This can be useful for functions that run only on the
|
||||
* local gpu when disco is enabled.
|
||||
* \return The array on the worker or local gpu.
|
||||
*/
|
||||
ObjectRef CopyToWorker0(const Tensor& host_array, String buffer_cache_key,
|
||||
Shape max_reserved_shape, bool local_only = false);
|
||||
|
||||
void DebugCallFuncOnAllAllWorker(const String& func_name, Optional<String> func_args) const;
|
||||
|
||||
bool use_disco = false;
|
||||
Device local_gpu_device;
|
||||
Session sess{nullptr};
|
||||
Optional<DRef> disco_mod = std::nullopt;
|
||||
Optional<Map<String, ObjectRef>> cached_buffers = std::nullopt;
|
||||
Optional<tvm::ffi::Module> local_vm = std::nullopt;
|
||||
tvm::ffi::json::Object model_config;
|
||||
|
||||
TypedFunction<Function(const std::string&)> mod_get_func;
|
||||
TypedFunction<Function(const std::string&)> get_global_func;
|
||||
|
||||
ModelMetadata model_metadata_;
|
||||
|
||||
Function embed_func_;
|
||||
Function image_embed_func_;
|
||||
Function single_batch_prefill_func_;
|
||||
Function single_batch_decode_func_;
|
||||
Function single_batch_extend_func_;
|
||||
Function prefill_func_;
|
||||
Function decode_func_;
|
||||
Function extend_func_;
|
||||
Function verify_func_;
|
||||
Function single_batch_prefill_to_last_hidden_func_;
|
||||
Function single_batch_decode_to_last_hidden_func_;
|
||||
Function prefill_to_last_hidden_func_;
|
||||
Function decode_to_last_hidden_func_;
|
||||
Function verify_to_last_hidden_func_;
|
||||
Function fuse_embed_hidden_func_;
|
||||
Function get_logits_func_;
|
||||
Function batch_get_logits_func_;
|
||||
Function batch_select_last_hidden_func_;
|
||||
Function softmax_func_;
|
||||
Function apply_logit_bias_func_;
|
||||
Function apply_penalty_func_;
|
||||
Function apply_bitmask_func_;
|
||||
Function alloc_embedding_tensor_func_;
|
||||
Function cuda_graph_alloc_init_func_;
|
||||
Function create_kv_cache_func_;
|
||||
Function create_rnn_state_func_;
|
||||
Function reset_kv_cache_func_;
|
||||
bool support_backtracking_kv_;
|
||||
Function kv_cache_add_sequence_func_;
|
||||
Function kv_cache_fork_sequence_func_;
|
||||
Function kv_cache_enable_sliding_window_for_seq_;
|
||||
Function kv_cache_remove_sequence_func_;
|
||||
Function kv_cache_begin_forward_func_;
|
||||
Function kv_cache_end_forward_func_;
|
||||
Function kv_cache_disagg_prepare_recv_func_;
|
||||
Function kv_cache_disagg_mark_send_func_;
|
||||
Function kv_cache_popn_func_;
|
||||
Function kv_cache_commit_accepted_token_tree_nodes_func_;
|
||||
Function kv_cache_get_num_available_pages_func_;
|
||||
Function kv_cache_get_total_sequence_length_func_;
|
||||
Function gpu_multinomial_from_uniform_func_;
|
||||
Function gpu_argsort_probs_func_;
|
||||
Function gpu_sample_with_top_p_func_;
|
||||
Function gpu_sampler_take_probs_func_;
|
||||
Function gpu_verify_draft_tokens_func_;
|
||||
Function gpu_renormalize_by_top_p_func_;
|
||||
Function nd_view_func_;
|
||||
Function nd_get_shape_func_;
|
||||
Function nd_copy_embedding_to_offset_func_;
|
||||
Function tuple_getitem_func_;
|
||||
Function last_group_send_to_worker_0_;
|
||||
// Auxiliary functions for speculative decoding.
|
||||
Function gather_probs_func_;
|
||||
Function scatter_probs_func_;
|
||||
Function gather_hidden_states_func_;
|
||||
Function scatter_hidden_states_func_;
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_FUNCTION_TABLE_H_
|
||||
@@ -0,0 +1,509 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/logit_processor.cc
|
||||
* \brief The implementation of logit processor.
|
||||
*/
|
||||
#include "logit_processor.h"
|
||||
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/runtime/device_api.h>
|
||||
#include <tvm/support/cuda/nvtx.h>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::support::NVTXScopedRange;
|
||||
|
||||
inline void CopyArray(Tensor src, Tensor dst, TVMStreamHandle copy_stream) {
|
||||
DLTensor dl_dst = *(dst.operator->());
|
||||
Tensor::CopyFromTo(src.operator->(), &dl_dst, copy_stream);
|
||||
}
|
||||
|
||||
inline void SyncCopyStream(Device device, TVMStreamHandle compute_stream,
|
||||
TVMStreamHandle copy_stream) {
|
||||
// - If there is no particular copy stream, no action is needed.
|
||||
if (copy_stream == nullptr) {
|
||||
return;
|
||||
}
|
||||
// - Sync two streams.
|
||||
DeviceAPI::Get(device)->SyncStreamFromTo(device, copy_stream, compute_stream);
|
||||
}
|
||||
|
||||
/***************** LogitProcessor Implementation *****************/
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() { LogitProcessorObj::RegisterReflection(); }
|
||||
|
||||
class LogitProcessorImpl : public LogitProcessorObj {
|
||||
public:
|
||||
/*! * \brief Constructor of LogitProcessorImpl. */
|
||||
explicit LogitProcessorImpl(int max_num_token, int vocab_size, FunctionTable* ft, DLDevice device,
|
||||
Optional<EventTraceRecorder> trace_recorder)
|
||||
: max_num_token_(max_num_token),
|
||||
vocab_size_(vocab_size),
|
||||
bitmask_size_((vocab_size + 31) / 32),
|
||||
softmax_func_(ft->softmax_func_),
|
||||
device_(device),
|
||||
apply_logit_bias_func_(ft->apply_logit_bias_func_),
|
||||
apply_penalty_func_(ft->apply_penalty_func_),
|
||||
apply_bitmask_func_(ft->apply_bitmask_func_),
|
||||
trace_recorder_(std::move(trace_recorder)) {
|
||||
Device preferred_host_device = GetPreferredHostDevice(device);
|
||||
// Initialize auxiliary arrays on CPU.
|
||||
seq_ids_host_ = Tensor::Empty({max_num_token}, dtype_i32_, preferred_host_device);
|
||||
pos2seq_id_host_ =
|
||||
Tensor::Empty({max_num_token * vocab_size}, dtype_i32_, preferred_host_device);
|
||||
token_ids_host_ =
|
||||
Tensor::Empty({max_num_token * vocab_size}, dtype_i32_, preferred_host_device);
|
||||
token_cnt_host_ =
|
||||
Tensor::Empty({max_num_token * vocab_size}, dtype_i32_, preferred_host_device);
|
||||
token_logit_bias_host_ =
|
||||
Tensor::Empty({max_num_token * vocab_size}, dtype_f32_, preferred_host_device);
|
||||
penalties_host_ = Tensor::Empty({max_num_token, 3}, dtype_f32_, preferred_host_device);
|
||||
bitmask_host_ =
|
||||
Tensor::Empty({max_num_token, bitmask_size_}, dtype_i32_, preferred_host_device);
|
||||
temperature_host_ = Tensor::Empty({max_num_token}, dtype_f32_, preferred_host_device);
|
||||
// Initialize auxiliary arrays on GPU.
|
||||
seq_ids_device_ = Tensor::Empty({max_num_token}, dtype_i32_, device);
|
||||
pos2seq_id_device_ = Tensor::Empty({max_num_token * vocab_size}, dtype_i32_, device);
|
||||
token_ids_device_ = Tensor::Empty({max_num_token * vocab_size}, dtype_i32_, device);
|
||||
token_cnt_device_ = Tensor::Empty({max_num_token * vocab_size}, dtype_i32_, device);
|
||||
token_logit_bias_device_ = Tensor::Empty({max_num_token * vocab_size}, dtype_f32_, device);
|
||||
penalties_device_ = Tensor::Empty({max_num_token, 3}, dtype_f32_, device);
|
||||
bitmask_device_ = Tensor::Empty({max_num_token, bitmask_size_}, dtype_i32_, device);
|
||||
temperature_device_ = Tensor::Empty({max_num_token}, dtype_f32_, device);
|
||||
|
||||
TVM_FFI_ICHECK(apply_logit_bias_func_.defined())
|
||||
<< "Function \"apply_logit_bias_inplace\" not found in model";
|
||||
TVM_FFI_ICHECK(apply_penalty_func_.defined())
|
||||
<< "Function \"apply_penalty_inplace\" not found in model";
|
||||
TVM_FFI_ICHECK(apply_bitmask_func_.defined())
|
||||
<< "Function \"apply_bitmask_inplace\" not found in model";
|
||||
|
||||
// If the device is CUDA/ROCm, we create a standalone copy stream, in
|
||||
// purpose to hide the latency of auxiliary stream copy.
|
||||
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);
|
||||
copy_stream_ = DeviceAPI::Get(device)->CreateStream(device);
|
||||
}
|
||||
}
|
||||
|
||||
~LogitProcessorImpl() {
|
||||
// Free the copy stream if defined.
|
||||
if (copy_stream_ != nullptr) {
|
||||
DeviceAPI::Get(device_)->FreeStream(device_, copy_stream_);
|
||||
}
|
||||
}
|
||||
|
||||
void InplaceUpdateLogits(Tensor logits, //
|
||||
const Array<GenerationConfig>& generation_cfg, //
|
||||
const Array<RequestModelState>& mstates, //
|
||||
const Array<String>& request_ids, //
|
||||
const std::vector<int>* cum_num_token, //
|
||||
const Array<RequestModelState>* draft_mstates, //
|
||||
const std::vector<std::vector<int>>* draft_token_indices) final {
|
||||
NVTXScopedRange nvtx_scope("Logit inplace update");
|
||||
TVM_FFI_ICHECK_EQ(logits->ndim, 2);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[1], vocab_size_);
|
||||
TVM_FFI_ICHECK(logits.DataType() == (DLDataType{kDLFloat, 32, 1}));
|
||||
TVM_FFI_ICHECK_EQ(generation_cfg.size(), mstates.size());
|
||||
TVM_FFI_ICHECK_LE(logits->shape[0], max_num_token_);
|
||||
int num_total_token = logits->shape[0];
|
||||
int num_sequence = generation_cfg.size();
|
||||
|
||||
TVM_FFI_ICHECK((draft_mstates == nullptr) == (draft_token_indices == nullptr));
|
||||
if (cum_num_token != nullptr) {
|
||||
TVM_FFI_ICHECK(draft_mstates != nullptr);
|
||||
TVM_FFI_ICHECK_EQ(cum_num_token->size(), num_sequence + 1);
|
||||
TVM_FFI_ICHECK_EQ(cum_num_token->back(), num_total_token);
|
||||
} else {
|
||||
TVM_FFI_ICHECK_EQ(num_sequence, num_total_token);
|
||||
}
|
||||
|
||||
if (draft_mstates != nullptr) {
|
||||
TVM_FFI_ICHECK_EQ(draft_mstates->size(), num_sequence);
|
||||
TVM_FFI_ICHECK_EQ(draft_token_indices->size(), num_sequence);
|
||||
}
|
||||
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start update logits");
|
||||
|
||||
// Update 1. logit bias
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start apply logit bias");
|
||||
UpdateWithLogitBias(logits, generation_cfg, cum_num_token);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish apply logit bias");
|
||||
|
||||
// Update 2. penalties
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start apply penalty");
|
||||
UpdateWithPenalty(logits, generation_cfg, mstates, cum_num_token, draft_mstates,
|
||||
draft_token_indices);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish apply penalty");
|
||||
|
||||
// Update 3. Vocabulary mask.
|
||||
// Note: The mask application must be placed as the last step in logit processor.
|
||||
// This is because the masked logits are set to the minimal value.
|
||||
// Further logit subtraction may cause issue such as underflow.
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start apply logit mask");
|
||||
UpdateWithMask(logits, mstates, cum_num_token, draft_mstates, draft_token_indices);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish apply logit mask");
|
||||
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish update logits");
|
||||
}
|
||||
|
||||
Tensor ComputeProbsFromLogits(Tensor logits, const Array<GenerationConfig>& generation_cfg,
|
||||
const Array<String>& request_ids,
|
||||
const std::vector<int>* cum_num_token) final {
|
||||
NVTXScopedRange nvtx_scope("Compute probs from logits");
|
||||
// logits: (n, v)
|
||||
TVM_FFI_ICHECK_EQ(logits->ndim, 2);
|
||||
TVM_FFI_ICHECK_LE(logits->shape[0], max_num_token_);
|
||||
TVM_FFI_ICHECK_EQ(logits->shape[1], vocab_size_);
|
||||
TVM_FFI_ICHECK(logits.DataType() == (DLDataType{kDLFloat, 32, 1}));
|
||||
int num_total_token = logits->shape[0];
|
||||
int num_sequence = generation_cfg.size();
|
||||
|
||||
if (cum_num_token != nullptr) {
|
||||
TVM_FFI_ICHECK_EQ(cum_num_token->size(), num_sequence + 1);
|
||||
TVM_FFI_ICHECK_EQ(cum_num_token->back(), num_total_token);
|
||||
} else {
|
||||
TVM_FFI_ICHECK_EQ(num_sequence, num_total_token);
|
||||
}
|
||||
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start softmax");
|
||||
|
||||
// Construct:
|
||||
// - temperature (max_num_token,) float32
|
||||
float* p_temperature = static_cast<float*>(temperature_host_->data);
|
||||
|
||||
// - Set arrays.
|
||||
for (int i = 0; i < num_sequence; ++i) {
|
||||
int num_token_to_process =
|
||||
cum_num_token == nullptr ? 1 : (cum_num_token->at(i + 1) - cum_num_token->at(i));
|
||||
int token_offset = cum_num_token == nullptr ? i : cum_num_token->at(i);
|
||||
for (int j = 0; j < num_token_to_process; ++j) {
|
||||
p_temperature[token_offset + j] = std::max(generation_cfg[i]->temperature, 0.0);
|
||||
}
|
||||
}
|
||||
|
||||
// - View arrays.
|
||||
Tensor temperature_host = temperature_host_.CreateView({num_total_token}, dtype_f32_);
|
||||
Tensor temperature_device = temperature_device_.CreateView({num_total_token}, dtype_f32_);
|
||||
|
||||
// - Copy arrays to GPU.
|
||||
CopyArray(/*src=*/temperature_host, /*dst=*/temperature_device, copy_stream_);
|
||||
SyncCopyStream(device_, compute_stream_, copy_stream_);
|
||||
|
||||
// - Call kernel.
|
||||
Tensor probs = softmax_func_(logits.CreateView({num_total_token, 1, vocab_size_}, dtype_f32_),
|
||||
temperature_device)
|
||||
.cast<Tensor>();
|
||||
TVM_FFI_ICHECK_EQ(probs->ndim, 3);
|
||||
TVM_FFI_ICHECK_EQ(probs->shape[0], num_total_token);
|
||||
TVM_FFI_ICHECK_EQ(probs->shape[1], 1);
|
||||
TVM_FFI_ICHECK_EQ(probs->shape[2], vocab_size_);
|
||||
if (trace_recorder_.has_value()) {
|
||||
DeviceAPI::Get(device_)->StreamSync(device_, /*stream=*/nullptr);
|
||||
}
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish softmax");
|
||||
return probs.CreateView({num_total_token, vocab_size_}, probs->dtype);
|
||||
}
|
||||
|
||||
private:
|
||||
void UpdateWithLogitBias(Tensor logits, const Array<GenerationConfig>& generation_cfg,
|
||||
const std::vector<int>* cum_num_token) {
|
||||
NVTXScopedRange nvtx_scope("UpdateWithLogitBias");
|
||||
// Construct:
|
||||
// - pos2seq_id (max_num_token * vocab_size,) int32
|
||||
// - token_ids (max_num_token * vocab_size,) int32
|
||||
// - token_logit_bias (max_num_token * vocab_size,) float32
|
||||
int* p_pos2seq_id = static_cast<int*>(pos2seq_id_host_->data);
|
||||
int* p_token_ids = static_cast<int*>(token_ids_host_->data);
|
||||
float* p_token_logit_bias = static_cast<float*>(token_logit_bias_host_->data);
|
||||
|
||||
// - Set arrays.
|
||||
int num_token_for_bias = 0;
|
||||
int num_bias_token = 0;
|
||||
for (int i = 0; i < static_cast<int>(generation_cfg.size()); ++i) {
|
||||
int num_token_to_process =
|
||||
cum_num_token == nullptr ? 1 : (cum_num_token->at(i + 1) - cum_num_token->at(i));
|
||||
int token_offset = cum_num_token == nullptr ? i : cum_num_token->at(i);
|
||||
for (int j = 0; j < num_token_to_process; ++j) {
|
||||
if (!generation_cfg[i]->logit_bias.empty()) {
|
||||
for (auto [token_id, bias] : generation_cfg[i]->logit_bias) {
|
||||
p_pos2seq_id[num_bias_token] = token_offset + j;
|
||||
p_token_ids[num_bias_token] = token_id;
|
||||
p_token_logit_bias[num_bias_token] = bias;
|
||||
++num_bias_token;
|
||||
}
|
||||
++num_token_for_bias;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (num_token_for_bias == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// - View arrays.
|
||||
int num_token = num_bias_token;
|
||||
Tensor pos2seq_id_host = pos2seq_id_host_.CreateView({num_token}, dtype_i32_);
|
||||
Tensor pos2seq_id_device = pos2seq_id_device_.CreateView({num_token}, dtype_i32_);
|
||||
Tensor token_ids_host = token_ids_host_.CreateView({num_token}, dtype_i32_);
|
||||
Tensor token_ids_device = token_ids_device_.CreateView({num_token}, dtype_i32_);
|
||||
Tensor token_logit_bias_host = token_logit_bias_host_.CreateView({num_token}, dtype_f32_);
|
||||
Tensor token_logit_bias_device = token_logit_bias_device_.CreateView({num_token}, dtype_f32_);
|
||||
|
||||
// - Copy arrays to GPU.
|
||||
CopyArray(/*src=*/pos2seq_id_host, /*dst=*/pos2seq_id_device, copy_stream_);
|
||||
CopyArray(/*src=*/token_ids_host, /*dst=*/token_ids_device, copy_stream_);
|
||||
CopyArray(/*src=*/token_logit_bias_host, /*dst=*/token_logit_bias_device, copy_stream_);
|
||||
SyncCopyStream(device_, compute_stream_, copy_stream_);
|
||||
|
||||
// - Call kernel.
|
||||
apply_logit_bias_func_(logits, pos2seq_id_device, token_ids_device, token_logit_bias_device);
|
||||
if (trace_recorder_.has_value()) {
|
||||
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
|
||||
}
|
||||
}
|
||||
|
||||
void UpdateWithPenalty(Tensor logits, const Array<GenerationConfig>& generation_cfg,
|
||||
const Array<RequestModelState>& mstates,
|
||||
const std::vector<int>* cum_num_token,
|
||||
const Array<RequestModelState>* draft_mstates,
|
||||
const std::vector<std::vector<int>>* draft_token_indices) {
|
||||
NVTXScopedRange nvtx_scope("UpdateWithPenalty");
|
||||
// Construct:
|
||||
// - seq_ids (max_num_token,) int32
|
||||
// - pos2seq_id (max_num_token * vocab_size,) int32
|
||||
// - token_ids (max_num_token * vocab_size,) int32
|
||||
// - token_cnt (max_num_token * vocab_size,) int32
|
||||
// - penalties (max_num_token, 3) float32
|
||||
int* p_seq_ids = static_cast<int*>(seq_ids_host_->data);
|
||||
int* p_pos2seq_id = static_cast<int*>(pos2seq_id_host_->data);
|
||||
int* p_token_ids = static_cast<int*>(token_ids_host_->data);
|
||||
int* p_token_cnt = static_cast<int*>(token_cnt_host_->data);
|
||||
float* p_penalties = static_cast<float*>(penalties_host_->data);
|
||||
|
||||
// - Set arrays.
|
||||
int num_token_for_penalty = 0;
|
||||
int num_penalty_appeared_token = 0;
|
||||
for (int i = 0; i < static_cast<int>(generation_cfg.size()); ++i) {
|
||||
if (generation_cfg[i]->frequency_penalty != 0.0 ||
|
||||
generation_cfg[i]->presence_penalty != 0.0 ||
|
||||
generation_cfg[i]->repetition_penalty != 1.0) {
|
||||
int num_token_to_process =
|
||||
cum_num_token == nullptr ? 1 : (cum_num_token->at(i + 1) - cum_num_token->at(i));
|
||||
int token_offset = cum_num_token == nullptr ? i : cum_num_token->at(i);
|
||||
TVM_FFI_ICHECK(num_token_to_process == 1 || mstates[i]->draft_output_tokens.empty());
|
||||
TVM_FFI_ICHECK(draft_token_indices == nullptr ||
|
||||
draft_token_indices->at(i).size() == num_token_to_process);
|
||||
for (int j = 0; j < num_token_to_process; ++j) {
|
||||
p_seq_ids[num_token_for_penalty] = token_offset + j;
|
||||
|
||||
std::vector<SampleResult> draft_token_seq;
|
||||
// Update appeared_token_ids with draft tokens
|
||||
if (draft_token_indices != nullptr) {
|
||||
int cur_draft_token_index = draft_token_indices->at(i)[j];
|
||||
while (cur_draft_token_index != -1) {
|
||||
draft_token_seq.push_back(
|
||||
(*draft_mstates)[i]->draft_output_tokens[cur_draft_token_index]);
|
||||
cur_draft_token_index =
|
||||
(*draft_mstates)[i]->draft_token_parent_idx[cur_draft_token_index];
|
||||
}
|
||||
}
|
||||
auto& appeared_token_ids = mstates[i]->appeared_token_ids;
|
||||
for (const auto& token : draft_token_seq) {
|
||||
appeared_token_ids[token.GetTokenId()] += 1;
|
||||
}
|
||||
for (auto [token_id, cnt] : appeared_token_ids) {
|
||||
p_pos2seq_id[num_penalty_appeared_token] = num_token_for_penalty;
|
||||
p_token_ids[num_penalty_appeared_token] = token_id;
|
||||
p_token_cnt[num_penalty_appeared_token] = cnt;
|
||||
++num_penalty_appeared_token;
|
||||
}
|
||||
for (const auto& token : draft_token_seq) {
|
||||
if ((--appeared_token_ids[token.GetTokenId()]) == 0) {
|
||||
appeared_token_ids.erase(token.GetTokenId());
|
||||
}
|
||||
}
|
||||
p_penalties[num_token_for_penalty * 3] = generation_cfg[i]->presence_penalty;
|
||||
p_penalties[num_token_for_penalty * 3 + 1] = generation_cfg[i]->frequency_penalty;
|
||||
p_penalties[num_token_for_penalty * 3 + 2] = generation_cfg[i]->repetition_penalty;
|
||||
++num_token_for_penalty;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (num_token_for_penalty == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// - View arrays.
|
||||
int num_seq = num_token_for_penalty;
|
||||
int num_token = num_penalty_appeared_token;
|
||||
Tensor seq_ids_host = seq_ids_host_.CreateView({num_seq}, dtype_i32_);
|
||||
Tensor seq_ids_device = seq_ids_device_.CreateView({num_seq}, dtype_i32_);
|
||||
Tensor pos2seq_id_host = pos2seq_id_host_.CreateView({num_token}, dtype_i32_);
|
||||
Tensor pos2seq_id_device = pos2seq_id_device_.CreateView({num_token}, dtype_i32_);
|
||||
Tensor token_ids_host = token_ids_host_.CreateView({num_token}, dtype_i32_);
|
||||
Tensor token_ids_device = token_ids_device_.CreateView({num_token}, dtype_i32_);
|
||||
Tensor token_cnt_host = token_cnt_host_.CreateView({num_token}, dtype_i32_);
|
||||
Tensor token_cnt_device = token_cnt_device_.CreateView({num_token}, dtype_i32_);
|
||||
Tensor penalties_host = penalties_host_.CreateView({num_seq, 3}, dtype_f32_);
|
||||
Tensor penalties_device = penalties_device_.CreateView({num_seq, 3}, dtype_f32_);
|
||||
|
||||
// - Copy arrays to GPU.
|
||||
CopyArray(/*src=*/seq_ids_host, /*dst=*/seq_ids_device, copy_stream_);
|
||||
CopyArray(/*src=*/pos2seq_id_host, /*dst=*/pos2seq_id_device, copy_stream_);
|
||||
CopyArray(/*src=*/token_ids_host, /*dst=*/token_ids_device, copy_stream_);
|
||||
CopyArray(/*src=*/token_cnt_host, /*dst=*/token_cnt_device, copy_stream_);
|
||||
CopyArray(/*src=*/penalties_host, /*dst=*/penalties_device, copy_stream_);
|
||||
SyncCopyStream(device_, compute_stream_, copy_stream_);
|
||||
|
||||
// - Call kernel.
|
||||
apply_penalty_func_(logits, seq_ids_device, pos2seq_id_device, token_ids_device,
|
||||
token_cnt_device, penalties_device);
|
||||
if (trace_recorder_.has_value()) {
|
||||
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
|
||||
}
|
||||
}
|
||||
|
||||
void UpdateWithMask(Tensor logits, const Array<RequestModelState>& mstates,
|
||||
const std::vector<int>* cum_num_token,
|
||||
const Array<RequestModelState>* draft_mstates,
|
||||
const std::vector<std::vector<int>>* draft_token_indices) {
|
||||
NVTXScopedRange nvtx_scope("UpdateWithMask");
|
||||
// Construct:
|
||||
// - seq_ids (max_num_token,) int32
|
||||
// - bitmask (max_num_token, ceildiv(vocab_size, 32)), int32
|
||||
int32_t* p_seq_ids = static_cast<int32_t*>(seq_ids_host_->data);
|
||||
uint32_t* p_bitmask = static_cast<uint32_t*>(bitmask_host_->data);
|
||||
|
||||
// - Set arrays.
|
||||
int batch_size = logits->shape[0];
|
||||
TVM_FFI_ICHECK((cum_num_token == nullptr && batch_size == mstates.size()) ||
|
||||
(cum_num_token != nullptr && batch_size == cum_num_token->back()));
|
||||
|
||||
std::memset(p_seq_ids, 0, batch_size * sizeof(int32_t));
|
||||
|
||||
for (int i = 0; i < static_cast<int>(mstates.size()); ++i) {
|
||||
int token_start_offset = cum_num_token == nullptr ? i : cum_num_token->at(i);
|
||||
int token_number =
|
||||
cum_num_token == nullptr ? 1 : (cum_num_token->at(i + 1) - cum_num_token->at(i));
|
||||
bool require_mask = mstates[i]->RequireNextTokenBitmask();
|
||||
TVM_FFI_ICHECK(draft_token_indices == nullptr ||
|
||||
draft_token_indices->at(i).size() == token_number);
|
||||
for (int j = 0; j < token_number; ++j) {
|
||||
if (require_mask) {
|
||||
std::vector<SampleResult> draft_token_seq;
|
||||
if (draft_token_indices != nullptr) {
|
||||
int cur_draft_token_index = draft_token_indices->at(i)[j];
|
||||
while (cur_draft_token_index != -1) {
|
||||
draft_token_seq.push_back(
|
||||
(*draft_mstates)[i]->draft_output_tokens[cur_draft_token_index]);
|
||||
cur_draft_token_index =
|
||||
(*draft_mstates)[i]->draft_token_parent_idx[cur_draft_token_index];
|
||||
}
|
||||
for (auto it = draft_token_seq.rbegin(); it != draft_token_seq.rend(); ++it) {
|
||||
mstates[i]->grammar_matcher.value().AcceptToken(it->GetTokenId());
|
||||
}
|
||||
}
|
||||
// Find a slice of bitmask_host_: bitmask_host_[num_token_for_mask, :]
|
||||
DLTensor bitmask_dltensor = *bitmask_host_.operator->();
|
||||
int64_t bitmask_shape[] = {bitmask_size_};
|
||||
bitmask_dltensor.data = p_bitmask + (token_start_offset + j) * bitmask_size_;
|
||||
bitmask_dltensor.shape = bitmask_shape;
|
||||
bitmask_dltensor.ndim = 1;
|
||||
|
||||
mstates[i]->GetNextTokenBitmask(&bitmask_dltensor);
|
||||
p_seq_ids[token_start_offset + j] = 1;
|
||||
|
||||
if (draft_token_seq.size() > 0) {
|
||||
mstates[i]->grammar_matcher.value().Rollback(draft_token_seq.size());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int num_token_for_mask = 0;
|
||||
for (int i = 0; i < batch_size; ++i) {
|
||||
if (p_seq_ids[i] == 1) {
|
||||
p_seq_ids[num_token_for_mask] = i;
|
||||
++num_token_for_mask;
|
||||
}
|
||||
}
|
||||
|
||||
if (num_token_for_mask == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// - View arrays.
|
||||
int num_seq = num_token_for_mask;
|
||||
Tensor seq_ids_host = seq_ids_host_.CreateView({num_seq}, dtype_i32_);
|
||||
Tensor seq_ids_device = seq_ids_device_.CreateView({num_seq}, dtype_i32_);
|
||||
Tensor bitmask_host = bitmask_host_.CreateView({batch_size, bitmask_size_}, dtype_i32_);
|
||||
Tensor bitmask_device = bitmask_device_.CreateView({batch_size, bitmask_size_}, dtype_i32_);
|
||||
|
||||
// - Copy arrays to GPU.
|
||||
CopyArray(/*src=*/seq_ids_host, /*dst=*/seq_ids_device, copy_stream_);
|
||||
CopyArray(/*src=*/bitmask_host, /*dst=*/bitmask_device, copy_stream_);
|
||||
SyncCopyStream(device_, compute_stream_, copy_stream_);
|
||||
|
||||
// - Call kernel.
|
||||
apply_bitmask_func_(logits, seq_ids_device, bitmask_device);
|
||||
if (trace_recorder_.has_value()) {
|
||||
DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
|
||||
}
|
||||
}
|
||||
|
||||
// Model configurations
|
||||
const int max_num_token_;
|
||||
const int vocab_size_;
|
||||
const int bitmask_size_;
|
||||
const DLDataType dtype_i32_ = DLDataType{kDLInt, 32, 1};
|
||||
const DLDataType dtype_u32_ = DLDataType{kDLUInt, 32, 1};
|
||||
const DLDataType dtype_f32_ = DLDataType{kDLFloat, 32, 1};
|
||||
// Packed functions.
|
||||
Device device_;
|
||||
Function softmax_func_;
|
||||
Function apply_logit_bias_func_;
|
||||
Function apply_penalty_func_;
|
||||
Function apply_bitmask_func_;
|
||||
// Auxiliary Tensors on CPU
|
||||
Tensor seq_ids_host_;
|
||||
Tensor pos2seq_id_host_;
|
||||
Tensor token_ids_host_;
|
||||
Tensor token_cnt_host_;
|
||||
Tensor token_logit_bias_host_;
|
||||
Tensor penalties_host_;
|
||||
Tensor bitmask_host_;
|
||||
Tensor temperature_host_;
|
||||
// Auxiliary Tensors on GPU
|
||||
Tensor seq_ids_device_;
|
||||
Tensor pos2seq_id_device_;
|
||||
Tensor token_ids_device_;
|
||||
Tensor token_cnt_device_;
|
||||
Tensor token_logit_bias_device_;
|
||||
Tensor penalties_device_;
|
||||
Tensor bitmask_device_;
|
||||
Tensor temperature_device_;
|
||||
// Event trace recorder.
|
||||
Optional<EventTraceRecorder> trace_recorder_;
|
||||
// The device stream for the default computation operations.
|
||||
TVMStreamHandle compute_stream_ = nullptr;
|
||||
// The device stream for copying auxiliary data structure to GPU.
|
||||
TVMStreamHandle copy_stream_ = nullptr;
|
||||
// A small epsilon.
|
||||
const double eps_ = 1e-5;
|
||||
};
|
||||
|
||||
LogitProcessor::LogitProcessor(int max_num_token, int vocab_size, FunctionTable* ft,
|
||||
DLDevice device, Optional<EventTraceRecorder> trace_recorder) {
|
||||
data_ = tvm::ffi::make_object<LogitProcessorImpl>(max_num_token, vocab_size, ft, device,
|
||||
std::move(trace_recorder));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,104 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/logit_processor.h
|
||||
* \brief The header for logit processor.
|
||||
*/
|
||||
|
||||
#ifndef MLC_LLM_SERVE_LOGIT_PROCESSOR_H_
|
||||
#define MLC_LLM_SERVE_LOGIT_PROCESSOR_H_
|
||||
|
||||
#include <tvm/ffi/extra/module.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
|
||||
#include "../base.h"
|
||||
#include "config.h"
|
||||
#include "event_trace_recorder.h"
|
||||
#include "function_table.h"
|
||||
#include "request_state.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::Device;
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
|
||||
/*!
|
||||
* \brief The logit processor class that updates logits with regard
|
||||
* presence/frequency penalties, logit bias, etc..
|
||||
*/
|
||||
class LogitProcessorObj : public Object {
|
||||
public:
|
||||
/*!
|
||||
* \brief In-place update a batch of logits with regard to the given
|
||||
* generation config and request states.
|
||||
* \param logits The batch of raw logits, in shape (num_total_token, vocab_size),
|
||||
* where `num_total_token` may be larger than the number of sequences
|
||||
* indicated by `generation_cfg`, in which case some sequences may have
|
||||
* more than one token.
|
||||
* \param generation_cfg The generation config of each sequence in the batch.
|
||||
* \param mstates The request states of each sequence in the batch.
|
||||
* \param request_ids The ids of each request.
|
||||
* \param cum_num_token The pointer to the cumulative token length of the sequences.
|
||||
* If the pointer is nullptr, it means each sequence has only one token.
|
||||
* \param draft_mstates Optional. The draft request states of each sequence.
|
||||
* \param draft_token_indices Optional. The pointer to the draft token indices of each draft token
|
||||
* in the model state (-1 indicates the token is not a draft token) when speculation is enabled.
|
||||
* This is used to compute the sequence state with the draft tokens considered (the saved sequence
|
||||
* state is not updated with the draft tokens).
|
||||
*/
|
||||
virtual void InplaceUpdateLogits(
|
||||
Tensor logits, const Array<GenerationConfig>& generation_cfg,
|
||||
const Array<RequestModelState>& mstates, const Array<String>& request_ids,
|
||||
const std::vector<int>* cum_num_token = nullptr,
|
||||
const Array<RequestModelState>* draft_mstates = nullptr,
|
||||
const std::vector<std::vector<int>>* draft_token_indices = nullptr) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Compute probability distributions for the input batch of logits.
|
||||
* \param logits The batch of updated logits.
|
||||
* \param generation_cfg The generation config of each sequence in the batch.
|
||||
* \param request_ids The ids of each request.
|
||||
* \param cum_num_token The pointer to the cumulative token length of the sequences.
|
||||
* If the pointer is nullptr, it means each sequence has only one token.
|
||||
* \return The batch of computed probability distributions on GPU.
|
||||
*/
|
||||
virtual Tensor ComputeProbsFromLogits(Tensor logits,
|
||||
const Array<GenerationConfig>& generation_cfg,
|
||||
const Array<String>& request_ids,
|
||||
const std::vector<int>* cum_num_token = nullptr) = 0;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<LogitProcessorObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.LogitProcessor", LogitProcessorObj, Object);
|
||||
};
|
||||
|
||||
class LogitProcessor : public ObjectRef {
|
||||
public:
|
||||
/*!
|
||||
* \brief Constructor.
|
||||
* \param max_num_token The max number of tokens in the token processor.
|
||||
* \param vocab_size The model's vocabulary size.
|
||||
* \param ft The packed function table.
|
||||
* \param device The device that the model runs on.
|
||||
* \param trace_recorder The event trace recorder.
|
||||
*/
|
||||
explicit LogitProcessor(int max_num_token, int vocab_size, FunctionTable* ft, DLDevice device,
|
||||
Optional<EventTraceRecorder> trace_recorder);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(LogitProcessor, ObjectRef, LogitProcessorObj);
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_LOGIT_PROCESSOR_H_
|
||||
@@ -0,0 +1,185 @@
|
||||
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/metrics.cc
|
||||
*/
|
||||
#include "metrics.h"
|
||||
|
||||
#include <tvm/runtime/logging.h>
|
||||
|
||||
#include <sstream>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
tvm::ffi::json::Object TimeCost::AsJSON() const {
|
||||
tvm::ffi::json::Object config;
|
||||
config.Set("count", count);
|
||||
if (count != 0) {
|
||||
config.Set("mean", sum / count);
|
||||
}
|
||||
return config;
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object SpecDecodeMetrics::AsJSON() const {
|
||||
tvm::ffi::json::Object metrics;
|
||||
auto f_vector_to_array = [](const std::vector<int64_t>& vec) {
|
||||
tvm::ffi::json::Array arr;
|
||||
for (int64_t v : vec) {
|
||||
arr.push_back(v);
|
||||
}
|
||||
return tvm::ffi::json::Value(arr);
|
||||
};
|
||||
metrics.Set("draft_count", f_vector_to_array(draft_count));
|
||||
metrics.Set("accept_count", f_vector_to_array(accept_count));
|
||||
|
||||
TVM_FFI_ICHECK_EQ(draft_count.size(), accept_count.size());
|
||||
// NOTE: label follows prometheus with full context
|
||||
// so it can be flattened and used in metrics reoorting end point
|
||||
tvm::ffi::json::Object accept_prob_metrics;
|
||||
tvm::ffi::json::Object accept_rate_metrics;
|
||||
tvm::ffi::json::Object accept_len_metrics;
|
||||
|
||||
double accept_len_value = 0;
|
||||
|
||||
for (size_t i = 0; i < draft_count.size(); ++i) {
|
||||
std::ostringstream accept_prob_label;
|
||||
accept_prob_label << "accept_prob{step=" << i << "}";
|
||||
double accept_prob_value =
|
||||
(static_cast<double>(accept_count[i]) / static_cast<double>(draft_count[i]));
|
||||
accept_prob_metrics.Set(accept_prob_label.str(), accept_prob_value);
|
||||
accept_len_value += accept_prob_value;
|
||||
|
||||
std::ostringstream accept_len_label;
|
||||
accept_len_label << "accept_len{step=" << i << "}";
|
||||
accept_len_metrics.Set(accept_len_label.str(), accept_len_value);
|
||||
|
||||
if (i != 0) {
|
||||
std::ostringstream accept_rate_label;
|
||||
accept_rate_label << "accept_rate{step=" << i << "}";
|
||||
double accept_rate_value =
|
||||
accept_count[i - 1] == 0
|
||||
? 0.0f
|
||||
: (static_cast<double>(accept_count[i]) / static_cast<double>(accept_count[i - 1]));
|
||||
accept_rate_metrics.Set(accept_rate_label.str(), accept_rate_value);
|
||||
}
|
||||
}
|
||||
metrics.Set("accept_prob", accept_prob_metrics);
|
||||
metrics.Set("accept_rate", accept_rate_metrics);
|
||||
metrics.Set("accept_len", accept_len_metrics);
|
||||
|
||||
return metrics;
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object RequestMetrics::AsJSON() const {
|
||||
tvm::ffi::json::Object metrics;
|
||||
metrics.Set("prompt_tokens", prompt_tokens);
|
||||
metrics.Set("completion_tokens", completion_tokens);
|
||||
metrics.Set("prefill_tokens", prefill_tokens);
|
||||
metrics.Set("decode_tokens", decode_tokens);
|
||||
metrics.Set("jump_forward_tokens", jump_forward_tokens);
|
||||
|
||||
if (prefill_tokens != 0) {
|
||||
metrics.Set("prefill_tokens_per_s", prefill_tokens / this->GetPrefillTime());
|
||||
}
|
||||
if (decode_tokens != 0) {
|
||||
metrics.Set("decode_tokens_per_s", decode_tokens / this->GetDecodeTime());
|
||||
}
|
||||
metrics.Set("end_to_end_latency_s", this->GetTotalTime());
|
||||
metrics.Set("ttft_s", this->GetTTFT());
|
||||
metrics.Set("inter_token_latency_s", this->GetInterTokenLatency());
|
||||
return metrics;
|
||||
}
|
||||
|
||||
std::string RequestMetrics::AsUsageJSONStr(bool include_extra) const {
|
||||
tvm::ffi::json::Object usage;
|
||||
usage.Set("prompt_tokens", prompt_tokens);
|
||||
usage.Set("completion_tokens", completion_tokens);
|
||||
usage.Set("total_tokens", prompt_tokens + completion_tokens);
|
||||
if (include_extra) {
|
||||
usage.Set("extra", this->AsJSON());
|
||||
}
|
||||
return tvm::ffi::json::Stringify(usage);
|
||||
}
|
||||
|
||||
tvm::ffi::json::Object EngineMetrics::AsJSON() const {
|
||||
tvm::ffi::json::Object metrics;
|
||||
metrics.Set("engine_prefill_time_sum", engine_prefill_time_sum);
|
||||
metrics.Set("engine_decode_time_sum", engine_decode_time_sum);
|
||||
metrics.Set("engine_jump_forward_time_sum", engine_jump_forward_time_sum);
|
||||
metrics.Set("prompt_tokens_sum", prompt_tokens_sum);
|
||||
metrics.Set("completion_tokens_sum", completion_tokens_sum);
|
||||
metrics.Set("prefill_tokens_sum", prefill_tokens_sum);
|
||||
metrics.Set("decode_tokens_sum", decode_tokens_sum);
|
||||
metrics.Set("jump_forward_tokens_sum", jump_forward_tokens_sum);
|
||||
|
||||
if (prefill_tokens_sum != 0) {
|
||||
metrics.Set("prefill_tokens_per_s", prefill_tokens_sum / engine_prefill_time_sum);
|
||||
}
|
||||
if (engine_decode_time_sum != 0) {
|
||||
metrics.Set("decode_tokens_per_s", decode_tokens_sum / engine_decode_time_sum);
|
||||
}
|
||||
|
||||
metrics.Set("last_finished_request", last_finished_request.AsJSON());
|
||||
if (!spec_decode.IsEmpty()) {
|
||||
metrics.Set("spec_decode", spec_decode.AsJSON());
|
||||
}
|
||||
|
||||
auto f_create_time_list = [](const std::vector<TimeCost>& time_list) {
|
||||
tvm::ffi::json::Object result;
|
||||
for (size_t i = 1; i < time_list.size(); ++i) {
|
||||
const TimeCost& item = time_list[i];
|
||||
if (item.count == 0) continue;
|
||||
std::ostringstream label_mean;
|
||||
label_mean << "mean{batch_size=" << i << "}";
|
||||
double mean = item.sum / item.count;
|
||||
result.Set(label_mean.str(), mean);
|
||||
std::ostringstream label_count;
|
||||
label_count << "count{batch_size=" << i << "}";
|
||||
result.Set(label_count.str(), item.count);
|
||||
}
|
||||
return tvm::ffi::json::Value(result);
|
||||
};
|
||||
|
||||
metrics.Set("decode_time_by_batch_size", f_create_time_list(decode_time_by_batch_size));
|
||||
metrics.Set("draft_time_by_batch_size", f_create_time_list(draft_time_by_batch_size));
|
||||
metrics.Set("verify_time_by_batch_size", f_create_time_list(verify_time_by_batch_size));
|
||||
|
||||
return metrics;
|
||||
}
|
||||
|
||||
std::string EngineMetrics::AsUsageJSONStr() const {
|
||||
tvm::ffi::json::Object usage;
|
||||
// We return engine usage as a usage field according to the OpenAI API.
|
||||
// To comply with the API, just set prompt_tokens, completion_tokens, and total_tokens to 0.
|
||||
// And store the information in the extra field.
|
||||
usage.Set("prompt_tokens", static_cast<int64_t>(0));
|
||||
usage.Set("completion_tokens", static_cast<int64_t>(0));
|
||||
usage.Set("total_tokens", static_cast<int64_t>(0));
|
||||
usage.Set("extra", this->AsJSON());
|
||||
return tvm::ffi::json::Stringify(usage);
|
||||
}
|
||||
|
||||
void EngineMetrics::Reset() {
|
||||
engine_prefill_time_sum = 0.0;
|
||||
engine_decode_time_sum = 0.0;
|
||||
engine_jump_forward_time_sum = 0;
|
||||
prompt_tokens_sum = 0;
|
||||
completion_tokens_sum = 0;
|
||||
prefill_tokens_sum = 0;
|
||||
decode_tokens_sum = 0;
|
||||
jump_forward_tokens_sum = 0;
|
||||
last_finished_request.Reset();
|
||||
spec_decode.Reset();
|
||||
decode_time_by_batch_size.clear();
|
||||
draft_time_by_batch_size.clear();
|
||||
verify_time_by_batch_size.clear();
|
||||
decode_time_by_batch_size.resize(kEndFineGrainedTrackingBatchSize);
|
||||
draft_time_by_batch_size.resize(kEndFineGrainedTrackingBatchSize);
|
||||
verify_time_by_batch_size.resize(kEndFineGrainedTrackingBatchSize);
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,264 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/metric.h
|
||||
* \brief Metrics of serving engine/requests.
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_METRICS_H_
|
||||
#define MLC_LLM_SERVE_METRICS_H_
|
||||
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
#include <tvm/runtime/logging.h>
|
||||
|
||||
#include <chrono>
|
||||
#include <string>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
// We keep all metrics containers in this header (instead of in Engine and Request State)
|
||||
// so we have a single central place to define all metrics across the engine.
|
||||
// Conceptually, these statistics are derived from engine/request behaviors.
|
||||
|
||||
/*!
|
||||
* \brief The class for tracking mean time cost.
|
||||
* - We maintain the number of updates (`count`) and the sum of updated values (`sum`).
|
||||
* - We support warmup. When `warmup` is false, the first update will be discarded.
|
||||
*/
|
||||
struct TimeCost {
|
||||
/*! \brief the total amount of cost excluding warm up time */
|
||||
double sum = 0.0;
|
||||
/*! \brief the total count of events excluding warmup */
|
||||
int64_t count = 0;
|
||||
/*! \brief Whether we warmed up already, assuming one hit is enough */
|
||||
bool warmed_up = false;
|
||||
|
||||
/*! \brief Update the metric with given value. */
|
||||
void Update(double value) {
|
||||
if (warmed_up) {
|
||||
sum += value;
|
||||
count += 1;
|
||||
} else {
|
||||
warmed_up = true;
|
||||
}
|
||||
}
|
||||
|
||||
/*! \brief Reset the metric. */
|
||||
void Reset() {
|
||||
// NOTE: no need to redo warmup
|
||||
// assuming we are measuring the same thing
|
||||
this->sum = 0.0;
|
||||
this->count = 0;
|
||||
}
|
||||
|
||||
/*! \brief Dump the metric as JSON. */
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
/*! \brief Runtime metrics for speculative decoding */
|
||||
struct SpecDecodeMetrics {
|
||||
/*! \brief The number of draft tokens in speculative decoding, per step */
|
||||
std::vector<int64_t> draft_count;
|
||||
/*! \brief The number of accepted tokens in speculative decoding, per step */
|
||||
std::vector<int64_t> accept_count;
|
||||
|
||||
/*!
|
||||
* \brief Update the metrics of speculative decoding.
|
||||
* \param draft_length The number of draft tokens (including the last prediction by the base
|
||||
* model)
|
||||
* \param accept_length The number of accepted tokens in the speculative decoding.
|
||||
*/
|
||||
void Update(int draft_length, int accept_length) {
|
||||
TVM_FFI_ICHECK_GE(accept_length, 1);
|
||||
if (accept_count.size() < draft_length) {
|
||||
this->accept_count.resize(draft_length, 0);
|
||||
this->draft_count.resize(draft_length, 0);
|
||||
}
|
||||
for (int j = 0; j < draft_length; ++j) {
|
||||
if (j < accept_length) {
|
||||
++this->accept_count[j];
|
||||
}
|
||||
++this->draft_count[j];
|
||||
}
|
||||
}
|
||||
|
||||
bool IsEmpty() const { return draft_count.size() == 0; }
|
||||
|
||||
void Reset() {
|
||||
accept_count.clear();
|
||||
draft_count.clear();
|
||||
}
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Metrics attached to each request
|
||||
*
|
||||
* Sometimes requests can involve tree decode(e.g. parallel n).
|
||||
* The metrics is collected across all branches of the tree.
|
||||
*/
|
||||
struct RequestMetrics {
|
||||
/*! \brief Request input tokens. */
|
||||
int64_t prompt_tokens = 0;
|
||||
/*! \brief Total number of output tokens. */
|
||||
int64_t completion_tokens = 0;
|
||||
/*! \brief Total number of tokens that needs to be prefilled */
|
||||
int64_t prefill_tokens = 0;
|
||||
/*! \brief The number of processed tokens (including tokens rolled back later) in decode. */
|
||||
int64_t decode_tokens = 0;
|
||||
/*! \brief The number of tokens predicted by jump-forward decoding. */
|
||||
int64_t jump_forward_tokens = 0;
|
||||
|
||||
/*! \brief The time of adding the request to engine. */
|
||||
std::chrono::high_resolution_clock::time_point add_time_point;
|
||||
/*! \brief The time of finishing prefill stage. */
|
||||
std::chrono::high_resolution_clock::time_point prefill_end_time_point;
|
||||
/*! \brief The time of finishing all decode. */
|
||||
std::chrono::high_resolution_clock::time_point finish_time_point;
|
||||
|
||||
/*! \brief check whether the request metrics is a completed request */
|
||||
bool IsComplete() const { return prompt_tokens != 0 && completion_tokens != 0; }
|
||||
|
||||
/*! \return the prefill time in seconds */
|
||||
double GetPrefillTime() const {
|
||||
return static_cast<double>((prefill_end_time_point - add_time_point).count()) / 1e9;
|
||||
}
|
||||
|
||||
/*! \return the decode time in seconds */
|
||||
double GetDecodeTime() const {
|
||||
return static_cast<double>((finish_time_point - prefill_end_time_point).count()) / 1e9;
|
||||
}
|
||||
|
||||
/*! \return the time to first token (TTFT) in seconds */
|
||||
double GetTTFT() const {
|
||||
return static_cast<double>((prefill_end_time_point - add_time_point).count()) / 1e9;
|
||||
}
|
||||
|
||||
/*! \return the prefill time in seconds */
|
||||
double GetTotalTime() const {
|
||||
return static_cast<double>((finish_time_point - add_time_point).count()) / 1e9;
|
||||
}
|
||||
|
||||
/*! \return the inter token latency (ITL) in seconds */
|
||||
double GetInterTokenLatency() const {
|
||||
return completion_tokens > 0 ? GetTotalTime() / completion_tokens : 0.0;
|
||||
}
|
||||
|
||||
/*! \brief Reset the metric. */
|
||||
void Reset() {
|
||||
this->prompt_tokens = 0;
|
||||
this->prefill_tokens = 0;
|
||||
this->completion_tokens = 0;
|
||||
}
|
||||
/*!
|
||||
* \brief Return the request metrics in JSON.
|
||||
* \return The metrics in JSON
|
||||
*/
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
/*!
|
||||
* \brief Return OpenAI compatible usage metrics
|
||||
* \param include_extra Whether to include extra set of metrics
|
||||
*
|
||||
* \return The usage metrics in json.
|
||||
*/
|
||||
std::string AsUsageJSONStr(bool include_extra) const;
|
||||
};
|
||||
|
||||
/*! \brief Runtime metrics of engine. */
|
||||
struct EngineMetrics {
|
||||
/*! \brief The total engine time on prefill, including warmup */
|
||||
double engine_prefill_time_sum = 0;
|
||||
/*! \brief The total engine time on decode/draft/verify, including warmup */
|
||||
double engine_decode_time_sum = 0;
|
||||
/*! \brief The total engine time on jump-forward prediction. */
|
||||
double engine_jump_forward_time_sum = 0;
|
||||
/*! \brief The total number of request input tokens. */
|
||||
int64_t prompt_tokens_sum = 0;
|
||||
/*! \brief The total number of request output tokens */
|
||||
int64_t completion_tokens_sum = 0;
|
||||
/*! \brief The total number of processed tokens (excluding the prefix-cached length) in prefill */
|
||||
int64_t prefill_tokens_sum = 0;
|
||||
/*! \brief The total number of processed tokens (including tokens rolled back later) in decode. */
|
||||
int64_t decode_tokens_sum = 0;
|
||||
/*! \brief The total number of tokens predicted by jump-forward decoding. */
|
||||
int64_t jump_forward_tokens_sum = 0;
|
||||
/*! \brief metrics from last finished request. */
|
||||
RequestMetrics last_finished_request;
|
||||
/*! \brief speculative decoding metrics */
|
||||
SpecDecodeMetrics spec_decode;
|
||||
|
||||
/*! \brief The maximum batch size we track for batch decode time. */
|
||||
static constexpr const int64_t kEndFineGrainedTrackingBatchSize = 65;
|
||||
/*! \brief The list of batch decode time under different batch size. */
|
||||
std::vector<TimeCost> decode_time_by_batch_size =
|
||||
std::vector<TimeCost>(kEndFineGrainedTrackingBatchSize);
|
||||
/*! \brief The list of batch draft time (a single decode step) under different batch size. */
|
||||
std::vector<TimeCost> draft_time_by_batch_size =
|
||||
std::vector<TimeCost>(kEndFineGrainedTrackingBatchSize);
|
||||
/*! \brief The list of batch verification time under different effective batch size. */
|
||||
std::vector<TimeCost> verify_time_by_batch_size =
|
||||
std::vector<TimeCost>(kEndFineGrainedTrackingBatchSize);
|
||||
|
||||
// NOTE: we keep most update function in header
|
||||
// so they can be inlined effectively
|
||||
/*!
|
||||
* \brief Update the batch decode time for the given batch size.
|
||||
* The time will be ignored if the batch size is greater than `kMaxBatchSizeForTracking`.
|
||||
*/
|
||||
void UpdateDecodeTimeByBatchSize(int batch_size, double time) {
|
||||
if (batch_size < kEndFineGrainedTrackingBatchSize) {
|
||||
decode_time_by_batch_size[batch_size].Update(time);
|
||||
}
|
||||
}
|
||||
/*!
|
||||
* \brief Update the single-step batch draft time for the given batch size.
|
||||
* The time will be ignored if the batch size is greater than `kMaxBatchSizeForTracking`.
|
||||
*/
|
||||
void UpdateDraftTimeByBatchSize(int batch_size, double time) {
|
||||
if (batch_size < kEndFineGrainedTrackingBatchSize) {
|
||||
draft_time_by_batch_size[batch_size].Update(time);
|
||||
}
|
||||
}
|
||||
/*!
|
||||
* \brief Update the batch decode time for the given effective batch sizPe.
|
||||
* The time will be ignored if the effective batch size is greater than
|
||||
* `kMaxBatchSizeForTracking`.
|
||||
*/
|
||||
void UpdateVerifyTimeByBatchSize(int effective_batch_size, double time) {
|
||||
if (effective_batch_size < kEndFineGrainedTrackingBatchSize) {
|
||||
verify_time_by_batch_size[effective_batch_size].Update(time);
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Update global engine metrics as we finish a request
|
||||
* by including the information from the finished request.
|
||||
*/
|
||||
void RequestFinishUpdate(const RequestMetrics& request_metrics) {
|
||||
prompt_tokens_sum += request_metrics.prompt_tokens;
|
||||
prefill_tokens_sum += request_metrics.prefill_tokens;
|
||||
completion_tokens_sum += request_metrics.completion_tokens;
|
||||
decode_tokens_sum += request_metrics.decode_tokens;
|
||||
jump_forward_tokens_sum += request_metrics.jump_forward_tokens;
|
||||
last_finished_request = request_metrics;
|
||||
}
|
||||
/*!
|
||||
* \brief Return the engine runtime metrics in JSON.
|
||||
* \return The metrics in JSON
|
||||
*/
|
||||
tvm::ffi::json::Object AsJSON() const;
|
||||
|
||||
/*!
|
||||
* \brief return engine metrics as usage json string.
|
||||
* \return The resulting usage json string.
|
||||
*/
|
||||
std::string AsUsageJSONStr() const;
|
||||
|
||||
/*! \brief Reset all the metrics. */
|
||||
void Reset();
|
||||
};
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_METRIC_H_
|
||||
+1293
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,418 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/model.h
|
||||
* \brief The header for runtime module of LLM functions (prefill/decode/etc.)
|
||||
*/
|
||||
|
||||
#ifndef MLC_LLM_SERVE_MODEL_H_
|
||||
#define MLC_LLM_SERVE_MODEL_H_
|
||||
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
#include "../base.h"
|
||||
#include "../support/result.h"
|
||||
#include "config.h"
|
||||
#include "draft_token_workspace_manager.h"
|
||||
#include "event_trace_recorder.h"
|
||||
#include "function_table.h"
|
||||
#include "logit_processor.h"
|
||||
#include "sampler/sampler.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::Device;
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
using tvm::ffi::Shape;
|
||||
|
||||
// Declare the sampler class for `Model::CreateSampler`.
|
||||
class Sampler;
|
||||
|
||||
/*!
|
||||
* \brief The workspace tensors that may be shared across different
|
||||
* calls to Model. For example, the prefill action use the `embeddings`
|
||||
* workspace for the concatenated embeddings of different sequences.
|
||||
* The workspace tensor is created by Model but owned by engine.
|
||||
*/
|
||||
struct ModelWorkspace {
|
||||
/*!
|
||||
* \brief The embedding tensor. It can be either an Tensor when tensor
|
||||
* model parallelism is not enabled, or a DRef when using tensor model parallelism.
|
||||
*/
|
||||
ObjectRef embeddings{nullptr};
|
||||
/*!
|
||||
* \brief The hidden_states tensor for the current batch. It can be either an Tensor when tensor
|
||||
* model parallelism is not enabled, or a DRef when using tensor model parallelism.
|
||||
*/
|
||||
ObjectRef hidden_states{nullptr};
|
||||
|
||||
/*!
|
||||
* \brief The draft token probabilities tensor for the current batch.
|
||||
*/
|
||||
Tensor draft_probs{nullptr};
|
||||
|
||||
/*!
|
||||
* \brief The hidden_states tensor storing the hidden_states of draft tokens of all requests.
|
||||
*/
|
||||
ObjectRef draft_hidden_states_storage{nullptr};
|
||||
|
||||
/*!
|
||||
* \brief The draft token probabilities tensor storing the probabilities of draft tokens of all
|
||||
* requests.
|
||||
*/
|
||||
Tensor draft_probs_storage{nullptr};
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief The model module for LLM functions.
|
||||
* It runs an LLM, and has an internal KV cache that maintains
|
||||
* the history KV values of all processed tokens.
|
||||
*
|
||||
* It contains the following functions:
|
||||
*
|
||||
* Model related:
|
||||
* - "token_embed": take token ids as input and return the embeddings,
|
||||
* - "batch_prefill": take embedding of a single sequence
|
||||
* as input, forward the embedding through LLM and return the logits,
|
||||
* - "decode": take the embeddings of the last-committed token of an
|
||||
* entire batch as input, forward through LLM and return the logits
|
||||
* for all sequences in the batch,
|
||||
* - "softmax_with_temperature": take logits and temperatures, return
|
||||
* probabilities.
|
||||
*
|
||||
* KV cache related:
|
||||
* - "create_kv_cache": create the KV cache for this module,
|
||||
* - "add_new_sequence": add (declare) a new sequence in the KV cache,
|
||||
* - "remove_sequence": remove a sequence from KV cache.
|
||||
*
|
||||
* ... and some other auxiliary functions.
|
||||
*/
|
||||
class ModelObj : public Object {
|
||||
public:
|
||||
/*********************** Model Computation ***********************/
|
||||
|
||||
/*!
|
||||
* \brief Compute embeddings for the input token ids.
|
||||
* When the input destination pointer is defined, it in-place writes the
|
||||
* embedding into the input destination array at the given offset.
|
||||
* Otherwise, the embeddings will be directly returned back.
|
||||
* \param token_ids The token ids to compute embedding for.
|
||||
* \param dst The destination array of the embedding lookup.
|
||||
* \param offset The token offset where the computed embeddings will be written
|
||||
* into the destination array.
|
||||
* \return The updated destination embedding array or the computed embeddings.
|
||||
* \note When `dst` is undefined, we require `offset` to be 0.
|
||||
*/
|
||||
virtual ObjectRef TokenEmbed(Shape batch_token_ids, ObjectRef* dst = nullptr, int offset = 0) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Compute embeddings for the input image.
|
||||
* \param image The image to compute embedding for.
|
||||
* \return The computed embeddings.
|
||||
*/
|
||||
virtual ObjectRef ImageEmbed(const Tensor& image, ObjectRef* dst = nullptr, int offset = 0) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Fuse the embeddings and hidden_states.
|
||||
* \param embeddings The embedding of the input to be prefilled.
|
||||
* \param previous_hidden_states The hidden_states from previous base model.
|
||||
* \param batch_size Batch size.
|
||||
* \param seq_len Sequence length.
|
||||
* \return The fused hidden_states.
|
||||
*/
|
||||
virtual ObjectRef FuseEmbedHidden(const ObjectRef& embeddings,
|
||||
const ObjectRef& previous_hidden_states, int batch_size,
|
||||
int seq_len) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Return if the model has lm_head so that we can get logits.
|
||||
*/
|
||||
virtual bool CanGetLogits() = 0;
|
||||
|
||||
/*!
|
||||
* \brief Compute logits for last hidden_states.
|
||||
* \param last_hidden_states The last hidden_states to compute logits for.
|
||||
* \return The computed logits.
|
||||
*/
|
||||
virtual Tensor GetLogits(const ObjectRef& last_hidden_states) = 0;
|
||||
|
||||
virtual Array<Tensor> GetMultiStepLogits(const ObjectRef& last_hidden_states) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Batch prefill function. Embedding in, logits out.
|
||||
* The embedding order of sequences in `embedding_arr` follows
|
||||
* the order of `seq_ids`.
|
||||
* \param embeddings The embedding of the input to be prefilled.
|
||||
* \param seq_id The id of the sequence in the KV cache.
|
||||
* \param lengths The length of each sequence to prefill.
|
||||
* \return The logits for the next token.
|
||||
*/
|
||||
virtual Tensor BatchPrefill(const ObjectRef& embeddings, const std::vector<int64_t>& seq_ids,
|
||||
const std::vector<int>& lengths) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Batch prefill function. Input hidden_states are computed from
|
||||
* input embeddings and previous hidden_states, output last hidden_states.
|
||||
* \param hidden_states The hidden_states of the input to be prefilled.
|
||||
* \param seq_id The id of the sequence in the KV cache.
|
||||
* \param lengths The length of each sequence to prefill.
|
||||
* \return The hidden_states for the next token.
|
||||
*/
|
||||
virtual ObjectRef BatchPrefillToLastHidden(const ObjectRef& hidden_states,
|
||||
const std::vector<int64_t>& seq_ids,
|
||||
const std::vector<int>& lengths) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Batch decode function. Embedding in, logits out.
|
||||
* The embedding order of sequences in `embeddings` follows
|
||||
* the order of `seq_ids`.
|
||||
* \param embeddings The embedding of last generated token in the entire batch.
|
||||
* \param seq_id The id of the sequence in the KV cache.
|
||||
* \return The logits for the next token for each sequence in the batch.
|
||||
*/
|
||||
virtual Tensor BatchDecode(const ObjectRef& embeddings, const std::vector<int64_t>& seq_ids) = 0;
|
||||
|
||||
virtual 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) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Batch decode function. Input hidden_states are computed from
|
||||
* input embeddings and previous hidden_states, output last hidden_states.
|
||||
* \param hidden_states The hidden_states of last generated token in the entire batch.
|
||||
* \param seq_id The id of the sequence in the KV cache.
|
||||
* \return The hidden_states for the next token for each sequence in the batch.
|
||||
*/
|
||||
virtual ObjectRef BatchDecodeToLastHidden(const ObjectRef& hidden_states,
|
||||
const std::vector<int64_t>& seq_ids) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Batch verify function. Embedding in, logits out.
|
||||
* \param embeddings The embedding of the input to be verified.
|
||||
* \param seq_id The id of the sequence in the KV cache.
|
||||
* \param lengths The length of each sequence to verify.
|
||||
* \param token_tree_parent_ptr The parent pointers of the token tree.
|
||||
* It's size is the sum of "lengths". It contains a batch of independent trees,
|
||||
* one for each sequence. Parent being "-1" means the node is a root.
|
||||
* \return The logits for the draft token for each sequence in the batch.
|
||||
* \note The function runs for **every** sequence in the batch.
|
||||
* That is to say, it does not accept "running a verify step for a subset
|
||||
* of the full batch".
|
||||
*/
|
||||
virtual 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) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Batch verify function. Input hidden_states are computed from
|
||||
* input embeddings and previous hidden_states, output last hidden_states.
|
||||
* \param hidden_states The hidden_states of the input to be verified.
|
||||
* \param seq_id The id of the sequence in the KV cache.
|
||||
* \param lengths The length of each sequence to verify.
|
||||
* \param token_tree_parent_ptr The parent pointers of the token tree.
|
||||
* It's size is the sum of "lengths". It contains a batch of independent trees,
|
||||
* one for each sequence. Parent being "-1" means the node is a root.
|
||||
* \return The hidden_states for the draft token for each sequence in the batch.
|
||||
* \note The function runs for **every** sequence in the batch.
|
||||
* That is to say, it does not accept "running a verify step for a subset
|
||||
* of the full batch".
|
||||
*/
|
||||
virtual ObjectRef BatchVerifyToLastHidden(const ObjectRef& hidden_states,
|
||||
const std::vector<int64_t>& seq_ids,
|
||||
const std::vector<int>& lengths,
|
||||
const std::vector<int64_t>& token_tree_parent_ptr) = 0;
|
||||
|
||||
/*********************** KV Cache Management ***********************/
|
||||
|
||||
/*!
|
||||
* \brief Create the KV cache inside the model with regard to the input config.
|
||||
* \param page_size The number of consecutive tokens handled in each page in paged KV cache.
|
||||
* \param max_num_sequence The maximum number of sequences that are allowed to be
|
||||
* processed by the KV cache at any time.
|
||||
* \param max_total_sequence_length The maximum length allowed for a single sequence
|
||||
* in the engine.
|
||||
* \param prefill_chunk_size The maximum total number of tokens whose KV data
|
||||
* are allowed to exist in the KV cache at any time.
|
||||
* \param max_history_size The maximum history size for RNN state to roll back.
|
||||
* The KV cache does not need this.
|
||||
* \param prefix_cache_max_num_recycling_seqs The maximum number of recycling
|
||||
* sequences kept by prefix cache. For hybrid models, the RNN state needs
|
||||
* extra slots beyond max_num_sequence to hold these recycling sequences
|
||||
* simultaneously with active sequences.
|
||||
*/
|
||||
virtual 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) = 0;
|
||||
|
||||
/*! \brief Add a new sequence with the given sequence id to the KV cache. */
|
||||
virtual void AddNewSequence(int64_t seq_id) = 0;
|
||||
|
||||
/*! \brief Fork a sequence from a given parent sequence. */
|
||||
virtual void ForkSequence(int64_t parent_seq_id, int64_t child_seq_id, int64_t fork_pos = -1) = 0;
|
||||
|
||||
/*! \brief Remove the given sequence from the KV cache in the model. */
|
||||
virtual void RemoveSequence(int64_t seq_id) = 0;
|
||||
|
||||
/*! \brief Pop out N pages from KV cache. */
|
||||
virtual void PopNFromKVCache(int64_t seq_id, int num_tokens) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Commit the accepted token tree nodes to KV cache.
|
||||
* The unaccepted token tree node will be removed from KV cache.
|
||||
* This is usually used in the verification stage of speculative decoding.
|
||||
*/
|
||||
virtual void CommitAcceptedTokenTreeNodesToKVCache(
|
||||
const std::vector<int64_t>& seq_ids, const std::vector<int64_t>& accepted_leaf_indices) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Enabling sliding window for the given sequence.
|
||||
* It is a no-op if the model does not support sliding window.
|
||||
* \note Given this operation is tied with the underlying KV cache,
|
||||
* we add the function in Model interface to expose this for Engine.
|
||||
* This may be optimized with decoupling KV cache and Model in the future.
|
||||
*/
|
||||
virtual void EnableSlidingWindowForSeq(int64_t seq_id) = 0;
|
||||
|
||||
/*! \brief Prepare for the disaggregation KV data receive for the specified sequence and length.*/
|
||||
virtual Shape DisaggPrepareKVRecv(int64_t seq_id, int length) = 0;
|
||||
|
||||
/*! \brief Prepare for the disaggregation KV data send for the specified sequence and length.*/
|
||||
virtual void DisaggMarkKVSend(int64_t seq_id, int begin_pos, Shape compressed_kv_append_metadata,
|
||||
int dst_group_offset) = 0;
|
||||
|
||||
/************** Raw Info Query **************/
|
||||
|
||||
/*! \brief Return the metadata JSON object of the model. */
|
||||
virtual ModelMetadata GetMetadata() const = 0;
|
||||
|
||||
/*! \brief Get the number of available pages in KV cache. */
|
||||
virtual int GetNumAvailablePages() const = 0;
|
||||
|
||||
/*! \brief Get the current total sequence length in the KV cache. */
|
||||
virtual int GetCurrentTotalSequenceLength() const = 0;
|
||||
|
||||
/*********************** Utilities ***********************/
|
||||
|
||||
/*! \brief Load the model's weight parameters, which is not loaded at construction time. */
|
||||
virtual void LoadParams() = 0;
|
||||
|
||||
/*!
|
||||
* \brief Set the maximum number of sequences to be processed for the model,
|
||||
* which is not initialized at construction time.
|
||||
*/
|
||||
virtual void SetMaxNumSequence(int max_num_sequence) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Set the prefill chunk size for the model,
|
||||
* which is not initialized at construction time.
|
||||
*/
|
||||
virtual void SetPrefillChunkSize(int prefill_chunk_size) = 0;
|
||||
|
||||
/*! \brief Create a logit processor from this model. */
|
||||
virtual LogitProcessor CreateLogitProcessor(int max_num_token,
|
||||
Optional<EventTraceRecorder> trace_recorder) = 0;
|
||||
|
||||
/*! \brief Create a sampler from this model. */
|
||||
virtual Sampler CreateSampler(int max_num_sample, int num_models,
|
||||
Optional<EventTraceRecorder> trace_recorder) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Estimate number of CPU units required to drive the model
|
||||
* executing during TP.
|
||||
* \note This normally equals to the number of TP shards (or 0 if
|
||||
* the model does not use TP) and can be used to hint runtime to
|
||||
* avoid overuse cores in other places.
|
||||
*/
|
||||
virtual int EstimateHostCPURequirement() const = 0;
|
||||
|
||||
/*! \brief Get the sliding window size of the model. "-1" means sliding window is not enabled. */
|
||||
virtual int GetSlidingWindowSize() const = 0;
|
||||
|
||||
/*! \brief Get the attention sink size of the model. */
|
||||
virtual int GetAttentionSinkSize() const = 0;
|
||||
|
||||
/*! \brief Allocate an embedding tensor with the prefill chunk size. */
|
||||
virtual ObjectRef AllocEmbeddingTensor() = 0;
|
||||
|
||||
/*! \brief Allocate an hidden_states tensor with the prefill chunk size. */
|
||||
virtual ObjectRef AllocHiddenStatesTensor() = 0;
|
||||
|
||||
/*! \brief Reset the model KV cache and other metrics. */
|
||||
virtual void Reset() = 0;
|
||||
|
||||
/*********************** Utilities for speculative decoding. ***********************/
|
||||
|
||||
virtual DraftTokenWorkspaceManager CreateDraftTokenWorkspaceManager(int max_num_token) = 0;
|
||||
|
||||
/*! \brief Gather the hidden_states of the given indices and in-place update the dst tensor. */
|
||||
virtual ObjectRef GatherHiddenStates(const ObjectRef& input, const std::vector<int>& indices,
|
||||
ObjectRef* dst) = 0;
|
||||
|
||||
/*! \brief Scatter the hidden_states of the given indices to the dst tensor. */
|
||||
virtual void ScatterHiddenStates(const ObjectRef& input, const std::vector<int>& indices,
|
||||
ObjectRef* dst) = 0;
|
||||
|
||||
/*! \brief Gather the draft token probabilities of the given indices and in-place update the dst
|
||||
* tensor. */
|
||||
virtual Tensor GatherDraftProbs(const Tensor& input, const std::vector<int>& indices,
|
||||
Tensor* dst) = 0;
|
||||
|
||||
/*! \brief Scatter the draft token probabilities of the given indices to the dst tensor. */
|
||||
virtual void ScatterDraftProbs(const Tensor& input, const std::vector<int>& indices,
|
||||
Tensor* dst) = 0;
|
||||
|
||||
/************** Debug/Profile **************/
|
||||
|
||||
/*! \brief Call the given global function on all workers. Only for debug purpose. */
|
||||
virtual void DebugCallFuncOnAllAllWorker(const String& func_name, Optional<String> func_args) = 0;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<ModelObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.Model", ModelObj, Object);
|
||||
};
|
||||
|
||||
class Model : public ObjectRef {
|
||||
public:
|
||||
/*!
|
||||
* \brief Create the runtime module for LLM functions.
|
||||
* \param reload_lib_path The model library path.
|
||||
* \param model_path The path to the model weight parameters.
|
||||
* \param model_config The model config json object.
|
||||
* \param device The device to run the model on.
|
||||
* \param session The session to run the model on.
|
||||
* \param num_shards The number of tensor parallel shards of the model.
|
||||
* \param num_stages The number of pipeline parallel stages of the model.
|
||||
* \param trace_enabled A boolean indicating whether tracing is enabled.
|
||||
* \return The created runtime module.
|
||||
*/
|
||||
static 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);
|
||||
|
||||
/*!
|
||||
* Load the model config from the given model path.
|
||||
* \param model_path The path to the model weight parameters.
|
||||
* \return The model config json object.
|
||||
*/
|
||||
static Result<tvm::ffi::json::Object> LoadModelConfig(const String& model_path);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Model, ObjectRef, ModelObj);
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_MODEL_H_
|
||||
@@ -0,0 +1,442 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/prefix_cache.cc
|
||||
*/
|
||||
#include "prefix_cache.h"
|
||||
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/support/cuda/nvtx.h>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::support::NVTXScopedRange;
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() { PrefixCacheObj::RegisterReflection(); }
|
||||
|
||||
/*!
|
||||
* \brief The implementation of prefix cache.
|
||||
*/
|
||||
class PrefixCacheImpl : public PrefixCacheObj {
|
||||
public:
|
||||
/*!
|
||||
* \brief Constructor of paged radix tree.
|
||||
* \param max_num_recycling_seqs The maximum number of sequences in prefix cache.
|
||||
* \param remove_callback The optional callback function to call when removing a sequence.
|
||||
*/
|
||||
explicit PrefixCacheImpl(size_t max_num_recycling_seqs, PrefixCacheRemoveCallback remove_callback)
|
||||
: radix_tree_(PagedRadixTree::Create()),
|
||||
max_num_recycling_seqs_(max_num_recycling_seqs),
|
||||
remove_callback_(std::move(remove_callback)) {
|
||||
recycling_seq_lrus_.clear();
|
||||
reversed_recycling_seq_lrus_.clear();
|
||||
seq_states_.clear();
|
||||
seq_sliding_window_infos_.clear();
|
||||
lru_counter_ = 0;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Insert a new tokenized sequence into Prefix Cache.
|
||||
* \param seq_id The sequence ID.
|
||||
* \param tokens The tokens of tokenized sequence.
|
||||
* \param sliding_window_size The sliding window size for the sequence, -1 as sliding window
|
||||
* disabled.
|
||||
* \param attention_sink_size The attention sink size for the sequence, 0 by default.
|
||||
* \return The matched result.
|
||||
*/
|
||||
PrefixCacheMatchedResult InsertSequence(int64_t seq_id, std::vector<int32_t> tokens,
|
||||
int sliding_window_size, int attention_sink_size) final {
|
||||
TVM_FFI_ICHECK_NE(sliding_window_size, 0);
|
||||
TVM_FFI_ICHECK_GE(attention_sink_size, 0);
|
||||
TVM_FFI_ICHECK(seq_states_.find(seq_id) == seq_states_.end());
|
||||
TVM_FFI_ICHECK(seq_sliding_window_infos_.find(seq_id) == seq_sliding_window_infos_.end());
|
||||
TVM_FFI_ICHECK(!tokens.empty());
|
||||
CommitSequenceExtention();
|
||||
tokens.pop_back();
|
||||
auto [matched_offset, matched_seqs] = radix_tree_->MatchPrefix(tokens);
|
||||
std::pair<int, size_t> sliding_window_info{sliding_window_size, attention_sink_size};
|
||||
// No prefix matched, directly adding new sequence.
|
||||
if (!matched_offset) {
|
||||
radix_tree_->AddSequence(seq_id);
|
||||
seq_states_.emplace(seq_id, SequenceState::kActive);
|
||||
seq_sliding_window_infos_.emplace(seq_id, sliding_window_info);
|
||||
return PrefixCacheMatchedResult{0, -1, -1, 0};
|
||||
}
|
||||
|
||||
TVM_FFI_ICHECK(!matched_seqs.empty());
|
||||
|
||||
// The reusage of recycling sequences logic is different between with/without sliding window
|
||||
// enabled.
|
||||
if (sliding_window_size != -1) {
|
||||
// If sliding window enabled, the reusage of recycling sequences should be limited to exactly
|
||||
// matched. And no rolling back is allowed due to the sliding window.
|
||||
for (int64_t matched_seq_id : matched_seqs) {
|
||||
if (seq_states_.at(matched_seq_id) == SequenceState::kRecycling &&
|
||||
seq_sliding_window_infos_.at(matched_seq_id) == sliding_window_info) {
|
||||
size_t matched_seq_length = radix_tree_->GetSequenceLength(matched_seq_id);
|
||||
if (matched_seq_length == matched_offset) {
|
||||
ReuseRecyclingSequence(matched_seq_id);
|
||||
return PrefixCacheMatchedResult{matched_offset, -1, matched_seq_id, 0};
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// If sliding window is not enabled, we can greedily reuse the shortest recycling sequence
|
||||
// without sliding window, so that the loss or roll back of trailing tokens will be minimum.
|
||||
size_t shortest_recycling_seq_length = 0;
|
||||
int64_t shortest_recycling_seq_id = -1;
|
||||
|
||||
for (int64_t matched_seq_id : matched_seqs) {
|
||||
if (seq_states_.at(matched_seq_id) == SequenceState::kRecycling &&
|
||||
seq_sliding_window_infos_.at(matched_seq_id) == sliding_window_info) {
|
||||
size_t matched_seq_length = radix_tree_->GetSequenceLength(matched_seq_id);
|
||||
if (shortest_recycling_seq_id == -1 ||
|
||||
matched_seq_length < shortest_recycling_seq_length) {
|
||||
shortest_recycling_seq_id = matched_seq_id;
|
||||
shortest_recycling_seq_length = matched_seq_length;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (shortest_recycling_seq_id != -1 && matched_offset > shortest_recycling_seq_length * 0.9) {
|
||||
ReuseRecyclingSequence(shortest_recycling_seq_id);
|
||||
if (shortest_recycling_seq_length > matched_offset) {
|
||||
// Recycling sequence is longer than new sequence, rolling back the redundant trailing
|
||||
// tokens, to match the new sequence.
|
||||
radix_tree_->RollBackSequence(shortest_recycling_seq_id,
|
||||
shortest_recycling_seq_length - matched_offset);
|
||||
}
|
||||
return PrefixCacheMatchedResult{matched_offset, -1, shortest_recycling_seq_id,
|
||||
shortest_recycling_seq_length - matched_offset};
|
||||
}
|
||||
// No reusage of recycling sequence, fallback to forking matched sequence. Currently, we only
|
||||
// fork from sequence without sliding window, due to current paged KVCache implementation.
|
||||
size_t longest_forking_offset = 0;
|
||||
int64_t longest_forking_seq_id = -1;
|
||||
for (int64_t matched_seq_id : matched_seqs) {
|
||||
auto [matched_seq_sliding_window_size, matched_seq_attention_sink_size] =
|
||||
seq_sliding_window_infos_.at(matched_seq_id);
|
||||
if (matched_seq_sliding_window_size != -1) {
|
||||
continue;
|
||||
}
|
||||
// If the matched is not enabled with sliding window, we can fork within matched offset
|
||||
// tokens arbitrarily.
|
||||
if (matched_offset > longest_forking_offset) {
|
||||
longest_forking_offset = matched_offset;
|
||||
longest_forking_seq_id = matched_seq_id;
|
||||
}
|
||||
}
|
||||
if (longest_forking_offset > 0) {
|
||||
radix_tree_->ForkSequence(seq_id, longest_forking_seq_id, longest_forking_offset);
|
||||
seq_states_.emplace(seq_id, SequenceState::kActive);
|
||||
seq_sliding_window_infos_.emplace(seq_id, sliding_window_info);
|
||||
return PrefixCacheMatchedResult{longest_forking_offset, longest_forking_seq_id, -1, 0};
|
||||
}
|
||||
}
|
||||
// No forking from matched sequence, fallback to adding new sequence.
|
||||
radix_tree_->AddSequence(seq_id);
|
||||
seq_states_.emplace(seq_id, SequenceState::kActive);
|
||||
seq_sliding_window_infos_.emplace(seq_id, sliding_window_info);
|
||||
return PrefixCacheMatchedResult{0, -1, -1, 0};
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Extend a sequence with new tokenized sequence suffix.
|
||||
* \param seq_id The sequence to be extended.
|
||||
* \param tokens The tokens of tokenized sequence suffix to extend.
|
||||
* \throw Error if the given sequence id is not valid or active.
|
||||
*/
|
||||
void ExtendSequence(int64_t seq_id, const std::vector<int32_t>& tokens) final {
|
||||
uncommitted_extended_token_ids_.emplace_back(seq_id, tokens);
|
||||
}
|
||||
|
||||
void CommitSequenceExtention() final {
|
||||
if (uncommitted_extended_token_ids_.empty()) {
|
||||
return;
|
||||
}
|
||||
NVTXScopedRange nvtx_scope("PrefixCache commit sequence extension");
|
||||
for (const auto& [seq_id, uncommitted_token_ids] : uncommitted_extended_token_ids_) {
|
||||
if (!HasSequence(seq_id)) {
|
||||
// The sequence has been removed. Hence no action is needed.
|
||||
continue;
|
||||
}
|
||||
const auto& it = seq_states_.find(seq_id);
|
||||
TVM_FFI_ICHECK(it == seq_states_.end() || it->second == SequenceState::kActive);
|
||||
radix_tree_->ExtendSequence(seq_id, uncommitted_token_ids);
|
||||
}
|
||||
uncommitted_extended_token_ids_.clear();
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Roll back a sequence by number of tokens.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \param num_tokens The number of tokens to be rolled back.
|
||||
* \throw Error if the given sequence id is not valid or active.
|
||||
*/
|
||||
void RollBackSequence(int64_t seq_id, size_t num_tokens) final {
|
||||
CommitSequenceExtention();
|
||||
TVM_FFI_ICHECK(seq_states_.at(seq_id) == SequenceState::kActive);
|
||||
radix_tree_->RollBackSequence(seq_id, num_tokens);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Recycle a sequence. The recycled sequence will not be removed immediately, as long as
|
||||
* memory is sufficient and the number of sequence in prefix cache belows the maximum number of
|
||||
* sequence. And it will be reused again in the future request.
|
||||
* \param seq_id The sequence to be recycled.
|
||||
* \param lazy The flag if the sequence should be removed lazily or intermediary.
|
||||
* \throw Error if the given sequence id is not valid.
|
||||
*/
|
||||
void RecycleSequence(int64_t seq_id, bool lazy = true) final {
|
||||
CommitSequenceExtention();
|
||||
TVM_FFI_ICHECK(seq_states_.at(seq_id) == SequenceState::kActive);
|
||||
TVM_FFI_ICHECK(recycling_seq_lrus_.find(seq_id) == recycling_seq_lrus_.end());
|
||||
if (lazy && max_num_recycling_seqs_ != 0) {
|
||||
// Remove the sequence lazily.
|
||||
if (recycling_seq_lrus_.size() == max_num_recycling_seqs_) {
|
||||
// If prefix cache has reached maximum number of recycling sequences, try to pop one
|
||||
// recycling sequence.
|
||||
TVM_FFI_ICHECK(TryFreeMemory());
|
||||
TVM_FFI_ICHECK_EQ(recycling_seq_lrus_.size(), max_num_recycling_seqs_ - 1);
|
||||
}
|
||||
seq_states_.at(seq_id) = SequenceState::kRecycling;
|
||||
++lru_counter_;
|
||||
recycling_seq_lrus_.emplace(seq_id, lru_counter_);
|
||||
reversed_recycling_seq_lrus_.emplace(lru_counter_, seq_id);
|
||||
} else {
|
||||
// Remove the sequence intermediately.
|
||||
radix_tree_->RemoveSequence(seq_id);
|
||||
if (remove_callback_ != nullptr) {
|
||||
remove_callback_(seq_id);
|
||||
}
|
||||
TVM_FFI_ICHECK(seq_states_.erase(seq_id));
|
||||
TVM_FFI_ICHECK(seq_sliding_window_infos_.erase(seq_id));
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Try to remove recycling sequence to free up memory. It will remove the oldest recycling
|
||||
sequence.
|
||||
* \return The flag if there is a sequence removed. In other word, return true when memory is
|
||||
freed successfully.
|
||||
* \throw Error if the given sequence id is not valid.
|
||||
*/
|
||||
bool TryFreeMemory() final {
|
||||
NVTXScopedRange nvtx_scope("PrefixCache TryFreeMemory");
|
||||
if (reversed_recycling_seq_lrus_.empty()) {
|
||||
// There is no recycling sequence. No memory can be freed.
|
||||
return false;
|
||||
}
|
||||
auto [lru, seq_id] = *reversed_recycling_seq_lrus_.begin();
|
||||
TVM_FFI_ICHECK(seq_states_.at(seq_id) == SequenceState::kRecycling);
|
||||
TVM_FFI_ICHECK_EQ(recycling_seq_lrus_.at(seq_id), lru);
|
||||
radix_tree_->RemoveSequence(seq_id);
|
||||
if (remove_callback_ != nullptr) {
|
||||
remove_callback_(seq_id);
|
||||
}
|
||||
TVM_FFI_ICHECK(seq_states_.erase(seq_id));
|
||||
TVM_FFI_ICHECK(recycling_seq_lrus_.erase(seq_id));
|
||||
TVM_FFI_ICHECK(reversed_recycling_seq_lrus_.erase(lru));
|
||||
TVM_FFI_ICHECK(seq_sliding_window_infos_.erase(seq_id));
|
||||
return true;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Check if a sequence exists.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \return The sequence existence.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
bool HasSequence(int64_t seq_id) final { return radix_tree_->HasSequence(seq_id); }
|
||||
|
||||
/*!
|
||||
* \brief Reset the prefix cache to initial status.
|
||||
*/
|
||||
void Reset() final {
|
||||
radix_tree_->Reset();
|
||||
recycling_seq_lrus_.clear();
|
||||
reversed_recycling_seq_lrus_.clear();
|
||||
seq_states_.clear();
|
||||
seq_sliding_window_infos_.clear();
|
||||
uncommitted_extended_token_ids_.clear();
|
||||
lru_counter_ = 0;
|
||||
}
|
||||
|
||||
PrefixCacheMode Mode() final { return PrefixCacheMode::kRadix; }
|
||||
|
||||
private:
|
||||
void ReuseRecyclingSequence(int64_t seq_id) {
|
||||
TVM_FFI_ICHECK(seq_states_.at(seq_id) == SequenceState::kRecycling);
|
||||
size_t lru = recycling_seq_lrus_.at(seq_id);
|
||||
TVM_FFI_ICHECK_EQ(reversed_recycling_seq_lrus_.at(lru), seq_id);
|
||||
seq_states_.at(seq_id) = SequenceState::kActive;
|
||||
TVM_FFI_ICHECK(recycling_seq_lrus_.erase(seq_id));
|
||||
TVM_FFI_ICHECK(reversed_recycling_seq_lrus_.erase(lru));
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief The sequence states.
|
||||
*/
|
||||
enum class SequenceState : int {
|
||||
/*!
|
||||
* \brief The state of active sequence. In this state, the sequence can be forked only. When
|
||||
* recycling a sequence, it will transfer to kRecycling.
|
||||
*/
|
||||
kActive = 0,
|
||||
/*!
|
||||
* \brief The state of recycling sequence. In this state, the sequence can be forked or be
|
||||
* reused. And it will transfer to kActive only when reused.
|
||||
*/
|
||||
kRecycling = 1,
|
||||
};
|
||||
/*!
|
||||
* \brief The core data structure radix tree.
|
||||
*/
|
||||
PagedRadixTree radix_tree_;
|
||||
/*!
|
||||
* \brief The map from sequence to LRU time stamps.
|
||||
*/
|
||||
std::unordered_map<int64_t, size_t> recycling_seq_lrus_;
|
||||
/*!
|
||||
* \brief The map from LRU time stamps to sequence, used to find the sequence with earliest LRU
|
||||
* time stamp.
|
||||
*/
|
||||
std::unordered_map<size_t, int64_t> reversed_recycling_seq_lrus_;
|
||||
/*!
|
||||
* \brief The maximum number of recycling sequences in prefix cache. Set -1 as infinite prefix
|
||||
* cache.
|
||||
*/
|
||||
int max_num_recycling_seqs_ = -1;
|
||||
/*!
|
||||
* \brief The LRU counter.
|
||||
*/
|
||||
size_t lru_counter_ = 0;
|
||||
/*!
|
||||
* \brief The callback function to call when removing a sequence. This can be used to
|
||||
* removing sequence in KVCache and return sequence ID to ID manager lazily
|
||||
*/
|
||||
PrefixCacheRemoveCallback remove_callback_ = nullptr;
|
||||
/*!
|
||||
* \brief The map from sequence to its sequence states.
|
||||
*/
|
||||
std::unordered_map<int64_t, SequenceState> seq_states_;
|
||||
/*!
|
||||
* \brief The map from sequence to its sliding window information. The sliding window information
|
||||
* is a pair of sliding window size and attention sink size. The sliding window size is -1 for
|
||||
* sliding window disabled, or positive for sliding window size. The attention sink size is
|
||||
* non-negative and used when sliding window size is positive.
|
||||
*/
|
||||
std::unordered_map<int64_t, std::pair<int, size_t>> seq_sliding_window_infos_;
|
||||
/*!
|
||||
* \brief The collection of uncommitted extended token ids of sequences.
|
||||
* The "ExtendSequence" method only lazily add token ids into this collection,
|
||||
* and these uncommitted token ids will be committed when needed.
|
||||
*
|
||||
* Note: Since the tokens stored are references, CommitSequenceExtention should be called after
|
||||
* each action, to avoid the uncaught changes of uncomitted extended token ids.
|
||||
*/
|
||||
std::vector<std::pair<int64_t, const std::vector<int32_t>&>> uncommitted_extended_token_ids_;
|
||||
}; // namespace serve
|
||||
|
||||
/*!
|
||||
* \brief The implementation of no prefix cache.
|
||||
*/
|
||||
class NoPrefixCache : public PrefixCacheObj {
|
||||
public:
|
||||
/*!
|
||||
* \brief Insert a new tokenized sequence into Prefix Cache.
|
||||
* \param seq_id The sequence ID.
|
||||
* \param tokens The tokens of tokenized sequence.
|
||||
* \param sliding_window_size The sliding window size for the sequence, -1 as sliding window
|
||||
* disabled.
|
||||
* \param attention_sink_size The attention sink size for the sequence, 0 by default.
|
||||
* \return The matched result.
|
||||
*/
|
||||
PrefixCacheMatchedResult InsertSequence(int64_t seq_id, std::vector<int32_t> tokens,
|
||||
int sliding_window_size, int attention_sink_size) final {
|
||||
// Since there is no prefix cache, always return as new sequence.
|
||||
return PrefixCacheMatchedResult{0, -1, -1, 0};
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Extend a sequence with new tokenized sequence suffix.
|
||||
* \param seq_id The sequence to be extended.
|
||||
* \param tokens The tokens of tokenized sequence suffix to extend.
|
||||
* \throw Error if called since this should never be called.
|
||||
*/
|
||||
void ExtendSequence(int64_t seq_id, const std::vector<int32_t>& tokens) final {
|
||||
// No-op;
|
||||
}
|
||||
|
||||
void CommitSequenceExtention() final {
|
||||
// No-op;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Roll back a sequence by number of tokens.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \param num_tokens The number of tokens to be rolled back.
|
||||
* \throw Error if called since this should never be called.
|
||||
*/
|
||||
void RollBackSequence(int64_t seq_id, size_t num_tokens) final {
|
||||
// Since there is no prefix cache, this method should never be called.
|
||||
LOG(FATAL) << "Unreachable code.";
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Recycle a sequence. The recycled sequence will not be removed immediately, as long as
|
||||
* memory is sufficient and the number of sequence in prefix cache belows the maximum number of
|
||||
* sequence. And it will be reused again in the future request.
|
||||
* \param seq_id The sequence to be recycled.
|
||||
* \param lazy The flag if the sequence should be removed lazily or intermediary.
|
||||
* \throw Error if the given sequence id is not valid.
|
||||
*/
|
||||
void RecycleSequence(int64_t seq_id, bool lazy = true) final {
|
||||
// Since there is no prefix cache, this method should never be called.
|
||||
LOG(FATAL) << "Unreachable code.";
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Try to remove recycling sequence to free up memory. It will remove the oldest
|
||||
recycling sequence.
|
||||
* \return Always return false as no sequence stored.
|
||||
*/
|
||||
bool TryFreeMemory() final {
|
||||
// Since there is no prefix cache, always return false.
|
||||
return false;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Check if a sequence exists.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \return Always return false as no sequence stored.
|
||||
*/
|
||||
bool HasSequence(int64_t seq_id) final {
|
||||
// Since there is no prefix cache, always return false.
|
||||
return false;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Reset the prefix cache to initial status. Do nothing and return.
|
||||
*/
|
||||
void Reset() final {}
|
||||
|
||||
PrefixCacheMode Mode() final { return PrefixCacheMode::kDisable; }
|
||||
};
|
||||
|
||||
PrefixCache PrefixCache::CreateRadixPrefixCache(size_t max_num_recycling_seqs,
|
||||
PrefixCacheRemoveCallback remove_callback) {
|
||||
ObjectPtr<PrefixCacheImpl> n =
|
||||
tvm::ffi::make_object<PrefixCacheImpl>(max_num_recycling_seqs, std::move(remove_callback));
|
||||
return PrefixCache(std::move(n));
|
||||
}
|
||||
|
||||
PrefixCache PrefixCache::CreateNoPrefixCache() {
|
||||
ObjectPtr<NoPrefixCache> n = tvm::ffi::make_object<NoPrefixCache>();
|
||||
return PrefixCache(std::move(n));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,158 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/prefix_cache.h
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_PREFIX_CACHE_H_
|
||||
#define MLC_LLM_SERVE_PREFIX_CACHE_H_
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
|
||||
#include <functional>
|
||||
#include <optional>
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
|
||||
#include "model.h"
|
||||
#include "radix_tree.h"
|
||||
#include "request_state.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
|
||||
/*!
|
||||
* \brief The signature of callback removing function.
|
||||
*/
|
||||
using PrefixCacheRemoveCallback = std::function<void(int64_t)>;
|
||||
|
||||
/*!
|
||||
* \brief The matched result from prefix cache. This result describes how to pre-process the new
|
||||
* sequence, to leverage the existing data in KVCache by reusing past sequences or forking from
|
||||
* other sequences.
|
||||
*/
|
||||
class PrefixCacheMatchedResult {
|
||||
public:
|
||||
/*!
|
||||
* \brief The matched and prefilled prefix offset.
|
||||
*/
|
||||
size_t prefilled_offset = 0;
|
||||
/*!
|
||||
* \brief The sequence ID to fork from.
|
||||
*/
|
||||
int64_t forked_seq_id = -1;
|
||||
/*!
|
||||
* \brief The finished sequence ID to reuse.
|
||||
*/
|
||||
int64_t reused_seq_id = -1;
|
||||
/*!
|
||||
* \brief The number of tailing tokens to be popped from the reused sequence.
|
||||
*/
|
||||
size_t reused_seq_pop_last_tokens = 0;
|
||||
};
|
||||
|
||||
class PrefixCacheObj : public Object {
|
||||
public:
|
||||
/*!
|
||||
* \brief Insert a new tokenized sequence into Prefix Cache.
|
||||
* \param seq_id The sequence ID.
|
||||
* \param tokens The tokens of tokenized sequence.
|
||||
* \param sliding_window_size The sliding window size for the sequence, -1 as sliding window
|
||||
* disabled.
|
||||
* \param attention_sink_size The attention sink size for the sequence, 0 by default.
|
||||
* \return The matched result.
|
||||
*/
|
||||
virtual PrefixCacheMatchedResult InsertSequence(int64_t seq_id, std::vector<int32_t> tokens,
|
||||
int sliding_window_size = -1,
|
||||
int attention_sink_size = 0) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Extend a sequence with new tokenized sequence suffix.
|
||||
* This extension might be cached and lazily committed later.
|
||||
* \param seq_id The sequence to be extended.
|
||||
* \param tokens The tokens of tokenized sequence suffix to extend.
|
||||
* \throw Error if the given sequence id is not valid or active.
|
||||
*/
|
||||
virtual void ExtendSequence(int64_t seq_id, const std::vector<int32_t>& tokens) = 0;
|
||||
|
||||
/*! \brief Commit the cached sequence extension from "ExtendSequence". */
|
||||
virtual void CommitSequenceExtention() = 0;
|
||||
|
||||
/*!
|
||||
* \brief Roll back a sequence by number of tokens.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \param num_tokens The number of tokens to be rolled back.
|
||||
* \throw Error if the given sequence id is not valid or active.
|
||||
*/
|
||||
virtual void RollBackSequence(int64_t seq_id, size_t num_tokens) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Recycle a sequence. The recycled sequence will not be removed immediately, as long as
|
||||
* memory is sufficient and the number of sequence in prefix cache belows the maximum number of
|
||||
* sequence. And it will be reused again in the future request.
|
||||
* \param seq_id The sequence to be recycled.
|
||||
* \param lazy The flag if the sequence should be removed lazily or intermediary.
|
||||
* \throw Error if the given sequence id is not valid.
|
||||
*/
|
||||
virtual void RecycleSequence(int64_t seq_id, bool lazy = true) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Try to remove recycling sequence to free up memory. It will remove the oldest recycling
|
||||
sequence.
|
||||
* \return The flag if there is a sequence removed. In other word, return true when memory is
|
||||
freed successfully.
|
||||
* \throw Error if the given sequence id is not valid.
|
||||
*/
|
||||
virtual bool TryFreeMemory() = 0;
|
||||
|
||||
/*!
|
||||
* \brief Check if a sequence exists.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \return The sequence existence.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
virtual bool HasSequence(int64_t seq_id) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Reset the prefix cache to initial status.
|
||||
*/
|
||||
virtual void Reset() = 0;
|
||||
|
||||
/*! \brief Return the prefix cache mode. */
|
||||
virtual PrefixCacheMode Mode() = 0;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<PrefixCacheObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.PrefixCache", PrefixCacheObj, Object);
|
||||
};
|
||||
|
||||
class PrefixCache : public ObjectRef {
|
||||
public:
|
||||
/*!
|
||||
* \brief Initialization of prefix cache.
|
||||
* \param max_recycling_seqs The maximum number of recycling sequences in prefix cache.
|
||||
* \param remove_callback The optional callback function to call when removing a sequence.
|
||||
*/
|
||||
static PrefixCache CreateRadixPrefixCache(size_t max_recycling_seqs,
|
||||
PrefixCacheRemoveCallback remove_callback = nullptr);
|
||||
/*!
|
||||
* \brief Initialization of no prefix cache.
|
||||
*/
|
||||
static PrefixCache CreateNoPrefixCache();
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrefixCache, ObjectRef, PrefixCacheObj);
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_PREFIX_CACHE_H_
|
||||
@@ -0,0 +1,845 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/radix_tree.cc
|
||||
*/
|
||||
#include "radix_tree.h"
|
||||
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/runtime/logging.h>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() { PagedRadixTreeObj::RegisterReflection(); }
|
||||
|
||||
/*!
|
||||
* \brief The sequence ID linked list structure in paged radix tree node.
|
||||
*/
|
||||
struct SequenceIDNode {
|
||||
/*! \brief The stored sequence ID. */
|
||||
int64_t id = 0;
|
||||
/*! \brief The pointer to the next sequence ID. */
|
||||
SequenceIDNode* next = nullptr;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief The sequence ID node pool.
|
||||
*
|
||||
* The sequence ID node pool allocates a block of sequence ID nodes when pool is full,
|
||||
* and frees all when destruction, to avoid frequent memory operation.
|
||||
*/
|
||||
class SequenceIDNodePool {
|
||||
public:
|
||||
/*! \brief The constructor of sequence ID node pool, allocating a new sequence ID node block. */
|
||||
SequenceIDNodePool() {
|
||||
NewNodeBlock_();
|
||||
used_nodes_.clear();
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get a sequence ID node from pool, and assign the fields.
|
||||
* If there is no available node, it will allocate a new sequence ID node block.
|
||||
* \param seq_id The assigned sequence ID of allocated sequence ID node.
|
||||
* \param node The next sequence ID node pointer of allocated sequence ID node.
|
||||
* \return The allocated radix page.
|
||||
*/
|
||||
SequenceIDNode* Allocate(int64_t seq_id, SequenceIDNode* next) {
|
||||
if (free_node_indices_.empty()) {
|
||||
NewNodeBlock_();
|
||||
TVM_FFI_ICHECK(!free_node_indices_.empty());
|
||||
}
|
||||
size_t id = free_node_indices_.back();
|
||||
free_node_indices_.pop_back();
|
||||
SequenceIDNode* node = nodes_[id];
|
||||
used_nodes_[node] = id;
|
||||
node->id = seq_id;
|
||||
node->next = next;
|
||||
return node;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Free a sequence ID node to pool.
|
||||
* \param node The sequence ID node to free.
|
||||
*/
|
||||
void Free(SequenceIDNode* node) {
|
||||
TVM_FFI_ICHECK(used_nodes_.find(node) != used_nodes_.end());
|
||||
free_node_indices_.push_back(used_nodes_[node]);
|
||||
used_nodes_.erase(node);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Reset the sequence ID node pool to initial status.
|
||||
*/
|
||||
void Reset() {
|
||||
used_nodes_.clear();
|
||||
free_node_indices_.reserve(nodes_.size());
|
||||
for (size_t i = 0; i < nodes_.size(); ++i) {
|
||||
nodes_[i]->id = 0;
|
||||
nodes_[i]->next = nullptr;
|
||||
free_node_indices_[i] = i;
|
||||
}
|
||||
}
|
||||
|
||||
/*! \brief The destructor of sequence ID node pool, freeing memory for each node. */
|
||||
~SequenceIDNodePool() {
|
||||
for (SequenceIDNode* node_block : node_blocks_) {
|
||||
delete[] node_block;
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
/*! \brief The size of each node pool block. */
|
||||
static constexpr size_t kNodeBlockSize_ = 64;
|
||||
/*! \brief The raw sequence ID node block pool, each element is a sequence ID node array. */
|
||||
std::vector<SequenceIDNode*> node_blocks_;
|
||||
/*! \brief The sequence ID node pool, each element is a sequence ID node pointer. */
|
||||
std::vector<SequenceIDNode*> nodes_;
|
||||
/*! \brief The indices of free sequence ID node in node pool. */
|
||||
std::vector<size_t> free_node_indices_;
|
||||
/*! \brief The map from used paged sequence ID node to its index in node pool. */
|
||||
std::unordered_map<SequenceIDNode*, size_t> used_nodes_;
|
||||
|
||||
/*! \brief Allocate a new node pool block. */
|
||||
void NewNodeBlock_() {
|
||||
size_t node_id_offset = node_blocks_.size() * kNodeBlockSize_;
|
||||
node_blocks_.push_back(new SequenceIDNode[kNodeBlockSize_]);
|
||||
nodes_.reserve(nodes_.size() + kNodeBlockSize_);
|
||||
free_node_indices_.reserve(free_node_indices_.size() + kNodeBlockSize_);
|
||||
for (size_t i = 0; i < kNodeBlockSize_; ++i) {
|
||||
nodes_.push_back(&node_blocks_.back()[i]);
|
||||
free_node_indices_.push_back(i + node_id_offset);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief The paged radix tree node data structure.
|
||||
*
|
||||
* The paged radix tree node is similar to original radix tree node, but with the limited length for
|
||||
* prefix in page, so that the memory usage in each page is the same and is fixed once allocated.
|
||||
* Since the page only consists of pointers and int tokens, the page memory layout is int array
|
||||
* indeed. The lower offset is the pointers and page information, while the higher offset is the
|
||||
* stored prefix tokens.
|
||||
*
|
||||
* And since the vocabulary size may be very large, the paged Radix tree is represented
|
||||
* as left-child, right-sibling binary tree.
|
||||
*
|
||||
* Also, due to possible pop/push front/back tokens in page, the page is designed as circular
|
||||
* buffer, to make full use of each page.
|
||||
*
|
||||
* Each page records the sequence exactly ends with the prefix tokens stored in page. In other word,
|
||||
* all sequences locate in the boundary of each page, or the end of each page.
|
||||
*/
|
||||
struct RadixPage {
|
||||
/*! \brief The parent page. */
|
||||
RadixPage* parent;
|
||||
/*! \brief The first child page. */
|
||||
RadixPage* first_child;
|
||||
/*! \brief The sibling page sharing the same parent page. */
|
||||
RadixPage* next_sibling;
|
||||
/*! \brief The head of sequence ID linked list. */
|
||||
SequenceIDNode* seq_ids;
|
||||
/*! \brief The capacity of maximum stored prefix tokens. */
|
||||
size_t capacity;
|
||||
/*! \brief The start offset of stored prefix tokens. The legal value is of [0, capacity). */
|
||||
size_t offset;
|
||||
/*! \brief The length of stored prefix tokens. The legal value is of [0, capacity). */
|
||||
size_t length;
|
||||
/*! \brief The offset of first prefix token in memory layout. */
|
||||
static constexpr int kDataOffset = (sizeof(RadixPage*) * 3 + sizeof(SequenceIDNode*) +
|
||||
sizeof(size_t) * 3 + sizeof(int32_t) - 1) /
|
||||
sizeof(int32_t);
|
||||
|
||||
/*!
|
||||
* \brief Overload operator [] to get the prefix tokens by index as simple int array.
|
||||
* \param i The prefix token index.
|
||||
* \return The value of i-th prefix token.
|
||||
*/
|
||||
int32_t& operator[](size_t i) {
|
||||
return reinterpret_cast<int32_t*>(this)[kDataOffset + (i + offset) % capacity];
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Extend or push back a suffix tokens in page.
|
||||
* \param suffix The suffix tokens array.
|
||||
* \param suffix_length The suffix length to extend.
|
||||
* \throw Error if suffix length is larger than current vacant space.
|
||||
*/
|
||||
void Extend(const int32_t* suffix, size_t suffix_length) {
|
||||
TVM_FFI_ICHECK_LE(suffix_length + length, capacity);
|
||||
for (int i = 0; i < suffix_length; ++i) {
|
||||
(*this)[i + length] = suffix[i];
|
||||
}
|
||||
length += suffix_length;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Add a sequence ID in page.
|
||||
* \param pool The sequence ID node pool to allocate new node.
|
||||
* \param id The sequence ID to add.
|
||||
*/
|
||||
void AddSequence(SequenceIDNodePool* pool, int64_t id) { seq_ids = pool->Allocate(id, seq_ids); }
|
||||
|
||||
/*!
|
||||
* \brief Pop a sequence ID in page.
|
||||
* \param pool The sequence ID node pool to free popped node.
|
||||
* \param id The sequence ID to pop.
|
||||
* \throw Error if no such sequence ID in page.
|
||||
*/
|
||||
void PopSequence(SequenceIDNodePool* pool, int64_t id) {
|
||||
if (seq_ids->id == id) {
|
||||
// If the popped sequence ID is the first node in linked list,
|
||||
// directly skip from head and free it.
|
||||
SequenceIDNode* next = seq_ids->next;
|
||||
pool->Free(seq_ids);
|
||||
seq_ids = next;
|
||||
} else {
|
||||
// If the popped sequence ID is not the first node in linked list,
|
||||
// skip it from previous node and free it.
|
||||
SequenceIDNode* last = seq_ids;
|
||||
SequenceIDNode* cur = seq_ids->next;
|
||||
while (cur) {
|
||||
if (cur->id == id) {
|
||||
last->next = cur->next;
|
||||
pool->Free(cur);
|
||||
return;
|
||||
}
|
||||
last = cur;
|
||||
cur = cur->next;
|
||||
}
|
||||
LOG(FATAL) << "Sequence ID = " << id << " not found.";
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get all sequence ID in page.
|
||||
* \return The std::vector of sequence ID in page.
|
||||
*/
|
||||
std::vector<int64_t> GetLocalSequence() {
|
||||
std::vector<int64_t> output;
|
||||
for (SequenceIDNode* node = seq_ids; node; node = node->next) {
|
||||
output.push_back(node->id);
|
||||
}
|
||||
return output;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get any sequence ID in current page or child pages.
|
||||
* Since there is always a sequence in leaf pages, it only check first child if no sequence ID in
|
||||
* current page.
|
||||
* \return The any sequence ID in current page or child pages.
|
||||
*/
|
||||
int32_t FindAnyChildSequence() {
|
||||
if (seq_ids) return seq_ids->id;
|
||||
return first_child->FindAnyChildSequence();
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get all sequence ID in current page and child pages, using Iterate method with lambda
|
||||
* expression as callback to avoid frequently memory allocation of std::vector.
|
||||
* \return The std::vector of all sequence ID in current page and child pages.
|
||||
*/
|
||||
std::vector<int64_t> FindAllChildSequence() {
|
||||
std::vector<int64_t> output = GetLocalSequence();
|
||||
if (first_child) {
|
||||
first_child->Iterate([&output](const RadixPage* page) {
|
||||
for (SequenceIDNode* node = page->seq_ids; node; node = node->next) {
|
||||
output.push_back(node->id);
|
||||
}
|
||||
});
|
||||
}
|
||||
return output;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief The iteration method for tree or sub-tree traverse.
|
||||
* \param f The callback function to invoke at each radix page visited.
|
||||
*/
|
||||
template <class CallbackFunc>
|
||||
void Iterate(CallbackFunc f) {
|
||||
f(this);
|
||||
if (next_sibling) next_sibling->Iterate(f);
|
||||
if (first_child) first_child->Iterate(f);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get the last sibling of current page.
|
||||
* \return The page whose next_sibling is current page, or nullptr if current is the first_child
|
||||
* of its parent page.
|
||||
*/
|
||||
RadixPage* GetLastSibling() {
|
||||
if (parent == nullptr) return nullptr;
|
||||
if (parent->first_child == this) return nullptr;
|
||||
for (RadixPage* child = parent->first_child; child; child = child->next_sibling) {
|
||||
if (child->next_sibling == this) return child;
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Find the child indexed by first token.
|
||||
* \return The child page started with first token, or nullptr if no such child page.
|
||||
*/
|
||||
RadixPage* FindChild(int64_t first_token) {
|
||||
int32_t casted = first_token;
|
||||
// Iterate all child radix pages, as the child radix pages are stored unorderly.
|
||||
for (RadixPage* child = first_child; child; child = child->next_sibling) {
|
||||
if ((*child)[0] == casted) return child;
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
/*! \brief Insert a new child page. */
|
||||
void InsertChild(RadixPage* child) {
|
||||
child->parent = this;
|
||||
child->next_sibling = first_child;
|
||||
first_child = child;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Remove a child page.
|
||||
* \throw Error if page to be removed is not child page.
|
||||
*/
|
||||
void RemoveChild(RadixPage* child) {
|
||||
TVM_FFI_ICHECK(child->parent == this);
|
||||
if (first_child == child) {
|
||||
first_child = child->next_sibling;
|
||||
} else {
|
||||
child->GetLastSibling()->next_sibling = child->next_sibling;
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Check current page is mergable with its child page.
|
||||
* The page is mergable if and only if
|
||||
* 1. No sequence ID in current page, as sequence ID is not allowed to exist within page.
|
||||
* 2. The current page has child page.
|
||||
* 3. The current page has only one child page.
|
||||
* 4. The current page prefix and the child page prefix can be concatenated into one page.
|
||||
* \return True if current page is mergable, or false.
|
||||
*/
|
||||
bool Mergeable() {
|
||||
if (seq_ids) return false;
|
||||
if (!first_child) return false;
|
||||
if (first_child->next_sibling) return false;
|
||||
if (length + first_child->length > capacity) return false;
|
||||
return true;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Match the given prefix within page.
|
||||
* \param prefix The prefix token array.
|
||||
* \param prefix_length The length of prefix token array.
|
||||
* \return The matched prefix offset within page, or the first mismatched token position. The
|
||||
* possible return value is [0, page->length], where page->length means the page is completely the
|
||||
* prefix of given prefix.
|
||||
*/
|
||||
size_t MatchPrefix(const int32_t* prefix, size_t prefix_length) {
|
||||
size_t n = std::min(length, prefix_length);
|
||||
for (int i = 0; i < n; ++i) {
|
||||
if ((*this)[i] != prefix[i]) return i;
|
||||
}
|
||||
return n;
|
||||
}
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief The paged radix tree page pool.
|
||||
*
|
||||
* The paged radix tree page pool allocates a block of radix tree pages when pool is full,
|
||||
* and frees all when destruction, to avoid frequent memory operation.
|
||||
*/
|
||||
class RadixPagePool {
|
||||
public:
|
||||
/*! \brief The constructor of paged radix tree page pool, allocating memory for each page. */
|
||||
RadixPagePool() {
|
||||
NewPageBlock_();
|
||||
used_pages_.clear();
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get a radix page from pool.
|
||||
* If there is no available page, it will allocate a new radix page block.
|
||||
* \return The allocated radix page.
|
||||
*/
|
||||
RadixPage* Allocate() {
|
||||
if (free_page_indices_.empty()) {
|
||||
NewPageBlock_();
|
||||
TVM_FFI_ICHECK(!free_page_indices_.empty());
|
||||
}
|
||||
int id = free_page_indices_.back();
|
||||
free_page_indices_.pop_back();
|
||||
RadixPage* page = pages_[id];
|
||||
used_pages_[page] = id;
|
||||
page->parent = page->first_child = page->next_sibling = nullptr;
|
||||
page->capacity = kPageCapacity_;
|
||||
page->offset = page->length = 0;
|
||||
page->seq_ids = nullptr;
|
||||
return page;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Free a radix page to pool.
|
||||
* \param page The radix page to free.
|
||||
*/
|
||||
void Free(RadixPage* page) {
|
||||
TVM_FFI_ICHECK_EQ(page->seq_ids, nullptr);
|
||||
TVM_FFI_ICHECK(used_pages_.find(page) != used_pages_.end());
|
||||
free_page_indices_.push_back(used_pages_[page]);
|
||||
TVM_FFI_ICHECK(used_pages_.erase(page));
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get the token capacity of free pages.
|
||||
* \return The the token capacity of free pages.
|
||||
*/
|
||||
size_t FreeCapacity() { return free_page_indices_.size() * kPageCapacity_; }
|
||||
|
||||
/*!
|
||||
* \brief Reset the paged radix tree page pool to initial status.
|
||||
*/
|
||||
void Reset() {
|
||||
used_pages_.clear();
|
||||
free_page_indices_.reserve(pages_.size());
|
||||
for (int i = 0; i < pages_.size(); ++i) {
|
||||
pages_[i]->parent = pages_[i]->first_child = pages_[i]->next_sibling = nullptr;
|
||||
pages_[i]->capacity = kPageCapacity_;
|
||||
pages_[i]->offset = pages_[i]->length = 0;
|
||||
pages_[i]->seq_ids = nullptr;
|
||||
free_page_indices_[i] = i;
|
||||
}
|
||||
}
|
||||
|
||||
/*! \brief The destructor of paged radix tree page pool, freeing memory for each page. */
|
||||
~RadixPagePool() {
|
||||
for (int32_t* page_block : page_blocks_) {
|
||||
delete[] page_block;
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
/*! \brief The size of each page pool block. */
|
||||
static constexpr size_t kPageBlockSize_ = 64;
|
||||
/*! \brief The page capacity of each paged radix tree page. */
|
||||
static constexpr size_t kPageCapacity_ = 64;
|
||||
/*! \brief The page size of each paged radix tree page. */
|
||||
static constexpr size_t kPageSize_ = kPageCapacity_ + RadixPage::kDataOffset;
|
||||
/*! \brief The raw paged radix tree page block pool,
|
||||
each element is a raw paged radix tree page array. */
|
||||
std::vector<int32_t*> page_blocks_;
|
||||
/*! \brief The paged radix tree page pool,
|
||||
each element is a raw paged radix tree page pointer. */
|
||||
std::vector<RadixPage*> pages_;
|
||||
/*! \brief The indices of free paged radix page in page pool. */
|
||||
std::vector<size_t> free_page_indices_;
|
||||
/*! \brief The map from used paged radix tree page to its index in page pool. */
|
||||
std::unordered_map<RadixPage*, size_t> used_pages_;
|
||||
|
||||
/*! \brief Allocate a new page pool block. */
|
||||
void NewPageBlock_() {
|
||||
size_t page_id_offset = page_blocks_.size() * kPageBlockSize_;
|
||||
page_blocks_.push_back(new int32_t[kPageBlockSize_ * kPageSize_]);
|
||||
pages_.reserve(pages_.size() + kPageBlockSize_);
|
||||
free_page_indices_.reserve(free_page_indices_.size() + kPageBlockSize_);
|
||||
for (size_t i = 0; i < kPageBlockSize_; ++i) {
|
||||
pages_.push_back(reinterpret_cast<RadixPage*>(page_blocks_.back() + i * kPageSize_));
|
||||
free_page_indices_.push_back(i + page_id_offset);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// PagedRadixTree
|
||||
|
||||
/*!
|
||||
* \brief The paged radix tree data structure.
|
||||
*/
|
||||
class PagedRadixTreeImpl : public PagedRadixTreeObj {
|
||||
public:
|
||||
/*! \brief The map from sequence to paged radix tree node it is stored. */
|
||||
std::unordered_map<int32_t, RadixPage*> seq2page;
|
||||
/*! \brief The sequence ID node pool. */
|
||||
SequenceIDNodePool* seq_id_node_pool = nullptr;
|
||||
/*! \brief The radix page pool. */
|
||||
RadixPagePool* radix_page_pool = nullptr;
|
||||
/*! \brief The root page of paged radix tree. */
|
||||
RadixPage* root = nullptr;
|
||||
|
||||
explicit PagedRadixTreeImpl() {
|
||||
seq_id_node_pool = new SequenceIDNodePool();
|
||||
radix_page_pool = new RadixPagePool();
|
||||
|
||||
root = reinterpret_cast<RadixPage*>(new int32_t[RadixPage::kDataOffset]);
|
||||
root->parent = root->first_child = root->next_sibling = nullptr;
|
||||
root->offset = root->length = root->capacity = 0;
|
||||
root->seq_ids = nullptr;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Check if a sequence exists.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \return The sequence existence.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
bool HasSequence(int64_t seq_id) { return seq2page.find(seq_id) != seq2page.end(); }
|
||||
|
||||
/*!
|
||||
* \brief Get a sequence's all tokens.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \return The sequence tokens.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
Shape GetSequence(int64_t seq_id) {
|
||||
TVM_FFI_ICHECK(seq2page.find(seq_id) != seq2page.end());
|
||||
size_t length = GetSequenceLength(seq_id);
|
||||
std::vector<int64_t> output(length);
|
||||
size_t offset = length;
|
||||
for (RadixPage* page = seq2page[seq_id]; page; page = page->parent) {
|
||||
offset -= page->length;
|
||||
for (int i = 0; i < page->length; ++i) {
|
||||
output[offset + i] = (*page)[i];
|
||||
}
|
||||
}
|
||||
return Shape(output);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get all sequences with longest common prefix with give prefix tokens.
|
||||
* \param tokens The prefix tokens for reference.
|
||||
* \return The pair of matched prefix length and the array of matched sequences indices.
|
||||
*/
|
||||
std::pair<size_t, std::vector<int64_t>> MatchPrefix(const std::vector<int32_t>& tokens) {
|
||||
const int32_t* prefix = tokens.data();
|
||||
size_t length = tokens.size();
|
||||
auto [page, offset, in_page_offset] = MatchSequence(root, prefix, length);
|
||||
if (!offset) return std::make_pair(0, std::vector<int64_t>());
|
||||
return std::make_pair(offset, page->FindAllChildSequence());
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get a sequence's length.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \return The sequence length.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
size_t GetSequenceLength(int64_t seq_id) {
|
||||
TVM_FFI_ICHECK(seq2page.find(seq_id) != seq2page.end());
|
||||
size_t length = 0;
|
||||
for (RadixPage* page = seq2page[seq_id]; page; page = page->parent) {
|
||||
length += page->length;
|
||||
}
|
||||
return length;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Fork a sequence from parent sequence at given position.
|
||||
* \param seq_id The new sequence ID.
|
||||
* \param parent_seq_id The parent sequence ID to fork from.
|
||||
* \param forked_offset The position of parent sequence to fork at.
|
||||
* The valid value is [1, length of forked sequence]. If the position equals the length of forked
|
||||
* sequence, the new sequence will copy the entire forked sequence.
|
||||
* \throw Error if sequence ID or
|
||||
* forked postion is not valid.
|
||||
*/
|
||||
void ForkSequence(int64_t seq_id, int64_t parent_seq_id, size_t forked_offset) {
|
||||
TVM_FFI_ICHECK(seq2page.find(seq_id) == seq2page.end());
|
||||
TVM_FFI_ICHECK(seq2page.find(parent_seq_id) != seq2page.end());
|
||||
TVM_FFI_ICHECK_GT(forked_offset, 0);
|
||||
size_t length = GetSequenceLength(parent_seq_id);
|
||||
TVM_FFI_ICHECK_LE(forked_offset, length);
|
||||
for (RadixPage* page = seq2page[parent_seq_id]; page; page = page->parent) {
|
||||
if (forked_offset > length - page->length) {
|
||||
if (forked_offset < length) {
|
||||
// Split radix page if forked position is within page
|
||||
page = SplitPage(page, forked_offset + page->length - length);
|
||||
}
|
||||
page->AddSequence(seq_id_node_pool, seq_id);
|
||||
seq2page[seq_id] = page;
|
||||
return;
|
||||
}
|
||||
length -= page->length;
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Add an empty sequence at root.
|
||||
* \param seq_id The new sequence ID.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
void AddSequence(int64_t seq_id) {
|
||||
TVM_FFI_ICHECK(seq2page.find(seq_id) == seq2page.end())
|
||||
<< "Sequence ID = " << seq_id << " has been added.";
|
||||
root->AddSequence(seq_id_node_pool, seq_id);
|
||||
seq2page[seq_id] = root;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Extend a sequence with given tokens.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \param tokens The given tokens to extend.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
void ExtendSequence(int64_t seq_id, const std::vector<int32_t>& tokens) {
|
||||
TVM_FFI_ICHECK(seq2page.find(seq_id) != seq2page.end());
|
||||
const int32_t* suffix = tokens.data();
|
||||
size_t length = tokens.size();
|
||||
RadixPage* original_page = seq2page[seq_id];
|
||||
original_page->PopSequence(seq_id_node_pool, seq_id);
|
||||
auto [page, offset, in_page_offset] = MatchSequence(original_page, suffix, length);
|
||||
if (in_page_offset < page->length) {
|
||||
// Split page if extended sequence mismatches within page
|
||||
page = SplitPage(page, in_page_offset);
|
||||
}
|
||||
if (offset < length && !page->seq_ids && !page->first_child && page->capacity > page->length) {
|
||||
// Extend in the existing leaf page first if possible.
|
||||
size_t suffix_length = std::min(page->capacity - page->length, length - offset);
|
||||
page->Extend(suffix + offset, suffix_length);
|
||||
offset += suffix_length;
|
||||
}
|
||||
while (offset < length) {
|
||||
// Allocate new radix page and extend tokens
|
||||
RadixPage* new_page = radix_page_pool->Allocate();
|
||||
page->InsertChild(new_page);
|
||||
page = new_page;
|
||||
size_t suffix_length = std::min(page->capacity - page->length, length - offset);
|
||||
page->Extend(suffix + offset, suffix_length);
|
||||
offset += suffix_length;
|
||||
}
|
||||
page->AddSequence(seq_id_node_pool, seq_id);
|
||||
seq2page[seq_id] = page;
|
||||
if (original_page->Mergeable()) {
|
||||
// The original page may be mergeable, as the sequence ID changes
|
||||
MergePage(original_page);
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Roll back a sequence by number of tokens.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \param num_tokens The number of tokens to be rolled back.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
void RollBackSequence(int64_t seq_id, size_t num_tokens) {
|
||||
size_t length = GetSequenceLength(seq_id);
|
||||
TVM_FFI_ICHECK_GT(num_tokens, 0);
|
||||
TVM_FFI_ICHECK_LE(num_tokens, length);
|
||||
if (num_tokens == length) {
|
||||
// If rolling back whole sequence, just remove the sequence and add it again equivalently.
|
||||
RemoveSequence(seq_id);
|
||||
AddSequence(seq_id);
|
||||
return;
|
||||
}
|
||||
RadixPage* page = seq2page[seq_id];
|
||||
// Remove the sequence temporarily, but keeping the data and starting rolling back.
|
||||
page->PopSequence(seq_id_node_pool, seq_id);
|
||||
seq2page.erase(seq_id);
|
||||
while (page->length <= num_tokens) {
|
||||
// Roll back entire page
|
||||
num_tokens -= page->length;
|
||||
RadixPage* parent = page->parent;
|
||||
if (page->seq_ids == nullptr && page->first_child == nullptr) {
|
||||
// The leaf page is removable
|
||||
parent->RemoveChild(page);
|
||||
radix_page_pool->Free(page);
|
||||
}
|
||||
page = parent;
|
||||
}
|
||||
if (page->seq_ids == nullptr && page->first_child == nullptr) {
|
||||
// The page is leaf page, directly roll back in page length
|
||||
page->length -= num_tokens;
|
||||
// Update the mapping from sequence to page
|
||||
page->AddSequence(seq_id_node_pool, seq_id);
|
||||
seq2page[seq_id] = page;
|
||||
return;
|
||||
}
|
||||
// Split page for rolled back sequence
|
||||
if (num_tokens) {
|
||||
page = SplitPage(page, page->length - num_tokens);
|
||||
}
|
||||
// Update the mapping from sequence to page
|
||||
page->AddSequence(seq_id_node_pool, seq_id);
|
||||
seq2page[seq_id] = page;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Remove a sequence.
|
||||
* \param seq_id The sequence ID to remove.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
void RemoveSequence(int64_t seq_id) {
|
||||
RadixPage* page = seq2page[seq_id];
|
||||
page->PopSequence(seq_id_node_pool, seq_id);
|
||||
seq2page.erase(seq_id);
|
||||
while (page->parent && !page->seq_ids && !page->first_child) {
|
||||
RadixPage* parent = page->parent;
|
||||
parent->RemoveChild(page);
|
||||
radix_page_pool->Free(page);
|
||||
page = parent;
|
||||
}
|
||||
if (page && page->Mergeable()) {
|
||||
// The remaining page may be mergeable, as the sequence ID changes
|
||||
MergePage(page);
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get the remaining token capacity of the paged radix tree.
|
||||
* \return The the remaining token capacity of the paged radix tree.
|
||||
*/
|
||||
size_t FreeCapacity() { return radix_page_pool->FreeCapacity(); }
|
||||
|
||||
void Reset() {
|
||||
radix_page_pool->Reset();
|
||||
seq_id_node_pool->Reset();
|
||||
seq2page.clear();
|
||||
root->parent = root->first_child = root->next_sibling = nullptr;
|
||||
root->offset = root->length = root->capacity = 0;
|
||||
root->seq_ids = nullptr;
|
||||
}
|
||||
|
||||
/*! \brief The destructor to free root page. */
|
||||
~PagedRadixTreeImpl() {
|
||||
delete[] reinterpret_cast<int32_t*>(root);
|
||||
delete seq_id_node_pool;
|
||||
delete radix_page_pool;
|
||||
}
|
||||
|
||||
private:
|
||||
/*!
|
||||
* \brief Merge a radix tree page with its child radix tree page, to save radix tree page.
|
||||
* e.g. MergePage([1, 2, _, _, _] -> [3, 4, 5, _, _]) = [1, 2, 3, 4, 5].
|
||||
* And the page to be merged should be page->Mergeable().
|
||||
* \param page The parent radix tree page.
|
||||
*/
|
||||
void MergePage(RadixPage* page) {
|
||||
TVM_FFI_ICHECK(page->Mergeable());
|
||||
RadixPage* child = page->first_child;
|
||||
for (int i = 0; i < child->length; ++i) {
|
||||
(*page)[i + page->length] = (*child)[i];
|
||||
}
|
||||
page->length += child->length;
|
||||
page->first_child = child->first_child;
|
||||
for (RadixPage* p = child->first_child; p; p = p->next_sibling) {
|
||||
p->parent = page;
|
||||
}
|
||||
page->seq_ids = child->seq_ids;
|
||||
std::vector<int64_t> seq_ids = page->GetLocalSequence();
|
||||
for (int64_t id : seq_ids) seq2page[id] = page;
|
||||
child->seq_ids = nullptr;
|
||||
radix_page_pool->Free(child);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Split a radix tree page at given position, to accept new sequence.
|
||||
* e.g. SplitPage([1, 2, 3, 4, 5], 2) = [1, 2, _, _, _] -> [3, 4, 5, _, _].
|
||||
* \param page The radix tree page to split.
|
||||
* \param offset The position to split the radix tree page.
|
||||
* \return The splitted radix tree page. It can be different from the input radix tree page, as
|
||||
* there may be implicit radix tree page merge.
|
||||
*/
|
||||
RadixPage* SplitPage(RadixPage* page, size_t offset) {
|
||||
TVM_FFI_ICHECK_LT(offset, page->length);
|
||||
RadixPage* child = radix_page_pool->Allocate();
|
||||
child->parent = page;
|
||||
child->first_child = page->first_child;
|
||||
for (RadixPage* p = page->first_child; p; p = p->next_sibling) {
|
||||
p->parent = child;
|
||||
}
|
||||
page->first_child = child;
|
||||
for (int i = offset; i < page->length; ++i) {
|
||||
(*child)[i - offset] = (*page)[i];
|
||||
}
|
||||
child->length = page->length - offset;
|
||||
page->length = offset;
|
||||
child->seq_ids = page->seq_ids;
|
||||
std::vector<int64_t> seq_ids = page->GetLocalSequence();
|
||||
for (int64_t id : seq_ids) seq2page[id] = child;
|
||||
page->seq_ids = nullptr;
|
||||
if (child->Mergeable()) {
|
||||
// The child page may be mergeable
|
||||
MergePage(child);
|
||||
}
|
||||
if (page->parent && page->parent->Mergeable()) {
|
||||
// The parent page may be mergeable
|
||||
page = page->parent;
|
||||
MergePage(page);
|
||||
}
|
||||
return page;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Match with given token from a radix tree page, stopping at first mismatch.
|
||||
* \param page The radix tree page to start matching.
|
||||
* \param tokens The given tokens to match.
|
||||
* \param length The length of given tokens.
|
||||
*/
|
||||
std::tuple<RadixPage*, size_t, size_t> MatchSequence(RadixPage* page, const int32_t* tokens,
|
||||
size_t length) {
|
||||
size_t offset = 0;
|
||||
while (offset < length) {
|
||||
if (RadixPage* child = page->FindChild(tokens[offset])) {
|
||||
// If child page starts with offset-th token, common prefix at least ends with child page
|
||||
size_t matched_offset = child->MatchPrefix(tokens + offset, length - offset);
|
||||
offset += matched_offset;
|
||||
if (matched_offset < child->length) {
|
||||
// Common prefix ends within child page
|
||||
return std::make_tuple(child, offset, matched_offset);
|
||||
}
|
||||
page = child;
|
||||
} else {
|
||||
// No child page starts with offset-th token, common prefix ends with current page
|
||||
return std::make_tuple(page, offset, page->length);
|
||||
}
|
||||
}
|
||||
return std::make_tuple(page, length, page->length);
|
||||
}
|
||||
};
|
||||
|
||||
PagedRadixTree PagedRadixTree::Create() {
|
||||
return PagedRadixTree(tvm::ffi::make_object<PagedRadixTreeImpl>());
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef()
|
||||
.def("mlc.serve.PagedRadixTree", []() { return PagedRadixTree::Create(); })
|
||||
.def("mlc.serve.PagedRadixTreeMatchPrefix",
|
||||
[](PagedRadixTree paged_radix_tree, Shape tokens) {
|
||||
std::vector<int32_t> token_ids{tokens.begin(), tokens.end()};
|
||||
auto [offset, seq_ids] = paged_radix_tree->MatchPrefix(token_ids);
|
||||
seq_ids.insert(seq_ids.begin(), offset);
|
||||
return Shape(seq_ids);
|
||||
})
|
||||
.def("mlc.serve.PagedRadixTreeExtendSequence",
|
||||
[](PagedRadixTree paged_radix_tree, int64_t seq_id, Shape tokens) {
|
||||
std::vector<int32_t> token_ids{tokens.begin(), tokens.end()};
|
||||
paged_radix_tree->ExtendSequence(seq_id, std::move(token_ids));
|
||||
})
|
||||
.def("mlc.serve.PagedRadixTreeRollBackSequence",
|
||||
[](PagedRadixTree paged_radix_tree, int64_t seq_id, int64_t num_tokens) {
|
||||
paged_radix_tree->RollBackSequence(seq_id, num_tokens);
|
||||
})
|
||||
.def("mlc.serve.PagedRadixTreeForkSequence",
|
||||
[](PagedRadixTree paged_radix_tree, int64_t seq_id, int64_t parent_seq_id,
|
||||
uint64_t forked_offset) {
|
||||
paged_radix_tree->ForkSequence(seq_id, parent_seq_id, forked_offset);
|
||||
})
|
||||
.def_method("mlc.serve.PagedRadixTreeHasSequence", &PagedRadixTreeObj::HasSequence)
|
||||
.def_method("mlc.serve.PagedRadixTreeAddSequence", &PagedRadixTreeObj::AddSequence)
|
||||
.def_method("mlc.serve.PagedRadixTreeRemoveSequence", &PagedRadixTreeObj::RemoveSequence)
|
||||
.def_method("mlc.serve.PagedRadixTreeGetSequence", &PagedRadixTreeObj::GetSequence)
|
||||
.def("mlc.serve.PagedRadixTreeGetSequenceLength",
|
||||
[](PagedRadixTree paged_radix_tree, int64_t seq_id) {
|
||||
return static_cast<int64_t>(paged_radix_tree->GetSequenceLength(seq_id));
|
||||
})
|
||||
.def("mlc.serve.PagedRadixTreeFreeCapacity", [](PagedRadixTree paged_radix_tree) {
|
||||
return static_cast<int64_t>(paged_radix_tree->FreeCapacity());
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,134 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/radix_tree.h
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_RADIX_TREE_H_
|
||||
#define MLC_LLM_SERVE_RADIX_TREE_H_
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
using tvm::ffi::Shape;
|
||||
|
||||
/*!
|
||||
* \brief The paged radix tree data structure.
|
||||
*/
|
||||
class PagedRadixTreeObj : public Object {
|
||||
public:
|
||||
/*!
|
||||
* \brief Check if a sequence exists.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \return The sequence existence.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
virtual bool HasSequence(int64_t seq_id) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Get a sequence's all tokens.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \return The sequence tokens.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
virtual Shape GetSequence(int64_t seq_id) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Get all sequences with longest common prefix with give prefix tokens.
|
||||
* \param tokens The prefix tokens for reference.
|
||||
* \return The pair of matched prefix length and the array of matched sequences indices.
|
||||
*/
|
||||
virtual std::pair<size_t, std::vector<int64_t>> MatchPrefix(
|
||||
const std::vector<int32_t>& tokens) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Get a sequence's length.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \return The sequence length.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
virtual size_t GetSequenceLength(int64_t seq_id) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Fork a sequence from parent sequence at given position.
|
||||
* \param seq_id The new sequence ID.
|
||||
* \param parent_seq_id The parent sequence ID to fork from.
|
||||
* \param forked_offset The position of parent sequence to fork at.
|
||||
* The valid value is [1, length of forked sequence]. If the position equals the length of forked
|
||||
* sequence, the new sequence will copy the entire forked sequence.
|
||||
* \throw Error if sequence ID or
|
||||
* forked postion is not valid.
|
||||
*/
|
||||
virtual void ForkSequence(int64_t seq_id, int64_t parent_seq_id, size_t forked_offset) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Add an empty sequence at root.
|
||||
* \param seq_id The new sequence ID.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
virtual void AddSequence(int64_t seq_id) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Extend a sequence with given tokens.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \param tokens The given tokens to extend.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
virtual void ExtendSequence(int64_t seq_id, const std::vector<int32_t>& tokens) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Roll back a sequence by number of tokens.
|
||||
* \param seq_id The sequence ID for index.
|
||||
* \param num_tokens The number of tokens to be rolled back.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
virtual void RollBackSequence(int64_t seq_id, size_t num_tokens) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Remove a sequence.
|
||||
* \param seq_id The sequence ID to remove.
|
||||
* \throw Error if sequence ID is not valid.
|
||||
*/
|
||||
virtual void RemoveSequence(int64_t seq_id) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Get the remaining token capacity of the paged radix tree.
|
||||
* \return The the remaining token capacity of the paged radix tree.
|
||||
*/
|
||||
virtual size_t FreeCapacity() = 0;
|
||||
|
||||
/*!
|
||||
* \brief Reset the paged radix tree to initial status.
|
||||
*/
|
||||
virtual void Reset() = 0;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<PagedRadixTreeObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.PagedRadixTree", PagedRadixTreeObj, Object);
|
||||
};
|
||||
|
||||
class PagedRadixTree : public ObjectRef {
|
||||
public:
|
||||
/*!
|
||||
* \brief Construct a paged radix tree.
|
||||
* \return The constructed paged radix tree. */
|
||||
static PagedRadixTree Create();
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PagedRadixTree, ObjectRef, PagedRadixTreeObj);
|
||||
};
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_RADIX_TREE_H_
|
||||
@@ -0,0 +1,82 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/request.cc
|
||||
*/
|
||||
|
||||
#include "request.h"
|
||||
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
|
||||
#include "data.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
/****************** Request ******************/
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() { RequestNode::RegisterReflection(); }
|
||||
|
||||
Request::Request(String id, Array<Data> inputs, GenerationConfig generation_cfg) {
|
||||
if (generation_cfg->debug_config.special_request == SpecialRequestKind::kNone) {
|
||||
TVM_FFI_ICHECK(!inputs.empty()) << "No input data is given.";
|
||||
}
|
||||
// Compute the total input length, or fall back to "-1" which means
|
||||
// unknown due to the existence of untokenized data.
|
||||
int prompt_tokens = 0;
|
||||
for (Data input : inputs) {
|
||||
if (const auto* token_data = input.as<TokenDataNode>()) {
|
||||
prompt_tokens += token_data->token_ids.size();
|
||||
} else if (const auto* image_data = input.as<ImageDataNode>()) {
|
||||
prompt_tokens += image_data->GetLength();
|
||||
} else {
|
||||
prompt_tokens = -1;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
ObjectPtr<RequestNode> n = tvm::ffi::make_object<RequestNode>();
|
||||
n->id = std::move(id);
|
||||
n->inputs = std::move(inputs);
|
||||
n->prompt_tokens = prompt_tokens;
|
||||
n->generation_cfg = std::move(generation_cfg);
|
||||
data_ = std::move(n);
|
||||
}
|
||||
|
||||
Request Request::FromUntokenized(const Request& request, const Tokenizer& tokenizer) {
|
||||
bool has_untokenized_input = false;
|
||||
Array<Data> inputs;
|
||||
inputs.reserve(request->inputs.size());
|
||||
// Tokenize all text inputs.
|
||||
for (Data input : request->inputs) {
|
||||
if (const auto* text_data = input.as<TextDataNode>()) {
|
||||
has_untokenized_input = true;
|
||||
std::vector<int> token_ids = tokenizer->Encode(text_data->text);
|
||||
inputs.push_back(TokenData(token_ids));
|
||||
} else {
|
||||
inputs.push_back(input);
|
||||
}
|
||||
}
|
||||
|
||||
// If there is no untokenized input, we don't need to create a new request.
|
||||
if (!has_untokenized_input) {
|
||||
TVM_FFI_ICHECK_NE(request->prompt_tokens, -1);
|
||||
return request;
|
||||
} else {
|
||||
return Request(request->id, std::move(inputs), request->generation_cfg);
|
||||
}
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef()
|
||||
.def("mlc.serve.RequestGetInputs", [](Request request) { return request->inputs; })
|
||||
.def("mlc.serve.RequestGetGenerationConfigJSON", [](Request request) {
|
||||
return tvm::ffi::json::Stringify(request->generation_cfg->AsJSON());
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,92 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/request.h
|
||||
* \brief Implementation of llm chat.
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_REQUEST_H_
|
||||
#define MLC_LLM_SERVE_REQUEST_H_
|
||||
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
|
||||
#include "../tokenizers/tokenizers.h"
|
||||
#include "config.h"
|
||||
#include "data.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
|
||||
/****************** Request ******************/
|
||||
|
||||
/*!
|
||||
* \brief The user submitted text-generation request, which contains
|
||||
* a unique request id, a list of multi-modal inputs, a set of
|
||||
* generation configuration parameters.
|
||||
* \note Request is immutable and can be re-dispatched to another
|
||||
* node and restart the request handling on the new one.
|
||||
*/
|
||||
class RequestNode : public Object {
|
||||
public:
|
||||
/*!
|
||||
* \brief The unique identifier of the request.
|
||||
* Different requests should have different ids.
|
||||
*/
|
||||
String id;
|
||||
/*!
|
||||
* \brief The user inputs of a request. Input may have multi-modality.
|
||||
* \sa data.h
|
||||
*/
|
||||
Array<Data> inputs;
|
||||
/*!
|
||||
* \brief The equivalent input sequence length of the request.
|
||||
* "-1" means the input length is unknown due to the existence
|
||||
* of untokenized text data.
|
||||
*/
|
||||
int prompt_tokens = -1;
|
||||
/*!
|
||||
* \brief The sampling configuration which may contain temperature,
|
||||
* top_p, repetition_penalty, max_gen_len, etc.
|
||||
*/
|
||||
GenerationConfig generation_cfg;
|
||||
/*! \brief Backward reference to the request state. */
|
||||
Object* rstate = nullptr;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<RequestNode>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.Request", RequestNode, Object);
|
||||
};
|
||||
|
||||
class Request : public ObjectRef {
|
||||
public:
|
||||
explicit Request(String id, Array<Data> inputs, GenerationConfig generation_cfg);
|
||||
|
||||
/*!
|
||||
* \brief Return a request object with all text data tokenized,
|
||||
* and the request ID kept the same as the input one.
|
||||
* \param request The request to be tokenized.
|
||||
* \param tokenizer The tokenizer to tokenize the input data of the given request.
|
||||
* \return The request object whose data are tokenized.
|
||||
*/
|
||||
static Request FromUntokenized(const Request& request, const Tokenizer& tokenizer);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Request, ObjectRef, RequestNode);
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_REQUEST_H_
|
||||
@@ -0,0 +1,298 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/request_state.cc
|
||||
*/
|
||||
|
||||
#include "request_state.h"
|
||||
|
||||
#include <unordered_set>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
RequestModelStateNode::RegisterReflection();
|
||||
RequestStateEntryNode::RegisterReflection();
|
||||
RequestStateNode::RegisterReflection();
|
||||
}
|
||||
|
||||
/****************** RequestModelState ******************/
|
||||
|
||||
RequestModelState::RequestModelState(
|
||||
Request request, int model_id, int64_t internal_id, Array<Data> inputs,
|
||||
const std::optional<xgrammar::CompiledGrammar>& compiled_grammar) {
|
||||
ObjectPtr<RequestModelStateNode> n = tvm::ffi::make_object<RequestModelStateNode>();
|
||||
n->model_id = model_id;
|
||||
n->internal_id = internal_id;
|
||||
n->inputs = std::move(inputs);
|
||||
|
||||
if (compiled_grammar.has_value()) {
|
||||
// TODO(yixin): set rollback limit to a configurable value.
|
||||
n->grammar_matcher =
|
||||
xgrammar::GrammarMatcher(compiled_grammar.value(), std::nullopt, false, std::nullopt, 10);
|
||||
}
|
||||
|
||||
n->request = std::move(request);
|
||||
data_ = std::move(n);
|
||||
}
|
||||
|
||||
int RequestModelStateNode::GetInputLength() const {
|
||||
int total_length = 0;
|
||||
for (Data input : inputs) {
|
||||
total_length += input->GetLength();
|
||||
}
|
||||
return total_length;
|
||||
}
|
||||
|
||||
bool RequestModelStateNode::RequireNextTokenBitmask() { return grammar_matcher.has_value(); }
|
||||
|
||||
void RequestModelStateNode::GetNextTokenBitmask(DLTensor* bitmask) {
|
||||
TVM_FFI_ICHECK(grammar_matcher.has_value());
|
||||
|
||||
grammar_matcher->GetNextTokenBitmask(bitmask);
|
||||
}
|
||||
|
||||
void RequestModelStateNode::CommitToken(SampleResult sampled_token) {
|
||||
committed_tokens.push_back(std::move(sampled_token));
|
||||
appeared_token_ids[sampled_token.GetTokenId()] += 1;
|
||||
// There will be one more token that will be processed in the next decoding.
|
||||
++num_tokens_for_next_decode;
|
||||
|
||||
// Update the grammar matcher state if it exists.
|
||||
if (grammar_matcher) {
|
||||
bool accepted = grammar_matcher->AcceptToken(sampled_token.GetTokenId());
|
||||
TVM_FFI_ICHECK(accepted) << "Token id " << sampled_token.GetTokenId()
|
||||
<< " is not accepted by the grammar state matcher.";
|
||||
}
|
||||
}
|
||||
|
||||
void RequestModelStateNode::RollbackTokens(int count) {
|
||||
TVM_FFI_ICHECK(count <= static_cast<int>(committed_tokens.size()));
|
||||
for (int i = 0; i < count; ++i) {
|
||||
auto it = appeared_token_ids.find(committed_tokens.back().GetTokenId());
|
||||
TVM_FFI_ICHECK(it != appeared_token_ids.end());
|
||||
if (--it->second == 0) {
|
||||
appeared_token_ids.erase(it);
|
||||
}
|
||||
committed_tokens.pop_back();
|
||||
if (grammar_matcher) {
|
||||
grammar_matcher->Rollback(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void RequestModelStateNode::AddDraftToken(SampleResult sampled_token, int draft_token_slot,
|
||||
int64_t parent_idx) {
|
||||
draft_output_tokens.push_back(std::move(sampled_token));
|
||||
draft_token_slots.push_back(draft_token_slot);
|
||||
draft_token_parent_idx.push_back(parent_idx);
|
||||
draft_token_first_child_idx.push_back(-1);
|
||||
if (parent_idx != -1) {
|
||||
if (draft_token_first_child_idx[parent_idx] == -1) {
|
||||
draft_token_first_child_idx[parent_idx] = static_cast<int>(draft_output_tokens.size()) - 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void RequestModelStateNode::RemoveAllDraftTokens(std::vector<int>* removed_draft_token_slots) {
|
||||
if (removed_draft_token_slots != nullptr) {
|
||||
std::unordered_set<int> dedup;
|
||||
removed_draft_token_slots->clear();
|
||||
for (auto slot : draft_token_slots) {
|
||||
bool inserted = dedup.insert(slot).second;
|
||||
if (inserted) {
|
||||
removed_draft_token_slots->push_back(slot);
|
||||
}
|
||||
}
|
||||
}
|
||||
draft_token_slots.clear();
|
||||
draft_token_parent_idx.clear();
|
||||
draft_token_first_child_idx.clear();
|
||||
draft_output_tokens.clear();
|
||||
}
|
||||
|
||||
/****************** RequestActionPostProcWorkspace ******************/
|
||||
|
||||
RequestStreamOutput RequestActionPostProcWorkspace::GetStreamOutput() {
|
||||
for (const RequestStreamOutput& stream_output : stream_outputs) {
|
||||
if (stream_output->unpacked) {
|
||||
return stream_output;
|
||||
}
|
||||
}
|
||||
|
||||
TVM_FFI_ICHECK(!stream_outputs.empty());
|
||||
int num_response = stream_outputs[0]->group_delta_token_ids.size();
|
||||
std::vector<std::vector<int64_t>> group_delta_token_ids;
|
||||
std::vector<std::vector<String>> group_delta_logprob_json_strs;
|
||||
std::vector<Optional<String>> group_finish_reason;
|
||||
std::vector<String> group_extra_prefix_string;
|
||||
group_delta_token_ids.resize(num_response);
|
||||
group_finish_reason.resize(num_response);
|
||||
group_extra_prefix_string.resize(num_response);
|
||||
if (stream_outputs[0]->group_delta_logprob_json_strs.has_value()) {
|
||||
group_delta_logprob_json_strs.resize(num_response);
|
||||
}
|
||||
RequestStreamOutput stream_output(stream_outputs[0]->request_id, std::move(group_delta_token_ids),
|
||||
stream_outputs[0]->group_delta_logprob_json_strs.has_value()
|
||||
? std::make_optional(group_delta_logprob_json_strs)
|
||||
: std::nullopt,
|
||||
std::move(group_finish_reason),
|
||||
std::move(group_extra_prefix_string));
|
||||
stream_outputs.push_back(stream_output);
|
||||
return stream_output;
|
||||
}
|
||||
|
||||
/****************** RequestStateEntry ******************/
|
||||
|
||||
RequestStateEntry::RequestStateEntry(
|
||||
Request request, int num_models, int64_t internal_id, int rng_seed,
|
||||
const std::vector<std::string>& token_table,
|
||||
const std::optional<xgrammar::CompiledGrammar>& compiled_grammar, int parent_idx) {
|
||||
ObjectPtr<RequestStateEntryNode> n = tvm::ffi::make_object<RequestStateEntryNode>();
|
||||
Array<RequestModelState> mstates;
|
||||
Array<Data> inputs;
|
||||
if (parent_idx == -1) {
|
||||
inputs = request->inputs;
|
||||
}
|
||||
mstates.reserve(num_models);
|
||||
for (int i = 0; i < num_models; ++i) {
|
||||
mstates.push_back(RequestModelState(request, i, internal_id, inputs, compiled_grammar));
|
||||
}
|
||||
n->status = RequestStateStatus::kPending;
|
||||
n->rng = RandomGenerator(rng_seed);
|
||||
n->stop_str_handler = StopStrHandler(!request->generation_cfg->debug_config.ignore_eos
|
||||
? request->generation_cfg->stop_strs
|
||||
: Array<String>(),
|
||||
token_table);
|
||||
n->request = std::move(request);
|
||||
n->parent_idx = parent_idx;
|
||||
n->mstates = std::move(mstates);
|
||||
n->next_callback_token_pos = 0;
|
||||
data_ = std::move(n);
|
||||
}
|
||||
|
||||
void RequestStateEntryNode::GetDeltaRequestReturn(const Tokenizer& tokenizer,
|
||||
int64_t max_single_sequence_length,
|
||||
RequestStreamOutput* delta_stream_output,
|
||||
int idx) {
|
||||
TVM_FFI_ICHECK_NOTNULL(delta_stream_output);
|
||||
bool needs_logprobs = (*delta_stream_output)->group_delta_logprob_json_strs.has_value();
|
||||
(*delta_stream_output)->group_delta_token_ids[idx].clear();
|
||||
if (needs_logprobs) {
|
||||
(*delta_stream_output)->group_delta_logprob_json_strs.value()[idx].clear();
|
||||
}
|
||||
(*delta_stream_output)->group_finish_reason[idx] = std::nullopt;
|
||||
(*delta_stream_output)->group_extra_prefix_string[idx] = this->extra_prefix_string;
|
||||
this->extra_prefix_string.clear();
|
||||
|
||||
const std::vector<SampleResult>& committed_tokens = this->mstates[0]->committed_tokens;
|
||||
int num_committed_tokens = committed_tokens.size();
|
||||
TVM_FFI_ICHECK_LE(this->next_callback_token_pos, num_committed_tokens);
|
||||
|
||||
// Case 1. There is no new token ids.
|
||||
if (this->next_callback_token_pos == num_committed_tokens && extra_prefix_string.empty()) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Case 2. Any of the stop strings is matched.
|
||||
TVM_FFI_ICHECK(!stop_str_handler->StopTriggered());
|
||||
while (next_callback_token_pos < num_committed_tokens) {
|
||||
stop_str_handler->Put(committed_tokens[next_callback_token_pos].GetTokenId(),
|
||||
&(*delta_stream_output)->group_delta_token_ids[idx]);
|
||||
if (needs_logprobs) {
|
||||
(*delta_stream_output)
|
||||
->group_delta_logprob_json_strs.value()[idx]
|
||||
.push_back(committed_tokens[next_callback_token_pos].GetLogProbJSON(
|
||||
tokenizer, request->generation_cfg->logprobs));
|
||||
}
|
||||
++next_callback_token_pos;
|
||||
if (stop_str_handler->StopTriggered()) {
|
||||
(*delta_stream_output)->group_finish_reason[idx] = "stop";
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Case 3. Any of the stop tokens appears in the committed tokens ===> Finished
|
||||
// `stop_token_ids` includes the stop tokens from conversation template and user-provided tokens.
|
||||
// This check will be ignored when `ignore_eos` is set for the benchmarking purpose.
|
||||
if (!request->generation_cfg->debug_config.ignore_eos) {
|
||||
for (int i = 0; i < static_cast<int>((*delta_stream_output)->group_delta_token_ids[idx].size());
|
||||
++i) {
|
||||
if (std::any_of(request->generation_cfg->stop_token_ids.begin(),
|
||||
request->generation_cfg->stop_token_ids.end(),
|
||||
[delta_stream_output, idx, i](int32_t token) {
|
||||
return token == (*delta_stream_output)->group_delta_token_ids[idx][i];
|
||||
})) {
|
||||
// Stop token matched. Erase the stop token and all tokens after it.
|
||||
(*delta_stream_output)->group_finish_reason[idx] = "stop";
|
||||
while (static_cast<int>((*delta_stream_output)->group_delta_token_ids[idx].size()) > i) {
|
||||
(*delta_stream_output)->group_delta_token_ids[idx].pop_back();
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Case 4. When stop token is not detected (e.g. ignore_eos is set), but the grammar state is
|
||||
// terminated, stop the generation and pop the last token (used to trigger the termination).
|
||||
if ((*delta_stream_output)->group_finish_reason[idx] != "stop" &&
|
||||
this->mstates[0]->grammar_matcher.has_value() &&
|
||||
this->mstates[0]->grammar_matcher->IsTerminated()) {
|
||||
(*delta_stream_output)->group_delta_token_ids[idx].pop_back();
|
||||
(*delta_stream_output)->group_finish_reason[idx] = "stop";
|
||||
}
|
||||
|
||||
if ((*delta_stream_output)->group_finish_reason[idx].has_value()) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Case 5. Generation reaches the specified max generation length ==> Finished
|
||||
// `max_tokens` means the generation length is limited by model capacity.
|
||||
if (request->generation_cfg->max_tokens >= 0 &&
|
||||
num_committed_tokens >= request->generation_cfg->max_tokens) {
|
||||
stop_str_handler->Finish(&(*delta_stream_output)->group_delta_token_ids[idx]);
|
||||
(*delta_stream_output)->group_finish_reason[idx] = "length";
|
||||
return;
|
||||
}
|
||||
// Case 6. Total length of the request reaches the maximum single sequence length ==> Finished
|
||||
if (request->prompt_tokens + num_committed_tokens >= max_single_sequence_length) {
|
||||
stop_str_handler->Finish(&(*delta_stream_output)->group_delta_token_ids[idx]);
|
||||
(*delta_stream_output)->group_finish_reason[idx] = "length";
|
||||
}
|
||||
}
|
||||
|
||||
/****************** RequestState ******************/
|
||||
|
||||
RequestState::RequestState(std::vector<RequestStateEntry> entries, int num_response,
|
||||
std::chrono::high_resolution_clock::time_point add_time_point) {
|
||||
TVM_FFI_ICHECK(!entries.empty());
|
||||
ObjectPtr<RequestStateNode> n = tvm::ffi::make_object<RequestStateNode>();
|
||||
n->entries = std::move(entries);
|
||||
n->metrics.prompt_tokens = n->entries[0]->request->prompt_tokens;
|
||||
n->metrics.add_time_point = add_time_point;
|
||||
|
||||
std::vector<std::vector<int64_t>> group_delta_token_ids;
|
||||
std::vector<std::vector<String>> group_delta_logprob_json_strs;
|
||||
std::vector<Optional<String>> group_finish_reason;
|
||||
std::vector<String> group_extra_prefix_string;
|
||||
group_delta_token_ids.resize(num_response);
|
||||
group_finish_reason.resize(num_response);
|
||||
group_extra_prefix_string.resize(num_response);
|
||||
if (n->entries[0]->request->generation_cfg->logprobs) {
|
||||
group_delta_logprob_json_strs.resize(num_response);
|
||||
}
|
||||
RequestStreamOutput stream_output(n->entries[0]->request->id, std::move(group_delta_token_ids),
|
||||
n->entries[0]->request->generation_cfg->logprobs
|
||||
? std::make_optional(group_delta_logprob_json_strs)
|
||||
: std::nullopt,
|
||||
std::move(group_finish_reason),
|
||||
std::move(group_extra_prefix_string));
|
||||
stream_output->unpacked = true;
|
||||
n->postproc_states.stream_outputs = {std::move(stream_output)};
|
||||
data_ = std::move(n);
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,318 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/request_state.h
|
||||
* \brief The data structure maintaining the generation states of user requests.
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_REQUEST_STATE_H_
|
||||
#define MLC_LLM_SERVE_REQUEST_STATE_H_
|
||||
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
#include <xgrammar/xgrammar.h>
|
||||
|
||||
#include <optional>
|
||||
|
||||
#include "../support/random.h"
|
||||
#include "../tokenizers/streamer.h"
|
||||
#include "config.h"
|
||||
#include "metrics.h"
|
||||
#include "request.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
|
||||
/*!
|
||||
* \brief The state of a request with regard to some single model.
|
||||
* \details In MLC LLM, the serving engine may leverage multiple models
|
||||
* to fulfill a user generation request (e.g., use speculation decoding).
|
||||
* For each request, we isolate its states (e.g. the generated tokens)
|
||||
* on each model. This is to say, we use RequestModelState to store
|
||||
* the state of a user request on a single model (rather than all models).
|
||||
*/
|
||||
class RequestModelStateNode : public Object {
|
||||
public:
|
||||
/*! \brief The request that this state corresponds to. */
|
||||
Request request;
|
||||
/*!
|
||||
* \brief The internal request id of this state.
|
||||
* It is the **physical index** of the request in the running request queue.
|
||||
* If the request is on hold (not in the running queue), the request id
|
||||
* should be -1.
|
||||
*/
|
||||
int64_t internal_id = -1;
|
||||
/*! \brief The corresponding model id of this state. */
|
||||
int model_id = -1;
|
||||
/*!
|
||||
* \brief The committed generated token ids and related probability info.
|
||||
* A token is "committed" means it will no longer be updated (or changed).
|
||||
*/
|
||||
std::vector<SampleResult> committed_tokens;
|
||||
/*! \brief The list of input data yet for the model to prefill. */
|
||||
Array<Data> inputs;
|
||||
/*! \brief The list of prefilled input data, used to notify prefix cache. */
|
||||
std::vector<Data> prefilled_inputs;
|
||||
/*! \brief The number of tokens already cached in prefix cache. */
|
||||
int64_t cached_committed_tokens = 0;
|
||||
/*! \brief The number of tokens that is already prefilled from the inputs. */
|
||||
int64_t num_prefilled_tokens = 0;
|
||||
/*! \brief The number of tokens that need to be processed in the next decoding. */
|
||||
int num_tokens_for_next_decode = 0;
|
||||
/*! \brief Whether retokenization is needed in the next decoding. When the jump-forward decoding
|
||||
* is enabled, retokenization is needed after every jump-forward and decoding action. */
|
||||
bool require_retokenization_in_next_decode = false;
|
||||
|
||||
// NOTE: The following fields are reserved for future speculative inference
|
||||
// settings, and are produced by the speculative small models.
|
||||
/*!
|
||||
* \brief The draft generated token ids and related probability info,
|
||||
* which are usually generated by "small" speculative models.
|
||||
* These tokens will be fed to a "large" model to determine the final
|
||||
* result of speculation.
|
||||
*/
|
||||
std::vector<SampleResult> draft_output_tokens;
|
||||
/*! \brief The storage slots for the associated states of draft tokens. */
|
||||
std::vector<int> draft_token_slots;
|
||||
/*! \brief The parent indices of the draft tokens. */
|
||||
std::vector<int64_t> draft_token_parent_idx;
|
||||
/*! \brief The first child indices of the draft tokens. */
|
||||
std::vector<int64_t> draft_token_first_child_idx;
|
||||
|
||||
/*! \brief The appeared committed and draft tokens and their occurrence times. */
|
||||
std::unordered_map<int32_t, int32_t> appeared_token_ids;
|
||||
|
||||
/*!
|
||||
* \brief The current state of the generated token matching the grammar. Used in grammar-guided
|
||||
* generation, otherwise it's std::nullopt.
|
||||
*/
|
||||
std::optional<xgrammar::GrammarMatcher> grammar_matcher;
|
||||
|
||||
/*! \brief Return the total length of the input data. */
|
||||
int GetInputLength() const;
|
||||
/*!
|
||||
* \brief Return whether the next token bitmask is required, i.e. the grammar-guided generation is
|
||||
* enabled.
|
||||
*/
|
||||
bool RequireNextTokenBitmask();
|
||||
/*!
|
||||
* \brief Find the next token bitmask and store it in the given DLTensor.
|
||||
* \param bitmask The DLTensor to store the next token bitmask. The bitmask should be a tensor
|
||||
* with dtype uint32_t and shape (ceildiv(vocab_size, 32),).
|
||||
*/
|
||||
void GetNextTokenBitmask(DLTensor* bitmask);
|
||||
/*! \brief Commit a new token into committed_tokens. Does not effect the kv cache. Update
|
||||
* appeared_token_ids and the grammar state. */
|
||||
void CommitToken(SampleResult sampled_token);
|
||||
/*! \brief Roll back the last tokens back from committed_tokens. Does not effect the kv cache.
|
||||
* Also roll back appeared_token_ids and the grammar state. */
|
||||
void RollbackTokens(int count);
|
||||
|
||||
/*! \brief Add a draft token into draft_output_tokens. Update appeared_token_ids. */
|
||||
void AddDraftToken(SampleResult sampled_token, int draft_token_slot, int64_t parent_idx);
|
||||
/*! \brief Remove all draft tokens from draft_output_tokens. Update appeared_token_ids. */
|
||||
void RemoveAllDraftTokens(std::vector<int>* removed_draft_token_slots = nullptr);
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<RequestModelStateNode>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.RequestModelState", RequestModelStateNode, Object);
|
||||
};
|
||||
|
||||
class RequestModelState : public ObjectRef {
|
||||
public:
|
||||
explicit RequestModelState(Request request, int model_id, int64_t internal_id, Array<Data> inputs,
|
||||
const std::optional<xgrammar::CompiledGrammar>& compiled_grammar);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(RequestModelState, ObjectRef, RequestModelStateNode);
|
||||
};
|
||||
|
||||
struct DeltaRequestReturn {
|
||||
std::vector<int64_t> delta_token_ids;
|
||||
std::vector<String> delta_logprob_json_strs;
|
||||
Optional<String> finish_reason;
|
||||
/*! \brief The extra string to prepend the delta output. The delta output should be
|
||||
* extra_prefix_string + detokenize(delta_token_ids). */
|
||||
String extra_prefix_string = "";
|
||||
};
|
||||
|
||||
/****************** Request States ******************/
|
||||
|
||||
/*!
|
||||
* \brief For each request, we maintain its "request state" in the
|
||||
* engine. Generally, the state of a request contains the information
|
||||
* of the request's generation at the current moment, including
|
||||
* the generated token ids, the grammar handler, etc.
|
||||
*
|
||||
* When a request has multiple parallel generations (e.g., the field
|
||||
* `n` of its generation config is more than 1), each generation will
|
||||
* have different states all the time.
|
||||
*
|
||||
* Therefore, to better support parallel generations, we denote the
|
||||
* state of a single generation as a "RequestStateEntry" instance,
|
||||
* and denote the state of a request's all generations using a vector,
|
||||
* named as a "RequestState" instance.
|
||||
*
|
||||
* A request's all state entries are organized as a tree structure
|
||||
* when there are parallel generations.
|
||||
* - the request input has the root status entry,
|
||||
* - each parallel generation is a child of the root.
|
||||
* This tree structure may be further extended to more complicated
|
||||
* cases in the future. As of now, for the case of `n > 1`, there
|
||||
* will be (n + 1) entries in total. In a "RequestState", the root
|
||||
* entry always has index 0. And we guarantee that the entry order
|
||||
* from the vector begin to the end is always a topological order
|
||||
* of the tree.
|
||||
*/
|
||||
|
||||
/*! \brief Request state status. */
|
||||
enum class RequestStateStatus : int {
|
||||
kPending = 0,
|
||||
kAlive = 1,
|
||||
kFinished = 2,
|
||||
};
|
||||
|
||||
/*! \brief The data structures for each request used in the action post-process. */
|
||||
struct RequestActionPostProcWorkspace {
|
||||
std::vector<RequestStreamOutput> stream_outputs;
|
||||
|
||||
RequestStreamOutput GetStreamOutput();
|
||||
};
|
||||
|
||||
// forward declare request state node.
|
||||
class RequestStateNode;
|
||||
|
||||
/*!
|
||||
* \brief A request's state entry. It contains the state of a single
|
||||
* generation of a request, or the state of a prompt prefix of a request.
|
||||
*/
|
||||
class RequestStateEntryNode : public Object {
|
||||
public:
|
||||
/*! \brief The status of the request state entry. */
|
||||
RequestStateStatus status;
|
||||
/*! \brief The request that this state corresponds to. */
|
||||
Request request;
|
||||
/*!
|
||||
* \brief The idx of the parent request state entry of this state.
|
||||
* Being -1 means the state has no parent and is the foremost
|
||||
* "prefix" entry or the only entry.
|
||||
*/
|
||||
int parent_idx = -1;
|
||||
/*! \brief The children indices of the request state entry. */
|
||||
std::vector<int> child_indices;
|
||||
|
||||
/*!
|
||||
* \brief The state with regard to each model.
|
||||
* \sa RequestModelState
|
||||
*/
|
||||
Array<RequestModelState> mstates;
|
||||
/*! \brief The random number generator of this request state entry. */
|
||||
RandomGenerator rng;
|
||||
/*! \brief The stop string handler of this request state entry. */
|
||||
StopStrHandler stop_str_handler;
|
||||
/*!
|
||||
* \brief The start position of the committed tokens in the
|
||||
* next request stream callback invocation.
|
||||
*/
|
||||
int next_callback_token_pos;
|
||||
|
||||
/*! \brief The extra string to prepend the output. */
|
||||
std::string extra_prefix_string;
|
||||
|
||||
std::vector<int32_t> token_ids_for_prefix_cache_update;
|
||||
|
||||
/*!
|
||||
* \brief Back reference to the request state.
|
||||
* Use ObjectRef to avoid circulate reference.
|
||||
*/
|
||||
RequestStateNode* rstate = nullptr;
|
||||
|
||||
/*!
|
||||
* \brief Get the delta token ids and the logprob JSON strings for this request to return since
|
||||
* the last time calling into this function, and return the finish reason if the request
|
||||
* generation has finished.
|
||||
* \note This function follows the destination passing style, which means it writes the
|
||||
* output into the "idx"-th slot in "delta_stream_output".
|
||||
* We adopt the destination passing style to reduce the CPU data structure allocation and
|
||||
* construction overhead.
|
||||
* \param tokenizer The tokenizer for logprob process.
|
||||
* \param max_single_sequence_length The maximum allowed single sequence length.
|
||||
* \param delta_stream_output The delta token ids to return, the logprob JSON strings
|
||||
* of each delta token id, and the optional finish reason.
|
||||
* \param idx The index denoting which slot to write results in "delta_request_return".
|
||||
*/
|
||||
void GetDeltaRequestReturn(const Tokenizer& tokenizer, int64_t max_single_sequence_length,
|
||||
RequestStreamOutput* delta_stream_output, int idx);
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<RequestStateEntryNode>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.RequestStateEntry", RequestStateEntryNode, Object);
|
||||
};
|
||||
|
||||
class RequestStateEntry : public ObjectRef {
|
||||
public:
|
||||
explicit RequestStateEntry(Request request, int num_models, int64_t internal_id, int rng_seed,
|
||||
const std::vector<std::string>& token_table,
|
||||
const std::optional<xgrammar::CompiledGrammar>& compiled_grammar,
|
||||
int parent_idx = -1);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(RequestStateEntry, ObjectRef, RequestStateEntryNode);
|
||||
};
|
||||
|
||||
/*! \brief A request's state, which groups all the request state entries. */
|
||||
class RequestStateNode : public Object {
|
||||
public:
|
||||
/*! \brief the request state entries */
|
||||
std::vector<RequestStateEntry> entries;
|
||||
/*! \brief tracks the request metrics. */
|
||||
RequestMetrics metrics;
|
||||
/*!
|
||||
* \brief The post-process data structures.
|
||||
* We make it a state to avoid repetitive memory allocation/free in the action post process.
|
||||
*/
|
||||
RequestActionPostProcWorkspace postproc_states;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<RequestStateNode>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.RequestState", RequestStateNode, Object);
|
||||
};
|
||||
|
||||
class RequestState : public ObjectRef {
|
||||
public:
|
||||
/*!
|
||||
* \brief Request state constructor. We take the number of response (namely "n" in OpenAI
|
||||
* API) to pre-allocate all the data structure, in order to reduce the CPU data structure
|
||||
* allocation overhead when updating the request state.
|
||||
*/
|
||||
explicit RequestState(std::vector<RequestStateEntry> entries, int num_response,
|
||||
std::chrono::high_resolution_clock::time_point add_time_point);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(RequestState, ObjectRef, RequestStateNode);
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_REQUEST_STATE_H_
|
||||
@@ -0,0 +1,588 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/sampler/cpu_sampler.cc
|
||||
* \brief The implementation for CPU sampler functions.
|
||||
*/
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
|
||||
#include "../../support/random.h"
|
||||
#include "../../support/threading_backend.h"
|
||||
#include "sampler.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() { SamplerObj::RegisterReflection(); }
|
||||
|
||||
/*!
|
||||
* \brief Sample a value from the input probability distribution with top-p.
|
||||
* The input is a batch of distributions, and we use `unit_offset` to specify
|
||||
* which distribution to sample from.
|
||||
* \param prob The input batch of probability distributions.
|
||||
* \param unit_offset The offset specifying which distribution to output
|
||||
* \param input_prob_offset The offset specifying which distribution to sample from.
|
||||
* \param top_p The top-p value of sampling.
|
||||
* \param uniform_sample The random number in [0, 1] for sampling.
|
||||
* \return The sampled value and probability.
|
||||
* \note This function is an enhancement of SampleTopPFromProb in TVM Unity.
|
||||
* We will upstream the enhancement after it gets stable.
|
||||
*/
|
||||
TokenProbPair SampleTopPFromProb(Tensor prob, int unit_offset, int input_prob_offset, double top_p,
|
||||
double uniform_sample) {
|
||||
// prob: (*, v)
|
||||
// The prob array may have arbitrary ndim and shape.
|
||||
// The last dimension corresponds to the prob distribution size.
|
||||
// We use the `unit_offset` parameter to determine which slice
|
||||
// of the prob array we sample from.
|
||||
|
||||
TVM_FFI_ICHECK(prob.IsContiguous());
|
||||
TVM_FFI_ICHECK(prob.DataType() == (DLDataType{kDLFloat, 32, 1}));
|
||||
TVM_FFI_ICHECK_EQ(prob->device.device_type, DLDeviceType::kDLCPU);
|
||||
|
||||
int64_t ndata = prob->shape[prob->ndim - 1];
|
||||
const float* __restrict p_prob =
|
||||
static_cast<float*>(__builtin_assume_aligned(prob->data, 4)) + (input_prob_offset * ndata);
|
||||
constexpr double one = 1.0f - 1e-5f;
|
||||
|
||||
if (top_p == 0) {
|
||||
// Specially handle case where top_p == 0.
|
||||
// This case is equivalent to doing argmax.
|
||||
int argmax_pos = -1;
|
||||
float max_prob = 0.0;
|
||||
float sum_prob = 0.0;
|
||||
for (int i = 0; i < ndata; ++i) {
|
||||
if (p_prob[i] > max_prob) {
|
||||
max_prob = p_prob[i];
|
||||
argmax_pos = i;
|
||||
}
|
||||
// Early exit.
|
||||
sum_prob += p_prob[i];
|
||||
if (1 - sum_prob <= max_prob) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
return {argmax_pos, 1.0};
|
||||
}
|
||||
|
||||
if (top_p >= one) {
|
||||
// Specially handle case where top_p == 1.
|
||||
double prob_sum = 0.0f;
|
||||
for (int64_t i = 0; i < ndata; ++i) {
|
||||
prob_sum += p_prob[i];
|
||||
if (prob_sum >= uniform_sample) {
|
||||
return {i, p_prob[i]};
|
||||
}
|
||||
}
|
||||
TVM_FFI_ICHECK(false) << "Possibly prob distribution contains NAN.";
|
||||
}
|
||||
|
||||
// Key observation: when we are doing top_p sampling
|
||||
// usually we only need to preserve some of the elements with
|
||||
// high probabilities before we do sort
|
||||
thread_local std::vector<std::pair<float, int>> data;
|
||||
|
||||
auto sample_top_p_with_filter = [&](float cuttoff) -> std::pair<float, int64_t> {
|
||||
data.clear();
|
||||
// filter the data with cuttoff
|
||||
float cutoff_sum = 0.0f;
|
||||
for (int64_t i = 0; i < ndata; ++i) {
|
||||
if (p_prob[i] >= cuttoff) {
|
||||
cutoff_sum += p_prob[i];
|
||||
data.emplace_back(std::make_pair(p_prob[i], static_cast<int>(i)));
|
||||
if (cutoff_sum > 1 - cuttoff) {
|
||||
// Short cut. When the remaining parts cannot have total
|
||||
// probability larger than cutoff, we can quit.
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (data.size() == 0) return std::make_pair(-1, -1);
|
||||
auto fcmp = [](const std::pair<float, int>& lhs, const std::pair<float, int>& rhs) {
|
||||
return lhs.first > rhs.first;
|
||||
};
|
||||
std::sort(data.begin(), data.end(), fcmp);
|
||||
|
||||
// short cut, if we know that
|
||||
// uniform sample < p[0] / top_p
|
||||
// we know that unform_sample < p[0] / top_p_sum
|
||||
// because top_p_sum guarantees to be smaller than top_p
|
||||
// so we can simply return the argmax sample
|
||||
// without computing anything
|
||||
if (uniform_sample < data[0].first / top_p) {
|
||||
return std::make_pair(data[0].first, data[0].second);
|
||||
}
|
||||
|
||||
// compute top_p_sum
|
||||
float cum_sum_prob = 0.0f;
|
||||
float top_p_sum = 0.0f;
|
||||
for (auto it = data.begin(); it != data.end(); ++it) {
|
||||
float prob = it->first;
|
||||
if (cum_sum_prob < top_p) {
|
||||
top_p_sum += prob;
|
||||
} else {
|
||||
// we get to the right cutoff pt
|
||||
break;
|
||||
}
|
||||
cum_sum_prob += prob;
|
||||
it->first = cum_sum_prob;
|
||||
}
|
||||
// we find that the current total sum by the given cutoff
|
||||
// is not sufficient to cover everything
|
||||
// this means we might need to retry a smaller cutoff pt.
|
||||
if (cum_sum_prob < top_p && cuttoff != 0.0f) return std::make_pair(-1, -1);
|
||||
|
||||
float last_cum_sum_prob = 0.0;
|
||||
for (auto it = data.begin(); it != data.end(); ++it) {
|
||||
if (uniform_sample < it->first / top_p_sum) {
|
||||
return std::make_pair(it->first - last_cum_sum_prob, it->second);
|
||||
}
|
||||
last_cum_sum_prob = it->first;
|
||||
}
|
||||
return std::make_pair(data[static_cast<int64_t>(data.size()) - 1].first - last_cum_sum_prob,
|
||||
data[static_cast<int64_t>(data.size()) - 1].second);
|
||||
};
|
||||
|
||||
if (top_p < 1) {
|
||||
// sample through cutoff by a number
|
||||
// by pigeonhole principle we will get at most 1024 elements
|
||||
// usually it is much less by applying this filtering(order of 10 - 20)
|
||||
data.reserve(256);
|
||||
std::pair<float, int64_t> sampled_index = sample_top_p_with_filter(top_p / 1024);
|
||||
if (sampled_index.second >= 0) return {sampled_index.second, sampled_index.first};
|
||||
}
|
||||
// fallback via full prob, rare case
|
||||
data.reserve(ndata);
|
||||
std::pair<float, int64_t> sampled_index = sample_top_p_with_filter(0.0f);
|
||||
TVM_FFI_ICHECK_GE(sampled_index.second, 0);
|
||||
return {sampled_index.second, sampled_index.first};
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Renormalize the probability distribution by the top p value.
|
||||
* \param prob The input batch of probability distributions.
|
||||
* \param unit_offset The offset specifying which distribution to output
|
||||
* \param top_p The top p value for renormalization.
|
||||
* \param eps A small epsilon value for comparison stability.
|
||||
*/
|
||||
void RenormalizeProbByTopP(Tensor prob, int unit_offset, double top_p, double eps) {
|
||||
// prob: (*, v)
|
||||
// The prob array may have arbitrary ndim and shape.
|
||||
// The last dimension corresponds to the prob distribution size.
|
||||
// We use the `unit_offset` parameter to determine which slice
|
||||
// of the prob array we will renormalize.
|
||||
TVM_FFI_ICHECK(prob.IsContiguous());
|
||||
TVM_FFI_ICHECK(prob.DataType() == (DLDataType{kDLFloat, 32, 1}));
|
||||
TVM_FFI_ICHECK_EQ(prob->device.device_type, DLDeviceType::kDLCPU);
|
||||
|
||||
if (top_p == 1.0) {
|
||||
// No renormalization is needed if top_p is 1.
|
||||
return;
|
||||
}
|
||||
|
||||
int vocab_size = prob->shape[prob->ndim - 1];
|
||||
float* __restrict p_prob =
|
||||
static_cast<float*>(__builtin_assume_aligned(prob->data, 4)) + (unit_offset * vocab_size);
|
||||
|
||||
// We manually choice the cutoff values of "top_p / 256" and "top_p / 8192".
|
||||
// In most of the cases, only one round is needed.
|
||||
std::vector<double> cutoff_values{top_p / 256, top_p / 8192, 0.0f};
|
||||
|
||||
// Create the upper partition vector and the lower partition rolling vectors.
|
||||
std::vector<float> upper_partition;
|
||||
std::vector<float> lower_partitions[2];
|
||||
upper_partition.reserve(vocab_size);
|
||||
lower_partitions[0].reserve(vocab_size);
|
||||
lower_partitions[1].reserve(vocab_size);
|
||||
float upper_partition_sum = 0.0;
|
||||
for (int round = 0; round < static_cast<int>(cutoff_values.size()); ++round) {
|
||||
const float* lower_partition_begin;
|
||||
const float* lower_partition_end;
|
||||
if (round == 0) {
|
||||
lower_partition_begin = p_prob;
|
||||
lower_partition_end = p_prob + vocab_size;
|
||||
} else {
|
||||
int idx = (round - 1) & 1;
|
||||
lower_partition_begin = lower_partitions[idx].data();
|
||||
lower_partition_end = lower_partitions[idx].data() + lower_partitions[idx].size();
|
||||
}
|
||||
|
||||
// - Partition the last round lower partition into upper and lower
|
||||
// based on the new cutoff value.
|
||||
std::vector<float>& lower_partition = lower_partitions[round & 1];
|
||||
lower_partition.clear();
|
||||
for (const float* ptr = lower_partition_begin; ptr != lower_partition_end; ++ptr) {
|
||||
if (*ptr >= cutoff_values[round]) {
|
||||
upper_partition.push_back(*ptr);
|
||||
upper_partition_sum += *ptr;
|
||||
} else {
|
||||
lower_partition.push_back(*ptr);
|
||||
}
|
||||
}
|
||||
// - If the upper partition sum is at least top p, exit the loop.
|
||||
if (upper_partition_sum >= top_p - eps) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// - Sort the upper partition in descending order.
|
||||
std::sort(upper_partition.begin(), upper_partition.end(), std::greater<>());
|
||||
// - Find the top p boundary prob value.
|
||||
float boundary_value = -1.0;
|
||||
upper_partition_sum = 0.0;
|
||||
for (float upper_value : upper_partition) {
|
||||
upper_partition_sum += upper_value;
|
||||
if (upper_partition_sum >= top_p - eps) {
|
||||
boundary_value = upper_value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
// - Mask all values smaller than the boundary to 0.
|
||||
float renormalize_sum = 0.0;
|
||||
std::vector<int> upper_partition_indices;
|
||||
upper_partition_indices.reserve(vocab_size);
|
||||
for (int i = 0; i < vocab_size; ++i) {
|
||||
if (p_prob[i] >= boundary_value) {
|
||||
upper_partition_indices.push_back(i);
|
||||
renormalize_sum += p_prob[i];
|
||||
} else {
|
||||
p_prob[i] = 0.0;
|
||||
}
|
||||
}
|
||||
// - Renormalize.
|
||||
for (int idx : upper_partition_indices) {
|
||||
p_prob[idx] /= renormalize_sum;
|
||||
}
|
||||
}
|
||||
|
||||
namespace detail {
|
||||
|
||||
/*! \brief Implementation of getting top probs on CPU. */
|
||||
template <int num_top_probs>
|
||||
std::vector<TokenProbPair> ComputeTopProbsImpl(const float* p_prob, int ndata) {
|
||||
std::vector<TokenProbPair> top_probs;
|
||||
top_probs.reserve(num_top_probs);
|
||||
for (int i = 0; i < num_top_probs; ++i) {
|
||||
top_probs.emplace_back(-1, -1.0f);
|
||||
}
|
||||
|
||||
float sum_prob = 0.0;
|
||||
// Selection argsort.
|
||||
for (int p = 0; p < ndata; ++p) {
|
||||
int i = num_top_probs - 1;
|
||||
for (; i >= 0; --i) {
|
||||
if (p_prob[p] > top_probs[i].second) {
|
||||
if (i != num_top_probs - 1) {
|
||||
top_probs[i + 1] = top_probs[i];
|
||||
}
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (i != num_top_probs - 1) {
|
||||
top_probs[i + 1] = {p, p_prob[p]};
|
||||
}
|
||||
|
||||
// Early exit.
|
||||
sum_prob += p_prob[p];
|
||||
if (1 - sum_prob <= top_probs[num_top_probs - 1].second) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
return top_probs;
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
/*! \brief Get the probs of a few number of tokens with top probabilities. */
|
||||
inline std::vector<TokenProbPair> ComputeTopProbs(Tensor prob, int unit_offset, int num_top_probs) {
|
||||
TVM_FFI_ICHECK_LE(num_top_probs, 5);
|
||||
TVM_FFI_ICHECK_EQ(prob->ndim, 2);
|
||||
int ndata = prob->shape[1];
|
||||
const float* __restrict p_prob =
|
||||
static_cast<float*>(__builtin_assume_aligned(prob->data, 4)) + (unit_offset * ndata);
|
||||
switch (num_top_probs) {
|
||||
case 0:
|
||||
return {};
|
||||
case 1:
|
||||
return detail::ComputeTopProbsImpl<1>(p_prob, ndata);
|
||||
case 2:
|
||||
return detail::ComputeTopProbsImpl<2>(p_prob, ndata);
|
||||
case 3:
|
||||
return detail::ComputeTopProbsImpl<3>(p_prob, ndata);
|
||||
case 4:
|
||||
return detail::ComputeTopProbsImpl<4>(p_prob, ndata);
|
||||
case 5:
|
||||
return detail::ComputeTopProbsImpl<5>(p_prob, ndata);
|
||||
}
|
||||
throw;
|
||||
}
|
||||
|
||||
/********************* CPU Sampler *********************/
|
||||
|
||||
class CPUSampler : public SamplerObj {
|
||||
public:
|
||||
explicit CPUSampler(Optional<EventTraceRecorder> trace_recorder)
|
||||
: trace_recorder_(std::move(trace_recorder)) {}
|
||||
|
||||
Tensor BatchRenormalizeProbsByTopP(Tensor probs_on_device, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<String>& request_ids, //
|
||||
const Array<GenerationConfig>& generation_cfg) final {
|
||||
// probs_on_device: (n, v)
|
||||
TVM_FFI_ICHECK_EQ(probs_on_device->ndim, 2);
|
||||
// - Copy probs to CPU
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start copy probs to CPU");
|
||||
Tensor probs_on_host = CopyProbsToCPU(probs_on_device);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish copy probs to CPU");
|
||||
int num_samples = sample_indices.size();
|
||||
int num_probs = probs_on_device->shape[0];
|
||||
int vocab_size = probs_on_device->shape[1];
|
||||
TVM_FFI_ICHECK_EQ(request_ids.size(), num_samples);
|
||||
TVM_FFI_ICHECK_EQ(generation_cfg.size(), num_samples);
|
||||
|
||||
std::vector<int> top_p_indices;
|
||||
std::vector<double> top_p_values;
|
||||
for (int i = 0; i < num_samples; ++i) {
|
||||
if (top_p_indices.empty() || top_p_indices.back() != sample_indices[i]) {
|
||||
top_p_indices.push_back(sample_indices[i]);
|
||||
top_p_values.push_back(generation_cfg[i]->top_p);
|
||||
} else {
|
||||
TVM_FFI_ICHECK(fabs(top_p_values.back() - generation_cfg[i]->top_p) < eps_)
|
||||
<< "Sampler requires the top_p values for each prob distribution are the same.";
|
||||
}
|
||||
}
|
||||
if (top_p_indices.empty()) {
|
||||
// Return if no top p needs to apply.
|
||||
return probs_on_host;
|
||||
}
|
||||
|
||||
tvm::runtime::parallel_for_with_threading_backend(
|
||||
[this, &probs_on_host, &request_ids, &top_p_indices, &top_p_values](int i) {
|
||||
RECORD_EVENT(this->trace_recorder_, request_ids[i], "start renormalize by top p");
|
||||
RenormalizeProbByTopP(probs_on_host, top_p_indices[i], top_p_values[i], eps_);
|
||||
RECORD_EVENT(this->trace_recorder_, request_ids[i], "finish renormalize by top p");
|
||||
},
|
||||
0, static_cast<int64_t>(top_p_indices.size()));
|
||||
|
||||
return probs_on_host;
|
||||
}
|
||||
|
||||
std::vector<SampleResult> BatchSampleTokensWithProbBeforeTopP(
|
||||
Tensor probs_on_device, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<String>& request_ids, //
|
||||
const Array<GenerationConfig>& generation_cfg, //
|
||||
const std::vector<RandomGenerator*>& rngs) final {
|
||||
// probs_on_device: (n, v)
|
||||
TVM_FFI_ICHECK_EQ(probs_on_device->ndim, 2);
|
||||
// - Copy probs to CPU
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start copy probs to CPU");
|
||||
Tensor probs_on_host = CopyProbsToCPU(probs_on_device);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish copy probs to CPU");
|
||||
|
||||
return BatchSampleTokensImpl(probs_on_host, sample_indices, request_ids, generation_cfg, rngs,
|
||||
/*top_p_applied=*/false);
|
||||
}
|
||||
|
||||
std::vector<SampleResult> BatchSampleTokensWithProbAfterTopP(
|
||||
Tensor probs_on_host, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<String>& request_ids, //
|
||||
const Array<GenerationConfig>& generation_cfg, //
|
||||
const std::vector<RandomGenerator*>& rngs) final {
|
||||
return BatchSampleTokensImpl(probs_on_host, sample_indices, request_ids, generation_cfg, rngs,
|
||||
/*top_p_applied=*/true);
|
||||
}
|
||||
|
||||
std::pair<std::vector<std::vector<SampleResult>>, std::vector<int>>
|
||||
BatchVerifyDraftTokensWithProbAfterTopP(
|
||||
Tensor probs_on_host, const Array<String>& request_ids,
|
||||
const std::vector<int>& cum_verify_lengths, const Array<GenerationConfig>& generation_cfg,
|
||||
const std::vector<RandomGenerator*>& rngs,
|
||||
const std::vector<std::vector<SampleResult>>& draft_output_tokens,
|
||||
const std::vector<int64_t>& token_tree_parent_ptr, Tensor draft_probs_on_device) final {
|
||||
// probs_on_host: (n, v)
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start draft verification");
|
||||
TVM_FFI_ICHECK_EQ(probs_on_host->ndim, 2);
|
||||
|
||||
int num_sequence = static_cast<int>(cum_verify_lengths.size()) - 1;
|
||||
TVM_FFI_ICHECK_EQ(rngs.size(), num_sequence);
|
||||
TVM_FFI_ICHECK_EQ(draft_output_tokens.size(), num_sequence);
|
||||
|
||||
Tensor draft_probs_on_host = draft_probs_on_device.CopyTo(DLDevice{kDLCPU, 0});
|
||||
std::vector<std::vector<SampleResult>> sample_results;
|
||||
sample_results.resize(num_sequence);
|
||||
|
||||
float* __restrict global_p_probs =
|
||||
static_cast<float*>(__builtin_assume_aligned(probs_on_host->data, 4));
|
||||
int vocab_size = probs_on_host->shape[1];
|
||||
|
||||
std::vector<int> last_accepted_tree_node(num_sequence, 0);
|
||||
tvm::runtime::parallel_for_with_threading_backend(
|
||||
[&](int i) {
|
||||
int verify_start = cum_verify_lengths[i];
|
||||
int verify_end = cum_verify_lengths[i + 1];
|
||||
|
||||
TVM_FFI_ICHECK_EQ(token_tree_parent_ptr[verify_start], -1);
|
||||
for (int j = verify_start + 1; j < verify_end; ++j) {
|
||||
TVM_FFI_ICHECK_EQ(token_tree_parent_ptr[j], j - verify_start - 1)
|
||||
<< "CPU sampler only supports chain-style draft tokens.";
|
||||
}
|
||||
|
||||
int cur_token_idx = 0;
|
||||
// Sub 1 to ignore the last prediction.
|
||||
for (; cur_token_idx < verify_end - verify_start - 1; ++cur_token_idx) {
|
||||
float* p_probs = global_p_probs + (verify_start + cur_token_idx) * vocab_size;
|
||||
int cur_token = draft_output_tokens[i][cur_token_idx].GetTokenId();
|
||||
float q_value = draft_output_tokens[i][cur_token_idx].sampled_token_id.second;
|
||||
float p_value = p_probs[cur_token];
|
||||
|
||||
if (p_value >= q_value) {
|
||||
sample_results[i].push_back(
|
||||
SampleResult{{cur_token, p_value},
|
||||
ComputeTopProbs(probs_on_host, verify_start + cur_token_idx,
|
||||
generation_cfg[i]->top_logprobs)});
|
||||
continue;
|
||||
}
|
||||
float r = rngs[i]->GetRandomNumber();
|
||||
if (r < p_value / (q_value + eps_)) {
|
||||
sample_results[i].push_back(
|
||||
SampleResult{{cur_token, p_value},
|
||||
ComputeTopProbs(probs_on_host, verify_start + cur_token_idx,
|
||||
generation_cfg[i]->top_logprobs)});
|
||||
continue;
|
||||
}
|
||||
|
||||
// normalize a new probability distribution
|
||||
double sum_v = 0.0;
|
||||
const float* __restrict p_qdist =
|
||||
static_cast<float*>(__builtin_assume_aligned(draft_probs_on_host->data, 4)) +
|
||||
(verify_start + cur_token_idx + 1) * vocab_size;
|
||||
|
||||
for (int j = 0; j < vocab_size; ++j) {
|
||||
p_probs[j] = std::max(p_probs[j] - p_qdist[j], 0.0f);
|
||||
sum_v += p_probs[j];
|
||||
}
|
||||
for (int j = 0; j < vocab_size; ++j) {
|
||||
p_probs[j] /= sum_v;
|
||||
}
|
||||
|
||||
// sample a new token from the new distribution
|
||||
SampleResult sample_result;
|
||||
sample_result.sampled_token_id = SampleTopPFromProb(
|
||||
probs_on_host, verify_start + cur_token_idx, verify_start + cur_token_idx,
|
||||
/*top_p=*/1.0f, rngs[i]->GetRandomNumber());
|
||||
sample_result.top_prob_tokens = ComputeTopProbs(
|
||||
probs_on_host, verify_start + cur_token_idx, generation_cfg[i]->top_logprobs);
|
||||
sample_results[i].push_back(sample_result);
|
||||
break;
|
||||
}
|
||||
last_accepted_tree_node[i] = cur_token_idx;
|
||||
// if cur_token_idx == verify_end - verify_start - 1
|
||||
// all draft tokens are accepted
|
||||
// we sample a new token
|
||||
if (cur_token_idx == verify_end - verify_start - 1) {
|
||||
SampleResult sample_result;
|
||||
// sample a new token from the original distribution
|
||||
sample_result.sampled_token_id = SampleTopPFromProb(
|
||||
probs_on_host, verify_start + cur_token_idx, verify_start + cur_token_idx,
|
||||
/*top_p=*/1.0f, rngs[i]->GetRandomNumber());
|
||||
sample_result.top_prob_tokens = ComputeTopProbs(
|
||||
probs_on_host, verify_start + cur_token_idx, generation_cfg[i]->top_logprobs);
|
||||
sample_results[i].push_back(sample_result);
|
||||
}
|
||||
},
|
||||
0, num_sequence);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish draft verification");
|
||||
return {sample_results, last_accepted_tree_node};
|
||||
}
|
||||
|
||||
private:
|
||||
std::vector<SampleResult> BatchSampleTokensImpl(Tensor probs_on_host, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<String>& request_ids, //
|
||||
const Array<GenerationConfig>& generation_cfg, //
|
||||
const std::vector<RandomGenerator*>& rngs, //
|
||||
bool top_p_applied) {
|
||||
// probs_on_host: (n, v)
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start sampling");
|
||||
TVM_FFI_ICHECK_EQ(probs_on_host->ndim, 2);
|
||||
TVM_FFI_ICHECK_EQ(probs_on_host->device.device_type, DLDeviceType::kDLCPU);
|
||||
|
||||
// - Sample tokens from probabilities.
|
||||
int n = request_ids.size();
|
||||
TVM_FFI_ICHECK_EQ(generation_cfg.size(), n);
|
||||
TVM_FFI_ICHECK_EQ(rngs.size(), n);
|
||||
|
||||
std::vector<SampleResult> sample_results;
|
||||
sample_results.resize(n);
|
||||
|
||||
tvm::runtime::parallel_for_with_threading_backend(
|
||||
[this, &sample_results, &probs_on_host, &generation_cfg, &rngs, &request_ids, top_p_applied,
|
||||
sample_indices](int i) {
|
||||
RECORD_EVENT(this->trace_recorder_, request_ids[i], "start sample token");
|
||||
// Sample top p from probability.
|
||||
double top_p =
|
||||
top_p_applied
|
||||
? 1.0f
|
||||
: (generation_cfg[i]->temperature < eps_ ? 0.0 : generation_cfg[i]->top_p);
|
||||
sample_results[i].sampled_token_id = SampleTopPFromProb(
|
||||
probs_on_host, i, sample_indices[i], top_p, rngs[i]->GetRandomNumber());
|
||||
sample_results[i].top_prob_tokens =
|
||||
ComputeTopProbs(probs_on_host, i, generation_cfg[i]->top_logprobs);
|
||||
RECORD_EVENT(this->trace_recorder_, request_ids[i], "finish sample token");
|
||||
},
|
||||
0, n);
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish sampling");
|
||||
return sample_results;
|
||||
}
|
||||
|
||||
/*! \brief Copy prob distributions from device to CPU. */
|
||||
Tensor CopyProbsToCPU(Tensor probs_on_device) {
|
||||
// probs_on_device: (n, v)
|
||||
if (probs_on_device->device.device_type == kDLCPU) {
|
||||
return probs_on_device;
|
||||
}
|
||||
|
||||
TVM_FFI_ICHECK(probs_on_device->device.device_type != kDLCPU);
|
||||
if (probs_host_.defined()) {
|
||||
TVM_FFI_ICHECK_EQ(probs_host_->shape[1], probs_on_device->shape[1]);
|
||||
}
|
||||
|
||||
int64_t init_size = probs_host_.defined() ? probs_host_->shape[0] : 32;
|
||||
int64_t num_tokens = probs_on_device->shape[0];
|
||||
int64_t vocab_size = probs_on_device->shape[1];
|
||||
while (init_size < num_tokens) {
|
||||
init_size *= 2;
|
||||
}
|
||||
if (!probs_host_.defined() || init_size != probs_host_->shape[0]) {
|
||||
probs_host_ =
|
||||
Tensor::Empty({init_size, vocab_size}, probs_on_device->dtype, DLDevice{kDLCPU, 0});
|
||||
}
|
||||
TVM_FFI_ICHECK_LE(num_tokens, probs_host_->shape[0]);
|
||||
Tensor view = probs_host_.CreateView({num_tokens, vocab_size}, probs_on_device->dtype);
|
||||
view.CopyFrom(probs_on_device);
|
||||
return view;
|
||||
}
|
||||
|
||||
/*! \brief The event trace recorder for requests. */
|
||||
Optional<EventTraceRecorder> trace_recorder_;
|
||||
/*! \brief Customized function which computes prob distribution from logits */
|
||||
Function flogits_to_probs_inplace_;
|
||||
/*! \brief Probability distribution array on CPU. */
|
||||
Tensor probs_host_{nullptr};
|
||||
const float eps_ = 1e-5;
|
||||
};
|
||||
|
||||
Sampler Sampler::CreateCPUSampler(Optional<EventTraceRecorder> trace_recorder) {
|
||||
return Sampler(tvm::ffi::make_object<CPUSampler>(std::move(trace_recorder)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,756 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/sampler/gpu_sampler.cc
|
||||
* \brief The implementation for GPU sampler functions.
|
||||
*/
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/runtime/device_api.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
#include <tvm/support/cuda/nvtx.h>
|
||||
|
||||
#include "../../support/random.h"
|
||||
#include "sampler.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::support::NVTXScopedRange;
|
||||
|
||||
inline bool FlashInferSamplingAvailable(Device device) {
|
||||
// Device must be CUDA, and FlashInfer must be enabled.
|
||||
if (device.device_type != DLDeviceType::kDLCUDA ||
|
||||
!Function::GetGlobal("flashinfer.sampling.parallel_sampling_from_prob").has_value()) {
|
||||
return false;
|
||||
}
|
||||
// Compute version must be at least 8.0
|
||||
Any rv;
|
||||
DeviceAPI::Get(device)->GetAttr(device, kComputeVersion, &rv);
|
||||
std::string compute_version = rv.cast<std::string>();
|
||||
std::string major_version = compute_version.substr(0, compute_version.find('.'));
|
||||
return std::stoi(major_version) >= 8;
|
||||
}
|
||||
|
||||
inline void CopyArray(Tensor src, Tensor dst, TVMStreamHandle copy_stream) {
|
||||
DLTensor dl_dst = *(dst.operator->());
|
||||
Tensor::CopyFromTo(src.operator->(), &dl_dst, copy_stream);
|
||||
}
|
||||
|
||||
inline void SyncCopyStream(Device device, TVMStreamHandle compute_stream,
|
||||
TVMStreamHandle copy_stream) {
|
||||
// - If there is no particular copy stream, no action is needed.
|
||||
if (copy_stream == nullptr) {
|
||||
return;
|
||||
}
|
||||
// - Sync two streams.
|
||||
DeviceAPI::Get(device)->SyncStreamFromTo(device, copy_stream, compute_stream);
|
||||
}
|
||||
|
||||
/*********************** GPU Sampler ***********************/
|
||||
|
||||
class GPUSampler : public SamplerObj {
|
||||
public:
|
||||
explicit GPUSampler(int max_num_sample, int vocab_size, FunctionTable* ft, DLDevice device,
|
||||
Optional<EventTraceRecorder> trace_recorder)
|
||||
: max_num_sample_(max_num_sample),
|
||||
vocab_size_(vocab_size),
|
||||
flashinfer_sampling_available_(FlashInferSamplingAvailable(device)),
|
||||
device_(device),
|
||||
gpu_multinomial_from_uniform_func_(ft->gpu_multinomial_from_uniform_func_),
|
||||
gpu_argsort_probs_func_(ft->gpu_argsort_probs_func_),
|
||||
gpu_sample_with_top_p_func_(ft->gpu_sample_with_top_p_func_),
|
||||
gpu_sampler_take_probs_func_(ft->gpu_sampler_take_probs_func_),
|
||||
gpu_verify_draft_tokens_func_(ft->gpu_verify_draft_tokens_func_),
|
||||
gpu_renormalize_by_top_p_func_(ft->gpu_renormalize_by_top_p_func_),
|
||||
trace_recorder_(std::move(trace_recorder)) {
|
||||
TVM_FFI_ICHECK(gpu_multinomial_from_uniform_func_.defined());
|
||||
TVM_FFI_ICHECK(gpu_argsort_probs_func_.defined());
|
||||
TVM_FFI_ICHECK(gpu_sample_with_top_p_func_.defined());
|
||||
TVM_FFI_ICHECK(gpu_sampler_take_probs_func_.defined());
|
||||
|
||||
flashinfer_multinomial_sample_func_ =
|
||||
Function::GetGlobal("flashinfer.sampling.parallel_sampling_from_prob");
|
||||
|
||||
Device preferred_host_device = GetPreferredHostDevice(device);
|
||||
// We support at most 5 top prob results for each sequence.
|
||||
// Initialize auxiliary arrays on CPU.
|
||||
uniform_samples_host_ = Tensor::Empty({max_num_sample}, dtype_f32_, preferred_host_device);
|
||||
sample_indices_host_ = Tensor::Empty({max_num_sample}, dtype_i32_, preferred_host_device);
|
||||
top_p_host_ = Tensor::Empty({max_num_sample}, dtype_f32_, preferred_host_device);
|
||||
top_p_init_pivots_host_ = Tensor::Empty({max_num_sample, num_top_p_cutoff_pivots_}, dtype_f32_,
|
||||
preferred_host_device);
|
||||
top_prob_offsets_host_ = Tensor::Empty({max_num_sample * 5}, dtype_i32_, preferred_host_device);
|
||||
draft_tokens_host_ = Tensor::Empty({max_num_sample}, dtype_i32_, preferred_host_device);
|
||||
token_tree_first_child_host_ =
|
||||
Tensor::Empty({max_num_sample}, dtype_i32_, preferred_host_device);
|
||||
token_tree_next_sibling_host_ =
|
||||
Tensor::Empty({max_num_sample}, dtype_i32_, preferred_host_device);
|
||||
token_tree_parent_ptr_host_ =
|
||||
Tensor::Empty({max_num_sample}, dtype_i32_, preferred_host_device);
|
||||
sampled_token_ids_host_ = Tensor::Empty({max_num_sample}, dtype_i32_, preferred_host_device);
|
||||
sampled_probs_host_ = Tensor::Empty({max_num_sample}, dtype_f32_, preferred_host_device);
|
||||
top_prob_probs_host_ = Tensor::Empty({max_num_sample * 5}, dtype_f32_, preferred_host_device);
|
||||
top_prob_indices_host_ = Tensor::Empty({max_num_sample * 5}, dtype_i32_, preferred_host_device);
|
||||
// Initialize auxiliary arrays on GPU.
|
||||
uniform_samples_device_ = Tensor::Empty({max_num_sample}, dtype_f32_, device);
|
||||
sample_indices_device_ = Tensor::Empty({max_num_sample}, dtype_i32_, device);
|
||||
top_p_device_ = Tensor::Empty({max_num_sample}, dtype_f32_, device);
|
||||
top_p_init_pivots_device_ =
|
||||
Tensor::Empty({max_num_sample, num_top_p_cutoff_pivots_}, dtype_f32_, device);
|
||||
top_prob_offsets_device_ = Tensor::Empty({max_num_sample * 5}, dtype_i32_, device);
|
||||
draft_tokens_device_ = Tensor::Empty({max_num_sample}, dtype_i32_, device);
|
||||
token_tree_first_child_device_ = Tensor::Empty({max_num_sample}, dtype_i32_, device);
|
||||
token_tree_next_sibling_device_ = Tensor::Empty({max_num_sample}, dtype_i32_, device);
|
||||
token_tree_parent_ptr_device_ = Tensor::Empty({max_num_sample}, dtype_i32_, device);
|
||||
sampled_token_ids_device_ = Tensor::Empty({max_num_sample}, dtype_i32_, device);
|
||||
|
||||
// If the device is CUDA/ROCm, we create a standalone copy stream, in
|
||||
// purpose to hide the latency of auxiliary stream copy.
|
||||
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);
|
||||
copy_stream_ = DeviceAPI::Get(device)->CreateStream(device);
|
||||
}
|
||||
}
|
||||
|
||||
~GPUSampler() {
|
||||
// Free the copy stream if defined.
|
||||
if (copy_stream_ != nullptr) {
|
||||
DeviceAPI::Get(device_)->FreeStream(device_, copy_stream_);
|
||||
}
|
||||
}
|
||||
|
||||
Tensor BatchRenormalizeProbsByTopP(Tensor probs_on_device, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<String>& request_ids, //
|
||||
const Array<GenerationConfig>& generation_cfg) final {
|
||||
NVTXScopedRange nvtx_scope("BatchRenormalizeProbsByTopP");
|
||||
// probs_on_device: (n, v)
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start renormalization by top p");
|
||||
TVM_FFI_ICHECK_EQ(probs_on_device->ndim, 2);
|
||||
int num_samples = sample_indices.size();
|
||||
int num_probs = probs_on_device->shape[0];
|
||||
int vocab_size = probs_on_device->shape[1];
|
||||
TVM_FFI_ICHECK_LE(num_probs, max_num_sample_);
|
||||
TVM_FFI_ICHECK_EQ(generation_cfg.size(), num_samples);
|
||||
|
||||
// - Check if there is need for applying top p.
|
||||
bool need_top_p = CheckTopP(generation_cfg, sample_indices, num_probs, num_samples, vocab_size);
|
||||
if (!need_top_p) {
|
||||
return probs_on_device;
|
||||
}
|
||||
|
||||
// - Copy auxiliary array for top-p and initial pivots.
|
||||
Tensor top_p_host = top_p_host_.CreateView({num_probs}, dtype_f32_);
|
||||
Tensor top_p_device = top_p_device_.CreateView({num_probs}, dtype_f32_);
|
||||
CopyArray(/*src=*/top_p_host, /*dst=*/top_p_device, copy_stream_);
|
||||
|
||||
Tensor top_p_init_pivots_host =
|
||||
top_p_init_pivots_host_.CreateView({num_probs, num_top_p_cutoff_pivots_}, dtype_f32_);
|
||||
Tensor top_p_init_pivots_device =
|
||||
top_p_init_pivots_device_.CreateView({num_probs, num_top_p_cutoff_pivots_}, dtype_f32_);
|
||||
const float* p_top_p = static_cast<const float*>(top_p_host->data);
|
||||
float* p_top_p_init_pivots = static_cast<float*>(top_p_init_pivots_host->data);
|
||||
for (int i = 0; i < num_probs; ++i) {
|
||||
if (1 - p_top_p[i] >= 0.02) {
|
||||
p_top_p_init_pivots[i * num_top_p_cutoff_pivots_] =
|
||||
std::min(1 - p_top_p[i], static_cast<float>(0.5));
|
||||
p_top_p_init_pivots[i * num_top_p_cutoff_pivots_ + 1] = 0.02;
|
||||
p_top_p_init_pivots[i * num_top_p_cutoff_pivots_ + 2] = 0.01;
|
||||
} else {
|
||||
p_top_p_init_pivots[i * num_top_p_cutoff_pivots_] = 1 - p_top_p[i];
|
||||
p_top_p_init_pivots[i * num_top_p_cutoff_pivots_ + 1] = (1 - p_top_p[i]) / 2;
|
||||
p_top_p_init_pivots[i * num_top_p_cutoff_pivots_ + 2] = (1 - p_top_p[i]) / 4;
|
||||
}
|
||||
}
|
||||
CopyArray(/*src=*/top_p_init_pivots_host, /*dst=*/top_p_init_pivots_device, copy_stream_);
|
||||
SyncCopyStream(device_, compute_stream_, copy_stream_);
|
||||
|
||||
// - Renormalize the prob with top p.
|
||||
Tensor renormed_probs_on_device =
|
||||
gpu_renormalize_by_top_p_func_(probs_on_device, top_p_device, top_p_init_pivots_device)
|
||||
.cast<Tensor>();
|
||||
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish renormalization by top p");
|
||||
return renormed_probs_on_device;
|
||||
}
|
||||
|
||||
std::vector<SampleResult> BatchSampleTokensWithProbBeforeTopP(
|
||||
Tensor probs_on_device, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<String>& request_ids, //
|
||||
const Array<GenerationConfig>& generation_cfg, //
|
||||
const std::vector<RandomGenerator*>& rngs) final {
|
||||
NVTXScopedRange nvtx_scope("BatchSampleTokensWithProbBeforeTopP");
|
||||
return BatchSampleTokensImpl(std::move(probs_on_device), sample_indices, request_ids,
|
||||
generation_cfg, rngs, /*top_p_applied=*/false);
|
||||
}
|
||||
|
||||
std::vector<SampleResult> BatchSampleTokensWithProbAfterTopP(
|
||||
Tensor probs_on_device, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<String>& request_ids, //
|
||||
const Array<GenerationConfig>& generation_cfg, //
|
||||
const std::vector<RandomGenerator*>& rngs) final {
|
||||
NVTXScopedRange nvtx_scope("BatchSampleTokensWithProbAfterTopP");
|
||||
return BatchSampleTokensImpl(std::move(probs_on_device), sample_indices, request_ids,
|
||||
generation_cfg, rngs, /*top_p_applied=*/true);
|
||||
}
|
||||
|
||||
std::pair<std::vector<std::vector<SampleResult>>, std::vector<int>>
|
||||
BatchVerifyDraftTokensWithProbAfterTopP(
|
||||
Tensor probs_on_device, const Array<String>& request_ids,
|
||||
const std::vector<int>& cum_verify_lengths, const Array<GenerationConfig>& generation_cfg,
|
||||
const std::vector<RandomGenerator*>& rngs,
|
||||
const std::vector<std::vector<SampleResult>>& draft_output_tokens,
|
||||
const std::vector<int64_t>& token_tree_parent_ptr, Tensor draft_probs_on_device) final {
|
||||
NVTXScopedRange nvtx_scope("BatchVerifyDraftTokensWithProbAfterTopP");
|
||||
std::vector<std::vector<SampleResult>> sample_results;
|
||||
// probs_on_device: (n, v)
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start draft verification");
|
||||
TVM_FFI_ICHECK_EQ(probs_on_device->ndim, 2);
|
||||
|
||||
int num_sequence = static_cast<int>(cum_verify_lengths.size()) - 1;
|
||||
TVM_FFI_ICHECK_EQ(rngs.size(), num_sequence);
|
||||
TVM_FFI_ICHECK_EQ(draft_output_tokens.size(), num_sequence);
|
||||
sample_results.resize(num_sequence);
|
||||
|
||||
int num_nodes = cum_verify_lengths.back();
|
||||
TVM_FFI_ICHECK(num_nodes <= max_num_sample_);
|
||||
TVM_FFI_ICHECK_EQ(draft_probs_on_device->shape[0], num_nodes);
|
||||
Tensor uniform_samples_device = GenerateUniformSamples(rngs, cum_verify_lengths);
|
||||
Tensor draft_tokens_host = draft_tokens_host_.CreateView({num_nodes}, dtype_i32_);
|
||||
Tensor draft_tokens_device = draft_tokens_device_.CreateView({num_nodes}, dtype_i32_);
|
||||
|
||||
// Copy draft tokens to GPU
|
||||
int* p_draft_tokens_host = static_cast<int*>(draft_tokens_host->data);
|
||||
for (int i = 0; i < num_sequence; i++) {
|
||||
const std::vector<SampleResult>& draft_output_tokens_i = draft_output_tokens[i];
|
||||
int start = cum_verify_lengths[i];
|
||||
int end = cum_verify_lengths[i + 1];
|
||||
// start/end is the range of the sequence i in probs_on_device, which includes the prob dist
|
||||
// of the draft tokens and the last committed token
|
||||
TVM_FFI_ICHECK_EQ(draft_output_tokens_i.size() + 1, end - start);
|
||||
for (int j = 0; j < end - start - 1; j++) {
|
||||
// Copy sampled token id
|
||||
p_draft_tokens_host[start + j + 1] = draft_output_tokens_i[j].GetTokenId();
|
||||
}
|
||||
}
|
||||
CopyArray(draft_tokens_host, draft_tokens_device, copy_stream_);
|
||||
|
||||
Tensor token_tree_first_child_host =
|
||||
token_tree_first_child_host_.CreateView({num_nodes}, dtype_i32_);
|
||||
Tensor token_tree_first_child_device =
|
||||
token_tree_first_child_device_.CreateView({num_nodes}, dtype_i32_);
|
||||
Tensor token_tree_next_sibling_host =
|
||||
token_tree_next_sibling_host_.CreateView({num_nodes}, dtype_i32_);
|
||||
Tensor token_tree_next_sibling_device =
|
||||
token_tree_next_sibling_device_.CreateView({num_nodes}, dtype_i32_);
|
||||
Tensor token_tree_parent_ptr_host =
|
||||
token_tree_parent_ptr_host_.CreateView({num_sequence}, dtype_i32_);
|
||||
Tensor token_tree_parent_ptr_device =
|
||||
token_tree_parent_ptr_device_.CreateView({num_sequence}, dtype_i32_);
|
||||
std::vector<int> token_tree_child_to_parent(/*n=*/num_nodes);
|
||||
|
||||
int* token_tree_first_child_ptr_host = static_cast<int*>(token_tree_first_child_host->data);
|
||||
int* token_tree_next_sibling_ptr_host = static_cast<int*>(token_tree_next_sibling_host->data);
|
||||
// Build the tree structure on CPU
|
||||
for (int i = 0; i < num_sequence; i++) {
|
||||
// Assuming no tree structure for now
|
||||
int start = cum_verify_lengths[i];
|
||||
int end = cum_verify_lengths[i + 1];
|
||||
TVM_FFI_ICHECK_GE(end - start, 2);
|
||||
for (int j = 0; j < end - start; j++) {
|
||||
int cur_node = j + start;
|
||||
int parent_node =
|
||||
token_tree_parent_ptr[cur_node] != -1 ? token_tree_parent_ptr[cur_node] + start : -1;
|
||||
token_tree_first_child_ptr_host[cur_node] = -1;
|
||||
if (parent_node != -1 && token_tree_first_child_ptr_host[parent_node] == -1) {
|
||||
token_tree_first_child_ptr_host[parent_node] = cur_node;
|
||||
}
|
||||
token_tree_child_to_parent[cur_node] = parent_node;
|
||||
if (cur_node + 1 < end && token_tree_parent_ptr[cur_node - start + 1] ==
|
||||
token_tree_parent_ptr[cur_node - start]) {
|
||||
token_tree_next_sibling_ptr_host[cur_node] = cur_node + 1;
|
||||
} else {
|
||||
token_tree_next_sibling_ptr_host[cur_node] = -1;
|
||||
}
|
||||
}
|
||||
static_cast<int*>(token_tree_parent_ptr_host->data)[i] = start; // point to the root
|
||||
}
|
||||
// Copy token tree structure to GPU
|
||||
CopyArray(token_tree_first_child_host, token_tree_first_child_device, copy_stream_);
|
||||
CopyArray(token_tree_next_sibling_host, token_tree_next_sibling_device, copy_stream_);
|
||||
CopyArray(token_tree_parent_ptr_host, token_tree_parent_ptr_device, copy_stream_);
|
||||
|
||||
SyncCopyStream(device_, compute_stream_, copy_stream_);
|
||||
|
||||
gpu_verify_draft_tokens_func_(draft_probs_on_device, draft_tokens_device, probs_on_device,
|
||||
token_tree_first_child_device, token_tree_next_sibling_device,
|
||||
uniform_samples_device, token_tree_parent_ptr_device);
|
||||
|
||||
DeviceAPI::Get(device_)->SyncStreamFromTo(device_, compute_stream_, copy_stream_);
|
||||
CopyArray(token_tree_parent_ptr_device, token_tree_parent_ptr_host, copy_stream_);
|
||||
|
||||
std::vector<SampleResult> additional_sample_result;
|
||||
{
|
||||
additional_sample_result.reserve(num_sequence);
|
||||
// Sample one additional token for each sequence using the probablity at the last accepted
|
||||
// token.
|
||||
uniform_samples_device = GenerateUniformSamples(rngs, num_sequence);
|
||||
const Tensor& sample_indices_device = token_tree_parent_ptr_device;
|
||||
// Check need_prob_values
|
||||
bool need_prob_values = false;
|
||||
for (int i = 0; i < num_sequence; i++) {
|
||||
need_prob_values |= generation_cfg[i]->logprobs;
|
||||
}
|
||||
std::vector<int> top_prob_offset_indptr;
|
||||
if (!need_prob_values) {
|
||||
top_prob_offset_indptr.resize(num_sequence + 1, 0);
|
||||
} else {
|
||||
// Slow path: if any of the generation config requires prob values, we need to copy
|
||||
// sample_indices to host to compute top_prob_offset_indptr.
|
||||
DeviceAPI::Get(device_)->StreamSync(device_, copy_stream_);
|
||||
std::vector<int> sample_indices;
|
||||
sample_indices.reserve(num_sequence);
|
||||
const int* p_token_tree_parent_ptr = static_cast<int*>(token_tree_parent_ptr_host->data);
|
||||
for (int i = 0; i < num_sequence; i++) {
|
||||
sample_indices.push_back(p_token_tree_parent_ptr[i]);
|
||||
}
|
||||
CheckProbValues(generation_cfg, sample_indices, num_nodes, num_sequence, vocab_size_,
|
||||
&top_prob_offset_indptr);
|
||||
}
|
||||
auto device_arrays =
|
||||
SampleOnGPU(probs_on_device, uniform_samples_device, sample_indices_device,
|
||||
/*need_top_p=*/false, need_prob_values, num_nodes, top_prob_offset_indptr);
|
||||
auto host_arrays = CopyArraysToCPU(device_arrays, num_sequence, need_prob_values,
|
||||
top_prob_offset_indptr.back());
|
||||
additional_sample_result =
|
||||
CollectSampleResult(host_arrays, num_sequence, need_prob_values, top_prob_offset_indptr);
|
||||
}
|
||||
|
||||
std::vector<int> last_accepted_tree_node;
|
||||
last_accepted_tree_node.reserve(num_sequence);
|
||||
for (int i = 0; i < num_sequence; i++) {
|
||||
int start = cum_verify_lengths[i];
|
||||
int end = cum_verify_lengths[i + 1];
|
||||
int last_accepted = static_cast<int*>(token_tree_parent_ptr_host->data)[i];
|
||||
last_accepted_tree_node.push_back(last_accepted - start);
|
||||
int num_accepted = 0;
|
||||
for (int cur_node = last_accepted; cur_node != start;
|
||||
cur_node = token_tree_child_to_parent[cur_node]) {
|
||||
sample_results[i].push_back(draft_output_tokens[i][cur_node - start - 1]);
|
||||
num_accepted++;
|
||||
}
|
||||
std::reverse(sample_results[i].rbegin(), sample_results[i].rbegin() + num_accepted);
|
||||
}
|
||||
|
||||
// Append the additional sample result to the sample_results
|
||||
TVM_FFI_ICHECK_EQ(additional_sample_result.size(), num_sequence);
|
||||
for (int i = 0; i < num_sequence; i++) {
|
||||
sample_results[i].push_back(additional_sample_result[i]);
|
||||
}
|
||||
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish draft verification");
|
||||
return {sample_results, last_accepted_tree_node};
|
||||
}
|
||||
|
||||
private:
|
||||
std::vector<SampleResult> BatchSampleTokensImpl(Tensor probs_on_device, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<String>& request_ids, //
|
||||
const Array<GenerationConfig>& generation_cfg, //
|
||||
const std::vector<RandomGenerator*>& rngs, //
|
||||
bool top_p_applied) {
|
||||
// probs_on_device: (n, v)
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "start sampling");
|
||||
TVM_FFI_ICHECK_EQ(probs_on_device->ndim, 2);
|
||||
TVM_FFI_ICHECK_EQ(probs_on_device->device.device_id, device_.device_id);
|
||||
TVM_FFI_ICHECK_EQ(probs_on_device->device.device_type, device_.device_type);
|
||||
int num_samples = sample_indices.size();
|
||||
int num_probs = probs_on_device->shape[0];
|
||||
int vocab_size = probs_on_device->shape[1];
|
||||
if (num_samples == 0) {
|
||||
// This synchronization is necessary for making sure that this round
|
||||
// of model forward is finished.
|
||||
DeviceAPI::Get(device_)->StreamSync(device_, compute_stream_);
|
||||
return {};
|
||||
}
|
||||
TVM_FFI_ICHECK_EQ(request_ids.size(), num_samples);
|
||||
TVM_FFI_ICHECK_EQ(generation_cfg.size(), num_samples);
|
||||
TVM_FFI_ICHECK_EQ(rngs.size(), num_samples);
|
||||
|
||||
// Since `num_samples` may be larger than `max_num_sample_` in some cases,
|
||||
// we apply chunking to support large `num_samples`.
|
||||
std::vector<SampleResult> sample_results;
|
||||
if (num_samples <= max_num_sample_) {
|
||||
sample_results = ChunkSampleTokensImpl(probs_on_device, sample_indices, generation_cfg, rngs,
|
||||
top_p_applied);
|
||||
} else {
|
||||
for (int chunk_start = 0; chunk_start < num_samples; chunk_start += max_num_sample_) {
|
||||
int chunk_end = std::min(chunk_start + max_num_sample_, num_samples);
|
||||
std::vector<int> sample_indices_chunk(sample_indices.begin() + chunk_start,
|
||||
sample_indices.begin() + chunk_end);
|
||||
Array<GenerationConfig> generation_cfg_chunk(generation_cfg.begin() + chunk_start,
|
||||
generation_cfg.begin() + chunk_end);
|
||||
std::vector<RandomGenerator*> rngs_chunk(rngs.begin() + chunk_start,
|
||||
rngs.begin() + chunk_end);
|
||||
std::vector<SampleResult> sample_results_chunk = ChunkSampleTokensImpl(
|
||||
probs_on_device, sample_indices_chunk, generation_cfg_chunk, rngs_chunk, top_p_applied);
|
||||
sample_results.insert(sample_results.end(), sample_results_chunk.begin(),
|
||||
sample_results_chunk.end());
|
||||
}
|
||||
}
|
||||
|
||||
RECORD_EVENT(trace_recorder_, request_ids, "finish sampling");
|
||||
return sample_results;
|
||||
}
|
||||
|
||||
/*! \brief Collect the sampling results from the computed Tensor results. */
|
||||
std::vector<SampleResult> CollectSampleResult(const std::vector<Tensor>& host_arrays,
|
||||
int num_samples, bool need_prob_values,
|
||||
const std::vector<int> top_prob_offset_indptr) {
|
||||
const int* p_sampled_token_ids = static_cast<const int*>(host_arrays[0]->data);
|
||||
const float* p_sampled_probs = nullptr;
|
||||
const float* p_top_prob_probs = nullptr;
|
||||
const int* p_top_prob_indices = nullptr;
|
||||
if (need_prob_values) {
|
||||
p_sampled_probs = static_cast<const float*>(host_arrays[1]->data);
|
||||
p_top_prob_probs = static_cast<const float*>(host_arrays[2]->data);
|
||||
p_top_prob_indices = static_cast<const int*>(host_arrays[3]->data);
|
||||
}
|
||||
std::vector<SampleResult> sample_results;
|
||||
sample_results.reserve(num_samples);
|
||||
TVM_FFI_ICHECK_EQ(top_prob_offset_indptr.size(), num_samples + 1);
|
||||
for (int i = 0; i < num_samples; ++i) {
|
||||
// Note: we set the probability in SampleResult to 1.0 since prob value is not needed.
|
||||
float sampled_prob = need_prob_values ? p_sampled_probs[i] : 1.0;
|
||||
std::vector<TokenProbPair> top_prob_tokens;
|
||||
top_prob_tokens.reserve(top_prob_offset_indptr[i + 1] - top_prob_offset_indptr[i]);
|
||||
for (int j = top_prob_offset_indptr[i]; j < top_prob_offset_indptr[i + 1]; ++j) {
|
||||
top_prob_tokens.emplace_back(p_top_prob_indices[j], p_top_prob_probs[j]);
|
||||
}
|
||||
sample_results.push_back(
|
||||
SampleResult{{p_sampled_token_ids[i], sampled_prob}, top_prob_tokens});
|
||||
}
|
||||
return sample_results;
|
||||
}
|
||||
|
||||
std::vector<SampleResult> ChunkSampleTokensImpl(Tensor probs_on_device, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<GenerationConfig>& generation_cfg, //
|
||||
const std::vector<RandomGenerator*>& rngs, //
|
||||
bool top_p_applied) {
|
||||
// probs_on_device: (n, v)
|
||||
int num_samples = sample_indices.size();
|
||||
int num_probs = probs_on_device->shape[0];
|
||||
int vocab_size = probs_on_device->shape[1];
|
||||
|
||||
// - Generate random numbers.
|
||||
// Copy the random numbers and sample indices.
|
||||
auto uniform_samples_device = GenerateUniformSamples(rngs, num_samples);
|
||||
auto sample_indices_device = CopySampleIndicesToGPU(sample_indices);
|
||||
|
||||
// - Check if there is need for applying top p or prob values,
|
||||
// so that argsort is needed.
|
||||
bool need_top_p = false;
|
||||
if (!top_p_applied) {
|
||||
need_top_p = CheckTopP(generation_cfg, sample_indices, num_probs, num_samples, vocab_size);
|
||||
}
|
||||
// The indptr array of the number of top probs for each sample.
|
||||
std::vector<int> top_prob_offset_indptr;
|
||||
bool need_prob_values = CheckProbValues(generation_cfg, sample_indices, num_probs, num_samples,
|
||||
vocab_size, &top_prob_offset_indptr);
|
||||
|
||||
// - Sample tokens on GPU, and take out the probability values if needed.
|
||||
std::vector<Tensor> device_arrays =
|
||||
SampleOnGPU(probs_on_device, uniform_samples_device, sample_indices_device, need_top_p,
|
||||
need_prob_values, num_probs, top_prob_offset_indptr);
|
||||
|
||||
// - Copy the GPU sampling function results to CPU.
|
||||
std::vector<Tensor> host_arrays = CopyArraysToCPU(device_arrays, num_samples, need_prob_values,
|
||||
top_prob_offset_indptr.back());
|
||||
|
||||
// - Collect the sampling results.
|
||||
return CollectSampleResult(host_arrays, num_samples, need_prob_values, top_prob_offset_indptr);
|
||||
}
|
||||
|
||||
/*! \brief Generate num_samples uniform random numbers, and copy them to GPU. */
|
||||
Tensor GenerateUniformSamples(const std::vector<RandomGenerator*>& rngs, int num_samples) {
|
||||
float* p_uniform_samples = static_cast<float*>(uniform_samples_host_->data);
|
||||
for (int i = 0; i < num_samples; ++i) {
|
||||
p_uniform_samples[i] = rngs[i]->GetRandomNumber();
|
||||
}
|
||||
Tensor uniform_samples_host = uniform_samples_host_.CreateView({num_samples}, dtype_f32_);
|
||||
Tensor uniform_samples_device = uniform_samples_device_.CreateView({num_samples}, dtype_f32_);
|
||||
CopyArray(/*src=*/uniform_samples_host, /*dst=*/uniform_samples_device, copy_stream_);
|
||||
return uniform_samples_device;
|
||||
}
|
||||
|
||||
/*! \brief Generate uniform random numbers, and copy the numbers and sample indices to GPU. The
|
||||
* number of samples for each random generator is given by `cum_num_samples`. */
|
||||
Tensor GenerateUniformSamples(const std::vector<RandomGenerator*>& rngs,
|
||||
const std::vector<int>& cum_num_samples) {
|
||||
float* p_uniform_samples = static_cast<float*>(uniform_samples_host_->data);
|
||||
int total_samples = cum_num_samples.back();
|
||||
for (int i = 0; i + 1 < static_cast<int>(cum_num_samples.size()); ++i) {
|
||||
for (int j = cum_num_samples[i]; j < cum_num_samples[i + 1]; ++j) {
|
||||
p_uniform_samples[j] = rngs[i]->GetRandomNumber();
|
||||
}
|
||||
}
|
||||
Tensor uniform_samples_host = uniform_samples_host_.CreateView({total_samples}, dtype_f32_);
|
||||
Tensor uniform_samples_device = uniform_samples_device_.CreateView({total_samples}, dtype_f32_);
|
||||
CopyArray(/*src=*/uniform_samples_host, /*dst=*/uniform_samples_device, copy_stream_);
|
||||
return uniform_samples_device;
|
||||
}
|
||||
|
||||
/*! \brief Generate uniform random numbers, and copy the numbers and sample indices to GPU. */
|
||||
Tensor CopySampleIndicesToGPU(const std::vector<int>& sample_indices) {
|
||||
int* p_sample_indices = static_cast<int*>(sample_indices_host_->data);
|
||||
std::copy(sample_indices.begin(), sample_indices.end(), p_sample_indices);
|
||||
// Copy the sample indices to GPU.
|
||||
int num_samples = static_cast<int>(sample_indices.size());
|
||||
Tensor sample_indices_host = sample_indices_host_.CreateView({num_samples}, dtype_i32_);
|
||||
Tensor sample_indices_device = sample_indices_device_.CreateView({num_samples}, dtype_i32_);
|
||||
CopyArray(/*src=*/sample_indices_host, /*dst=*/sample_indices_device, copy_stream_);
|
||||
return sample_indices_device;
|
||||
}
|
||||
|
||||
/*! \brief Check if top p is needed. Update host top p array in place. */
|
||||
bool CheckTopP(const Array<GenerationConfig>& generation_cfg,
|
||||
const std::vector<int>& sample_indices, int num_probs, int num_samples,
|
||||
int vocab_size) {
|
||||
// Initialize top p values with -1.
|
||||
float* p_top_p = static_cast<float*>(top_p_host_->data);
|
||||
for (int i = 0; i < num_probs; ++i) {
|
||||
p_top_p[i] = -1.0;
|
||||
}
|
||||
bool need_top_p = false;
|
||||
for (int i = 0; i < num_samples; ++i) {
|
||||
if (p_top_p[sample_indices[i]] == -1.0) {
|
||||
p_top_p[sample_indices[i]] = generation_cfg[i]->top_p;
|
||||
need_top_p |= generation_cfg[i]->top_p != 1.0;
|
||||
} else {
|
||||
TVM_FFI_ICHECK(fabs(p_top_p[sample_indices[i]] - generation_cfg[i]->top_p) < eps_)
|
||||
<< "GPU sampler requires the top_p values for each prob distribution are the same.";
|
||||
}
|
||||
}
|
||||
for (int i = 0; i < num_probs; ++i) {
|
||||
p_top_p[i] = std::max(p_top_p[i], eps_);
|
||||
}
|
||||
return need_top_p;
|
||||
}
|
||||
|
||||
/*! \brief Check whether prob values are needed, and collect info when necessary. */
|
||||
bool CheckProbValues(const Array<GenerationConfig>& generation_cfg,
|
||||
const std::vector<int>& sample_indices, int num_probs, int num_samples,
|
||||
int vocab_size, std::vector<int>* top_prob_offset_indptr) {
|
||||
top_prob_offset_indptr->reserve(num_samples + 1);
|
||||
top_prob_offset_indptr->push_back(0);
|
||||
int* p_top_prob_offsets = static_cast<int*>(top_prob_offsets_host_->data);
|
||||
int num_top_probs = 0;
|
||||
bool need_prob_values = false;
|
||||
for (int i = 0; i < num_samples; ++i) {
|
||||
need_prob_values |= generation_cfg[i]->logprobs;
|
||||
for (int j = 0; j < generation_cfg[i]->top_logprobs; ++j) {
|
||||
p_top_prob_offsets[num_top_probs++] = sample_indices[i] * vocab_size + j;
|
||||
}
|
||||
top_prob_offset_indptr->push_back(top_prob_offset_indptr->back() +
|
||||
generation_cfg[i]->top_logprobs);
|
||||
}
|
||||
TVM_FFI_ICHECK_EQ(num_top_probs, top_prob_offset_indptr->back());
|
||||
return need_prob_values;
|
||||
}
|
||||
|
||||
/*! \brief Sample tokens on GPU. Take out the probability values when needed. */
|
||||
std::vector<Tensor> SampleOnGPU(Tensor probs_on_device, Tensor uniform_samples_device,
|
||||
Tensor sample_indices_device, //
|
||||
bool need_top_p, bool need_prob_values, int num_probs,
|
||||
const std::vector<int>& top_prob_offset_indptr) {
|
||||
Tensor sampled_token_ids_device{nullptr};
|
||||
Tensor sampled_probs_device{nullptr};
|
||||
Tensor top_prob_probs_device{nullptr};
|
||||
Tensor top_prob_indices_device{nullptr};
|
||||
|
||||
if (!need_top_p && !need_prob_values) {
|
||||
// - Short path: If top_p and prob values are not needed, we directly sample from multinomial.
|
||||
SyncCopyStream(device_, compute_stream_, copy_stream_);
|
||||
if (flashinfer_sampling_available_) {
|
||||
sampled_token_ids_device =
|
||||
sampled_token_ids_device_.CreateView({sample_indices_device->shape[0]}, dtype_i32_);
|
||||
flashinfer_multinomial_sample_func_.value()(probs_on_device, uniform_samples_device,
|
||||
sample_indices_device,
|
||||
sampled_token_ids_device);
|
||||
} else {
|
||||
sampled_token_ids_device =
|
||||
gpu_multinomial_from_uniform_func_(probs_on_device, uniform_samples_device,
|
||||
sample_indices_device)
|
||||
.cast<Tensor>();
|
||||
}
|
||||
return {sampled_token_ids_device, sampled_probs_device, top_prob_probs_device,
|
||||
top_prob_indices_device};
|
||||
}
|
||||
|
||||
// - Argsort the probability.
|
||||
Array<Tensor> argsort_results = gpu_argsort_probs_func_(probs_on_device).cast<Array<Tensor>>();
|
||||
TVM_FFI_ICHECK_EQ(argsort_results.size(), 2);
|
||||
Tensor sorted_probs_on_device = argsort_results[0];
|
||||
Tensor sorted_indices_on_device = argsort_results[1];
|
||||
|
||||
// - Copy auxiliary array for top-p and prob values in ahead.
|
||||
Tensor top_p_device;
|
||||
Tensor top_prob_offsets_device;
|
||||
if (need_top_p) {
|
||||
Tensor top_p_host = top_p_host_.CreateView({num_probs}, dtype_f32_);
|
||||
top_p_device = top_p_device_.CreateView({num_probs}, dtype_f32_);
|
||||
CopyArray(/*src=*/top_p_host, /*dst=*/top_p_device, copy_stream_);
|
||||
}
|
||||
if (need_prob_values) {
|
||||
int num_top_probs = top_prob_offset_indptr.back();
|
||||
Tensor top_prob_offsets_host = top_prob_offsets_host_.CreateView({num_top_probs}, dtype_i32_);
|
||||
top_prob_offsets_device = top_prob_offsets_device_.CreateView({num_top_probs}, dtype_i32_);
|
||||
CopyArray(/*src=*/top_prob_offsets_host, /*dst=*/top_prob_offsets_device, copy_stream_);
|
||||
}
|
||||
SyncCopyStream(device_, compute_stream_, copy_stream_);
|
||||
|
||||
if (need_top_p) {
|
||||
// - Sample with top_p applied.
|
||||
sampled_token_ids_device =
|
||||
gpu_sample_with_top_p_func_(sorted_probs_on_device, sorted_indices_on_device,
|
||||
uniform_samples_device, sample_indices_device, top_p_device)
|
||||
.cast<Tensor>();
|
||||
} else {
|
||||
// - Sample without top_p.
|
||||
if (flashinfer_sampling_available_) {
|
||||
sampled_token_ids_device =
|
||||
sampled_token_ids_device_.CreateView({sample_indices_device->shape[0]}, dtype_i32_);
|
||||
flashinfer_multinomial_sample_func_
|
||||
.value()(probs_on_device, uniform_samples_device, sample_indices_device,
|
||||
sampled_token_ids_device)
|
||||
.cast<Tensor>();
|
||||
} else {
|
||||
sampled_token_ids_device =
|
||||
gpu_multinomial_from_uniform_func_(probs_on_device, uniform_samples_device,
|
||||
sample_indices_device)
|
||||
.cast<Tensor>();
|
||||
}
|
||||
}
|
||||
|
||||
if (need_prob_values) {
|
||||
// - Take the probability values.
|
||||
Array<Tensor> prob_value_results =
|
||||
gpu_sampler_take_probs_func_(probs_on_device, sorted_indices_on_device,
|
||||
sample_indices_device, sampled_token_ids_device,
|
||||
top_prob_offsets_device)
|
||||
.cast<Array<Tensor>>();
|
||||
sampled_probs_device = prob_value_results[0];
|
||||
top_prob_probs_device = prob_value_results[1];
|
||||
top_prob_indices_device = prob_value_results[2];
|
||||
}
|
||||
|
||||
return {sampled_token_ids_device, sampled_probs_device, top_prob_probs_device,
|
||||
top_prob_indices_device};
|
||||
}
|
||||
|
||||
/*! \brief Copy the results of GPU sampling functions back to CPU. */
|
||||
std::vector<Tensor> CopyArraysToCPU(const std::vector<Tensor>& device_arrays, //
|
||||
int num_samples, bool need_prob_values, int num_top_probs) {
|
||||
Tensor sampled_token_ids_device = device_arrays[0];
|
||||
Tensor sampled_probs_device = device_arrays[1];
|
||||
Tensor top_prob_probs_device = device_arrays[2];
|
||||
Tensor top_prob_indices_device = device_arrays[3];
|
||||
TVM_FFI_ICHECK(sampled_token_ids_device.defined());
|
||||
TVM_FFI_ICHECK_EQ(sampled_token_ids_device->ndim, 1);
|
||||
TVM_FFI_ICHECK_EQ(sampled_token_ids_device->shape[0], num_samples);
|
||||
Tensor sampled_token_ids_host = sampled_token_ids_host_.CreateView({num_samples}, dtype_i32_);
|
||||
CopyArray(/*src=*/sampled_token_ids_device, /*dst=*/sampled_token_ids_host, compute_stream_);
|
||||
|
||||
Tensor sampled_probs_host{nullptr};
|
||||
Tensor top_prob_probs_host{nullptr};
|
||||
Tensor top_prob_indices_host{nullptr};
|
||||
if (need_prob_values) {
|
||||
TVM_FFI_ICHECK(sampled_probs_device.defined());
|
||||
TVM_FFI_ICHECK(top_prob_probs_device.defined());
|
||||
TVM_FFI_ICHECK(top_prob_indices_device.defined());
|
||||
TVM_FFI_ICHECK_EQ(sampled_probs_device->ndim, 1);
|
||||
TVM_FFI_ICHECK_EQ(top_prob_probs_device->ndim, 1);
|
||||
TVM_FFI_ICHECK_EQ(top_prob_indices_device->ndim, 1);
|
||||
TVM_FFI_ICHECK_EQ(sampled_probs_device->shape[0], num_samples);
|
||||
TVM_FFI_ICHECK_EQ(top_prob_probs_device->shape[0], num_top_probs);
|
||||
TVM_FFI_ICHECK_EQ(top_prob_indices_device->shape[0], num_top_probs);
|
||||
sampled_probs_host = sampled_probs_host_.CreateView({num_samples}, dtype_i32_);
|
||||
top_prob_probs_host = top_prob_probs_host_.CreateView({num_top_probs}, dtype_f32_);
|
||||
top_prob_indices_host = top_prob_indices_host_.CreateView({num_top_probs}, dtype_i32_);
|
||||
CopyArray(/*src=*/sampled_probs_device, /*dst=*/sampled_probs_host, compute_stream_);
|
||||
if (num_top_probs > 0) {
|
||||
CopyArray(/*src=*/top_prob_probs_device, /*dst=*/top_prob_probs_host, compute_stream_);
|
||||
CopyArray(/*src=*/top_prob_indices_device, /*dst=*/top_prob_indices_host, compute_stream_);
|
||||
}
|
||||
}
|
||||
|
||||
// Synchronize for CPU to get the correct array results.
|
||||
DeviceAPI::Get(device_)->StreamSync(device_, compute_stream_);
|
||||
|
||||
return {sampled_token_ids_host, sampled_probs_host, top_prob_probs_host, top_prob_indices_host};
|
||||
}
|
||||
|
||||
// Model configurations
|
||||
const int max_num_sample_;
|
||||
const int vocab_size_;
|
||||
const DLDataType dtype_i32_ = DLDataType{kDLInt, 32, 1};
|
||||
const DLDataType dtype_f32_ = DLDataType{kDLFloat, 32, 1};
|
||||
const bool flashinfer_sampling_available_;
|
||||
// Functions for sampling on GPU.
|
||||
Device device_;
|
||||
Function gpu_multinomial_from_uniform_func_;
|
||||
Function gpu_argsort_probs_func_;
|
||||
Function gpu_sample_with_top_p_func_;
|
||||
Function gpu_sampler_take_probs_func_;
|
||||
Function gpu_verify_draft_tokens_func_;
|
||||
Function gpu_renormalize_by_top_p_func_;
|
||||
Optional<Function> flashinfer_multinomial_sample_func_;
|
||||
// Auxiliary Tensors on CPU
|
||||
Tensor uniform_samples_host_;
|
||||
Tensor sample_indices_host_;
|
||||
Tensor top_p_host_;
|
||||
Tensor top_p_init_pivots_host_;
|
||||
Tensor top_prob_offsets_host_;
|
||||
Tensor draft_tokens_host_;
|
||||
Tensor token_tree_first_child_host_;
|
||||
Tensor token_tree_next_sibling_host_;
|
||||
Tensor token_tree_parent_ptr_host_;
|
||||
Tensor sampled_token_ids_host_;
|
||||
Tensor sampled_probs_host_;
|
||||
Tensor top_prob_probs_host_;
|
||||
Tensor top_prob_indices_host_;
|
||||
// Auxiliary Tensors on GPU
|
||||
Tensor uniform_samples_device_;
|
||||
Tensor sample_indices_device_;
|
||||
Tensor top_p_device_;
|
||||
Tensor top_p_init_pivots_device_;
|
||||
Tensor top_prob_offsets_device_;
|
||||
Tensor draft_tokens_device_;
|
||||
Tensor token_tree_first_child_device_;
|
||||
Tensor token_tree_next_sibling_device_;
|
||||
Tensor token_tree_parent_ptr_device_;
|
||||
Tensor sampled_token_ids_device_;
|
||||
// The event trace recorder for requests. */
|
||||
Optional<EventTraceRecorder> trace_recorder_;
|
||||
// The device stream for the default computation operations.
|
||||
TVMStreamHandle compute_stream_ = nullptr;
|
||||
// The device stream for copying auxiliary data structure to GPU.
|
||||
TVMStreamHandle copy_stream_ = nullptr;
|
||||
const float eps_ = 1e-5;
|
||||
const int num_top_p_cutoff_pivots_ = 3;
|
||||
};
|
||||
|
||||
Sampler Sampler::CreateGPUSampler(int max_num_sample, int vocab_size, FunctionTable* ft,
|
||||
DLDevice device, Optional<EventTraceRecorder> trace_recorder) {
|
||||
return Sampler(tvm::ffi::make_object<GPUSampler>(max_num_sample, vocab_size, ft, device,
|
||||
std::move(trace_recorder)));
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,165 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/sampler/sampler.h
|
||||
* \brief The header for runtime module of sampler functions.
|
||||
*/
|
||||
|
||||
#ifndef MLC_LLM_SERVE_SAMPLER_SAMPLER_H_
|
||||
#define MLC_LLM_SERVE_SAMPLER_SAMPLER_H_
|
||||
|
||||
#include <tvm/ffi/extra/module.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
|
||||
#include "../../base.h"
|
||||
#include "../../support/random.h"
|
||||
#include "../data.h"
|
||||
#include "../event_trace_recorder.h"
|
||||
#include "../model.h"
|
||||
#include "../request_state.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::Device;
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
|
||||
/*!
|
||||
* \brief The base class of runtime sampler.
|
||||
* Its main function is `BatchSampleTokensWithProbBeforeTopP`, which takes a batch of
|
||||
* logits and corresponding configuration, and sample one token
|
||||
* for each instance of the batch.
|
||||
*/
|
||||
class SamplerObj : public Object {
|
||||
public:
|
||||
/*!
|
||||
* \brief Renormalize the input batch of probability distributions with top p values.
|
||||
* \param probs_on_device The batch of prob distributions before normalization.
|
||||
* \param sample_indices Specifying which request we will sample for
|
||||
* in i-th output for the sampling later on.
|
||||
* The output result of the sampling will be as follow:
|
||||
* result[i] = sample_from(prob_on_device[sample_indices[i],:], generation_config[i]));
|
||||
* For renormalization, the sample indices are used for determine the top-p grouping.
|
||||
* \param request_ids The id of each request.
|
||||
* \param generation_cfg The generation config of each request in the input batch.
|
||||
* \return The renormalized probability distributions, residing on device
|
||||
* if the sampler is GPU sampler, or on host if the sampler is CPU sampler.
|
||||
*/
|
||||
virtual Tensor BatchRenormalizeProbsByTopP(Tensor probs_on_device, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<String>& request_ids, //
|
||||
const Array<GenerationConfig>& generation_cfg) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Sample tokens from the input batch of prob distribution on device.
|
||||
* The input prob distributions are not yet applied with top-p.
|
||||
* \param probs_on_device The prob distributions on GPU to sample tokens from.
|
||||
* \param sample_indices Specifying which request we should sample for
|
||||
* in i-th output. The output result is sample as follow:
|
||||
* result[i] = sample_from(prob_on_device[sample_indices[i],:], generation_config[i]));
|
||||
* \param request_ids The id of each request.
|
||||
* \param generation_cfg The generation config of each request
|
||||
* in the input batch.
|
||||
* \param rngs The random number generator of each sequence.
|
||||
* \return The batch of sampling results, which contain the sampled token id
|
||||
* and other probability info.
|
||||
*/
|
||||
virtual std::vector<SampleResult> BatchSampleTokensWithProbBeforeTopP(
|
||||
Tensor probs_on_device, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<String>& request_ids, //
|
||||
const Array<GenerationConfig>& generation_cfg, //
|
||||
const std::vector<RandomGenerator*>& rngs) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Sample tokens from the input batch of prob distribution on device.
|
||||
* The input prob distributions are already applied with top-p.
|
||||
* \param probs The prob distributions.
|
||||
* It resides on GPU if the sampler is GPU sampler, or on host if hte sampler is CPU sampler.
|
||||
* \param sample_indices Specifying which request we should sample for
|
||||
* in i-th output. The output result is sample as follow:
|
||||
* result[i] = sample_from(prob_on_device[sample_indices[i],:], generation_config[i]));
|
||||
* \param request_ids The id of each request.
|
||||
* \param generation_cfg The generation config of each request
|
||||
* in the input batch.
|
||||
* \param rngs The random number generator of each sequence.
|
||||
* \return The batch of sampling results, which contain the sampled token id
|
||||
* and other probability info.
|
||||
*/
|
||||
virtual std::vector<SampleResult> BatchSampleTokensWithProbAfterTopP(
|
||||
Tensor probs, //
|
||||
const std::vector<int>& sample_indices, //
|
||||
const Array<String>& request_ids, //
|
||||
const Array<GenerationConfig>& generation_cfg, //
|
||||
const std::vector<RandomGenerator*>& rngs) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Verify draft tokens generated by small models in the large model
|
||||
* in speculative decoding. The input corresponds to a batch of sequences.
|
||||
* The input prob distributions are already applied with top-p.
|
||||
* \param probs The prob distributions on GPU to sample tokens from.
|
||||
* It resides on GPU if the sampler is GPU sampler, or on host if hte sampler is CPU sampler.
|
||||
* \param request_ids The id of each request.
|
||||
* \param cum_verify_lengths The cumulative draft lengths to verify of all sequences.
|
||||
* \param generation_cfg The generation config of each request
|
||||
* in the input batch.
|
||||
* \param rngs The random number generator of each sequence.
|
||||
* \param draft_output_tokens The draft tokens generated by the small model for
|
||||
* each sequence.
|
||||
* \param token_tree_parent_ptr The parent pointer of the token tree.
|
||||
* \param draft_probs_on_device The probability distribution computed from the
|
||||
* small model for each sequence. Concatenated tensor of shape (total_verify_length, vocab_size).
|
||||
* It includes the slot for the last committed token that has undefined probablity value.
|
||||
* \return The list of accepted tokens for each request and the index of the last accepted tree
|
||||
* node for each request.
|
||||
*/
|
||||
virtual std::pair<std::vector<std::vector<SampleResult>>, std::vector<int>>
|
||||
BatchVerifyDraftTokensWithProbAfterTopP(
|
||||
Tensor probs, const Array<String>& request_ids, const std::vector<int>& cum_verify_lengths,
|
||||
const Array<GenerationConfig>& generation_cfg, const std::vector<RandomGenerator*>& rngs,
|
||||
const std::vector<std::vector<SampleResult>>& draft_output_tokens,
|
||||
const std::vector<int64_t>& token_tree_parent_ptr, Tensor draft_probs_on_device) = 0;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<SamplerObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.Sampler", SamplerObj, Object);
|
||||
};
|
||||
|
||||
class Sampler : public ObjectRef {
|
||||
public:
|
||||
/*! * \brief Create a CPU sampler. */
|
||||
static Sampler CreateCPUSampler(Optional<EventTraceRecorder> trace_recorder);
|
||||
/*!
|
||||
* \brief Create a GPU sampler.
|
||||
* \param max_num_sample The max number of samples to sample at a time.
|
||||
* \param vocab_size The model's vocabulary size.
|
||||
* \param ft The packed function table.
|
||||
* \param device The device that the model runs on.
|
||||
* \param trace_recorder The event trace recorder.
|
||||
*/
|
||||
static Sampler CreateGPUSampler(int max_num_sample, int vocab_size, FunctionTable* ft,
|
||||
DLDevice device, Optional<EventTraceRecorder> trace_recorder);
|
||||
|
||||
/*! \brief Check if the given device supports GPU sampling. */
|
||||
static bool SupportGPUSampler(Device device) {
|
||||
return device.device_type == DLDeviceType::kDLCUDA ||
|
||||
device.device_type == DLDeviceType::kDLVulkan ||
|
||||
device.device_type == DLDeviceType::kDLMetal;
|
||||
}
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Sampler, ObjectRef, SamplerObj);
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_SAMPLER_SAMPLER_H_
|
||||
@@ -0,0 +1,417 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/threaded_engine.cc
|
||||
* \brief The implementation for threaded serving engine in MLC LLM.
|
||||
*/
|
||||
#include "threaded_engine.h"
|
||||
|
||||
#include <tvm/ffi/extra/module.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
|
||||
#include <atomic>
|
||||
#include <condition_variable>
|
||||
#include <mutex>
|
||||
#include <optional>
|
||||
|
||||
#include "../support/json_parser.h"
|
||||
#include "../support/module_vtable.h"
|
||||
#include "../support/result.h"
|
||||
#include "engine.h"
|
||||
#include "request.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using tvm::Device;
|
||||
using namespace tvm::runtime;
|
||||
|
||||
/*! \brief The threaded engine instruction kind. */
|
||||
enum class InstructionKind : int {
|
||||
kAddRequest = 0,
|
||||
kAbortRequest = 1,
|
||||
kUnloadEngine = 2,
|
||||
kReloadEngine = 3,
|
||||
kResetEngine = 4,
|
||||
kDebugCallFuncOnAllAllWorker = 5,
|
||||
};
|
||||
|
||||
/*! \brief The implementation of ThreadedEngine. */
|
||||
class ThreadedEngineImpl : public ThreadedEngine {
|
||||
public:
|
||||
void InitThreadedEngine(Device device, Optional<Function> request_stream_callback,
|
||||
Optional<EventTraceRecorder> trace_recorder) final {
|
||||
device_ = device;
|
||||
TVM_FFI_ICHECK(request_stream_callback.has_value())
|
||||
<< "ThreadedEngine requires request stream callback function, but it is not given.";
|
||||
request_stream_callback_ = request_stream_callback.value();
|
||||
trace_recorder_ = trace_recorder;
|
||||
}
|
||||
|
||||
void Reload(String engine_config_json_str) final {
|
||||
// NOTE: important to set this before, we send out
|
||||
// reload instruction to the other threads
|
||||
// otherwise there can be deadlocks
|
||||
reload_finished_ = false;
|
||||
bool need_notify = false;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
||||
instruction_queue_.emplace_back(InstructionKind::kReloadEngine,
|
||||
std::move(engine_config_json_str));
|
||||
++pending_request_operation_cnt_;
|
||||
need_notify = engine_waiting_;
|
||||
}
|
||||
if (need_notify) {
|
||||
background_loop_cv_.notify_one();
|
||||
}
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(reload_unload_mutex_);
|
||||
reload_unload_cv_.wait(lock, [this] { return reload_finished_; });
|
||||
}
|
||||
}
|
||||
|
||||
void Unload() final {
|
||||
// NOTE: important to set this before, we send out
|
||||
// reload instruction to the other threads
|
||||
// otherwise there can be deadlocks
|
||||
// e.g. the other thread finish unload job and set the flag to true
|
||||
// then we set it back to false
|
||||
unload_finished_ = false;
|
||||
bool need_notify = false;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
||||
instruction_queue_.emplace_back(InstructionKind::kUnloadEngine, ObjectRef(nullptr));
|
||||
++pending_request_operation_cnt_;
|
||||
need_notify = engine_waiting_;
|
||||
}
|
||||
if (need_notify) {
|
||||
background_loop_cv_.notify_one();
|
||||
}
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(reload_unload_mutex_);
|
||||
reload_unload_cv_.wait(lock, [this] { return unload_finished_; });
|
||||
}
|
||||
}
|
||||
|
||||
void Reset() final {
|
||||
bool need_notify = false;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
||||
instruction_queue_.emplace_back(InstructionKind::kResetEngine, ObjectRef(nullptr));
|
||||
++pending_request_operation_cnt_;
|
||||
need_notify = engine_waiting_;
|
||||
}
|
||||
if (need_notify) {
|
||||
background_loop_cv_.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
void AddRequest(Request request) final {
|
||||
bool need_notify = false;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
||||
instruction_queue_.emplace_back(InstructionKind::kAddRequest, request);
|
||||
++pending_request_operation_cnt_;
|
||||
need_notify = engine_waiting_;
|
||||
}
|
||||
if (need_notify) {
|
||||
background_loop_cv_.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
void AbortRequest(const String& request_id) final {
|
||||
bool need_notify = false;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
||||
instruction_queue_.emplace_back(InstructionKind::kAbortRequest, request_id);
|
||||
++pending_request_operation_cnt_;
|
||||
need_notify = engine_waiting_;
|
||||
}
|
||||
if (need_notify) {
|
||||
background_loop_cv_.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
void RunBackgroundLoop() final {
|
||||
// The local vectors that load the requests from critical regions.
|
||||
std::vector<std::pair<InstructionKind, Any>> local_instruction_queue;
|
||||
|
||||
while (!exit_now_.load(std::memory_order_relaxed)) {
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(background_loop_mutex_);
|
||||
engine_waiting_ = true;
|
||||
background_loop_cv_.wait(lock, [this] {
|
||||
return (background_engine_ != nullptr && !background_engine_->Empty()) ||
|
||||
pending_request_operation_cnt_.load() > 0 ||
|
||||
exit_now_.load(std::memory_order_relaxed);
|
||||
});
|
||||
engine_waiting_ = false;
|
||||
local_instruction_queue = instruction_queue_;
|
||||
instruction_queue_.clear();
|
||||
pending_request_operation_cnt_ = 0;
|
||||
}
|
||||
for (const auto& [kind, arg] : local_instruction_queue) {
|
||||
if (kind == InstructionKind::kAddRequest) {
|
||||
TVM_FFI_ICHECK(background_engine_ != nullptr) << "Background engine is not loaded.";
|
||||
background_engine_->AddRequest(arg.as_or_throw<Request>());
|
||||
} else if (kind == InstructionKind::kAbortRequest) {
|
||||
// in a rare case, abort request can happen after unloading
|
||||
// aka background engine is nullptr
|
||||
// this happens when the on going generation was interrupted
|
||||
// the engine get unloaded, and then abort was called.
|
||||
// it is safe to ignore these abort in such case
|
||||
if (background_engine_ != nullptr) {
|
||||
background_engine_->AbortRequest(arg.as_or_throw<String>());
|
||||
}
|
||||
} else if (kind == InstructionKind::kUnloadEngine) {
|
||||
EngineUnloadImpl();
|
||||
} else if (kind == InstructionKind::kReloadEngine) {
|
||||
EngineUnloadImpl();
|
||||
EngineReloadImpl(arg.as_or_throw<String>());
|
||||
} else if (kind == InstructionKind::kResetEngine) {
|
||||
if (background_engine_ != nullptr) {
|
||||
background_engine_->Reset();
|
||||
}
|
||||
} else if (kind == InstructionKind::kDebugCallFuncOnAllAllWorker) {
|
||||
TVM_FFI_ICHECK(background_engine_ != nullptr) << "Background engine is not loaded.";
|
||||
Array<Any> packed_args = arg.as_or_throw<Array<Any>>();
|
||||
background_engine_->DebugCallFuncOnAllAllWorker(
|
||||
packed_args[0].as_or_throw<String>(), packed_args[1].as_or_throw<Optional<String>>());
|
||||
} else {
|
||||
LOG(FATAL) << "Cannot reach here";
|
||||
}
|
||||
}
|
||||
if (background_engine_ != nullptr) {
|
||||
background_engine_->Step();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void RunBackgroundStreamBackLoop() final {
|
||||
// The local vectors that load the request stream callback inputs from critical regions.
|
||||
std::vector<Array<RequestStreamOutput>> local_request_stream_callback_inputs;
|
||||
std::vector<RequestStreamOutput> flattened_callback_inputs;
|
||||
|
||||
while (!exit_now_.load(std::memory_order_relaxed)) {
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(request_stream_callback_mutex_);
|
||||
stream_callback_waiting_ = true;
|
||||
request_stream_callback_cv_.wait(lock, [this] {
|
||||
return pending_request_stream_callback_cnt_.load() > 0 ||
|
||||
exit_now_.load(std::memory_order_relaxed);
|
||||
});
|
||||
stream_callback_waiting_ = false;
|
||||
|
||||
local_request_stream_callback_inputs = request_stream_callback_inputs_;
|
||||
request_stream_callback_inputs_.clear();
|
||||
pending_request_stream_callback_cnt_ = 0;
|
||||
}
|
||||
for (const Array<RequestStreamOutput>& callback_inputs :
|
||||
local_request_stream_callback_inputs) {
|
||||
for (const RequestStreamOutput& callback_input : callback_inputs) {
|
||||
flattened_callback_inputs.push_back(callback_input);
|
||||
}
|
||||
}
|
||||
if (!flattened_callback_inputs.empty()) {
|
||||
request_stream_callback_(Array<RequestStreamOutput>(flattened_callback_inputs));
|
||||
}
|
||||
flattened_callback_inputs.clear();
|
||||
}
|
||||
}
|
||||
|
||||
void ExitBackgroundLoop() final {
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
||||
exit_now_.store(true);
|
||||
}
|
||||
background_loop_cv_.notify_one();
|
||||
request_stream_callback_cv_.notify_one();
|
||||
}
|
||||
|
||||
/************** Query/Profile/Debug **************/
|
||||
|
||||
GenerationConfig GetDefaultGenerationConfig() const final {
|
||||
TVM_FFI_ICHECK(default_generation_config_.has_value())
|
||||
<< "The default generation config has not been set.";
|
||||
return default_generation_config_.value();
|
||||
}
|
||||
|
||||
Request CreateRequest(String id, Array<Data> inputs, String generation_cfg_json_str) const {
|
||||
json::Object config = json::ParseToJSONObject(generation_cfg_json_str);
|
||||
auto gen_config = GenerationConfig::FromJSON(config, GetDefaultGenerationConfig());
|
||||
TVM_FFI_ICHECK(gen_config.IsOk()) << gen_config.UnwrapErr();
|
||||
return Request(std::move(id), std::move(inputs), gen_config.Unwrap());
|
||||
}
|
||||
|
||||
EngineConfig GetCompleteEngineConfig() const final {
|
||||
TVM_FFI_ICHECK(complete_engine_config_.has_value()) << "The engine config has not been set.";
|
||||
return complete_engine_config_.value();
|
||||
}
|
||||
|
||||
String GetCompleteEngineConfigJSONString() const {
|
||||
return GetCompleteEngineConfig()->AsJSONString();
|
||||
}
|
||||
|
||||
void DebugCallFuncOnAllAllWorker(const String& func_name, Optional<String> func_args) final {
|
||||
bool need_notify = false;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(background_loop_mutex_);
|
||||
instruction_queue_.emplace_back(InstructionKind::kDebugCallFuncOnAllAllWorker,
|
||||
Array<Any>{func_name, func_args});
|
||||
++pending_request_operation_cnt_;
|
||||
need_notify = engine_waiting_;
|
||||
}
|
||||
if (need_notify) {
|
||||
background_loop_cv_.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
void EngineReloadImpl(const std::string& engine_config_json_str) {
|
||||
auto frequest_stream_callback_wrapper = [this](Array<RequestStreamOutput> delta_outputs) {
|
||||
bool need_notify = false;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(request_stream_callback_mutex_);
|
||||
request_stream_callback_inputs_.push_back(std::move(delta_outputs));
|
||||
++pending_request_stream_callback_cnt_;
|
||||
need_notify = stream_callback_waiting_;
|
||||
}
|
||||
if (need_notify) {
|
||||
request_stream_callback_cv_.notify_one();
|
||||
}
|
||||
};
|
||||
|
||||
FRequestStreamCallback request_stream_callback(frequest_stream_callback_wrapper);
|
||||
Result<EngineCreationOutput> output_res =
|
||||
Engine::Create(engine_config_json_str, device_, request_stream_callback, trace_recorder_);
|
||||
TVM_FFI_ICHECK(output_res.IsOk()) << output_res.UnwrapErr();
|
||||
EngineCreationOutput output = output_res.Unwrap();
|
||||
background_engine_ = std::move(output.reloaded_engine);
|
||||
default_generation_config_ = output.default_generation_cfg;
|
||||
complete_engine_config_ = output.completed_engine_config;
|
||||
{
|
||||
// Wake up the thread waiting for reload finish.
|
||||
std::lock_guard<std::mutex> lock(reload_unload_mutex_);
|
||||
reload_finished_ = true;
|
||||
}
|
||||
reload_unload_cv_.notify_one();
|
||||
}
|
||||
|
||||
void EngineUnloadImpl() {
|
||||
if (background_engine_ != nullptr) {
|
||||
background_engine_->AbortAllRequests();
|
||||
background_engine_ = nullptr;
|
||||
// Clear the allocated memory in cached memory pool.
|
||||
static Function fclear_memory_manager =
|
||||
Function::GetGlobalRequired("vm.builtin.memory_manager.clear");
|
||||
fclear_memory_manager();
|
||||
default_generation_config_ = std::nullopt;
|
||||
complete_engine_config_ = std::nullopt;
|
||||
}
|
||||
{
|
||||
// Wake up the thread waiting for unload finish.
|
||||
std::lock_guard<std::mutex> lock(reload_unload_mutex_);
|
||||
unload_finished_ = true;
|
||||
}
|
||||
reload_unload_cv_.notify_one();
|
||||
}
|
||||
|
||||
/*! \brief The device to run models on. */
|
||||
Device device_;
|
||||
/*! \brief The background normal engine for request processing. */
|
||||
std::unique_ptr<Engine> background_engine_;
|
||||
/*! \brief The request stream callback. */
|
||||
Function request_stream_callback_;
|
||||
/*! \brief Event trace recorder. */
|
||||
Optional<EventTraceRecorder> trace_recorder_;
|
||||
|
||||
/*! \brief complete engine config. */
|
||||
Optional<EngineConfig> complete_engine_config_;
|
||||
/*! \brief The default generation config. */
|
||||
Optional<GenerationConfig> default_generation_config_;
|
||||
|
||||
/*! \brief The mutex ensuring only one thread can access critical regions. */
|
||||
std::mutex background_loop_mutex_;
|
||||
std::mutex request_stream_callback_mutex_;
|
||||
std::mutex reload_unload_mutex_;
|
||||
/*! \brief The condition variable preventing threaded engine from spinning. */
|
||||
std::condition_variable background_loop_cv_;
|
||||
std::condition_variable request_stream_callback_cv_;
|
||||
std::condition_variable reload_unload_cv_;
|
||||
/*! \brief A boolean flag denoting if the engine needs to exit background loop. */
|
||||
std::atomic<bool> exit_now_ = false;
|
||||
|
||||
/************** Critical Regions **************/
|
||||
/*!
|
||||
* \brief The instruction queue for the threaded engine.
|
||||
* The instructions include:
|
||||
* - requests to add into the background engine,
|
||||
* - requests to abort from the background engine,
|
||||
* - engine unload/reload,
|
||||
* - and other debugging instructions.
|
||||
* Elements are sended from other threads and consumed by
|
||||
* the threaded engine in the background loop.
|
||||
*/
|
||||
std::vector<std::pair<InstructionKind, Any>> instruction_queue_;
|
||||
/*!
|
||||
* \brief The delta outputs to pass through callback.
|
||||
* Elements are sended from the background loop thread and
|
||||
* consumed by the foreground thread.
|
||||
*/
|
||||
std::vector<Array<RequestStreamOutput>> request_stream_callback_inputs_;
|
||||
/*!
|
||||
* \brief Number of pending request operations, should be the size of
|
||||
* `requests_to_add_` and `requests_to_abort_`.
|
||||
*/
|
||||
std::atomic<int> pending_request_operation_cnt_ = 0;
|
||||
/*!
|
||||
* \brief Number of pending request stream callback invocations.
|
||||
* It should be the size of `request_stream_callback_inputs_`.
|
||||
*/
|
||||
std::atomic<int> pending_request_stream_callback_cnt_ = 0;
|
||||
/*! \brief A boolean flag indicating if the engine is waiting for new requests/aborts. */
|
||||
bool engine_waiting_ = false;
|
||||
/*! \brief A boolean flag indicating if the stream callback loop is waiting. */
|
||||
bool stream_callback_waiting_ = false;
|
||||
/*! \brief A boolean indicating if the engine reload has finished. */
|
||||
bool reload_finished_ = false;
|
||||
/*! \brief A boolean indicating if the engine unload has finished. */
|
||||
bool unload_finished_ = false;
|
||||
};
|
||||
|
||||
/*! \brief The implementation of ThreadedEngine. */
|
||||
class ThreadedEngineModule : public ThreadedEngineImpl, public ffi::ModuleObj {
|
||||
public:
|
||||
TVM_MODULE_VTABLE_BEGIN("mlc.serve.async_threaded_engine");
|
||||
TVM_MODULE_VTABLE_ENTRY("init_threaded_engine", &ThreadedEngineImpl::InitThreadedEngine);
|
||||
TVM_MODULE_VTABLE_ENTRY("reload", &ThreadedEngineImpl::Reload);
|
||||
TVM_MODULE_VTABLE_ENTRY("add_request", &ThreadedEngineImpl::AddRequest);
|
||||
TVM_MODULE_VTABLE_ENTRY("create_request", &ThreadedEngineImpl::CreateRequest);
|
||||
TVM_MODULE_VTABLE_ENTRY("abort_request", &ThreadedEngineImpl::AbortRequest);
|
||||
TVM_MODULE_VTABLE_ENTRY("run_background_loop", &ThreadedEngineImpl::RunBackgroundLoop);
|
||||
TVM_MODULE_VTABLE_ENTRY("run_background_stream_back_loop",
|
||||
&ThreadedEngineImpl::RunBackgroundStreamBackLoop);
|
||||
TVM_MODULE_VTABLE_ENTRY("exit_background_loop", &ThreadedEngineImpl::ExitBackgroundLoop);
|
||||
TVM_MODULE_VTABLE_ENTRY("get_complete_engine_config",
|
||||
&ThreadedEngineImpl::GetCompleteEngineConfigJSONString);
|
||||
TVM_MODULE_VTABLE_ENTRY("reset", &ThreadedEngineImpl::Reset);
|
||||
TVM_MODULE_VTABLE_ENTRY("debug_call_func_on_all_worker",
|
||||
&ThreadedEngineImpl::DebugCallFuncOnAllAllWorker);
|
||||
TVM_MODULE_VTABLE_END();
|
||||
};
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("mlc.serve.create_threaded_engine",
|
||||
[]() { return Module(tvm::ffi::make_object<ThreadedEngineModule>()); });
|
||||
}
|
||||
|
||||
std::unique_ptr<ThreadedEngine> ThreadedEngine::Create() {
|
||||
std::unique_ptr<ThreadedEngineImpl> threaded_engine = std::make_unique<ThreadedEngineImpl>();
|
||||
return std::move(threaded_engine);
|
||||
}
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,90 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file serve/threaded_engine.h
|
||||
* \brief The header of threaded serving engine in MLC LLM.
|
||||
*/
|
||||
#ifndef MLC_LLM_SERVE_THREADED_ENGINE_H_
|
||||
#define MLC_LLM_SERVE_THREADED_ENGINE_H_
|
||||
|
||||
#include "data.h"
|
||||
#include "engine.h"
|
||||
#include "request.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace serve {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
|
||||
/*!
|
||||
* \brief The interface threaded engine in MLC LLM.
|
||||
* The threaded engine keeps running a background request processing
|
||||
* loop on a standalone thread. Ensuring thread safety, it exposes
|
||||
* `AddRequest` and `AbortRequest` to receive new requests or
|
||||
* abortions from other threads, and the internal request processing
|
||||
* is backed by a normal engine wrapped inside.
|
||||
*/
|
||||
class ThreadedEngine {
|
||||
public:
|
||||
/*! \brief Create a ThreadedEngine. */
|
||||
static std::unique_ptr<ThreadedEngine> Create();
|
||||
|
||||
virtual ~ThreadedEngine() = default;
|
||||
|
||||
/*!
|
||||
* \brief Initialize the threaded engine from packed arguments in PackedArgs.
|
||||
* \param device The device where to run models.
|
||||
* \param request_stream_callback The request stream callback function to.
|
||||
* \param trace_recorder Event trace recorder for requests.
|
||||
*/
|
||||
virtual void InitThreadedEngine(Device device, Optional<Function> request_stream_callback,
|
||||
Optional<EventTraceRecorder> trace_recorder) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Reload the engine with the new engine config.
|
||||
* \param engine_config_json_str The engine config JSON string.
|
||||
*/
|
||||
virtual void Reload(String engine_config_json_str) = 0;
|
||||
|
||||
/*! \brief Unload the background engine. */
|
||||
virtual void Unload() = 0;
|
||||
|
||||
/*! \brief Reset the engine to the initial state. */
|
||||
virtual void Reset() = 0;
|
||||
|
||||
/*! \brief Starts the background request processing loop. */
|
||||
virtual void RunBackgroundLoop() = 0;
|
||||
|
||||
/*! \brief Starts the request stream callback loop. */
|
||||
virtual void RunBackgroundStreamBackLoop() = 0;
|
||||
|
||||
/*!
|
||||
* \brief Notify the ThreadedEngine to exit the background
|
||||
* request processing loop. This method is invoked by threads
|
||||
* other than the engine-driving thread.
|
||||
*/
|
||||
virtual void ExitBackgroundLoop() = 0;
|
||||
|
||||
/*! \brief Add a new request to the engine. */
|
||||
virtual void AddRequest(Request request) = 0;
|
||||
|
||||
/*! \brief Abort the input request (specified by id string) from engine. */
|
||||
virtual void AbortRequest(const String& request_id) = 0;
|
||||
|
||||
/************** Query/Profile/Debug **************/
|
||||
|
||||
/*! \brief Return the default generation config. */
|
||||
virtual GenerationConfig GetDefaultGenerationConfig() const = 0;
|
||||
|
||||
/*! \brief Return the complete engine config. */
|
||||
virtual EngineConfig GetCompleteEngineConfig() const = 0;
|
||||
|
||||
/*! \brief Call the given global function on all workers. Only for debug purpose. */
|
||||
virtual void DebugCallFuncOnAllAllWorker(const String& func_name, Optional<String> func_args) = 0;
|
||||
};
|
||||
|
||||
} // namespace serve
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SERVE_THREADED_ENGINE_H_
|
||||
@@ -0,0 +1,37 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file support/debug_utils.h
|
||||
* \brief Tools for debug purposes.
|
||||
*/
|
||||
#ifndef MLC_LLM_SUPPORT_DEBUG_UTILS_H_
|
||||
#define MLC_LLM_SUPPORT_DEBUG_UTILS_H_
|
||||
|
||||
#include "../tokenizers/tokenizers.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
/*! \brief A registry for debug information. */
|
||||
class DebugRegistry {
|
||||
public:
|
||||
static DebugRegistry* Global() {
|
||||
static DebugRegistry reg;
|
||||
return ®
|
||||
}
|
||||
|
||||
// Tokenizer information, helpful for converting token id to token string in debugging
|
||||
Tokenizer tokenizer;
|
||||
};
|
||||
|
||||
/*! \brief Register the tokenizer to the global tokenizer registry. */
|
||||
inline void DebugRegisterTokenizer(const Tokenizer& tokenizer) {
|
||||
DebugRegistry::Global()->tokenizer = tokenizer;
|
||||
}
|
||||
|
||||
/*! \brief Get the registered tokenizer from the global tokenizer registry. */
|
||||
inline Tokenizer DebugGetTokenizer() { return DebugRegistry::Global()->tokenizer; }
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_DEBUG_UTILS_H_
|
||||
@@ -0,0 +1,147 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file support/dynamic_bitset.h
|
||||
* \brief The header for utilities used in grammar-guided generation.
|
||||
*/
|
||||
#ifndef MLC_LLM_SUPPORT_DYNAMIC_BITSET_H_
|
||||
#define MLC_LLM_SUPPORT_DYNAMIC_BITSET_H_
|
||||
|
||||
#include <tvm/runtime/logging.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <vector>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
/*!
|
||||
* \brief A bitset whose length is specified at runtime. Note the size cannot be changed after
|
||||
* construction.
|
||||
* \details The buffer of the bitset is a uint32_t array. There are two uses for this class:
|
||||
* - When passing nullptr to data, it maintains an internal buffer for the bitset.
|
||||
* - When passing a pointer to a buffer with enough size, it uses the external buffer for the
|
||||
* bitset.
|
||||
*/
|
||||
class DynamicBitset {
|
||||
public:
|
||||
/*!
|
||||
* \brief Calculate the minimal size of the uint32_t buffer for the bitset with the given size.
|
||||
* \param element_size The size of the bitset.
|
||||
* \return The minimal buffer size.
|
||||
*/
|
||||
static int CalculateBufferSize(int element_size) { return (element_size + 31) / 32; }
|
||||
|
||||
/*!
|
||||
* \brief Construct a empty bitset. This object should be assigned to a valid bitset before using.
|
||||
*/
|
||||
DynamicBitset() : size_(0), buffer_size_(0), data_(nullptr), is_internal_(true) {}
|
||||
|
||||
/*!
|
||||
* \brief Construct a bitset with the given size.
|
||||
* \param size The size of the bitset.
|
||||
* \param data The buffer for the bitset. If nullptr, the bitset will maintain an internal buffer.
|
||||
*/
|
||||
DynamicBitset(int size, uint32_t* data = nullptr)
|
||||
: size_(size), buffer_size_(CalculateBufferSize(size)) {
|
||||
if (data == nullptr) {
|
||||
internal_buffer_.resize(buffer_size_, 0);
|
||||
data_ = internal_buffer_.data();
|
||||
is_internal_ = true;
|
||||
} else {
|
||||
data_ = data;
|
||||
is_internal_ = false;
|
||||
}
|
||||
}
|
||||
|
||||
/*! \brief Copy assignment. */
|
||||
DynamicBitset& operator=(const DynamicBitset& other) {
|
||||
TVM_FFI_DCHECK(is_internal_ || size_ >= other.size_)
|
||||
<< "Expanding bitset size is not allowed when the "
|
||||
"memory of the bitset is externally managed";
|
||||
size_ = other.size_;
|
||||
buffer_size_ = other.buffer_size_;
|
||||
if (is_internal_) {
|
||||
internal_buffer_.reserve(buffer_size_);
|
||||
data_ = internal_buffer_.data();
|
||||
}
|
||||
if (data_ != other.data_) {
|
||||
std::memcpy(data_, other.data_, buffer_size_ * sizeof(uint32_t));
|
||||
}
|
||||
return *this;
|
||||
}
|
||||
|
||||
/*! \brief Move assignment. */
|
||||
DynamicBitset& operator=(DynamicBitset&& other) {
|
||||
size_ = other.size_;
|
||||
buffer_size_ = other.buffer_size_;
|
||||
is_internal_ = other.is_internal_;
|
||||
if (is_internal_) {
|
||||
internal_buffer_ = std::move(other.internal_buffer_);
|
||||
data_ = internal_buffer_.data();
|
||||
} else {
|
||||
data_ = other.data_;
|
||||
}
|
||||
return *this;
|
||||
}
|
||||
|
||||
/*! \brief Get the value of the bit at the given index. */
|
||||
bool operator[](int index) const {
|
||||
TVM_FFI_DCHECK(data_ && index >= 0 && index < size_);
|
||||
return (data_[index / 32] >> (index % 32)) & 1;
|
||||
}
|
||||
|
||||
/*! \brief Get the size of the bitset. */
|
||||
int Size() const { return size_; }
|
||||
|
||||
/*! \brief Set the whole bitset to true. */
|
||||
void Set() {
|
||||
TVM_FFI_DCHECK(data_);
|
||||
std::memset(data_, 0xFF, buffer_size_ * sizeof(uint32_t));
|
||||
}
|
||||
|
||||
/*! \brief Set the bit at the given index to the given value. */
|
||||
void Set(int index, bool value = true) {
|
||||
TVM_FFI_DCHECK(data_ && index >= 0 && index < size_);
|
||||
if (value) {
|
||||
data_[index / 32] |= 1 << (index % 32);
|
||||
} else {
|
||||
data_[index / 32] &= ~(1 << (index % 32));
|
||||
}
|
||||
}
|
||||
|
||||
/*! \brief Set the whole bitset to false. */
|
||||
void Reset() {
|
||||
TVM_FFI_DCHECK(data_);
|
||||
std::memset(data_, 0, buffer_size_ * sizeof(uint32_t));
|
||||
}
|
||||
|
||||
/*! \brief Set the bit at the given index to false. */
|
||||
void Reset(int index) { Set(index, false); }
|
||||
|
||||
/*! \brief Perform a bitwise OR operation between the current bitset and another bitset. */
|
||||
DynamicBitset& operator|=(const DynamicBitset& other) {
|
||||
TVM_FFI_DCHECK(buffer_size_ <= other.buffer_size_);
|
||||
for (int i = 0; i < buffer_size_; ++i) {
|
||||
data_[i] |= other.data_[i];
|
||||
}
|
||||
return *this;
|
||||
}
|
||||
|
||||
private:
|
||||
// The size of the bitset.
|
||||
int size_;
|
||||
// The size of the buffer.
|
||||
int buffer_size_;
|
||||
// The buffer for the bitset.
|
||||
uint32_t* data_;
|
||||
// The internal buffer. It is empty if not needed.
|
||||
std::vector<uint32_t> internal_buffer_;
|
||||
// Whether the buffer is internally managed.
|
||||
bool is_internal_;
|
||||
};
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_DYNAMIC_BITSET_H_
|
||||
@@ -0,0 +1,211 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file support/encoding.cc
|
||||
*/
|
||||
#include "encoding.h"
|
||||
|
||||
#include <tvm/runtime/logging.h>
|
||||
|
||||
#include <array>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
std::string PrintAsUTF8(TCodepoint codepoint) {
|
||||
TVM_FFI_ICHECK(codepoint <= 0x10FFFF) << "Invalid codepoint: " << codepoint;
|
||||
std::string utf8;
|
||||
if (codepoint <= 0x7F) {
|
||||
// 1-byte sequence
|
||||
utf8 += static_cast<char>(codepoint);
|
||||
} else if (codepoint <= 0x7FF) {
|
||||
// 2-byte sequence
|
||||
utf8 += static_cast<char>(0xC0 | ((codepoint >> 6) & 0x1F));
|
||||
utf8 += static_cast<char>(0x80 | (codepoint & 0x3F));
|
||||
} else if (codepoint <= 0xFFFF) {
|
||||
// 3-byte sequence
|
||||
utf8 += static_cast<char>(0xE0 | ((codepoint >> 12) & 0x0F));
|
||||
utf8 += static_cast<char>(0x80 | ((codepoint >> 6) & 0x3F));
|
||||
utf8 += static_cast<char>(0x80 | (codepoint & 0x3F));
|
||||
} else {
|
||||
// 4-byte sequence
|
||||
utf8 += static_cast<char>(0xF0 | ((codepoint >> 18) & 0x07));
|
||||
utf8 += static_cast<char>(0x80 | ((codepoint >> 12) & 0x3F));
|
||||
utf8 += static_cast<char>(0x80 | ((codepoint >> 6) & 0x3F));
|
||||
utf8 += static_cast<char>(0x80 | (codepoint & 0x3F));
|
||||
}
|
||||
return utf8;
|
||||
}
|
||||
|
||||
std::string PrintAsEscaped(
|
||||
TCodepoint codepoint,
|
||||
const std::unordered_map<TCodepoint, std::string>& additional_escape_map) {
|
||||
static const std::unordered_map<TCodepoint, std::string> kCodepointToEscape = {
|
||||
{'\'', "\\\'"}, {'\"', "\\\""}, {'\?', "\\\?"}, {'\\', "\\\\"}, {'\a', "\\a"},
|
||||
{'\b', "\\b"}, {'\f', "\\f"}, {'\n', "\\n"}, {'\r', "\\r"}, {'\t', "\\t"},
|
||||
{'\v', "\\v"}, {'\0', "\\0"}, {'\x1B', "\\e"}};
|
||||
|
||||
if (auto it = additional_escape_map.find(codepoint); it != additional_escape_map.end()) {
|
||||
return it->second;
|
||||
}
|
||||
|
||||
if (auto it = kCodepointToEscape.find(codepoint); it != kCodepointToEscape.end()) {
|
||||
return it->second;
|
||||
}
|
||||
|
||||
if (codepoint >= 0x20 && codepoint <= 0x7E) {
|
||||
return std::string({static_cast<char>(codepoint)});
|
||||
}
|
||||
|
||||
// convert codepoint to hex
|
||||
char prefix = codepoint <= 0xFF ? 'x' : codepoint <= 0xFFFF ? 'u' : 'U';
|
||||
int width = codepoint <= 0xFF ? 2 : codepoint <= 0xFFFF ? 4 : 8;
|
||||
std::stringstream ss;
|
||||
ss << std::setfill('0') << std::setw(width) << std::hex << codepoint;
|
||||
auto hex = ss.str();
|
||||
return std::string("\\") + prefix + hex;
|
||||
}
|
||||
|
||||
std::string PrintAsEscaped(uint8_t raw_char) { return PrintAsEscaped(raw_char); }
|
||||
|
||||
std::string PrintAsEscaped(std::string raw_str) {
|
||||
std::string res;
|
||||
auto codepoints = ParseUTF8(raw_str.c_str(), UTF8ErrorPolicy::kReturnByte);
|
||||
for (auto c : codepoints) {
|
||||
res += PrintAsEscaped(c);
|
||||
}
|
||||
return res;
|
||||
}
|
||||
|
||||
std::tuple<bool, int, TCodepoint> HandleUTF8FirstByte(uint8_t byte) {
|
||||
static const std::array<int8_t, 5> kFirstByteMask = {0x00, 0x7F, 0x1F, 0x0F, 0x07};
|
||||
// clang-format off
|
||||
static const std::array<int, 256> kUtf8Bytes = {
|
||||
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
|
||||
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
|
||||
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
|
||||
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
|
||||
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
|
||||
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
|
||||
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
|
||||
3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
|
||||
4, 4, 4, 4, 4, 4, 4, 4, -1, -1, -1, -1, -1, -1, -1, -1,
|
||||
};
|
||||
// clang-format on
|
||||
auto num_bytes = kUtf8Bytes[static_cast<uint8_t>(byte)];
|
||||
if (num_bytes == -1) {
|
||||
return {false, 0, 0};
|
||||
}
|
||||
return {true, num_bytes, byte & kFirstByteMask[num_bytes]};
|
||||
}
|
||||
|
||||
std::pair<TCodepoint, const char*> ParseNextUTF8(const char* utf8, UTF8ErrorPolicy error_policy) {
|
||||
auto [accepted, num_bytes, res] = HandleUTF8FirstByte(utf8[0]);
|
||||
if (accepted) {
|
||||
for (int i = 1; i < num_bytes; ++i) {
|
||||
if (utf8[i] == 0 || (static_cast<uint8_t>(utf8[i]) & 0xC0) != 0x80) {
|
||||
// invalid utf8
|
||||
accepted = false;
|
||||
break;
|
||||
}
|
||||
res = (res << 6) | (static_cast<uint8_t>(utf8[i]) & 0x3F);
|
||||
}
|
||||
}
|
||||
|
||||
if (!accepted) {
|
||||
// invalid utf8
|
||||
if (error_policy == UTF8ErrorPolicy::kReturnInvalid) {
|
||||
return {CharHandlingError::kInvalidUTF8, utf8};
|
||||
} else {
|
||||
return {static_cast<unsigned char>(utf8[0]), utf8 + 1};
|
||||
}
|
||||
}
|
||||
|
||||
return {res, utf8 + num_bytes};
|
||||
}
|
||||
|
||||
std::vector<TCodepoint> ParseUTF8(const char* utf8, UTF8ErrorPolicy error_policy) {
|
||||
std::vector<TCodepoint> codepoints;
|
||||
while (*utf8 != 0) {
|
||||
TCodepoint codepoint;
|
||||
std::tie(codepoint, utf8) = ParseNextUTF8(utf8, error_policy);
|
||||
if (codepoint == CharHandlingError::kInvalidUTF8) {
|
||||
return {codepoint};
|
||||
}
|
||||
codepoints.push_back(codepoint);
|
||||
}
|
||||
return codepoints;
|
||||
}
|
||||
|
||||
inline int HexCharToInt(char c) {
|
||||
if (c >= '0' && c <= '9') {
|
||||
return c - '0';
|
||||
} else if (c >= 'a' && c <= 'f') {
|
||||
return c - 'a' + 10;
|
||||
} else if (c >= 'A' && c <= 'F') {
|
||||
return c - 'A' + 10;
|
||||
} else {
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
std::pair<TCodepoint, const char*> ParseNextUTF8OrEscaped(
|
||||
const char* utf8, const std::unordered_map<std::string, TCodepoint>& additional_escape_map) {
|
||||
static const std::unordered_map<std::string, TCodepoint> kEscapeToCodepoint = {
|
||||
{"\\\'", '\''}, {"\\\"", '\"'}, {"\\\?", '\?'}, {"\\\\", '\\'}, {"\\a", '\a'},
|
||||
{"\\b", '\b'}, {"\\f", '\f'}, {"\\n", '\n'}, {"\\r", '\r'}, {"\\t", '\t'},
|
||||
{"\\v", '\v'}, {"\\0", '\0'}, {"\\e", '\x1B'}};
|
||||
if (utf8[0] != '\\') {
|
||||
return ParseNextUTF8(utf8, UTF8ErrorPolicy::kReturnInvalid);
|
||||
}
|
||||
|
||||
auto escape_sequence = std::string(utf8, 2);
|
||||
if (auto it = additional_escape_map.find(escape_sequence); it != additional_escape_map.end()) {
|
||||
return {it->second, utf8 + 2};
|
||||
}
|
||||
if (auto it = kEscapeToCodepoint.find(escape_sequence); it != kEscapeToCodepoint.end()) {
|
||||
return {it->second, utf8 + 2};
|
||||
}
|
||||
|
||||
if (utf8[1] == 'x') {
|
||||
// arbitrary length hex
|
||||
int len = 0;
|
||||
int32_t codepoint = 0;
|
||||
while (true) {
|
||||
auto digit = HexCharToInt(utf8[2 + len]);
|
||||
if (digit == -1) {
|
||||
break;
|
||||
}
|
||||
codepoint = codepoint * 16 + digit;
|
||||
++len;
|
||||
}
|
||||
if (len == 0) {
|
||||
return {CharHandlingError::kInvalidEscape, utf8};
|
||||
}
|
||||
return {codepoint, utf8 + len + 2};
|
||||
} else if (utf8[1] == 'u' || utf8[1] == 'U') {
|
||||
// 4- or 8-digit hex
|
||||
int len = utf8[1] == 'u' ? 4 : 8;
|
||||
int32_t codepoint = 0;
|
||||
|
||||
for (int i = 0; i < len; ++i) {
|
||||
auto digit = HexCharToInt(utf8[i + 2]);
|
||||
if (digit == -1) {
|
||||
return {CharHandlingError::kInvalidEscape, utf8};
|
||||
}
|
||||
codepoint = codepoint * 16 + digit;
|
||||
}
|
||||
return {codepoint, utf8 + len + 2};
|
||||
} else {
|
||||
return {CharHandlingError::kInvalidEscape, utf8};
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,114 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file support/encoding.h
|
||||
* \brief Encoding and decoding from/to UTF-8 and escape sequence to/from codepoints.
|
||||
*/
|
||||
#ifndef MLC_LLM_SUPPORT_ENCODING_H_
|
||||
#define MLC_LLM_SUPPORT_ENCODING_H_
|
||||
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
/*! \brief Represents a unicode codepoint. */
|
||||
using TCodepoint = int32_t;
|
||||
|
||||
/*!
|
||||
* \brief Handle the utf-8 first byte.
|
||||
* \returns (is_valid, total_number_of_bytes, initial_codepoint).
|
||||
*/
|
||||
std::tuple<bool, int, TCodepoint> HandleUTF8FirstByte(uint8_t byte);
|
||||
|
||||
/*!
|
||||
* \brief Print a codepoint to a UTF-8 string.
|
||||
* \param codepoint The codepoint.
|
||||
* \return The UTF-8 string.
|
||||
*/
|
||||
std::string PrintAsUTF8(TCodepoint codepoint);
|
||||
|
||||
/*!
|
||||
* \brief Print a codepoint to a escaped string. If the codepoint is not printable, it will be
|
||||
* escaped. By default the function support escape sequences in C ("\n", "\t", "\u0123"). User can
|
||||
* specify more escape sequences using additional_escape_map.
|
||||
* \param codepoint The codepoint.
|
||||
* \param additional_escape_map A map from codepoint to escape sequence. If the codepoint is in the
|
||||
* map, it will be escaped using the corresponding escape sequence. e.g. {{'-', "\\-"}}. \return The
|
||||
* printable string.
|
||||
*/
|
||||
std::string PrintAsEscaped(
|
||||
TCodepoint codepoint,
|
||||
const std::unordered_map<TCodepoint, std::string>& additional_escape_map = {});
|
||||
|
||||
/*!
|
||||
* \brief Print the given char to a escaped string that can be printed.
|
||||
* \return The escaped string.
|
||||
*/
|
||||
std::string PrintAsEscaped(uint8_t raw_char);
|
||||
|
||||
/*!
|
||||
* \brief Print the given string to a escaped string that can be printed.
|
||||
* \return The escaped string.
|
||||
*/
|
||||
std::string PrintAsEscaped(std::string raw_str);
|
||||
|
||||
/*!
|
||||
* \brief Represents an error when handling characters. Will be returned as a special TCodepoint
|
||||
* value.
|
||||
*/
|
||||
enum CharHandlingError : TCodepoint {
|
||||
/*! \brief The UTF-8 string is invalid. */
|
||||
kInvalidUTF8 = -10,
|
||||
/*! \brief The escape sequence is invalid. */
|
||||
kInvalidEscape = -11,
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief The method to handle invalid UTF-8 sequence.
|
||||
*/
|
||||
enum class UTF8ErrorPolicy {
|
||||
/*! \brief Return an error codepoint when an error is encountered. */
|
||||
kReturnInvalid,
|
||||
/*! \brief Skip the error and continue parsing. */
|
||||
kReturnByte,
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Parse the first codepoint in a UTF-8 string.
|
||||
* \param utf8 The UTF-8 string.
|
||||
* \return The codepoint and new pointer. If the UTF-8 string is invalid, and the error policy is
|
||||
* kReturnInvalid, the function returns (CharHandlingError::kInvalidUTF8, input char pointer).
|
||||
*/
|
||||
std::pair<TCodepoint, const char*> ParseNextUTF8(
|
||||
const char* utf8, UTF8ErrorPolicy error_policy = UTF8ErrorPolicy::kReturnInvalid);
|
||||
|
||||
/*!
|
||||
* \brief Parse all codepoints in a UTF-8 string.
|
||||
* \param utf8 The UTF-8 string.
|
||||
* \return All codepoints. If the UTF-8 string is invalid, and the error policy is
|
||||
* kReturnInvalid, the function returns {CharHandlingError::kInvalidUTF8}.
|
||||
*/
|
||||
std::vector<TCodepoint> ParseUTF8(const char* utf8,
|
||||
UTF8ErrorPolicy error_policy = UTF8ErrorPolicy::kReturnInvalid);
|
||||
|
||||
/*!
|
||||
* \brief Parse the first codepoint from a UTF-8 string. Also checks escape sequences and converts
|
||||
* the escaped char to its original value.
|
||||
* \param utf8 The UTF-8 string or the escape sequence.
|
||||
* \param additional_escape_map A map from escape sequence to codepoint. If the escape sequence is
|
||||
* in the map, it will be converted to the corresponding codepoint. e.g. {{"\\-", '-'}}.
|
||||
* \return The codepoint and the new pointer. If the UTF-8 string or the escape sequence is
|
||||
* invalid, and the error policy is kReturnInvalid, the function returns
|
||||
* (CharHandlingError::kInvalidUTF8, input char pointer).
|
||||
*/
|
||||
std::pair<TCodepoint, const char*> ParseNextUTF8OrEscaped(
|
||||
const char* utf8,
|
||||
const std::unordered_map<std::string, TCodepoint>& additional_escape_map = {});
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_ENCODING_H_
|
||||
@@ -0,0 +1,284 @@
|
||||
/*!
|
||||
* \file support/json_parser.h
|
||||
* \brief Helps to parse JSON strings and objects.
|
||||
*/
|
||||
#ifndef MLC_LLM_SUPPORT_JSON_PARSER_H_
|
||||
#define MLC_LLM_SUPPORT_JSON_PARSER_H_
|
||||
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/dtype.h>
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
#include <tvm/runtime/logging.h>
|
||||
|
||||
#include <optional>
|
||||
|
||||
#include "result.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace json {
|
||||
|
||||
using ::tvm::ffi::json::Array;
|
||||
using ::tvm::ffi::json::Object;
|
||||
using ::tvm::ffi::json::Value;
|
||||
|
||||
/*!
|
||||
* \brief Parse a JSON string to a JSON object.
|
||||
* \param json_str The JSON string to parse.
|
||||
* \return The parsed JSON object.
|
||||
*/
|
||||
inline Object ParseToJSONObject(const std::string& json_str) {
|
||||
tvm::ffi::String err;
|
||||
Value result = ::tvm::ffi::json::Parse(json_str, &err);
|
||||
TVM_FFI_CHECK(err.empty(), ValueError)
|
||||
<< "Failed to parse JSON: err. The JSON string is:" << json_str;
|
||||
auto opt = result.try_cast<Object>();
|
||||
TVM_FFI_CHECK(opt.has_value(), ValueError)
|
||||
<< "The given string is not a JSON object: " << json_str;
|
||||
return *opt;
|
||||
}
|
||||
/*!
|
||||
* \brief Parse a JSON string to a JSON object.
|
||||
* \param json_str The JSON string to parse.
|
||||
* \return The parsed JSON object, or the error message.
|
||||
*/
|
||||
inline Result<Object> ParseToJSONObjectWithResultReturn(const std::string& json_str) {
|
||||
using TResult = Result<Object>;
|
||||
tvm::ffi::String err;
|
||||
Value result = ::tvm::ffi::json::Parse(json_str, &err);
|
||||
if (!err.empty()) {
|
||||
return TResult::Error("Failed to parse JSON: err. The JSON string is: " + json_str +
|
||||
". The error is " + std::string(err));
|
||||
}
|
||||
auto opt = result.try_cast<Object>();
|
||||
if (!opt.has_value()) {
|
||||
return TResult::Error("ValueError: The given string is not a JSON object: " + json_str);
|
||||
}
|
||||
return TResult::Ok(*opt);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
||||
* \param json The JSON object to look up.
|
||||
* \param key The key to look up.
|
||||
* \tparam ValueType The type to be converted to.
|
||||
* \return The converted value.
|
||||
*/
|
||||
template <typename ValueType>
|
||||
ValueType Lookup(const Object& json, const std::string& key);
|
||||
/*!
|
||||
* \brief Lookup a JSON array by an index, and convert it to a given type.
|
||||
* \param json The JSON array to look up.
|
||||
* \param index The index to look up.
|
||||
* \tparam ValueType The type to be converted to.
|
||||
* \return The converted value.
|
||||
*/
|
||||
template <typename ValueType>
|
||||
ValueType Lookup(const Array& json, int index);
|
||||
/*!
|
||||
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
||||
* If the key doesn't exist or has null value, the default value is returned.
|
||||
* \param json The JSON object to look up.
|
||||
* \param key The key to look up.
|
||||
* \tparam ValueType The type to be converted to.
|
||||
* \return The converted value, or the default value if the key doesn't exist or has null value.
|
||||
*/
|
||||
template <typename ValueType>
|
||||
inline ValueType LookupOrDefault(const Object& json, const std::string& key,
|
||||
const ValueType& default_value) {
|
||||
if (json.count(key) == 0 || json.at(key) == nullptr) {
|
||||
return default_value;
|
||||
}
|
||||
auto opt = json.at(key).try_cast<ValueType>();
|
||||
TVM_FFI_CHECK(opt.has_value(), ValueError) << "key `" << key << "` has unexpected type";
|
||||
return *opt;
|
||||
}
|
||||
/*!
|
||||
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
||||
* If the key doesn't exist or has null value, return std::nullopt.
|
||||
* \param json The JSON object to look up.
|
||||
* \param key The key to look up.
|
||||
* \tparam ValueType The type to be converted to.
|
||||
* \return The converted value, or std::nullopt if the value doesn't exist or has null value.
|
||||
*/
|
||||
template <typename ValueType>
|
||||
inline std::optional<ValueType> LookupOptional(const Object& json, const std::string& key) {
|
||||
if (json.count(key) == 0 || json.at(key) == nullptr) {
|
||||
return std::nullopt;
|
||||
}
|
||||
auto opt = json.at(key).try_cast<ValueType>();
|
||||
TVM_FFI_CHECK(opt.has_value(), ValueError) << "key `" << key << "` has unexpected type";
|
||||
return *opt;
|
||||
}
|
||||
/*!
|
||||
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
||||
* \param json The JSON object to look up.
|
||||
* \param key The key to look up.
|
||||
* \tparam ValueType The type to be converted to.
|
||||
* \return The converted value, or the error message.
|
||||
*/
|
||||
template <typename ValueType>
|
||||
inline Result<ValueType> LookupWithResultReturn(const Object& json, const std::string& key) {
|
||||
using TResult = Result<ValueType>;
|
||||
if (json.count(key) == 0) {
|
||||
return TResult::Error("ValueError: key \"" + key + "\" not found in the JSON object");
|
||||
}
|
||||
auto opt = json.at(key).try_cast<ValueType>();
|
||||
if (!opt.has_value()) {
|
||||
return TResult::Error("ValueError: key \"" + key + "\" has unexpected value type.");
|
||||
}
|
||||
return TResult::Ok(*opt);
|
||||
}
|
||||
/*!
|
||||
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
||||
* If the key doesn't exist or has null value, the default value is returned.
|
||||
* \param json The JSON object to look up.
|
||||
* \param key The key to look up.
|
||||
* \tparam ValueType The type to be converted to.
|
||||
* \return The converted value, or the default value if the key doesn't exist or has null value
|
||||
* , or the error message.
|
||||
*/
|
||||
template <typename ValueType>
|
||||
inline Result<ValueType> LookupOrDefaultWithResultReturn(const Object& json, const std::string& key,
|
||||
const ValueType& default_value) {
|
||||
using TResult = Result<ValueType>;
|
||||
if (json.count(key) == 0 || json.at(key) == nullptr) {
|
||||
return TResult::Ok(default_value);
|
||||
}
|
||||
auto opt = json.at(key).try_cast<ValueType>();
|
||||
if (!opt.has_value()) {
|
||||
return TResult::Error("ValueError: key \"" + key + "\" has unexpected value type.");
|
||||
}
|
||||
return TResult::Ok(*opt);
|
||||
}
|
||||
/*!
|
||||
* \brief Lookup a JSON object by a key, and convert it to a given type.
|
||||
* If the key doesn't exist or has null value, return std::nullopt.
|
||||
* \param json The JSON object to look up.
|
||||
* \param key The key to look up.
|
||||
* \tparam ValueType The type to be converted to.
|
||||
* \return The converted value, or std::nullopt if the value doesn't exist or has null value,
|
||||
* , or the error message.
|
||||
*/
|
||||
template <typename ValueType>
|
||||
inline Result<std::optional<ValueType>> LookupOptionalWithResultReturn(const Object& json,
|
||||
const std::string& key) {
|
||||
using TResult = Result<std::optional<ValueType>>;
|
||||
if (json.count(key) == 0 || json.at(key) == nullptr) {
|
||||
return TResult::Ok(std::nullopt);
|
||||
}
|
||||
auto opt = json.at(key).try_cast<ValueType>();
|
||||
if (!opt.has_value()) {
|
||||
return TResult::Error("ValueError: key \"" + key + "\" has unexpected value type.");
|
||||
}
|
||||
return TResult::Ok(*opt);
|
||||
}
|
||||
|
||||
// Implementation details
|
||||
|
||||
/*! \brief Shape extension to incorporate symbolic shapes. */
|
||||
struct SymShapeTuple {
|
||||
tvm::ffi::Shape shape_values;
|
||||
std::vector<std::string> sym_names;
|
||||
|
||||
/*! \brief Convert symbolic shape tuple to static shape tuple with model config. */
|
||||
tvm::ffi::Shape ToStatic(const Object& model_config) {
|
||||
std::vector<int64_t> shape;
|
||||
shape.reserve(shape_values.size());
|
||||
for (int i = 0; i < static_cast<int>(shape_values.size()); ++i) {
|
||||
if (shape_values[i] != -1) {
|
||||
shape.push_back(shape_values[i]);
|
||||
} else {
|
||||
auto opt = model_config.at(sym_names[i]).try_cast<int64_t>();
|
||||
TVM_FFI_CHECK(opt.has_value(), ValueError)
|
||||
<< "model config is expected to contain \"" << sym_names[i]
|
||||
<< "\" as an integer. However, the given config has unexpected type for \""
|
||||
<< sym_names[i] << "\".";
|
||||
shape.push_back(*opt);
|
||||
}
|
||||
}
|
||||
return tvm::ffi::Shape(std::move(shape));
|
||||
}
|
||||
};
|
||||
|
||||
namespace details {
|
||||
|
||||
inline DLDataType DTypeFromString(const std::string& s) { return tvm::ffi::StringToDLDataType(s); }
|
||||
|
||||
inline SymShapeTuple SymShapeTupleFromArray(const Array& shape) {
|
||||
std::vector<int64_t> result;
|
||||
std::vector<std::string> sym_names;
|
||||
result.reserve(shape.size());
|
||||
sym_names.reserve(shape.size());
|
||||
for (int i = 0; i < static_cast<int>(shape.size()); ++i) {
|
||||
const auto& dim = shape[i];
|
||||
auto str_opt = dim.try_cast<std::string>();
|
||||
if (str_opt.has_value()) {
|
||||
result.push_back(-1);
|
||||
sym_names.push_back(*str_opt);
|
||||
} else {
|
||||
auto int_opt = dim.try_cast<int64_t>();
|
||||
TVM_FFI_CHECK(int_opt.has_value(), ValueError) << "shape has unexpected type";
|
||||
result.push_back(*int_opt);
|
||||
sym_names.push_back("");
|
||||
}
|
||||
}
|
||||
return SymShapeTuple{tvm::ffi::Shape(std::move(result)), sym_names};
|
||||
}
|
||||
|
||||
} // namespace details
|
||||
|
||||
template <typename ValueType>
|
||||
inline ValueType Lookup(const Object& json, const std::string& key) {
|
||||
TVM_FFI_CHECK(json.count(key) != 0, ValueError)
|
||||
<< "key `" << key << "` not found in the JSON object";
|
||||
auto opt = json.at(key).try_cast<ValueType>();
|
||||
TVM_FFI_CHECK(opt.has_value(), ValueError) << "key `" << key << "` has unexpected type";
|
||||
return *opt;
|
||||
}
|
||||
|
||||
template <typename ValueType>
|
||||
inline ValueType Lookup(const Array& json, int index) {
|
||||
TVM_FFI_ICHECK(index < static_cast<int>(json.size()))
|
||||
<< "IndexError: json::array index out of range";
|
||||
auto opt = json[index].try_cast<ValueType>();
|
||||
TVM_FFI_ICHECK(opt.has_value()) << "ValueError: value at index `" << index
|
||||
<< "` has unexpected type";
|
||||
return *opt;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline DLDataType Lookup(const Object& json, const std::string& key) {
|
||||
return details::DTypeFromString(Lookup<std::string>(json, key));
|
||||
}
|
||||
|
||||
template <>
|
||||
inline DLDataType Lookup(const Array& json, int index) {
|
||||
return details::DTypeFromString(Lookup<std::string>(json, index));
|
||||
}
|
||||
|
||||
template <>
|
||||
inline SymShapeTuple Lookup(const Object& json, const std::string& key) {
|
||||
return details::SymShapeTupleFromArray(Lookup<Array>(json, key));
|
||||
}
|
||||
|
||||
template <>
|
||||
inline SymShapeTuple LookupOrDefault(const Object& json, const std::string& key,
|
||||
const SymShapeTuple& default_value) {
|
||||
if (json.count(key) == 0 || json.at(key) == nullptr) {
|
||||
return default_value;
|
||||
}
|
||||
return details::SymShapeTupleFromArray(Lookup<Array>(json, key));
|
||||
}
|
||||
|
||||
template <>
|
||||
inline SymShapeTuple Lookup(const Array& json, int index) {
|
||||
return details::SymShapeTupleFromArray(Lookup<Array>(json, index));
|
||||
}
|
||||
|
||||
} // namespace json
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_JSON_PARSER_H_
|
||||
@@ -0,0 +1,32 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file support/load_bytes_from_file.h
|
||||
* \brief Utility methods to load from files.
|
||||
*/
|
||||
#ifndef MLC_LLM_SUPPORT_LOAD_BYTES_FROM_FILE_H_
|
||||
#define MLC_LLM_SUPPORT_LOAD_BYTES_FROM_FILE_H_
|
||||
|
||||
#include <tvm/runtime/logging.h>
|
||||
|
||||
#include <fstream>
|
||||
#include <string>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
inline std::string LoadBytesFromFile(const std::string& path) {
|
||||
std::ifstream fs(path, std::ios::in | std::ios::binary);
|
||||
TVM_FFI_ICHECK(!fs.fail()) << "Cannot open " << path;
|
||||
std::string data;
|
||||
fs.seekg(0, std::ios::end);
|
||||
size_t size = static_cast<size_t>(fs.tellg());
|
||||
fs.seekg(0, std::ios::beg);
|
||||
data.resize(size);
|
||||
fs.read(data.data(), size);
|
||||
return data;
|
||||
}
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_LOAD_BYTES_FROM_FILE_H_
|
||||
@@ -0,0 +1,107 @@
|
||||
/*!
|
||||
* \file support/module_vtable.h
|
||||
* \brief Compatibility shim providing the TVM_MODULE_VTABLE_* macros that
|
||||
* previously lived in <tvm/runtime/module.h>. After the TVM runtime
|
||||
* refactor the macros were moved into TVM's private
|
||||
* src/runtime/vm/module_utils.h, so we vendor the surface mlc-llm uses
|
||||
* here to keep the existing call sites unchanged.
|
||||
*/
|
||||
#ifndef MLC_LLM_SUPPORT_MODULE_VTABLE_H_
|
||||
#define MLC_LLM_SUPPORT_MODULE_VTABLE_H_
|
||||
|
||||
#include <tvm/ffi/cast.h>
|
||||
#include <tvm/ffi/extra/module.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
|
||||
#include <utility>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
namespace module_vtable_detail {
|
||||
|
||||
template <typename T>
|
||||
struct EntryHelper {};
|
||||
|
||||
template <typename T, typename R, typename... Args>
|
||||
struct EntryHelper<R (T::*)(Args...) const> {
|
||||
using MemFnType = R (T::*)(Args...) const;
|
||||
TVM_FFI_INLINE static void Call(::tvm::ffi::Any* rv, T* self, MemFnType f,
|
||||
::tvm::ffi::PackedArgs args) {
|
||||
auto wrapped = [self, f](Args... args) -> R { return (self->*f)(std::forward<Args>(args)...); };
|
||||
::tvm::ffi::details::unpack_call<R>(std::make_index_sequence<sizeof...(Args)>{}, nullptr,
|
||||
wrapped, args.data(), args.size(), rv);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, typename R, typename... Args>
|
||||
struct EntryHelper<R (T::*)(Args...)> {
|
||||
using MemFnType = R (T::*)(Args...);
|
||||
TVM_FFI_INLINE static void Call(::tvm::ffi::Any* rv, T* self, MemFnType f,
|
||||
::tvm::ffi::PackedArgs args) {
|
||||
auto wrapped = [self, f](Args... args) -> R { return (self->*f)(std::forward<Args>(args)...); };
|
||||
::tvm::ffi::details::unpack_call<R>(std::make_index_sequence<sizeof...(Args)>{}, nullptr,
|
||||
wrapped, args.data(), args.size(), rv);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, typename... Args>
|
||||
struct EntryHelper<void (T::*)(Args...) const> {
|
||||
using MemFnType = void (T::*)(Args...) const;
|
||||
TVM_FFI_INLINE static void Call(::tvm::ffi::Any* rv, T* self, MemFnType f,
|
||||
::tvm::ffi::PackedArgs args) {
|
||||
auto wrapped = [self, f](Args... args) -> void { (self->*f)(std::forward<Args>(args)...); };
|
||||
::tvm::ffi::details::unpack_call<void>(std::make_index_sequence<sizeof...(Args)>{}, nullptr,
|
||||
wrapped, args.data(), args.size(), rv);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, typename... Args>
|
||||
struct EntryHelper<void (T::*)(Args...)> {
|
||||
using MemFnType = void (T::*)(Args...);
|
||||
TVM_FFI_INLINE static void Call(::tvm::ffi::Any* rv, T* self, MemFnType f,
|
||||
::tvm::ffi::PackedArgs args) {
|
||||
auto wrapped = [self, f](Args... args) -> void { (self->*f)(std::forward<Args>(args)...); };
|
||||
::tvm::ffi::details::unpack_call<void>(std::make_index_sequence<sizeof...(Args)>{}, nullptr,
|
||||
wrapped, args.data(), args.size(), rv);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace module_vtable_detail
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#define TVM_MODULE_VTABLE_BEGIN(TypeKey) \
|
||||
const char* kind() const final { return TypeKey; } \
|
||||
::tvm::ffi::Optional<::tvm::ffi::Function> GetFunction(const ::tvm::ffi::String& _name) \
|
||||
override { \
|
||||
using SelfPtr = std::remove_cv_t<decltype(this)>; \
|
||||
::tvm::ffi::ObjectPtr<::tvm::ffi::Object> _self = \
|
||||
::tvm::ffi::GetObjectPtr<::tvm::ffi::Object>(this);
|
||||
#define TVM_MODULE_VTABLE_END() \
|
||||
return std::nullopt; \
|
||||
}
|
||||
#define TVM_MODULE_VTABLE_END_WITH_DEFAULT(MemFunc) \
|
||||
{ \
|
||||
auto f = (MemFunc); \
|
||||
return (this->*f)(_name); \
|
||||
} \
|
||||
}
|
||||
#define TVM_MODULE_VTABLE_ENTRY(Name, MemFunc) \
|
||||
if (_name == Name) { \
|
||||
return ::tvm::ffi::Function::FromPacked( \
|
||||
[_self](::tvm::ffi::PackedArgs args, ::tvm::ffi::Any* rv) -> void { \
|
||||
using Helper = ::mlc::llm::module_vtable_detail::EntryHelper<decltype(MemFunc)>; \
|
||||
SelfPtr self = static_cast<SelfPtr>(_self.get()); \
|
||||
Helper::Call(rv, self, MemFunc, args); \
|
||||
}); \
|
||||
}
|
||||
#define TVM_MODULE_VTABLE_ENTRY_PACKED(Name, MemFunc) \
|
||||
if (_name == Name) { \
|
||||
return ::tvm::ffi::Function( \
|
||||
[_self](::tvm::ffi::PackedArgs args, ::tvm::ffi::Any* rv) -> void { \
|
||||
(static_cast<SelfPtr>(_self.get())->*(MemFunc))(args, rv); \
|
||||
}); \
|
||||
}
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_MODULE_VTABLE_H_
|
||||
@@ -0,0 +1,49 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file support/progress_bar.h
|
||||
* \brief A simple progress bar in C++.
|
||||
*/
|
||||
#ifndef MLC_LLM_SUPPORT_PROGRESS_BAR_H_
|
||||
#define MLC_LLM_SUPPORT_PROGRESS_BAR_H_
|
||||
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
class ProgressBar {
|
||||
public:
|
||||
explicit ProgressBar(int total, int width = 100) : total(total), width(width), cur(0) {}
|
||||
|
||||
void Progress() {
|
||||
if (cur < total) {
|
||||
++cur;
|
||||
}
|
||||
int bar_width = width - 2; // Adjust for borders
|
||||
int completed = static_cast<int>(static_cast<float>(cur) / total * bar_width);
|
||||
int remaining = bar_width - completed;
|
||||
std::cout << "[" //
|
||||
<< std::string(completed, '=') //
|
||||
<< ">" //
|
||||
<< std::string(remaining, ' ') //
|
||||
<< "] " //
|
||||
<< " [" << cur << "/" << total << "]";
|
||||
if (cur < total) {
|
||||
std::cout << "\r";
|
||||
std::cout.flush();
|
||||
} else {
|
||||
std::cout << std::endl; // Move to the next line after the progress bar is complete
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
int total;
|
||||
int width;
|
||||
int cur;
|
||||
};
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_PROGRESS_BAR_H_
|
||||
@@ -0,0 +1,37 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file support/random.h
|
||||
* \brief Header of random number generator.
|
||||
*/
|
||||
|
||||
#ifndef MLC_LLM_SUPPORT_RANDOM_H_
|
||||
#define MLC_LLM_SUPPORT_RANDOM_H_
|
||||
|
||||
#include <random>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
// Random number generator
|
||||
class RandomGenerator {
|
||||
private:
|
||||
std::mt19937 gen;
|
||||
std::uniform_real_distribution<> dis;
|
||||
|
||||
public:
|
||||
RandomGenerator(int seed = std::random_device{}()) : gen(seed), dis(0.0, 1.0) {}
|
||||
|
||||
static RandomGenerator& GetInstance(int seed = std::random_device{}()) {
|
||||
static RandomGenerator instance(seed);
|
||||
return instance;
|
||||
}
|
||||
|
||||
double GetRandomNumber() { return dis(gen); }
|
||||
|
||||
void SetSeed(int seed) { gen.seed(seed); }
|
||||
};
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_RANDOM_H_
|
||||
@@ -0,0 +1,77 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file support/result.h
|
||||
* \brief The header for the Result class in MLC LLM.
|
||||
*/
|
||||
#ifndef MLC_LLM_SUPPORT_RESULT_H_
|
||||
#define MLC_LLM_SUPPORT_RESULT_H_
|
||||
|
||||
#include <tvm/runtime/logging.h>
|
||||
|
||||
#include <optional>
|
||||
#include <string>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
/*!
|
||||
* \brief The result class in MLC LLM.
|
||||
* Each instance is either an okay value or an error.
|
||||
* \tparam T The okay value type of the result.
|
||||
* \tparam E The error type of the result.
|
||||
*/
|
||||
template <typename T, typename E = std::string>
|
||||
class Result {
|
||||
public:
|
||||
/*! \brief Create a result with an okay value. */
|
||||
static Result Ok(T value) {
|
||||
Result result;
|
||||
result.ok_value_ = std::move(value);
|
||||
return result;
|
||||
}
|
||||
/*! \brief Create a result with an error value. */
|
||||
static Result Error(E error) {
|
||||
Result result;
|
||||
result.err_value_ = std::move(error);
|
||||
return result;
|
||||
}
|
||||
/*! \brief Check if the result is okay or not. */
|
||||
bool IsOk() const { return ok_value_.has_value(); }
|
||||
/*! \brief Check if the result is an error or not. */
|
||||
bool IsErr() const { return err_value_.has_value(); }
|
||||
/*!
|
||||
* \brief Unwrap the result and return the okay value.
|
||||
* Throwing exception if it is an error.
|
||||
* \note This function returns the ok value by moving, so a Result can be unwrapped only once.
|
||||
*/
|
||||
T Unwrap() {
|
||||
TVM_FFI_ICHECK(ok_value_.has_value()) << "Cannot unwrap result on an error value.";
|
||||
TVM_FFI_ICHECK(!unwrapped_) << "Cannot unwrap a Result instance twice.";
|
||||
unwrapped_ = true;
|
||||
return std::move(ok_value_.value());
|
||||
}
|
||||
/*!
|
||||
* \brief Unwrap the result and return the error value.
|
||||
* Throwing exception if it is an okay value.
|
||||
* \note This function returns the error value by moving, so a Result can be unwrapped only once.
|
||||
*/
|
||||
E UnwrapErr() {
|
||||
TVM_FFI_ICHECK(err_value_.has_value()) << "Cannot unwrap result on an okay value.";
|
||||
TVM_FFI_ICHECK(!unwrapped_) << "Cannot unwrap a Result instance twice.";
|
||||
unwrapped_ = true;
|
||||
return std::move(err_value_.value());
|
||||
}
|
||||
|
||||
private:
|
||||
/*! \brief A boolean flag indicating if the result is okay or error. */
|
||||
bool unwrapped_ = false;
|
||||
/*! \brief The internal optional okay value. */
|
||||
std::optional<T> ok_value_;
|
||||
/*! \brief The internal optional error value. */
|
||||
std::optional<E> err_value_;
|
||||
};
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_RESULT_H_
|
||||
@@ -0,0 +1,73 @@
|
||||
/*!
|
||||
* \file support/threading_backend.h
|
||||
* \brief Compatibility shim providing the threading helpers that used to live
|
||||
* at <tvm/runtime/threading_backend.h>. The public TVM header was
|
||||
* removed in a runtime refactor, but the underlying symbols
|
||||
* (TVMBackendParallelLaunch, threading::MaxConcurrency,
|
||||
* threading::SetMaxConcurrency) are still exported from libtvm_runtime.
|
||||
* We re-declare the inline parallel-for template here so existing
|
||||
* mlc-llm call sites keep compiling.
|
||||
*/
|
||||
#ifndef MLC_LLM_SUPPORT_THREADING_BACKEND_H_
|
||||
#define MLC_LLM_SUPPORT_THREADING_BACKEND_H_
|
||||
|
||||
#include <tvm/runtime/base.h>
|
||||
#include <tvm/runtime/c_backend_api.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
namespace threading {
|
||||
|
||||
/*! \return the maximum number of effective workers for this system. */
|
||||
int MaxConcurrency();
|
||||
/*! \brief Setting the maximum number of available cores. */
|
||||
void SetMaxConcurrency(int value);
|
||||
|
||||
} // namespace threading
|
||||
|
||||
namespace detail {
|
||||
|
||||
template <typename T>
|
||||
struct ParallelForWithThreadingBackendLambdaInvoker {
|
||||
static int TVMParallelLambdaInvoke(int task_id, TVMParallelGroupEnv* penv, void* cdata) {
|
||||
int num_task = penv->num_task;
|
||||
T* lambda_ptr = static_cast<T*>(cdata);
|
||||
(*lambda_ptr)(task_id, num_task);
|
||||
return 0;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
inline void parallel_launch_with_threading_backend(T flambda) {
|
||||
void* cdata = &flambda;
|
||||
TVMBackendParallelLaunch(ParallelForWithThreadingBackendLambdaInvoker<T>::TVMParallelLambdaInvoke,
|
||||
cdata, /*num_task=*/0);
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
template <typename T>
|
||||
inline void parallel_for_with_threading_backend(T flambda, int64_t begin, int64_t end) {
|
||||
if (end - begin == 1) {
|
||||
flambda(begin);
|
||||
return;
|
||||
}
|
||||
auto flaunch = [begin, end, flambda](int task_id, int num_task) {
|
||||
int64_t total_len = end - begin;
|
||||
int64_t step = (total_len + num_task - 1) / num_task;
|
||||
int64_t local_begin = std::min(begin + step * task_id, end);
|
||||
int64_t local_end = std::min(local_begin + step, end);
|
||||
for (int64_t i = local_begin; i < local_end; ++i) {
|
||||
flambda(i);
|
||||
}
|
||||
};
|
||||
detail::parallel_launch_with_threading_backend(flaunch);
|
||||
}
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_THREADING_BACKEND_H_
|
||||
@@ -0,0 +1,79 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file support/utils.h
|
||||
* \brief Utility functions.
|
||||
*/
|
||||
#ifndef MLC_LLM_SUPPORT_UTILS_H_
|
||||
#define MLC_LLM_SUPPORT_UTILS_H_
|
||||
|
||||
#include <tvm/support/io.h>
|
||||
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "../../3rdparty/tvm/src/support/base64.h"
|
||||
#include "../../3rdparty/tvm/src/support/bytes_io.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
/*! \brief Split the input string by the given delimiter character. */
|
||||
inline std::vector<std::string> Split(const std::string& str, char delim) {
|
||||
std::string item;
|
||||
std::istringstream is(str);
|
||||
std::vector<std::string> ret;
|
||||
while (std::getline(is, item, delim)) {
|
||||
ret.push_back(item);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Check whether the string starts with a given prefix.
|
||||
* \param str The given string.
|
||||
* \param prefix The given prefix.
|
||||
* \return Whether the prefix matched.
|
||||
*/
|
||||
inline bool StartsWith(const std::string& str, const char* prefix) {
|
||||
size_t n = str.length();
|
||||
for (size_t i = 0; i < n; i++) {
|
||||
if (prefix[i] == '\0') return true;
|
||||
if (str.data()[i] != prefix[i]) return false;
|
||||
}
|
||||
// return true if the str is equal to the prefix
|
||||
return prefix[n] == '\0';
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get the base64 encoded result of a string.
|
||||
* \param str The string to encode.
|
||||
* \return The base64 encoded string.
|
||||
*/
|
||||
inline std::string Base64Encode(std::string str) {
|
||||
std::string result;
|
||||
tvm::support::BytesOutStream m_stream(&result);
|
||||
tvm::support::Base64OutStream b64stream(&m_stream);
|
||||
static_cast<tvm::support::Stream*>(&b64stream)->Write(str);
|
||||
b64stream.Finish();
|
||||
return result;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get the base64 decoded result of a string.
|
||||
* \param str The string to decode.
|
||||
* \return The base64 decoded string.
|
||||
*/
|
||||
inline std::string Base64Decode(std::string str) {
|
||||
std::string result;
|
||||
tvm::support::BytesInStream m_stream(str);
|
||||
tvm::support::Base64InStream b64stream(&m_stream);
|
||||
b64stream.InitPosition();
|
||||
static_cast<tvm::support::Stream*>(&b64stream)->Read(&result);
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_UTILS_H_
|
||||
@@ -0,0 +1,59 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file support/image_utils.cc
|
||||
*/
|
||||
#include "vlm_utils.h"
|
||||
|
||||
#include <cmath>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
void CalculateResizeShape(tvm::runtime::Tensor image_data, std::string model_type,
|
||||
int* p_target_height, int* p_target_width) {
|
||||
TVM_FFI_ICHECK_EQ(image_data->shape[3], 3) << "Image format must be NHWC";
|
||||
int height = image_data->shape[1];
|
||||
int width = image_data->shape[2];
|
||||
|
||||
if ("phi3_v" == model_type) {
|
||||
const int hd_num = 4;
|
||||
double ratio = static_cast<double>(width) / height;
|
||||
int scale = 1;
|
||||
while (scale * std::ceil(scale / ratio) <= hd_num) {
|
||||
scale += 1;
|
||||
}
|
||||
scale -= 1;
|
||||
*p_target_width = static_cast<int>(scale * 336);
|
||||
*p_target_height = static_cast<int>(*p_target_width / ratio);
|
||||
}
|
||||
}
|
||||
|
||||
void CalculatePadShape(tvm::runtime::Tensor image_data, std::string model_type, int* p_pad_height,
|
||||
int* p_pad_width) {
|
||||
TVM_FFI_ICHECK_EQ(image_data->shape[3], 3) << "Image format must be NHWC";
|
||||
if ("phi3_v" == model_type) {
|
||||
int resized_height = 0, resized_width = 0;
|
||||
CalculateResizeShape(image_data, model_type, &resized_height, &resized_width);
|
||||
int tar = (int)(ceil(resized_height / 336.0) * 336);
|
||||
int top_padding = (int)((tar - resized_height) / 2);
|
||||
int bottom_padding = tar - resized_height - top_padding;
|
||||
TVM_FFI_ICHECK_EQ(tar, resized_height + top_padding + bottom_padding)
|
||||
<< "Padding size not equal!";
|
||||
*p_pad_height = tar;
|
||||
*p_pad_width = resized_width;
|
||||
}
|
||||
}
|
||||
|
||||
void CalculateCropShape(tvm::runtime::Tensor image_data, std::string model_type, int* p_crop_height,
|
||||
int* p_crop_width) {
|
||||
TVM_FFI_ICHECK_EQ(image_data->shape[3], 3) << "Image format must be NHWC";
|
||||
if ("phi3_v" == model_type) {
|
||||
int pad_h = 0, pad_w = 0;
|
||||
CalculatePadShape(image_data, model_type, &pad_h, &pad_w);
|
||||
*p_crop_height = pad_h / 336;
|
||||
*p_crop_width = pad_w / 336;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,48 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file support/vlm_utils.h
|
||||
* \brief Tools for debug purposes.
|
||||
*/
|
||||
#ifndef MLC_LLM_SUPPORT_VLM_UTILS_H_
|
||||
#define MLC_LLM_SUPPORT_VLM_UTILS_H_
|
||||
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
#include <string>
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
/*!
|
||||
* \brief Calculate the target height and width for resizing an image based on the input data and
|
||||
* model type. \param image_data The input image data as a TVM Tensor. \param model_type The type
|
||||
* of the model influencing the resizing parameters (e.g., phi3v). \param target_height Reference to
|
||||
* the variable where the calculated target height will be stored. \param target_width Reference to
|
||||
* the variable where the calculated target width will be stored.
|
||||
*/
|
||||
void CalculateResizeShape(tvm::runtime::Tensor image_data, std::string model_type,
|
||||
int* p_target_height, int* p_target_width);
|
||||
/*!
|
||||
* \brief Calculate the padding height and width for an image based on the input data and model
|
||||
* type. \param image_data The input image data as a TVM Tensor. \param model_type The type of the
|
||||
* model influencing the padding parameters (e.g., phi3v). \param pad_height Reference to the
|
||||
* variable where the calculated padding height will be stored. \param pad_width Reference to the
|
||||
* variable where the calculated padding width will be stored.
|
||||
*/
|
||||
void CalculatePadShape(tvm::runtime::Tensor image_data, std::string model_type, int* p_pad_height,
|
||||
int* p_pad_width);
|
||||
|
||||
/*!
|
||||
* \brief Calculate the cropping height and width for an image based on the input data and model
|
||||
* type. \param image_data The input image data as a TVM Tensor. \param model_type The type of the
|
||||
* model influencing the cropping parameters (e.g., phi3v). \param crop_height Reference to the
|
||||
* variable where the calculated cropping height will be stored. \param crop_width Reference to the
|
||||
* variable where the calculated cropping width will be stored.
|
||||
*/
|
||||
void CalculateCropShape(tvm::runtime::Tensor image_data, std::string model_type, int* p_crop_height,
|
||||
int* p_crop_width);
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_SUPPORT_IMAGE_UTILS_H_
|
||||
@@ -0,0 +1,292 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file streamer.cc
|
||||
*/
|
||||
|
||||
#include "streamer.h"
|
||||
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <string>
|
||||
|
||||
#include "tokenizers.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
TextStreamerObj::RegisterReflection();
|
||||
StopStrHandlerObj::RegisterReflection();
|
||||
}
|
||||
|
||||
/****************** TextStreamer ******************/
|
||||
|
||||
TextStreamerObj::TextStreamerObj(Tokenizer tokenizer) : tokenizer_(std::move(tokenizer)) {}
|
||||
|
||||
TextStreamer::TextStreamer(Tokenizer tokenizer) {
|
||||
data_ = tvm::ffi::make_object<TextStreamerObj>(std::move(tokenizer));
|
||||
}
|
||||
|
||||
std::string TextStreamerObj::Put(const std::vector<int32_t>& delta_tokens) {
|
||||
TVM_FFI_ICHECK(!finished_) << "`put` is not expected to be invoked after finish.";
|
||||
if (delta_tokens.empty()) {
|
||||
return "";
|
||||
}
|
||||
|
||||
std::string ret;
|
||||
// We process delta tokens one by one.
|
||||
for (int32_t delta_token : delta_tokens) {
|
||||
// push to pending tokens.
|
||||
pending_tokens_.push_back(delta_token);
|
||||
|
||||
// all_tokens = prefix_tokens_ + pending_tokens_
|
||||
std::vector<int32_t> all_tokens;
|
||||
all_tokens.reserve(prefix_tokens_.size() + pending_tokens_.size());
|
||||
all_tokens.insert(all_tokens.end(), prefix_tokens_.begin(), prefix_tokens_.end());
|
||||
all_tokens.insert(all_tokens.end(), pending_tokens_.begin(), pending_tokens_.end());
|
||||
|
||||
// Decode prefix_tokens_ and all_tokens.
|
||||
std::string prefix_str = prefix_tokens_.empty() ? "" : tokenizer_->Decode(prefix_tokens_);
|
||||
std::string full_str = tokenizer_->Decode(all_tokens);
|
||||
|
||||
std::string validated_str;
|
||||
std::vector<int32_t> new_pending_tokens;
|
||||
if (full_str.compare(0, prefix_str.length(), prefix_str) == 0) {
|
||||
// Case 1. prefix_str is a prefix of `full_str`.
|
||||
// validated_str = full_str[len(prefix_str):]
|
||||
validated_str = full_str.substr(prefix_str.length());
|
||||
// Pop UTF-8 replacement character from the back of pending tokens.
|
||||
// - The UTF-8 replacement character take 3 chars.
|
||||
// - A valid UTF-8 has 4 chars at most.
|
||||
// So there will be at most 3 tokens popped.
|
||||
while (!pending_tokens_.empty() && //
|
||||
static_cast<int>(new_pending_tokens.size()) < 3 && //
|
||||
validated_str.length() >= 3 && //
|
||||
validated_str.compare(validated_str.length() - 3, /*n=*/3, kReplacementCharacter) ==
|
||||
0) {
|
||||
new_pending_tokens.push_back(pending_tokens_.back());
|
||||
pending_tokens_.pop_back();
|
||||
all_tokens.pop_back();
|
||||
validated_str = tokenizer_->Decode(all_tokens).substr(prefix_str.length());
|
||||
}
|
||||
} else {
|
||||
// Case 2. prefix_str is not a prefix of `full_str`.
|
||||
// Pop pending tokens from the back.
|
||||
// - Pop until prefix_str is indeed a prefix of full_str.
|
||||
// - A valid UTF-8 has 4 chars at most.
|
||||
// So there will be at most 3 tokens popped.
|
||||
// - If there are no more than 3 pending tokens, skip popping.
|
||||
// This is because it is impossible to make full_str contain
|
||||
// prefix_str without popping all the pending tokens.
|
||||
if (static_cast<int>(pending_tokens_.size()) < 3) {
|
||||
continue;
|
||||
}
|
||||
bool get_valid_full_str = false;
|
||||
while (!pending_tokens_.empty() && static_cast<int>(new_pending_tokens.size()) < 3) {
|
||||
new_pending_tokens.push_back(pending_tokens_.back());
|
||||
pending_tokens_.pop_back();
|
||||
all_tokens.pop_back();
|
||||
full_str = tokenizer_->Decode(all_tokens);
|
||||
if (full_str.compare(0, prefix_str.length(), prefix_str) == 0) {
|
||||
get_valid_full_str = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (get_valid_full_str) {
|
||||
// We find a full_str which starts from prefix_str.
|
||||
// So we return the sliced full string without the prefix.
|
||||
validated_str = full_str.substr(prefix_str.length());
|
||||
} else {
|
||||
// We cannot find a full_str which starts from prefix_str by
|
||||
// popping 3 tokens.
|
||||
// In this case, the remaining pending tokens are invalid UTF-8
|
||||
// characters already, so we return the decoded pending tokens.
|
||||
validated_str = tokenizer_->Decode(pending_tokens_);
|
||||
}
|
||||
}
|
||||
|
||||
if (!pending_tokens_.empty()) {
|
||||
// Set the new prefix.
|
||||
prefix_tokens_ = pending_tokens_;
|
||||
}
|
||||
std::reverse(new_pending_tokens.begin(), new_pending_tokens.end());
|
||||
pending_tokens_ = new_pending_tokens;
|
||||
ret += validated_str;
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
std::string TextStreamerObj::Finish() {
|
||||
// all_tokens = prefix_tokens_ + pending_tokens_
|
||||
std::vector<int32_t> all_tokens;
|
||||
all_tokens.reserve(prefix_tokens_.size() + pending_tokens_.size());
|
||||
all_tokens.insert(all_tokens.end(), prefix_tokens_.begin(), prefix_tokens_.end());
|
||||
all_tokens.insert(all_tokens.end(), pending_tokens_.begin(), pending_tokens_.end());
|
||||
|
||||
// Decode prefix_tokens_ and all_tokens.
|
||||
std::string prefix_str = prefix_tokens_.empty() ? "" : tokenizer_->Decode(prefix_tokens_);
|
||||
std::string full_str = all_tokens.empty() ? "" : tokenizer_->Decode(all_tokens);
|
||||
|
||||
finished_ = true;
|
||||
if (full_str.compare(0, prefix_str.length(), prefix_str) == 0) {
|
||||
// Case 1. prefix_str is a prefix of `full_str`.
|
||||
return full_str.substr(prefix_str.length());
|
||||
} else {
|
||||
// Case 2. prefix_str is not a prefix of `full_str`.
|
||||
// In this case, the remaining pending tokens are invalid UTF-8
|
||||
// characters already, so we return the decoded pending tokens.
|
||||
return tokenizer_->Decode(pending_tokens_);
|
||||
}
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef()
|
||||
.def("mlc.tokenizers.TextStreamer",
|
||||
[](Tokenizer tokenizer) { return TextStreamer(std::move(tokenizer)); })
|
||||
.def("mlc.tokenizers.TextStreamerPut",
|
||||
[](TextStreamer text_streamer, const Shape& delta_tokens) {
|
||||
return text_streamer->Put(
|
||||
{delta_tokens->data, delta_tokens->data + delta_tokens->size});
|
||||
})
|
||||
.def_method("mlc.tokenizers.TextStreamerFinish", &TextStreamerObj::Finish);
|
||||
}
|
||||
|
||||
/****************** StopStrHandler ******************/
|
||||
|
||||
/*! \brief Create the KMP partial match table for the input string. */
|
||||
inline std::vector<int> CreatePartialMatchTable(const String& str) {
|
||||
int length = str.length();
|
||||
std::vector<int> partial_match_table = {-1};
|
||||
partial_match_table.reserve(length);
|
||||
for (int i = 1; i < length; ++i) {
|
||||
int ptr = partial_match_table[i - 1];
|
||||
while (ptr != -1 && str.at(ptr) != str.at(i - 1)) {
|
||||
ptr = partial_match_table[ptr];
|
||||
}
|
||||
partial_match_table.push_back(ptr + 1);
|
||||
}
|
||||
return partial_match_table;
|
||||
}
|
||||
|
||||
StopStrHandlerObj::StopStrHandlerObj(Array<String> stop_strs,
|
||||
const std::vector<std::string>& token_table)
|
||||
: stop_strs_(std::move(stop_strs)), token_table_(token_table) {
|
||||
int num_stop_strs = stop_strs_.size();
|
||||
cur_match_lengths_.resize(num_stop_strs, 0);
|
||||
|
||||
// Create the KMP partial match table for each stop string.
|
||||
partial_match_tables_.reserve(num_stop_strs);
|
||||
for (const String& stop_str : stop_strs_) {
|
||||
TVM_FFI_ICHECK(!stop_str.empty()) << "Stop string cannot be empty.";
|
||||
partial_match_tables_.push_back(CreatePartialMatchTable(stop_str));
|
||||
}
|
||||
}
|
||||
|
||||
void StopStrHandlerObj::Put(int32_t token_id, std::vector<int64_t>* return_token_ids) {
|
||||
TVM_FFI_ICHECK_NOTNULL(return_token_ids);
|
||||
|
||||
// Return the input token id if there is no stop string.
|
||||
if (stop_strs_.empty()) {
|
||||
return_token_ids->push_back(token_id);
|
||||
return;
|
||||
}
|
||||
|
||||
TVM_FFI_ICHECK(!stop_triggered_) << "Cannot put new token when already stopped.";
|
||||
|
||||
TVM_FFI_ICHECK_LT(token_id, static_cast<int>(token_table_.size()));
|
||||
const std::string& token = token_table_[token_id];
|
||||
pending_token_ids_.push_back(token_id);
|
||||
pending_token_lengths_.push_back(token.length());
|
||||
|
||||
for (char ch : token) {
|
||||
// The earliest starting point of stop string.
|
||||
int stop_starting_pos = std::numeric_limits<int>::max();
|
||||
// The cutoff length that can be safely return.
|
||||
int cutoff_length = std::numeric_limits<int>::max();
|
||||
// The maximum matched length.
|
||||
int max_match_length = 0;
|
||||
|
||||
for (int str_id = 0; str_id < static_cast<int>(stop_strs_.size()); ++str_id) {
|
||||
// - Run one step of KMP algorithm.
|
||||
const std::vector<int>& partial_match_table = partial_match_tables_[str_id];
|
||||
int& cur_match_length = cur_match_lengths_[str_id];
|
||||
while (cur_match_length != -1 && ch != stop_strs_[str_id].at(cur_match_length)) {
|
||||
cur_match_length = partial_match_table[cur_match_length];
|
||||
}
|
||||
++cur_match_length;
|
||||
|
||||
// Case 1. The stop string is matched.
|
||||
if (cur_match_length == stop_strs_[str_id].length()) {
|
||||
stop_triggered_ = true;
|
||||
stop_starting_pos =
|
||||
std::min(stop_starting_pos,
|
||||
pending_string_len_ + 1 - static_cast<int>(stop_strs_[str_id].length()));
|
||||
continue;
|
||||
}
|
||||
|
||||
// Case 2. The stop string is not matched.
|
||||
// - Get the cutoff length that can be safely return.
|
||||
TVM_FFI_ICHECK_GE(pending_string_len_ + 1, cur_match_length);
|
||||
cutoff_length = std::min(cutoff_length, pending_string_len_ + 1 - cur_match_length);
|
||||
// - Get the updated pending string length.
|
||||
max_match_length = std::max(max_match_length, cur_match_length);
|
||||
}
|
||||
|
||||
// Collect the token ids that can be safely cut off and returned.
|
||||
if (stop_triggered_) {
|
||||
cutoff_length = stop_starting_pos;
|
||||
}
|
||||
TVM_FFI_ICHECK_NE(cutoff_length, std::numeric_limits<int>::max());
|
||||
TVM_FFI_ICHECK_GE(cutoff_length, 0);
|
||||
int cum_length = 0;
|
||||
while (!pending_token_ids_.empty() &&
|
||||
cum_length + pending_token_lengths_.front() <= cutoff_length) {
|
||||
cum_length += pending_token_lengths_.front();
|
||||
return_token_ids->push_back(pending_token_ids_.front());
|
||||
pending_token_ids_.erase(pending_token_ids_.begin());
|
||||
pending_token_lengths_.erase(pending_token_lengths_.begin());
|
||||
}
|
||||
if (stop_triggered_) {
|
||||
return;
|
||||
}
|
||||
|
||||
TVM_FFI_ICHECK_LE(cum_length, cutoff_length);
|
||||
// `cum_length` is the prefix length what we actually cut off.
|
||||
pending_string_len_ = (cutoff_length - cum_length) + max_match_length;
|
||||
}
|
||||
}
|
||||
|
||||
StopStrHandler::StopStrHandler(Array<String> stop_strs,
|
||||
const std::vector<std::string>& token_table) {
|
||||
data_ = tvm::ffi::make_object<StopStrHandlerObj>(std::move(stop_strs), token_table);
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef()
|
||||
.def("mlc.tokenizers.StopStrHandler",
|
||||
[](Array<String> stop_strs, const Tokenizer& tokenizer) {
|
||||
return StopStrHandler(std::move(stop_strs), tokenizer->PostProcessedTokenTable());
|
||||
})
|
||||
.def("mlc.tokenizers.StopStrHandlerPut",
|
||||
[](StopStrHandler handler, int token_id) {
|
||||
std::vector<int64_t> delta_tokens;
|
||||
handler->Put(token_id, &delta_tokens);
|
||||
return Shape(std::move(delta_tokens));
|
||||
})
|
||||
.def("mlc.tokenizers.StopStringHandlerFinish",
|
||||
[](StopStrHandler handler) {
|
||||
std::vector<int64_t> remaining_token_ids;
|
||||
handler->Finish(&remaining_token_ids);
|
||||
return Shape(std::move(remaining_token_ids));
|
||||
})
|
||||
.def_method("mlc.tokenizers.StopStrHandlerStopTriggered", &StopStrHandlerObj::StopTriggered);
|
||||
}
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,160 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file streamer.h
|
||||
* \brief Header of streamers in MLC LLM.
|
||||
*/
|
||||
|
||||
#ifndef MLC_LLM_STREAMER_H_
|
||||
#define MLC_LLM_STREAMER_H_
|
||||
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
|
||||
#include "tokenizers.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectRef;
|
||||
|
||||
/****************** TextStreamer ******************/
|
||||
|
||||
/*!
|
||||
* \brief The class that streams back validated utf-8 text strings
|
||||
* that generated by tokenizer.
|
||||
*/
|
||||
class TextStreamerObj : public Object {
|
||||
public:
|
||||
explicit TextStreamerObj(Tokenizer tokenizer);
|
||||
|
||||
/*!
|
||||
* \brief Put new delta tokens into the streamer, and get the UTF-8-valid
|
||||
* delta string. The text streamer may hold some of the input delta tokens
|
||||
* which cannot decode into valid UTF-8 strings. The returned string
|
||||
* is always guaranteed to be UTF-8 valid.
|
||||
* \param delta_tokens The new tokens to put into the streamer.
|
||||
* \return The decoded delta string after putting the input new tokens.
|
||||
*/
|
||||
std::string Put(const std::vector<int32_t>& delta_tokens);
|
||||
|
||||
/*! \brief Return the string decoded by remaining tokens. */
|
||||
std::string Finish();
|
||||
|
||||
// REPLACEMENT CHARACTER (U+FFFD) in UTF-8.
|
||||
static constexpr const char* kReplacementCharacter = "\xef\xbf\xbd";
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<TextStreamerObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.TextStreamer", TextStreamerObj, Object);
|
||||
|
||||
private:
|
||||
Tokenizer tokenizer_;
|
||||
std::vector<int32_t> prefix_tokens_;
|
||||
std::vector<int32_t> pending_tokens_;
|
||||
bool finished_ = false;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Managed reference to TextStreamerObj
|
||||
* \sa TextStreamerObj
|
||||
*/
|
||||
class TextStreamer : public ObjectRef {
|
||||
public:
|
||||
/*! \brief Construct a text streamer with tokenizer. */
|
||||
explicit TextStreamer(Tokenizer tokenizer);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TextStreamer, ObjectRef, TextStreamerObj);
|
||||
};
|
||||
|
||||
/****************** StopStrHandler ******************/
|
||||
|
||||
/*!
|
||||
* \brief The stop string handler in MLC LLM, which takes input delta tokens
|
||||
* one at a time, and return the output delta token before stopping due to
|
||||
* stop strings.
|
||||
*/
|
||||
class StopStrHandlerObj : public Object {
|
||||
public:
|
||||
explicit StopStrHandlerObj(Array<String> stop_strs, const std::vector<std::string>& token_table);
|
||||
|
||||
/*!
|
||||
* \brief Add new input delta token to the handler, push the output
|
||||
* delta tokens before stopping into the given vector.
|
||||
* The stop string handler may hold some of the input delta token
|
||||
* which may be part of a stop string.
|
||||
* The returned tokens are always guaranteed not to be part of stop string.
|
||||
*/
|
||||
void Put(int32_t token_id, std::vector<int64_t>* return_token_ids);
|
||||
|
||||
/*!
|
||||
* \brief Stop string handling has finished, append the remaining
|
||||
* cached token ids into the given vector.
|
||||
*/
|
||||
void Finish(std::vector<int64_t>* return_token_ids) const {
|
||||
return_token_ids->insert(return_token_ids->end(), pending_token_ids_.begin(),
|
||||
pending_token_ids_.end());
|
||||
};
|
||||
|
||||
/*! \brief Check if the generation has stopped due to stop string. */
|
||||
bool StopTriggered() const { return stop_triggered_; }
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<StopStrHandlerObj>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.StopStrHandler", StopStrHandlerObj, Object);
|
||||
|
||||
private:
|
||||
/*! \brief The stop strings. */
|
||||
Array<String> stop_strs_;
|
||||
/*! \brief The partial match table for each stop string in the KMP algorithm. */
|
||||
std::vector<std::vector<int>> partial_match_tables_;
|
||||
/*! \brief The tokenizer token table for token id lookup. */
|
||||
const std::vector<std::string>& token_table_;
|
||||
|
||||
/************ Global states across all stop strings. ************/
|
||||
|
||||
/*! \brief The globally pending string length. */
|
||||
int pending_string_len_ = 0;
|
||||
/*! \brief The globally pending token ids. */
|
||||
std::vector<int32_t> pending_token_ids_;
|
||||
/*! \brief The token string length of each pending token id. */
|
||||
std::vector<int> pending_token_lengths_;
|
||||
/*! \brief A boolean flag indicating if stop has been triggered. */
|
||||
bool stop_triggered_ = false;
|
||||
|
||||
/************ Per-stop-string states. ************/
|
||||
|
||||
/*! \brief The current match position of the pending string to each stop string. */
|
||||
std::vector<int> cur_match_lengths_;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Managed reference to StopStrHandlerObj
|
||||
* \sa StopStrHandlerObj
|
||||
*/
|
||||
class StopStrHandler : public ObjectRef {
|
||||
public:
|
||||
explicit StopStrHandler(Array<String> stop_strs, const std::vector<std::string>& token_table);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(StopStrHandler, ObjectRef, StopStrHandlerObj);
|
||||
};
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_STREAMER_H_
|
||||
@@ -0,0 +1,535 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file tokenizer.cc
|
||||
*/
|
||||
|
||||
#include "tokenizers.h"
|
||||
|
||||
#include <tokenizers_cpp.h>
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/extra/json.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/runtime/logging.h>
|
||||
|
||||
#include <array>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
|
||||
#include "./../support/encoding.h"
|
||||
#include "./../support/load_bytes_from_file.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
TokenizerInfoNode::RegisterReflection();
|
||||
TokenizerObj::RegisterReflection();
|
||||
}
|
||||
|
||||
#ifndef COMPILE_MLC_WASM_RUNTIME
|
||||
|
||||
String TokenizerInfoNode::AsJSONString() const {
|
||||
tvm::ffi::json::Object obj;
|
||||
obj.Set("token_postproc_method", token_postproc_method);
|
||||
obj.Set("prepend_space_in_encode", prepend_space_in_encode);
|
||||
obj.Set("strip_space_in_decode", strip_space_in_decode);
|
||||
return tvm::ffi::json::Stringify(obj);
|
||||
}
|
||||
|
||||
TokenizerInfo TokenizerInfo::FromJSONString(String json_string) {
|
||||
tvm::ffi::String err;
|
||||
auto v = tvm::ffi::json::Parse(json_string, &err);
|
||||
TVM_FFI_ICHECK(err.empty()) << "Failed to parse JSON: " << err;
|
||||
|
||||
TVM_FFI_ICHECK(v.try_cast<tvm::ffi::json::Object>().has_value()) << "JSON must be an object.";
|
||||
const auto& obj = v.cast<tvm::ffi::json::Object>();
|
||||
|
||||
ObjectPtr<TokenizerInfoNode> n = tvm::ffi::make_object<TokenizerInfoNode>();
|
||||
if (obj.count("token_postproc_method")) {
|
||||
TVM_FFI_ICHECK(obj.at("token_postproc_method").try_cast<tvm::ffi::String>().has_value());
|
||||
n->token_postproc_method = obj.at("token_postproc_method").cast<tvm::ffi::String>();
|
||||
}
|
||||
if (obj.count("prepend_space_in_encode")) {
|
||||
TVM_FFI_ICHECK(obj.at("prepend_space_in_encode").try_cast<bool>().has_value());
|
||||
n->prepend_space_in_encode = obj.at("prepend_space_in_encode").cast<bool>();
|
||||
}
|
||||
if (obj.count("strip_space_in_decode")) {
|
||||
TVM_FFI_ICHECK(obj.at("strip_space_in_decode").try_cast<bool>().has_value());
|
||||
n->strip_space_in_decode = obj.at("strip_space_in_decode").cast<bool>();
|
||||
}
|
||||
|
||||
return TokenizerInfo(n);
|
||||
}
|
||||
|
||||
Tokenizer::Tokenizer(std::unique_ptr<tokenizers::Tokenizer> tokenizer, TokenizerInfo info) {
|
||||
ObjectPtr<TokenizerObj> n = tvm::ffi::make_object<TokenizerObj>();
|
||||
n->tokenizer = std::move(tokenizer);
|
||||
n->info_ = std::move(info);
|
||||
data_ = std::move(n);
|
||||
}
|
||||
|
||||
std::vector<int32_t> TokenizerObj::Encode(const std::string& text) const {
|
||||
return tokenizer->Encode(text);
|
||||
}
|
||||
|
||||
std::vector<int32_t> TokenizerObj::EncodeNoPrependSpace(const std::string& text) const {
|
||||
// TODO(yixin): now this only supports tokenizers with tokenizer.json
|
||||
// other tokenizers should be supported.
|
||||
static const constexpr char* kPaddingPrefix = "\x01";
|
||||
if (!info_->prepend_space_in_encode) {
|
||||
return tokenizer->Encode(text);
|
||||
}
|
||||
|
||||
auto result = tokenizer->Encode(kPaddingPrefix + text);
|
||||
// remove the first two tokens: "▁" and "<0x01>"
|
||||
result.erase(result.begin(), result.begin() + 2);
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<std::vector<int32_t>> TokenizerObj::EncodeBatch(const Array<String>& texts) const {
|
||||
std::vector<std::string> texts_vec;
|
||||
for (const String& text : texts) {
|
||||
texts_vec.push_back(text);
|
||||
}
|
||||
return tokenizer->EncodeBatch(texts_vec);
|
||||
}
|
||||
|
||||
std::string TokenizerObj::Decode(const std::vector<int32_t>& token_ids) const {
|
||||
return tokenizer->Decode(token_ids);
|
||||
}
|
||||
|
||||
const DynamicBitset& TokenizerObj::GetPrefixTokenMask() {
|
||||
if (prefix_token_mask_.Size() != 0) {
|
||||
return prefix_token_mask_;
|
||||
}
|
||||
|
||||
int vocab_size = GetVocabSize();
|
||||
prefix_token_mask_ = DynamicBitset(vocab_size);
|
||||
|
||||
// Sort all tokens
|
||||
const auto& token_table = PostProcessedTokenTable();
|
||||
std::vector<std::pair<std::string, int>> sorted_tokens;
|
||||
for (int32_t token_id = 0; token_id < vocab_size; ++token_id) {
|
||||
sorted_tokens.emplace_back(token_table[token_id], token_id);
|
||||
}
|
||||
std::sort(sorted_tokens.begin(), sorted_tokens.end());
|
||||
|
||||
// Check every token if it is a prefix of another token
|
||||
for (int idx = 0; idx < vocab_size - 1; ++idx) {
|
||||
auto cur_token = sorted_tokens[idx].first;
|
||||
auto nxt_token = sorted_tokens[idx + 1].first;
|
||||
if (cur_token.length() <= nxt_token.length() &&
|
||||
std::string_view(nxt_token).substr(0, cur_token.length()) == cur_token) {
|
||||
prefix_token_mask_.Set(sorted_tokens[idx].second);
|
||||
}
|
||||
}
|
||||
|
||||
return prefix_token_mask_;
|
||||
}
|
||||
|
||||
size_t TokenizerObj::GetVocabSize() const { return tokenizer->GetVocabSize(); }
|
||||
|
||||
std::string TokenizerObj::IdToToken(int32_t token_id) const {
|
||||
return tokenizer->IdToToken(token_id);
|
||||
}
|
||||
|
||||
int32_t TokenizerObj::TokenToId(const std::string& token) const {
|
||||
return tokenizer->TokenToId(token);
|
||||
}
|
||||
|
||||
Tokenizer Tokenizer::FromPath(const String& _path, std::optional<TokenizerInfo> info) {
|
||||
TokenizerInfo info_value = info.value_or(DetectTokenizerInfo(_path));
|
||||
std::filesystem::path path{std::string(_path)};
|
||||
std::filesystem::path sentencepiece;
|
||||
std::filesystem::path huggingface;
|
||||
std::filesystem::path rwkvworld;
|
||||
TVM_FFI_ICHECK(std::filesystem::exists(path)) << "Cannot find tokenizer via path: " << _path;
|
||||
if (std::filesystem::is_directory(path)) {
|
||||
sentencepiece = path / "tokenizer.model";
|
||||
huggingface = path / "tokenizer.json";
|
||||
rwkvworld = path / "tokenizer_model";
|
||||
} else {
|
||||
sentencepiece = path.parent_path() / "tokenizer.model";
|
||||
huggingface = path.parent_path() / "tokenizer.json";
|
||||
rwkvworld = path.parent_path() / "tokenizer_model";
|
||||
}
|
||||
if (std::filesystem::exists(huggingface)) {
|
||||
// Check HuggingFace
|
||||
return Tokenizer(tokenizers::Tokenizer::FromBlobJSON(LoadBytesFromFile(huggingface.string())),
|
||||
info_value);
|
||||
}
|
||||
if (std::filesystem::exists(sentencepiece)) {
|
||||
// Check SentencePiece
|
||||
LOG(WARNING)
|
||||
<< "Using `tokenizer.model` since we cannot locate `tokenizer.json`.\n"
|
||||
<< "It is recommended to use `tokenizer.json` to ensure all token mappings are included, "
|
||||
<< "since currently, files like `added_tokens.json`, `tokenizer_config.json` are ignored.\n"
|
||||
<< "Consider converting `tokenizer.model` to `tokenizer.json` by compiling the model "
|
||||
<< "with MLC again, or see if MLC's huggingface provides this file.";
|
||||
return Tokenizer(
|
||||
tokenizers::Tokenizer::FromBlobSentencePiece(LoadBytesFromFile(sentencepiece.string())),
|
||||
info_value);
|
||||
}
|
||||
{
|
||||
// Check ByteLevelBPE
|
||||
std::filesystem::path merges_path = path / "merges.txt";
|
||||
std::filesystem::path vocab_path = path / "vocab.json";
|
||||
std::filesystem::path added_tokens_path = path / "added_tokens.json";
|
||||
if (std::filesystem::exists(merges_path) && std::filesystem::exists(vocab_path) &&
|
||||
std::filesystem::exists(added_tokens_path)) {
|
||||
std::string vocab = LoadBytesFromFile(vocab_path.string());
|
||||
std::string merges = LoadBytesFromFile(merges_path.string());
|
||||
std::string added_tokens = LoadBytesFromFile(added_tokens_path.string());
|
||||
return Tokenizer(tokenizers::Tokenizer::FromBlobByteLevelBPE(vocab, merges, added_tokens),
|
||||
info_value);
|
||||
}
|
||||
}
|
||||
if (std::filesystem::exists(rwkvworld)) {
|
||||
// Check RWKV
|
||||
return Tokenizer(tokenizers::Tokenizer::FromBlobRWKVWorld(rwkvworld.string()), info_value);
|
||||
}
|
||||
LOG(FATAL) << "Cannot find any tokenizer under: " << _path;
|
||||
}
|
||||
|
||||
TokenizerInfo Tokenizer::DetectTokenizerInfo(const String& path_str) {
|
||||
std::filesystem::path path{std::string(path_str)};
|
||||
TVM_FFI_ICHECK(std::filesystem::exists(path)) << "Cannot find tokenizer via path: " << path_str;
|
||||
if (!std::filesystem::is_directory(path)) {
|
||||
path = path.parent_path();
|
||||
}
|
||||
path = path / "tokenizer.json";
|
||||
if (!std::filesystem::exists(path)) {
|
||||
LOG(WARNING) << "Tokenizer info is not detected as tokenizer.json is not found. The default "
|
||||
<< "tokenizer info will be used.";
|
||||
return TokenizerInfo(tvm::ffi::make_object<TokenizerInfoNode>());
|
||||
}
|
||||
|
||||
std::string tokenizer_json = LoadBytesFromFile(path.string());
|
||||
tvm::ffi::String err;
|
||||
auto v = tvm::ffi::json::Parse(tokenizer_json, &err);
|
||||
TVM_FFI_ICHECK(err.empty()) << "Failed to parse JSON: " << err;
|
||||
TVM_FFI_ICHECK(v.try_cast<tvm::ffi::json::Object>().has_value()) << "JSON must be an object.";
|
||||
const auto& obj = v.cast<tvm::ffi::json::Object>();
|
||||
|
||||
ObjectPtr<TokenizerInfoNode> n = tvm::ffi::make_object<TokenizerInfoNode>();
|
||||
|
||||
// Step 1. Detect token_postproc_method: byte_fallback or byte_level
|
||||
// Detect {"type": "ByteLevel"} or {"type": "ByteFallback"} in "decoder" field of the tokenizer
|
||||
if (!obj.count("decoder") || !obj.at("decoder").try_cast<tvm::ffi::json::Object>().has_value()) {
|
||||
LOG(WARNING) << "Decoder field is not found in tokenizer.json. Use ByteFallback as default.";
|
||||
n->token_postproc_method = "byte_fallback";
|
||||
} else {
|
||||
auto decoder_obj = obj.at("decoder").cast<tvm::ffi::json::Object>();
|
||||
TVM_FFI_ICHECK(decoder_obj.count("type") &&
|
||||
decoder_obj.at("type").try_cast<tvm::ffi::String>().has_value());
|
||||
auto type = decoder_obj.at("type").cast<tvm::ffi::String>();
|
||||
|
||||
auto f_detect_decoder_type = [](ObjectPtr<TokenizerInfoNode> n,
|
||||
const tvm::ffi::json::Value& decoder_json) {
|
||||
TVM_FFI_ICHECK(decoder_json.try_cast<tvm::ffi::json::Object>().has_value());
|
||||
TVM_FFI_ICHECK(decoder_json.cast<tvm::ffi::json::Object>().count("type") &&
|
||||
decoder_json.cast<tvm::ffi::json::Object>()
|
||||
.at("type")
|
||||
.try_cast<tvm::ffi::String>()
|
||||
.has_value());
|
||||
auto type = decoder_json.cast<tvm::ffi::json::Object>().at("type").cast<tvm::ffi::String>();
|
||||
if (type == "ByteLevel") {
|
||||
n->token_postproc_method = "byte_level";
|
||||
return true;
|
||||
} else if (type == "ByteFallback") {
|
||||
n->token_postproc_method = "byte_fallback";
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
};
|
||||
|
||||
bool found = false;
|
||||
|
||||
// For sequence, examine every decoder
|
||||
if (type == "Sequence") {
|
||||
TVM_FFI_ICHECK(decoder_obj.count("decoders") &&
|
||||
decoder_obj.at("decoders").try_cast<tvm::ffi::json::Array>().has_value());
|
||||
for (const tvm::ffi::json::Value& decoder :
|
||||
decoder_obj.at("decoders").cast<tvm::ffi::json::Array>()) {
|
||||
if (f_detect_decoder_type(n, decoder)) {
|
||||
found = true;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (f_detect_decoder_type(n, obj.at("decoder"))) {
|
||||
found = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (!found) {
|
||||
LOG(WARNING) << "Neither ByteLevel nor ByteFallback decoder is detected in tokenizer.json. "
|
||||
<< "Use ByteFallback as default.";
|
||||
n->token_postproc_method = "byte_fallback";
|
||||
}
|
||||
}
|
||||
|
||||
// Step 2. Detect prepend_space_in_encode
|
||||
// Find {"type": "Prepend", "prepend": "▁"} in "normalizer" field of the tokenizer
|
||||
if (obj.count("normalizer") &&
|
||||
obj.at("normalizer").try_cast<tvm::ffi::json::Object>().has_value()) {
|
||||
const tvm::ffi::json::Value& normalizer_json = obj.at("normalizer");
|
||||
|
||||
auto f_handle_normalizer = [](ObjectPtr<TokenizerInfoNode> n,
|
||||
const tvm::ffi::json::Value& normalizer_json) {
|
||||
TVM_FFI_ICHECK(normalizer_json.try_cast<tvm::ffi::json::Object>().has_value());
|
||||
auto obj = normalizer_json.cast<tvm::ffi::json::Object>();
|
||||
TVM_FFI_ICHECK(obj.count("type") && obj.at("type").try_cast<tvm::ffi::String>().has_value());
|
||||
if (obj.at("type").cast<tvm::ffi::String>() == "Prepend" && obj.count("prepend") &&
|
||||
obj.at("prepend").try_cast<tvm::ffi::String>().has_value() &&
|
||||
obj.at("prepend").cast<tvm::ffi::String>() == "\xe2\x96\x81") {
|
||||
n->prepend_space_in_encode = true;
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
};
|
||||
|
||||
auto type = normalizer_json.cast<tvm::ffi::json::Object>().at("type").cast<tvm::ffi::String>();
|
||||
if (type == "Sequence") {
|
||||
TVM_FFI_ICHECK(normalizer_json.cast<tvm::ffi::json::Object>().count("normalizers") &&
|
||||
normalizer_json.cast<tvm::ffi::json::Object>()
|
||||
.at("normalizers")
|
||||
.try_cast<tvm::ffi::json::Array>()
|
||||
.has_value());
|
||||
for (const tvm::ffi::json::Value& normalizer : normalizer_json.cast<tvm::ffi::json::Object>()
|
||||
.at("normalizers")
|
||||
.cast<tvm::ffi::json::Array>()) {
|
||||
if (f_handle_normalizer(n, normalizer)) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
f_handle_normalizer(n, normalizer_json);
|
||||
}
|
||||
}
|
||||
|
||||
// Step 3. Detect strip_space_in_decode
|
||||
// Find {"type": "Strip", "content": " ", "start": 1, "stop": 0} in "decoder" field of the
|
||||
// tokenizer
|
||||
if (obj.count("decoder") && obj.at("decoder").try_cast<tvm::ffi::json::Object>().has_value()) {
|
||||
const tvm::ffi::json::Value& decoders_json = obj.at("decoder");
|
||||
|
||||
auto f_handle_decoder = [](ObjectPtr<TokenizerInfoNode> n,
|
||||
const tvm::ffi::json::Value& decoder_json) {
|
||||
TVM_FFI_ICHECK(decoder_json.try_cast<tvm::ffi::json::Object>().has_value());
|
||||
auto obj = decoder_json.cast<tvm::ffi::json::Object>();
|
||||
TVM_FFI_ICHECK(obj.count("type") && obj.at("type").try_cast<tvm::ffi::String>().has_value());
|
||||
if (obj.at("type").cast<tvm::ffi::String>() == "Strip" && obj.count("content") &&
|
||||
obj.at("content").try_cast<tvm::ffi::String>().has_value() &&
|
||||
obj.at("content").cast<tvm::ffi::String>() == " " && obj.count("start") &&
|
||||
obj.at("start").try_cast<int64_t>().has_value() && obj.at("start").cast<int64_t>() == 1 &&
|
||||
obj.count("stop") && obj.at("stop").try_cast<int64_t>().has_value() &&
|
||||
obj.at("stop").cast<int64_t>() == 0) {
|
||||
n->strip_space_in_decode = true;
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
};
|
||||
|
||||
auto type = decoders_json.cast<tvm::ffi::json::Object>().at("type").cast<tvm::ffi::String>();
|
||||
if (type == "Sequence") {
|
||||
TVM_FFI_ICHECK(decoders_json.cast<tvm::ffi::json::Object>().count("decoders") &&
|
||||
decoders_json.cast<tvm::ffi::json::Object>()
|
||||
.at("decoders")
|
||||
.try_cast<tvm::ffi::json::Array>()
|
||||
.has_value());
|
||||
for (const tvm::ffi::json::Value& decoder : decoders_json.cast<tvm::ffi::json::Object>()
|
||||
.at("decoders")
|
||||
.cast<tvm::ffi::json::Array>()) {
|
||||
if (f_handle_decoder(n, decoder)) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
f_handle_decoder(n, decoders_json);
|
||||
}
|
||||
}
|
||||
|
||||
return TokenizerInfo(n);
|
||||
}
|
||||
#endif
|
||||
|
||||
/*! \brief ByteFallback decoder: transform tokens like <0x1B> to hex char byte 1B */
|
||||
inline std::string ByteFallbackDecoder(const std::string& token) {
|
||||
if (token.length() == 6 && token.substr(0, 3) == "<0x" && token.back() == '>') {
|
||||
int byte = 0;
|
||||
for (int i = 0; i < 2; ++i) {
|
||||
byte *= 16;
|
||||
byte +=
|
||||
token[3 + i] >= '0' && token[3 + i] <= '9' ? token[3 + i] - '0' : token[3 + i] - 'A' + 10;
|
||||
}
|
||||
TVM_FFI_ICHECK(byte >= 0 && byte < 256);
|
||||
return std::string(/*n=*/1, static_cast<char>(byte));
|
||||
}
|
||||
return token;
|
||||
}
|
||||
|
||||
/*! \brief SpaceReplacer decoder: transform "\u2581" back to space */
|
||||
inline std::string SpaceReplacerDecoder(const std::string& token) {
|
||||
// \u2581 is the unicode for "lower one eighth block"
|
||||
// UTF8 encoding for \u2581 is 0xE2 0x96 0x81
|
||||
std::string result;
|
||||
for (size_t i = 0; i < token.size(); ++i) {
|
||||
if (i + 2 < token.size() && token[i] == char(0xE2) && token[i + 1] == char(0x96) &&
|
||||
token[i + 2] == char(0x81)) {
|
||||
result += ' ';
|
||||
i += 2;
|
||||
} else {
|
||||
result += token[i];
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
/*! \brief ByteLevel decoder: inverses the bytes-to-unicode transformation in the encoding
|
||||
* process as in
|
||||
* https://github.com/huggingface/transformers/blob/87be06ca77166e6a6215eee5a990ab9f07238a18/src/transformers/models/gpt2/tokenization_gpt2.py#L38-L59
|
||||
*/
|
||||
inline std::string ByteLevelDecoder(const std::string& token) {
|
||||
// clang-format off
|
||||
// The inverse map of bytes_to_unicode. -1 means there is no mapping to this unicode.
|
||||
static const std::array<int, 324> char_to_byte_map = {
|
||||
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
|
||||
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45,
|
||||
46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68,
|
||||
69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91,
|
||||
92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
|
||||
112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, -1, -1, -1, -1,
|
||||
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
|
||||
-1, -1, -1, -1, -1, -1, -1, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, -1,
|
||||
174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191,
|
||||
192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209,
|
||||
210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227,
|
||||
228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245,
|
||||
246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12,
|
||||
13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 127, 128,
|
||||
129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146,
|
||||
147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 173
|
||||
};
|
||||
// clang-format on
|
||||
|
||||
auto unicode_codepoints = ParseUTF8(token.c_str(), UTF8ErrorPolicy::kReturnInvalid);
|
||||
if (unicode_codepoints.size() == 1 && unicode_codepoints[0] == kInvalidUTF8) {
|
||||
return token;
|
||||
}
|
||||
|
||||
std::string decoded;
|
||||
|
||||
for (auto unicode_codepoint : unicode_codepoints) {
|
||||
TVM_FFI_ICHECK(unicode_codepoint >= 0);
|
||||
if (unicode_codepoint >= static_cast<int>(char_to_byte_map.size()) ||
|
||||
char_to_byte_map[unicode_codepoint] == -1) {
|
||||
// If there is no mapping, return the original token
|
||||
return token;
|
||||
}
|
||||
decoded += static_cast<char>(char_to_byte_map[unicode_codepoint]);
|
||||
}
|
||||
return decoded;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Post-process a raw token to the actual token with the given post-processing method.
|
||||
*/
|
||||
inline std::string PostProcessToken(const std::string& token,
|
||||
const std::string& token_postproc_method) {
|
||||
if (token_postproc_method == "byte_fallback") {
|
||||
return SpaceReplacerDecoder(ByteFallbackDecoder(token));
|
||||
} else if (token_postproc_method == "byte_level") {
|
||||
return ByteLevelDecoder(token);
|
||||
} else {
|
||||
LOG(FATAL) << "Unknown post-processing method: " << token_postproc_method;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::string> Tokenizer::PostProcessTokenTable(
|
||||
const std::vector<std::string>& token_table, const std::string& token_postproc_method) {
|
||||
std::vector<std::string> post_processed_token_table;
|
||||
post_processed_token_table.reserve(token_table.size());
|
||||
for (const std::string& token : token_table) {
|
||||
post_processed_token_table.push_back(PostProcessToken(token, token_postproc_method));
|
||||
}
|
||||
return post_processed_token_table;
|
||||
}
|
||||
|
||||
#ifndef COMPILE_MLC_WASM_RUNTIME
|
||||
const std::vector<std::string>& TokenizerObj::PostProcessedTokenTable() {
|
||||
if (!post_processed_token_table_.empty()) {
|
||||
return post_processed_token_table_;
|
||||
}
|
||||
|
||||
std::vector<std::string> raw_token_table;
|
||||
int vocab_size = tokenizer->GetVocabSize();
|
||||
raw_token_table.reserve(vocab_size);
|
||||
for (int32_t token_id = 0; token_id < vocab_size; ++token_id) {
|
||||
raw_token_table.push_back(tokenizer->IdToToken(token_id));
|
||||
}
|
||||
post_processed_token_table_ =
|
||||
Tokenizer::PostProcessTokenTable(raw_token_table, info_->token_postproc_method);
|
||||
return post_processed_token_table_;
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef()
|
||||
.def("mlc.tokenizers.Tokenizer", [](const String& path) { return Tokenizer::FromPath(path); })
|
||||
.def("mlc.tokenizers.TokenizerEncode",
|
||||
[](const Tokenizer& tokenizer, const String& text) {
|
||||
std::vector<int32_t> token_ids = tokenizer->Encode(text);
|
||||
return Shape{token_ids.begin(), token_ids.end()};
|
||||
})
|
||||
.def("mlc.tokenizers.TokenizerEncodeBatch",
|
||||
[](const Tokenizer& tokenizer, const Array<String>& texts) {
|
||||
std::vector<std::vector<int32_t>> results = tokenizer->EncodeBatch(texts);
|
||||
Array<Shape> ret;
|
||||
ret.reserve(results.size());
|
||||
for (const auto& result : results) {
|
||||
ret.push_back(Shape{result.begin(), result.end()});
|
||||
}
|
||||
return ret;
|
||||
})
|
||||
.def("mlc.tokenizers.TokenizerDecode",
|
||||
[](const Tokenizer& tokenizer, const Shape& token_ids) {
|
||||
return tokenizer->Decode({token_ids->data, token_ids->data + token_ids->size});
|
||||
})
|
||||
.def("mlc.tokenizers.DetectTokenizerInfo",
|
||||
[](const String& path) { return Tokenizer::DetectTokenizerInfo(path)->AsJSONString(); });
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef()
|
||||
.def_packed("mlc.tokenizers.PostProcessTokenTable",
|
||||
[](tvm::ffi::PackedArgs args, tvm::ffi::Any* rv) {
|
||||
Array<String> token_table_arr = args[0].cast<Array<String>>();
|
||||
std::string token_postproc_method = args[args.size() - 1].cast<String>();
|
||||
std::vector<std::string> token_table;
|
||||
for (int i = 0; i < token_table_arr.size(); ++i) {
|
||||
token_table.push_back(token_table_arr[i]);
|
||||
}
|
||||
std::vector<std::string> processed_token_table =
|
||||
Tokenizer::PostProcessTokenTable(token_table, token_postproc_method);
|
||||
|
||||
// Convert std::vector<std::string> to Array<String>
|
||||
Array<String> processed_token_table_tvm;
|
||||
for (int i = 0; i < processed_token_table.size(); ++i) {
|
||||
processed_token_table_tvm.push_back(processed_token_table[i]);
|
||||
}
|
||||
*rv = processed_token_table_tvm;
|
||||
})
|
||||
.def("mlc.tokenizers.PostProcessToken",
|
||||
[](const String& token, const String& token_postproc_method) {
|
||||
return PostProcessToken(token, token_postproc_method);
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
@@ -0,0 +1,171 @@
|
||||
/*!
|
||||
* Copyright (c) 2023-2025 by Contributors
|
||||
* \file tokenizers.h
|
||||
* \brief Header of tokenizer related functions.
|
||||
*/
|
||||
|
||||
#ifndef MLC_LLM_TOKENIZER_H_
|
||||
#define MLC_LLM_TOKENIZER_H_
|
||||
|
||||
#include <tokenizers_cpp.h>
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
#include <tvm/ffi/reflection/registry.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
|
||||
#include <optional>
|
||||
#include <unordered_map>
|
||||
|
||||
#include "../base.h"
|
||||
#include "../support/dynamic_bitset.h"
|
||||
|
||||
namespace mlc {
|
||||
namespace llm {
|
||||
|
||||
using namespace tvm::runtime;
|
||||
using tvm::ffi::Array;
|
||||
using tvm::ffi::Object;
|
||||
using tvm::ffi::ObjectPtr;
|
||||
using tvm::ffi::ObjectRef;
|
||||
using tvm::ffi::Shape;
|
||||
using tvm::ffi::String;
|
||||
|
||||
/*! \brief Useful information of the tokenizer during generation. */
|
||||
class TokenizerInfoNode : public Object {
|
||||
public:
|
||||
/*! \brief The method to post-process the tokens to their original strings.
|
||||
* Possible values (each refers to a kind of tokenizer):
|
||||
* - "byte_fallback": The same as the byte-fallback BPE tokenizer, including LLaMA-2,
|
||||
* Mixtral-7b, etc. E.g. "▁of" -> " of", "<0x1B>" -> "\x1B".
|
||||
* This method:
|
||||
* 1) Transform tokens like <0x1B> to hex char byte 1B. (so-called byte-fallback)
|
||||
* 2) Replace \\u2581 "▁" with space.
|
||||
* - "byte_level": The same as the byte-level BPE tokenizer, including LLaMA-3, GPT-2,
|
||||
* Phi-2, etc. E.g. "Ġin" -> " in", "ě" -> "\x1B"
|
||||
* This method inverses the bytes-to-unicode transformation in the encoding process in
|
||||
* https://github.com/huggingface/transformers/blob/87be06ca77166e6a6215eee5a990ab9f07238a18/src/transformers/models/gpt2/tokenization_gpt2.py#L38-L59
|
||||
*/
|
||||
String token_postproc_method = "byte_fallback";
|
||||
/*! \brief Whether to prepend a space during encoding. */
|
||||
bool prepend_space_in_encode = false;
|
||||
/*! \brief Whether to strip the first space during decoding. */
|
||||
bool strip_space_in_decode = false;
|
||||
|
||||
String AsJSONString() const;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<TokenizerInfoNode>();
|
||||
}
|
||||
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.TokenizerInfo", TokenizerInfoNode, Object);
|
||||
};
|
||||
|
||||
class TokenizerInfo : public ObjectRef {
|
||||
public:
|
||||
/*! \brief Create a TokenizerInfo object from a dumped string. */
|
||||
static TokenizerInfo FromJSONString(String json_string);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TokenizerInfo, ObjectRef, TokenizerInfoNode);
|
||||
};
|
||||
|
||||
/*! \brief A wrapper object class for tokenizer. */
|
||||
class TokenizerObj : public Object {
|
||||
public:
|
||||
/*! \brief The underlying tokenizer. */
|
||||
std::unique_ptr<tokenizers::Tokenizer> tokenizer;
|
||||
|
||||
/*! \brief Encode text into ids. */
|
||||
std::vector<int32_t> Encode(const std::string& text) const;
|
||||
|
||||
/*! \brief Encode text into ids. Some tokenizers may prepend a space in encoding, this method
|
||||
* guarantees the space is not prepended. */
|
||||
std::vector<int32_t> EncodeNoPrependSpace(const std::string& text) const;
|
||||
|
||||
/*! \brief Encode texts into ids. */
|
||||
std::vector<std::vector<int32_t>> EncodeBatch(const Array<String>& texts) const;
|
||||
|
||||
/*! \brief Decode token ids into text. */
|
||||
std::string Decode(const std::vector<int32_t>& token_ids) const;
|
||||
|
||||
/*! \brief Return the post-processed token table of the tokenizer. Special tokens are included. */
|
||||
const std::vector<std::string>& PostProcessedTokenTable();
|
||||
|
||||
/*! \brief Get the prefix token mask as a bitset. The tokens which is a prefix of another token
|
||||
* are set to true, and others are set to false in the bitset. */
|
||||
const DynamicBitset& GetPrefixTokenMask();
|
||||
|
||||
/*!
|
||||
* \brief Returns the vocabulary size. Special tokens are considered. This may be smaller than the
|
||||
* `vocab_size` in config.json (length of logits), see https://github.com/QwenLM/Qwen2/issues/147
|
||||
* and https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/discussions/47.
|
||||
*/
|
||||
size_t GetVocabSize() const;
|
||||
|
||||
/*!
|
||||
* \brief Convert the given id to its corresponding token if it exists. If not, return an
|
||||
* empty string.
|
||||
*/
|
||||
std::string IdToToken(int32_t token_id) const;
|
||||
|
||||
/*!
|
||||
* \brief Convert the given token to its corresponding id if it exists. If not, return -1.
|
||||
*/
|
||||
int32_t TokenToId(const std::string& token) const;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<TokenizerObj>();
|
||||
}
|
||||
|
||||
friend class Tokenizer;
|
||||
static constexpr const bool _type_has_method_sequal_reduce = false;
|
||||
static constexpr const bool _type_has_method_shash_reduce = false;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.Tokenizer", TokenizerObj, Object);
|
||||
|
||||
private:
|
||||
/*! \brief Useful information of the tokenizer during generation. */
|
||||
TokenizerInfo info_;
|
||||
/*! \brief The cached token table. */
|
||||
std::vector<std::string> post_processed_token_table_;
|
||||
/*! \brief The cached prefix token mask. */
|
||||
DynamicBitset prefix_token_mask_;
|
||||
};
|
||||
|
||||
class Tokenizer : public ObjectRef {
|
||||
public:
|
||||
/*!
|
||||
* \brief Create a tokenizer from a directory path on disk.
|
||||
* \param path The path to the tokenizer or the tokenizer directory.
|
||||
* \param info The tokenizer info. If not provided, the info will be detected automatically.
|
||||
*/
|
||||
MLC_LLM_DLL static Tokenizer FromPath(const String& path,
|
||||
std::optional<TokenizerInfo> info = std::nullopt);
|
||||
|
||||
/*! \brief Detect the tokenizer info from the given path of the tokenizer. */
|
||||
MLC_LLM_DLL static TokenizerInfo DetectTokenizerInfo(const String& path);
|
||||
|
||||
/*!
|
||||
* \brief Post-process the token table to their original strings.
|
||||
* \param token_table The raw token table.
|
||||
* \param postproc_method The postprocessing method to use.
|
||||
* \returns The postprocessed token table containing the original strings.
|
||||
*/
|
||||
static std::vector<std::string> PostProcessTokenTable(const std::vector<std::string>& token_table,
|
||||
const std::string& token_postproc_method);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Tokenizer, ObjectRef, TokenizerObj);
|
||||
|
||||
private:
|
||||
explicit Tokenizer(std::unique_ptr<tokenizers::Tokenizer> tokenizer, TokenizerInfo info);
|
||||
};
|
||||
|
||||
} // namespace llm
|
||||
} // namespace mlc
|
||||
|
||||
#endif // MLC_LLM_TOKENIZER_H_
|
||||
Reference in New Issue
Block a user