vllm.v1.attention.backends.mla.rocm_aiter_mla

Classes:

• AiterMLADCPVerifyMetadata –

  One paged-KV row per verify token for causal 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 causal DCP verification.

Segmented MLA and the gfx942 Triton path both consume this view. Which kernel runs is AiterMLADecodeMetadata.dcp_route.

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

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

    Segmented MLA and the gfx942 Triton path both consume this view. Which
    kernel runs is ``AiterMLADecodeMetadata.dcp_route``.
    """

    # 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.

• is_valid_fp8_prefill_num_heads –

  Whether the FP8 PS asm prefill should run at num_heads per rank.

• 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 is_valid_fp8_prefill_num_heads(num_heads: int) -> bool:
        """Whether the FP8 PS asm prefill should run at ``num_heads`` per rank.

        Rejects counts that padding to the kernel width would at least double
        (1-8 heads). At 8 heads (Kimi-K2.x at TP8) the doubled kernel plus its
        reduce is slower than bf16 flash attention, so those counts keep the
        bf16 prefill.
        """
        return AiterMLAHelper.is_valid_num_heads(
            num_heads
        ) and 2 * num_heads > AiterMLAHelper.get_fp8_prefill_num_heads(num_heads)

    @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 q.shape[1] == m:
            return q
        if m % num_heads == 0:
            return q.repeat_interleave(m // num_heads, dim=1)
        if q.stride(0) == m * q.stride(1):
            return q.as_strided((q.shape[0], m, q.shape[2]), q.stride())
        # 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

is_valid_fp8_prefill_num_heads(num_heads) staticmethod

Whether the FP8 PS asm prefill should run at num_heads per rank.

Rejects counts that padding to the kernel width would at least double (1-8 heads). At 8 heads (Kimi-K2.x at TP8) the doubled kernel plus its reduce is slower than bf16 flash attention, so those counts keep the bf16 prefill.

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

@staticmethod
def is_valid_fp8_prefill_num_heads(num_heads: int) -> bool:
    """Whether the FP8 PS asm prefill should run at ``num_heads`` per rank.

    Rejects counts that padding to the kernel width would at least double
    (1-8 heads). At 8 heads (Kimi-K2.x at TP8) the doubled kernel plus its
    reduce is slower than bf16 flash attention, so those counts keep the
    bf16 prefill.
    """
    return AiterMLAHelper.is_valid_num_heads(
        num_heads
    ) and 2 * num_heads > AiterMLAHelper.get_fp8_prefill_num_heads(num_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)
        self._sm_count = current_platform.num_compute_units()
        self._triton_verify_warmed_up = False

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

        self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and (
            is_quantized_kv_cache(kv_cache_dtype)
            and AiterMLAHelper.is_valid_fp8_prefill_num_heads(self.num_heads)
        )
        self._fused_prefill_qkv = False
        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
            self._fp8_prefill_one_scale: torch.Tensor | None = None
            head_dims = (self.qk_nope_head_dim, self.qk_rope_head_dim, self.v_head_dim)
            self._fused_prefill_qkv = (
                rocm_aiter_ops.is_enabled()
                and _fused_mla_prefill_qkv_supported()
                and not any(d & (d - 1) for d in head_dims)
            )

        self._use_fused_mla_kv_concat = (
            hasattr(torch.ops._C, "fused_kimi_k3_mla_kv_concat")
            and self.qk_nope_head_dim == 128
            and self.qk_rope_head_dim == 64
        )

    def _concat_k_nope_k_pe(
        self, k_nope: torch.Tensor, k_pe: torch.Tensor
    ) -> torch.Tensor:
        """Build [k_nope | k_pe] with one kernel, fall back to the base copies."""
        if not (
            self._use_fused_mla_kv_concat
            and k_nope.dim() == 3
            and k_nope.dtype in (torch.bfloat16, torch.float16)
            and k_pe.dtype == k_nope.dtype
        ):
            return super()._concat_k_nope_k_pe(k_nope, k_pe)
        k = torch.empty(
            (*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]),
            dtype=k_nope.dtype,
            device=k_nope.device,
        )
        torch.ops._C.fused_kimi_k3_mla_kv_concat(
            k_nope, k_pe.reshape(k_pe.shape[0], k_pe.shape[-1]), k
        )
        return k

    def _can_fuse_context_gather(
        self,
        q: torch.Tensor,
        prefill_metadata: MLACommonPrefillMetadata,
        kv_c_and_k_pe_cache: torch.Tensor,
    ) -> bool:
        from vllm.platforms import current_platform

        weight = getattr(self.kv_b_proj, "weight", None)
        return (
            self.kv_cache_dtype != "fp8_ds_mla"
            and self.kv_lora_rank == 512
            and prefill_metadata.q_data_type != current_platform.fp8_dtype()
            and q.dtype in (torch.bfloat16, torch.float16)
            and kv_c_and_k_pe_cache.is_contiguous()
            and weight is not None
            and weight.dtype in (torch.bfloat16, torch.float16)
            and weight.is_contiguous()
            and getattr(self.kv_b_proj, "weight_scale", None) is None
        )

    def _gather_context_chunk(
        self,
        chunk,
        kv_c_and_k_pe_cache: torch.Tensor,
        kv_indices: torch.Tensor,
        k_scale: torch.Tensor,
        out_dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Gather one context chunk and expand it into (k, v)."""
        from aiter.ops.triton.gather_kv_b_proj import gather_kv_b_proj

        num_blocks, block_size, cache_width = kv_c_and_k_pe_cache.shape
        rope_dim = cache_width - self.kv_lora_rank
        num_rows = chunk.num_context_tokens
        k = torch.empty(
            (num_rows, self.num_heads, self.qk_nope_head_dim + rope_dim),
            dtype=out_dtype,
            device=kv_c_and_k_pe_cache.device,
        )
        v = torch.empty(
            (num_rows, self.num_heads, self.v_head_dim),
            dtype=out_dtype,
            device=kv_c_and_k_pe_cache.device,
        )
        if num_rows == 0:
            return k, v

        gather_kv_b_proj(
            kv_c_and_k_pe_cache.view(num_blocks * block_size, 1, cache_width),
            k_scale,
            chunk.cu_seq_lens,
            kv_indices,
            chunk.cu_seq_lens,
            self.kv_b_proj.weight,
            None,
            k,
            v,
        )
        return k, v

    def _compute_prefill_context(
        self,
        q: torch.Tensor,
        kv_c_and_k_pe_cache: torch.Tensor,
        attn_metadata: MLACommonMetadata,
        k_scale: torch.Tensor,
    ):
        """Attend the cached prefix, gathering and expanding it in one kernel"""
        assert isinstance(attn_metadata, AiterMLAMetadata)
        assert attn_metadata.prefill is not None
        prefill_metadata = attn_metadata.prefill

        if not self._can_fuse_context_gather(q, prefill_metadata, kv_c_and_k_pe_cache):
            return super()._compute_prefill_context(
                q, kv_c_and_k_pe_cache, attn_metadata, k_scale
            )
        kv_indices = attn_metadata.context_chunk_kv_indices
        if kv_indices is None:
            return super()._compute_prefill_context(
                q, kv_c_and_k_pe_cache, attn_metadata, k_scale
            )

        assert prefill_metadata.prefill_backend is not None
        chunked_context = prefill_metadata.chunked_context
        assert chunked_context is not None

        output = None
        output_lse = None
        for chunk in chunked_context.chunks:
            k, v = self._gather_context_chunk(
                chunk,
                kv_c_and_k_pe_cache,
                kv_indices[chunk.index],
                k_scale,
                q.dtype,
            )

            attn_output, attn_softmax_lse = (
                prefill_metadata.prefill_backend.run_prefill_context_chunk(
                    chunk=chunk,
                    q=q[chunk.token_slice],
                    k=k,
                    v=v,
                )
            )

            if output is None:
                if (
                    len(chunked_context.chunks) == 1
                    and not chunked_context.empty_token_slices
                ):
                    return attn_output, attn_softmax_lse
                output, output_lse = init_mla_context_partial(
                    chunked_context,
                    attn_output,
                    attn_softmax_lse,
                    num_tokens=q.shape[0],
                )
            accumulate_mla_context_chunk(
                chunk, attn_output, attn_softmax_lse, output, output_lse
            )
        return output, output_lse

    def _context_parallel_compute_prefill_context(
        self,
        q: torch.Tensor,
        kv_c_and_k_pe_cache: torch.Tensor,
        attn_metadata: MLACommonMetadata,
        k_scale: torch.Tensor,
        dcp_world_size: int,
        fused_mla_kv_concat_fn: Callable[
            [torch.Tensor, torch.Tensor, bool], tuple[torch.Tensor, torch.Tensor]
        ]
        | None = None,
    ):
        attn_metadata = cast(AiterMLAMetadata, attn_metadata)
        assert attn_metadata.prefill is not None
        prefill = attn_metadata.prefill
        weight = getattr(self.kv_b_proj, "weight", None)
        eligible = (
            _supports_fused_dcp_prefill_config(
                self.kv_cache_dtype,
                prefill.q_data_type,
                self.kv_lora_rank,
                self.qk_nope_head_dim,
                self.qk_rope_head_dim,
                self.v_head_dim,
            )
            and q.dtype == torch.bfloat16
            and weight is not None
            and weight.dtype == torch.bfloat16
            and weight.shape == (self.num_heads * 256, 512)
            and weight.is_contiguous()
            and getattr(self.kv_b_proj, "bias", None) is None
            and getattr(self.kv_b_proj, "weight_scale", None) is None
            and kv_c_and_k_pe_cache.is_contiguous()
            and attn_metadata.dcp_context_row_indices is not None
        )
        if not eligible:
            return super()._context_parallel_compute_prefill_context(
                q,
                kv_c_and_k_pe_cache,
                attn_metadata,
                k_scale,
                dcp_world_size,
                fused_mla_kv_concat_fn,
            )
        from vllm.platforms import current_platform

        assert weight is not None
        assert attn_metadata.dcp_context_row_indices is not None
        chunked = prefill.chunked_context
        assert chunked is not None and chunked.dcp_manager is not None
        assert prefill.prefill_backend is not None
        assert len(chunked.chunks) == len(attn_metadata.dcp_context_row_indices)
        output = output_lse = None
        for chunk, row_indices in zip(
            chunked.chunks, attn_metadata.dcp_context_row_indices
        ):
            gathered = gather_compressed_context(
                kv_c_and_k_pe_cache,
                chunked.workspace,
                prefill.block_table[chunk.request_slice],
                chunk,
                chunked.dcp_manager.kv_gather,
                current_platform.fp8_dtype(),
            )
            k, v = expand_context(
                gathered,
                k_scale,
                row_indices,
                chunk.cu_seq_lens,
                weight,
                self.num_heads,
                self.qk_nope_head_dim,
                self.qk_rope_head_dim,
                self.v_head_dim,
            )
            attn_output, attn_lse = prefill.prefill_backend.run_prefill_context_chunk(
                chunk=chunk,
                q=q[chunk.token_slice],
                k=k,
                v=v,
            )
            if output is None:
                if len(chunked.chunks) == 1 and not chunked.empty_token_slices:
                    return attn_output, attn_lse
                output, output_lse = init_mla_context_partial(
                    chunked,
                    attn_output,
                    attn_lse,
                    num_tokens=q.shape[0],
                )
            accumulate_mla_context_chunk(
                chunk, attn_output, attn_lse, output, output_lse
            )
        return output, output_lse

    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 every padded head a duplicate of
        # a real head, so the real heads are bit-identical to the unpadded
        # result. get_mla_padded_q repeat_interleaves when the target is a
        # multiple of num_heads (real heads at stride target // num_heads) and
        # tiles otherwise (real heads [0:num_heads]). 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)

        if self._fp8_prefill_one_scale is None:
            self._fp8_prefill_one_scale = torch.ones(
                (), dtype=torch.float32, device=q.device
            )
        one_scale = self._fp8_prefill_one_scale

        # 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:
            real_out = (
                out_3d[:, :: nhead // _real_nhead, :]
                if nhead % _real_nhead == 0
                else out_3d[:, :_real_nhead, :]
            )
            out.view(total_q, _real_nhead, v_head_dim).copy_(real_out)

    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, the selected prefill backend supersedes
        it, 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.
        """
        prefill_metadata = attn_metadata.prefill
        if (
            not self._fp8_prefill_enabled
            or not isinstance(attn_metadata, AiterMLAMetadata)
            or prefill_metadata is None
            or _asm_fp8_prefill_superseded(prefill_metadata.prefill_backend)
            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,
            )

        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)
        if self._fused_prefill_qkv:
            from aiter.ops.triton.fusions.fused_mla_prefill_qkv_fp8 import (
                fused_mla_prefill_qkv_fp8,
            )

            from vllm.platforms import current_platform

            # The k_pe concat, the head padding and the FP8 cast in one launch.
            q, k, v = fused_mla_prefill_qkv_fp8(
                q,
                k_nope,
                v,
                k_pe,
                AiterMLAHelper.get_fp8_prefill_num_heads(self.num_heads),
                current_platform.fp8_dtype(),
            )
        else:
            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_triton_dcp_verify(
        self,
        q_mla: torch.Tensor,
        verify: AiterMLADCPVerifyMetadata,
        kv_c_and_k_pe_cache: torch.Tensor,
        layer: AttentionLayer,
        out_dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """gfx942 causal verify using the LDS-safe generic split-KV kernel.

        Each verify token is its own row whose causal window is its
        ``row_lens`` entry. A zero-length row (an empty DCP shard or graph
        padding) comes back as a NaN output with an LSE of -inf; the DCP
        combine gives it zero weight and never reads the output.
        """
        num_rows, num_heads = q_mla.shape[:2]
        # Stage 2 stores every output and LSE element, empty rows included.
        output = torch.empty(
            num_rows, num_heads, self.kv_lora_rank, dtype=out_dtype, device=q_mla.device
        )
        # Matches merge_mla_segments_triton and the fp32 LSE the DCP combine
        # kernel is warmed up for.
        lse = torch.empty(num_rows, num_heads, dtype=torch.float32, device=q_mla.device)
        if envs.VLLM_BATCH_INVARIANT:
            num_kv_splits = 1
        else:
            num_kv_splits = _triton_dcp_verify_num_kv_splits(
                verify.max_kv_seq_len, num_rows, num_heads, self._sm_count
            )
        logits_shape = (num_rows, num_heads, num_kv_splits, self.kv_lora_rank + 1)
        if is_workspace_manager_initialized():
            (attn_logits,) = current_workspace_manager().get_simultaneous(
                (logits_shape, torch.float32),
            )
        else:
            attn_logits = torch.empty(
                logits_shape, dtype=torch.float32, device=q_mla.device
            )

        paged_kv = kv_c_and_k_pe_cache.unsqueeze(2)
        decode_attention_fwd(
            q_mla,
            paged_kv,
            paged_kv[..., : self.kv_lora_rank],
            output,
            lse,
            verify.block_table,
            verify.row_lens,
            attn_logits,
            num_kv_splits,
            self.scale,
            paged_kv.size(1),
            k_scale=layer._k_scale,
            v_scale=layer._k_scale,
            is_mla=True,
        )
        return output, lse

    def _forward_triton_verify(
        self,
        q_mla: torch.Tensor,
        decode: AiterMLADecodeMetadata,
        kv_c_and_k_pe_cache: torch.Tensor,
        layer: AttentionLayer,
    ) -> torch.Tensor:
        """Causal verify past qlen 4 where no asm decode kernel is correct.

        One Triton program covers a block of a request's (verify token, head)
        rows, so each KV tile is read once for all of them, and each row is
        bounded by its own causal window. A graph padding request (or a
        varlen tail past its real length) comes back as zeros.
        """
        assert decode.max_qo_len is not None
        num_rows, num_heads = q_mla.shape[:2]
        output = torch.empty(
            num_rows,
            num_heads,
            self.kv_lora_rank,
            dtype=decode.attn_out_dtype,
            device=q_mla.device,
        )
        # Only the DCP merge would read the LSE; the kernel writes it anyway.
        lse = torch.empty(num_rows, num_heads, dtype=torch.float32, device=q_mla.device)
        num_kv_splits = decode.triton_verify_num_kv_splits
        logits_shape = (num_rows, num_heads, num_kv_splits, self.kv_lora_rank + 1)
        # A single split writes the output directly and needs no scratch.
        attn_logits = None
        if num_kv_splits > 1 and is_workspace_manager_initialized():
            (attn_logits,) = current_workspace_manager().get_simultaneous(
                (logits_shape, torch.float32),
            )
        elif num_kv_splits > 1:
            attn_logits = torch.empty(
                logits_shape, dtype=torch.float32, device=q_mla.device
            )
        paged_kv = kv_c_and_k_pe_cache.unsqueeze(2)
        if (
            not self._triton_verify_warmed_up
            and not torch.cuda.is_current_stream_capturing()
        ):
            self._triton_verify_warmed_up = True
            warmup_decode_attention_fwd_verify(
                q_mla,
                paged_kv,
                output,
                lse,
                decode.block_table,
                decode.triton_verify_real_seq_lens is not None,
                decode.triton_verify_qlens,
                attn_logits,
                self.scale,
                paged_kv.size(1),
                k_scale=layer._k_scale,
            )
        decode_attention_fwd_verify(
            q_mla,
            paged_kv,
            paged_kv[..., : self.kv_lora_rank],
            output,
            lse,
            decode.block_table,
            decode.seq_lens,
            decode.triton_verify_real_seq_lens,
            decode.max_qo_len,
            attn_logits,
            num_kv_splits,
            self.scale,
            paged_kv.size(1),
            k_scale=layer._k_scale,
        )
        return output

    def forward_mqa(
        self,
        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        kv_c_and_k_pe_cache: torch.Tensor,
        attn_metadata: AiterMLAMetadata,
        layer: AttentionLayer,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        decode = attn_metadata.decode
        if decode is None or decode.varlen_q_rows is None:
            return self._forward_mqa(q, kv_c_and_k_pe_cache, attn_metadata, layer)
        assert decode.varlen_o_rows is not None
        q_rows = decode.varlen_q_rows
        if isinstance(q, tuple):
            num_tokens = q[0].shape[0]
            q = (q[0].index_select(0, q_rows), q[1].index_select(0, q_rows))
        else:
            num_tokens = q.shape[0]
            q = q.index_select(0, q_rows)
        o, lse = self._forward_mqa(q, kv_c_and_k_pe_cache, attn_metadata, layer)
        o_rows = decode.varlen_o_rows[:num_tokens]
        assert lse is None or self.dcp_world_size > 1
        if lse is not None:
            lse = lse.index_select(0, o_rows)
        return o.index_select(0, o_rows), lse

    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

        if decode.triton_verify_num_kv_splits:
            q_mla = torch.cat(q, dim=-1) if isinstance(q, tuple) else q
            if (
                is_quantized_kv_cache(self.kv_cache_dtype)
                and q_mla.dtype != torch.bfloat16
            ):
                q_mla = q_mla.to(torch.bfloat16) * layer._q_scale
            return (
                self._forward_triton_verify(q_mla, decode, kv_c_and_k_pe_cache, layer),
                None,
            )

        dcp_route = decode.dcp_route
        if dcp_route is _DCPDecodeRoute.TRITON:
            verify = decode.dcp_verify
            assert verify is not None
            q_mla = torch.cat(q, dim=-1) if isinstance(q, tuple) else q
            if (
                is_quantized_kv_cache(self.kv_cache_dtype)
                and q_mla.dtype != torch.bfloat16
            ):
                q_mla = q_mla.to(torch.bfloat16) * layer._q_scale
            return self._forward_triton_dcp_verify(
                q_mla,
                verify,
                kv_c_and_k_pe_cache,
                layer,
                decode.attn_out_dtype,
            )

        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 decode.fp8_v_scale_fix:
                o.mul_(layer._k_scale)

        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.

_compute_prefill_context(q, kv_c_and_k_pe_cache, attn_metadata, k_scale)

Attend the cached prefix, gathering and expanding it in one kernel

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

def _compute_prefill_context(
    self,
    q: torch.Tensor,
    kv_c_and_k_pe_cache: torch.Tensor,
    attn_metadata: MLACommonMetadata,
    k_scale: torch.Tensor,
):
    """Attend the cached prefix, gathering and expanding it in one kernel"""
    assert isinstance(attn_metadata, AiterMLAMetadata)
    assert attn_metadata.prefill is not None
    prefill_metadata = attn_metadata.prefill

    if not self._can_fuse_context_gather(q, prefill_metadata, kv_c_and_k_pe_cache):
        return super()._compute_prefill_context(
            q, kv_c_and_k_pe_cache, attn_metadata, k_scale
        )
    kv_indices = attn_metadata.context_chunk_kv_indices
    if kv_indices is None:
        return super()._compute_prefill_context(
            q, kv_c_and_k_pe_cache, attn_metadata, k_scale
        )

    assert prefill_metadata.prefill_backend is not None
    chunked_context = prefill_metadata.chunked_context
    assert chunked_context is not None

    output = None
    output_lse = None
    for chunk in chunked_context.chunks:
        k, v = self._gather_context_chunk(
            chunk,
            kv_c_and_k_pe_cache,
            kv_indices[chunk.index],
            k_scale,
            q.dtype,
        )

        attn_output, attn_softmax_lse = (
            prefill_metadata.prefill_backend.run_prefill_context_chunk(
                chunk=chunk,
                q=q[chunk.token_slice],
                k=k,
                v=v,
            )
        )

        if output is None:
            if (
                len(chunked_context.chunks) == 1
                and not chunked_context.empty_token_slices
            ):
                return attn_output, attn_softmax_lse
            output, output_lse = init_mla_context_partial(
                chunked_context,
                attn_output,
                attn_softmax_lse,
                num_tokens=q.shape[0],
            )
        accumulate_mla_context_chunk(
            chunk, attn_output, attn_softmax_lse, output, output_lse
        )
    return output, output_lse

_concat_k_nope_k_pe(k_nope, k_pe)

Build [k_nope | k_pe] with one kernel, fall back to the base copies.

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

def _concat_k_nope_k_pe(
    self, k_nope: torch.Tensor, k_pe: torch.Tensor
) -> torch.Tensor:
    """Build [k_nope | k_pe] with one kernel, fall back to the base copies."""
    if not (
        self._use_fused_mla_kv_concat
        and k_nope.dim() == 3
        and k_nope.dtype in (torch.bfloat16, torch.float16)
        and k_pe.dtype == k_nope.dtype
    ):
        return super()._concat_k_nope_k_pe(k_nope, k_pe)
    k = torch.empty(
        (*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]),
        dtype=k_nope.dtype,
        device=k_nope.device,
    )
    torch.ops._C.fused_kimi_k3_mla_kv_concat(
        k_nope, k_pe.reshape(k_pe.shape[0], k_pe.shape[-1]), k
    )
    return k

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

_forward_triton_dcp_verify(q_mla, verify, kv_c_and_k_pe_cache, layer, out_dtype)

gfx942 causal verify using the LDS-safe generic split-KV kernel.

Each verify token is its own row whose causal window is its row_lens entry. A zero-length row (an empty DCP shard or graph padding) comes back as a NaN output with an LSE of -inf; the DCP combine gives it zero weight and never reads the output.

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

def _forward_triton_dcp_verify(
    self,
    q_mla: torch.Tensor,
    verify: AiterMLADCPVerifyMetadata,
    kv_c_and_k_pe_cache: torch.Tensor,
    layer: AttentionLayer,
    out_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
    """gfx942 causal verify using the LDS-safe generic split-KV kernel.

    Each verify token is its own row whose causal window is its
    ``row_lens`` entry. A zero-length row (an empty DCP shard or graph
    padding) comes back as a NaN output with an LSE of -inf; the DCP
    combine gives it zero weight and never reads the output.
    """
    num_rows, num_heads = q_mla.shape[:2]
    # Stage 2 stores every output and LSE element, empty rows included.
    output = torch.empty(
        num_rows, num_heads, self.kv_lora_rank, dtype=out_dtype, device=q_mla.device
    )
    # Matches merge_mla_segments_triton and the fp32 LSE the DCP combine
    # kernel is warmed up for.
    lse = torch.empty(num_rows, num_heads, dtype=torch.float32, device=q_mla.device)
    if envs.VLLM_BATCH_INVARIANT:
        num_kv_splits = 1
    else:
        num_kv_splits = _triton_dcp_verify_num_kv_splits(
            verify.max_kv_seq_len, num_rows, num_heads, self._sm_count
        )
    logits_shape = (num_rows, num_heads, num_kv_splits, self.kv_lora_rank + 1)
    if is_workspace_manager_initialized():
        (attn_logits,) = current_workspace_manager().get_simultaneous(
            (logits_shape, torch.float32),
        )
    else:
        attn_logits = torch.empty(
            logits_shape, dtype=torch.float32, device=q_mla.device
        )

    paged_kv = kv_c_and_k_pe_cache.unsqueeze(2)
    decode_attention_fwd(
        q_mla,
        paged_kv,
        paged_kv[..., : self.kv_lora_rank],
        output,
        lse,
        verify.block_table,
        verify.row_lens,
        attn_logits,
        num_kv_splits,
        self.scale,
        paged_kv.size(1),
        k_scale=layer._k_scale,
        v_scale=layer._k_scale,
        is_mla=True,
    )
    return output, lse

_forward_triton_verify(q_mla, decode, kv_c_and_k_pe_cache, layer)

Causal verify past qlen 4 where no asm decode kernel is correct.

One Triton program covers a block of a request's (verify token, head) rows, so each KV tile is read once for all of them, and each row is bounded by its own causal window. A graph padding request (or a varlen tail past its real length) comes back as zeros.

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

def _forward_triton_verify(
    self,
    q_mla: torch.Tensor,
    decode: AiterMLADecodeMetadata,
    kv_c_and_k_pe_cache: torch.Tensor,
    layer: AttentionLayer,
) -> torch.Tensor:
    """Causal verify past qlen 4 where no asm decode kernel is correct.

    One Triton program covers a block of a request's (verify token, head)
    rows, so each KV tile is read once for all of them, and each row is
    bounded by its own causal window. A graph padding request (or a
    varlen tail past its real length) comes back as zeros.
    """
    assert decode.max_qo_len is not None
    num_rows, num_heads = q_mla.shape[:2]
    output = torch.empty(
        num_rows,
        num_heads,
        self.kv_lora_rank,
        dtype=decode.attn_out_dtype,
        device=q_mla.device,
    )
    # Only the DCP merge would read the LSE; the kernel writes it anyway.
    lse = torch.empty(num_rows, num_heads, dtype=torch.float32, device=q_mla.device)
    num_kv_splits = decode.triton_verify_num_kv_splits
    logits_shape = (num_rows, num_heads, num_kv_splits, self.kv_lora_rank + 1)
    # A single split writes the output directly and needs no scratch.
    attn_logits = None
    if num_kv_splits > 1 and is_workspace_manager_initialized():
        (attn_logits,) = current_workspace_manager().get_simultaneous(
            (logits_shape, torch.float32),
        )
    elif num_kv_splits > 1:
        attn_logits = torch.empty(
            logits_shape, dtype=torch.float32, device=q_mla.device
        )
    paged_kv = kv_c_and_k_pe_cache.unsqueeze(2)
    if (
        not self._triton_verify_warmed_up
        and not torch.cuda.is_current_stream_capturing()
    ):
        self._triton_verify_warmed_up = True
        warmup_decode_attention_fwd_verify(
            q_mla,
            paged_kv,
            output,
            lse,
            decode.block_table,
            decode.triton_verify_real_seq_lens is not None,
            decode.triton_verify_qlens,
            attn_logits,
            self.scale,
            paged_kv.size(1),
            k_scale=layer._k_scale,
        )
    decode_attention_fwd_verify(
        q_mla,
        paged_kv,
        paged_kv[..., : self.kv_lora_rank],
        output,
        lse,
        decode.block_table,
        decode.seq_lens,
        decode.triton_verify_real_seq_lens,
        decode.max_qo_len,
        attn_logits,
        num_kv_splits,
        self.scale,
        paged_kv.size(1),
        k_scale=layer._k_scale,
    )
    return output

_gather_context_chunk(chunk, kv_c_and_k_pe_cache, kv_indices, k_scale, out_dtype)

Gather one context chunk and expand it into (k, v).

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

def _gather_context_chunk(
    self,
    chunk,
    kv_c_and_k_pe_cache: torch.Tensor,
    kv_indices: torch.Tensor,
    k_scale: torch.Tensor,
    out_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Gather one context chunk and expand it into (k, v)."""
    from aiter.ops.triton.gather_kv_b_proj import gather_kv_b_proj

    num_blocks, block_size, cache_width = kv_c_and_k_pe_cache.shape
    rope_dim = cache_width - self.kv_lora_rank
    num_rows = chunk.num_context_tokens
    k = torch.empty(
        (num_rows, self.num_heads, self.qk_nope_head_dim + rope_dim),
        dtype=out_dtype,
        device=kv_c_and_k_pe_cache.device,
    )
    v = torch.empty(
        (num_rows, self.num_heads, self.v_head_dim),
        dtype=out_dtype,
        device=kv_c_and_k_pe_cache.device,
    )
    if num_rows == 0:
        return k, v

    gather_kv_b_proj(
        kv_c_and_k_pe_cache.view(num_blocks * block_size, 1, cache_width),
        k_scale,
        chunk.cu_seq_lens,
        kv_indices,
        chunk.cu_seq_lens,
        self.kv_b_proj.weight,
        None,
        k,
        v,
    )
    return k, v

_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 every padded head a duplicate of
    # a real head, so the real heads are bit-identical to the unpadded
    # result. get_mla_padded_q repeat_interleaves when the target is a
    # multiple of num_heads (real heads at stride target // num_heads) and
    # tiles otherwise (real heads [0:num_heads]). 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)

    if self._fp8_prefill_one_scale is None:
        self._fp8_prefill_one_scale = torch.ones(
            (), dtype=torch.float32, device=q.device
        )
    one_scale = self._fp8_prefill_one_scale

    # 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:
        real_out = (
            out_3d[:, :: nhead // _real_nhead, :]
            if nhead % _real_nhead == 0
            else out_3d[:, :_real_nhead, :]
        )
        out.view(total_q, _real_nhead, v_head_dim).copy_(real_out)

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, the selected prefill backend supersedes it, 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, the selected prefill backend supersedes
    it, 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.
    """
    prefill_metadata = attn_metadata.prefill
    if (
        not self._fp8_prefill_enabled
        or not isinstance(attn_metadata, AiterMLAMetadata)
        or prefill_metadata is None
        or _asm_fp8_prefill_superseded(prefill_metadata.prefill_backend)
        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,
        )

    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)
    if self._fused_prefill_qkv:
        from aiter.ops.triton.fusions.fused_mla_prefill_qkv_fp8 import (
            fused_mla_prefill_qkv_fp8,
        )

        from vllm.platforms import current_platform

        # The k_pe concat, the head padding and the FP8 cast in one launch.
        q, k, v = fused_mla_prefill_qkv_fp8(
            q,
            k_nope,
            v,
            k_pe,
            AiterMLAHelper.get_fp8_prefill_num_heads(self.num_heads),
            current_platform.fp8_dtype(),
        )
    else:
        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 % num_reqs != 0:
            return 0

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

        # qo_len is derived from query_start_loc_cpu, so it is already
        # host-side and tiny (one entry per request). Reading it out once and
        # deciding on Python ints avoids eight separate tensor dispatches --
        # sum/item, boolean-mask index, all, eq, nonzero, item, any -- on the
        # metadata-build path, which runs between decode steps while the GPUs
        # are idle.
        lens = qo_len.tolist()
        if num_decode_tokens <= sum(lens):
            return 0

        first_zero = -1
        num_positive = 0
        for i, qlen in enumerate(lens):
            if qlen > 0:
                num_positive += 1
                # A positive entry after a zero breaks the contiguous live
                # prefix that the padded layout requires.
                if first_zero >= 0:
                    return 0
            elif first_zero < 0:
                first_zero = i

        if num_positive == num_reqs:
            return 0
        if num_positive > 0:
            if max_qo_len != uniform_qo_len:
                return 0
            if any(qlen != uniform_qo_len for qlen in lens if qlen > 0):
                return 0

        return uniform_qo_len

    @property
    def _varlen_right_align(self) -> bool:
        return self.dcp_world_size > 1

    @classmethod
    def get_varlen_cudagraph_max_query_len(
        cls: type["AiterMLAMetadataBuilder"],
        vllm_config: VllmConfig,
        kv_cache_spec: KVCacheSpec,
    ) -> int | None:
        if not _adaptive_varlen_verify_enabled(vllm_config):
            return None
        if (
            cls.get_cudagraph_support(vllm_config, kv_cache_spec)
            != AttentionCGSupport.UNIFORM_BATCH
        ):
            return None
        return max_decode_query_len(vllm_config)

    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,
        )
        supports_triton_dcp_verify = _triton_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,
        )
        super().__init__(
            kv_cache_spec,
            layer_names,
            vllm_config,
            device,
            AiterMLAMetadata,
            supports_dcp_with_varlen=(
                supports_segmented_dcp_verify
                or asm_dcp_verify_config
                or supports_triton_dcp_verify
            ),
        )
        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 self._asm_dcp_verify_heads:
                self._asm_dcp_verify = True
            elif not supports_segmented_dcp_verify:
                raise ValueError(
                    "ROCM_AITER_MLA DCP multi-token verify has no kernel for "
                    f"{self.num_heads * self.dcp_world_size} DCP-gathered heads: "
                    "the round-robin asm decode only has native counts "
                    f"{_NATIVE_CPRR_HEADS}, and this AITER build lacks segmented "
                    "MLA decode."
                )
        # The group's cache dtype can differ from cache_config.cache_dtype
        # (a speculative kv_cache_dtype override, kv_cache_dtype_skip_layers).
        # Use the group's original format: decode normalizes E5M2 to "fp8".
        kv_cache_dtype = getattr(kv_cache_spec, "cache_dtype_str", None)
        if kv_cache_dtype is None:
            kv_cache_dtype = vllm_config.cache_config.cache_dtype
        self._supports_segmented_dcp_verify = (
            supports_segmented_dcp_verify
            and _segmented_mla_fits_lds(
                self.num_heads * self.dcp_world_size,
                kv_cache_spec.block_size,
                is_quantized_kv_cache(kv_cache_dtype),
            )
        )
        self._use_triton_dcp_verify = (
            supports_triton_dcp_verify and not self._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._supports_fused_dcp_prefill = _supports_fused_dcp_prefill_config(
            kv_cache_dtype,
            self.q_data_type,
            self.mla_dims.kv_lora_rank,
            self.mla_dims.qk_nope_head_dim,
            self.mla_dims.qk_rope_head_dim,
            self.mla_dims.v_head_dim,
        )

        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

        self._varlen_decode_qlen = 0
        if _adaptive_varlen_verify_enabled(vllm_config):
            assert self._mtp_decode_qlen == max_decode_query_len(vllm_config), (
                self._mtp_decode_qlen,
                max_decode_query_len(vllm_config),
            )
            self._varlen_decode_qlen = self._mtp_decode_qlen
            self.query_len_support = QueryLenSupport.VARLEN  # type: ignore[misc]
            max_rows = vllm_config.scheduler_config.max_num_seqs * (
                self._varlen_decode_qlen
            )
            row_ids = torch.arange(max_rows, dtype=torch.int64, device=device)
            self._varlen_row_req = row_ids // self._varlen_decode_qlen
            self._varlen_row_pos = row_ids % self._varlen_decode_qlen
            self._varlen_row_ids = row_ids
            self._varlen_q_rows = torch.zeros_like(row_ids)
            self._varlen_o_rows = torch.zeros(
                max_rows + 1, dtype=torch.int64, device=device
            )

        # 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
        # Generic Triton indexes physical KV blocks; segmented MLA indexes
        # TILE subpages of them.
        self._dcp_verify_page_size = (
            kv_cache_spec.block_size
            if self._use_triton_dcp_verify
            else _segmented_mla_page_size(kv_cache_spec.block_size)
        )
        # A DCP rank's shard of the longest sequence bounds every verify row.
        dcp_local_max_seq_len = (
            cdiv(
                vllm_config.model_config.max_model_len,
                self.dcp_world_size * parallel_config.cp_kv_cache_interleave_size,
            )
            * parallel_config.cp_kv_cache_interleave_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
        speculative_config = vllm_config.speculative_config
        if _adaptive_varlen_verify_enabled(vllm_config) and not (
            self._varlen_right_align
        ):
            assert speculative_config is not None
            max_num_pages += max_num_reqs * (
                speculative_config.num_speculative_tokens or 0
            )

        # 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 = kv_cache_dtype
        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
        )
        from vllm.platforms.rocm import on_gfx950

        self._supports_fused_gather = (
            self._kv_cache_dtype_str != "fp8_ds_mla" and on_gfx950()
        )
        # 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
        # The metadata is built for every qlen up to _mtp_decode_qlen, and
        # aiter's buffer needs are not monotonic in qlen: on gfx950 a bf16
        # block of fewer than 128 packed rows is tiled by 64 instead of 128.
        (
            (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),
        ) = _max_mla_metadata_info(
            get_mla_metadata_info_v1(
                max_num_reqs,
                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 {}
                ),
            )
            for qlen in range(1, self._mtp_decode_qlen + 1)
        )
        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.
        # A FlyDSL FP8 prefill backend supersedes the assembly kernel, so the
        # PS workspace buffers below are not allocated either.
        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_fp8_prefill_num_heads(self.num_heads)
            and not _asm_fp8_prefill_superseded(self._prefill_backend)
        )
        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 or self._use_triton_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. Full graphs need the row bound to be constant.
                max_verify_rows = max_num_reqs * self._mtp_decode_qlen
                max_local_pages = cdiv(
                    dcp_local_max_seq_len, self._dcp_verify_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._dcp_verify_page_size,
                    max_kv_seq_len=dcp_local_max_seq_len,
                )
        # gfx942 DCP verify launches the generic split-KV kernel, which draws
        # its partials from the shared workspace. Reserve it before warmup
        # locks the pool.
        if self._use_triton_dcp_verify and self._mtp_decode_qlen > 1:
            _reserve_triton_dcp_verify_workspace(
                max_num_reqs * self._mtp_decode_qlen,
                self._decode_num_heads,
                dcp_local_max_seq_len,
                self.mla_dims.kv_lora_rank,
                current_platform.num_compute_units(),
            )

        # FP8 verify lengths whose asm kernel leaves V unscaled.
        self._fp8_v_scale_fix_qlens: frozenset[int] = frozenset()
        if (
            is_quantized_kv_cache(self._kv_cache_dtype_str)
            and self.dcp_world_size == 1
            and self._mtp_decode_qlen > 1
        ):
            heads = AiterMLAHelper.get_actual_mla_num_heads(self._decode_num_heads)
            self._fp8_v_scale_fix_qlens = frozenset(
                qlen
                for qlen in range(2, min(self._mtp_decode_qlen, 4) + 1)
                if _aiter_fp8_mla_decode_drops_v_scale(
                    heads, qlen, self._kv_cache_dtype_str == "fp8"
                )
            )

        # Verify past qlen 4 with no correct asm kernel takes the Triton one.
        self._use_triton_verify = _asm_causal_verify_unsupported(
            self._decode_num_heads,
            self._mtp_decode_qlen,
            self._kv_cache_dtype_str,
            self.dcp_world_size,
        )
        # Split cap bound; must be static for full CUDA graphs.
        self._triton_verify_max_seq_len = vllm_config.model_config.max_model_len
        self._sm_count = current_platform.num_compute_units()
        self._triton_verify_qlens: tuple[int, ...] = ()
        if self._use_triton_verify:
            self._triton_verify_qlens = tuple(
                range(
                    AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN + 1,
                    self._mtp_decode_qlen + 1,
                )
            )
            logger.info_once(
                "ROCM_AITER_MLA: causal verify past %d tokens uses the Triton "
                "verify kernel; the AITER asm decode has no correct kernel "
                "for it on this GPU.",
                AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN,
            )
            _reserve_triton_verify_workspace(
                max_num_reqs,
                self._triton_verify_qlens,
                self._decode_num_heads,
                self._triton_verify_max_seq_len,
                self.mla_dims.kv_lora_rank,
                self._sm_count,
            )

    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 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._dcp_verify_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, page_size
        )
        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,
        page_size: int,
    ) -> None:
        """Fill one page list per request and broadcast it across qlen rows.

        ``page_size`` is the unit the kernel indexes. Segmented MLA uses TILE
        subpages; generic Triton uses the physical KV block. When those are
        equal, pages_per_block is 1 and this copies block IDs as-is.

        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
        assert kernel_block_size is not None
        assert page_size > 0
        assert kernel_block_size % page_size == 0, (
            f"DCP verify page_size {page_size} must divide kernel_block_size "
            f"{kernel_block_size}"
        )
        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]}"
        )
        if pages_per_block == 1:
            per_req_page_table = block_table[:num_reqs, :max_local_pages]
        else:
            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]
        varlen_decode = self._varlen_decode_qlen > 0 and causal
        if varlen_decode:
            max_qo_len = self._varlen_decode_qlen
            pad_uniform_mtp = True
        else:
            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),
                )

        varlen_q_rows = varlen_o_rows = None
        varlen_left_align = varlen_decode and not self._varlen_right_align
        if varlen_left_align:
            seq_lens_for_kernel = seq_lens_device + (max_qo_len - qo_lens_device)
        if varlen_decode:
            varlen_q_rows, varlen_o_rows = self._fill_varlen_rows(
                query_start_loc_device, qo_lens_device, num_reqs
            )

        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(
                supports_segmented=self._supports_segmented_dcp_verify,
                supports_triton=self._use_triton_dcp_verify,
                causal=causal,
                max_qo_len=int(max_qo_len),
                asm_selected=self._asm_dcp_verify,
            )

        # forward_mqa serves Gluon verify first, and Gluon needs the flat
        # page indices this route skips.
        use_triton_verify = (
            self._use_triton_verify
            and causal
            and not use_gluon_verify
            and max_qo_len > AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN
        )
        triton_verify_num_kv_splits = 0
        if use_triton_verify:
            # The kernel assumes every request carries max_qo_len query rows;
            # varlen verify gathers q to that layout in forward_mqa.
            assert varlen_decode or num_decode_tokens == num_reqs * max_qo_len, (
                num_decode_tokens,
                num_reqs,
                max_qo_len,
            )
            triton_verify_num_kv_splits = _triton_verify_num_kv_splits(
                (
                    self._triton_verify_max_seq_len
                    if self.compilation_config.cudagraph_mode.has_full_cudagraphs()
                    else max_seq_len
                ),
                num_reqs,
                int(max_qo_len),
                self._decode_num_heads,
                self._sm_count,
            )

        # Row-view DCP verify and the folded Triton verify read block tables
        # directly, so the flat per-token view is dead work for them. 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 not _is_row_view_dcp_route(dcp_route) and not use_triton_verify:
            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
            if varlen_left_align:
                max_tokens_per_req += max_qo_len - 1
            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,
                seq_lens_device,
                KERNEL_BLOCK_SIZE=self.kernel_block_size,
                BLOCK_SIZE=1024,
                CLAMP_TAIL=varlen_left_align,
            )
            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. Gluon and the row-view DCP routes
        # (segmented and gfx942 Triton) cover the non-ASM decode/verify paths.
        use_persistent_metadata = (
            not use_gluon_decode
            and not use_gluon_verify
            and not _is_row_view_dcp_route(dcp_route)
            and not use_triton_verify
            # A padded rank has no bf16 persistent kernel past qlen 4 where the
            # gfx950 fold is absent; the non-persistent entry covers it on
            # gfx950, and CDNA3 takes the Triton verify above. 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 _is_row_view_dcp_route(dcp_route):
            assert causal
            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,
            triton_verify_num_kv_splits=triton_verify_num_kv_splits,
            triton_verify_qlens=self._triton_verify_qlens,
            fp8_v_scale_fix=(
                has_persistent_metadata
                and self.dcp_world_size == 1
                and int(max_qo_len) in self._fp8_v_scale_fix_qlens
            ),
            triton_verify_real_seq_lens=(
                seq_lens_device if use_triton_verify and varlen_decode else None
            ),
            attn_out_dtype=self.decode_attn_out_dtype,
            has_persistent_metadata=has_persistent_metadata,
            varlen_q_rows=varlen_q_rows,
            varlen_o_rows=varlen_o_rows,
        )

        return attn_metadata

    def _build_context_chunk_kv_indices(self, prefill) -> list[torch.Tensor]:
        """Flatten each context chunk's block table into per-token KV indices.

        Built here so the expansion runs once per step instead of once per
        MLA layer.
        """
        indices = []
        for chunk in prefill.chunked_context.chunks:
            block_table = prefill.block_table[chunk.request_slice]
            cu_seq_lens = chunk.cu_seq_lens
            kv_indices = torch.empty(
                chunk.num_context_tokens,
                dtype=torch.int32,
                device=block_table.device,
            )
            num_chunks = max(1, cdiv(chunk.max_seq_len, 1024))
            _expand_page_indices_kernel[(cu_seq_lens.shape[0] - 1, num_chunks)](
                kv_indices,
                block_table,
                block_table.stride(0),
                block_table.stride(1),
                cu_seq_lens,
                None,
                KERNEL_BLOCK_SIZE=self.kernel_block_size,
                BLOCK_SIZE=1024,
                HAS_START_OFFSETS=True,
                start_offsets=chunk.starts,
            )
            indices.append(kv_indices)
        return indices

    def _fill_varlen_rows(
        self,
        query_start_loc_device: torch.Tensor,
        qo_lens_device: torch.Tensor,
        num_reqs: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        width = self._varlen_decode_qlen
        num_rows = num_reqs * width
        row_req = self._varlen_row_req[:num_rows]
        row_pos = self._varlen_row_pos[:num_rows]
        row_qlen = qo_lens_device.to(torch.int64).index_select(0, row_req)
        row_start = query_start_loc_device.to(torch.int64).index_select(0, row_req)
        if self._varlen_right_align:
            row_pos = row_pos - (width - row_qlen)
        is_real = (row_pos >= 0) & (row_pos < row_qlen)
        packed_row = row_start + row_pos
        q_rows = self._varlen_q_rows[:num_rows]
        q_rows.copy_(torch.where(is_real, packed_row, 0))
        o_rows = self._varlen_o_rows
        o_rows.zero_()
        dump_slot = o_rows.size(0) - 1
        o_rows.scatter_(
            0,
            torch.where(is_real, packed_row, dump_slot),
            self._varlen_row_ids[:num_rows],
        )
        return q_rows, o_rows

    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
        )
        prefill = attn_metadata.prefill
        if (
            prefill is not None
            and prefill.chunked_context is not None
            and self.dcp_world_size == 1
            and self._supports_fused_gather
        ):
            attn_metadata.context_chunk_kv_indices = (
                self._build_context_chunk_kv_indices(prefill)
            )
        if (
            prefill is not None
            and prefill.chunked_context is not None
            and self.dcp_world_size > 1
            and self._supports_fused_dcp_prefill
        ):
            attn_metadata.dcp_context_row_indices = [
                context_row_indices(chunk, prefill.block_table.device)
                for chunk in prefill.chunked_context.chunks
            ]
        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_context_chunk_kv_indices(prefill)

