[ROCm][Spec Decode] Add Aiter MLA decode support non-causal draft block (#55966)

Signed-off-by: Poovaiah Palangappa <poovaiah.palangappa@amd.com>
Co-authored-by: Poovaiah Palangappa <poovaiah.palangappa@amd.com>
Co-authored-by: John Qin <zhengsheng1997@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com>

tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py

@@ -30,7 +30,10 @@ def reset_aiter_mla_support_cache(monkeypatch: pytest.MonkeyPatch) -> None:
 
 
 def _install_fake_aiter_modules(
-    monkeypatch: pytest.MonkeyPatch, *, supports_fp8: bool
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    supports_fp8: bool,
+    supports_causal: bool = False,
 ) -> None:
     aiter_mod: Any = types.ModuleType("aiter")
     mla_mod: Any = types.ModuleType("aiter.mla")
@@ -65,6 +68,14 @@ def _install_fake_aiter_modules(
 
         mla_decode_fwd = mla_decode_fwd_without_fp8
 
+    if supports_causal:
+        inner = mla_decode_fwd
+
+        def mla_decode_fwd_with_causal(*args, causal=True, **kwargs):
+            return inner(*args, **kwargs)
+
+        mla_decode_fwd = mla_decode_fwd_with_causal
+
     mla_mod.mla_decode_fwd = mla_decode_fwd
     aiter_mod.mla = mla_mod
 
@@ -135,3 +146,107 @@ def test_aiter_mla_fp8_support_result_is_cached(monkeypatch):
         assert _check_aiter_mla_fp8_support() is True
         assert _check_aiter_mla_fp8_support() is True
         assert signature_mock.call_count == 1
+
+
+def test_missing_causal_arg_fails_closed(monkeypatch):
+    """Old aiter wheels have no causal=; do not report that as a valid draft path."""
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
+
+    _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=False)
+    with pytest.raises(RuntimeError, match="causal-only"):
+        rocm_aiter_ops.mla_decode_supports_non_causal()
+    with pytest.raises(RuntimeError, match="causal-only"):
+        AiterMLABackend.supports_non_causal()
+
+
+@pytest.mark.parametrize("has_kernels", [False, True])
+def test_supports_non_causal_requires_arch_kernels(monkeypatch, has_kernels):
+    """The Python argument alone does not guarantee an arch has the kernel."""
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.v1.attention.backends.mla import rocm_aiter_mla
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
+
+    _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True)
+    monkeypatch.setattr(
+        rocm_aiter_mla,
+        "_aiter_mla_non_causal_asm_kernels",
+        lambda: has_kernels,
+    )
+    assert rocm_aiter_ops.mla_decode_supports_non_causal() is True
+    assert AiterMLABackend.supports_non_causal() is has_kernels
+
+
+@pytest.mark.parametrize("causal", [True, False])
+def test_decode_fwd_forwards_causal_to_aiter(monkeypatch, causal):
+    """The decode op must pass the block mask through; omitting it is causal-only."""
+    seen: list = []
+
+    def mla_decode_fwd(*args, causal=None, **kwargs):
+        seen.append(causal)
+
+    mla_mod: Any = types.ModuleType("aiter.mla")
+    mla_mod.mla_decode_fwd = mla_decode_fwd
+    aiter_mod: Any = types.ModuleType("aiter")
+    aiter_mod.mla = mla_mod
+    monkeypatch.setitem(sys.modules, "aiter", aiter_mod)
+    monkeypatch.setitem(sys.modules, "aiter.mla", mla_mod)
+
+    import torch
+
+    import vllm._aiter_ops as aiter_ops
+
+    aiter_ops._rocm_aiter_mla_decode_fwd_impl(
+        torch.zeros(2, 1, 8),
+        torch.zeros(2, 1, 1, 8),
+        torch.zeros(2, 1, 8),
+        torch.zeros(3, dtype=torch.int32),
+        1,
+        causal=causal,
+    )
+    assert seen == [causal]
+
+
+def _aiter_mla_validate(
+    monkeypatch,
+    *,
+    dtype,
+    kv_cache_dtype,
+    use_non_causal,
+):
+    from vllm.platforms.interface import DeviceCapability
+    from vllm.v1.attention.backend import AttentionType
+    from vllm.v1.attention.backends.mla import rocm_aiter_mla
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
+
+    _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True)
+    monkeypatch.setattr(
+        rocm_aiter_mla, "_aiter_mla_non_causal_asm_kernels", lambda: True
+    )
+    return AiterMLABackend.validate_configuration(
+        head_size=576,
+        dtype=dtype,
+        kv_cache_dtype=kv_cache_dtype,
+        block_size=1,
+        use_mla=True,
+        has_sink=False,
+        use_sparse=False,
+        use_mm_prefix=False,
+        use_per_head_quant_scales=False,
+        device_capability=DeviceCapability(9, 5),
+        attn_type=AttentionType.DECODER,
+        use_non_causal=use_non_causal,
+    )
+
+
+def test_non_causal_fp16_auto_cache_is_rejected(monkeypatch):
+    """Pinned AITER aborts fp16 Q; auto-select must fall through to Triton."""
+    import torch
+
+    reasons = _aiter_mla_validate(
+        monkeypatch,
+        dtype=torch.float16,
+        kv_cache_dtype="auto",
+        use_non_causal=True,
+    )
+    assert any("non-causal fp16" in reason for reason in reasons)

tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py

@@ -176,3 +176,8 @@ def test_pad_unpad_round_trip_preserves_head_order(num_heads):
 
     assert unpadded.shape == q.shape
     torch.testing.assert_close(unpadded, q)
+
+
+def test_a_non_causal_block_never_routes_to_gluon(gluon_available):
+    """A small-head block that normally uses Gluon must use ASM when non-causal."""
+    assert not AiterMLAHelper.use_gluon_verify(12, 8, "auto", causal=False)

tests/v1/attention/test_rocm_aiter_mla_mtp_split.py

@@ -130,6 +130,7 @@ def _builder(
         _mla_q_dtype=torch.bfloat16,
         _mla_kv_dtype=torch.bfloat16,
         decode_attn_out_dtype=torch.bfloat16,
+        _decode_causal=True,
     )
     # Bound to the stub rather than faked: the verify flatten's per-row view is
     # part of what _build_decode is being tested for.
@@ -155,6 +156,25 @@ def test_backend_declares_uniform_batch_support():
     )
 
 
+@pytest.mark.parametrize("causal", [True, False])
+def test_build_hands_common_causal_flag_to_aiter_builder(monkeypatch, causal):
+    seen = []
+
+    def fake_common_build(self, *args):
+        seen.append(self._decode_causal)
+        return SimpleNamespace(decode=None, prefill=None)
+
+    monkeypatch.setattr(
+        rocm_aiter_mla.MLACommonMetadataBuilder, "build", fake_common_build
+    )
+    builder = object.__new__(AiterMLAMetadataBuilder)
+    builder._fp8_prefill_enabled = False
+
+    builder.build(0, SimpleNamespace(causal=causal))
+
+    assert seen == [causal]
+
+
 def test_dcp_verify_row_view_is_causal_per_row():
     """Row t must cover this rank's share of positions [0, seq_len - qlen + t].
 
@@ -949,3 +969,81 @@ def test_persistent_metadata_gate_without_gluon_build(
 
     assert metadata.has_persistent_metadata is expect_persistent
     assert get_mla_metadata_v1.called is expect_persistent
+
+
+def _build_non_causal(monkeypatch, *, num_heads, kv_cache_dtype, qlen, mtp_qlen):
+    """Drive _build_decode with causal=False and hand back the aiter mock."""
+    get_mla_metadata_v1 = mock.MagicMock()
+    monkeypatch.setitem(
+        sys.modules,
+        "aiter",
+        SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
+    )
+    monkeypatch.setattr(
+        rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
+    )
+    builder = _builder(
+        mtp_decode_qlen=mtp_qlen,
+        kv_cache_dtype=kv_cache_dtype,
+        num_heads=num_heads,
+    )
+    builder._decode_causal = False
+    metadata = AiterMLAMetadataBuilder._build_decode(
+        builder,
+        block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
+        seq_lens_device=torch.tensor([7, 5], dtype=torch.int32),
+        max_seq_len=7,
+        query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
+        query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
+        num_decode_tokens=2 * qlen,
+        dcp_tot_seq_lens_device=None,
+    )
+    return metadata, get_mla_metadata_v1
+
+
+def test_non_causal_build_hands_the_mask_to_aiter(monkeypatch):
+    """The schedule carries the mask, so is_causal must follow the block."""
+    metadata, get_mla_metadata_v1 = _build_non_causal(
+        monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=8, mtp_qlen=8
+    )
+    assert metadata.has_persistent_metadata
+    # 6th positional arg of get_mla_metadata_v1 is is_causal.
+    assert get_mla_metadata_v1.call_args.args[5] is False
+
+
+@pytest.mark.parametrize("num_heads", [12, 16, 48])
+def test_a_two_token_fp8_block_is_refused_without_a_fold(monkeypatch, num_heads):
+    """16-head (and padded-to-16) qlen-2 fp8 has no non-causal kernel.
+
+    48/80/112 (and 33 padded to 48) fold to H16 while keeping Q2, so they
+    must be refused the same way as native 16-head qlen 2.
+    """
+    with pytest.raises(ValueError, match="2-token"):
+        _build_non_causal(
+            monkeypatch,
+            num_heads=num_heads,
+            kv_cache_dtype="fp8",
+            qlen=2,
+            mtp_qlen=8,
+        )
+
+
+def test_a_two_token_fp8_block_is_allowed_when_folded(monkeypatch):
+    """32 heads at qlen 2 folds onto the 16-head / 4-token kernel."""
+    metadata, get_mla_metadata_v1 = _build_non_causal(
+        monkeypatch,
+        num_heads=32,
+        kv_cache_dtype="fp8",
+        qlen=2,
+        mtp_qlen=8,
+    )
+    assert metadata.has_persistent_metadata
+    assert get_mla_metadata_v1.called
+
+
+def test_a_two_token_bf16_block_is_allowed(monkeypatch):
+    """The bf16 fold has no such hole, so the guard must not fire there."""
+    metadata, _ = _build_non_causal(
+        monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=2, mtp_qlen=8
+    )
+    assert metadata.has_persistent_metadata

vllm/_aiter_ops.py

@@ -583,10 +583,11 @@ def _rocm_aiter_mla_decode_fwd_impl(
     reduce_indptr: torch.Tensor | None = None,
     reduce_final_map: torch.Tensor | None = None,
     reduce_partial_map: torch.Tensor | None = None,
+    causal: bool = True,
 ) -> None:
     from aiter.mla import mla_decode_fwd
 
-    kwargs: dict[str, float | torch.Tensor | None] = {
+    kwargs: dict[str, float | torch.Tensor | None | bool] = {
         "sm_scale": sm_scale,
         "logit_cap": logit_cap,
     }
@@ -596,6 +597,8 @@ def _rocm_aiter_mla_decode_fwd_impl(
         kwargs["q_scale"] = q_scale
         kwargs["kv_scale"] = kv_scale
 
+    kwargs["causal"] = causal
+
     if work_meta_data is not None:
         assert work_indptr is not None, (
             "work_indptr must be provided with work_meta_data"
@@ -2036,6 +2039,26 @@ class rocm_aiter_ops:
     def is_mla_enabled(cls) -> bool:
         return cls._AITER_ENABLED and cls._MLA_ENABLED
 
+    @classmethod
+    @if_aiter_supported
+    def mla_decode_supports_non_causal(cls) -> bool:
+        """Whether installed aiter.mla.mla_decode_fwd accepts `causal`.
+
+        Added in aiter v0.1.20. A missing argument is a hard error rather
+        than a boolean the caller can use to pick another MLA backend.
+        """
+        import inspect
+
+        from aiter.mla import mla_decode_fwd
+
+        if "causal" not in inspect.signature(mla_decode_fwd).parameters:
+            raise RuntimeError(
+                "ROCM_AITER_MLA requires aiter.mla.mla_decode_fwd(..., causal=). "
+                "The installed aiter is causal-only; upgrade aiter rather than "
+                "falling back to another MLA backend."
+            )
+        return True
+
     @classmethod
     @if_aiter_supported
     def is_mha_enabled(cls) -> bool:
@@ -2858,6 +2881,7 @@ class rocm_aiter_ops:
         reduce_indptr: torch.Tensor | None = None,
         reduce_final_map: torch.Tensor | None = None,
         reduce_partial_map: torch.Tensor | None = None,
+        causal: bool = True,
     ):
         torch.ops.vllm.rocm_aiter_mla_decode_fwd(
             q,
@@ -2878,6 +2902,7 @@ class rocm_aiter_ops:
             reduce_indptr=reduce_indptr,
             reduce_final_map=reduce_final_map,
             reduce_partial_map=reduce_partial_map,
+            causal=causal,
         )
 
     @staticmethod

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

@@ -4,7 +4,7 @@
 import functools
 from dataclasses import dataclass
 from pathlib import Path
-from typing import ClassVar, Final
+from typing import TYPE_CHECKING, ClassVar, Final
 
 import torch
 
@@ -38,6 +38,9 @@ from vllm.v1.attention.ops.rocm_aiter_mla_merge import (
 )
 from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
 
+if TYPE_CHECKING:
+    from vllm.platforms.interface import DeviceCapability
+
 logger = init_logger(__name__)
 
 
@@ -144,6 +147,20 @@ def _aiter_mla_native_h24_supported() -> bool:
     )
 
 
+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())
+
+
 @functools.lru_cache(maxsize=1)
 def _gluon_mla_decode_supported() -> bool:
     """The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build.
