[SM90] Use native block32 DeepGEMM for DeepSeek-V4.1

python/pyproject.toml

@@ -77,7 +77,7 @@ dependencies = [
   "sentencepiece==0.2.1",  # 0.2.2 rejects null pieces in InternVL tokenizers.
   "setproctitle",
   "sgl-deep-ep==0.1.2.post1",
-  "sgl-deep-gemm==0.2.1.post1",
+  "sgl-deep-gemm==0.2.2",
   "sglang-kernel==0.4.9",
   "smg-grpc-servicer>=0.9.0",
   "soundfile==0.13.1",

python/sglang/kernels/ops/gemm/fp8_kernel.py

@@ -33,8 +33,11 @@ def deep_gemm_fp8_fp8_bf16_nt(
     B: torch.Tensor,
     Bs: torch.Tensor,
     C: torch.Tensor,
+    block_n: int = 128,
+    block_k: int = 128,
 ) -> None:
-    deep_gemm_wrapper.gemm_nt_f8f8bf16((A, As), (B, Bs), C)
+    recipe = (1, block_n, block_k) if (block_n, block_k) != (128, 128) else None
+    deep_gemm_wrapper.gemm_nt_f8f8bf16((A, As), (B, Bs), C, recipe=recipe)
 
 
 @register_custom_op(mutates_args=["C"])
@@ -746,7 +749,7 @@ def w8a8_block_fp8_matmul_deepgemm(
     # Deepgemm only supports output tensor type as bfloat16
     assert C.dtype == torch.bfloat16 and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
 
-    deep_gemm_fp8_fp8_bf16_nt(A, As, B, Bs, C)
+    deep_gemm_fp8_fp8_bf16_nt(A, As, B, Bs, C, block_size[0], block_size[1])
 
     return C
 

python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py

@@ -5,6 +5,7 @@
 import time
 from contextlib import contextmanager, nullcontext
 from enum import IntEnum, auto
+from functools import partial
 from pathlib import Path
 from typing import Dict, List, Tuple
 
@@ -111,6 +112,7 @@ class DeepGemmKernelType(IntEnum):
     GROUPED_GEMM_NT_BF16_MASKED = auto()
     GROUPED_GEMM_NT_BF16_CONTIG = auto()
     GEMM_NT_F8F8BF16 = auto()
+    GEMM_NT_F8F8BF16_BLOCK32 = auto()
     GEMM_NT_BF16BF16F32 = auto()
 
 
@@ -265,6 +267,9 @@ class _BaseWarmupExecutor:
     def create(kernel_type: DeepGemmKernelType, **kwargs):
         return {
             DeepGemmKernelType.GEMM_NT_F8F8BF16: _NormalWarmupExecutor,
+            DeepGemmKernelType.GEMM_NT_F8F8BF16_BLOCK32: partial(
+                _NormalWarmupExecutor, block_size=32
+            ),
             DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG: _GroupedContWarmupExecutor,
             DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED: _GroupedMaskedWarmupExecutor,
             DeepGemmKernelType.GEMM_NT_BF16BF16F32: _BF16F32WarmupExecutor,
@@ -278,7 +283,10 @@ def get_memory_requirement(
     ) -> int:
         # Return the required memory space in GB for warmup executor
         _GB = 1 << 30
-        if kernel_type == DeepGemmKernelType.GEMM_NT_F8F8BF16:
+        if kernel_type in (
+            DeepGemmKernelType.GEMM_NT_F8F8BF16,
+            DeepGemmKernelType.GEMM_NT_F8F8BF16_BLOCK32,
+        ):
             return (max_m * k + n * k + max_m * n * 2) / _GB
         elif kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG:
             return (max_m * k + num_groups * n * k + max_m * 4 + max_m * n * 2) / _GB
@@ -310,22 +318,22 @@ def execute(self, m):
         raise NotImplementedError
 
 
-def _empty_token_fp8(size):
+def _empty_token_fp8(size, block_size=128):
     *dims, k = size
     return (
         torch.empty(size, device="cuda", dtype=torch.float8_e4m3fn),
         torch.ones(
-            (*dims, ceil_div(k, _BLOCK_SIZE)), device="cuda", dtype=torch.float32
+            (*dims, ceil_div(k, block_size)), device="cuda", dtype=torch.float32
         ),
     )
 
 
-def _empty_block_fp8(size):
+def _empty_block_fp8(size, block_size=128):
     *dims, n, k = size
     return (
         torch.empty(size, device="cuda", dtype=torch.float8_e4m3fn),
         torch.ones(
-            (*dims, ceil_div(n, _BLOCK_SIZE), ceil_div(k, _BLOCK_SIZE)),
+            (*dims, ceil_div(n, block_size), ceil_div(k, block_size)),
             device="cuda",
             dtype=torch.float32,
         ),
@@ -336,16 +344,20 @@ def _empty_block_fp8(size):
 
 
 class _NormalWarmupExecutor(_BaseWarmupExecutor):
-    def __init__(self, max_m: int, n: int, k: int, num_groups: int):
-        self.lhs_q, self.lhs_s = _empty_token_fp8((max_m, k))
-        self.rhs_q, self.rhs_s = _empty_block_fp8((n, k))
+    def __init__(
+        self, max_m: int, n: int, k: int, num_groups: int, block_size: int = 128
+    ):
+        self.recipe = (1, block_size, block_size)
+        self.lhs_q, self.lhs_s = _empty_token_fp8((max_m, k), block_size)
+        self.rhs_q, self.rhs_s = _empty_block_fp8((n, k), block_size)
         self.out = torch.empty((max_m, n), device="cuda", dtype=torch.bfloat16)
 
     def execute(self, m):
         deep_gemm.fp8_gemm_nt(
             (self.lhs_q[:m], self.lhs_s[:m]),
             (self.rhs_q, self.rhs_s),
             self.out[:m],
+            recipe=self.recipe,
         )
 
 

python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py

@@ -220,11 +220,16 @@ def gemm_nt_f8f8bf16(
     lhs: Tuple[torch.Tensor, torch.Tensor],
     rhs: Tuple[torch.Tensor, torch.Tensor],
     out: torch.Tensor,
+    recipe: Optional[Tuple[int, int, int]] = None,
 ):
     m, k = lhs[0].shape
     n, _ = rhs[0].shape
     num_groups = 1
-    kernel_type = compile_utils.DeepGemmKernelType.GEMM_NT_F8F8BF16
+    kernel_type = (
+        compile_utils.DeepGemmKernelType.GEMM_NT_F8F8BF16_BLOCK32
+        if recipe == (1, 32, 32)
+        else compile_utils.DeepGemmKernelType.GEMM_NT_F8F8BF16
+    )
 
     _sanity_check_input(lhs)
     _sanity_check_input(rhs)
@@ -234,6 +239,7 @@ def gemm_nt_f8f8bf16(
             lhs,
             rhs,
             out,
+            **({"recipe": recipe} if recipe is not None else {}),
         )
 
 

python/sglang/srt/layers/quantization/fp8.py

@@ -53,6 +53,7 @@
     QuantizeMethodBase,
 )
 from sglang.srt.layers.quantization.fp8_utils import (
+    _deepgemm_block_fp8_supported,
     _use_aiter_bpreshuffle_gfx95,
     apply_fp8_linear,
     block_fp8_scale_to_mxfp8_e8m0,
@@ -729,6 +730,10 @@ def create_weights(
         )
 
     def process_weights_after_loading_block_quant(self, layer: Module) -> None:
+        use_deepgemm_runner = (
+            getattr(self.w8a8_block_fp8_linear, "func", self.w8a8_block_fp8_linear)
+            is deepgemm_w8a8_block_fp8_linear_with_fallback
+        )
         if self.convert_mxfp8_to_block:
             from sglang.srt.layers.quantization.mxfp8_block_convert import (
                 convert_mxfp8_weight_to_block_fp8,
@@ -810,10 +815,6 @@ def process_weights_after_loading_block_quant(self, layer: Module) -> None:
             return
         else:
             # Requantize block scales to UE8M0 when DeepGEMM is the active runner.
-            use_deepgemm_runner = (
-                self.w8a8_block_fp8_linear
-                is deepgemm_w8a8_block_fp8_linear_with_fallback
-            )
             requant_block_scale_ue8m0_for_deepgemm(
                 layer.weight,
                 layer.weight_scale_inv,
@@ -833,6 +834,14 @@ def process_weights_after_loading_block_quant(self, layer: Module) -> None:
             _is_cuda
             and get_platform().is_sm90
             and envs.SGLANG_OPT_HOPPER_BLOCK_FP8_BF16.get()
+            and not (
+                use_deepgemm_runner
+                and _deepgemm_block_fp8_supported(
+                    weight.shape,
+                    self.weight_block_size,
+                    getattr(layer, "orig_dtype", None),
+                )
+            )
             and weight.is_cuda
             and weight.dtype == torch.float8_e4m3fn
             and self.weight_block_size == [32, 32]

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -614,19 +614,32 @@ def dispatch_w8a8_block_fp8_linear(
     1. The --fp8-gemm-backend server argument (preferred)
     2. Auto-detection based on hardware capabilities
     """
-    # Only Triton reads the block size at launch; DeepGEMM, the FlashInfer
-    # groupwise kernels and CUTLASS take 128-wide K blocks only.
+    backend = get_fp8_gemm_runner_backend()
+    if (
+        weight_block_size == [32, 32]
+        and get_platform().is_sm90
+        and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+        and (backend.is_auto() or backend.is_deep_gemm())
+    ):
+        return partial(
+            deepgemm_w8a8_block_fp8_linear_with_fallback,
+            act_scale_ue8m0=act_scale_ue8m0,
+        )
+
+    # Other non-128 layouts still require Triton or the separate MXFP8 route.
     if weight_block_size is not None and weight_block_size[1] != 128:
         return partial(triton_w8a8_block_fp8_linear, act_scale_ue8m0=act_scale_ue8m0)
 
-    backend = get_fp8_gemm_runner_backend()
-
     # Handle explicit backend selection via --fp8-gemm-backend
     if not backend.is_auto():
-        return _dispatch_explicit_backend(backend)
-
-    # Auto mode: Select based purely on hardware/backend availability
-    return _dispatch_auto_backend()
+        linear = _dispatch_explicit_backend(backend)
+    else:
+        # Auto mode: Select based purely on hardware/backend availability
+        linear = _dispatch_auto_backend()
+    # Preserve the unwrapped callable when no activation policy needs binding.
+    if act_scale_ue8m0 and linear is deepgemm_w8a8_block_fp8_linear_with_fallback:
+        return partial(linear, act_scale_ue8m0=True)
+    return linear
 
 
 def torch_w8a8_block_fp8_linear(
@@ -1238,50 +1251,34 @@ def cutlass_w8a8_block_fp8_linear_with_fallback(
     return output.to(dtype=input_2d.dtype).view(*output_shape)
 
 
+def _deepgemm_block_fp8_supported(
+    weight_shape: Tuple[int, int], block_size: List[int], output_dtype: torch.dtype
+) -> bool:
+    if output_dtype != torch.bfloat16:
+        return False
+    n, k = weight_shape
+    if list(block_size) == [32, 32] and get_platform().is_sm90:
+        return n % 8 == 0 and k % 32 == 0
+    return list(block_size) == [128, 128] and n % 64 == 0 and k % 128 == 0
+
+
 def deepgemm_w8a8_block_fp8_linear_with_fallback(
     input: torch.Tensor,
     weight: torch.Tensor,
     block_size: List[int],
     weight_scale: torch.Tensor,
     input_scale: Optional[torch.Tensor] = None,
     bias: Optional[torch.Tensor] = None,
+    act_scale_ue8m0: bool = False,
+    weight_bf16: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     if input_scale is not None:
-        # Pre-quantized activation (SGLANG_OPT_MOE_QUANT_ONCE): ``input`` is
-        # the fp8 per-token-group-128 q with rows padded to a multiple of 4
-        # and ``input_scale`` the matching column-major fp32 scales
-        # (stride == (1, padded_rows)) -- identical to the MN-major
-        # TMA-aligned layout this path's own quant would produce below.
-        # Output keeps the padded row count; the caller slices back.
-        # UE8M0 packed scales (Blackwell DeepGEMM) use a different layout;
-        # the caller gates on it.
         assert not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
         assert input.dtype == torch.float8_e4m3fn
-        assert weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0, (
-            "pre-quantized fp8 input requires DeepGEMM-supported weight shapes "
-            f"(got {tuple(weight.shape)})"
-        )
-        input_2d = input.view(-1, input.shape[-1])
-        output = w8a8_block_fp8_matmul_deepgemm(
-            input_2d,
-            weight,
-            input_scale,
-            weight_scale,
-            block_size,
-            output_dtype=torch.bfloat16,
-        )
-        if bias is not None:
-            output += bias
-        return output.view(*input.shape[:-1], weight.shape[0])
-
-    output_dtype = input.dtype
-    dtype_supported = output_dtype == torch.bfloat16
+    output_dtype = torch.bfloat16 if input_scale is not None else input.dtype
 
     # TODO: https://github.com/sgl-project/sglang/pull/6890#issuecomment-2943395737
-    shape_supported = weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0
-    block_supported = list(block_size) == [128, 128]
-
-    if not (shape_supported and dtype_supported and block_supported):
+    if not _deepgemm_block_fp8_supported(weight.shape, block_size, output_dtype):
         # fall back to triton
         # If weight_scale is in UE8M0 packed format (int32), convert back to float32
         # UE8M0 format has shape (N, K//block_k//4) with dtype int32
@@ -1291,13 +1288,42 @@ def deepgemm_w8a8_block_fp8_linear_with_fallback(
                 weight_scale, weight.shape, block_size
             )
         return triton_w8a8_block_fp8_linear(
-            input, weight, block_size, weight_scale, input_scale, bias
+            input,
+            weight,
+            block_size,
+            weight_scale,
+            input_scale,
+            bias,
+            act_scale_ue8m0=act_scale_ue8m0,
+            weight_bf16=weight_bf16,
         )
 
+    if input_scale is not None:
+        # Pre-quantized activations retain their per-token group size and
+        # fp32 scales. DeepGEMM aligns scale storage for TMA when needed.
+        input_2d = input.view(-1, input.shape[-1])
+        output = w8a8_block_fp8_matmul_deepgemm(
+            input_2d,
+            weight,
+            input_scale,
+            weight_scale,
+            block_size,
+            output_dtype=output_dtype,
+        )
+        if bias is not None:
+            output += bias
+        return output.view(*input.shape[:-1], weight.shape[0])
+
     input_2d = input.view(-1, input.shape[-1])
     output_shape = [*input.shape[:-1], weight.shape[0]]
 
-    if not _is_musa:
+    if get_platform().is_sm90 and act_scale_ue8m0:
+        # Keep the checkpoint's power-of-two activation quantization. SM90
+        # consumes fp32 scales; DeepGEMM aligns their storage for TMA.
+        q_input, x_scale = sglang_per_token_group_quant_fp8(
+            input_2d, block_size[1], scale_ue8m0=True
+        )
+    elif not _is_musa:
         q_input, x_scale = sglang_per_token_group_quant_fp8(
             input_2d,
             block_size[1],

python/sglang/srt/models/deepseek_v2.py

@@ -1775,7 +1775,8 @@ def _compute_moe_quant_once_enabled(self) -> Tuple[bool, str]:
             if gate_up.weight.shape[0] % 64 != 0 or gate_up.weight.shape[1] % 128 != 0:
                 return False, "gate_up weight shape unsupported by deepgemm"
         else:
-            return False, f"w8a8 linear backend {linear_fn.__name__} unsupported"
+            backend_name = getattr(linear_fn, "__name__", type(linear_fn).__name__)
+            return False, f"w8a8 linear backend {backend_name} unsupported"
         # Routed side: standard dispatcher + triton fused func with dynamic
         # per-token-group-128 fp8 activation quant.
         experts = self.experts