[Bugfix][MoE] Tune FlashInfer experts to scheduler token limit (#52989)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Misha Goin committed
Aug 20, 2026 at 16:26 UTC
bfb6c134997aace3e801c9ae3251728bd5312003
6 files changed
+12
-1
tests/kernels/moe/test_flashinfer_moe.py
+2
@@ -72,6 +72,7 @@ def test_flashinfer_swigluoai_params_are_forwarded(activation, monkeypatch):
72
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
73
in_dtype=torch.bfloat16,
74
routing_method=RoutingMethodType.TopK,
75
+ max_num_tokens=16_384,
76
)
77
quant_config = FusedMoEQuantConfig.make(
78
gemm1_alpha=1.702,
@@ -110,6 +111,7 @@ def test_flashinfer_swigluoai_params_are_forwarded(activation, monkeypatch):
111
112
assert experts._supports_activation(activation)
113
assert call_args["activation_type"] == ActivationType.Swiglu
114
+ assert call_args["tune_max_num_tokens"] == 16_384
115
for name, value in (
116
("swiglu_alpha", 1.702),
117
("swiglu_beta", 1.0),
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
+2
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.fused_moe.config import (
13
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
14
TopKWeightAndReduceNoOP,
15
)
16
+from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket
17
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
18
activation_to_flashinfer_type,
19
)
@@ -388,6 +389,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
389
use_deepseek_fp8_block_scale=self.use_deepseek_fp8_block_scale,
390
use_mxfp8_act_scaling=use_mxfp8_act_scaling,
391
use_w4_group_scaling=use_w4_group_scaling,
392
+ tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
393
)
394
395
def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+1
@@ -219,6 +219,7 @@ class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular)
219
weight_layout=WeightLayout.BlockMajorK,
220
do_finalize=True,
221
activation_type=activation_to_flashinfer_int(activation),
222
+ tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
223
)
224
# FlashInfer's BF16 routed wrapper does not expose an output= argument.
225
output.copy_(result[0] if isinstance(result, list) else result)
vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py
+2
@@ -39,6 +39,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
39
TopKWeightAndReduceNoOP,
40
)
41
from vllm.model_executor.layers.fused_moe.utils import (
42
+ fi_moe_largest_bucket,
43
trtllm_moe_pack_topk_ids_weights,
44
)
45
from vllm.platforms import current_platform
@@ -549,6 +550,7 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase):
550
routing_method_type=self.routing_method_type,
551
do_finalize=do_finalize,
552
output=output if do_finalize else None,
553
+ tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
554
)
555
if not do_finalize:
556
# [gemm2_output, expert_weights, expanded_idx, gemm1_activation]
vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_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.quant_utils import (
16
QuantKey,
17
kInt4Static32,
@@ -169,6 +170,7 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
170
e_score_correction_bias=e_score_correction_bias,
171
routing_method_type=self.routing_method,
172
routing_replay_out=routing_replay_out,
173
+ tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
174
)
175
self._maybe_dispatch_routing_replay(
176
routing_replay_out, num_tokens=hidden_states.shape[0]
vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py
+3
-1
@@ -191,6 +191,7 @@ def flashinfer_trtllm_mxint4_moe(
191
e_score_correction_bias: torch.Tensor | None = None,
192
routing_method_type: int | None = None,
193
routing_replay_out: torch.Tensor | None = None,
194
+ tune_max_num_tokens: int = 8192,
195
) -> torch.Tensor:
196
"""
197
Apply FlashInfer TensorRT-LLM MxInt4 MoE kernel.
@@ -211,6 +212,7 @@ def flashinfer_trtllm_mxint4_moe(
212
topk_group: Top-k within groups (default: None -> 0)
213
e_score_correction_bias: Optional routing bias. dtype: bfloat16
214
routing_method_type: FlashInfer RoutingMethodType enum value
215
+ tune_max_num_tokens: Maximum token count covered by autotuning.
216
217
Returns:
218
Output tensor from MoE layer. dtype: same as x (bfloat16)
@@ -262,7 +264,7 @@ def flashinfer_trtllm_mxint4_moe(
264
enable_pdl=None,
265
do_finalize=True,
266
output=None,
265
- tune_max_num_tokens=8192,
267
+ tune_max_num_tokens=tune_max_num_tokens,
268
routing_replay_out=routing_replay_out,
269
)
270
if isinstance(out, (tuple, list)):