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

vllm/model_executor/warmup/kernel_warmup.py

@@ -330,13 +330,16 @@ def _flashinfer_deferred_moe_token_counts(
     return tuple(dict.fromkeys(token_counts))
 
 
-def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ...]:
+def _flashinfer_autotune_token_counts(
+    runner: "GPUModelRunner", *, include_bf16: bool = True
+) -> 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"
+        include_bf16
+        and linear_backend == "flashinfer_cutedsl"
         and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
     ):
         token_counts.append(_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS)
@@ -350,7 +353,7 @@ def _run_flashinfer_autotune_dummy_runs(
     import vllm.utils.flashinfer as fi_utils
 
     dummy_run_kwargs = {"skip_attn": True} if skip_attn else {}
-    for num_tokens in _flashinfer_autotune_token_counts(runner):
+    for num_tokens in _flashinfer_autotune_token_counts(runner, include_bf16=False):
         tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens)
         logger.info(
             "Running FlashInfer autotune with %d tokens and token buckets %s.",
@@ -367,6 +370,36 @@ def _run_flashinfer_autotune_dummy_runs(
             )
 
 
+def _run_flashinfer_bf16_autotune_dummy_run(
+    runner: "GPUModelRunner",
+    *,
+    skip_ops: set[str] | None = None,
+    skip_attn: bool = False,
+) -> None:
+    import vllm.utils.flashinfer as fi_utils
+
+    if (
+        runner.vllm_config.kernel_config.linear_backend != "flashinfer_cutedsl"
+        or runner.scheduler_config.max_num_batched_tokens
+        <= _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
+    ):
+        return
+
+    num_tokens = _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
+    tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens)
+    logger.info("Running FlashInfer BF16-only autotune with %d tokens.", num_tokens)
+    # Other ops execute from their existing cache or fallback, without the
+    # BF16 bucket cap. No full-model autotune context may enclose this pass.
+    with fi_utils.autotune_bf16_only(tuning_buckets, skip_ops=skip_ops):
+        runner._dummy_run(
+            num_tokens=num_tokens,
+            skip_eplb=True,
+            is_profile=True,
+            randomize_inputs=True,
+            **({"skip_attn": True} if skip_attn else {}),
+        )
+
+
 def flashinfer_autotune(runner: "GPUModelRunner") -> None:
     """Autotune FlashInfer operations.
     FlashInfer have many implementations for the same operation,
@@ -435,6 +468,10 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
             _run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled)
             replayssm_autotune_warmup(runner)
             _autotune_kimi_k3_kda_qkvg(runner.get_model())
+        with torch.inference_mode():
+            _run_flashinfer_bf16_autotune_dummy_run(
+                runner, skip_ops=skip_ops, skip_attn=hisparse_enabled
+            )
     finally:
         set_autotune_process_group(None)
 

vllm/utils/flashinfer.py

@@ -11,7 +11,8 @@ import importlib
 import importlib.util
 import os
 import shutil
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
+from contextvars import ContextVar
 from typing import Any, NoReturn
 
 import requests
@@ -26,6 +27,24 @@ from vllm.utils.torch_utils import PIN_MEMORY
 logger = init_logger(__name__)
 
 
+_bf16_autotune_buckets: ContextVar[tuple[int, ...] | None] = ContextVar(
+    "flashinfer_bf16_autotune_buckets", default=None
+)
+
+
+@contextlib.contextmanager
+def autotune_bf16_only(
+    tuning_buckets: tuple[int, ...], *, skip_ops: set[str] | None = None
+) -> Iterator[None]:
+    """Tune BF16 calls with bounded buckets, outside full-model autotuning."""
+    token = _bf16_autotune_buckets.set(tuning_buckets)
+    try:
+        with autotune(tune_mode=False, skip_ops=skip_ops):
+            yield
+    finally:
+        _bf16_autotune_buckets.reset(token)
+
+
 def flashinfer_bf16_mm(
     a: torch.Tensor,
     b: torch.Tensor,
@@ -35,14 +54,21 @@ def flashinfer_bf16_mm(
 ) -> torch.Tensor:
     from flashinfer import mm_bf16
 
-    return mm_bf16(
-        a,
-        b,
-        bias=bias,
-        pdl=pdl,
-        out_dtype=torch.bfloat16,
-        backend=backend,
+    tuning_buckets = _bf16_autotune_buckets.get()
+    tuning = (
+        autotune(tune_mode=True, tuning_buckets=tuning_buckets)
+        if tuning_buckets is not None
+        else contextlib.nullcontext()
     )
+    with tuning:
+        return mm_bf16(
+            a,
+            b,
+            bias=bias,
+            pdl=pdl,
+            out_dtype=torch.bfloat16,
+            backend=backend,
+        )
 
 
 # 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:
 __all__ = [
     "has_flashinfer",
     "flashinfer_bf16_mm",
+    "autotune_bf16_only",
     "has_flashinfer_bf16_gemm",
     "is_flashinfer_bf16_gemm_supported",
     "is_flashinfer_cutedsl_bf16_gemm_supported",