@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 315 lines 10 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import itertools
5 from typing import Any
6
7 import torch
8 from torch._higher_order_ops.auto_functionalize import auto_functionalized
9 from torch._ops import OpOverload
10
11 from vllm.config import VllmConfig
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.quantization.utils.quant_utils import (
14 QuantKey,
15 kFp8Dynamic64Sym,
16 kFp8Dynamic128Sym,
17 kFp8StaticTensorSym,
18 kNvfp4Dynamic,
19 )
20 from vllm.platforms import current_platform
21
22 from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement
23 from .matcher_utils import MatcherQuantFP8, MatcherSiluAndMul
24 from .rms_quant_fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
25
26 logger = init_logger(__name__)
27
28 FP8_DTYPE = current_platform.fp8_dtype()
29 FP4_DTYPE = torch.uint8
30
31 SILU_MUL_OP = torch.ops._C.silu_and_mul.default
32
33 FUSED_OPS: dict[QuantKey, OpOverload] = {
34 kFp8StaticTensorSym: torch.ops._C.silu_and_mul_quant.default, # noqa: E501
35 }
36 silu_and_mul_nvfp4_quant_supported = current_platform.is_cuda() and hasattr(
37 torch.ops._C, "silu_and_mul_nvfp4_quant"
38 )
39 if silu_and_mul_nvfp4_quant_supported:
40 FUSED_OPS[kNvfp4Dynamic] = torch.ops._C.silu_and_mul_nvfp4_quant.default # noqa: E501
41
42 if current_platform.is_cuda_alike():
43 FUSED_OPS[kFp8Dynamic128Sym] = torch.ops._C.silu_and_mul_per_block_quant.default
44 FUSED_OPS[kFp8Dynamic64Sym] = torch.ops._C.silu_and_mul_per_block_quant.default
45
46
47 class ActivationQuantPattern(VllmPatternReplacement):
48 """Base class for Activation+Quant fusions.
49 Should not be used directly.
50 """
51
52 def __init__(
53 self,
54 quant_key: QuantKey,
55 ) -> None:
56 self.quant_key = quant_key
57 self.quant_dtype = quant_key.dtype
58
59 assert self.quant_key in QUANT_OPS, (
60 f"unsupported quantization scheme {self.quant_key}"
61 )
62 self.QUANT_OP = QUANT_OPS[self.quant_key]
63
64 assert self.quant_key in FUSED_OPS, (
65 f"unsupported fusion scheme {self.quant_key}"
66 )
67 self.FUSED_OP = FUSED_OPS[self.quant_key]
68
69 self.silu_and_mul_matcher = MatcherSiluAndMul()
70
71 def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor:
72 kwargs = {
73 "dtype": self.quant_dtype,
74 "device": current_platform.device_type,
75 **kwargs,
76 }
77 return torch.empty(*args, **kwargs)
78
79
80 class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):
81 """Fusion for SiluMul+Fp8StaticQuant Pattern."""
82
83 def __init__(self) -> None:
84 super().__init__(kFp8StaticTensorSym)
85 self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
86
87 def get_inputs(self) -> list[torch.Tensor]:
88 scale = self.quant_matcher.inputs()[1]
89 return [
90 *self.silu_and_mul_matcher.inputs(), # input
91 scale,
92 ]
93
94 @property
95 def pattern(self):
96 def _pattern(
97 input: torch.Tensor,
98 scale: torch.Tensor,
99 ) -> torch.Tensor:
100 result_silu_mul = self.silu_and_mul_matcher(input)
101 result_quant = self.quant_matcher(result_silu_mul, scale)
102 return result_quant[0]
103
104 return _pattern
105
106 @property
107 def replacement(self):
108 def _replacement(
109 input: torch.Tensor,
110 scale: torch.Tensor,
111 ) -> torch.Tensor:
112 d = input.shape[-1] // 2
113 output_shape = input.shape[:-1] + (d,)
114 result = torch.empty(
115 output_shape, device=input.device, dtype=self.quant_dtype
116 )
117 at = auto_functionalized(
118 self.FUSED_OP, result=result, input=input, scale=scale
119 )
120 return at[1]
121
122 return _replacement
123
124
125 class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
126 """Fusion for SiluMul+Nvfp4Quant Pattern."""
127
128 def __init__(self) -> None:
129 super().__init__(kNvfp4Dynamic)
130
131 def get_inputs(self) -> list[torch.Tensor]:
132 result = self.empty_quant(5, 32)
133 output_scale = empty_i32(128, 4)
134 input_ = empty_bf16(5, 64)
135 scale = empty_fp32(1, 1)
136 return [result, output_scale, input_, scale]
137
138 @property
139 def pattern(self):
140 def _pattern(
141 result: torch.Tensor,
142 output_scale: torch.Tensor,
143 input: torch.Tensor,
144 scale: torch.Tensor,
145 ) -> tuple[torch.Tensor, torch.Tensor]:
146 result_silu_mul = self.silu_and_mul_matcher(input)
147 at = auto_functionalized(
148 self.QUANT_OP,
149 input=result_silu_mul,
150 input_scale=scale,
151 is_sf_swizzled_layout=True,
152 output=result,
153 output_scale=output_scale,
154 )
155 return at[1], at[2]
156
157 return _pattern
158
159 @property
160 def replacement(self):
161 def _replacement(
162 result: torch.Tensor,
163 output_scale: torch.Tensor,
164 input: torch.Tensor,
165 scale: torch.Tensor,
166 ) -> tuple[torch.Tensor, torch.Tensor]:
167 at = auto_functionalized(
168 self.FUSED_OP,
169 result=result,
170 result_block_scale=output_scale,
171 input=input,
172 input_global_scale=scale,
173 )
174 return at[1], at[2]
175
176 return _replacement
177
178
179 class SiluMulBlockQuantPattern(ActivationQuantPattern):
180 """Fusion for SiluMul+BlockQuant (FP8 dynamic per-group) Pattern.
181 Supports group_size 128 and 64 via QuantKey.
182 Parameterized on is_scale_transposed for different scale layouts.
183 """
184
185 def __init__(
186 self,
187 quant_key: QuantKey,
188 is_scale_transposed: bool = False,
189 is_e8m0: bool = False,
190 is_tma_aligned: bool = False,
191 match_aiter: bool = False,
192 ) -> None:
193 super().__init__(quant_key)
194 self.quant_matcher = MatcherQuantFP8(
195 quant_key,
196 has_col_major_scales=is_scale_transposed,
197 is_e8m0=is_e8m0,
198 is_tma_aligned=is_tma_aligned,
199 )
200 self.group_size = quant_key.scale.group_shape[1]
201 self.is_scale_transposed = is_scale_transposed
202 self.is_e8m0 = is_e8m0
203 self.is_tma_aligned = is_tma_aligned
204
205 def get_inputs(self) -> list[torch.Tensor]:
206 scale = self.quant_matcher.empty_f32(1, 1)
207 return self.silu_and_mul_matcher.inputs() + [scale]
208
209 @property
210 def pattern(self):
211 def _pattern(
212 input: torch.Tensor,
213 scale: torch.Tensor,
214 ) -> tuple[torch.Tensor, torch.Tensor]:
215 silu_out = self.silu_and_mul_matcher(input)
216 result = torch.empty(
217 silu_out.shape,
218 device=silu_out.device,
219 dtype=self.quant_dtype,
220 )
221 assert scale is not None
222 finfo = torch.finfo(self.quant_dtype)
223 _, result, scale = auto_functionalized(
224 self.quant_matcher.QUANT_OP,
225 input=silu_out,
226 output_q=result,
227 output_s=scale,
228 group_size=self.group_size,
229 eps=1e-10,
230 fp8_min=finfo.min,
231 fp8_max=finfo.max,
232 scale_ue8m0=self.is_e8m0,
233 dummy_is_scale_transposed=self.is_scale_transposed,
234 dummy_is_tma_aligned=self.is_tma_aligned,
235 )
236 return result, scale
237
238 return _pattern
239
240 @property
241 def replacement(self):
242 def _replacement(
243 input: torch.Tensor,
244 scale: torch.Tensor,
245 ) -> tuple[torch.Tensor, torch.Tensor]:
246 d = input.shape[-1] // 2
247 output_shape = input.shape[:-1] + (d,)
248 result = torch.empty(
249 output_shape, device=input.device, dtype=self.quant_dtype
250 )
251 if self.is_scale_transposed:
252 scale = torch.empty(
253 (d // self.group_size, input.shape[0]),
254 device=input.device,
255 dtype=torch.float32,
256 ).permute(-1, -2)
257 else:
258 scale = torch.empty(
259 (input.shape[0], d // self.group_size),
260 device=input.device,
261 dtype=torch.float32,
262 )
263 at = auto_functionalized(
264 self.FUSED_OP,
265 out=result,
266 input=input,
267 scales=scale,
268 group_size=self.group_size,
269 scale_ub=None,
270 is_scale_transposed=self.is_scale_transposed,
271 )
272 return at[1], at[2]
273
274 return _replacement
275
276
277 class ActivationQuantFusionPass(VllmFusionPatternMatcherPass):
278 """This pass fuses a pre-defined set of custom ops into fused ops.
279 It uses the torch pattern matcher to find the patterns and replace them.
280
281 Because patterns can only be registered once, the pass is a singleton.
282 This will be addressed in a future version of PyTorch:
283 https://github.com/pytorch/pytorch/pull/139321#issuecomment-2452354980
284 """
285
286 def __init__(self, config: VllmConfig) -> None:
287 super().__init__(config, "activation_quant_fusion_pass")
288
289 self.register(SiluMulFp8StaticQuantPattern())
290
291 if silu_and_mul_nvfp4_quant_supported:
292 self.register(SiluMulNvfp4QuantPattern())
293
294 if current_platform.is_cuda():
295 for (
296 quant_key,
297 is_scale_transposed,
298 is_e8m0,
299 is_tma_aligned,
300 ) in itertools.product(
301 [kFp8Dynamic128Sym, kFp8Dynamic64Sym],
302 [False, True],
303 [True, False],
304 [False, True],
305 ):
306 self.register(
307 SiluMulBlockQuantPattern(
308 quant_key,
309 is_scale_transposed=is_scale_transposed,
310 is_e8m0=is_e8m0,
311 is_tma_aligned=is_tma_aligned,
312 )
313 )
314
315 self.dump_patterns(config, self.pm_pass)