vllm.model_executor.layers.sparse_attn_indexer

Custom Sparse Attention Indexer layers.

Classes:

• PrefillRowShard –

  One TP rank's contiguous slice of the replicated prefill rows.

• SparseAttnIndexer –

  Sparse Attention Indexer Custom Op Layer. This layer is extracted as a

Functions:

• as_fp32_logits –

  Upcast SM100's BF16 MQA logits for the FP32 top-k kernels.

• dcp_gather_kv_rows –

  All-gather this rank's KV shard and return it in global token order.

• kv_cache_as_quant_view –

  4D [num_blocks, block_size, 1, head_width] view expected by

PrefillRowShard

One TP rank's contiguous slice of the replicated prefill rows.

Rows are independent, so each rank scores only its slice and the group exchanges final int32 top-k indices instead of the logits.

Methods:

• exchange_topk –

  Gather per-rank final indices back into global row order.

• narrow –

  Row window and narrowed causal bounds of chunk for this rank.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

class PrefillRowShard:
    """One TP rank's contiguous slice of the replicated prefill rows.

    Rows are independent, so each rank scores only its slice and the group
    exchanges final int32 top-k indices instead of the logits."""

    def __init__(self, num_decode_tokens: int, sizes: list[int], rank: int):
        self.num_decode_tokens = num_decode_tokens
        self.sizes = sizes
        self.start = num_decode_tokens + sum(sizes[:rank])
        self.stop = self.start + sizes[rank]

    @classmethod
    def from_metadata(
        cls, prefill_metadata, num_decode_tokens: int
    ) -> "PrefillRowShard | None":
        sizes = getattr(prefill_metadata, "row_shard_sizes", None)
        if sizes is None:
            return None
        return cls(num_decode_tokens, sizes, get_tensor_model_parallel_rank())

    def narrow(self, chunk) -> "tuple[int, int, torch.Tensor, torch.Tensor] | None":
        """Row window and narrowed causal bounds of ``chunk`` for this rank."""
        start = max(chunk.token_start, self.start)
        end = min(chunk.token_end, self.stop)
        if start >= end:
            return None
        lo, hi = start - chunk.token_start, end - chunk.token_start
        return start, end, chunk.cu_seqlen_ks[lo:hi], chunk.cu_seqlen_ke[lo:hi]

    def exchange_topk(self, topk_buffer: torch.Tensor, width: int) -> None:
        """Gather per-rank final indices back into global row order."""
        # The local input must outlive the async PyNCCL gather's enqueued copy.
        local = topk_buffer[self.start : self.stop, :width].contiguous()
        end = self.num_decode_tokens + sum(self.sizes)
        topk_buffer[self.num_decode_tokens : end, :width] = get_tp_group().all_gatherv(
            local, dim=0, sizes=self.sizes
        )

exchange_topk(topk_buffer, width)

Gather per-rank final indices back into global row order.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

def exchange_topk(self, topk_buffer: torch.Tensor, width: int) -> None:
    """Gather per-rank final indices back into global row order."""
    # The local input must outlive the async PyNCCL gather's enqueued copy.
    local = topk_buffer[self.start : self.stop, :width].contiguous()
    end = self.num_decode_tokens + sum(self.sizes)
    topk_buffer[self.num_decode_tokens : end, :width] = get_tp_group().all_gatherv(
        local, dim=0, sizes=self.sizes
    )

narrow(chunk)

Row window and narrowed causal bounds of chunk for this rank.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

def narrow(self, chunk) -> "tuple[int, int, torch.Tensor, torch.Tensor] | None":
    """Row window and narrowed causal bounds of ``chunk`` for this rank."""
    start = max(chunk.token_start, self.start)
    end = min(chunk.token_end, self.stop)
    if start >= end:
        return None
    lo, hi = start - chunk.token_start, end - chunk.token_start
    return start, end, chunk.cu_seqlen_ks[lo:hi], chunk.cu_seqlen_ke[lo:hi]

SparseAttnIndexer

Bases: CustomOp

Sparse Attention Indexer Custom Op Layer. This layer is extracted as a separate custom op since it involves heavy custom kernels like mqa_logits, paged_mqa_logits and top_k_per_row, etc. Those kernels maybe requires specific memory layout or implementation for different hardware backends to achieve optimal performance.

For now, the default native path will use CUDA backend path. Other platform may requires add the corresponding Custom Op name sparse_attn_indexer to custom_ops in CompilationConfig to enable the platform specific path.

Methods:

• forward_cpu –

  CPU sparse attention indexer: cache write stays eager Python glue

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

