[MoE] Align TRTLLM MXFP4 autotune buckets (#50942)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: OpenAI Codex <noreply@openai.com>
Thien Tran committed
Aug 6, 2026 at 09:17 UTC
9f3169960a37f230f9fe9629b6bc8e1bbf5969eb
1 file changed
+6
-5
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+6
-5
@@ -14,7 +14,10 @@ from vllm.model_executor.layers.fused_moe.config import (
14
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
15
TopKWeightAndReduceNoOP,
16
)
17
-from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights
17
+from vllm.model_executor.layers.fused_moe.utils import (
18
+ fi_moe_largest_bucket,
19
+ trtllm_moe_pack_topk_ids_weights,
20
+)
21
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
22
activation_to_flashinfer_int,
23
has_flashinfer_situ_activation,
@@ -113,8 +116,6 @@ class TrtLlmMxfp4ExpertsBase:
116
)
117
self.gemm1_clamp_limit = None
118
116
- self.max_capture_size = moe_config.max_capture_size
117
-
119
@staticmethod
120
def _supports_current_device() -> bool:
121
p = current_platform
@@ -263,7 +264,7 @@ class TrtLlmMxfp4ExpertsMonolithic(
264
routing_method_type=self.routing_method_type,
265
do_finalize=True,
266
activation_type=self._flashinfer_activation_type(activation),
266
- tune_max_num_tokens=max(self.max_capture_size, 1),
267
+ tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
268
output=output,
269
routing_replay_out=routing_replay_out,
270
)
@@ -382,7 +383,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
383
enable_pdl=True,
384
activation_type=self._flashinfer_activation_type(activation),
385
output=output,
385
- tune_max_num_tokens=max(self.max_capture_size, 1),
386
+ tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
387
)
388
389
def apply(