vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90

FlashInfer sparse MLA backend for SM90 (Hopper) NoPE models.

Wraps FlashInfer's BatchMLAPagedAttentionWrapper (FA2/FA3 paths), which as of FlashInfer 0.6.18 supports head_dim_kpe=0 (GLM-5.3-Flash NoPE MLA) and FP8 E4M3 KV caches on SM90 with in-kernel dequantization: the FP8 cache is read directly (half the bf16 HBM traffic) and converted to BF16 in shared memory, while queries stay BF16 (no query quantization).

Sparsity rides the same trick the FA-based sparse backend uses: with page_size=1 the per-token top-k slot indices ARE the page table, so each query token becomes one varlen batch row whose kv_indices slice is its top-k row and whose kv_len is its valid count. Causality is already encoded by the indexer's selection, so causal=False.

CUDA-graph handling: plan() copies its inputs to host unconditionally, so it must stay outside graph capture. Each metadata builder owns a wrapper, reserved capture-stable device buffers, and the plan parameters. The wrapper bakes the per-row kv_len into its int schedule at plan() time — run() never reads the device-side buffer. The builder plans (outside capture) from sync-free host upper bounds and clamps each scheduled work item's kv_end on device to the exact valid count, so the kernel never reads past a row's valid prefix (the -1 tail of the converted index buffer) and async scheduling needs no D2H sync. Per-step content (top-k slots) is written into the reserved buffers by kernels inside the captured forward, and captured runs read the refreshed plan buffers on replay.

KV cache format: plain contiguous E4M3 [num_blocks, block_size, 512] (uint8 storage) with a per-tensor k_scale; BF16 caches also work. The per-token x 128-channel-group ckv_scale_arr layout is supported by the kernel but not wired yet (it needs a group-quantizing cache-write op).

Classes:

• FlashInferMLASparseSM90Builder –

  Reuse the common sparse metadata (req ids, topk buffer access).

FlashInferMLASparseSM90Builder

Bases: FlashInferMLASparseMetadataBuilder

Reuse the common sparse metadata (req ids, topk buffer access).

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