@CustomOp.register("sparse_attn_indexer")
class SparseAttnIndexer(CustomOp):
    """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a
    separate custom op since it involves heavy custom kernels like `mqa_logits`,
    `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires
    specific memory layout or implementation for different hardware backends to
    achieve optimal performance.

    For now, the default native path will use CUDA backend path. Other platform
    may requires add the corresponding Custom Op name `sparse_attn_indexer` to
    `custom_ops` in `CompilationConfig` to enable the platform specific path.
    """

    def __init__(
        self,
        k_cache,
        quant_block_size: int,
        scale_fmt: str,
        topk_tokens: int,
        head_dim: int,
        max_model_len: int,
        max_total_seq_len: int,
        topk_indices_buffer: torch.Tensor,
        skip_k_cache_insert: bool = False,
        use_fp4_cache: bool = False,
        compress_ratio: int = 1,
        candidate_blocks: torch.Tensor | None = None,
        candidate_block_size: int = 0,
        candidate_write: bool = False,
    ):
        super().__init__()
        self.k_cache = k_cache
        self.quant_block_size = quant_block_size
        self.scale_fmt = scale_fmt
        self.topk_tokens = topk_tokens
        self.head_dim = head_dim
        self.max_model_len = max_model_len
        self.max_total_seq_len = max_total_seq_len
        self.topk_indices_buffer = topk_indices_buffer
        self.skip_k_cache_insert = skip_k_cache_insert
        self.use_fp4_cache = use_fp4_cache
        self.compress_ratio = compress_ratio
        # v4.1 two-level selection: the candidate source indexer writes the
        # top candidate blocks here; later indexers mask their scores with it.
        self.candidate_blocks = candidate_blocks
        self.candidate_block_size = candidate_block_size
        self.candidate_write = candidate_write
        self.dense_mha_metadata_layer_name = ""
        # DCP scalars are constant for the run; resolve them here (config is set
        # during model construction) and pass them into the custom op, rather
        # than threading them through per-step metadata.
        vllm_config = get_current_vllm_config()
        parallel_config = vllm_config.parallel_config
        self.topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend
        self._parallel_config = parallel_config
        self.dcp_world_size = parallel_config.decode_context_parallel_size
        self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
        self.use_pcp = parallel_config.prefill_context_parallel_size > 1
        self.pcp_shard_decode_requests = parallel_config.pcp_shard_decode_requests
        if current_platform.is_cuda() and not has_deep_gemm():
            raise RuntimeError(
                "Sparse Attention Indexer CUDA op requires DeepGEMM support in "
                "the current vLLM environment."
            )

        if vllm_config.kernel_config.enable_jit_warmup:
            from vllm.v1.attention.ops.common import (
                _PACK_SEQ_TRITON_KERNEL,
                _UNPACK_SEQ_TRITON_KERNEL,
            )

            pack_dtype = torch.uint8 if use_fp4_cache else current_platform.fp8_dtype()
            _PACK_SEQ_TRITON_KERNEL.register_warmup(
                dtype=pack_dtype,
                pad_value=0 if use_fp4_cache else -float("inf"),
            )
            _UNPACK_SEQ_TRITON_KERNEL.register_warmup()

            if self.dcp_world_size > 1 and current_platform.is_cuda() and has_cutedsl():
                from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import (  # noqa: E501
                    _PACK_DCP_TOPK_CANDIDATES_KERNEL,
                    _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL,
                )

                _PACK_DCP_TOPK_CANDIDATES_KERNEL.register_warmup()
                _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL.register_warmup()

    def forward_native(
        self,
        hidden_states: torch.Tensor,
        q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        k: torch.Tensor | None,
        weights: torch.Tensor,
    ):
        if current_platform.is_cuda() or current_platform.is_xpu():
            return self.forward_cuda(hidden_states, q_quant, k, weights)
        elif current_platform.is_rocm():
            if self.use_pcp and self.dcp_world_size > 1:
                raise NotImplementedError(
                    "The ROCm sparse-indexer path does not support PCP+DCP."
                )
            return self.forward_hip(hidden_states, q_quant, k, weights)
        elif current_platform.is_cpu():
            return self.forward_cpu(hidden_states, q_quant, k, weights)
        else:
            raise NotImplementedError(
                "SparseAttnIndexer native forward is only implemented for "
                "CUDA, ROCm, XPU and CPU platforms."
            )

    def forward_cuda(
        self,
        hidden_states: torch.Tensor,
        q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        k: torch.Tensor | None,
        weights: torch.Tensor,
    ):
        # FP8 path: single tensor (per-token scale is folded into `weights`).
        # FP4 path: (values, scales) tuple with scales required by the kernel.
        if isinstance(q_quant, tuple):
            q_values, q_scale = q_quant
        else:
            q_values, q_scale = q_quant, None
        return torch.ops.vllm.sparse_attn_indexer(
            hidden_states,
            _encode_layer_name(self.k_cache.prefix),
            self.k_cache.kv_cache,
            q_values,
            q_scale,
            k,
            weights,
            self.quant_block_size,
            self.scale_fmt,
            self.topk_tokens,
            self.head_dim,
            self.max_model_len,
            self.max_total_seq_len,
            self.topk_indices_buffer,
            self.skip_k_cache_insert,
            self.use_pcp,
            self.pcp_shard_decode_requests,
            _encode_layer_name(self.dense_mha_metadata_layer_name),
            self.use_fp4_cache,
            self.dcp_rank,
            self.dcp_world_size,
            self._parallel_config.cp_kv_cache_interleave_size,
            False,
            candidate_blocks=self.candidate_blocks,
            candidate_block_size=self.candidate_block_size,
            candidate_write=self.candidate_write,
            topk_backend=self.topk_backend,
        )

    def forward_xpu(
        self,
        hidden_states: torch.Tensor,
        q_fp8: torch.Tensor,
        k: torch.Tensor | None,
        weights: torch.Tensor,
    ):
        return self.forward_cuda(hidden_states, q_fp8, k, weights)

    def forward_hip(
        self,
        hidden_states: torch.Tensor,
        q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        k: torch.Tensor | None,
        weights: torch.Tensor,
    ):
        assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"
        assert isinstance(q_quant, torch.Tensor), (
            "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
        )
        from vllm.platforms.rocm import on_gfx950

        if (
            rocm_aiter_ops.is_enabled()
            or rocm_aiter_ops.is_rdna_aiter_enabled()
            # The so-called AITER sparse indexer op has a native gfx950 path:
            # its cache insert, MQA logits, and top-k fallbacks are implemented
            # by local Triton/C++ kernels and do not require the aiter package.
            or on_gfx950()
        ):
            return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
                hidden_states,
                _encode_layer_name(self.k_cache.prefix),
                self.k_cache.kv_cache,
                q_quant,
                k,
                weights,
                self.quant_block_size,
                self.scale_fmt,
                self.topk_tokens,
                self.head_dim,
                self.max_model_len,
                self.max_total_seq_len,
                self.topk_indices_buffer,
                skip_k_cache_insert=self.skip_k_cache_insert,
                compress_ratio=self.compress_ratio,
                candidate_blocks=self.candidate_blocks,
                candidate_block_size=self.candidate_block_size,
                candidate_write=self.candidate_write,
            )
        raise RuntimeError(
            "Sparse attention indexer ROCm path requires AITER or a supported "
            "native architecture (gfx950)."
        )

    def forward_cpu(
        self,
        hidden_states: torch.Tensor,
        q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        k: torch.Tensor | None,
        weights: torch.Tensor,
    ):
        """CPU sparse attention indexer: cache write stays eager Python glue
        (own K-cache layout, not shared with the main attention cache
        write). PREFILL and DECODE both call the ported
        ``fp8_paged_mqa_logits_cpu``/``topk_transform_512_cpu`` kernels,
        which read the paged K-cache directly via ``page_table`` -- no
        eager gather step, no per-request Python loop.

        ``prefill_metadata.chunks`` always has exactly one entry here:
        ``DeepseekV4CPUIndexerMetadataBuilder`` overrides the base chunk
        split to always return the whole step's prefill batch as one
        chunk, since the base chunking only bounds CUDA/XPU's dense M*N
        logits tensor and flat K-gather workspace, neither of which this
        paged kernel allocates.
        """
        assert not self.use_fp4_cache, (
            "CPU sparse indexer doesn't support fp4 cache yet"
        )
        assert isinstance(q_quant, torch.Tensor), (
            "CPU sparse_attn_indexer expects a single FP8 q_quant tensor"
        )
        assert self.dcp_world_size <= 1 and not self.use_pcp, (
            "CPU sparse indexer doesn't support decode/prefill context parallelism yet."
        )

        forward_context = get_forward_context()
        attn_metadata = forward_context.attn_metadata
        attn_metadata_narrowed: DeepseekV32IndexerMetadata | None = None
        if isinstance(attn_metadata, dict):
            metadata = attn_metadata[self.k_cache.prefix]
            assert isinstance(metadata, DeepseekV32IndexerMetadata)
            attn_metadata_narrowed = metadata
        if attn_metadata_narrowed is None:
            # Profiling/dummy run: no real metadata to act on.
            return self.topk_indices_buffer

        kv_cache = self.k_cache.kv_cache
        topk_tokens = self.topk_tokens
        topk_indices_buffer = self.topk_indices_buffer
        slot_mapping = attn_metadata_narrowed.slot_mapping
        has_decode = attn_metadata_narrowed.num_decodes > 0
        has_prefill = attn_metadata_narrowed.num_prefills > 0
        num_decode_tokens = attn_metadata_narrowed.num_decode_tokens

        num_tokens = slot_mapping.shape[0]
        if k is not None:
            k = k[:num_tokens]

        if not self.skip_k_cache_insert:
            # Only reachable via DeepseekV32Attention with
            # prefill_context_parallel_size > 1 on CPU -- set_k_cpu/set_s_cpu
            # (the kernels this used to call) have been removed as unused/
            # untested (csrc/cpu/sgl-kernels/store_cache.cpp).
            raise NotImplementedError(
                "SparseAttnIndexer.forward_cpu: skip_k_cache_insert=False "
                "(prefill context parallel on CPU) is not supported."
            )

        topk_indices_buffer[: hidden_states.shape[0]] = -1

        if has_prefill:
            assert topk_tokens == 512, (
                "topk_transform_512_cpu only supports index_topk == 512."
            )
            prefill_metadata = attn_metadata_narrowed.prefill
            assert prefill_metadata is not None
            assert len(prefill_metadata.chunks) == 1, (
                "forward_cpu expects the prefill metadata builder to always "
                "produce a single chunk -- see "
                "DeepseekV4CPUIndexerMetadataBuilder._split_indexer_prefill_chunks."
            )
            chunk = prefill_metadata.chunks[0]
            # kv_cache is a per-layer view into vLLM's shared multi-layer
            # cache allocation, so its block stride generally exceeds
            # block_size * page_width; fp8_paged_mqa_logits_cpu reads
            # kv_view.stride(0) explicitly, so no copy is needed here.
            kv_view = kv_cache.view(kv_cache.shape[0], -1)
            block_size = kv_cache.shape[1]
            q_slice = q_quant[chunk.token_start : chunk.token_end]
            topk_indices = topk_indices_buffer[
                chunk.token_start : chunk.token_end, :topk_tokens
            ]
            if chunk.local_total_seq_lens == 0:
                topk_indices.fill_(-1)
            else:
                assert chunk.local_cu_seq_lens is not None
                # Each token's own (per-request, DCP-local) causal length.
                local_seq_lens = chunk.cu_seqlen_ke - chunk.cu_seqlen_ks
                # Recover each token's owning request from its row-start
                # tag. Ties (from zero-length requests) are harmless: those
                # tokens have local length 0 and never dereference
                # page_table.
                req_idx = (
                    torch.searchsorted(
                        chunk.local_cu_seq_lens, chunk.cu_seqlen_ks, right=True
                    )
                    - 1
                )
                page_table = chunk.block_table[req_idx]
                # The true max over this chunk's own rows, NOT
                # chunk.max_local_total_seq_lens (that field sums every
                # request's length in the chunk, bounding the old
                # flat-gather buffer this paged path no longer allocates).
                max_seq_len = int(local_seq_lens.max().item())

                logits = ops.fp8_paged_mqa_logits_cpu(
                    q_slice,
                    kv_view,
                    weights[chunk.token_start : chunk.token_end],
                    local_seq_lens,
                    page_table,
                    block_size,
                    max_seq_len,
                )
                out_page_scratch = torch.empty(
                    (q_slice.shape[0], topk_tokens),
                    dtype=torch.int32,
                    device=kv_cache.device,
                )
                ops.topk_transform_512_cpu(
                    logits,
                    local_seq_lens,
                    page_table,
                    out_page_scratch,
                    block_size,
                    topk_indices,
                )

        if has_decode:
            decode_metadata = attn_metadata_narrowed.decode
            assert decode_metadata is not None
            assert not decode_metadata.requires_padding, (
                "CPU sparse indexer decode path does not support speculative "
                "decoding (native MTP) yet."
            )
            batch_size = decode_metadata.decode_lens.shape[0]
            if batch_size > 0:
                # No native MTP on CPU (asserted above) => exactly one
                # query token per decode request, so the flat slice below
                # is already the batch-major layout
                # fp8_paged_mqa_logits_cpu wants.
                assert num_decode_tokens == batch_size, (
                    "CPU sparse indexer decode path expects exactly one query "
                    "token per decode request."
                )
                assert topk_tokens == 512, (
                    "topk_transform_512_cpu only supports index_topk == 512."
                )
                seq_lens = decode_metadata.seq_lens
                seq_lens = (
                    seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens
                )
                block_table = decode_metadata.block_table[:batch_size]
                block_size = kv_cache.shape[1]
                kv_view = kv_cache.view(kv_cache.shape[0], -1)

                logits = ops.fp8_paged_mqa_logits_cpu(
                    q_quant[:num_decode_tokens],
                    kv_view,
                    weights[:num_decode_tokens],
                    seq_lens,
                    block_table,
                    block_size,
                    attn_metadata_narrowed.max_seq_len,
                )

                # out_page_indices is a required kernel output but unused:
                # the indexer's topk output must stay local/compressed-
                # context positions (resolved later by
                # DeepseekV4CPUAttention.forward_mqa via
                # map_local_to_global_slots_cpu).
                out_page_scratch = torch.empty(
                    (batch_size, topk_tokens),
                    dtype=torch.int32,
                    device=kv_cache.device,
                )
                ops.topk_transform_512_cpu(
                    logits,
                    seq_lens,
                    block_table,
                    out_page_scratch,
                    block_size,
                    topk_indices_buffer[:num_decode_tokens, :topk_tokens],
                )

        return topk_indices_buffer

forward_cpu(hidden_states, q_quant, k, weights)

CPU sparse attention indexer: cache write stays eager Python glue (own K-cache layout, not shared with the main attention cache write). PREFILL and DECODE both call the ported fp8_paged_mqa_logits_cpu/topk_transform_512_cpu kernels, which read the paged K-cache directly via page_table -- no eager gather step, no per-request Python loop.

prefill_metadata.chunks always has exactly one entry here: DeepseekV4CPUIndexerMetadataBuilder overrides the base chunk split to always return the whole step's prefill batch as one chunk, since the base chunking only bounds CUDA/XPU's dense M*N logits tensor and flat K-gather workspace, neither of which this paged kernel allocates.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

def forward_cpu(
    self,
    hidden_states: torch.Tensor,
    q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
    k: torch.Tensor | None,
    weights: torch.Tensor,
):
    """CPU sparse attention indexer: cache write stays eager Python glue
    (own K-cache layout, not shared with the main attention cache
    write). PREFILL and DECODE both call the ported
    ``fp8_paged_mqa_logits_cpu``/``topk_transform_512_cpu`` kernels,
    which read the paged K-cache directly via ``page_table`` -- no
    eager gather step, no per-request Python loop.

    ``prefill_metadata.chunks`` always has exactly one entry here:
    ``DeepseekV4CPUIndexerMetadataBuilder`` overrides the base chunk
    split to always return the whole step's prefill batch as one
    chunk, since the base chunking only bounds CUDA/XPU's dense M*N
    logits tensor and flat K-gather workspace, neither of which this
    paged kernel allocates.
    """
    assert not self.use_fp4_cache, (
        "CPU sparse indexer doesn't support fp4 cache yet"
    )
    assert isinstance(q_quant, torch.Tensor), (
        "CPU sparse_attn_indexer expects a single FP8 q_quant tensor"
    )
    assert self.dcp_world_size <= 1 and not self.use_pcp, (
        "CPU sparse indexer doesn't support decode/prefill context parallelism yet."
    )

    forward_context = get_forward_context()
    attn_metadata = forward_context.attn_metadata
    attn_metadata_narrowed: DeepseekV32IndexerMetadata | None = None
    if isinstance(attn_metadata, dict):
        metadata = attn_metadata[self.k_cache.prefix]
        assert isinstance(metadata, DeepseekV32IndexerMetadata)
        attn_metadata_narrowed = metadata
    if attn_metadata_narrowed is None:
        # Profiling/dummy run: no real metadata to act on.
        return self.topk_indices_buffer

    kv_cache = self.k_cache.kv_cache
    topk_tokens = self.topk_tokens
    topk_indices_buffer = self.topk_indices_buffer
    slot_mapping = attn_metadata_narrowed.slot_mapping
    has_decode = attn_metadata_narrowed.num_decodes > 0
    has_prefill = attn_metadata_narrowed.num_prefills > 0
    num_decode_tokens = attn_metadata_narrowed.num_decode_tokens

    num_tokens = slot_mapping.shape[0]
    if k is not None:
        k = k[:num_tokens]

    if not self.skip_k_cache_insert:
        # Only reachable via DeepseekV32Attention with
        # prefill_context_parallel_size > 1 on CPU -- set_k_cpu/set_s_cpu
        # (the kernels this used to call) have been removed as unused/
        # untested (csrc/cpu/sgl-kernels/store_cache.cpp).
        raise NotImplementedError(
            "SparseAttnIndexer.forward_cpu: skip_k_cache_insert=False "
            "(prefill context parallel on CPU) is not supported."
        )

    topk_indices_buffer[: hidden_states.shape[0]] = -1

    if has_prefill:
        assert topk_tokens == 512, (
            "topk_transform_512_cpu only supports index_topk == 512."
        )
        prefill_metadata = attn_metadata_narrowed.prefill
        assert prefill_metadata is not None
        assert len(prefill_metadata.chunks) == 1, (
            "forward_cpu expects the prefill metadata builder to always "
            "produce a single chunk -- see "
            "DeepseekV4CPUIndexerMetadataBuilder._split_indexer_prefill_chunks."
        )
        chunk = prefill_metadata.chunks[0]
        # kv_cache is a per-layer view into vLLM's shared multi-layer
        # cache allocation, so its block stride generally exceeds
        # block_size * page_width; fp8_paged_mqa_logits_cpu reads
        # kv_view.stride(0) explicitly, so no copy is needed here.
        kv_view = kv_cache.view(kv_cache.shape[0], -1)
        block_size = kv_cache.shape[1]
        q_slice = q_quant[chunk.token_start : chunk.token_end]
        topk_indices = topk_indices_buffer[
            chunk.token_start : chunk.token_end, :topk_tokens
        ]
        if chunk.local_total_seq_lens == 0:
            topk_indices.fill_(-1)
        else:
            assert chunk.local_cu_seq_lens is not None
            # Each token's own (per-request, DCP-local) causal length.
            local_seq_lens = chunk.cu_seqlen_ke - chunk.cu_seqlen_ks
            # Recover each token's owning request from its row-start
            # tag. Ties (from zero-length requests) are harmless: those
            # tokens have local length 0 and never dereference
            # page_table.
            req_idx = (
                torch.searchsorted(
                    chunk.local_cu_seq_lens, chunk.cu_seqlen_ks, right=True
                )
                - 1
            )
            page_table = chunk.block_table[req_idx]
            # The true max over this chunk's own rows, NOT
            # chunk.max_local_total_seq_lens (that field sums every
            # request's length in the chunk, bounding the old
            # flat-gather buffer this paged path no longer allocates).
            max_seq_len = int(local_seq_lens.max().item())

            logits = ops.fp8_paged_mqa_logits_cpu(
                q_slice,
                kv_view,
                weights[chunk.token_start : chunk.token_end],
                local_seq_lens,
                page_table,
                block_size,
                max_seq_len,
            )
            out_page_scratch = torch.empty(
                (q_slice.shape[0], topk_tokens),
                dtype=torch.int32,
                device=kv_cache.device,
            )
            ops.topk_transform_512_cpu(
                logits,
                local_seq_lens,
                page_table,
                out_page_scratch,
                block_size,
                topk_indices,
            )

    if has_decode:
        decode_metadata = attn_metadata_narrowed.decode
        assert decode_metadata is not None
        assert not decode_metadata.requires_padding, (
            "CPU sparse indexer decode path does not support speculative "
            "decoding (native MTP) yet."
        )
        batch_size = decode_metadata.decode_lens.shape[0]
        if batch_size > 0:
            # No native MTP on CPU (asserted above) => exactly one
            # query token per decode request, so the flat slice below
            # is already the batch-major layout
            # fp8_paged_mqa_logits_cpu wants.
            assert num_decode_tokens == batch_size, (
                "CPU sparse indexer decode path expects exactly one query "
                "token per decode request."
            )
            assert topk_tokens == 512, (
                "topk_transform_512_cpu only supports index_topk == 512."
            )
            seq_lens = decode_metadata.seq_lens
            seq_lens = (
                seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens
            )
            block_table = decode_metadata.block_table[:batch_size]
            block_size = kv_cache.shape[1]
            kv_view = kv_cache.view(kv_cache.shape[0], -1)

            logits = ops.fp8_paged_mqa_logits_cpu(
                q_quant[:num_decode_tokens],
                kv_view,
                weights[:num_decode_tokens],
                seq_lens,
                block_table,
                block_size,
                attn_metadata_narrowed.max_seq_len,
            )

            # out_page_indices is a required kernel output but unused:
            # the indexer's topk output must stay local/compressed-
            # context positions (resolved later by
            # DeepseekV4CPUAttention.forward_mqa via
            # map_local_to_global_slots_cpu).
            out_page_scratch = torch.empty(
                (batch_size, topk_tokens),
                dtype=torch.int32,
                device=kv_cache.device,
            )
            ops.topk_transform_512_cpu(
                logits,
                seq_lens,
                block_table,
                out_page_scratch,
                block_size,
                topk_indices_buffer[:num_decode_tokens, :topk_tokens],
            )

    return topk_indices_buffer

_decode_chunk_ranges(rows, width, budget)

Partition independent Q rows under a BF16 payload budget.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

def _decode_chunk_ranges(rows: int, width: int, budget: int) -> list[tuple[int, int]]:
    """Partition independent Q rows under a BF16 payload budget."""
    if budget < 0:
        raise ValueError("Decode logits budget must be nonnegative")
    if not rows or not budget or rows * width * 2 <= budget:
        return []
    limit = budget // (2 * width) // 16 * 16
    if limit < 16:
        raise ValueError("Decode logits budget must fit at least 16 BF16 rows")
    units = (rows + 15) // 16
    count = (units + limit // 16 - 1) // (limit // 16)
    base, extra = divmod(units, count)
    ranges = []
    start = 0
    # Allocate the largest scratch first so later graph nodes can reuse it.
    for i in range(count):
        end = min(rows, start + 16 * (base + (i < extra)))
        ranges.append((start, end))
        start = end
    return ranges

_gather_workspace_shapes(total_seq_lens, head_dim, fp8_dtype, use_fp4_cache)

Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32 scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 + (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

def _gather_workspace_shapes(
    total_seq_lens: int,
    head_dim: int,
    fp8_dtype: torch.dtype,
    use_fp4_cache: bool,
) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]:
    """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for
    the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32
    scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 +
    (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales."""
    if use_fp4_cache:
        return (
            ((total_seq_lens, head_dim // 2), torch.uint8),
            ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8),
        )
    return (
        ((total_seq_lens, head_dim), fp8_dtype),
        ((total_seq_lens, 4), torch.uint8),
    )

_merge_dcp_topk_global(logits, topk_indices, topk_tokens, dcp_rank, dcp_world_size, cp_interleave, row_starts=None)

Merge each DCP rank's local top-K into the global top-K.

topk_indices are this rank's local top-K positions into its 1/N KV shard. A token in the global top-K must also be in its owning rank's local top-K (at most topk_tokens - 1 tokens rank globally above it, hence at most that many on its own rank), so exchanging only the per-rank local candidates is exact -- equivalent to all-gathering the full logit matrix, but it ships dcp_world_size * topk_tokens candidates instead of the whole score row. Overwrites topk_indices with global token ids (-1 for padding); the attention backend localizes them back to physical slots per rank.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

def _merge_dcp_topk_global(
    logits: torch.Tensor,
    topk_indices: torch.Tensor,
    topk_tokens: int,
    dcp_rank: int,
    dcp_world_size: int,
    cp_interleave: int,
    row_starts: torch.Tensor | None = None,
) -> None:
    """Merge each DCP rank's local top-K into the global top-K.

    ``topk_indices`` are this rank's local top-K positions into its 1/N KV
    shard. A token in the global top-K must also be in its owning rank's local
    top-K (at most ``topk_tokens - 1`` tokens rank globally above it, hence at
    most that many on its own rank), so exchanging only the per-rank local
    candidates is exact -- equivalent to all-gathering the full logit matrix,
    but it ships ``dcp_world_size * topk_tokens`` candidates instead of the whole
    score row. Overwrites ``topk_indices`` with global token ids (``-1`` for
    padding); the attention backend localizes them back to physical slots per
    rank.
    """
    if dcp_world_size <= 1:
        return

    # CuteDSL-only path (no PyTorch fallback): Triton-pack each rank's
    # (score, global_id) candidates on-device, all-gather, then the CuteDSL
    # stable-topk selector.
    _assert_cutedsl_dcp_merge_supported(logits, topk_indices, topk_tokens)
    from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import (
        pack_dcp_topk_candidates_cutedsl,
        stable_topk_from_gathered_candidates_cutedsl,
    )

    packed = torch.empty(
        (*topk_indices.shape, 2),
        dtype=torch.float32,
        device=topk_indices.device,
    )
    pack_dcp_topk_candidates_cutedsl(
        logits,
        topk_indices,
        packed,
        dcp_rank,
        dcp_world_size,
        cp_interleave,
        row_starts,
    )
    # Gather along dim 0 and let the selector read the rank-major
    # [world_size, rows, topk, 2] result: a dim-1 gather would reorder it into
    # [rows, world_size * topk, 2] with a full extra copy.
    gathered = get_dcp_group().all_gather(packed, dim=0)
    stable_topk_from_gathered_candidates_cutedsl(
        gathered.view(dcp_world_size, *packed.shape),
        topk_tokens,
        out=topk_indices,
    )

_reserve_row_shard_exchange_transient(topk_indices_buffer, device)

Profiling/dummy runs return before exchange_topk executes; reserve its worst-case transient so the memory profile accounts for it.

The exchange stages this rank's contiguous slice and the all_gatherv wrapper allocates the full gathered tensor; each is bounded by the buffer's row capacity times its (tile-padded) width.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

def _reserve_row_shard_exchange_transient(
    topk_indices_buffer: torch.Tensor, device: torch.device
) -> None:
    """Profiling/dummy runs return before ``exchange_topk`` executes; reserve
    its worst-case transient so the memory profile accounts for it.

    The exchange stages this rank's contiguous slice and the all_gatherv
    wrapper allocates the full gathered tensor; each is bounded by the
    buffer's row capacity times its (tile-padded) width.
    """
    cfg = get_current_vllm_config_or_none()
    if cfg is None:
        return
    parallel = cfg.parallel_config
    if not tp_prefill_row_sharding_supported(
        cfg,
        parallel.decode_context_parallel_size,
        parallel.prefill_context_parallel_size > 1,
        get_tp_group().world_size,
    ):
        return
    rows, width = topk_indices_buffer.shape[:2]
    _ = torch.empty(
        2 * rows * width * topk_indices_buffer.element_size(),
        dtype=torch.uint8,
        device=device,
    )

_varlen_decode_rows(q, q_scale, weights, kv_cache, decode, max_model_len, max_seq_len, topk_backend, topk_tokens, topk_indices_buffer, candidate_blocks, candidate_block_size, candidate_write, start, end)

The varlen (SM100) decode logits, candidate selection and top-k of rows [start, end), with the slice's own DeepGEMM schedule.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

def _varlen_decode_rows(
    q: torch.Tensor,
    q_scale: torch.Tensor | None,
    weights: torch.Tensor,
    kv_cache: torch.Tensor,
    decode: DeepSeekV32IndexerDecodeMetadata,
    max_model_len: int,
    max_seq_len: int,
    topk_backend: str,
    topk_tokens: int,
    topk_indices_buffer: torch.Tensor,
    candidate_blocks: torch.Tensor | None,
    candidate_block_size: int,
    candidate_write: bool,
    start: int,
    end: int,
) -> None:
    """The varlen (SM100) decode logits, candidate selection and top-k of rows
    ``[start, end)``, with the slice's own DeepGEMM schedule."""
    assert decode.indices is not None
    seq_lens = decode.seq_lens[start:end]
    indices = decode.indices[start:end]
    schedule = decode.schedule_metadata
    if end - start < decode.seq_lens.shape[0]:
        schedule = get_paged_mqa_logits_metadata(
            seq_lens, kv_cache.shape[1], schedule.shape[0] - 1, indices=indices
        )
    logits = fp8_fp4_paged_mqa_logits(
        (q[start:end], q_scale[start:end] if q_scale is not None else None),
        kv_cache,
        weights[start:end],
        seq_lens,
        decode.block_table[start:end],
        schedule,
        max_model_len=max_model_len,
        indices=indices,
    )
    row_ends = seq_lens.view(-1)
    selector = get_indexer_topk(topk_backend)
    if not selector.takes_bf16_logits(logits, topk_tokens, end - start):
        logits = as_fp32_logits(logits, row_ends)
    if candidate_blocks is not None:
        candidates = candidate_blocks[start:end]
        if candidate_write:
            _select_candidate_blocks(
                logits,
                None,
                row_ends,
                candidates.shape[1],
                candidate_block_size,
                candidates,
                1,
                score_dtype=(
                    torch.float32 if topk_backend in ("auto", "persistent") else None
                ),
            )
        else:
            _apply_candidate_mask(
                logits, None, row_ends, candidates, candidate_block_size, 1
            )
    selector(
        logits,
        seq_lens,
        1,
        topk_indices_buffer[start:end, :topk_tokens],
        topk_tokens,
        max_seq_len,
    )

as_fp32_logits(logits, row_ends)

Upcast SM100's BF16 MQA logits for the FP32 top-k kernels.

Only columns below row_ends (int32, one per row) are converted, the span the consumers read; DeepGEMM's aligned row stride is kept.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

def as_fp32_logits(logits: torch.Tensor, row_ends: torch.Tensor) -> torch.Tensor:
    """Upcast SM100's BF16 MQA logits for the FP32 top-k kernels.

    Only columns below ``row_ends`` (int32, one per row) are converted, the
    span the consumers read; DeepGEMM's aligned row stride is kept.
    """
    if logits.dtype == torch.float32:
        return logits
    out = torch.empty_strided(
        logits.shape, logits.stride(), dtype=torch.float32, device=logits.device
    )
    num_rows, width = logits.shape
    if num_rows == 0 or width == 0:
        return out
    assert row_ends.numel() == num_rows and row_ends.dtype == torch.int32
    block = 1024
    splits = min(triton.cdiv(width, block), max(1, 1024 // num_rows))
    _upcast_logits_kernel[(num_rows, splits)](
        logits, out, row_ends.contiguous(), logits.stride(0), BLOCK=block
    )
    return out

dcp_gather_kv_rows(local_padded, deinterleave_idx, gathered_buf, out_buf)

All-gather this rank's KV shard and return it in global token order.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

def dcp_gather_kv_rows(
    local_padded: torch.Tensor,
    deinterleave_idx: torch.Tensor,
    gathered_buf: torch.Tensor,
    out_buf: torch.Tensor,
) -> torch.Tensor:
    """All-gather this rank's KV shard and return it in global token order."""
    gathered = gathered_buf[: get_dcp_group().world_size * local_padded.shape[0]]
    dist.all_gather_into_tensor(
        gathered, local_padded.contiguous(), group=get_dcp_group().device_group
    )
    return torch.index_select(
        gathered, 0, deinterleave_idx, out=out_buf[: deinterleave_idx.shape[0]]
    )

kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)

4D [num_blocks, block_size, 1, head_width] view expected by DeepGEMM, from the 3D indexer kv-cache allocation.

Source code in vllm/model_executor/layers/sparse_attn_indexer.py

def kv_cache_as_quant_view(
    kv_cache: torch.Tensor,
    head_dim: int,
    use_fp4_cache: bool,
) -> torch.Tensor:
    """4D ``[num_blocks, block_size, 1, head_width]`` view expected by
    DeepGEMM, from the 3D indexer kv-cache allocation."""
    if use_fp4_cache:
        assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8
        num_blocks, block_size, _ = kv_cache.shape
        page_bytes = int(kv_cache.stride(0))
        fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE
        return torch.as_strided(
            kv_cache,
            size=(num_blocks, block_size, 1, fp4_bytes),
            stride=(page_bytes, fp4_bytes, fp4_bytes, 1),
        )
    return kv_cache.unsqueeze(-2)