diff --git a/tests/kernels/moe/test_flashinfer_moe.py b/tests/kernels/moe/test_flashinfer_moe.py index ece11c371c..a98844d5f8 100644 --- a/tests/kernels/moe/test_flashinfer_moe.py +++ b/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 test_flashinfer_swigluoai_params_are_forwarded(activation, monkeypatch): 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), diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py index 543968b671..1614e41b33 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py @@ -13,7 +13,6 @@ 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 fi_moe_largest_bucket from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_type, ) @@ -389,7 +388,6 @@ class FlashInferExperts(mk.FusedMoEExpertsModular): 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: 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 555fc59360..6b29a6f69d 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 @@ -220,7 +220,6 @@ class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular) 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) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py index 2ecda740c0..1a947756ce 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py @@ -39,7 +39,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, trtllm_moe_pack_topk_ids_weights, ) from vllm.platforms import current_platform @@ -550,7 +549,6 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase): 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] diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py index 12b0dce12e..a9e725f296 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -11,7 +11,6 @@ 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.quant_utils import ( QuantKey, kInt4Static32, @@ -170,7 +169,6 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): 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] diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py index af15b206f0..1b5320615f 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py +++ b/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)):