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 000000000000..53e0b8c757e1 --- /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/kernels/aot/benchmark/bench_fp8_gemm.py b/python/sglang/kernels/aot/benchmark/bench_fp8_gemm.py index 2205d8e392a9..0132395001b0 100644 --- a/python/sglang/kernels/aot/benchmark/bench_fp8_gemm.py +++ b/python/sglang/kernels/aot/benchmark/bench_fp8_gemm.py @@ -8,9 +8,7 @@ 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 @@ -106,7 +104,7 @@ def sglang_scaled_fp8_quant( if IS_CI: batch_sizes = [1] # Single batch size for CI else: - batch_sizes = [1, 16, 64, 128, 256, 512, 1024, 2048] + batch_sizes = [1, 2, 8, 16, 64, 128, 256, 512, 1024, 2048] # Filter line_vals based on vLLM availability if VLLM_AVAILABLE: @@ -115,24 +113,39 @@ 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", "--")] + styles = [("blue", "-"), ("blue", "--"), ("red", "-"), ("red", "--")] @triton.testing.perf_report( @@ -174,8 +187,9 @@ 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-fp8" in provider: - a_fp8, scale_a_fp8 = sglang_scaled_fp8_quant(a, scale_a) + 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) 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 ca3946764350..0a3895e48af8 100644 --- a/python/sglang/kernels/aot/csrc/gemm/fp8_gemm_kernel.cu +++ b/python/sglang/kernels/aot/csrc/gemm/fp8_gemm_kernel.cu @@ -448,19 +448,22 @@ template < typename MainloopScheduleType, typename EpilogueScheduleType, typename TileSchedulerType = void, - bool WithBias = false> + bool WithBias = false, + bool ScalarA = 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 ScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast< + using VectorScaleA = 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, @@ -551,7 +554,11 @@ struct DeviceGemmFp8RowwiseSm100 { auto* data_ptr = static_cast(tensor.data_ptr()); static_assert( std::is_same_v || std::is_same_v || std::is_same_v); - return Arguments{data_ptr}; + if constexpr (std::is_same_v) { + return Arguments{{}, {data_ptr}, {}}; + } else { + return Arguments{data_ptr}; + } } public: @@ -657,7 +664,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, @@ -695,7 +702,8 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - true>; + true, + ScalarA>; using BiasGemm256 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -705,7 +713,8 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - true>; + true, + ScalarA>; using BiasGemm64 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -715,7 +724,8 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - true>; + true, + ScalarA>; using BiasGemm16 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -725,7 +735,8 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - true>; + true, + ScalarA>; // Gemm type without bias using GemmDefault = DeviceGemmFp8RowwiseSm100< @@ -737,7 +748,8 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - false>; + false, + ScalarA>; using Gemm256 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -747,7 +759,8 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - false>; + false, + ScalarA>; using Gemm64 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -757,7 +770,8 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - false>; + false, + ScalarA>; using Gemm16 = DeviceGemmFp8RowwiseSm100< ElementInput, ElementOutput, @@ -767,7 +781,8 @@ void sm100_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - false>; + false, + ScalarA>; // next power of 2 (minimum 16) uint32_t const m = a.size(0); @@ -811,7 +826,10 @@ void sm100_fp8_dispatch_shape( const torch::Tensor& scales_a, const torch::Tensor& scales_b, const c10::optional& bias) { - return sm100_fp8_dispatch_bias(out, a, b, scales_a, scales_b, 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); } template < @@ -823,19 +841,22 @@ template < typename MainloopScheduleType, typename EpilogueScheduleType, typename TileSchedulerType = void, - bool WithBias = false> + bool WithBias = false, + bool ScalarA = 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 ScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast< + using VectorScaleA = 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, @@ -926,7 +947,11 @@ struct DeviceGemmFp8RowwiseSm120 { auto* data_ptr = static_cast(tensor.data_ptr()); static_assert( std::is_same_v || std::is_same_v || std::is_same_v); - return Arguments{data_ptr}; + if constexpr (std::is_same_v) { + return Arguments{{}, {data_ptr}, {}}; + } else { + return Arguments{data_ptr}; + } } public: @@ -1032,7 +1057,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, @@ -1060,7 +1085,8 @@ void sm120_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - true>; + true, + ScalarA>; using GemmDefault = DeviceGemmFp8RowwiseSm120< ElementInput, @@ -1071,7 +1097,8 @@ void sm120_fp8_dispatch_bias( MainloopScheduleType, EpilogueScheduleType, TileSchedulerType, - false>; + false, + ScalarA>; if (bias) { return launch_sm120_fp8_scaled_mm(out, a, b, scales_a, scales_b, bias); @@ -1088,7 +1115,10 @@ void sm120_fp8_dispatch_shape( const torch::Tensor& scales_a, const torch::Tensor& scales_b, const c10::optional& bias) { - return sm120_fp8_dispatch_bias(out, a, b, scales_a, scales_b, 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); } #endif @@ -1115,7 +1145,26 @@ 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"); - TORCH_CHECK(scales_a.numel() == mat_a.size(0), "size of scales_a is not matched"); + 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_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"); @@ -1131,8 +1180,6 @@ 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 5f8fbdbf4a02..0d6ec145e9cd 100644 --- a/python/sglang/kernels/aot/tests/test_fp8_gemm.py +++ b/python/sglang/kernels/aot/tests/test_fp8_gemm.py @@ -5,6 +5,22 @@ 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) @@ -38,6 +54,40 @@ 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]) @@ -92,6 +142,43 @@ 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 67661147f35d..a126d39f1356 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -55,6 +55,7 @@ _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 @@ 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, @@ -157,6 +169,7 @@ 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 @@ -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(MultiPlatformOp): def __init__( self, @@ -407,6 +471,7 @@ 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: @@ -436,6 +501,20 @@ 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 ( @@ -493,6 +572,7 @@ 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: @@ -508,6 +588,7 @@ 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 @@ -584,6 +665,7 @@ 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: @@ -623,6 +705,7 @@ 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) @@ -646,6 +729,7 @@ 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() @@ -696,6 +780,7 @@ 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: @@ -716,6 +801,7 @@ 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) @@ -769,6 +855,65 @@ 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 edc720f2ae23..a829de606a93 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 @@ 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 3e142d29d967..243cc0dfde6f 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -139,10 +139,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 ( @@ -1028,6 +1025,24 @@ 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, @@ -1824,9 +1839,7 @@ 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 38a524beaf98..20dca003e012 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -59,6 +59,7 @@ _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() @@ -312,7 +313,9 @@ 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 @@ -1550,9 +1553,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 @@ -1841,6 +1842,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. @@ -1857,10 +1859,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. @@ -1889,13 +1923,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 @@ -1916,13 +1956,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( @@ -1930,7 +1969,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) @@ -1963,7 +2002,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 @@ -1979,7 +2018,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, @@ -1993,7 +2032,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, @@ -2022,7 +2061,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 9209979e7721..c3349251ded5 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -348,9 +348,13 @@ def forward( # 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 @@ def forward( ) # 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 30dc09b7cec4..5b7b95d47301 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: 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 0fb1592cec09..228ac550d64d 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -291,9 +291,13 @@ def forward( # 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 @@ def forward( ) # 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 017e75bc8784..20c4a152a556 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: 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 000000000000..3cf4d581afc6 --- /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 34dbc9ccd9fd..49885ee4be08 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.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 @@ 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()