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
python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py
+3 -3
View File
@@ -7,7 +7,7 @@ import triton
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
from sglang.srt.utils import is_sm100_supported
from sglang.srt.utils import is_sm100_supported, is_sm120_supported
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
@@ -15,7 +15,7 @@ register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
BLOCK_SIZE = 16
_NVFP4_SUPPORTED = is_sm100_supported()
_NVFP4_SUPPORTED = is_sm100_supported() or is_sm120_supported()
K_E2M1_TO_FLOAT = [
0.0,
@@ -178,7 +178,7 @@ def benchmark(m, n, k, provider):
if __name__ == "__main__":
if not _NVFP4_SUPPORTED:
print("[skip] NVFP4 scaled_mm benchmark requires sm100+ with CUDA 12.8+.")
print("[skip] NVFP4 scaled_mm benchmark requires sm100/sm120 with CUDA 12.8+.")
sys.exit(0)
if not _AOT_SCALED_MM_AVAILABLE:
print(