[Bugfix] Autotune FlashInfer deferred MoE decode kernels before CUDA graph capture (#55377)
Signed-off-by: ming <ming@inferact.ai> Co-authored-by: Kimi Code <kimi-code@moonshot.cn> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
mingg26 committed
Sep 7, 2026 at 21:02 UTC
7fa2c637968e2f7f7f815f3acc9ceb847a8a462d
3 files changed
+159
-8
tests/model_executor/test_flashinfer_autotune_warmup.py
new
+116
@@ -0,0 +1,116 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+from typing import Any
6
+from unittest.mock import Mock, call, patch
7
+
8
+import pytest
9
+
10
+from vllm.model_executor.warmup.kernel_warmup import (
11
+ _flashinfer_autotune_token_counts,
12
+ _run_flashinfer_autotune_dummy_runs,
13
+)
14
+
15
+pytestmark = pytest.mark.cpu_test
16
+
17
+
18
+class _FakeMoERunner:
19
+ """Minimal MoERunner stand-in for token-count discovery tests."""
20
+
21
+ moe_config: Any
22
+
23
+
24
+def _make_moe(*, max_deferred_tokens: int = 128, enabled: bool = True):
25
+ """Create a fake MoE layer with a deferred-finalize token limit."""
26
+ moe = _FakeMoERunner()
27
+ moe.moe_config = SimpleNamespace(
28
+ use_deferred_moe_finalize=enabled,
29
+ defer_moe_finalize_max_num_tokens=max_deferred_tokens,
30
+ )
31
+ return moe
32
+
33
+
34
+def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto"):
35
+ """Create a runner carrying only the state used by FlashInfer warmup."""
36
+ return SimpleNamespace(
37
+ scheduler_config=SimpleNamespace(max_num_batched_tokens=max_tokens),
38
+ vllm_config=SimpleNamespace(
39
+ kernel_config=SimpleNamespace(linear_backend=linear_backend)
40
+ ),
41
+ get_model=Mock(
42
+ return_value=SimpleNamespace(modules=Mock(return_value=modules))
43
+ ),
44
+ _dummy_run=Mock(),
45
+ )
46
+
47
+
48
+def test_flashinfer_autotune_token_counts_include_deferred_moe_limits():
49
+ runner = _make_runner(
50
+ [
51
+ object(),
52
+ _make_moe(max_deferred_tokens=128),
53
+ _make_moe(max_deferred_tokens=128),
54
+ _make_moe(max_deferred_tokens=64, enabled=False),
55
+ _make_moe(max_deferred_tokens=-1),
56
+ ],
57
+ linear_backend="flashinfer_cutedsl",
58
+ )
59
+
60
+ with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner):
61
+ token_counts = _flashinfer_autotune_token_counts(runner)
62
+
63
+ assert token_counts == (8192, 32, 128)
64
+
65
+
66
+def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated():
67
+ runner = _make_runner(
68
+ [_make_moe(max_deferred_tokens=4096)],
69
+ max_tokens=32,
70
+ linear_backend="flashinfer_cutedsl",
71
+ )
72
+
73
+ with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner):
74
+ token_counts = _flashinfer_autotune_token_counts(runner)
75
+
76
+ assert token_counts == (32,)
77
+
78
+
79
+def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run():
80
+ runner = _make_runner([])
81
+ max_buckets = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 8192)
82
+ deferred_buckets = (1, 2, 4, 8, 16, 32, 64, 128)
83
+
84
+ with (
85
+ patch(
86
+ "vllm.model_executor.warmup.kernel_warmup."
87
+ "_flashinfer_autotune_token_counts",
88
+ return_value=(8192, 128),
89
+ ),
90
+ patch(
91
+ "vllm.utils.flashinfer.flashinfer_get_hybrid_num_tokens_buckets",
92
+ side_effect=(max_buckets, deferred_buckets),
93
+ ) as get_buckets,
94
+ patch("vllm.utils.flashinfer.autotune") as autotune,
95
+ ):
96
+ _run_flashinfer_autotune_dummy_runs(runner)
97
+
98
+ assert get_buckets.call_args_list == [call(8192), call(128)]
99
+ assert autotune.call_args_list == [
100
+ call(tuning_buckets=max_buckets),
101
+ call(tuning_buckets=deferred_buckets),
102
+ ]
103
+ assert runner._dummy_run.call_args_list == [
104
+ call(
105
+ num_tokens=8192,
106
+ skip_eplb=True,
107
+ is_profile=True,
108
+ randomize_inputs=True,
109
+ ),
110
+ call(
111
+ num_tokens=128,
112
+ skip_eplb=True,
113
+ is_profile=True,
114
+ randomize_inputs=True,
115
+ ),
116
+ ]
vllm/model_executor/warmup/kernel_warmup.py
+39
-8
@@ -274,26 +274,57 @@ def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None:
274
_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS = 32
275
276
277
+def _flashinfer_deferred_moe_token_counts(
278
+ runner: "GPUModelRunner",
279
+) -> tuple[int, ...]:
280
+ """Return bounded token counts that exercise deferred MoE finalization."""
281
+ from vllm.model_executor.layers.fused_moe import MoERunner
282
+
283
+ max_tokens = runner.scheduler_config.max_num_batched_tokens
284
+ token_counts: list[int] = []
285
+ for module in runner.get_model().modules():
286
+ if not isinstance(module, MoERunner):
287
+ continue
288
+
289
+ moe_config = module.moe_config
290
+ max_deferred_tokens = moe_config.defer_moe_finalize_max_num_tokens
291
+ if moe_config.use_deferred_moe_finalize and max_deferred_tokens > 0:
292
+ token_counts.append(min(max_tokens, max_deferred_tokens))
293
+
294
+ return tuple(dict.fromkeys(token_counts))
295
+
296
+
297
def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ...]:
298
max_tokens = runner.scheduler_config.max_num_batched_tokens
299
+ # Tune the widest bucket set first so bounded passes reuse its configs.
300
+ token_counts = [max_tokens]
301
linear_backend = runner.vllm_config.kernel_config.linear_backend
302
if (
303
linear_backend == "flashinfer_cutedsl"
304
and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
305
):
284
- return max_tokens, _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
285
- return (max_tokens,)
306
+ token_counts.append(_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS)
307
+ token_counts.extend(_flashinfer_deferred_moe_token_counts(runner))
308
+ return tuple(dict.fromkeys(token_counts))
309
310
311
def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None:
312
+ import vllm.utils.flashinfer as fi_utils
313
+
314
for num_tokens in _flashinfer_autotune_token_counts(runner):
290
- logger.info("Running FlashInfer autotune with %d tokens.", num_tokens)
291
- runner._dummy_run(
292
- num_tokens=num_tokens,
293
- skip_eplb=True,
294
- is_profile=True,
295
- randomize_inputs=True,
315
+ tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens)
316
+ logger.info(
317
+ "Running FlashInfer autotune with %d tokens and token buckets %s.",
318
+ num_tokens,
319
+ tuning_buckets,
320
)
321
+ with fi_utils.autotune(tuning_buckets=tuning_buckets):
322
+ runner._dummy_run(
323
+ num_tokens=num_tokens,
324
+ skip_eplb=True,
325
+ is_profile=True,
326
+ randomize_inputs=True,
327
+ )
328
329
330
def flashinfer_autotune(runner: "GPUModelRunner") -> None:
vllm/utils/flashinfer.py
+4
@@ -247,6 +247,9 @@ flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper(
247
flashinfer_b12x_fused_moe = _lazy_import_wrapper(
248
"flashinfer.fused_moe", "b12x_fused_moe"
249
)
250
+flashinfer_get_hybrid_num_tokens_buckets = _lazy_import_wrapper(
251
+ "flashinfer.fused_moe.utils", "get_hybrid_num_tokens_buckets"
252
+)
253
trtllm_fp4_block_scale_moe = _lazy_import_wrapper(
254
"flashinfer", "trtllm_fp4_block_scale_moe"
255
)
@@ -1286,6 +1289,7 @@ __all__ = [
1289
"nvfp4_block_scale_interleave",
1290
"flashinfer_cute_dsl_fused_moe_nvfp4",
1291
"flashinfer_b12x_fused_moe",
1292
+ "flashinfer_get_hybrid_num_tokens_buckets",
1293
"flashinfer_convert_sf_to_mma_layout",
1294
"trtllm_fp4_block_scale_moe",
1295
"flashinfer_trtllm_batch_decode_with_kv_cache_mla",