From 7ab548ef641b41ee36da82b67b68761987b39b7c Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Thu, 27 Nov 2025 09:00:26 -0500 Subject: [PATCH] [2/2] Refactor DeepGeem requant for FP8 FusedMoE on Blackwell (#13960) --- python/sglang/srt/layers/quantization/fp8.py | 31 ++++++++++++++--- python/sglang/srt/models/deepseek_v2.py | 35 -------------------- python/sglang/srt/server_args.py | 5 +++ 3 files changed, 32 insertions(+), 39 deletions(-) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 5f0962769..878f70619 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -764,8 +764,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): w2_weight_scale, requires_grad=False ) layer.w2_input_scale = None - - if _use_aiter: + elif _use_aiter: # Pre-shuffle weights layer.w13_weight.data = shuffle_weight( layer.w13_weight.contiguous(), (16, 16) @@ -773,13 +772,37 @@ class Fp8MoEMethod(FusedMoEMethodBase): layer.w2_weight.data = shuffle_weight( layer.w2_weight.contiguous(), (16, 16) ) - - if _is_cpu: + elif _is_cpu: assert ( _is_cpu_amx_available ), "Fp8MoEMethod on CPU requires that CPU has AMX support" _amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"]) + else: + # For fp8 moe run with deepgemm, the expert weights and scales need be requantized to ue8m0 + from sglang.srt.layers.moe import get_moe_runner_backend + from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE + 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 get_moe_runner_backend().is_deep_gemm() + ): + assert isinstance( + layer, DeepEPMoE + ), "DeepGemm MoE is only supported with DeepEPMoE" + weight_block_size = self.quant_config.weight_block_size + requant_weight_ue8m0_inplace( + layer.w13_weight, layer.w13_weight_scale_inv, weight_block_size + ) + requant_weight_ue8m0_inplace( + layer.w2_weight, layer.w2_weight_scale_inv, weight_block_size + ) return # If checkpoint is fp16 or bfloat16, quantize in place. diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index e09881293..163d12845 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -114,7 +114,6 @@ from sglang.srt.layers.quantization.fp8_utils import ( inverse_transform_scale_ue8m0, normalize_e4m3fn_to_e4m3fnuz, quant_weight_ue8m0, - requant_weight_ue8m0_inplace, transform_scale_ue8m0_inplace, ) from sglang.srt.layers.quantization.int8_utils import ( @@ -3684,43 +3683,9 @@ 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 - # Requant the weights and scales of MoE layers - if get_moe_runner_backend().is_deep_gemm(): - self._maybe_moe_weight_requant_ue8m0(is_nextn) if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config): self._transform_scale_nextn_moe_ue8m0() - 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( - range( - self.config.first_k_dense_replace, - self.config.num_hidden_layers, - self.config.moe_layer_freq, - ) - ) - - 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] - - if layer_id in moe_layers or is_nextn: - 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) - # TODO avoid code dup (currently combine from weight_requant_ue8m0 and transform_scale_ue8m0) def _transform_scale_nextn_moe_ue8m0(self): layer = self.model.decoder diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 5b13ac801..07a31229a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1517,6 +1517,11 @@ class ServerArgs: self.ep_size == 1 ), "FP8 Cutlass MoE is only supported with ep_size == 1" + if self.moe_runner_backend == "deep_gemm": + assert ( + self.ep_size > 1 and self.moe_a2a_backend == "deepep" + ), "DeepGemm MoE runner is only supported when ep is enabled and moe_a2a_backend is deepep" + def _handle_a2a_moe(self): if self.moe_a2a_backend == "deepep": if self.deepep_mode == "normal":