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

[Deepseek V4] Keep fp32 routing weights in the mxfp4 trtllm MoE (#33608)

Browse Source 
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
Brayden Zhong
2026-09-07 00:12:55 +00:00
committed by GitHub
co-authored by Brayden Zhong
parent 2e8c03e2c7
commit 30705c004c
1 changed files with 2 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py
+2 -5
View File
@@ -7,7 +7,6 @@ import torch
from torch.nn import Module from torch.nn import Module
from torch.nn.parameter import Parameter from torch.nn.parameter import Parameter
from sglang.kernels.ops.moe.pack_topk_ids import PackTopkIds
from sglang.srt.distributed import get_tp_group from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
@@ -295,8 +294,6 @@ class Mxfp4FlashinferTrtllmMoEMethod:
else: else:
raise ValueError(f"Unsupported topk output format: {topk_output.format}") raise ValueError(f"Unsupported topk output format: {topk_output.format}")
packed_topk = PackTopkIds.execute(topk_ids, topk_weights)
precision = self.flashinfer_mxfp4_moe_precision precision = self.flashinfer_mxfp4_moe_precision
if precision == "bf16": if precision == "bf16":
assert hidden_states.dtype == torch.bfloat16 assert hidden_states.dtype == torch.bfloat16
@@ -345,7 +342,7 @@ class Mxfp4FlashinferTrtllmMoEMethod:
) )
output = trtllm_fp4_block_scale_routed_moe( output = trtllm_fp4_block_scale_routed_moe(
topk_ids=packed_topk, topk_ids=(topk_ids, topk_weights),
routing_bias=None, routing_bias=None,
hidden_states=x_quant, hidden_states=x_quant,
hidden_states_scale=x_scale, hidden_states_scale=x_scale,
@@ -362,7 +359,7 @@ class Mxfp4FlashinferTrtllmMoEMethod:
output1_scale_gate_scalar=layer.output1_scale_gate_scalar, output1_scale_gate_scalar=layer.output1_scale_gate_scalar,
output2_scale_scalar=layer.output2_scale_scalar, output2_scale_scalar=layer.output2_scale_scalar,
num_experts=layer.num_experts, num_experts=layer.num_experts,
top_k=packed_topk.shape[1], top_k=topk_ids.shape[1],
n_group=1, n_group=1,
topk_group=1, topk_group=1,
intermediate_size=intermediate_size, intermediate_size=intermediate_size,