@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fb7fdc49c4
Fork 0

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 7, 09:00 UTC

[ROCm] [CI] Add new fusion test cases that are relevant to vLLM IR Ops (#34307)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com>

TJian committed Mar 3, 2026 at 22:24 UTC fb7fdc49c4a0c629fd92a5e49c08ec86f5dd8ff9
10 files changed +217 -61
.buildkite/test-amd.yaml
+107 -40
@@ -610,6 +610,8 @@ steps:
610 --ignore=lora/test_qwen3moe_tp.py
611 parallelism: 4
612
613 +##### .buildkite/test_areas/pytorch.yaml #####
614 +# corresponds to .buildkite/test_areas/pytorch.yaml
615 - label: PyTorch Compilation Unit Tests # 15min
616 timeout_in_minutes: 30
617 mirror_hardwares: [amdexperimental, amdproduction]
@@ -627,6 +629,20 @@ steps:
629 # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
630 - "find compile/ -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\\\;"
631
632 +# corresponds to .buildkite/test_areas/pytorch.yaml
633 +- label: PyTorch Compilation Passes Unit Tests
634 + timeout_in_minutes: 20
635 + mirror_hardwares: [amdexperimental, amdproduction]
636 + agent_pool: mi325_1
637 + source_file_dependencies:
638 + - vllm/
639 + - tests/compile/passes
640 + commands:
641 + # TODO: clean up this comment if not needed. It is used to
642 + # keep track of the tests changes during vLLM IR Ops refactoring.
643 + # Use `find` to launch multiple instances of pytest.
644 + - "find compile/passes -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\\\;"
645 +
646 - label: PyTorch Fullgraph Smoke Test # 15min
647 timeout_in_minutes: 30
648 mirror_hardwares: [amdexperimental, amdproduction]
@@ -1211,41 +1227,6 @@ steps:
1227 - pytest -v -s tests/kernels/moe/test_flashinfer.py
1228 - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
1229
1214 -- label: Blackwell Fusion and Compile Tests # 30 min
1215 - timeout_in_minutes: 40
1216 - working_dir: "/vllm-workspace/"
1217 - gpu: b200
1218 - source_file_dependencies:
1219 - - csrc/quantization/fp4/
1220 - - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
1221 - - vllm/v1/attention/backends/flashinfer.py
1222 - - vllm/v1/worker/
1223 - - vllm/v1/cudagraph_dispatcher.py
1224 - - vllm/compilation/
1225 - # can affect pattern matching
1226 - - vllm/model_executor/layers/layernorm.py
1227 - - vllm/model_executor/layers/activation.py
1228 - - vllm/model_executor/layers/quantization/input_quant_fp8.py
1229 - - tests/compile/passes/test_fusion_attn.py
1230 - - tests/compile/passes/test_silu_mul_quant_fusion.py
1231 - - tests/compile/passes/distributed/test_fusion_all_reduce.py
1232 - - tests/compile/fullgraph/test_full_graph.py
1233 - commands:
1234 - - nvidia-smi
1235 - - pytest -v -s tests/compile/passes/test_fusion_attn.py
1236 - - pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py
1237 - # this runner has 2 GPUs available even though num_gpus=2 is not set
1238 - - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
1239 -
1240 - # # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
1241 - # # Wrap with quotes to escape yaml
1242 - # - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
1243 - # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
1244 - # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
1245 -
1246 - # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
1247 - - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
1248 -
1230 - label: Blackwell GPT-OSS Eval
1231 timeout_in_minutes: 60
1232 working_dir: "/vllm-workspace/"
@@ -1371,7 +1352,6 @@ steps:
1352 - pytest -v -s ./compile/test_wrapper.py
1353 - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
1354 - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
1374 - - pytest -v -s compile/correctness_e2e/test_sequence_parallel.py
1355 - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
1356 - pytest -v -s v1/worker/test_worker_memory_snapshot.py
1357
@@ -1601,16 +1581,16 @@ steps:
1581 commands:
1582 - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
1583 - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
1604 - - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
1584 + # TODO: this test is not supported on ROCm, there are aiter kernels for this.
1585 + # - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
1586 #- pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm
1587 # - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
1588 # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
1589 # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
1609 -
1610 - - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
1590 - pytest -v -s tests/distributed/test_context_parallel.py
1591 - HIP_VISIBLE_DEVICES=0,1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
1613 - - pytest -v -s tests/v1/distributed/test_dbo.py
1592 + # this test is not supported on ROCm
1593 + # - pytest -v -s tests/v1/distributed/test_dbo.py
1594
1595 ##### B200 test #####
1596 - label: Distributed Tests (B200) # optional
@@ -1721,6 +1701,93 @@ steps:
1701 commands:
1702 - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040
1703
1704 +##### .buildkite/test_areas/compile.yaml #####
1705 +# Slowly setting up the tests so that it is also easier for the
1706 +# CI team to review and upstream to the pipelinev2.
1707 +# The following tests are important for vLLM IR Ops refactoring,
1708 +# which affects fusion passes on ROCm. So we have to
1709 +# enable them as as soon as possible.
1710 +
1711 +## TODO: Enable the test in this group
1712 +# # corresponds to .buildkite/test_areas/compile.yaml
1713 +# - label: Fusion and Compile Unit Tests (2xMI325 GPUs)
1714 +# timeout_in_minutes: 20
1715 +# working_dir: "/vllm-workspace/"
1716 +# mirror_hardwares: [amdexperimental, amdproduction, tj]
1717 +# agent_pool: mi325_1 # changed to 1 GPU until the fusion all reduce is enabled then only revert back to 2 GPUs
1718 +# source_file_dependencies:
1719 +# - csrc/quantization/fp4/
1720 +# - vllm/model_executor/layers/quantization/
1721 +# - vllm/model_executor/layers/layernorm.py
1722 +# - vllm/model_executor/layers/activation.py
1723 +# - vllm/model_executor/layers/attention/attention.py
1724 +# - vllm/v1/attention/backends/flashinfer.py
1725 +# - vllm/compilation/ # TODO(luka) limit to vllm/compilation/passes
1726 +# - tests/compile/test_fusion_attn.py
1727 +# - tests/compile/test_silu_mul_quant_fusion.py
1728 +# - tests/compile/distributed/test_fusion_all_reduce.py
1729 +# - tests/compile/fullgraph/test_full_graph.py
1730 +# commands:
1731 +# - rocm-smi
1732 +# # we run all backend tests on ROCm
1733 +# # These two tests are covered in "PyTorch Compilation Passes Unit Tests"
1734 +# # - "pytest -v -s tests/compile/passes/test_fusion_attn.py"
1735 +# # - "pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py"
1736 +# # TODO: this test is not supported on ROCm, there are aiter kernels for this.
1737 +# # - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
1738 +# # TODO: find out more details
1739 +# # - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
1740 +
1741 +# corresponds to .buildkite/test_areas/compile.yaml
1742 +- label: Fusion E2E Quick (MI325)
1743 + timeout_in_minutes: 15
1744 + working_dir: "/vllm-workspace/"
1745 + mirror_hardwares: [amdexperimental, amdproduction]
1746 + agent_pool: mi325_1
1747 + num_devices: 1
1748 + source_file_dependencies:
1749 + - csrc/quantization/
1750 + - vllm/model_executor/
1751 + - vllm/v1/attention/
1752 + - vllm/compilation/
1753 + - tests/compile/fusions_e2e/
1754 + commands:
1755 + - rocm-smi
1756 + # Run all models and attn backends but only Inductor partition and native custom ops
1757 + - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and not +rms_norm and not +quant_fp8'"
1758 + # Different from CUDA, Qwen requires +rms_norm and +quant_fp8 as rms+quant fusion is only supported on AITER
1759 + - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and +rms_norm and +quant_fp8 and qwen3'"
1760 +
1761 +# corresponds to .buildkite/test_areas/compile.yaml
1762 +- label: Fusion E2E Config Sweep (MI325)
1763 + timeout_in_minutes: 30
1764 + working_dir: "/vllm-workspace/"
1765 + mirror_hardwares: [amdexperimental, amdproduction]
1766 + agent_pool: mi325_1
1767 + num_devices: 1
1768 + source_file_dependencies:
1769 + - csrc/quantization/
1770 + - vllm/compilation/
1771 + # can affect pattern matching
1772 + - vllm/model_executor/layers/layernorm.py
1773 + - vllm/model_executor/layers/activation.py
1774 + - vllm/model_executor/layers/attention/attention.py
1775 + - vllm/model_executor/layers/quantization/input_quant_fp8.py
1776 + - tests/compile/fusions_e2e/
1777 + commands:
1778 + - rocm-smi
1779 + # Run just llama3 (fp8) for all config combinations
1780 + - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3"
1781 +
1782 +## There are no ops on ROCm for these tests.
1783 +## The test still passes but the logs are not useful.
1784 +## fused ops just call torch.ops.symm_mem which
1785 +## exists in ROCm even though they don't work
1786 +# - label: AsyncTP Correctness Tests (2xMI325 GPUs)
1787 +# - label: Fusion E2E TP2 Quick (MI325)
1788 +# - label: Fusion E2E TP2 AsyncTP Config Sweep (MI325)
1789 +# - label: Fusion E2E TP2 (MI325)
1790 +# - label: Sequence Parallel Correctness Tests (2xMI325 GPUs)
1791
1792
1793 #####################################################################################################################################
tests/compile/fusions_e2e/common.py
+4
@@ -13,6 +13,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
13
14 class Matches(NamedTuple):
15 # simple pointwise
16 + aiter_rms_quant_fusion: int = 0
17 rms_quant_fusion: int = 0
18 act_quant_fusion: int = 0
19 norm_rope_fusion: int = 0
@@ -82,6 +83,9 @@ INDUCTOR_GRAPH_PARTITION = [
83 ]
84
85 FUSION_LOG_PATTERNS: dict[str, re.Pattern] = {
86 + "aiter_rms_quant_fusion": re.compile(
87 + r"RocmAiterRMSNormQuantFusionPass Replaced (\d+) patterns"
88 + ),
89 "rms_quant_fusion": re.compile(r"rms_quant_fusion.py:\d+] Replaced (\d+) patterns"),
90 "act_quant_fusion": re.compile(r"act_quant_fusion.py:\d+] Replaced (\d+) patterns"),
91 "norm_rope_fusion": re.compile(
tests/compile/fusions_e2e/conftest.py
+5
@@ -63,9 +63,14 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
63 compilation_config: dict,
64 matches_check: list[str],
65 use_deepgemm: bool = False,
66 + use_aiter: bool = False,
67 tp_size: int = 1,
68 ):
69 monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1" if use_deepgemm else "0")
70 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
71 + from vllm._aiter_ops import rocm_aiter_ops
72 +
73 + rocm_aiter_ops.refresh_env_variables()
74
75 # Disable, compile cache to make sure custom passes run.
76 # Otherwise, we can't verify fusion happened through the logs.
tests/compile/fusions_e2e/models.py
+20 -2
@@ -2,6 +2,8 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4
5 +from vllm._aiter_ops import is_aiter_found_and_supported
6 +from vllm.platforms import current_platform
7 from vllm.utils.flashinfer import has_flashinfer
8 from vllm.v1.attention.backends.registry import AttentionBackendEnum
9
@@ -24,6 +26,24 @@ TRITON_ATTN = pytest.param(
26 AttentionBackendCase(backend=AttentionBackendEnum.TRITON_ATTN), id="TRITON_ATTN"
27 )
28
29 +ROCM_ATTN = pytest.param(
30 + AttentionBackendCase(backend=AttentionBackendEnum.ROCM_ATTN),
31 + id="ROCM_ATTN",
32 + marks=pytest.mark.skipif(
33 + not current_platform.is_rocm(),
34 + reason="ROCm attention only for AMD",
35 + ),
36 +)
37 +
38 +ROCM_AITER_UNIFIED_ATTN = pytest.param(
39 + AttentionBackendCase(backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN),
40 + id="ROCM_AITER_UNIFIED_ATTN",
41 + marks=pytest.mark.skipif(
42 + not is_aiter_found_and_supported(),
43 + reason="ROCM_AITER_UNIFIED_ATTN only for AMD when AITER is installed",
44 + ),
45 +)
46 +
47 # Models
48 llama3_8b = ModelFusionInfo(
49 model_name="meta-llama/Llama-3.1-8B-Instruct",
@@ -49,7 +69,6 @@ llama3_8b_fp8 = ModelFusionInfo(
69 llama3_8b_fp4 = ModelFusionInfo(
70 model_name="nvidia/Llama-3.1-8B-Instruct-FP4",
71 matches=lambda n_layers: Matches(
52 - rms_quant_fusion=0,
72 act_quant_fusion=n_layers,
73 attn_quant_fusion=n_layers,
74 ar_rms_fusion=n_layers * 2 + 1,
@@ -79,7 +98,6 @@ llama4_scout_fp4 = ModelFusionInfo(
98 model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-NVFP4",
99 hf_overrides=lambda n_layers: {"text_config": {"num_hidden_layers": n_layers}},
100 matches=lambda n_layers: Matches(
82 - rms_quant_fusion=0,
101 attn_quant_fusion=n_layers,
102 ar_rms_fusion=n_layers * 2,
103 sequence_parallel=n_layers * 2,
tests/compile/fusions_e2e/test_tp1_quant.py
+39 -3
@@ -5,6 +5,7 @@ from collections.abc import Callable
5 import pytest
6
7 from vllm.config import PassConfig
8 +from vllm.platforms import current_platform
9 from vllm.utils.flashinfer import is_flashinfer_fp8_blockscale_gemm_supported
10
11 from .common import (
@@ -16,6 +17,8 @@ from .common import (
17 )
18 from .models import (
19 FLASHINFER_ATTN,
20 + ROCM_AITER_UNIFIED_ATTN,
21 + ROCM_ATTN,
22 TRITON_ATTN,
23 llama3_8b_fp4,
24 llama3_8b_fp8,
@@ -29,12 +32,33 @@ from .models import (
32 "model_name, matches_fn, model_kwargs, hf_overrides, use_deepgemm",
33 [
34 (*llama3_8b_fp8, False),
32 - (*llama4_scout_fp8, False),
35 (*qwen3_a3b_fp8, False),
34 - (*qwen3_a3b_fp8, True),
36 + pytest.param(
37 + *llama4_scout_fp8,
38 + False,
39 + marks=pytest.mark.skipif(
40 + not current_platform.is_cuda(),
41 + reason="Llama4 Scout FP8 only supported on CUDA",
42 + ),
43 + ),
44 + pytest.param(
45 + *qwen3_a3b_fp8,
46 + True,
47 + marks=pytest.mark.skipif(
48 + not current_platform.is_cuda(), reason="DeepGemm only supported on CUDA"
49 + ),
50 + ),
51 + ],
52 +)
53 +@pytest.mark.parametrize(
54 + "attn_backend",
55 + [
56 + TRITON_ATTN,
57 + FLASHINFER_ATTN,
58 + ROCM_ATTN,
59 + ROCM_AITER_UNIFIED_ATTN,
60 ],
61 )
37 -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN])
62 @pytest.mark.parametrize("n_layers", [6])
63 @pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
64 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
@@ -81,6 +105,8 @@ def test_tp1_fp8_fusions(
105 ),
106 )
107
108 + use_aiter = current_platform.is_rocm() and ("qwen" in model_name.lower())
109 +
110 matches_check = [
111 "rms_quant_fusion",
112 "act_quant_fusion",
@@ -88,6 +114,15 @@ def test_tp1_fp8_fusions(
114 "attn_quant_fusion",
115 ]
116
117 + if use_aiter:
118 + matches_check[0] = "aiter_rms_quant_fusion"
119 +
120 + matches = matches._replace(aiter_rms_quant_fusion=matches.rms_quant_fusion)
121 + # TODO: enable the `norm_rope_fusion` test,
122 + # On ROCm norm_rope_fusion is only supported without
123 + # enabling AITER.
124 + matches_check.remove("norm_rope_fusion")
125 +
126 run_e2e_fusion_test(
127 model_name,
128 matches,
@@ -96,6 +131,7 @@ def test_tp1_fp8_fusions(
131 compilation_config,
132 matches_check,
133 use_deepgemm=use_deepgemm,
134 + use_aiter=use_aiter,
135 )
136
137
tests/compile/fusions_e2e/test_tp2_ar_rms.py
+3
@@ -5,6 +5,7 @@ from collections.abc import Callable
5 import pytest
6
7 from vllm.config import PassConfig
8 +from vllm.platforms import current_platform
9
10 from ...utils import multi_gpu_test
11 from .common import (
@@ -26,6 +27,8 @@ from .models import (
27 qwen3_a3b_fp8,
28 )
29
30 +pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
31 +
32
33 @multi_gpu_test(num_gpus=2)
34 @pytest.mark.parametrize(
tests/compile/fusions_e2e/test_tp2_async_tp.py
+3
@@ -5,6 +5,7 @@ from collections.abc import Callable
5 import pytest
6
7 from vllm.config import PassConfig
8 +from vllm.platforms import current_platform
9
10 from ...utils import multi_gpu_test
11 from .common import (
@@ -23,6 +24,8 @@ from .models import (
24 qwen3_a3b,
25 )
26
27 +pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
28 +
29
30 @multi_gpu_test(num_gpus=2)
31 @pytest.mark.parametrize(
tests/compile/passes/distributed/test_sequence_parallelism.py
+2
@@ -36,6 +36,8 @@ from vllm.platforms import current_platform
36 from vllm.utils.system_utils import update_environment_variables
37 from vllm.utils.torch_utils import set_random_seed
38
39 +pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
40 +
41 FP8_DTYPE = current_platform.fp8_dtype()
42 prompts = [
43 "Hello, my name is",
tests/compile/passes/test_silu_mul_quant_fusion.py
+24 -4
@@ -182,8 +182,24 @@ TEST_KERNELS = ROCM_KERNELS if current_platform.is_rocm() else CUDA_KERNELS
182 "model_class, enable_quant_fp8_custom_op, force_kernel",
183 list(itertools.product([TestSiluMulFp8QuantModel], [True, False], TEST_KERNELS))
184 + [
185 - (TestSiluMulNvfp4QuantModel, False, None),
186 - (TestSiluMulGroupFp8QuantModel, False, None),
185 + pytest.param(
186 + TestSiluMulNvfp4QuantModel,
187 + False,
188 + None,
189 + marks=pytest.mark.skipif(
190 + not current_platform.is_cuda(), reason="CUDA only"
191 + ),
192 + ),
193 + # GroupFP8Quant fusion only works with AITER on ROCm.
194 + # and the enable_quant_fp8_custom_op must be True.
195 + pytest.param(
196 + TestSiluMulGroupFp8QuantModel,
197 + True,
198 + None,
199 + marks=pytest.mark.skipif(
200 + not current_platform.is_rocm(), reason="ROCm only"
201 + ),
202 + ),
203 ],
204 )
205 @pytest.mark.skipif(
@@ -201,6 +217,7 @@ def test_fusion_silu_and_mul_quant(
217 enable_silu_mul_custom_op: bool,
218 enable_quant_fp8_custom_op: bool,
219 force_kernel: FP8ScaledMMLinearKernel | None,
220 + monkeypatch: pytest.MonkeyPatch,
221 ):
222 if model_class is TestSiluMulNvfp4QuantModel and not is_nvfp4_supported():
223 pytest.skip("NVFP4 is not supported on this GPU.")
@@ -227,13 +244,16 @@ def test_fusion_silu_and_mul_quant(
244 ),
245 )
246
230 - with set_current_vllm_config(config):
247 + with set_current_vllm_config(config), monkeypatch.context() as m:
248 fusion_passes = [ActivationQuantFusionPass(config)]
232 - if IS_AITER_FOUND:
249 + if IS_AITER_FOUND and model_class is TestSiluMulGroupFp8QuantModel:
250 + from vllm._aiter_ops import rocm_aiter_ops
251 from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
252 RocmAiterSiluMulFp8GroupQuantFusionPass,
253 )
254
255 + m.setenv("VLLM_ROCM_USE_AITER", "1")
256 + rocm_aiter_ops.refresh_env_variables()
257 fusion_passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]
258
259 passes = [NoOpEliminationPass(config), *fusion_passes, PostCleanupPass(config)]
vllm/compilation/passes/fusion/rocm_aiter_fusion.py
+10 -12
@@ -5,7 +5,6 @@ import torch
5 import torch._inductor.pattern_matcher as pm
6 from torch import fx
7 from torch._inductor.pattern_matcher import PatternMatcherPass
8 -from torch._ops import OpOverload
8
9 import vllm.model_executor.layers.quantization.utils.fp8_utils # noqa: F401
10 from vllm._aiter_ops import rocm_aiter_ops
@@ -15,6 +14,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
14 GroupShape,
15 QuantKey,
16 ScaleDesc,
17 + kFp8Dynamic128Sym,
18 )
19 from vllm.platforms import current_platform
20
@@ -312,7 +312,9 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
312 @VllmInductorPass.time_and_log
313 def __call__(self, graph: fx.Graph) -> None:
314 self.matched_count = self.patterns.apply(graph)
315 - logger.debug("Replaced %s patterns", self.matched_count)
315 + logger.debug(
316 + "%s Replaced %s patterns", self.__class__.__name__, self.matched_count
317 + )
318
319 def uuid(self) -> str:
320 fusion_patterns = [
@@ -332,9 +334,11 @@ class AiterSiluMulFp8GroupQuantPattern(ActivationQuantPattern):
334
335 FUSED_SILU_MUL_QUANT_OP = rocm_aiter_ops.get_act_mul_fused_fp8_group_quant_op()
336
335 - def __init__(self, quant_op: OpOverload) -> None:
337 + def __init__(self) -> None:
338 self.silu_and_mul_matcher = MatcherSiluAndMul()
337 - self.quant_op = quant_op
339 + self.quant_matcher = MatcherQuantFP8(
340 + quant_key=kFp8Dynamic128Sym, match_rocm_aiter=True
341 + )
342
343 def get_inputs(self) -> list[torch.Tensor]:
344 return [
@@ -346,7 +350,7 @@ class AiterSiluMulFp8GroupQuantPattern(ActivationQuantPattern):
350 input: torch.Tensor,
351 ) -> tuple[torch.Tensor, torch.Tensor]:
352 at1 = self.silu_and_mul_matcher(input)
349 - at2 = self.quant_op(at1, 128)
353 + at2 = self.quant_matcher(at1)
354 return at2[0], at2[1]
355
356 def replacement(
@@ -370,11 +374,6 @@ class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
374 https://github.com/pytorch/pytorch/pull/139321#issuecomment-2452354980
375 """
376
373 - AITER_GROUP_FP8_QUANT_OP = rocm_aiter_ops.get_group_quant_op()
374 - TRITON_GROUP_FP8_QUANT_OP = torch.ops.vllm.triton_per_token_group_quant_fp8.default
375 -
376 - QUANT_OPS = [AITER_GROUP_FP8_QUANT_OP, TRITON_GROUP_FP8_QUANT_OP]
377 -
377 @enable_fake_mode
378 def __init__(self, config: VllmConfig) -> None:
379 super().__init__(config)
@@ -383,8 +382,7 @@ class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
382 pass_name="rocm_aiter_silu_mul_fp8_group_quant_fusion_pass"
383 )
384
386 - for quant_op in self.QUANT_OPS:
387 - AiterSiluMulFp8GroupQuantPattern(quant_op).register(self.patterns)
385 + AiterSiluMulFp8GroupQuantPattern().register(self.patterns)
386
387 self.dump_patterns(config, self.patterns)
388