class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
    """Reuse the common sparse metadata (req ids, topk buffer access)."""

    metadata_cls = FlashInferMLASparseSM90Metadata

    def __init__(
        self,
        kv_cache_spec: "AttentionSpec",
        layer_names: list[str],
        vllm_config: "VllmConfig",
        device: torch.device,
    ) -> None:
        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
        attention_layer = vllm_config.compilation_config.static_forward_context[
            layer_names[0]
        ]
        self._attention_layer = attention_layer
        impl = attention_layer.impl
        if not isinstance(impl, FlashInferMLASparseSM90Impl):
            raise TypeError(
                "FlashInferMLASparseSM90Builder requires an SM90 FlashInfer "
                f"implementation, got {type(impl).__name__}."
            )
        topk_indices_buffer = impl.topk_indices_buffer
        assert topk_indices_buffer is not None
        hf_config = vllm_config.model_config.hf_text_config
        assert hf_config.index_topk is not None
        # fp8 KV cache is stored as uint8; plan() needs the logical dtype.
        kv_plan_dtype = (
            torch.float8_e4m3fn
            if kv_cache_spec.dtype == torch.uint8
            else kv_cache_spec.dtype
        )
        self.state = _SM90State(
            device,
            impl.num_heads,
            kv_plan_dtype,
            vllm_config.scheduler_config.max_num_batched_tokens,
            topk_indices_buffer.shape[1],
            kv_lora_rank=impl.kv_lora_rank,
            qk_rope_head_dim=impl.qk_rope_head_dim,
            sm_scale=impl.scale,
            index_topk=int(hf_config.index_topk),
            index_kpool=int(getattr(hf_config, "index_kpool", 1) or 1),
        )
        spec_config = vllm_config.speculative_config
        self._adaptive_verification = bool(
            spec_config is not None and spec_config.enable_adaptive_verification
        )

    def _kv_lens_host(
        self, cam: CommonAttentionMetadata
    ) -> tuple[int, torch.Tensor, bool]:
        """Host per-row KV lengths for the plan and whether they are exact.

        A row for the j-th query token of request i attends
        ``seq_lens[i] - q_len[i] + j + 1`` tokens. The indexer's selection
        then bounds the valid count: contexts up to ``index_topk`` select
        everything (valid == context); longer contexts keep the top
        ``index_topk`` pool-expanded tokens plus the trailing incomplete
        pool (valid == ``index_topk + context % index_kpool``). Both match
        the count of non -1 entries the convert kernel produces.

        With ``seq_lens_cpu_upper_bound`` (optimistic under async spec
        decode) returns sync-free upper bounds, which the plan clamps on
        device. Otherwise returns exact lengths via a D2H copy.
        """
        num_reqs = cam.num_reqs
        qsl = cam.query_start_loc_cpu[: num_reqs + 1]
        num_rows = int(qsl[-1])
        if num_rows == 0:
            return 0, torch.zeros(0, dtype=torch.int32), True
        state = self.state
        sl_host = cam.seq_lens_cpu_upper_bound
        if sl_host is None:
            ctx = self._row_contexts(qsl, cam.seq_lens[:num_reqs].cpu())
            topk, kpool = state.index_topk, state.index_kpool
            lens = torch.where(ctx <= topk, ctx, topk + ctx % kpool)
            return num_rows, lens.to(torch.int32), True
        if self._adaptive_verification:
            # The host query split may not match the device one, so bound
            # every row by the longest request.
            ctx = torch.full((num_rows,), int(sl_host[:num_reqs].max()))
        else:
            ctx = self._row_contexts(qsl, sl_host[:num_reqs])
        return num_rows, torch.clamp(ctx, max=state.max_valid), False

    @staticmethod
    def _row_contexts(qsl: torch.Tensor, seq_lens: torch.Tensor) -> torch.Tensor:
        num_reqs = qsl.shape[0] - 1
        q_lens = (qsl[1:] - qsl[:-1]).to(torch.int64)
        first_ctx = seq_lens.to(torch.int64) - q_lens + 1
        req_of_row = torch.repeat_interleave(
            torch.arange(num_reqs, dtype=torch.int64), q_lens
        )
        rows = torch.arange(int(qsl[-1]), dtype=torch.int64)
        return first_ctx[req_of_row] + rows - qsl.to(torch.int64)[req_of_row]

    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> FlashInferMLASparseSM90Metadata:
        metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)
        assert isinstance(metadata, FlashInferMLASparseSM90Metadata)
        # Plan outside any CUDA graph capture; captured runs read the
        # refreshed buffers.
        num_rows, kv_lens, exact = self._kv_lens_host(common_attn_metadata)
        # MHA prefills are removed from q before forward_mqa. Plan only the
        # decode prefix in that case, using the same routing decision as MLA.
        if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha(
            metadata
        ):
            num_rows = metadata.num_decode_tokens
            kv_lens = kv_lens[:num_rows]
        if exact:
            self.state.plan(num_rows, kv_lens, None, None)
        else:
            self.state.plan(
                num_rows, kv_lens, common_attn_metadata, metadata.req_id_per_token
            )
        metadata.state = self.state
        return metadata

_kv_lens_host(cam)

Host per-row KV lengths for the plan and whether they are exact.

A row for the j-th query token of request i attends seq_lens[i] - q_len[i] + j + 1 tokens. The indexer's selection then bounds the valid count: contexts up to index_topk select everything (valid == context); longer contexts keep the top index_topk pool-expanded tokens plus the trailing incomplete pool (valid == index_topk + context % index_kpool). Both match the count of non -1 entries the convert kernel produces.

With seq_lens_cpu_upper_bound (optimistic under async spec decode) returns sync-free upper bounds, which the plan clamps on device. Otherwise returns exact lengths via a D2H copy.

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

