vllm.models.deepseek_v41.common.ops.indexer_k_store ¶

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.
        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.
    """
    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,
        SHUFFLE=shuffle,
        BLOCK_TILE_SIZE=_BLOCK_TILE_SIZE,
        HEAD_TILE_SIZE=_HEAD_TILE_SIZE,
        num_warps=1,
        **launch_kwargs,
    )