diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py index 64b786375..f178c5cc8 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py @@ -7,7 +7,6 @@ import torch from torch.nn import Module 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.device_communicators.pynccl_allocator import ( use_symmetric_memory, @@ -295,8 +294,6 @@ class Mxfp4FlashinferTrtllmMoEMethod: else: raise ValueError(f"Unsupported topk output format: {topk_output.format}") - packed_topk = PackTopkIds.execute(topk_ids, topk_weights) - precision = self.flashinfer_mxfp4_moe_precision if precision == "bf16": assert hidden_states.dtype == torch.bfloat16 @@ -345,7 +342,7 @@ class Mxfp4FlashinferTrtllmMoEMethod: ) output = trtllm_fp4_block_scale_routed_moe( - topk_ids=packed_topk, + topk_ids=(topk_ids, topk_weights), routing_bias=None, hidden_states=x_quant, hidden_states_scale=x_scale, @@ -362,7 +359,7 @@ class Mxfp4FlashinferTrtllmMoEMethod: output1_scale_gate_scalar=layer.output1_scale_gate_scalar, output2_scale_scalar=layer.output2_scale_scalar, num_experts=layer.num_experts, - top_k=packed_topk.shape[1], + top_k=topk_ids.shape[1], n_group=1, topk_group=1, intermediate_size=intermediate_size,