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 11, 16:15 UTC
[Bugfix] Fix grouped_topk pytorch impl when num_experts can't be grouped properly (#29439)
Signed-off-by: Divakar Verma <divakar.verma@amd.com> Co-authored-by: Gregory Shtrasberg <156009573+gshtras@users.noreply.github.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Divakar Verma committed
Dec 10, 2025 at 21:47 UTC
d1e1fb4363c61080b7cd20469d5a751e88a1cdb3
1 file changed
+9
-1
vllm/model_executor/layers/fused_moe/layer.py
+9
-1
@@ -1556,6 +1556,14 @@ class FusedMoE(CustomOp):
1556
f"EPLB is not supported for {self.quant_method.method_name}."
1557
)
1558
1559
+ def valid_grouping() -> bool:
1560
+ # Check if num_experts is greater than num_expert_group
1561
+ # and is divisible by num_expert_group
1562
+ num_experts = router_logits.shape[-1]
1563
+ if num_experts <= self.num_expert_group:
1564
+ return False
1565
+ return num_experts % self.num_expert_group == 0
1566
+
1567
indices_type = self.quant_method.topk_indices_dtype
1568
1569
# Check if we should use a routing simulation strategy
@@ -1570,7 +1578,7 @@ class FusedMoE(CustomOp):
1578
)
1579
1580
# DeepSeekv2 uses grouped_top_k
1573
- elif self.use_grouped_topk:
1581
+ elif self.use_grouped_topk and valid_grouping():
1582
assert self.topk_group is not None
1583
assert self.num_expert_group is not None
1584
if rocm_aiter_ops.is_fused_moe_enabled():