From 4edb2401129887a6c3e400868ad3ee199292092a Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Thu, 13 Nov 2025 00:18:58 +0800 Subject: [PATCH] Fuse routed_scaling_factor to fused_marlin_moe (#12998) --- .../compressed_tensors/compressed_tensors_moe.py | 1 + python/sglang/srt/models/deepseek_v2.py | 12 +----------- 2 files changed, 2 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 30883e391..5c66a472f 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -653,5 +653,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): num_bits=self.num_bits, is_k_full=self.is_k_full, expert_map=torch.empty(1, device=x.device), + routed_scaling_factor=self.moe_runner_config.routed_scaling_factor, ) return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 5411052e9..ed19fe474 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -81,9 +81,6 @@ from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat from sglang.srt.layers.quantization.base_config import QuantizationConfig -from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors_moe import ( - CompressedTensorsWNA16MoEMethod, -) from sglang.srt.layers.quantization.fp8 import Fp8Config from sglang.srt.layers.quantization.fp8_kernel import ( is_fp8_fnuz, @@ -783,13 +780,7 @@ class DeepseekV2MoE(nn.Module): router_logits = self.gate(hidden_states, gemm_output_zero_allocator) topk_output = self.topk(hidden_states, router_logits) final_hidden_states = self.experts(hidden_states, topk_output) - if ( - not _is_cuda - or isinstance(self.experts.quant_method, KTEPWrapperMethod) - or isinstance( - self.experts.quant_method, CompressedTensorsWNA16MoEMethod - ) - ): + if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod): final_hidden_states *= self.routed_scaling_factor current_stream.wait_stream(self.alt_stream) @@ -852,7 +843,6 @@ class DeepseekV2MoE(nn.Module): not _is_cuda and not _use_aiter or isinstance(self.experts.quant_method, KTEPWrapperMethod) - or isinstance(self.experts.quant_method, CompressedTensorsWNA16MoEMethod) ): # fused in biased_grouped_topk so we can skip here final_hidden_states *= self.routed_scaling_factor