Flatten each context chunk's block table into per-token KV indices.

Built here so the expansion runs once per step instead of once per MLA layer.

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

def _build_context_chunk_kv_indices(self, prefill) -> list[torch.Tensor]:
    """Flatten each context chunk's block table into per-token KV indices.

    Built here so the expansion runs once per step instead of once per
    MLA layer.
    """
    indices = []
    for chunk in prefill.chunked_context.chunks:
        block_table = prefill.block_table[chunk.request_slice]
        cu_seq_lens = chunk.cu_seq_lens
        kv_indices = torch.empty(
            chunk.num_context_tokens,
            dtype=torch.int32,
            device=block_table.device,
        )
        num_chunks = max(1, cdiv(chunk.max_seq_len, 1024))
        _expand_page_indices_kernel[(cu_seq_lens.shape[0] - 1, num_chunks)](
            kv_indices,
            block_table,
            block_table.stride(0),
            block_table.stride(1),
            cu_seq_lens,
            None,
            KERNEL_BLOCK_SIZE=self.kernel_block_size,
            BLOCK_SIZE=1024,
            HAS_START_OFFSETS=True,
            start_offsets=chunk.starts,
        )
        indices.append(kv_indices)
    return indices

_build_dcp_verify_row_view(qlen, block_table, dcp_tot_seq_lens)

