vllm.v1.attention.backends.mla.rocm_aiter_mla ¶

class AiterMLAHelper:
    """
    AITER MLA persistent (asm) decode requires a multiple of 16 heads. Unaligned
    head counts through 128 are padded to the next multiple of 16 by tiling the
    query heads and slicing to the padded size. Native H24 AITER builds bypass
    that padding. Small divisors of 16 retain the existing repeat_interleave and
    strided-unpad behavior. Native and aligned counts pass through without
    copies.
    """

    _AITER_MIN_MLA_HEADS: Final = 16
    _AITER_MAX_PADDED_MLA_HEADS: Final = 128
    # Largest qlen the padded gqa=16 asm decode has a bf16 persistent kernel
    # for. Above it only the non-persistent qseqlen=8 entry exists, and the
    # fold that reaches a persistent one is gfx950-only.
    _ASM_PADDED_MAX_PS_QLEN: Final = 4
    _AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = ()

    @staticmethod
    def check_num_heads_validity(num_heads: int):
        assert AiterMLAHelper.is_valid_num_heads(num_heads), (
            "ROCM AITER MLA requires a positive multiple of 16 heads, or an "
            "unaligned head count up to 128 (padded to the next multiple of "
            f"16), but got {num_heads}.\n"
            f"Try adjusting tensor_parallel_size value."
        )

    @staticmethod
    def is_valid_num_heads(num_heads: int) -> bool:
        return (
            num_heads > 0
            and num_heads not in AiterMLAHelper._AITER_UNSUPPORTED_HEADS
            and (
                num_heads <= AiterMLAHelper._AITER_MAX_PADDED_MLA_HEADS
                or num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0
            )
        )

    @staticmethod
    def get_actual_mla_num_heads(num_heads: int) -> int:
        if num_heads == 24 and _aiter_mla_native_h24_supported():
            return num_heads
        m = AiterMLAHelper._AITER_MIN_MLA_HEADS
        return -(-num_heads // m) * m

    @staticmethod
    def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor:
        m = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
        if num_heads == m:
            return q
        if m % num_heads == 0:
            return q.repeat_interleave(m // num_heads, dim=1)
        # Non-divisor head counts cannot be padded by repeat_interleave. Tile
        # the query heads and slice to exactly m. MLA attention is independent
        # per query head over the shared KV, so padding heads cannot affect
        # heads [0:num_heads]; they are sliced back off the output.
        reps = -(-m // num_heads)  # ceil(m / num_heads)
        # Slicing a tiled tensor yields a non-contiguous view. The asm decode
        # reads q as packed [tokens, m, head_dim], so materialize it.
        return q.repeat(1, reps, 1)[:, :m, :].contiguous()

    @staticmethod
    def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor:
        m = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
        if num_heads == m:
            return o
        if m % num_heads == 0:
            return o[:, :: m // num_heads, :]
        # Undo the tile-padding from get_mla_padded_q: the real heads are the
        # first num_heads.
        return o[:, :num_heads, :]

    @staticmethod
    def use_gluon_decode(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
        # Small-head (<16) single-token decode takes either the Gluon kernel or
        # the padded asm persistent decode, selected by
        # VLLM_ROCM_AITER_MLA_ASM_PADDING and the arch (Gluon is gfx950 only).
        m = AiterMLAHelper._AITER_MIN_MLA_HEADS
        if num_heads >= m or max_qo_len != 1:
            return False
        # Gluon's only fp8-KV regime, bh16bn128, is a bf16-query kernel with a
        # hardcoded scale that asserts batch_size == 1, so it cannot serve a
        # decode batch. Checked before the mode knob: an explicit "gluon"
        # request under fp8 would assert immediately.
        if is_quantized_kv_cache(kv_cache_dtype):
            return False
        mode = _aiter_mla_small_head_mode()
        if mode == "asm":
            return False
        gluon_supported = _gluon_mla_decode_supported()
        if mode == "gluon":
            return gluon_supported
        return m % num_heads == 0 and gluon_supported

    @staticmethod
    def use_gluon_verify(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
        """Whether a small-head multi-token verify is flattened onto Gluon.

        bf16 has no gqa<16, qseqlen>1 asm kernel, so the verify is flattened
        into per-token Gluon decodes. fp8 has one via the q-row fold and must
        not come here: the flatten hands Gluon the batch size its fp8 regime
        asserts against. A predicate rather than inline in forward_mqa so the
        builder sees the same answer the impl acts on.
        """
        if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= 1:
            return False
        if is_quantized_kv_cache(kv_cache_dtype):
            return False
        # Same arch and mode gating as use_gluon_decode.
        return _aiter_mla_small_head_mode() != "asm" and _gluon_mla_decode_supported()

class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: int,
        alibi_slopes: list[float] | None,
        sliding_window: int | None,
        kv_cache_dtype: str,
        logits_soft_cap: float | None,
        attn_type: str,
        kv_sharing_target_layer_name: str | None,
        # MLA Specific Arguments
        **mla_args,
    ) -> None:
        super().__init__(
            num_heads,
            head_size,
            scale,
            num_kv_heads,
            alibi_slopes,
            sliding_window,
            kv_cache_dtype,
            logits_soft_cap,
            attn_type,
            kv_sharing_target_layer_name,
            **mla_args,
        )
        AiterMLAHelper.check_num_heads_validity(num_heads)

        unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
        if any(unsupported_features):
            raise NotImplementedError(
                "Aiter MLA does not support one of the following: "
                "alibi_slopes, sliding_window, logits_soft_cap"
            )

        from aiter import flash_attn_varlen_func

        self.flash_attn_varlen_func = flash_attn_varlen_func

        # FP8 MLA prefill kernel imports (lazy, only when enabled).
        # Auto-enabled on gfx950 when AITER ships the kernels.
        # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads.
        self._fp8_prefill_enabled = (
            _fp8_mla_prefill_supported() and self.num_heads % 16 == 0
        )
        if self._fp8_prefill_enabled:
            from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1

            self._mla_prefill_ps_asm_fwd = mla_prefill_ps_asm_fwd
            self._mla_reduce_v1 = mla_reduce_v1

    def _flash_attn_varlen_diff_headdims(
        self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs
    ):
        output = self.flash_attn_varlen_func(  # type: ignore[call-arg]
            q=q,
            k=k,
            v=v,
            softmax_scale=softmax_scale,
            return_lse=return_softmax_lse,
            **kwargs,
        )

        return output

    def _mla_fp8_prefill_attn(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        attn_metadata: AiterMLAMetadata,
        out: torch.Tensor,
    ) -> None:
        """Run FP8 MLA prefill via mla_prefill_ps_asm_fwd + mla_reduce_v1.

        Q, K, V are already decompressed (post-kv_b_proj), so K and V have
        ``num_heads`` heads (same as Q) and gqa_ratio=1.  Writes the
        result in-place to ``out``, which is the [total_q, nhead * v_head_dim]
        output buffer supplied by ``forward_mha``; no extra allocation or
        copy is required.
        """
        from vllm.platforms import current_platform
        from vllm.v1.worker.workspace import current_workspace_manager

        fp8_dtype = current_platform.fp8_dtype()
        total_q = q.shape[0]
        nhead = self.num_heads
        v_head_dim = self.v_head_dim
        tile_q = _FP8_PREFILL_TILE_Q

        # The FP8 ASM kernel expects FP8 inputs; the q_scale/k_scale/v_scale
        # parameters select per-tensor dequant scales.  Q/K/V arrive as
        # bf16 from kv_b_proj, so cast here (one_scale=1.0 disables scaling).
        if q.dtype != fp8_dtype:
            q = q.to(fp8_dtype)
        if k.dtype != fp8_dtype:
            k = k.to(fp8_dtype)
        if v.dtype != fp8_dtype:
            v = v.to(fp8_dtype)

        one_scale = torch.ones((), dtype=torch.float32, device=q.device)

        # num_partial_tiles is resolved during metadata build to avoid an
        # in-forward .item() sync that would prevent CUDA Graph capture.
        # forward_mha gates the FP8 path on fp8_prefill_qo_indptr being set,
        # and the builder always sets every fp8_prefill_* field together, so
        # num_partial_tiles is non-None here.
        num_partial_tiles = attn_metadata.fp8_prefill_num_partial_tiles
        assert num_partial_tiles is not None

        # Reuse the caller's output buffer to skip the per-call alloc + copy.
        # The ASM and reduce kernels both write to a [total_q, nhead, v_head_dim]
        # view, which aliases the [total_q, nhead * v_head_dim] storage of out.
        out_3d = out.view(total_q, nhead, v_head_dim)

        # Per-call scratch (logits, attn_lse, final_lse) is served from the
        # workspace manager so allocator churn in the prefill hot path is
        # bounded after warmup, matching the pattern in PR #41002.
        logits, attn_lse, final_lse = current_workspace_manager().get_simultaneous(
            ((num_partial_tiles * tile_q, nhead, v_head_dim), torch.float32),
            ((num_partial_tiles * tile_q, nhead), torch.float32),
            ((total_q, nhead), torch.float32),
        )

        # Phase 1: persistent-scheduling assembly prefill kernel.
        self._mla_prefill_ps_asm_fwd(
            q,
            k,
            v,
            attn_metadata.fp8_prefill_qo_indptr,
            attn_metadata.fp8_prefill_kv_indptr,
            attn_metadata.fp8_prefill_kv_indices,
            attn_metadata.fp8_prefill_work_indptr,
            attn_metadata.fp8_prefill_work_info_set,
            attn_metadata.fp8_prefill_max_q_len,
            self.scale,
            True,  # is_causal
            logits,
            attn_lse,
            out_3d,
            one_scale,
            one_scale,
            one_scale,
        )

        # Phase 2: reduction across KV splits.
        self._mla_reduce_v1(
            logits,
            attn_lse,
            attn_metadata.fp8_prefill_reduce_indptr,
            attn_metadata.fp8_prefill_reduce_final_map,
            attn_metadata.fp8_prefill_reduce_partial_map,
            tile_q,
            # num_kv_splits added by ROCm/aiter#3391; 0 selects the kernel
            # default max(cu_num, 0) == cu_num, matching pre-#3391 behavior.
            0,
            out_3d,
            final_lse,
        )

    def forward_mha(
        self,
        q: torch.Tensor,
        kv_c_normed: torch.Tensor,
        k_pe: torch.Tensor,
        kv_c_and_k_pe_cache: torch.Tensor,
        attn_metadata: MLACommonMetadata,
        k_scale: torch.Tensor,
        output: torch.Tensor,
        output_scale: torch.Tensor | None = None,
    ) -> None:
        """Dispatch prefill to the FP8 ASM kernel when available.

        Falls back to the parent (``flash_attn_varlen_func``) when FP8
        MLA prefill is disabled, PS metadata is missing, or chunked
        context requires two-pass merge.

        The annotation uses the base ``MLACommonMetadata`` to honour LSP
        with ``MLACommonImpl.forward_mha``; the AITER builder always
        produces ``AiterMLAMetadata`` instances at runtime, so we narrow
        with ``isinstance`` before reading the AITER-specific FP8 fields.
        """
        if (
            not self._fp8_prefill_enabled
            or not isinstance(attn_metadata, AiterMLAMetadata)
            or attn_metadata.fp8_prefill_qo_indptr is None
        ):
            return super().forward_mha(
                q,
                kv_c_normed,
                k_pe,
                kv_c_and_k_pe_cache,
                attn_metadata,
                k_scale,
                output,
                output_scale,
            )

        assert attn_metadata.prefill is not None
        prefill_metadata = attn_metadata.prefill
        has_context = prefill_metadata.chunked_context is not None

        if has_context:
            return super().forward_mha(
                q,
                kv_c_normed,
                k_pe,
                kv_c_and_k_pe_cache,
                attn_metadata,
                k_scale,
                output,
                output_scale,
            )

        assert output_scale is None, (
            "fused FP8 output not supported by the AITER FP8 MLA prefill path"
        )

        kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
            -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
        )
        k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
        k = self._concat_k_nope_k_pe(k_nope, k_pe)

        self._mla_fp8_prefill_attn(q, k, v, attn_metadata, output)

    def forward_mqa(
        self,
        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        kv_c_and_k_pe_cache: torch.Tensor,
        attn_metadata: AiterMLAMetadata,
        layer: AttentionLayer,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        assert kv_c_and_k_pe_cache.numel() > 0
        assert attn_metadata.decode is not None

        decode = attn_metadata.decode
        assert decode.max_qo_len is not None
        assert decode.paged_kv_indptr is not None
        assert decode.paged_kv_indices is not None
        if decode.use_gluon_decode:
            if type(q) is tuple:
                q_nope, q_pe = q
            else:
                q_nope, q_pe = torch.split(
                    q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
                )
            B, num_q_heads, _ = q_nope.shape
            o = torch.empty(
                B,
                num_q_heads,
                self.kv_lora_rank,
                dtype=decode.attn_out_dtype,
                device=q_nope.device,
            )
            kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])
            mla_gluon = _get_mla_gluon()
            mla_gluon(
                q_nope=q_nope,
                q_pe=q_pe,
                kv_c=kv_buffer,
                o=o,
                page_table=decode.paged_kv_indices,
                seq_info=decode.paged_kv_indptr,
                sm_scale=self.scale,
                k_pe=None,
                kv_pe_offset=self.kv_lora_rank,
                use_2d_view=False,
                kv_scale=1.0,
                min_kv_seq_len=decode.min_kv_seq_len,
            )
            return o, None

        # 12-head (<16) multi-token verify (DSpark): the asm path has no
        # gqa<16, qseqlen>1 kernel. Flatten each verify token to its own
        # qseqlen=1 gluon decode, mirroring the TRITON_MLA / sparse-backend
        # flatten but on the fast gluon kernel. The block is causal -- the
        # target is checking draft tokens, so position t must not see t+1 --
        # and attention rows are independent, so giving row t the KV range
        # [0, context + t] is exactly causal multi-token attention.
        # Arch, mode and dtype gating all live in use_gluon_verify, so that the
        # builder -- which has to know whether the asm decode will run -- sees
        # the same answer as this branch.
        if AiterMLAHelper.use_gluon_verify(
            self.num_heads, int(decode.max_qo_len), self.kv_cache_dtype
        ):
            qlen = int(decode.max_qo_len)
            if type(q) is tuple:
                q_nope, q_pe = q
            else:
                q_nope, q_pe = torch.split(
                    q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
                )
            B, num_q_heads, _ = q_nope.shape
            o = torch.empty(
                B,
                num_q_heads,
                self.kv_lora_rank,
                dtype=decode.attn_out_dtype,
                device=q_nope.device,
            )
            kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])
            # Expand per-request paged-KV to per-verify-token. Row r*qlen+t is
            # request r's verify token t, and seq_lens counts the tokens
            # scheduled in this step, so a request's KV range already spans its
            # whole verify block and context_r = seq_len_r - qlen. Token t may
            # attend to [0, context_r + t], i.e. seq_len_r - (qlen - 1) + t
            # entries. paged_kv_indices lists a request's pages in ascending
            # position order, so each row's causal window is a prefix of that
            # request's slice and only the row length changes. Rows clamp to
            # zero for cudagraph padding requests, whose seq_len is 0. Fully
            # vectorized (no host loop).
            old_indptr = decode.paged_kv_indptr
            per_req_len = old_indptr[1:] - old_indptr[:-1]
            dev = q_nope.device
            row_req = torch.arange(per_req_len.shape[0], device=dev).repeat_interleave(
                qlen
            )
            row_len = (
                (
                    per_req_len.unsqueeze(1)
                    - (qlen - 1)
                    + torch.arange(qlen, device=dev, dtype=per_req_len.dtype)
                )
                .clamp_(min=0)
                .flatten()
            )
            new_indptr = torch.cat([old_indptr.new_zeros(1), row_len.cumsum(0)]).to(
                torch.int32
            )
            total = int(new_indptr[-1].item())
            within = torch.arange(total, device=dev, dtype=torch.int64) - new_indptr[
                :-1
            ].to(torch.int64).repeat_interleave(row_len)
            src = (
                old_indptr[row_req].to(torch.int64).repeat_interleave(row_len) + within
            )
            new_indices = decode.paged_kv_indices[src]
            mla_gluon = _get_mla_gluon()
            mla_gluon(
                q_nope=q_nope,
                q_pe=q_pe,
                kv_c=kv_buffer,
                o=o,
                page_table=new_indices,
                seq_info=new_indptr,
                sm_scale=self.scale,
                k_pe=None,
                kv_pe_offset=self.kv_lora_rank,
                use_2d_view=False,
                kv_scale=1.0,
                min_kv_seq_len=int(row_len.min()),
            )
            return o, None

        if type(q) is tuple:
            q = torch.cat(q, dim=-1)

        assert isinstance(q, torch.Tensor)
        B = q.shape[0]

        mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)
        mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads)
        o = torch.empty(
            B,
            mla_num_heads,
            self.kv_lora_rank,
            dtype=attn_metadata.decode.attn_out_dtype,
            device=q.device,
        )
        if decode.max_qo_len > 1 and not decode.has_persistent_metadata:
            # MTP verification can call the AITER MLA decode kernel with
            # qlen > 1. If that path is running without persistent metadata,
            # zero-fill so unwritten lanes cannot leak into logits.
            o.zero_()

        kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)

        # Build kwargs for mla_decode_fwd. Pass persistent metadata only
        # when it was successfully computed.
        mla_kwargs = dict(
            q_scale=layer._q_scale,
            kv_scale=layer._k_scale,
        )
        if attn_metadata.work_meta_data is not None:
            mla_kwargs.update(
                work_meta_data=attn_metadata.work_meta_data,
                work_indptr=attn_metadata.work_indptr,
                work_info_set=attn_metadata.work_info_set,
                reduce_indptr=attn_metadata.reduce_indptr,
                reduce_final_map=attn_metadata.reduce_final_map,
                reduce_partial_map=attn_metadata.reduce_partial_map,
            )

        rocm_aiter_ops.mla_decode_fwd(
            mla_padded_q,
            kv_buffer,
            o,
            self.scale,
            decode.qo_indptr,
            decode.max_qo_len,
            decode.paged_kv_indptr,
            decode.paged_kv_indices,
            decode.paged_kv_last_page_len,
            **mla_kwargs,
        )

        return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, o), None

