[ROCm] Fix AITER all-reduce fusion coverage (#50802)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Aug 6, 2026 at 01:52 UTC
7e724dca897dda9e3fd05abcfe5b0f6d77c564b1
2 files changed
+95
-46
tests/compile/passes/distributed/test_fusion_all_reduce.py
+60
-46
@@ -43,6 +43,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
43
kFp8StaticTensorSym,
44
)
45
from vllm.platforms import current_platform
46
+from vllm.utils.network_utils import get_open_port
47
from vllm.utils.system_utils import update_environment_variables
48
from vllm.utils.torch_utils import set_random_seed
49
@@ -257,10 +258,9 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
258
modeling the DSv3.2 indexer fan-out)
259
-> ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern``
260
260
- The chain feeds the next AllReduce by dequantizing the FP8 output (FP8
261
- cast back to bf16 multiplied by the per-group scale), which is enough to
262
- keep the matmul chain bf16 without depending on a real FP8 block-scaled
263
- GEMM kernel.
261
+ Each site starts from the same bf16 activation and is returned separately.
262
+ This avoids compounding four quantize/dequantize errors while still testing
263
+ all four pattern firings without depending on a real FP8 block-scaled GEMM.
264
"""
265
266
quant_group_size = 128
@@ -276,6 +276,7 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
276
super().__init__()
277
self.hidden_size = hidden_size
278
self.eps = eps
279
+ self.dtype = dtype
280
assert hidden_size % self.quant_group_size == 0, (
281
f"hidden_size ({hidden_size}) must be a multiple of "
282
f"quant_group_size ({self.quant_group_size}) for per-group FP8 quant"
@@ -291,44 +292,39 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
292
rms, self.quant_group_size
293
)
294
294
- def _dequantize_to_bf16(
295
- self, q: torch.Tensor, s: torch.Tensor, ref: torch.Tensor
296
- ) -> torch.Tensor:
295
+ def _dequantize_to_bf16(self, q: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
296
# Broadcast the per-group scale across each group of `quant_group_size`
298
- # so we can chain the FP8 output back into a bf16 matmul. This avoids
299
- # depending on a real FP8 block-scaled GEMM kernel in the test.
300
- s_full = s.repeat_interleave(self.quant_group_size, dim=-1).to(ref.dtype)
301
- return q.to(ref.dtype) * s_full
297
+ # to validate the quantized result without depending on a real FP8
298
+ # block-scaled GEMM kernel in the test.
299
+ s_full = s.repeat_interleave(self.quant_group_size, dim=-1).to(self.dtype)
300
+ return q.to(self.dtype) * s_full
301
302
def forward(self, hidden_states):
303
z = torch.relu(hidden_states)
305
- x = resid = tensor_model_parallel_all_reduce(z)
306
- rms = self.norm[0](x)
307
- q0, s0 = self._group_quant(rms)
308
- y = self._dequantize_to_bf16(q0, s0, rms)
309
-
310
- z2 = torch.mm(y, self.w[0])
311
- x2 = tensor_model_parallel_all_reduce(z2)
312
- rms2, resid = self.norm[1](x2, resid)
313
- q1, s1 = self._group_quant(rms2)
314
- y2 = self._dequantize_to_bf16(q1, s1, rms2)
315
-
316
- z3 = torch.mm(y2, self.w[1])
317
- x3 = tensor_model_parallel_all_reduce(z3)
318
- rms3, resid = self.norm[2](x3, resid)
319
- q2, s2 = self._group_quant(rms3)
304
+ x0 = tensor_model_parallel_all_reduce(z)
305
+ rms0 = self.norm[0](x0)
306
+ q0, s0 = self._group_quant(rms0)
307
+ y0 = self._dequantize_to_bf16(q0, s0)
308
+
309
+ x1 = tensor_model_parallel_all_reduce(torch.mm(z, self.w[0]))
310
+ rms1, resid1 = self.norm[1](x1, z.clone())
311
+ q1, s1 = self._group_quant(rms1)
312
+ y1 = self._dequantize_to_bf16(q1, s1)
313
+
314
+ x2 = tensor_model_parallel_all_reduce(torch.mm(z, self.w[1]))
315
+ rms2, resid2 = self.norm[2](x2, z.clone())
316
+ q2, s2 = self._group_quant(rms2)
317
+ # Second consumer of ``rms2``: forces the with-indexer pattern.
318
+ idx2 = torch.ops.vllm.rocm_unquantized_gemm(rms2, self.indexer_w[0], None)
319
+ y2 = self._dequantize_to_bf16(q2, s2)
320
+
321
+ x3 = tensor_model_parallel_all_reduce(torch.mm(z, self.w[2]))
322
+ rms3, resid3 = self.norm[3](x3, z.clone())
323
+ q3, s3 = self._group_quant(rms3)
324
# Second consumer of ``rms3``: forces the with-indexer pattern.
321
- idx2 = torch.ops.vllm.rocm_unquantized_gemm(rms3, self.indexer_w[0], None)
322
- y3 = self._dequantize_to_bf16(q2, s2, rms3)
323
-
324
- z4 = torch.mm(y3, self.w[2])
325
- x4 = tensor_model_parallel_all_reduce(z4)
326
- rms4, resid = self.norm[3](x4, resid)
327
- q3, s3 = self._group_quant(rms4)
328
- # Second consumer of ``rms4``: forces the with-indexer pattern.
329
- idx3 = torch.ops.vllm.rocm_unquantized_gemm(rms4, self.indexer_w[1], None)
330
- y4 = self._dequantize_to_bf16(q3, s3, rms4)
331
- return y4, idx2, idx3
325
+ idx3 = torch.ops.vllm.rocm_unquantized_gemm(rms3, self.indexer_w[1], None)
326
+ y3 = self._dequantize_to_bf16(q3, s3)
327
+ return y0, y1, y2, y3, resid1, resid2, resid3, idx2, idx3
328
329
def ops_in_model_before(self):
330
return [
@@ -505,10 +501,12 @@ def test_all_reduce_fusion_pass_replace(
501
)
502
503
def run_torch_spawn(fn, nprocs):
504
+ master_port = get_open_port()
505
torch.multiprocessing.spawn(
506
fn,
507
args=(
508
num_processes,
509
+ master_port,
510
test_model,
511
batch_size,
512
seq_len,
@@ -529,6 +527,7 @@ def test_all_reduce_fusion_pass_replace(
527
def all_reduce_fusion_pass_on_test_model(
528
local_rank: int,
529
world_size: int,
530
+ master_port: int,
531
test_model_cls: torch.nn.Module,
532
batch_size: int,
533
seq_len: int,
@@ -553,7 +552,7 @@ def all_reduce_fusion_pass_on_test_model(
552
"LOCAL_RANK": str(local_rank),
553
"WORLD_SIZE": str(world_size),
554
"MASTER_ADDR": "localhost",
556
- "MASTER_PORT": "12345",
555
+ "MASTER_PORT": str(master_port),
556
"VLLM_FLASHINFER_ALLREDUCE_BACKEND": flashinfer_allreduce_backend,
557
"VLLM_ROCM_USE_AITER": str(int(use_aiter)),
558
"VLLM_ROCM_USE_AITER_CUSTOM_AR": str(int(use_aiter)),
@@ -681,10 +680,12 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
680
num_processes = 2
681
682
def run_torch_spawn(fn, nprocs):
683
+ master_port = get_open_port()
684
torch.multiprocessing.spawn(
685
fn,
686
args=(
687
num_processes,
688
+ master_port,
689
TestAiterAllReduceRMSNormGroupQuantFP8Model,
690
batch_size,
691
seq_len,
@@ -702,6 +703,7 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
703
def rocm_aiter_group_quant_fusion_pass_on_test_model(
704
local_rank: int,
705
world_size: int,
706
+ master_port: int,
707
test_model_cls: torch.nn.Module,
708
batch_size: int,
709
seq_len: int,
@@ -723,7 +725,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model(
725
"LOCAL_RANK": str(local_rank),
726
"WORLD_SIZE": str(world_size),
727
"MASTER_ADDR": "localhost",
726
- "MASTER_PORT": "12345",
728
+ "MASTER_PORT": str(master_port),
729
"VLLM_ROCM_USE_AITER": "1",
730
"VLLM_ROCM_USE_AITER_CUSTOM_AR": "1",
731
}
@@ -775,17 +777,29 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model(
777
778
results_unfused = model(hidden_states)
779
results_fused = compiled_model(hidden_states)
778
- # The fused per-group AR+RMS+QUANT op is bit-equivalent to the unfused
779
- # chain modulo the small AllReduce + RMSNorm reordering inside aiter.
780
- # Per-group FP8 quant introduces step noise <=1 per group; use the
781
- # same tolerance as the sibling FP8 static test.
782
- torch.testing.assert_close(results_unfused, results_fused, atol=1e-2, rtol=1e-2)
780
+ # AITER's kernel-level oracle permits up to 5% isolated FP8 values to
781
+ # fall outside a 5% allclose envelope after collective reordering.
782
+ for site, (reference, actual) in enumerate(
783
+ zip(results_unfused[:4], results_fused[:4])
784
+ ):
785
+ close = torch.isclose(reference, actual, atol=5e-2, rtol=5e-2)
786
+ mismatch_ratio = float((~close).float().mean().item())
787
+ assert mismatch_ratio <= 0.05, (
788
+ f"FP8 site {site} mismatch ratio {mismatch_ratio:.2%} exceeds 5%"
789
+ )
790
+
791
+ # Residual and bf16 indexer outputs do not incur FP8 quantization noise.
792
+ torch.testing.assert_close(
793
+ results_unfused[4:], results_fused[4:], atol=1e-2, rtol=1e-2
794
+ )
795
796
# Four pattern firings: norm[0] (no-add quant), norm[1] (add quant,
797
# single ``rms`` consumer), norm[2..3] (add quant + indexer fan-out).
798
assert all_reduce_fusion_pass.matched_count == 4, (
799
f"{all_reduce_fusion_pass.matched_count=}"
800
)
789
- backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
790
- backend.check_after_ops(model.ops_in_model_after())
801
+ backend.check_before_ops(model.ops_in_model_before(), fully_replaced=True)
802
+ fused_quant_op, fused_indexer_op = model.ops_in_model_after()
803
+ assert backend.op_count(fused_quant_op) == 2
804
+ assert backend.op_count(fused_indexer_op) == 2
805
del all_reduce_fusion_pass
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+35
@@ -1245,6 +1245,34 @@ class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement):
1245
return _replacement
1246
1247
1248
+class AiterAllreduceFusedAddRMSNormOutputOnlyPattern(
1249
+ AiterAllreduceFusedAddRMSNormPattern
1250
+):
1251
+ """Match the add-RMSNorm form when its residual output is dead."""
1252
+
1253
+ @property
1254
+ def pattern(self):
1255
+ pattern = super().pattern
1256
+
1257
+ def _pattern(
1258
+ residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
1259
+ ) -> torch.Tensor:
1260
+ return pattern(residual, input, weight)[0]
1261
+
1262
+ return _pattern
1263
+
1264
+ @property
1265
+ def replacement(self):
1266
+ replacement = super().replacement
1267
+
1268
+ def _replacement(
1269
+ residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
1270
+ ) -> torch.Tensor:
1271
+ return replacement(residual, input, weight)[0]
1272
+
1273
+ return _replacement
1274
+
1275
+
1276
class AiterAllreduceFusedRMSNormGroupQuantFP8Pattern(
1277
BasePattern, VllmPatternReplacement
1278
):
@@ -1624,6 +1652,13 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
1652
self.device,
1653
)
1654
)
1655
+ self.register(
1656
+ AiterAllreduceFusedAddRMSNormOutputOnlyPattern(
1657
+ epsilon,
1658
+ self.model_dtype,
1659
+ self.device,
1660
+ )
1661
+ )
1662
1663
# WARNING: This is a hack to clear the pattern matcher cache
1664
# and allow multiple values of epsilon.