@@ -252,6 +269,72 @@ class AiterMLABackend(MLACommonBackend):
     def get_name() -> str:
         return "ROCM_AITER_MLA"
 
+    @classmethod
+    def supports_non_causal(cls) -> bool:
+        # causal= on mla_decode_fwd is necessary but not sufficient: pinned
+        # AITER v0.1.21.post2 still exports that argument on gfx942, whose ASM
+        # manifest has no non-causal decode kernels. Selecting this backend
+        # there reaches `cannot get heuristic kernel!` instead of TRITON_MLA.
+        return bool(
+            rocm_aiter_ops.mla_decode_supports_non_causal()
+            and _aiter_mla_non_causal_asm_kernels()
+        )
+
+    @classmethod
+    def validate_configuration(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        use_mm_prefix: bool,
+        use_per_head_quant_scales: bool,
+        device_capability: "DeviceCapability",
+        attn_type: str,
+        has_sliding_window: bool = False,
+        use_non_causal: bool = False,
+        use_batch_invariant: bool = False,
+        use_kv_connector: bool = False,
+        use_pcp: bool = False,
+        use_adaptive_verification: bool = False,
+        use_dcp: bool = False,
+        use_rswa: bool = False,
+    ) -> list[str]:
+        invalid_reasons = super().validate_configuration(
+            head_size=head_size,
+            dtype=dtype,
+            kv_cache_dtype=kv_cache_dtype,
+            block_size=block_size,
+            use_mla=use_mla,
+            has_sink=has_sink,
+            use_sparse=use_sparse,
+            use_mm_prefix=use_mm_prefix,
+            use_per_head_quant_scales=use_per_head_quant_scales,
+            device_capability=device_capability,
+            attn_type=attn_type,
+            has_sliding_window=has_sliding_window,
+            use_non_causal=use_non_causal,
+            use_batch_invariant=use_batch_invariant,
+            use_kv_connector=use_kv_connector,
+            use_pcp=use_pcp,
+            use_adaptive_verification=use_adaptive_verification,
+            use_dcp=use_dcp,
+            use_rswa=use_rswa,
+        )
+        # Arch/signature cannot express this: the backend still advertises
+        # fp16, so a DSpark draft with kv_cache_dtype="auto" would select
+        # AITER and abort on fp16 query. Causal fp16 is fine.
+        if (
+            use_non_causal
+            and dtype == torch.float16
+            and kv_cache_dtype in (None, "auto", "float16")
+        ):
+            invalid_reasons.append("non-causal fp16 MLA decode not supported")
+        return invalid_reasons
+
     @staticmethod
     def get_impl_cls() -> type["AiterMLAImpl"]:
         return AiterMLAImpl
