@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 339 lines 9.9 KB
Raw
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 from vllm.platforms import current_platform
9
10 from ...utils import multi_gpu_test
11 from .common import (
12 INDUCTOR_GRAPH_PARTITION,
13 AttentionBackendCase,
14 Matches,
15 custom_ops_combos,
16 is_blackwell,
17 nvfp4_kernel_exposes_input_quant_key,
18 )
19 from .models import (
20 FLASHINFER_ATTN,
21 TRITON_ATTN,
22 llama3_8b,
23 llama3_8b_fp4,
24 llama3_8b_fp8,
25 llama4_scout_fp8,
26 qwen3_a3b,
27 )
28
29 pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
30
31
32 @multi_gpu_test(num_gpus=2)
33 @pytest.mark.parametrize(
34 "model_name, matches_fn, model_kwargs, hf_overrides",
35 [llama3_8b_fp8, llama4_scout_fp8],
36 )
37 @pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN])
38 @pytest.mark.parametrize("n_layers", [4])
39 @pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
40 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
41 def test_tp2_async_tp_fp8_fusions(
42 model_name: str,
43 matches_fn: Callable[[int], Matches],
44 model_kwargs: dict,
45 hf_overrides: Callable[[int], dict],
46 attn_backend: AttentionBackendCase,
47 n_layers: int,
48 custom_ops: str,
49 inductor_graph_partition: bool,
50 run_e2e_fusion_test,
51 ):
52 matches = matches_fn(n_layers)
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 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
59
60 compilation_config = dict(
61 use_inductor_graph_partition=inductor_graph_partition,
62 custom_ops=custom_ops.split(","),
63 pass_config=PassConfig(
64 fuse_norm_quant=True,
65 fuse_act_quant=False,
66 fuse_attn_quant=True,
67 enable_qk_norm_rope_fusion=True,
68 enable_sp=True,
69 fuse_gemm_comms=True,
70 fuse_allreduce_rms=False,
71 # Override threshold for testing (models have small hidden_size)
72 sp_min_token_num=512,
73 ),
74 )
75
76 matches_check = [
77 "rms_quant_fusion",
78 "norm_rope_fusion",
79 "attn_quant_fusion",
80 "sequence_parallel",
81 "async_tp",
82 ]
83
84 run_e2e_fusion_test(
85 model_name,
86 matches,
87 model_kwargs,
88 attn_backend,
89 compilation_config,
90 matches_check,
91 tp_size=2,
92 )
93
94
95 @multi_gpu_test(num_gpus=2)
96 @pytest.mark.parametrize(
97 "model_name, matches_fn, model_kwargs, hf_overrides",
98 [llama3_8b_fp4],
99 )
100 @pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN])
101 @pytest.mark.parametrize("n_layers", [4])
102 @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
103 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
104 @pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4")
105 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
106 def test_tp2_async_tp_nvfp4_fusions(
107 model_name: str,
108 matches_fn: Callable[[int], Matches],
109 model_kwargs: dict,
110 hf_overrides: Callable[[int], dict],
111 attn_backend: AttentionBackendCase,
112 n_layers: int,
113 custom_ops: str,
114 inductor_graph_partition: bool,
115 run_e2e_fusion_test,
116 ):
117 if nvfp4_kernel_exposes_input_quant_key():
118 pytest.skip(
119 "NVFP4 kernel exposes input_quant_key; manual fusion fires "
120 "instead of compiler pass-based fusion"
121 )
122
123 # NVFP4 currently wires the all-gather + GEMM path only.
124 matches = matches_fn(n_layers)._replace(async_tp=n_layers * 2)
125
126 # Reduce size of model and skip weight loading time
127 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
128 model_kwargs["load_format"] = "dummy"
129 model_kwargs["max_model_len"] = 1024
130 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
131
132 compilation_config = dict(
133 use_inductor_graph_partition=inductor_graph_partition,
134 custom_ops=custom_ops.split(","),
135 pass_config=PassConfig(
136 fuse_act_quant=False,
137 fuse_attn_quant=True,
138 enable_sp=True,
139 fuse_gemm_comms=True,
140 fuse_allreduce_rms=False,
141 # Override threshold for testing (models have small hidden_size)
142 sp_min_token_num=512,
143 ),
144 )
145
146 matches_check = [
147 "attn_quant_fusion",
148 "sequence_parallel",
149 "async_tp",
150 ]
151
152 run_e2e_fusion_test(
153 model_name,
154 matches,
155 model_kwargs,
156 attn_backend,
157 compilation_config,
158 matches_check,
159 tp_size=2,
160 )
161
162
163 @multi_gpu_test(num_gpus=2)
164 @pytest.mark.parametrize(
165 "model_name, matches_fn, model_kwargs, hf_overrides",
166 [llama3_8b, qwen3_a3b],
167 )
168 @pytest.mark.parametrize("attn_backend", [TRITON_ATTN])
169 @pytest.mark.parametrize("n_layers", [4])
170 @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
171 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
172 def test_tp2_async_tp_fusions(
173 model_name: str,
174 matches_fn: Callable[[int], Matches],
175 model_kwargs: dict,
176 hf_overrides: Callable[[int], dict],
177 attn_backend: AttentionBackendCase,
178 n_layers: int,
179 custom_ops: str,
180 inductor_graph_partition: bool,
181 run_e2e_fusion_test,
182 ):
183 matches = matches_fn(n_layers)
184
185 # Reduce size of model and skip weight loading time
186 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
187 model_kwargs["load_format"] = "dummy"
188 model_kwargs["max_model_len"] = 1024
189 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
190
191 compilation_config = dict(
192 use_inductor_graph_partition=inductor_graph_partition,
193 custom_ops=custom_ops.split(","),
194 pass_config=PassConfig(
195 enable_qk_norm_rope_fusion=True,
196 enable_sp=True,
197 fuse_gemm_comms=True,
198 fuse_allreduce_rms=False,
199 # Override threshold for testing (models have small hidden_size)
200 sp_min_token_num=512,
201 ),
202 )
203
204 matches_check = [
205 "norm_rope_fusion",
206 "sequence_parallel",
207 "async_tp",
208 ]
209
210 run_e2e_fusion_test(
211 model_name,
212 matches,
213 model_kwargs,
214 attn_backend,
215 compilation_config,
216 matches_check,
217 tp_size=2,
218 )
219
220
221 @multi_gpu_test(num_gpus=2)
222 @pytest.mark.parametrize(
223 "model_name, matches_fn, model_kwargs, hf_overrides",
224 [llama3_8b_fp8, llama4_scout_fp8],
225 )
226 @pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN])
227 @pytest.mark.parametrize("n_layers", [4])
228 @pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
229 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
230 def test_tp2_sp_ar_rms_fp8_fusions(
231 model_name: str,
232 matches_fn: Callable[[int], Matches],
233 model_kwargs: dict,
234 hf_overrides: Callable[[int], dict],
235 attn_backend: AttentionBackendCase,
236 n_layers: int,
237 custom_ops: str,
238 inductor_graph_partition: bool,
239 run_e2e_fusion_test,
240 ):
241 matches = matches_fn(n_layers)
242
243 # Reduce size of model and skip weight loading time
244 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
245 model_kwargs["load_format"] = "dummy"
246 model_kwargs["max_model_len"] = 1024
247
248 compilation_config = dict(
249 use_inductor_graph_partition=inductor_graph_partition,
250 custom_ops=custom_ops.split(","),
251 pass_config=PassConfig(
252 fuse_norm_quant=True,
253 fuse_act_quant=False,
254 fuse_attn_quant=True,
255 enable_qk_norm_rope_fusion=True,
256 enable_sp=True,
257 fuse_gemm_comms=True,
258 fuse_allreduce_rms=True,
259 # Override threshold for testing (models have small hidden_size)
260 sp_min_token_num=512,
261 ),
262 )
263
264 matches_check = [
265 "rms_quant_fusion",
266 "norm_rope_fusion",
267 "attn_quant_fusion",
268 "ar_rms_fusion",
269 "sequence_parallel",
270 "async_tp",
271 ]
272
273 run_e2e_fusion_test(
274 model_name,
275 matches,
276 model_kwargs,
277 attn_backend,
278 compilation_config,
279 matches_check,
280 tp_size=2,
281 )
282
283
284 @multi_gpu_test(num_gpus=2)
285 @pytest.mark.parametrize(
286 "model_name, matches_fn, model_kwargs, hf_overrides",
287 [llama3_8b, qwen3_a3b],
288 )
289 @pytest.mark.parametrize("attn_backend", [TRITON_ATTN])
290 @pytest.mark.parametrize("n_layers", [4])
291 @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
292 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
293 def test_tp2_sp_ar_rms_fusions(
294 model_name: str,
295 matches_fn: Callable[[int], Matches],
296 model_kwargs: dict,
297 hf_overrides: Callable[[int], dict],
298 attn_backend: AttentionBackendCase,
299 n_layers: int,
300 custom_ops: str,
301 inductor_graph_partition: bool,
302 run_e2e_fusion_test,
303 ):
304 matches = matches_fn(n_layers)
305
306 # Reduce size of model and skip weight loading time
307 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
308 model_kwargs["load_format"] = "dummy"
309 model_kwargs["max_model_len"] = 1024
310
311 compilation_config = dict(
312 use_inductor_graph_partition=inductor_graph_partition,
313 custom_ops=custom_ops.split(","),
314 pass_config=PassConfig(
315 enable_qk_norm_rope_fusion=True,
316 enable_sp=True,
317 fuse_gemm_comms=True,
318 fuse_allreduce_rms=True,
319 # Override threshold for testing (models have small hidden_size)
320 sp_min_token_num=512,
321 ),
322 )
323
324 matches_check = [
325 "norm_rope_fusion",
326 "ar_rms_fusion",
327 "sequence_parallel",
328 "async_tp",
329 ]
330
331 run_e2e_fusion_test(
332 model_name,
333 matches,
334 model_kwargs,
335 attn_backend,
336 compilation_config,
337 matches_check,
338 tp_size=2,
339 )