A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 14:45 UTC
[Bugfix][MXFP4] Call `trtllm_fp4_block_scale_moe` with kwargs (#33104)
Signed-off-by: Pengchao Wang <wpc@fb.com>
Pengchao Wang committed
Jan 26, 2026 at 15:13 UTC
8caffd92dfb011733b3c71c76756a687f64eb3ae
1 file changed
+26
-26
vllm/model_executor/layers/quantization/mxfp4.py
+26
-26
@@ -1053,32 +1053,32 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1053
x_scale = x_scale.view(torch.float8_e4m3fn).reshape(*x.shape[:-1], -1)
1054
1055
trtllm_gen_output = trtllm_fp4_block_scale_moe(
1056
- router_logits.to(torch.bfloat16),
1057
- None, # routing_bias
1058
- x_quant,
1059
- x_scale,
1060
- layer.w13_weight, # uint8 (e2m1 x 2)
1061
- layer.w13_weight_scale, # uint8 (e4m3 x 2)
1062
- layer.w13_bias, # fp32 per expert per channel
1063
- layer.gemm1_alpha, # fp32 per expert
1064
- layer.gemm1_beta, # fp32 per expert
1065
- layer.gemm1_clamp_limit, # fp32 per expert
1066
- layer.w2_weight, # uint8 (e2m1 x 2)
1067
- layer.w2_weight_scale, # ue8m0
1068
- layer.w2_bias, # fp32 per expert per channel
1069
- None, # output1_scale_scalar
1070
- None, # output1_scale_gate_scalar
1071
- None, # output2_scale_scalar
1072
- layer.global_num_experts,
1073
- layer.top_k,
1074
- None, # n_group
1075
- None, # topk_group
1076
- self.intermediate_size, # padded to multiple of 256
1077
- layer.ep_rank * layer.local_num_experts, # local_expert_offset
1078
- self.num_experts, # local num experts
1079
- None, # routed_scaling_factor
1080
- 1 if layer.renormalize else 0, # routing_method_type, renormalize
1081
- True, # do finalize
1056
+ routing_logits=router_logits.to(torch.bfloat16),
1057
+ routing_bias=None,
1058
+ hidden_states=x_quant,
1059
+ hidden_states_scale=x_scale,
1060
+ gemm1_weights=layer.w13_weight, # uint8 (e2m1 x 2)
1061
+ gemm1_weights_scale=layer.w13_weight_scale, # uint8 (e4m3 x 2)
1062
+ gemm1_bias=layer.w13_bias, # fp32 per expert per channel
1063
+ gemm1_alpha=layer.gemm1_alpha, # fp32 per expert
1064
+ gemm1_beta=layer.gemm1_beta, # fp32 per expert
1065
+ gemm1_clamp_limit=layer.gemm1_clamp_limit, # fp32 per expert
1066
+ gemm2_weights=layer.w2_weight, # uint8 (e2m1 x 2)
1067
+ gemm2_weights_scale=layer.w2_weight_scale, # ue8m0
1068
+ gemm2_bias=layer.w2_bias, # fp32 per expert per channel
1069
+ output1_scale_scalar=None,
1070
+ output1_scale_gate_scalar=None,
1071
+ output2_scale_scalar=None,
1072
+ num_experts=layer.global_num_experts,
1073
+ top_k=layer.top_k,
1074
+ n_group=None,
1075
+ topk_group=None,
1076
+ intermediate_size=self.intermediate_size, # padded to multiple of 256
1077
+ local_expert_offset=layer.ep_rank * layer.local_num_experts,
1078
+ local_num_experts=self.num_experts,
1079
+ routed_scaling_factor=None,
1080
+ routing_method_type=1 if layer.renormalize else 0,
1081
+ do_finalize=True,
1082
tune_max_num_tokens=max(self.max_capture_size, 1),
1083
)[0]
1084
return trtllm_gen_output