diff --git a/tests/compile/fusions_e2e/common.py b/tests/compile/fusions_e2e/common.py index 2c6dc2b3eb..c239e93785 100644 --- a/tests/compile/fusions_e2e/common.py +++ b/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" ), diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index 3a06087472..9f34d25c46 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/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 diff --git a/tests/compile/fusions_e2e/models.py b/tests/compile/fusions_e2e/models.py index 32f1ea3506..847783d703 100644 --- a/tests/compile/fusions_e2e/models.py +++ b/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, ), diff --git a/tests/compile/fusions_e2e/test_tp2_ar_rms.py b/tests/compile/fusions_e2e/test_tp2_ar_rms.py index b5e2b2dc07..b18c41658f 100644 --- a/tests/compile/fusions_e2e/test_tp2_ar_rms.py +++ b/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, diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py index 33e050d776..7d1320bcca 100644 --- a/tests/compile/passes/distributed/test_async_tp.py +++ b/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, } ) diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index 5406f611e8..569fac667e 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/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 diff --git a/vllm/compilation/passes/fusion/collective_fusion.py b/vllm/compilation/passes/fusion/collective_fusion.py index 29d79c9b92..3658877c67 100644 --- a/vllm/compilation/passes/fusion/collective_fusion.py +++ b/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"):