vllm.v1.attention.backends.mla.rocm_aiter_mla

Classes:

• AiterMLADCPVerifyMetadata –

  One paged-KV row per verify token, for segmented DCP verification.

• AiterMLAHelper –

  AITER MLA persistent (asm) decode requires a multiple of 16 heads. Unaligned

• AiterMLAImpl –
• AiterMLAMetadataBuilder –

AiterMLADCPVerifyMetadata dataclass

One paged-KV row per verify token, for segmented DCP verification.

These are produced together or not at all, so they travel as one value: its presence on the decode metadata is the routing decision the builder made, and the impl does not re-derive it.

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

@dataclass
class AiterMLADCPVerifyMetadata:
    """One paged-KV row per verify token, for segmented DCP verification.

    These are produced together or not at all, so they travel as one value: its
    presence on the decode metadata *is* the routing decision the builder made,
    and the impl does not re-derive it.
    """

    # Local KV length of each verify row, already causally bounded.
    row_lens: torch.Tensor
    # Subpage IDs each row reads, shape [num_rows, max_local_pages].
    block_table: torch.Tensor
    # One query per row, so this is an arange.
    qo_indptr: torch.Tensor
    # Subpage size the block table is expressed in, and the segmented kernel's
    # TILE_SIZE. Carried so the impl reinterprets the cache exactly the way the
    # block table was built, instead of re-deriving it from the cache shape.
    page_size: int
    # Kernel-visible KV bound. Whenever full graphs are enabled this is the
    # configuration's maximum for every batch, not just during capture, so the
    # page table keeps one shape across replays.
    max_kv_seq_len: int

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.

Methods:

• dcp_local_verify_row_lens –

  Local KV length of every verify row, in global-causal order.

• get_fp8_prefill_num_heads –

  Head count the FP8 PS asm prefill runs at: the next multiple of 16.

• has_fp8_non_causal_qlen2_kernel –

  Whether fp8 (num_heads, qlen=2) folds onto a non-causal ASM kernel.

• qo_indptr_for_uniform_qlen –

  Build [0, qlen, 2*qlen, ..., num_reqs*qlen].

• use_gluon_verify –

  Whether a small-head multi-token verify uses native Gluon MTP.

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

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, ...]] = ()
    # Pinned AITER v0.1.21.post2 folds these fp8 qlen-2 counts onto the
    # non-causal 16-head / 4-token kernel. Other multiples of 16 (48/80/112)
    # fold to H16 while keeping Q2, which has no non-causal fp8 entry.
    _AITER_FP8_NON_CAUSAL_QLEN2_HEADS: ClassVar[tuple[int, ...]] = (32, 64, 96, 128)

    @staticmethod
    def qo_indptr_for_uniform_qlen(
        num_reqs: int,
        qlen: int,
        device: torch.device,
        dtype: torch.dtype = torch.int32,
    ) -> torch.Tensor:
        """Build ``[0, qlen, 2*qlen, ..., num_reqs*qlen]``."""
        return torch.arange(
            0,
            (num_reqs + 1) * qlen,
            step=qlen,
            dtype=dtype,
            device=device,
        )

    @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 has_fp8_non_causal_qlen2_kernel(num_heads: int) -> bool:
        """Whether fp8 (num_heads, qlen=2) folds onto a non-causal ASM kernel.

        Pinned AITER v0.1.21.post2 folds 32/64/96/128 heads at qlen 2 onto the
        non-causal 16-head / 4-token kernel. A 16-head (or padded-to-16) qlen-2
        block has no matching entry, and padded 48/80/112 keep Q2 after the
        H16 fold.
        """
        kernel_heads = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
        return kernel_heads in AiterMLAHelper._AITER_FP8_NON_CAUSAL_QLEN2_HEADS

    @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_fp8_prefill_num_heads(num_heads: int) -> int:
        """Head count the FP8 PS asm prefill runs at: the next multiple of 16.

        Deliberately *not* ``get_actual_mla_num_heads``. That one carves out
        native H24 when ``_aiter_mla_native_h24_supported()``, which probes the
        asm *decode* reducer and metadata. The prefill is a different kernel
        pair (``mla_prefill_ps_asm_fwd`` + ``mla_reduce_v1``) with no such
        probe, so 24 heads pad to 32 here even on a native-H24 build.

        The PS metadata in ``_init_fp8_prefill_ps_buffers``/``build()`` must be
        sized with this same function, or the work/reduce maps describe a
        different head count than the kernel is handed.

        This function itself has no upper bound; the ceiling comes from the
        ``is_valid_num_heads`` gate, which rejects counts above
        ``_AITER_MAX_PADDED_MLA_HEADS`` (128) unless they are already
        16-aligned. That bound was established for the asm *decode* padding, so
        an architecture with, say, 136 heads per rank would be refused the
        prefill here for a reason that was never measured against this kernel
        pair. Revisit the constant rather than special-casing prefill.
        """
        m = AiterMLAHelper._AITER_MIN_MLA_HEADS
        return -(-num_heads // m) * m

    @staticmethod
    def get_mla_padded_q(
        num_heads: int, q: torch.Tensor, target_heads: int | None = None
    ) -> torch.Tensor:
        m = (
            target_heads
            if target_heads is not None
            else 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:
        return AiterMLAHelper._get_mla_unpadded_heads(num_heads, o)

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

    @staticmethod
    def get_mla_unpadded_lse(num_heads: int, lse: torch.Tensor) -> torch.Tensor:
        return AiterMLAHelper._get_mla_unpadded_heads(num_heads, lse)

    @staticmethod
    def use_gluon_decode(
        num_heads: int,
        max_qo_len: int,
        kv_cache_dtype: str,
        kv_cache_bytes: int | None = None,
    ) -> 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
        if not _gluon_mla_decode_supported():
            return False
        if mode != "gluon" and m % num_heads != 0:
            return False
        # Last, so the size warning only fires where Gluon would have run.
        return _gluon_kv_cache_in_bounds(kv_cache_bytes)

    @staticmethod
    def use_gluon_verify(
        num_heads: int,
        max_qo_len: int,
        kv_cache_dtype: str,
        dcp_world_size: int = 1,
        causal: bool = True,
        kv_cache_bytes: int | None = None,
    ) -> bool:
        """Whether a small-head multi-token verify uses native Gluon MTP.

        bf16 has no gqa<16, qseqlen>1 asm kernel, so verify goes through
        ``mla_gluon``'s 4-D MTP entry (``q`` shaped ``[batch, qlen, nhead, dim]``)
        with ``use_2d_view=False``. fp8 has one via the q-row fold and must not
        come here: the MTP path 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.

        DCP verify is excluded: its per-row causal windows are served by the
        segmented path, which Gluon's MTP entry cannot express.

        Gluon masks the block causally with no way to turn it off, so a
        non-causal block takes the padded asm decode whatever the head count
        -- padding is what gives it a kernel there.
        """
        if not causal:
            return False
        if max_qo_len <= 1 or dcp_world_size > 1:
            return False
        if is_quantized_kv_cache(kv_cache_dtype):
            return False
        if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS:
            return False
        if not _gluon_mla_decode_supported():
            return False
        # Same arch, mode and cache-size gating as use_gluon_decode.
        if _aiter_mla_small_head_mode() == "asm":
            return False
        return _gluon_kv_cache_in_bounds(kv_cache_bytes)

    @staticmethod
    def dcp_local_verify_row_lens(
        tot_seq_lens: torch.Tensor,
        qlen: int,
        dcp_world_size: int,
        dcp_rank: int,
        cp_interleave: int,
    ) -> torch.Tensor:
        """Local KV length of every verify row, in global-causal order.

        Row ``t`` of a qlen-token verify attends global positions
        ``[0, seq_len - qlen + t]``, its own token included, so its local length
        is the round-robin count evaluated at that bound. Truncating causally in
        global coordinates has to happen before the shard mapping: the tokens a
        row drops sit on ``qlen - 1 - t`` different ranks, while subtracting the
        offset from a request's local length takes one from every rank.
        """
        offsets = torch.arange(
            1,
            qlen + 1,
            device=tot_seq_lens.device,
            dtype=tot_seq_lens.dtype,
        )
        visible = (tot_seq_lens.unsqueeze(1) - qlen + offsets).clamp_(min=0)
        return get_dcp_local_seq_lens(
            visible,
            dcp_world_size,
            dcp_rank,
            cp_interleave,
        ).flatten()

dcp_local_verify_row_lens(tot_seq_lens, qlen, dcp_world_size, dcp_rank, cp_interleave) staticmethod

Local KV length of every verify row, in global-causal order.

Row t of a qlen-token verify attends global positions [0, seq_len - qlen + t], its own token included, so its local length is the round-robin count evaluated at that bound. Truncating causally in global coordinates has to happen before the shard mapping: the tokens a row drops sit on qlen - 1 - t different ranks, while subtracting the offset from a request's local length takes one from every rank.

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

@staticmethod
def dcp_local_verify_row_lens(
    tot_seq_lens: torch.Tensor,
    qlen: int,
    dcp_world_size: int,
    dcp_rank: int,
    cp_interleave: int,
) -> torch.Tensor:
    """Local KV length of every verify row, in global-causal order.

    Row ``t`` of a qlen-token verify attends global positions
    ``[0, seq_len - qlen + t]``, its own token included, so its local length
    is the round-robin count evaluated at that bound. Truncating causally in
    global coordinates has to happen before the shard mapping: the tokens a
    row drops sit on ``qlen - 1 - t`` different ranks, while subtracting the
    offset from a request's local length takes one from every rank.
    """
    offsets = torch.arange(
        1,
        qlen + 1,
        device=tot_seq_lens.device,
        dtype=tot_seq_lens.dtype,
    )
    visible = (tot_seq_lens.unsqueeze(1) - qlen + offsets).clamp_(min=0)
    return get_dcp_local_seq_lens(
        visible,
        dcp_world_size,
        dcp_rank,
        cp_interleave,
    ).flatten()

get_fp8_prefill_num_heads(num_heads) staticmethod

Head count the FP8 PS asm prefill runs at: the next multiple of 16.

Deliberately not get_actual_mla_num_heads. That one carves out native H24 when _aiter_mla_native_h24_supported(), which probes the asm decode reducer and metadata. The prefill is a different kernel pair (mla_prefill_ps_asm_fwd + mla_reduce_v1) with no such probe, so 24 heads pad to 32 here even on a native-H24 build.

The PS metadata in _init_fp8_prefill_ps_buffers/build() must be sized with this same function, or the work/reduce maps describe a different head count than the kernel is handed.

This function itself has no upper bound; the ceiling comes from the is_valid_num_heads gate, which rejects counts above _AITER_MAX_PADDED_MLA_HEADS (128) unless they are already 16-aligned. That bound was established for the asm decode padding, so an architecture with, say, 136 heads per rank would be refused the prefill here for a reason that was never measured against this kernel pair. Revisit the constant rather than special-casing prefill.

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

@staticmethod
def get_fp8_prefill_num_heads(num_heads: int) -> int:
    """Head count the FP8 PS asm prefill runs at: the next multiple of 16.

    Deliberately *not* ``get_actual_mla_num_heads``. That one carves out
    native H24 when ``_aiter_mla_native_h24_supported()``, which probes the
    asm *decode* reducer and metadata. The prefill is a different kernel
    pair (``mla_prefill_ps_asm_fwd`` + ``mla_reduce_v1``) with no such
    probe, so 24 heads pad to 32 here even on a native-H24 build.

    The PS metadata in ``_init_fp8_prefill_ps_buffers``/``build()`` must be
    sized with this same function, or the work/reduce maps describe a
    different head count than the kernel is handed.

    This function itself has no upper bound; the ceiling comes from the
    ``is_valid_num_heads`` gate, which rejects counts above
    ``_AITER_MAX_PADDED_MLA_HEADS`` (128) unless they are already
    16-aligned. That bound was established for the asm *decode* padding, so
    an architecture with, say, 136 heads per rank would be refused the
    prefill here for a reason that was never measured against this kernel
    pair. Revisit the constant rather than special-casing prefill.
    """
    m = AiterMLAHelper._AITER_MIN_MLA_HEADS
    return -(-num_heads // m) * m

has_fp8_non_causal_qlen2_kernel(num_heads) staticmethod

Whether fp8 (num_heads, qlen=2) folds onto a non-causal ASM kernel.

Pinned AITER v0.1.21.post2 folds 32/64/96/128 heads at qlen 2 onto the non-causal 16-head / 4-token kernel. A 16-head (or padded-to-16) qlen-2 block has no matching entry, and padded 48/80/112 keep Q2 after the H16 fold.

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

@staticmethod
def has_fp8_non_causal_qlen2_kernel(num_heads: int) -> bool:
    """Whether fp8 (num_heads, qlen=2) folds onto a non-causal ASM kernel.

    Pinned AITER v0.1.21.post2 folds 32/64/96/128 heads at qlen 2 onto the
    non-causal 16-head / 4-token kernel. A 16-head (or padded-to-16) qlen-2
    block has no matching entry, and padded 48/80/112 keep Q2 after the
    H16 fold.
    """
    kernel_heads = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
    return kernel_heads in AiterMLAHelper._AITER_FP8_NON_CAUSAL_QLEN2_HEADS

qo_indptr_for_uniform_qlen(num_reqs, qlen, device, dtype=torch.int32) staticmethod

Build [0, qlen, 2*qlen, ..., num_reqs*qlen].

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

@staticmethod
def qo_indptr_for_uniform_qlen(
    num_reqs: int,
    qlen: int,
    device: torch.device,
    dtype: torch.dtype = torch.int32,
) -> torch.Tensor:
    """Build ``[0, qlen, 2*qlen, ..., num_reqs*qlen]``."""
    return torch.arange(
        0,
        (num_reqs + 1) * qlen,
        step=qlen,
        dtype=dtype,
        device=device,
    )

use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype, dcp_world_size=1, causal=True, kv_cache_bytes=None) staticmethod

Whether a small-head multi-token verify uses native Gluon MTP.

bf16 has no gqa<16, qseqlen>1 asm kernel, so verify goes through mla_gluon's 4-D MTP entry (q shaped [batch, qlen, nhead, dim]) with use_2d_view=False. fp8 has one via the q-row fold and must not come here: the MTP path 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.

DCP verify is excluded: its per-row causal windows are served by the segmented path, which Gluon's MTP entry cannot express.

Gluon masks the block causally with no way to turn it off, so a non-causal block takes the padded asm decode whatever the head count -- padding is what gives it a kernel there.

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

@staticmethod
def use_gluon_verify(
    num_heads: int,
    max_qo_len: int,
    kv_cache_dtype: str,
    dcp_world_size: int = 1,
    causal: bool = True,
    kv_cache_bytes: int | None = None,
) -> bool:
    """Whether a small-head multi-token verify uses native Gluon MTP.

    bf16 has no gqa<16, qseqlen>1 asm kernel, so verify goes through
    ``mla_gluon``'s 4-D MTP entry (``q`` shaped ``[batch, qlen, nhead, dim]``)
    with ``use_2d_view=False``. fp8 has one via the q-row fold and must not
    come here: the MTP path 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.

    DCP verify is excluded: its per-row causal windows are served by the
    segmented path, which Gluon's MTP entry cannot express.

    Gluon masks the block causally with no way to turn it off, so a
    non-causal block takes the padded asm decode whatever the head count
    -- padding is what gives it a kernel there.
    """
    if not causal:
        return False
    if max_qo_len <= 1 or dcp_world_size > 1:
        return False
    if is_quantized_kv_cache(kv_cache_dtype):
        return False
    if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS:
        return False
    if not _gluon_mla_decode_supported():
        return False
    # Same arch, mode and cache-size gating as use_gluon_decode.
    if _aiter_mla_small_head_mode() == "asm":
        return False
    return _gluon_kv_cache_in_bounds(kv_cache_bytes)

AiterMLAImpl

Bases: MLACommonImpl[AiterMLAMetadata]

Methods:

• forward_mha –

  Dispatch prefill to the FP8 ASM kernel when available.

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

class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
    # DCP decode paths return natural-log softmax LSE for the cross-rank merge.
    can_return_lse_for_decode: bool = True
    supports_dcp: bool = True
    # Measured on gfx950: aiter.mla.mla_decode_fwd(return_lse=True) matches
    # logsumexp to fp32 exactly, and merge_mla_segments_triton converts AITER's
    # base-2 segment statistics with LOGE2. Stated rather than inherited because
    # the DCP combine silently corrupts the softmax denominator if it disagrees.
    lse_base_on_e: bool = True

    @property
    def _decode_num_heads(self) -> int:
        """Return the query-head count after DCP gathering."""
        return self.num_heads * self.dcp_world_size

    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)
        AiterMLAHelper.check_num_heads_validity(self._decode_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. Only runs when the
        # KV cache is FP8. Head counts that are not a multiple of 16 are
        # replicate-padded up to one (see _mla_fp8_prefill_attn).
        from vllm.utils.torch_utils import is_quantized_kv_cache

        self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and (
            is_quantized_kv_cache(kv_cache_dtype)
            and AiterMLAHelper.is_valid_num_heads(self.num_heads)
        )
        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]
        # PS asm prefill + mla_reduce_v1 require 16-aligned heads, and the PS
        # metadata is built for get_fp8_prefill_num_heads(num_heads). For head
        # counts that are not a multiple of 16 (K3 = 12/rank at TP8)
        # replicate-pad q/k/v up to that count, then slice the output back to
        # the real head count.
        #
        # Counts above 16 that are not multiples of 16 (24, 40, ...) are
        # handled by the same code but are not reached by any current model
        # and TP that fits: 96 heads would need TP4, whose weights exceed a
        # 288 GiB GPU, and 128-head models land on 128/64/32/16/8. The path is
        # kept general for future architectures and its numerics are covered by
        # test_fp8_prefill_matches_reference[num_heads=24]. Note the cost is
        # (padded - real)/real extra FLOPs and q/k/v bytes, which is worst just
        # above a multiple of 16 (17 heads pad to 32); a future arch landing
        # there should measure against the flash_attn_varlen_func fallback
        # rather than assume the asm path wins.
        #
        # Exact, not approximate: after kv_b_proj gqa_ratio is 1, so q, k and v
        # all carry num_heads heads and attention is independent per head.
        # Padding all three identically makes padded head j a duplicate of real
        # head j % num_heads, so the real heads [0:num_heads] are bit-identical
        # to the unpadded result. Same argument as the decode path; only the
        # target width differs.
        _real_nhead = self.num_heads
        nhead = AiterMLAHelper.get_fp8_prefill_num_heads(_real_nhead)
        _pad = nhead != _real_nhead
        if _pad:
            q = AiterMLAHelper.get_mla_padded_q(_real_nhead, q, nhead)
            k = AiterMLAHelper.get_mla_padded_q(_real_nhead, k, nhead)
            v = AiterMLAHelper.get_mla_padded_q(_real_nhead, v, nhead)
        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

        # Per-call scratch is served from the workspace manager so allocator
        # churn in the prefill hot path is bounded after warmup, matching the
        # pattern in PR #41002.  The builder reserves the maximum shape of every
        # tensor requested here before the workspace is locked.
        scratch: list[tuple[tuple[int, ...], torch.dtype]] = [
            ((num_partial_tiles * tile_q, nhead, v_head_dim), torch.float32),
            ((num_partial_tiles * tile_q, nhead), torch.float32),
            ((total_q, nhead), torch.float32),
        ]
        if _pad:
            # The ASM and reduce kernels write a [total_q, nhead, v_head_dim]
            # buffer.  With unpadded heads that aliases the caller's
            # [total_q, nhead * v_head_dim] output, so write straight into it;
            # padded heads do not fit that storage and need their own buffer.
            scratch.append(((total_q, nhead, v_head_dim), out.dtype))

        workspace = current_workspace_manager()
        logits, attn_lse, final_lse, *pad_out = workspace.get_simultaneous(*scratch)
        out_3d = pad_out[0] if _pad else out.view(total_q, nhead, v_head_dim)

        # 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,
        )

        if _pad:
            out.view(total_q, _real_nhead, v_head_dim).copy_(out_3d[:, :_real_nhead, :])

    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_segmented_dcp_verify(
        self,
        q_nope: torch.Tensor,
        q_pe: torch.Tensor,
        verify: AiterMLADCPVerifyMetadata,
        kv_c_and_k_pe_cache: torch.Tensor,
        layer: AttentionLayer,
        out_dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Run segmented attention over this rank's shard of every verify row.

        Each row's length already covers the current tokens this rank holds, so
        the cross-rank LSE merge in the MLA layer completes the causal block.
        """
        q_mla = torch.cat([q_nope, q_pe], dim=-1)
        # skip_reduce with NUM_SEGMENTS>1 returns the partials and never writes
        # `out`. Pass None rather than aliasing q_mla: if that ever stops
        # holding, this fails loudly instead of scribbling over the query.
        segment_partials = _get_segmented_mla_decode()(
            q_mla,
            kv_c_and_k_pe_cache.view(
                -1,
                verify.page_size,
                1,
                kv_c_and_k_pe_cache.shape[-1],
            ),
            None,
            verify.qo_indptr,
            verify.row_lens,
            verify.max_kv_seq_len,
            verify.block_table,
            self.scale,
            self.kv_lora_rank,
            self.qk_rope_head_dim,
            causal=True,
            q_descale=None,
            kv_descale=layer._k_scale,
            skip_reduce=True,
        )
        assert isinstance(segment_partials, tuple) and len(segment_partials) == 3, (
            "AITER segmented MLA verify must return segment partials "
            "when skip_reduce=True."
        )
        segm_output, segm_max, segm_expsum = segment_partials
        return merge_mla_segments_triton(
            segm_output,
            segm_max,
            segm_expsum,
            verify.row_lens,
            verify.page_size,
            out_dtype,
        )

    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
        if decode.use_gluon_decode:
            assert decode.paged_kv_indptr is not None
            assert decode.paged_kv_indices is not None
            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, so the block goes to the gluon kernel's 4-D
        # MTP entry, which serves a whole (1 + num_spec) block in one launch.
        # The block is causal -- the target is checking draft tokens, so
        # position t must not see t+1 -- and the kernel bounds each query
        # position's scores itself.
        # Arch, mode and dtype gating all live in use_gluon_verify, which the
        # builder evaluates and records on the metadata, so this branch acts on
        # the same answer the builder used to decide whether the asm decode
        # would run. DCP verify is routed to the segmented path below instead.
        if decode.use_gluon_verify:
            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])
            assert attn_metadata.causal, (
                "AITER MLA small-head verify MTP is causal-only"
            )
            # Hand mla_gluon its 4-D MTP entry instead of an expanded
            # per-verify-token paged-KV view. The flat query layout is already
            # row-major (r * qlen + t), so unflatten is a free view. mla_gluon
            # applies the per-position causal bound
            # score_end = min(split_kv_end, seq_len - qlen + q_pos + 1)
            # in-kernel, and seq_lens already spans the verify block, so that
            # bound is context_r + q_pos + 1 -- the same window the expanded
            # view supplied by truncating each row's page list.
            qlen = int(decode.max_qo_len)
            num_reqs = B // qlen
            if num_reqs * qlen != B:
                raise ValueError(
                    f"verify block {B} rows is not a multiple of qlen {qlen}"
                )
            assert decode.paged_kv_indptr is not None
            assert decode.paged_kv_indices is not None
            mla_gluon = _get_mla_gluon()
            mla_gluon(
                q_nope=q_nope.unflatten(0, (num_reqs, qlen)),
                q_pe=q_pe.unflatten(0, (num_reqs, qlen)),
                kv_c=kv_buffer,
                o=o.unflatten(0, (num_reqs, qlen)),
                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

        dcp_route = decode.dcp_route
        if dcp_route is _DCPDecodeRoute.SEGMENTED:
            verify = decode.dcp_verify
            assert verify is not None
            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
                )
            if (
                is_quantized_kv_cache(self.kv_cache_dtype)
                and q_nope.dtype != torch.bfloat16
            ):
                q_nope = q_nope.to(torch.bfloat16) * layer._q_scale
                q_pe = q_pe.to(torch.bfloat16) * layer._q_scale
            return self._forward_segmented_dcp_verify(
                q_nope,
                q_pe,
                verify,
                kv_c_and_k_pe_cache,
                layer,
                decode.attn_out_dtype,
            )

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

        assert isinstance(q, torch.Tensor)
        assert decode.paged_kv_indptr is not None
        assert decode.paged_kv_indices is not None
        B = q.shape[0]

        assert q.shape[1] == self._decode_num_heads, (
            "ROCM_AITER_MLA decode expected the DCP-gathered query head count "
            f"{self._decode_num_heads}, got {q.shape[1]}"
        )
        # cprr has kernels only at native head counts, so pad q up (e.g. 96 ->
        # 128). Heads are independent in MLA; the extras are sliced off below.
        asm_dcp_heads = decode.asm_decode_num_heads
        if asm_dcp_heads > self._decode_num_heads:
            mla_num_heads = asm_dcp_heads
            mla_padded_q = AiterMLAHelper.get_mla_padded_q(
                self._decode_num_heads, q, target_heads=mla_num_heads
            )
        else:
            mla_padded_q = AiterMLAHelper.get_mla_padded_q(self._decode_num_heads, q)
            mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(
                self._decode_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,
            )

        lse = None
        if self.dcp_world_size > 1:
            if dcp_route is _DCPDecodeRoute.CPRR:
                # cprr only: global positions for the in-kernel causal window.
                assert decode.g_kv_indptr is not None
                mla_kwargs["g_kv_indptr"] = decode.g_kv_indptr
                mla_kwargs["cp_world_size"] = decode.cp_world_size
                mla_kwargs["cp_rank"] = decode.cp_rank
                # Must match the schedule get_mla_metadata_v1 was built with.
                assert decode.mla_num_kv_splits > 0
                mla_kwargs["num_kv_splits"] = decode.mla_num_kv_splits
                mla_kwargs["intra_batch_mode"] = False
            # The vLLM custom-op wrapper exposes only the in-place output and
            # drops aiter's final LSE, which the cross-shard merge needs, so go
            # through aiter's native entry point on the DCP path.
            _, lse = _get_aiter_mla_decode()(
                mla_padded_q,
                kv_buffer.view(-1, 1, 1, mla_padded_q.shape[-1]),
                o,
                decode.qo_indptr,
                decode.paged_kv_indptr,
                decode.paged_kv_indices,
                decode.paged_kv_last_page_len,
                decode.max_qo_len,
                sm_scale=self.scale,
                return_lse=True,
                causal=attn_metadata.causal,
                **mla_kwargs,
            )
            assert lse is not None, (
                "aiter mla_decode_fwd(return_lse=True) returned no LSE; upgrade "
                "aiter to a build with decode LSE support."
            )
        else:
            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,
                causal=attn_metadata.causal,
                **mla_kwargs,
            )

        if asm_dcp_heads > self._decode_num_heads:
            # Keep a view: a post-kernel copy races the symm-mem a2a combine.
            output = o[:, : self._decode_num_heads]
            if lse is not None:
                lse = lse[:, : self._decode_num_heads, ...]
        else:
            output = AiterMLAHelper.get_mla_unpadded_o(self._decode_num_heads, o)
            if lse is not None:
                lse = AiterMLAHelper.get_mla_unpadded_lse(self._decode_num_heads, lse)
        return output, lse

_decode_num_heads property

Return the query-head count after DCP gathering.

_forward_segmented_dcp_verify(q_nope, q_pe, verify, kv_c_and_k_pe_cache, layer, out_dtype)

Run segmented attention over this rank's shard of every verify row.

Each row's length already covers the current tokens this rank holds, so the cross-rank LSE merge in the MLA layer completes the causal block.

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

def _forward_segmented_dcp_verify(
    self,
    q_nope: torch.Tensor,
    q_pe: torch.Tensor,
    verify: AiterMLADCPVerifyMetadata,
    kv_c_and_k_pe_cache: torch.Tensor,
    layer: AttentionLayer,
    out_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Run segmented attention over this rank's shard of every verify row.

    Each row's length already covers the current tokens this rank holds, so
    the cross-rank LSE merge in the MLA layer completes the causal block.
    """
    q_mla = torch.cat([q_nope, q_pe], dim=-1)
    # skip_reduce with NUM_SEGMENTS>1 returns the partials and never writes
    # `out`. Pass None rather than aliasing q_mla: if that ever stops
    # holding, this fails loudly instead of scribbling over the query.
    segment_partials = _get_segmented_mla_decode()(
        q_mla,
        kv_c_and_k_pe_cache.view(
            -1,
            verify.page_size,
            1,
            kv_c_and_k_pe_cache.shape[-1],
        ),
        None,
        verify.qo_indptr,
        verify.row_lens,
        verify.max_kv_seq_len,
        verify.block_table,
        self.scale,
        self.kv_lora_rank,
        self.qk_rope_head_dim,
        causal=True,
        q_descale=None,
        kv_descale=layer._k_scale,
        skip_reduce=True,
    )
    assert isinstance(segment_partials, tuple) and len(segment_partials) == 3, (
        "AITER segmented MLA verify must return segment partials "
        "when skip_reduce=True."
    )
    segm_output, segm_max, segm_expsum = segment_partials
    return merge_mla_segments_triton(
        segm_output,
        segm_max,
        segm_expsum,
        verify.row_lens,
        verify.page_size,
        out_dtype,
    )

_mla_fp8_prefill_attn(q, k, v, attn_metadata, out)

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.

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

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]
    # PS asm prefill + mla_reduce_v1 require 16-aligned heads, and the PS
    # metadata is built for get_fp8_prefill_num_heads(num_heads). For head
    # counts that are not a multiple of 16 (K3 = 12/rank at TP8)
    # replicate-pad q/k/v up to that count, then slice the output back to
    # the real head count.
    #
    # Counts above 16 that are not multiples of 16 (24, 40, ...) are
    # handled by the same code but are not reached by any current model
    # and TP that fits: 96 heads would need TP4, whose weights exceed a
    # 288 GiB GPU, and 128-head models land on 128/64/32/16/8. The path is
    # kept general for future architectures and its numerics are covered by
    # test_fp8_prefill_matches_reference[num_heads=24]. Note the cost is
    # (padded - real)/real extra FLOPs and q/k/v bytes, which is worst just
    # above a multiple of 16 (17 heads pad to 32); a future arch landing
    # there should measure against the flash_attn_varlen_func fallback
    # rather than assume the asm path wins.
    #
    # Exact, not approximate: after kv_b_proj gqa_ratio is 1, so q, k and v
    # all carry num_heads heads and attention is independent per head.
    # Padding all three identically makes padded head j a duplicate of real
    # head j % num_heads, so the real heads [0:num_heads] are bit-identical
    # to the unpadded result. Same argument as the decode path; only the
    # target width differs.
    _real_nhead = self.num_heads
    nhead = AiterMLAHelper.get_fp8_prefill_num_heads(_real_nhead)
    _pad = nhead != _real_nhead
    if _pad:
        q = AiterMLAHelper.get_mla_padded_q(_real_nhead, q, nhead)
        k = AiterMLAHelper.get_mla_padded_q(_real_nhead, k, nhead)
        v = AiterMLAHelper.get_mla_padded_q(_real_nhead, v, nhead)
    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

    # Per-call scratch is served from the workspace manager so allocator
    # churn in the prefill hot path is bounded after warmup, matching the
    # pattern in PR #41002.  The builder reserves the maximum shape of every
    # tensor requested here before the workspace is locked.
    scratch: list[tuple[tuple[int, ...], torch.dtype]] = [
        ((num_partial_tiles * tile_q, nhead, v_head_dim), torch.float32),
        ((num_partial_tiles * tile_q, nhead), torch.float32),
        ((total_q, nhead), torch.float32),
    ]
    if _pad:
        # The ASM and reduce kernels write a [total_q, nhead, v_head_dim]
        # buffer.  With unpadded heads that aliases the caller's
        # [total_q, nhead * v_head_dim] output, so write straight into it;
        # padded heads do not fit that storage and need their own buffer.
        scratch.append(((total_q, nhead, v_head_dim), out.dtype))

    workspace = current_workspace_manager()
    logits, attn_lse, final_lse, *pad_out = workspace.get_simultaneous(*scratch)
    out_3d = pad_out[0] if _pad else out.view(total_q, nhead, v_head_dim)

    # 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,
    )

    if _pad:
        out.view(total_q, _real_nhead, v_head_dim).copy_(out_3d[:, :_real_nhead, :])

