vllm.models.deepseek_v41.common.ops.cache_utils

Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index preparation.

• quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into the paged cache.
• dequantize_and_gather_k_cache: gather and dequantize FP8 K from the paged cache for sparse/SWA prefill.
• compute_global_topk_indices_and_lens: map local topk indices to global KV cache slots and count valid entries.
• combine_topk_swa_indices: concatenate topk compressed indices with SWA window indices for sparse prefill.

Functions:

• build_flashinfer_mixed_sparse_indices –

  Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.

• compute_global_topk_indices_and_lens –

  Map local topk indices to global KV cache slots and count valid entries.

• dequantize_and_gather_k_cache –

  Dequantize and gather a paged DSv4 K cache.

• gather_num_workers –

  Grid width for a gather of at most max_gather_len tokens.

• quantize_and_insert_k_cache –

  Quantize K tensor and insert into paged K cache.

_combine_topk_swa_warmup_inputs(vllm_config)

One warmup row per v4.1 layer type (compress ratio) in the model.

Every layer slices the shared index_topk-wide indices buffer, so the padded width is the same for all rows; SWA-only layers (ratio 0) pass TOP_K=0.

Source code in vllm/models/deepseek_v41/common/ops/cache_utils.py

def _combine_topk_swa_warmup_inputs(vllm_config: Any) -> _WarmupInputRows:
    """One warmup row per v4.1 layer type (compress ratio) in the model.

    Every layer slices the shared ``index_topk``-wide indices buffer, so the
    padded width is the same for all rows; SWA-only layers (ratio 0) pass
    ``TOP_K=0``.
    """
    hf_config = vllm_config.model_config.hf_config
    compress_ratios = get_compress_ratios(hf_config) or [0]
    topk = hf_config.index_topk
    return zip_inputs(
        *(
            dict(compress_ratio=ratio, topk=topk if ratio > 0 else 0, topk_width=topk)
            for ratio in sorted({int(ratio) for ratio in compress_ratios})
        )
    )

_dequantize_and_gather_k_mxfp8_kernel(out_ptr, out_stride0, out_stride1, k_cache_ptr, seq_lens_ptr, block_table_ptr, offset, gather_lens_ptr, max_blocks_per_seq, head_dim, scale_dim, quant_block, cache_block_size, block_stride, use_fnuz=False)

Gather V4.1 MXFP8 rows into a bf16 workspace, dequantizing in place.

Source code in vllm/models/deepseek_v41/common/ops/cache_utils.py

@triton.jit
def _dequantize_and_gather_k_mxfp8_kernel(
    out_ptr,
    out_stride0,
    out_stride1,
    k_cache_ptr,
    seq_lens_ptr,
    block_table_ptr,
    offset,
    gather_lens_ptr,
    max_blocks_per_seq: tl.constexpr,
    head_dim: tl.constexpr,  # 512
    scale_dim: tl.constexpr,  # 16
    quant_block: tl.constexpr,  # 32
    cache_block_size: tl.constexpr,
    block_stride: tl.constexpr,
    use_fnuz: tl.constexpr = False,
):
    """Gather V4.1 MXFP8 rows into a bf16 workspace, dequantizing in place."""
    batch_idx = tl.program_id(0)
    worker_id = tl.program_id(1)
    num_workers = tl.num_programs(1)

    seq_len = tl.load(seq_lens_ptr + batch_idx)
    if gather_lens_ptr is not None:  # noqa: SIM108
        gather_len = tl.load(gather_lens_ptr + batch_idx)
    else:
        gather_len = seq_len
    start_pos = seq_len - gather_len

    d = tl.arange(0, head_dim)
    s = tl.arange(0, scale_dim)
    for i in range(worker_id, gather_len, num_workers):
        pos = start_pos + i
        block_in_seq = pos // cache_block_size
        pos_in_block = pos % cache_block_size

        block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
        physical_block_idx = tl.load(block_table_row_ptr + block_in_seq)  # int32

        # int64: physical_block_idx * block_stride can exceed 2^31.
        cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
        token_data_ptr = cache_block_ptr + pos_in_block * (head_dim + scale_dim)
        token_scale_ptr = token_data_ptr + head_dim

        x_uint8 = tl.load(token_data_ptr + d)
        if use_fnuz:
            x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
        else:
            x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
        tiles = tl.reshape(x_fp8.to(tl.float32), (scale_dim, quant_block))

        # UE8M0: scale = 2^(stored_value - 127).
        encoded = tl.load(token_scale_ptr + s)
        scale = tl.exp2(encoded.to(tl.float32) - 127.0)
        dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,))

        output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
        tl.store(output_row_ptr + d, dequant.to(tl.bfloat16))

