From 2ab93e89e9a79aad3db4193e15abcf75180445d7 Mon Sep 17 00:00:00 2001 From: Vladislav Nosivskoy Date: Mon, 5 Oct 2026 12:42:13 +0300 Subject: [PATCH 1/5] Use native DeepGEMM block32 FP8 GEMM on SM90 --- python/pyproject.toml | 2 +- python/sglang/kernels/ops/gemm/fp8_kernel.py | 7 ++- .../layers/deep_gemm_wrapper/compile_utils.py | 25 +++++--- .../layers/deep_gemm_wrapper/entrypoint.py | 8 ++- .../srt/layers/quantization/fp8_utils.py | 58 +++++++++++++------ .../test_fp8_blockwise_linear_backends.py | 32 ++++++++++ 6 files changed, 102 insertions(+), 30 deletions(-) diff --git a/python/pyproject.toml b/python/pyproject.toml index 2f9a5200609c..e53c281a31a1 100755 --- a/python/pyproject.toml +++ b/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", diff --git a/python/sglang/kernels/ops/gemm/fp8_kernel.py b/python/sglang/kernels/ops/gemm/fp8_kernel.py index 46234dc6eb3d..ff32551d33f1 100644 --- a/python/sglang/kernels/ops/gemm/fp8_kernel.py +++ b/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 diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 2db04ae05d98..db878de14f5a 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/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,7 @@ 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 +315,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,9 +341,12 @@ 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): @@ -346,6 +354,7 @@ def execute(self, m): (self.lhs_q[:m], self.lhs_s[:m]), (self.rhs_q, self.rhs_s), self.out[:m], + recipe=self.recipe, ) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index 6134374bc332..7fff7d00bb16 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/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 {}), ) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 2c4a70be9fff..b7a147cb799c 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -608,13 +608,22 @@ 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) @@ -1236,19 +1245,22 @@ def deepgemm_w8a8_block_fp8_linear_with_fallback( 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: + block32 = list(block_size) == [32, 32] and get_platform().is_sm90 + shape_supported = ( + weight.shape[0] % 8 == 0 and weight.shape[1] % 32 == 0 + if block32 + else weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0 + ) + block_supported = block32 or list(block_size) == [128, 128] 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. + # Pre-quantized activations retain their per-token group size and + # fp32 scales. DeepGEMM aligns scale storage for TMA when needed. 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, ( + assert shape_supported and block_supported, ( "pre-quantized fp8 input requires DeepGEMM-supported weight shapes " f"(got {tuple(weight.shape)})" ) @@ -1269,9 +1281,6 @@ def deepgemm_w8a8_block_fp8_linear_with_fallback( dtype_supported = output_dtype == torch.bfloat16 # 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): # fall back to triton # If weight_scale is in UE8M0 packed format (int32), convert back to float32 @@ -1282,13 +1291,26 @@ 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, ) input_2d = input.view(-1, input.shape[-1]) output_shape = [*input.shape[:-1], weight.shape[0]] - if not _is_musa: + if block32 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], diff --git a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py index 56317e09fa00..bc37b7545620 100644 --- a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py +++ b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py @@ -339,6 +339,38 @@ def test_mxfp8_view_and_swizzled_input(self): plain_layer.quant_method.apply(plain_layer, Mxfp8SwizzledInput(q, s)) +@unittest.skipUnless(get_device_sm() == 90, "Native block32 DeepGEMM requires SM90") +class TestSm90Block32Linear(_LinearBackendCheck): + _build_layer = staticmethod(_build_block32_layer) + _shapes = [(1, 96, 288), (6, 5120, 288), (65, 1792, 5120), (512, 512, 512)] + + def test_deepgemm_and_auto_use_native_blocks(self): + for backend in ("deep_gemm", "auto"): + with self.subTest(backend=backend), mock.patch.object( + fp8_utils, "triton_w8a8_block_fp8_linear", + side_effect=AssertionError("native block32 unexpectedly fell back to Triton"), + ): + self._check_backend(backend, [backend], self._shapes, self._build_layer) + + def test_triton_override(self): + self._check_backend("triton", ["triton"], self._shapes, self._build_layer) + + def test_preserves_checkpoint_and_prequantized_input(self): + torch.manual_seed(7) + with mock.patch.object(fp8_utils, "FP8_GEMM_RUNNER_BACKEND", Fp8GemmRunnerBackend.DEEP_GEMM): + layer, _ = self._build_layer(96, 288) + original_weight = layer.weight.detach().clone() + layer.quant_method.process_weights_after_loading(layer) + torch.testing.assert_close(layer.weight.float(), original_weight.float(), rtol=0, atol=0) + self.assertEqual(tuple(layer.weight_scale_inv.shape), (3, 9)) + x = torch.randn(5, 288, device="cuda", dtype=torch.bfloat16) / 10 + q, scales = fp8_utils.sglang_per_token_group_quant_fp8(x, 32, scale_ue8m0=True) + bias = torch.randn(96, device="cuda", dtype=torch.bfloat16) + plain = layer.quant_method.apply(layer, x, bias) + quantized = layer.quant_method.apply(layer, (q, scales), bias) + torch.testing.assert_close(plain, quantized, rtol=0, atol=0) + + @unittest.skipUnless( "flashinfer_cutedsl" in _block32_backends(), "block-fp8-as-MXFP8 prefill tuning needs the FlashInfer CuTe-DSL kernel", From 7818b247701ba08db485458e1a8546ba176ced56 Mon Sep 17 00:00:00 2001 From: Vladislav Nosivskoy Date: Mon, 5 Oct 2026 12:59:21 +0300 Subject: [PATCH 2/5] Avoid unused BF16 caches for native block32 DeepGEMM --- python/sglang/srt/layers/quantization/fp8.py | 9 +++++---- .../quantization/test_fp8_blockwise_linear_backends.py | 1 + 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 3d62fd5e1eb2..87e68e103ec3 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -725,6 +725,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, @@ -806,10 +810,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, @@ -829,6 +829,7 @@ 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 weight.is_cuda and weight.dtype == torch.float8_e4m3fn and self.weight_block_size == [32, 32] diff --git a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py index bc37b7545620..ac53b7756e6e 100644 --- a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py +++ b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py @@ -361,6 +361,7 @@ def test_preserves_checkpoint_and_prequantized_input(self): layer, _ = self._build_layer(96, 288) original_weight = layer.weight.detach().clone() layer.quant_method.process_weights_after_loading(layer) + self.assertIsNone(getattr(layer, "_block_fp8_bf16_weight", None)) torch.testing.assert_close(layer.weight.float(), original_weight.float(), rtol=0, atol=0) self.assertEqual(tuple(layer.weight_scale_inv.shape), (3, 9)) x = torch.randn(5, 288, device="cuda", dtype=torch.bfloat16) / 10 From b9aa964f9467c80878cfaaaed8c164d6ad67dbb6 Mon Sep 17 00:00:00 2001 From: Vladislav Nosivskoy Date: Mon, 5 Oct 2026 16:46:47 +0300 Subject: [PATCH 3/5] Format native block32 DeepGEMM integration --- .../layers/deep_gemm_wrapper/compile_utils.py | 5 +++- .../test_fp8_blockwise_linear_backends.py | 24 ++++++++++++++----- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index db878de14f5a..dc00ac7cc4c1 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -283,7 +283,10 @@ def get_memory_requirement( ) -> int: # Return the required memory space in GB for warmup executor _GB = 1 << 30 - if kernel_type in (DeepGemmKernelType.GEMM_NT_F8F8BF16, DeepGemmKernelType.GEMM_NT_F8F8BF16_BLOCK32): + 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 diff --git a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py index ac53b7756e6e..9f6aeaea9a0e 100644 --- a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py +++ b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py @@ -346,9 +346,15 @@ class TestSm90Block32Linear(_LinearBackendCheck): def test_deepgemm_and_auto_use_native_blocks(self): for backend in ("deep_gemm", "auto"): - with self.subTest(backend=backend), mock.patch.object( - fp8_utils, "triton_w8a8_block_fp8_linear", - side_effect=AssertionError("native block32 unexpectedly fell back to Triton"), + with ( + self.subTest(backend=backend), + mock.patch.object( + fp8_utils, + "triton_w8a8_block_fp8_linear", + side_effect=AssertionError( + "native block32 unexpectedly fell back to Triton" + ), + ), ): self._check_backend(backend, [backend], self._shapes, self._build_layer) @@ -357,15 +363,21 @@ def test_triton_override(self): def test_preserves_checkpoint_and_prequantized_input(self): torch.manual_seed(7) - with mock.patch.object(fp8_utils, "FP8_GEMM_RUNNER_BACKEND", Fp8GemmRunnerBackend.DEEP_GEMM): + with mock.patch.object( + fp8_utils, "FP8_GEMM_RUNNER_BACKEND", Fp8GemmRunnerBackend.DEEP_GEMM + ): layer, _ = self._build_layer(96, 288) original_weight = layer.weight.detach().clone() layer.quant_method.process_weights_after_loading(layer) self.assertIsNone(getattr(layer, "_block_fp8_bf16_weight", None)) - torch.testing.assert_close(layer.weight.float(), original_weight.float(), rtol=0, atol=0) + torch.testing.assert_close( + layer.weight.float(), original_weight.float(), rtol=0, atol=0 + ) self.assertEqual(tuple(layer.weight_scale_inv.shape), (3, 9)) x = torch.randn(5, 288, device="cuda", dtype=torch.bfloat16) / 10 - q, scales = fp8_utils.sglang_per_token_group_quant_fp8(x, 32, scale_ue8m0=True) + q, scales = fp8_utils.sglang_per_token_group_quant_fp8( + x, 32, scale_ue8m0=True + ) bias = torch.randn(96, device="cuda", dtype=torch.bfloat16) plain = layer.quant_method.apply(layer, x, bias) quantized = layer.quant_method.apply(layer, (q, scales), bias) From 848b1f3926586cf1e7afb15983a9b54fc89954e7 Mon Sep 17 00:00:00 2001 From: Vladislav Nosivskoy Date: Tue, 6 Oct 2026 16:49:54 +0300 Subject: [PATCH 4/5] Fix native block32 DeepGEMM fallbacks --- python/sglang/srt/layers/quantization/fp8.py | 10 +++- .../srt/layers/quantization/fp8_utils.py | 57 ++++++++++--------- .../test_fp8_blockwise_linear_backends.py | 31 ++++++++++ 3 files changed, 69 insertions(+), 29 deletions(-) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 87e68e103ec3..93c79d1ce55b 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/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, @@ -829,7 +830,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 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] diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index b7a147cb799c..c539604b522c 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1238,6 +1238,17 @@ 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, @@ -1249,39 +1260,13 @@ def deepgemm_w8a8_block_fp8_linear_with_fallback( weight_bf16: Optional[torch.Tensor] = None, ) -> torch.Tensor: block32 = list(block_size) == [32, 32] and get_platform().is_sm90 - shape_supported = ( - weight.shape[0] % 8 == 0 and weight.shape[1] % 32 == 0 - if block32 - else weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0 - ) - block_supported = block32 or list(block_size) == [128, 128] 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. assert not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 assert input.dtype == torch.float8_e4m3fn - assert shape_supported and block_supported, ( - "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 - 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 @@ -1301,6 +1286,22 @@ def deepgemm_w8a8_block_fp8_linear_with_fallback( 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]] diff --git a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py index 9f6aeaea9a0e..f8daeda7e318 100644 --- a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py +++ b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py @@ -361,6 +361,37 @@ def test_deepgemm_and_auto_use_native_blocks(self): def test_triton_override(self): self._check_backend("triton", ["triton"], self._shapes, self._build_layer) + def test_unsupported_shape_keeps_triton_and_bf16_fallback(self): + torch.manual_seed(7) + with ( + mock.patch.object( + fp8_utils, "FP8_GEMM_RUNNER_BACKEND", Fp8GemmRunnerBackend.DEEP_GEMM + ), + mock.patch.object( + fp8_utils.envs.SGLANG_OPT_HOPPER_BLOCK_FP8_BF16, + "get", + return_value=True, + ), + mock.patch.object( + fp8_utils, + "w8a8_block_fp8_matmul_deepgemm", + side_effect=AssertionError("unsupported shape reached DeepGEMM"), + ), + ): + layer, w_dequant = self._build_layer(32, 64) + layer.weight.data = layer.weight.data[:10].contiguous() + layer.quant_method.process_weights_after_loading(layer) + self.assertIsNotNone(getattr(layer, "_block_fp8_bf16_weight", None)) + x = torch.randn(65, 64, device="cuda", dtype=torch.bfloat16) / 10 + q, scales = fp8_utils.sglang_per_token_group_quant_fp8( + x, 32, scale_ue8m0=True + ) + ref = fp8_utils.dequant_group_fp8_to_bf16(q, scales).float() + ref = ref @ w_dequant[:10].T + for inputs in (x, (q, scales)): + out = layer.quant_method.apply(layer, inputs) + assert_output_close(self, out, ref, rtol=5e-2, atol=1e-1) + def test_preserves_checkpoint_and_prequantized_input(self): torch.manual_seed(7) with mock.patch.object( From b093f33e6c21097c6e833c55ad232d2e7b75893e Mon Sep 17 00:00:00 2001 From: Vladislav Nosivskoy Date: Tue, 6 Oct 2026 17:44:10 +0300 Subject: [PATCH 5/5] Preserve UE8M0 activation policy in DeepGEMM --- .../srt/layers/quantization/fp8_utils.py | 15 +++++++----- python/sglang/srt/models/deepseek_v2.py | 3 ++- .../test_fp8_blockwise_linear_backends.py | 24 +++++++++++++++++++ 3 files changed, 35 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index c539604b522c..843a757d527a 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -626,10 +626,14 @@ def dispatch_w8a8_block_fp8_linear( # 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( @@ -1259,7 +1263,6 @@ def deepgemm_w8a8_block_fp8_linear_with_fallback( act_scale_ue8m0: bool = False, weight_bf16: Optional[torch.Tensor] = None, ) -> torch.Tensor: - block32 = list(block_size) == [32, 32] and get_platform().is_sm90 if input_scale is not None: assert not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 assert input.dtype == torch.float8_e4m3fn @@ -1305,7 +1308,7 @@ def deepgemm_w8a8_block_fp8_linear_with_fallback( input_2d = input.view(-1, input.shape[-1]) output_shape = [*input.shape[:-1], weight.shape[0]] - if block32 and act_scale_ue8m0: + 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( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 737037f2b8bb..8a3c3223006f 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1845,7 +1845,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 diff --git a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py index f8daeda7e318..536f3b5b37be 100644 --- a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py +++ b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py @@ -192,6 +192,30 @@ def test_triton(self): def test_deep_gemm(self): self._run("deep_gemm") + @unittest.skipUnless( + get_device_sm() == 90, "SM90 UE8M0 activation-scale regression" + ) + def test_deep_gemm_ue8m0_activation_scales(self): + torch.manual_seed(7) + x = torch.randn(3, 128, device="cuda", dtype=torch.bfloat16) + weight = torch.eye(128, device="cuda").to(torch.float8_e4m3fn) + weight_scale = torch.ones(1, 1, device="cuda") + q, scale = fp8_utils.sglang_per_token_group_quant_fp8(x, 128, scale_ue8m0=True) + ref = (q.float() * scale).to(torch.bfloat16) + for backend in (Fp8GemmRunnerBackend.AUTO, Fp8GemmRunnerBackend.DEEP_GEMM): + with mock.patch.object(fp8_utils, "FP8_GEMM_RUNNER_BACKEND", backend): + self.assertIs( + fp8_utils.dispatch_w8a8_block_fp8_linear([128, 128]), + fp8_utils.deepgemm_w8a8_block_fp8_linear_with_fallback, + ) + linear = fp8_utils.dispatch_w8a8_block_fp8_linear( + [128, 128], act_scale_ue8m0=True + ) + for n in (128, 96): + with self.subTest(backend=backend, n=n): + out = linear(x, weight[:n], [128, 128], weight_scale) + torch.testing.assert_close(out, ref[:, :n], rtol=0, atol=0) + def test_flashinfer_trtllm(self): self._run("flashinfer_trtllm")