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

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

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

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

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

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

    @staticmethod
    def 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 get_actual_mla_num_heads(num_heads: int) -> int:
        if num_heads == 24 and _aiter_mla_native_h24_supported():
            return num_heads
        m = AiterMLAHelper._AITER_MIN_MLA_HEADS
        return -(-num_heads // m) * m

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

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

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

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

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

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

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

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

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

    @staticmethod
    def use_gluon_verify(
        num_heads: int,
        max_qo_len: int,
        kv_cache_dtype: str,
        dcp_world_size: int = 1,
    ) -> 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.
        """
        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 and mode gating as use_gluon_decode.
        return _aiter_mla_small_head_mode() != "asm"

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

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

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

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

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

        from aiter import flash_attn_varlen_func

        self.flash_attn_varlen_func = flash_attn_varlen_func

        # FP8 MLA prefill kernel imports (lazy, only when enabled).
        # Auto-enabled on gfx950 when AITER ships the kernels. Only runs when the
        # KV cache is FP8. Head counts that are not a multiple of 16 are
        # replicate-padded up to one (see _mla_fp8_prefill_attn).
        from vllm.utils.torch_utils import is_quantized_kv_cache

        self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and (
            is_quantized_kv_cache(kv_cache_dtype)
            and AiterMLAHelper.is_valid_num_heads(self.num_heads)
        )
        if self._fp8_prefill_enabled:
            from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1

            self._mla_prefill_ps_asm_fwd = mla_prefill_ps_asm_fwd
            self._mla_reduce_v1 = mla_reduce_v1

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

        return output

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        # 12-head (<16) multi-token verify (DSpark): the asm path has no
        # gqa<16, qseqlen>1 kernel, so the block goes to the gluon kernel's 4-D
        # MTP entry, which serves a whole (1 + num_spec) block in one launch.
        # The block is causal -- the target is checking draft tokens, so
        # position t must not see t+1 -- and the kernel bounds each query
        # position's scores itself.
        # Arch, mode and dtype gating all live in use_gluon_verify, which the
        # builder evaluates and records on the metadata, so this branch acts on
        # the same answer the builder used to decide whether the asm decode
        # would run. DCP verify is routed to the segmented path below instead.
        if decode.use_gluon_verify:
            if type(q) is tuple:
                q_nope, q_pe = q
            else:
                q_nope, q_pe = torch.split(
                    q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
                )
            B, num_q_heads, _ = q_nope.shape
            o = torch.empty(
                B,
                num_q_heads,
                self.kv_lora_rank,
                dtype=decode.attn_out_dtype,
                device=q_nope.device,
            )
            kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])
            assert attn_metadata.causal, (
                "AITER MLA small-head verify MTP is causal-only"
            )
            # Hand mla_gluon its 4-D MTP entry instead of an expanded
            # per-verify-token paged-KV view. The flat query layout is already
            # row-major (r * qlen + t), so unflatten is a free view. mla_gluon
            # applies the per-position causal bound
            # score_end = min(split_kv_end, seq_len - qlen + q_pos + 1)
            # in-kernel, and seq_lens already spans the verify block, so that
            # bound is context_r + q_pos + 1 -- the same window the expanded
            # view supplied by truncating each row's page list.
            qlen = int(decode.max_qo_len)
            num_reqs = B // qlen
            if num_reqs * qlen != B:
                raise ValueError(
                    f"verify block {B} rows is not a multiple of qlen {qlen}"
                )
            assert decode.paged_kv_indptr is not None
            assert decode.paged_kv_indices is not None
            mla_gluon = _get_mla_gluon()
            mla_gluon(
                q_nope=q_nope.unflatten(0, (num_reqs, qlen)),
                q_pe=q_pe.unflatten(0, (num_reqs, qlen)),
                kv_c=kv_buffer,
                o=o.unflatten(0, (num_reqs, qlen)),
                page_table=decode.paged_kv_indices,
                seq_info=decode.paged_kv_indptr,
                sm_scale=self.scale,
                k_pe=None,
                kv_pe_offset=self.kv_lora_rank,
                use_2d_view=False,
                kv_scale=1.0,
                min_kv_seq_len=decode.min_kv_seq_len,
            )
            return o, None

        verify = decode.dcp_verify
        if 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 self.dcp_world_size > 1 and int(decode.max_qo_len) > 1:
            raise RuntimeError(
                "ROCM_AITER_MLA DCP multi-token verify requires segmented MLA."
            )

        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]}"
        )
        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:
            # 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,
                **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,
                **mla_kwargs,
            )

        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