[Kernel] Enable optimized FlashInfer add-RMSNorm NVFP4 fusion (#51925)
Signed-off-by: Qidong Su <qidongs@nvidia.com> Signed-off-by: Qidong Su <soodoshll@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Qidong Su committed
Sep 8, 2026 at 20:20 UTC
9e137d8d230c9d62ba15c2a85f248ae7b0e20e2d
3 files changed
+286
tests/compile/passes/test_fusion.py
+103
@@ -9,14 +9,20 @@ import vllm.config
9
import vllm.ir.ops
10
import vllm.plugins
11
from tests.compile.backend import TestBackend
12
+from tests.kernels.quantization.nvfp4_utils import quant_nvfp4_tensor
13
from tests.utils import TestFP8Layer
14
from vllm._aiter_ops import IS_AITER_FOUND, rocm_aiter_ops
15
+from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
16
from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS
17
from vllm.compilation.passes.fusion.rms_quant_fusion import (
18
FUSED_OPS,
19
FusedRMSQuantKey,
20
RMSNormQuantFusionPass,
21
)
22
+from vllm.compilation.passes.fx_utils import find_auto_fn_maybe
23
+from vllm.compilation.passes.utility.fix_functionalization import (
24
+ FixFunctionalizationPass,
25
+)
26
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
27
from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
28
from vllm.config import (
@@ -244,6 +250,31 @@ class TestModel(torch.nn.Module):
250
]
251
252
253
+class AddRMSNormNvfp4Model(torch.nn.Module):
254
+ def __init__(self, hidden_size: int, eps: float):
255
+ super().__init__()
256
+ self.norm = RMSNorm(hidden_size, eps)
257
+ self.activation_scale = torch.rand((1, 1), dtype=torch.float32)
258
+
259
+ weight = torch.rand((hidden_size, hidden_size))
260
+ self.weight, self.weight_scale, weight_global_scale = quant_nvfp4_tensor(weight)
261
+ self.alpha = (1 / (weight_global_scale * self.activation_scale)).reshape(1)
262
+
263
+ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
264
+ x = residual = torch.relu(x)
265
+ x, residual = self.norm(x, residual)
266
+ x, block_scale = scaled_fp4_quant(x, self.activation_scale)
267
+ x = cutlass_scaled_fp4_mm(
268
+ x,
269
+ self.weight,
270
+ block_scale,
271
+ self.weight_scale,
272
+ self.alpha,
273
+ out_dtype=torch.get_default_dtype(),
274
+ )
275
+ return x, residual
276
+
277
+
278
def _run_fusion_test(
279
model,
280
fusion_pass,
@@ -384,6 +415,78 @@ def test_fusion_rmsnorm_quant(
415
)
416
417
418
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
419
+@pytest.mark.parametrize("eps", [1e-5, 1e-6])
420
+@pytest.mark.parametrize(
421
+ ("num_tokens", "hidden_size"),
422
+ [(1, 256), (32, 1024), (257, 4096)],
423
+)
424
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
425
+def test_fusion_add_rmsnorm_nvfp4_quant(
426
+ dtype: torch.dtype, eps: float, num_tokens: int, hidden_size: int
427
+):
428
+ fused_op = getattr(
429
+ torch.ops.vllm,
430
+ "flashinfer_fused_add_rms_norm_nvfp4_quant",
431
+ None,
432
+ )
433
+ if not current_platform.has_device_capability(100) or fused_op is None:
434
+ pytest.skip("FlashInfer add-RMSNorm NVFP4 fusion is not available")
435
+ assert fused_op is not None
436
+ fused_op_default = fused_op.default
437
+
438
+ vllm_config = VllmConfig(
439
+ model_config=ModelConfig(dtype=dtype),
440
+ compilation_config=CompilationConfig(
441
+ mode=CompilationMode.VLLM_COMPILE,
442
+ pass_config=PassConfig(
443
+ fuse_norm_quant=True,
444
+ eliminate_noops=True,
445
+ ),
446
+ ),
447
+ )
448
+
449
+ with (
450
+ vllm.config.set_current_vllm_config(vllm_config),
451
+ vllm_config.kernel_config.ir_op_priority.set_priority(),
452
+ ):
453
+ torch.set_default_device("cuda")
454
+ torch.set_default_dtype(dtype)
455
+ torch.manual_seed(1)
456
+
457
+ fusion_pass = RMSNormQuantFusionPass(vllm_config)
458
+ fused_backend = TestBackend(
459
+ NoOpEliminationPass(vllm_config),
460
+ fusion_pass,
461
+ PostCleanupPass(vllm_config),
462
+ FixFunctionalizationPass(vllm_config),
463
+ )
464
+ unfused_backend = TestBackend(
465
+ NoOpEliminationPass(vllm_config),
466
+ PostCleanupPass(vllm_config),
467
+ FixFunctionalizationPass(vllm_config),
468
+ )
469
+
470
+ model = AddRMSNormNvfp4Model(hidden_size=hidden_size, eps=eps)
471
+ x = torch.rand(num_tokens, hidden_size)
472
+ torch._dynamo.mark_dynamic(x, 0)
473
+
474
+ result_fused = torch.compile(model, backend=fused_backend)(x)
475
+ result_unfused = torch.compile(model, backend=unfused_backend)(x)
476
+
477
+ torch.testing.assert_close(result_fused, result_unfused, atol=2e-1, rtol=2e-1)
478
+ assert fusion_pass.matched_count == 1
479
+ fused_backend.check_before_ops([torch.ops._C.scaled_fp4_quant.out])
480
+ fused_backend.check_after_ops([fused_op_default])
481
+ assert (
482
+ find_auto_fn_maybe(
483
+ fused_backend.graph_post_pass.nodes,
484
+ fused_op_default,
485
+ )
486
+ is None
487
+ )
488
+
489
+
490
@pytest.mark.parametrize("dtype", [torch.bfloat16])
491
@pytest.mark.parametrize("hidden_size", [256])
492
@pytest.mark.parametrize("num_tokens", [257])
vllm/compilation/passes/fusion/rms_quant_fusion.py
+169
@@ -26,6 +26,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
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
@@ -40,6 +42,73 @@ FP4_DTYPE = torch.uint8
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,
@@ -610,6 +679,97 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
679
)
680
681
682
+class FusedAddRMSNormNvfp4QuantPattern:
683
+ """Fuse add-RMSNorm with NVFP4 quantization for either scale layout."""
684
+
685
+ def __init__(self, epsilon: float, is_sf_swizzled_layout: bool) -> None:
686
+ assert _FLASHINFER_NVFP4_RMS_QUANT_OP is not None
687
+ self.epsilon = epsilon
688
+ self.is_sf_swizzled_layout = is_sf_swizzled_layout
689
+ self.FUSED_OP = _FLASHINFER_NVFP4_RMS_QUANT_OP
690
+
691
+ def register(self, pm_pass: PatternMatcherPass) -> None:
692
+ def pattern(
693
+ result: torch.Tensor,
694
+ result_block_scale: torch.Tensor,
695
+ input: torch.Tensor,
696
+ weight: torch.Tensor,
697
+ residual: torch.Tensor,
698
+ input_global_scale: torch.Tensor,
699
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
700
+ result_rms, updated_residual = vllm.ir.ops.fused_add_rms_norm(
701
+ input, residual, weight, self.epsilon
702
+ )
703
+ at = auto_functionalized(
704
+ torch.ops._C.scaled_fp4_quant.out,
705
+ input=result_rms,
706
+ input_scale=input_global_scale,
707
+ is_sf_swizzled_layout=self.is_sf_swizzled_layout,
708
+ output=result,
709
+ output_scale=result_block_scale,
710
+ )
711
+ return at[1], updated_residual, at[2]
712
+
713
+ def replacement(
714
+ result: torch.Tensor,
715
+ result_block_scale: torch.Tensor,
716
+ input: torch.Tensor,
717
+ weight: torch.Tensor,
718
+ residual: torch.Tensor,
719
+ input_global_scale: torch.Tensor,
720
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
721
+ hidden_size = input.shape[-1]
722
+ num_tokens = input.numel() // hidden_size
723
+ # This full-size dummy is required by FlashInfer's TVM-FFI tensor
724
+ # validation even though output_both_sf_layouts=False leaves it untouched.
725
+ block_scale_unswizzled = torch.empty(
726
+ (num_tokens, hidden_size // 16),
727
+ dtype=torch.float8_e4m3fn,
728
+ device=input.device,
729
+ )
730
+ at = auto_functionalized(
731
+ self.FUSED_OP,
732
+ result=result,
733
+ result_block_scale=result_block_scale,
734
+ residual=residual,
735
+ input=input,
736
+ weight=weight,
737
+ input_global_scale=input_global_scale,
738
+ block_scale_unswizzled=block_scale_unswizzled,
739
+ is_sf_swizzled_layout=self.is_sf_swizzled_layout,
740
+ epsilon=self.epsilon,
741
+ )
742
+ # result, updated residual, block scale in the requested layout
743
+ return at[1], at[3], at[2]
744
+
745
+ inputs = [
746
+ torch.empty(
747
+ (5, 32), dtype=torch.uint8, device=current_platform.device_type
748
+ ),
749
+ (
750
+ empty_i32(128, 4)
751
+ if self.is_sf_swizzled_layout
752
+ else torch.empty(
753
+ (5, 4),
754
+ dtype=torch.uint8,
755
+ device=current_platform.device_type,
756
+ )
757
+ ),
758
+ empty_bf16(5, 64),
759
+ empty_bf16(64),
760
+ empty_bf16(5, 64),
761
+ empty_fp32(1),
762
+ ]
763
+ pm.register_replacement(
764
+ pattern,
765
+ replacement,
766
+ inputs,
767
+ pm.fwd_only,
768
+ pm_pass,
769
+ extra_check=_rms_input_weight_dtype_match,
770
+ )
771
+
772
+
773
class RMSNormQuantFusionPass(VllmPatternMatcherPass):
774
"""
775
This pass fuses rms_norm & quant custom ops into a fused rms_norm_quant op.
@@ -627,6 +787,14 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):
787
# Make sure fused add patterns are before simple rms norm,
788
# as the latter is a subset of the former in torch ops
789
for epsilon in [1e-5, 1e-6]:
790
+ if _FLASHINFER_NVFP4_RMS_QUANT_OP is not None and (
791
+ current_platform.has_device_capability(100)
792
+ ):
793
+ for is_sf_swizzled_layout in (True, False):
794
+ FusedAddRMSNormNvfp4QuantPattern(
795
+ epsilon, is_sf_swizzled_layout
796
+ ).register(self.patterns)
797
+
798
# Fuse fused_add_rms_norm + static fp8 quant
799
FusedAddRMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(
800
self.patterns
@@ -685,4 +853,5 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):
853
FusedAddRMSNormStaticQuantPattern,
854
FusedAddRMSNormDynamicQuantPattern,
855
FusedAddRMSNormGroupQuantPattern,
856
+ FusedAddRMSNormNvfp4QuantPattern,
857
)
vllm/compilation/passes/utility/fix_functionalization.py
+14
@@ -124,6 +124,20 @@ class FixFunctionalizationPass(VllmInductorPass):
124
5: "scale_out",
125
}
126
self.defunctionalize(graph, node, mutated_args)
127
+ elif (
128
+ hasattr(
129
+ torch.ops.vllm,
130
+ "flashinfer_fused_add_rms_norm_nvfp4_quant",
131
+ )
132
+ and at_target
133
+ == torch.ops.vllm.flashinfer_fused_add_rms_norm_nvfp4_quant.default
134
+ ):
135
+ mutated_args = {
136
+ 1: "result",
137
+ 2: "result_block_scale",
138
+ 3: "residual",
139
+ }
140
+ self.defunctionalize(graph, node, mutated_args)
141
# For some reason we need to specify the args for both
142
# silu_and_mul and silu_and_mul_quant. The kwargs
143
# pathway gets the wrong answer.