[MoE] FI autotuning: max bucket = max token count [e.g. `DP_size*MNBT`] (#47427)
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Netanel Haber committed
Jul 7, 2026 at 12:08 UTC
066f02ae94c75d7e3ad7f0faaa2df348fce98c8d
4 files changed
+38
-2
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+2
@@ -11,6 +11,7 @@ from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEQuantConfig,
12
RoutingMethodType,
13
)
14
+from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket
15
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
16
activation_to_flashinfer_int,
17
)
@@ -145,4 +146,5 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
146
routed_scaling_factor=routed_scaling_factor,
147
routing_method_type=self.routing_method_type,
148
activation_type=activation_to_flashinfer_int(activation),
149
+ tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
150
)
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+7
-1
@@ -15,7 +15,10 @@ from vllm.model_executor.layers.fused_moe.config import (
15
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
16
TopKWeightAndReduceNoOP,
17
)
18
-from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights
18
+from vllm.model_executor.layers.fused_moe.utils import (
19
+ fi_moe_largest_bucket,
20
+ trtllm_moe_pack_topk_ids_weights,
21
+)
22
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
23
activation_to_flashinfer_int,
24
)
@@ -249,6 +252,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
252
weight_layout=weight_layout,
253
fp8_quantization_type=fp8_quant_type,
254
output=output,
255
+ tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
256
)
257
258
@@ -419,6 +423,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
423
use_shuffled_weight=use_shuffled_weight,
424
weight_layout=weight_layout,
425
fp8_quantization_type=fp8_quant_type,
426
+ tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
427
)
428
if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL:
429
kwargs["activation_type"] = activation_type
@@ -475,6 +480,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
480
use_routing_scales_on_input=apply_router_weight_on_input,
481
routing_method_type=self.routing_method_type,
482
activation_type=activation_type,
483
+ tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
484
)
485
return out
486
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+8
-1
@@ -16,7 +16,10 @@ from vllm.model_executor.layers.fused_moe.config import (
16
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
17
TopKWeightAndReduceNoOP,
18
)
19
-from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights
19
+from vllm.model_executor.layers.fused_moe.utils import (
20
+ fi_moe_largest_bucket,
21
+ trtllm_moe_pack_topk_ids_weights,
22
+)
23
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
24
activation_to_flashinfer_int,
25
)
@@ -319,6 +322,9 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
322
do_finalize=True,
323
activation_type=activation_to_flashinfer_int(activation),
324
output=output,
325
+ tune_max_num_tokens=min(
326
+ fi_moe_largest_bucket(self.moe_config), self._get_chunk_size()
327
+ ),
328
)
329
330
def apply(
@@ -479,4 +485,5 @@ class TrtLlmNvFp4ExpertsMonolithic(
485
routing_method_type=self.routing_method_type,
486
do_finalize=True,
487
activation_type=activation_to_flashinfer_int(activation),
488
+ tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
489
)[0]
vllm/model_executor/layers/fused_moe/utils.py
+21
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import functools
4
from math import prod
5
+from typing import TYPE_CHECKING
6
7
import torch
8
import torch.nn.functional as F
@@ -33,6 +34,9 @@ from vllm.platforms import current_platform
34
from vllm.triton_utils import tl, triton
35
from vllm.utils.math_utils import cdiv
36
37
+if TYPE_CHECKING:
38
+ from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
39
+
40
41
@triton.jit
42
def _count_expert_num_tokens(
@@ -403,6 +407,23 @@ def _pack_topk_ids_weights_kernel(
407
tl.store(output_ptr + offsets, packed, mask=mask)
408
409
410
+def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int:
411
+ """Estimate FlashInfer's MoE autotuning maximum token count.
412
+
413
+ All DP ranks may contribute `max_num_tokens` to one invocation.
414
+ Keep FlashInfer's default moe `tune_max_num_tokens=8192`
415
+ floor to avoid over-underestimation.
416
+ DeepEP, SP, or PCP may make this underestimate, however overestimation
417
+ may be dangerous, increasing tuning- cost and memory use.
418
+
419
+ NOTE: The DP factor applies even when EP is disabled:
420
+ > Without `--enable-expert-parallel`, MoE layers would use tensor parallelism.
421
+
422
+ For a detailed explanation, see: `docs/serving/data_parallel_deployment.md`
423
+ """
424
+ return max(moe_config.max_num_tokens * moe_config.dp_size, 8192)
425
+
426
+
427
def trtllm_moe_pack_topk_ids_weights(
428
topk_ids: torch.Tensor,
429
topk_weights: torch.Tensor,