_dequantize_and_gather_k_nvfp4_kernel(out_ptr, out_stride0, out_stride1, k_cache_ptr, seq_lens_ptr, block_table_ptr, offset, gather_lens_ptr, max_blocks_per_seq, head_dim, scale_dim, quant_block, cache_block_size, block_stride, TILE)

Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace, TILE tokens per iteration. A nibble placed as fp16 bits (sign at 15, exponent and mantissa at 11:9) reads as e2m1 * 2^-14; times scale * 2^14 the fp32 product is exact, so the bf16 result matches a direct decode.

Source code in vllm/models/deepseek_v41/common/ops/cache_utils.py

@triton.jit
def _dequantize_and_gather_k_nvfp4_kernel(
    out_ptr,
    out_stride0,
    out_stride1,
    k_cache_ptr,
    seq_lens_ptr,
    block_table_ptr,
    offset,
    gather_lens_ptr,
    max_blocks_per_seq: tl.constexpr,
    head_dim: tl.constexpr,  # 512
    scale_dim: tl.constexpr,  # 32
    quant_block: tl.constexpr,  # 16
    cache_block_size: tl.constexpr,
    block_stride: tl.constexpr,
    TILE: tl.constexpr,  # tokens per iteration
):
    """Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace, TILE tokens
    per iteration. A nibble placed as fp16 bits (sign at 15, exponent and
    mantissa at 11:9) reads as e2m1 * 2^-14; times scale * 2^14 the fp32 product
    is exact, so the bf16 result matches a direct decode."""
    batch_idx = tl.program_id(0)
    worker_id = tl.program_id(1)
    num_workers = tl.num_programs(1)

    seq_len = tl.load(seq_lens_ptr + batch_idx)
    if gather_lens_ptr is not None:  # noqa: SIM108
        gather_len = tl.load(gather_lens_ptr + batch_idx)
    else:
        gather_len = seq_len
    start_pos = seq_len - gather_len

    packed_bytes: tl.constexpr = head_dim // 2
    group_bytes: tl.constexpr = quant_block // 2
    # Byte b of a record packs elements 2b (low nibble) and 2b + 1, which share
    # scale b // group_bytes: index the bytes as [scale, byte within group].
    byte = (
        tl.arange(0, scale_dim)[:, None] * group_bytes
        + tl.arange(0, group_bytes)[None, :]
    )
    t = tl.arange(0, TILE)
    d = tl.arange(0, head_dim)
    block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
    out_row_ptr = out_ptr + batch_idx.to(tl.int64) * out_stride0
    for tile_start in range(worker_id * TILE, gather_len, num_workers * TILE):
        i = tile_start + t
        valid = i < gather_len
        pos = start_pos + i
        physical_block_idx = tl.load(
            block_table_row_ptr + pos // cache_block_size, mask=valid, other=0
        )
        page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
        token_ptr = page + (pos % cache_block_size) * (packed_bytes + scale_dim)

        packed = tl.load(
            token_ptr[:, None, None] + byte[None, :, :],
            mask=valid[:, None, None],
            other=0,
        ).to(tl.int32)
        sf = tl.load(
            token_ptr[:, None] + packed_bytes + tl.arange(0, scale_dim)[None, :],
            mask=valid[:, None],
            other=0,
        )
        scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32) * 16384.0
        scale = scale[:, :, None]

        lo = packed & 0xF
        hi = packed >> 4
        lo = ((lo & 7) << 9) | ((lo & 8) << 12)
        hi = ((hi & 7) << 9) | ((hi & 8) << 12)
        lo = lo.to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32) * scale
        hi = hi.to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32) * scale
        dequant = tl.reshape(tl.join(lo, hi), (TILE, head_dim))

        out_ptrs = (
            out_row_ptr + (offset + i).to(tl.int64)[:, None] * out_stride1 + d[None, :]
        )
        tl.store(out_ptrs, dequant.to(tl.bfloat16), mask=valid[:, None])

