[AMD][CI][BugFix] Fix Distributed Compile Unit Tests (2xH100-2xMI300) group (#43120)

Signed-off-by: Randall Smith <[email protected]>

tests/compile/fusions_e2e/common.py

@@ -20,6 +20,7 @@ class Matches(NamedTuple):
     attn_quant_fusion: int = 0
     # distributed
     ar_rms_fusion: int = 0
+    aiter_ar_rms_fusion: int = 0
     sequence_parallel: int = 0
     async_tp: int = 0
 
@@ -97,6 +98,9 @@ FUSION_LOG_PATTERNS: dict[str, re.Pattern] = {
     "ar_rms_fusion": re.compile(
         r"allreduce_rms_fusion.py:\d+] Replaced (\d+) patterns"
     ),
+    "aiter_ar_rms_fusion": re.compile(
+        r"RocmAiterAllReduceFusionPass Replaced (\d+) patterns"
+    ),
     "sequence_parallel": re.compile(
         r"sequence_parallelism.py:\d+] Replaced (\d+) patterns"
     ),

tests/compile/fusions_e2e/conftest.py

@@ -97,6 +97,11 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
                 f"attention backend '{attn_backend.backend.name}'"
             )
 
+        if backend_name == "rocm_attn" and model_name == "openai/gpt-oss-20b":
+            pytest.skip(
+                "ROCM_ATTN does not support attention sinks (required by gpt-oss-20b)"
+            )
+
         if attn_backend.backend.name == "FLASHINFER":
             from vllm.utils.flashinfer import supports_trtllm_attention
 

tests/compile/fusions_e2e/models.py

@@ -75,6 +75,7 @@ llama3_8b = ModelFusionInfo(
     model_name="meta-llama/Llama-3.1-8B-Instruct",
     matches=lambda n_layers: Matches(
         ar_rms_fusion=n_layers * 2 + 1,
+        aiter_ar_rms_fusion=n_layers * 2,
         sequence_parallel=n_layers * 2 + 1,
         async_tp=n_layers * 4,
     ),
@@ -136,6 +137,7 @@ qwen3_a3b = ModelFusionInfo(
     matches=lambda n_layers: Matches(
         norm_rope_fusion=n_layers,
         ar_rms_fusion=n_layers * 2 + 1,
+        aiter_ar_rms_fusion=n_layers * 2,
         sequence_parallel=n_layers * 2 + 1,
         async_tp=n_layers * 2,
     ),
@@ -211,6 +213,7 @@ gpt_oss_20b = ModelFusionInfo(
     model_name="openai/gpt-oss-20b",
     matches=lambda n_layers: Matches(
         ar_rms_fusion=n_layers * 2 + 1,
+        aiter_ar_rms_fusion=n_layers + 1,
         sequence_parallel=n_layers * 2 + 1,
         async_tp=n_layers * 2,
     ),

tests/compile/fusions_e2e/test_tp2_ar_rms.py

@@ -84,6 +84,7 @@ def test_tp2_ar_rms_fp8_fusions(
     model_kwargs["load_format"] = "dummy"
     model_kwargs["max_model_len"] = 1024
     model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
+    model_kwargs["disable_custom_all_reduce"] = False
 
     compilation_config = dict(
         use_inductor_graph_partition=inductor_graph_partition,
@@ -149,6 +150,7 @@ def test_tp2_ar_rms_fp4_fusions(
     model_kwargs["load_format"] = "dummy"
     model_kwargs["max_model_len"] = 1024
     model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
+    model_kwargs["disable_custom_all_reduce"] = False
 
     compilation_config = dict(
         use_inductor_graph_partition=inductor_graph_partition,
@@ -213,6 +215,7 @@ def test_tp2_ar_rms_fusions(
     model_kwargs["load_format"] = "dummy"
     model_kwargs["max_model_len"] = 1024
     model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
+    model_kwargs["disable_custom_all_reduce"] = False
 
     compilation_config = dict(
         use_inductor_graph_partition=inductor_graph_partition,
@@ -225,9 +228,13 @@ def test_tp2_ar_rms_fusions(
 
     matches_check = [
         "norm_rope_fusion",
-        "ar_rms_fusion",
     ]
 
+    if current_platform.is_rocm():
+        matches_check.append("aiter_ar_rms_fusion")
+    else:
+        matches_check.append("ar_rms_fusion")
+
     run_e2e_fusion_test(
         model_name,
         matches,

tests/compile/passes/distributed/test_async_tp.py

@@ -29,6 +29,7 @@ from vllm.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
@@ -234,8 +235,20 @@ class TestAGCutlassScaledMMModel(_BaseScaledMMModel):
         TestAGMMModel,
         TestScaledMMRSModel,
         TestAGScaledMMModel,
-        TestCutlassScaledMMRSModel,
-        TestAGCutlassScaledMMModel,
+        pytest.param(
+            TestCutlassScaledMMRSModel,
+            marks=pytest.mark.skipif(
+                not hasattr(torch.ops._C, "cutlass_scaled_mm"),
+                reason="Requires cutlass_scaled_mm",
+            ),
+        ),
+        pytest.param(
+            TestAGCutlassScaledMMModel,
+            marks=pytest.mark.skipif(
+                not hasattr(torch.ops._C, "cutlass_scaled_mm"),
+                reason="Requires cutlass_scaled_mm",
+            ),
+        ),
     ],
 )
 @pytest.mark.parametrize("batch_size", [8])
@@ -268,6 +281,7 @@ def test_async_tp_pass_replace(
         )
 
     num_processes = 2
+    master_port = str(get_open_port())
 
     def run_torch_spawn(fn, nprocs):
         # need to use torch.mp.spawn otherwise will have problems with
@@ -282,6 +296,7 @@ def test_async_tp_pass_replace(
                 hidden_size,
                 dtype,
                 dynamic,
+                master_port,
             ),
             nprocs=nprocs,
         )
@@ -314,6 +329,7 @@ def async_tp_pass_on_test_model(
     hidden_size: int,
     dtype: torch.dtype,
     dynamic: bool,
+    master_port: str = "0",
 ):
     set_random_seed(0)
 
@@ -328,7 +344,7 @@ def async_tp_pass_on_test_model(
             "LOCAL_RANK": str(local_rank),
             "WORLD_SIZE": str(world_size),
             "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "12345",
+            "MASTER_PORT": master_port,
         }
     )
 

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

@@ -1101,6 +1101,14 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
             return False
         return bool(compile_range.end <= self.max_token_num)
 
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None:
+        self.matched_count = self.pm_pass.apply(graph)
+        VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count
+        logger.debug(
+            "%s Replaced %s patterns", self.__class__.__name__, self.matched_count
+        )
+
     def __del__(self) -> None:
         if getattr(self, "disabled", True):
             return

vllm/compilation/passes/fusion/collective_fusion.py

@@ -917,13 +917,13 @@ class AsyncTPPass(VllmFusionPatternMatcherPass):
             AllGatherScaledMMPattern(self.model_dtype, self.device).register(
                 self.pm_pass
             )
-
-            CutlassScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
-                self.pm_pass
-            )
-            AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register(
-                self.pm_pass
-            )
+            if hasattr(torch.ops._C, "cutlass_scaled_mm"):
+                CutlassScaledMMReduceScatterPattern(
+                    self.model_dtype, self.device
+                ).register(self.pm_pass)
+                AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register(
+                    self.pm_pass
+                )
             with suppress(ImportError):
                 import vllm.utils.flashinfer  # noqa: F401
             if hasattr(torch.ops.vllm, "bmm_fp8"):