vllm.models.deepseek_v41.common.ops ¶

def build_flashinfer_mixed_sparse_indices(
    decode_swa_indices: torch.Tensor,
    decode_compressed_indices: torch.Tensor | None,
    decode_compressed_topk_lens: torch.Tensor | None,
    prefill_topk_indices: torch.Tensor,
    query_start_loc: torch.Tensor,
    seq_lens: torch.Tensor,
    token_to_req_indices: torch.Tensor,
    swa_block_table: torch.Tensor,
    swa_block_size: int,
    compressed_block_table: torch.Tensor | None,
    compressed_block_size: int,
    window_size: int,
    compress_ratio: int,
    topk: int,
    decode_compressed_indices_are_local: bool = False,
    decode_is_valid_token: torch.Tensor | None = None,
    swa_block_span: int | None = None,
    compressed_block_span: int | None = None,
    prefill_left_visible: torch.Tensor | None = None,
    prefill_right_visible: torch.Tensor | None = None,
    max_image_tokens: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.

    Produces ``sparse_indices`` of shape ``[num_tokens, swa_total_width +
    padded_topk]`` (the first ``swa_total_width`` columns are SWA slot ids, the
    rest are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length
    per token). Decode tokens read precomputed SWA/compressed indices; prefill
    tokens derive their SWA window from the position and translate local
    compressed indices to global slots via the block tables.

    When ``prefill_left_visible``/``prefill_right_visible`` are given (vision
    variant), the SWA column region widens by ``max_image_tokens`` and prefill
    tokens inside an image span get a bidirectionally widened window; decode
    rows are padded with -1 across the extra columns.
    """
    assert decode_swa_indices.dtype == torch.int32
    assert decode_swa_indices.dim() == 2
    swa_index_width = decode_swa_indices.shape[-1]
    assert swa_index_width >= window_size
    has_image = prefill_left_visible is not None
    image_width = max_image_tokens if has_image else 0
    swa_total_width = swa_index_width + image_width
    if has_image:
        assert prefill_right_visible is not None
    if decode_compressed_topk_lens is not None:
        assert decode_compressed_topk_lens.dtype == torch.int32
    assert prefill_topk_indices.dtype == torch.int32
    assert prefill_topk_indices.dim() == 2
    assert query_start_loc.dtype == torch.int32
    assert seq_lens.dtype == torch.int32
    assert token_to_req_indices.dtype == torch.int32
    assert swa_block_table.dtype == torch.int32

    num_decode_tokens = decode_swa_indices.shape[0]
    num_prefill_tokens = prefill_topk_indices.shape[0]
    num_tokens = num_decode_tokens + num_prefill_tokens
    assert token_to_req_indices.shape[0] >= num_tokens
    if decode_compressed_topk_lens is not None:
        assert decode_compressed_topk_lens.shape[0] >= num_decode_tokens

    decode_compressed_topk = 0
    if decode_compressed_indices is None:
        decode_compressed_indices = prefill_topk_indices
    else:
        assert decode_compressed_indices.dtype == torch.int32
        assert decode_compressed_indices.dim() == 2
        assert decode_compressed_indices.shape[0] == num_decode_tokens
        decode_compressed_topk = decode_compressed_indices.shape[-1]
    if decode_compressed_topk > 0 and decode_compressed_indices_are_local:
        assert decode_is_valid_token is not None
        assert decode_is_valid_token.dtype == torch.bool
        assert decode_is_valid_token.shape[0] >= num_decode_tokens
    else:
        decode_is_valid_token = token_to_req_indices

    if compressed_block_table is None:
        compressed_block_table = swa_block_table
    assert compressed_block_table.dtype == torch.int32
    has_decode_compressed_lens = decode_compressed_topk_lens is not None
    if decode_compressed_topk_lens is None:
        decode_compressed_topk_lens = token_to_req_indices

    # The FlashInfer TRTLLM-gen sparse-MLA kernels require every per-token topk
    # index row to start on a 16-byte boundary: the kernel loads the compressed
    # indices with 128-bit (16-byte) vectorized loads, so a misaligned row would
    # fault or read across rows. 16 bytes = 4 int32 indices, so round the topk
    # width (and hence the row stride, since the SWA columns are fixed-width) up
    # to a multiple of 4. The extra columns are filled with -1 (invalid) and bounded
    # by ``sparse_topk_lens``, so padding never changes the attention result.
    padded_topk = max(topk, decode_compressed_topk)
    padded_topk = (padded_topk + 3) // 4 * 4
    sparse_indices = torch.empty(
        (num_tokens, swa_total_width + padded_topk),
        dtype=torch.int32,
        device=decode_swa_indices.device,
    )
    sparse_topk_lens = torch.empty(
        num_tokens, dtype=torch.int32, device=decode_swa_indices.device
    )
    if num_tokens == 0:
        return sparse_indices, sparse_topk_lens

    window_block_size = triton.next_power_of_2(max(swa_index_width, 1))
    topk_block_size = triton.next_power_of_2(max(padded_topk, 1))
    max_block_size = max(window_block_size, topk_block_size)
    num_warps = 4 if max_block_size >= 256 else 1

    # block_span = page_stride / token_stride; == block_size (no-op) for unpacked KV.
    swa_span = swa_block_size if swa_block_span is None else swa_block_span
    compressed_span = (
        compressed_block_size
        if compressed_block_span is None
        else compressed_block_span
    )
    _build_flashinfer_mixed_sparse_indices_kernel[(num_tokens,)](
        sparse_indices,
        sparse_indices.stride(0),
        sparse_topk_lens,
        decode_swa_indices,
        decode_swa_indices.stride(0),
        decode_compressed_indices,
        decode_compressed_indices.stride(0),
        decode_compressed_topk_lens,
        decode_is_valid_token,
        prefill_topk_indices,
        prefill_topk_indices.stride(0),
        prefill_left_visible if has_image else token_to_req_indices,
        prefill_right_visible if has_image else token_to_req_indices,
        query_start_loc,
        seq_lens,
        token_to_req_indices,
        swa_block_table,
        swa_block_table.stride(0),
        swa_block_size,
        swa_span,
        compressed_block_table,
        compressed_block_table.stride(0),
        compressed_block_size,
        compressed_span,
        NUM_DECODE_TOKENS=num_decode_tokens,
        WINDOW_SIZE=window_size,
        SWA_INDEX_WIDTH=swa_index_width,
        IMAGE_WIDTH=image_width,
        COMPRESS_RATIO=compress_ratio,
        TOP_K=topk,
        PADDED_TOP_K=padded_topk,
        PREFILL_TOPK_STRIDE=prefill_topk_indices.shape[-1],
        DECODE_COMPRESSED_TOPK=decode_compressed_topk,
        DECODE_COMPRESSED_INDICES_ARE_LOCAL=decode_compressed_indices_are_local,
        HAS_DECODE_COMPRESSED_LENS=has_decode_compressed_lens,
        WINDOW_BLOCK_SIZE=window_block_size,
        TOPK_BLOCK_SIZE=topk_block_size,
        num_warps=num_warps,
    )
    return sparse_indices, sparse_topk_lens

def compute_global_topk_indices_and_lens(
    topk_indices: torch.Tensor,
    token_to_req_indices: torch.Tensor,
    block_table: torch.Tensor,
    block_size: int,
    is_valid_token: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Map local topk indices to global KV cache slots and count valid entries.

    Fuses three operations into a single kernel:
    1. Block-table lookup (local index → global slot id)
    2. Valid-entry counting (topk_lens per token)
    3. Masking padding tokens to length 0
    """
    num_tokens = topk_indices.shape[0]
    global_topk_indices = torch.empty_like(topk_indices)
    topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
    _compute_global_topk_indices_and_lens_kernel[(num_tokens,)](
        global_topk_indices,
        global_topk_indices.stride(0),
        topk_lens,
        topk_indices,
        topk_indices.stride(0),
        topk_indices.shape[-1],
        token_to_req_indices,
        block_table,
        block_table.stride(0),
        block_size,
        is_valid_token,
        TRITON_BLOCK_SIZE=1024,
    )
    return global_topk_indices, topk_lens

def dequantize_and_gather_k_cache(
    # [num_reqs, max_num_tokens, head_size]
    out: torch.Tensor,
    # [num_blocks, block_size, head_bytes]
    k_cache: torch.Tensor,
    # [num_reqs]
    seq_lens: torch.Tensor,
    # [num_reqs]
    gather_lens: torch.Tensor | None,
    # [num_reqs, max_blocks_per_seq]
    block_table: torch.Tensor,
    block_size: int,
    offset: int,
    use_fnuz: bool = False,
) -> None:
    """Dequantize and gather a paged DSv4 K cache.

    ``use_fnuz`` MUST match the encoder of the specific cache being read:
    ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere),
    ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder
    writes FNUZ on gfx942 and OCP on gfx950).
    """
    if has_cutedsl():
        # lazily import, otherwise some tests fail due to CUDA driver init failure.
        from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (
            _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL,
        )

        _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL(
            out=out,
            k_cache=k_cache,
            seq_lens=seq_lens,
            gather_lens=gather_lens,
            block_table=block_table,
            block_size=block_size,
            offset=offset,
        )
        return

    dequantize_and_gather_k_cache_triton(
        out,
        k_cache,
        seq_lens,
        gather_lens,
        block_table,
        block_size,
        offset,
        use_fnuz=use_fnuz,
    )

def fused_indexer_q_rope_quant(
    positions: torch.Tensor,
    index_q: torch.Tensor,
    index_q_cos_sin_cache: torch.Tensor,
    # Index weights
    index_weights: torch.Tensor,
    index_weights_softmax_scale: float,
    index_weights_head_scale: float,
    use_fp4: bool = False,
) -> tuple[
    torch.Tensor | tuple[torch.Tensor, torch.Tensor],
    torch.Tensor,
]:
    """Fused RoPE + quantize Q for the sparse indexer.

    Weight-fold semantics (important — the two paths differ):

    FP8 path (use_fp4=False, default):
        q_fp8      : (T, H, HEAD_DIM) platform fp8 (e4m3fnuz on gfx942,
                     e4m3fn elsewhere); per-token-per-head scalar scale
                     (NOT stored — folded into weights below)
        weights_out = weights * q_scale * softmax_scale * head_scale
        Rationale: a single per-token q_scale is a scalar the downstream FP8
        logits kernel would otherwise multiply in. Folding it into `weights`
        avoids emitting a separate tensor and is free for the logits kernel.

    MXFP4 path (use_fp4=True):
        q_packed   : (T, H, HEAD_DIM // 2) uint8 (2 E2M1 nibbles per byte)
        q_scale    : (T, H, HEAD_DIM // MXFP4_BLOCK_SIZE) uint8 ue8m0 bytes
        weights_out = weights * softmax_scale * head_scale
        Rationale: MXFP4 has PER-BLOCK (32-element) scales that live with
        the Q values — they cannot be folded into a per-token weight
        scalar, so `weights` carries only the softmax and head scales.

    Returns (q_quant, weights_out) where q_quant is either a Tensor (FP8) or
    a (values, scales) tuple (MXFP4). This matches the union type accepted
    by `SparseAttnIndexer.forward_*`.
    """
    assert positions.ndim == 1
    assert index_q.ndim == 3
    assert index_q_cos_sin_cache.ndim == 2

    num_tokens = positions.shape[0]
    num_index_q_heads = index_q.shape[1]
    index_q_head_dim = index_q.shape[2]

    index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)

    if use_fp4:
        assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, (
            f"head_dim={index_q_head_dim} must be a multiple of MXFP4 block "
            f"size {MXFP4_BLOCK_SIZE}"
        )
        num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE
        index_q_packed = torch.empty(
            (num_tokens, num_index_q_heads, index_q_head_dim // 2),
            dtype=torch.uint8,
            device=index_q.device,
        )
        index_q_scale = torch.empty(
            (num_tokens, num_index_q_heads, num_scale_blocks),
            dtype=torch.uint8,
            device=index_q.device,
        )
        if has_cutedsl():
            # lazily import, otherwise some tests fail due to CUDA driver init failure.
            from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
                _INDEXER_Q_MXFP4_KERNEL,
            )

            _INDEXER_Q_MXFP4_KERNEL(
                positions=positions,
                q=index_q,
                cos_sin_cache=index_q_cos_sin_cache,
                weights=index_weights,
                weights_softmax_scale=index_weights_softmax_scale,
                weights_head_scale=index_weights_head_scale,
                q_packed=index_q_packed,
                q_scale=index_q_scale,
                weights_out=index_weights_out,
            )
        elif current_platform.is_xpu():
            torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4(
                index_q,
                positions,
                index_q_cos_sin_cache,
                index_weights,
                index_weights_softmax_scale,
                index_weights_head_scale,
                index_q_packed,
                index_q_scale,
                index_weights_out,
            )
        else:
            _FUSED_INDEXER_Q_ROPE_MXFP4_TRITON_KERNEL(
                positions,
                index_q,
                index_q_cos_sin_cache,
                index_weights,
                index_weights_softmax_scale,
                index_weights_head_scale,
                index_q_packed,
                index_q_scale,
                index_weights_out,
            )

        # Values stay uint8 (2 E2M1 nibbles per byte). Scales are 4 ue8m0
        # bytes per (token, head) reinterpreted as one int32, then squeezed
        # from (T, H, 1) to (T, H) to match DeepGEMM's expected q_sf rank
        # (prefill wants 2-D (seq_len, num_heads); decode reshapes this to
        # 3-D (batch, next_n, num_heads)).
        return (
            index_q_packed,
            index_q_scale.view(torch.int32).squeeze(-1),
        ), index_weights_out

    fp8_dtype = current_platform.fp8_dtype()
    use_fnuz = fp8_dtype == torch.float8_e4m3fnuz
    fp8_max = 224.0 if use_fnuz else 448.0
    index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype)
    if has_cutedsl():
        # lazily import, otherwise some tests fail due to CUDA driver init failure.
        from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
            _INDEXER_Q_FP8_KERNEL,
        )

        _INDEXER_Q_FP8_KERNEL(
            positions=positions,
            q=index_q,
            cos_sin_cache=index_q_cos_sin_cache,
            weights=index_weights,
            weights_softmax_scale=index_weights_softmax_scale,
            weights_head_scale=index_weights_head_scale,
            q_fp8=index_q_fp8.view(torch.uint8),
            weights_out=index_weights_out,
        )
    elif current_platform.is_xpu():
        torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_fp8(
            index_q,
            positions,
            index_q_cos_sin_cache,
            index_weights,
            index_weights_softmax_scale,
            index_weights_head_scale,
            index_q_fp8,
            index_weights_out,
        )
    elif current_platform.is_cpu():
        from vllm._custom_ops import fused_indexer_q_rope_quant_cpu

        fused_indexer_q_rope_quant_cpu(
            positions.contiguous(),
            index_q.to(torch.float32),
            index_q_cos_sin_cache.to(torch.float32).contiguous(),
            index_q_fp8,
            index_weights.to(torch.float32),
            index_weights_softmax_scale,
            index_weights_head_scale,
            index_weights_out,
        )
    else:
        _FUSED_INDEXER_Q_ROPE_QUANT_TRITON_KERNEL(
            positions,
            index_q,
            index_q_cos_sin_cache,
            index_weights,
            index_weights_softmax_scale,
            index_weights_head_scale,
            index_q_fp8,
            index_weights_out,
            fp8_max=fp8_max,
            use_fnuz=use_fnuz,
        )
    return index_q_fp8, index_weights_out

