From 11d03eaeefd87c867f3fcbdf63f58cc0ae04de39 Mon Sep 17 00:00:00 2001 From: DevashishLal-CB Date: Sun, 9 Aug 2026 05:15:03 -0700 Subject: [PATCH] runtime: Add flashinfer rmsnorm + quant fusion support SM90, SM100, SM120- #32994 (#33471) Signed-off-by: Devashish Lal Co-authored-by: Devashish Lal Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- .../kernels/bench_fused_rmsnorm_fp8_quant.py | 185 ++++++++++++ python/sglang/srt/layers/layernorm.py | 145 ++++++++++ .../schemes/compressed_tensors_w8a8_fp8.py | 17 ++ python/sglang/srt/layers/quantization/fp8.py | 27 +- .../srt/layers/quantization/fp8_utils.py | 72 +++-- python/sglang/srt/models/llama.py | 12 +- python/sglang/srt/models/llama_eagle.py | 2 +- python/sglang/srt/models/qwen2.py | 12 +- python/sglang/srt/models/qwen2_eagle.py | 2 +- .../layers/test_layernorm_fusion.py | 142 ++++++++++ test/registered/quant/test_fp8_utils.py | 268 +++++++++++++++++- 11 files changed, 851 insertions(+), 33 deletions(-) create mode 100644 benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py create mode 100644 test/registered/layers/test_layernorm_fusion.py diff --git a/benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py b/benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py new file mode 100644 index 000000000..53e0b8c75 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 92b26a592..73ac933c7 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/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__( diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py index edc720f2a..a829de606 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index acc40d05c..0498e6be7 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/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() diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 3a31c3905..b772079c5 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/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, ) diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py index 9209979e7..c3349251d 100644 --- a/python/sglang/srt/models/llama.py +++ b/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 diff --git a/python/sglang/srt/models/llama_eagle.py b/python/sglang/srt/models/llama_eagle.py index 30dc09b7c..5b7b95d47 100644 --- a/python/sglang/srt/models/llama_eagle.py +++ b/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): diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 0fb1592ce..228ac550d 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/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 diff --git a/python/sglang/srt/models/qwen2_eagle.py b/python/sglang/srt/models/qwen2_eagle.py index 017e75bc8..20c4a152a 100644 --- a/python/sglang/srt/models/qwen2_eagle.py +++ b/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): diff --git a/test/registered/layers/test_layernorm_fusion.py b/test/registered/layers/test_layernorm_fusion.py new file mode 100644 index 000000000..3cf4d581a --- /dev/null +++ b/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() diff --git a/test/registered/quant/test_fp8_utils.py b/test/registered/quant/test_fp8_utils.py index 34dbc9ccd..49885ee4b 100644 --- a/test/registered/quant/test_fp8_utils.py +++ b/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()