[1/2] Refactor DeepGeem requant for FP8 Linear on Blackwell (#13601)

Co-authored-by: fy1214

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

@@ -44,6 +44,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
     dispatch_w8a8_block_fp8_linear,
     input_to_float8,
     normalize_e4m3fn_to_e4m3fnuz,
+    requant_weight_ue8m0_inplace,
 )
 from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
 from sglang.srt.layers.quantization.marlin_utils_fp8 import (
@@ -261,6 +262,7 @@ class Fp8LinearMethod(LinearMethodBase):
         layer.input_size_per_partition = input_size_per_partition
         layer.output_size_per_partition = output_size_per_partition
         layer.orig_dtype = params_dtype
+        layer.executed_weight_requant_ue8m0 = False
 
         # WEIGHT
         weight_dtype = (
@@ -347,7 +349,34 @@ class Fp8LinearMethod(LinearMethodBase):
                 )
                 return
             else:
+                # For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0
+                from sglang.srt.layers.quantization.fp8_utils import (
+                    deepgemm_w8a8_block_fp8_linear_with_fallback,
+                )
+                from sglang.srt.model_loader.utils import (
+                    should_deepgemm_weight_requant_ue8m0,
+                )
+
+                if (
+                    should_deepgemm_weight_requant_ue8m0(
+                        weight_block_size=getattr(
+                            self.quant_config, "weight_block_size", None
+                        ),
+                    )
+                    and (
+                        self.w8a8_block_fp8_linear
+                        is deepgemm_w8a8_block_fp8_linear_with_fallback
+                    )
+                    and (not layer.executed_weight_requant_ue8m0)
+                ):
+                    requant_weight_ue8m0_inplace(
+                        layer.weight,
+                        layer.weight_scale_inv,
+                        self.quant_config.weight_block_size,
+                    )
+                    layer.executed_weight_requant_ue8m0 = True
                 weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data
+
             layer.weight.data = weight.data
             layer.weight_scale_inv.data = weight_scale.data
         else:

python/sglang/srt/models/deepseek_nextn.py

@@ -220,7 +220,6 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
             use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
         )
         self.logits_processor = LogitsProcessor(config)
-        self._executed_weight_requant_ue8m0 = False
 
     @torch.no_grad()
     def forward(

python/sglang/srt/models/deepseek_v2.py

@@ -108,7 +108,6 @@ from sglang.srt.layers.quantization.fp8_kernel import (
     per_token_group_quant_mla_deep_gemm_masked_fp8,
 )
 from sglang.srt.layers.quantization.fp8_utils import (
-    ENABLE_FLASHINFER_FP8_GEMM,
     block_quant_dequant,
     block_quant_to_tensor_quant,
     channel_quant_to_tensor_quant,
@@ -3419,7 +3418,6 @@ class DeepseekV2ForCausalLM(nn.Module):
             }
         )
         self.capture_aux_hidden_states = False
-        self._executed_weight_requant_ue8m0 = False
 
         self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
         if self.nsa_enable_prefill_cp:
@@ -3594,13 +3592,14 @@ class DeepseekV2ForCausalLM(nn.Module):
                         weight = w
                         weight_scale = self_attn.kv_b_proj.weight_scale_inv
 
+                    # In multiple weight loading scenarios (e.g. RL), we need to inverse the scale of the weights after the requantization happened at the first loading.
                     if (
                         should_deepgemm_weight_requant_ue8m0(
                             weight_block_size=getattr(
                                 self.quant_config, "weight_block_size", None
                             )
                         )
-                        and self._executed_weight_requant_ue8m0
+                        and self_attn.kv_b_proj.executed_weight_requant_ue8m0
                     ):
                         weight_scale = inverse_transform_scale_ue8m0(
                             weight_scale, mn=weight.shape[-2]
@@ -3716,27 +3715,15 @@ class DeepseekV2ForCausalLM(nn.Module):
                 self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc.contiguous())
                 self_attn.use_deep_gemm_bmm = True
 