def _kv_lens_host(
    self, cam: CommonAttentionMetadata
) -> tuple[int, torch.Tensor, bool]:
    """Host per-row KV lengths for the plan and whether they are exact.

    A row for the j-th query token of request i attends
    ``seq_lens[i] - q_len[i] + j + 1`` tokens. The indexer's selection
    then bounds the valid count: contexts up to ``index_topk`` select
    everything (valid == context); longer contexts keep the top
    ``index_topk`` pool-expanded tokens plus the trailing incomplete
    pool (valid == ``index_topk + context % index_kpool``). Both match
    the count of non -1 entries the convert kernel produces.

    With ``seq_lens_cpu_upper_bound`` (optimistic under async spec
    decode) returns sync-free upper bounds, which the plan clamps on
    device. Otherwise returns exact lengths via a D2H copy.
    """
    num_reqs = cam.num_reqs
    qsl = cam.query_start_loc_cpu[: num_reqs + 1]
    num_rows = int(qsl[-1])
    if num_rows == 0:
        return 0, torch.zeros(0, dtype=torch.int32), True
    state = self.state
    sl_host = cam.seq_lens_cpu_upper_bound
    if sl_host is None:
        ctx = self._row_contexts(qsl, cam.seq_lens[:num_reqs].cpu())
        topk, kpool = state.index_topk, state.index_kpool
        lens = torch.where(ctx <= topk, ctx, topk + ctx % kpool)
        return num_rows, lens.to(torch.int32), True
    if self._adaptive_verification:
        # The host query split may not match the device one, so bound
        # every row by the longest request.
        ctx = torch.full((num_rows,), int(sl_host[:num_reqs].max()))
    else:
        ctx = self._row_contexts(qsl, sl_host[:num_reqs])
    return num_rows, torch.clamp(ctx, max=state.max_valid), False

_SM90State

Builder-owned wrapper, capture-stable buffers, and plan parameters.

One instance serves every MLA layer in an attention group because the plan depends only on the batch shape, not the layer.

Methods:

• plan –

  Plan per-row KV lengths (CPU int32, [num_tokens]).

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

