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>

benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py

@@ -0,0 +1,185 @@
+"""Microbenchmark: fused RMSNorm + static per-tensor FP8 quant, comparing the
+flashinfer default kernels against the CuTe-DSL kernels and the unfused
+baseline (RMSNorm followed by a separate static FP8 quant).
+
+Providers:
+  unfused     RMSNorm.forward_cuda  +  static_quant_fp8
+  fused       flashinfer rmsnorm_quant / fused_add_rmsnorm_quant (default)
+  fused_cute  flashinfer rmsnorm_quant_cute / fused_add_rmsnorm_quant_cute
+
+All fused providers produce an ``(fp8, scale)`` activation (and updated residual
+when a residual is supplied), matching what a downstream FP8 static-per-tensor
+linear consumes. Covers the no-residual and residual (fused-add) cases across a
+few hidden sizes so you can pick the fastest kernel per shape.
+
+Run:
+    python benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py
+"""
+
+import itertools
+
+import numpy as np
+import torch
+import triton
+from flashinfer.norm import fused_add_rmsnorm_quant, rmsnorm_quant
+from flashinfer.testing import bench_gpu_time
+
+from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
+from sglang.srt.layers.layernorm import RMSNorm, _flashinfer_rmsnorm_quant_available
+
+if not torch.cuda.is_available():
+    raise RuntimeError("CUDA is required for this benchmark")
+if not _flashinfer_rmsnorm_quant_available:
+    raise RuntimeError(
+        "flashinfer rmsnorm_quant / fused_add_rmsnorm_quant is not available; "
+        "install flashinfer to benchmark the fused path"
+    )
+
+try:
+    from flashinfer.norm import fused_add_rmsnorm_quant_cute, rmsnorm_quant_cute
+
+    _CUTE_AVAILABLE = True
+except ImportError:
+    _CUTE_AVAILABLE = False
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+FP8_DTYPE = torch.float8_e4m3fn
+HIDDEN_SIZES = [4096, 8192]
+# Per-tensor reciprocal scale (q = normed / scale); 0.05 keeps normed/scale well
+# within the e4m3 range for unit-scale activations.
+SCALE_VALUE = 0.05
+
+
+def make_layer(hidden_size):
+    layer = RMSNorm(hidden_size).to(device=DEVICE, dtype=DTYPE)
+    layer.weight.data.normal_(mean=1.0, std=0.1)
+    return layer
+
+
+def make_inputs(num_tokens, hidden_size, add_residual):
+    x = torch.randn(num_tokens, hidden_size, device=DEVICE, dtype=DTYPE)
+    residual = torch.randn_like(x) if add_residual else None
+    scale = torch.tensor([SCALE_VALUE], device=DEVICE, dtype=torch.float32)
+    return x, residual, scale
+
+
+def run_unfused(layer, x, residual, scale):
+    out = layer(x, residual)
+    if residual is not None:
+        normed, residual_out = out
+        q, q_scale = static_quant_fp8(normed, scale)
+        return (q, q_scale), residual_out
+    q, q_scale = static_quant_fp8(out, scale)
+    return q, q_scale
+
+
+def _run_fused(kernel, add_kernel, layer, x, residual, scale):
+    out = torch.empty_like(x, dtype=FP8_DTYPE)
+    if residual is not None:
+        # In-place: residual += x, then out = quant(rmsnorm(residual) * w).
+        add_kernel(out, x, residual, layer.weight.data, scale, layer.variance_epsilon)
+        return (out, scale), residual
+    kernel(out, x, layer.weight.data, scale, layer.variance_epsilon)
+    return out, scale
+
+
+def run_fused_default(layer, x, residual, scale):
+    return _run_fused(rmsnorm_quant, fused_add_rmsnorm_quant, layer, x, residual, scale)
+
+
+def run_fused_cute(layer, x, residual, scale):
+    return _run_fused(
+        rmsnorm_quant_cute, fused_add_rmsnorm_quant_cute, layer, x, residual, scale
+    )
+
+
+RUNNERS = {
+    "unfused": run_unfused,
+    "fused": run_fused_default,
+    "fused_cute": run_fused_cute,
+}
+
+# (provider key, plot label, style)
+_PROVIDERS = [
+    ("unfused", "rmsnorm + static_quant_fp8 (unfused)", ("blue", "-")),
+    ("fused", "rmsnorm_quant (fused, default)", ("green", "-")),
+]
+if _CUTE_AVAILABLE:
+    _PROVIDERS.append(
+        ("fused_cute", "rmsnorm_quant_cute (fused, cute-dsl)", ("red", "-"))
+    )
+
+
+def _bench_ms(fn, args, quantiles=(0.5, 0.2, 0.8)):
+    # Pass the GPU tensors as input_args so flashinfer's cold_l2_cache flush can
+    # find them; a zero-arg callable trips its "no GPU tensors found" warning and
+    # silently disables cold-L2 timing.
+    times = bench_gpu_time(
+        fn=fn,
+        input_args=args,
+        use_cuda_graph=True,
+        dry_run_time_ms=25,
+        repeat_time_ms=100,
+    )
+    return tuple(float(np.percentile(times, q * 100)) for q in quantiles)
+
+
+def _check_correctness():
+    """One-shot sanity check that every fused provider agrees with the unfused
+    baseline within FP8 precision."""
+    fused_providers = [p for p in RUNNERS if p != "unfused"]
+    for hidden_size, add_residual in itertools.product(HIDDEN_SIZES, [False, True]):
+        layer = make_layer(hidden_size)
+        x, residual, scale = make_inputs(64, hidden_size, add_residual)
+        with torch.inference_mode():
+            ref = run_unfused(
+                layer, x.clone(), residual.clone() if add_residual else None, scale
+            )
+        (uq, _), _ = ref if add_residual else (ref, None)
+        ref_deq = uq.float() * scale
+        for provider in fused_providers:
+            if provider == "fused_cute" and not _CUTE_AVAILABLE:
+                continue
+            with torch.inference_mode():
+                out = RUNNERS[provider](
+                    layer, x.clone(), residual.clone() if add_residual else None, scale
+                )
+            (q, _), _ = out if add_residual else (out, None)
+            cos = torch.nn.functional.cosine_similarity(
+                (q.float() * scale).flatten(), ref_deq.flatten(), dim=0
+            ).item()
+            assert (
+                cos > 0.99
+            ), f"{provider} h={hidden_size} residual={add_residual} cos={cos:.4f}"
+    print("correctness check passed (all fused providers vs unfused within FP8)")
+
+
+configs = [
+    triton.testing.Benchmark(
+        x_names=["num_tokens"],
+        x_vals=[512, 1024, 2048, 4096, 8192, 16384],
+        x_log=False,
+        line_arg="provider",
+        line_vals=[p[0] for p in _PROVIDERS],
+        line_names=[p[1] for p in _PROVIDERS],
+        styles=[p[2] for p in _PROVIDERS],
+        ylabel="latency (ms)",
+        plot_name=f"rmsnorm_fp8_quant_h{hidden_size}_residual{add_residual}",
+        args={"hidden_size": hidden_size, "add_residual": add_residual},
+    )
+    for hidden_size, add_residual in itertools.product(HIDDEN_SIZES, [False, True])
+]
+
+
+@triton.testing.perf_report(configs)
+def benchmark(num_tokens, hidden_size, add_residual, provider):
+    layer = make_layer(hidden_size)
+    x, residual, scale = make_inputs(num_tokens, hidden_size, add_residual)
+    return _bench_ms(RUNNERS[provider], (layer, x, residual, scale))
+
+
+if __name__ == "__main__":
+    torch.manual_seed(0)
+    _check_correctness()
+    benchmark.run(print_data=True, show_plots=False)

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):

test/registered/layers/test_layernorm_fusion.py

@@ -0,0 +1,142 @@
+import itertools
+import unittest
+
+import torch
+
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-large")
+
+
+class TestRMSNormFp8QuantFusion(CustomTestCase):
+    DTYPES = [torch.bfloat16, torch.half]
+    NUM_TOKENS = [7, 83, 512]
+    HIDDEN_SIZES = [512, 4096]
+    ADD_RESIDUAL = [False, True]
+    SEED = 0
+    FP8_DTYPE = torch.float8_e4m3fn
+
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA is not available")
+        from sglang.srt.layers.layernorm import _flashinfer_rmsnorm_quant_available
+
+        if not _flashinfer_rmsnorm_quant_available:
+            raise unittest.SkipTest("flashinfer rmsnorm_quant is not available")
+        torch.set_default_device("cuda")
+
+    def _run_fusion_test(self, num_tokens, hidden_size, add_residual, dtype):
+        torch.manual_seed(self.SEED)
+
+        layer = RMSNorm(hidden_size).to(dtype=dtype)
+        layer.weight.data.normal_(mean=1.0, std=0.1)
+        x = torch.randn(num_tokens, hidden_size, dtype=dtype)
+        residual = torch.randn_like(x) if add_residual else None
+        # Per-tensor reciprocal scale (as carried by a static FP8 linear).
+        scale = torch.tensor([0.05], dtype=torch.float32)
+
+        with torch.inference_mode():
+            ref = layer.forward_native(
+                x.clone(), residual.clone() if add_residual else None
+            )
+            normed_ref = ref[0] if add_residual else ref
+            residual_ref = ref[1] if add_residual else None
+
+            result = layer.forward_with_per_tensor_quant_fusion(
+                x.clone(), scale, residual.clone() if add_residual else None
+            )
+
+        if add_residual:
+            (q, s, out_dtype), r = result
+        else:
+            q, s, out_dtype = result
+            r = None
+
+        # Output contract.
+        self.assertEqual(q.dtype, self.FP8_DTYPE)
+        self.assertIs(s, scale)
+        self.assertEqual(out_dtype, dtype)
+        self.assertEqual(tuple(q.shape), (num_tokens, hidden_size))
+        if add_residual:
+            self.assertEqual(r.dtype, dtype)
+            self.assertTrue(
+                torch.allclose(r.float(), residual_ref.float(), atol=1e-2, rtol=1e-2)
+            )
+
+        # Numerical: dequantized (q * scale) matches the reference normed output
+        # within FP8 e4m3 precision.
+        deq = q.float() * scale
+        ref_flat = normed_ref.float().flatten()
+        cos = torch.nn.functional.cosine_similarity(deq.flatten(), ref_flat, dim=0)
+        self.assertGreater(cos.item(), 0.99)
+        rel_err = (
+            deq.flatten() - ref_flat
+        ).abs().mean() / ref_flat.abs().mean().clamp_min(1e-6)
+        self.assertLess(rel_err.item(), 0.1)
+
+    def test_rms_norm_fp8_quant_fusion(self):
+        for params in itertools.product(
+            self.NUM_TOKENS,
+            self.HIDDEN_SIZES,
+            self.ADD_RESIDUAL,
+            self.DTYPES,
+        ):
+            with self.subTest(
+                num_tokens=params[0],
+                hidden_size=params[1],
+                add_residual=params[2],
+                dtype=params[3],
+            ):
+                self._run_fusion_test(*params)
+
+    def test_forward_cuda_quant_linear_dispatch(self):
+        """forward_cuda routes to the fused path only when applicable."""
+        import sglang.srt.layers.layernorm as ln_mod
+
+        torch.manual_seed(self.SEED)
+        hidden_size, num_tokens = 512, 32
+        x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16)
+        residual = torch.randn_like(x)
+        scale = torch.tensor([0.05], dtype=torch.float32)
+
+        orig_static_scale = ln_mod._fp8_static_input_scale
+        ln_mod._fp8_static_input_scale = lambda linear: scale
+        try:
+            plain = RMSNorm(hidden_size).to(dtype=torch.bfloat16)
+            plain.weight.data.normal_(mean=1.0, std=0.1)
+
+            with torch.inference_mode():
+                # Plain norm -> fused (fp8, scale, dtype) + bf16 residual.
+                (q, s, out_dtype), r = plain(
+                    x.clone(), residual.clone(), quant_linear=object()
+                )
+
+                # variance_size_override is incompatible -> must not fuse.
+                var_layer = RMSNorm(hidden_size, var_hidden_size=hidden_size // 2).to(
+                    dtype=torch.bfloat16
+                )
+                var_out = var_layer(x.clone(), residual.clone(), quant_linear=object())
+                # cast_x_before_out_mul (HF semantics) is incompatible -> must not fuse.
+                cast_layer = RMSNorm(hidden_size, cast_x_before_out_mul=True).to(
+                    dtype=torch.bfloat16
+                )
+                cast_out = cast_layer(
+                    x.clone(), residual.clone(), quant_linear=object()
+                )
+        finally:
+            ln_mod._fp8_static_input_scale = orig_static_scale
+
+        self.assertEqual(q.dtype, self.FP8_DTYPE)
+        self.assertIs(s, scale)
+        self.assertEqual(out_dtype, torch.bfloat16)
+        self.assertEqual(r.dtype, torch.bfloat16)
+
+        self.assertEqual(var_out[0].dtype, torch.bfloat16)
+        self.assertEqual(cast_out[0].dtype, torch.bfloat16)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/quant/test_fp8_utils.py

