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, 21:30 UTC
[CI][torch.compile] Reduce e2e fusion test time (#33293)
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 4, 2026 at 19:09 UTC
4d9513537d00a9b6678a2b1ed3c3566a81f7dd77
17 files changed
+1065
-818
.buildkite/test-amd.yaml
+16
-29
@@ -604,9 +604,11 @@ steps:
604
- tests/compile
605
commands:
606
- pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
607
- # Limit to no custom ops to reduce running time
608
- # Wrap with quotes to escape yaml and avoid starting -k string with a -
609
- - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
607
+ # # Limit to no custom ops to reduce running time
608
+ # # Wrap with quotes to escape yaml and avoid starting -k string with a -
609
+ # - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
610
+ # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
611
+ # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
612
613
- label: Cudagraph test
614
timeout_in_minutes: 20
@@ -1181,7 +1183,6 @@ steps:
1183
- tests/compile/test_fusion_attn.py
1184
- tests/compile/test_silu_mul_quant_fusion.py
1185
- tests/compile/distributed/test_fusion_all_reduce.py
1184
- - tests/compile/distributed/test_fusions_e2e.py
1186
- tests/compile/fullgraph/test_full_graph.py
1187
commands:
1188
- nvidia-smi
@@ -1189,33 +1190,16 @@ steps:
1190
- pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
1191
# this runner has 2 GPUs available even though num_gpus=2 is not set
1192
- pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
1192
- # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
1193
- # Wrap with quotes to escape yaml
1194
- - "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'"
1193
+
1194
+ # # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
1195
+ # # Wrap with quotes to escape yaml
1196
+ # - "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'"
1197
+ # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
1198
+ # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
1199
+
1200
# test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
1201
- pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
1202
1198
-- label: Blackwell Fusion E2E Tests # 30 min
1199
- timeout_in_minutes: 40
1200
- working_dir: "/vllm-workspace/"
1201
- gpu: b200
1202
- optional: true
1203
- num_gpus: 2
1204
- source_file_dependencies:
1205
- - csrc/quantization/fp4/
1206
- - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
1207
- - vllm/v1/attention/backends/flashinfer.py
1208
- - vllm/compilation/
1209
- # can affect pattern matching
1210
- - vllm/model_executor/layers/layernorm.py
1211
- - vllm/model_executor/layers/activation.py
1212
- - vllm/model_executor/layers/quantization/input_quant_fp8.py
1213
- - tests/compile/distributed/test_fusions_e2e.py
1214
- commands:
1215
- - nvidia-smi
1216
- # Run all e2e fusion tests
1217
- - pytest -v -s tests/compile/distributed/test_fusions_e2e.py
1218
-
1203
- label: Blackwell GPT-OSS Eval
1204
timeout_in_minutes: 60
1205
working_dir: "/vllm-workspace/"
@@ -1566,7 +1550,10 @@ steps:
1550
- pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
1551
- pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
1552
#- pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm
1569
- - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
1553
+ # - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
1554
+ # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
1555
+ # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
1556
+
1557
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
1558
- pytest -v -s tests/distributed/test_context_parallel.py
1559
- 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
.buildkite/test-pipeline.yaml
+11
-70
@@ -537,9 +537,11 @@ steps:
537
commands:
538
# fp8 kv scales not supported on sm89, tested on Blackwell instead
539
- pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
540
- # Limit to no custom ops to reduce running time
541
- # Wrap with quotes to escape yaml and avoid starting -k string with a -
542
- - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
540
+ # # Limit to no custom ops to reduce running time
541
+ # # Wrap with quotes to escape yaml and avoid starting -k string with a -
542
+ # - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
543
+ # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
544
+ # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
545
546
- label: Cudagraph test
547
timeout_in_minutes: 20
@@ -1069,7 +1071,6 @@ steps:
1071
- tests/compile/test_fusion_attn.py
1072
- tests/compile/test_silu_mul_quant_fusion.py
1073
- tests/compile/distributed/test_fusion_all_reduce.py
1072
- - tests/compile/distributed/test_fusions_e2e.py
1074
- tests/compile/fullgraph/test_full_graph.py
1075
commands:
1076
- nvidia-smi
@@ -1077,75 +1078,15 @@ steps:
1078
- pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
1079
# this runner has 2 GPUs available even though num_gpus=2 is not set
1080
- pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
1080
- # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
1081
- # Wrap with quotes to escape yaml
1082
- - "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'"
1081
+ # # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
1082
+ # # Wrap with quotes to escape yaml
1083
+ # - "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'"
1084
+ # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
1085
+ # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
1086
+
1087
# test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
1088
- pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
1089
1086
-- label: Blackwell Fusion E2E Tests # 30 min
1087
- timeout_in_minutes: 40
1088
- working_dir: "/vllm-workspace/"
1089
- gpu: b200
1090
- optional: true
1091
- num_gpus: 2
1092
- source_file_dependencies:
1093
- - csrc/quantization/fp4/
1094
- - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
1095
- - vllm/v1/attention/backends/flashinfer.py
1096
- - vllm/compilation/
1097
- # can affect pattern matching
1098
- - vllm/model_executor/layers/layernorm.py
1099
- - vllm/model_executor/layers/activation.py
1100
- - vllm/model_executor/layers/quantization/input_quant_fp8.py
1101
- - tests/compile/distributed/test_fusions_e2e.py
1102
- commands:
1103
- - nvidia-smi
1104
- # Run all e2e fusion tests
1105
- - pytest -v -s tests/compile/distributed/test_fusions_e2e.py
1106
-
1107
-- label: Hopper Fusion E2E Tests (H100) # 10min
1108
- timeout_in_minutes: 70
1109
- working_dir: "/vllm-workspace/"
1110
- gpu: h100
1111
- optional: true
1112
- source_file_dependencies:
1113
- - csrc/quantization/fp4/
1114
- - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
1115
- - vllm/v1/attention/backends/flashinfer.py
1116
- - vllm/compilation/
1117
- # can affect pattern matching
1118
- - vllm/model_executor/layers/layernorm.py
1119
- - vllm/model_executor/layers/activation.py
1120
- - vllm/model_executor/layers/quantization/input_quant_fp8.py
1121
- - tests/compile/test_fusion_attn.py
1122
- commands:
1123
- - export VLLM_TEST_CLEAN_GPU_MEMORY=1
1124
- # skip Llama-4 since it does not fit on this device
1125
- - pytest -v -s tests/compile/test_fusion_attn.py -k 'not Llama-4'
1126
-
1127
-- label: Hopper Fusion Distributed E2E Tests (2xH100) # 70min
1128
- timeout_in_minutes: 70
1129
- working_dir: "/vllm-workspace/"
1130
- gpu: h100
1131
- optional: true
1132
- num_gpus: 2
1133
- source_file_dependencies:
1134
- - csrc/quantization/fp4/
1135
- - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
1136
- - vllm/v1/attention/backends/flashinfer.py
1137
- - vllm/compilation/
1138
- # can affect pattern matching
1139
- - vllm/model_executor/layers/layernorm.py
1140
- - vllm/model_executor/layers/activation.py
1141
- - vllm/model_executor/layers/quantization/input_quant_fp8.py
1142
- - tests/compile/distributed/test_fusions_e2e.py
1143
- commands:
1144
- - export VLLM_TEST_CLEAN_GPU_MEMORY=1
1145
- # Run all e2e fusion tests
1146
- - pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'
1147
- - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
1148
-
1090
- label: Blackwell GPT-OSS Eval
1091
timeout_in_minutes: 60
1092
working_dir: "/vllm-workspace/"
.buildkite/test_areas/compile.yaml
+166
-26
@@ -2,56 +2,196 @@ group: Compile
2
depends_on:
3
- image-build
4
steps:
5
-- label: Fusion and Compile Tests (B200)
6
- timeout_in_minutes: 40
5
+- label: Sequence Parallel Tests (2 GPUs)
6
+ timeout_in_minutes: 50
7
working_dir: "/vllm-workspace/"
8
- device: b200
8
+ num_devices: 2
9
source_file_dependencies:
10
- - csrc/quantization/fp4/
11
- - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
12
- - vllm/v1/attention/backends/flashinfer.py
10
+ - vllm/model_executor/layers/
11
+ - vllm/compilation/
12
- vllm/v1/worker/
13
- vllm/v1/cudagraph_dispatcher.py
14
+ - tests/distributed/test_sequence_parallel.py
15
+ commands:
16
+ - export VLLM_TEST_CLEAN_GPU_MEMORY=1
17
+ - pytest -v -s tests/distributed/test_sequence_parallel.py
18
+
19
+- label: Sequence Parallel Tests (2xH100)
20
+ timeout_in_minutes: 50
21
+ working_dir: "/vllm-workspace/"
22
+ device: h100
23
+ optional: true
24
+ num_devices: 2
25
+ commands:
26
+ - export VLLM_TEST_CLEAN_GPU_MEMORY=1
27
+ - pytest -v -s tests/distributed/test_sequence_parallel.py
28
+
29
+- label: Distributed Compile Unit Tests (2xH100)
30
+ timeout_in_minutes: 40
31
+ working_dir: "/vllm-workspace/"
32
+ device: h100
33
+ num_devices: 2
34
+ source_file_dependencies:
35
- vllm/compilation/
16
- # can affect pattern matching
36
+ - 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
40
+ commands:
41
+ - 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
45
+
46
+- label: Fusion and Compile Unit Tests (B200)
47
+ timeout_in_minutes: 20
48
+ working_dir: "/vllm-workspace/"
49
+ device: b200
50
+ source_file_dependencies:
51
+ - csrc/quantization/fp4/
52
+ - vllm/model_executor/layers/quantization/
53
- vllm/model_executor/layers/layernorm.py
54
- vllm/model_executor/layers/activation.py
19
- - vllm/model_executor/layers/quantization/input_quant_fp8.py
55
+ - vllm/model_executor/layers/attention/attention.py
56
+ - vllm/v1/attention/backends/flashinfer.py
57
+ - 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
23
- - tests/compile/distributed/test_fusions_e2e.py
61
- tests/compile/fullgraph/test_full_graph.py
62
commands:
63
+ # b200 runners are limited, so we limit the tests to the minimum set only supported on Blackwell
64
- nvidia-smi
27
- - pytest -v -s tests/compile/test_fusion_attn.py
65
+ - pytest -v -s tests/compile/test_fusion_attn.py -k FLASHINFER
66
- pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
67
# 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
31
- # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
32
- # Wrap with quotes to escape yaml
33
- - "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'"
69
# test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
70
+ # TODO(luka) move to H100 once pass tests run on H100
71
- pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
72
37
-- label: Fusion E2E (2 GPUs)(B200)
38
- timeout_in_minutes: 40
73
+- label: Fusion E2E Quick (H100)
74
+ timeout_in_minutes: 15
75
+ working_dir: "/vllm-workspace/"
76
+ device: h100
77
+ num_devices: 1
78
+ source_file_dependencies:
79
+ - csrc/quantization/
80
+ - vllm/model_executor/
81
+ - vllm/v1/attention/
82
+ - vllm/compilation/
83
+ - tests/compile/fusions_e2e/
84
+ commands:
85
+ - nvidia-smi
86
+ # Run all models and attn backends but only Inductor partition and native custom ops
87
+ - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
88
+ # Qwen requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
89
+ - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3"
90
+
91
+- label: Fusion E2E Config Sweep (H100)
92
+ timeout_in_minutes: 30
93
+ working_dir: "/vllm-workspace/"
94
+ device: h100
95
+ num_devices: 1
96
+ source_file_dependencies:
97
+ - csrc/quantization/
98
+ - vllm/compilation/
99
+ # can affect pattern matching
100
+ - vllm/model_executor/layers/layernorm.py
101
+ - vllm/model_executor/layers/activation.py
102
+ - vllm/model_executor/layers/attention/attention.py
103
+ - vllm/model_executor/layers/quantization/input_quant_fp8.py
104
+ - tests/compile/fusions_e2e/
105
+ commands:
106
+ - nvidia-smi
107
+ # Run just llama3 (fp8) for all config combinations
108
+ - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3"
109
+
110
+- label: Fusion E2E Config Sweep (B200)
111
+ timeout_in_minutes: 30
112
working_dir: "/vllm-workspace/"
113
device: b200
114
+ num_devices: 1
115
optional: true
116
+ commands:
117
+ - nvidia-smi
118
+ # Run all models and attn backends but only Inductor partition and native custom ops
119
+ # -k "inductor_partition and not +rms_norm and not +quant_fp8"
120
+ # Qwen requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
121
+ # -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3"
122
+ # Run just llama3 (fp8 & fp4) for all config combinations
123
+ # -k "llama-3"
124
+ - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and not +quant_fp8" -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3" -k "llama-3"
125
+
126
+- label: Fusion E2E TP2 Quick (H100)
127
+ timeout_in_minutes: 20
128
+ working_dir: "/vllm-workspace/"
129
+ device: h100
130
num_devices: 2
131
source_file_dependencies:
44
- - csrc/quantization/fp4/
45
- - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
46
- - vllm/v1/attention/backends/flashinfer.py
47
- - vllm/compilation/
48
- # can affect pattern matching
49
- - vllm/model_executor/layers/layernorm.py
50
- - vllm/model_executor/layers/activation.py
51
- - vllm/model_executor/layers/quantization/input_quant_fp8.py
52
- - tests/compile/distributed/test_fusions_e2e.py
132
+ - csrc/quantization/
133
+ - vllm/model_executor/
134
+ - vllm/v1/attention/
135
+ - vllm/compilation/
136
+ - tests/compile/fusions_e2e/
137
+ commands:
138
+ - nvidia-smi
139
+ # Run all models and attn backends but only Inductor partition and native custom ops
140
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
141
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
142
+
143
+- label: Fusion E2E TP2 AR-RMS Config Sweep (H100)
144
+ timeout_in_minutes: 40
145
+ working_dir: "/vllm-workspace/"
146
+ device: h100
147
+ num_devices: 2
148
+ source_file_dependencies:
149
+ - csrc/quantization/
150
+ - vllm/compilation/
151
+ # can affect pattern matching
152
+ - vllm/model_executor/layers/layernorm.py
153
+ - vllm/model_executor/layers/activation.py
154
+ - vllm/model_executor/layers/attention/attention.py
155
+ - vllm/model_executor/layers/quantization/input_quant_fp8.py
156
+ - tests/compile/fusions_e2e/
157
+ commands:
158
+ - nvidia-smi
159
+ # Run just llama3 (fp4 & fp8 & bf16) for all config combinations
160
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3"
161
+
162
+- label: Fusion E2E TP2 AsyncTP Config Sweep (H100)
163
+ timeout_in_minutes: 40
164
+ working_dir: "/vllm-workspace/"
165
+ device: h100
166
+ num_devices: 2
167
+ source_file_dependencies:
168
+ - csrc/quantization/
169
+ - vllm/compilation/
170
+ # can affect pattern matching
171
+ - vllm/model_executor/layers/layernorm.py
172
+ - vllm/model_executor/layers/activation.py
173
+ - vllm/model_executor/layers/attention/attention.py
174
+ - vllm/model_executor/layers/quantization/input_quant_fp8.py
175
+ - tests/compile/fusions_e2e/
176
commands:
177
- nvidia-smi
55
- # Run all e2e fusion tests
56
- - pytest -v -s tests/compile/distributed/test_fusions_e2e.py
178
+ # Run just llama3 (fp8 & bf16) for all config combinations
179
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3"
180
181
+- label: Fusion E2E TP2 (B200)
182
+ timeout_in_minutes: 20
183
+ working_dir: "/vllm-workspace/"
184
+ device: b200
185
+ num_devices: 2
186
+ source_file_dependencies:
187
+ - csrc/quantization/
188
+ - vllm/model_executor/
189
+ - vllm/v1/attention/
190
+ - vllm/compilation/
191
+ - tests/compile/fusions_e2e/
192
+ commands:
193
+ - nvidia-smi
194
+ # Run all models and attn backends but only Inductor partition and native custom ops
195
+ # for ar-rms-quant-fp4, also sweep llama3
196
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and not +quant_fp8" -k "Llama-3.1-8B-Instruct-FP4"
197
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
.buildkite/test_areas/distributed.yaml
+2
-57
@@ -16,7 +16,7 @@ steps:
16
- pytest -v -s distributed/test_shm_storage.py
17
18
- label: Distributed (2 GPUs)
19
- timeout_in_minutes: 90
19
+ timeout_in_minutes: 60
20
working_dir: "/vllm-workspace/tests"
21
num_devices: 2
22
source_file_dependencies:
@@ -47,7 +47,6 @@ steps:
47
- pytest -v -s ./compile/test_wrapper.py
48
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
49
- 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'
50
- - pytest -v -s distributed/test_sequence_parallel.py
50
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
51
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
52
@@ -133,25 +132,13 @@ steps:
132
- TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
133
- pytest -v -s -x lora/test_mixtral.py
134
136
-- label: Sequence Parallel Tests (H100)
137
- timeout_in_minutes: 60
138
- working_dir: "/vllm-workspace/"
139
- device: h100
140
- optional: true
141
- num_devices: 2
142
- commands:
143
- - export VLLM_TEST_CLEAN_GPU_MEMORY=1
144
- # Run sequence parallel tests
145
- - pytest -v -s tests/distributed/test_sequence_parallel.py
146
- - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
147
-
135
- label: Distributed Tests (2 GPUs)(H100)
136
+ timeout_in_minutes: 15
137
device: h100
138
optional: true
139
working_dir: "/vllm-workspace/"
140
num_devices: 2
141
commands:
154
- - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py
142
- pytest -v -s tests/distributed/test_context_parallel.py
143
- 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
144
- pytest -v -s tests/v1/distributed/test_dbo.py
@@ -217,45 +204,3 @@ steps:
204
commands:
205
- pytest -v -s distributed/test_pp_cudagraph.py
206
- pytest -v -s distributed/test_pipeline_parallel.py
220
-
221
-- label: Hopper Fusion E2E Tests (H100)
222
- timeout_in_minutes: 70
223
- working_dir: "/vllm-workspace/"
224
- device: h100
225
- optional: true
226
- source_file_dependencies:
227
- - csrc/quantization/fp4/
228
- - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
229
- - vllm/v1/attention/backends/flashinfer.py
230
- - vllm/compilation/
231
- # can affect pattern matching
232
- - vllm/model_executor/layers/layernorm.py
233
- - vllm/model_executor/layers/activation.py
234
- - vllm/model_executor/layers/quantization/input_quant_fp8.py
235
- - tests/compile/test_fusion_attn.py
236
- commands:
237
- - export VLLM_TEST_CLEAN_GPU_MEMORY=1
238
- # skip Llama-4 since it does not fit on this device
239
- - pytest -v -s tests/compile/test_fusion_attn.py -k 'not Llama-4'
240
-
241
-- label: Hopper Fusion Distributed E2E Tests (2xH100)
242
- timeout_in_minutes: 70
243
- working_dir: "/vllm-workspace/"
244
- device: h100
245
- optional: true
246
- num_devices: 2
247
- source_file_dependencies:
248
- - csrc/quantization/fp4/
249
- - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
250
- - vllm/v1/attention/backends/flashinfer.py
251
- - vllm/compilation/
252
- # can affect pattern matching
253
- - vllm/model_executor/layers/layernorm.py
254
- - vllm/model_executor/layers/activation.py
255
- - vllm/model_executor/layers/quantization/input_quant_fp8.py
256
- - tests/compile/distributed/test_fusions_e2e.py
257
- commands:
258
- - export VLLM_TEST_CLEAN_GPU_MEMORY=1
259
- # Run all e2e fusion tests
260
- - pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'
261
- - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
.buildkite/test_areas/pytorch.yaml
+2
-5
@@ -18,7 +18,7 @@ steps:
18
- "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
19
20
- label: PyTorch Fullgraph Smoke Test
21
- timeout_in_minutes: 30
21
+ timeout_in_minutes: 35
22
source_file_dependencies:
23
- vllm/
24
- tests/compile
@@ -30,16 +30,13 @@ steps:
30
- "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\;"
31
32
- label: PyTorch Fullgraph
33
- timeout_in_minutes: 40
33
+ timeout_in_minutes: 30
34
source_file_dependencies:
35
- vllm/
36
- tests/compile
37
commands:
38
# fp8 kv scales not supported on sm89, tested on Blackwell instead
39
- pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
40
- # Limit to no custom ops to reduce running time
41
- # Wrap with quotes to escape yaml and avoid starting -k string with a -
42
- - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
40
41
- label: Pytorch Nightly Dependency Override Check # 2min
42
# if this test fails, it means the nightly torch version is not compatible with some
tests/compile/distributed/test_fusions_e2e.py
deleted
-321
@@ -1,321 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from __future__ import annotations
5
-
6
-import logging
7
-from typing import Any
8
-
9
-import pytest
10
-import regex as re
11
-
12
-from tests.compile.fusion_test_utils import (
13
- CUSTOM_OPS_FP8,
14
- CUSTOM_OPS_QUANT_RMS_NORM,
15
- CUSTOM_OPS_RMS_NORM,
16
- MODELS,
17
- MODELS_FP4,
18
- MODELS_FP8,
19
- MODELS_GROUP_FP8,
20
- Matches,
21
- custom_ops_product,
22
- is_blackwell,
23
- run_model,
24
-)
25
-from tests.v1.attention.utils import AttentionBackendEnum
26
-from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig
27
-from vllm.platforms import current_platform
28
-from vllm.utils.flashinfer import has_flashinfer
29
-from vllm.utils.torch_utils import is_torch_equal_or_newer
30
-
31
-from ...utils import flat_product, multi_gpu_test
32
-
33
-
34
-@multi_gpu_test(num_gpus=2)
35
-@pytest.mark.parametrize(
36
- "model_name, model_kwargs, backend, matches, custom_ops",
37
- # Toggle RMSNorm and QuantFP8 for FP8 models
38
- list(
39
- flat_product(
40
- MODELS_FP8, custom_ops_product(CUSTOM_OPS_FP8, CUSTOM_OPS_RMS_NORM)
41
- )
42
- )
43
- # Toggle RMSNorm for FP4 models and unquant models
44
- + list(flat_product(MODELS_FP4 + MODELS, CUSTOM_OPS_RMS_NORM)),
45
-)
46
-@pytest.mark.parametrize("inductor_graph_partition", [True, False])
47
-@pytest.mark.skipif(
48
- not current_platform.is_cuda()
49
- or not has_flashinfer()
50
- or not current_platform.has_device_capability(90),
51
- reason="allreduce+rmsnorm fusion requires flashinfer",
52
-)
53
-def test_tp2_attn_quant_allreduce_rmsnorm(
54
- model_name: str,
55
- model_kwargs: dict,
56
- backend: AttentionBackendEnum,
57
- matches: Matches,
58
- custom_ops: str,
59
- inductor_graph_partition: bool,
60
- caplog_mp_spawn,
61
- monkeypatch,
62
-):
63
- if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
64
- pytest.skip("Inductor graph partition requires torch>=2.9")
65
-
66
- if "fp4" in model_name.lower() and not is_blackwell():
67
- pytest.skip("NVFP4 quant requires Blackwell")
68
-
69
- if backend == AttentionBackendEnum.FLASHINFER and not is_blackwell():
70
- # FlashInfer attn fusion requires Blackwell
71
- matches = matches._replace(attention_fusion=0)
72
-
73
- custom_ops_list = custom_ops.split(",") if custom_ops else []
74
-
75
- if inductor_graph_partition:
76
- mode = CUDAGraphMode.FULL_AND_PIECEWISE
77
- splitting_ops: list[str] | None = None
78
- else:
79
- mode = CUDAGraphMode.FULL_DECODE_ONLY
80
- splitting_ops = []
81
-
82
- # Disable, compile cache to make sure custom passes run.
83
- # Otherwise, we can't verify fusion happened through the logs.
84
- monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
85
-
86
- # To capture subprocess logs, we need to know whether spawn or fork is used.
87
- # Force spawn as it is more general.
88
- monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
89
-
90
- model_kwargs["attention_config"] = {"backend": backend.name}
91
-
92
- compilation_config = CompilationConfig(
93
- # Testing properties
94
- use_inductor_graph_partition=inductor_graph_partition,
95
- cudagraph_mode=mode,
96
- custom_ops=custom_ops_list,
97
- splitting_ops=splitting_ops,
98
- # Common
99
- mode=CompilationMode.VLLM_COMPILE,
100
- pass_config=PassConfig(
101
- fuse_attn_quant=True,
102
- eliminate_noops=True,
103
- fuse_allreduce_rms=True,
104
- ),
105
- # Inductor caches custom passes by default as well via uuid
106
- inductor_compile_config={"force_disable_caches": True},
107
- )
108
-
109
- with caplog_mp_spawn(logging.DEBUG) as log_holder:
110
- run_model(
111
- compilation_config, model_name, tensor_parallel_size=2, **model_kwargs
112
- )
113
- log_matches = re.findall(
114
- r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes",
115
- log_holder.text,
116
- )
117
- # 2 for each compile range
118
- # (global compile range can be split due to fuse_allreduce_rmsnorm)
119
- num_compile_ranges = len(compilation_config.get_compile_ranges())
120
- assert num_compile_ranges in [1, 2]
121
-
122
- assert len(log_matches) == 2 * num_compile_ranges, log_holder.text
123
-
124
- assert all(int(log_match) == matches.attention_fusion for log_match in log_matches)
125
-
126
- log_matches = re.findall(
127
- r"collective_fusion.py:\d+] Replaced (\d+) patterns",
128
- log_holder.text,
129
- )
130
- assert len(log_matches) == 2, log_holder.text
131
-
132
- assert int(log_matches[0]) == matches.allreduce_fusion
133
- assert int(log_matches[1]) == matches.allreduce_fusion
134
-
135
- log_matches = re.findall(
136
- r"pass_manager.py:\d+] Skipping .*AllReduceFusionPass.* with compile range",
137
- log_holder.text,
138
- )
139
- assert len(log_matches) == 2 * (num_compile_ranges - 1), log_holder.text
140
-
141
-
142
-@multi_gpu_test(num_gpus=2)
143
-@pytest.mark.parametrize(
144
- "model_name, model_kwargs, backend, matches, custom_ops",
145
- # Toggle RMSNorm and QuantFP8 for FP8 models
146
- list(
147
- flat_product(
148
- MODELS_FP8, custom_ops_product(CUSTOM_OPS_FP8, CUSTOM_OPS_RMS_NORM)
149
- )
150
- )
151
- # Toggle RMSNorm for FP4 models and unquant models
152
- + list(flat_product(MODELS_FP4 + MODELS, CUSTOM_OPS_RMS_NORM)),
153
-)
154
-@pytest.mark.parametrize("inductor_graph_partition", [True, False])
155
-@pytest.mark.skipif(
156
- not current_platform.is_cuda(),
157
- reason="sequence parallel only tested on CUDA",
158
-)
159
-def test_tp2_attn_quant_async_tp(
160
- model_name: str,
161
- model_kwargs: dict,
162
- backend: AttentionBackendEnum,
163
- matches: Matches,
164
- custom_ops: str,
165
- inductor_graph_partition: bool,
166
- caplog_mp_spawn,
167
- monkeypatch,
168
-):
169
- if is_blackwell():
170
- # TODO: https://github.com/vllm-project/vllm/issues/27893
171
- pytest.skip("Blackwell is not supported for AsyncTP pass")
172
-
173
- if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
174
- pytest.skip("Inductor graph partition requires torch>=2.9")
175
-
176
- if "fp4" in model_name.lower() and not is_blackwell():
177
- pytest.skip("NVFP4 quant requires Blackwell")
178
-
179
- if backend == AttentionBackendEnum.FLASHINFER:
180
- if not has_flashinfer():
181
- pytest.skip("FlashInfer backend requires flashinfer installed")
182
- if not is_blackwell():
183
- # FlashInfer attn fusion requires Blackwell
184
- matches = matches._replace(attention_fusion=0)
185
-
186
- custom_ops_list = custom_ops.split(",") if custom_ops else []
187
-
188
- if inductor_graph_partition:
189
- mode = CUDAGraphMode.FULL_AND_PIECEWISE
190
- splitting_ops: list[str] | None = None
191
- else:
192
- mode = CUDAGraphMode.FULL_DECODE_ONLY
193
- splitting_ops = []
194
-
195
- # Disable, compile cache to make sure custom passes run.
196
- # Otherwise, we can't verify fusion happened through the logs.
197
- monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
198
-
199
- # To capture subprocess logs, we need to know whether spawn or fork is used.
200
- # Force spawn as it is more general.
201
- monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
202
-
203
- model_kwargs["attention_config"] = {"backend": backend.name}
204
-
205
- compilation_config = CompilationConfig(
206
- # Testing properties
207
- use_inductor_graph_partition=inductor_graph_partition,
208
- cudagraph_mode=mode,
209
- custom_ops=custom_ops_list,
210
- splitting_ops=splitting_ops,
211
- # Common
212
- mode=CompilationMode.VLLM_COMPILE,
213
- pass_config=PassConfig(
214
- fuse_attn_quant=True,
215
- eliminate_noops=True,
216
- enable_sp=True,
217
- fuse_gemm_comms=True,
218
- ),
219
- # Inductor caches custom passes by default as well via uuid
220
- inductor_compile_config={"force_disable_caches": True},
221
- )
222
-
223
- with caplog_mp_spawn(logging.DEBUG) as log_holder:
224
- run_model(
225
- compilation_config, model_name, tensor_parallel_size=2, **model_kwargs
226
- )
227
- log_matches = re.findall(
228
- r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes",
229
- log_holder.text,
230
- )
231
- assert len(log_matches) == 2, log_holder.text
232
-
233
- assert int(log_matches[0]) == matches.attention_fusion
234
- assert int(log_matches[1]) == matches.attention_fusion
235
-
236
- log_matches = re.findall(
237
- r"sequence_parallelism.py:\d+] Replaced (\d+) patterns",
238
- log_holder.text,
239
- )
240
- assert len(log_matches) == 2, log_holder.text
241
-
242
- assert int(log_matches[0]) == matches.sequence_parallel
243
- assert int(log_matches[1]) == matches.sequence_parallel
244
-
245
- log_matches = re.findall(
246
- r"collective_fusion.py:\d+] Replaced (\d+) patterns",
247
- log_holder.text,
248
- )
249
- assert len(log_matches) == 2, log_holder.text
250
-
251
- assert int(log_matches[0]) == matches.async_tp
252
- assert int(log_matches[1]) == matches.async_tp
253
-
254
-
255
-@pytest.mark.parametrize(
256
- "model_name, model_kwargs, backend, matches, custom_ops",
257
- # Test rms norm+group quant_fp8 fusion
258
- list[tuple[Any, ...]](flat_product(MODELS_GROUP_FP8, CUSTOM_OPS_QUANT_RMS_NORM)),
259
-)
260
-@pytest.mark.parametrize("inductor_graph_partition", [True, False])
261
-# TODO: remove skip after we fix the fusion thoroughly
262
-@pytest.mark.skipif(is_blackwell(), reason="Temporarily disabled on Blackwell")
263
-def test_rms_group_quant(
264
- model_name: str,
265
- model_kwargs: dict[str, Any],
266
- backend: AttentionBackendEnum,
267
- matches: Matches,
268
- custom_ops: str,
269
- inductor_graph_partition: bool,
270
- caplog_mp_spawn,
271
- monkeypatch,
272
-):
273
- if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
274
- pytest.skip("Inductor graph partition requires torch>=2.9")
275
-
276
- custom_ops_list = custom_ops.split(",") if custom_ops else []
277
-
278
- if inductor_graph_partition:
279
- mode = CUDAGraphMode.FULL_AND_PIECEWISE
280
- splitting_ops: list[str] | None = None
281
- else:
282
- mode = CUDAGraphMode.FULL_DECODE_ONLY
283
- splitting_ops = []
284
-
285
- # Disable, compile cache to make sure custom passes run.
286
- # Otherwise, we can't verify fusion happened through the logs.
287
- monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
288
-
289
- # To capture subprocess logs, we need to know whether spawn or fork is used.
290
- # Force spawn as it is more general.
291
- monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
292
-
293
- # TODO: remove this after fusion is fixed
294
- monkeypatch.setenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "0")
295
-
296
- model_kwargs["attention_config"] = {"backend": backend.name}
297
-
298
- compilation_config = CompilationConfig(
299
- # Testing properties
300
- custom_ops=custom_ops_list,
301
- use_inductor_graph_partition=inductor_graph_partition,
302
- cudagraph_mode=mode,
303
- splitting_ops=splitting_ops,
304
- # Common
305
- mode=CompilationMode.VLLM_COMPILE,
306
- pass_config=PassConfig(
307
- fuse_norm_quant=True, fuse_act_quant=True, eliminate_noops=True
308
- ),
309
- # Inductor caches custom passes by default as well via uuid
310
- inductor_compile_config={"force_disable_caches": True},
311
- )
312
-
313
- with caplog_mp_spawn(logging.DEBUG) as log_holder:
314
- run_model(compilation_config, model_name, **model_kwargs)
315
-
316
- log_matches = re.findall(
317
- r"\[fusion.py:\d+] Replaced (\d+) patterns",
318
- log_holder.text,
319
- )
320
- assert len(log_matches) == 1, log_holder.text
321
- assert int(log_matches[0]) == matches.rms_quant_norm_fusion
tests/compile/fusion_test_utils.py
deleted
-208
@@ -1,208 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Shared utilities for fusion tests (e.g. test_fusion_attn.py)."""
4
-
5
-from __future__ import annotations
6
-
7
-import itertools
8
-from collections.abc import Iterable
9
-from typing import Any, NamedTuple
10
-
11
-from tests.v1.attention.utils import AttentionBackendEnum
12
-from vllm import LLM, SamplingParams
13
-from vllm.config import CompilationConfig, CUDAGraphMode
14
-from vllm.platforms import current_platform
15
-
16
-is_blackwell = lambda: current_platform.is_device_capability_family(100)
17
-"""Are we running on Blackwell, a lot of tests depend on it"""
18
-
19
-
20
-def has_cuda_graph_wrapper_metadata() -> bool:
21
- from importlib import import_module
22
-
23
- try:
24
- module = import_module("torch._inductor.utils")
25
- module.CUDAGraphWrapperMetadata # noqa B018
26
- except AttributeError:
27
- return False
28
- return True
29
-
30
-
31
-class Matches(NamedTuple):
32
- attention_fusion: int = 0
33
- allreduce_fusion: int = 0
34
- sequence_parallel: int = 0
35
- async_tp: int = 0
36
- rms_quant_norm_fusion: int = 0
37
-
38
-
39
-class ModelBackendTestCase(NamedTuple):
40
- model_name: str
41
- model_kwargs: dict[str, Any]
42
- backend: AttentionBackendEnum
43
- matches: Matches
44
-
45
-
46
-# E2E model test cases
47
-MODELS_FP8: list[ModelBackendTestCase] = []
48
-MODELS_FP4: list[ModelBackendTestCase] = []
49
-MODELS: list[ModelBackendTestCase] = [] # tp-only (unquantized)
50
-MODELS_GROUP_FP8: list[ModelBackendTestCase] = []
51
-
52
-if current_platform.is_cuda():
53
- MODELS_FP8 = [
54
- ModelBackendTestCase(
55
- # Use smaller model for L40s in CI
56
- model_name="RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
57
- model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
58
- backend=AttentionBackendEnum.TRITON_ATTN,
59
- matches=Matches(
60
- attention_fusion=32,
61
- allreduce_fusion=65,
62
- sequence_parallel=65,
63
- async_tp=128,
64
- ),
65
- ),
66
- ModelBackendTestCase(
67
- model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
68
- model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
69
- # TODO FlashInfer attn broken on Hopper with kvcache=fp8:
70
- # https://github.com/vllm-project/vllm/issues/28568
71
- backend=AttentionBackendEnum.FLASHINFER
72
- if is_blackwell()
73
- else AttentionBackendEnum.TRITON_ATTN,
74
- matches=Matches(
75
- attention_fusion=48,
76
- allreduce_fusion=96,
77
- sequence_parallel=96,
78
- async_tp=95, # mlp is moe, no fusion there
79
- ),
80
- ),
81
- ]
82
-
83
- MODELS_FP4 = [
84
- ModelBackendTestCase(
85
- model_name="nvidia/Llama-3.1-8B-Instruct-FP4",
86
- model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
87
- backend=AttentionBackendEnum.FLASHINFER,
88
- matches=Matches(
89
- attention_fusion=32,
90
- allreduce_fusion=65,
91
- sequence_parallel=65,
92
- async_tp=128,
93
- ),
94
- ),
95
- ]
96
-
97
- # TP only (unquantized models)
98
- MODELS = [
99
- ModelBackendTestCase(
100
- model_name="meta-llama/Llama-3.1-8B-Instruct",
101
- model_kwargs=dict(max_model_len=1024),
102
- backend=AttentionBackendEnum.TRITON_ATTN,
103
- matches=Matches(
104
- attention_fusion=0,
105
- allreduce_fusion=65,
106
- sequence_parallel=65,
107
- async_tp=128,
108
- ),
109
- ),
110
- ModelBackendTestCase(
111
- model_name="Qwen/Qwen3-30B-A3B",
112
- model_kwargs=dict(max_model_len=1024),
113
- backend=AttentionBackendEnum.TRITON_ATTN,
114
- matches=Matches(
115
- attention_fusion=0,
116
- allreduce_fusion=97,
117
- sequence_parallel=97,
118
- async_tp=96, # MLP is MoE, half the fusions of dense
119
- ),
120
- ),
121
- ]
122
-
123
- MODELS_GROUP_FP8 = [
124
- ModelBackendTestCase(
125
- model_name="Qwen/Qwen3-30B-A3B-FP8",
126
- model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
127
- backend=AttentionBackendEnum.TRITON_ATTN,
128
- matches=Matches(
129
- rms_quant_norm_fusion=48,
130
- ),
131
- ),
132
- ]
133
-
134
-elif current_platform.is_rocm():
135
- MODELS_FP8 = [
136
- ModelBackendTestCase(
137
- model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
138
- model_kwargs=dict(max_model_len=1024),
139
- backend=AttentionBackendEnum.TRITON_ATTN,
140
- matches=Matches(attention_fusion=32),
141
- ),
142
- ModelBackendTestCase(
143
- model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
144
- model_kwargs=dict(max_model_len=1024),
145
- backend=AttentionBackendEnum.ROCM_ATTN,
146
- matches=Matches(attention_fusion=32),
147
- ),
148
- ModelBackendTestCase(
149
- model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
150
- model_kwargs=dict(max_model_len=1024),
151
- backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
152
- matches=Matches(attention_fusion=32),
153
- ),
154
- ]
155
-
156
-
157
-# Custom ops toggle lists for parametrization
158
-CUSTOM_OPS_FP8 = ["-quant_fp8", "+quant_fp8"]
159
-CUSTOM_OPS_RMS_NORM = ["-rms_norm", "+rms_norm"]
160
-CUSTOM_OPS_QUANT_RMS_NORM = ["+quant_fp8,+rms_norm"]
161
-
162
-
163
-def custom_ops_product(*custom_ops_lists: list[str]) -> Iterable[str]:
164
- """Generate all combinations of custom ops for parametrization."""
165
- for op_list in itertools.product(*custom_ops_lists):
166
- yield ",".join(op_list)
167
-
168
-
169
-def run_model(compile_config: int | CompilationConfig, model: str, **model_kwargs):
170
- """Run a model with the given compilation config for E2E fusion tests."""
171
- compilation_config = (
172
- compile_config
173
- if isinstance(compile_config, CompilationConfig)
174
- else CompilationConfig(mode=compile_config)
175
- )
176
-
177
- prompts = [
178
- "Hello, my name is",
179
- "The president of the United States is",
180
- "The capital of France is",
181
- "The future of AI is",
182
- ]
183
- sampling_params = SamplingParams(temperature=0)
184
- # Allow override from model_kwargs
185
- model_kwargs = {"tensor_parallel_size": 1, **model_kwargs}
186
- model_kwargs = {"disable_custom_all_reduce": True, **model_kwargs}
187
-
188
- # No cudagraphs by default
189
- if compilation_config.cudagraph_mode is None:
190
- compilation_config.cudagraph_mode = CUDAGraphMode.NONE
191
- llm = LLM(
192
- model=model,
193
- compilation_config=compilation_config,
194
- **model_kwargs,
195
- )
196
- outputs = llm.generate(prompts, sampling_params)
197
-
198
- # Print the outputs.
199
- for output in outputs:
200
- prompt = output.prompt
201
- generated_text = output.outputs[0].text
202
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
203
-
204
- # Get the compile ranges split points after vllm config post init
205
- # in order to compute compile ranges correctly
206
- compilation_config.compile_ranges_split_points = (
207
- llm.llm_engine.vllm_config.compilation_config.compile_ranges_split_points
208
- )
tests/compile/fusions_e2e/__init__.py
tests/compile/fusions_e2e/common.py
new
+102
@@ -0,0 +1,102 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import itertools
4
+from collections.abc import Callable, Iterable
5
+from typing import Any, NamedTuple
6
+
7
+import pytest
8
+import regex as re
9
+
10
+from vllm.platforms import current_platform
11
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
12
+
13
+
14
+class Matches(NamedTuple):
15
+ # simple pointwise
16
+ rms_quant_fusion: int = 0
17
+ act_quant_fusion: int = 0
18
+ norm_rope_fusion: int = 0
19
+ attn_quant_fusion: int = 0
20
+ # distributed
21
+ ar_rms_fusion: int = 0
22
+ sequence_parallel: int = 0
23
+ async_tp: int = 0
24
+
25
+
26
+class ModelFusionInfo(NamedTuple):
27
+ model_name: str
28
+ matches: Callable[[int], Matches]
29
+ """Given number of hidden layers, produces the matches object"""
30
+ model_kwargs: dict[str, Any] = {}
31
+ hf_overrides: Callable[[int], dict] = lambda n: {"num_hidden_layers": n}
32
+
33
+
34
+class AttentionBackendCase(NamedTuple):
35
+ backend: AttentionBackendEnum
36
+ model_kwargs: dict[str, Any] = {}
37
+ """Additional args required for attn+quant fusion"""
38
+
39
+
40
+is_blackwell = lambda: current_platform.is_device_capability_family(100)
41
+"""Are we running on Blackwell, a lot of tests depend on it"""
42
+
43
+
44
+def custom_ops_combos(*custom_ops: str) -> Iterable[str]:
45
+ """Generate all combinations of custom ops for parametrization."""
46
+ custom_ops_lists = [[f"-{op}", f"+{op}"] for op in custom_ops]
47
+ for op_list in itertools.product(*custom_ops_lists):
48
+ yield ",".join(op_list)
49
+
50
+
51
+# Quick inline validation
52
+assert list(custom_ops_combos("silu_and_mul")) == ["-silu_and_mul", "+silu_and_mul"]
53
+assert list(custom_ops_combos("quant_fp8", "rms_norm")) == [
54
+ "-quant_fp8,-rms_norm",
55
+ "-quant_fp8,+rms_norm",
56
+ "+quant_fp8,-rms_norm",
57
+ "+quant_fp8,+rms_norm",
58
+]
59
+
60
+
61
+def has_cuda_graph_wrapper_metadata() -> bool:
62
+ from importlib import import_module
63
+
64
+ try:
65
+ module = import_module("torch._inductor.utils")
66
+ module.CUDAGraphWrapperMetadata # noqa B018
67
+ except AttributeError:
68
+ return False
69
+ return True
70
+
71
+
72
+INDUCTOR_GRAPH_PARTITION = [
73
+ pytest.param(
74
+ True,
75
+ marks=pytest.mark.skipif(
76
+ not has_cuda_graph_wrapper_metadata(),
77
+ reason="torch version does not support Inductor partition",
78
+ ),
79
+ id="inductor_partition",
80
+ ),
81
+ pytest.param(False, id="dynamo_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
+ ),
91
+ "norm_rope_fusion": re.compile(
92
+ r"qk_norm_rope_fusion.py:\d+] Fused QK Norm\+RoPE on (\d+) sites"
93
+ ),
94
+ "attn_quant_fusion": re.compile(
95
+ r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes"
96
+ ),
97
+ "ar_rms_fusion": re.compile(r"collective_fusion.py:\d+] Replaced (\d+) patterns"),
98
+ "sequence_parallel": re.compile(
99
+ r"sequence_parallelism.py:\d+] Replaced (\d+) patterns"
100
+ ),
101
+ "async_tp": re.compile(r"collective_fusion.py:\d+] Replaced (\d+) patterns"),
102
+}
tests/compile/fusions_e2e/conftest.py
new
+158
@@ -0,0 +1,158 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import logging
4
+
5
+import pytest
6
+import regex as re
7
+
8
+from vllm import LLM, SamplingParams
9
+from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
10
+
11
+from .common import FUSION_LOG_PATTERNS, AttentionBackendCase, Matches
12
+
13
+
14
+def run_model(compile_config: int | CompilationConfig, model: str, **model_kwargs):
15
+ """Run a model with the given compilation config for E2E fusion tests."""
16
+ compilation_config = (
17
+ compile_config
18
+ if isinstance(compile_config, CompilationConfig)
19
+ else CompilationConfig(mode=compile_config)
20
+ )
21
+
22
+ prompts = [
23
+ "Hello, my name is",
24
+ "The president of the United States is",
25
+ "The capital of France is",
26
+ "The future of AI is",
27
+ ]
28
+ sampling_params = SamplingParams(temperature=0)
29
+ # Allow override from model_kwargs
30
+ model_kwargs = {"tensor_parallel_size": 1, **model_kwargs}
31
+ model_kwargs = {"disable_custom_all_reduce": True, **model_kwargs}
32
+
33
+ # No cudagraphs by default
34
+ if compilation_config.cudagraph_mode is None:
35
+ compilation_config.cudagraph_mode = CUDAGraphMode.NONE
36
+ llm = LLM(
37
+ model=model,
38
+ compilation_config=compilation_config,
39
+ **model_kwargs,
40
+ )
41
+ outputs = llm.generate(prompts, sampling_params)
42
+
43
+ # Print the outputs.
44
+ for output in outputs:
45
+ prompt = output.prompt
46
+ generated_text = output.outputs[0].text
47
+ print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
48
+
49
+ # Get the compile ranges split points after vllm config post init
50
+ # in order to compute compile ranges correctly
51
+ compilation_config.compile_ranges_split_points = (
52
+ llm.llm_engine.vllm_config.compilation_config.compile_ranges_split_points
53
+ )
54
+
55
+
56
+@pytest.fixture
57
+def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
58
+ def run(
59
+ model_name: str,
60
+ matches: Matches,
61
+ model_kwargs: dict,
62
+ attn_backend: AttentionBackendCase,
63
+ compilation_config: dict,
64
+ matches_check: list[str],
65
+ use_deepgemm: bool = False,
66
+ tp_size: int = 1,
67
+ ):
68
+ monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1" if use_deepgemm else "0")
69
+
70
+ # Disable, compile cache to make sure custom passes run.
71
+ # Otherwise, we can't verify fusion happened through the logs.
72
+ monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
73
+
74
+ # To capture subprocess logs, we need to know whether spawn or fork is used.
75
+ # Force spawn as it is more general.
76
+ monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
77
+
78
+ model_kwargs = {**attn_backend.model_kwargs, **model_kwargs}
79
+ model_kwargs["attention_config"] = {"backend": attn_backend.backend.name}
80
+ model_kwargs["tensor_parallel_size"] = tp_size
81
+
82
+ # Always compile the full graph instead of piecewise
83
+ if not compilation_config["use_inductor_graph_partition"]:
84
+ compilation_config["splitting_ops"] = []
85
+
86
+ full_compilation_config = CompilationConfig(
87
+ cudagraph_mode=CUDAGraphMode.NONE,
88
+ mode=CompilationMode.VLLM_COMPILE,
89
+ inductor_compile_config={"force_disable_caches": True},
90
+ **compilation_config,
91
+ )
92
+
93
+ with caplog_mp_spawn(logging.DEBUG) as log_holder:
94
+ run_model(full_compilation_config, model_name, **model_kwargs)
95
+
96
+ num_compile_ranges = len(full_compilation_config.get_compile_ranges())
97
+ assert num_compile_ranges in [1, 2]
98
+
99
+ print(f"Compile ranges: {full_compilation_config.get_compile_ranges()}")
100
+ print("Fusion results:")
101
+
102
+ # Iterate through all so printing happens before asserting
103
+ log_matches_dict = {}
104
+ for match_name, pattern in FUSION_LOG_PATTERNS.items():
105
+ log_matches_dict[match_name] = list(pattern.findall(log_holder.text))
106
+ print(f"- {match_name}={','.join(log_matches_dict[match_name])}")
107
+
108
+ # Now check the matches
109
+ for match_name in matches_check:
110
+ num_ranges_activated = (
111
+ 1 if match_name == "ar_rms_fusion" else num_compile_ranges
112
+ )
113
+ n_expected = tp_size * num_ranges_activated
114
+
115
+ log_matches = list(int(ms) for ms in log_matches_dict[match_name])
116
+ assert len(log_matches) == n_expected, (
117
+ f"Could not find {n_expected} {match_name} "
118
+ f"(found {len(log_matches)}) in:\n {log_holder.text}"
119
+ )
120
+
121
+ expected_matches = getattr(matches, match_name)
122
+
123
+ if match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check:
124
+ # AR+rms+quant takes precedence over rms+quant if activated.
125
+ # That means we get full matching where ar+rms+quant was not activated,
126
+ # and less where it was
127
+ assert sum(m == expected_matches for m in log_matches) == tp_size * (
128
+ num_ranges_activated - 1
129
+ ), "Expecting full rms+quant fusion where ar+rms+quant not activated"
130
+
131
+ assert all(
132
+ expected_matches - matches.ar_rms_fusion <= m <= expected_matches
133
+ for m in log_matches
134
+ ), (
135
+ f"Expecting at least {expected_matches - matches.ar_rms_fusion} "
136
+ f"where ar+rms+quant was activated"
137
+ )
138
+ else:
139
+ expected_matches_list = [expected_matches] * n_expected
140
+ assert sorted(log_matches) == expected_matches_list, (
141
+ f"{match_name} expected: {expected_matches_list}, "
142
+ f"found: {sorted(log_matches)}"
143
+ )
144
+
145
+ if match_name == "ar_rms_fusion":
146
+ log_matches = re.findall(
147
+ r"pass_manager.py:\d+] Skipping "
148
+ r".*AllReduceFusionPass.* with compile range",
149
+ log_holder.text,
150
+ )
151
+
152
+ n_expected = tp_size * (num_compile_ranges - num_ranges_activated)
153
+ assert len(log_matches) == n_expected, (
154
+ f'Could not find {n_expected} "Skipping AllReduceFusionPass" '
155
+ f"(found {len(log_matches)}) in:\n {log_holder.text}"
156
+ )
157
+
158
+ return run
tests/compile/fusions_e2e/models.py
new
+112
@@ -0,0 +1,112 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+
5
+from vllm.utils.flashinfer import has_flashinfer
6
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
7
+
8
+from .common import AttentionBackendCase, Matches, ModelFusionInfo, is_blackwell
9
+
10
+# Attn backends
11
+FLASHINFER_ATTN = pytest.param(
12
+ AttentionBackendCase(
13
+ backend=AttentionBackendEnum.FLASHINFER,
14
+ model_kwargs=dict(kv_cache_dtype="fp8"),
15
+ ),
16
+ id="FLASHINFER",
17
+ marks=pytest.mark.skipif(
18
+ not is_blackwell() or not has_flashinfer(),
19
+ reason="FI backend requires Blackwell and FlashInfer",
20
+ ),
21
+)
22
+
23
+TRITON_ATTN = pytest.param(
24
+ AttentionBackendCase(backend=AttentionBackendEnum.TRITON_ATTN), id="TRITON_ATTN"
25
+)
26
+
27
+# Models
28
+llama3_8b = ModelFusionInfo(
29
+ model_name="meta-llama/Llama-3.1-8B-Instruct",
30
+ matches=lambda n_layers: Matches(
31
+ ar_rms_fusion=n_layers * 2 + 1,
32
+ sequence_parallel=n_layers * 2 + 1,
33
+ async_tp=n_layers * 4,
34
+ ),
35
+)
36
+
37
+llama3_8b_fp8 = ModelFusionInfo(
38
+ model_name="RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
39
+ matches=lambda n_layers: Matches(
40
+ rms_quant_fusion=n_layers * 2,
41
+ act_quant_fusion=n_layers,
42
+ attn_quant_fusion=n_layers,
43
+ ar_rms_fusion=n_layers * 2 + 1,
44
+ sequence_parallel=n_layers * 2 + 1,
45
+ async_tp=n_layers * 4,
46
+ ),
47
+)
48
+
49
+llama3_8b_fp4 = ModelFusionInfo(
50
+ model_name="nvidia/Llama-3.1-8B-Instruct-FP4",
51
+ matches=lambda n_layers: Matches(
52
+ rms_quant_fusion=0,
53
+ act_quant_fusion=n_layers,
54
+ attn_quant_fusion=n_layers,
55
+ ar_rms_fusion=n_layers * 2 + 1,
56
+ sequence_parallel=n_layers * 2 + 1,
57
+ async_tp=n_layers * 4,
58
+ ),
59
+)
60
+
61
+# MoEs cannot do act+quant fusion because those ops are hidden from torch.compile.
62
+# MoEs also only expose 1 rms+quant fusion because the quant for up_proj is hidden.
63
+# TODO(luka): https://github.com/vllm-project/vllm/issues/31985
64
+# Also, for MoEs, gemm+collective fusion only happens for dense GEMMs (o_proj/qkv proj)
65
+
66
+llama4_scout_fp8 = ModelFusionInfo(
67
+ model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",
68
+ hf_overrides=lambda n_layers: {"text_config": {"num_hidden_layers": n_layers}},
69
+ matches=lambda n_layers: Matches(
70
+ rms_quant_fusion=n_layers,
71
+ attn_quant_fusion=n_layers,
72
+ ar_rms_fusion=n_layers * 2,
73
+ sequence_parallel=n_layers * 2,
74
+ async_tp=n_layers * 2 - 1,
75
+ ),
76
+)
77
+
78
+llama4_scout_fp4 = ModelFusionInfo(
79
+ model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-NVFP4",
80
+ hf_overrides=lambda n_layers: {"text_config": {"num_hidden_layers": n_layers}},
81
+ matches=lambda n_layers: Matches(
82
+ rms_quant_fusion=0,
83
+ attn_quant_fusion=n_layers,
84
+ ar_rms_fusion=n_layers * 2,
85
+ sequence_parallel=n_layers * 2,
86
+ async_tp=n_layers * 2 - 1,
87
+ ),
88
+)
89
+
90
+qwen3_a3b = ModelFusionInfo(
91
+ model_name="Qwen/Qwen3-30B-A3B",
92
+ matches=lambda n_layers: Matches(
93
+ norm_rope_fusion=n_layers,
94
+ ar_rms_fusion=n_layers * 2 + 1,
95
+ sequence_parallel=n_layers * 2 + 1,
96
+ async_tp=n_layers * 2,
97
+ ),
98
+)
99
+
100
+qwen3_a3b_fp8 = ModelFusionInfo(
101
+ model_name="Qwen/Qwen3-30B-A3B-FP8",
102
+ matches=lambda n_layers: Matches(
103
+ rms_quant_fusion=n_layers,
104
+ # TODO broken on Blackwell:
105
+ # https://github.com/vllm-project/vllm/issues/33295
106
+ norm_rope_fusion=0 if is_blackwell() else n_layers,
107
+ attn_quant_fusion=0, # attn + group quant not supported
108
+ ar_rms_fusion=n_layers * 2 + 1,
109
+ sequence_parallel=n_layers * 2 + 1,
110
+ async_tp=n_layers * 2,
111
+ ),
112
+)
tests/compile/fusions_e2e/test_tp1_quant.py
new
+146
@@ -0,0 +1,146 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import pytest
6
+
7
+from vllm.config import PassConfig
8
+
9
+from .common import (
10
+ INDUCTOR_GRAPH_PARTITION,
11
+ AttentionBackendCase,
12
+ Matches,
13
+ custom_ops_combos,
14
+ is_blackwell,
15
+)
16
+from .models import (
17
+ FLASHINFER_ATTN,
18
+ TRITON_ATTN,
19
+ llama3_8b_fp4,
20
+ llama3_8b_fp8,
21
+ llama4_scout_fp4,
22
+ llama4_scout_fp8,
23
+ qwen3_a3b_fp8,
24
+)
25
+
26
+
27
+@pytest.mark.parametrize(
28
+ "model_name, matches_fn, model_kwargs, hf_overrides, use_deepgemm",
29
+ [
30
+ (*llama3_8b_fp8, False),
31
+ (*llama4_scout_fp8, False),
32
+ (*qwen3_a3b_fp8, False),
33
+ (*qwen3_a3b_fp8, True),
34
+ ],
35
+)
36
+@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN])
37
+@pytest.mark.parametrize("n_layers", [6])
38
+@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
39
+@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
40
+def test_tp1_fp8_fusions(
41
+ model_name: str,
42
+ matches_fn: Callable[[int], Matches],
43
+ model_kwargs: dict,
44
+ hf_overrides: Callable[[int], dict],
45
+ attn_backend: AttentionBackendCase,
46
+ n_layers: int,
47
+ custom_ops: str,
48
+ inductor_graph_partition: bool,
49
+ use_deepgemm: bool,
50
+ run_e2e_fusion_test,
51
+ monkeypatch,
52
+):
53
+ if use_deepgemm:
54
+ # TODO(luka/eliza) DeepGEMM uses different quants, matching not supported
55
+ # - on Blackwell, uses a special quant fp8, currently not supported
56
+ # - on Hopper, tma-aligned scales inhibit matching (fix WIP)
57
+ pytest.skip("DeepGEMM & quant matching not currently supported")
58
+
59
+ matches = matches_fn(n_layers)
60
+
61
+ if "qwen" in model_name.lower() and "-quant_fp8" in custom_ops:
62
+ # This is why config forces +quant_fp8 by default
63
+ pytest.skip("native QuantFP8 matching not supported for group quant")
64
+
65
+ # Reduce size of model and skip weight loading time
66
+ model_kwargs["hf_overrides"] = hf_overrides(n_layers)
67
+ model_kwargs["load_format"] = "dummy"
68
+ model_kwargs["max_model_len"] = 1024
69
+
70
+ compilation_config = dict(
71
+ use_inductor_graph_partition=inductor_graph_partition,
72
+ custom_ops=custom_ops.split(","),
73
+ pass_config=PassConfig(
74
+ fuse_norm_quant=True,
75
+ fuse_act_quant=True,
76
+ fuse_attn_quant=True,
77
+ enable_qk_norm_rope_fusion=True,
78
+ ),
79
+ )
80
+
81
+ matches_check = [
82
+ "rms_quant_fusion",
83
+ "act_quant_fusion",
84
+ "norm_rope_fusion",
85
+ "attn_quant_fusion",
86
+ ]
87
+
88
+ run_e2e_fusion_test(
89
+ model_name,
90
+ matches,
91
+ model_kwargs,
92
+ attn_backend,
93
+ compilation_config,
94
+ matches_check,
95
+ use_deepgemm=use_deepgemm,
96
+ )
97
+
98
+
99
+@pytest.mark.parametrize(
100
+ "model_name, matches_fn, model_kwargs, hf_overrides",
101
+ [llama3_8b_fp4, llama4_scout_fp4],
102
+)
103
+@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN])
104
+@pytest.mark.parametrize("n_layers", [6])
105
+@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
106
+@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
107
+@pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4")
108
+def test_tp1_fp4_fusions(
109
+ model_name: str,
110
+ matches_fn: Callable[[int], Matches],
111
+ model_kwargs: dict,
112
+ hf_overrides: Callable[[int], dict],
113
+ attn_backend: AttentionBackendCase,
114
+ n_layers: int,
115
+ custom_ops: str,
116
+ inductor_graph_partition: bool,
117
+ run_e2e_fusion_test,
118
+):
119
+ matches = matches_fn(n_layers)
120
+
121
+ # Reduce size of model and skip weight loading time
122
+ model_kwargs["hf_overrides"] = hf_overrides(n_layers)
123
+ model_kwargs["load_format"] = "dummy"
124
+ model_kwargs["max_model_len"] = 1024
125
+
126
+ compilation_config = dict(
127
+ use_inductor_graph_partition=inductor_graph_partition,
128
+ custom_ops=custom_ops.split(","),
129
+ pass_config=PassConfig(
130
+ fuse_norm_quant=True,
131
+ fuse_act_quant=True,
132
+ fuse_attn_quant=True,
133
+ enable_qk_norm_rope_fusion=True,
134
+ ),
135
+ )
136
+
137
+ matches_check = ["act_quant_fusion", "attn_quant_fusion", "norm_rope_fusion"]
138
+
139
+ run_e2e_fusion_test(
140
+ model_name,
141
+ matches,
142
+ model_kwargs,
143
+ attn_backend,
144
+ compilation_config,
145
+ matches_check,
146
+ )
tests/compile/fusions_e2e/test_tp2_ar_rms.py
new
+199
@@ -0,0 +1,199 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import pytest
6
+
7
+from vllm.config import PassConfig
8
+
9
+from ...utils import multi_gpu_test
10
+from .common import (
11
+ INDUCTOR_GRAPH_PARTITION,
12
+ AttentionBackendCase,
13
+ Matches,
14
+ custom_ops_combos,
15
+ is_blackwell,
16
+)
17
+from .models import (
18
+ FLASHINFER_ATTN,
19
+ TRITON_ATTN,
20
+ llama3_8b,
21
+ llama3_8b_fp4,
22
+ llama3_8b_fp8,
23
+ llama4_scout_fp4,
24
+ llama4_scout_fp8,
25
+ qwen3_a3b,
26
+ qwen3_a3b_fp8,
27
+)
28
+
29
+
30
+@multi_gpu_test(num_gpus=2)
31
+@pytest.mark.parametrize(
32
+ "model_name, matches_fn, model_kwargs, hf_overrides",
33
+ # qwen3-fp8 should still fuse AR+rms even though group quant is not yet supported
34
+ [llama3_8b_fp8, llama4_scout_fp8, qwen3_a3b_fp8],
35
+)
36
+@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN])
37
+@pytest.mark.parametrize("n_layers", [4])
38
+@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
39
+@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
40
+def test_tp2_ar_rms_fp8_fusions(
41
+ model_name: str,
42
+ matches_fn: Callable[[int], Matches],
43
+ model_kwargs: dict,
44
+ hf_overrides: Callable[[int], dict],
45
+ attn_backend: AttentionBackendCase,
46
+ n_layers: int,
47
+ custom_ops: str,
48
+ inductor_graph_partition: bool,
49
+ run_e2e_fusion_test,
50
+ monkeypatch,
51
+):
52
+ matches = matches_fn(n_layers)
53
+
54
+ if "qwen" in model_name.lower() and "-quant_fp8" in custom_ops:
55
+ # This is why config forces +quant_fp8 by default
56
+ pytest.skip("native QuantFP8 matching not supported for group quant")
57
+
58
+ # Reduce size of model and skip weight loading time
59
+ model_kwargs["hf_overrides"] = hf_overrides(n_layers)
60
+ model_kwargs["load_format"] = "dummy"
61
+ model_kwargs["max_model_len"] = 1024
62
+
63
+ compilation_config = dict(
64
+ use_inductor_graph_partition=inductor_graph_partition,
65
+ custom_ops=custom_ops.split(","),
66
+ pass_config=PassConfig(
67
+ fuse_norm_quant=True,
68
+ fuse_act_quant=True,
69
+ fuse_attn_quant=True,
70
+ enable_qk_norm_rope_fusion=True,
71
+ fuse_allreduce_rms=True,
72
+ ),
73
+ )
74
+
75
+ matches_check = [
76
+ "rms_quant_fusion",
77
+ "act_quant_fusion",
78
+ "norm_rope_fusion",
79
+ "attn_quant_fusion",
80
+ "ar_rms_fusion",
81
+ ]
82
+
83
+ run_e2e_fusion_test(
84
+ model_name,
85
+ matches,
86
+ model_kwargs,
87
+ attn_backend,
88
+ compilation_config,
89
+ matches_check,
90
+ tp_size=2,
91
+ )
92
+
93
+
94
+@multi_gpu_test(num_gpus=2)
95
+@pytest.mark.parametrize(
96
+ "model_name, matches_fn, model_kwargs, hf_overrides",
97
+ [llama3_8b_fp4, llama4_scout_fp4],
98
+)
99
+@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN])
100
+@pytest.mark.parametrize("n_layers", [4])
101
+@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
102
+@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
103
+@pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4")
104
+def test_tp2_ar_rms_fp4_fusions(
105
+ model_name: str,
106
+ matches_fn: Callable[[int], Matches],
107
+ model_kwargs: dict,
108
+ hf_overrides: Callable[[int], dict],
109
+ attn_backend: AttentionBackendCase,
110
+ n_layers: int,
111
+ custom_ops: str,
112
+ inductor_graph_partition: bool,
113
+ run_e2e_fusion_test,
114
+ monkeypatch,
115
+):
116
+ matches = matches_fn(n_layers)
117
+
118
+ # Reduce size of model and skip weight loading time
119
+ model_kwargs["hf_overrides"] = hf_overrides(n_layers)
120
+ model_kwargs["load_format"] = "dummy"
121
+ model_kwargs["max_model_len"] = 1024
122
+
123
+ compilation_config = dict(
124
+ use_inductor_graph_partition=inductor_graph_partition,
125
+ custom_ops=custom_ops.split(","),
126
+ pass_config=PassConfig(
127
+ fuse_act_quant=True,
128
+ fuse_attn_quant=True,
129
+ fuse_allreduce_rms=True,
130
+ ),
131
+ )
132
+
133
+ matches_check = [
134
+ "act_quant_fusion",
135
+ "attn_quant_fusion",
136
+ "ar_rms_fusion",
137
+ ]
138
+
139
+ run_e2e_fusion_test(
140
+ model_name,
141
+ matches,
142
+ model_kwargs,
143
+ attn_backend,
144
+ compilation_config,
145
+ matches_check,
146
+ tp_size=2,
147
+ )
148
+
149
+
150
+@multi_gpu_test(num_gpus=2)
151
+@pytest.mark.parametrize(
152
+ "model_name, matches_fn, model_kwargs, hf_overrides",
153
+ [llama3_8b, qwen3_a3b],
154
+)
155
+@pytest.mark.parametrize("attn_backend", [TRITON_ATTN])
156
+@pytest.mark.parametrize("n_layers", [4])
157
+@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
158
+@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
159
+def test_tp2_ar_rms_fusions(
160
+ model_name: str,
161
+ matches_fn: Callable[[int], Matches],
162
+ model_kwargs: dict,
163
+ hf_overrides: Callable[[int], dict],
164
+ attn_backend: AttentionBackendCase,
165
+ n_layers: int,
166
+ custom_ops: str,
167
+ inductor_graph_partition: bool,
168
+ run_e2e_fusion_test,
169
+):
170
+ matches = matches_fn(n_layers)
171
+
172
+ # Reduce size of model and skip weight loading time
173
+ model_kwargs["hf_overrides"] = hf_overrides(n_layers)
174
+ model_kwargs["load_format"] = "dummy"
175
+ model_kwargs["max_model_len"] = 1024
176
+
177
+ compilation_config = dict(
178
+ use_inductor_graph_partition=inductor_graph_partition,
179
+ custom_ops=custom_ops.split(","),
180
+ pass_config=PassConfig(
181
+ enable_qk_norm_rope_fusion=True,
182
+ fuse_allreduce_rms=True,
183
+ ),
184
+ )
185
+
186
+ matches_check = [
187
+ "norm_rope_fusion",
188
+ "ar_rms_fusion",
189
+ ]
190
+
191
+ run_e2e_fusion_test(
192
+ model_name,
193
+ matches,
194
+ model_kwargs,
195
+ attn_backend,
196
+ compilation_config,
197
+ matches_check,
198
+ tp_size=2,
199
+ )
tests/compile/fusions_e2e/test_tp2_async_tp.py
new
+143
@@ -0,0 +1,143 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import pytest
6
+
7
+from vllm.config import PassConfig
8
+
9
+from ...utils import multi_gpu_test
10
+from .common import (
11
+ INDUCTOR_GRAPH_PARTITION,
12
+ AttentionBackendCase,
13
+ Matches,
14
+ custom_ops_combos,
15
+ is_blackwell,
16
+)
17
+from .models import (
18
+ FLASHINFER_ATTN,
19
+ TRITON_ATTN,
20
+ llama3_8b,
21
+ llama3_8b_fp8,
22
+ llama4_scout_fp8,
23
+ qwen3_a3b,
24
+)
25
+
26
+
27
+@multi_gpu_test(num_gpus=2)
28
+@pytest.mark.parametrize(
29
+ "model_name, matches_fn, model_kwargs, hf_overrides",
30
+ [llama3_8b_fp8, llama4_scout_fp8],
31
+)
32
+@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN])
33
+@pytest.mark.parametrize("n_layers", [4])
34
+@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
35
+@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
36
+def test_tp2_async_tp_fp8_fusions(
37
+ model_name: str,
38
+ matches_fn: Callable[[int], Matches],
39
+ model_kwargs: dict,
40
+ hf_overrides: Callable[[int], dict],
41
+ attn_backend: AttentionBackendCase,
42
+ n_layers: int,
43
+ custom_ops: str,
44
+ inductor_graph_partition: bool,
45
+ run_e2e_fusion_test,
46
+ monkeypatch,
47
+):
48
+ matches = matches_fn(n_layers)
49
+
50
+ if is_blackwell():
51
+ # Disable FlashInfer scaled_mm FP8 as it's not supported in async tp patterns
52
+ monkeypatch.setenv("VLLM_DISABLED_KERNELS", "FlashInferFP8ScaledMMLinearKernel")
53
+
54
+ # Reduce size of model and skip weight loading time
55
+ model_kwargs["hf_overrides"] = hf_overrides(n_layers)
56
+ model_kwargs["load_format"] = "dummy"
57
+ model_kwargs["max_model_len"] = 1024
58
+
59
+ compilation_config = dict(
60
+ use_inductor_graph_partition=inductor_graph_partition,
61
+ custom_ops=custom_ops.split(","),
62
+ pass_config=PassConfig(
63
+ fuse_norm_quant=True,
64
+ fuse_act_quant=True,
65
+ fuse_attn_quant=True,
66
+ enable_qk_norm_rope_fusion=True,
67
+ enable_sp=True,
68
+ fuse_gemm_comms=True,
69
+ ),
70
+ )
71
+
72
+ matches_check = [
73
+ "rms_quant_fusion",
74
+ "act_quant_fusion",
75
+ "norm_rope_fusion",
76
+ "attn_quant_fusion",
77
+ "sequence_parallel",
78
+ "async_tp",
79
+ ]
80
+
81
+ run_e2e_fusion_test(
82
+ model_name,
83
+ matches,
84
+ model_kwargs,
85
+ attn_backend,
86
+ compilation_config,
87
+ matches_check,
88
+ tp_size=2,
89
+ )
90
+
91
+
92
+@multi_gpu_test(num_gpus=2)
93
+@pytest.mark.parametrize(
94
+ "model_name, matches_fn, model_kwargs, hf_overrides",
95
+ [llama3_8b, qwen3_a3b],
96
+)
97
+@pytest.mark.parametrize("attn_backend", [TRITON_ATTN])
98
+@pytest.mark.parametrize("n_layers", [4])
99
+@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
100
+@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
101
+def test_tp2_async_tp_fusions(
102
+ model_name: str,
103
+ matches_fn: Callable[[int], Matches],
104
+ model_kwargs: dict,
105
+ hf_overrides: Callable[[int], dict],
106
+ attn_backend: AttentionBackendCase,
107
+ n_layers: int,
108
+ custom_ops: str,
109
+ inductor_graph_partition: bool,
110
+ run_e2e_fusion_test,
111
+):
112
+ matches = matches_fn(n_layers)
113
+
114
+ # Reduce size of model and skip weight loading time
115
+ model_kwargs["hf_overrides"] = hf_overrides(n_layers)
116
+ model_kwargs["load_format"] = "dummy"
117
+ model_kwargs["max_model_len"] = 1024
118
+
119
+ compilation_config = dict(
120
+ use_inductor_graph_partition=inductor_graph_partition,
121
+ custom_ops=custom_ops.split(","),
122
+ pass_config=PassConfig(
123
+ enable_qk_norm_rope_fusion=True,
124
+ enable_sp=True,
125
+ fuse_gemm_comms=True,
126
+ ),
127
+ )
128
+
129
+ matches_check = [
130
+ "norm_rope_fusion",
131
+ "sequence_parallel",
132
+ "async_tp",
133
+ ]
134
+
135
+ run_e2e_fusion_test(
136
+ model_name,
137
+ matches,
138
+ model_kwargs,
139
+ attn_backend,
140
+ compilation_config,
141
+ matches_check,
142
+ tp_size=2,
143
+ )
tests/compile/test_fusion_attn.py
-99
@@ -1,23 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import copy
4
-import logging
5
-from typing import Any
4
5
import pytest
8
-import regex as re
6
import torch._dynamo
7
8
from tests.compile.backend import LazyInitPass, TestBackend
12
-from tests.compile.fusion_test_utils import (
13
- CUSTOM_OPS_FP8,
14
- MODELS_FP4,
15
- MODELS_FP8,
16
- Matches,
17
- has_cuda_graph_wrapper_metadata,
18
- is_blackwell,
19
- run_model,
20
-)
9
from tests.utils import 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
@@ -31,7 +19,6 @@ from vllm.config import (
19
CacheConfig,
20
CompilationConfig,
21
CompilationMode,
34
- CUDAGraphMode,
22
ModelConfig,
23
PassConfig,
24
SchedulerConfig,
@@ -47,7 +34,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
34
)
35
from vllm.platforms import current_platform
36
from vllm.utils.flashinfer import has_flashinfer
50
-from vllm.utils.torch_utils import is_torch_equal_or_newer
37
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
@@ -501,88 +487,3 @@ def test_attention_quant_pattern(
487
488
# Check that results are close
489
torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2)
504
-
505
-
506
-@pytest.mark.parametrize(
507
- "model_name, model_kwargs, backend, matches, custom_ops",
508
- # Test attention+quant_fp8 fusion with custom and torch impls of QuantFP8
509
- list(flat_product(MODELS_FP8, CUSTOM_OPS_FP8))
510
- # quant_fp4 only has the custom impl
511
- + list(flat_product(MODELS_FP4, [""])),
512
-)
513
-@pytest.mark.parametrize(
514
- "inductor_graph_partition",
515
- [
516
- pytest.param(
517
- True,
518
- marks=pytest.mark.skipif(
519
- not has_cuda_graph_wrapper_metadata(),
520
- reason="This test requires"
521
- "torch._inductor.utils.CUDAGraphWrapperMetadata to run",
522
- ),
523
- ),
524
- False,
525
- ],
526
-)
527
-def test_attn_quant(
528
- model_name: str,
529
- model_kwargs: dict[str, Any],
530
- backend: AttentionBackendEnum,
531
- matches: Matches,
532
- custom_ops: str,
533
- inductor_graph_partition: bool,
534
- caplog_mp_spawn,
535
- monkeypatch,
536
-):
537
- if not current_platform.has_device_capability(90):
538
- pytest.skip("test_attn_quant requires H100 (SM90) or B200 (SM100) GPU")
539
- if backend == AttentionBackendEnum.FLASHINFER and (
540
- not is_blackwell() or not has_flashinfer()
541
- ):
542
- pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")
543
- if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
544
- pytest.skip("Inductor graph partition requires torch>=2.9")
545
-
546
- custom_ops_list = custom_ops.split(",") if custom_ops else []
547
-
548
- if inductor_graph_partition:
549
- mode = CUDAGraphMode.FULL_AND_PIECEWISE
550
- splitting_ops: list[str] | None = None
551
- else:
552
- # FIXME: Llama-4-Scout-17B-16E-Instruct-FP8 + FlashInfer + Blackwell end at
553
- # CUDAGraphMode.NONE here because it derives an attention backend that
554
- # does not support full cudagraphs
555
- mode = CUDAGraphMode.FULL_DECODE_ONLY
556
- splitting_ops = []
557
-
558
- # Disable, compile cache to make sure custom passes run.
559
- # Otherwise, we can't verify fusion happened through the logs.
560
- monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
561
-
562
- # To capture subprocess logs, we need to know whether spawn or fork is used.
563
- # Force spawn as it is more general.
564
- monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
565
- model_kwargs["attention_config"] = {"backend": backend.name}
566
-
567
- compilation_config = CompilationConfig(
568
- # Testing properties
569
- custom_ops=custom_ops_list,
570
- use_inductor_graph_partition=inductor_graph_partition,
571
- cudagraph_mode=mode,
572
- splitting_ops=splitting_ops,
573
- # Common
574
- mode=CompilationMode.VLLM_COMPILE,
575
- pass_config=PassConfig(fuse_attn_quant=True, eliminate_noops=True),
576
- # Inductor caches custom passes by default as well via uuid
577
- inductor_compile_config={"force_disable_caches": True},
578
- )
579
-
580
- with caplog_mp_spawn(logging.DEBUG) as log_holder:
581
- run_model(compilation_config, model_name, **model_kwargs)
582
-
583
- log_matches = re.findall(
584
- r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes",
585
- log_holder.text,
586
- )
587
- assert len(log_matches) == 1, log_holder.text
588
- assert int(log_matches[0]) == matches.attention_fusion
tests/test_config.py
+2
-2
@@ -1002,7 +1002,7 @@ def test_vllm_config_explicit_overrides():
1002
assert config.compilation_config.pass_config.fuse_attn_quant is True
1003
1004
# Explicit cudagraph mode override on quantized model at O2
1005
- pass_config = PassConfig(fuse_gemm_comms=True)
1005
+ pass_config = PassConfig(enable_qk_norm_rope_fusion=True)
1006
compilation_config = CompilationConfig(
1007
cudagraph_mode=CUDAGraphMode.NONE, pass_config=pass_config
1008
)
@@ -1012,7 +1012,7 @@ def test_vllm_config_explicit_overrides():
1012
compilation_config=compilation_config,
1013
)
1014
assert config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE
1015
- assert config.compilation_config.pass_config.fuse_gemm_comms is True
1015
+ assert config.compilation_config.pass_config.enable_qk_norm_rope_fusion is True
1016
# Mode should still use default for O2
1017
assert config.compilation_config.mode == CompilationMode.VLLM_COMPILE
1018
vllm/config/vllm.py
+6
-1
@@ -766,7 +766,12 @@ class VllmConfig:
766
if self.compilation_config.pass_config.fuse_gemm_comms:
767
self.compilation_config.pass_config.enable_sp = True
768
if self.compilation_config.pass_config.enable_sp:
769
- if "-rms_norm" in self.compilation_config.custom_ops:
769
+ if self.parallel_config.tensor_parallel_size == 1:
770
+ logger.warning("Sequence Parallelism requires TP>1, disabling")
771
+ self.compilation_config.pass_config.enable_sp = False
772
+ self.compilation_config.pass_config.fuse_gemm_comms = False
773
+
774
+ elif "-rms_norm" in self.compilation_config.custom_ops:
775
logger.warning(
776
"RMS norm force disabled, sequence parallelism might break"
777
)