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

Revert "Support CuteDSL mm_fp4 backend" (#21077)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-03-20 22:47:47 -07:00
committed by GitHub
parent c076968c52
commit 67cad3e69e
5 changed files with 219 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

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