_dequantize_and_gather_k_one_pass_kernel(out_ptr, out_stride0, out_stride1, k_cache_ptr, seq_lens_ptr, block_table_ptr, offset, gather_lens_ptr, max_blocks_per_seq, fp8_dim, bf16_dim, scale_dim, quant_block, cache_block_size, token_data_size, block_stride, output_dim, fp8_max, n_quant_blocks, use_fnuz=False)

_dequantize_and_gather_k_kernel with the token read in one tile.

The seven quantization blocks tile the fp8 region exactly, so the per-block form's bounds never bite and its twenty-two memory operations per token reduce to four. Selected on ROCm; CUDA keeps the per-block kernel.

Source code in vllm/models/deepseek_v41/common/ops/cache_utils.py

@triton.jit
def _dequantize_and_gather_k_one_pass_kernel(
    out_ptr,
    out_stride0,
    out_stride1,
    k_cache_ptr,
    seq_lens_ptr,
    block_table_ptr,
    offset,
    gather_lens_ptr,
    # Constants
    max_blocks_per_seq: tl.constexpr,
    fp8_dim: tl.constexpr,  # 448
    bf16_dim: tl.constexpr,  # 64
    scale_dim: tl.constexpr,  # 8
    quant_block: tl.constexpr,  # 64 (quantization block size)
    cache_block_size: tl.constexpr,  # 64 or 128 (paged cache block size)
    token_data_size: tl.constexpr,  # 576 bytes per token data
    block_stride: tl.constexpr,  # total bytes per block (padded) int32
    output_dim: tl.constexpr,  # 512
    fp8_max: tl.constexpr,
    n_quant_blocks: tl.constexpr,  # 7 real blocks; the per-block form only
    use_fnuz: tl.constexpr = False,
):
    """``_dequantize_and_gather_k_kernel`` with the token read in one tile.

    The seven quantization blocks tile the fp8 region exactly, so the per-block
    form's bounds never bite and its twenty-two memory operations per token
    reduce to four. Selected on ROCm; CUDA keeps the per-block kernel.
    """
    batch_idx = tl.program_id(0)
    worker_id = tl.program_id(1)
    num_workers = tl.num_programs(1)

    seq_len = tl.load(seq_lens_ptr + batch_idx)
    if gather_lens_ptr is not None:  # noqa: SIM108
        gather_len = tl.load(gather_lens_ptr + batch_idx)
    else:
        gather_len = seq_len
    start_pos = seq_len - gather_len

    for i in range(worker_id, gather_len, num_workers):
        pos = start_pos + i
        block_in_seq = pos // cache_block_size
        pos_in_block = pos % cache_block_size

        block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
        physical_block_idx = tl.load(block_table_row_ptr + block_in_seq)
        cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride

        token_fp8_ptr = cache_block_ptr + pos_in_block * token_data_size
        token_bf16_ptr = token_fp8_ptr + fp8_dim
        token_scale_ptr = (
            cache_block_ptr
            + cache_block_size * token_data_size
            + pos_in_block * scale_dim
        )
        output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1

        # The tile is scale_dim * quant_block = 512 against an fp8 region of
        # 448: the tail reads the token's own bf16 bytes, is scaled by a scale
        # nothing reads, and is dropped by the store mask.
        tile_dim: tl.constexpr = scale_dim * quant_block
        tl.static_assert(tile_dim <= token_data_size, "tile outruns the token")
        fp8_offsets = tl.arange(0, tile_dim)
        x_uint8 = tl.load(token_fp8_ptr + fp8_offsets)
        if use_fnuz:
            x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
        else:
            x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)

        encoded_scale = tl.load(token_scale_ptr + tl.arange(0, scale_dim))
        scale = tl.exp2(encoded_scale.to(tl.float32) - 127.0)  # UE8M0
        tiles = tl.reshape(x_fp8.to(tl.float32), (scale_dim, quant_block))
        x_dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (tile_dim,))
        tl.store(
            output_row_ptr + fp8_offsets,
            x_dequant.to(tl.bfloat16),
            mask=fp8_offsets < fp8_dim,
        )

        bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16))
        bf16_offsets = tl.arange(0, bf16_dim)
        tl.store(
            output_row_ptr + fp8_dim + bf16_offsets,
            tl.load(bf16_cache_ptr + bf16_offsets),
        )