@@ -347,6 +430,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
     #  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
 
     @staticmethod
     def _uniform_padded_mtp_qo_len(
@@ -919,6 +1008,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         num_decode_tokens: 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]
@@ -970,6 +1060,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
             int(max_qo_len),
             self._kv_cache_dtype_str,
             self.dcp_world_size,
+            causal,
         )
         use_segmented_dcp_verify = (
             self._supports_segmented_dcp_verify and max_qo_len > 1
@@ -1047,14 +1138,36 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
             # A padded rank has no bf16 persistent kernel past qlen 4 where the
             # gfx950 fold is absent; the non-persistent entry covers it. fp8
             # keeps the schedule -- its fold rejects non-persistent outright.
+            # A non-causal block keeps it too: what the fold drops past qlen 4
+            # is the block's causal staircase, which a non-causal block does
+            # not have, and the schedule is the only thing carrying its mask.
             and (
-                self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
+                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
 
@@ -1067,7 +1180,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
                 paged_kv_last_page_len,
                 self._num_attention_heads,
                 1,
-                True,
+                causal,
                 self._mla_work_meta_data,
                 self._mla_work_info_set,
                 self._mla_work_indptr,
@@ -1137,6 +1250,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         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
         )
@@ -1226,6 +1344,10 @@ class AiterMLAHelper:
     # 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(
@@ -1263,6 +1385,18 @@ class AiterMLAHelper:
             )
         )
 
+    @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():
@@ -1364,6 +1498,7 @@ class AiterMLAHelper:
         max_qo_len: int,
         kv_cache_dtype: str,
         dcp_world_size: int = 1,
+        causal: bool = True,
     ) -> bool:
         """Whether a small-head multi-token verify uses native Gluon MTP.
 
@@ -1376,7 +1511,13 @@ class AiterMLAHelper:
 
         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):
@@ -1974,6 +2115,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
                 decode.paged_kv_indptr,
                 decode.paged_kv_indices,
                 decode.paged_kv_last_page_len,
+                causal=attn_metadata.causal,
                 **mla_kwargs,
             )