From 4e5ffda19d897f3c34455167dc0503620f1b3e2a Mon Sep 17 00:00:00 2001 From: "Jiahan Chang (Cyrus)" <173873397+jiahanc@users.noreply.github.com> Date: Thu, 17 Sep 2026 15:38:58 +0800 Subject: [PATCH] [Bugfix] Isolate supplemental FlashInfer BF16 autotuning (#57285) Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: OpenAI Codex --- vllm/model_executor/warmup/kernel_warmup.py | 43 +++++++++++++++++++-- vllm/utils/flashinfer.py | 43 +++++++++++++++++---- 2 files changed, 75 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index bf2c1ad0b6..81eca7ac50 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/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) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 7647179134..abc4d75500 100644 --- a/vllm/utils/flashinfer.py +++ b/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",