@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 248 lines 7.58 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 from vllm.utils.flashinfer import is_flashinfer_fp8_blockscale_gemm_supported
10
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 FLASHINFER_MLA_ATTN,
22 ROCM_AITER_UNIFIED_ATTN,
23 ROCM_ATTN,
24 TRITON_ATTN,
25 TRITON_MLA_ATTN,
26 deepseek_coder_v2_lite_fp8,
27 deepseek_r1_fp4,
28 deepseek_v3_fp8,
29 llama3_8b_fp4,
30 llama3_8b_fp8,
31 llama4_scout_fp4,
32 llama4_scout_fp8,
33 qwen3_8_27b_mixed_fp4,
34 qwen3_a3b_fp8,
35 )
36
37
38 @pytest.mark.parametrize(
39 "model_name, matches_fn, model_kwargs, hf_overrides, use_deepgemm",
40 [
41 (*llama3_8b_fp8, False),
42 (*qwen3_a3b_fp8, False),
43 (*qwen3_a3b_fp8, True),
44 (*deepseek_coder_v2_lite_fp8, False),
45 (*deepseek_v3_fp8, False),
46 (*deepseek_v3_fp8, True),
47 pytest.param(
48 *llama4_scout_fp8,
49 False,
50 marks=pytest.mark.skipif(
51 not current_platform.is_cuda(),
52 reason="Llama4 Scout FP8 only supported on CUDA",
53 ),
54 ),
55 ],
56 )
57 @pytest.mark.parametrize(
58 "attn_backend",
59 [
60 TRITON_ATTN,
61 FLASHINFER_ATTN,
62 ROCM_ATTN,
63 ROCM_AITER_UNIFIED_ATTN,
64 FLASHINFER_MLA_ATTN,
65 TRITON_MLA_ATTN,
66 ],
67 )
68 @pytest.mark.parametrize("n_layers", [6])
69 @pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
70 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
71 def test_tp1_fp8_fusions(
72 model_name: str,
73 matches_fn: Callable[[int], Matches],
74 model_kwargs: dict,
75 hf_overrides: Callable[[int], dict],
76 attn_backend: AttentionBackendCase,
77 n_layers: int,
78 custom_ops: str,
79 inductor_graph_partition: bool,
80 use_deepgemm: bool,
81 run_e2e_fusion_test,
82 ):
83 if use_deepgemm and not current_platform.is_cuda():
84 pytest.skip("DeepGemm only supported on CUDA")
85
86 if use_deepgemm and is_flashinfer_fp8_blockscale_gemm_supported():
87 # Flashinfer block FP8 GEMM has internal quantization, so it can't
88 # be fused with other ops.
89 pytest.skip("FlashInfer block FP8 GEMM not supported")
90 if use_deepgemm and is_blackwell():
91 # TODO(luka) DeepGEMM uses different quants, matching not supported
92 # - on Blackwell, uses a special quant fp8, currently not supported
93 pytest.skip("DeepGEMM & quant matching not currently supported")
94
95 matches = matches_fn(n_layers)
96
97 block_fp8 = "qwen" in model_name.lower() or "deepseek" in model_name.lower()
98 if block_fp8 and "-quant_fp8" in custom_ops:
99 # This is why config forces +quant_fp8 by default
100 pytest.skip("native QuantFP8 matching not supported for group quant")
101
102 # Reduce size of model and skip weight loading time
103 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
104 model_kwargs["load_format"] = "dummy"
105 model_kwargs["max_model_len"] = 1024
106 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
107
108 compilation_config = dict(
109 use_inductor_graph_partition=inductor_graph_partition,
110 custom_ops=custom_ops.split(","),
111 pass_config=PassConfig(
112 fuse_norm_quant=True,
113 fuse_act_quant=False,
114 fuse_attn_quant=True,
115 enable_qk_norm_rope_fusion=True,
116 ),
117 )
118
119 use_aiter = current_platform.is_rocm() and ("qwen" in model_name.lower())
120
121 matches_check = [
122 "rms_quant_fusion",
123 "norm_rope_fusion",
124 "attn_quant_fusion",
125 ]
126
127 if use_aiter:
128 matches_check[0] = "aiter_rms_quant_fusion"
129
130 matches = matches._replace(aiter_rms_quant_fusion=matches.rms_quant_fusion)
131 # TODO: enable the `norm_rope_fusion` test,
132 # On ROCm norm_rope_fusion is only supported without
133 # enabling AITER.
134 matches_check.remove("norm_rope_fusion")
135
136 run_e2e_fusion_test(
137 model_name,
138 matches,
139 model_kwargs,
140 attn_backend,
141 compilation_config,
142 matches_check,
143 use_deepgemm=use_deepgemm,
144 use_aiter=use_aiter,
145 )
146
147
148 @pytest.mark.parametrize(
149 "model_name, matches_fn, model_kwargs, hf_overrides",
150 [llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4],
151 )
152 @pytest.mark.parametrize(
153 "attn_backend",
154 [FLASHINFER_ATTN, FLASHINFER_MLA_ATTN],
155 )
156 @pytest.mark.parametrize("n_layers", [6])
157 @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
158 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
159 @pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4")
160 def test_tp1_fp4_fusions(
161 model_name: str,
162 matches_fn: Callable[[int], Matches],
163 model_kwargs: dict,
164 hf_overrides: Callable[[int], dict],
165 attn_backend: AttentionBackendCase,
166 n_layers: int,
167 custom_ops: str,
168 inductor_graph_partition: bool,
169 run_e2e_fusion_test,
170 ):
171 if nvfp4_kernel_exposes_input_quant_key():
172 pytest.skip(
173 "NVFP4 kernel exposes input_quant_key; manual fusion fires "
174 "instead of compiler pass-based fusion"
175 )
176
177 matches = matches_fn(n_layers)
178
179 # Reduce size of model and skip weight loading time
180 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
181 model_kwargs["load_format"] = "dummy"
182 model_kwargs["max_model_len"] = 1024
183 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
184
185 compilation_config = dict(
186 use_inductor_graph_partition=inductor_graph_partition,
187 custom_ops=custom_ops.split(","),
188 pass_config=PassConfig(
189 fuse_norm_quant=True,
190 fuse_act_quant=False,
191 fuse_attn_quant=True,
192 enable_qk_norm_rope_fusion=True,
193 ),
194 )
195
196 matches_check = ["attn_quant_fusion", "norm_rope_fusion"]
197
198 run_e2e_fusion_test(
199 model_name,
200 matches,
201 model_kwargs,
202 attn_backend,
203 compilation_config,
204 matches_check,
205 )
206
207
208 @pytest.mark.parametrize(
209 "model_name, matches_fn, model_kwargs, hf_overrides", [qwen3_8_27b_mixed_fp4]
210 )
211 @pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN])
212 @pytest.mark.parametrize("n_layers", [4])
213 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
214 @pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4")
215 def test_tp1_fp4_mixed_precision_act_fusion(
216 model_name: str,
217 matches_fn: Callable[[int], Matches],
218 model_kwargs: dict,
219 hf_overrides: Callable[[int], dict],
220 attn_backend: AttentionBackendCase,
221 n_layers: int,
222 inductor_graph_partition: bool,
223 run_e2e_fusion_test,
224 ):
225 """Act+quant fusion matches on a MIXED_PRECISION NVFP4 checkpoint."""
226 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
227 model_kwargs["load_format"] = "dummy"
228 model_kwargs["max_model_len"] = 1024
229 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
230
231 compilation_config = dict(
232 use_inductor_graph_partition=inductor_graph_partition,
233 pass_config=PassConfig(
234 fuse_norm_quant=False,
235 fuse_act_quant=True,
236 fuse_attn_quant=False,
237 enable_qk_norm_rope_fusion=False,
238 ),
239 )
240
241 run_e2e_fusion_test(
242 model_name,
243 matches_fn(n_layers),
244 model_kwargs,
245 attn_backend,
246 compilation_config,
247 ["act_quant_fusion"],
248 )