minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Drop the routing bias casts in flashinfer trtllm MoE (#33591)

Browse Source 
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
Brayden Zhong
2026-09-07 09:04:33 -07:00
committed by GitHub
co-authored by Brayden Zhong
parent 62bca081a3
commit c99d906eff
2 changed files with 2 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+1 -4
View File
@@ -905,9 +905,6 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
assert quant_info.output2_scales_scalar is not None assert quant_info.output2_scales_scalar is not None
a_q, _ = scaled_fp8_quant(hidden_states, quant_info.w13_input_scale) a_q, _ = scaled_fp8_quant(hidden_states, quant_info.w13_input_scale)
routing_bias_cast = (
None if correction_bias is None else correction_bias.to(torch.bfloat16)
)
# Allocate output inside symmetric memory context # Allocate output inside symmetric memory context
with use_symmetric_memory( with use_symmetric_memory(
@@ -928,7 +925,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
output = trtllm_fp8_per_tensor_scale_moe_wrapper( output = trtllm_fp8_per_tensor_scale_moe_wrapper(
routing_logits=router_logits, routing_logits=router_logits,
routing_bias=routing_bias_cast, routing_bias=correction_bias,
hidden_states=a_q, hidden_states=a_q,
gemm1_weights=quant_info.w13_weight, gemm1_weights=quant_info.w13_weight,
output1_scales_scalar=quant_info.output1_scales_scalar, output1_scales_scalar=quant_info.output1_scales_scalar,
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py
+1 -5
View File
@@ -310,11 +310,7 @@ class CompressedTensorsMxInt4MoE(CompressedTensorsMoEScheme):
router_logits = topk_output.router_logits router_logits = topk_output.router_logits
topk_config = topk_output.topk_config topk_config = topk_output.topk_config
correction_bias = ( correction_bias = topk_config.correction_bias
None
if topk_config.correction_bias is None
else topk_config.correction_bias.to(x.dtype)
)
local_num_experts = self.moe_runner_config.num_local_experts local_num_experts = self.moe_runner_config.num_local_experts
routing_method_type = layer.routing_method_type routing_method_type = layer.routing_method_type