Implementation:Mlc ai Mlc llm KV Cache

Overview

The KV Cache module defines the PagedKVCache class used for efficient key-value cache management in LLM batching and attention computation. It is located at python/mlc_llm/nn/kv_cache.py (94 lines).

This module extends TVM's built-in PagedKVCache with a generic factory method (create_generic) that produces a cache instance suitable for multiple attention types including multi-head attention (MHA), multi-latent attention (MLA), and sliding-window MHA. The factory method emits a call to the mlc.create_paged_kv_cache_generic packed function, which is resolved during the compilation pipeline.

Source File

• File: python/mlc_llm/nn/kv_cache.py
• Lines: 94
• Module: mlc_llm.nn.kv_cache

Dependencies

Import	Purpose
json	Serializes rope_scaling dictionary to JSON string for the packed function call
numpy	Creates constant arrays for rope_ext_factors
tvm.relax	Relax IR types: call_pure_packed, ShapeExpr, PrimValue, StringImm, DataTypeImm, ObjectStructInfo, const
tvm.tir	TIR variable types for symbolic shape parameters
tvm.relax.frontend.nn.llm.kv_cache.PagedKVCache	Base TVM paged KV cache class (aliased as TVMPagedKVCache)
tvm.relax.frontend.nn.llm.kv_cache.RopeMode	Enumeration for rotary position embedding modes

Class: PagedKVCache

class PagedKVCache(TVMPagedKVCache):

Extends TVMPagedKVCache with the create_generic static factory method.

Static Method: create_generic

@staticmethod
def create_generic(
    attn_kind: Union[Literal["mha", "mla"], List[Literal["mha", "mla", "mha_sliding"]]],
    max_batch_size: tir.Var,
    max_total_seq_len: tir.Var,
    prefill_chunk_size: tir.Var,
    page_size: tir.Var,
    support_sliding_window: tir.Var,
    num_hidden_layers: int,
    num_attention_heads: int,
    num_key_value_heads: int,
    qk_head_dim: int,
    v_head_dim: int,
    rope_mode: RopeMode,
    rope_scale: int,
    rope_theta: int,
    dtype: str,
    mla_original_qk_head_dim: int = 0,
    mla_original_v_head_dim: int = 0,
    rotary_dim: Optional[int] = None,
    rope_scaling: Optional[Dict[str, Any]] = None,
    rope_ext_factors: Optional[List[int]] = None,
    layer_partition: Optional[List[int]] = None,
    enable_disaggregation: bool = False,
    name: str = "paged_kv_cache",
) -> "PagedKVCache":

Creates a PagedKVCache instance by emitting a call_pure_packed to the runtime function mlc.create_paged_kv_cache_generic.

Parameters

Parameter	Type	Description
attn_kind	str or List[str]	Attention type per layer: "mha", "mla", or "mha_sliding". Can be a single string (same for all layers) or a per-layer list.
max_batch_size	tir.Var	Maximum batch size (symbolic)
max_total_seq_len	tir.Var	Maximum total sequence length across all batches (symbolic)
prefill_chunk_size	tir.Var	Prefill chunk size (symbolic)
page_size	tir.Var	Number of tokens per KV cache page (symbolic)
support_sliding_window	tir.Var	Whether sliding window attention is supported (symbolic)
num_hidden_layers	int	Number of transformer layers
num_attention_heads	int	Number of query attention heads
num_key_value_heads	int	Number of key/value heads (for grouped-query attention)
qk_head_dim	int	Dimension of each Q/K head
v_head_dim	int	Dimension of each V head
rope_mode	RopeMode	Rotary position embedding mode
rope_scale	int	RoPE scaling factor
rope_theta	int	RoPE theta base frequency
dtype	str	Data type for the KV cache storage
mla_original_qk_head_dim	int	Original QK head dim for MLA (default 0)
mla_original_v_head_dim	int	Original V head dim for MLA (default 0)
rotary_dim	Optional[int]	Rotary dimension; defaults to qk_head_dim
rope_scaling	Optional[Dict]	Additional RoPE scaling configuration, serialized as JSON
rope_ext_factors	Optional[List[int]]	RoPE extension factors, passed as a float32 constant array
layer_partition	Optional[List[int]]	Partition boundaries for pipeline parallelism; defaults to [0, num_hidden_layers]
enable_disaggregation	bool	Enables disaggregated inference (default False)

Packed Function Call

The method constructs a Relax IR expression that calls mlc.create_paged_kv_cache_generic with all parameters encoded as Relax IR values:

return PagedKVCache(
    _expr=rx.call_pure_packed(
        "mlc.create_paged_kv_cache_generic",
        rx_attn_kind,
        rx.ShapeExpr([max_batch_size, max_total_seq_len, prefill_chunk_size, page_size, support_sliding_window]),
        rx.ShapeExpr(layer_partition),
        rx.PrimValue(num_hidden_layers),
        rx.PrimValue(num_attention_heads),
        rx.PrimValue(num_key_value_heads),
        rx.PrimValue(qk_head_dim),
        rx.PrimValue(v_head_dim),
        # ... additional parameters ...
        rx.DataTypeImm(dtype),
        sinfo_args=rx.ObjectStructInfo(),
    ),
    _name=name,
)

Attention Kind Handling

When attn_kind is a list, each element is converted to a rx.StringImm to support per-layer attention type configuration:

if isinstance(attn_kind, List):
    rx_attn_kind = [rx.StringImm(layer_kind) for layer_kind in attn_kind]
else:
    rx_attn_kind = rx.StringImm(attn_kind)

Optional Value Encoding

Since Relax does not have an "Optional" type, rope_ext_factors uses rx.PrimValue(0) as a sentinel for "undefined":

(
    rx.const(np.array(rope_ext_factors, "float32"))
    if rope_ext_factors is not None
    else rx.PrimValue(0)
)

Design Notes

• All symbolic shape parameters (max_batch_size, max_total_seq_len, etc.) remain as tir.Var, allowing the cache to be configured at runtime.
• The actual KV cache creation logic is deferred to the mlc.create_paged_kv_cache_generic runtime function, which is resolved during the MLC compilation pipeline.
• The layer_partition parameter supports pipeline parallelism by specifying which layers belong to each partition.

Categories

• KV Cache
• Attention
• Paged Memory
• TVM Relax
• LLM Inference