forward_mha(q, kv_c_normed, k_pe, kv_c_and_k_pe_cache, attn_metadata, k_scale, output, output_scale=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.

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

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)

AiterMLAMetadataBuilder

Bases: MLACommonMetadataBuilder[AiterMLAMetadata]

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

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
    # Served by passing the mask to the kernel; _build_decode turns away the
    # shapes AITER has no non-causal kernel for.
    supports_non_causal_multi_token_decode: ClassVar[bool] = True
    # Set from the common metadata every build; a batch is causal unless the
    # drafter says otherwise.
    _decode_causal: bool = True
    # A non-causal draft needs no cross-shard causal window: the ordinary
    # mask0 decode returns per-row LSE for the DCP merge.
    supports_non_causal_multi_token_dcp: ClassVar[bool] = True

    @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,
    ):
        parallel_config = vllm_config.parallel_config
        supports_segmented_dcp_verify = _segmented_dcp_verify_supported(
            parallel_config.decode_context_parallel_size,
            parallel_config.cp_kv_cache_interleave_size,
        )
        asm_dcp_verify_config = (
            _asm_dcp_verify_configured(
                parallel_config.decode_context_parallel_size,
                parallel_config.cp_kv_cache_interleave_size,
                multi_token_decode=vllm_config.speculative_config is not None,
            )
            and envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "asm"
        )
        super().__init__(
            kv_cache_spec,
            layer_names,
            vllm_config,
            device,
            AiterMLAMetadata,
            supports_dcp_with_varlen=(
                supports_segmented_dcp_verify or asm_dcp_verify_config
            ),
        )
        self._asm_dcp_verify = False
        self._asm_dcp_verify_heads = 0
        # The head-count half of the route decision needs self.num_heads, so only
        # the config half is known before super().__init__().
        if asm_dcp_verify_config:
            self._asm_dcp_verify_heads = _asm_dcp_verify_heads(
                self.num_heads * self.dcp_world_size
            )
            if not self._asm_dcp_verify_heads:
                raise ValueError(
                    "VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm, but the round-robin asm "
                    f"decode has no kernel for {self.num_heads * self.dcp_world_size} "
                    f"DCP-gathered heads (native counts: {_NATIVE_CPRR_HEADS}). "
                    "Set VLLM_ROCM_AITER_MLA_DCP_VERIFY=segmented."
                )
            self._asm_dcp_verify = True
        self._supports_segmented_dcp_verify = supports_segmented_dcp_verify
        self._mla_max_split_per_batch = 0
        if self._asm_dcp_verify:
            # Cap on KV splits per batch; at low batch the KV axis is the only
            # parallelism. Buffer sizing and every runtime call must see this value.
            self._mla_max_split_per_batch = torch.cuda.get_device_properties(
                device
            ).multi_processor_count

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

        # Needed to place a verify row's causal window on this rank's KV shard.
        self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0

        # 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
        self._segmented_page_size = _segmented_mla_page_size(self.kernel_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

        # Decode kernels consume the DCP-gathered query heads.
        self._decode_num_heads = self.num_heads * self.dcp_world_size
        # Keep metadata sizing consistent with the padded tensor shape passed
        # to mla_decode_fwd, including native 24-head AITER builds.
        self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads(
            self._decode_num_heads
        )
        # the cprr kernel exists only at the native head
        # counts, so run (and size the persistent schedule) at the padded count.
        # 96 is 16-aligned, so get_actual_mla_num_heads leaves it alone and the
        # pad has to be applied here explicitly.
        if self._asm_dcp_verify:
            self._num_attention_heads = self._asm_dcp_verify_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
        # Sized per layer, matching the flat KV tensor the kernel is handed.
        num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks
        self._kv_cache_bytes = (
            None
            if num_gpu_blocks is None
            else num_gpu_blocks * kv_cache_spec.page_size_bytes
        )
        # 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,
            # Must match max_split_per_batch in get_mla_metadata_v1; without it
            # aiter sizes the reduce scratch from capacity (~9 GiB).
            **(
                dict(max_split_per_batch=self._mla_max_split_per_batch)
                if self._asm_dcp_verify
                else {}
            ),
        )
        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,
        )

        # The assembly prefill requires FP8 KV, bf16 output, and 16-aligned
        # heads. It writes bf16 through a raw output pointer, so fp16 must use
        # the standard prefill path. Head counts that are not a multiple of 16
        # are replicate-padded up to one in _mla_fp8_prefill_attn, so the gate
        # is the same head-count predicate the decode path uses.
        self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and (
            kv_cache_dtype_str == "fp8"
            and vllm_config.model_config.dtype == torch.bfloat16
            and AiterMLAHelper.is_valid_num_heads(self.num_heads)
        )
        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,
                vllm_config.model_config.dtype,
                device,
            )

        # Full cudagraphs need a stable address for the cprr global page indptr.
        # Zero-initialized; element 0 is the indptr base and is never rewritten.
        self._g_kv_indptr_buf: torch.Tensor | None = None
        if self._asm_dcp_verify:
            self._g_kv_indptr_buf = torch.zeros(
                max_num_reqs + 1, dtype=torch.int32, device=device
            )

        # Persistent buffers for segmented DCP verification. Captured graphs
        # require stable addresses while row lengths vary between replays.
        self._dcp_verify_buffers: AiterMLADCPVerifyMetadata | None = None
        self._graph_seq_lens: torch.Tensor | None = None
        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
            )

            # Full graphs require a stable address after uniform-MTP padding.
            self._graph_seq_lens = torch.zeros(
                max_num_reqs, dtype=torch.int32, device=device
            )

            if self._supports_segmented_dcp_verify and self._mtp_decode_qlen > 1:
                # Allocate even when CPRR is the preferred route: a later
                # replay can fall below _MIN_CPRR_QLEN and needs these
                # addresses. A DCP rank's shard of the longest sequence
                # bounds every verify row, and full graphs need that bound
                # to be constant.
                num_dcp_partitions = (
                    self.dcp_world_size * self.cp_kv_cache_interleave_size
                )
                graph_max_kv_seq_len = (
                    cdiv(vllm_config.model_config.max_model_len, num_dcp_partitions)
                    * self.cp_kv_cache_interleave_size
                )
                max_verify_rows = max_num_reqs * self._mtp_decode_qlen
                max_local_pages = cdiv(graph_max_kv_seq_len, self._segmented_page_size)
                self._dcp_verify_buffers = AiterMLADCPVerifyMetadata(
                    row_lens=torch.zeros(
                        max_verify_rows, dtype=torch.int32, device=device
                    ),
                    block_table=torch.zeros(
                        (max_verify_rows, max_local_pages),
                        dtype=torch.int32,
                        device=device,
                    ),
                    qo_indptr=torch.arange(
                        max_verify_rows + 1, dtype=torch.int32, device=device
                    ),
                    page_size=self._segmented_page_size,
                    max_kv_seq_len=graph_max_kv_seq_len,
                )

    def _init_fp8_prefill_ps_buffers(
        self,
        max_num_reqs: int,
        max_prefill_qlen: int,
        max_num_batched_tokens: int,
        attn_out_dtype: torch.dtype,
        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``.
            attn_out_dtype: Dtype of the attention output buffer, used to size
                the padded-head output scratch (small head counts only).
            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.
        # Head counts that are not a multiple of 16 (K3: 12/rank at TP8) are
        # replicate-padded up to one in _mla_fp8_prefill_attn; build the PS
        # metadata for that same padded count so the work/reduce maps and the
        # scratch reservations describe the width the kernel is handed.
        #
        # This was previously max(16, num_heads), which agrees with the helper
        # at every head count reachable today (12 and 16 both give 16) and
        # differs only above 16: at 24 it leaves num_head_k=24 while the
        # forward pads to 32. That is not a correctness bug -- the 24-wide work
        # maps still cover head-tiles 0..23, which are the real heads -- but it
        # is expensive, because a lower head alignment yields more partial
        # tiles: gcd-driven, 24 heads -> 64 tiles vs 32 heads -> 16, i.e.
        # 193.6 MiB of reservations instead of 68.6 MiB at batch=1/qlen=512
        # (measured on gfx950). Sizing both from one helper keeps the widths
        # equal and takes the cheaper tiling.
        num_head_k = AiterMLAHelper.get_fp8_prefill_num_heads(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,
        )

        # get_ps_metadata_v1 builds the plan on the host and writes its outputs with
        # blocking hipMemcpy that is not ordered on the current stream. Writing the
        # device buffers directly lets step N+1's build (async scheduling) overwrite a
        # plan that step N's prefill kernels are still reading. Build into pinned host
        # staging instead and copy on the current stream. Two staging slots, each
        # reused only after its previous H2D copy has completed.
        self._fp8_ps_device_outputs = (
            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,
        )
        self._fp8_ps_staging = [
            tuple(
                torch.empty(t.shape, dtype=t.dtype, device="cpu", pin_memory=True)
                for t in self._fp8_ps_device_outputs
            )
            for _ in range(2)
        ]
        # One event per slot, recorded after that slot's H2D copy. synchronize()
        # on an event that has not been recorded yet returns immediately.
        self._fp8_ps_staging_free = [torch.cuda.Event(), torch.cuda.Event()]
        self._fp8_ps_slot = 0

        from vllm.platforms import current_platform
        from vllm.v1.worker.workspace import current_workspace_manager

        # AITER metadata sizing assumes all requests can carry max_prefill_qlen tokens,
        # which is a loose worst case in the number of QO tiles. Rather, the sum of
        # qlens is bounded by max_num_batched_tokens. Manually compute the number of
        # qo tiles to avoid OOM at startup.
        # TODO: AITER should give us this budget constrained value
        qo_tile_cnt = (
            cdiv(max_num_batched_tokens, _FP8_PREFILL_TILE_Q) + max_num_reqs - 1
        )
        max_num_partial_tiles = qo_tile_cnt + current_platform.num_compute_units()
        reservations: list[tuple[tuple[int, ...], torch.dtype]] = [
            (
                (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),
        ]
        if self.num_heads < num_head_k:
            # Padded head counts also take their kernel output buffer from the
            # workspace: the caller's [total_q, num_heads * v_head_dim] output
            # cannot back a num_head_k-head view (see _mla_fp8_prefill_attn).
            reservations.append(
                ((max_num_batched_tokens, num_head_k, v_head_dim), attn_out_dtype)
            )
        current_workspace_manager().get_simultaneous(*reservations)

        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)

        # Head counts that are not a multiple of 16 (K3: 12/rank at TP8) are
        # replicate-padded up to one in _mla_fp8_prefill_attn; build the PS
        # metadata for that same padded count so the work/reduce maps and the
        # scratch reservations describe the width the kernel is handed.
        #
        # This was previously max(16, num_heads), which agrees with the helper
        # at every head count reachable today (12 and 16 both give 16) and
        # differs only above 16: at 24 it leaves num_head_k=24 while the
        # forward pads to 32. That is not a correctness bug -- the 24-wide work
        # maps still cover head-tiles 0..23, which are the real heads -- but it
        # is expensive, because a lower head alignment yields more partial
        # tiles: gcd-driven, 24 heads -> 64 tiles vs 32 heads -> 16, i.e.
        # 193.6 MiB of reservations instead of 68.6 MiB at batch=1/qlen=512
        # (measured on gfx950). Sizing both from one helper keeps the widths
        # equal and takes the cheaper tiling.
        num_head_k = AiterMLAHelper.get_fp8_prefill_num_heads(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

        slot = self._fp8_ps_slot
        self._fp8_ps_slot ^= 1
        staging_free = self._fp8_ps_staging_free[slot]
        with gpu_sync_allowed():
            staging_free.synchronize()
        (
            work_indptr_host,
            work_info_host,
            reduce_indptr_host,
            reduce_final_map_host,
            reduce_partial_map_host,
        ) = self._fp8_ps_staging[slot]

        # work_metadata is not written by the host planner; it stays on device.
        get_ps_metadata_v1(
            qo_indptr_cpu,
            kv_indptr_cpu,
            seq_lens_cpu,
            gqa_ratio,
            num_head_k,
            self.fp8_ps_work_metadata,
            work_indptr_host,
            work_info_host,
            reduce_indptr_host,
            reduce_final_map_host,
            reduce_partial_map_host,
            qhead_granularity=qhead_granularity,
            qlen_granularity=qlen_granularity,
            kvlen_granularity=kvlen_granularity,
            block_size=block_size,
            is_causal=True,
        )
        for dst, src in zip(self._fp8_ps_device_outputs, self._fp8_ps_staging[slot]):
            dst.copy_(src, non_blocking=True)
        staging_free.record()

        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, read from the host plan (no GPU sync).
        num_partial_tiles = int(reduce_indptr_host[-1])

        # 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_dcp_verify_row_view(
        self,
        qlen: int,
        block_table: torch.Tensor,
        dcp_tot_seq_lens: torch.Tensor,
    ) -> AiterMLADCPVerifyMetadata:
        """Build one paged-KV row per verify token for segmented DCP verification.

        Every row is a single query (``qo_indptr`` is an arange), so the whole
        causal structure is carried by ``dcp_local_verify_row_lens`` and the
        kernel applies no tail of its own.
        """
        assert self.dcp_world_size > 1
        num_reqs = dcp_tot_seq_lens.numel()
        row_lens = AiterMLAHelper.dcp_local_verify_row_lens(
            dcp_tot_seq_lens,
            qlen,
            self.dcp_world_size,
            self.dcp_rank,
            self.cp_kv_cache_interleave_size,
        )
        num_rows = row_lens.numel()
        buffers = self._dcp_verify_buffers
        # Persistent buffers exist exactly when full graphs are on, and they
        # already carry the static bound they were sized for.
        max_kv_seq_len = (
            buffers.max_kv_seq_len
            if buffers is not None
            else max(1, int(row_lens.max().item()))
        )
        page_size = self._segmented_page_size
        max_local_pages = cdiv(max_kv_seq_len, page_size)
        if buffers is not None:
            row_block_table = buffers.block_table[:num_rows]
            buffers.row_lens[:num_rows].copy_(row_lens, non_blocking=True)
            row_lens = buffers.row_lens[:num_rows]
            qo_indptr = buffers.qo_indptr[: num_rows + 1]
        else:
            row_block_table = torch.empty(
                (num_rows, max_local_pages),
                dtype=torch.int32,
                device=block_table.device,
            )
            qo_indptr = torch.arange(
                num_rows + 1,
                dtype=torch.int32,
                device=block_table.device,
            )
        self._fill_dcp_verify_page_table(row_block_table, block_table, num_reqs, qlen)
        return AiterMLADCPVerifyMetadata(
            row_lens=row_lens,
            block_table=row_block_table,
            qo_indptr=qo_indptr,
            page_size=page_size,
            max_kv_seq_len=max_kv_seq_len,
        )

    def _fill_dcp_verify_page_table(
        self,
        row_block_table: torch.Tensor,
        block_table: torch.Tensor,
        num_reqs: int,
        qlen: int,
    ) -> None:
        """Expand each request's blocks into the subpages one verify row reads.

        Every row of a request shares the request's shard, so the page list is
        built once per request and repeated; only ``row_lens`` distinguishes the
        rows. A DCP rank holds ``1/dcp_world_size`` of the sequence, so the
        request's block table is always wider than the pages a row can reach.
        """
        kernel_block_size = self.kernel_block_size
        page_size = self._segmented_page_size
        assert kernel_block_size is not None
        assert page_size is not None
        pages_per_block = kernel_block_size // page_size
        max_local_pages = row_block_table.shape[1]
        max_local_blocks = cdiv(max_local_pages, pages_per_block)
        assert max_local_blocks <= block_table.shape[1], (
            f"DCP verify needs {max_local_blocks} blocks per request but the "
            f"block table only has {block_table.shape[1]}"
        )
        subpage_offsets = torch.arange(
            pages_per_block,
            dtype=block_table.dtype,
            device=block_table.device,
        )
        per_req_page_table = (
            block_table[:num_reqs, :max_local_blocks, None] * pages_per_block
            + subpage_offsets
        ).flatten(1)[:, :max_local_pages]
        # Broadcast the request's page list across its rows instead of
        # materializing a repeat_interleave copy.
        row_block_table.unflatten(0, (num_reqs, qlen)).copy_(
            per_req_page_table.unsqueeze(1)
        )

    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,
        max_query_len: int,
        dcp_tot_seq_lens_device: torch.Tensor | None,
    ) -> AiterMLADecodeMetadata:
        causal = self._decode_causal
        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
        # the global lengths need the same dummy-row
        # treatment as the local ones below, but only for g_kv_indptr -- the
        # metadata field keeps the true values, which is what the segmented
        # route wants (it derives its own row lens and gives a dummy row
        # length 0, correctly).
        g_tot_seq_lens = dcp_tot_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),
            )
            # A dummy row just got a LOCAL length of max_qo_len, so its GLOBAL
            # length has to say the same thing or the cprr mask is degenerate:
            # global_len 0 puts the row query positions at {-3,-2,-1} while its
            # local shard sits at {rank, W+rank, 2W+rank}, so every entry is
            # masked and the row softmaxes over nothing. max_qo_len * W is
            # exactly the global length whose round-robin shard is max_qo_len
            # on every rank.
            if g_tot_seq_lens is not None:
                g_tot_seq_lens = torch.where(
                    qo_lens_device > 0,
                    g_tot_seq_lens,
                    g_tot_seq_lens.new_full((), max_qo_len * self.dcp_world_size),
                )

        if self._graph_seq_lens is not None:
            self._graph_seq_lens[:num_kernel_reqs].copy_(
                seq_lens_for_kernel, non_blocking=True
            )
            seq_lens_for_kernel = self._graph_seq_lens[:num_kernel_reqs]

        # 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._decode_num_heads,
            int(max_qo_len),
            self._kv_cache_dtype_str,
            self._kv_cache_bytes,
        )
        use_gluon_verify = AiterMLAHelper.use_gluon_verify(
            self._decode_num_heads,
            int(max_qo_len),
            self._kv_cache_dtype_str,
            self.dcp_world_size,
            causal,
            self._kv_cache_bytes,
        )
        dcp_route = _DCPDecodeRoute.PLAIN
        if self.dcp_world_size > 1:
            dcp_route = _select_dcp_decode_route(
                self._supports_segmented_dcp_verify,
                causal,
                int(max_qo_len),
                self._asm_dcp_verify,
            )

        # Segmented DCP verify carries its own per-row subpage table, so the
        # flat per-token view is dead work for it. Leave the buffer alone and
        # hand the metadata None, so a future reader cannot pick up whatever
        # the previous batch left behind.
        paged_kv_indices = None
        if dcp_route is not _DCPDecodeRoute.SEGMENTED:
            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).
            # Chunk count comes from the block table width, an upper bound on
            # tokens per request that is available host-side, so this adds no
            # device synchronisation. Programs whose chunk lies past
            # num_tokens mask out entirely.
            max_tokens_per_req = block_table_tensor.shape[1] * self.kernel_block_size
            num_chunks = max(1, cdiv(max_tokens_per_req, 1024))
            _expand_page_indices_kernel[(num_reqs, num_chunks)](
                self.paged_kv_indices,
                block_table_tensor,
                block_table_tensor.stride(0),
                block_table_tensor.stride(1),
                paged_kv_indptr,
                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 = AiterMLAHelper.qo_indptr_for_uniform_qlen(
                    num_kernel_reqs, int(max_qo_len), 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 = AiterMLAHelper.qo_indptr_for_uniform_qlen(
                    num_kernel_reqs, 1, device
                )
            else:
                if pad_uniform_mtp:
                    qo_indptr = AiterMLAHelper.qo_indptr_for_uniform_qlen(
                        num_kernel_reqs, int(max_qo_len), device
                    )
                else:
                    qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs]

        g_kv_indptr = None
        if dcp_route is _DCPDecodeRoute.CPRR and g_tot_seq_lens is not None:
            assert self._g_kv_indptr_buf is not None
            _ngk = g_tot_seq_lens.shape[0]
            g_kv_indptr = self._g_kv_indptr_buf[: _ngk + 1]
            g_kv_indptr[1:].copy_(g_tot_seq_lens.cumsum(dim=0, dtype=torch.int32))
            # A FULL cudagraph replays at its capture-time request count, so a
            # smaller real batch leaves the tail holding the PREVIOUS step's
            # cumsum. Flatten it the way paged_kv_indptr is flattened below, so
            # padding rows read global_len == 0 and agree with their zero-length
            # local shard instead of a stale global length.
            self._g_kv_indptr_buf[_ngk + 1 :].fill_(g_kv_indptr[-1])

        has_persistent_metadata = False
        # 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 predicates are disjoint -- decode
        # is qlen==1, verify is qlen>1 -- and cover both Gluon entries plus the
        # segmented DCP verify.
        use_persistent_metadata = (
            not use_gluon_decode
            and not use_gluon_verify
            and dcp_route is not _DCPDecodeRoute.SEGMENTED
            # 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.
            # A non-causal block keeps it too: what the fold drops past qlen 4
            # is the block's causal staircase, which a non-causal block does
            # not have, and the schedule is the only thing carrying its mask.
            and (
                not causal
                or self._decode_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 (
            not causal
            and max_qo_len == 2
            and is_quantized_kv_cache(self._kv_cache_dtype_str)
            and not AiterMLAHelper.has_fp8_non_causal_qlen2_kernel(
                self._decode_num_heads
            )
        ):
            # AITER's fp8 dispatch folds (gqa 16, qlen 3 or 4) onto the
            # qseqlen-4 kernel but never lists 2. Only 32/64/96/128 heads
            # at qlen 2 fold onto that 16-head / 4-token non-causal kernel;
            # 16-head qlen 2, and padded 48/80/112, keep Q2 and abort the
            # process rather than raising. The bf16 fold has no such hole.
            raise ValueError(
                "AITER has no non-causal fp8 MLA kernel for this 2-token "
                "query block. Pin the draft to TRITON_MLA for this "
                "speculative config."
            )
        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
            )
            cprr_kwargs: dict = {}
            is_causal = causal
            if self._asm_dcp_verify:
                cprr_kwargs = dict(
                    max_split_per_batch=self._mla_max_split_per_batch,
                    intra_batch_mode=False,
                )
            if dcp_route is _DCPDecodeRoute.CPRR:
                is_causal = False
                cprr_kwargs["is_cp_round_robin"] = True
            metadata_num_heads = (
                self._asm_dcp_verify_heads
                if dcp_route is _DCPDecodeRoute.CPRR
                else AiterMLAHelper.get_actual_mla_num_heads(self._decode_num_heads)
            )
            get_mla_metadata_v1(
                qo_indptr,
                paged_kv_indptr,
                paged_kv_last_page_len,
                metadata_num_heads,
                1,
                is_causal,
                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,
                **cprr_kwargs,
            )
            has_persistent_metadata = True

        # Small-head multi-token verify uses mla_gluon's 4-D MTP entry over the
        # ordinary per-request paged-KV view, so there is no expanded per-token
        # buffer to build here. mla_gluon still wants a lower bound on the KV
        # length it is asked to split; on the verify path that bound is over
        # active requests, not cudagraph padding rows pinned to max_qo_len.
        # The .item() below runs in the builder, outside the captured region,
        # so it does not abort HIP graph capture the way the per-layer syncs
        # in forward_mqa did. DCP verify never reaches here -- it carries its
        # own per-row lengths instead of a single split bound.
        min_kv_seq_len = 1
        if use_gluon_verify:
            per_req_len = paged_kv_indptr[1:] - paged_kv_indptr[:-1]
            if pad_uniform_mtp:
                active = qo_lens_device > 0
                if active.any():
                    min_kv_seq_len = int(per_req_len[active].min().item())
            else:
                min_kv_seq_len = int(per_req_len.min().item())

        dcp_verify = None
        if dcp_route is _DCPDecodeRoute.SEGMENTED:
            assert dcp_tot_seq_lens_device is not None
            dcp_verify = self._build_dcp_verify_row_view(
                int(max_qo_len),
                block_table_tensor,
                dcp_tot_seq_lens_device,
            )

        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,
            min_kv_seq_len=min_kv_seq_len,
            dcp_route=dcp_route,
            dcp_verify=dcp_verify,
            g_kv_indptr=g_kv_indptr,
            cp_world_size=self.dcp_world_size,
            cp_rank=self.dcp_rank,
            asm_decode_num_heads=(
                self._asm_dcp_verify_heads if dcp_route is _DCPDecodeRoute.CPRR else 0
            ),
            mla_num_kv_splits=(
                self._mla_max_split_per_batch
                if dcp_route is _DCPDecodeRoute.CPRR
                else 0
            ),
            use_gluon_decode=use_gluon_decode,
            use_gluon_verify=use_gluon_verify,
            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:
        # The common MLA builder owns the causal flag but its _build_decode
        # contract is shared with CUDA backends. Keep this ROCm-specific state
        # on the per-ubatch Aiter builder instead of changing that shared API;
        # super().build invokes _build_decode synchronously.
        self._decode_causal = common_attn_metadata.causal
        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

_build_dcp_verify_row_view(qlen, block_table, dcp_tot_seq_lens)

Build one paged-KV row per verify token for segmented DCP verification.

Every row is a single query (qo_indptr is an arange), so the whole causal structure is carried by dcp_local_verify_row_lens and the kernel applies no tail of its own.

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

def _build_dcp_verify_row_view(
    self,
    qlen: int,
    block_table: torch.Tensor,
    dcp_tot_seq_lens: torch.Tensor,
) -> AiterMLADCPVerifyMetadata:
    """Build one paged-KV row per verify token for segmented DCP verification.

    Every row is a single query (``qo_indptr`` is an arange), so the whole
    causal structure is carried by ``dcp_local_verify_row_lens`` and the
    kernel applies no tail of its own.
    """
    assert self.dcp_world_size > 1
    num_reqs = dcp_tot_seq_lens.numel()
    row_lens = AiterMLAHelper.dcp_local_verify_row_lens(
        dcp_tot_seq_lens,
        qlen,
        self.dcp_world_size,
        self.dcp_rank,
        self.cp_kv_cache_interleave_size,
    )
    num_rows = row_lens.numel()
    buffers = self._dcp_verify_buffers
    # Persistent buffers exist exactly when full graphs are on, and they
    # already carry the static bound they were sized for.
    max_kv_seq_len = (
        buffers.max_kv_seq_len
        if buffers is not None
        else max(1, int(row_lens.max().item()))
    )
    page_size = self._segmented_page_size
    max_local_pages = cdiv(max_kv_seq_len, page_size)
    if buffers is not None:
        row_block_table = buffers.block_table[:num_rows]
        buffers.row_lens[:num_rows].copy_(row_lens, non_blocking=True)
        row_lens = buffers.row_lens[:num_rows]
        qo_indptr = buffers.qo_indptr[: num_rows + 1]
    else:
        row_block_table = torch.empty(
            (num_rows, max_local_pages),
            dtype=torch.int32,
            device=block_table.device,
        )
        qo_indptr = torch.arange(
            num_rows + 1,
            dtype=torch.int32,
            device=block_table.device,
        )
    self._fill_dcp_verify_page_table(row_block_table, block_table, num_reqs, qlen)
    return AiterMLADCPVerifyMetadata(
        row_lens=row_lens,
        block_table=row_block_table,
        qo_indptr=qo_indptr,
        page_size=page_size,
        max_kv_seq_len=max_kv_seq_len,
    )

_build_fp8_prefill_ps_metadata(metadata, common_attn_metadata)

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()).

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

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)

    # Head counts that are not a multiple of 16 (K3: 12/rank at TP8) are
    # replicate-padded up to one in _mla_fp8_prefill_attn; build the PS
    # metadata for that same padded count so the work/reduce maps and the
    # scratch reservations describe the width the kernel is handed.
    #
    # This was previously max(16, num_heads), which agrees with the helper
    # at every head count reachable today (12 and 16 both give 16) and
    # differs only above 16: at 24 it leaves num_head_k=24 while the
    # forward pads to 32. That is not a correctness bug -- the 24-wide work
    # maps still cover head-tiles 0..23, which are the real heads -- but it
    # is expensive, because a lower head alignment yields more partial
    # tiles: gcd-driven, 24 heads -> 64 tiles vs 32 heads -> 16, i.e.
    # 193.6 MiB of reservations instead of 68.6 MiB at batch=1/qlen=512
    # (measured on gfx950). Sizing both from one helper keeps the widths
    # equal and takes the cheaper tiling.
    num_head_k = AiterMLAHelper.get_fp8_prefill_num_heads(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

    slot = self._fp8_ps_slot
    self._fp8_ps_slot ^= 1
    staging_free = self._fp8_ps_staging_free[slot]
    with gpu_sync_allowed():
        staging_free.synchronize()
    (
        work_indptr_host,
        work_info_host,
        reduce_indptr_host,
        reduce_final_map_host,
        reduce_partial_map_host,
    ) = self._fp8_ps_staging[slot]

    # work_metadata is not written by the host planner; it stays on device.
    get_ps_metadata_v1(
        qo_indptr_cpu,
        kv_indptr_cpu,
        seq_lens_cpu,
        gqa_ratio,
        num_head_k,
        self.fp8_ps_work_metadata,
        work_indptr_host,
        work_info_host,
        reduce_indptr_host,
        reduce_final_map_host,
        reduce_partial_map_host,
        qhead_granularity=qhead_granularity,
        qlen_granularity=qlen_granularity,
        kvlen_granularity=kvlen_granularity,
        block_size=block_size,
        is_causal=True,
    )
    for dst, src in zip(self._fp8_ps_device_outputs, self._fp8_ps_staging[slot]):
        dst.copy_(src, non_blocking=True)
    staging_free.record()

    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, read from the host plan (no GPU sync).
    num_partial_tiles = int(reduce_indptr_host[-1])

    # 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

_fill_dcp_verify_page_table(row_block_table, block_table, num_reqs, qlen)

Expand each request's blocks into the subpages one verify row reads.

Every row of a request shares the request's shard, so the page list is built once per request and repeated; only row_lens distinguishes the rows. A DCP rank holds 1/dcp_world_size of the sequence, so the request's block table is always wider than the pages a row can reach.

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

def _fill_dcp_verify_page_table(
    self,
    row_block_table: torch.Tensor,
    block_table: torch.Tensor,
    num_reqs: int,
    qlen: int,
) -> None:
    """Expand each request's blocks into the subpages one verify row reads.

    Every row of a request shares the request's shard, so the page list is
    built once per request and repeated; only ``row_lens`` distinguishes the
    rows. A DCP rank holds ``1/dcp_world_size`` of the sequence, so the
    request's block table is always wider than the pages a row can reach.
    """
    kernel_block_size = self.kernel_block_size
    page_size = self._segmented_page_size
    assert kernel_block_size is not None
    assert page_size is not None
    pages_per_block = kernel_block_size // page_size
    max_local_pages = row_block_table.shape[1]
    max_local_blocks = cdiv(max_local_pages, pages_per_block)
    assert max_local_blocks <= block_table.shape[1], (
        f"DCP verify needs {max_local_blocks} blocks per request but the "
        f"block table only has {block_table.shape[1]}"
    )
    subpage_offsets = torch.arange(
        pages_per_block,
        dtype=block_table.dtype,
        device=block_table.device,
    )
    per_req_page_table = (
        block_table[:num_reqs, :max_local_blocks, None] * pages_per_block
        + subpage_offsets
    ).flatten(1)[:, :max_local_pages]
    # Broadcast the request's page list across its rows instead of
    # materializing a repeat_interleave copy.
    row_block_table.unflatten(0, (num_reqs, qlen)).copy_(
        per_req_page_table.unsqueeze(1)
    )

_init_fp8_prefill_ps_buffers(max_num_reqs, max_prefill_qlen, max_num_batched_tokens, attn_out_dtype, device)

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.

Parameters:

• max_num_reqs

  (int) –

  Maximum number of concurrent requests.

• max_prefill_qlen

  (int) –

  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

  (int) –

  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.

• attn_out_dtype

  (dtype) –

  Dtype of the attention output buffer, used to size the padded-head output scratch (small head counts only).

• device

  (device) –

  Target device for the buffers.

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

def _init_fp8_prefill_ps_buffers(
    self,
    max_num_reqs: int,
    max_prefill_qlen: int,
    max_num_batched_tokens: int,
    attn_out_dtype: torch.dtype,
    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``.
        attn_out_dtype: Dtype of the attention output buffer, used to size
            the padded-head output scratch (small head counts only).
        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.
    # Head counts that are not a multiple of 16 (K3: 12/rank at TP8) are
    # replicate-padded up to one in _mla_fp8_prefill_attn; build the PS
    # metadata for that same padded count so the work/reduce maps and the
    # scratch reservations describe the width the kernel is handed.
    #
    # This was previously max(16, num_heads), which agrees with the helper
    # at every head count reachable today (12 and 16 both give 16) and
    # differs only above 16: at 24 it leaves num_head_k=24 while the
    # forward pads to 32. That is not a correctness bug -- the 24-wide work
    # maps still cover head-tiles 0..23, which are the real heads -- but it
    # is expensive, because a lower head alignment yields more partial
    # tiles: gcd-driven, 24 heads -> 64 tiles vs 32 heads -> 16, i.e.
    # 193.6 MiB of reservations instead of 68.6 MiB at batch=1/qlen=512
    # (measured on gfx950). Sizing both from one helper keeps the widths
    # equal and takes the cheaper tiling.
    num_head_k = AiterMLAHelper.get_fp8_prefill_num_heads(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,
    )

    # get_ps_metadata_v1 builds the plan on the host and writes its outputs with
    # blocking hipMemcpy that is not ordered on the current stream. Writing the
    # device buffers directly lets step N+1's build (async scheduling) overwrite a
    # plan that step N's prefill kernels are still reading. Build into pinned host
    # staging instead and copy on the current stream. Two staging slots, each
    # reused only after its previous H2D copy has completed.
    self._fp8_ps_device_outputs = (
        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,
    )
    self._fp8_ps_staging = [
        tuple(
            torch.empty(t.shape, dtype=t.dtype, device="cpu", pin_memory=True)
            for t in self._fp8_ps_device_outputs
        )
        for _ in range(2)
    ]
    # One event per slot, recorded after that slot's H2D copy. synchronize()
    # on an event that has not been recorded yet returns immediately.
    self._fp8_ps_staging_free = [torch.cuda.Event(), torch.cuda.Event()]
    self._fp8_ps_slot = 0

    from vllm.platforms import current_platform
    from vllm.v1.worker.workspace import current_workspace_manager

    # AITER metadata sizing assumes all requests can carry max_prefill_qlen tokens,
    # which is a loose worst case in the number of QO tiles. Rather, the sum of
    # qlens is bounded by max_num_batched_tokens. Manually compute the number of
    # qo tiles to avoid OOM at startup.
    # TODO: AITER should give us this budget constrained value
    qo_tile_cnt = (
        cdiv(max_num_batched_tokens, _FP8_PREFILL_TILE_Q) + max_num_reqs - 1
    )
    max_num_partial_tiles = qo_tile_cnt + current_platform.num_compute_units()
    reservations: list[tuple[tuple[int, ...], torch.dtype]] = [
        (
            (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),
    ]
    if self.num_heads < num_head_k:
        # Padded head counts also take their kernel output buffer from the
        # workspace: the caller's [total_q, num_heads * v_head_dim] output
        # cannot back a num_head_k-head view (see _mla_fp8_prefill_attn).
        reservations.append(
            ((max_num_batched_tokens, num_head_k, v_head_dim), attn_out_dtype)
        )
    current_workspace_manager().get_simultaneous(*reservations)

    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,
    )

_aiter_mla_native_h24_metadata_supported() cached

Whether AITER's fast MLA metadata planner accepts native H24.

The reducer and metadata planner have independent shape dispatch. Checking only the reducer can route H24 into a planner that rejects it before the attention kernel launches. Until AITER exposes a capability API, inspect the shipped JIT source for an explicit native-H24 planner branch.

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

@functools.lru_cache(maxsize=1)
def _aiter_mla_native_h24_metadata_supported() -> bool:
    """Whether AITER's fast MLA metadata planner accepts native H24.

    The reducer and metadata planner have independent shape dispatch. Checking
    only the reducer can route H24 into a planner that rejects it before the
    attention kernel launches. Until AITER exposes a capability API, inspect
    the shipped JIT source for an explicit native-H24 planner branch.
    """
    try:
        from aiter.jit.core import AITER_CSRC_DIR

        metadata_source = (
            Path(AITER_CSRC_DIR) / "kernels" / "mla" / "metadata" / "v1_2_device.cuh"
        )
        source = "".join(metadata_source.read_text(encoding="utf-8").split())
    except (ImportError, OSError):
        return False
    return "num_heads==24" in source

_aiter_mla_native_h24_reducer_supported() cached

Whether AITER's JIT reducer supports the native H24/512 shape.

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

@functools.lru_cache(maxsize=1)
def _aiter_mla_native_h24_reducer_supported() -> bool:
    """Whether AITER's JIT reducer supports the native H24/512 shape."""
    try:
        from aiter.jit.core import AITER_CSRC_DIR

        reduce_source = Path(AITER_CSRC_DIR) / "kernels" / "mla" / "reduce.cu"
        source = "".join(reduce_source.read_text(encoding="utf-8").split())
    except (ImportError, OSError):
        return False
    return "MLA_REDUCE_CASE_EF(NUM_HEAD,24,HEAD_DIM,512," in source

_aiter_mla_native_h24_supported()

Whether the complete AITER decode path supports native H24.

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

def _aiter_mla_native_h24_supported() -> bool:
    """Whether the complete AITER decode path supports native H24."""
    return (
        _aiter_mla_native_h24_reducer_supported()
        and _aiter_mla_native_h24_metadata_supported()
    )

_aiter_mla_non_causal_asm_kernels()

Whether this arch ships non-causal MLA decode ASM kernels.

The Python causal= probe cannot see the per-arch manifest. gfx950 has the kernels; gfx942 and everything else must fall through to another backend.

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

def _aiter_mla_non_causal_asm_kernels() -> bool:
    """Whether this arch ships non-causal MLA decode ASM kernels.

    The Python `causal=` probe cannot see the per-arch manifest. gfx950 has
    the kernels; gfx942 and everything else must fall through to another
    backend.
    """
    try:
        from vllm.platforms.rocm import on_gfx950
    except Exception:  # noqa: BLE001
        return False
    return bool(on_gfx950())

_aiter_mla_small_head_mode()

Small-head (<16) MLA decode kernel selection.

Controlled by VLLM_ROCM_AITER_MLA_ASM_PADDING:

• "auto" (default): let the arch decide -- divisor head counts keep the Gluon decode where a build exists (gfx950), everything else (non-divisor counts and all counts on gfx942) uses the padded persistent-scheduling ASM decode.
• "gluon": prefer the Gluon path wherever a build exists.
• "asm": force the padded persistent-scheduling ASM decode.

On gfx942 (no Gluon build) the ASM path is always used regardless of this setting; "gluon" there falls back to ASM with a one-time warning.

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

def _aiter_mla_small_head_mode() -> str:
    """Small-head (<16) MLA decode kernel selection.

    Controlled by ``VLLM_ROCM_AITER_MLA_ASM_PADDING``:

    - ``"auto"`` (default): let the arch decide -- divisor head counts keep the
      Gluon decode where a build exists (gfx950), everything else (non-divisor
      counts and all counts on gfx942) uses the padded persistent-scheduling
      ASM decode.
    - ``"gluon"``: prefer the Gluon path wherever a build exists.
    - ``"asm"``: force the padded persistent-scheduling ASM decode.

    On gfx942 (no Gluon build) the ASM path is always used regardless of this
    setting; ``"gluon"`` there falls back to ASM with a one-time warning.
    """
    import vllm.envs as envs

    mode = (envs.VLLM_ROCM_AITER_MLA_ASM_PADDING or "auto").lower()
    if mode == "gluon" and not _gluon_mla_decode_supported():
        logger.warning_once(
            "VLLM_ROCM_AITER_MLA_ASM_PADDING=gluon requested, but this device "
            "has no Gluon MLA decode build (Gluon requires gfx950); using the "
            "padded persistent-scheduling ASM decode instead."
        )
    return mode

_asm_dcp_verify_configured(dcp_world_size, cp_interleave, multi_token_decode)

Whether asm cprr DCP verify is reachable in this process.

Requires DCP with interleave size 1, speculative decoding, and gfx950 (the only arch AITER builds cprr kernels for).

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

def _asm_dcp_verify_configured(
    dcp_world_size: int, cp_interleave: int, multi_token_decode: bool
) -> bool:
    """Whether asm cprr DCP verify is reachable in this process.

    Requires DCP with interleave size 1, speculative decoding, and gfx950 (the
    only arch AITER builds cprr kernels for).
    """
    if not (dcp_world_size > 1 and cp_interleave == 1 and multi_token_decode):
        return False
    from vllm.platforms.rocm import on_gfx950

    return bool(on_gfx950())

_asm_dcp_verify_heads(decode_num_heads)

Head count the cprr asm decode runs at, or 0 if it cannot serve this shape.

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

def _asm_dcp_verify_heads(decode_num_heads: int) -> int:
    """Head count the cprr asm decode runs at, or 0 if it cannot serve this shape."""
    if decode_num_heads in _NATIVE_CPRR_HEADS:
        return decode_num_heads
    for nat in _NATIVE_CPRR_HEADS:
        if nat >= decode_num_heads:
            return nat
    return 0

_expand_page_indices_kernel(page_indices, block_table, block_table_stride_0, block_table_stride_1, cu_num_tokens, KERNEL_BLOCK_SIZE, BLOCK_SIZE)

Expand block table entries into per-token flat page indices.

The aiter MLA kernel always operates with page_size=1 internally (kv_buffer is flattened via .view(-1, 1, 1, H)). This kernel converts block-level indices from the block table into individual token positions in the flattened KV buffer.

When KERNEL_BLOCK_SIZE=1: block_idx=t, offset=0, flat=block_id (equivalent to a direct copy -- no regression from the original kernel).

When KERNEL_BLOCK_SIZE=K: block table entry b (covering K tokens) is expanded to flat indices bK, bK+1, ..., b*K+(K-1).

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

@triton.jit
def _expand_page_indices_kernel(
    page_indices,
    block_table,
    block_table_stride_0,
    block_table_stride_1,
    cu_num_tokens,
    KERNEL_BLOCK_SIZE: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    """Expand block table entries into per-token flat page indices.

    The aiter MLA kernel always operates with page_size=1 internally
    (kv_buffer is flattened via .view(-1, 1, 1, H)). This kernel converts
    block-level indices from the block table into individual token positions
    in the flattened KV buffer.

    When KERNEL_BLOCK_SIZE=1: block_idx=t, offset=0, flat=block_id
    (equivalent to a direct copy -- no regression from the original kernel).

    When KERNEL_BLOCK_SIZE=K: block table entry b (covering K tokens)
    is expanded to flat indices b*K, b*K+1, ..., b*K+(K-1).
    """
    # One program per (request, token-chunk). Parallelising over requests
    # alone degenerates at low concurrency: with num_reqs == 1 a single
    # workgroup walked the whole sequence in a serial loop.
    req_idx = tl.program_id(0)
    chunk_idx = tl.program_id(1)
    row_ptr = block_table + req_idx * block_table_stride_0
    start_idx = tl.load(cu_num_tokens + req_idx)
    num_tokens = tl.load(cu_num_tokens + req_idx + 1) - start_idx

    # The grid is sized from the block table width, an upper bound over all
    # requests, so a ragged batch launches chunks past a short request's end.
    # Returning here keeps those programs from issuing masked-out loads and
    # stores at all.
    chunk_start = chunk_idx * BLOCK_SIZE
    if chunk_start >= num_tokens:
        return

    token_offsets = chunk_start + tl.arange(0, BLOCK_SIZE)
    mask = token_offsets < num_tokens

    # Which block in the block table does this token belong to?
    block_idx = token_offsets // KERNEL_BLOCK_SIZE
    # Offset within that block
    offset_in_block = token_offsets % KERNEL_BLOCK_SIZE

    # Load the block ID from the block table
    # Both strides are taken from the caller: the block table is a view owned
    # elsewhere, so a unit column stride must not be assumed.
    block_ids = tl.load(row_ptr + block_idx * block_table_stride_1, mask=mask)

    # Compute flat index in the flattened kv_buffer
    flat_indices = block_ids * KERNEL_BLOCK_SIZE + offset_in_block

    tl.store(
        page_indices + start_idx + token_offsets,
        flat_indices,
        mask=mask,
    )

_fp8_mla_prefill_supported() cached

Auto-detect FP8 MLA prefill via mla_prefill_ps_asm_fwd + mla_reduce_v1.

Requires gfx950 plus an AITER build that exports both kernels. When either is missing we silently fall back to flash_attn_varlen_func.

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

@functools.lru_cache(maxsize=1)
def _fp8_mla_prefill_supported() -> bool:
    """Auto-detect FP8 MLA prefill via mla_prefill_ps_asm_fwd + mla_reduce_v1.

    Requires gfx950 plus an AITER build that exports both kernels.  When
    either is missing we silently fall back to ``flash_attn_varlen_func``.
    """
    try:
        from vllm.platforms.rocm import on_gfx950
    except Exception:  # noqa: BLE001
        return False
    if not on_gfx950():
        return False
    try:
        from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1  # noqa: F401
    except Exception:  # noqa: BLE001
        return False
    return True

_get_mla_gluon() cached

Load the small-head Gluon MLA entry point.

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

@functools.lru_cache(maxsize=1)
def _get_mla_gluon():
    """Load the small-head Gluon MLA entry point."""
    unified_module = "aiter.ops.triton.gluon.mla_gluon"
    try:
        from aiter.ops.triton.gluon.mla_gluon import mla_gluon

        return mla_gluon
    except ModuleNotFoundError as unified_import_error:
        if not unified_module.startswith(unified_import_error.name or ""):
            raise
        legacy_module = "aiter.ops.triton.gluon.mla_decode_gluon"
        try:
            from aiter.ops.triton.gluon.mla_decode_gluon import mla_decode_gluon

            return mla_decode_gluon
        except ModuleNotFoundError as legacy_import_error:
            if not legacy_module.startswith(legacy_import_error.name or ""):
                raise
            raise RuntimeError(
                "ROCM_AITER_MLA requires an AITER build with the small-head "
                "Gluon MLA kernel (mla_gluon or mla_decode_gluon) when decode "
                "heads are fewer than 16."
            ) from unified_import_error

_get_segmented_mla_decode() cached

Load AITER's segmented MLA decode with unreduced partial output.

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

@functools.lru_cache(maxsize=1)
def _get_segmented_mla_decode():
    """Load AITER's segmented MLA decode with unreduced partial output."""
    from aiter.ops.triton.attention.mla import mla_decode_fwd

    return mla_decode_fwd

_gluon_kv_cache_in_bounds(kv_cache_bytes)

Whether this layer's KV cache keeps Gluon on its bounds-checked path.

The cache is one flat tensor per layer, so the bound disqualifies the layer outright; no batch shape brings an oversized cache back in range. None is the not-yet-sized case during profiling, before any kernel runs.

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

def _gluon_kv_cache_in_bounds(kv_cache_bytes: int | None) -> bool:
    """Whether this layer's KV cache keeps Gluon on its bounds-checked path.

    The cache is one flat tensor per layer, so the bound disqualifies the layer
    outright; no batch shape brings an oversized cache back in range. ``None``
    is the not-yet-sized case during profiling, before any kernel runs.
    """
    if kv_cache_bytes is None or kv_cache_bytes <= _GLUON_MAX_KV_CACHE_BYTES:
        return True
    logger.warning_once(
        "KV cache is %.1f GiB per layer, past the 2 GiB where the Gluon MLA "
        "kernel drops its KV bounds mask; using the padded ASM decode instead. "
        "Lower --gpu-memory-utilization or --max-model-len to get Gluon back.",
        kv_cache_bytes / (1 << 30),
    )
    return False

_gluon_mla_decode_supported() cached

The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build.

Its tiling needs ~160 KiB of LDS, which exceeds CDNA3's 64 KiB, so on gfx942 there is no kernel to fall through to and selecting it asserts (mla_gluon requires gfx950). Restrict Gluon decode to gfx950; other archs use the asm persistent decode, which get_mla_padded_q makes correct for any 1..15 heads.

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

@functools.lru_cache(maxsize=1)
def _gluon_mla_decode_supported() -> bool:
    """The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build.

    Its tiling needs ~160 KiB of LDS, which exceeds CDNA3's 64 KiB, so on
    gfx942 there is no kernel to fall through to and selecting it asserts
    (``mla_gluon requires gfx950``). Restrict Gluon decode to gfx950; other
    archs use the asm persistent decode, which ``get_mla_padded_q`` makes
    correct for any 1..15 heads.
    """
    try:
        from vllm.platforms.rocm import on_gfx950
    except Exception:  # noqa: BLE001
        return False
    return on_gfx950()

_segmented_dcp_verify_supported(dcp_world_size, cp_interleave)

Whether this configuration can serve DCP verify on segmented MLA.

Configuration only -- whether a given batch takes the route additionally depends on its query length. Round-robin interleaving other than 1 is excluded because the per-row causal window has not been validated there.

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

def _segmented_dcp_verify_supported(dcp_world_size: int, cp_interleave: int) -> bool:
    """Whether this configuration can serve DCP verify on segmented MLA.

    Configuration only -- whether a given batch takes the route additionally
    depends on its query length. Round-robin interleaving other than 1 is
    excluded because the per-row causal window has not been validated there.
    """
    return (
        dcp_world_size > 1 and cp_interleave == 1 and _segmented_mla_decode_supported()
    )

_segmented_mla_decode_supported() cached

Whether AITER exposes the segmented MLA decode used by DCP verify.

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

@functools.lru_cache(maxsize=1)
def _segmented_mla_decode_supported() -> bool:
    """Whether AITER exposes the segmented MLA decode used by DCP verify."""
    try:
        _get_segmented_mla_decode()
    except Exception:  # noqa: BLE001
        return False
    return True

_segmented_mla_page_size(block_size)

Largest supported power-of-two subpage dividing a physical KV block.

The subpage size becomes the segmented kernel's TILE_SIZE: the cache is reinterpreted as (-1, page_size, 1, head_dim) before the call. It must divide the physical block and be a power of two, since the kernel walks a tile with tl.arange(0, TILE_SIZE). 128 is the largest tile that tiling supports -- a tile holds TILE_SIZE x kv_lora_rank keys and its scores are BLOCK_M x TILE_SIZE -- so a larger manager block is split into several subpages instead of widening the tile.

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

def _segmented_mla_page_size(block_size: int) -> int:
    """Largest supported power-of-two subpage dividing a physical KV block.

    The subpage size becomes the segmented kernel's ``TILE_SIZE``: the cache is
    reinterpreted as ``(-1, page_size, 1, head_dim)`` before the call. It must
    divide the physical block and be a power of two, since the kernel walks a
    tile with ``tl.arange(0, TILE_SIZE)``. 128 is the largest tile that tiling
    supports -- a tile holds ``TILE_SIZE x kv_lora_rank`` keys and its scores
    are ``BLOCK_M x TILE_SIZE`` -- so a larger manager block is split into
    several subpages instead of widening the tile.
    """
    assert block_size > 0
    return min(128, largest_power_of_2_divisor(block_size))

_select_dcp_decode_route(supports_segmented, causal, max_qo_len, asm_selected)

Select one DCP decode route for this batch; see DCP decode routing.

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

def _select_dcp_decode_route(
    supports_segmented: bool,
    causal: bool,
    max_qo_len: int,
    asm_selected: bool,
) -> _DCPDecodeRoute:
    """Select one DCP decode route for this batch; see DCP decode routing."""
    if not causal or max_qo_len <= 1:
        return _DCPDecodeRoute.PLAIN
    if asm_selected and max_qo_len >= _MIN_CPRR_QLEN:
        return _DCPDecodeRoute.CPRR
    if supports_segmented:
        return _DCPDecodeRoute.SEGMENTED
    raise RuntimeError(
        "ROCM_AITER_MLA DCP multi-token verify requires either segmented "
        "MLA or the round-robin asm decode."
    )