390 lines
11 KiB
Python
Executable File
390 lines
11 KiB
Python
Executable File
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,
|
||
is_sm120_supported,
|
||
)
|
||
from sglang.utils import is_in_ci
|
||
|
||
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"
|
||
|
||
# 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),
|
||
],
|
||
}
|
||
|
||
DEEPSEEK_R1_WEIGHT_SHAPES = {
|
||
4: [[1024, 3584], [7168, 256], [7168, 2304], [9216, 3584]],
|
||
8: [[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
|
||
|
||
|
||
if IS_CI:
|
||
batch_sizes = [1, 8]
|
||
else:
|
||
batch_sizes = [
|
||
1,
|
||
2,
|
||
4,
|
||
8,
|
||
16,
|
||
32,
|
||
64,
|
||
128,
|
||
256,
|
||
512,
|
||
1024,
|
||
2048,
|
||
3072,
|
||
4096,
|
||
8192,
|
||
16384,
|
||
]
|
||
|
||
|
||
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_log=False,
|
||
line_arg="provider",
|
||
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={},
|
||
)
|
||
)
|
||
def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
|
||
M = batch_size
|
||
packed_k = K
|
||
K = 2 * packed_k
|
||
a_dtype = torch.randn((M, K), dtype=dtype, device="cuda")
|
||
b_dtype = torch.randn((N, K), dtype=dtype, device="cuda")
|
||
a_global_scale = (
|
||
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
|
||
).to(torch.float32)
|
||
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)
|
||
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")
|
||
|
||
if provider == "sglang_cutlass":
|
||
times_ms = bench_gpu_time(
|
||
fn=cutlass_scaled_fp4_mm,
|
||
input_args=(
|
||
a_fp4,
|
||
b_fp4,
|
||
a_scale_interleaved,
|
||
b_scale_interleaved,
|
||
alpha,
|
||
dtype,
|
||
),
|
||
use_cuda_graph=True,
|
||
)
|
||
elif provider == "cutlass":
|
||
times_ms = bench_gpu_time(
|
||
fn=partial(_run_mm_fp4, backend="cutlass"),
|
||
input_args=(
|
||
a_fp4,
|
||
b_fp4_T,
|
||
a_scale_interleaved,
|
||
b_sf_T,
|
||
alpha,
|
||
dtype,
|
||
res_fi,
|
||
),
|
||
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,
|
||
a_scale_interleaved,
|
||
b_sf_T,
|
||
alpha,
|
||
dtype,
|
||
res_fi,
|
||
backend="cudnn",
|
||
)
|
||
assert torch.allclose(
|
||
res_fi, res_cutlass, atol=1e-3, rtol=1e-3
|
||
), "cudnn fp4 doesn't match cutlass fp4"
|
||
mm_fp4(
|
||
a_fp4,
|
||
b_fp4_T,
|
||
a_scale_interleaved,
|
||
b_sf_T,
|
||
alpha,
|
||
dtype,
|
||
res_fi,
|
||
backend="trtllm",
|
||
)
|
||
assert torch.allclose(
|
||
res_fi, res_cutlass, atol=1e-3, rtol=1e-3
|
||
), "trtllm fp4 doesn't match cutlass fp4"
|
||
|
||
if csv_file:
|
||
with open(csv_file, "a", newline="") as f:
|
||
writer = csv.writer(f)
|
||
writer.writerow([provider, M, N, K, ms, bandwidth_gbs])
|
||
|
||
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="+",
|
||
type=int,
|
||
default=[1],
|
||
help="List of tensor parallel sizes",
|
||
)
|
||
parser.add_argument(
|
||
"--dtype",
|
||
type=torch.dtype,
|
||
default=torch.bfloat16,
|
||
help="Output data type",
|
||
)
|
||
parser.add_argument(
|
||
"--correctness",
|
||
action="store_true",
|
||
help="Check correctness",
|
||
)
|
||
parser.add_argument(
|
||
"--csv",
|
||
type=str,
|
||
default="results_cutlass_cudnn.csv",
|
||
help="CSV file to save results",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
if IS_CI:
|
||
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", "bandwidth_gbs"])
|
||
|
||
major, minor = get_device_capability()
|
||
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+, but found sm{major}{minor}")
|
||
else:
|
||
print("Could not determine device capability")
|
||
else:
|
||
NKs = get_weight_shapes(args)
|
||
|
||
if IS_CI:
|
||
NKs = NKs[:2]
|
||
|
||
for N, K, model_name in NKs:
|
||
print(f"{model_name} N={N} packed_k={K}: ")
|
||
benchmark.run(
|
||
print_data=True,
|
||
N=N,
|
||
K=K,
|
||
dtype=args.dtype,
|
||
correctness=args.correctness,
|
||
csv_file=args.csv,
|
||
)
|
||
print("Benchmark finished!")
|