diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py new file mode 100644 index 0000000000..a8ecd1276e --- /dev/null +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock, call, patch + +import pytest + +from vllm.model_executor.warmup.kernel_warmup import ( + _flashinfer_autotune_token_counts, + _run_flashinfer_autotune_dummy_runs, +) + +pytestmark = pytest.mark.cpu_test + + +class _FakeMoERunner: + """Minimal MoERunner stand-in for token-count discovery tests.""" + + moe_config: Any + + +def _make_moe(*, max_deferred_tokens: int = 128, enabled: bool = True): + """Create a fake MoE layer with a deferred-finalize token limit.""" + moe = _FakeMoERunner() + moe.moe_config = SimpleNamespace( + use_deferred_moe_finalize=enabled, + defer_moe_finalize_max_num_tokens=max_deferred_tokens, + ) + return moe + + +def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto"): + """Create a runner carrying only the state used by FlashInfer warmup.""" + return SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=max_tokens), + vllm_config=SimpleNamespace( + kernel_config=SimpleNamespace(linear_backend=linear_backend) + ), + get_model=Mock( + return_value=SimpleNamespace(modules=Mock(return_value=modules)) + ), + _dummy_run=Mock(), + ) + + +def test_flashinfer_autotune_token_counts_include_deferred_moe_limits(): + runner = _make_runner( + [ + object(), + _make_moe(max_deferred_tokens=128), + _make_moe(max_deferred_tokens=128), + _make_moe(max_deferred_tokens=64, enabled=False), + _make_moe(max_deferred_tokens=-1), + ], + linear_backend="flashinfer_cutedsl", + ) + + with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): + token_counts = _flashinfer_autotune_token_counts(runner) + + assert token_counts == (8192, 32, 128) + + +def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated(): + runner = _make_runner( + [_make_moe(max_deferred_tokens=4096)], + max_tokens=32, + linear_backend="flashinfer_cutedsl", + ) + + with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): + token_counts = _flashinfer_autotune_token_counts(runner) + + assert token_counts == (32,) + + +def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(): + runner = _make_runner([]) + max_buckets = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 8192) + deferred_buckets = (1, 2, 4, 8, 16, 32, 64, 128) + + 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", + side_effect=(max_buckets, deferred_buckets), + ) as get_buckets, + patch("vllm.utils.flashinfer.autotune") as autotune, + ): + _run_flashinfer_autotune_dummy_runs(runner) + + 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), + ] + assert runner._dummy_run.call_args_list == [ + call( + num_tokens=8192, + skip_eplb=True, + is_profile=True, + randomize_inputs=True, + ), + call( + num_tokens=128, + skip_eplb=True, + is_profile=True, + randomize_inputs=True, + ), + ] diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index d7a73ef2f6..7158e23456 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -274,26 +274,57 @@ def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS = 32 +def _flashinfer_deferred_moe_token_counts( + runner: "GPUModelRunner", +) -> tuple[int, ...]: + """Return bounded token counts that exercise deferred MoE finalization.""" + from vllm.model_executor.layers.fused_moe import MoERunner + + max_tokens = runner.scheduler_config.max_num_batched_tokens + token_counts: list[int] = [] + for module in runner.get_model().modules(): + if not isinstance(module, MoERunner): + continue + + moe_config = module.moe_config + max_deferred_tokens = moe_config.defer_moe_finalize_max_num_tokens + if moe_config.use_deferred_moe_finalize and max_deferred_tokens > 0: + token_counts.append(min(max_tokens, max_deferred_tokens)) + + return tuple(dict.fromkeys(token_counts)) + + def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> 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" and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS ): - return max_tokens, _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS - return (max_tokens,) + token_counts.append(_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS) + token_counts.extend(_flashinfer_deferred_moe_token_counts(runner)) + return tuple(dict.fromkeys(token_counts)) def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None: + import vllm.utils.flashinfer as fi_utils + for num_tokens in _flashinfer_autotune_token_counts(runner): - logger.info("Running FlashInfer autotune with %d tokens.", num_tokens) - runner._dummy_run( - num_tokens=num_tokens, - skip_eplb=True, - is_profile=True, - randomize_inputs=True, + tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_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): + runner._dummy_run( + num_tokens=num_tokens, + skip_eplb=True, + is_profile=True, + randomize_inputs=True, + ) def flashinfer_autotune(runner: "GPUModelRunner") -> None: diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 8c158dedb7..a3c24a3fb9 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -247,6 +247,9 @@ flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper( flashinfer_b12x_fused_moe = _lazy_import_wrapper( "flashinfer.fused_moe", "b12x_fused_moe" ) +flashinfer_get_hybrid_num_tokens_buckets = _lazy_import_wrapper( + "flashinfer.fused_moe.utils", "get_hybrid_num_tokens_buckets" +) trtllm_fp4_block_scale_moe = _lazy_import_wrapper( "flashinfer", "trtllm_fp4_block_scale_moe" ) @@ -1286,6 +1289,7 @@ __all__ = [ "nvfp4_block_scale_interleave", "flashinfer_cute_dsl_fused_moe_nvfp4", "flashinfer_b12x_fused_moe", + "flashinfer_get_hybrid_num_tokens_buckets", "flashinfer_convert_sf_to_mma_layout", "trtllm_fp4_block_scale_moe", "flashinfer_trtllm_batch_decode_with_kv_cache_mla",