diff --git a/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py b/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py index f608be4fb..84ecb4c59 100644 --- a/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py +++ b/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index edc72f0a7..fe78b4fcc 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 67eed9672..f99edc07a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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, ), diff --git a/test/registered/kernels/ops/gemm/test_cutedsl_bf16_gemm.py b/test/registered/kernels/ops/gemm/test_cutedsl_bf16_gemm.py index f52b1137e..fd8d7b93e 100644 --- a/test/registered/kernels/ops/gemm/test_cutedsl_bf16_gemm.py +++ b/test/registered/kernels/ops/gemm/test_cutedsl_bf16_gemm.py @@ -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")