_quantize_and_insert_k_mxfp8_kernel(k_ptr, slot_mapping_ptr, k_cache_ptr, num_tokens, head_dim, scale_dim, quant_block, cache_block_size, block_stride, fp8_max, use_fnuz=False)

Quantize a bf16 K row to MXFP8 and insert it into a V4.1 paged cache.

One program per token. Unlike the V4 record there is no bf16 tail, so every dim is scaled.

Source code in vllm/models/deepseek_v41/common/ops/cache_utils.py

@triton.jit
def _quantize_and_insert_k_mxfp8_kernel(
    k_ptr,  # [num_tokens, 512] bf16
    slot_mapping_ptr,  # [num_tokens] int64
    k_cache_ptr,  # [num_blocks, block_bytes] as uint8 (flattened view)
    num_tokens,
    head_dim: tl.constexpr,  # 512, RoPE dims included
    scale_dim: tl.constexpr,  # 16
    quant_block: tl.constexpr,  # 32 dims per UE8M0 scale
    cache_block_size: tl.constexpr,  # paged cache block size, in tokens
    block_stride: tl.constexpr,  # total bytes per block (padded)
    fp8_max: tl.constexpr,
    use_fnuz: tl.constexpr = False,
):
    """Quantize a bf16 K row to MXFP8 and insert it into a V4.1 paged cache.

    One program per token. Unlike the V4 record there is no bf16 tail, so every
    dim is scaled.
    """
    pid = tl.program_id(0)
    if pid >= num_tokens:
        return

    slot_idx = tl.load(slot_mapping_ptr + pid)
    if slot_idx == -1:
        return

    block_idx = slot_idx // cache_block_size
    pos_in_block = slot_idx % cache_block_size

    # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks.
    cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride
    token_data_ptr = cache_block_ptr + pos_in_block * (head_dim + scale_dim)
    token_scale_ptr = token_data_ptr + head_dim

    d = tl.arange(0, head_dim)
    x = tl.load(k_ptr + pid.to(tl.int64) * head_dim + d).to(tl.float32)

    tiles = tl.reshape(x, (scale_dim, quant_block))
    # Floor the absmax like the CUDA encoder so an all-zero tile still lands on
    # a finite scale.
    amax = tl.maximum(tl.max(tl.abs(tiles), 1), 1e-4)
    exponent = tl.ceil(tl.log2(amax * (1.0 / fp8_max)))
    scaled = tl.clamp(
        tiles * tl.reshape(tl.exp2(-exponent), (scale_dim, 1)), -fp8_max, fp8_max
    )
    if use_fnuz:  # noqa: SIM108
        fp8 = scaled.to(tl.float8e4b8)
    else:
        fp8 = scaled.to(tl.float8e4nv)
    tl.store(
        token_data_ptr + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (head_dim,))
    )

    # UE8M0 encoding: stored_value = exponent + 127 (bias).
    encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), 255.0)
    tl.store(token_scale_ptr + tl.arange(0, scale_dim), encoded.to(tl.uint8))

build_flashinfer_mixed_sparse_indices(decode_swa_indices, decode_compressed_indices, decode_compressed_topk_lens, prefill_topk_indices, query_start_loc, seq_lens, token_to_req_indices, swa_block_table, swa_block_size, compressed_block_table, compressed_block_size, window_size, compress_ratio, topk, decode_compressed_indices_are_local=False, decode_is_valid_token=None, swa_block_span=None, compressed_block_span=None, *, replay_start)

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.

