From 67cad3e69e3c5c001644c4e04af905eeefc6ce80 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Fri, 20 Mar 2026 22:47:47 -0700 Subject: [PATCH] Revert "Support CuteDSL `mm_fp4` backend" (#21077) --- .../srt/layers/quantization/fp4_utils.py | 7 - python/sglang/srt/server_args.py | 4 +- sgl-kernel/benchmark/bench_fp4_gemm.py | 120 +++-------- .../benchmark/bench_nvfp4_scaled_gemm.py | 192 ++++++++++++++++++ test/registered/quant/test_nvfp4_gemm.py | 5 - 5 files changed, 219 insertions(+), 109 deletions(-) create mode 100644 sgl-kernel/benchmark/bench_nvfp4_scaled_gemm.py diff --git a/python/sglang/srt/layers/quantization/fp4_utils.py b/python/sglang/srt/layers/quantization/fp4_utils.py index 1ec712fc9..3e913e137 100644 --- a/python/sglang/srt/layers/quantization/fp4_utils.py +++ b/python/sglang/srt/layers/quantization/fp4_utils.py @@ -18,7 +18,6 @@ class Fp4GemmRunnerBackend(Enum): AUTO = "auto" FLASHINFER_CUDNN = "flashinfer_cudnn" - FLASHINFER_CUTEDSL = "flashinfer_cutedsl" FLASHINFER_CUTLASS = "flashinfer_cutlass" FLASHINFER_TRTLLM = "flashinfer_trtllm" @@ -34,9 +33,6 @@ 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. @@ -45,10 +41,7 @@ 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 11e2fcd6b..8a011005d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -213,7 +213,6 @@ FP8_GEMM_RUNNER_BACKEND_CHOICES = [ FP4_GEMM_RUNNER_BACKEND_CHOICES = [ "auto", "flashinfer_cudnn", - "flashinfer_cutedsl", "flashinfer_cutlass", "flashinfer_trtllm", ] @@ -4605,8 +4604,7 @@ 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_cutedsl' (FlashInfer CuTe DSL backend), " - "'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling), " + "'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 d5d358611..f8f0bd666 100755 --- a/sgl-kernel/benchmark/bench_fp4_gemm.py +++ b/sgl-kernel/benchmark/bench_fp4_gemm.py @@ -1,14 +1,12 @@ 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 from sglang.utils import is_in_ci @@ -17,68 +15,25 @@ IS_CI = is_in_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" -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), - ], -} +def get_weight_shapes(args): + models_tps = args.tp_sizes -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 == [4]: + return [[1024, 3584], [7168, 256], [7168, 2304], [9216, 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 + 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], + ] # CI environment uses simplified parameters @@ -112,13 +67,12 @@ else: # x_vals = [64], x_log=False, line_arg="provider", - line_vals=["sglang_cutlass", "cutlass", "cudnn", "trtllm", "cute-dsl", "auto"], + line_vals=["sglang_cutlass", "cutlass", "cudnn", "trtllm", "auto"], line_names=[ "sglang cutlass fp4", "flashinfer cutlass fp4", "cudnn fp4", "trtllm fp4", - "cute-dsl fp4", "auto fp4 (cudnn/cutlass)", ], styles=[ @@ -126,7 +80,6 @@ else: ("orange", "solid"), ("blue", "solid"), ("green", "solid"), - ("brown", "solid"), ("purple", "solid"), ], ylabel="latency (ms)", @@ -155,14 +108,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 = _bench_cudagraph_with_cupti( + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( 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 = _bench_cudagraph_with_cupti( + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( lambda: mm_fp4( a_fp4, b_fp4.T, @@ -176,7 +129,7 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): quantiles=quantiles, ) if provider == "cudnn": - ms, min_ms, max_ms = _bench_cudagraph_with_cupti( + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( lambda: mm_fp4( a_fp4, b_fp4.T, @@ -192,7 +145,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 = _bench_cudagraph_with_cupti( + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( lambda: mm_fp4( a_fp4, b_fp4.T, @@ -205,22 +158,8 @@ 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 = _bench_cudagraph_with_cupti( + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( lambda: mm_fp4( a_fp4, b_fp4.T, @@ -273,13 +212,6 @@ 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="+", @@ -291,7 +223,7 @@ if __name__ == "__main__": "--dtype", type=torch.dtype, default=torch.bfloat16, - help="Output data type", + help="Data type", ) parser.add_argument( "--correctness", @@ -332,8 +264,8 @@ if __name__ == "__main__": if IS_CI: NKs = NKs[:2] # Only test first 2 shapes in CI - for N, K, model_name in NKs: - print(f"{model_name} N={N} packed_k={K}: ") + for N, K in NKs: + print(f"DeepSeek-R1-0528-FP4 N={N} 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 new file mode 100644 index 000000000..eeb5842ed --- /dev/null +++ b/sgl-kernel/benchmark/bench_nvfp4_scaled_gemm.py @@ -0,0 +1,192 @@ +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 a18c94c92..1a94b6b48 100644 --- a/test/registered/quant/test_nvfp4_gemm.py +++ b/test/registered/quant/test_nvfp4_gemm.py @@ -81,10 +81,5 @@ 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()