@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 855 lines 29.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from typing import Any, NamedTuple
4
5 import torch
6 import torch._inductor.pattern_matcher as pm
7 from torch import fx
8 from torch._higher_order_ops.auto_functionalize import auto_functionalized
9 from torch._inductor.pattern_matcher import PatternMatcherPass
10 from torch._ops import OpOverload
11
12 import vllm.ir.ops
13 from vllm.config import VllmConfig, get_current_vllm_config
14 from vllm.logger import init_logger
15 from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 GroupShape,
17 QuantKey,
18 ScaleDesc,
19 get_fp8_min_max,
20 kFp8Dynamic64Sym,
21 kFp8Dynamic128Sym,
22 kFp8DynamicTensorSym,
23 kFp8DynamicTokenSym,
24 kFp8StaticTensorSym,
25 kNvfp4Dynamic,
26 kStaticTensorScale,
27 )
28 from vllm.platforms import current_platform
29 from vllm.utils.flashinfer import has_flashinfer
30 from vllm.utils.torch_utils import direct_register_custom_op
31
32 from ..inductor_pass import enable_fake_mode
33 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
34 from .matcher_utils import (
35 MatcherQuantFP8,
36 )
37
38 logger = init_logger(__name__)
39 FP8_DTYPE = current_platform.fp8_dtype()
40 FP4_DTYPE = torch.uint8
41
42
43 _RMS_NORM_OP = torch.ops.vllm_ir.rms_norm.default
44 _FUSED_ADD_RMS_NORM_OP = torch.ops.vllm_ir.fused_add_rms_norm.default
45 _FLASHINFER_ADD_RMSNORM_FP4QUANT: Any | None = None
46
47
48 def _flashinfer_fused_add_rms_norm_nvfp4_quant(
49 result: torch.Tensor,
50 result_block_scale: torch.Tensor,
51 residual: torch.Tensor,
52 input: torch.Tensor,
53 weight: torch.Tensor,
54 input_global_scale: torch.Tensor,
55 block_scale_unswizzled: torch.Tensor,
56 is_sf_swizzled_layout: bool,
57 epsilon: float,
58 ) -> None:
59 """FlashInfer fused add + RMSNorm + NVFP4 quantization."""
60 assert _FLASHINFER_ADD_RMSNORM_FP4QUANT is not None
61 _FLASHINFER_ADD_RMSNORM_FP4QUANT(
62 input,
63 residual,
64 weight,
65 y_fp4=result.view(torch.float4_e2m1fn_x2),
66 block_scale=result_block_scale.view(torch.float8_e4m3fn),
67 global_scale=input_global_scale.reshape(1),
68 eps=epsilon,
69 block_size=16,
70 scale_format="e4m3",
71 is_sf_swizzled_layout=is_sf_swizzled_layout,
72 output_both_sf_layouts=False,
73 block_scale_unswizzled=block_scale_unswizzled,
74 )
75
76
77 def _flashinfer_fused_add_rms_norm_nvfp4_quant_fake(
78 result: torch.Tensor,
79 result_block_scale: torch.Tensor,
80 residual: torch.Tensor,
81 input: torch.Tensor,
82 weight: torch.Tensor,
83 input_global_scale: torch.Tensor,
84 block_scale_unswizzled: torch.Tensor,
85 is_sf_swizzled_layout: bool,
86 epsilon: float,
87 ) -> None:
88 return None
89
90
91 _FLASHINFER_NVFP4_RMS_QUANT_OP: OpOverload | None = None
92 if (
93 current_platform.is_cuda()
94 and hasattr(torch, "float4_e2m1fn_x2")
95 and has_flashinfer()
96 ):
97 try:
98 from flashinfer import add_rmsnorm_fp4quant
99 except ImportError:
100 pass
101 else:
102 _FLASHINFER_ADD_RMSNORM_FP4QUANT = add_rmsnorm_fp4quant
103 direct_register_custom_op(
104 op_name="flashinfer_fused_add_rms_norm_nvfp4_quant",
105 op_func=_flashinfer_fused_add_rms_norm_nvfp4_quant,
106 mutates_args=["result", "result_block_scale", "residual"],
107 fake_impl=_flashinfer_fused_add_rms_norm_nvfp4_quant_fake,
108 )
109 _FLASHINFER_NVFP4_RMS_QUANT_OP = (
110 torch.ops.vllm.flashinfer_fused_add_rms_norm_nvfp4_quant.default
111 )
112
113
114 # TODO: extend rmsnorm quant kernels to support mixed input/weight dtypes,
115 # and remove this check.
116 def _rms_input_weight_dtype_match(match: pm.Match) -> bool:
117 """Prevent fusion when rms_norm input and weight dtypes differ."""
118 for node in match.nodes:
119 if node.target == _RMS_NORM_OP:
120 # rms_norm(x, weight, epsilon, variance_size)
121 x, weight = node.args[0], node.args[1]
122 elif node.target == _FUSED_ADD_RMS_NORM_OP:
123 # fused_add_rms_norm(x, residual, weight, epsilon, variance_size)
124 x, weight = node.args[0], node.args[2]
125 else:
126 continue
127 if isinstance(x, fx.Node) and isinstance(weight, fx.Node):
128 return x.meta["val"].dtype == weight.meta["val"].dtype
129 return True
130
131
132 def empty_bf16(*args: Any, **kwargs: Any) -> torch.Tensor:
133 return torch.empty(
134 *args, **kwargs, dtype=torch.bfloat16, device=current_platform.device_type
135 )
136
137
138 def empty_fp32(*args: Any, **kwargs: Any) -> torch.Tensor:
139 return torch.empty(
140 *args, **kwargs, dtype=torch.float32, device=current_platform.device_type
141 )
142
143
144 def empty_i32(*args: Any, **kwargs: Any) -> torch.Tensor:
145 return torch.empty(
146 *args, **kwargs, dtype=torch.int32, device=current_platform.device_type
147 )
148
149
150 def empty_i64(*args: Any, **kwargs: Any) -> torch.Tensor:
151 return torch.empty(
152 *args, **kwargs, dtype=torch.int64, device=current_platform.device_type
153 )
154
155
156 RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
157
158 QUANT_OPS: dict[QuantKey, OpOverload] = {
159 kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default, # noqa: E501
160 kFp8DynamicTensorSym: torch.ops._C.dynamic_scaled_fp8_quant.default, # noqa: E501
161 kFp8DynamicTokenSym: torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501
162 }
163 if hasattr(torch.ops._C, "per_token_group_fp8_quant"):
164 QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501
165 QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501
166 if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
167 QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out
168
169
170 class FusedRMSQuantKey(NamedTuple):
171 """Named tuple for identifying the type of RMSNorm + quant fusion.
172 quant: type of quantization
173 fused_add: does the op also perform the residual add
174 """
175
176 quant: QuantKey
177 fused_add: bool
178
179 def __str__(self) -> str:
180 return (
181 f"FusedQuantKey({self.quant}, with"
182 f"{'' if self.fused_add else 'out'} residual)"
183 )
184
185
186 FUSED_OPS: dict[FusedRMSQuantKey, OpOverload] = {
187 FusedRMSQuantKey(
188 kFp8StaticTensorSym, False
189 ): torch.ops._C.rms_norm_static_fp8_quant.default, # noqa: E501
190 FusedRMSQuantKey(
191 kFp8StaticTensorSym, True
192 ): torch.ops._C.fused_add_rms_norm_static_fp8_quant.default, # noqa: E501
193 FusedRMSQuantKey(
194 kFp8DynamicTokenSym, False
195 ): torch.ops._C.rms_norm_dynamic_per_token_quant.default, # noqa: E501
196 FusedRMSQuantKey(
197 kFp8DynamicTokenSym, True
198 ): torch.ops._C.rms_norm_dynamic_per_token_quant.default, # noqa: E501
199 }
200 # rms_norm_per_block_quant is CUDA-only; guard it like per_token_group_fp8_quant above.
201 if hasattr(torch.ops._C, "rms_norm_per_block_quant"):
202 _rms_norm_per_block_quant = torch.ops._C.rms_norm_per_block_quant.default
203 FUSED_OPS[FusedRMSQuantKey(kFp8Dynamic128Sym, False)] = _rms_norm_per_block_quant # noqa: E501
204 FUSED_OPS[FusedRMSQuantKey(kFp8Dynamic128Sym, True)] = _rms_norm_per_block_quant # noqa: E501
205 FUSED_OPS[FusedRMSQuantKey(kFp8Dynamic64Sym, False)] = _rms_norm_per_block_quant # noqa: E501
206 FUSED_OPS[FusedRMSQuantKey(kFp8Dynamic64Sym, True)] = _rms_norm_per_block_quant # noqa: E501
207
208
209 class RMSNormQuantPattern:
210 def __init__(
211 self,
212 epsilon: float,
213 key: FusedRMSQuantKey,
214 has_col_major_scales: bool = False,
215 is_e8m0: bool = False,
216 is_tma_aligned: bool = False,
217 ) -> None:
218 self.epsilon = epsilon
219 self.quant_dtype = key.quant.dtype
220 config = get_current_vllm_config()
221 self.model_dtype = config.model_config.dtype if config.model_config else None
222
223 assert key in FUSED_OPS, f"unsupported fused rmsnorm+quant op for {key}"
224 self.FUSED_OP = FUSED_OPS[key]
225
226 self.quant_matcher = MatcherQuantFP8(
227 key.quant,
228 has_col_major_scales=has_col_major_scales,
229 is_e8m0=is_e8m0,
230 is_tma_aligned=is_tma_aligned,
231 )
232
233
234 class RMSNormStaticQuantPattern(RMSNormQuantPattern):
235 def __init__(
236 self, epsilon: float, quant_dtype: torch.dtype, symmetric: bool = True
237 ) -> None:
238 fused_key = FusedRMSQuantKey(
239 fused_add=False,
240 quant=QuantKey(
241 dtype=quant_dtype, scale=kStaticTensorScale, symmetric=symmetric
242 ),
243 )
244 super().__init__(epsilon, fused_key)
245
246 def register(self, pm_pass: PatternMatcherPass) -> None:
247 # Cannot use methods, as the self argument affects tracing
248 def pattern(
249 input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
250 ) -> torch.Tensor:
251 result_rms = vllm.ir.ops.rms_norm(input, weight, self.epsilon)
252 return self.quant_matcher(result_rms, scale)[0]
253
254 def replacement(
255 input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
256 ) -> torch.Tensor:
257 result = torch.empty(
258 input.shape, device=input.device, dtype=self.quant_dtype
259 )
260 at = auto_functionalized(
261 self.FUSED_OP,
262 result=result,
263 input=input,
264 weight=weight,
265 scale=scale,
266 epsilon=self.epsilon,
267 )
268
269 # result
270 return at[1]
271
272 inputs = [
273 empty_bf16(5, 16), # input
274 empty_bf16(16), # weight
275 self.quant_matcher.inputs()[1], # scale
276 ]
277 pattern(*inputs)
278
279 pm.register_replacement(
280 pattern,
281 replacement,
282 inputs,
283 pm.fwd_only,
284 pm_pass,
285 extra_check=_rms_input_weight_dtype_match,
286 )
287
288
289 class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
290 def __init__(
291 self, epsilon: float, quant_dtype: torch.dtype, symmetric: bool = True
292 ) -> None:
293 key = FusedRMSQuantKey(
294 fused_add=True,
295 quant=QuantKey(
296 dtype=quant_dtype, scale=kStaticTensorScale, symmetric=symmetric
297 ),
298 )
299 super().__init__(epsilon, key)
300
301 def register(self, pm_pass: PatternMatcherPass) -> None:
302 def pattern(
303 input: torch.Tensor,
304 weight: torch.Tensor,
305 residual: torch.Tensor,
306 scale: torch.Tensor,
307 ) -> tuple[torch.Tensor, torch.Tensor]:
308 result_rms, residual = vllm.ir.ops.fused_add_rms_norm(
309 input, residual, weight, self.epsilon
310 )
311 result, _ = self.quant_matcher(result_rms, scale)
312
313 return result, residual
314
315 def replacement(
316 input: torch.Tensor,
317 weight: torch.Tensor,
318 residual: torch.Tensor,
319 scale: torch.Tensor,
320 ) -> tuple[torch.Tensor, torch.Tensor]:
321 # In case we're matching native rms-norm, conversions might be
322 # optimized out. We convert here just to be safe.
323 input = input.to(dtype=self.model_dtype)
324
325 result = torch.empty_like(input, dtype=self.quant_dtype)
326 at = auto_functionalized(
327 self.FUSED_OP,
328 result=result,
329 input=input,
330 residual=residual,
331 weight=weight,
332 scale=scale,
333 epsilon=self.epsilon,
334 )
335
336 # result, residual
337 return at[1], at[2]
338
339 inputs = [
340 empty_bf16(5, 16), # input
341 empty_bf16(16), # weight
342 empty_bf16(5, 16), # residual
343 self.quant_matcher.inputs()[1], # scale
344 ]
345
346 pm.register_replacement(
347 pattern,
348 replacement,
349 inputs,
350 pm.fwd_only,
351 pm_pass,
352 extra_check=_rms_input_weight_dtype_match,
353 )
354
355
356 class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
357 def __init__(
358 self,
359 epsilon: float,
360 quant_dtype: torch.dtype,
361 group_shape: GroupShape,
362 symmetric: bool = True,
363 is_e8m0: bool = False,
364 has_col_major_scales: bool = True,
365 is_tma_aligned: bool = True,
366 ) -> None:
367 scale = ScaleDesc(torch.float32, False, group_shape)
368 key = FusedRMSQuantKey(
369 fused_add=True,
370 quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
371 )
372 self.group_shape = group_shape
373 self.is_e8m0 = is_e8m0
374 self.has_col_major_scales = has_col_major_scales
375 self.is_tma_aligned = is_tma_aligned
376 super().__init__(
377 epsilon,
378 key,
379 has_col_major_scales=has_col_major_scales,
380 is_e8m0=is_e8m0,
381 is_tma_aligned=is_tma_aligned,
382 )
383
384 def register(self, pm_pass: PatternMatcherPass) -> None:
385 def pattern(
386 input: torch.Tensor,
387 weight: torch.Tensor,
388 residual: torch.Tensor,
389 scale: torch.Tensor,
390 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
391 result_rms, residual = vllm.ir.ops.fused_add_rms_norm(
392 input, residual, weight, self.epsilon
393 )
394 result = torch.empty(
395 result_rms.shape,
396 device=result_rms.device,
397 dtype=self.quant_matcher.quant_key.dtype,
398 )
399 assert scale is not None
400 fp8_min, fp8_max = get_fp8_min_max()
401
402 _, result, scale = auto_functionalized(
403 self.quant_matcher.QUANT_OP,
404 input=result_rms,
405 output_q=result,
406 output_s=scale,
407 group_size=self.quant_matcher.quant_key.scale.group_shape[1],
408 eps=1e-10,
409 fp8_min=fp8_min,
410 fp8_max=fp8_max,
411 scale_ue8m0=self.quant_matcher.is_e8m0,
412 dummy_is_scale_transposed=self.has_col_major_scales,
413 dummy_is_tma_aligned=self.is_tma_aligned,
414 )
415
416 return result, residual, scale
417
418 def replacement(
419 input: torch.Tensor,
420 weight: torch.Tensor,
421 residual: torch.Tensor,
422 scale: torch.Tensor,
423 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
424 # In case we're matching native rms-norm, conversions might be
425 # optimized out. We convert here just to be safe.
426 input = input.to(dtype=self.model_dtype)
427
428 result = torch.empty_like(input, dtype=self.quant_dtype)
429
430 at = auto_functionalized(
431 self.FUSED_OP,
432 result=result,
433 input=input,
434 weight=weight,
435 scale=scale,
436 epsilon=self.epsilon,
437 scale_ub=None,
438 residual=residual,
439 group_size=self.group_shape[1],
440 is_scale_transposed=self.has_col_major_scales,
441 )
442
443 # result, residual, scale
444 return at[1], at[3], at[2]
445
446 inputs = [
447 empty_bf16(5, 16), # input
448 empty_bf16(16), # weight
449 empty_bf16(5, 16), # residual
450 self.quant_matcher.empty_f32(1, 1), # scale
451 ]
452
453 pm.register_replacement(
454 pattern,
455 replacement,
456 inputs,
457 pm.fwd_only,
458 pm_pass,
459 extra_check=_rms_input_weight_dtype_match,
460 )
461
462
463 class RMSNormGroupQuantPattern(RMSNormQuantPattern):
464 def __init__(
465 self,
466 epsilon: float,
467 quant_dtype: torch.dtype,
468 group_shape: GroupShape,
469 symmetric: bool = True,
470 is_e8m0: bool = False,
471 has_col_major_scales: bool = True,
472 is_tma_aligned: bool = True,
473 ) -> None:
474 scale = ScaleDesc(torch.float32, False, group_shape)
475 key = FusedRMSQuantKey(
476 fused_add=False,
477 quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
478 )
479 self.group_shape = group_shape
480 self.has_col_major_scales = has_col_major_scales
481 self.is_tma_aligned = is_tma_aligned
482 super().__init__(
483 epsilon,
484 key,
485 has_col_major_scales=self.has_col_major_scales,
486 is_e8m0=is_e8m0,
487 is_tma_aligned=is_tma_aligned,
488 )
489
490 def register(self, pm_pass: PatternMatcherPass) -> None:
491 def pattern(
492 input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
493 ) -> tuple[torch.Tensor, torch.Tensor]:
494 result_rms = vllm.ir.ops.rms_norm(input, weight, self.epsilon)
495 result = torch.empty(
496 result_rms.shape,
497 device=result_rms.device,
498 dtype=self.quant_matcher.quant_key.dtype,
499 )
500 assert scale is not None
501 fp8_min, fp8_max = get_fp8_min_max()
502
503 _, result, scale = auto_functionalized(
504 self.quant_matcher.QUANT_OP,
505 input=result_rms,
506 output_q=result,
507 output_s=scale,
508 group_size=self.quant_matcher.quant_key.scale.group_shape[1],
509 eps=1e-10,
510 fp8_min=fp8_min,
511 fp8_max=fp8_max,
512 scale_ue8m0=self.quant_matcher.is_e8m0,
513 dummy_is_scale_transposed=self.has_col_major_scales,
514 dummy_is_tma_aligned=self.is_tma_aligned,
515 )
516
517 return result, scale
518
519 def replacement(
520 input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
521 ) -> tuple[torch.Tensor, torch.Tensor]:
522 # In case we're matching native rms-norm, conversions might be
523 # optimized out. We convert here just to be safe.
524 input = input.to(dtype=self.model_dtype)
525
526 result = torch.empty_like(input, dtype=self.quant_dtype)
527 at = auto_functionalized(
528 self.FUSED_OP,
529 result=result,
530 input=input,
531 weight=weight,
532 scale=scale,
533 epsilon=self.epsilon,
534 scale_ub=None,
535 residual=None,
536 group_size=self.group_shape[1],
537 is_scale_transposed=self.has_col_major_scales,
538 )
539
540 # result, scale
541 return at[1], at[2]
542
543 pm.register_replacement(
544 pattern,
545 replacement,
546 [
547 empty_bf16(5, 16), # input
548 empty_bf16(16), # weight
549 self.quant_matcher.empty_f32(1, 1), # scale
550 ],
551 pm.fwd_only,
552 pm_pass,
553 extra_check=_rms_input_weight_dtype_match,
554 )
555
556
557 class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
558 def __init__(
559 self,
560 epsilon: float,
561 quant_dtype: torch.dtype,
562 group_shape: GroupShape = GroupShape.PER_TOKEN,
563 symmetric: bool = True,
564 ) -> None:
565 scale = ScaleDesc(torch.float32, False, group_shape)
566 key = FusedRMSQuantKey(
567 fused_add=False,
568 quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
569 )
570 super().__init__(epsilon, key)
571
572 def register(self, pm_pass: PatternMatcherPass) -> None:
573 def pattern(
574 input: torch.Tensor, weight: torch.Tensor
575 ) -> tuple[torch.Tensor, torch.Tensor]:
576 result_rms = vllm.ir.ops.rms_norm(input, weight, self.epsilon)
577 # result, scale
578 return self.quant_matcher(result_rms) # type: ignore[no-any-return]
579
580 def replacement(
581 input: torch.Tensor, weight: torch.Tensor
582 ) -> tuple[torch.Tensor, torch.Tensor]:
583 # In case we're matching native rms-norm, conversions might be
584 # optimized out. We convert here just to be safe.
585 input = input.to(dtype=self.model_dtype)
586
587 result = torch.empty_like(input, dtype=self.quant_dtype)
588 scale = self.quant_matcher.make_scale(input)
589 at = auto_functionalized(
590 self.FUSED_OP,
591 result=result,
592 input=input,
593 weight=weight,
594 scale=scale,
595 epsilon=self.epsilon,
596 scale_ub=None,
597 residual=None,
598 )
599
600 # result, scale
601 return at[1], at[2]
602
603 pm.register_replacement(
604 pattern,
605 replacement,
606 [
607 empty_bf16(5, 16), # input
608 empty_bf16(16), # weight
609 ],
610 pm.fwd_only,
611 pm_pass,
612 extra_check=_rms_input_weight_dtype_match,
613 )
614
615
616 class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
617 def __init__(
618 self,
619 epsilon: float,
620 quant_dtype: torch.dtype,
621 group_shape: GroupShape = GroupShape.PER_TOKEN,
622 symmetric: bool = True,
623 ) -> None:
624 scale = ScaleDesc(torch.float32, False, group_shape)
625 key = FusedRMSQuantKey(
626 fused_add=True,
627 quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
628 )
629 super().__init__(epsilon, key)
630
631 def register(self, pm_pass: PatternMatcherPass) -> None:
632 def pattern(
633 input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
634 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
635 result_rms, residual = vllm.ir.ops.fused_add_rms_norm(
636 input, residual, weight, self.epsilon
637 )
638 result, scale = self.quant_matcher(result_rms)
639
640 return result, residual, scale
641
642 def replacement(
643 input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
644 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
645 # In case we're matching native rms-norm, conversions might be
646 # optimized out. We convert here just to be safe.
647 input = input.to(dtype=self.model_dtype)
648
649 result = torch.empty_like(input, dtype=self.quant_dtype)
650 scale = self.quant_matcher.make_scale(input)
651 at = auto_functionalized(
652 self.FUSED_OP,
653 result=result,
654 input=input,
655 weight=weight,
656 scale=scale,
657 epsilon=self.epsilon,
658 scale_ub=None,
659 residual=residual,
660 )
661
662 # result, residual, scale
663 return at[1], at[3], at[2]
664
665 inputs = [
666 empty_bf16(5, 16), # input
667 empty_bf16(16), # weight
668 empty_bf16(5, 16), # residual
669 ]
670
671 pm.register_replacement(
672 pattern,
673 replacement,
674 inputs,
675 pm.fwd_only,
676 pm_pass,
677 extra_check=_rms_input_weight_dtype_match,
678 )
679
680
681 class FusedAddRMSNormNvfp4QuantPattern:
682 """Fuse add-RMSNorm with NVFP4 quantization for either scale layout."""
683
684 def __init__(self, epsilon: float, is_sf_swizzled_layout: bool) -> None:
685 assert _FLASHINFER_NVFP4_RMS_QUANT_OP is not None
686 self.epsilon = epsilon
687 self.is_sf_swizzled_layout = is_sf_swizzled_layout
688 self.FUSED_OP = _FLASHINFER_NVFP4_RMS_QUANT_OP
689
690 def register(self, pm_pass: PatternMatcherPass) -> None:
691 def pattern(
692 result: torch.Tensor,
693 result_block_scale: torch.Tensor,
694 input: torch.Tensor,
695 weight: torch.Tensor,
696 residual: torch.Tensor,
697 input_global_scale: torch.Tensor,
698 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
699 result_rms, updated_residual = vllm.ir.ops.fused_add_rms_norm(
700 input, residual, weight, self.epsilon
701 )
702 at = auto_functionalized(
703 torch.ops._C.scaled_fp4_quant.out,
704 input=result_rms,
705 input_scale=input_global_scale,
706 is_sf_swizzled_layout=self.is_sf_swizzled_layout,
707 output=result,
708 output_scale=result_block_scale,
709 )
710 return at[1], updated_residual, at[2]
711
712 def replacement(
713 result: torch.Tensor,
714 result_block_scale: torch.Tensor,
715 input: torch.Tensor,
716 weight: torch.Tensor,
717 residual: torch.Tensor,
718 input_global_scale: torch.Tensor,
719 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
720 hidden_size = input.shape[-1]
721 num_tokens = input.numel() // hidden_size
722 # This full-size dummy is required by FlashInfer's TVM-FFI tensor
723 # validation even though output_both_sf_layouts=False leaves it untouched.
724 block_scale_unswizzled = torch.empty(
725 (num_tokens, hidden_size // 16),
726 dtype=torch.float8_e4m3fn,
727 device=input.device,
728 )
729 at = auto_functionalized(
730 self.FUSED_OP,
731 result=result,
732 result_block_scale=result_block_scale,
733 residual=residual,
734 input=input,
735 weight=weight,
736 input_global_scale=input_global_scale,
737 block_scale_unswizzled=block_scale_unswizzled,
738 is_sf_swizzled_layout=self.is_sf_swizzled_layout,
739 epsilon=self.epsilon,
740 )
741 # result, updated residual, block scale in the requested layout
742 return at[1], at[3], at[2]
743
744 inputs = [
745 torch.empty(
746 (5, 32), dtype=torch.uint8, device=current_platform.device_type
747 ),
748 (
749 empty_i32(128, 4)
750 if self.is_sf_swizzled_layout
751 else torch.empty(
752 (5, 4),
753 dtype=torch.uint8,
754 device=current_platform.device_type,
755 )
756 ),
757 empty_bf16(5, 64),
758 empty_bf16(64),
759 empty_bf16(5, 64),
760 empty_fp32(1),
761 ]
762 pm.register_replacement(
763 pattern,
764 replacement,
765 inputs,
766 pm.fwd_only,
767 pm_pass,
768 extra_check=_rms_input_weight_dtype_match,
769 )
770
771
772 class RMSNormQuantFusionPass(VllmPatternMatcherPass):
773 """This pass fuses rms_norm & quant custom ops into a fused rms_norm_quant op.
774 It also supports fused_add_rms_norm.
775 """
776
777 @enable_fake_mode
778 def __init__(self, config: VllmConfig) -> None:
779 super().__init__(config)
780
781 self.patterns: PatternMatcherPass = PatternMatcherPass(
782 pass_name="rmsnorm_quant_fusion_pass"
783 )
784
785 # Make sure fused add patterns are before simple rms norm,
786 # as the latter is a subset of the former in torch ops
787 for epsilon in [1e-5, 1e-6]:
788 if _FLASHINFER_NVFP4_RMS_QUANT_OP is not None and (
789 current_platform.has_device_capability(100)
790 ):
791 for is_sf_swizzled_layout in (True, False):
792 FusedAddRMSNormNvfp4QuantPattern(
793 epsilon, is_sf_swizzled_layout
794 ).register(self.patterns)
795
796 # Fuse fused_add_rms_norm + static fp8 quant
797 FusedAddRMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(
798 self.patterns
799 )
800
801 # Fuse rms_norm + static fp8 quant
802 RMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(self.patterns)
803
804 # Fuse fused_add_rms_norm + dynamic per-token fp8 quant
805 FusedAddRMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(
806 self.patterns
807 )
808
809 # Fuse rms_norm + dynamic per-token fp8 quant
810 RMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(self.patterns)
811
812 # Only register group quant patterns on CUDA/ROCm where the C++ op exists
813 for group_shape in [GroupShape(1, 128), GroupShape(1, 64)]:
814 for has_col_major_scales in [True, False]:
815 for is_e8m0 in [True, False]:
816 for is_tma_aligned in [False, True]:
817 # Fuse fused_add_rms_norm + fp8 group quant
818 FusedAddRMSNormGroupQuantPattern(
819 epsilon,
820 FP8_DTYPE,
821 group_shape=group_shape,
822 is_e8m0=is_e8m0,
823 has_col_major_scales=has_col_major_scales,
824 is_tma_aligned=is_tma_aligned,
825 ).register(self.patterns)
826
827 # Fuse rms_norm + fp8 group quant
828 RMSNormGroupQuantPattern(
829 epsilon,
830 FP8_DTYPE,
831 group_shape=group_shape,
832 is_e8m0=is_e8m0,
833 has_col_major_scales=has_col_major_scales,
834 is_tma_aligned=is_tma_aligned,
835 ).register(self.patterns)
836
837 self.dump_patterns(config, self.patterns)
838
839 @VllmInductorPass.time_and_log
840 def __call__(self, graph: fx.Graph) -> None:
841 self.matched_count = self.patterns.apply(graph)
842 logger.debug("Replaced %s patterns", self.matched_count)
843
844 def uuid(self) -> str:
845 return self.hash_source(
846 self,
847 RMSNormGroupQuantPattern,
848 RMSNormQuantPattern,
849 RMSNormStaticQuantPattern,
850 RMSNormDynamicQuantPattern,
851 FusedAddRMSNormStaticQuantPattern,
852 FusedAddRMSNormDynamicQuantPattern,
853 FusedAddRMSNormGroupQuantPattern,
854 FusedAddRMSNormNvfp4QuantPattern,
855 )