vllm.v1.attention.backends.mla.sparse_utils

Utility functions for sparse MLA backends.

Functions:

• flat_kv_row_view –

  Flat [row, head_dim] view of a paged cache and its physical rows per block.

• prepare_sparse_mla_safe_lengths –

  Install dummy slots for empty queries and return nonzero kernel lengths.

• request_row_bounds –

  Bounds of the runs of adjacent rows that belong to one request: run

• triton_convert_req_index_to_global_index –

  out[token_id, indice_id] =

• triton_filter_and_convert_dcp_index –

  Filter global per-request indices to this DCP rank's local slots.

_remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, count_valid)

Pick the column tiling for the index remap kernel.

Counting the valid slots per row is the only reason the column tiles have to talk to each other, so when counting give one program the whole row: the count becomes an in-register reduction plus a plain store, needing neither atomics nor a zero-initialized counter. Pad modest non-power-of-two widths (including GLM's 2176 entries) to one tile; larger widths stay tiled and atomic.

Returns:

• tuple[bool, int, int, int] –

  (single_tile, block_n, tiles_per_row, num_warps)

Source code in vllm/v1/attention/backends/mla/sparse_utils.py

def _remap_tiling(
    NUM_TOPK_TOKENS: int, BLOCK_N: int, count_valid: bool
) -> tuple[bool, int, int, int]:
    """Pick the column tiling for the index remap kernel.

    Counting the valid slots per row is the only reason the column tiles have to
    talk to each other, so when counting give one program the whole row: the
    count becomes an in-register reduction plus a plain store, needing neither
    atomics nor a zero-initialized counter. Pad modest non-power-of-two widths
    (including GLM's 2176 entries) to one tile; larger widths stay tiled and atomic.

    Returns:
        (single_tile, block_n, tiles_per_row, num_warps)

    """
    padded_width = triton.next_power_of_2(NUM_TOPK_TOKENS)
    single_tile = count_valid and (
        padded_width == NUM_TOPK_TOKENS or padded_width <= 4096
    )
    if single_tile:
        return True, padded_width, 1, 8
    return False, BLOCK_N, NUM_TOPK_TOKENS // BLOCK_N, 4

flat_kv_row_view(kv_cache, block_size)

Flat [row, head_dim] view of a paged cache and its physical rows per block.

Token offset is block_idx * block_stride_rows + offset_in_block. When other layers' pages sit between consecutive blocks of this cache, block_stride_rows exceeds block_size; those in-between rows are never indexed (triton_convert_req_index_to_global_index ensures this).

Source code in vllm/v1/attention/backends/mla/sparse_utils.py

def flat_kv_row_view(
    kv_cache: torch.Tensor,  # [num_blocks, block_size, head_dim]
    block_size: int,
) -> tuple[torch.Tensor, int]:
    """Flat [row, head_dim] view of a paged cache and its physical rows per block.

    Token offset is  ``block_idx * block_stride_rows + offset_in_block``.
    When other layers' pages sit between consecutive blocks of this cache,
    ``block_stride_rows`` exceeds ``block_size``; those in-between rows are never
    indexed (`triton_convert_req_index_to_global_index` ensures this).
    """
    num_blocks, _, head_dim = kv_cache.shape
    assert kv_cache.stride(0) % head_dim == 0, (
        "block stride is not a whole number of rows; flat row indexing would "
        "silently misaddress"
    )
    block_stride_rows = kv_cache.stride(0) // head_dim
    num_rows = (num_blocks - 1) * block_stride_rows + block_size
    rows = kv_cache.as_strided((num_rows, head_dim), (head_dim, 1))
    return rows, block_stride_rows

prepare_sparse_mla_safe_lengths(physical_indices, valid_counts)

Install dummy slots for empty queries and return nonzero kernel lengths.

Preserve the raw counts so empty outputs and LSE can be neutralized later.

Source code in vllm/v1/attention/backends/mla/sparse_utils.py

def prepare_sparse_mla_safe_lengths(
    physical_indices: torch.Tensor, valid_counts: torch.Tensor
) -> torch.Tensor:
    """Install dummy slots for empty queries and return nonzero kernel lengths.

    Preserve the raw counts so empty outputs and LSE can be neutralized later.
    """
    num_tokens = valid_counts.numel()
    safe_lengths = torch.empty(
        (num_tokens,), dtype=valid_counts.dtype, device=valid_counts.device
    )
    if num_tokens:
        sparse_mla_prepare_safe_lengths_kernel[(triton.cdiv(num_tokens, 256),)](
            physical_indices,
            valid_counts,
            safe_lengths,
            num_tokens,
            physical_indices.stride(0),
            valid_counts.stride(0),
            BLOCK=256,
        )
    return safe_lengths

request_row_bounds(req_idx)

Bounds of the runs of adjacent rows that belong to one request: run r is rows [bounds[r], bounds[r + 1]).

Under PCP a rank holds two adjacent chunk rows of a split prefill; the sparse backends give such a run one KV region.

Source code in vllm/v1/attention/backends/mla/sparse_utils.py

def request_row_bounds(req_idx: np.ndarray) -> np.ndarray:
    """Bounds of the runs of adjacent rows that belong to one request: run
    ``r`` is rows ``[bounds[r], bounds[r + 1])``.

    Under PCP a rank holds two adjacent chunk rows of a split prefill; the
    sparse backends give such a run one KV region.
    """
    assert req_idx.size > 0
    bounds = np.flatnonzero(np.r_[True, req_idx[1:] != req_idx[:-1], True])
    assert bounds.size - 1 == np.unique(req_idx).size, (
        "rows of one request must be adjacent"
    )
    return bounds

triton_convert_req_index_to_global_index(req_id, block_table, token_indices, BLOCK_SIZE=64, BLOCK_STRIDE_ROWS=None, NUM_TOPK_TOKENS=2048, BLOCK_N=128, HAS_PREFILL_WORKSPACE=False, prefill_workspace_request_ids=None, prefill_workspace_starts=None, prefill_workspace_rank_stride=None, dcp_size=1, dcp_rank=0, cp_kv_cache_interleave_size=1, return_valid_counts=False, out=None, valid_counts_out=None)

out[token_id, indice_id] = block_table[req_id[token_id], token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE + token_indices[token_id, indice_id] % BLOCK_SIZE

Only when token_indices[token_id, indice_id] == -1 do we output -1. For safety, we also output -1 if the derived block_id would be out-of-bounds.

When HAS_PREFILL_WORKSPACE is True, prefill tokens are mapped to workspace offsets instead of global cache slots. prefill_workspace_request_ids and prefill_workspace_starts must be provided.

int32 [num_tokens], -1 for decode else

prefill request index (maps to prefill_workspace_starts)

prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace starts for each prefill request

When return_valid_counts is True, also returns the count of valid (non -1) indices per row, computed during the same kernel pass (no extra overhead).

Source code in vllm/v1/attention/backends/mla/sparse_utils.py

def triton_convert_req_index_to_global_index(
    req_id: torch.Tensor,  # int32 [num_tokens]
    block_table: torch.Tensor,  # int32 [num_requests, max_num_blocks_per_req]
    token_indices: torch.Tensor,  # int32 [num_tokens, NUM_TOPK_TOKENS]
    BLOCK_SIZE: int = 64,
    BLOCK_STRIDE_ROWS: int | None = None,
    NUM_TOPK_TOKENS: int = 2048,
    BLOCK_N: int = 128,  # tile width along columns
    HAS_PREFILL_WORKSPACE: bool = False,
    prefill_workspace_request_ids: torch.Tensor | None = None,
    prefill_workspace_starts: torch.Tensor | None = None,
    prefill_workspace_rank_stride: int | None = None,
    dcp_size: int = 1,
    dcp_rank: int = 0,
    cp_kv_cache_interleave_size: int = 1,
    return_valid_counts: bool = False,
    out: torch.Tensor | None = None,
    valid_counts_out: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """out[token_id, indice_id] =
        block_table[req_id[token_id],
            token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE
        + token_indices[token_id, indice_id] % BLOCK_SIZE

    Only when token_indices[token_id, indice_id] == -1 do we output -1.
    For safety, we also output -1 if the derived block_id would be
        out-of-bounds.

    When HAS_PREFILL_WORKSPACE is True, prefill tokens are mapped to workspace offsets
    instead of global cache slots. prefill_workspace_request_ids and
    prefill_workspace_starts must be provided.

    prefill_workspace_request_ids: int32 [num_tokens], -1 for decode else
        prefill request index (maps to prefill_workspace_starts)
    prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace
        starts for each prefill request

    When return_valid_counts is True, also returns the count of valid (non -1)
    indices per row, computed during the same kernel pass (no extra overhead).
    """
    assert req_id.dtype == torch.int32
    assert block_table.dtype == torch.int32
    assert token_indices.dtype == torch.int32
    assert req_id.shape[0] == token_indices.shape[0], (
        f"req_id ({req_id.shape[0]}) and token_indices ({token_indices.shape[0]}) "
        "must cover the same tokens; the grid is sized by req_id but the output "
        "is allocated like token_indices, so a longer req_id writes out of bounds"
    )
    assert token_indices.shape[1] == NUM_TOPK_TOKENS
    assert NUM_TOPK_TOKENS % BLOCK_N == 0, (
        f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})"
    )

    if HAS_PREFILL_WORKSPACE:
        assert prefill_workspace_request_ids is not None
        assert prefill_workspace_starts is not None
        assert prefill_workspace_request_ids.dtype == torch.int32
        assert prefill_workspace_starts.dtype == torch.int32

    num_tokens = req_id.shape[0]
    max_num_blocks_per_req = block_table.shape[1]

    single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts)

    # Ensure contiguous tensors on the same device
    req_id_c = req_id.contiguous()
    block_table_c = block_table.contiguous()
    token_indices_c = token_indices.contiguous()
    # Only multi-tile compaction needs separate padding initialization.
    if out is None:
        out = (
            torch.full_like(token_indices_c, -1)
            if return_valid_counts and not single_tile
            else torch.empty_like(token_indices_c)
        )
    else:
        assert out.dtype == token_indices_c.dtype
        assert out.device == token_indices_c.device
        assert out.shape == token_indices_c.shape
        assert out.is_contiguous()
        if return_valid_counts and not single_tile:
            out.fill_(-1)

    valid_counts: torch.Tensor | None = None
    if return_valid_counts:
        if valid_counts_out is None:
            # Zero-init only matters for the atomic accumulation path.
            alloc = torch.empty if single_tile else torch.zeros
            valid_counts = alloc(
                num_tokens, dtype=torch.int32, device=token_indices.device
            )
        else:
            assert valid_counts_out.dtype == torch.int32
            assert valid_counts_out.device == token_indices.device
            assert valid_counts_out.shape == (num_tokens,)
            assert valid_counts_out.is_contiguous()
            valid_counts = valid_counts_out
            if not single_tile:
                valid_counts.zero_()

    # Prepare prefill pointers
    if HAS_PREFILL_WORKSPACE:
        assert prefill_workspace_request_ids is not None  # for mypy
        assert prefill_workspace_starts is not None  # for mypy
        assert prefill_workspace_request_ids.is_contiguous()
        assert prefill_workspace_starts.is_contiguous()

    _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL(
        req_id_c,
        block_table_c,
        token_indices_c,
        out,
        valid_counts,
        prefill_workspace_request_ids,
        prefill_workspace_starts,
        # shapes / constexprs
        workspace_rank_stride=prefill_workspace_rank_stride or 0,
        max_num_blocks_per_req=max_num_blocks_per_req,
        BLOCK_SIZE=BLOCK_SIZE,
        BLOCK_STRIDE_ROWS=(
            BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE
        ),
        BLOCK_N=BLOCK_N,
        NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
        HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE,
        COUNT_VALID=return_valid_counts,
        COMPACT_TO_FRONT=return_valid_counts,
        DCP_SIZE=dcp_size,
        DCP_RANK=dcp_rank,
        DCP_INTERLEAVE=cp_kv_cache_interleave_size,
    )

    if return_valid_counts:
        assert valid_counts is not None
        return out, valid_counts
    return out

