Implementation:Mlc ai Mlc llm Function Table

Overview

File: cpp/serve/function_table.h

Purpose: Declares the FunctionTable struct, which serves as the central registry of all TVM packed functions used by the MLC LLM serving engine for batched and distributed inference. It manages function pointers for model operations (embedding, prefill, decode, verify), KV cache management, GPU sampling, and auxiliary speculative decoding operations. The struct also handles initialization across both local and Disco (distributed) execution modes.

Namespace: mlc::llm::serve

Design Context

As noted in the source code comments, this function table is mostly identical to the one used for single-sequence distributed inference in llm_chat.cc, but uses a different set of packed functions tailored for batched serving. The duplicate was intentionally maintained to keep the batching/serving development independent from single-sequence inference, with plans to merge once the batching implementation stabilizes.

Struct: FunctionTable

Initialization and Loading Methods

static Function SessionFuncAsPackedFunc(Session sess, DRef sess_func, String name);

Wraps a Disco session function reference (DRef) as a standard TVM packed function, enabling uniform invocation regardless of whether the function runs locally or is distributed across workers.

void Init(String reload_lib_path, Device device, picojson::object model_config,
          Optional<Session> session, int num_shards, int num_stages);

Initializes the function table by loading the model library, setting up the execution device, and populating all function pointers. When a Session is provided, Disco distributed mode is activated. The num_shards and num_stages parameters configure the parallelism topology.

ObjectRef LoadParams(const std::string& model_path, Device device);

Loads model parameters (weights) from disk into the specified device.

void _InitFunctions();

Internal method that populates all the function member variables from the loaded model library module.

Utility Methods

ObjectRef Empty(Shape shape, DataType dtype, Device device, bool worker0_only) const;

Allocates an empty tensor with the given shape and data type. When worker0_only is true in Disco mode, the allocation is performed only on the first worker.

ObjectRef CopyToWorker0(const Tensor& host_array, String buffer_cache_key,
                        Shape max_reserved_shape, bool local_only = false);

Copies a host-side array to GPU memory (either local or worker 0 in Disco mode). Uses a buffer cache keyed by buffer_cache_key with a maximum reserved shape to avoid repeated allocations. The local_only parameter bypasses Disco distribution for functions that run only on the local GPU.

void DebugCallFuncOnAllAllWorker(const String& func_name, Optional<String> func_args) const;

Invokes a named function on all Disco workers for debugging purposes.

Configuration Members

Member	Type	Description
use_disco	bool	Whether distributed execution (Disco) is enabled
local_gpu_device	Device	The local GPU device handle
sess	Session	Disco session (null if not using Disco)
disco_mod	Optional<DRef>	Disco module reference
cached_buffers	Optional<Map<String, ObjectRef>>	Cache of pre-allocated buffers
local_vm	Optional<tvm::ffi::Module>	Local virtual machine module
model_config	picojson::object	Model configuration in JSON
model_metadata_	ModelMetadata	Parsed model metadata

Function Lookup Members

TypedFunction<Function(const std::string&)> mod_get_func;
TypedFunction<Function(const std::string&)> get_global_func;

These two typed functions provide the ability to look up other functions by name -- mod_get_func from the loaded model module and get_global_func from the TVM global function registry.

Model Execution Functions

Function	Purpose
embed_func_	Token embedding
image_embed_func_	Image embedding (multimodal models)
single_batch_prefill_func_	Single-batch prefill
single_batch_decode_func_	Single-batch decode
single_batch_extend_func_	Single-batch extend
prefill_func_	Batched prefill
decode_func_	Batched decode
extend_func_	Batched extend
verify_func_	Speculative decoding verification
single_batch_prefill_to_last_hidden_func_	Single-batch prefill returning last hidden state
single_batch_decode_to_last_hidden_func_	Single-batch decode returning last hidden state
prefill_to_last_hidden_func_	Batched prefill returning last hidden state
decode_to_last_hidden_func_	Batched decode returning last hidden state
verify_to_last_hidden_func_	Verification returning last hidden state
fuse_embed_hidden_func_	Fuse embeddings with hidden states (EAGLE)
get_logits_func_	Compute logits from hidden states
batch_get_logits_func_	Batched logit computation
batch_select_last_hidden_func_	Select last hidden state from batch

Logit Processing Functions

Function	Purpose
softmax_func_	Softmax computation
apply_logit_bias_func_	Apply logit bias from generation config
apply_penalty_func_	Apply repetition/frequency/presence penalties
apply_bitmask_func_	Apply vocabulary bitmask (grammar constraints)

KV Cache Management Functions

Function	Purpose
alloc_embedding_tensor_func_	Allocate embedding tensor
cuda_graph_alloc_init_func_	Initialize CUDA graph allocations
create_kv_cache_func_	Create a new KV cache
reset_kv_cache_func_	Reset an existing KV cache
support_backtracking_kv_	Boolean flag: whether KV cache supports backtracking
kv_cache_add_sequence_func_	Add a new sequence to the KV cache
kv_cache_fork_sequence_func_	Fork a sequence in the KV cache (for prefix sharing)
kv_cache_enable_sliding_window_for_seq_	Enable sliding window attention for a sequence
kv_cache_remove_sequence_func_	Remove a sequence from the KV cache
kv_cache_begin_forward_func_	Begin a forward pass on the KV cache
kv_cache_end_forward_func_	End a forward pass on the KV cache
kv_cache_disagg_prepare_recv_func_	Prepare to receive KV cache in disaggregated mode
kv_cache_disagg_mark_send_func_	Mark KV cache for sending in disaggregated mode
kv_cache_popn_func_	Pop N entries from the KV cache (rollback)
kv_cache_commit_accepted_token_tree_nodes_func_	Commit accepted nodes from a speculative token tree
kv_cache_get_num_available_pages_func_	Query available KV cache pages
kv_cache_get_total_sequence_length_func_	Query total sequence length across all sequences

GPU Sampling Functions

Function	Purpose
gpu_multinomial_from_uniform_func_	GPU multinomial sampling from uniform random values
gpu_argsort_probs_func_	GPU argsort of probabilities
gpu_sample_with_top_p_func_	GPU top-p (nucleus) sampling
gpu_sampler_take_probs_func_	Take probabilities at specific indices
gpu_verify_draft_tokens_func_	GPU verification of speculative draft tokens
gpu_renormalize_by_top_p_func_	GPU probability renormalization after top-p filtering

Tensor Utility Functions

Function	Purpose
nd_view_func_	Create a view of an NDArray
nd_get_shape_func_	Get the shape of an NDArray
nd_copy_embedding_to_offset_func_	Copy embedding data to a specific offset
tuple_getitem_func_	Extract an item from a tuple
last_group_send_to_worker_0_	Send data from the last pipeline group to worker 0

Speculative Decoding Auxiliary Functions

Function	Purpose
gather_probs_func_	Gather probability distributions for draft tokens
scatter_probs_func_	Scatter probability distributions to draft token slots
gather_hidden_states_func_	Gather hidden states for draft token generation
scatter_hidden_states_func_	Scatter hidden states to draft token positions

Design Notes

• The function table acts as a vtable-like mechanism, decoupling the engine actions from specific model implementations. Functions are resolved by name at initialization time and invoked as opaque packed functions thereafter.
• The Disco support enables transparent distributed execution -- the same engine action code works whether running on a single GPU or across multiple GPUs/nodes.
• The large number of KV cache functions reflects the complexity of paged attention with features like forking (prefix sharing), sliding window, backtracking (rollback), and disaggregated inference.
• The support_backtracking_kv_ boolean flag (rather than a function) indicates a compile-time capability of the KV cache implementation.