[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>

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

@@ -11,6 +11,7 @@ from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
     RoutingMethodType,
 )
+from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     activation_to_flashinfer_int,
 )
@@ -145,4 +146,5 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
             routed_scaling_factor=routed_scaling_factor,
             routing_method_type=self.routing_method_type,
             activation_type=activation_to_flashinfer_int(activation),
+            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
         )

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

@@ -15,7 +15,10 @@ from vllm.model_executor.layers.fused_moe.config import (
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
-from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights
+from vllm.model_executor.layers.fused_moe.utils import (
+    fi_moe_largest_bucket,
+    trtllm_moe_pack_topk_ids_weights,
+)
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     activation_to_flashinfer_int,
 )
@@ -249,6 +252,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
             weight_layout=weight_layout,
             fp8_quantization_type=fp8_quant_type,
             output=output,
+            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
         )
 
 
@@ -419,6 +423,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
             use_shuffled_weight=use_shuffled_weight,
             weight_layout=weight_layout,
             fp8_quantization_type=fp8_quant_type,
+            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
         )
         if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL:
             kwargs["activation_type"] = activation_type
@@ -475,6 +480,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
             use_routing_scales_on_input=apply_router_weight_on_input,
             routing_method_type=self.routing_method_type,
             activation_type=activation_type,
+            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
         )
         return out
 

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

@@ -16,7 +16,10 @@ from vllm.model_executor.layers.fused_moe.config import (
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
-from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights
+from vllm.model_executor.layers.fused_moe.utils import (
+    fi_moe_largest_bucket,
+    trtllm_moe_pack_topk_ids_weights,
+)
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     activation_to_flashinfer_int,
 )
@@ -319,6 +322,9 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
             do_finalize=True,
             activation_type=activation_to_flashinfer_int(activation),
             output=output,
+            tune_max_num_tokens=min(
+                fi_moe_largest_bucket(self.moe_config), self._get_chunk_size()
+            ),
         )
 
     def apply(
@@ -479,4 +485,5 @@ class TrtLlmNvFp4ExpertsMonolithic(
             routing_method_type=self.routing_method_type,
             do_finalize=True,
             activation_type=activation_to_flashinfer_int(activation),
+            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
         )[0]

vllm/model_executor/layers/fused_moe/utils.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import functools
 from math import prod
+from typing import TYPE_CHECKING
 
 import torch
 import torch.nn.functional as F
@@ -33,6 +34,9 @@ from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
+
 
 @triton.jit
 def _count_expert_num_tokens(
@@ -403,6 +407,23 @@ def _pack_topk_ids_weights_kernel(
     tl.store(output_ptr + offsets, packed, mask=mask)
 
 
+def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int:
+    """Estimate FlashInfer's MoE autotuning maximum token count.
+
+    All DP ranks may contribute `max_num_tokens` to one invocation.
+    Keep FlashInfer's default moe `tune_max_num_tokens=8192`
+    floor to avoid over-underestimation.
+    DeepEP, SP, or PCP may make this underestimate, however overestimation
+    may be dangerous, increasing tuning- cost and memory use.
+
+    NOTE: The DP factor applies even when EP is disabled:
+    > Without `--enable-expert-parallel`, MoE layers would use tensor parallelism.
+
+    For a detailed explanation, see: `docs/serving/data_parallel_deployment.md`
+    """
+    return max(moe_config.max_num_tokens * moe_config.dp_size, 8192)
+
+
 def trtllm_moe_pack_topk_ids_weights(
     topk_ids: torch.Tensor,
     topk_weights: torch.Tensor,