@@ -1,4 +1,6 @@
 import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
 
 import torch
 
@@ -10,7 +12,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=12, stage="base-b", runner_config="1-gpu-large")
 
 
 class TestInverseTransformScaleUe8m0(CustomTestCase):
@@ -43,5 +45,269 @@ class TestInverseTransformScaleUe8m0(CustomTestCase):
             ), f"{sf_fp32_original=} {sf_fp32_recreated}"
 
 
+class TestApplyFp8LinearScaleDispatch(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA is not available")
+        torch.set_default_device("cuda")
+
+    @staticmethod
+    def _make_inputs(dtype=torch.bfloat16):
+        M, K, N = 8, 16, 32
+        input = torch.randn(M, K, dtype=dtype)
+        qinput = input.to(torch.float8_e4m3fn)
+        weight = torch.randn(N, K).to(torch.float8_e4m3fn).t()
+        input_scale = torch.tensor([0.05], dtype=torch.float32)
+        weight_scale = torch.linspace(0.01, 0.03, N, dtype=torch.float32)
+        return input, qinput, weight, input_scale, weight_scale
+
+    def test_native_scalar_a_static_prequant_and_dynamic_scale_shapes(self):
+        import sglang.srt.layers.quantization.fp8_utils as fp8_utils
+
+        exec_config = SimpleNamespace(
+            graph=SimpleNamespace(
+                cuda_graph_config=SimpleNamespace(
+                    prefill=SimpleNamespace(tc_compiler="none")
+                )
+            )
+        )
+        for capability in (
+            "_is_sm90_supported",
+            "_is_sm100_supported",
+            "_is_sm120_supported",
+        ):
+            with self.subTest(capability=capability):
+                input, qinput, weight, input_scale, weight_scale = self._make_inputs()
+                seen_scales = []
+
+                def fake_fp8_scaled_mm(
+                    mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None
+                ):
+                    seen_scales.append(scales_a)
+                    return torch.empty(
+                        (mat_a.shape[0], mat_b.shape[1]),
+                        dtype=out_dtype,
+                        device=mat_a.device,
+                    )
+
+                capabilities = {
+                    "_is_sm90_supported": False,
+                    "_is_sm100_supported": False,
+                    "_is_sm120_supported": False,
+                }
+                capabilities[capability] = True
+                with patch.multiple(fp8_utils, **capabilities), patch.object(
+                    fp8_utils, "fp8_scaled_mm", side_effect=fake_fp8_scaled_mm
+                ), patch.object(fp8_utils, "get_exec", return_value=exec_config):
+                    fp8_utils.apply_fp8_linear(
+                        input,
+                        weight,
+                        weight_scale,
+                        input_scale=input_scale,
+                        cutlass_fp8_supported=True,
+                    )
+                    fp8_utils.apply_fp8_linear(
+                        input,
+                        weight,
+                        weight_scale,
+                        input_scale=input_scale,
+                        cutlass_fp8_supported=True,
+                        use_per_token_if_dynamic=True,
+                        compressed_tensor_quant=True,
+                    )
+                    fp8_utils.apply_fp8_linear(
+                        qinput,
+                        weight,
+                        weight_scale,
+                        input_scale=input_scale,
+                        cutlass_fp8_supported=True,
+                        pre_quant_output_dtype=input.dtype,
+                    )
+                    fp8_utils.apply_fp8_linear(
+                        input,
+                        weight,
+                        weight_scale,
+                        input_scale=None,
+                        cutlass_fp8_supported=True,
+                        use_per_token_if_dynamic=True,
+                        compressed_tensor_quant=True,
+                    )
+
+                self.assertEqual(seen_scales[0].numel(), 1)
+                self.assertEqual(seen_scales[1].numel(), 1)
+                self.assertIs(seen_scales[2], input_scale)
+                self.assertEqual(tuple(seen_scales[3].shape), (input.shape[0], 1))
+
+    def test_without_native_scalar_a_static_scale_is_repeated(self):
+        import sglang.srt.layers.quantization.fp8_utils as fp8_utils
+
+        input, qinput, weight, input_scale, weight_scale = self._make_inputs()
+        seen_scales = []
+
+        def fake_fp8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):
+            seen_scales.append(scales_a)
+            return torch.empty(
+                (mat_a.shape[0], mat_b.shape[1]), dtype=out_dtype, device=mat_a.device
+            )
+
+        with patch.multiple(
+            fp8_utils,
+            _is_sm90_supported=False,
+            _is_sm100_supported=False,
+            _is_sm120_supported=False,
+        ), patch.object(fp8_utils, "fp8_scaled_mm", side_effect=fake_fp8_scaled_mm):
+            fp8_utils.apply_fp8_linear(
+                input,
+                weight,
+                weight_scale,
+                input_scale=input_scale,
+                cutlass_fp8_supported=True,
+            )
+            fp8_utils.apply_fp8_linear(
+                qinput,
+                weight,
+                weight_scale,
+                input_scale=input_scale,
+                cutlass_fp8_supported=True,
+                pre_quant_output_dtype=input.dtype,
+            )
+
+        self.assertEqual(tuple(seen_scales[0].shape), (input.shape[0], 1))
+        self.assertEqual(tuple(seen_scales[1].shape), (input.shape[0], 1))
+
+    def test_linear_methods_forward_fused_scalar_tuple(self):
+        import sglang.srt.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_fp8 as compressed_fp8
+        import sglang.srt.layers.quantization.fp8 as native_fp8
+
+        input, qinput, weight, input_scale, weight_scale = self._make_inputs(
+            torch.float16
+        )
+
+        class Layer:
+            pass
+
+        layer = Layer()
+        layer.weight = weight
+        layer.weight_scale = weight_scale
+        layer.input_scale = input_scale
+
+        native_method = native_fp8.Fp8LinearMethod.__new__(native_fp8.Fp8LinearMethod)
+        native_method.use_marlin = False
+        native_method.use_mxfp8 = False
+        native_method.block_quant = False
+        native_method.cutlass_fp8_supported = True
+        native_method.use_per_token_if_dynamic = False
+
+        compressed_method = compressed_fp8.CompressedTensorsW8A8Fp8.__new__(
+            compressed_fp8.CompressedTensorsW8A8Fp8
+        )
+        compressed_method.weight_block_size = None
+
+        fused_input = (qinput, input_scale, input.dtype)
+        with patch.object(native_fp8, "apply_fp8_linear") as native_apply:
+            native_apply.return_value = torch.empty(
+                (qinput.shape[0], weight.shape[1]), dtype=input.dtype
+            )
+            native_method.apply(layer, fused_input)
+            self.assertIs(native_apply.call_args.kwargs["input_scale"], input_scale)
+            self.assertEqual(
+                native_apply.call_args.kwargs["pre_quant_output_dtype"], input.dtype
+            )
+
+        with patch.object(compressed_fp8, "apply_fp8_linear") as compressed_apply:
+            compressed_apply.return_value = torch.empty(
+                (qinput.shape[0], weight.shape[1]), dtype=input.dtype
+            )
+            compressed_method.apply_weights(layer, fused_input)
+            self.assertIs(compressed_apply.call_args.kwargs["input_scale"], input_scale)
+            self.assertEqual(
+                compressed_apply.call_args.kwargs["pre_quant_output_dtype"],
+                input.dtype,
+            )
+
+
+class TestApplyFp8LinearPrequantOutputDtype(CustomTestCase):
+    """apply_fp8_linear with a pre-quantized fp8 activation must emit the
+    caller-supplied ``pre_quant_output_dtype`` (the model's activation dtype),
+    not the fp8 input dtype. Regression test for FP16 models where hardcoding
+    bf16 caused a query/key dtype mismatch in attention."""
+
+    DTYPES = [torch.float16, torch.bfloat16]
+    FP8_DTYPE = torch.float8_e4m3fn
+
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA is not available")
+        torch.set_default_device("cuda")
+
+    def _run(self, dtype):
+        from sglang.srt.layers.quantization.fp8_utils import (
+            apply_fp8_linear,
+            cutlass_fp8_supported,
+        )
+
+        torch.manual_seed(0)
+        M, K, N = 33, 512, 256
+        cf = cutlass_fp8_supported()
+        fp8_info = torch.finfo(self.FP8_DTYPE)
+
+        normed = torch.randn(M, K, dtype=dtype)
+        input_scale = torch.tensor([0.05], dtype=torch.float32)
+        # Per-channel fp8 weight in column-major (K, N) layout.
+        w = torch.randn(N, K, dtype=dtype) * 0.05
+        w_scale = (w.abs().amax(dim=1) / fp8_info.max).float()
+        weight = (
+            (w.float() / w_scale[:, None])
+            .clamp(fp8_info.min, fp8_info.max)
+            .to(self.FP8_DTYPE)
+            .t()
+        )
+
+        # Reference: non-pre-quantized input -> output dtype == input dtype.
+        ref = apply_fp8_linear(
+            input=normed,
+            weight=weight,
+            weight_scale=w_scale,
+            input_scale=input_scale,
+            cutlass_fp8_supported=cf,
+        )
+        self.assertEqual(ref.dtype, dtype)
+
+        qinput = (
+            (normed.float() * input_scale.reciprocal())
+            .clamp(fp8_info.min, fp8_info.max)
+            .to(self.FP8_DTYPE)
+        )
+
+        # Pre-quantized input with the dtype propagated -> output matches dtype.
+        out = apply_fp8_linear(
+            input=qinput,
+            weight=weight,
+            weight_scale=w_scale,
+            input_scale=input_scale,
+            cutlass_fp8_supported=cf,
+            pre_quant_output_dtype=dtype,
+        )
+        self.assertEqual(out.dtype, dtype)
+        self.assertTrue(torch.allclose(out.float(), ref.float(), atol=2e-2, rtol=2e-2))
+
+        # Without the dtype hint, the pre-quantized path falls back to bf16.
+        out_default = apply_fp8_linear(
+            input=qinput,
+            weight=weight,
+            weight_scale=w_scale,
+            input_scale=input_scale,
+            cutlass_fp8_supported=cf,
+        )
+        self.assertEqual(out_default.dtype, torch.bfloat16)
+
+    def test_prequant_output_dtype(self):
+        for dtype in self.DTYPES:
+            with self.subTest(dtype=dtype):
+                self._run(dtype)
+
+
 if __name__ == "__main__":
     unittest.main()