vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90 ¶

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
        self.state = _SM90State(
            device,
            impl.num_heads,
            kv_cache_spec.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,
        )
        # seq_lens_cpu_upper_bound is optimistic on decode rows under async
        # spec decode, so the fast sync-free path is only safe without it;
        # under async scheduling the exact (device) positions are used at
        # the cost of one D2H sync per metadata build.
        self._async_scheduling = bool(vllm_config.scheduler_config.async_scheduling)
        hf_config = vllm_config.model_config.hf_text_config
        assert hf_config.index_topk is not None
        self._index_topk = int(hf_config.index_topk)
        self._index_kpool = int(kv_cache_spec.tokens_per_state)

    def _kv_lens_host(self, cam: CommonAttentionMetadata) -> tuple[int, torch.Tensor]:
        """Exact per-row KV lengths, host-side (the flashinfer wrapper bakes
        them into its schedule at plan time; there is no device-side path).

        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.
        """
        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)
        # Row context == position + 1. Without async scheduling the
        # maintained host upper bound is exact, giving a sync-free path;
        # under async scheduling it is optimistic on decode rows, so fall
        # back to the exact device positions (one D2H sync per build). The
        # seq_lens derivation equals position + 1 because positions are
        # contiguous per request.
        sl_host = cam.seq_lens_cpu_upper_bound
        positions = cam.positions
        if not self._async_scheduling and sl_host is not None:
            seq_lens = sl_host[:num_reqs].to(torch.int32)
            q_lens = qsl[1:] - qsl[:-1]
            first_pos = seq_lens - q_lens
            req_of_row = torch.repeat_interleave(
                torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64)
            )
            rows = torch.arange(num_rows, dtype=torch.int32)
            ctx = (
                first_pos.to(torch.int64)[req_of_row]
                + rows.to(torch.int64)
                - qsl.to(torch.int64)[req_of_row]
                + 1
            )
        elif positions is not None and num_rows <= positions.shape[0]:
            ctx = positions[:num_rows].cpu().to(torch.int64) + 1
        else:
            seq_lens = cam.seq_lens[:num_reqs].cpu().to(torch.int32)
            q_lens = qsl[1:] - qsl[:-1]
            first_pos = seq_lens - q_lens
            req_of_row = torch.repeat_interleave(
                torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64)
            )
            rows = torch.arange(num_rows, dtype=torch.int32)
            ctx = (
                first_pos.to(torch.int64)[req_of_row]
                + rows.to(torch.int64)
                - qsl.to(torch.int64)[req_of_row]
                + 1
            )
        topk = self._index_topk
        kpool = max(self._index_kpool, 1)
        lens = torch.where(ctx <= topk, ctx, topk + ctx % kpool)
        return num_rows, lens.to(torch.int32)

    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)
        # Replan every step outside any CUDA graph capture with this step's
        # exact per-row lengths; captured runs read the refreshed buffers.
        num_rows, kv_lens = 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]
        self.state.plan(num_rows, kv_lens)
        metadata.state = self.state
        return metadata

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,
    ) -> 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
        # 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.wrapper = BatchMLAPagedAttentionWrapper(
            self.workspace,
            qo_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
            kv_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
            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
        )

    def plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None:
        """Replan with exact per-row KV lengths (CPU int32, ``[num_tokens]``).

        The wrapper bakes kv_len into its int schedule from host values;
        run() never reads the device kv_len_arr buffer. Scheduling with the
        full buffer width while kv_indices rows carry a -1 tail past each
        row's valid count makes the kernel compute ``-1 * ckv_stride_page``
        (illegal address), so the lengths must be exact at plan time and
        replanned every step as contexts grow. 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."
            )
        # 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. Padded rows keep the full width lens; the
        # value is never dereferenced.
        torch.clamp(self._arange_cpu, max=num_tokens, out=self._qo_cpu)
        torch.mul(self._qo_cpu, self.topk_width, out=self._kv_cpu)
        self._lens_cpu.fill_(self.topk_width)
        self._lens_cpu[:num_tokens] = kv_lens.to(torch.int32)
        self.wrapper.plan(
            self._qo_cpu,
            self._kv_cpu,
            self.kv_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,
        )