diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 550d6b5341d..bd1b9ecaa8a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/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), ) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index a7faa5f6e17..a4cce79741e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 518c87ce4df..f046dfeaf26 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/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] diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index fce74346d62..ed512e1ff8f 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/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,