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

[NVIDIA] Enable CuTe DSL BF16 GEMM on SM107 (#33617)

Browse Source 
Co-authored-by: Lee Nau <lnau@nvidia.com>
This commit is contained in:
YAMY
2026-08-06 02:06:08 -07:00
committed by GitHub
co-authored by Lee Nau
parent dfe53232d7
commit 8b29c90218
4 changed files with 20 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py
+9 -7
View File
@@ -7,7 +7,7 @@
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
"""CuTe DSL TGV BF16 GEMM (low-latency Blackwell GEMM, SM100/SM103 only).
"""CuTe DSL TGV BF16 GEMM (low-latency SM10x GEMM).
Computes ``out[M, N] = x[M, K] @ weight[N, K].T (+ bias[N])`` for bf16 inputs,
fp32 accumulation, bf16 output. The kernel writes M-contiguous output, so the
@@ -39,7 +39,7 @@ from cutlass.cute.nvgpu import tcgen05
from cutlass.cute.runtime import from_dlpack, make_fake_stream
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.utils import get_device_sm
from sglang.srt.utils import is_sm100_supported
from sglang.srt.utils.common import direct_register_custom_op
# Tuple format: (cta_m, cta_n, num_ab_stage, use_2cta); cta_k is fixed at
@@ -1205,7 +1205,9 @@ def _to_cute_swap(
M_ce = c_swap.shape[1] # == PyTorch N
N_ce = c_swap.shape[2] # == PyTorch M
bias_3d = bias_pt.as_strided(size=(L, M_ce, N_ce), stride=(0, 1, 0))
bias_ = from_dlpack(bias_3d, assumed_align=2).mark_layout_dynamic(leading_dim=1)
bias_ = from_dlpack(bias_3d.detach(), assumed_align=2).mark_layout_dynamic(
leading_dim=1
)
return a_, b_, c_, bias_, layout
@@ -1352,8 +1354,8 @@ def _tgv_bf16_gemm_run(
weight: torch.Tensor,
bias: Optional[torch.Tensor],
) -> torch.Tensor:
if get_device_sm() not in (100, 103):
raise RuntimeError("cutedsl_bf16_gemm requires SM100/SM103 (Blackwell)")
if not is_sm100_supported():
raise RuntimeError("cutedsl_bf16_gemm requires an SM10x GPU")
assert x.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16
assert x.stride(-1) == 1, "x must be K-major [M, K]"
assert weight.stride(-1) == 1, "weight must be K-major [N, K]"
@@ -1380,8 +1382,8 @@ def _tgv_bf16_gemm_out_run(
out: torch.Tensor,
bias: Optional[torch.Tensor],
) -> None:
if get_device_sm() not in (100, 103):
raise RuntimeError("cutedsl_bf16_gemm requires SM100/SM103 (Blackwell)")
if not is_sm100_supported():
raise RuntimeError("cutedsl_bf16_gemm requires an SM10x GPU")
assert x.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16
assert out.dtype == torch.bfloat16 and out.device == x.device
assert x.ndim == 2 and weight.ndim == 2 and out.ndim == 2
python/sglang/srt/layers/quantization/unquant.py
+1 -3
View File
@@ -95,9 +95,7 @@ def initialize_bf16_gemm_config(server_args: ServerArgs) -> None:
if backend.is_cutedsl():
if not is_sm100_supported():
raise ValueError(
"--bf16-gemm-backend cutedsl requires SM100/SM103 (Blackwell)"
)
raise ValueError("--bf16-gemm-backend cutedsl requires an SM10x GPU")
from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import (
cutedsl_bf16_gemm,
python/sglang/srt/server_args.py
+1 -1
View File
@@ -1729,7 +1729,7 @@ class ServerArgs:
bf16_gemm_backend: A[
str,
Arg(
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).",
help="Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; selects 'cutedsl' on SM10x GPUs, 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).",
cli_name="--bf16-gemm-backend",
choices=BF16_GEMM_BACKEND_CHOICES,
),
test/registered/kernels/ops/gemm/test_cutedsl_bf16_gemm.py
+9 -8
View File
@@ -23,32 +23,33 @@ from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import ( # noqa: E402
use_cutedsl_bf16_gemm,
)
N_VALUES = [1024, 2624, 6144]
K_VALUES = [2048, 6144]
SHAPES = [(n, k) for n in [1024, 2624, 6144] for k in [2048, 6144]] + [(2048, 4096)]
NUM_TOKENS = get_ci_test_range(list(range(1, 33)), [1, 15, 16, 32])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize("has_bias", [False, True])
@pytest.mark.parametrize("n", N_VALUES)
@pytest.mark.parametrize("k", K_VALUES)
@pytest.mark.parametrize("n,k", SHAPES)
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
def test_cutedsl_bf16_gemm(num_tokens, k, n, has_bias):
if is_hip_runtime() or get_jit_cuda_arch().major != 10:
pytest.skip("SM100/SM103 required")
pytest.skip("SM10x required")
torch.manual_seed(num_tokens)
x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
bias = torch.randn(n, dtype=torch.bfloat16, device="cuda") if has_bias else None
if bias is not None:
bias.requires_grad_(True)
out = cutedsl_bf16_gemm(x, weight, bias)
with torch.no_grad():
out = cutedsl_bf16_gemm(x, weight, bias)
assert out.shape == (num_tokens, n)
assert out.dtype == torch.bfloat16
ref = x.float() @ weight.float().T
if bias is not None:
ref = ref + bias.float()
ref = ref + bias.detach().float()
torch.testing.assert_close(out, ref.bfloat16(), rtol=2e-2, atol=2.5)
@@ -65,7 +66,7 @@ def test_cutedsl_bf16_gemm_empty_batch(has_bias):
"""Empty input must yield the empty [0, N] output, mirroring F.linear —
launching TGV with a 0-CTA grid fails with CUDA_ERROR_INVALID_VALUE."""
if is_hip_runtime() or get_jit_cuda_arch().major != 10:
pytest.skip("SM100/SM103 required")
pytest.skip("SM10x required")
n, k = 6144, 2048
x = torch.empty(0, k, dtype=torch.bfloat16, device="cuda")