Build one paged-KV row per verify token for 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 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._dcp_verify_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, page_size
    )
    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, page_size)

Fill one page list per request and broadcast it across qlen rows.

page_size is the unit the kernel indexes. Segmented MLA uses TILE subpages; generic Triton uses the physical KV block. When those are equal, pages_per_block is 1 and this copies block IDs as-is.

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,
    page_size: int,
) -> None:
    """Fill one page list per request and broadcast it across qlen rows.

    ``page_size`` is the unit the kernel indexes. Segmented MLA uses TILE
    subpages; generic Triton uses the physical KV block. When those are
    equal, pages_per_block is 1 and this copies block IDs as-is.

    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
    assert kernel_block_size is not None
    assert page_size > 0
    assert kernel_block_size % page_size == 0, (
        f"DCP verify page_size {page_size} must divide kernel_block_size "
        f"{kernel_block_size}"
    )
    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]}"
    )
    if pages_per_block == 1:
        per_req_page_table = block_table[:num_reqs, :max_local_pages]
    else:
        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_fp8_mla_decode_drops_v_scale(num_heads, qlen, q_fp8) cached

Whether AITER's FP8 persistent MLA decode leaves V unscaled at a shape.

gfx942's (gqa 16, qseqlen 4) FP8 kernel applies kv_scale to the scores but not to V, so its output is short by exactly kv_scale. Measured on a tiny causal block with kv_scale 2 rather than assumed per arch, so a fixed AITER build turns the correction off by itself.

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

@functools.cache
def _aiter_fp8_mla_decode_drops_v_scale(num_heads: int, qlen: int, q_fp8: bool) -> bool:
    """Whether AITER's FP8 persistent MLA decode leaves V unscaled at a shape.

    gfx942's (gqa 16, qseqlen 4) FP8 kernel applies kv_scale to the scores but
    not to V, so its output is short by exactly kv_scale. Measured on a tiny
    causal block with kv_scale 2 rather than assumed per arch, so a fixed
    AITER build turns the correction off by itself.
    """
    try:
        from aiter import dtypes, get_mla_metadata_info_v1, get_mla_metadata_v1
        from aiter.mla import mla_decode_fwd

        device = torch.device("cuda", torch.accelerator.current_device_index())
        gen = torch.Generator(device=device).manual_seed(0)
        kv_len, head_dim, kv_lora_rank, scale = 64, 576, 512, 576**-0.5
        q_dtype = dtypes.fp8 if q_fp8 else dtypes.bf16
        q = torch.randn(qlen, num_heads, head_dim, device=device, generator=gen)
        q = (q * 0.5).to(q_dtype)
        kv = torch.randn(kv_len, head_dim, device=device, generator=gen)
        kv = (kv * 0.5).to(dtypes.fp8)
        qo_indptr = torch.tensor([0, qlen], dtype=torch.int32, device=device)
        kv_indptr = torch.tensor([0, kv_len], dtype=torch.int32, device=device)
        kv_indices = torch.arange(kv_len, dtype=torch.int32, device=device)
        last_page_len = torch.ones(1, dtype=torch.int32, device=device)
        work = [
            torch.empty(size, dtype=dtype, device=device)
            for size, dtype in get_mla_metadata_info_v1(
                1, qlen, num_heads, q_dtype, dtypes.fp8, is_sparse=False, fast_mode=True
            )
        ]
        get_mla_metadata_v1(
            qo_indptr,
            kv_indptr,
            last_page_len,
            num_heads,
            1,
            True,
            work[0],
            work[2],
            work[1],
            work[3],
            work[4],
            work[5],
            page_size=1,
            kv_granularity=16,
            max_seqlen_qo=qlen,
            uni_seqlen_qo=qlen,
            fast_mode=True,
            dtype_q=q_dtype,
            dtype_kv=dtypes.fp8,
        )
        out = torch.empty(
            qlen, num_heads, kv_lora_rank, dtype=torch.bfloat16, device=device
        )
        one = torch.ones(1, dtype=torch.float32, device=device)
        mla_decode_fwd(
            q,
            kv.view(-1, 1, 1, head_dim),
            out,
            qo_indptr,
            kv_indptr,
            kv_indices,
            last_page_len,
            qlen,
            sm_scale=scale,
            logit_cap=0.0,
            q_scale=one,
            kv_scale=one * 2,
            causal=True,
            work_meta_data=work[0],
            work_indptr=work[1],
            work_info_set=work[2],
            reduce_indptr=work[3],
            reduce_final_map=work[4],
            reduce_partial_map=work[5],
        )
        k = kv.float() * 2
        scores = torch.einsum("qhd,kd->qhk", q.float(), k) * scale
        visible = torch.arange(kv_len, device=device)[None, :] <= (
            torch.arange(qlen, device=device)[:, None] + kv_len - qlen
        )
        scores = scores.masked_fill(~visible[:, None, :], float("-inf"))
        ref = torch.einsum("qhk,kd->qhd", scores.softmax(-1), k[:, :kv_lora_rank])
        err_scaled = (out.float() - ref).abs().max().item()
        err_unscaled = (out.float() - ref / 2).abs().max().item()
    except Exception as exc:  # noqa: BLE001
        logger.warning(
            "AITER FP8 MLA V-scale probe failed at %d heads, qlen %d, so its "
            "V-scale correction stays off: %s",
            num_heads,
            qlen,
            exc,
        )
        return False
    return err_unscaled < 0.05 < err_scaled

_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_causal_verify_unsupported(decode_num_heads, mtp_decode_qlen, kv_cache_dtype, dcp_world_size)

Whether CDNA3's asm decode has no correct kernel for causal verify past qlen 4: its persistent kernels stop at qseqlen 4 (bf16, and FP8 below 128 heads) and its non-persistent entries are not usable here (ROCm/aiter#5679). DCP verify keeps its row-view and cprr routes.

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

def _asm_causal_verify_unsupported(
    decode_num_heads: int,
    mtp_decode_qlen: int,
    kv_cache_dtype: str,
    dcp_world_size: int,
) -> bool:
    """Whether CDNA3's asm decode has no correct kernel for causal verify past
    qlen 4: its persistent kernels stop at qseqlen 4 (bf16, and FP8 below 128
    heads) and its non-persistent entries are not usable here
    (ROCm/aiter#5679). DCP verify keeps its row-view and cprr routes.
    """
    if dcp_world_size > 1:
        return False
    if mtp_decode_qlen <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN:
        return False
    from vllm.platforms.rocm import get_cdna_version

    if get_cdna_version() != 3:
        return False
    if is_quantized_kv_cache(kv_cache_dtype):
        return AiterMLAHelper.get_actual_mla_num_heads(decode_num_heads) != 128
    return True

_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

_asm_fp8_prefill_superseded(prefill_backend)

Whether the selected MLA prefill backend supersedes the ASM FP8 prefill.

AiterFlyDSLFP8PrefillBackend runs FP8 attention itself and also serves context-free prefills, so the assembly prefill kernel must not claim them.

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

def _asm_fp8_prefill_superseded(prefill_backend: object) -> bool:
    """Whether the selected MLA prefill backend supersedes the ASM FP8 prefill.

    ``AiterFlyDSLFP8PrefillBackend`` runs FP8 attention itself and also serves
    context-free prefills, so the assembly prefill kernel must not claim them.
    """
    return isinstance(prefill_backend, AiterFlyDSLFP8PrefillBackend)

_expand_page_indices_kernel(page_indices, block_table, block_table_stride_0, block_table_stride_1, cu_num_tokens, num_real_tokens, KERNEL_BLOCK_SIZE, BLOCK_SIZE, CLAMP_TAIL=False, HAS_START_OFFSETS=False, start_offsets=None)

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

With CLAMP_TAIL, entries past num_real_tokens[req] (padding a request's page list out to its kernel length) repeat the request's first page, so they never index past its block-table row.

With HAS_START_OFFSETS the request's rows begin start_offsets[req] tokens into its sequence rather than at 0, which is what a chunked-context prefill needs.

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,
    num_real_tokens,
    KERNEL_BLOCK_SIZE: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    CLAMP_TAIL: tl.constexpr = False,
    HAS_START_OFFSETS: tl.constexpr = False,
    start_offsets=None,
):
    """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).

    With CLAMP_TAIL, entries past num_real_tokens[req] (padding a request's
    page list out to its kernel length) repeat the request's first page, so
    they never index past its block-table row.

    With HAS_START_OFFSETS the request's rows begin start_offsets[req] tokens
    into its sequence rather than at 0, which is what a chunked-context prefill
    needs.
    """
    # 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
    if CLAMP_TAIL:
        real_tokens = tl.load(num_real_tokens + req_idx)

    start = tl.load(start_offsets + req_idx) if HAS_START_OFFSETS else 0

    # 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
    source_offsets = token_offsets
    if CLAMP_TAIL:
        source_offsets = tl.where(token_offsets < real_tokens, token_offsets, 0)
    seq_positions = source_offsets + start

    # Which block in the block table does this token belong to?
    block_idx = seq_positions // KERNEL_BLOCK_SIZE
    # Offset within that block
    offset_in_block = seq_positions % 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

_fused_mla_prefill_qkv_supported() cached

Gate fused_mla_prefill_qkv_fp8 to CDNA4 (gfx950).

Older AITER builds without the kernel fall back to the eager concat/cast.

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

@functools.lru_cache(maxsize=1)
def _fused_mla_prefill_qkv_supported() -> bool:
    """Gate ``fused_mla_prefill_qkv_fp8`` to CDNA4 (gfx950).

    Older AITER builds without the kernel fall back to the eager concat/cast.
    """
    from vllm.platforms import current_platform

    if not current_platform.is_rocm():
        return False
    from vllm.platforms.rocm import get_cdna_version

    if get_cdna_version() != 4:
        return False
    try:
        from aiter.ops.triton.fusions.fused_mla_prefill_qkv_fp8 import (  # noqa: F401
            fused_mla_prefill_qkv_fp8,
        )
    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()

_is_row_view_dcp_route(route)

Routes that attend one causal window per verify token.

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

def _is_row_view_dcp_route(route: _DCPDecodeRoute) -> bool:
    """Routes that attend one causal window per verify token."""
    return route in (_DCPDecodeRoute.SEGMENTED, _DCPDecodeRoute.TRITON)

_max_mla_metadata_info(infos)

Per-dimension maximum of several get_mla_metadata_info_v1 results.

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

def _max_mla_metadata_info(
    infos: Iterable[tuple[tuple[int | tuple[int, ...], torch.dtype], ...]],
) -> list[tuple[tuple[int, ...], torch.dtype]]:
    """Per-dimension maximum of several ``get_mla_metadata_info_v1`` results."""
    merged = []
    for entries in zip(*infos):
        shapes = [
            shape if isinstance(shape, tuple) else (shape,) for shape, _ in entries
        ]
        merged.append((tuple(max(dims) for dims in zip(*shapes)), entries[0][1]))
    return merged

_reserve_triton_dcp_verify_workspace(max_rows, num_heads, max_seq_len, kv_lora_rank, sm_count)

Reserve split-KV scratch before warmup locks the workspace manager.

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

def _reserve_triton_dcp_verify_workspace(
    max_rows: int,
    num_heads: int,
    max_seq_len: int,
    kv_lora_rank: int,
    sm_count: int,
) -> None:
    """Reserve split-KV scratch before warmup locks the workspace manager."""
    if not is_workspace_manager_initialized():
        return
    # Fewer rows may take more splits, so size for the largest product.
    max_row_splits = max(
        rows * _triton_dcp_verify_num_kv_splits(max_seq_len, rows, num_heads, sm_count)
        for rows in range(1, max_rows + 1)
    )
    current_workspace_manager().get_simultaneous(
        ((max_row_splits, num_heads, kv_lora_rank + 1), torch.float32),
    )

_reserve_triton_verify_workspace(max_num_reqs, qlens, num_heads, max_seq_len, kv_lora_rank, sm_count)

Reserve the folded verify kernel's split scratch before warmup.

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

def _reserve_triton_verify_workspace(
    max_num_reqs: int,
    qlens: tuple[int, ...],
    num_heads: int,
    max_seq_len: int,
    kv_lora_rank: int,
    sm_count: int,
) -> None:
    """Reserve the folded verify kernel's split scratch before warmup."""
    if not is_workspace_manager_initialized():
        return
    # Fewer requests or verify tokens may take more splits, so size for the
    # largest launch. A single split needs no scratch.
    max_rows = 0
    for qlen in qlens:
        for reqs in range(1, max_num_reqs + 1):
            splits = _triton_verify_num_kv_splits(
                max_seq_len, reqs, qlen, num_heads, sm_count
            )
            if splits > 1:
                max_rows = max(max_rows, reqs * splits * qlen)
    if max_rows:
        current_workspace_manager().get_simultaneous(
            ((max_rows, num_heads, kv_lora_rank + 1), torch.float32),
        )

_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_fits_lds(num_heads, block_size, fp8_kv)

Whether AITER segmented MLA compiles for this DCP verify shape.

On gfx942 it needs 128 KiB of LDS, over the 64 KiB available, for bf16 KV at TILE-128 and for fp8 KV past 64 heads, where its query block widens to 128 heads. num_heads is the DCP-gathered count.

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

def _segmented_mla_fits_lds(num_heads: int, block_size: int, fp8_kv: bool) -> bool:
    """Whether AITER segmented MLA compiles for this DCP verify shape.

    On gfx942 it needs 128 KiB of LDS, over the 64 KiB available, for bf16 KV
    at TILE-128 and for fp8 KV past 64 heads, where its query block widens to
    128 heads. ``num_heads`` is the DCP-gathered count.
    """
    from vllm.platforms.rocm import on_gfx942

    if not on_gfx942():
        return True
    if fp8_kv:
        return num_heads <= 64
    return _segmented_mla_page_size(block_size) < 128

_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, supports_triton, 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,
    supports_triton: 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_triton:
        return _DCPDecodeRoute.TRITON
    if supports_segmented:
        return _DCPDecodeRoute.SEGMENTED
    raise RuntimeError(
        "ROCM_AITER_MLA DCP multi-token verify requires either segmented "
        "MLA, the round-robin asm decode, or the gfx942 Triton verify path."
    )