triton_filter_and_convert_dcp_index(req_id, block_table, token_indices, dcp_size, dcp_rank, cp_kv_cache_interleave_size=1, BLOCK_SIZE=64, BLOCK_STRIDE_ROWS=None, NUM_TOPK_TOKENS=2048, BLOCK_N=128, return_valid_counts=False, compact_valid_to_front=True)

Filter global per-request indices to this DCP rank's local slots.

With compact_valid_to_front (default), the conversion kernel scatters this rank's owned slots to a contiguous prefix [0, valid_count) and leaves the rest -1. DCP filtering marks non-owned slots -1 and so creates interior gaps; the trtllm-gen sparse kernel reads the first valid_count entries of each row, so they must be a contiguous prefix. Compaction is fused into the kernel (atomic slot allocator) rather than a separate sort/gather pass. Prefix order is unspecified (only the set matters).

Source code in vllm/v1/attention/backends/mla/sparse_utils.py

def triton_filter_and_convert_dcp_index(
    req_id: torch.Tensor,
    block_table: torch.Tensor,
    token_indices: torch.Tensor,
    dcp_size: int,
    dcp_rank: int,
    cp_kv_cache_interleave_size: int = 1,
    BLOCK_SIZE: int = 64,
    BLOCK_STRIDE_ROWS: int | None = None,
    NUM_TOPK_TOKENS: int = 2048,
    BLOCK_N: int = 128,
    return_valid_counts: bool = False,
    compact_valid_to_front: bool = True,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Filter global per-request indices to this DCP rank's local slots.

    With ``compact_valid_to_front`` (default), the conversion kernel scatters
    this rank's owned slots to a contiguous prefix ``[0, valid_count)`` and
    leaves the rest ``-1``. DCP filtering marks non-owned slots ``-1`` and so
    creates interior gaps; the trtllm-gen sparse kernel reads the first
    ``valid_count`` entries of each row, so they must be a contiguous prefix.
    Compaction is fused into the kernel (atomic slot allocator) rather than a
    separate sort/gather pass. Prefix order is unspecified (only the set matters).
    """
    assert dcp_size >= 1
    assert 0 <= dcp_rank < dcp_size
    # Interleave groups must align to KV blocks (globally enforced by
    # VllmConfig: block_size % cp_kv_cache_interleave_size == 0); assert the
    # local invariant so local_idx // BLOCK_SIZE never straddles a group.
    assert BLOCK_SIZE % cp_kv_cache_interleave_size == 0, (
        f"BLOCK_SIZE ({BLOCK_SIZE}) must be divisible by "
        f"cp_kv_cache_interleave_size ({cp_kv_cache_interleave_size})."
    )
    assert req_id.dtype == torch.int32
    assert block_table.dtype == torch.int32
    assert token_indices.dtype == torch.int32
    assert token_indices.shape[1] == NUM_TOPK_TOKENS
    assert NUM_TOPK_TOKENS % BLOCK_N == 0

    if dcp_size == 1:
        return triton_convert_req_index_to_global_index(
            req_id,
            block_table,
            token_indices,
            BLOCK_SIZE=BLOCK_SIZE,
            BLOCK_STRIDE_ROWS=BLOCK_STRIDE_ROWS,
            NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
            BLOCK_N=BLOCK_N,
            return_valid_counts=return_valid_counts,
        )

    num_tokens = req_id.shape[0]
    max_num_blocks_per_req = block_table.shape[1]

    req_id_c = req_id.contiguous()
    block_table_c = block_table.contiguous()
    token_indices_c = token_indices.contiguous()

    # The compaction uses the valid-count buffer as a slot allocator, so it
    # requires counting. Only the multi-tile path needs pre-filled padding.
    count_valid = return_valid_counts or compact_valid_to_front

    # The compaction builds on the counting, so it shares the tiling.
    single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, count_valid)

    if compact_valid_to_front and not single_tile:
        out = torch.full_like(token_indices_c, -1)
    else:
        out = torch.empty_like(token_indices_c)

    valid_counts: torch.Tensor | None = None
    if count_valid:
        # Zero-init only matters for the atomic accumulation path.
        alloc = torch.empty if single_tile else torch.zeros
        valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)

    _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL(
        req_id_c,
        block_table_c,
        token_indices_c,
        out,
        valid_counts,
        # No prefill workspace on the DCP decode path.
        None,
        None,
        workspace_rank_stride=0,
        max_num_blocks_per_req=max_num_blocks_per_req,
        BLOCK_SIZE=BLOCK_SIZE,
        BLOCK_STRIDE_ROWS=(
            BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE
        ),
        BLOCK_N=BLOCK_N,
        NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
        HAS_PREFILL_WORKSPACE=False,
        COUNT_VALID=count_valid,
        COMPACT_TO_FRONT=compact_valid_to_front,
        DCP_SIZE=dcp_size,
        DCP_RANK=dcp_rank,
        DCP_INTERLEAVE=cp_kv_cache_interleave_size,
    )

    if return_valid_counts:
        assert valid_counts is not None
        return out, valid_counts
    return out