[ROCm][CI] Mirror generic GEMM-RS/AR on MI355 (#58433)
Signed-off-by: Andreas Karatzas <andreas.karatzas@protonmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Andreas Karatzas committed
Sep 29, 2026 at 05:28 UTC
05d89636fca80d806732ced636ef505f42e6fe35
4 files changed
+289
-4
.buildkite/test-amd.yaml
+29
-2
@@ -3419,6 +3419,33 @@ steps:
3419
- pytest -v -s distributed/test_sharded_rdt_producer.py
3420
- pytest -v -s distributed/test_sharded_rdt_trainer.py
3421
3422
+- label: ":amd: (MI355) GEMM-RS-AR"
3423
+ timeout_in_minutes: 20
3424
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3425
+ dind: false
3426
+ agent_pool: mi355_4
3427
+ num_gpus: 4
3428
+ optional: true
3429
+ working_dir: "/vllm-workspace/"
3430
+ source_file_dependencies:
3431
+ - .buildkite/test-amd.yaml
3432
+ - tests/kernels/test_gemm_rs_ar.py
3433
+ - tests/kernels/test_rocm_gemm_rs_ar.py
3434
+ - tests/conftest.py
3435
+ - vllm/models/common/ops/sequence_parallel.py
3436
+ - vllm/model_executor/layers/linear.py
3437
+ - vllm/model_executor/layers/utils.py
3438
+ - vllm/model_executor/layers/quantization/online/
3439
+ - vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
3440
+ - vllm/model_executor/kernels/linear/
3441
+ - vllm/distributed/
3442
+ - vllm/_aiter_ops.py
3443
+ - vllm/_custom_ops.py
3444
+ - csrc/rocm/skinny_gemms.cu
3445
+ - vllm/platforms/rocm.py
3446
+ commands:
3447
+ - pytest -v -s tests/kernels/test_rocm_gemm_rs_ar.py
3448
+
3449
- label: ":amd: (MI355) Distributed DP Basic" # TBD
3450
timeout_in_minutes: 45
3451
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -4533,7 +4560,7 @@ steps:
4560
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
4561
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi355-deepseek-r1-dp.txt
4562
4536
-- label: ":amd: (MI355) LM Eval DeepSeek V3.2 TP"
4563
+- label: ":amd: (MI355) LM Eval DeepSeek V3-2 TP"
4564
timeout_in_minutes: 65
4565
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
4566
dind: false
@@ -4566,7 +4593,7 @@ steps:
4593
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
4594
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi355-deepseek-v32-tp.txt
4595
4569
-- label: ":amd: (MI355) LM Eval DeepSeek V3.2 DP"
4596
+- label: ":amd: (MI355) LM Eval DeepSeek V3-2 DP"
4597
timeout_in_minutes: 65
4598
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
4599
dind: false
.buildkite/test_areas/distributed.yaml
+28
@@ -473,6 +473,34 @@ steps:
473
- tests/kernels/test_gemm_rs_ar.py
474
commands:
475
- pytest -v -s tests/kernels/test_gemm_rs_ar.py
476
+ mirror:
477
+ amd:
478
+ label: ":amd: (MI355) GEMM-RS-AR"
479
+ dind: false
480
+ device: mi355_4
481
+ num_devices: 4
482
+ timeout_in_minutes: 20
483
+ working_dir: /vllm-workspace/
484
+ depends_on:
485
+ - image-build-amd
486
+ source_file_dependencies:
487
+ - .buildkite/test_areas/distributed.yaml
488
+ - tests/kernels/test_gemm_rs_ar.py
489
+ - tests/kernels/test_rocm_gemm_rs_ar.py
490
+ - tests/conftest.py
491
+ - vllm/models/common/ops/sequence_parallel.py
492
+ - vllm/model_executor/layers/linear.py
493
+ - vllm/model_executor/layers/utils.py
494
+ - vllm/model_executor/layers/quantization/online/
495
+ - vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
496
+ - vllm/model_executor/kernels/linear/
497
+ - vllm/distributed/
498
+ - vllm/_aiter_ops.py
499
+ - vllm/_custom_ops.py
500
+ - csrc/rocm/skinny_gemms.cu
501
+ - vllm/platforms/rocm.py
502
+ commands:
503
+ - pytest -v -s tests/kernels/test_rocm_gemm_rs_ar.py
504
505
- label: ":nvidia: (L4) Distributed 2-Node"
506
key: 2-node-test-4-gpus
.buildkite/test_areas/lm_eval.yaml
+2
-2
@@ -578,7 +578,7 @@ steps:
578
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-v32-tp.txt
579
mirror:
580
amd:
581
- label: ":amd: (MI355) LM Eval DeepSeek V3.2 TP"
581
+ label: ":amd: (MI355) LM Eval DeepSeek V3-2 TP"
582
dind: false
583
device: mi355_4
584
num_devices: 4
@@ -621,7 +621,7 @@ steps:
621
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-v32-dp.txt
622
mirror:
623
amd:
624
- label: ":amd: (MI355) LM Eval DeepSeek V3.2 DP"
624
+ label: ":amd: (MI355) LM Eval DeepSeek V3-2 DP"
625
dind: false
626
device: mi355_4
627
num_devices: 4
tests/kernels/test_rocm_gemm_rs_ar.py
new
+230
@@ -0,0 +1,230 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ROCm projection/collective contract corresponding to the SM100 GEMM-RS/AR.
4
+
5
+Exercise the production RowParallelLinear and sequence-parallel reduction
6
+paths: native BF16/MXFP8 GEMM, AITER all-reduce, and RCCL reduce-scatter. This
7
+checks the same observable contract, not NVIDIA's fused GEMM/communication
8
+implementation or its performance.
9
+"""
10
+
11
+import pytest
12
+import torch
13
+import torch.distributed as dist
14
+import torch.multiprocessing as mp
15
+
16
+from tests.kernels.test_gemm_rs_ar import _N, _SHAPES
17
+from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
18
+from vllm.config.quantization import QuantizationConfigArgs
19
+from vllm.distributed import cleanup_dist_env_and_memory, get_tp_group
20
+from vllm.distributed.parallel_state import (
21
+ graph_capture,
22
+ init_distributed_environment,
23
+ initialize_model_parallel,
24
+)
25
+from vllm.model_executor.layers.linear import RowParallelLinear
26
+from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
27
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
28
+ _mxfp8_e4m3_quantize_torch,
29
+ dequant_mxfp8_to_bf16,
30
+)
31
+from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter
32
+from vllm.platforms import current_platform
33
+from vllm.utils.network_utils import get_open_port
34
+from vllm.utils.system_utils import update_environment_variables
35
+
36
+# Retain the expanded parent's shape transitions, plus fully padded RS ranks.
37
+_ROCM_SHAPES = ((1, 512), (7, 768), *_SHAPES)
38
+
39
+
40
+def _reference(
41
+ x: torch.Tensor,
42
+ layer: RowParallelLinear,
43
+ all_reduce: bool,
44
+) -> torch.Tensor:
45
+ weight = layer.weight
46
+ if hasattr(layer, "weight_scale"):
47
+ # The production ROCm path uses the Triton quantizer and dot_scaled;
48
+ # use the independent torch quantizer and FP32 matmul as the oracle.
49
+ x_q, x_scale = _mxfp8_e4m3_quantize_torch(x)
50
+ x = dequant_mxfp8_to_bf16(x_q, x_scale)
51
+ weight = dequant_mxfp8_to_bf16(weight, layer.weight_scale)
52
+ partial = (x.float() @ weight.float().T).to(torch.bfloat16).float()
53
+ group = get_tp_group()
54
+ # Always use FP32 RCCL AR for the oracle, including the RS candidate.
55
+ dist.all_reduce(partial, group=group.device_group)
56
+ expected = partial.to(torch.bfloat16)
57
+ if all_reduce:
58
+ return expected
59
+ padding = -x.shape[0] % group.world_size
60
+ expected = torch.nn.functional.pad(expected, (0, 0, 0, padding))
61
+ return expected.chunk(group.world_size)[group.rank_in_group]
62
+
63
+
64
+def _apply(x: torch.Tensor, layer: RowParallelLinear) -> torch.Tensor:
65
+ output, _ = layer(x)
66
+ return output if layer.reduce_results else sp_reduce_scatter(output)
67
+
68
+
69
+def _assert_output(
70
+ actual: torch.Tensor, expected: torch.Tensor, M: int, all_reduce: bool
71
+) -> None:
72
+ assert actual.dtype == torch.bfloat16
73
+ assert actual.shape == expected.shape
74
+ assert torch.isfinite(actual).all()
75
+ # Match the parent's tolerance at 0.1 weight scale, then bound total error.
76
+ # BF16 partial/reduction rounding can dominate near a cancelled output.
77
+ torch.testing.assert_close(actual, expected, rtol=5e-2, atol=0.4)
78
+ error = (actual.float() - expected.float()).norm()
79
+ assert error <= 5e-3 * expected.float().norm() + 1e-6
80
+ if not all_reduce:
81
+ rank_start = get_tp_group().rank_in_group * actual.shape[0]
82
+ valid_rows = min(max(M - rank_start, 0), actual.shape[0])
83
+ assert not actual[valid_rows:].any(), "RS padding must remain zero"
84
+
85
+
86
+def _run_mode(
87
+ device: torch.device, backend: str, all_reduce: bool, world_size: int
88
+) -> None:
89
+ group = get_tp_group()
90
+ generator = torch.Generator(device=device)
91
+ quant_config = (
92
+ OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8"))
93
+ if backend == "mxfp8"
94
+ else None
95
+ )
96
+ layers = {}
97
+ for K in sorted({K for _, K in _ROCM_SHAPES}):
98
+ with torch.device(device):
99
+ layer = RowParallelLinear(
100
+ K * world_size,
101
+ _N,
102
+ bias=False,
103
+ params_dtype=torch.bfloat16,
104
+ quant_config=quant_config,
105
+ reduce_results=all_reduce,
106
+ )
107
+ generator.manual_seed(1000 + group.rank_in_group * 100 + K)
108
+ weight = (
109
+ torch.randn(_N, K, dtype=torch.bfloat16, device=device, generator=generator)
110
+ * 0.1
111
+ )
112
+ layer.weight = torch.nn.Parameter(weight, requires_grad=False)
113
+ layer.quant_method.process_weights_after_loading(layer)
114
+ if backend == "mxfp8":
115
+ from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
116
+ RocmDotScaledMxfp8LinearKernel,
117
+ )
118
+
119
+ assert isinstance(layer.quant_method.kernel, RocmDotScaledMxfp8LinearKernel)
120
+ assert layer.weight.dtype == torch.float8_e4m3fn
121
+ layers[K] = layer
122
+
123
+ def make_input(M: int, K: int, seed: int) -> torch.Tensor:
124
+ generator.manual_seed(seed + group.rank_in_group * 100)
125
+ return torch.randn(
126
+ M, K, dtype=torch.bfloat16, device=device, generator=generator
127
+ )
128
+
129
+ comm = group.device_communicator
130
+ assert comm is not None
131
+ assert comm.use_aiter_allreduce
132
+ assert comm.aiter_ar_comm is not None and not comm.aiter_ar_comm.disabled
133
+ assert comm.qr_comm is None or comm.qr_comm.disabled
134
+ assert comm.pynccl_comm is not None and not comm.pynccl_comm.disabled
135
+ for M, K in (*_ROCM_SHAPES, *_ROCM_SHAPES[::-1]):
136
+ x = make_input(M, K, 2000 + M + K)
137
+ if all_reduce:
138
+ probe = torch.empty(M, _N, dtype=x.dtype, device=device)
139
+ assert comm.aiter_ar_comm.should_custom_ar(probe)
140
+ expected = _reference(x, layers[K], all_reduce)
141
+ actual = _apply(x, layers[K])
142
+ _assert_output(actual, expected, M, all_reduce)
143
+
144
+ # Retained outputs must survive the next call's collective buffer reuse.
145
+ x = make_input(257, 768, 2500)
146
+ first = _apply(x, layers[768])
147
+ snapshot = first.clone()
148
+ second = _apply(-x, layers[768])
149
+ torch.accelerator.synchronize(device)
150
+ assert first.data_ptr() != second.data_ptr()
151
+ torch.testing.assert_close(first, snapshot, rtol=0, atol=0)
152
+
153
+ M, K = 1025, 4224
154
+ graph_x = make_input(M, K, 3000)
155
+ for _ in range(3):
156
+ _apply(graph_x, layers[K])
157
+ torch.accelerator.synchronize(device)
158
+ dist.barrier(group=group.device_group)
159
+ graph = torch.cuda.CUDAGraph()
160
+ # vLLM's context registers the AITER IPC addresses used by graph replay.
161
+ with (
162
+ graph_capture(device=device) as context,
163
+ torch.cuda.graph(graph, stream=context.stream),
164
+ ):
165
+ graph_output = _apply(graph_x, layers[K])
166
+ torch.cuda.current_stream().wait_stream(context.stream)
167
+ dist.barrier(group=group.device_group)
168
+ for replay in range(3):
169
+ graph_x.copy_(make_input(M, K, 3001 + replay))
170
+ expected = _reference(graph_x, layers[K], all_reduce)
171
+ graph.replay()
172
+ torch.accelerator.synchronize(device)
173
+ _assert_output(graph_output, expected, M, all_reduce)
174
+ dist.barrier(group=group.device_group)
175
+
176
+
177
+def _worker(rank: int, world_size: int, port: int, backend: str) -> None:
178
+ device = torch.device("cuda", rank)
179
+ torch.accelerator.set_device_index(device)
180
+ update_environment_variables(
181
+ {
182
+ "RANK": str(rank),
183
+ "LOCAL_RANK": str(rank),
184
+ "WORLD_SIZE": str(world_size),
185
+ "MASTER_ADDR": "localhost",
186
+ "MASTER_PORT": str(port),
187
+ }
188
+ )
189
+ config = VllmConfig()
190
+ config.model_config = ModelConfig(dtype="bfloat16")
191
+ try:
192
+ with set_current_vllm_config(config), torch.inference_mode():
193
+ init_distributed_environment()
194
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
195
+ for all_reduce in (False, True):
196
+ _run_mode(device, backend, all_reduce, world_size)
197
+ finally:
198
+ cleanup_dist_env_and_memory()
199
+
200
+
201
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm GEMM-RS/AR contract")
202
+@pytest.mark.parametrize("world_size", [2, 4])
203
+@pytest.mark.parametrize(
204
+ "backend",
205
+ [
206
+ "bf16",
207
+ pytest.param(
208
+ "mxfp8",
209
+ marks=pytest.mark.skipif(
210
+ not current_platform.supports_mx(), reason="native MXFP8 requires CDNA4"
211
+ ),
212
+ ),
213
+ ],
214
+)
215
+@pytest.mark.distributed(num_gpus=4)
216
+def test_rocm_gemm_rs_ar(
217
+ monkeypatch: pytest.MonkeyPatch, world_size: int, backend: str
218
+) -> None:
219
+ """Preserve rank reduction, tails, shape reuse, lifetime, and graph replay."""
220
+ if torch.accelerator.device_count() < world_size:
221
+ pytest.skip(f"GEMM-RS/AR requires {world_size} GPUs")
222
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
223
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1")
224
+ monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE")
225
+ try:
226
+ mp.spawn(
227
+ _worker, args=(world_size, get_open_port(), backend), nprocs=world_size
228
+ )
229
+ finally:
230
+ cleanup_dist_env_and_memory()