minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

CUTLASS NVFP4 GEMM improvement of SM120 (#21314)

Browse Source
This commit is contained in:
Brayden Zhong
2026-04-01 09:04:34 +08:00
committed by GitHub
parent 5bbf347bb3
commit 6a9b09847c
12 changed files with 845 additions and 901 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/benchmark/bench_fp4_gemm.py
+220 -108
View File
@@ -1,13 +1,20 @@
import argparse
import csv
import os
from functools import partial
from typing import List, Tuple
import torch
import triton
from flashinfer import mm_fp4
from flashinfer.testing import bench_gpu_time
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.srt.utils import (
get_device_capability,
is_sm100_supported,
is_sm120_supported,
)
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
@@ -15,30 +22,102 @@ IS_CI = is_in_ci()
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
DEEPSEEK_R1_MODEL = "deepseek-ai/DeepSeek-R1-0528-FP4"
def get_weight_shapes(args):
models_tps = args.tp_sizes
# 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.
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),
],
"Qwen/Qwen3.5-27B": [
([5120, 8192], 1),
([6144, 5120], 0),
([5120, 34816], 1),
([17408, 5120], 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),
],
}
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 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
if IS_CI:
batch_sizes = [1, 8] # Simplified for CI
batch_sizes = [1, 8]
else:
batch_sizes = [
1,
@@ -60,29 +139,54 @@ else:
]
def _run_mm_fp4(a_fp4, b_fp4_T, a_sf, b_sf_T, alpha, dtype, res_fi, backend):
return mm_fp4(a_fp4, b_fp4_T, a_sf, b_sf_T, alpha, dtype, res_fi, backend=backend)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size"],
x_vals=batch_sizes,
# x_vals = [64],
x_log=False,
line_arg="provider",
line_vals=["sglang_cutlass", "cutlass", "cudnn", "trtllm", "auto"],
line_names=[
"sglang cutlass fp4",
"flashinfer cutlass fp4",
"cudnn fp4",
"trtllm fp4",
"auto fp4 (cudnn/cutlass)",
],
styles=[
("red", "solid"),
("orange", "solid"),
("blue", "solid"),
("green", "solid"),
("purple", "solid"),
],
ylabel="latency (ms)",
line_vals=(
["sglang_cutlass", "cutlass", "cudnn", "trtllm", "auto"]
if is_sm100_supported()
else ["sglang_cutlass", "cutlass", "cudnn", "auto"]
),
line_names=(
[
"sglang cutlass fp4",
"flashinfer cutlass fp4",
"cudnn fp4",
"trtllm fp4",
"auto fp4 (cudnn/cutlass)",
]
if is_sm100_supported()
else [
"sglang cutlass fp4",
"flashinfer cutlass fp4",
"cudnn fp4",
"auto fp4",
]
),
styles=(
[
("red", "solid"),
("orange", "solid"),
("blue", "solid"),
("green", "solid"),
("purple", "solid"),
]
if is_sm100_supported()
else [
("red", "solid"),
("orange", "solid"),
("blue", "solid"),
("purple", "solid"),
]
),
ylabel="bandwidth (GB/s)",
plot_name="fp4_gemm_benchmark",
args={},
)
@@ -99,87 +203,93 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
b_global_scale = (
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
).to(torch.float32)
alpha = 1.0 / (a_global_scale * b_global_scale)
a_fp4, a_scale_interleaved = scaled_fp4_quant(a_dtype, a_global_scale)
# print("a_fp4", a_fp4)
b_fp4, b_scale_interleaved = scaled_fp4_quant(b_dtype, b_global_scale)
b_fp4_T = b_fp4.T
b_sf_T = b_scale_interleaved.T
res_fi = torch.empty((M, N), dtype=dtype, device="cuda")
quantiles = [0.5, 0.2, 0.8]
if provider == "sglang_cutlass":
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 = triton.testing.do_bench_cudagraph(
lambda: mm_fp4(
times_ms = bench_gpu_time(
fn=cutlass_scaled_fp4_mm,
input_args=(
a_fp4,
b_fp4.T,
b_fp4,
a_scale_interleaved,
b_scale_interleaved.T,
b_scale_interleaved,
alpha,
dtype,
res_fi,
backend="cutlass",
),
quantiles=quantiles,
use_cuda_graph=True,
)
if provider == "cudnn":
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
lambda: mm_fp4(
elif provider == "cutlass":
times_ms = bench_gpu_time(
fn=partial(_run_mm_fp4, backend="cutlass"),
input_args=(
a_fp4,
b_fp4.T,
b_fp4_T,
a_scale_interleaved,
b_scale_interleaved.T,
alpha,
dtype,
res_fi,
backend="cudnn",
),
quantiles=quantiles,
)
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(
lambda: mm_fp4(
a_fp4,
b_fp4.T,
a_scale_interleaved,
b_scale_interleaved.T,
alpha,
dtype,
res_fi,
backend="trtllm",
),
quantiles=quantiles,
)
if provider == "auto":
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
lambda: mm_fp4(
a_fp4,
b_fp4.T,
a_scale_interleaved,
b_scale_interleaved.T,
b_sf_T,
alpha,
dtype,
res_fi,
),
quantiles=quantiles,
use_cuda_graph=True,
)
elif provider == "cudnn":
times_ms = bench_gpu_time(
fn=partial(_run_mm_fp4, backend="cudnn"),
input_args=(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
),
use_cuda_graph=True,
)
elif provider == "trtllm":
a_sf_u8 = a_scale_interleaved.to(torch.uint8)
b_sf_u8_T = b_sf_T.to(torch.uint8)
times_ms = bench_gpu_time(
fn=partial(_run_mm_fp4, backend="trtllm"),
input_args=(a_fp4, b_fp4_T, a_sf_u8, b_sf_u8_T, alpha, dtype, res_fi),
use_cuda_graph=True,
)
elif provider == "auto":
times_ms = bench_gpu_time(
fn=partial(_run_mm_fp4, backend="auto"),
input_args=(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
),
use_cuda_graph=True,
)
ms = torch.tensor(times_ms).median().item()
# A: M×packed_k bytes (fp4 packed), B: N×packed_k bytes, C: M×N×element_size bytes
element_size = torch.finfo(dtype).bits // 8
total_bytes = M * packed_k + N * packed_k + M * N * element_size
bandwidth_gbs = total_bytes / (ms * 1e-3) / 1e9
if correctness:
res_cutlass = cutlass_scaled_fp4_mm(
a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype
)
mm_fp4(
a_fp4,
b_fp4.T,
b_fp4_T,
a_scale_interleaved,
b_scale_interleaved.T,
b_sf_T,
alpha,
dtype,
res_fi,
@@ -190,9 +300,9 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
), "cudnn fp4 doesn't match cutlass fp4"
mm_fp4(
a_fp4,
b_fp4.T,
b_fp4_T,
a_scale_interleaved,
b_scale_interleaved.T,
b_sf_T,
alpha,
dtype,
res_fi,
@@ -205,13 +315,20 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
if csv_file:
with open(csv_file, "a", newline="") as f:
writer = csv.writer(f)
writer.writerow([provider, M, N, K, ms])
writer.writerow([provider, M, N, K, ms, bandwidth_gbs])
return ms, min_ms, max_ms
return bandwidth_gbs
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, Qwen, Mistral, DeepSeek.",
)
parser.add_argument(
"--tp-sizes",
nargs="+",
@@ -223,7 +340,7 @@ if __name__ == "__main__":
"--dtype",
type=torch.dtype,
default=torch.bfloat16,
help="Data type",
help="Output data type",
)
parser.add_argument(
"--correctness",
@@ -238,34 +355,29 @@ if __name__ == "__main__":
)
args = parser.parse_args()
# Simplify for CI environment
if IS_CI:
args.tp_sizes = [args.tp_sizes[0]] # Use only first TP size
args.tp_sizes = [args.tp_sizes[0]]
if args.csv:
with open(args.csv, "w", newline="") as f:
writer = csv.writer(f)
writer.writerow(["provider", "m", "n", "k", "time_ms"])
writer.writerow(["provider", "m", "n", "k", "time_ms", "bandwidth_gbs"])
# FP4 operations require Blackwell SM100 support
major, minor = get_device_capability()
if not is_sm100_supported():
if not (is_sm100_supported() or is_sm120_supported()):
print("Skipping FP4 GEMM benchmark")
if major is not None:
print(
f"FP4 operations require SM100 (Blackwell), but found sm{major}{minor}"
)
print(f"FP4 operations require sm100+, but found sm{major}{minor}")
else:
print("Could not determine device capability")
else:
NKs = get_weight_shapes(args)
# Limit iterations in CI
if IS_CI:
NKs = NKs[:2] # Only test first 2 shapes in CI
NKs = NKs[:2]
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,
sgl-kernel/benchmark/bench_nvfp4_scaled_gemm.py
-192
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!")