iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Compile] Support BF16 AsyncTP fusion (#58098)

Browse Source 
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
Andreas Karatzas
2026-09-22 20:36:38 -05:00
committed by GitHub
parent f9dce295c9
commit 0eb42dbbfc
9 changed files with 310 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test-amd.yaml
+30
View File
@@ -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
.buildkite/test_areas/compile.yaml
+27
View File
@@ -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
tests/compile/correctness_e2e/test_async_tp.py
+76
View File
@@ -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",
tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py
+58
View File
@@ -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,
)
tests/compile/passes/distributed/test_async_tp.py
+14 -1
View File
@@ -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
tests/compile/rocm/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/compile/rocm/test_async_tp.py
+75
View File
@@ -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,
)
vllm/compilation/passes/fusion/collective_fusion.py
+27 -3
View File
@@ -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
vllm/compilation/passes/pass_manager.py
+1 -1
View File
@@ -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