From 9b3d4963ea626a4dcaeb33bf67cf8fdb65cf1f7f Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Mon, 3 Aug 2026 19:03:23 -0700 Subject: [PATCH] =?UTF-8?q?Revert=20"Add=20flashinfer=20rmsnorm=20+=20quan?= =?UTF-8?q?t=20fusion=20support=20SM90,=20SM100,=20SM120=20(#32=E2=80=A6"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 39609837537fbae03eb148339d2406651bac8ba4. --- .../kernels/bench_fused_rmsnorm_fp8_quant.py | 185 ------------ .../kernels/aot/benchmark/bench_fp8_gemm.py | 30 +- .../kernels/aot/csrc/gemm/fp8_gemm_kernel.cu | 93 ++---- .../sglang/kernels/aot/tests/test_fp8_gemm.py | 87 ------ 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 | 75 ++--- 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 +----------------- 14 files changed, 65 insertions(+), 1032 deletions(-) delete mode 100644 benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py delete 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 deleted file mode 100644 index 53e0b8c757e1..000000000000 --- a/benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py +++ /dev/null @@ -1,185 +0,0 @@ -"""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/kernels/aot/benchmark/bench_fp8_gemm.py b/python/sglang/kernels/aot/benchmark/bench_fp8_gemm.py index 0132395001b0..2205d8e392a9 100644 --- a/python/sglang/kernels/aot/benchmark/bench_fp8_gemm.py +++ b/python/sglang/kernels/aot/benchmark/bench_fp8_gemm.py @@ -8,7 +8,9 @@ import triton from sgl_kernel import fp8_scaled_mm as sgl_scaled_mm -from sglang.kernels.ops.quantization.per_tensor_quant_fp8 import per_tensor_quant_fp8 +from sglang.kernels.ops.quantization.per_tensor_quant_fp8 import ( + per_tensor_quant_fp8, +) from sglang.utils import is_in_ci # Optional vLLM import @@ -104,7 +106,7 @@ def sglang_scaled_fp8_quant( if IS_CI: batch_sizes = [1] # Single batch size for CI else: - batch_sizes = [1, 2, 8, 16, 64, 128, 256, 512, 1024, 2048] + batch_sizes = [1, 16, 64, 128, 256, 512, 1024, 2048] # Filter line_vals based on vLLM availability if VLLM_AVAILABLE: @@ -113,39 +115,24 @@ def sglang_scaled_fp8_quant( "vllm-fp8-bf16", "sglang-fp8-fp16", "sglang-fp8-bf16", - "sglang-scalar-a-fp8-fp16", - "sglang-scalar-a-fp8-bf16", ] line_names = [ "vllm-fp8-fp16", "vllm-fp8-bf16", "sglang-fp8-fp16", "sglang-fp8-bf16", - "sglang-scalar-a-fp8-fp16", - "sglang-scalar-a-fp8-bf16", - ] - styles = [ - ("green", "-"), - ("green", "--"), - ("blue", "-"), - ("blue", "--"), - ("red", "-"), - ("red", "--"), ] + styles = [("green", "-"), ("green", "--"), ("blue", "-"), ("blue", "--")] else: line_vals = [ "sglang-fp8-fp16", "sglang-fp8-bf16", - "sglang-scalar-a-fp8-fp16", - "sglang-scalar-a-fp8-bf16", ] line_names = [ "sglang-fp8-fp16", "sglang-fp8-bf16", - "sglang-scalar-a-fp8-fp16", - "sglang-scalar-a-fp8-bf16", ] - styles = [("blue", "-"), ("blue", "--"), ("red", "-"), ("red", "--")] + styles = [("blue", "-"), ("blue", "--")] @triton.testing.perf_report( @@ -187,9 +174,8 @@ def benchmark(batch_size, provider, N, K): lambda: vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype), quantiles=quantiles, ) - elif "sglang" in provider: - a_scale = scale_a_scalar if "scalar-a" in provider else scale_a - a_fp8, scale_a_fp8 = sglang_scaled_fp8_quant(a, a_scale) + elif "sglang-fp8" in provider: + a_fp8, scale_a_fp8 = sglang_scaled_fp8_quant(a, scale_a) b_fp8, scale_b_fp8 = sglang_scaled_fp8_quant(b, scale_b) b_fp8 = b_fp8.t() ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( diff --git a/python/sglang/kernels/aot/csrc/gemm/fp8_gemm_kernel.cu b/python/sglang/kernels/aot/csrc/gemm/fp8_gemm_kernel.cu index 0a3895e48af8..ca3946764350 100644 --- a/python/sglang/kernels/aot/csrc/gemm/fp8_gemm_kernel.cu +++ b/python/sglang/kernels/aot/csrc/gemm/fp8_gemm_kernel.cu @@ -448,22 +448,19 @@ template < typename MainloopScheduleType, typename EpilogueScheduleType, typename TileSchedulerType = void, - bool WithBias = false, - bool ScalarA = false> + bool WithBias = false> struct DeviceGemmFp8RowwiseSm100 { static_assert(std::is_same_v, "ElementType must be FP8(e4m3)"); using TileShape = CTAShape; using Accum = cutlass::epilogue::fusion::Sm90AccFetch; using ElementComputeEpilogue = float; - using VectorScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast< + using ScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast< 0, TileShape, ElementComputeEpilogue, ElementComputeEpilogue, cute::Stride, cute::Int<0>, cute::Int<0>>>; - using ScalarScaleA = cutlass::epilogue::fusion::Sm90ScalarBroadcast; - using ScaleA = std::conditional_t; using ScaleB = cutlass::epilogue::fusion::Sm90RowBroadcast< 0, @@ -554,11 +551,7 @@ struct DeviceGemmFp8RowwiseSm100 { auto* data_ptr = static_cast(tensor.data_ptr()); static_assert( std::is_same_v || std::is_same_v || std::is_same_v); - if constexpr (std::is_same_v) { - return Arguments{{}, {data_ptr}, {}}; - } else { - return Arguments{data_ptr}; - } + return Arguments{data_ptr}; } public: @@ -664,7 +657,7 @@ void launch_sm100_fp8_scaled_mm( TORCH_CHECK(status == cutlass::Status::kSuccess) } -template +template void sm100_fp8_dispatch_bias( torch::Tensor& out, const torch::Tensor& a, @@ -702,8 +695,7 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - true, - ScalarA>; + true>; using BiasGemm256 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -713,8 +705,7 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - true, - ScalarA>; + true>; using BiasGemm64 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -724,8 +715,7 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - true, - ScalarA>; + true>; using BiasGemm16 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -735,8 +725,7 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - true, - ScalarA>; + true>; // Gemm type without bias using GemmDefault = DeviceGemmFp8RowwiseSm100< @@ -748,8 +737,7 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - false, - ScalarA>; + false>; using Gemm256 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -759,8 +747,7 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - false, - ScalarA>; + false>; using Gemm64 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -770,8 +757,7 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - false, - ScalarA>; + false>; using Gemm16 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -781,8 +767,7 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - false, - ScalarA>; + false>; // next power of 2 (minimum 16) uint32_t const m = a.size(0); @@ -826,10 +811,7 @@ void sm100_fp8_dispatch_shape( const torch::Tensor& scales_a, const torch::Tensor& scales_b, const c10::optional& bias) { - if (scales_a.numel() == 1) { - return sm100_fp8_dispatch_bias(out, a, b, scales_a, scales_b, bias); - } - return sm100_fp8_dispatch_bias(out, a, b, scales_a, scales_b, bias); + return sm100_fp8_dispatch_bias(out, a, b, scales_a, scales_b, bias); } template < @@ -841,22 +823,19 @@ template < typename MainloopScheduleType, typename EpilogueScheduleType, typename TileSchedulerType = void, - bool WithBias = false, - bool ScalarA = false> + bool WithBias = false> struct DeviceGemmFp8RowwiseSm120 { static_assert(std::is_same_v, "ElementType must be FP8(e4m3)"); using TileShape = CTAShape; using Accum = cutlass::epilogue::fusion::Sm90AccFetch; using ElementComputeEpilogue = float; - using VectorScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast< + using ScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast< 0, TileShape, ElementComputeEpilogue, ElementComputeEpilogue, cute::Stride, cute::Int<0>, cute::Int<0>>>; - using ScalarScaleA = cutlass::epilogue::fusion::Sm90ScalarBroadcast; - using ScaleA = std::conditional_t; using ScaleB = cutlass::epilogue::fusion::Sm90RowBroadcast< 0, @@ -947,11 +926,7 @@ struct DeviceGemmFp8RowwiseSm120 { auto* data_ptr = static_cast(tensor.data_ptr()); static_assert( std::is_same_v || std::is_same_v || std::is_same_v); - if constexpr (std::is_same_v) { - return Arguments{{}, {data_ptr}, {}}; - } else { - return Arguments{data_ptr}; - } + return Arguments{data_ptr}; } public: @@ -1057,7 +1032,7 @@ void launch_sm120_fp8_scaled_mm( TORCH_CHECK(status == cutlass::Status::kSuccess) } -template +template void sm120_fp8_dispatch_bias( torch::Tensor& out, const torch::Tensor& a, @@ -1085,8 +1060,7 @@ void sm120_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - true, - ScalarA>; + true>; using GemmDefault = DeviceGemmFp8RowwiseSm120< ElementInput, @@ -1097,8 +1071,7 @@ void sm120_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - false, - ScalarA>; + false>; if (bias) { return launch_sm120_fp8_scaled_mm(out, a, b, scales_a, scales_b, bias); @@ -1115,10 +1088,7 @@ void sm120_fp8_dispatch_shape( const torch::Tensor& scales_a, const torch::Tensor& scales_b, const c10::optional& bias) { - if (scales_a.numel() == 1) { - return sm120_fp8_dispatch_bias(out, a, b, scales_a, scales_b, bias); - } - return sm120_fp8_dispatch_bias(out, a, b, scales_a, scales_b, bias); + return sm120_fp8_dispatch_bias(out, a, b, scales_a, scales_b, bias); } #endif @@ -1145,26 +1115,7 @@ torch::Tensor fp8_scaled_mm( TORCH_CHECK(mat_b.scalar_type() == torch::kFloat8_e4m3fn, "mat_b must be Float8_e4m3fn"); TORCH_CHECK(out_dtype == torch::kHalf || out_dtype == torch::kBFloat16, "out_dtype must be Half or BFloat16"); - auto sm_version = getSMVersion(); - TORCH_CHECK( - scales_a.numel() == 1 || scales_a.numel() == mat_a.size(0), - "scales_a must contain either one scalar scale or one scale per row; got ", - scales_a.numel(), - " elements for M=", - mat_a.size(0)); - bool scalar_a_scale_supported = false; -#if defined CUDA_VERSION && CUDA_VERSION >= 12000 - scalar_a_scale_supported = sm_version == 90; -#endif -#if defined CUDA_VERSION && CUDA_VERSION >= 12080 - scalar_a_scale_supported = scalar_a_scale_supported || sm_version >= 100; -#endif - TORCH_CHECK( - scales_a.numel() != 1 || mat_a.size(0) == 1 || scalar_a_scale_supported, - "scalar scales_a with M > 1 is unsupported on SM", - sm_version, - " for this build; got M=", - mat_a.size(0)); + TORCH_CHECK(scales_a.numel() == mat_a.size(0), "size of scales_a is not matched"); TORCH_CHECK(scales_b.numel() == mat_b.size(1), "size of scales_b is not matched"); TORCH_CHECK(scales_a.is_contiguous(), "scales_a must be contiguous"); TORCH_CHECK(scales_b.is_contiguous(), "scales_b msut be contiguous"); @@ -1180,6 +1131,8 @@ torch::Tensor fp8_scaled_mm( torch::Tensor out = torch::empty({mat_a.size(0), mat_b.size(1)}, mat_a.options().dtype(out_dtype)); TORCH_CHECK((out.size(1) * out.element_size()) % 16 == 0, "out must be multiple of 16 bytes for memory alignment"); + auto sm_version = getSMVersion(); + #if defined CUDA_VERSION && CUDA_VERSION >= 12080 if (sm_version >= 120) { if (out_dtype == torch::kBFloat16) { diff --git a/python/sglang/kernels/aot/tests/test_fp8_gemm.py b/python/sglang/kernels/aot/tests/test_fp8_gemm.py index 0d6ec145e9cd..5f8fbdbf4a02 100644 --- a/python/sglang/kernels/aot/tests/test_fp8_gemm.py +++ b/python/sglang/kernels/aot/tests/test_fp8_gemm.py @@ -5,22 +5,6 @@ from sgl_kernel import fp8_scaled_mm -def _cuda_version_at_least(major, minor): - if torch.version.cuda is None: - return False - version = tuple(int(component) for component in torch.version.cuda.split(".")[:2]) - return version >= (major, minor) - - -def _native_scalar_a_supported(): - if not torch.cuda.is_available(): - return False - capability = torch.cuda.get_device_capability() - if capability == (9, 0): - return _cuda_version_at_least(12, 0) - return capability[0] in (10, 12) and _cuda_version_at_least(12, 8) - - def torch_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias): o = torch.matmul(a.to(torch.float32), b.to(torch.float32)) o = o.to(torch.float32) @@ -54,40 +38,6 @@ def _test_accuracy_once(M, N, K, with_bias, out_dtype, device): print(f"M={M}, N={N}, K={K}, with_bias={with_bias}, out_dtype={out_dtype}: OK") -def _test_scalar_a_accuracy_once(M, N, K, with_bias, out_dtype, device): - fp8_info = torch.finfo(torch.float8_e4m3fn) - a_fp8 = ( - torch.randn(M, K, dtype=torch.float32, device=device) - .clamp(min=fp8_info.min, max=fp8_info.max) - .to(torch.float8_e4m3fn) - ) - b_fp8 = ( - torch.randn(N, K, dtype=torch.float32, device=device) - .clamp(min=fp8_info.min, max=fp8_info.max) - .to(torch.float8_e4m3fn) - .t() - ) - scale_a = torch.tensor([0.03125], device=device, dtype=torch.float32) - scale_a_repeated = scale_a.repeat(M) - - # Resemble merged projections whose component matrices were quantized with - # different tensorwise scales before concatenation. - scale_b = torch.empty(N, device=device, dtype=torch.float32) - first_boundary = N // 3 - second_boundary = 2 * N // 3 - scale_b[:first_boundary] = 0.015625 - scale_b[first_boundary:second_boundary] = 0.03125 - scale_b[second_boundary:] = 0.0625 - - bias = torch.randn(N, device=device, dtype=out_dtype) if with_bias else None - expected = torch_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias) - actual = fp8_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias) - repeated = fp8_scaled_mm(a_fp8, b_fp8, scale_a_repeated, scale_b, out_dtype, bias) - - torch.testing.assert_close(expected, actual, rtol=0.02, atol=1) - torch.testing.assert_close(repeated, actual, rtol=0, atol=0) - - @pytest.mark.parametrize("M", [1, 128, 512, 1024, 4096]) @pytest.mark.parametrize("N", [16, 128, 512, 1024, 4096]) @pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 16384]) @@ -142,43 +92,6 @@ def test_accuracy_sm90_swap_ab(shape_mn, K, with_bias, out_dtype): _test_accuracy_once(M, N, K, with_bias, out_dtype, "cuda") -@pytest.mark.skipif( - not _native_scalar_a_supported(), - reason="native scalar A scales require a compatible SM90, SM100, or SM120 build", -) -@pytest.mark.parametrize("M", [1, 2, 8, 16, 64, 189]) -@pytest.mark.parametrize("with_bias", [True, False]) -@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16]) -def test_scalar_a_channelwise_b(M, with_bias, out_dtype): - _test_scalar_a_accuracy_once(M, 6144, 4096, with_bias, out_dtype, "cuda") - - -def test_rejects_invalid_a_scale_count(): - M, N, K = 8, 128, 512 - a = torch.randn(M, K, device="cuda").to(torch.float8_e4m3fn) - b = torch.randn(N, K, device="cuda").to(torch.float8_e4m3fn).t() - scale_a = torch.ones(2, device="cuda", dtype=torch.float32) - scale_b = torch.ones(N, device="cuda", dtype=torch.float32) - - with pytest.raises(RuntimeError, match="scales_a must contain either"): - fp8_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16, None) - - -@pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (8, 9), - reason="SM89-specific scalar A validation", -) -def test_rejects_scalar_a_with_multiple_rows_on_sm89(): - M, N, K = 8, 128, 512 - a = torch.randn(M, K, device="cuda").to(torch.float8_e4m3fn) - b = torch.randn(N, K, device="cuda").to(torch.float8_e4m3fn).t() - scale_a = torch.ones(1, device="cuda", dtype=torch.float32) - scale_b = torch.ones(N, device="cuda", dtype=torch.float32) - - with pytest.raises(RuntimeError, match="scalar scales_a with M > 1 is unsupported"): - fp8_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16, None) - - PRODUCTION_LIKE_FP8_GEMM_CASES = [ (189, 4608, 8192, False, torch.bfloat16), (3330, 256, 8192, False, torch.bfloat16), diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index a126d39f1356..67661147f35d 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -55,7 +55,6 @@ _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: @@ -84,19 +83,8 @@ def layernorm( _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, @@ -169,7 +157,6 @@ def is_supported_rmsnorm_hf_hidden_size(d: int) -> bool: logger = logging.getLogger(__name__) - if _is_npu: import torch_npu from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm @@ -367,57 +354,6 @@ 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(MultiPlatformOp): def __init__( self, @@ -471,7 +407,6 @@ def forward_cuda( 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: @@ -501,20 +436,6 @@ def forward_cuda( 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 ( @@ -572,7 +493,6 @@ def forward_npu( 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: @@ -588,7 +508,6 @@ def forward_aiter( 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 @@ -665,7 +584,6 @@ def forward_hip( 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: @@ -705,7 +623,6 @@ def forward_musa( 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) @@ -729,7 +646,6 @@ def forward_native( 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() @@ -780,7 +696,6 @@ def forward_cpu( 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: @@ -801,7 +716,6 @@ def forward_xpu( 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) @@ -855,65 +769,6 @@ def forward_with_allreduce_fusion_quant_per_group( 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(MultiPlatformOp): 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 a829de606a93..edc720f2ae23 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,23 +231,6 @@ def apply_weights( 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 243cc0dfde6f..3e142d29d967 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -139,7 +139,10 @@ 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 ( @@ -1025,24 +1028,6 @@ def apply( 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, @@ -1839,7 +1824,9 @@ def _quantize_with_flashinfer_trtllm(weight: torch.Tensor): ) 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 20dca003e012..38a524beaf98 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -59,7 +59,6 @@ _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() @@ -313,9 +312,7 @@ def is_aiter(self) -> bool: if is_blackwell_supported() and is_flashinfer_available(): - from flashinfer import ( - SfLayout, - ) + from flashinfer import SfLayout from flashinfer import bmm_fp8 as _raw_flashinfer_bmm_fp8 from flashinfer import mm_mxfp8 as _raw_flashinfer_mm_mxfp8 from flashinfer import mxfp8_quantize as _raw_flashinfer_mxfp8_quantize @@ -1553,7 +1550,9 @@ 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 @@ -1842,7 +1841,6 @@ 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. @@ -1859,42 +1857,10 @@ def apply_fp8_linear( input_2d = input.view(-1, input.shape[-1]) output_shape = [*input.shape[:-1], weight.shape[1]] - # 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: + if compressed_tensor_quant: # Maybe apply padding to output, see comment in __init__ num_token_padding = output_padding - if channelwise_cutlass: + if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]: 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. @@ -1923,19 +1889,13 @@ 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=channelwise_cutlass and not native_scalar_a_scale, + input_2d, input_scale, repeat_scale=cutlass_fp8_supported ) else: # default use per-token quantization if dynamic @@ -1956,12 +1916,13 @@ def apply_fp8_linear( input_2d, group_size=input_2d.shape[1] ) - if channelwise_cutlass: - if not use_cutlass_channelwise_gemm: + 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: # Massage the input to be 2D qinput = qinput.view(-1, qinput.shape[-1]) output = triton_scaled_mm( - qinput, weight, x_scale, weight_scale, output_dtype, bias + qinput, weight, x_scale, weight_scale, input.dtype, bias ) else: output = fp8_scaled_mm( @@ -1969,7 +1930,7 @@ def apply_fp8_linear( weight, x_scale, weight_scale, - out_dtype=output_dtype, + out_dtype=input.dtype, bias=bias, ) return output.view(*output_shape) @@ -2002,7 +1963,7 @@ def apply_fp8_linear( WQ=weight.T, x_scale=x_scale, w_scale=weight_scale, - dtype=output_dtype, + dtype=input.dtype, ) if bias is not None: output += bias @@ -2018,7 +1979,7 @@ def apply_fp8_linear( output = torch._scaled_mm( qinput, weight, - out_dtype=output_dtype, + out_dtype=input.dtype, scale_a=x_scale, scale_b=weight_scale.t(), bias=bias, @@ -2032,7 +1993,7 @@ def apply_fp8_linear( output = torch._scaled_mm( qinput, weight, - out_dtype=output_dtype, + out_dtype=input.dtype, scale_a=x_scale, scale_b=weight_scale, bias=bias, @@ -2061,7 +2022,7 @@ def apply_fp8_linear( input_2d.shape, output_shape, bias, - output_dtype, + input.dtype, ) diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py index c3349251ded5..9209979e7721 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -348,13 +348,9 @@ def forward( # Self Attention if residual is None: residual = hidden_states - hidden_states = self.input_layernorm( - hidden_states, quant_linear=self.self_attn.qkv_proj - ) + hidden_states = self.input_layernorm(hidden_states) else: - hidden_states, residual = self.input_layernorm( - hidden_states, residual, quant_linear=self.self_attn.qkv_proj - ) + hidden_states, residual = self.input_layernorm(hidden_states, residual) hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, @@ -362,9 +358,7 @@ def forward( ) # Fully Connected - hidden_states, residual = self.post_attention_layernorm( - hidden_states, residual, quant_linear=self.mlp.gate_up_proj - ) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) 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 5b7b95d47301..30dc09b7cec4 100644 --- a/python/sglang/srt/models/llama_eagle.py +++ b/python/sglang/srt/models/llama_eagle.py @@ -50,7 +50,7 @@ def __init__( # 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, quant_linear=None: x) + setattr(self, "input_layernorm", lambda x: x) class LlamaModel(nn.Module): diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 228ac550d64d..0fb1592cec09 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -291,13 +291,9 @@ def forward( # Self Attention if residual is None: residual = hidden_states - hidden_states = self.input_layernorm( - hidden_states, quant_linear=self.self_attn.qkv_proj - ) + hidden_states = self.input_layernorm(hidden_states) else: - hidden_states, residual = self.input_layernorm( - hidden_states, residual, quant_linear=self.self_attn.qkv_proj - ) + hidden_states, residual = self.input_layernorm(hidden_states, residual) hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, @@ -305,9 +301,7 @@ def forward( ) # Fully Connected - hidden_states, residual = self.post_attention_layernorm( - hidden_states, residual, quant_linear=self.mlp.gate_up_proj - ) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) 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 20c4a152a556..017e75bc8784 100644 --- a/python/sglang/srt/models/qwen2_eagle.py +++ b/python/sglang/srt/models/qwen2_eagle.py @@ -51,7 +51,7 @@ def __init__( # 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, quant_linear=None: x) + setattr(self, "input_layernorm", lambda x: x) class Qwen2Model(nn.Module): diff --git a/test/registered/layers/test_layernorm_fusion.py b/test/registered/layers/test_layernorm_fusion.py deleted file mode 100644 index 3cf4d581afc6..000000000000 --- a/test/registered/layers/test_layernorm_fusion.py +++ /dev/null @@ -1,142 +0,0 @@ -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 49885ee4be08..34dbc9ccd9fd 100644 --- a/test/registered/quant/test_fp8_utils.py +++ b/test/registered/quant/test_fp8_utils.py @@ -1,6 +1,4 @@ import unittest -from types import SimpleNamespace -from unittest.mock import patch import torch @@ -12,7 +10,7 @@ from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase -register_cuda_ci(est_time=12, stage="base-b", runner_config="1-gpu-large") +register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-large") class TestInverseTransformScaleUe8m0(CustomTestCase): @@ -45,269 +43,5 @@ def test_round_trip(self): ), 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()