[Test] Add ROCm AITER MLA op registration and env gating tests (#50930)
Signed-off-by: Aarushi Jain <aarushi@amd.com> Co-authored-by: Aarushi Jain <aarushi@amd.com>
Aarushi Jain committed
Aug 7, 2026 at 14:07 UTC
a0056e103e214d5b0b73ca9e24003e09f8cbec63
2 files changed
+83
.buildkite/test-amd.yaml
+4
@@ -1543,10 +1543,12 @@ steps:
1543
- vllm/_aiter_ops.py
1544
- tests/kernels/attention/test_rocm_aiter_mla_decode.py
1545
- tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
1546
+ - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
1547
commands:
1548
- rocm-smi
1549
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
1550
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py
1551
+ - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
1552
1553
- label: Kernels Attention Test %N # TBD
1554
timeout_in_minutes: 180
@@ -3441,10 +3443,12 @@ steps:
3443
- vllm/_aiter_ops.py
3444
- tests/kernels/attention/test_rocm_aiter_mla_decode.py
3445
- tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
3446
+ - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
3447
commands:
3448
- rocm-smi
3449
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
3450
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py
3451
+ - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
3452
3453
- label: Kernels Attention Test %N # TBD
3454
timeout_in_minutes: 180
tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
new
+79
@@ -0,0 +1,79 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ROCm custom op schema test for AITER MLA decode.
4
+
5
+A single ``opcheck`` call on ``torch.ops.vllm.rocm_aiter_mla_decode_fwd``
6
+verifies that the custom op is registered and that its schema and fake
7
+implementation are consistent with the real kernel: fake-tensor support for
8
+torch.compile tracing and the ``mutates_args=["o"]`` in-place output aliasing.
9
+"""
10
+
11
+import pytest
12
+import torch
13
+
14
+from tests.kernels.utils import opcheck
15
+from vllm.platforms import current_platform
16
+
17
+_SKIP_NON_MI3XX = True
18
+if current_platform.is_rocm():
19
+ from vllm.platforms.rocm import on_mi3xx
20
+
21
+ _SKIP_NON_MI3XX = not on_mi3xx()
22
+
23
+pytestmark = [
24
+ pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"),
25
+ pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"),
26
+]
27
+
28
+Q_HEAD_DIM = 576 # kv_lora_rank + qk_rope_head_dim
29
+V_HEAD_DIM = 512 # kv_lora_rank
30
+
31
+
32
+def _require_aiter():
33
+ from vllm._aiter_ops import is_aiter_found_and_supported
34
+
35
+ if not is_aiter_found_and_supported():
36
+ pytest.skip("aiter is required on supported ROCm hardware for this test")
37
+
38
+
39
+@torch.inference_mode()
40
+def test_mla_decode_fwd_op_schema() -> None:
41
+ """opcheck validates registration, schema, fake-tensor, and ``o`` aliasing.
42
+
43
+ A single opcheck call covers that the op is registered/callable, that its
44
+ fake implementation matches the real op (torch.compile tracing), and that
45
+ the ``mutates_args=["o"]`` in-place output aliasing is declared correctly.
46
+ """
47
+ _require_aiter()
48
+ # Import ensures the custom op is registered.
49
+ from vllm._aiter_ops import rocm_aiter_ops # noqa: F401
50
+
51
+ batch_size, nhead = 4, 128
52
+
53
+ q = torch.randn(batch_size, nhead, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
54
+ kv_buffer = torch.randn(64, 1, 1, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
55
+ o = torch.zeros(batch_size, nhead, V_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
56
+ qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda")
57
+ kv_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") * 16
58
+ kv_indices = torch.arange(0, 64, dtype=torch.int32, device="cuda")
59
+ kv_last_page_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda")
60
+
61
+ opcheck(
62
+ torch.ops.vllm.rocm_aiter_mla_decode_fwd,
63
+ (q, kv_buffer, o, qo_indptr, 1),
64
+ {
65
+ "kv_indptr": kv_indptr,
66
+ "kv_indices": kv_indices,
67
+ "kv_last_page_lens": kv_last_page_lens,
68
+ "sm_scale": Q_HEAD_DIM**-0.5,
69
+ "logit_cap": 0.0,
70
+ "q_scale": None,
71
+ "kv_scale": None,
72
+ "work_meta_data": None,
73
+ "work_indptr": None,
74
+ "work_info_set": None,
75
+ "reduce_indptr": None,
76
+ "reduce_final_map": None,
77
+ "reduce_partial_map": None,
78
+ },
79
+ )