Skip to content

vllm.models.deepseek_v41.common.ops.indexer_k_store

Indexer K production for DeepSeek V4.1 kv-source layers.

In v4.1 the index key is derived from the main compressor's latent: k = k_norm(wk(latent)) (reference model.py Indexer.forward), then RoPE'd at the group's first-token position and MXFP4/FP8-quantized into the paged indexer K cache. The wk(latent) GEMM runs in torch; this kernel fuses the remaining k_norm → RoPE → quant → paged store, one program per token.

Unlike the legacy (v4.0) indexer path there is no per-token pooling from a compressor state cache: the latent already stands for a whole group, so only group-boundary tokens (position + 1) % compress_ratio == 0 produce a key.

Functions:

indexer_k_norm_rope_store(k_pre, positions, cos_sin_cache, rms_norm_weight, rms_norm_eps, k_cache, kv_slot_mapping, compress_ratio, use_fp4_cache)

k_norm → RoPE → quant → paged store for indexer keys.

Parameters:

  • k_pre

    (Tensor) –

    [num_tokens, 128] bf16, the wk(latent) projection. Only group-boundary rows are read.

  • positions

    (Tensor) –

    [num_tokens] int64 token positions.

  • cos_sin_cache

    (Tensor) –

    [max_pos, rope_head_dim] GPT-J layout (cos half, then sin half), from the layer's compress-RoPE instance.

  • rms_norm_weight

    (Tensor) –

    [128] k_norm weight.

  • rms_norm_eps

    (float) –

    Epsilon of the k_norm RMSNorm.

  • k_cache

    (Tensor) –

    uint8 paged indexer cache [num_blocks, block_size, row_bytes].

  • kv_slot_mapping

    (Tensor) –

    [num_tokens] slots in the indexer cache (-1 = skip).

  • compress_ratio

    (int) –

    group size; keys are emitted at group boundaries.

  • use_fp4_cache

    (bool) –

    MXFP4 (2 nibbles/byte + ue8m0 per 32) when True, else per-token FP8 with a single fp32 scale (four packed ue8m0 bytes under mqa_logits_use_mx_scales).

Source code in vllm/models/deepseek_v41/common/ops/indexer_k_store.py 
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def indexer_k_norm_rope_store(
    k_pre: torch.Tensor,
    positions: torch.Tensor,
    cos_sin_cache: torch.Tensor,
    rms_norm_weight: torch.Tensor,
    rms_norm_eps: float,
    k_cache: torch.Tensor,
    kv_slot_mapping: torch.Tensor,
    compress_ratio: int,
    use_fp4_cache: bool,
) -> None:
    """k_norm → RoPE → quant → paged store for indexer keys.

    Args:
        k_pre: [num_tokens, 128] bf16, the ``wk(latent)`` projection. Only
            group-boundary rows are read.
        positions: [num_tokens] int64 token positions.
        cos_sin_cache: [max_pos, rope_head_dim] GPT-J layout (cos half, then
            sin half), from the layer's compress-RoPE instance.
        rms_norm_weight: [128] k_norm weight.
        rms_norm_eps: Epsilon of the k_norm RMSNorm.
        k_cache: uint8 paged indexer cache [num_blocks, block_size, row_bytes].
        kv_slot_mapping: [num_tokens] slots in the indexer cache (-1 = skip).
        compress_ratio: group size; keys are emitted at group boundaries.
        use_fp4_cache: MXFP4 (2 nibbles/byte + ue8m0 per 32) when True, else
            per-token FP8 with a single fp32 scale (four packed ue8m0 bytes
            under `mqa_logits_use_mx_scales`).

    """
    num_tokens = kv_slot_mapping.numel()
    assert k_pre.ndim == 2 and k_pre.shape[1] == 128
    assert k_pre.dtype == torch.bfloat16 and k_pre.stride(1) == 1
    assert num_tokens <= k_pre.shape[0] and num_tokens <= positions.numel()
    assert compress_ratio in (1, 2)
    if num_tokens == 0:
        return

    head_dim = k_pre.shape[1]
    if use_fp4_cache:
        token_stride = head_dim // 2
        scale_dim = head_dim // MXFP4_BLOCK_SIZE
    else:
        token_stride = head_dim
        scale_dim = 4  # single float32 scale

    # ROCm reads this cache back with the 16x16-tiled ("SHUFFLE") value
    # layout: cp_gather_indexer_k_quant_cache_triton on prefill and aiter's
    # deepgemm_fp8_paged_mqa_logits(Preshuffle=True) on decode both select it
    # from ``block_size > 1``, matching the v4.0 writer
    # (indexer_k_quant_and_cache_triton). Writing row-major here would hand
    # both readers permuted key bytes.
    block_size = k_cache.shape[1]
    shuffle = current_platform.is_rocm() and block_size > 1
    if shuffle:
        if use_fp4_cache:
            raise NotImplementedError(
                "MXFP4 indexer K cache has no tiled ROCm layout; "
                "the ROCm readers only implement the FP8 one."
            )
        if block_size % _BLOCK_TILE_SIZE != 0 or head_dim % _HEAD_TILE_SIZE != 0:
            raise ValueError(
                f"ROCm tiled indexer K cache needs block_size "
                f"({block_size}) % {_BLOCK_TILE_SIZE} == 0 and head_dim "
                f"({head_dim}) % {_HEAD_TILE_SIZE} == 0."
            )

    launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {}
    _indexer_k_norm_rope_quant_store_kernel[(num_tokens,)](
        k_pre,
        k_pre.stride(0),
        positions,
        rms_norm_weight,
        rms_norm_eps,
        cos_sin_cache,
        cos_sin_cache.stride(0),
        k_cache,
        kv_slot_mapping,
        k_cache.shape[1],
        HEAD_SIZE=head_dim,
        ROPE_HEAD_DIM=64,
        COMPRESS_RATIO=compress_ratio,
        TOKEN_STRIDE=token_stride,
        SCALE_DIM=scale_dim,
        KV_BLOCK_STRIDE=k_cache.stride(0),
        FP8_MAX=448.0,
        USE_FP4=use_fp4_cache,
        MX_SCALE=not use_fp4_cache and mqa_logits_use_mx_scales(),
        SHUFFLE=shuffle,
        BLOCK_TILE_SIZE=_BLOCK_TILE_SIZE,
        HEAD_TILE_SIZE=_HEAD_TILE_SIZE,
        num_warps=1,
        **launch_kwargs,
    )