@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 35dc273072

[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>

ppalanga committed Sep 15, 2026 at 10:18 UTC 35dc273072ff133743bc5fcfa306bf030fa8730b
5 files changed +390 -5
tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
+116 -1
@@ -30,7 +30,10 @@ def reset_aiter_mla_support_cache(monkeypatch: pytest.MonkeyPatch) -> None:
30
31
32 def _install_fake_aiter_modules(
33 - monkeypatch: pytest.MonkeyPatch, *, supports_fp8: bool
33 + monkeypatch: pytest.MonkeyPatch,
34 + *,
35 + supports_fp8: bool,
36 + supports_causal: bool = False,
37 ) -> None:
38 aiter_mod: Any = types.ModuleType("aiter")
39 mla_mod: Any = types.ModuleType("aiter.mla")
@@ -65,6 +68,14 @@ def _install_fake_aiter_modules(
68
69 mla_decode_fwd = mla_decode_fwd_without_fp8
70
71 + if supports_causal:
72 + inner = mla_decode_fwd
73 +
74 + def mla_decode_fwd_with_causal(*args, causal=True, **kwargs):
75 + return inner(*args, **kwargs)
76 +
77 + mla_decode_fwd = mla_decode_fwd_with_causal
78 +
79 mla_mod.mla_decode_fwd = mla_decode_fwd
80 aiter_mod.mla = mla_mod
81
@@ -135,3 +146,107 @@ def test_aiter_mla_fp8_support_result_is_cached(monkeypatch):
146 assert _check_aiter_mla_fp8_support() is True
147 assert _check_aiter_mla_fp8_support() is True
148 assert signature_mock.call_count == 1
149 +
150 +
151 +def test_missing_causal_arg_fails_closed(monkeypatch):
152 + """Old aiter wheels have no causal=; do not report that as a valid draft path."""
153 + from vllm._aiter_ops import rocm_aiter_ops
154 + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
155 +
156 + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=False)
157 + with pytest.raises(RuntimeError, match="causal-only"):
158 + rocm_aiter_ops.mla_decode_supports_non_causal()
159 + with pytest.raises(RuntimeError, match="causal-only"):
160 + AiterMLABackend.supports_non_causal()
161 +
162 +
163 +@pytest.mark.parametrize("has_kernels", [False, True])
164 +def test_supports_non_causal_requires_arch_kernels(monkeypatch, has_kernels):
165 + """The Python argument alone does not guarantee an arch has the kernel."""
166 + from vllm._aiter_ops import rocm_aiter_ops
167 + from vllm.v1.attention.backends.mla import rocm_aiter_mla
168 + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
169 +
170 + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True)
171 + monkeypatch.setattr(
172 + rocm_aiter_mla,
173 + "_aiter_mla_non_causal_asm_kernels",
174 + lambda: has_kernels,
175 + )
176 + assert rocm_aiter_ops.mla_decode_supports_non_causal() is True
177 + assert AiterMLABackend.supports_non_causal() is has_kernels
178 +
179 +
180 +@pytest.mark.parametrize("causal", [True, False])
181 +def test_decode_fwd_forwards_causal_to_aiter(monkeypatch, causal):
182 + """The decode op must pass the block mask through; omitting it is causal-only."""
183 + seen: list = []
184 +
185 + def mla_decode_fwd(*args, causal=None, **kwargs):
186 + seen.append(causal)
187 +
188 + mla_mod: Any = types.ModuleType("aiter.mla")
189 + mla_mod.mla_decode_fwd = mla_decode_fwd
190 + aiter_mod: Any = types.ModuleType("aiter")
191 + aiter_mod.mla = mla_mod
192 + monkeypatch.setitem(sys.modules, "aiter", aiter_mod)
193 + monkeypatch.setitem(sys.modules, "aiter.mla", mla_mod)
194 +
195 + import torch
196 +
197 + import vllm._aiter_ops as aiter_ops
198 +
199 + aiter_ops._rocm_aiter_mla_decode_fwd_impl(
200 + torch.zeros(2, 1, 8),
201 + torch.zeros(2, 1, 1, 8),
202 + torch.zeros(2, 1, 8),
203 + torch.zeros(3, dtype=torch.int32),
204 + 1,
205 + causal=causal,
206 + )
207 + assert seen == [causal]
208 +
209 +
210 +def _aiter_mla_validate(
211 + monkeypatch,
212 + *,
213 + dtype,
214 + kv_cache_dtype,
215 + use_non_causal,
216 +):
217 + from vllm.platforms.interface import DeviceCapability
218 + from vllm.v1.attention.backend import AttentionType
219 + from vllm.v1.attention.backends.mla import rocm_aiter_mla
220 + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
221 +
222 + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True)
223 + monkeypatch.setattr(
224 + rocm_aiter_mla, "_aiter_mla_non_causal_asm_kernels", lambda: True
225 + )
226 + return AiterMLABackend.validate_configuration(
227 + head_size=576,
228 + dtype=dtype,
229 + kv_cache_dtype=kv_cache_dtype,
230 + block_size=1,
231 + use_mla=True,
232 + has_sink=False,
233 + use_sparse=False,
234 + use_mm_prefix=False,
235 + use_per_head_quant_scales=False,
236 + device_capability=DeviceCapability(9, 5),
237 + attn_type=AttentionType.DECODER,
238 + use_non_causal=use_non_causal,
239 + )
240 +
241 +
242 +def test_non_causal_fp16_auto_cache_is_rejected(monkeypatch):
243 + """Pinned AITER aborts fp16 Q; auto-select must fall through to Triton."""
244 + import torch
245 +
246 + reasons = _aiter_mla_validate(
247 + monkeypatch,
248 + dtype=torch.float16,
249 + kv_cache_dtype="auto",
250 + use_non_causal=True,
251 + )
252 + assert any("non-causal fp16" in reason for reason in reasons)
tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py
+5
@@ -176,3 +176,8 @@ def test_pad_unpad_round_trip_preserves_head_order(num_heads):
176
177 assert unpadded.shape == q.shape
178 torch.testing.assert_close(unpadded, q)
179 +
180 +
181 +def test_a_non_causal_block_never_routes_to_gluon(gluon_available):
182 + """A small-head block that normally uses Gluon must use ASM when non-causal."""
183 + assert not AiterMLAHelper.use_gluon_verify(12, 8, "auto", causal=False)
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+98
@@ -130,6 +130,7 @@ def _builder(
130 _mla_q_dtype=torch.bfloat16,
131 _mla_kv_dtype=torch.bfloat16,
132 decode_attn_out_dtype=torch.bfloat16,
133 + _decode_causal=True,
134 )
135 # Bound to the stub rather than faked: the verify flatten's per-row view is
136 # part of what _build_decode is being tested for.
@@ -155,6 +156,25 @@ def test_backend_declares_uniform_batch_support():
156 )
157
158
159 +@pytest.mark.parametrize("causal", [True, False])
160 +def test_build_hands_common_causal_flag_to_aiter_builder(monkeypatch, causal):
161 + seen = []
162 +
163 + def fake_common_build(self, *args):
164 + seen.append(self._decode_causal)
165 + return SimpleNamespace(decode=None, prefill=None)
166 +
167 + monkeypatch.setattr(
168 + rocm_aiter_mla.MLACommonMetadataBuilder, "build", fake_common_build
169 + )
170 + builder = object.__new__(AiterMLAMetadataBuilder)
171 + builder._fp8_prefill_enabled = False
172 +
173 + builder.build(0, SimpleNamespace(causal=causal))
174 +
175 + assert seen == [causal]
176 +
177 +
178 def test_dcp_verify_row_view_is_causal_per_row():
179 """Row t must cover this rank's share of positions [0, seq_len - qlen + t].
180
@@ -949,3 +969,81 @@ def test_persistent_metadata_gate_without_gluon_build(
969
970 assert metadata.has_persistent_metadata is expect_persistent
971 assert get_mla_metadata_v1.called is expect_persistent
972 +
973 +
974 +def _build_non_causal(monkeypatch, *, num_heads, kv_cache_dtype, qlen, mtp_qlen):
975 + """Drive _build_decode with causal=False and hand back the aiter mock."""
976 + get_mla_metadata_v1 = mock.MagicMock()
977 + monkeypatch.setitem(
978 + sys.modules,
979 + "aiter",
980 + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
981 + )
982 + monkeypatch.setattr(
983 + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
984 + )
985 + builder = _builder(
986 + mtp_decode_qlen=mtp_qlen,
987 + kv_cache_dtype=kv_cache_dtype,
988 + num_heads=num_heads,
989 + )
990 + builder._decode_causal = False
991 + metadata = AiterMLAMetadataBuilder._build_decode(
992 + builder,
993 + block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
994 + seq_lens_device=torch.tensor([7, 5], dtype=torch.int32),
995 + max_seq_len=7,
996 + query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
997 + query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
998 + num_decode_tokens=2 * qlen,
999 + dcp_tot_seq_lens_device=None,
1000 + )
1001 + return metadata, get_mla_metadata_v1
1002 +
1003 +
1004 +def test_non_causal_build_hands_the_mask_to_aiter(monkeypatch):
1005 + """The schedule carries the mask, so is_causal must follow the block."""
1006 + metadata, get_mla_metadata_v1 = _build_non_causal(
1007 + monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=8, mtp_qlen=8
1008 + )
1009 + assert metadata.has_persistent_metadata
1010 + # 6th positional arg of get_mla_metadata_v1 is is_causal.
1011 + assert get_mla_metadata_v1.call_args.args[5] is False
1012 +
1013 +
1014 +@pytest.mark.parametrize("num_heads", [12, 16, 48])
1015 +def test_a_two_token_fp8_block_is_refused_without_a_fold(monkeypatch, num_heads):
1016 + """16-head (and padded-to-16) qlen-2 fp8 has no non-causal kernel.
1017 +
1018 + 48/80/112 (and 33 padded to 48) fold to H16 while keeping Q2, so they
1019 + must be refused the same way as native 16-head qlen 2.
1020 + """
1021 + with pytest.raises(ValueError, match="2-token"):
1022 + _build_non_causal(
1023 + monkeypatch,
1024 + num_heads=num_heads,
1025 + kv_cache_dtype="fp8",
1026 + qlen=2,
1027 + mtp_qlen=8,
1028 + )
1029 +
1030 +
1031 +def test_a_two_token_fp8_block_is_allowed_when_folded(monkeypatch):
1032 + """32 heads at qlen 2 folds onto the 16-head / 4-token kernel."""
1033 + metadata, get_mla_metadata_v1 = _build_non_causal(
1034 + monkeypatch,
1035 + num_heads=32,
1036 + kv_cache_dtype="fp8",
1037 + qlen=2,
1038 + mtp_qlen=8,
1039 + )
1040 + assert metadata.has_persistent_metadata
1041 + assert get_mla_metadata_v1.called
1042 +
1043 +
1044 +def test_a_two_token_bf16_block_is_allowed(monkeypatch):
1045 + """The bf16 fold has no such hole, so the guard must not fire there."""
1046 + metadata, _ = _build_non_causal(
1047 + monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=2, mtp_qlen=8
1048 + )
1049 + assert metadata.has_persistent_metadata
vllm/_aiter_ops.py
+26 -1
@@ -583,10 +583,11 @@ def _rocm_aiter_mla_decode_fwd_impl(
583 reduce_indptr: torch.Tensor | None = None,
584 reduce_final_map: torch.Tensor | None = None,
585 reduce_partial_map: torch.Tensor | None = None,
586 + causal: bool = True,
587 ) -> None:
588 from aiter.mla import mla_decode_fwd
589
589 - kwargs: dict[str, float | torch.Tensor | None] = {
590 + kwargs: dict[str, float | torch.Tensor | None | bool] = {
591 "sm_scale": sm_scale,
592 "logit_cap": logit_cap,
593 }
@@ -596,6 +597,8 @@ def _rocm_aiter_mla_decode_fwd_impl(
597 kwargs["q_scale"] = q_scale
598 kwargs["kv_scale"] = kv_scale
599
600 + kwargs["causal"] = causal
601 +
602 if work_meta_data is not None:
603 assert work_indptr is not None, (
604 "work_indptr must be provided with work_meta_data"
@@ -2036,6 +2039,26 @@ class rocm_aiter_ops:
2039 def is_mla_enabled(cls) -> bool:
2040 return cls._AITER_ENABLED and cls._MLA_ENABLED
2041
2042 + @classmethod
2043 + @if_aiter_supported
2044 + def mla_decode_supports_non_causal(cls) -> bool:
2045 + """Whether installed aiter.mla.mla_decode_fwd accepts `causal`.
2046 +
2047 + Added in aiter v0.1.20. A missing argument is a hard error rather
2048 + than a boolean the caller can use to pick another MLA backend.
2049 + """
2050 + import inspect
2051 +
2052 + from aiter.mla import mla_decode_fwd
2053 +
2054 + if "causal" not in inspect.signature(mla_decode_fwd).parameters:
2055 + raise RuntimeError(
2056 + "ROCM_AITER_MLA requires aiter.mla.mla_decode_fwd(..., causal=). "
2057 + "The installed aiter is causal-only; upgrade aiter rather than "
2058 + "falling back to another MLA backend."
2059 + )
2060 + return True
2061 +
2062 @classmethod
2063 @if_aiter_supported
2064 def is_mha_enabled(cls) -> bool:
@@ -2858,6 +2881,7 @@ class rocm_aiter_ops:
2881 reduce_indptr: torch.Tensor | None = None,
2882 reduce_final_map: torch.Tensor | None = None,
2883 reduce_partial_map: torch.Tensor | None = None,
2884 + causal: bool = True,
2885 ):
2886 torch.ops.vllm.rocm_aiter_mla_decode_fwd(
2887 q,
@@ -2878,6 +2902,7 @@ class rocm_aiter_ops:
2902 reduce_indptr=reduce_indptr,
2903 reduce_final_map=reduce_final_map,
2904 reduce_partial_map=reduce_partial_map,
2905 + causal=causal,
2906 )
2907
2908 @staticmethod
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+145 -3
@@ -4,7 +4,7 @@
4 import functools
5 from dataclasses import dataclass
6 from pathlib import Path
7 -from typing import ClassVar, Final
7 +from typing import TYPE_CHECKING, ClassVar, Final
8
9 import torch
10
@@ -38,6 +38,9 @@ from vllm.v1.attention.ops.rocm_aiter_mla_merge import (
38 )
39 from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
40
41 +if TYPE_CHECKING:
42 + from vllm.platforms.interface import DeviceCapability
43 +
44 logger = init_logger(__name__)
45
46
@@ -144,6 +147,20 @@ def _aiter_mla_native_h24_supported() -> bool:
147 )
148
149
150 +def _aiter_mla_non_causal_asm_kernels() -> bool:
151 + """Whether this arch ships non-causal MLA decode ASM kernels.
152 +
153 + The Python `causal=` probe cannot see the per-arch manifest. gfx950 has
154 + the kernels; gfx942 and everything else must fall through to another
155 + backend.
156 + """
157 + try:
158 + from vllm.platforms.rocm import on_gfx950
159 + except Exception: # noqa: BLE001
160 + return False
161 + return bool(on_gfx950())
162 +
163 +
164 @functools.lru_cache(maxsize=1)
165 def _gluon_mla_decode_supported() -> bool:
166 """The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build.
@@ -252,6 +269,72 @@ class AiterMLABackend(MLACommonBackend):
269 def get_name() -> str:
270 return "ROCM_AITER_MLA"
271
272 + @classmethod
273 + def supports_non_causal(cls) -> bool:
274 + # causal= on mla_decode_fwd is necessary but not sufficient: pinned
275 + # AITER v0.1.21.post2 still exports that argument on gfx942, whose ASM
276 + # manifest has no non-causal decode kernels. Selecting this backend
277 + # there reaches `cannot get heuristic kernel!` instead of TRITON_MLA.
278 + return bool(
279 + rocm_aiter_ops.mla_decode_supports_non_causal()
280 + and _aiter_mla_non_causal_asm_kernels()
281 + )
282 +
283 + @classmethod
284 + def validate_configuration(
285 + cls,
286 + head_size: int,
287 + dtype: torch.dtype,
288 + kv_cache_dtype: CacheDType | None,
289 + block_size: int | None,
290 + use_mla: bool,
291 + has_sink: bool,
292 + use_sparse: bool,
293 + use_mm_prefix: bool,
294 + use_per_head_quant_scales: bool,
295 + device_capability: "DeviceCapability",
296 + attn_type: str,
297 + has_sliding_window: bool = False,
298 + use_non_causal: bool = False,
299 + use_batch_invariant: bool = False,
300 + use_kv_connector: bool = False,
301 + use_pcp: bool = False,
302 + use_adaptive_verification: bool = False,
303 + use_dcp: bool = False,
304 + use_rswa: bool = False,
305 + ) -> list[str]:
306 + invalid_reasons = super().validate_configuration(
307 + head_size=head_size,
308 + dtype=dtype,
309 + kv_cache_dtype=kv_cache_dtype,
310 + block_size=block_size,
311 + use_mla=use_mla,
312 + has_sink=has_sink,
313 + use_sparse=use_sparse,
314 + use_mm_prefix=use_mm_prefix,
315 + use_per_head_quant_scales=use_per_head_quant_scales,
316 + device_capability=device_capability,
317 + attn_type=attn_type,
318 + has_sliding_window=has_sliding_window,
319 + use_non_causal=use_non_causal,
320 + use_batch_invariant=use_batch_invariant,
321 + use_kv_connector=use_kv_connector,
322 + use_pcp=use_pcp,
323 + use_adaptive_verification=use_adaptive_verification,
324 + use_dcp=use_dcp,
325 + use_rswa=use_rswa,
326 + )
327 + # Arch/signature cannot express this: the backend still advertises
328 + # fp16, so a DSpark draft with kv_cache_dtype="auto" would select
329 + # AITER and abort on fp16 query. Causal fp16 is fine.
330 + if (
331 + use_non_causal
332 + and dtype == torch.float16
333 + and kv_cache_dtype in (None, "auto", "float16")
334 + ):
335 + invalid_reasons.append("non-causal fp16 MLA decode not supported")
336 + return invalid_reasons
337 +
338 @staticmethod
339 def get_impl_cls() -> type["AiterMLAImpl"]:
340 return AiterMLAImpl
@@ -347,6 +430,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
430 # https://github.com/vllm-project/vllm/issues/22945
431 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
432 query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
433 + # Served by passing the mask to the kernel; _build_decode turns away the
434 + # shapes AITER has no non-causal kernel for.
435 + supports_non_causal_multi_token_decode: ClassVar[bool] = True
436 + # Set from the common metadata every build; a batch is causal unless the
437 + # drafter says otherwise.
438 + _decode_causal: bool = True
439
440 @staticmethod
441 def _uniform_padded_mtp_qo_len(
@@ -919,6 +1008,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1008 num_decode_tokens: int,
1009 dcp_tot_seq_lens_device: torch.Tensor | None,
1010 ) -> AiterMLADecodeMetadata:
1011 + causal = self._decode_causal
1012 device = self.device
1013 num_reqs = seq_lens_device.size(0)
1014 qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
@@ -970,6 +1060,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1060 int(max_qo_len),
1061 self._kv_cache_dtype_str,
1062 self.dcp_world_size,
1063 + causal,
1064 )
1065 use_segmented_dcp_verify = (
1066 self._supports_segmented_dcp_verify and max_qo_len > 1
@@ -1047,14 +1138,36 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1138 # A padded rank has no bf16 persistent kernel past qlen 4 where the
1139 # gfx950 fold is absent; the non-persistent entry covers it. fp8
1140 # keeps the schedule -- its fold rejects non-persistent outright.
1141 + # A non-causal block keeps it too: what the fold drops past qlen 4
1142 + # is the block's causal staircase, which a non-causal block does
1143 + # not have, and the schedule is the only thing carrying its mask.
1144 and (
1051 - self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
1145 + not causal
1146 + or self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
1147 or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN
1148 or is_quantized_kv_cache(self._kv_cache_dtype_str)
1149 )
1150 and max_qo_len >= 1
1151 and max_qo_len <= self._mtp_decode_qlen
1152 )
1153 + if (
1154 + not causal
1155 + and max_qo_len == 2
1156 + and is_quantized_kv_cache(self._kv_cache_dtype_str)
1157 + and not AiterMLAHelper.has_fp8_non_causal_qlen2_kernel(
1158 + self._decode_num_heads
1159 + )
1160 + ):
1161 + # AITER's fp8 dispatch folds (gqa 16, qlen 3 or 4) onto the
1162 + # qseqlen-4 kernel but never lists 2. Only 32/64/96/128 heads
1163 + # at qlen 2 fold onto that 16-head / 4-token non-causal kernel;
1164 + # 16-head qlen 2, and padded 48/80/112, keep Q2 and abort the
1165 + # process rather than raising. The bf16 fold has no such hole.
1166 + raise ValueError(
1167 + "AITER has no non-causal fp8 MLA kernel for this 2-token "
1168 + "query block. Pin the draft to TRITON_MLA for this "
1169 + "speculative config."
1170 + )
1171 if use_persistent_metadata:
1172 from aiter import get_mla_metadata_v1
1173
@@ -1067,7 +1180,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1180 paged_kv_last_page_len,
1181 self._num_attention_heads,
1182 1,
1070 - True,
1183 + causal,
1184 self._mla_work_meta_data,
1185 self._mla_work_info_set,
1186 self._mla_work_indptr,
@@ -1137,6 +1250,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1250 common_attn_metadata: CommonAttentionMetadata,
1251 fast_build: bool = False,
1252 ) -> AiterMLAMetadata:
1253 + # The common MLA builder owns the causal flag but its _build_decode
1254 + # contract is shared with CUDA backends. Keep this ROCm-specific state
1255 + # on the per-ubatch Aiter builder instead of changing that shared API;
1256 + # super().build invokes _build_decode synchronously.
1257 + self._decode_causal = common_attn_metadata.causal
1258 attn_metadata = super().build(
1259 common_prefix_len, common_attn_metadata, fast_build
1260 )
@@ -1226,6 +1344,10 @@ class AiterMLAHelper:
1344 # fold that reaches a persistent one is gfx950-only.
1345 _ASM_PADDED_MAX_PS_QLEN: Final = 4
1346 _AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = ()
1347 + # Pinned AITER v0.1.21.post2 folds these fp8 qlen-2 counts onto the
1348 + # non-causal 16-head / 4-token kernel. Other multiples of 16 (48/80/112)
1349 + # fold to H16 while keeping Q2, which has no non-causal fp8 entry.
1350 + _AITER_FP8_NON_CAUSAL_QLEN2_HEADS: ClassVar[tuple[int, ...]] = (32, 64, 96, 128)
1351
1352 @staticmethod
1353 def qo_indptr_for_uniform_qlen(
@@ -1263,6 +1385,18 @@ class AiterMLAHelper:
1385 )
1386 )
1387
1388 + @staticmethod
1389 + def has_fp8_non_causal_qlen2_kernel(num_heads: int) -> bool:
1390 + """Whether fp8 (num_heads, qlen=2) folds onto a non-causal ASM kernel.
1391 +
1392 + Pinned AITER v0.1.21.post2 folds 32/64/96/128 heads at qlen 2 onto the
1393 + non-causal 16-head / 4-token kernel. A 16-head (or padded-to-16) qlen-2
1394 + block has no matching entry, and padded 48/80/112 keep Q2 after the
1395 + H16 fold.
1396 + """
1397 + kernel_heads = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
1398 + return kernel_heads in AiterMLAHelper._AITER_FP8_NON_CAUSAL_QLEN2_HEADS
1399 +
1400 @staticmethod
1401 def get_actual_mla_num_heads(num_heads: int) -> int:
1402 if num_heads == 24 and _aiter_mla_native_h24_supported():
@@ -1364,6 +1498,7 @@ class AiterMLAHelper:
1498 max_qo_len: int,
1499 kv_cache_dtype: str,
1500 dcp_world_size: int = 1,
1501 + causal: bool = True,
1502 ) -> bool:
1503 """Whether a small-head multi-token verify uses native Gluon MTP.
1504
@@ -1376,7 +1511,13 @@ class AiterMLAHelper:
1511
1512 DCP verify is excluded: its per-row causal windows are served by the
1513 segmented path, which Gluon's MTP entry cannot express.
1514 +
1515 + Gluon masks the block causally with no way to turn it off, so a
1516 + non-causal block takes the padded asm decode whatever the head count
1517 + -- padding is what gives it a kernel there.
1518 """
1519 + if not causal:
1520 + return False
1521 if max_qo_len <= 1 or dcp_world_size > 1:
1522 return False
1523 if is_quantized_kv_cache(kv_cache_dtype):
@@ -1974,6 +2115,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
2115 decode.paged_kv_indptr,
2116 decode.paged_kv_indices,
2117 decode.paged_kv_last_page_len,
2118 + causal=attn_metadata.causal,
2119 **mla_kwargs,
2120 )
2121