[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>

tests/model_executor/test_flashinfer_autotune_warmup.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import Mock, call, patch
+
+import pytest
+
+from vllm.model_executor.warmup.kernel_warmup import (
+    _flashinfer_autotune_token_counts,
+    _run_flashinfer_autotune_dummy_runs,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+
+class _FakeMoERunner:
+    """Minimal MoERunner stand-in for token-count discovery tests."""
+
+    moe_config: Any
+
+
+def _make_moe(*, max_deferred_tokens: int = 128, enabled: bool = True):
+    """Create a fake MoE layer with a deferred-finalize token limit."""
+    moe = _FakeMoERunner()
+    moe.moe_config = SimpleNamespace(
+        use_deferred_moe_finalize=enabled,
+        defer_moe_finalize_max_num_tokens=max_deferred_tokens,
+    )
+    return moe
+
+
+def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto"):
+    """Create a runner carrying only the state used by FlashInfer warmup."""
+    return SimpleNamespace(
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=max_tokens),
+        vllm_config=SimpleNamespace(
+            kernel_config=SimpleNamespace(linear_backend=linear_backend)
+        ),
+        get_model=Mock(
+            return_value=SimpleNamespace(modules=Mock(return_value=modules))
+        ),
+        _dummy_run=Mock(),
+    )
+
+
+def test_flashinfer_autotune_token_counts_include_deferred_moe_limits():
+    runner = _make_runner(
+        [
+            object(),
+            _make_moe(max_deferred_tokens=128),
+            _make_moe(max_deferred_tokens=128),
+            _make_moe(max_deferred_tokens=64, enabled=False),
+            _make_moe(max_deferred_tokens=-1),
+        ],
+        linear_backend="flashinfer_cutedsl",
+    )
+
+    with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner):
+        token_counts = _flashinfer_autotune_token_counts(runner)
+
+    assert token_counts == (8192, 32, 128)
+
+
+def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated():
+    runner = _make_runner(
+        [_make_moe(max_deferred_tokens=4096)],
+        max_tokens=32,
+        linear_backend="flashinfer_cutedsl",
+    )
+
+    with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner):
+        token_counts = _flashinfer_autotune_token_counts(runner)
+
+    assert token_counts == (32,)
+
+
+def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run():
+    runner = _make_runner([])
+    max_buckets = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 8192)
+    deferred_buckets = (1, 2, 4, 8, 16, 32, 64, 128)
+
+    with (
+        patch(
+            "vllm.model_executor.warmup.kernel_warmup."
+            "_flashinfer_autotune_token_counts",
+            return_value=(8192, 128),
+        ),
+        patch(
+            "vllm.utils.flashinfer.flashinfer_get_hybrid_num_tokens_buckets",
+            side_effect=(max_buckets, deferred_buckets),
+        ) as get_buckets,
+        patch("vllm.utils.flashinfer.autotune") as autotune,
+    ):
+        _run_flashinfer_autotune_dummy_runs(runner)
+
+    assert get_buckets.call_args_list == [call(8192), call(128)]
+    assert autotune.call_args_list == [
+        call(tuning_buckets=max_buckets),
+        call(tuning_buckets=deferred_buckets),
+    ]
+    assert runner._dummy_run.call_args_list == [
+        call(
+            num_tokens=8192,
+            skip_eplb=True,
+            is_profile=True,
+            randomize_inputs=True,
+        ),
+        call(
+            num_tokens=128,
+            skip_eplb=True,
+            is_profile=True,
+            randomize_inputs=True,
+        ),
+    ]

vllm/model_executor/warmup/kernel_warmup.py

@@ -274,26 +274,57 @@ def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None:
 _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS = 32
 
 
+def _flashinfer_deferred_moe_token_counts(
+    runner: "GPUModelRunner",
+) -> tuple[int, ...]:
+    """Return bounded token counts that exercise deferred MoE finalization."""
+    from vllm.model_executor.layers.fused_moe import MoERunner
+
+    max_tokens = runner.scheduler_config.max_num_batched_tokens
+    token_counts: list[int] = []
+    for module in runner.get_model().modules():
+        if not isinstance(module, MoERunner):
+            continue
+
+        moe_config = module.moe_config
+        max_deferred_tokens = moe_config.defer_moe_finalize_max_num_tokens
+        if moe_config.use_deferred_moe_finalize and max_deferred_tokens > 0:
+            token_counts.append(min(max_tokens, max_deferred_tokens))
+
+    return tuple(dict.fromkeys(token_counts))
+
+
 def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ...]:
     max_tokens = runner.scheduler_config.max_num_batched_tokens
+    # Tune the widest bucket set first so bounded passes reuse its configs.
+    token_counts = [max_tokens]
     linear_backend = runner.vllm_config.kernel_config.linear_backend
     if (
         linear_backend == "flashinfer_cutedsl"
         and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
     ):
-        return max_tokens, _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
-    return (max_tokens,)
+        token_counts.append(_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS)
+    token_counts.extend(_flashinfer_deferred_moe_token_counts(runner))
+    return tuple(dict.fromkeys(token_counts))
 
 
 def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None:
+    import vllm.utils.flashinfer as fi_utils
+
     for num_tokens in _flashinfer_autotune_token_counts(runner):
-        logger.info("Running FlashInfer autotune with %d tokens.", num_tokens)
-        runner._dummy_run(
-            num_tokens=num_tokens,
-            skip_eplb=True,
-            is_profile=True,
-            randomize_inputs=True,
+        tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens)
+        logger.info(
+            "Running FlashInfer autotune with %d tokens and token buckets %s.",
+            num_tokens,
+            tuning_buckets,
         )
+        with fi_utils.autotune(tuning_buckets=tuning_buckets):
+            runner._dummy_run(
+                num_tokens=num_tokens,
+                skip_eplb=True,
+                is_profile=True,
+                randomize_inputs=True,
+            )
 
 
 def flashinfer_autotune(runner: "GPUModelRunner") -> None:

vllm/utils/flashinfer.py

@@ -247,6 +247,9 @@ flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper(
 flashinfer_b12x_fused_moe = _lazy_import_wrapper(
     "flashinfer.fused_moe", "b12x_fused_moe"
 )
+flashinfer_get_hybrid_num_tokens_buckets = _lazy_import_wrapper(
+    "flashinfer.fused_moe.utils", "get_hybrid_num_tokens_buckets"
+)
 trtllm_fp4_block_scale_moe = _lazy_import_wrapper(
     "flashinfer", "trtllm_fp4_block_scale_moe"
 )
@@ -1286,6 +1289,7 @@ __all__ = [
     "nvfp4_block_scale_interleave",
     "flashinfer_cute_dsl_fused_moe_nvfp4",
     "flashinfer_b12x_fused_moe",
+    "flashinfer_get_hybrid_num_tokens_buckets",
     "flashinfer_convert_sf_to_mma_layout",
     "trtllm_fp4_block_scale_moe",
     "flashinfer_trtllm_batch_decode_with_kv_cache_mla",