chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support CuteDSL mm_fp4 backend (#18801)

Browse Source
This commit is contained in:
Brayden Zhong
2026-03-19 17:20:01 -04:00
committed by GitHub
parent d8ece7fb22
commit b42b9f6e1a
5 changed files with 109 additions and 219 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/srt/layers/quantization/fp4_utils.py
View File

@@ -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:

4
python/sglang/srt/server_args.py
View File

@@ -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.",
)

120
sgl-kernel/benchmark/bench_fp4_gemm.py
View File

@@ -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,

192
sgl-kernel/benchmark/bench_nvfp4_scaled_gemm.py
View File

@@ -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!")

5
test/registered/quant/test_nvfp4_gemm.py
View File

@@ -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()