_supports_fused_dcp_prefill_config(kv_cache_dtype, query_dtype, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim)

Whether the configuration supports fused DCP context expansion.

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

def _supports_fused_dcp_prefill_config(
    kv_cache_dtype: str | None,
    query_dtype: torch.dtype,
    kv_lora_rank: int,
    qk_nope_head_dim: int,
    qk_rope_head_dim: int,
    v_head_dim: int,
) -> bool:
    """Whether the configuration supports fused DCP context expansion."""
    return (
        kv_cache_dtype in ("fp8", "fp8_e4m3")
        and query_dtype == torch.bfloat16
        and kv_lora_rank == 512
        and qk_nope_head_dim == 128
        and qk_rope_head_dim == 64
        and v_head_dim == 128
        and is_aiter_found_and_supported()
    )

_triton_dcp_verify_supported(dcp_world_size, cp_interleave)

Whether gfx942 can fall back to generic split-KV causal target verify.

The batch shape is causal qlen>1 over a DCP shard. Eagle MTP and DSpark target verify are that same shape, so this is not limited to one speculative method. gfx942 has no cprr kernel, so this serves the shapes segmented MLA cannot (see _segmented_mla_fits_lds). qlen==1 and non-causal stay on the plain ASM decode.

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

def _triton_dcp_verify_supported(dcp_world_size: int, cp_interleave: int) -> bool:
    """Whether gfx942 can fall back to generic split-KV causal target verify.

    The batch shape is causal qlen>1 over a DCP shard. Eagle MTP and DSpark
    target verify are that same shape, so this is not limited to one
    speculative method. gfx942 has no cprr kernel, so this serves the shapes
    segmented MLA cannot (see _segmented_mla_fits_lds). qlen==1 and non-causal
    stay on the plain ASM decode.
    """
    from vllm.platforms.rocm import on_gfx942

    return on_gfx942() and dcp_world_size > 1 and cp_interleave == 1

_triton_verify_num_kv_splits(max_seq_len, num_reqs, qlen, num_heads, sm_count)

KV split cap for the folded verify kernel at this launch shape.

Depends only on values that are static for a captured graph; the kernel lengthens the splits of shorter requests at runtime.

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

def _triton_verify_num_kv_splits(
    max_seq_len: int, num_reqs: int, qlen: int, num_heads: int, sm_count: int
) -> int:
    """KV split cap for the folded verify kernel at this launch shape.

    Depends only on values that are static for a captured graph; the kernel
    lengthens the splits of shorter requests at runtime.
    """
    if envs.VLLM_BATCH_INVARIANT:
        return 1
    # Must agree with the launch config decode_attention_fwd_verify selects.
    cfg = select_verify_config(num_reqs, qlen, num_heads)
    return verify_num_kv_splits(
        max_seq_len,
        num_reqs,
        qlen,
        num_heads,
        sm_count,
        block_h=cfg["BLOCK_H"],
        occupancy=cfg.get("occupancy", 1.0),
        min_work_per_split=cfg["MIN_SPLIT_LEN"],
    )