replay_start ([num_reqs], SWA bounded replay) lower-bounds every prefill token's window: positions below it hold no window KV.

Source code in vllm/models/deepseek_v41/common/ops/cache_utils.py

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,
    *,
    replay_start: torch.Tensor,
) -> 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.

    ``replay_start`` ([num_reqs], SWA bounded replay) lower-bounds every
    prefill token's window: positions below it hold no window KV.
    """
    assert decode_swa_indices.dtype == torch.int32
    assert decode_swa_indices.dim() == 2
    swa_total_width = decode_swa_indices.shape[-1]
    assert swa_total_width >= window_size
    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_total_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),
        query_start_loc,
        seq_lens,
        token_to_req_indices,
        replay_start,
        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_total_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

compute_global_topk_indices_and_lens(topk_indices, token_to_req_indices, block_table, block_size, is_valid_token)

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

Source code in vllm/models/deepseek_v41/common/ops/cache_utils.py

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

dequantize_and_gather_k_cache(out, k_cache, seq_lens, gather_lens, block_table, block_size, offset, use_fnuz=False, max_gather_len=None)

Dequantize and gather a paged DSv4 K cache.

The record is read off k_cache.shape[-1]; see the module header. Only the fp8 records have a CuteDSL gather, so NVFP4 always takes the Triton path.

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

Source code in vllm/models/deepseek_v41/common/ops/cache_utils.py

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,
    max_gather_len: int | None = None,
) -> None:
    """Dequantize and gather a paged DSv4 K cache.

    The record is read off ``k_cache.shape[-1]``; see the module header. Only
    the fp8 records have a CuteDSL gather, so NVFP4 always takes the Triton
    path.

    ``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() and k_cache.shape[-1] != V41_NVFP4_BYTES_PER_TOKEN:
        # 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,
        )

        mxfp8 = k_cache.shape[-1] == V41_BYTES_PER_TOKEN
        _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,
            fp8_dim=512 if mxfp8 else 448,
            group_size=V41_QUANT_BLOCK if mxfp8 else 64,
        )
        return

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

gather_num_workers(max_gather_len)

Grid width for a gather of at most max_gather_len tokens.

Source code in vllm/models/deepseek_v41/common/ops/cache_utils.py

def gather_num_workers(max_gather_len: int | None) -> int:
    """Grid width for a gather of at most ``max_gather_len`` tokens."""
    if max_gather_len is None:
        return _GATHER_MIN_WORKERS
    workers = cdiv(max_gather_len, _GATHER_STEPS_PER_WORKER)
    return max(_GATHER_MIN_WORKERS, min(_GATHER_MAX_WORKERS, workers))

quantize_and_insert_k_cache(k, k_cache, slot_mapping, block_size=64, is_ue8m0=True, use_fnuz=False, bytes_per_token=V4_BYTES_PER_TOKEN)

Quantize K tensor and insert into paged K cache.

bytes_per_token picks the record (see the module header): the V4 one, or V4.1's all-dims MXFP8 one.

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.

Source code in vllm/models/deepseek_v41/common/ops/cache_utils.py

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,
    bytes_per_token: int = V4_BYTES_PER_TOKEN,
):
    """Quantize K tensor and insert into paged K cache.

    ``bytes_per_token`` picks the record (see the module header): the V4 one,
    or V4.1's all-dims MXFP8 one.

    ``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

    if bytes_per_token == V41_BYTES_PER_TOKEN:
        if use_fnuz and not current_platform.is_fp8_fnuz():
            raise ValueError("use_fnuz=True requires a platform using FNUZ FP8")
        _quantize_and_insert_k_mxfp8_kernel[(num_tokens,)](
            k,
            slot_mapping,
            k_cache,
            num_tokens,
            head_dim=512,
            scale_dim=V41_NUM_SCALES,
            quant_block=V41_QUANT_BLOCK,
            cache_block_size=block_size,
            block_stride=block_stride,
            fp8_max=get_fp8_min_max()[1]
            if use_fnuz
            else torch.finfo(torch.float8_e4m3fn).max,
            use_fnuz=use_fnuz,
        )
        return

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