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 8, 19:00 UTC
[torch.compile] Reorganize vllm/compilation and tests/compile (0/N for vLLM IR) (#33731)
Signed-off-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: ProExpertProg <luka.govedic@gmail.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Luka Govedič committed
Feb 6, 2026 at 07:19 UTC
ac32e66cf95c40502ad7b00b2e80dfb0315bfee4
47 files changed
+717
-651
.buildkite/test-amd.yaml
+19
-19
@@ -551,7 +551,7 @@ steps:
551
- label: LoRA Test %N # 20min each
552
timeout_in_minutes: 30
553
mirror_hardwares: [amdexperimental]
554
- agent_pool: mi325_1
554
+ agent_pool: mi325_8
555
# grade: Blocking
556
source_file_dependencies:
557
- vllm/lora
@@ -647,7 +647,7 @@ steps:
647
- label: Kernels Attention Test %N # 23min
648
timeout_in_minutes: 35
649
mirror_hardwares: [amdexperimental, amdproduction]
650
- agent_pool: mi325_1
650
+ agent_pool: mi325_8
651
# grade: Blocking
652
source_file_dependencies:
653
- csrc/attention/
@@ -662,7 +662,7 @@ steps:
662
- label: Kernels Quantization Test %N # 64min
663
timeout_in_minutes: 90
664
mirror_hardwares: [amdexperimental]
665
- agent_pool: mi325_1
665
+ agent_pool: mi325_8
666
# grade: Blocking
667
source_file_dependencies:
668
- csrc/quantization/
@@ -675,7 +675,7 @@ steps:
675
- label: Kernels MoE Test %N # 40min
676
timeout_in_minutes: 60
677
mirror_hardwares: [amdexperimental, amdproduction]
678
- agent_pool: mi325_1
678
+ agent_pool: mi325_8
679
# grade: Blocking
680
source_file_dependencies:
681
- csrc/quantization/cutlass_w8a8/moe/
@@ -753,7 +753,7 @@ steps:
753
- label: Benchmarks # 11min
754
timeout_in_minutes: 20
755
mirror_hardwares: [amdexperimental, amdproduction]
756
- agent_pool: mi325_1
756
+ agent_pool: mi325_8
757
# grade: Blocking
758
working_dir: "/vllm-workspace/.buildkite"
759
source_file_dependencies:
@@ -764,7 +764,7 @@ steps:
764
- label: Benchmarks CLI Test # 7min
765
timeout_in_minutes: 20
766
mirror_hardwares: [amdexperimental, amdproduction]
767
- agent_pool: mi325_1
767
+ agent_pool: mi325_8
768
# grade: Blocking
769
source_file_dependencies:
770
- vllm/
@@ -838,7 +838,7 @@ steps:
838
- label: Basic Models Tests (Extra Initialization) %N
839
timeout_in_minutes: 45
840
mirror_hardwares: [amdexperimental, amdproduction]
841
- agent_pool: mi325_1
841
+ agent_pool: mi325_8
842
# grade: Blocking
843
torch_nightly: true
844
source_file_dependencies:
@@ -900,7 +900,7 @@ steps:
900
- label: Language Models Tests (Extra Standard) %N
901
timeout_in_minutes: 45
902
mirror_hardwares: [amdexperimental]
903
- agent_pool: mi325_1
903
+ agent_pool: mi325_8
904
# grade: Blocking
905
torch_nightly: true
906
source_file_dependencies:
@@ -921,7 +921,7 @@ steps:
921
- label: Language Models Tests (Hybrid) %N
922
timeout_in_minutes: 75
923
mirror_hardwares: [amdexperimental]
924
- agent_pool: mi325_1
924
+ agent_pool: mi325_8
925
# grade: Blocking
926
torch_nightly: true
927
source_file_dependencies:
@@ -1190,16 +1190,16 @@ steps:
1190
- vllm/model_executor/layers/layernorm.py
1191
- vllm/model_executor/layers/activation.py
1192
- vllm/model_executor/layers/quantization/input_quant_fp8.py
1193
- - tests/compile/test_fusion_attn.py
1194
- - tests/compile/test_silu_mul_quant_fusion.py
1195
- - tests/compile/distributed/test_fusion_all_reduce.py
1193
+ - tests/compile/passes/test_fusion_attn.py
1194
+ - tests/compile/passes/test_silu_mul_quant_fusion.py
1195
+ - tests/compile/passes/distributed/test_fusion_all_reduce.py
1196
- tests/compile/fullgraph/test_full_graph.py
1197
commands:
1198
- nvidia-smi
1199
- - pytest -v -s tests/compile/test_fusion_attn.py
1200
- - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
1199
+ - pytest -v -s tests/compile/passes/test_fusion_attn.py
1200
+ - pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py
1201
# this runner has 2 GPUs available even though num_gpus=2 is not set
1202
- - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
1202
+ - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
1203
1204
# # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
1205
# # Wrap with quotes to escape yaml
@@ -1556,15 +1556,15 @@ steps:
1556
working_dir: "/vllm-workspace/"
1557
num_gpus: 2
1558
commands:
1559
- - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py
1560
- - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
1561
- - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
1559
+ - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
1560
+ - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
1561
+ - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
1562
#- pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm
1563
# - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
1564
# Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
1565
# in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
1566
1567
- - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
1567
+ - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
1568
- pytest -v -s tests/distributed/test_context_parallel.py
1569
- 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
1570
- pytest -v -s tests/v1/distributed/test_dbo.py
.buildkite/test-pipeline.yaml
+6
-5
@@ -519,6 +519,7 @@ steps:
519
# However, find does not normally propagate error codes, so we combine it with xargs
520
# (using -0 for proper path handling)
521
- "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
522
+ - pytest -s -v compile/passes --ignore compile/passes/distributed
523
524
- label: PyTorch Fullgraph Smoke Test # 15min
525
timeout_in_minutes: 30
@@ -1080,14 +1081,14 @@ steps:
1081
- vllm/model_executor/layers/quantization/input_quant_fp8.py
1082
- tests/compile/test_fusion_attn.py
1083
- tests/compile/test_silu_mul_quant_fusion.py
1083
- - tests/compile/distributed/test_fusion_all_reduce.py
1084
+ - tests/compile/passes/distributed/test_fusion_all_reduce.py
1085
- tests/compile/fullgraph/test_full_graph.py
1086
commands:
1087
- nvidia-smi
1088
- pytest -v -s tests/compile/test_fusion_attn.py
1089
- pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
1090
# this runner has 2 GPUs available even though num_gpus=2 is not set
1090
- - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
1091
+ - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
1092
# # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
1093
# # Wrap with quotes to escape yaml
1094
# - "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'"
@@ -1421,8 +1422,8 @@ steps:
1422
commands:
1423
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
1424
# Run sequence parallel tests
1424
- - pytest -v -s tests/distributed/test_sequence_parallel.py
1425
- - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
1425
+ - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
1426
+ - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
1427
1428
- label: Distributed Tests (H100) # optional
1429
gpu: h100
@@ -1430,7 +1431,7 @@ steps:
1431
working_dir: "/vllm-workspace/"
1432
num_gpus: 2
1433
commands:
1433
- - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py
1434
+ - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
1435
- pytest -v -s tests/distributed/test_context_parallel.py
1436
- VLLM_USE_DEEP_GEMM=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=deepep_high_throughput
1437
- pytest -v -s tests/v1/distributed/test_dbo.py
.buildkite/test_areas/compile.yaml
+24
-18
@@ -2,7 +2,7 @@ group: Compile
2
depends_on:
3
- image-build
4
steps:
5
-- label: Sequence Parallel Tests (2 GPUs)
5
+- label: Sequence Parallel Correctness Tests (2 GPUs)
6
timeout_in_minutes: 50
7
working_dir: "/vllm-workspace/"
8
num_devices: 2
@@ -11,12 +11,12 @@ steps:
11
- vllm/compilation/
12
- vllm/v1/worker/
13
- vllm/v1/cudagraph_dispatcher.py
14
- - tests/distributed/test_sequence_parallel.py
14
+ - tests/compile/correctness_e2e/test_sequence_parallel.py
15
commands:
16
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
17
- - pytest -v -s tests/distributed/test_sequence_parallel.py
17
+ - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
18
19
-- label: Sequence Parallel Tests (2xH100)
19
+- label: Sequence Parallel Correctness Tests (2xH100)
20
timeout_in_minutes: 50
21
working_dir: "/vllm-workspace/"
22
device: h100
@@ -24,24 +24,30 @@ steps:
24
num_devices: 2
25
commands:
26
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
27
- - pytest -v -s tests/distributed/test_sequence_parallel.py
27
+ - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
28
+
29
+- label: AsyncTP Correctness Tests (2xH100)
30
+ timeout_in_minutes: 50
31
+ working_dir: "/vllm-workspace/"
32
+ device: h100
33
+ optional: true
34
+ num_devices: 2
35
+ commands:
36
+ - export VLLM_TEST_CLEAN_GPU_MEMORY=1
37
+ - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
38
39
- label: Distributed Compile Unit Tests (2xH100)
30
- timeout_in_minutes: 40
40
+ timeout_in_minutes: 20
41
working_dir: "/vllm-workspace/"
42
device: h100
43
num_devices: 2
44
source_file_dependencies:
45
- vllm/compilation/
46
- vllm/model_executor/layers
37
- - tests/compile/distributed/test_fusion_all_reduce.py
38
- - tests/compile/distributed/test_sequence_parallelism.py
39
- - tests/compile/distributed/test_async_tp.py
47
+ - tests/compile/passes/distributed/
48
commands:
49
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
42
- - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
43
- - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
44
- - pytest -v -s tests/compile/distributed/test_async_tp.py
50
+ - pytest -s -v tests/compile/passes/distributed
51
52
- label: Fusion and Compile Unit Tests (B200)
53
timeout_in_minutes: 20
@@ -55,17 +61,17 @@ steps:
61
- vllm/model_executor/layers/attention/attention.py
62
- vllm/v1/attention/backends/flashinfer.py
63
- vllm/compilation/ # TODO(luka) limit to vllm/compilation/passes
58
- - tests/compile/test_fusion_attn.py
59
- - tests/compile/test_silu_mul_quant_fusion.py
60
- - tests/compile/distributed/test_fusion_all_reduce.py
64
+ - tests/compile/passes/test_fusion_attn.py
65
+ - tests/compile/passes/test_silu_mul_quant_fusion.py
66
+ - tests/compile/passes/distributed/test_fusion_all_reduce.py
67
- tests/compile/fullgraph/test_full_graph.py
68
commands:
69
# b200 runners are limited, so we limit the tests to the minimum set only supported on Blackwell
70
- nvidia-smi
65
- - pytest -v -s tests/compile/test_fusion_attn.py -k FLASHINFER
66
- - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
71
+ - pytest -v -s tests/compile/passes/test_fusion_attn.py -k FLASHINFER
72
+ - pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py
73
# this runner has 2 GPUs available even though num_devices=2 is not set
68
- - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
74
+ - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
75
# test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
76
# TODO(luka) move to H100 once pass tests run on H100
77
- pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
.buildkite/test_areas/pytorch.yaml
+9
-1
@@ -3,7 +3,7 @@ depends_on:
3
- image-build
4
steps:
5
- label: PyTorch Compilation Unit Tests
6
- timeout_in_minutes: 30
6
+ timeout_in_minutes: 10
7
source_file_dependencies:
8
- vllm/
9
- tests/compile
@@ -17,6 +17,14 @@ steps:
17
# (using -0 for proper path handling)
18
- "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
19
20
+- label: PyTorch Compilation Passes Unit Tests
21
+ timeout_in_minutes: 20
22
+ source_file_dependencies:
23
+ - vllm/
24
+ - tests/compile/passes
25
+ commands:
26
+ - pytest -s -v compile/passes --ignore compile/passes/distributed
27
+
28
- label: PyTorch Fullgraph Smoke Test
29
timeout_in_minutes: 35
30
source_file_dependencies:
tests/compile/backend.py
+4
-4
@@ -11,10 +11,10 @@ from torch import fx
11
from torch._ops import OpOverload
12
from torch.fx._utils import lazy_format_graph_code
13
14
-from vllm.compilation.fx_utils import find_op_nodes
15
-from vllm.compilation.inductor_pass import InductorPass
16
-from vllm.compilation.pass_manager import with_pattern_match_debug
17
-from vllm.compilation.vllm_inductor_pass import VllmInductorPass
14
+from vllm.compilation.passes.fx_utils import find_op_nodes
15
+from vllm.compilation.passes.inductor_pass import InductorPass
16
+from vllm.compilation.passes.pass_manager import with_pattern_match_debug
17
+from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
18
from vllm.config import VllmConfig, get_current_vllm_config
19
from vllm.logger import init_logger
20
tests/compile/correctness_e2e/__init__.py
renamed
tests/compile/correctness_e2e/test_async_tp.py
new
+79
@@ -0,0 +1,79 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import json
5
+
6
+import pytest
7
+
8
+from tests.models.registry import HF_EXAMPLE_MODELS
9
+from tests.utils import (
10
+ compare_two_settings,
11
+ create_new_process_for_each_test,
12
+)
13
+from vllm.config import (
14
+ CompilationMode,
15
+)
16
+
17
+
18
+@create_new_process_for_each_test()
19
+@pytest.mark.parametrize(
20
+ "model_id",
21
+ ["meta-llama/Llama-3.2-1B-Instruct", "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8"],
22
+)
23
+@pytest.mark.parametrize("tp_size", [2])
24
+@pytest.mark.parametrize("async_tp_enabled", [True])
25
+@pytest.mark.parametrize("distributed_backend", ["mp"])
26
+@pytest.mark.parametrize("eager_mode", [False, True])
27
+def test_async_tp_pass_correctness(
28
+ model_id: str,
29
+ tp_size: int,
30
+ async_tp_enabled: bool,
31
+ distributed_backend: str,
32
+ eager_mode: bool,
33
+ num_gpus_available: int,
34
+):
35
+ model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
36
+ model_info.check_transformers_version(on_fail="skip")
37
+ model_info.check_available_online(on_fail="skip")
38
+
39
+ pp_size = 1
40
+ if num_gpus_available < tp_size:
41
+ pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs")
42
+
43
+ common_args = [
44
+ "--dtype",
45
+ "bfloat16",
46
+ "--max-model-len",
47
+ "2048",
48
+ "--max-num-seqs",
49
+ "8",
50
+ ]
51
+ if eager_mode:
52
+ common_args.append("--enforce-eager")
53
+
54
+ compilation_config = {
55
+ "mode": CompilationMode.VLLM_COMPILE,
56
+ "compile_sizes": [2, 4, 8],
57
+ "splitting_ops": [],
58
+ "pass_config": {"fuse_gemm_comms": async_tp_enabled},
59
+ }
60
+
61
+ async_tp_args = [
62
+ *common_args,
63
+ "--tensor-parallel-size",
64
+ str(tp_size),
65
+ "--distributed-executor-backend",
66
+ distributed_backend,
67
+ "--compilation_config",
68
+ json.dumps(compilation_config),
69
+ ]
70
+
71
+ tp_args = [
72
+ *common_args,
73
+ "--tensor-parallel-size",
74
+ str(tp_size),
75
+ "--distributed-executor-backend",
76
+ "mp",
77
+ ]
78
+
79
+ compare_two_settings(model_id, async_tp_args, tp_args, method="generate")
tests/compile/correctness_e2e/test_sequence_parallel.py
renamed
+2
-2
@@ -21,8 +21,8 @@ from vllm.logger import init_logger
21
from vllm.platforms import current_platform
22
from vllm.utils.torch_utils import is_torch_equal_or_newer
23
24
-from ..models.registry import HF_EXAMPLE_MODELS
25
-from ..utils import compare_two_settings, create_new_process_for_each_test
24
+from ...models.registry import HF_EXAMPLE_MODELS
25
+from ...utils import compare_two_settings, create_new_process_for_each_test
26
27
logger = init_logger("test_sequence_parallel")
28
tests/compile/fusions_e2e/common.py
+6
-8
@@ -82,19 +82,17 @@ INDUCTOR_GRAPH_PARTITION = [
82
]
83
84
FUSION_LOG_PATTERNS: dict[str, re.Pattern] = {
85
- "rms_quant_fusion": re.compile(
86
- r"\[(?:compilation/)?fusion.py:\d+] Replaced (\d+) patterns"
87
- ),
88
- "act_quant_fusion": re.compile(
89
- r"activation_quant_fusion.py:\d+] Replaced (\d+) patterns"
90
- ),
85
+ "rms_quant_fusion": re.compile(r"rms_quant_fusion.py:\d+] Replaced (\d+) patterns"),
86
+ "act_quant_fusion": re.compile(r"act_quant_fusion.py:\d+] Replaced (\d+) patterns"),
87
"norm_rope_fusion": re.compile(
88
r"qk_norm_rope_fusion.py:\d+] Fused QK Norm\+RoPE on (\d+) sites"
89
),
90
"attn_quant_fusion": re.compile(
95
- r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes"
91
+ r"attn_quant_fusion.py:\d+] Fused quant onto (\d+) attention nodes"
92
+ ),
93
+ "ar_rms_fusion": re.compile(
94
+ r"allreduce_rms_fusion.py:\d+] Replaced (\d+) patterns"
95
),
97
- "ar_rms_fusion": re.compile(r"collective_fusion.py:\d+] Replaced (\d+) patterns"),
96
"sequence_parallel": re.compile(
97
r"sequence_parallelism.py:\d+] Replaced (\d+) patterns"
98
),
tests/compile/passes/__init__.py
tests/compile/passes/distributed/__init__.py
tests/compile/passes/distributed/test_async_tp.py
renamed
+5
-75
@@ -1,16 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import json
4
5
import pytest
6
import torch
7
8
import vllm.envs as envs
10
-from vllm.compilation.collective_fusion import AsyncTPPass
9
+from tests.compile.backend import TestBackend
10
+from tests.utils import (
11
+ multi_gpu_test,
12
+)
13
+from vllm.compilation.passes.fusion.collective_fusion import AsyncTPPass
14
from vllm.config import (
15
CompilationConfig,
13
- CompilationMode,
16
DeviceConfig,
17
ModelConfig,
18
PassConfig,
@@ -29,14 +31,6 @@ from vllm.platforms import current_platform
31
from vllm.utils.system_utils import update_environment_variables
32
from vllm.utils.torch_utils import set_random_seed
33
32
-from ...models.registry import HF_EXAMPLE_MODELS
33
-from ...utils import (
34
- compare_two_settings,
35
- create_new_process_for_each_test,
36
- multi_gpu_test,
37
-)
38
-from ..backend import TestBackend
39
-
34
FP8_DTYPE = current_platform.fp8_dtype()
35
36
prompts = [
@@ -377,67 +371,3 @@ def async_tp_pass_on_test_model(
371
# In post-nodes, fused_matmul_reduce_scatter or \
372
# fused_all_gather_matmul should exist
373
backend.check_after_ops(model.ops_in_model_after())
380
-
381
-
382
-@create_new_process_for_each_test()
383
-@pytest.mark.parametrize(
384
- "model_id",
385
- ["meta-llama/Llama-3.2-1B-Instruct", "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8"],
386
-)
387
-@pytest.mark.parametrize("tp_size", [2])
388
-@pytest.mark.parametrize("async_tp_enabled", [True])
389
-@pytest.mark.parametrize("distributed_backend", ["mp"])
390
-@pytest.mark.parametrize("eager_mode", [False, True])
391
-def test_async_tp_pass_correctness(
392
- model_id: str,
393
- tp_size: int,
394
- async_tp_enabled: bool,
395
- distributed_backend: str,
396
- eager_mode: bool,
397
- num_gpus_available: int,
398
-):
399
- model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
400
- model_info.check_transformers_version(on_fail="skip")
401
- model_info.check_available_online(on_fail="skip")
402
-
403
- pp_size = 1
404
- if num_gpus_available < tp_size:
405
- pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs")
406
-
407
- common_args = [
408
- "--dtype",
409
- "bfloat16",
410
- "--max-model-len",
411
- "2048",
412
- "--max-num-seqs",
413
- "8",
414
- ]
415
- if eager_mode:
416
- common_args.append("--enforce-eager")
417
-
418
- compilation_config = {
419
- "mode": CompilationMode.VLLM_COMPILE,
420
- "compile_sizes": [2, 4, 8],
421
- "splitting_ops": [],
422
- "pass_config": {"fuse_gemm_comms": async_tp_enabled},
423
- }
424
-
425
- async_tp_args = [
426
- *common_args,
427
- "--tensor-parallel-size",
428
- str(tp_size),
429
- "--distributed-executor-backend",
430
- distributed_backend,
431
- "--compilation_config",
432
- json.dumps(compilation_config),
433
- ]
434
-
435
- tp_args = [
436
- *common_args,
437
- "--tensor-parallel-size",
438
- str(tp_size),
439
- "--distributed-executor-backend",
440
- "mp",
441
- ]
442
-
443
- compare_two_settings(model_id, async_tp_args, tp_args, method="generate")
tests/compile/passes/distributed/test_fusion_all_reduce.py
renamed
+8
-7
@@ -6,11 +6,15 @@ import pytest
6
import torch
7
8
import vllm.envs as envs
9
+from tests.compile.backend import TestBackend
10
+from tests.utils import TestFP8Layer, has_module_attribute, multi_gpu_test
11
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
10
-from vllm.compilation.collective_fusion import AllReduceFusionPass
11
-from vllm.compilation.fix_functionalization import FixFunctionalizationPass
12
-from vllm.compilation.noop_elimination import NoOpEliminationPass
13
-from vllm.compilation.post_cleanup import PostCleanupPass
12
+from vllm.compilation.passes.fusion.allreduce_rms_fusion import AllReduceFusionPass
13
+from vllm.compilation.passes.utility.fix_functionalization import (
14
+ FixFunctionalizationPass,
15
+)
16
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
17
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
18
from vllm.config import (
19
CompilationConfig,
20
CompilationMode,
@@ -33,9 +37,6 @@ from vllm.platforms import current_platform
37
from vllm.utils.system_utils import update_environment_variables
38
from vllm.utils.torch_utils import set_random_seed
39
36
-from ...utils import TestFP8Layer, has_module_attribute, multi_gpu_test
37
-from ..backend import TestBackend
38
-
40
41
class TestAllReduceRMSNormModel(torch.nn.Module):
42
def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
tests/compile/passes/distributed/test_sequence_parallelism.py
renamed
+8
-9
@@ -5,12 +5,14 @@ import pytest
5
import torch
6
7
import vllm.envs as envs
8
-from vllm.compilation.fusion import RMSNormQuantFusionPass
9
-from vllm.compilation.fx_utils import find_auto_fn
10
-from vllm.compilation.noop_elimination import NoOpEliminationPass
11
-from vllm.compilation.post_cleanup import PostCleanupPass
12
-from vllm.compilation.sequence_parallelism import SequenceParallelismPass
13
-from vllm.compilation.vllm_inductor_pass import VllmInductorPass
8
+from tests.compile.backend import TestBackend
9
+from tests.utils import TestFP8Layer, multi_gpu_test
10
+from vllm.compilation.passes.fusion.rms_quant_fusion import RMSNormQuantFusionPass
11
+from vllm.compilation.passes.fusion.sequence_parallelism import SequenceParallelismPass
12
+from vllm.compilation.passes.fx_utils import find_auto_fn
13
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
14
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
15
+from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
16
from vllm.config import (
17
CompilationConfig,
18
CUDAGraphMode,
@@ -34,9 +36,6 @@ 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
37
-from ...utils import TestFP8Layer, multi_gpu_test
38
-from ..backend import TestBackend
39
-
39
FP8_DTYPE = current_platform.fp8_dtype()
40
prompts = [
41
"Hello, my name is",
tests/compile/passes/test_functionalization.py
renamed
+12
-9
@@ -5,12 +5,18 @@ import pytest
5
import torch
6
7
import vllm.envs as envs
8
-from vllm.compilation.activation_quant_fusion import ActivationQuantFusionPass
9
-from vllm.compilation.fix_functionalization import FixFunctionalizationPass
10
-from vllm.compilation.fusion import RMSNormQuantFusionPass
11
-from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func
12
-from vllm.compilation.noop_elimination import NoOpEliminationPass
13
-from vllm.compilation.post_cleanup import PostCleanupPass
8
+from tests.compile.backend import TestBackend
9
+from tests.utils import TestFP8Layer
10
+from vllm.compilation.passes.fusion.act_quant_fusion import (
11
+ ActivationQuantFusionPass,
12
+)
13
+from vllm.compilation.passes.fusion.rms_quant_fusion import RMSNormQuantFusionPass
14
+from vllm.compilation.passes.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func
15
+from vllm.compilation.passes.utility.fix_functionalization import (
16
+ FixFunctionalizationPass,
17
+)
18
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
19
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
20
from vllm.config import (
21
CompilationConfig,
22
ModelConfig,
@@ -26,9 +32,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
32
from vllm.model_executor.layers.rotary_embedding import get_rope
33
from vllm.platforms import current_platform
34
29
-from ..utils import TestFP8Layer
30
-from .backend import TestBackend
31
-
35
TEST_FP8 = current_platform.supports_fp8()
36
FP8_DTYPE = current_platform.fp8_dtype()
37
tests/compile/passes/test_fuse_act_padding.py
renamed
+4
-5
@@ -6,9 +6,10 @@ import pytest
6
import torch
7
8
import vllm.config
9
+from tests.compile.backend import TestBackend
10
from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
10
-from vllm.compilation.noop_elimination import NoOpEliminationPass
11
-from vllm.compilation.post_cleanup import PostCleanupPass
11
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
12
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
13
from vllm.config import (
14
CompilationConfig,
15
CompilationMode,
@@ -19,8 +20,6 @@ from vllm.config import (
20
from vllm.model_executor.layers.layernorm import RMSNorm
21
from vllm.model_executor.layers.utils import rocm_unquantized_gemm
22
22
-from .backend import TestBackend
23
-
23
24
class TestModel(torch.nn.Module):
25
def __init__(
@@ -95,7 +94,7 @@ def test_fuse_act_padding(
94
)
95
96
with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
98
- from vllm.compilation.rocm_aiter_fusion import (
97
+ from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
98
RocmAiterTritonAddRMSNormPadFusionPass,
99
)
100
tests/compile/passes/test_fusion.py
renamed
+16
-11
@@ -7,12 +7,18 @@ import torch
7
8
import vllm.config
9
import vllm.plugins
10
+from tests.compile.backend import TestBackend
11
+from tests.utils import TestBlockFP8Layer, TestFP8Layer
12
from vllm._aiter_ops import IS_AITER_FOUND, rocm_aiter_ops
11
-from vllm.compilation.fusion import FUSED_OPS, FusedRMSQuantKey, RMSNormQuantFusionPass
12
-from vllm.compilation.fx_utils import find_op_nodes
13
-from vllm.compilation.matcher_utils import QUANT_OPS
14
-from vllm.compilation.noop_elimination import NoOpEliminationPass
15
-from vllm.compilation.post_cleanup import PostCleanupPass
13
+from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS
14
+from vllm.compilation.passes.fusion.rms_quant_fusion import (
15
+ FUSED_OPS,
16
+ FusedRMSQuantKey,
17
+ RMSNormQuantFusionPass,
18
+)
19
+from vllm.compilation.passes.fx_utils import find_op_nodes
20
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
21
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
22
from vllm.config import (
23
CompilationConfig,
24
CompilationMode,
@@ -51,9 +57,6 @@ from vllm.utils.deep_gemm import (
57
is_deep_gemm_supported,
58
)
59
54
-from ..utils import TestBlockFP8Layer, TestFP8Layer
55
-from .backend import TestBackend
56
-
60
FP8_DTYPE = current_platform.fp8_dtype()
61
62
RMS_OP = torch.ops._C.rms_norm.default
@@ -223,7 +226,7 @@ class TestModel(torch.nn.Module):
226
if self.use_aiter_fusion:
227
if self.group_shape.is_per_group():
228
# Blockwise aiter fusion
226
- from vllm.compilation.rocm_aiter_fusion import (
229
+ from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
230
AiterFusedAddRMSFp8GroupQuantPattern,
231
AiterRMSFp8GroupQuantPattern,
232
)
@@ -234,7 +237,7 @@ class TestModel(torch.nn.Module):
237
]
238
else:
239
# Per-token aiter fusion
237
- from vllm.compilation.rocm_aiter_fusion import (
240
+ from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
241
AiterFusedAddRMSNormDynamicQuantPattern,
242
AiterRMSNormDynamicQuantPattern,
243
)
@@ -410,7 +413,9 @@ def test_aiter_fusion_rmsnorm_quant(
413
)
414
415
with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
413
- from vllm.compilation.rocm_aiter_fusion import RocmAiterRMSNormQuantFusionPass
416
+ from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
417
+ RocmAiterRMSNormQuantFusionPass,
418
+ )
419
420
m.setenv("VLLM_ROCM_USE_AITER", "1")
421
tests/compile/passes/test_fusion_attn.py
renamed
+6
-8
@@ -6,14 +6,14 @@ import pytest
6
import torch._dynamo
7
8
from tests.compile.backend import LazyInitPass, TestBackend
9
-from tests.utils import flat_product
9
+from tests.utils import TestFP8Layer, flat_product
10
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
11
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
12
-from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass
13
-from vllm.compilation.fx_utils import find_op_nodes
14
-from vllm.compilation.matcher_utils import QUANT_OPS
15
-from vllm.compilation.noop_elimination import NoOpEliminationPass
16
-from vllm.compilation.post_cleanup import PostCleanupPass
12
+from vllm.compilation.passes.fusion.attn_quant_fusion import ATTN_OP, AttnFusionPass
13
+from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS
14
+from vllm.compilation.passes.fx_utils import find_op_nodes
15
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
16
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
17
from vllm.config import (
18
AttentionConfig,
19
CacheConfig,
@@ -38,8 +38,6 @@ from vllm.v1.attention.backend import AttentionMetadata
38
from vllm.v1.attention.backends.registry import AttentionBackendEnum
39
from vllm.v1.kv_cache_interface import AttentionSpec
40
41
-from ..utils import TestFP8Layer
42
-
41
FP8_DTYPE = current_platform.fp8_dtype()
42
FP4_DTYPE = torch.uint8
43
tests/compile/passes/test_noop_elimination.py
renamed
+2
-3
@@ -5,11 +5,10 @@ import pytest
5
import torch
6
7
import vllm
8
-from vllm.compilation.noop_elimination import NoOpEliminationPass
8
+from tests.compile.backend import TestBackend
9
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
10
from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig
11
11
-from .backend import TestBackend
12
-
12
13
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
14
# Important edge case is when `num_tokens == buffer_size`
tests/compile/passes/test_pass_manager.py
renamed
+2
-2
@@ -5,12 +5,12 @@ import copy
5
import pytest
6
import torch
7
8
-from vllm.compilation.inductor_pass import (
8
+from vllm.compilation.passes.inductor_pass import (
9
CallableInductorPass,
10
InductorPass,
11
pass_context,
12
)
13
-from vllm.compilation.pass_manager import PostGradPassManager
13
+from vllm.compilation.passes.pass_manager import PostGradPassManager
14
from vllm.config import ModelConfig, VllmConfig
15
from vllm.config.utils import Range
16
tests/compile/passes/test_qk_norm_rope_fusion.py
renamed
+8
-4
@@ -5,13 +5,17 @@ import pytest
5
import torch
6
7
from tests.compile.backend import TestBackend
8
-from vllm.compilation.matcher_utils import FLASHINFER_ROTARY_OP, RMS_OP, ROTARY_OP
9
-from vllm.compilation.noop_elimination import NoOpEliminationPass
10
-from vllm.compilation.post_cleanup import PostCleanupPass
11
-from vllm.compilation.qk_norm_rope_fusion import (
8
+from vllm.compilation.passes.fusion.matcher_utils import (
9
+ FLASHINFER_ROTARY_OP,
10
+ RMS_OP,
11
+ ROTARY_OP,
12
+)
13
+from vllm.compilation.passes.fusion.qk_norm_rope_fusion import (
14
FUSED_QK_ROPE_OP,
15
QKNormRoPEFusionPass,
16
)
17
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
18
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
19
from vllm.config import (
20
CompilationConfig,
21
CompilationMode,
tests/compile/passes/test_silu_mul_quant_fusion.py
renamed
+8
-9
@@ -6,17 +6,19 @@ import pytest
6
import torch
7
8
import vllm.envs as envs
9
+from tests.compile.backend import TestBackend
10
from tests.kernels.quantization.nvfp4_utils import quant_nvfp4_tensor
11
+from tests.utils import TestFP8Layer
12
from vllm._aiter_ops import IS_AITER_FOUND
13
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
12
-from vllm.compilation.activation_quant_fusion import (
14
+from vllm.compilation.passes.fusion.act_quant_fusion import (
15
FUSED_OPS,
16
SILU_MUL_OP,
17
ActivationQuantFusionPass,
18
)
17
-from vllm.compilation.fusion import QUANT_OPS
18
-from vllm.compilation.noop_elimination import NoOpEliminationPass
19
-from vllm.compilation.post_cleanup import PostCleanupPass
19
+from vllm.compilation.passes.fusion.rms_quant_fusion import QUANT_OPS
20
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
21
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
22
from vllm.config import (
23
CompilationConfig,
24
CompilationMode,
@@ -48,9 +50,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
50
)
51
from vllm.platforms import current_platform
52
51
-from ..utils import TestFP8Layer
52
-from .backend import TestBackend
53
-
53
FP8_DTYPE = current_platform.fp8_dtype()
54
FP4_DTYPE = torch.uint8
55
@@ -100,7 +99,7 @@ class TestSiluMulFp8QuantModel(torch.nn.Module):
99
class TestSiluMulNvfp4QuantModel(torch.nn.Module):
100
def __init__(self, hidden_size: int, x: torch.Tensor, **kwargs):
101
super().__init__()
103
- from vllm.compilation.activation_quant_fusion import (
102
+ from vllm.compilation.passes.fusion.act_quant_fusion import (
103
silu_and_mul_nvfp4_quant_supported,
104
)
105
@@ -239,7 +238,7 @@ def test_fusion_silu_and_mul_quant(
238
with set_current_vllm_config(config):
239
fusion_passes = [ActivationQuantFusionPass(config)]
240
if IS_AITER_FOUND:
242
- from vllm.compilation.rocm_aiter_fusion import (
241
+ from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
242
RocmAiterSiluMulFp8GroupQuantFusionPass,
243
)
244
tests/compile/test_compile_ranges.py
+1
-1
@@ -10,7 +10,7 @@ from torch import nn
10
import tests.compile.silly_attention # noqa
11
from vllm.compilation.counter import compilation_counter
12
from vllm.compilation.decorators import support_torch_compile
13
-from vllm.compilation.inductor_pass import (
13
+from vllm.compilation.passes.inductor_pass import (
14
InductorPass,
15
get_pass_context,
16
)
tests/compile/test_config.py
+3
-1
@@ -8,7 +8,9 @@ import pytest
8
from pydantic import ValidationError
9
10
from vllm.compilation.counter import compilation_counter
11
-from vllm.compilation.fix_functionalization import FixFunctionalizationPass
11
+from vllm.compilation.passes.utility.fix_functionalization import (
12
+ FixFunctionalizationPass,
13
+)
14
from vllm.config import (
15
CompilationConfig,
16
CUDAGraphMode,
tests/compile/test_graph_partition.py
+1
-1
@@ -8,7 +8,7 @@ import torch
8
from torch.fx.experimental.proxy_tensor import make_fx
9
10
from vllm.compilation.backends import split_graph
11
-from vllm.compilation.fx_utils import find_op_nodes
11
+from vllm.compilation.passes.fx_utils import find_op_nodes
12
13
# This import automatically registers `torch.ops.silly.attention`
14
from . import silly_attention # noqa: F401
vllm/compilation/backends.py
+6
-7
@@ -22,11 +22,6 @@ from torch._dispatch.python import enable_python_dispatcher
22
from torch._logging._internal import trace_structured
23
24
import vllm.envs as envs
25
-from vllm.compilation.inductor_pass import pass_context
26
-from vllm.compilation.partition_rules import (
27
- inductor_partition_rule_context,
28
- should_split,
29
-)
25
from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
26
from vllm.config.compilation import DynamicShapesType
27
from vllm.config.utils import Range, hash_factors
@@ -44,8 +39,12 @@ from .compiler_interface import (
39
is_compile_cache_enabled,
40
)
41
from .counter import compilation_counter
47
-from .inductor_pass import InductorPass
48
-from .pass_manager import PostGradPassManager
42
+from .partition_rules import (
43
+ inductor_partition_rule_context,
44
+ should_split,
45
+)
46
+from .passes.inductor_pass import InductorPass, pass_context
47
+from .passes.pass_manager import PostGradPassManager
48
49
logger = init_logger(__name__)
50
vllm/compilation/passes/__init__.py
vllm/compilation/passes/fusion/__init__.py
vllm/compilation/passes/fusion/act_quant_fusion.py
renamed
+3
-3
@@ -22,10 +22,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
22
)
23
from vllm.platforms import current_platform
24
25
-from .fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
26
-from .inductor_pass import enable_fake_mode
25
+from ..inductor_pass import enable_fake_mode
26
+from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
27
from .matcher_utils import MatcherQuantFP8, MatcherSiluAndMul
28
-from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
28
+from .rms_quant_fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
29
30
logger = init_logger(__name__)
31
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
renamed
+13
-403
@@ -8,7 +8,6 @@ import torch._inductor.pattern_matcher as pm
8
import torch.fx as fx
9
from torch._higher_order_ops.auto_functionalize import auto_functionalized
10
from torch._inductor.pattern_matcher import PatternMatcherPass
11
-from torch.distributed._symmetric_memory import enable_symm_mem_for_group
11
12
from vllm.config import VllmConfig
13
from vllm.config.utils import Range
@@ -24,12 +23,14 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
from vllm.platforms import current_platform
24
from vllm.utils.torch_utils import direct_register_custom_op
25
27
-from .inductor_pass import enable_fake_mode
26
+from ..inductor_pass import enable_fake_mode
27
+from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
28
from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNorm
29
-from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
29
30
FP8_DTYPE = current_platform.fp8_dtype()
31
32
+logger = init_logger(__name__)
33
+
34
flashinfer_comm: ModuleType | None = None
35
if find_spec("flashinfer"):
36
try:
@@ -45,406 +46,6 @@ logger = init_logger(__name__)
46
if hasattr(torch.ops._C, "scaled_fp4_quant"):
47
STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.default
48
48
-
49
-class BasePattern:
50
- def __init__(self, dtype: torch.dtype, device: str | None) -> None:
51
- self.dtype = dtype
52
- self.device = device
53
- self.tp = get_tp_group()
54
- self.tp_size = get_tensor_model_parallel_world_size()
55
-
56
-
57
-class GEMMReduceScatterPattern(BasePattern):
58
- def get_inputs(self) -> list[torch.Tensor]:
59
- mul = torch.empty([16, 4], device=self.device, dtype=self.dtype)
60
- mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
61
- return [mul, mm_weight]
62
-
63
- def register(self, pm_pass: PatternMatcherPass) -> None:
64
- def pattern(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
65
- mm = torch.ops.aten.mm.default(mul, mm_weight)
66
- reduce_scatter = torch.ops.vllm.reduce_scatter.default(
67
- mm,
68
- dim=0,
69
- world_size=self.tp_size,
70
- group_name=self.tp.unique_name,
71
- )
72
- return reduce_scatter
73
-
74
- def replacement(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
75
- gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter(
76
- mul,
77
- mm_weight,
78
- "avg",
79
- scatter_dim=0,
80
- group_name=self.tp.device_group.group_name,
81
- )
82
-
83
- return gemm_rs
84
-
85
- pm.register_replacement(
86
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
87
- )
88
-
89
-
90
-class AllGatherGEMMPattern(BasePattern):
91
- def get_inputs(self) -> list[torch.Tensor]:
92
- x = torch.empty([4, 4], device=self.device, dtype=self.dtype)
93
- weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
94
-
95
- return [x, weight]
96
-
97
- def register(self, pm_pass: PatternMatcherPass) -> None:
98
- def pattern(
99
- x: torch.Tensor,
100
- weight: torch.Tensor,
101
- ) -> torch.Tensor:
102
- all_gather = torch.ops.vllm.all_gather.default(
103
- x,
104
- dim=0,
105
- world_size=self.tp_size,
106
- group_name=self.tp.unique_name,
107
- )
108
-
109
- return torch.ops.aten.mm.default(all_gather, weight)
110
-
111
- def replacement(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
112
- ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul(
113
- x,
114
- [weight],
115
- gather_dim=0,
116
- group_name=self.tp.device_group.group_name,
117
- )
118
- return mm_outputs
119
-
120
- pm.register_replacement(
121
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
122
- )
123
-
124
-
125
-class ScaledMMReduceScatterPattern(BasePattern):
126
- def get_inputs(self) -> list[torch.Tensor]:
127
- input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
128
- mm_weight = (
129
- torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
130
- .contiguous()
131
- .transpose(0, 1)
132
- )
133
- scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32)
134
- scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
135
- return [input, mm_weight, scale_a, scale_b]
136
-
137
- def register(self, pm_pass: PatternMatcherPass) -> None:
138
- def pattern(
139
- input: torch.Tensor,
140
- mat2: torch.Tensor,
141
- scale_a: torch.Tensor,
142
- scale_b: torch.Tensor,
143
- ) -> torch.Tensor:
144
- scaled_mm = torch.ops.aten._scaled_mm.default(
145
- input,
146
- mat2=mat2,
147
- scale_a=scale_a,
148
- scale_b=scale_b,
149
- bias=None,
150
- scale_result=None,
151
- out_dtype=self.dtype,
152
- )
153
- reduce_scatter = torch.ops.vllm.reduce_scatter.default(
154
- scaled_mm,
155
- dim=0,
156
- world_size=self.tp_size,
157
- group_name=self.tp.unique_name,
158
- )
159
- return reduce_scatter
160
-
161
- def replacement(
162
- input: torch.Tensor,
163
- mat2: torch.Tensor,
164
- scale_a: torch.Tensor,
165
- scale_b: torch.Tensor,
166
- ) -> torch.Tensor:
167
- # Calculate output shape: input @ mat2 with scatter_dim reduced
168
- output_shape = [*input.shape[:-1], mat2.shape[1]]
169
- scatter_dim = 0
170
- gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter(
171
- input,
172
- mat2,
173
- scale_a,
174
- scale_b,
175
- "avg",
176
- scatter_dim, # orig_scatter_dim
177
- scatter_dim, # scatter_dim_after_maybe_reshape
178
- self.tp.device_group.group_name,
179
- output_shape,
180
- None, # bias
181
- None, # result_scale
182
- self.dtype, # out_dtype
183
- False, # use_fast_accum
184
- )
185
-
186
- return gemm_rs
187
-
188
- pm.register_replacement(
189
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
190
- )
191
-
192
-
193
-class AllGatherScaledMMPattern(BasePattern):
194
- def get_inputs(self) -> list[torch.Tensor]:
195
- x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
196
- weight = (
197
- torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
198
- .contiguous()
199
- .transpose(0, 1)
200
- )
201
-
202
- s1 = x.shape[0] * self.tp_size
203
-
204
- scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32)
205
- scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
206
-
207
- return [x, weight, scale_a, scale_b]
208
-
209
- def register(self, pm_pass: PatternMatcherPass) -> None:
210
- def pattern(
211
- x: torch.Tensor,
212
- weight: torch.Tensor,
213
- scale_a: torch.Tensor,
214
- scale_b: torch.Tensor,
215
- ) -> torch.Tensor:
216
- all_gather = torch.ops.vllm.all_gather.default(
217
- x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name
218
- )
219
-
220
- return torch.ops.aten._scaled_mm.default(
221
- all_gather,
222
- mat2=weight,
223
- scale_a=scale_a,
224
- scale_b=scale_b,
225
- bias=None,
226
- scale_result=None,
227
- out_dtype=self.dtype,
228
- )
229
-
230
- def replacement(
231
- x: torch.Tensor,
232
- weight: torch.Tensor,
233
- scale_a: torch.Tensor,
234
- scale_b: torch.Tensor,
235
- ) -> torch.Tensor:
236
- ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa
237
- x,
238
- [weight],
239
- scale_a,
240
- [scale_b],
241
- gather_dim=0,
242
- biases=[None],
243
- result_scales=[None],
244
- out_dtypes=[self.dtype],
245
- use_fast_accum=[False],
246
- group_name=self.tp.device_group.group_name,
247
- )
248
- return mm_outputs
249
-
250
- pm.register_replacement(
251
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
252
- )
253
-
254
-
255
-class CutlassScaledMMReduceScatterPattern(BasePattern):
256
- def get_inputs(self) -> list[torch.Tensor]:
257
- input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
258
- mm_weight = (
259
- torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
260
- .contiguous()
261
- .transpose(0, 1)
262
- )
263
- scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32)
264
- scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
265
-
266
- cutlass_mm_output = torch.empty([16, 16], device=self.device, dtype=self.dtype)
267
- return [input, mm_weight, scale_a, scale_b, cutlass_mm_output]
268
-
269
- def register(self, pm_pass: PatternMatcherPass) -> None:
270
- def pattern(
271
- input: torch.Tensor,
272
- weight: torch.Tensor,
273
- scale_a: torch.Tensor,
274
- scale_b: torch.Tensor,
275
- cutlass_mm_output: torch.Tensor,
276
- ) -> torch.Tensor:
277
- cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized(
278
- torch.ops._C.cutlass_scaled_mm.default,
279
- out=cutlass_mm_output,
280
- a=input,
281
- b=weight,
282
- a_scales=scale_a,
283
- b_scales=scale_b,
284
- bias=None,
285
- )
286
-
287
- reduce_scatter = torch.ops.vllm.reduce_scatter.default(
288
- cutlass_scaled_mm[1],
289
- dim=0,
290
- world_size=self.tp_size,
291
- group_name=self.tp.unique_name,
292
- )
293
- return reduce_scatter
294
-
295
- def replacement(
296
- input: torch.Tensor,
297
- mat2: torch.Tensor,
298
- scale_a: torch.Tensor,
299
- scale_b: torch.Tensor,
300
- cutlass_mm_output: torch.Tensor,
301
- ) -> torch.Tensor:
302
- # Calculate output shape: input @ mat2 with scatter_dim reduced
303
- output_shape = [*input.shape[:-1], mat2.shape[1]]
304
- scatter_dim = 0
305
- gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter(
306
- input,
307
- mat2,
308
- scale_a,
309
- scale_b,
310
- "avg",
311
- scatter_dim, # orig_scatter_dim
312
- scatter_dim, # scatter_dim_after_maybe_reshape
313
- self.tp.device_group.group_name,
314
- output_shape,
315
- None, # bias
316
- None, # result_scale
317
- self.dtype, # out_dtype
318
- False, # use_fast_accum
319
- )
320
-
321
- return gemm_rs
322
-
323
- pm.register_replacement(
324
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
325
- )
326
-
327
-
328
-class AllGatherCutlassScaledMMPattern(BasePattern):
329
- def get_inputs(self) -> list[torch.Tensor]:
330
- x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
331
- weight = (
332
- torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
333
- .contiguous()
334
- .transpose(0, 1)
335
- )
336
-
337
- s1 = x.shape[0] * self.tp_size
338
-
339
- scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32)
340
- scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
341
-
342
- s2 = weight.shape[1]
343
- output = torch.empty([s1, s2], device=self.device, dtype=self.dtype)
344
-
345
- return [x, weight, scale_a, scale_b, output]
346
-
347
- def register(self, pm_pass: PatternMatcherPass) -> None:
348
- def pattern(
349
- x: torch.Tensor,
350
- weight: torch.Tensor,
351
- scale_a: torch.Tensor,
352
- scale_b: torch.Tensor,
353
- output: torch.Tensor,
354
- ) -> torch.Tensor:
355
- all_gather = torch.ops.vllm.all_gather.default(
356
- x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name
357
- )
358
-
359
- cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized(
360
- torch.ops._C.cutlass_scaled_mm.default,
361
- out=output,
362
- a=all_gather,
363
- b=weight,
364
- a_scales=scale_a,
365
- b_scales=scale_b,
366
- bias=None,
367
- )
368
- return cutlass_scaled_mm[1]
369
-
370
- def replacement(
371
- x: torch.Tensor,
372
- weight: torch.Tensor,
373
- scale_a: torch.Tensor,
374
- scale_b: torch.Tensor,
375
- output: torch.Tensor,
376
- ) -> torch.Tensor:
377
- ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa
378
- x,
379
- [weight],
380
- scale_a,
381
- [scale_b],
382
- gather_dim=0,
383
- biases=[None],
384
- result_scales=[None],
385
- out_dtypes=[self.dtype],
386
- use_fast_accum=[False],
387
- group_name=self.tp.device_group.group_name,
388
- )
389
- return mm_outputs
390
-
391
- pm.register_replacement(
392
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
393
- )
394
-
395
-
396
-class AsyncTPPass(VllmPatternMatcherPass):
397
- @enable_fake_mode
398
- def __init__(self, config: VllmConfig) -> None:
399
- super().__init__(config)
400
-
401
- # Enable symmetric memory for the TP process group
402
- enable_symm_mem_for_group(get_tp_group().device_group.group_name)
403
- self.patterns: PatternMatcherPass = PatternMatcherPass(
404
- pass_name="async_tp_pass"
405
- )
406
- GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.patterns)
407
-
408
- AllGatherGEMMPattern(self.model_dtype, self.device).register(self.patterns)
409
-
410
- # These fusions are enabled only for bfloat16 models because
411
- # `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling
412
- # only supports bfloat16 as the output dtype.
413
- if self.model_dtype == torch.bfloat16:
414
- ScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
415
- self.patterns
416
- )
417
- AllGatherScaledMMPattern(self.model_dtype, self.device).register(
418
- self.patterns
419
- )
420
-
421
- CutlassScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
422
- self.patterns
423
- )
424
- AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register(
425
- self.patterns
426
- )
427
-
428
- self.dump_patterns(config, self.patterns)
429
-
430
- def is_applicable_for_range(self, compile_range: Range) -> bool:
431
- # This pass is applied on top of the sequence parallelism pass.
432
- # It inherits the same applicability condition as `SequenceParallelismPass`.
433
- # See `SequenceParallelismPass.is_applicable` for more details.
434
- if (
435
- not self.compilation_config.splitting_ops
436
- or self.compilation_config.use_inductor_graph_partition
437
- ):
438
- return True
439
- tp_size = get_tensor_model_parallel_world_size()
440
- return bool(compile_range.is_single_size() and compile_range.end % tp_size == 0)
441
-
442
- @VllmInductorPass.time_and_log
443
- def __call__(self, graph: fx.Graph) -> None:
444
- self.matched_count = self.patterns.apply(graph)
445
- logger.debug("Replaced %s patterns", self.matched_count)
446
-
447
-
49
# Max size of the input tensor per world size per device capability
50
# to use flashinfer fused allreduce
51
FI_ALLREDUCE_FUSION_MAX_SIZE_MB: dict[int, dict[int, float]] = {
@@ -623,6 +224,15 @@ class FlashInferFusedAllReduceParams:
224
}
225
226
227
+# TODO(luka): unify
228
+class BasePattern:
229
+ def __init__(self, dtype: torch.dtype, device: str | None) -> None:
230
+ self.dtype = dtype
231
+ self.device = device
232
+ self.tp = get_tp_group()
233
+ self.tp_size = get_tensor_model_parallel_world_size()
234
+
235
+
236
class AllReduceRMSNormPattern(BasePattern):
237
"""
238
This pattern replaces the allreduce + rms norm (without residual)
vllm/compilation/passes/fusion/attn_quant_fusion.py
renamed
+4
-4
@@ -22,11 +22,11 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
22
from vllm.platforms import current_platform
23
from vllm.utils.math_utils import round_up
24
25
-from .fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
26
-from .fx_utils import is_func
27
-from .inductor_pass import enable_fake_mode
25
+from ..fx_utils import is_func
26
+from ..inductor_pass import enable_fake_mode
27
+from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
28
from .matcher_utils import MatcherQuantFP8
29
-from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
29
+from .rms_quant_fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
30
31
logger = init_logger(__name__)
32
P = ParamSpec("P")
vllm/compilation/passes/fusion/collective_fusion.py
new
+423
@@ -0,0 +1,423 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+import torch._inductor.pattern_matcher as pm
6
+import torch.fx as fx
7
+from torch._inductor.pattern_matcher import PatternMatcherPass
8
+from torch.distributed._symmetric_memory import enable_symm_mem_for_group
9
+
10
+from vllm.config import VllmConfig
11
+from vllm.config.utils import Range
12
+from vllm.distributed import get_tp_group
13
+from vllm.distributed.parallel_state import (
14
+ get_tensor_model_parallel_world_size,
15
+)
16
+from vllm.logger import init_logger
17
+from vllm.platforms import current_platform
18
+
19
+from ..inductor_pass import enable_fake_mode
20
+from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
21
+
22
+FP8_DTYPE = current_platform.fp8_dtype()
23
+
24
+logger = init_logger(__name__)
25
+
26
+
27
+class BasePattern:
28
+ def __init__(self, dtype: torch.dtype, device: str | None) -> None:
29
+ self.dtype = dtype
30
+ self.device = device
31
+ self.tp = get_tp_group()
32
+ self.tp_size = get_tensor_model_parallel_world_size()
33
+
34
+
35
+class GEMMReduceScatterPattern(BasePattern):
36
+ def get_inputs(self) -> list[torch.Tensor]:
37
+ mul = torch.empty([16, 4], device=self.device, dtype=self.dtype)
38
+ mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
39
+ return [mul, mm_weight]
40
+
41
+ def register(self, pm_pass: PatternMatcherPass) -> None:
42
+ def pattern(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
43
+ mm = torch.ops.aten.mm.default(mul, mm_weight)
44
+ reduce_scatter = torch.ops.vllm.reduce_scatter.default(
45
+ mm,
46
+ dim=0,
47
+ world_size=self.tp_size,
48
+ group_name=self.tp.unique_name,
49
+ )
50
+ return reduce_scatter
51
+
52
+ def replacement(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
53
+ gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter(
54
+ mul,
55
+ mm_weight,
56
+ "avg",
57
+ scatter_dim=0,
58
+ group_name=self.tp.device_group.group_name,
59
+ )
60
+
61
+ return gemm_rs
62
+
63
+ pm.register_replacement(
64
+ pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
65
+ )
66
+
67
+
68
+class AllGatherGEMMPattern(BasePattern):
69
+ def get_inputs(self) -> list[torch.Tensor]:
70
+ x = torch.empty([4, 4], device=self.device, dtype=self.dtype)
71
+ weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
72
+
73
+ return [x, weight]
74
+
75
+ def register(self, pm_pass: PatternMatcherPass) -> None:
76
+ def pattern(
77
+ x: torch.Tensor,
78
+ weight: torch.Tensor,
79
+ ) -> torch.Tensor:
80
+ all_gather = torch.ops.vllm.all_gather.default(
81
+ x,
82
+ dim=0,
83
+ world_size=self.tp_size,
84
+ group_name=self.tp.unique_name,
85
+ )
86
+
87
+ return torch.ops.aten.mm.default(all_gather, weight)
88
+
89
+ def replacement(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
90
+ ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul(
91
+ x,
92
+ [weight],
93
+ gather_dim=0,
94
+ group_name=self.tp.device_group.group_name,
95
+ )
96
+ return mm_outputs
97
+
98
+ pm.register_replacement(
99
+ pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
100
+ )
101
+
102
+
103
+class ScaledMMReduceScatterPattern(BasePattern):
104
+ def get_inputs(self) -> list[torch.Tensor]:
105
+ input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
106
+ mm_weight = (
107
+ torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
108
+ .contiguous()
109
+ .transpose(0, 1)
110
+ )
111
+ scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32)
112
+ scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
113
+ return [input, mm_weight, scale_a, scale_b]
114
+
115
+ def register(self, pm_pass: PatternMatcherPass) -> None:
116
+ def pattern(
117
+ input: torch.Tensor,
118
+ mat2: torch.Tensor,
119
+ scale_a: torch.Tensor,
120
+ scale_b: torch.Tensor,
121
+ ) -> torch.Tensor:
122
+ scaled_mm = torch.ops.aten._scaled_mm.default(
123
+ input,
124
+ mat2=mat2,
125
+ scale_a=scale_a,
126
+ scale_b=scale_b,
127
+ bias=None,
128
+ scale_result=None,
129
+ out_dtype=self.dtype,
130
+ )
131
+ reduce_scatter = torch.ops.vllm.reduce_scatter.default(
132
+ scaled_mm,
133
+ dim=0,
134
+ world_size=self.tp_size,
135
+ group_name=self.tp.unique_name,
136
+ )
137
+ return reduce_scatter
138
+
139
+ def replacement(
140
+ input: torch.Tensor,
141
+ mat2: torch.Tensor,
142
+ scale_a: torch.Tensor,
143
+ scale_b: torch.Tensor,
144
+ ) -> torch.Tensor:
145
+ # Calculate output shape: input @ mat2 with scatter_dim reduced
146
+ output_shape = [*input.shape[:-1], mat2.shape[1]]
147
+ scatter_dim = 0
148
+ gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter(
149
+ input,
150
+ mat2,
151
+ scale_a,
152
+ scale_b,
153
+ "avg",
154
+ scatter_dim, # orig_scatter_dim
155
+ scatter_dim, # scatter_dim_after_maybe_reshape
156
+ self.tp.device_group.group_name,
157
+ output_shape,
158
+ None, # bias
159
+ None, # result_scale
160
+ self.dtype, # out_dtype
161
+ False, # use_fast_accum
162
+ )
163
+
164
+ return gemm_rs
165
+
166
+ pm.register_replacement(
167
+ pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
168
+ )
169
+
170
+
171
+class AllGatherScaledMMPattern(BasePattern):
172
+ def get_inputs(self) -> list[torch.Tensor]:
173
+ x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
174
+ weight = (
175
+ torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
176
+ .contiguous()
177
+ .transpose(0, 1)
178
+ )
179
+
180
+ s1 = x.shape[0] * self.tp_size
181
+
182
+ scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32)
183
+ scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
184
+
185
+ return [x, weight, scale_a, scale_b]
186
+
187
+ def register(self, pm_pass: PatternMatcherPass) -> None:
188
+ def pattern(
189
+ x: torch.Tensor,
190
+ weight: torch.Tensor,
191
+ scale_a: torch.Tensor,
192
+ scale_b: torch.Tensor,
193
+ ) -> torch.Tensor:
194
+ all_gather = torch.ops.vllm.all_gather.default(
195
+ x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name
196
+ )
197
+
198
+ return torch.ops.aten._scaled_mm.default(
199
+ all_gather,
200
+ mat2=weight,
201
+ scale_a=scale_a,
202
+ scale_b=scale_b,
203
+ bias=None,
204
+ scale_result=None,
205
+ out_dtype=self.dtype,
206
+ )
207
+
208
+ def replacement(
209
+ x: torch.Tensor,
210
+ weight: torch.Tensor,
211
+ scale_a: torch.Tensor,
212
+ scale_b: torch.Tensor,
213
+ ) -> torch.Tensor:
214
+ ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa
215
+ x,
216
+ [weight],
217
+ scale_a,
218
+ [scale_b],
219
+ gather_dim=0,
220
+ biases=[None],
221
+ result_scales=[None],
222
+ out_dtypes=[self.dtype],
223
+ use_fast_accum=[False],
224
+ group_name=self.tp.device_group.group_name,
225
+ )
226
+ return mm_outputs
227
+
228
+ pm.register_replacement(
229
+ pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
230
+ )
231
+
232
+
233
+class CutlassScaledMMReduceScatterPattern(BasePattern):
234
+ def get_inputs(self) -> list[torch.Tensor]:
235
+ input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
236
+ mm_weight = (
237
+ torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
238
+ .contiguous()
239
+ .transpose(0, 1)
240
+ )
241
+ scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32)
242
+ scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
243
+
244
+ cutlass_mm_output = torch.empty([16, 16], device=self.device, dtype=self.dtype)
245
+ return [input, mm_weight, scale_a, scale_b, cutlass_mm_output]
246
+
247
+ def register(self, pm_pass: PatternMatcherPass) -> None:
248
+ def pattern(
249
+ input: torch.Tensor,
250
+ weight: torch.Tensor,
251
+ scale_a: torch.Tensor,
252
+ scale_b: torch.Tensor,
253
+ cutlass_mm_output: torch.Tensor,
254
+ ) -> torch.Tensor:
255
+ cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized(
256
+ torch.ops._C.cutlass_scaled_mm.default,
257
+ out=cutlass_mm_output,
258
+ a=input,
259
+ b=weight,
260
+ a_scales=scale_a,
261
+ b_scales=scale_b,
262
+ bias=None,
263
+ )
264
+
265
+ reduce_scatter = torch.ops.vllm.reduce_scatter.default(
266
+ cutlass_scaled_mm[1],
267
+ dim=0,
268
+ world_size=self.tp_size,
269
+ group_name=self.tp.unique_name,
270
+ )
271
+ return reduce_scatter
272
+
273
+ def replacement(
274
+ input: torch.Tensor,
275
+ mat2: torch.Tensor,
276
+ scale_a: torch.Tensor,
277
+ scale_b: torch.Tensor,
278
+ cutlass_mm_output: torch.Tensor,
279
+ ) -> torch.Tensor:
280
+ # Calculate output shape: input @ mat2 with scatter_dim reduced
281
+ output_shape = [*input.shape[:-1], mat2.shape[1]]
282
+ scatter_dim = 0
283
+ gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter(
284
+ input,
285
+ mat2,
286
+ scale_a,
287
+ scale_b,
288
+ "avg",
289
+ scatter_dim, # orig_scatter_dim
290
+ scatter_dim, # scatter_dim_after_maybe_reshape
291
+ self.tp.device_group.group_name,
292
+ output_shape,
293
+ None, # bias
294
+ None, # result_scale
295
+ self.dtype, # out_dtype
296
+ False, # use_fast_accum
297
+ )
298
+
299
+ return gemm_rs
300
+
301
+ pm.register_replacement(
302
+ pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
303
+ )
304
+
305
+
306
+class AllGatherCutlassScaledMMPattern(BasePattern):
307
+ def get_inputs(self) -> list[torch.Tensor]:
308
+ x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE)
309
+ weight = (
310
+ torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
311
+ .contiguous()
312
+ .transpose(0, 1)
313
+ )
314
+
315
+ s1 = x.shape[0] * self.tp_size
316
+
317
+ scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32)
318
+ scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32)
319
+
320
+ s2 = weight.shape[1]
321
+ output = torch.empty([s1, s2], device=self.device, dtype=self.dtype)
322
+
323
+ return [x, weight, scale_a, scale_b, output]
324
+
325
+ def register(self, pm_pass: PatternMatcherPass) -> None:
326
+ def pattern(
327
+ x: torch.Tensor,
328
+ weight: torch.Tensor,
329
+ scale_a: torch.Tensor,
330
+ scale_b: torch.Tensor,
331
+ output: torch.Tensor,
332
+ ) -> torch.Tensor:
333
+ all_gather = torch.ops.vllm.all_gather.default(
334
+ x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name
335
+ )
336
+
337
+ cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized(
338
+ torch.ops._C.cutlass_scaled_mm.default,
339
+ out=output,
340
+ a=all_gather,
341
+ b=weight,
342
+ a_scales=scale_a,
343
+ b_scales=scale_b,
344
+ bias=None,
345
+ )
346
+ return cutlass_scaled_mm[1]
347
+
348
+ def replacement(
349
+ x: torch.Tensor,
350
+ weight: torch.Tensor,
351
+ scale_a: torch.Tensor,
352
+ scale_b: torch.Tensor,
353
+ output: torch.Tensor,
354
+ ) -> torch.Tensor:
355
+ ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa
356
+ x,
357
+ [weight],
358
+ scale_a,
359
+ [scale_b],
360
+ gather_dim=0,
361
+ biases=[None],
362
+ result_scales=[None],
363
+ out_dtypes=[self.dtype],
364
+ use_fast_accum=[False],
365
+ group_name=self.tp.device_group.group_name,
366
+ )
367
+ return mm_outputs
368
+
369
+ pm.register_replacement(
370
+ pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
371
+ )
372
+
373
+
374
+class AsyncTPPass(VllmPatternMatcherPass):
375
+ @enable_fake_mode
376
+ def __init__(self, config: VllmConfig) -> None:
377
+ super().__init__(config)
378
+
379
+ # Enable symmetric memory for the TP process group
380
+ enable_symm_mem_for_group(get_tp_group().device_group.group_name)
381
+ self.patterns: PatternMatcherPass = PatternMatcherPass(
382
+ pass_name="async_tp_pass"
383
+ )
384
+ GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.patterns)
385
+
386
+ AllGatherGEMMPattern(self.model_dtype, self.device).register(self.patterns)
387
+
388
+ # These fusions are enabled only for bfloat16 models because
389
+ # `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling
390
+ # only supports bfloat16 as the output dtype.
391
+ if self.model_dtype == torch.bfloat16:
392
+ ScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
393
+ self.patterns
394
+ )
395
+ AllGatherScaledMMPattern(self.model_dtype, self.device).register(
396
+ self.patterns
397
+ )
398
+
399
+ CutlassScaledMMReduceScatterPattern(self.model_dtype, self.device).register(
400
+ self.patterns
401
+ )
402
+ AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register(
403
+ self.patterns
404
+ )
405
+
406
+ self.dump_patterns(config, self.patterns)
407
+
408
+ def is_applicable_for_range(self, compile_range: Range) -> bool:
409
+ # This pass is applied on top of the sequence parallelism pass.
410
+ # It inherits the same applicability condition as `SequenceParallelismPass`.
411
+ # See `SequenceParallelismPass.is_applicable` for more details.
412
+ if (
413
+ not self.compilation_config.splitting_ops
414
+ or self.compilation_config.use_inductor_graph_partition
415
+ ):
416
+ return True
417
+ tp_size = get_tensor_model_parallel_world_size()
418
+ return bool(compile_range.is_single_size() and compile_range.end % tp_size == 0)
419
+
420
+ @VllmInductorPass.time_and_log
421
+ def __call__(self, graph: fx.Graph) -> None:
422
+ self.matched_count = self.patterns.apply(graph)
423
+ logger.debug("Replaced %s patterns", self.matched_count)
vllm/compilation/passes/fusion/matcher_utils.py
renamed
vllm/compilation/passes/fusion/qk_norm_rope_fusion.py
renamed
+3
-3
@@ -15,10 +15,10 @@ from vllm.logger import init_logger
15
from vllm.model_executor.layers.attention import Attention
16
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
17
18
-from .fusion import empty_bf16, empty_fp32, empty_i64
19
-from .inductor_pass import enable_fake_mode
18
+from ..inductor_pass import enable_fake_mode
19
+from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
20
from .matcher_utils import MatcherRMSNorm, MatcherRotaryEmbedding
21
-from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
21
+from .rms_quant_fusion import empty_bf16, empty_fp32, empty_i64
22
23
logger = init_logger(__name__)
24
vllm/compilation/passes/fusion/rms_quant_fusion.py
renamed
+2
-2
@@ -25,13 +25,13 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
25
)
26
from vllm.platforms import current_platform
27
28
-from .inductor_pass import enable_fake_mode
28
+from ..inductor_pass import enable_fake_mode
29
+from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
30
from .matcher_utils import (
31
MatcherFusedAddRMSNorm,
32
MatcherQuantFP8,
33
MatcherRMSNorm,
34
)
34
-from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
35
36
logger = init_logger(__name__)
37
FP8_DTYPE = current_platform.fp8_dtype()
vllm/compilation/passes/fusion/rocm_aiter_fusion.py
renamed
+6
-6
@@ -9,7 +9,6 @@ from torch._ops import OpOverload
9
10
import vllm.model_executor.layers.quantization.utils.fp8_utils # noqa: F401
11
from vllm._aiter_ops import rocm_aiter_ops
12
-from vllm.compilation.activation_quant_fusion import ActivationQuantPattern
12
from vllm.config import VllmConfig
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -19,17 +18,18 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
)
19
from vllm.platforms import current_platform
20
22
-from .fusion import (
23
- FusedRMSQuantKey,
24
-)
25
-from .inductor_pass import enable_fake_mode
21
+from ..activation_quant_fusion import ActivationQuantPattern
22
+from ..inductor_pass import enable_fake_mode
23
+from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
24
from .matcher_utils import (
25
MatcherFusedAddRMSNorm,
26
MatcherQuantFP8,
27
MatcherRMSNorm,
28
MatcherSiluAndMul,
29
)
32
-from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
30
+from .rms_quant_fusion import (
31
+ FusedRMSQuantKey,
32
+)
33
34
logger = init_logger(__name__)
35
FP8_DTYPE = current_platform.fp8_dtype()
vllm/compilation/passes/fusion/sequence_parallelism.py
renamed
+3
-3
@@ -20,10 +20,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
)
21
from vllm.platforms import current_platform
22
23
-from .inductor_pass import enable_fake_mode
23
+from ..inductor_pass import enable_fake_mode
24
+from ..utility.noop_elimination import NoOpEliminationPass
25
+from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
26
from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNorm
25
-from .noop_elimination import NoOpEliminationPass
26
-from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
27
28
logger = init_logger(__name__)
29
vllm/compilation/passes/fx_utils.py
renamed
vllm/compilation/passes/inductor_pass.py
renamed
vllm/compilation/passes/pass_manager.py
renamed
+11
-10
@@ -8,38 +8,39 @@ from torch import fx as fx
8
9
from vllm import envs
10
from vllm._aiter_ops import rocm_aiter_ops
11
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
12
from vllm.config import VllmConfig, set_current_vllm_config
13
from vllm.logger import init_logger
14
from vllm.platforms import current_platform
15
from vllm.utils.system_utils import set_env_var
16
16
-from .post_cleanup import PostCleanupPass
17
from .vllm_inductor_pass import VllmInductorPass
18
19
if rocm_aiter_ops.is_enabled():
20
- from vllm.compilation.rocm_aiter_fusion import (
20
+ from .fusion.rocm_aiter_fusion import (
21
RocmAiterRMSNormQuantFusionPass,
22
RocmAiterSiluMulFp8GroupQuantFusionPass,
23
RocmAiterTritonAddRMSNormPadFusionPass,
24
)
25
26
if current_platform.is_cuda_alike():
27
- from .activation_quant_fusion import ActivationQuantFusionPass
28
- from .fusion import RMSNormQuantFusionPass
29
- from .fusion_attn import AttnFusionPass
30
- from .qk_norm_rope_fusion import QKNormRoPEFusionPass
31
- from .sequence_parallelism import SequenceParallelismPass
27
+ from .fusion.act_quant_fusion import ActivationQuantFusionPass
28
+ from .fusion.attn_quant_fusion import AttnFusionPass
29
+ from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
30
+ from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
31
+ from .fusion.sequence_parallelism import SequenceParallelismPass
32
33
if current_platform.is_cuda():
34
- from .collective_fusion import AllReduceFusionPass, AsyncTPPass
34
+ from .fusion.allreduce_rms_fusion import AllReduceFusionPass
35
+ from .fusion.collective_fusion import AsyncTPPass
36
36
-from .fix_functionalization import FixFunctionalizationPass
37
from .inductor_pass import (
38
CustomGraphPass,
39
InductorPass,
40
get_pass_context,
41
)
42
-from .noop_elimination import NoOpEliminationPass
42
+from .utility.fix_functionalization import FixFunctionalizationPass
43
+from .utility.noop_elimination import NoOpEliminationPass
44
45
logger = init_logger(__name__)
46
vllm/compilation/passes/utility/__init__.py
vllm/compilation/passes/utility/fix_functionalization.py
renamed
+2
-2
@@ -10,8 +10,8 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized
10
from vllm.logger import init_logger
11
from vllm.platforms import current_platform
12
13
-from .fx_utils import is_func
14
-from .vllm_inductor_pass import VllmInductorPass
13
+from ..fx_utils import is_func
14
+from ..vllm_inductor_pass import VllmInductorPass
15
16
logger = init_logger(__name__)
17
vllm/compilation/passes/utility/noop_elimination.py
renamed
+2
-2
@@ -9,8 +9,8 @@ from torch.fx.experimental.symbolic_shapes import statically_known_true
9
10
from vllm.logger import init_logger
11
12
-from .fx_utils import is_func
13
-from .vllm_inductor_pass import VllmInductorPass
12
+from ..fx_utils import is_func
13
+from ..vllm_inductor_pass import VllmInductorPass
14
15
logger = init_logger(__name__)
16
vllm/compilation/passes/utility/post_cleanup.py
renamed
+1
-1
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from torch import fx
4
5
-from vllm.compilation.vllm_inductor_pass import VllmInductorPass
5
+from ..vllm_inductor_pass import VllmInductorPass
6
7
8
class PostCleanupPass(VllmInductorPass):
vllm/compilation/passes/vllm_inductor_pass.py
renamed
vllm/config/compilation.py
+4
-2
@@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Literal
11
from pydantic import Field, TypeAdapter, field_validator
12
13
import vllm.envs as envs
14
-from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass
14
+from vllm.compilation.passes.inductor_pass import CallableInductorPass, InductorPass
15
from vllm.config.utils import (
16
Range,
17
config,
@@ -170,7 +170,9 @@ class PassConfig:
170
171
@staticmethod
172
def default_fi_allreduce_fusion_max_size_mb() -> dict[int, float]:
173
- from vllm.compilation.collective_fusion import FI_ALLREDUCE_FUSION_MAX_SIZE_MB
173
+ from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
174
+ FI_ALLREDUCE_FUSION_MAX_SIZE_MB,
175
+ )
176
from vllm.platforms import current_platform
177
178
if not current_platform.is_cuda():
vllm/platforms/interface.py
+1
-1
@@ -191,7 +191,7 @@ class Platform:
191
Get the pass manager class for this platform.
192
It will be registered as a custom pass under the current_platform.pass_key.
193
"""
194
- return "vllm.compilation.pass_manager.PostGradPassManager"
194
+ return "vllm.compilation.passes.pass_manager.PostGradPassManager"
195
196
@classmethod
197
def get_compile_backend(cls) -> str: