runtime: Add flashinfer rmsnorm + quant fusion support SM90, SM100, SM120- #32994 (#33471)

Signed-off-by: Devashish Lal <devcode@fb.com>
Co-authored-by: Devashish Lal <devcode@fb.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

python/sglang/srt/layers/layernorm.py

@@ -55,6 +55,7 @@ _is_cpu_amx_available = cpu_has_amx_support()
 _is_cpu = is_cpu()
 _is_xpu = is_xpu()
 _flashinfer_layernorm_available = False
+_flashinfer_rmsnorm_quant_available = False
 
 if _is_cuda or _is_xpu or _is_musa:
     if _is_flashinfer_available:
@@ -83,8 +84,19 @@ if _is_cuda or _is_xpu or _is_musa:
             _flashinfer_layernorm_available = True
         except (ImportError, AttributeError):
             _flashinfer_layernorm_available = False
+
+        try:
+            from flashinfer.norm import (
+                fused_add_rmsnorm_quant as _flashinfer_fused_add_rmsnorm_quant,
+            )
+            from flashinfer.norm import rmsnorm_quant as _flashinfer_rmsnorm_quant
+
+            _flashinfer_rmsnorm_quant_available = True
+        except (ImportError, AttributeError):
+            _flashinfer_rmsnorm_quant_available = False
     else:
         _flashinfer_layernorm_available = False
+        _flashinfer_rmsnorm_quant_available = False
 
     from sgl_kernel import (
         fused_add_rmsnorm,
@@ -157,6 +169,7 @@ if _is_cuda:
 
 logger = logging.getLogger(__name__)
 
+
 if _is_npu:
     import torch_npu
     from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm
@@ -354,6 +367,57 @@ def _forward_with_allreduce_fusion_quant_per_group(
     return (bf16_out, fp8_out, scale_out), residual_out
 
 
+def _fp8_static_input_scale(linear) -> Optional[torch.Tensor]:
+    """Return the per-tensor static FP8 activation scale of ``linear`` if it is
+    an FP8 linear using static per-tensor activation scaling that can consume a
+    pre-quantized ``(fp8, scale)`` input; otherwise ``None``.
+
+    Recognizes both the native ``Fp8LinearMethod`` (non block/mxfp8/marlin) and
+    the compressed-tensors W8A8-FP8 scheme with a static per-tensor input scale
+    (e.g. RedHatAI ``*-FP8`` checkpoints). The flashinfer fused kernel only
+    supports per-tensor quant, hence the ``numel() == 1`` requirement.
+    """
+    if linear is None:
+        return None
+    quant_method = getattr(linear, "quant_method", None)
+    if quant_method is None:
+        return None
+    if not _is_static_per_tensor_fp8_linear(quant_method, linear):
+        return None
+    input_scale = getattr(linear, "input_scale", None)
+    if input_scale is None or input_scale.numel() != 1:
+        return None
+    return input_scale
+
+
+def _is_static_per_tensor_fp8_linear(quant_method, linear) -> bool:
+    try:
+        from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod
+    except ImportError:
+        Fp8LinearMethod = ()
+    if isinstance(quant_method, Fp8LinearMethod):
+        return not (
+            getattr(quant_method, "block_quant", False)
+            or getattr(quant_method, "use_mxfp8", False)
+            or getattr(quant_method, "use_marlin", False)
+        )
+    try:
+        from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
+            CompressedTensorsLinearMethod,
+        )
+        from sglang.srt.layers.quantization.compressed_tensors.schemes import (
+            CompressedTensorsW8A8Fp8,
+        )
+    except ImportError:
+        return False
+    if isinstance(quant_method, CompressedTensorsLinearMethod):
+        scheme = getattr(linear, "scheme", None)
+        return isinstance(scheme, CompressedTensorsW8A8Fp8) and getattr(
+            scheme, "is_static_input_scheme", False
+        )
+    return False
+
+
 class RMSNorm(BaseFusedOp):
     def __init__(
         self,
@@ -407,6 +471,7 @@ class RMSNorm(BaseFusedOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
+        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if x.numel() == 0:
             if residual is not None:
@@ -436,6 +501,20 @@ class RMSNorm(BaseFusedOp):
             if needs_reshape:
                 out = out.reshape(original_shape)
             return out
+        # Fuse the downstream FP8 static per-tensor activation quant into the
+        # norm when supported. Placed after the empty / variance-override /
+        # batch-invariant guards above (all incompatible with the fused kernel)
+        # and gated on not-HF-cast, so it only runs on the standard RMSNorm path.
+        if (
+            quant_linear is not None
+            and not self.cast_x_before_out_mul
+            and _flashinfer_rmsnorm_quant_available
+        ):
+            scale = _fp8_static_input_scale(quant_linear)
+            if scale is not None:
+                return self.forward_with_per_tensor_quant_fusion(
+                    x, scale, residual, post_residual_addition
+                )
         if self.cast_x_before_out_mul and residual is None:
             # Use HF-semantics kernel (cast to dtype before weight multiply).
             if (
@@ -493,6 +572,7 @@ class RMSNorm(BaseFusedOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
+        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if residual is not None:
             if post_residual_addition is not None:
@@ -508,6 +588,7 @@ class RMSNorm(BaseFusedOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
+        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         # Fix dsv4 dp attenton issue
         # the symptom is torch.AcceleratorError: HIP error: invalid configuration argument
@@ -584,6 +665,7 @@ class RMSNorm(BaseFusedOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
+        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         # Fallback to native implementation if vllm is not available
         if not _has_vllm_rms_norm:
@@ -623,6 +705,7 @@ class RMSNorm(BaseFusedOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
+        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
             return self.forward_native(x, residual, post_residual_addition)
@@ -646,6 +729,7 @@ class RMSNorm(BaseFusedOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
+        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if not x.is_contiguous():
             x = x.contiguous()
@@ -696,6 +780,7 @@ class RMSNorm(BaseFusedOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
+        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if _is_cpu_amx_available:
             if residual is not None:
@@ -716,6 +801,7 @@ class RMSNorm(BaseFusedOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
+        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if self.variance_size_override is not None:
             return self.forward_native(x, residual, post_residual_addition)
@@ -769,6 +855,65 @@ class RMSNorm(BaseFusedOp):
             self, x, residual, self.weight, group_size, use_attn_tp_group, keep_bf16
         )
 
+    def forward_with_per_tensor_quant_fusion(
+        self,
+        x: torch.Tensor,
+        scale: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+        post_residual_addition: Optional[torch.Tensor] = None,
+        fp8_dtype: torch.dtype = torch.float8_e4m3fn,
+    ) -> Union[
+        Tuple[torch.Tensor, torch.Tensor, torch.dtype],
+        Tuple[Tuple[torch.Tensor, torch.Tensor, torch.dtype], torch.Tensor],
+    ]:
+        """Fused RMSNorm + static per-tensor FP8 quantization.
+
+        The normed activation is quantized to ``fp8_dtype`` using the per-tensor
+        reciprocal ``scale`` (same convention as ``static_quant_fp8``:
+        ``q = normed / scale``), so a downstream FP8 linear carrying a matching
+        static ``input_scale`` can skip its own activation quant.
+
+        The quantized activation is emitted as a ``(fp8_out, scale, orig_dtype)``
+        tuple; ``orig_dtype`` (the un-quantized activation dtype) is carried so
+        the downstream FP8 GEMM produces its output in the model's dtype rather
+        than defaulting to bf16.
+
+        Return contract mirrors ``forward``:
+        * no residual -> ``(fp8_out, scale, orig_dtype)``
+        * w/ residual -> ``((fp8_out, scale, orig_dtype), residual_out)``
+        """
+        orig_dtype = x.dtype
+        needs_reshape = x.dim() != 2
+        if needs_reshape:
+            original_shape = x.shape
+            x = x.contiguous().reshape(-1, original_shape[-1])
+        elif not x.is_contiguous():
+            x = x.contiguous()
+
+        out = torch.empty_like(x, dtype=fp8_dtype)
+        if residual is not None:
+            if post_residual_addition is not None:
+                residual = residual + post_residual_addition
+            if residual.dim() != 2:
+                residual = residual.contiguous().reshape(-1, residual.shape[-1])
+            elif not residual.is_contiguous():
+                residual = residual.contiguous()
+            # In-place: residual += x, then out = quant(rmsnorm(residual) * w).
+            _flashinfer_fused_add_rmsnorm_quant(
+                out, x, residual, self.weight.data, scale, self.variance_epsilon
+            )
+            if needs_reshape:
+                out = out.reshape(original_shape)
+                residual = residual.reshape(original_shape)
+            return (out, scale, orig_dtype), residual
+
+        _flashinfer_rmsnorm_quant(
+            out, x, self.weight.data, scale, self.variance_epsilon
+        )
+        if needs_reshape:
+            out = out.reshape(original_shape)
+        return out, scale, orig_dtype
+
 
 class LayerNorm(BaseFusedOp):
     def __init__(

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py

@@ -231,6 +231,23 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsLinearScheme):
         x: torch.Tensor,
         bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
+        if isinstance(x, tuple):
+            # Pre-quantized activation from a fused RMSNorm+FP8 quant kernel:
+            # x = (fp8_input, per_tensor_input_scale[, orig_dtype]).
+            # apply_fp8_linear detects the fp8 dtype and skips re-quantizing;
+            # orig_dtype (when present) sets the GEMM output dtype.
+            qx, x_scale = x[0], x[1]
+            out_dtype = x[2] if len(x) > 2 else None
+            return apply_fp8_linear(
+                input=qx,
+                weight=layer.weight,
+                weight_scale=layer.weight_scale,
+                input_scale=x_scale,
+                bias=bias,
+                use_per_token_if_dynamic=True,
+                compressed_tensor_quant=True,
+                pre_quant_output_dtype=out_dtype,
+            )
         if self.weight_block_size is not None:
             return self.w8a8_block_fp8_linear(
                 input=x,

python/sglang/srt/layers/quantization/fp8.py

@@ -141,10 +141,7 @@ def _require_fp4_dtype():
 
 
 if _use_aiter or _use_hip_int4:
-    from aiter.ops.shuffle import (
-        shuffle_scale,
-        shuffle_weight,
-    )
+    from aiter.ops.shuffle import shuffle_scale, shuffle_weight
 
 if _use_aiter:
     from sglang.srt.layers.quantization.fp8_utils import (
@@ -1035,6 +1032,24 @@ class Fp8LinearMethod(LinearMethodBase):
                 bias=bias,
             )
 
+        if isinstance(x, tuple):
+            # Pre-quantized activation from a fused RMSNorm+FP8 quant kernel:
+            # x = (fp8_input, per_tensor_input_scale[, orig_dtype]).
+            # apply_fp8_linear detects the fp8 dtype and skips re-quantizing;
+            # orig_dtype (when present) sets the GEMM output dtype.
+            qx, x_scale = x[0], x[1]
+            out_dtype = x[2] if len(x) > 2 else None
+            return apply_fp8_linear(
+                input=qx,
+                weight=layer.weight,
+                weight_scale=layer.weight_scale,
+                input_scale=x_scale,
+                bias=bias,
+                cutlass_fp8_supported=self.cutlass_fp8_supported,
+                use_per_token_if_dynamic=self.use_per_token_if_dynamic,
+                pre_quant_output_dtype=out_dtype,
+            )
+
         return apply_fp8_linear(
             input=x,
             weight=layer.weight,
@@ -1838,9 +1853,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             )
             return qweight.view_as(weight), scale_u8
 
-        from sglang.srt.layers.quantization.mxfp8_block_convert import (
-            _ue8m0_to_fp32,
-        )
+        from sglang.srt.layers.quantization.mxfp8_block_convert import _ue8m0_to_fp32
 
         def _quantize_for_deepgemm(weight: torch.Tensor):
             weight = weight.contiguous()

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -57,7 +57,9 @@ logger = logging.getLogger(__name__)
 _is_hip = is_hip()
 _is_cuda = is_cuda()
 _is_fp8_fnuz = is_fp8_fnuz()
+_is_sm90_supported = is_sm90_supported()
 _is_sm100_supported = is_sm100_supported()
+_is_sm120_supported = is_sm120_supported()
 _is_gfx95_supported = is_gfx95_supported()
 _is_musa = is_musa()
 
@@ -1430,9 +1432,7 @@ def requant_block_scale_ue8m0_for_deepgemm(
     scales are not already UE8M0, and DeepGEMM can run the layer (bf16 output,
     aligned shape). Returns True when it requantizes.
     """
-    from sglang.srt.model_loader.utils import (
-        should_deepgemm_weight_requant_ue8m0,
-    )
+    from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0
 
     if (
         not use_deepgemm_runner
@@ -1721,6 +1721,7 @@ def apply_fp8_linear(
     use_per_token_if_dynamic: bool = False,
     pad_output: Optional[bool] = None,
     compressed_tensor_quant: bool = False,
+    pre_quant_output_dtype: Optional[torch.dtype] = None,
 ) -> torch.Tensor:
     # Note: we pad the input because torch._scaled_mm is more performant
     # for matrices with batch dimension > 16.
@@ -1737,10 +1738,42 @@ def apply_fp8_linear(
     input_2d = input.view(-1, input.shape[-1])
     output_shape = [*input.shape[:-1], weight.shape[1]]
 
-    if compressed_tensor_quant:
+    # A pre-quantized fp8 activation (e.g. from a fused RMSNorm+quant kernel)
+    # carries no original dtype: skip re-quant, reuse the supplied per-tensor
+    # input_scale, and emit ``pre_quant_output_dtype`` (the model's activation
+    # dtype, propagated by the producer) or bf16 if it was not provided.
+    input_prequantized = input_2d.dtype in (
+        torch.float8_e4m3fn,
+        torch.float8_e4m3fnuz,
+    )
+    if input_prequantized:
+        output_dtype = pre_quant_output_dtype or torch.bfloat16
+    else:
+        output_dtype = input.dtype
+
+    channelwise_cutlass = (
+        cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]
+    )
+    cutlass_compatible_b = weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
+    use_cutlass_channelwise_gemm = (
+        channelwise_cutlass and cutlass_compatible_b and not use_triton_w8a8_fp8_kernel
+    )
+    native_scalar_a_scale = use_cutlass_channelwise_gemm and (
+        _is_sm90_supported or _is_sm100_supported or _is_sm120_supported
+    )
+
+    if input_prequantized:
+        assert input_scale is not None and input_scale.numel() == 1
+        qinput = input_2d
+        if channelwise_cutlass and not native_scalar_a_scale:
+            # Unsupported CUTLASS epilogues require one A scale per row.
+            x_scale = input_scale.repeat(input_2d.shape[0]).view(-1, 1)
+        else:
+            x_scale = input_scale
+    elif compressed_tensor_quant:
         # Maybe apply padding to output, see comment in __init__
         num_token_padding = output_padding
-        if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]:
+        if channelwise_cutlass:
             num_token_padding = None
         # For static per-tensor activation scales when using inductor compiler,
         # use pure PyTorch ops instead of the opaque sgl_kernel quant kernel.
@@ -1769,13 +1802,19 @@ def apply_fp8_linear(
                 num_token_padding=num_token_padding,
                 use_per_token_if_dynamic=use_per_token_if_dynamic,
             )
+        if (
+            input_scale is not None
+            and channelwise_cutlass
+            and not native_scalar_a_scale
+        ):
+            x_scale = input_scale.repeat(input_2d.shape[0]).view(-1, 1)
     else:
-        # cutlass w8a8 fp8 sgl-kernel only supports per-token scale
         if input_scale is not None:
             assert input_scale.numel() == 1
-            # broadcast per-tensor scale to per-token scale when supporting cutlass
             qinput, x_scale = static_quant_fp8(
-                input_2d, input_scale, repeat_scale=cutlass_fp8_supported
+                input_2d,
+                input_scale,
+                repeat_scale=channelwise_cutlass and not native_scalar_a_scale,
             )
         else:
             # default use per-token quantization if dynamic
@@ -1796,13 +1835,12 @@ def apply_fp8_linear(
                         input_2d, group_size=input_2d.shape[1]
                     )
 
-    if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]:
-        cutlass_compatible_b = weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
-        if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel:
+    if channelwise_cutlass:
+        if not use_cutlass_channelwise_gemm:
             # Massage the input to be 2D
             qinput = qinput.view(-1, qinput.shape[-1])
             output = triton_scaled_mm(
-                qinput, weight, x_scale, weight_scale, input.dtype, bias
+                qinput, weight, x_scale, weight_scale, output_dtype, bias
             )
         else:
             output = fp8_scaled_mm(
@@ -1810,7 +1848,7 @@ def apply_fp8_linear(
                 weight,
                 x_scale,
                 weight_scale,
-                out_dtype=input.dtype,
+                out_dtype=output_dtype,
                 bias=bias,
             )
         return output.view(*output_shape)
@@ -1843,7 +1881,7 @@ def apply_fp8_linear(
                 WQ=weight.T,
                 x_scale=x_scale,
                 w_scale=weight_scale,
-                dtype=input.dtype,
+                dtype=output_dtype,
             )
             if bias is not None:
                 output += bias
@@ -1859,7 +1897,7 @@ def apply_fp8_linear(
             output = torch._scaled_mm(
                 qinput,
                 weight,
-                out_dtype=input.dtype,
+                out_dtype=output_dtype,
                 scale_a=x_scale,
                 scale_b=weight_scale.t(),
                 bias=bias,
@@ -1873,7 +1911,7 @@ def apply_fp8_linear(
         output = torch._scaled_mm(
             qinput,
             weight,
-            out_dtype=input.dtype,
+            out_dtype=output_dtype,
             scale_a=x_scale,
             scale_b=weight_scale,
             bias=bias,
@@ -1902,7 +1940,7 @@ def apply_fp8_linear(
         input_2d.shape,
         output_shape,
         bias,
-        input.dtype,
+        output_dtype,
     )
 
 

python/sglang/srt/models/llama.py

@@ -348,9 +348,13 @@ class LlamaDecoderLayer(nn.Module):
         # Self Attention
         if residual is None:
             residual = hidden_states
-            hidden_states = self.input_layernorm(hidden_states)
+            hidden_states = self.input_layernorm(
+                hidden_states, quant_linear=self.self_attn.qkv_proj
+            )
         else:
-            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+            hidden_states, residual = self.input_layernorm(
+                hidden_states, residual, quant_linear=self.self_attn.qkv_proj
+            )
         hidden_states = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,
@@ -358,7 +362,9 @@ class LlamaDecoderLayer(nn.Module):
         )
 
         # Fully Connected
-        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states, residual = self.post_attention_layernorm(
+            hidden_states, residual, quant_linear=self.mlp.gate_up_proj
+        )
         hidden_states = self.mlp(hidden_states)
         return hidden_states, residual
 

python/sglang/srt/models/llama_eagle.py

@@ -50,7 +50,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
         # https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427
         if layer_id == 0:
             del self.input_layernorm
-            setattr(self, "input_layernorm", lambda x: x)
+            setattr(self, "input_layernorm", lambda x, quant_linear=None: x)
 
 
 class LlamaModel(nn.Module):

python/sglang/srt/models/qwen2.py

@@ -291,9 +291,13 @@ class Qwen2DecoderLayer(nn.Module):
         # Self Attention
         if residual is None:
             residual = hidden_states
-            hidden_states = self.input_layernorm(hidden_states)
+            hidden_states = self.input_layernorm(
+                hidden_states, quant_linear=self.self_attn.qkv_proj
+            )
         else:
-            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+            hidden_states, residual = self.input_layernorm(
+                hidden_states, residual, quant_linear=self.self_attn.qkv_proj
+            )
         hidden_states = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,
@@ -301,7 +305,9 @@ class Qwen2DecoderLayer(nn.Module):
         )
 
         # Fully Connected
-        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states, residual = self.post_attention_layernorm(
+            hidden_states, residual, quant_linear=self.mlp.gate_up_proj
+        )
         hidden_states = self.mlp(hidden_states)
         return hidden_states, residual
 

python/sglang/srt/models/qwen2_eagle.py

@@ -51,7 +51,7 @@ class Qwen2DecoderLayer(Qwen2DecoderLayer):
         # https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427
         if layer_id == 0:
             del self.input_layernorm
-            setattr(self, "input_layernorm", lambda x: x)
+            setattr(self, "input_layernorm", lambda x, quant_linear=None: x)
 
 
 class Qwen2Model(nn.Module):