diff --git a/python/sglang/srt/layers/quantization/fp4_utils.py b/python/sglang/srt/layers/quantization/fp4_utils.py index 3e913e137..1ec712fc9 100644 --- a/python/sglang/srt/layers/quantization/fp4_utils.py +++ b/python/sglang/srt/layers/quantization/fp4_utils.py @@ -18,6 +18,7 @@ class Fp4GemmRunnerBackend(Enum): AUTO = "auto" FLASHINFER_CUDNN = "flashinfer_cudnn" + FLASHINFER_CUTEDSL = "flashinfer_cutedsl" FLASHINFER_CUTLASS = "flashinfer_cutlass" FLASHINFER_TRTLLM = "flashinfer_trtllm" @@ -33,6 +34,9 @@ class Fp4GemmRunnerBackend(Enum): def is_flashinfer_trtllm(self) -> bool: return self == Fp4GemmRunnerBackend.FLASHINFER_TRTLLM + def is_flashinfer_cutedsl(self) -> bool: + return self == Fp4GemmRunnerBackend.FLASHINFER_CUTEDSL + def get_flashinfer_backend(self) -> str: """Get the backend string to pass to FlashInfer's mm_fp4 API. @@ -41,7 +45,10 @@ class Fp4GemmRunnerBackend(Enum): 'flashinfer_trtllm' -> 'trtllm' 'flashinfer_cutlass' -> 'cutlass' 'flashinfer_cudnn' -> 'cudnn' + 'flashinfer_cutedsl' -> 'cute-dsl' """ + if self == Fp4GemmRunnerBackend.FLASHINFER_CUTEDSL: + return "cute-dsl" if self.value.startswith("flashinfer_"): return self.value.removeprefix("flashinfer_") else: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e0421138a..5faa2e2cc 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -213,6 +213,7 @@ FP8_GEMM_RUNNER_BACKEND_CHOICES = [ FP4_GEMM_RUNNER_BACKEND_CHOICES = [ "auto", "flashinfer_cudnn", + "flashinfer_cutedsl", "flashinfer_cutlass", "flashinfer_trtllm", ] @@ -4546,7 +4547,8 @@ class ServerArgs: "Options: 'auto' (default; selects flashinfer_cudnn on SM120, flashinfer_cutlass otherwise), " "'flashinfer_cutlass' (CUTLASS backend), " "'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), " - "'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling). " + "'flashinfer_cutedsl' (FlashInfer CuTe DSL backend), " + "'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling), " "NOTE: This replaces the deprecated environment variable " "SGLANG_FLASHINFER_FP4_GEMM_BACKEND.", ) diff --git a/sgl-kernel/benchmark/bench_fp4_gemm.py b/sgl-kernel/benchmark/bench_fp4_gemm.py index a7fc55e0b..278b965c2 100755 --- a/sgl-kernel/benchmark/bench_fp4_gemm.py +++ b/sgl-kernel/benchmark/bench_fp4_gemm.py @@ -1,12 +1,14 @@ import argparse import csv import os +from typing import List, Tuple import torch import triton from flashinfer import mm_fp4 +from flashinfer.testing import bench_gpu_time_with_cupti +from sgl_kernel import cutlass_scaled_fp4_mm, scaled_fp4_quant -from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant from sglang.srt.utils import get_device_capability, is_sm100_supported # CI environment detection @@ -18,25 +20,68 @@ IS_CI = ( FLOAT4_E2M1_MAX = 6.0 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max +# Weight shapes are in the format: ([K, N], TP_SPLIT_DIM) +# TP split dim 0 means split K by tp size; dim 1 means split N by tp size. +DEEPSEEK_R1_MODEL = "deepseek-ai/DeepSeek-R1-0528-FP4" -def get_weight_shapes(args): - models_tps = args.tp_sizes +WEIGHT_SHAPES = { + "meta-llama/Llama-3.1-8B-Instruct": [ + ([4096, 6144], 1), + ([4096, 4096], 0), + ([4096, 28672], 1), + ([14336, 4096], 0), + ], + "meta-llama/Llama-3.3-70B-Instruct": [ + ([8192, 10240], 1), + ([8192, 8192], 0), + ([8192, 57344], 1), + ([28672, 8192], 0), + ], +} - if models_tps == [4]: - return [[1024, 3584], [7168, 256], [7168, 2304], [9216, 3584]] +DEEPSEEK_R1_WEIGHT_SHAPES = { + 4: [[1024, 3584], [7168, 256], [7168, 2304], [9216, 3584]], + 8: [[512, 3584], [7168, 128], [7168, 1152], [4608, 3584]], +} - if models_tps == [8]: - return [[512, 3584], [7168, 128], [7168, 1152], [4608, 3584]] - return [ - [1024, 3584], - [7168, 256], - [7168, 2304], - [9216, 3584], - [512, 3584], - [7168, 128], - [7168, 1152], - [4608, 3584], - ] + +def _bench_cudagraph_with_cupti(fn, quantiles): + times_ms = bench_gpu_time_with_cupti(fn=fn, use_cuda_graph=True) + if not times_ms: + return 0.0, 0.0, 0.0 + quantiles_tensor = torch.tensor(quantiles, dtype=torch.float32) + times_tensor = torch.tensor(times_ms, dtype=torch.float32) + qs = torch.quantile(times_tensor, quantiles_tensor).tolist() + return qs[0], qs[1], qs[2] + + +def get_weight_shapes(args) -> List[Tuple[int, int, str]]: + shapes: List[Tuple[int, int, str]] = [] + for model in args.models: + if model == DEEPSEEK_R1_MODEL: + for tp_size in args.tp_sizes: + if tp_size in DEEPSEEK_R1_WEIGHT_SHAPES: + selected = DEEPSEEK_R1_WEIGHT_SHAPES[tp_size] + else: + selected = ( + DEEPSEEK_R1_WEIGHT_SHAPES[4] + DEEPSEEK_R1_WEIGHT_SHAPES[8] + ) + for n, packed_k in selected: + shapes.append((n, packed_k, model)) + continue + + if model not in WEIGHT_SHAPES: + raise ValueError(f"Unsupported model: {model}") + for tp_size in args.tp_sizes: + for k_n, tp_split_dim in WEIGHT_SHAPES[model]: + k, n = k_n + if tp_split_dim == 0: + k = k // tp_size + else: + n = n // tp_size + packed_k = k // 2 + shapes.append((n, packed_k, model)) + return shapes # CI environment uses simplified parameters @@ -70,12 +115,13 @@ else: # x_vals = [64], x_log=False, line_arg="provider", - line_vals=["sglang_cutlass", "cutlass", "cudnn", "trtllm", "auto"], + line_vals=["sglang_cutlass", "cutlass", "cudnn", "trtllm", "cute-dsl", "auto"], line_names=[ "sglang cutlass fp4", "flashinfer cutlass fp4", "cudnn fp4", "trtllm fp4", + "cute-dsl fp4", "auto fp4 (cudnn/cutlass)", ], styles=[ @@ -83,6 +129,7 @@ else: ("orange", "solid"), ("blue", "solid"), ("green", "solid"), + ("brown", "solid"), ("purple", "solid"), ], ylabel="latency (ms)", @@ -111,14 +158,14 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): quantiles = [0.5, 0.2, 0.8] if provider == "sglang_cutlass": - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + ms, min_ms, max_ms = _bench_cudagraph_with_cupti( lambda: cutlass_scaled_fp4_mm( a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype ), quantiles=quantiles, ) if provider == "cutlass": - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + ms, min_ms, max_ms = _bench_cudagraph_with_cupti( lambda: mm_fp4( a_fp4, b_fp4.T, @@ -132,7 +179,7 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): quantiles=quantiles, ) if provider == "cudnn": - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + ms, min_ms, max_ms = _bench_cudagraph_with_cupti( lambda: mm_fp4( a_fp4, b_fp4.T, @@ -148,7 +195,7 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): if provider == "trtllm": a_scale_interleaved = a_scale_interleaved.to(torch.uint8) b_scale_interleaved = b_scale_interleaved.to(torch.uint8) - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + ms, min_ms, max_ms = _bench_cudagraph_with_cupti( lambda: mm_fp4( a_fp4, b_fp4.T, @@ -161,8 +208,22 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): ), quantiles=quantiles, ) + if provider == "cute-dsl": + ms, min_ms, max_ms = _bench_cudagraph_with_cupti( + lambda: mm_fp4( + a_fp4, + b_fp4.T, + a_scale_interleaved, + b_scale_interleaved.T, + alpha, + dtype, + res_fi, + backend="cute-dsl", + ), + quantiles=quantiles, + ) if provider == "auto": - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + ms, min_ms, max_ms = _bench_cudagraph_with_cupti( lambda: mm_fp4( a_fp4, b_fp4.T, @@ -215,6 +276,13 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): if __name__ == "__main__": parser = argparse.ArgumentParser() + parser.add_argument( + "--models", + nargs="+", + type=str, + default=[DEEPSEEK_R1_MODEL], + help="List of models to benchmark. Supported: Llama 8B/70B and deepseek-ai/DeepSeek-R1-0528-FP4.", + ) parser.add_argument( "--tp-sizes", nargs="+", @@ -226,7 +294,7 @@ if __name__ == "__main__": "--dtype", type=torch.dtype, default=torch.bfloat16, - help="Data type", + help="Output data type", ) parser.add_argument( "--correctness", @@ -267,8 +335,8 @@ if __name__ == "__main__": if IS_CI: NKs = NKs[:2] # Only test first 2 shapes in CI - for N, K in NKs: - print(f"DeepSeek-R1-0528-FP4 N={N} K={K}: ") + for N, K, model_name in NKs: + print(f"{model_name} N={N} packed_k={K}: ") benchmark.run( print_data=True, N=N, diff --git a/sgl-kernel/benchmark/bench_nvfp4_scaled_gemm.py b/sgl-kernel/benchmark/bench_nvfp4_scaled_gemm.py deleted file mode 100644 index eeb5842ed..000000000 --- a/sgl-kernel/benchmark/bench_nvfp4_scaled_gemm.py +++ /dev/null @@ -1,192 +0,0 @@ -import argparse -import copy -import itertools -import os - -import torch -import triton - -from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant -from sglang.srt.utils import get_device_capability - -# CI environment detection -IS_CI = ( - os.getenv("CI", "false").lower() == "true" - or os.getenv("GITHUB_ACTIONS", "false").lower() == "true" -) - -FLOAT4_E2M1_MAX = 6.0 -FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max - -# Weight Shapes are in the format -# ([K, N], TP_SPLIT_DIM) -# Example: -# A shape of ([14336, 4096], 0) indicates the following GEMM shape, -# - TP1 : K = 14336, N = 4096 -# - TP2 : K = 7168, N = 4096 -# A shape of ([4096, 6144], 1) indicates the following GEMM shape, -# - TP1 : K = 4096, N = 6144 -# - TP4 : K = 4096, N = 1536 - -# TP1 shapes -WEIGHT_SHAPES = { - "meta-llama/Llama-3.1-8B-Instruct": [ - ([4096, 6144], 1), - ([4096, 4096], 0), - ([4096, 28672], 1), - ([14336, 4096], 0), - ], - "meta-llama/Llama-3.3-70B-Instruct": [ - ([8192, 10240], 1), - ([8192, 8192], 0), - ([8192, 57344], 1), - ([28672, 8192], 0), - ], - "mistralai/Mistral-Large-Instruct-2407": [ - ([12288, 14336], 1), - ([12288, 12288], 0), - ([12288, 57344], 1), - ([28672, 12288], 0), - ], - "Qwen/Qwen2.5-7B-Instruct": [ - ([3584, 4608], 1), - ([3584, 3584], 0), - ([3584, 37888], 1), - ([18944, 3584], 0), - ], - "Qwen/Qwen2.5-32B-Instruct": [ - ([5120, 7168], 1), - ([5120, 5120], 0), - ([5120, 55296], 1), - ([27648, 5120], 0), - ], - "Qwen/Qwen2.5-72B-Instruct": [ - ([8192, 10240], 1), - ([8192, 8192], 0), - ([8192, 59136], 1), - ([29568, 8192], 0), - ], - "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": [ - ([2048, 3072], 1), - ([2048, 4096], 1), - ([2048, 2048], 0), - ([2048, 576], 0), - ([2048, 21888], 1), - ([10944, 2048], 0), - ([2048, 2816], 1), - ([1408, 2048], 0), - ], -} - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["batch_size"], - x_vals=[1, 16, 64, 128, 256, 512, 1024, 2048], - x_log=False, - line_arg="provider", - line_vals=[ - "sglang-fp4-fp16", - "sglang-fp4-bf16", - ], - line_names=[ - "sglang-fp4-fp16", - "sglang-fp4-bf16", - ], - styles=[("green", "-"), ("blue", "-")], - ylabel="TFLOPS", - plot_name="fp4 block scaled matmul", - args={}, - ) -) -def benchmark(batch_size, provider, N, K): - # M, N, K = batch_size, 4096, 8192 - run_step = 100 - dtype = torch.float16 if "fp16" in provider else torch.bfloat16 - M = batch_size - a = torch.randn((M, K), dtype=dtype, device="cuda") - b = torch.randn((N, K), dtype=dtype, device="cuda") - a_global_scale = ( - (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1) - ).to(torch.float32) - b_global_scale = ( - (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b.flatten(), dim=-1) - ).to(torch.float32) - alpha = 1.0 / (a_global_scale * b_global_scale) - a_fp4, a_scale_interleaved = scaled_fp4_quant(a, a_global_scale) - b_fp4, b_scale_interleaved = scaled_fp4_quant(b, b_global_scale) - - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - - # Bridging the gap between CPU and GPU - for _ in range(25): - c = a @ b.t() - # Warmup - for _ in range(5): - cutlass_scaled_fp4_mm( - a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype - ) - start_event.record() - for _ in range(run_step): - cutlass_scaled_fp4_mm( - a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype - ) - end_event.record() - end_event.synchronize() - torch.cuda.synchronize() - ms = start_event.elapsed_time(end_event) / run_step - - tflops = lambda ms: (2 * M * N * K) * 1e-9 / ms - return tflops(ms) - - -def prepare_shapes(args): - KN_model_names = [] - models_tps = list(itertools.product(args.models, args.tp_sizes)) - for model, tp_size in models_tps: - assert model in WEIGHT_SHAPES - for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model]): - KN[tp_split_dim] = KN[tp_split_dim] // tp_size - KN.append(model) - KN_model_names.append(KN) - return KN_model_names - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--models", - nargs="+", - type=str, - default=["meta-llama/Llama-3.1-8B-Instruct"], - help="List of models to benchmark", - ) - parser.add_argument( - "--tp-sizes", - nargs="+", - type=int, - default=[1], - help="List of tensor parallel sizes", - ) - args = parser.parse_args() - - # Check architecture compatibility - FP4 operations require sm100a/sm103a - major, minor = get_device_capability() - if major is None or major < 10: # Requires compute capability 10.0+ (sm100a/sm103a) - print("Skipping NVIDIA FP4 scaled GEMM benchmark") - if major is not None: - print(f"FP4 operations require sm100a/sm103a, but found sm{major}{minor}") - else: - print("Could not determine device capability") - else: - KN_model_names = prepare_shapes(args) - - # Limit iterations in CI - if IS_CI: - KN_model_names = KN_model_names[:2] # Only test first 2 shapes in CI - - for K, N, model_name in KN_model_names: - print(f"{model_name} N={N} K={K}: ") - benchmark.run(print_data=True, N=N, K=K) - print("Benchmark finished!") diff --git a/test/registered/quant/test_nvfp4_gemm.py b/test/registered/quant/test_nvfp4_gemm.py index 1a94b6b48..a18c94c92 100644 --- a/test/registered/quant/test_nvfp4_gemm.py +++ b/test/registered/quant/test_nvfp4_gemm.py @@ -81,5 +81,10 @@ class TestFP4GemmFlashinferTrtllm(FP4GemmBase, unittest.TestCase): backend = "flashinfer_trtllm" +@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher") +class TestFP4GemmFlashinferCutedsl(FP4GemmBase, unittest.TestCase): + backend = "flashinfer_cutedsl" + + if __name__ == "__main__": unittest.main()