Skip to content

Commit 05c8d53

Browse files
fix(moe): scale FlashInfer b12x capacity for DP
Size the b12x wrapper for tokens contributed by all data-parallel ranks. Fixes #47982 Co-authored-by: Devin AI <noreply@devin.ai> Signed-off-by: Sidharth Rajmohan <dumko.raj@gmail.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
1 parent 1ff9429 commit 05c8d53

1 file changed

Lines changed: 2 additions & 1 deletion

File tree

vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
1616
TopKWeightAndReduceNoOP,
1717
)
18+
from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket
1819
from vllm.model_executor.layers.quantization.utils.quant_utils import (
1920
QuantKey,
2021
kNvfp4Dynamic,
@@ -73,7 +74,7 @@ def __init__(
7374
self.intermediate_size_per_partition = (
7475
moe_config.intermediate_size_per_partition
7576
)
76-
self.max_num_tokens = moe_config.max_num_tokens
77+
self.max_num_tokens = fi_moe_largest_bucket(moe_config)
7778
self.local_expert_offset = self.ep_rank * self.num_local_experts
7879

7980
activation = moe_config.activation

0 commit comments

Comments
 (0)