Skip to content

vllm.models.kimi_k3.common.dsa_cache

Kimi-K3.1 DSA indexer cache, shared by the NVIDIA and ROCm stacks.

Classes:

Functions:

KimiK3IndexerMetadataBuilder

Bases: DeepseekV32IndexerMetadataBuilder

Writes the indexer keys at slots derived from the block table, as the compressed indexers are written: in a packed block row the block table addresses 64-token kernel pages, which manager slots do not.

Source code in vllm/models/kimi_k3/common/dsa_cache.py 
 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
class KimiK3IndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder):
    """Writes the indexer keys at slots derived from the block table, as the
    compressed indexers are written: in a packed block row the block table
    addresses 64-token kernel pages, which manager slots do not."""

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        if self.use_pcp:
            raise NotImplementedError("Kimi-K3.1 packed DSA caches with PCP")
        # The K3.1 key format, validated as KimiK3DSAIndexer selects it.
        self.indexer_uses_fp4 = kimi_k3_indexer_uses_fp4(self.vllm_config)
        self.slot_mapping_buffer = torch.zeros(
            self.vllm_config.scheduler_config.max_num_batched_tokens,
            dtype=torch.int64,
            device=self.device,
        )

    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> DeepseekV32IndexerMetadata:
        slot_mapping = get_compressed_slot_mapping(
            common_attn_metadata.num_actual_tokens,
            common_attn_metadata.slot_mapping,
            common_attn_metadata.query_start_loc,
            common_attn_metadata.seq_lens,
            common_attn_metadata.block_table_tensor,
            self.kv_cache_spec.block_size,
            compress_ratio=1,
            out=self.slot_mapping_buffer,
            dcp_world_size=self.dcp_world_size,
            cp_interleave=self.cp_kv_cache_interleave_size,
        )
        metadata = super().build(
            common_prefix_len,
            replace(common_attn_metadata, slot_mapping=slot_mapping),
            fast_build,
        )
        return metadata

dense_mla_draft_bytes_per_token(vllm_config)

Per-token bytes of a dense MLA draft's caches, which share the DSA row: their full-attention MLA specs group with the target's (see configure_dense_draft).

Source code in vllm/models/kimi_k3/common/dsa_cache.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def dense_mla_draft_bytes_per_token(vllm_config: VllmConfig) -> int:
    """Per-token bytes of a dense MLA draft's caches, which share the DSA row:
    their full-attention MLA specs group with the target's (see
    ``configure_dense_draft``)."""
    spec_config = vllm_config.speculative_config
    if spec_config is None or spec_config.method not in ("eagle3", "dspark"):
        return 0
    draft_model_config = spec_config.draft_model_config
    draft_config = draft_model_config.hf_text_config
    kv_lora_rank = getattr(draft_config, "kv_lora_rank", None)
    if kv_lora_rank is None:
        return 0
    cache_dtype = spec_config.kv_cache_dtype or vllm_config.cache_config.cache_dtype
    page = MLAAttentionSpec(
        block_size=1,
        num_kv_heads=1,
        head_size=kv_lora_rank + draft_config.qk_rope_head_dim,
        dtype=kv_cache_dtype_str_to_dtype(cache_dtype, draft_model_config),
        cache_dtype_str=cache_dtype,
        kv_quant_mode=get_kv_quant_mode(cache_dtype),
    ).page_size_bytes
    return draft_config.num_hidden_layers * page

indexer_cache_entry_bytes(head_dim, use_fp4_cache)

Indexer key cache bytes per token: FP8 values + one 4 B scale, or MXFP4 values (two per byte) + one UE8M0 byte per 32 values.

Source code in vllm/models/kimi_k3/common/dsa_cache.py 
56
57
58
59
60
61
def indexer_cache_entry_bytes(head_dim: int, use_fp4_cache: bool) -> int:
    """Indexer key cache bytes per token: FP8 values + one 4 B scale, or MXFP4
    values (two per byte) + one UE8M0 byte per 32 values."""
    if use_fp4_cache:
        return head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE
    return head_dim + 4

kimi_k3_indexer_uses_fp4(vllm_config)

FP8 (132 B/token) indexer keys by default (indexer_kv_dtype='auto'); MXFP4 (68 B/token) only when requested with --attention-config '{"indexer_kv_dtype": "mxfp4"}'. MXFP4 needs SM10x (checked by dsa_indexer_uses_fp4), the fused key insert (layernorm k-norm, no indexer RoPE) and no PCP; anything else raises instead of writing another format.

