[ROCm][Compile] Support BF16 AsyncTP fusion (#58098)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Sep 22, 2026 at 20:36 UTC
0eb42dbbfce96477f4ca174980f4d68336fb1971
9 files changed
+310
-5
.buildkite/test-amd.yaml
+30
@@ -5069,3 +5069,33 @@ steps:
5069
commands:
5070
- pip install 'modelscope<1.38'
5071
- pytest -v -s test_regression.py
5072
+
5073
+- label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep'
5074
+ timeout_in_minutes: 30
5075
+ working_dir: /vllm-workspace/
5076
+ source_file_dependencies:
5077
+ - vllm/compilation/
5078
+ - vllm/model_executor/layers/linear.py
5079
+ - vllm/model_executor/layers/utils.py
5080
+ - vllm/distributed/
5081
+ - vllm/v1/attention/backends/triton_attn.py
5082
+ - tests/compile/fusions_e2e/
5083
+ - tests/compile/passes/distributed/test_async_tp.py
5084
+ - tests/compile/rocm/test_async_tp.py
5085
+ - tests/compile/correctness_e2e/test_async_tp.py
5086
+ - vllm/model_executor/kernels/linear/
5087
+ - vllm/model_executor/models/qwen3.py
5088
+ commands:
5089
+ - rocm-smi
5090
+ - pytest -v -s tests/compile/rocm/test_async_tp.py
5091
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py
5092
+ - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness
5093
+ dind: false
5094
+ agent_pool: mi355_2
5095
+ num_gpus: 2
5096
+ mirror_hardwares:
5097
+ - amdexperimental
5098
+ - amdproduction
5099
+ - amdgfx950nightly
5100
+ - amdmi355
5101
+ optional: true
.buildkite/test_areas/compile.yaml
+27
@@ -221,6 +221,33 @@ steps:
221
- nvidia-smi
222
# Run just llama3 (fp8 & bf16) for all config combinations
223
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3"
224
+ mirror:
225
+ amd:
226
+ label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep'
227
+ dind: false
228
+ device: mi355_2
229
+ num_devices: 2
230
+ working_dir: /vllm-workspace/
231
+ timeout_in_minutes: 30
232
+ depends_on:
233
+ - image-build-amd
234
+ source_file_dependencies:
235
+ - vllm/compilation/
236
+ - vllm/model_executor/layers/linear.py
237
+ - vllm/model_executor/layers/utils.py
238
+ - vllm/distributed/
239
+ - vllm/v1/attention/backends/triton_attn.py
240
+ - tests/compile/fusions_e2e/
241
+ - tests/compile/passes/distributed/test_async_tp.py
242
+ - tests/compile/rocm/test_async_tp.py
243
+ - tests/compile/correctness_e2e/test_async_tp.py
244
+ - vllm/model_executor/kernels/linear/
245
+ - vllm/model_executor/models/qwen3.py
246
+ commands:
247
+ - rocm-smi
248
+ - pytest -v -s tests/compile/rocm/test_async_tp.py
249
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py
250
+ - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness
251
252
- label: ":nvidia: (B200) Fusion E2E TP2"
253
key: fusion-e2e-tp2-b200
tests/compile/correctness_e2e/test_async_tp.py
+76
@@ -2,18 +2,22 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import json
5
+import logging
6
7
import pytest
8
9
+from tests.compile.fusions_e2e.common import FUSION_LOG_PATTERNS
10
from tests.models.registry import HF_EXAMPLE_MODELS
11
from tests.utils import (
12
compare_two_settings,
13
create_new_process_for_each_test,
14
+ multi_gpu_test,
15
)
16
from vllm.config import (
17
CompilationMode,
18
)
19
from vllm.platforms import current_platform
20
+from vllm.transformers_utils.repo_utils import hf_api
21
from vllm.utils.flashinfer import has_flashinfer
22
23
NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4"
@@ -26,6 +30,78 @@ NVFP4_HF_OVERRIDES = {
30
}
31
32
33
+@multi_gpu_test(num_gpus=2)
34
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract")
35
+def test_rocm_async_tp_bf16_output_correctness(caplog_mp_spawn):
36
+ """Trained-model outputs must agree with unfused TP, including padded decode."""
37
+ # A snapshot path pins both servers and the comparison helper's tokenizer.
38
+ model = hf_api().snapshot_download(
39
+ "Qwen/Qwen3-0.6B", revision="c1899de289a04d12100db370d81485cdf75e47ca"
40
+ )
41
+ common_args = [
42
+ "--dtype",
43
+ "bfloat16",
44
+ "--load-format",
45
+ "auto",
46
+ "--tensor-parallel-size",
47
+ "2",
48
+ "--distributed-executor-backend",
49
+ "mp",
50
+ "--max-model-len",
51
+ "512",
52
+ "--max-num-batched-tokens",
53
+ "512",
54
+ "--max-num-seqs",
55
+ "8",
56
+ "--kv-cache-memory-bytes",
57
+ str(64 * 1024**2),
58
+ "--gpu-memory-utilization",
59
+ "0.01",
60
+ "--attention-backend",
61
+ "TRITON_ATTN",
62
+ "--disable-custom-all-reduce",
63
+ ]
64
+
65
+ def settings(enabled):
66
+ return [
67
+ *common_args,
68
+ "--compilation-config",
69
+ json.dumps(
70
+ {
71
+ "mode": CompilationMode.VLLM_COMPILE,
72
+ "cudagraph_mode": "NONE",
73
+ "splitting_ops": [],
74
+ "inductor_compile_config": {"force_disable_caches": True},
75
+ "pass_config": {
76
+ "enable_sp": enabled,
77
+ "fuse_gemm_comms": enabled,
78
+ "fuse_allreduce_rms": False,
79
+ "sp_min_token_num": 1,
80
+ },
81
+ }
82
+ ),
83
+ ]
84
+
85
+ env = {"VLLM_ROCM_USE_AITER": "0", "VLLM_DISABLE_COMPILE_CACHE": "1"}
86
+ with caplog_mp_spawn(logging.DEBUG) as logs:
87
+ compare_two_settings(
88
+ model,
89
+ settings(False),
90
+ settings(True),
91
+ env1=env,
92
+ env2=env,
93
+ method="generate",
94
+ force_v1_runner=True,
95
+ )
96
+ matches = [
97
+ int(count) for count in FUSION_LOG_PATTERNS["async_tp"].findall(logs.text)
98
+ ]
99
+ print(f"AsyncTP rewrite counts: {matches}")
100
+ assert len(matches) >= 2 and all(count > 0 for count in matches), (
101
+ f"Expected positive AsyncTP rewrites on both TP workers, found {matches}"
102
+ )
103
+
104
+
105
@create_new_process_for_each_test()
106
@pytest.mark.parametrize(
107
"model_id",
tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py
new
+58
@@ -0,0 +1,58 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Bounded ROCm model coverage for actual sequence-parallel AsyncTP rewrites."""
4
+
5
+import pytest
6
+
7
+from tests.utils import multi_gpu_test
8
+from vllm.config import PassConfig
9
+from vllm.platforms import current_platform
10
+
11
+from .common import INDUCTOR_GRAPH_PARTITION, custom_ops_combos
12
+from .models import TRITON_ATTN, llama3_8b
13
+
14
+
15
+@multi_gpu_test(num_gpus=2)
16
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract")
17
+@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
18
+@pytest.mark.parametrize("custom_ops", tuple(custom_ops_combos("rms_norm")))
19
+@pytest.mark.parametrize("attn_backend", [TRITON_ATTN])
20
+def test_tp2_rocm_async_tp_bf16_fusions(
21
+ inductor_graph_partition, custom_ops, attn_backend, run_e2e_fusion_test
22
+):
23
+ n_layers = 4
24
+ model_kwargs = dict(
25
+ hf_overrides=dict(
26
+ num_hidden_layers=n_layers,
27
+ hidden_size=512,
28
+ intermediate_size=1024,
29
+ num_attention_heads=8,
30
+ num_key_value_heads=2,
31
+ head_dim=64,
32
+ ),
33
+ load_format="dummy",
34
+ max_model_len=1024,
35
+ kv_cache_memory_bytes=64 * 1024**2,
36
+ gpu_memory_utilization=0.01,
37
+ max_num_batched_tokens=1024,
38
+ max_num_seqs=32,
39
+ )
40
+ run_e2e_fusion_test(
41
+ llama3_8b.model_name,
42
+ llama3_8b.matches(n_layers),
43
+ model_kwargs,
44
+ attn_backend,
45
+ dict(
46
+ use_inductor_graph_partition=inductor_graph_partition,
47
+ custom_ops=custom_ops.split(","),
48
+ pass_config=PassConfig(
49
+ enable_qk_norm_rope_fusion=True,
50
+ enable_sp=True,
51
+ fuse_gemm_comms=True,
52
+ fuse_allreduce_rms=False,
53
+ sp_min_token_num=512,
54
+ ),
55
+ ),
56
+ ["norm_rope_fusion", "sequence_parallel", "async_tp"],
57
+ tp_size=2,
58
+ )
tests/compile/passes/distributed/test_async_tp.py
+14
-1
@@ -303,6 +303,7 @@ def async_tp_pass_on_test_model(
303
dtype: torch.dtype,
304
dynamic: bool,
305
distributed_init_method: str,
306
+ check_numerics: bool = False,
307
):
308
set_random_seed(0)
309
@@ -363,12 +364,24 @@ def async_tp_pass_on_test_model(
364
hidden_states = torch.randn(
365
(batch_size * seq_len, hidden_size), dtype=dtype, requires_grad=False
366
)
367
+ if check_numerics:
368
+ hidden_states += local_rank * 0.125
369
370
if dynamic:
371
torch._dynamo.mark_dynamic(hidden_states, 0)
372
373
compiled_model = torch.compile(model, backend=backend)
371
- compiled_model(hidden_states)
374
+ actual = compiled_model(hidden_states)
375
+
376
+ if check_numerics:
377
+ # These are new checks, so no existing tolerance is relaxed.
378
+ # Allow BF16 rounding vs. eager, but keep the lifetime check exact.
379
+ expected = model(hidden_states)
380
+ torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2)
381
+ snapshot = actual.clone()
382
+ changed = compiled_model(-hidden_states)
383
+ torch.testing.assert_close(changed, -expected, rtol=1e-2, atol=1e-2)
384
+ torch.testing.assert_close(actual, snapshot, rtol=0, atol=0)
385
386
assert async_tp_pass.matched_count == 1
387
tests/compile/rocm/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/compile/rocm/test_async_tp.py
new
+75
@@ -0,0 +1,75 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ROCm AsyncTP numerical coverage kept outside the shared distributed suite."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from tests.compile.passes.distributed.test_async_tp import (
9
+ TestAGMMModel,
10
+ TestMMRSModel,
11
+ async_tp_pass_on_test_model,
12
+)
13
+from tests.utils import multi_gpu_test
14
+from vllm.distributed import (
15
+ tensor_model_parallel_all_gather,
16
+ tensor_model_parallel_reduce_scatter,
17
+)
18
+from vllm.platforms import current_platform
19
+from vllm.utils.network_utils import get_file_store_init_method
20
+
21
+
22
+class TestRocmMMRSModel(TestMMRSModel):
23
+ def forward(self, hidden_states):
24
+ view = hidden_states.reshape(-1, self.hidden_size)
25
+ mm = torch.ops.vllm.rocm_unquantized_gemm(view, self.gate_proj, None)
26
+ return tensor_model_parallel_reduce_scatter(mm, dim=0)
27
+
28
+ def ops_in_model_before(self):
29
+ return super().ops_in_model_before() + [
30
+ torch.ops.vllm.rocm_unquantized_gemm.default
31
+ ]
32
+
33
+
34
+class TestRocmAGMMModel(TestAGMMModel):
35
+ def __init__(self, hidden_size=16, dtype=torch.bfloat16):
36
+ super().__init__(hidden_size, dtype)
37
+ self.weight = torch.nn.Parameter(
38
+ torch.empty((hidden_size * 2, hidden_size)), requires_grad=False
39
+ )
40
+ torch.nn.init.normal_(self.weight, std=0.02)
41
+
42
+ def forward(self, hidden_states):
43
+ view = hidden_states.reshape(-1, self.hidden_size)
44
+ gathered = tensor_model_parallel_all_gather(view, dim=0)
45
+ return torch.ops.vllm.rocm_unquantized_gemm(gathered, self.weight, None)
46
+
47
+ def ops_in_model_before(self):
48
+ return super().ops_in_model_before() + [
49
+ torch.ops.vllm.rocm_unquantized_gemm.default
50
+ ]
51
+
52
+
53
+@multi_gpu_test(num_gpus=2)
54
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend contract")
55
+@pytest.mark.parametrize(
56
+ "test_model", [TestMMRSModel, TestAGMMModel, TestRocmMMRSModel, TestRocmAGMMModel]
57
+)
58
+@pytest.mark.parametrize("dynamic", [False, True])
59
+def test_rocm_async_tp_bf16_rewrite_and_correctness(test_model, dynamic):
60
+ """Real symmetric-memory AG/GEMM and GEMM/RS must replace the collectives."""
61
+ torch.multiprocessing.spawn(
62
+ async_tp_pass_on_test_model,
63
+ args=(
64
+ 2,
65
+ test_model,
66
+ 8,
67
+ 16,
68
+ 128,
69
+ torch.bfloat16,
70
+ dynamic,
71
+ get_file_store_init_method(),
72
+ True,
73
+ ),
74
+ nprocs=2,
75
+ )
vllm/compilation/passes/fusion/collective_fusion.py
+27
-3
@@ -339,6 +339,8 @@ class BasePattern:
339
340
341
class GEMMReduceScatterPattern(BasePattern):
342
+ rocm_gemm = False
343
+
344
def get_inputs(self) -> list[torch.Tensor]:
345
mul = torch.empty([16, 4], device=self.device, dtype=self.dtype)
346
mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
@@ -346,7 +348,10 @@ class GEMMReduceScatterPattern(BasePattern):
348
349
def register(self, pm_pass: PatternMatcherPass) -> None:
350
def pattern(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
349
- mm = torch.ops.aten.mm.default(mul, mm_weight)
351
+ if self.rocm_gemm:
352
+ mm = torch.ops.vllm.rocm_unquantized_gemm(mul, mm_weight, None)
353
+ else:
354
+ mm = torch.ops.aten.mm.default(mul, mm_weight)
355
reduce_scatter = torch.ops.vllm.reduce_scatter.default(
356
mm,
357
dim=0,
@@ -358,7 +363,7 @@ class GEMMReduceScatterPattern(BasePattern):
363
def replacement(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
364
gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter(
365
mul,
361
- mm_weight,
366
+ mm_weight.t() if self.rocm_gemm else mm_weight,
367
"sum",
368
scatter_dim=0,
369
group_name=self.tp.device_group.group_name,
@@ -371,7 +376,13 @@ class GEMMReduceScatterPattern(BasePattern):
376
)
377
378
379
+class RocmGEMMReduceScatterPattern(GEMMReduceScatterPattern):
380
+ rocm_gemm = True
381
+
382
+
383
class AllGatherGEMMPattern(BasePattern):
384
+ rocm_gemm = False
385
+
386
def get_inputs(self) -> list[torch.Tensor]:
387
x = torch.empty([4, 4], device=self.device, dtype=self.dtype)
388
weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
@@ -390,12 +401,14 @@ class AllGatherGEMMPattern(BasePattern):
401
group_name=self.tp.unique_name,
402
)
403
404
+ if self.rocm_gemm:
405
+ return torch.ops.vllm.rocm_unquantized_gemm(all_gather, weight, None)
406
return torch.ops.aten.mm.default(all_gather, weight)
407
408
def replacement(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
409
ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul(
410
x,
398
- [weight],
411
+ [weight.t() if self.rocm_gemm else weight],
412
gather_dim=0,
413
group_name=self.tp.device_group.group_name,
414
)
@@ -406,6 +419,10 @@ class AllGatherGEMMPattern(BasePattern):
419
)
420
421
422
+class RocmAllGatherGEMMPattern(AllGatherGEMMPattern):
423
+ rocm_gemm = True
424
+
425
+
426
class ScaledMMReduceScatterPattern(BasePattern):
427
def get_inputs(self) -> list[torch.Tensor]:
428
input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
@@ -906,6 +923,13 @@ class AsyncTPPass(VllmFusionPatternMatcherPass):
923
GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.pm_pass)
924
925
AllGatherGEMMPattern(self.model_dtype, self.device).register(self.pm_pass)
926
+ if current_platform.is_rocm() and self.model_dtype == torch.bfloat16:
927
+ RocmGEMMReduceScatterPattern(self.model_dtype, self.device).register(
928
+ self.pm_pass
929
+ )
930
+ RocmAllGatherGEMMPattern(self.model_dtype, self.device).register(
931
+ self.pm_pass
932
+ )
933
934
# These fusions are enabled only for bfloat16 models because
935
# `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling
vllm/compilation/passes/pass_manager.py
+1
-1
@@ -41,6 +41,7 @@ if current_platform.is_cuda_alike() or current_platform.is_xpu():
41
if current_platform.is_cuda_alike():
42
from .fusion.act_quant_fusion import ActivationQuantFusionPass
43
from .fusion.attn_quant_fusion import AttnQuantFusionPass
44
+ from .fusion.collective_fusion import AsyncTPPass
45
from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass
46
from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass
47
from .fusion.qk_norm_rope_kvcache_fusion import QkNormRopeKvCacheFusionPass
@@ -50,7 +51,6 @@ if current_platform.is_cuda_alike():
51
52
if current_platform.is_cuda():
53
from .fusion.allreduce_rms_fusion import AllReduceFusionPass
53
- from .fusion.collective_fusion import AsyncTPPass
54
55
if current_platform.is_xpu():
56
from .fusion.act_quant_fusion import ActivationQuantFusionPass