def fused_inv_rope_fp8_quant(
    o: torch.Tensor,
    positions: torch.Tensor,
    cos_sin_cache: torch.Tensor,
    n_groups: int,
    heads_per_group: int,
    nope_dim: int = 448,
    rope_dim: int = 64,
    quant_group_size: int = 128,
    tma_aligned_scales: bool = False,
    quantize: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fused inverse RoPE + block-scaled FP8 quantization.

    Args:
        o: Attention output [num_tokens, num_heads, head_dim] bf16.
        positions: Token positions [num_tokens] int64.
        cos_sin_cache: Precomputed [max_pos, rope_dim] with cos||sin.
        n_groups: Number of output groups.
        heads_per_group: Heads per group.
        nope_dim: Non-RoPE dimensions per head (default 448).
        rope_dim: RoPE dimensions per head (default 64).
        quant_group_size: FP8 quantization block size (default 128).
        tma_aligned_scales: Output INT32 packed UE8M0 for SM100 (True)
                            or FP32 for SM90 (False).
        quantize: Quantize the rotated output to FP8 and return its scales.

    Returns:
        Rotated output in [T, G, D] and its FP8 scales. The scale tensor is
        empty when quantization is disabled.
    """
    from vllm.utils.deep_gemm import get_tma_aligned_size

    num_tokens, num_heads, head_dim = o.shape
    assert num_heads == n_groups * heads_per_group
    assert head_dim == nope_dim + rope_dim
    assert head_dim % quant_group_size == 0
    assert rope_dim % 2 == 0
    assert cos_sin_cache.shape[-1] == rope_dim
    assert cos_sin_cache.dtype == torch.float32

    d = heads_per_group * head_dim
    num_scale_blocks = d // quant_group_size
    chunks_per_head = head_dim // quant_group_size

    fp8_dtype = torch.float8_e4m3fn
    fp8_max = torch.finfo(fp8_dtype).max

    tma_aligned_T = get_tma_aligned_size(num_tokens, 4) if quantize else num_tokens
    if quantize and tma_aligned_scales:
        assert chunks_per_head % 4 == 0
        packed_sf_k = (num_scale_blocks + 3) // 4
        scale_inner = packed_sf_k
    elif quantize:
        scale_inner = num_scale_blocks
    else:
        scale_inner = 0

    # Run kernel through a custom op so inductor sees an opaque boundary.
    # It's a pytorch bug, see https://github.com/vllm-project/vllm/issues/41106
    out_buf, scale_buf = torch.ops.vllm.fused_inv_rope_fp8_quant_kernel(
        o,
        positions,
        cos_sin_cache,
        heads_per_group,
        quant_group_size,
        chunks_per_head,
        nope_dim,
        rope_dim // 2,
        tma_aligned_scales,
        fp8_max,
        tma_aligned_T,
        num_tokens,
        n_groups,
        d,
        scale_inner,
        quantize,
    )
    output = out_buf.transpose(0, 1)
    scales = scale_buf.transpose(0, 1) if quantize else scale_buf
    return output, scales

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,
    )

def quantize_and_insert_k_cache(
    k: torch.Tensor,  # [num_tokens, 512] bf16
    k_cache: torch.Tensor,  # [num_blocks, block_bytes] uint8
    slot_mapping: torch.Tensor,  # [num_tokens] int64
    block_size: int = 64,
    is_ue8m0: bool = True,
    use_fnuz: bool = False,
):
    """
    Quantize K tensor and insert into paged K cache.

    K Cache block layout (block_size=64 tokens):
    - First 64 * 576 = 36864 bytes: Token data
      - Each token: 448 bytes (fp8) + 128 bytes (bf16)
    - Next 64 * 8 = 512 bytes: Scales
      - Each token: 8 bytes (uint8 scales, 7 real + 1 padding)
    - Padded to multiple of 576

    ``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on
    platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3,
    which is used by OCP-encoded caches even on gfx942.
    """
    assert k.dim() == 2 and k.shape[1] == 512, (
        f"K must be [num_tokens, 512], got {k.shape}"
    )
    assert k.dtype == torch.bfloat16, f"K must be bf16, got {k.dtype}"
    assert is_ue8m0, "Only support ue8m0 quantization."

    # NOTE: When using DP, slot_mapping.shape[0] can be less than k.shape[0] due to
    # padding. Always use slot_mapping.shape[0] as the token count.
    num_tokens = slot_mapping.shape[0]
    block_stride = k_cache.stride(0)  # bytes per block

    TOKEN_FP8_DIM = 448
    TOKEN_BF16_DIM = 64
    TOKEN_SCALE_DIM = 8
    QUANT_BLOCK_SIZE = 64
    if use_fnuz:
        if not current_platform.is_fp8_fnuz():
            raise ValueError("use_fnuz=True requires a platform using FNUZ FP8")
        _, FP8_MAX = get_fp8_min_max()
    else:
        FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
    TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2

    grid = (num_tokens,)

    quantize_and_insert_k_kernel[grid](
        k,
        slot_mapping,
        k_cache,
        num_tokens,
        input_dim=512,
        fp8_dim=TOKEN_FP8_DIM,
        bf16_dim=TOKEN_BF16_DIM,
        scale_dim=TOKEN_SCALE_DIM,
        quant_block=QUANT_BLOCK_SIZE,
        cache_block_size=block_size,
        token_data_size=TOKEN_DATA_SIZE,
        block_stride=block_stride,
        fp8_max=FP8_MAX,
        n_quant_blocks=8,
        use_fnuz=use_fnuz,
    )