diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index e1914d287..d79a1e105 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -905,9 +905,6 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( assert quant_info.output2_scales_scalar is not None 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 with use_symmetric_memory( @@ -928,7 +925,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( output = trtllm_fp8_per_tensor_scale_moe_wrapper( routing_logits=router_logits, - routing_bias=routing_bias_cast, + routing_bias=correction_bias, hidden_states=a_q, gemm1_weights=quant_info.w13_weight, output1_scales_scalar=quant_info.output1_scales_scalar, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py index 3756ac26e..50e3ea6cf 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py @@ -310,11 +310,7 @@ class CompressedTensorsMxInt4MoE(CompressedTensorsMoEScheme): router_logits = topk_output.router_logits topk_config = topk_output.topk_config - correction_bias = ( - None - if topk_config.correction_bias is None - else topk_config.correction_bias.to(x.dtype) - ) + correction_bias = topk_config.correction_bias local_num_experts = self.moe_runner_config.num_local_experts routing_method_type = layer.routing_method_type