class _SM90State:
    """Builder-owned wrapper, capture-stable buffers, and plan parameters.

    One instance serves every MLA layer in an attention group because the plan
    depends only on the batch shape, not the layer.
    """

    def __init__(
        self,
        device: torch.device,
        num_heads: int,
        kv_dtype: torch.dtype,
        max_tokens: int,
        topk_width: int,
        kv_lora_rank: int,
        qk_rope_head_dim: int,
        sm_scale: float,
        index_topk: int,
        index_kpool: int,
    ) -> None:
        from flashinfer.mla import BatchMLAPagedAttentionWrapper

        self.workspace = torch.empty(_WORKSPACE_BYTES, dtype=torch.uint8, device=device)
        self.device = device
        self.num_heads = num_heads
        self.kv_dtype = kv_dtype
        self.max_tokens = max_tokens
        self.topk_width = topk_width
        self.kv_lora_rank = kv_lora_rank
        self.qk_rope_head_dim = qk_rope_head_dim
        self.sm_scale = sm_scale
        self.index_topk = index_topk
        self.index_kpool = max(index_kpool, 1)
        # Largest valid count the convert kernel can produce for any context.
        self.max_valid = min(self.index_topk + self.index_kpool - 1, topk_width)
        # User-reserved buffers: with use_cuda_graph=True plan() refreshes
        # these in place, so run()'s captured kernels always read them.
        self.kv_indices = torch.zeros(
            max_tokens * topk_width, dtype=torch.int32, device=device
        )
        self.kv_len_arr = torch.full(
            (max_tokens,), topk_width, dtype=torch.int32, device=device
        )
        self.kv_indptr = torch.zeros(max_tokens + 1, dtype=torch.int32, device=device)
        self.wrapper = BatchMLAPagedAttentionWrapper(
            self.workspace,
            qo_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
            kv_indptr=self.kv_indptr,
            kv_indices=self.kv_indices,
            kv_len_arr=self.kv_len_arr,
            use_cuda_graph=True,
            backend="fa3",
        )
        self._arange_cpu = torch.arange(self.max_tokens + 1, dtype=torch.int32)
        self._qo_cpu = torch.empty(self.max_tokens + 1, dtype=torch.int32)
        self._kv_cpu = torch.empty(self.max_tokens + 1, dtype=torch.int32)
        self._lens_cpu = torch.full(
            (self.max_tokens,), self.topk_width, dtype=torch.int32
        )
        # Planned per-work kv_end of the current plan, restored by each clamp;
        # sized on first plan to the planner's work capacity.
        self._saved_kv_end = torch.empty(0, dtype=torch.int32, device=device)
        self._planned_num_tokens = -1
        self._int_ws: torch.Tensor | None = None
        self._plan_offsets: tuple[int, int, int, int] | None = None

    def plan(
        self,
        num_tokens: int,
        kv_lens: torch.Tensor,
        cam: CommonAttentionMetadata | None,
        req_id_per_token: torch.Tensor | None,
    ) -> None:
        """Plan per-row KV lengths (CPU int32, ``[num_tokens]``).

        The wrapper bakes kv_len into its int schedule from host values, so
        rows past their valid count would send the kernel into the -1 tail of
        the converted index buffer. With ``cam`` and ``req_id_per_token``
        (device batch layout), ``kv_lens`` only needs to upper-bound the
        valid counts: the plan is padded by ``_PLAN_SLACK`` and reused while
        the bound still fits, and every call clamps each work item's kv_end
        on device to the exact count, so no D2H sync is needed. Without them
        the lengths must be exact. Must run outside CUDA graph capture: the
        in-place refreshed plan_info/indptr buffers are what captured runs
        read.
        """
        if torch.cuda.is_current_stream_capturing():
            raise RuntimeError(
                "FlashInferMLASparseSM90 plan() called inside CUDA graph "
                "capture; lengths must be planned host-side before capture."
            )
        kv_lens = kv_lens.to(torch.int32)
        if cam is None or req_id_per_token is None:
            self._plan(num_tokens, kv_lens)
            self._planned_num_tokens = -1
            return
        planned = self._lens_cpu[:num_tokens]
        if num_tokens != self._planned_num_tokens or not bool(
            ((planned >= kv_lens) & (planned <= kv_lens + _PLAN_SLACK)).all()
        ):
            self._plan(
                num_tokens, torch.clamp(kv_lens + _PLAN_SLACK, max=self.max_valid)
            )
            self._planned_num_tokens = num_tokens
        self._clamp_kv_end(cam, req_id_per_token)

    def _plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None:
        # CPU staging buffers are filled in place: plan() runs per step
        # (once per draft/verify metadata build), so per-call allocations and
        # device round trips are on the hot path. Passing CPU tensors lets
        # the wrapper's internal .to("cpu") no-op; its reserved-buffer
        # copy_ then performs the single H2D transfer per tensor.
        # use_cuda_graph=True makes the wrapper copy qo/kv indptr into its
        # fixed (max_tokens+1)-sized buffers with exact-size copy_, so the
        # indptr must always be full-size. Rows past num_tokens are padded
        # empty (qo_indptr flat at num_tokens) — zero-query rows read no q
        # and schedule no work. FlashInfer 0.7 requires packed CSR offsets
        # matching the exact lengths, including zero-length padded rows.
        torch.clamp(self._arange_cpu, max=num_tokens, out=self._qo_cpu)
        self._lens_cpu.zero_()
        self._lens_cpu[:num_tokens] = kv_lens
        self._kv_cpu[0] = 0
        torch.cumsum(self._lens_cpu, dim=0, out=self._kv_cpu[1:])
        # Only the live prefix of kv_indices: the wrapper snapshots and
        # copies whatever is passed.
        num_indices = max(int(self._kv_cpu[-1]), 1)
        self.wrapper.plan(
            self._qo_cpu,
            self._kv_cpu,
            self.kv_indices[:num_indices],
            self._lens_cpu,
            self.num_heads,
            self.kv_lora_rank,  # head_dim_ckv
            self.qk_rope_head_dim,  # 0 (NoPE) or 64 (rope MLA)
            1,  # page_size: top-k slots are the page table
            False,  # causal: encoded by the indexer's selection
            self.sm_scale,
            q_data_type=torch.bfloat16,
            kv_data_type=self.kv_dtype,
        )
        plan_info = [int(x) for x in self.wrapper._plan_info]
        assert len(plan_info) == _PLAN_INFO_LEN, plan_info
        # Per-work arrays are laid out back to back at the planner's capacity.
        max_works = (plan_info[_PI_KV_INDPTR] - plan_info[_PI_Q_INDPTR]) // 4
        if self._saved_kv_end.numel() != max_works:
            self._saved_kv_end = torch.empty(
                max_works, dtype=torch.int32, device=self.device
            )
        self._int_ws = self.wrapper._planned_backend._int_workspace_buffer.view(
            torch.int32
        )
        self._plan_offsets = (
            plan_info[_PI_WORK_INDPTR] // 4 + plan_info[_PI_NUM_BLKS_Y],
            plan_info[_PI_Q_INDPTR] // 4,
            plan_info[_PI_KV_START] // 4,
            plan_info[_PI_KV_END] // 4,
        )
        kv_end_off = self._plan_offsets[3]
        self._saved_kv_end.copy_(self._int_ws[kv_end_off : kv_end_off + max_works])

    def _clamp_kv_end(
        self, cam: CommonAttentionMetadata, req_id_per_token: torch.Tensor
    ) -> None:
        assert self._int_ws is not None and self._plan_offsets is not None
        block = 1024
        _clamp_work_kv_end[(triton.cdiv(self._saved_kv_end.numel(), block),)](
            self._int_ws,
            self._saved_kv_end,
            cam.seq_lens,
            cam.query_start_loc,
            req_id_per_token,
            *self._plan_offsets,
            TOPK=self.index_topk,
            KPOOL=self.index_kpool,
            BLOCK_SIZE=block,
        )

    def pack_indices(self, topk_slots: torch.Tensor) -> None:
        num_tokens, width = topk_slots.shape
        _pack_topk_indices[(num_tokens, triton.cdiv(width, 256))](
            topk_slots, self.kv_indptr, self.kv_indices, width, 256
        )

plan(num_tokens, kv_lens, cam, req_id_per_token)

Plan per-row KV lengths (CPU int32, [num_tokens]).

The wrapper bakes kv_len into its int schedule from host values, so rows past their valid count would send the kernel into the -1 tail of the converted index buffer. With cam and req_id_per_token (device batch layout), kv_lens only needs to upper-bound the valid counts: the plan is padded by _PLAN_SLACK and reused while the bound still fits, and every call clamps each work item's kv_end on device to the exact count, so no D2H sync is needed. Without them the lengths must be exact. Must run outside CUDA graph capture: the in-place refreshed plan_info/indptr buffers are what captured runs read.

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

def plan(
    self,
    num_tokens: int,
    kv_lens: torch.Tensor,
    cam: CommonAttentionMetadata | None,
    req_id_per_token: torch.Tensor | None,
) -> None:
    """Plan per-row KV lengths (CPU int32, ``[num_tokens]``).

    The wrapper bakes kv_len into its int schedule from host values, so
    rows past their valid count would send the kernel into the -1 tail of
    the converted index buffer. With ``cam`` and ``req_id_per_token``
    (device batch layout), ``kv_lens`` only needs to upper-bound the
    valid counts: the plan is padded by ``_PLAN_SLACK`` and reused while
    the bound still fits, and every call clamps each work item's kv_end
    on device to the exact count, so no D2H sync is needed. Without them
    the lengths must be exact. Must run outside CUDA graph capture: the
    in-place refreshed plan_info/indptr buffers are what captured runs
    read.
    """
    if torch.cuda.is_current_stream_capturing():
        raise RuntimeError(
            "FlashInferMLASparseSM90 plan() called inside CUDA graph "
            "capture; lengths must be planned host-side before capture."
        )
    kv_lens = kv_lens.to(torch.int32)
    if cam is None or req_id_per_token is None:
        self._plan(num_tokens, kv_lens)
        self._planned_num_tokens = -1
        return
    planned = self._lens_cpu[:num_tokens]
    if num_tokens != self._planned_num_tokens or not bool(
        ((planned >= kv_lens) & (planned <= kv_lens + _PLAN_SLACK)).all()
    ):
        self._plan(
            num_tokens, torch.clamp(kv_lens + _PLAN_SLACK, max=self.max_valid)
        )
        self._planned_num_tokens = num_tokens
    self._clamp_kv_end(cam, req_id_per_token)