diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 9a2116a16d..0b06aab230 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -5069,3 +5069,33 @@ steps: commands: - pip install 'modelscope<1.38' - pytest -v -s test_regression.py + +- label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep' + timeout_in_minutes: 30 + working_dir: /vllm-workspace/ + source_file_dependencies: + - vllm/compilation/ + - vllm/model_executor/layers/linear.py + - vllm/model_executor/layers/utils.py + - vllm/distributed/ + - vllm/v1/attention/backends/triton_attn.py + - tests/compile/fusions_e2e/ + - tests/compile/passes/distributed/test_async_tp.py + - tests/compile/rocm/test_async_tp.py + - tests/compile/correctness_e2e/test_async_tp.py + - vllm/model_executor/kernels/linear/ + - vllm/model_executor/models/qwen3.py + commands: + - rocm-smi + - pytest -v -s tests/compile/rocm/test_async_tp.py + - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py + - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness + dind: false + agent_pool: mi355_2 + num_gpus: 2 + mirror_hardwares: + - amdexperimental + - amdproduction + - amdgfx950nightly + - amdmi355 + optional: true diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index 219a4ca018..e7be6e3ac6 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -221,6 +221,33 @@ steps: - nvidia-smi # Run just llama3 (fp8 & bf16) for all config combinations - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3" + mirror: + amd: + label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep' + dind: false + device: mi355_2 + num_devices: 2 + working_dir: /vllm-workspace/ + timeout_in_minutes: 30 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/compilation/ + - vllm/model_executor/layers/linear.py + - vllm/model_executor/layers/utils.py + - vllm/distributed/ + - vllm/v1/attention/backends/triton_attn.py + - tests/compile/fusions_e2e/ + - tests/compile/passes/distributed/test_async_tp.py + - tests/compile/rocm/test_async_tp.py + - tests/compile/correctness_e2e/test_async_tp.py + - vllm/model_executor/kernels/linear/ + - vllm/model_executor/models/qwen3.py + commands: + - rocm-smi + - pytest -v -s tests/compile/rocm/test_async_tp.py + - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py + - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - label: ":nvidia: (B200) Fusion E2E TP2" key: fusion-e2e-tp2-b200 diff --git a/tests/compile/correctness_e2e/test_async_tp.py b/tests/compile/correctness_e2e/test_async_tp.py index 97248d153f..6b8c26cadc 100644 --- a/tests/compile/correctness_e2e/test_async_tp.py +++ b/tests/compile/correctness_e2e/test_async_tp.py @@ -2,18 +2,22 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +import logging import pytest +from tests.compile.fusions_e2e.common import FUSION_LOG_PATTERNS from tests.models.registry import HF_EXAMPLE_MODELS from tests.utils import ( compare_two_settings, create_new_process_for_each_test, + multi_gpu_test, ) from vllm.config import ( CompilationMode, ) from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.flashinfer import has_flashinfer NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4" @@ -26,6 +30,78 @@ NVFP4_HF_OVERRIDES = { } +@multi_gpu_test(num_gpus=2) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") +def test_rocm_async_tp_bf16_output_correctness(caplog_mp_spawn): + """Trained-model outputs must agree with unfused TP, including padded decode.""" + # A snapshot path pins both servers and the comparison helper's tokenizer. + model = hf_api().snapshot_download( + "Qwen/Qwen3-0.6B", revision="c1899de289a04d12100db370d81485cdf75e47ca" + ) + common_args = [ + "--dtype", + "bfloat16", + "--load-format", + "auto", + "--tensor-parallel-size", + "2", + "--distributed-executor-backend", + "mp", + "--max-model-len", + "512", + "--max-num-batched-tokens", + "512", + "--max-num-seqs", + "8", + "--kv-cache-memory-bytes", + str(64 * 1024**2), + "--gpu-memory-utilization", + "0.01", + "--attention-backend", + "TRITON_ATTN", + "--disable-custom-all-reduce", + ] + + def settings(enabled): + return [ + *common_args, + "--compilation-config", + json.dumps( + { + "mode": CompilationMode.VLLM_COMPILE, + "cudagraph_mode": "NONE", + "splitting_ops": [], + "inductor_compile_config": {"force_disable_caches": True}, + "pass_config": { + "enable_sp": enabled, + "fuse_gemm_comms": enabled, + "fuse_allreduce_rms": False, + "sp_min_token_num": 1, + }, + } + ), + ] + + env = {"VLLM_ROCM_USE_AITER": "0", "VLLM_DISABLE_COMPILE_CACHE": "1"} + with caplog_mp_spawn(logging.DEBUG) as logs: + compare_two_settings( + model, + settings(False), + settings(True), + env1=env, + env2=env, + method="generate", + force_v1_runner=True, + ) + matches = [ + int(count) for count in FUSION_LOG_PATTERNS["async_tp"].findall(logs.text) + ] + print(f"AsyncTP rewrite counts: {matches}") + assert len(matches) >= 2 and all(count > 0 for count in matches), ( + f"Expected positive AsyncTP rewrites on both TP workers, found {matches}" + ) + + @create_new_process_for_each_test() @pytest.mark.parametrize( "model_id", diff --git a/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py b/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py new file mode 100644 index 0000000000..3143143bec --- /dev/null +++ b/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bounded ROCm model coverage for actual sequence-parallel AsyncTP rewrites.""" + +import pytest + +from tests.utils import multi_gpu_test +from vllm.config import PassConfig +from vllm.platforms import current_platform + +from .common import INDUCTOR_GRAPH_PARTITION, custom_ops_combos +from .models import TRITON_ATTN, llama3_8b + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) +@pytest.mark.parametrize("custom_ops", tuple(custom_ops_combos("rms_norm"))) +@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) +def test_tp2_rocm_async_tp_bf16_fusions( + inductor_graph_partition, custom_ops, attn_backend, run_e2e_fusion_test +): + n_layers = 4 + model_kwargs = dict( + hf_overrides=dict( + num_hidden_layers=n_layers, + hidden_size=512, + intermediate_size=1024, + num_attention_heads=8, + num_key_value_heads=2, + head_dim=64, + ), + load_format="dummy", + max_model_len=1024, + kv_cache_memory_bytes=64 * 1024**2, + gpu_memory_utilization=0.01, + max_num_batched_tokens=1024, + max_num_seqs=32, + ) + run_e2e_fusion_test( + llama3_8b.model_name, + llama3_8b.matches(n_layers), + model_kwargs, + attn_backend, + dict( + use_inductor_graph_partition=inductor_graph_partition, + custom_ops=custom_ops.split(","), + pass_config=PassConfig( + enable_qk_norm_rope_fusion=True, + enable_sp=True, + fuse_gemm_comms=True, + fuse_allreduce_rms=False, + sp_min_token_num=512, + ), + ), + ["norm_rope_fusion", "sequence_parallel", "async_tp"], + tp_size=2, + ) diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py index ff0b048962..c0fd4c00a8 100644 --- a/tests/compile/passes/distributed/test_async_tp.py +++ b/tests/compile/passes/distributed/test_async_tp.py @@ -303,6 +303,7 @@ def async_tp_pass_on_test_model( dtype: torch.dtype, dynamic: bool, distributed_init_method: str, + check_numerics: bool = False, ): set_random_seed(0) @@ -363,12 +364,24 @@ def async_tp_pass_on_test_model( hidden_states = torch.randn( (batch_size * seq_len, hidden_size), dtype=dtype, requires_grad=False ) + if check_numerics: + hidden_states += local_rank * 0.125 if dynamic: torch._dynamo.mark_dynamic(hidden_states, 0) compiled_model = torch.compile(model, backend=backend) - compiled_model(hidden_states) + actual = compiled_model(hidden_states) + + if check_numerics: + # These are new checks, so no existing tolerance is relaxed. + # Allow BF16 rounding vs. eager, but keep the lifetime check exact. + expected = model(hidden_states) + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) + snapshot = actual.clone() + changed = compiled_model(-hidden_states) + torch.testing.assert_close(changed, -expected, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(actual, snapshot, rtol=0, atol=0) assert async_tp_pass.matched_count == 1 diff --git a/tests/compile/rocm/__init__.py b/tests/compile/rocm/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/tests/compile/rocm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/compile/rocm/test_async_tp.py b/tests/compile/rocm/test_async_tp.py new file mode 100644 index 0000000000..a4462265f8 --- /dev/null +++ b/tests/compile/rocm/test_async_tp.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm AsyncTP numerical coverage kept outside the shared distributed suite.""" + +import pytest +import torch + +from tests.compile.passes.distributed.test_async_tp import ( + TestAGMMModel, + TestMMRSModel, + async_tp_pass_on_test_model, +) +from tests.utils import multi_gpu_test +from vllm.distributed import ( + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, +) +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_file_store_init_method + + +class TestRocmMMRSModel(TestMMRSModel): + def forward(self, hidden_states): + view = hidden_states.reshape(-1, self.hidden_size) + mm = torch.ops.vllm.rocm_unquantized_gemm(view, self.gate_proj, None) + return tensor_model_parallel_reduce_scatter(mm, dim=0) + + def ops_in_model_before(self): + return super().ops_in_model_before() + [ + torch.ops.vllm.rocm_unquantized_gemm.default + ] + + +class TestRocmAGMMModel(TestAGMMModel): + def __init__(self, hidden_size=16, dtype=torch.bfloat16): + super().__init__(hidden_size, dtype) + self.weight = torch.nn.Parameter( + torch.empty((hidden_size * 2, hidden_size)), requires_grad=False + ) + torch.nn.init.normal_(self.weight, std=0.02) + + def forward(self, hidden_states): + view = hidden_states.reshape(-1, self.hidden_size) + gathered = tensor_model_parallel_all_gather(view, dim=0) + return torch.ops.vllm.rocm_unquantized_gemm(gathered, self.weight, None) + + def ops_in_model_before(self): + return super().ops_in_model_before() + [ + torch.ops.vllm.rocm_unquantized_gemm.default + ] + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend contract") +@pytest.mark.parametrize( + "test_model", [TestMMRSModel, TestAGMMModel, TestRocmMMRSModel, TestRocmAGMMModel] +) +@pytest.mark.parametrize("dynamic", [False, True]) +def test_rocm_async_tp_bf16_rewrite_and_correctness(test_model, dynamic): + """Real symmetric-memory AG/GEMM and GEMM/RS must replace the collectives.""" + torch.multiprocessing.spawn( + async_tp_pass_on_test_model, + args=( + 2, + test_model, + 8, + 16, + 128, + torch.bfloat16, + dynamic, + get_file_store_init_method(), + True, + ), + nprocs=2, + ) diff --git a/vllm/compilation/passes/fusion/collective_fusion.py b/vllm/compilation/passes/fusion/collective_fusion.py index 7a413665ff..f2aa6d7d2f 100644 --- a/vllm/compilation/passes/fusion/collective_fusion.py +++ b/vllm/compilation/passes/fusion/collective_fusion.py @@ -339,6 +339,8 @@ class BasePattern: class GEMMReduceScatterPattern(BasePattern): + rocm_gemm = False + def get_inputs(self) -> list[torch.Tensor]: mul = torch.empty([16, 4], device=self.device, dtype=self.dtype) mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype) @@ -346,7 +348,10 @@ class GEMMReduceScatterPattern(BasePattern): def register(self, pm_pass: PatternMatcherPass) -> None: def pattern(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor: - mm = torch.ops.aten.mm.default(mul, mm_weight) + if self.rocm_gemm: + mm = torch.ops.vllm.rocm_unquantized_gemm(mul, mm_weight, None) + else: + mm = torch.ops.aten.mm.default(mul, mm_weight) reduce_scatter = torch.ops.vllm.reduce_scatter.default( mm, dim=0, @@ -358,7 +363,7 @@ class GEMMReduceScatterPattern(BasePattern): def replacement(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor: gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter( mul, - mm_weight, + mm_weight.t() if self.rocm_gemm else mm_weight, "sum", scatter_dim=0, group_name=self.tp.device_group.group_name, @@ -371,7 +376,13 @@ class GEMMReduceScatterPattern(BasePattern): ) +class RocmGEMMReduceScatterPattern(GEMMReduceScatterPattern): + rocm_gemm = True + + class AllGatherGEMMPattern(BasePattern): + rocm_gemm = False + def get_inputs(self) -> list[torch.Tensor]: x = torch.empty([4, 4], device=self.device, dtype=self.dtype) weight = torch.empty([4, 4], device=self.device, dtype=self.dtype) @@ -390,12 +401,14 @@ class AllGatherGEMMPattern(BasePattern): group_name=self.tp.unique_name, ) + if self.rocm_gemm: + return torch.ops.vllm.rocm_unquantized_gemm(all_gather, weight, None) return torch.ops.aten.mm.default(all_gather, weight) def replacement(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul( x, - [weight], + [weight.t() if self.rocm_gemm else weight], gather_dim=0, group_name=self.tp.device_group.group_name, ) @@ -406,6 +419,10 @@ class AllGatherGEMMPattern(BasePattern): ) +class RocmAllGatherGEMMPattern(AllGatherGEMMPattern): + rocm_gemm = True + + class ScaledMMReduceScatterPattern(BasePattern): def get_inputs(self) -> list[torch.Tensor]: input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) @@ -906,6 +923,13 @@ class AsyncTPPass(VllmFusionPatternMatcherPass): GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.pm_pass) AllGatherGEMMPattern(self.model_dtype, self.device).register(self.pm_pass) + if current_platform.is_rocm() and self.model_dtype == torch.bfloat16: + RocmGEMMReduceScatterPattern(self.model_dtype, self.device).register( + self.pm_pass + ) + RocmAllGatherGEMMPattern(self.model_dtype, self.device).register( + self.pm_pass + ) # These fusions are enabled only for bfloat16 models because # `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index 8ee8f0354e..df0e58b210 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -41,6 +41,7 @@ if current_platform.is_cuda_alike() or current_platform.is_xpu(): if current_platform.is_cuda_alike(): from .fusion.act_quant_fusion import ActivationQuantFusionPass from .fusion.attn_quant_fusion import AttnQuantFusionPass + from .fusion.collective_fusion import AsyncTPPass from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass from .fusion.qk_norm_rope_kvcache_fusion import QkNormRopeKvCacheFusionPass @@ -50,7 +51,6 @@ if current_platform.is_cuda_alike(): if current_platform.is_cuda(): from .fusion.allreduce_rms_fusion import AllReduceFusionPass - from .fusion.collective_fusion import AsyncTPPass if current_platform.is_xpu(): from .fusion.act_quant_fusion import ActivationQuantFusionPass