wehub-resource-sync
118 lines
4.0 KiB
C++
118 lines
4.0 KiB
C++
/*!
|
|
* 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_
|