[Bugfix]Fix FlashInfer warmup crash: autotune(tuning_buckets=...) must round M up, not down (Invalid MXFP8 split-K tactic with spec-decode) (#58165)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

tests/model_executor/test_flashinfer_autotune_warmup.py

@@ -115,8 +115,8 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn):
 
     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),
+        call(tuning_buckets=max_buckets, round_up=True),
+        call(tuning_buckets=deferred_buckets, round_up=True),
     ]
     assert runner._dummy_run.call_args_list == [
         call(
@@ -136,6 +136,33 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn):
     ]
 
 
+def test_flashinfer_autotune_buckets_cover_drafter_tokens():
+    """The drafter's M can exceed the pass size; buckets must include it."""
+    runner = _make_runner([])
+    runner.max_num_reqs = 256
+    runner.max_num_tokens = 8192
+    runner.speculator = SimpleNamespace(num_query_per_req=6)
+
+    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"
+        ) as get_buckets,
+        patch("vllm.utils.flashinfer.autotune"),
+    ):
+        _run_flashinfer_autotune_dummy_runs(runner)
+
+    assert get_buckets.call_args_list == [call(8192), call(128 * 6)]
+    assert [c.kwargs["num_tokens"] for c in runner._dummy_run.call_args_list] == [
+        8192,
+        128,
+    ]
+
+
 class _AutotuneGroup:
     def __init__(self, run, ranks):
         self.run = run

vllm/model_executor/warmup/kernel_warmup.py

@@ -363,13 +363,32 @@ def _run_flashinfer_autotune_dummy_runs(
 
     dummy_run_kwargs = {"skip_attn": True} if skip_attn else {}
     for num_tokens in _flashinfer_autotune_token_counts(runner, include_bf16=False):
-        tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens)
+        # The drafter may use more tokens than this pass in the same dummy run.
+        # Include its max M when building the autotune buckets.
+        max_tuning_tokens = num_tokens
+        speculator = getattr(runner, "speculator", None)
+        if speculator is not None:
+            num_query_per_req = speculator.num_query_per_req
+            num_draft_reqs = min(
+                num_tokens,
+                runner.max_num_reqs,
+                runner.max_num_tokens // num_query_per_req,
+            )
+            max_tuning_tokens = max(
+                max_tuning_tokens, num_draft_reqs * num_query_per_req
+            )
+
+        tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(
+            max_tuning_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):
+        # Round M up to match serving-time bucket selection for non-bucket M.
+        # See https://github.com/flashinfer-ai/flashinfer/issues/5450
+        with fi_utils.autotune(tuning_buckets=tuning_buckets, round_up=True):
             runner._dummy_run(
                 num_tokens=num_tokens,
                 skip_eplb=True,