chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Fuse routed_scaling_factor to fused_marlin_moe (#12998)

Browse Source
This commit is contained in:
Ke Bao
2025-11-13 00:18:58 +08:00
committed by GitHub
parent 706502ff6c
commit 4edb240112
2 changed files with 2 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py
View File

@@ -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)

12
python/sglang/srt/models/deepseek_v2.py
View File

@@ -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