class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
    # TODO(luka, lucas): audit this as part of:
    #  https://github.com/vllm-project/vllm/issues/22945
    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
    query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM

    @staticmethod
    def _uniform_padded_mtp_qo_len(
        qo_len: torch.Tensor,
        max_qo_len: int,
        num_decode_tokens: int,
    ) -> int:
        num_reqs = qo_len.numel()
        if num_reqs == 0 or num_decode_tokens <= 0:
            return 0

        # Full-CG pads q to a captured token count while leaving
        # query_start_loc flat for dummy requests. Only synthesize dummy rows
        # when every padded request maps to the same qlen and the q buffer has
        # exactly that many rows.
        if num_decode_tokens <= int(qo_len.sum().item()):
            return 0
        if num_decode_tokens % num_reqs != 0:
            return 0

        uniform_qo_len = num_decode_tokens // num_reqs
        if uniform_qo_len <= 1:
            return 0

        positive_qo_len = qo_len[qo_len > 0]
        if positive_qo_len.numel() == qo_len.numel():
            return 0
        if positive_qo_len.numel() > 0:
            if max_qo_len != uniform_qo_len:
                return 0
            if not torch.all(positive_qo_len == uniform_qo_len):
                return 0

        zero_positions = torch.nonzero(qo_len == 0, as_tuple=False).flatten()
        if zero_positions.numel() > 0:
            first_zero = int(zero_positions[0].item())
            if torch.any(qo_len[first_zero:] > 0):
                return 0

        return uniform_qo_len

    def __init__(
        self,
        kv_cache_spec: AttentionSpec,
        layer_names: list[str],
        vllm_config: VllmConfig,
        device: torch.device,
    ):
        super().__init__(
            kv_cache_spec, layer_names, vllm_config, device, AiterMLAMetadata
        )

        self.compilation_config = vllm_config.compilation_config
        self.decode_attn_out_dtype = vllm_config.model_config.dtype

        # reorder_batch_threshold is the largest query length decode can be
        # handed, and already accounts for the drafting scheme. A method-name
        # whitelist sizes unlisted drafters for qlen=1, which closes the
        # persistent gate below and makes aiter raise a KeyError mid-run.
        self._mtp_decode_qlen = self.reorder_batch_threshold or 1

        # Store the kernel block size from the spec. When kernel_block_size=1
        # (no spec-dec), behavior is identical to the original. When > 1
        # (e.g. 16 with Eagle3), we expand block-level indices into per-token
        # flat indices since the aiter kernel always uses page_size=1 internally.
        self.kernel_block_size = kv_cache_spec.block_size

        # In the flat view (.view(-1,1,1,H)), each token is its own page,
        # so max_num_pages_per_req = max_model_len regardless of
        # kernel_block_size.
        max_num_pages_per_req = vllm_config.model_config.max_model_len
        max_num_reqs = vllm_config.scheduler_config.max_num_seqs
        max_num_pages = max_num_reqs * max_num_pages_per_req

        # Preparing persistent buffers
        # TODO: we can disambiguate between decode and mixed-prefill decode here
        # so we can only use the persistent buffer if a cudagraph is actually
        # being used.

        # paged_kv_last_page_len is always 1s (the aiter kernel always sees
        # page_size=1 after .view(-1,1,1,H) flattening), so we create it
        # once and reuse slices in both eager and cudagraph modes.
        self.paged_kv_last_page_len = torch.ones(
            max_num_reqs, dtype=torch.int32, device=device
        )

        # Persistent buffer for paged_kv_indices to avoid blocking boolean mask
        # indexing (block_table_tensor[mask]) which has data-dependent output size.
        self.paged_kv_indices = torch.zeros(
            max_num_pages, dtype=torch.int32, device=device
        )

        from aiter import dtypes, get_mla_metadata_info_v1

        # Keep metadata sizing consistent with the padded tensor shape passed
        # to mla_decode_fwd.
        self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads(
            self.num_heads
        )
        kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
        if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):
            kv_cache_dtype_str = "fp8"
            kv_dtype = dtypes.fp8
        else:
            kv_dtype = {
                torch.float16: dtypes.fp16,
                torch.bfloat16: dtypes.bf16,
            }[kv_cache_spec.dtype]
        # _build_decode needs the cache dtype to pick the decode kernel; keep
        # the normalized string instead of dropping it at the end of __init__.
        self._kv_cache_dtype_str = kv_cache_dtype_str
        # MLAAttention quantizes decode Q to FP8 before calling this backend
        # whenever the KV cache is FP8 and supports_quant_query_input is true.
        q_dtype = (
            dtypes.fp8 if kv_cache_dtype_str == "fp8" else self.decode_attn_out_dtype
        )
        # Persist for get_mla_metadata_v1 (decode build): omitting these causes
        # wrong split/reduce metadata for the gfx950 fp8 nhead=32 fold path.
        self._mla_q_dtype = q_dtype
        self._mla_kv_dtype = kv_dtype
        (
            (work_meta_data_size, work_meta_data_type),
            (work_indptr_size, work_indptr_type),
            (work_info_set_size, work_info_set_type),
            (reduce_indptr_size, reduce_indptr_type),
            (reduce_final_map_size, reduce_final_map_type),
            (reduce_partial_map_size, reduce_partial_map_type),
        ) = get_mla_metadata_info_v1(
            max_num_reqs,
            self._mtp_decode_qlen,
            self._num_attention_heads,
            q_dtype,
            kv_dtype,
            is_sparse=False,
            fast_mode=True,
        )
        self._mla_work_meta_data = torch.empty(
            work_meta_data_size, dtype=work_meta_data_type, device=device
        )
        self._mla_work_indptr = torch.empty(
            work_indptr_size, dtype=work_indptr_type, device=device
        )
        self._mla_work_info_set = torch.empty(
            work_info_set_size, dtype=work_info_set_type, device=device
        )
        self._mla_reduce_indptr = torch.empty(
            reduce_indptr_size, dtype=reduce_indptr_type, device=device
        )
        self._mla_reduce_final_map = torch.empty(
            reduce_final_map_size, dtype=reduce_final_map_type, device=device
        )
        self._mla_reduce_partial_map = torch.empty(
            reduce_partial_map_size,
            dtype=reduce_partial_map_type,
            device=device,
        )

        # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads.
        self._fp8_prefill_enabled = (
            _fp8_mla_prefill_supported() and self.num_heads % 16 == 0
        )
        if self._fp8_prefill_enabled:
            max_prefill_qlen = min(
                vllm_config.model_config.max_model_len,
                vllm_config.scheduler_config.max_num_batched_tokens,
            )
            self._init_fp8_prefill_ps_buffers(
                max_num_reqs,
                max_prefill_qlen,
                vllm_config.scheduler_config.max_num_batched_tokens,
                device,
            )

        if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
            self.paged_kv_indptr = torch.zeros(
                max_num_reqs + 1, dtype=torch.int32, device=device
            )

            self.qo_indptr = torch.zeros(
                max_num_reqs + 1, dtype=torch.int32, device=device
            )

    def _init_fp8_prefill_ps_buffers(
        self,
        max_num_reqs: int,
        max_prefill_qlen: int,
        max_num_batched_tokens: int,
        device: torch.device,
    ) -> None:
        """Pre-allocate persistent buffers for FP8 MLA prefill PS metadata.

        Uses ``get_ps_metadata_info_v1`` with max values so the buffers are
        large enough for any batch.  ``get_ps_metadata_v1`` fills them
        per-batch in ``build()``.  The FP8 prefill forward path also uses the
        global workspace manager for per-call scratch, so reserve its maximum
        shape here before the workspace manager is locked after warmup.

        Args:
            max_num_reqs: Maximum number of concurrent requests.
            max_prefill_qlen: Maximum Q-length for a single request in one
                prefill batch.  Should be ``min(max_model_len,
                max_num_batched_tokens)`` — a single request never exceeds
                ``max_model_len`` tokens, nor the per-batch token budget.
            max_num_batched_tokens: Maximum number of tokens scheduled in one
                batch.  The ``final_lse`` scratch is sized by ``total_q`` (the
                summed Q-length over all prefill requests in the batch), which
                is bounded by this budget rather than by a single request's
                ``max_prefill_qlen`` — concurrent requests can sum to more than
                ``max_model_len`` when ``max_model_len < max_num_batched_tokens``.
            device: Target device for the buffers.
        """
        from aiter import get_ps_metadata_info_v1

        # After kv_b_proj decompression, K has num_heads heads (same as Q).
        # So gqa_ratio=1 and num_head_k=num_heads for the PS kernel.
        num_head_k = self.num_heads
        v_head_dim = self.mla_dims.v_head_dim
        # gqa_ratio = 1
        # qlen_granularity = _FP8_PREFILL_TILE_Q // max(gqa_ratio, 1)
        qlen_granularity = _FP8_PREFILL_TILE_Q

        (
            (work_metadata_size, work_metadata_dtype),
            (work_indptr_size, work_indptr_dtype),
            (work_info_size, work_info_dtype),
            (reduce_indptr_size, reduce_indptr_dtype),
            (reduce_final_map_size, reduce_final_map_dtype),
            (reduce_partial_map_size, reduce_partial_map_dtype),
        ) = get_ps_metadata_info_v1(
            batch_size=max_num_reqs,
            num_head_k=num_head_k,
            max_qlen=max_prefill_qlen,
            qlen_granularity=qlen_granularity,
        )

        self.fp8_ps_work_metadata = torch.empty(
            work_metadata_size, dtype=work_metadata_dtype, device=device
        )
        self.fp8_ps_work_indptr = torch.empty(
            work_indptr_size, dtype=work_indptr_dtype, device=device
        )
        self.fp8_ps_work_info = torch.empty(
            *work_info_size, dtype=work_info_dtype, device=device
        )
        self.fp8_ps_reduce_indptr = torch.empty(
            reduce_indptr_size, dtype=reduce_indptr_dtype, device=device
        )
        self.fp8_ps_reduce_final_map = torch.empty(
            *reduce_final_map_size, dtype=reduce_final_map_dtype, device=device
        )
        self.fp8_ps_reduce_partial_map = torch.empty(
            reduce_partial_map_size,
            dtype=reduce_partial_map_dtype,
            device=device,
        )

        from vllm.v1.worker.workspace import current_workspace_manager

        max_num_partial_tiles = reduce_partial_map_size
        current_workspace_manager().get_simultaneous(
            (
                (max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k, v_head_dim),
                torch.float32,
            ),
            (
                (max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k),
                torch.float32,
            ),
            ((max_num_batched_tokens, num_head_k), torch.float32),
        )

        logger.info(
            "FP8 MLA prefill PS buffers allocated "
            "(max_batch=%d, max_qlen=%d, num_head_k=%d)",
            max_num_reqs,
            max_prefill_qlen,
            num_head_k,
        )

    def _build_fp8_prefill_ps_metadata(
        self,
        metadata: AiterMLAMetadata,
        common_attn_metadata: CommonAttentionMetadata,
    ) -> None:
        """Build per-batch FP8 MLA prefill PS metadata and attach to *metadata*.

        Called from ``build()`` when prefill tokens are present and
        FP8 MLA prefill is enabled (auto-detected via
        ``_fp8_mla_prefill_supported()``).
        """
        from aiter import get_ps_metadata_v1

        prefill = metadata.prefill
        # Caller (build()) only invokes this when prefill tokens exist, so
        # metadata.prefill is guaranteed non-None.  Assert to narrow for mypy.
        assert prefill is not None
        qo_indptr = prefill.query_start_loc
        kv_indptr = qo_indptr  # new tokens: KV length == Q length

        # Reuse the existing CPU view of query_start_loc instead of forcing a
        # device->host copy.  Prefill batches sit at the tail of the request
        # list, so we slice from num_decodes onwards and rebase to zero, the
        # same transform the parent build applies on device tensors.
        num_decodes = metadata.num_decodes
        qsl_cpu = common_attn_metadata.query_start_loc_cpu
        qo_indptr_cpu = (qsl_cpu[num_decodes:] - qsl_cpu[num_decodes]).to(torch.int32)
        kv_indptr_cpu = qo_indptr_cpu.clone()
        seq_lens_cpu = (qo_indptr_cpu[1:] - qo_indptr_cpu[:-1]).to(torch.int32)

        num_head_k = self.num_heads
        # gqa_ratio = 1
        # qhead_granularity = max(gqa_ratio, 1)
        # qlen_granularity = _FP8_PREFILL_TILE_Q // qhead_granularity
        gqa_ratio = 1
        qhead_granularity = 1
        qlen_granularity = _FP8_PREFILL_TILE_Q
        kvlen_granularity = 128
        block_size = 1  # non-paged: each "page" is one token

        get_ps_metadata_v1(
            qo_indptr_cpu,
            kv_indptr_cpu,
            seq_lens_cpu,
            gqa_ratio,
            num_head_k,
            self.fp8_ps_work_metadata,
            self.fp8_ps_work_indptr,
            self.fp8_ps_work_info,
            self.fp8_ps_reduce_indptr,
            self.fp8_ps_reduce_final_map,
            self.fp8_ps_reduce_partial_map,
            qhead_granularity=qhead_granularity,
            qlen_granularity=qlen_granularity,
            kvlen_granularity=kvlen_granularity,
            block_size=block_size,
            is_causal=True,
        )

        total_prefill_tokens = int(qo_indptr_cpu[-1].item())
        kv_indices = torch.arange(
            total_prefill_tokens, device=qo_indptr.device, dtype=torch.int32
        )

        # The actual number of active partial tiles for this batch is the
        # final value of reduce_indptr.  Resolving it here (during metadata
        # build) keeps it off the per-layer forward path where a sync would
        # break CUDA Graph capture.  Using the device-side reduce_indptr is
        # acceptable since build is allowed to incur an occasional sync.
        num_partial_tiles = int(self.fp8_ps_reduce_indptr[-1].item())

        # Attach PS metadata to the metadata object so forward_mha can read it.
        metadata.fp8_prefill_qo_indptr = qo_indptr
        metadata.fp8_prefill_kv_indptr = kv_indptr
        metadata.fp8_prefill_kv_indices = kv_indices
        metadata.fp8_prefill_work_indptr = self.fp8_ps_work_indptr
        metadata.fp8_prefill_work_info_set = self.fp8_ps_work_info
        metadata.fp8_prefill_reduce_indptr = self.fp8_ps_reduce_indptr
        metadata.fp8_prefill_reduce_final_map = self.fp8_ps_reduce_final_map
        metadata.fp8_prefill_reduce_partial_map = self.fp8_ps_reduce_partial_map
        metadata.fp8_prefill_max_q_len = prefill.max_query_len
        metadata.fp8_prefill_num_partial_tiles = num_partial_tiles

    def _build_decode(
        self,
        block_table_tensor: torch.Tensor,
        seq_lens_device: torch.Tensor,
        max_seq_len: int,
        query_start_loc_cpu: torch.Tensor,
        query_start_loc_device: torch.Tensor,
        num_decode_tokens: int,
        dcp_tot_seq_lens_device: torch.Tensor | None,
    ) -> AiterMLADecodeMetadata:
        device = self.device
        num_reqs = seq_lens_device.size(0)
        qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
        max_qo_len = qo_len.max().item()
        padded_mtp_qo_len = self._uniform_padded_mtp_qo_len(
            qo_len, max_qo_len, num_decode_tokens
        )
        if padded_mtp_qo_len > 0:
            max_qo_len = padded_mtp_qo_len
        pad_uniform_mtp = padded_mtp_qo_len > 0

        seq_lens_for_kernel = seq_lens_device
        num_kernel_reqs = num_reqs
        if pad_uniform_mtp:
            qo_lens_device = (
                query_start_loc_device[1 : num_reqs + 1]
                - query_start_loc_device[:num_reqs]
            ).to(torch.int32)
            seq_lens_for_kernel = torch.where(
                qo_lens_device > 0,
                seq_lens_for_kernel,
                seq_lens_for_kernel.new_full((), max_qo_len),
            )

        # The aiter kernel always operates with page_size=1 (the wrapper
        # flattens kv_buffer). last_page_len is always 1.
        paged_kv_last_page_len = self.paged_kv_last_page_len[:num_kernel_reqs]

        # indptr: cumsum of seq_lens (one page per token in the flat view)
        paged_kv_indptr = torch.cat(
            [
                torch.zeros(1, dtype=torch.int32, device=device),
                seq_lens_for_kernel.cumsum(dim=0, dtype=torch.int32),
            ]
        )
        use_gluon_decode = AiterMLAHelper.use_gluon_decode(
            self.num_heads, int(max_qo_len), self._kv_cache_dtype_str
        )

        if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
            self.paged_kv_indices.fill_(-1)

        # Expand block_table entries into per-token flat indices.
        # When kernel_block_size=1, this degrades to a direct copy (identical
        # to the original _copy_page_indices_kernel).
        # When kernel_block_size=K>1, block_table entry b covering K tokens
        # gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1).
        _expand_page_indices_kernel[(num_reqs,)](
            self.paged_kv_indices,
            block_table_tensor,
            block_table_tensor.stride(0),
            paged_kv_indptr,
            seq_lens_for_kernel,
            KERNEL_BLOCK_SIZE=self.kernel_block_size,
            BLOCK_SIZE=1024,
        )
        paged_kv_indices = self.paged_kv_indices

        if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
            self.paged_kv_indptr[: 1 + num_kernel_reqs].copy_(
                paged_kv_indptr, non_blocking=True
            )
            self.paged_kv_indptr[1 + num_kernel_reqs :].fill_(paged_kv_indptr[-1])
            paged_kv_indptr = self.paged_kv_indptr[: 1 + num_kernel_reqs]

            # paged_kv_last_page_len already uses the pre-initialized buffer slice
            # (set above), so no copy needed - buffer is always 1s.

            if pad_uniform_mtp:
                qo_indptr_src = torch.arange(
                    0,
                    (num_kernel_reqs + 1) * max_qo_len,
                    step=max_qo_len,
                    dtype=torch.int32,
                    device=device,
                )
            else:
                qo_indptr_src = query_start_loc_device[: 1 + num_kernel_reqs]
            self.qo_indptr[: 1 + num_kernel_reqs].copy_(
                qo_indptr_src, non_blocking=True
            )
            self.qo_indptr[1 + num_kernel_reqs :] = qo_indptr_src[-1]
            qo_indptr = self.qo_indptr[: 1 + num_kernel_reqs]

        else:
            if max_qo_len == 1:
                qo_indptr = torch.arange(
                    0,
                    num_kernel_reqs + 1,
                    step=1,
                    dtype=torch.int32,
                    device=device,
                )
            else:
                if pad_uniform_mtp:
                    qo_indptr = torch.arange(
                        0,
                        (num_kernel_reqs + 1) * max_qo_len,
                        step=max_qo_len,
                        dtype=torch.int32,
                        device=device,
                    )
                else:
                    qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs]

        # Only the asm decode consumes the schedule, so gate on the routing
        # rather than on num_heads >= 16, which denies it to a padded rank
        # running the same asm kernels. The two predicates are disjoint --
        # decode is qlen==1, verify is qlen>1 -- and cover both Gluon entries.
        has_persistent_metadata = False
        use_persistent_metadata = (
            not AiterMLAHelper.use_gluon_decode(
                self.num_heads, max_qo_len, self._kv_cache_dtype_str
            )
            and not AiterMLAHelper.use_gluon_verify(
                self.num_heads, max_qo_len, self._kv_cache_dtype_str
            )
            # A padded rank has no bf16 persistent kernel past qlen 4 where the
            # gfx950 fold is absent; the non-persistent entry covers it. fp8
            # keeps the schedule -- its fold rejects non-persistent outright.
            and (
                self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
                or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN
                or is_quantized_kv_cache(self._kv_cache_dtype_str)
            )
            and max_qo_len >= 1
            and max_qo_len <= self._mtp_decode_qlen
        )
        if use_persistent_metadata:
            from aiter import get_mla_metadata_v1

            uni_qo_len = (
                max_qo_len if pad_uniform_mtp or torch.all(qo_len == max_qo_len) else -1
            )
            get_mla_metadata_v1(
                qo_indptr,
                paged_kv_indptr,
                paged_kv_last_page_len,
                self._num_attention_heads,
                1,
                True,
                self._mla_work_meta_data,
                self._mla_work_info_set,
                self._mla_work_indptr,
                self._mla_reduce_indptr,
                self._mla_reduce_final_map,
                self._mla_reduce_partial_map,
                page_size=1,
                kv_granularity=16,
                max_seqlen_qo=max_qo_len,
                uni_seqlen_qo=uni_qo_len,
                fast_mode=True,
                dtype_q=self._mla_q_dtype,
                dtype_kv=self._mla_kv_dtype,
            )
            has_persistent_metadata = True

        attn_metadata = AiterMLADecodeMetadata(
            block_table=block_table_tensor,
            seq_lens=seq_lens_for_kernel,
            paged_kv_indptr=paged_kv_indptr,
            paged_kv_indices=paged_kv_indices,
            paged_kv_last_page_len=paged_kv_last_page_len,
            qo_indptr=qo_indptr,
            dcp_tot_seq_lens=dcp_tot_seq_lens_device,
            max_qo_len=max_qo_len,
            use_gluon_decode=use_gluon_decode,
            attn_out_dtype=self.decode_attn_out_dtype,
            has_persistent_metadata=has_persistent_metadata,
        )

        return attn_metadata

    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> AiterMLAMetadata:
        attn_metadata = super().build(
            common_prefix_len, common_attn_metadata, fast_build
        )
        if (
            attn_metadata.decode is not None
            and attn_metadata.decode.has_persistent_metadata
        ):
            attn_metadata.work_meta_data = self._mla_work_meta_data
            attn_metadata.work_indptr = self._mla_work_indptr
            attn_metadata.work_info_set = self._mla_work_info_set
            attn_metadata.reduce_indptr = self._mla_reduce_indptr
            attn_metadata.reduce_final_map = self._mla_reduce_final_map
            attn_metadata.reduce_partial_map = self._mla_reduce_partial_map
        if (
            self._fp8_prefill_enabled
            and attn_metadata.prefill is not None
            and attn_metadata.prefill.chunked_context is None
        ):
            self._build_fp8_prefill_ps_metadata(attn_metadata, common_attn_metadata)
        return attn_metadata