Revert "[Bugfix][MoE] Tune FlashInfer experts to scheduler token limit" (#52989)

tests/kernels/moe/test_flashinfer_moe.py

@@ -72,7 +72,6 @@ def test_flashinfer_swigluoai_params_are_forwarded(activation, monkeypatch):
         moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
         in_dtype=torch.bfloat16,
         routing_method=RoutingMethodType.TopK,
-        max_num_tokens=16_384,
     )
     quant_config = FusedMoEQuantConfig.make(
         gemm1_alpha=1.702,
@@ -111,7 +110,6 @@ def fake_flashinfer_cutlass_fused_moe(**kwargs):
 
     assert experts._supports_activation(activation)
     assert call_args["activation_type"] == ActivationType.Swiglu
-    assert call_args["tune_max_num_tokens"] == 16_384
     for name, value in (
         ("swiglu_alpha", 1.702),
         ("swiglu_beta", 1.0),

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

@@ -13,7 +13,6 @@
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
-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_type,
 )
@@ -389,7 +388,6 @@ def apply(
             use_deepseek_fp8_block_scale=self.use_deepseek_fp8_block_scale,
             use_mxfp8_act_scaling=use_mxfp8_act_scaling,
             use_w4_group_scaling=use_w4_group_scaling,
-            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
         )
 
     def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:

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

@@ -219,7 +219,6 @@ def apply(
             weight_layout=WeightLayout.BlockMajorK,
             do_finalize=True,
             activation_type=activation_to_flashinfer_int(activation),
-            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
         )
         # FlashInfer's BF16 routed wrapper does not expose an output= argument.
         output.copy_(result[0] if isinstance(result, list) else result)

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

@@ -39,7 +39,6 @@
     TopKWeightAndReduceNoOP,
 )
 from vllm.model_executor.layers.fused_moe.utils import (
-    fi_moe_largest_bucket,
     trtllm_moe_pack_topk_ids_weights,
 )
 from vllm.platforms import current_platform
@@ -550,7 +549,6 @@ def invoke_routed_moe(
             routing_method_type=self.routing_method_type,
             do_finalize=do_finalize,
             output=output if do_finalize else None,
-            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
         )
         if not do_finalize:
             # [gemm2_output, expert_weights, expanded_idx, gemm1_activation]

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

@@ -11,7 +11,6 @@
     FusedMoEQuantConfig,
     RoutingMethodType,
 )
-from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
     kInt4Static32,
@@ -170,7 +169,6 @@ def apply(
             e_score_correction_bias=e_score_correction_bias,
             routing_method_type=self.routing_method,
             routing_replay_out=routing_replay_out,
-            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
         )
         self._maybe_dispatch_routing_replay(
             routing_replay_out, num_tokens=hidden_states.shape[0]

vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py

@@ -191,7 +191,6 @@ def flashinfer_trtllm_mxint4_moe(
     e_score_correction_bias: torch.Tensor | None = None,
     routing_method_type: int | None = None,
     routing_replay_out: torch.Tensor | None = None,
-    tune_max_num_tokens: int = 8192,
 ) -> torch.Tensor:
     """
     Apply FlashInfer TensorRT-LLM MxInt4 MoE kernel.
@@ -212,7 +211,6 @@ def flashinfer_trtllm_mxint4_moe(
         topk_group: Top-k within groups (default: None -> 0)
         e_score_correction_bias: Optional routing bias. dtype: bfloat16
         routing_method_type: FlashInfer RoutingMethodType enum value
-        tune_max_num_tokens: Maximum token count covered by autotuning.
 
     Returns:
         Output tensor from MoE layer. dtype: same as x (bfloat16)
@@ -264,7 +262,7 @@ def flashinfer_trtllm_mxint4_moe(
         enable_pdl=None,
         do_finalize=True,
         output=None,
-        tune_max_num_tokens=tune_max_num_tokens,
+        tune_max_num_tokens=8192,
         routing_replay_out=routing_replay_out,
     )
     if isinstance(out, (tuple, list)):