update marlin M size for EP (#48144)
Signed-off-by: gnovack <novackgm@gmail.com>
gnovack committed
Jul 9, 2026 at 16:52 UTC
f1a5adddb815610db46cae81f2a1d5a4609bf99d
1 file changed
+7
-2
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+7
-2
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Fused MoE utilities for GPTQ."""
4
5
+import math
6
from collections.abc import Callable
7
8
import torch
@@ -312,6 +313,12 @@ def fused_marlin_moe(
313
assert num_bits in [4, 8]
314
assert topk_weights.dtype == torch.float32
315
316
+ if global_num_experts == -1:
317
+ global_num_experts = E
318
+ else:
319
+ # Set M to estimated valid tokens per rank
320
+ M = math.ceil(M * E / global_num_experts)
321
+
322
# M block size selection logic
323
# TODO: tune this further for specific models
324
for block_size_m in [8, 16, 32, 48, 64]:
@@ -321,8 +328,6 @@ def fused_marlin_moe(
328
if input_dtype is not None and input_dtype.itemsize == 1:
329
block_size_m = max(block_size_m, 16)
330
324
- if global_num_experts == -1:
325
- global_num_experts = E
331
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
332
topk_ids,
333
block_size_m,