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

Support CuteDSL GEMM BF16 on SM100 on by default when allowed by heuristic (#30567)

Browse Source
This commit is contained in:
Brayden Zhong
2026-07-22 14:13:12 -07:00
committed by GitHub
parent 98cc8d91cf
commit e7511141ea
3 changed files with 18 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/unquant.py
+10 -3
View File
@@ -67,6 +67,7 @@ if _use_aiter:
class Bf16GemmBackend(Enum):
AUTO = "auto"
CUTEDSL = "cutedsl"
TORCH = "torch"
def is_auto(self) -> bool:
return self == Bf16GemmBackend.AUTO
@@ -83,11 +84,15 @@ _use_cutedsl_bf16_gemm = None
def initialize_bf16_gemm_config(server_args: ServerArgs) -> None:
global _BF16_GEMM_BACKEND, _cutedsl_bf16_gemm, _use_cutedsl_bf16_gemm
backend = Bf16GemmBackend(server_args.bf16_gemm_backend)
from sglang.srt.utils import is_sm100_supported
backend_str = server_args.bf16_gemm_backend
if backend_str == "auto" and is_sm100_supported():
backend_str = "cutedsl"
backend = Bf16GemmBackend(backend_str)
if backend.is_cutedsl():
from sglang.srt.utils import is_sm100_supported
if not is_sm100_supported():
raise ValueError(
"--bf16-gemm-backend cutedsl requires SM100/SM103 (Blackwell)"
@@ -207,6 +212,8 @@ class UnquantizedLinearMethod(LinearMethodBase):
and x.dtype == torch.bfloat16
and layer.weight.dtype == torch.bfloat16
and (bias is None or bias.dtype == torch.bfloat16)
and not layer.weight.requires_grad
and (bias is None or not bias.requires_grad)
and _use_cutedsl_bf16_gemm(
x.numel() // x.shape[-1],
layer.weight.shape[0],
python/sglang/srt/server_args.py
+2 -2
View File
@@ -306,7 +306,7 @@ FP4_GEMM_RUNNER_BACKEND_CHOICES = [
"marlin",
]
BF16_GEMM_BACKEND_CHOICES = ["auto", "cutedsl"]
BF16_GEMM_BACKEND_CHOICES = ["auto", "cutedsl", "torch"]
RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu", "slru", "priority"]
RETRACTION_POLICY_CHOICES = ["length", "priority"]
@@ -1662,7 +1662,7 @@ class ServerArgs:
bf16_gemm_backend: A[
str,
Arg(
help="Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; uses cuBLAS via torch.nn.functional.linear), 'cutedsl' (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10X; dispatches between the CuTe DSL kernel and cuBLAS).",
help="Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; selects 'cutedsl' on SM100/SM103 (Blackwell), otherwise uses cuBLAS via torch.nn.functional.linear), 'cutedsl' (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10X; dispatches between the CuTe DSL kernel and cuBLAS), 'torch' (always uses cuBLAS via torch.nn.functional.linear, even on SM100/SM103).",
cli_name="--bf16-gemm-backend",
choices=BF16_GEMM_BACKEND_CHOICES,
),