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 14, 05:30 UTC
Fix Llama4 FlashInfer FP4 MoE issues (#22511)
Signed-off-by: Po-Han Huang <pohanh@nvidia.com>
Po-Han Huang (NVIDIA) committed
Aug 12, 2025 at 20:50 UTC
67c153b88a2129c3b6fb78af09901738f1034a68
3 files changed
+9
-5
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
-2
@@ -170,8 +170,6 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
170
"w1_scale and w2_scale must not "
171
"be None for FlashInferExperts")
172
173
- assert not apply_router_weight_on_input
174
-
173
quant_scales = [
174
a1_gscale,
175
w1_scale.view(torch.int32),
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+6
-1
@@ -60,7 +60,12 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
60
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
61
Optional[torch.Tensor], Optional[torch.Tensor]]:
62
63
- assert not apply_router_weight_on_input
63
+ if apply_router_weight_on_input:
64
+ topk = topk_ids.size(1)
65
+ # TODO: this only works for topK=1, will need to update for topK>1
66
+ assert topk == 1, \
67
+ "apply_router_weight_on_input is only implemented for topk=1"
68
+ a1.mul_(topk_weights.to(a1.dtype))
69
70
(a1_gscale, use_dp, local_tokens) = extract_required_args(
71
extra_prepare_args, ['a1_gscale', 'use_dp', 'local_tokens'])
vllm/model_executor/layers/quantization/modelopt.py
+3
-2
@@ -1299,8 +1299,9 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1299
output2_scale_scalar=layer.g2_alphas.data,
1300
num_experts=global_num_experts,
1301
top_k=top_k,
1302
- n_group=num_expert_group,
1303
- topk_group=topk_group,
1302
+ n_group=num_expert_group
1303
+ if num_expert_group is not None else 0,
1304
+ topk_group=topk_group if topk_group is not None else 0,
1305
intermediate_size=layer.intermediate_size_per_partition,
1306
local_expert_offset=layer.ep_rank * layer.local_num_experts,
1307
local_num_experts=layer.local_num_experts,