Source code in vllm/models/kimi_k3/common/dsa_cache.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def kimi_k3_indexer_uses_fp4(vllm_config: VllmConfig) -> bool:
    """FP8 (132 B/token) indexer keys by default (``indexer_kv_dtype='auto'``);
    MXFP4 (68 B/token) only when requested with ``--attention-config
    '{"indexer_kv_dtype": "mxfp4"}'``. MXFP4 needs SM10x (checked by
    ``dsa_indexer_uses_fp4``), the fused key insert (layernorm k-norm, no indexer
    RoPE) and no PCP; anything else raises instead of writing another format."""
    if not dsa_indexer_uses_fp4(vllm_config):
        return False
    text_config = getattr(vllm_config.model_config, "hf_text_config", None)
    dsa_config = getattr(text_config, "dsa_config", None)
    if dsa_config is not None and not (
        getattr(dsa_config, "index_k_norm_type", "rmsnorm") == "layernorm"
        and getattr(dsa_config, "index_rope_dim", 0) == 0
    ):
        raise ValueError(
            "indexer_kv_dtype='mxfp4' needs the fused Kimi K3 DSA key insert "
            "(index_k_norm_type='layernorm', index_rope_dim=0)."
        )
    parallel_config = getattr(vllm_config, "parallel_config", None)
    if getattr(parallel_config, "prefill_context_parallel_size", 1) > 1:
        raise ValueError(
            "indexer_kv_dtype='mxfp4' is not supported with prefill context "
            "parallelism for Kimi K3 DSA."
        )
    return True

packed_dsa_block_size(vllm_config, mla_bytes_per_token, state_bytes, alignment)

Block size of the packed cache of a KDA-hybrid DSA model, else None.

A block row holds every DSA layer's MLA and indexer pages, plus a dense MLA draft's pages; size it to fit as many KDA states as there are DSA layers. The KDA state shards with TP and the row does not, so the block is rounded at _PD_REFERENCE_TP and scaled by _PD_REFERENCE_TP / tp. Block sizes of different TP sizes stay divisible, and the rounding stays fine: a power of two would leave up to half the row empty once spec decoding grows the conv state past one.

Source code in vllm/models/kimi_k3/common/dsa_cache.py 
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
def packed_dsa_block_size(
    vllm_config: VllmConfig,
    mla_bytes_per_token: int,
    state_bytes: int,
    alignment: int,
) -> int | None:
    """Block size of the packed cache of a KDA-hybrid DSA model, else None.

    A block row holds every DSA layer's MLA and indexer pages, plus a dense
    MLA draft's pages; size it to fit as many KDA states as there are DSA
    layers. The KDA state shards with TP and the row does not, so the block is
    rounded at ``_PD_REFERENCE_TP`` and scaled by ``_PD_REFERENCE_TP / tp``.
    Block sizes of different TP sizes stay divisible, and the rounding stays
    fine: a power of two would leave up to half the row empty once spec
    decoding grows the conv state past one.
    """
    text_config = vllm_config.model_config.hf_text_config
    dsa_config = getattr(text_config, "dsa_config", None)
    linear_attn_config = getattr(text_config, "linear_attn_config", None)
    if dsa_config is None or linear_attn_config is None:
        return None
    num_dsa = len(linear_attn_config["full_attn_layers"])
    num_states = len(linear_attn_config["kda_layers"])
    states_per_row = cdiv(num_states, cdiv(num_states, num_dsa))
    # The indexer keeps fp8 values plus one fp32 scale per token, or MXFP4.
    index_bytes = indexer_cache_entry_bytes(
        dsa_config.index_head_dim, kimi_k3_indexer_uses_fp4(vllm_config)
    )
    row_bytes_per_token = num_dsa * (
        mla_bytes_per_token + index_bytes
    ) + dense_mla_draft_bytes_per_token(vllm_config)
    tp_size = vllm_config.parallel_config.tensor_parallel_size
    reference_tp = math.lcm(_PD_REFERENCE_TP, tp_size)
    # TP-invariant: the unsharded state bytes, in reference-TP block units.
    reference_units = cdiv(
        states_per_row * state_bytes * tp_size,
        alignment * reference_tp * row_bytes_per_token,
    )
    return alignment * reference_units * (reference_tp // tp_size)