Implementation:Mlc ai Mlc llm Radix Tree

Knowledge Sources	Mlc_ai_Mlc_llm
Domains	LLM Serving, Data Structures, KV Cache Management
Last Updated	2026-02-09 19:00 GMT

Overview

A paged radix tree implementation used for efficiently managing and sharing token sequences across multiple requests in the MLC LLM serving engine.

Description

The PagedRadixTree is a specialized radix tree data structure designed for LLM serving workloads where multiple concurrent requests may share common token prefixes (e.g., system prompts). Unlike a standard radix tree, this implementation uses fixed-size pages to store token prefixes, making memory usage predictable and manageable through pool-based allocation.

The data structure is organized around several key components:

RadixPage is the core node type. Each page stores a segment of a token sequence using a circular buffer layout. Pages are linked together in a left-child, right-sibling binary tree representation, which avoids the overhead of maintaining a child array proportional to the vocabulary size. Each page tracks which sequence IDs terminate at that node, enabling efficient prefix sharing.

RadixPagePool manages the allocation and deallocation of RadixPage instances using a block-based memory pool. Pages are allocated in blocks of 64, with each page having a capacity of 64 tokens. This avoids frequent memory allocation and deallocation operations during serving.

SequenceIDNodePool similarly manages linked-list nodes used to track which sequences are associated with each page. This pool also uses block-based allocation with blocks of size 64.

PagedRadixTreeImpl is the main tree implementation. It maintains a mapping from sequence IDs to their terminal pages, and provides operations for adding, forking, extending, rolling back, and removing sequences. The tree automatically performs page splitting and merging to maintain efficiency. When a sequence is extended, the tree tries to share existing prefix pages with other sequences. When pages become redundant (e.g., a page with no sequences and exactly one child), they are merged to conserve memory.

The implementation is registered with the TVM FFI system, exposing operations such as MatchPrefix, ExtendSequence, RollBackSequence, ForkSequence, AddSequence, RemoveSequence, and others for use from Python or other TVM-based components.

Usage

This radix tree is used by the MLC LLM serving engine to manage the KV cache prefix sharing among concurrent requests. It is particularly useful in scenarios where many requests share common prefixes (such as system prompts), allowing the engine to avoid redundant computation and memory usage. The tree supports operations needed for speculative decoding (forking and rolling back sequences) and prefix caching (matching existing prefixes).

Code Reference

Source Location

• Repository: Mlc_ai_Mlc_llm
• File: cpp/serve/radix_tree.cc

Signature

class PagedRadixTreeImpl : public PagedRadixTreeObj {
 public:
  bool HasSequence(int64_t seq_id);
  IntTuple GetSequence(int64_t seq_id);
  std::pair<size_t, std::vector<int64_t>> MatchPrefix(const std::vector<int32_t>& tokens);
  size_t GetSequenceLength(int64_t seq_id);
  void ForkSequence(int64_t seq_id, int64_t parent_seq_id, size_t forked_offset);
  void AddSequence(int64_t seq_id);
  void ExtendSequence(int64_t seq_id, const std::vector<int32_t>& tokens);
  void RollBackSequence(int64_t seq_id, size_t num_tokens);
  void RemoveSequence(int64_t seq_id);
  size_t FreeCapacity();
  void Reset();
};

Import

#include "radix_tree.h"

I/O Contract

Inputs

Name	Type	Required	Description
seq_id	int64_t	Yes	The unique identifier for a token sequence.
tokens	std::vector<int32_t>	Varies	Token IDs to extend or match against the tree.
parent_seq_id	int64_t	For ForkSequence	The parent sequence ID to fork from.
forked_offset	size_t	For ForkSequence	The position in the parent sequence at which to fork (range [1, parent_length]).
num_tokens	size_t	For RollBackSequence	Number of tokens to roll back from the end of a sequence.

Outputs

Name	Type	Description
HasSequence result	bool	Whether the given sequence ID exists in the tree.
GetSequence result	IntTuple	All token IDs of the specified sequence.
MatchPrefix result	(size_t, vector<int64_t>)	The matched prefix length and all sequence IDs sharing that prefix.
GetSequenceLength result	size_t	The total number of tokens in the specified sequence.
FreeCapacity result	size_t	The remaining token capacity in the page pool.

Usage Examples

// Create a paged radix tree
PagedRadixTree tree = PagedRadixTree::Create();

// Add a new empty sequence with ID 0
tree->AddSequence(0);

// Extend the sequence with tokens
std::vector<int32_t> tokens = {1, 2, 3, 4, 5};
tree->ExtendSequence(0, tokens);

// Fork sequence 0 into a new sequence 1 at offset 3
tree->ForkSequence(1, 0, 3);

// Extend the forked sequence with different tokens
std::vector<int32_t> new_tokens = {10, 11};
tree->ExtendSequence(1, new_tokens);

// Match prefix returns matched length and all sequences sharing the prefix
auto [matched_len, seq_ids] = tree->MatchPrefix({1, 2, 3});

// Roll back the last 2 tokens from sequence 1
tree->RollBackSequence(1, 2);

// Remove a sequence
tree->RemoveSequence(1);

Related Pages

• Mlc_ai_Mlc_llm_Request_State