A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 12:15 UTC
[AMD][CI][BugFix] Fix Distributed Compile Unit Tests (2xH100-2xMI300) group (#43120)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
May 28, 2026 at 16:39 UTC
9769e2df2ac94f28e8af7686b24ed62dba37c0fc
7 files changed
+54
-11
tests/compile/fusions_e2e/common.py
+4
@@ -20,6 +20,7 @@ class Matches(NamedTuple):
20
attn_quant_fusion: int = 0
21
# distributed
22
ar_rms_fusion: int = 0
23
+ aiter_ar_rms_fusion: int = 0
24
sequence_parallel: int = 0
25
async_tp: int = 0
26
@@ -97,6 +98,9 @@ FUSION_LOG_PATTERNS: dict[str, re.Pattern] = {
98
"ar_rms_fusion": re.compile(
99
r"allreduce_rms_fusion.py:\d+] Replaced (\d+) patterns"
100
),
101
+ "aiter_ar_rms_fusion": re.compile(
102
+ r"RocmAiterAllReduceFusionPass Replaced (\d+) patterns"
103
+ ),
104
"sequence_parallel": re.compile(
105
r"sequence_parallelism.py:\d+] Replaced (\d+) patterns"
106
),
tests/compile/fusions_e2e/conftest.py
+5
@@ -97,6 +97,11 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
97
f"attention backend '{attn_backend.backend.name}'"
98
)
99
100
+ if backend_name == "rocm_attn" and model_name == "openai/gpt-oss-20b":
101
+ pytest.skip(
102
+ "ROCM_ATTN does not support attention sinks (required by gpt-oss-20b)"
103
+ )
104
+
105
if attn_backend.backend.name == "FLASHINFER":
106
from vllm.utils.flashinfer import supports_trtllm_attention
107
tests/compile/fusions_e2e/models.py
+3
@@ -75,6 +75,7 @@ llama3_8b = ModelFusionInfo(
75
model_name="meta-llama/Llama-3.1-8B-Instruct",
76
matches=lambda n_layers: Matches(
77
ar_rms_fusion=n_layers * 2 + 1,
78
+ aiter_ar_rms_fusion=n_layers * 2,
79
sequence_parallel=n_layers * 2 + 1,
80
async_tp=n_layers * 4,
81
),
@@ -136,6 +137,7 @@ qwen3_a3b = ModelFusionInfo(
137
matches=lambda n_layers: Matches(
138
norm_rope_fusion=n_layers,
139
ar_rms_fusion=n_layers * 2 + 1,
140
+ aiter_ar_rms_fusion=n_layers * 2,
141
sequence_parallel=n_layers * 2 + 1,
142
async_tp=n_layers * 2,
143
),
@@ -211,6 +213,7 @@ gpt_oss_20b = ModelFusionInfo(
213
model_name="openai/gpt-oss-20b",
214
matches=lambda n_layers: Matches(
215
ar_rms_fusion=n_layers * 2 + 1,
216
+ aiter_ar_rms_fusion=n_layers + 1,
217
sequence_parallel=n_layers * 2 + 1,
218
async_tp=n_layers * 2,
219
),
tests/compile/fusions_e2e/test_tp2_ar_rms.py
+8
-1
@@ -84,6 +84,7 @@ def test_tp2_ar_rms_fp8_fusions(
84
model_kwargs["load_format"] = "dummy"
85
model_kwargs["max_model_len"] = 1024
86
model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
87
+ model_kwargs["disable_custom_all_reduce"] = False
88
89
compilation_config = dict(
90
use_inductor_graph_partition=inductor_graph_partition,
@@ -149,6 +150,7 @@ def test_tp2_ar_rms_fp4_fusions(
150
model_kwargs["load_format"] = "dummy"
151
model_kwargs["max_model_len"] = 1024
152
model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
153
+ model_kwargs["disable_custom_all_reduce"] = False
154
155
compilation_config = dict(
156
use_inductor_graph_partition=inductor_graph_partition,
@@ -213,6 +215,7 @@ def test_tp2_ar_rms_fusions(
215
model_kwargs["load_format"] = "dummy"
216
model_kwargs["max_model_len"] = 1024
217
model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
218
+ model_kwargs["disable_custom_all_reduce"] = False
219
220
compilation_config = dict(
221
use_inductor_graph_partition=inductor_graph_partition,
@@ -225,9 +228,13 @@ def test_tp2_ar_rms_fusions(
228
229
matches_check = [
230
"norm_rope_fusion",
228
- "ar_rms_fusion",
231
]
232
233
+ if current_platform.is_rocm():
234
+ matches_check.append("aiter_ar_rms_fusion")
235
+ else:
236
+ matches_check.append("ar_rms_fusion")
237
+
238
run_e2e_fusion_test(
239
model_name,
240
matches,
tests/compile/passes/distributed/test_async_tp.py
+19
-3
@@ -29,6 +29,7 @@ from vllm.distributed.parallel_state import (
29
initialize_model_parallel,
30
)
31
from vllm.platforms import current_platform
32
+from vllm.utils.network_utils import get_open_port
33
from vllm.utils.system_utils import update_environment_variables
34
from vllm.utils.torch_utils import set_random_seed
35
@@ -234,8 +235,20 @@ class TestAGCutlassScaledMMModel(_BaseScaledMMModel):
235
TestAGMMModel,
236
TestScaledMMRSModel,
237
TestAGScaledMMModel,
237
- TestCutlassScaledMMRSModel,
238
- TestAGCutlassScaledMMModel,
238
+ pytest.param(
239
+ TestCutlassScaledMMRSModel,
240
+ marks=pytest.mark.skipif(
241
+ not hasattr(torch.ops._C, "cutlass_scaled_mm"),
242
+ reason="Requires cutlass_scaled_mm",
243
+ ),
244
+ ),
245
+ pytest.param(
246
+ TestAGCutlassScaledMMModel,
247
+ marks=pytest.mark.skipif(
248
+ not hasattr(torch.ops._C, "cutlass_scaled_mm"),
249
+ reason="Requires cutlass_scaled_mm",
250
+ ),
251
+ ),
252
],
253
)
254
@pytest.mark.parametrize("batch_size", [8])
@@ -268,6 +281,7 @@ def test_async_tp_pass_replace(
281
)
282
283
num_processes = 2
284
+ master_port = str(get_open_port())
285
286
def run_torch_spawn(fn, nprocs):
287
# need to use torch.mp.spawn otherwise will have problems with
@@ -282,6 +296,7 @@ def test_async_tp_pass_replace(
296
hidden_size,
297
dtype,
298
dynamic,
299
+ master_port,
300
),
301
nprocs=nprocs,
302
)
@@ -314,6 +329,7 @@ def async_tp_pass_on_test_model(
329
hidden_size: int,
330
dtype: torch.dtype,
331
dynamic: bool,
332
+ master_port: str = "0",
333
):
334
set_random_seed(0)
335
@@ -328,7 +344,7 @@ def async_tp_pass_on_test_model(
344
"LOCAL_RANK": str(local_rank),
345
"WORLD_SIZE": str(world_size),
346
"MASTER_ADDR": "localhost",
331
- "MASTER_PORT": "12345",
347
+ "MASTER_PORT": master_port,
348
}
349
)
350
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+8
@@ -1101,6 +1101,14 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
1101
return False
1102
return bool(compile_range.end <= self.max_token_num)
1103
1104
+ @VllmInductorPass.time_and_log
1105
+ def __call__(self, graph: fx.Graph) -> None:
1106
+ self.matched_count = self.pm_pass.apply(graph)
1107
+ VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count
1108
+ logger.debug(
1109
+ "%s Replaced %s patterns", self.__class__.__name__, self.matched_count
1110
+ )
1111
+
1112
def __del__(self) -> None:
1113
if getattr(self, "disabled", True):
1114
return
vllm/compilation/passes/fusion/collective_fusion.py
+7
-7
@@ -917,13 +917,13 @@ class AsyncTPPass(VllmFusionPatternMatcherPass):
917
AllGatherScaledMMPattern(self.model_dtype, self.device).register(
918
self.pm_pass
919
)
920
-
921
- CutlassScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
922
- self.pm_pass
923
- )
924
- AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register(
925
- self.pm_pass
926
- )
920
+ if hasattr(torch.ops._C, "cutlass_scaled_mm"):
921
+ CutlassScaledMMReduceScatterPattern(
922
+ self.model_dtype, self.device
923
+ ).register(self.pm_pass)
924
+ AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register(
925
+ self.pm_pass
926
+ )
927
with suppress(ImportError):
928
import vllm.utils.flashinfer # noqa: F401
929
if hasattr(torch.ops.vllm, "bmm_fp8"):