[Fix] [SM120] quantization: Requantize block-FP8 MoE weights to UE8M0 for DeepGEMM

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

@@ -2205,6 +2205,7 @@ def process_weights_after_loading_block_quant(self, layer: Module) -> None:
                         use_deepgemm_runner=will_use_deepgemm,
                         output_dtype=torch.bfloat16,
                         weight_shape=weight.shape[-2:],
+                        grouped_moe=True,
                     )
 
     def _convert_mxfp8_moe_to_block_fp8(self, layer: Module) -> None:

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

@@ -2023,6 +2023,8 @@ def requant_block_scale_ue8m0_for_deepgemm(
     use_deepgemm_runner: bool,
     output_dtype: Optional[torch.dtype] = None,
     weight_shape=None,
+    *,
+    grouped_moe: bool = False,
 ) -> bool:
     """Requantize block-FP8 weight scales to UE8M0 in place for DeepGEMM.
 
@@ -2041,6 +2043,7 @@ def requant_block_scale_ue8m0_for_deepgemm(
             weight_block_size=weight_block_size,
             output_dtype=output_dtype,
             weight_shape=weight_shape,
+            grouped_moe=grouped_moe,
         )
     ):
         return False

python/sglang/srt/model_loader/utils.py

@@ -264,7 +264,7 @@ def get_architecture_class_name(model_config: ModelConfig) -> str:
 
 
 def should_deepgemm_weight_requant_ue8m0(
-    weight_block_size, output_dtype=None, weight_shape=None
+    weight_block_size, output_dtype=None, weight_shape=None, *, grouped_moe=False
 ):
     """Should we requant fp8 weights into UE8M0 format when loading the model.
 
@@ -281,7 +281,7 @@ def should_deepgemm_weight_requant_ue8m0(
         return False
     # SM120 routes dense block-FP8 GEMMs to CUTLASS/Triton (fp32 scales);
     # only the grouped MoE GEMM consumes DeepGEMM layouts there.
-    if get_device_sm() == 120:
+    if get_device_sm() == 120 and not grouped_moe:
         return False
     if output_dtype is not None and output_dtype != torch.bfloat16:
         return False

test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py

@@ -11,6 +11,7 @@
 from compressed_tensors.quantization import QuantizationStrategy
 
 import sglang.srt.layers.quantization.fp8_utils as fp8_utils
+import sglang.srt.model_loader.utils as model_loader_utils
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.quantization import fp8 as fp8_quant
 from sglang.srt.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_fp8 import (
@@ -65,6 +66,37 @@ def test_helper_requants_supported_deepgemm_bf16_once(self):
         self.assertTrue(weight_scale.format_ue8m0)
         requant.assert_called_once_with(weight, weight_scale, BLOCK_SIZE)
 
+    def test_sm120_requants_grouped_moe_but_not_dense_weights(self):
+        """SM120's DeepGEMM grouped MoE GEMM faults on FP32 block scales that are
+        not powers of two; dense block-FP8 GEMMs there keep FP32 scales."""
+        dense = _make_params()
+        experts = _make_params()
+
+        with (
+            self._enabled_deepgemm_ue8m0(),
+            patch.object(model_loader_utils, "get_device_sm", return_value=120),
+            patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant,
+        ):
+            dense_fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
+                *dense,
+                BLOCK_SIZE,
+                use_deepgemm_runner=True,
+                output_dtype=torch.bfloat16,
+                weight_shape=dense[0].shape,
+            )
+            experts_fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
+                *experts,
+                BLOCK_SIZE,
+                use_deepgemm_runner=True,
+                output_dtype=torch.bfloat16,
+                weight_shape=experts[0].shape,
+                grouped_moe=True,
+            )
+
+        self.assertFalse(dense_fired)
+        self.assertTrue(experts_fired)
+        requant.assert_called_once_with(*experts, BLOCK_SIZE)
+
     def test_helper_skips_non_bf16_output(self):
         weight, weight_scale = _make_params()
 
@@ -213,6 +245,7 @@ def _mark_ue8m0(weight, weight_scale, *args, **kwargs):
                     use_deepgemm_runner=True,
                     output_dtype=torch.bfloat16,
                     weight_shape=layer.w13_weight.shape[-2:],
+                    grouped_moe=True,
                 ),
                 call(
                     layer.w2_weight,
@@ -221,6 +254,7 @@ def _mark_ue8m0(weight, weight_scale, *args, **kwargs):
                     use_deepgemm_runner=True,
                     output_dtype=torch.bfloat16,
                     weight_shape=layer.w2_weight.shape[-2:],
+                    grouped_moe=True,
                 ),
             ],
         )