-        if (
+        # Requant the weights and scales of MoE layers
-            not ENABLE_FLASHINFER_FP8_GEMM
+        if get_moe_runner_backend().is_deep_gemm():
-            and should_deepgemm_weight_requant_ue8m0(
+            self._maybe_moe_weight_requant_ue8m0(is_nextn)
-                weight_block_size=getattr(self.quant_config, "weight_block_size", None)
-            )
-            and not self._executed_weight_requant_ue8m0
-        ):
-            self._executed_weight_requant_ue8m0 = True
-            self._weight_requant_ue8m0(is_nextn)
-
-        # TODO can move weight_requant_ue8m0 and transform_scale_ue8m0 into Fp8LinearMethod.process_weights_after_loading
-        if (
-            deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
-            and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
-            and get_bool_env_var("SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN")
-        ):
-            self._transform_scale_ue8m0(is_nextn)
         if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config):
             self._transform_scale_nextn_moe_ue8m0()
 
-    def _weight_requant_ue8m0(self, is_nextn=False):
+    def _maybe_moe_weight_requant_ue8m0(self, is_nextn=False):
+        # Dense fp8 layers will be processed in Fp8LinearMethod.process_weights_after_loading
+        # So we only need to process sparse MoE layers here
         weight_block_size = self.quant_config.weight_block_size
 
         moe_layers = list(
@@ -3755,70 +3742,15 @@ class DeepseekV2ForCausalLM(nn.Module):
             else:
                 layer = self.model.layers[layer_id]
 
-            module_list = [
-                layer.self_attn.kv_b_proj,
-                layer.self_attn.o_proj,
-            ]
-
-            if self.config.q_lora_rank is not None:
-                module_list.append(layer.self_attn.fused_qkv_a_proj_with_mqa)
-                module_list.append(layer.self_attn.q_b_proj)
-            else:
-                module_list.append(layer.self_attn.kv_a_proj_with_mqa)
-                module_list.append(layer.self_attn.q_proj)
-
-            for module in module_list:
-                requant_weight_ue8m0_inplace(
-                    module.weight, module.weight_scale_inv, weight_block_size
-                )
-
             if layer_id in moe_layers or is_nextn:
-                shared_experts = getattr(layer.mlp, "shared_experts", None)
-                if shared_experts is not None:
-                    for module in [
-                        shared_experts.gate_up_proj,
-                        shared_experts.down_proj,
-                    ]:
-                        requant_weight_ue8m0_inplace(
-                            module.weight, module.weight_scale_inv, weight_block_size
-                        )
-
                 experts = layer.mlp.experts
+                # TODO: move this logic to Fp8MoEMethod.process_weights_after_loading
                 if isinstance(experts, DeepEPMoE):
                     for w in [
                         (experts.w13_weight, experts.w13_weight_scale_inv),
                         (experts.w2_weight, experts.w2_weight_scale_inv),
                     ]:
                         requant_weight_ue8m0_inplace(w[0], w[1], weight_block_size)
-            else:
-                mlp = layer.mlp
-                assert isinstance(mlp, DeepseekV2MLP)
-                for module in [
-                    mlp.gate_up_proj,
-                    mlp.down_proj,
-                ]:
-                    requant_weight_ue8m0_inplace(
-                        module.weight, module.weight_scale_inv, weight_block_size
-                    )
-
-    # TODO can move weight_requant_ue8m0 and transform_scale_ue8m0 into Fp8LinearMethod.process_weights_after_loading
-    def _transform_scale_ue8m0(self, is_nextn=False):
-        num_hidden_layers = 1 if is_nextn else self.config.num_hidden_layers
-
-        for layer_id in range(num_hidden_layers):
-            if is_nextn:
-                layer = self.model.decoder
-            else:
-                layer = self.model.layers[layer_id]
-
-            module_list = []
-            if self.config.q_lora_rank is not None:
-                module_list.append(layer.self_attn.q_b_proj)
-
-            for module in module_list:
-                transform_scale_ue8m0_inplace(
-                    module.weight_scale_inv, mn=module.weight.shape[-2]
-                )
 
     # TODO avoid code dup (currently combine from weight_requant_ue8m0 and transform_scale_ue8m0)
     def _transform_scale_nextn_moe_ue8m0(self):

python/sglang/test/test_block_fp8_deep_gemm_blackwell.py

@@ -222,7 +222,7 @@ class TestDeepGemmBlackwell(CustomTestCase):
 
         with torch.inference_mode():
             ref_out = native_w8a8_block_fp8_matmul(
-                A_q, B_q, A_s, B_s, block_size, out_dtype
+                A_qu[0], B_qu[0], A_qu[1], B_qu[1], block_size, out_dtype
             )
             out = torch.empty_like(ref_out)
             fp8_gemm_nt(A_qu, B_qu, out)