[ROCm][CI] Add MLA decode accuracy and determinism tests (#50480)
Signed-off-by: aarushjain29 <Aarushi.Jain2@amd.com>
Aarushi Jain committed
Aug 5, 2026 at 20:35 UTC
71f975af437ae5316679752e8cdb72cd1c8b7b1d
2 files changed
+261
.buildkite/test-amd.yaml
+36
@@ -1530,6 +1530,24 @@ steps:
1530
- pytest -v -s tests/ir
1531
- pytest -v -s tests/kernels/ir
1532
1533
+- label: Kernels MLA (MI300) # TBD
1534
+ timeout_in_minutes: 180
1535
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
1536
+ dind: false
1537
+ agent_pool: mi300_1
1538
+ working_dir: "/vllm-workspace/"
1539
+ source_file_dependencies:
1540
+ - vllm/v1/attention/backends/mla/rocm_aiter_mla.py
1541
+ - vllm/v1/attention/selector.py
1542
+ - vllm/platforms/rocm.py
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
+ commands:
1547
+ - rocm-smi
1548
+ - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
1549
+ - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py
1550
+
1551
- label: Kernels Attention Test %N # TBD
1552
timeout_in_minutes: 180
1553
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -3387,6 +3405,24 @@ steps:
3405
- pytest -v -s tests/kernels/attention/test_attention_selector.py
3406
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
3407
3408
+- label: Kernels MLA (MI355) # TBD
3409
+ timeout_in_minutes: 180
3410
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3411
+ dind: false
3412
+ agent_pool: mi355_1
3413
+ working_dir: "/vllm-workspace/"
3414
+ source_file_dependencies:
3415
+ - vllm/v1/attention/backends/mla/rocm_aiter_mla.py
3416
+ - vllm/v1/attention/selector.py
3417
+ - vllm/platforms/rocm.py
3418
+ - vllm/_aiter_ops.py
3419
+ - tests/kernels/attention/test_rocm_aiter_mla_decode.py
3420
+ - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
3421
+ commands:
3422
+ - rocm-smi
3423
+ - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
3424
+ - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py
3425
+
3426
- label: Kernels Attention Test %N # TBD
3427
timeout_in_minutes: 180
3428
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
tests/kernels/attention/test_rocm_aiter_mla_decode.py
new
+225
@@ -0,0 +1,225 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ROCm kernel correctness tests for AITER MLA decode.
4
+
5
+Compares ``rocm_aiter_ops.mla_decode_fwd`` against a pure PyTorch reference
6
+under the absorbed MLA formulation (DeepSeek-V3/V4) with varied batch sizes,
7
+head counts, and sequence lengths.
8
+"""
9
+
10
+import pytest
11
+import torch
12
+
13
+from vllm.platforms import current_platform
14
+from vllm.utils.torch_utils import set_random_seed
15
+
16
+_SKIP_NON_MI3XX = True
17
+if current_platform.is_rocm():
18
+ from vllm.platforms.rocm import on_mi3xx
19
+
20
+ _SKIP_NON_MI3XX = not on_mi3xx()
21
+
22
+pytestmark = [
23
+ pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"),
24
+ pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"),
25
+]
26
+
27
+# DeepSeek-V3/V4 MLA config.
28
+Q_HEAD_DIM = 576 # kv_lora_rank + qk_rope_head_dim
29
+V_HEAD_DIM = 512 # kv_lora_rank
30
+SM_SCALE = Q_HEAD_DIM**-0.5
31
+
32
+NUM_HEADS = [16, 128]
33
+BATCH_SIZES = [1, 4, 16]
34
+KV_SEQ_LENS = [16, 256]
35
+
36
+ATOL, RTOL = 1.5e-2, 1e-2
37
+
38
+
39
+def _require_aiter():
40
+ from vllm._aiter_ops import is_aiter_found_and_supported
41
+
42
+ if not is_aiter_found_and_supported():
43
+ pytest.skip("aiter is required on supported ROCm hardware for this test")
44
+
45
+
46
+def _ref_mla_decode(
47
+ q: torch.Tensor,
48
+ kv_buffer: torch.Tensor,
49
+ kv_indptr: torch.Tensor,
50
+ kv_indices: torch.Tensor,
51
+) -> torch.Tensor:
52
+ """Pure PyTorch reference for absorbed MLA decode.
53
+
54
+ K uses full q_head_dim (576) for scoring, V uses first v_head_dim (512)
55
+ for output. This mirrors the production kernel's absorbed formulation.
56
+ """
57
+ batch_size, num_heads, _ = q.shape
58
+ output = torch.zeros(
59
+ batch_size, num_heads, V_HEAD_DIM, dtype=q.dtype, device=q.device
60
+ )
61
+
62
+ for b in range(batch_size):
63
+ start = kv_indptr[b].item()
64
+ end = kv_indptr[b + 1].item()
65
+ token_indices = kv_indices[start:end]
66
+
67
+ k = kv_buffer[token_indices].float()
68
+ v = kv_buffer[token_indices, :V_HEAD_DIM].float()
69
+
70
+ for h in range(num_heads):
71
+ q_h = q[b, h, :].float()
72
+ scores = torch.mv(k, q_h) * SM_SCALE
73
+ weights = torch.softmax(scores, dim=0)
74
+ output[b, h, :] = torch.mv(v.t(), weights).to(q.dtype)
75
+
76
+ return output
77
+
78
+
79
+def _make_inputs(
80
+ batch_size: int,
81
+ nhead: int,
82
+ kv_seq_len: int,
83
+ *,
84
+ contiguous_indices: bool = True,
85
+):
86
+ """Build valid MLA decode inputs on the current CUDA device.
87
+
88
+ When contiguous_indices=False, tokens are randomly scattered in a 2x pool
89
+ to simulate real paged allocation.
90
+ """
91
+ from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper
92
+
93
+ actual_nhead = AiterMLAHelper.get_actual_mla_num_heads(nhead)
94
+ total_kv_tokens = batch_size * kv_seq_len
95
+
96
+ pool_size = total_kv_tokens if contiguous_indices else total_kv_tokens * 2
97
+
98
+ q = torch.randn(
99
+ batch_size, actual_nhead, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
100
+ )
101
+ kv_buffer = torch.randn(
102
+ pool_size, 1, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
103
+ )
104
+ o = torch.zeros(
105
+ batch_size, actual_nhead, V_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
106
+ )
107
+
108
+ qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda")
109
+ kv_indptr = (
110
+ torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") * kv_seq_len
111
+ )
112
+ kv_last_page_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda")
113
+
114
+ if contiguous_indices:
115
+ kv_indices = torch.arange(0, total_kv_tokens, dtype=torch.int32, device="cuda")
116
+ else:
117
+ perm = torch.randperm(pool_size, device="cuda")[:total_kv_tokens]
118
+ kv_indices = perm.to(torch.int32)
119
+
120
+ return {
121
+ "q": q,
122
+ "kv_buffer": kv_buffer,
123
+ "o": o,
124
+ "qo_indptr": qo_indptr,
125
+ "kv_indptr": kv_indptr,
126
+ "kv_indices": kv_indices,
127
+ "kv_last_page_lens": kv_last_page_lens,
128
+ "actual_nhead": actual_nhead,
129
+ }
130
+
131
+
132
+def _run_kernel(inputs: dict) -> torch.Tensor:
133
+ """Run decode through the production rocm_aiter_ops path."""
134
+ from vllm._aiter_ops import rocm_aiter_ops
135
+
136
+ rocm_aiter_ops.mla_decode_fwd(
137
+ inputs["q"],
138
+ inputs["kv_buffer"],
139
+ inputs["o"],
140
+ SM_SCALE,
141
+ inputs["qo_indptr"],
142
+ 1,
143
+ inputs["kv_indptr"],
144
+ inputs["kv_indices"],
145
+ inputs["kv_last_page_lens"],
146
+ )
147
+ return inputs["o"]
148
+
149
+
150
+def _ref_output(inputs: dict) -> torch.Tensor:
151
+ """Compute reference output from the same inputs."""
152
+ kv_flat = inputs["kv_buffer"].squeeze(1)
153
+ return _ref_mla_decode(
154
+ inputs["q"], kv_flat, inputs["kv_indptr"], inputs["kv_indices"]
155
+ )
156
+
157
+
158
+@pytest.mark.parametrize("nhead", NUM_HEADS)
159
+@pytest.mark.parametrize("batch_size", BATCH_SIZES)
160
+@pytest.mark.parametrize("kv_seq_len", KV_SEQ_LENS)
161
+@torch.inference_mode()
162
+def test_mla_decode_accuracy(
163
+ nhead: int,
164
+ batch_size: int,
165
+ kv_seq_len: int,
166
+) -> None:
167
+ """BF16 decode accuracy vs PyTorch reference (contiguous indices)."""
168
+ _require_aiter()
169
+ set_random_seed(0)
170
+
171
+ inputs = _make_inputs(batch_size, nhead, kv_seq_len)
172
+ output = _run_kernel(inputs)
173
+ output_ref = _ref_output(inputs)
174
+
175
+ torch.testing.assert_close(output, output_ref, atol=ATOL, rtol=RTOL)
176
+
177
+
178
+@pytest.mark.parametrize("nhead", NUM_HEADS)
179
+@pytest.mark.parametrize("batch_size", [1, 4, 8])
180
+@pytest.mark.parametrize("kv_seq_len", [16, 128])
181
+@torch.inference_mode()
182
+def test_mla_decode_noncontiguous_indices(
183
+ nhead: int,
184
+ batch_size: int,
185
+ kv_seq_len: int,
186
+) -> None:
187
+ """BF16 decode accuracy with shuffled kv_indices (paged allocation)."""
188
+ _require_aiter()
189
+ set_random_seed(0)
190
+
191
+ inputs = _make_inputs(batch_size, nhead, kv_seq_len, contiguous_indices=False)
192
+ output = _run_kernel(inputs)
193
+ output_ref = _ref_output(inputs)
194
+
195
+ torch.testing.assert_close(output, output_ref, atol=ATOL, rtol=RTOL)
196
+
197
+
198
+@torch.inference_mode()
199
+def test_mla_decode_determinism() -> None:
200
+ """Repeated decode calls produce bitwise-identical results."""
201
+ _require_aiter()
202
+ set_random_seed(0)
203
+
204
+ inputs = _make_inputs(batch_size=4, nhead=128, kv_seq_len=16)
205
+ out_first = _run_kernel(inputs).clone()
206
+
207
+ for _ in range(3):
208
+ inputs["o"] = torch.zeros_like(inputs["o"])
209
+ out = _run_kernel(inputs)
210
+ torch.testing.assert_close(out, out_first, atol=0, rtol=0)
211
+
212
+
213
+@torch.inference_mode()
214
+def test_mla_decode_smoke() -> None:
215
+ """Basic sanity: correct shape, dtype, finite, non-zero."""
216
+ _require_aiter()
217
+ set_random_seed(0)
218
+
219
+ inputs = _make_inputs(batch_size=4, nhead=128, kv_seq_len=64)
220
+ output = _run_kernel(inputs)
221
+
222
+ assert output.shape == (4, 128, V_HEAD_DIM)
223
+ assert output.dtype == torch.bfloat16
224
+ assert torch.isfinite(output).all()
225
+ assert not torch.all(output == 0)