sglang/sgl-kernel/benchmark/bench_fp4_gemm.py

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.srt.utils import get_device_capability, is_sm100_supported

# 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)
# 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),
    ],
}

DEEPSEEK_R1_WEIGHT_SHAPES = {
    4: [[1024, 3584], [7168, 256], [7168, 2304], [9216, 3584]],
    8: [[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
if IS_CI:
    batch_sizes = [1, 8]  # Simplified for CI
else:
    batch_sizes = [
        1,
        2,
        4,
        8,
        16,
        32,
        64,
        128,
        256,
        512,
        1024,
        2048,
        3072,
        4096,
        8192,
        16384,
    ]


@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", "cute-dsl", "auto"],
        line_names=[
            "sglang cutlass fp4",
            "flashinfer cutlass fp4",
            "cudnn fp4",
            "trtllm fp4",
            "cute-dsl fp4",
            "auto fp4 (cudnn/cutlass)",
        ],
        styles=[
            ("red", "solid"),
            ("orange", "solid"),
            ("blue", "solid"),
            ("green", "solid"),
            ("brown", "solid"),
            ("purple", "solid"),
        ],
        ylabel="latency (ms)",
        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)
    # print("a_fp4", a_fp4)
    b_fp4, b_scale_interleaved = scaled_fp4_quant(b_dtype, b_global_scale)
    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 = _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 = _bench_cudagraph_with_cupti(
            lambda: mm_fp4(
                a_fp4,
                b_fp4.T,
                a_scale_interleaved,
                b_scale_interleaved.T,
                alpha,
                dtype,
                res_fi,
                backend="cutlass",
            ),
            quantiles=quantiles,
        )
    if provider == "cudnn":
        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="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 = _bench_cudagraph_with_cupti(
            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 == "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(
            lambda: mm_fp4(
                a_fp4,
                b_fp4.T,
                a_scale_interleaved,
                b_scale_interleaved.T,
                alpha,
                dtype,
                res_fi,
            ),
            quantiles=quantiles,
        )
    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_scale_interleaved.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_scale_interleaved.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])

    return ms, min_ms, max_ms


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="+",
        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()

    # Simplify for CI environment
    if IS_CI:
        args.tp_sizes = [args.tp_sizes[0]]  # Use only first TP size

    if args.csv:
        with open(args.csv, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["provider", "m", "n", "k", "time_ms"])

    # FP4 operations require Blackwell SM100 support
    major, minor = get_device_capability()
    if not is_sm100_supported():
        print("Skipping FP4 GEMM benchmark")
        if major is not None:
            print(
                f"FP4 operations require SM100 (Blackwell), 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

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