[Bugfix] Isolate supplemental FlashInfer BF16 autotuning (#57285)
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Jiahan Chang (Cyrus) committed
Sep 17, 2026 at 15:38 UTC
4e5ffda19d897f3c34455167dc0503620f1b3e2a
2 files changed
+75
-11
vllm/model_executor/warmup/kernel_warmup.py
+40
-3
@@ -330,13 +330,16 @@ def _flashinfer_deferred_moe_token_counts(
330
return tuple(dict.fromkeys(token_counts))
331
332
333
-def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ...]:
333
+def _flashinfer_autotune_token_counts(
334
+ runner: "GPUModelRunner", *, include_bf16: bool = True
335
+) -> tuple[int, ...]:
336
max_tokens = runner.scheduler_config.max_num_batched_tokens
337
# Tune the widest bucket set first so bounded passes reuse its configs.
338
token_counts = [max_tokens]
339
linear_backend = runner.vllm_config.kernel_config.linear_backend
340
if (
339
- linear_backend == "flashinfer_cutedsl"
341
+ include_bf16
342
+ and linear_backend == "flashinfer_cutedsl"
343
and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
344
):
345
token_counts.append(_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS)
@@ -350,7 +353,7 @@ def _run_flashinfer_autotune_dummy_runs(
353
import vllm.utils.flashinfer as fi_utils
354
355
dummy_run_kwargs = {"skip_attn": True} if skip_attn else {}
353
- for num_tokens in _flashinfer_autotune_token_counts(runner):
356
+ for num_tokens in _flashinfer_autotune_token_counts(runner, include_bf16=False):
357
tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens)
358
logger.info(
359
"Running FlashInfer autotune with %d tokens and token buckets %s.",
@@ -367,6 +370,36 @@ def _run_flashinfer_autotune_dummy_runs(
370
)
371
372
373
+def _run_flashinfer_bf16_autotune_dummy_run(
374
+ runner: "GPUModelRunner",
375
+ *,
376
+ skip_ops: set[str] | None = None,
377
+ skip_attn: bool = False,
378
+) -> None:
379
+ import vllm.utils.flashinfer as fi_utils
380
+
381
+ if (
382
+ runner.vllm_config.kernel_config.linear_backend != "flashinfer_cutedsl"
383
+ or runner.scheduler_config.max_num_batched_tokens
384
+ <= _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
385
+ ):
386
+ return
387
+
388
+ num_tokens = _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
389
+ tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens)
390
+ logger.info("Running FlashInfer BF16-only autotune with %d tokens.", num_tokens)
391
+ # Other ops execute from their existing cache or fallback, without the
392
+ # BF16 bucket cap. No full-model autotune context may enclose this pass.
393
+ with fi_utils.autotune_bf16_only(tuning_buckets, skip_ops=skip_ops):
394
+ runner._dummy_run(
395
+ num_tokens=num_tokens,
396
+ skip_eplb=True,
397
+ is_profile=True,
398
+ randomize_inputs=True,
399
+ **({"skip_attn": True} if skip_attn else {}),
400
+ )
401
+
402
+
403
def flashinfer_autotune(runner: "GPUModelRunner") -> None:
404
"""Autotune FlashInfer operations.
405
FlashInfer have many implementations for the same operation,
@@ -435,6 +468,10 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
468
_run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled)
469
replayssm_autotune_warmup(runner)
470
_autotune_kimi_k3_kda_qkvg(runner.get_model())
471
+ with torch.inference_mode():
472
+ _run_flashinfer_bf16_autotune_dummy_run(
473
+ runner, skip_ops=skip_ops, skip_attn=hisparse_enabled
474
+ )
475
finally:
476
set_autotune_process_group(None)
477
vllm/utils/flashinfer.py
+35
-8
@@ -11,7 +11,8 @@ import importlib
11
import importlib.util
12
import os
13
import shutil
14
-from collections.abc import Callable
14
+from collections.abc import Callable, Iterator
15
+from contextvars import ContextVar
16
from typing import Any, NoReturn
17
18
import requests
@@ -26,6 +27,24 @@ from vllm.utils.torch_utils import PIN_MEMORY
27
logger = init_logger(__name__)
28
29
30
+_bf16_autotune_buckets: ContextVar[tuple[int, ...] | None] = ContextVar(
31
+ "flashinfer_bf16_autotune_buckets", default=None
32
+)
33
+
34
+
35
+@contextlib.contextmanager
36
+def autotune_bf16_only(
37
+ tuning_buckets: tuple[int, ...], *, skip_ops: set[str] | None = None
38
+) -> Iterator[None]:
39
+ """Tune BF16 calls with bounded buckets, outside full-model autotuning."""
40
+ token = _bf16_autotune_buckets.set(tuning_buckets)
41
+ try:
42
+ with autotune(tune_mode=False, skip_ops=skip_ops):
43
+ yield
44
+ finally:
45
+ _bf16_autotune_buckets.reset(token)
46
+
47
+
48
def flashinfer_bf16_mm(
49
a: torch.Tensor,
50
b: torch.Tensor,
@@ -35,14 +54,21 @@ def flashinfer_bf16_mm(
54
) -> torch.Tensor:
55
from flashinfer import mm_bf16
56
38
- return mm_bf16(
39
- a,
40
- b,
41
- bias=bias,
42
- pdl=pdl,
43
- out_dtype=torch.bfloat16,
44
- backend=backend,
57
+ tuning_buckets = _bf16_autotune_buckets.get()
58
+ tuning = (
59
+ autotune(tune_mode=True, tuning_buckets=tuning_buckets)
60
+ if tuning_buckets is not None
61
+ else contextlib.nullcontext()
62
)
63
+ with tuning:
64
+ return mm_bf16(
65
+ a,
66
+ b,
67
+ bias=bias,
68
+ pdl=pdl,
69
+ out_dtype=torch.bfloat16,
70
+ backend=backend,
71
+ )
72
73
74
# This is the storage path for the cubins, it can be replaced
@@ -1237,6 +1263,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool:
1263
__all__ = [
1264
"has_flashinfer",
1265
"flashinfer_bf16_mm",
1266
+ "autotune_bf16_only",
1267
"has_flashinfer_bf16_gemm",
1268
"is_flashinfer_bf16_gemm_supported",
1269
"is_flashinfer_cutedsl_bf16_gemm_supported",