[ROCm][CI] Add MiniMax reduce RMS kernel coverage (#55057)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Sep 3, 2026 at 04:07 UTC
096d8e8ce61f43a7e142fba4b8ccb845d793c737
3 files changed
+173
-5
.buildkite/test-amd.yaml
+36
@@ -1603,6 +1603,24 @@ steps:
1603
commands:
1604
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
1605
1606
+- label: ":amd: (MI300) MiniMax Reduce RMS Kernels" # TBD
1607
+ timeout_in_minutes: 45
1608
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
1609
+ dind: false
1610
+ agent_pool: mi300_2
1611
+ num_gpus: 2
1612
+ working_dir: "/vllm-workspace/tests"
1613
+ source_file_dependencies:
1614
+ - vllm/model_executor/layers/minimax_rms_norm/
1615
+ - vllm/distributed/
1616
+ - vllm/_aiter_ops.py
1617
+ - vllm/envs.py
1618
+ - vllm/platforms/rocm.py
1619
+ - tests/kernels/core/test_minimax_reduce_rms.py
1620
+ - tests/utils.py
1621
+ commands:
1622
+ - pytest -v -s kernels/core/test_minimax_reduce_rms.py
1623
+
1624
- label: ":amd: (MI300) KDA Kernels" # TBD
1625
timeout_in_minutes: 30
1626
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -3455,6 +3473,24 @@ steps:
3473
commands:
3474
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
3475
3476
+- label: ":amd: (MI355) MiniMax Reduce RMS Kernels" # TBD
3477
+ timeout_in_minutes: 45
3478
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3479
+ dind: false
3480
+ agent_pool: mi355_2
3481
+ num_gpus: 2
3482
+ working_dir: "/vllm-workspace/tests"
3483
+ source_file_dependencies:
3484
+ - vllm/model_executor/layers/minimax_rms_norm/
3485
+ - vllm/distributed/
3486
+ - vllm/_aiter_ops.py
3487
+ - vllm/envs.py
3488
+ - vllm/platforms/rocm.py
3489
+ - tests/kernels/core/test_minimax_reduce_rms.py
3490
+ - tests/utils.py
3491
+ commands:
3492
+ - pytest -v -s kernels/core/test_minimax_reduce_rms.py
3493
+
3494
- label: ":amd: (MI355) Kernels" # TBD
3495
timeout_in_minutes: 25
3496
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
.buildkite/test_areas/kernels.yaml
+23
-4
@@ -33,13 +33,32 @@ steps:
33
num_devices: 2
34
device: h100
35
source_file_dependencies:
36
- - csrc/minimax_reduce_rms_kernel.cu
37
- - csrc/minimax_reduce_rms_kernel.h
38
- - vllm/model_executor/layers/mamba/linear_attn.py
39
- - vllm/model_executor/layers/mamba/lamport_workspace.py
36
+ - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu
37
+ - csrc/libtorch_stable/minimax_reduce_rms_kernel.h
38
+ - csrc/libtorch_stable/ops.h
39
+ - csrc/libtorch_stable/torch_bindings.cpp
40
+ - vllm/model_executor/layers/minimax_rms_norm/
41
+ - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py
42
- tests/kernels/core/test_minimax_reduce_rms.py
43
commands:
44
- pytest -v -s kernels/core/test_minimax_reduce_rms.py
45
+ mirror:
46
+ amd:
47
+ label: ":amd: (MI355) MiniMax Reduce RMS Kernels"
48
+ dind: false
49
+ device: mi355_2
50
+ timeout_in_minutes: 45
51
+ working_dir: "/vllm-workspace/tests"
52
+ depends_on:
53
+ - image-build-amd
54
+ source_file_dependencies:
55
+ - vllm/model_executor/layers/minimax_rms_norm/
56
+ - vllm/distributed/
57
+ - vllm/_aiter_ops.py
58
+ - vllm/envs.py
59
+ - vllm/platforms/rocm.py
60
+ - tests/kernels/core/test_minimax_reduce_rms.py
61
+ - tests/utils.py
62
63
- label: ":nvidia: (H100) DeepSeek V4 Kernels"
64
key: deepseek-v4-kernel-test-h100
tests/kernels/core/test_minimax_reduce_rms.py
+114
-1
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests for MiniMax QK RMS-norm: NCCL reference vs Lamport fused kernel."""
3
+"""Tests for MiniMax QK RMS-norm fused all-reduce kernels and fallbacks."""
4
+
5
+from unittest.mock import patch
6
7
import pytest
8
import torch
@@ -9,6 +11,7 @@ from torch.multiprocessing import spawn
11
12
from tests.kernels.utils import opcheck
13
from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
14
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
15
from vllm.distributed import cleanup_dist_env_and_memory
16
from vllm.model_executor.layers.minimax_rms_norm import (
17
MiniMaxText01RMSNormTP,
@@ -124,6 +127,116 @@ def _worker_forward_qk(
127
cleanup_dist_env_and_memory()
128
129
130
+@ensure_current_vllm_config()
131
+def _worker_forward_qk_rocm_aiter(
132
+ local_rank: int,
133
+ world_size: int,
134
+ port: str,
135
+ hidden_q_full: int,
136
+ hidden_k_full: int,
137
+ seed: int,
138
+ eps: float,
139
+) -> None:
140
+ """Compare the ROCm AITER fusion with the eager TP implementation."""
141
+ device = torch.device(f"cuda:{local_rank}")
142
+ torch.accelerator.set_device_index(device)
143
+ init_test_distributed_environment(
144
+ world_size, 1, local_rank, port, local_rank=local_rank
145
+ )
146
+
147
+ try:
148
+ aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
149
+ assert aiter_ar is not None, "AITER custom all-reduce was not initialized"
150
+ assert not aiter_ar.disabled, "AITER custom all-reduce is disabled"
151
+ assert hasattr(aiter_ar.aiter_ca, "custom_fused_qknorm_ar")
152
+
153
+ hq = hidden_q_full // world_size
154
+ hk = hidden_k_full // world_size
155
+ test_cases = (
156
+ (1, torch.bfloat16),
157
+ (128, torch.bfloat16),
158
+ (333, torch.bfloat16),
159
+ (1, torch.float16),
160
+ (128, torch.float16),
161
+ (333, torch.float16),
162
+ )
163
+
164
+ for num_tokens, dtype in test_cases:
165
+ q_norm = MiniMaxText01RMSNormTP(hidden_q_full, eps=eps).cuda()
166
+ k_norm = MiniMaxText01RMSNormTP(hidden_k_full, eps=eps).cuda()
167
+
168
+ set_random_seed(seed)
169
+ q_weight = torch.randn(hidden_q_full, dtype=dtype, device="cuda")
170
+ k_weight = torch.randn(hidden_k_full, dtype=dtype, device="cuda")
171
+ q_norm.weight = nn.Parameter(
172
+ q_weight[local_rank * hq : (local_rank + 1) * hq]
173
+ )
174
+ k_norm.weight = nn.Parameter(
175
+ k_weight[local_rank * hk : (local_rank + 1) * hk]
176
+ )
177
+
178
+ torch.manual_seed(seed + num_tokens + local_rank)
179
+ qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda")
180
+
181
+ ref_q, ref_k = rms_norm_tp._minimax_qk_norm_tp_eager(
182
+ qkv.clone(),
183
+ q_norm.weight,
184
+ k_norm.weight,
185
+ hq,
186
+ hk,
187
+ world_size,
188
+ eps,
189
+ )
190
+
191
+ with (
192
+ patch.object(rms_norm_tp, "_MINIMAX_FUSED_AR_RMS_QK", None),
193
+ patch.object(
194
+ rms_norm_tp,
195
+ "_minimax_qk_norm_tp_fallback",
196
+ side_effect=AssertionError(
197
+ "ROCm AITER test unexpectedly used the fallback"
198
+ ),
199
+ ),
200
+ ):
201
+ fused_q, fused_k, fused_v = MiniMaxText01RMSNormTP.forward_qkv(
202
+ q_norm, k_norm, qkv.clone(), hq, hk
203
+ )
204
+
205
+ torch.accelerator.synchronize()
206
+ torch.testing.assert_close(fused_q, ref_q, atol=3e-2, rtol=3e-2)
207
+ torch.testing.assert_close(fused_k, ref_k, atol=3e-2, rtol=3e-2)
208
+ assert torch.equal(fused_v, qkv[:, hq + hk :])
209
+ finally:
210
+ cleanup_dist_env_and_memory()
211
+
212
+
213
+@pytest.mark.distributed(num_gpus=2)
214
+@pytest.mark.skipif(
215
+ not current_platform.is_rocm(),
216
+ reason="ROCm required",
217
+)
218
+def test_minimax_reduce_rms_qk_rocm_aiter(
219
+ monkeypatch: pytest.MonkeyPatch,
220
+) -> None:
221
+ """The MiniMax TP QK path dispatches to AITER and matches eager RMSNorm."""
222
+ assert is_aiter_found_and_supported(), "AITER requires ROCm CDNA3 or newer"
223
+
224
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
225
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1")
226
+ monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE")
227
+
228
+ world_size = 2
229
+ if current_platform.device_count() < world_size:
230
+ pytest.skip("Need at least two GPUs")
231
+
232
+ spawn(
233
+ _worker_forward_qk_rocm_aiter,
234
+ args=(world_size, str(get_open_port()), 6144, 1024, 42, 1e-6),
235
+ nprocs=world_size,
236
+ join=True,
237
+ )
238
+
239
+
240
@pytest.mark.skipif(
241
not current_platform.is_cuda(),
242
reason="CUDA required",