From f311cfc7c41304dd911879f5dfef2d36101a5f2c Mon Sep 17 00:00:00 2001 From: Qidong Su Date: Sat, 18 Jul 2026 19:30:56 -0700 Subject: [PATCH] Fuse add-RMSNorm with NVFP4 quantization --- .../passes/fusion/rms_quant_fusion.py | 169 ++++++++++++++++++ 1 file changed, 169 insertions(+) diff --git a/vllm/compilation/passes/fusion/rms_quant_fusion.py b/vllm/compilation/passes/fusion/rms_quant_fusion.py index 670349a08b2a..88aa1993b5ed 100644 --- a/vllm/compilation/passes/fusion/rms_quant_fusion.py +++ b/vllm/compilation/passes/fusion/rms_quant_fusion.py @@ -26,6 +26,8 @@ kStaticTensorScale, ) from vllm.platforms import current_platform +from vllm.utils.flashinfer import has_flashinfer +from vllm.utils.torch_utils import direct_register_custom_op from ..inductor_pass import enable_fake_mode from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass @@ -42,6 +44,72 @@ _FUSED_ADD_RMS_NORM_OP = torch.ops.vllm_ir.fused_add_rms_norm.default +def _flashinfer_fused_add_rms_norm_nvfp4_quant( + result: torch.Tensor, + result_block_scale: torch.Tensor, + residual: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, + input_global_scale: torch.Tensor, + block_scale_unswizzled: torch.Tensor, + is_sf_swizzled_layout: bool, + epsilon: float, +) -> None: + """FlashInfer fused add + RMSNorm + NVFP4 quantization.""" + from flashinfer import add_rmsnorm_fp4quant + + add_rmsnorm_fp4quant( + input, + residual, + weight, + y_fp4=result.view(torch.float4_e2m1fn_x2), + block_scale=result_block_scale.view(torch.float8_e4m3fn), + global_scale=input_global_scale.reshape(1), + eps=epsilon, + block_size=16, + scale_format="e4m3", + is_sf_swizzled_layout=is_sf_swizzled_layout, + output_both_sf_layouts=False, + block_scale_unswizzled=block_scale_unswizzled, + ) + + +def _flashinfer_fused_add_rms_norm_nvfp4_quant_fake( + result: torch.Tensor, + result_block_scale: torch.Tensor, + residual: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, + input_global_scale: torch.Tensor, + block_scale_unswizzled: torch.Tensor, + is_sf_swizzled_layout: bool, + epsilon: float, +) -> None: + return None + + +_FLASHINFER_NVFP4_RMS_QUANT_OP: OpOverload | None = None +if ( + current_platform.is_cuda() + and hasattr(torch, "float4_e2m1fn_x2") + and has_flashinfer() +): + try: + from flashinfer import add_rmsnorm_fp4quant as _add_rmsnorm_fp4quant # noqa: F401 + except ImportError: + pass + else: + direct_register_custom_op( + op_name="flashinfer_fused_add_rms_norm_nvfp4_quant", + op_func=_flashinfer_fused_add_rms_norm_nvfp4_quant, + mutates_args=["result", "result_block_scale", "residual"], + fake_impl=_flashinfer_fused_add_rms_norm_nvfp4_quant_fake, + ) + _FLASHINFER_NVFP4_RMS_QUANT_OP = ( + torch.ops.vllm.flashinfer_fused_add_rms_norm_nvfp4_quant.default + ) + + # TODO: extend rmsnorm quant kernels to support mixed input/weight dtypes, # and remove this check. def _rms_input_weight_dtype_match(match: pm.Match) -> bool: @@ -615,6 +683,95 @@ def replacement( ) +class FusedAddRMSNormNvfp4QuantPattern: + """Fuse add-RMSNorm with NVFP4 quantization for either scale layout.""" + + def __init__(self, epsilon: float, is_sf_swizzled_layout: bool) -> None: + assert _FLASHINFER_NVFP4_RMS_QUANT_OP is not None + self.epsilon = epsilon + self.is_sf_swizzled_layout = is_sf_swizzled_layout + self.FUSED_OP = _FLASHINFER_NVFP4_RMS_QUANT_OP + + def register(self, pm_pass: PatternMatcherPass) -> None: + def pattern( + result: torch.Tensor, + result_block_scale: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + input_global_scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + result_rms, updated_residual = vllm.ir.ops.fused_add_rms_norm( + input, residual, weight, self.epsilon + ) + at = auto_functionalized( + torch.ops._C.scaled_fp4_quant.out, + input=result_rms, + input_scale=input_global_scale, + is_sf_swizzled_layout=self.is_sf_swizzled_layout, + output=result, + output_scale=result_block_scale, + ) + return at[1], updated_residual, at[2] + + def replacement( + result: torch.Tensor, + result_block_scale: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + input_global_scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + hidden_size = input.shape[-1] + num_tokens = input.numel() // hidden_size + block_scale_unswizzled = torch.empty( + (num_tokens, hidden_size // 16), + dtype=torch.float8_e4m3fn, + device=input.device, + ) + at = auto_functionalized( + self.FUSED_OP, + result=result, + result_block_scale=result_block_scale, + residual=residual, + input=input, + weight=weight, + input_global_scale=input_global_scale, + block_scale_unswizzled=block_scale_unswizzled, + is_sf_swizzled_layout=self.is_sf_swizzled_layout, + epsilon=self.epsilon, + ) + # result, updated residual, block scale in the requested layout + return at[1], at[3], at[2] + + inputs = [ + torch.empty( + (5, 32), dtype=torch.uint8, device=current_platform.device_type + ), + ( + empty_i32(128, 4) + if self.is_sf_swizzled_layout + else torch.empty( + (5, 4), + dtype=torch.uint8, + device=current_platform.device_type, + ) + ), + empty_bf16(5, 64), + empty_bf16(64), + empty_bf16(5, 64), + empty_fp32(1), + ] + pm.register_replacement( + pattern, + replacement, + inputs, + pm.fwd_only, + pm_pass, + extra_check=_rms_input_weight_dtype_match, + ) + + class RMSNormQuantFusionPass(VllmPatternMatcherPass): """ This pass fuses rms_norm & quant custom ops into a fused rms_norm_quant op. @@ -632,6 +789,15 @@ def __init__(self, config: VllmConfig) -> None: # Make sure fused add patterns are before simple rms norm, # as the latter is a subset of the former in torch ops for epsilon in [1e-5, 1e-6]: + if ( + _FLASHINFER_NVFP4_RMS_QUANT_OP is not None + and current_platform.has_device_capability(100) + ): + for is_sf_swizzled_layout in (True, False): + FusedAddRMSNormNvfp4QuantPattern( + epsilon, is_sf_swizzled_layout + ).register(self.patterns) + # Fuse fused_add_rms_norm + static fp8 quant FusedAddRMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register( self.patterns @@ -690,4 +856,7 @@ def uuid(self) -> str: FusedAddRMSNormStaticQuantPattern, FusedAddRMSNormDynamicQuantPattern, FusedAddRMSNormGroupQuantPattern, + FusedAddRMSNormNvfp4QuantPattern, + _flashinfer_fused_add_rms_norm_nvfp4_quant, + _flashinfer_fused_add_rms_norm_nvfp4_quant_fake, )