diff --git a/python/sglang/kernels/aot/benchmark/bench_fp8_gemm.py b/python/sglang/kernels/aot/benchmark/bench_fp8_gemm.py index 2205d8e39..013239500 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 torch 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 @@ if VLLM_AVAILABLE: "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 ca3946764..0a3895e48 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 5f8fbdbf4..0d6ec145e 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 @@ import torch 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),