diff --git a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py index d67df8ee2b..d5c37724aa 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py +++ b/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) diff --git a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py index 06684302cf..a86c7e617b 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py +++ b/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) diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index 112cdff324..efffcad927 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/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 diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 4cfb3b3884..9bea4f8b3a 100644 --- a/vllm/_aiter_ops.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index b850deb070..c9ba2b0a56 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/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, )