diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index d65f2a4ad1..fe89b09d09 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/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 diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 06f04b75d9..d6544333be 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/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,