[NVIDIA] Fix FP8 gemm performance with fp16 models (MInimax-M2.5) (#22300)

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

@@ -534,11 +534,16 @@ class Fp8LinearMethod(LinearMethodBase):
                 should_deepgemm_weight_requant_ue8m0,
             )
 
+            # Only requantize to UE8M0 if DeepGEMM can actually run
+            # this layer. If the dtype or shape is unsupported, the GEMM
+            # falls back to triton at runtime, which needs float32 scales.
             if (
                 should_deepgemm_weight_requant_ue8m0(
                     weight_block_size=getattr(
                         self.quant_config, "weight_block_size", None
                     ),
+                    output_dtype=getattr(layer, "orig_dtype", None),
+                    weight_shape=layer.weight.shape,
                 )
                 and (
                     self.w8a8_block_fp8_linear

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

@@ -496,8 +496,11 @@ def flashinfer_gemm_w8a8_block_fp8_linear_with_fallback(
 
     input_2d = input.view(-1, input.shape[-1])
     backend = _get_flashinfer_groupwise_backend()
-    # TRTLLM backend requires K dimension >= 256.
+    # TRTLLM backend requires K >= 256 and weight scales in UE8M0/R128c4
-    if backend == "trtllm" and input_2d.shape[1] < 256:
+    # packed format. Fall back to triton when scales are plain float32.
+    if backend == "trtllm" and (
+        input_2d.shape[1] < 256 or not getattr(weight_scale, "format_ue8m0", False)
+    ):
         return triton_w8a8_block_fp8_linear(
             input, weight, block_size, weight_scale, input_scale, bias
         )

python/sglang/srt/model_loader/utils.py

@@ -249,13 +249,41 @@ def get_architecture_class_name(model_config: ModelConfig) -> str:
     return get_model_architecture(model_config)[1]
 
 
-def should_deepgemm_weight_requant_ue8m0(weight_block_size):
+def post_load_weights(model: nn.Module, model_config: ModelConfig):
-    """Should we requant fp8 weights into UE8M0 format when loading the model"""
+    # Model weight loading consists of two stages:
-    return (
+    # 1. Initial weight loading.
+    # 2. Post-processing of weights, including assigning specific member variables.
+    # For `dummy_init`, only the second stage is required.
+    if hasattr(model, "post_load_weights"):
+        if model_config.hf_config.architectures[0] == "DeepseekV3ForCausalLMNextN":
+            model.post_load_weights(is_nextn=True)
+        else:
+            model.post_load_weights()
+
+
+def should_deepgemm_weight_requant_ue8m0(
+    weight_block_size, output_dtype=None, weight_shape=None
+):
+    """Should we requant fp8 weights into UE8M0 format when loading the model.
+
+    When output_dtype or weight_shape are provided, also checks that DeepGEMM
+    can actually run this layer at runtime (bf16 output, N%64==0, K%128==0).
+    Without these checks, scales would be converted to UE8M0 but the GEMM would
+    fall back to triton which expects float32 scales, causing wrong results.
+    """
+    if not (
         deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
         and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
         and weight_block_size is not None
-    )
+    ):
+        return False
+    if output_dtype is not None and output_dtype != torch.bfloat16:
+        return False
+    if weight_shape is not None and (
+        weight_shape[0] % 64 != 0 or weight_shape[1] % 128 != 0
+    ):
+        return False
+    return True
 
 
 def should_async_load(weight: torch.Tensor) -> bool: