From 0d0504b54c73119ed643c80b4ed56ef3cf80e209 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 28 Jul 2026 07:58:02 -0700 Subject: [PATCH] [Core] Warm up runner-owned Triton kernels before the first request (#49903) --- tests/v1/worker/test_kv_block_zeroer.py | 66 +++++++++- vllm/model_executor/warmup/kernel_warmup.py | 16 ++- .../warmup/qwen_triton_warmup.py | 114 ----------------- .../warmup/v1_block_table_warmup.py | 42 ++---- vllm/v1/worker/gpu/warmup.py | 121 ++++++++++++------ vllm/v1/worker/mamba_utils.py | 6 +- vllm/v1/worker/utils.py | 7 +- 7 files changed, 183 insertions(+), 189 deletions(-) diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index b212e3ae17..17aa1bf38d 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -4,7 +4,7 @@ import pytest import torch -from vllm.v1.worker.utils import KVBlockZeroer +from vllm.v1.worker.utils import KVBlockZeroer, _zero_kv_blocks_kernel @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -86,3 +86,67 @@ def test_non_uniform_page_sizes(): assert torch.all(storage[1] == 0) assert torch.all(storage[2] == 0) assert torch.all(storage[3] == 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_warmup_compiles_every_n_blocks_specialization(): + """After warmup, no launch should trigger a first-request JIT compile. + + ``n_blocks`` is ``do_not_specialize``, so a single warmup launch must + cover every block count. + """ + device = torch.device("cuda") + num_blocks = 64 + page_size_el = 4 + storage = torch.ones((num_blocks, page_size_el), dtype=torch.int32, device=device) + + zeroer = KVBlockZeroer.__new__(KVBlockZeroer) + zeroer.device = device + zeroer._meta = ( + torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device), + torch.tensor([page_size_el], dtype=torch.int64, device=device), + 1, # max_chunks + page_size_el, # blk_size + 1, # n_segs + ) + + def compiled_variants() -> set: + return { + key + for caches in _zero_kv_blocks_kernel.device_caches.values() + for key in caches[0] + } + + zeroer.warmup(num_blocks) + torch.accelerator.synchronize() + warmed = compiled_variants() + assert warmed + + for n_blocks in (1, 2, 3, 16, 32): + zeroer.zero_block_ids(list(range(n_blocks))) + torch.accelerator.synchronize() + + assert compiled_variants() == warmed + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_warmup_respects_available_block_count(): + """An empty KV cache must not be warmed with out-of-range block IDs.""" + device = torch.device("cuda") + page_size_el = 4 + storage = torch.ones((1, page_size_el), dtype=torch.int32, device=device) + + zeroer = KVBlockZeroer.__new__(KVBlockZeroer) + zeroer.device = device + zeroer._meta = ( + torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device), + torch.tensor([page_size_el], dtype=torch.int64, device=device), + 1, + page_size_el, + 1, + ) + + zeroer.warmup(0) + torch.accelerator.synchronize() + + assert torch.all(storage == 1) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index b2c989b295..afe47f39ce 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -98,12 +98,16 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): minimax_m3_msa_warmup, ) - # Pooling models do not use the generation slot-mapping path. - if not worker.use_v2_model_runner and not worker.model_runner.is_pooling_model: - warm_v1_block_table_kernels( - getattr(worker.model_runner, "device", torch.device("cuda")), - worker.scheduler_config.max_num_batched_tokens, - ) + if not worker.use_v2_model_runner: + # Pooling models do not use the generation slot-mapping path. + if not worker.model_runner.is_pooling_model: + warm_v1_block_table_kernels(worker.model_runner) + # The KV-block zeroing kernel is driven by the scheduler's + # `new_block_ids_to_zero`, so no dummy run ever reaches it. + zeroer = getattr(worker.model_runner, "_kv_block_zeroer", None) + if zeroer is not None: + zeroer.warmup(worker.model_runner.kv_cache_config.num_blocks) + qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config) # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py index 8cbfa539b7..06a016a5a0 100644 --- a/vllm/model_executor/warmup/qwen_triton_warmup.py +++ b/vllm/model_executor/warmup/qwen_triton_warmup.py @@ -24,24 +24,10 @@ _QWEN_MODEL_TYPES = frozenset( } ) -_ZERO_KV_N_BLOCKS = (1, 2) - -_SLOT_MAPPING_KV_BLOCK_SIZE = 16 -_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE = 1 -_SLOT_MAPPING_BLOCK_TABLE_STRIDES = (1, 3) - # Covers L=1 constexpr, non-divisible runtime L, and divisible runtime L. _FLA_POST_CONV_WARMUP_LENGTHS = (1, 2, 16) -@dataclass(frozen=True) -class _ZeroKvWarmupConfig: - seg_page_sizes: torch.Tensor - max_chunks: int - block_size: int - n_segs: int - - @dataclass(frozen=True) class _QwenGDNWarmupConfig: h: int @@ -148,96 +134,6 @@ def _qwen_gdn_warmup_config( return None -def _get_kv_block_zeroer(runner: object) -> object | None: - zeroer = getattr(runner, "kv_block_zeroer", None) - if zeroer is None: - zeroer = getattr(runner, "_kv_block_zeroer", None) - return zeroer - - -def _zero_kv_warmup_config(runner: object) -> _ZeroKvWarmupConfig | None: - zeroer = _get_kv_block_zeroer(runner) - meta = getattr(zeroer, "_meta", None) - if meta is None: - return None - - _, seg_page_sizes, max_chunks, block_size, n_segs = meta - return _ZeroKvWarmupConfig( - seg_page_sizes=seg_page_sizes, - max_chunks=int(max_chunks), - block_size=int(block_size), - n_segs=int(n_segs), - ) - - -def _warm_zero_kv_blocks_with_runner_zeroer(runner: object) -> bool: - zeroer = _get_kv_block_zeroer(runner) - zero_block_ids = getattr(zeroer, "zero_block_ids", None) - if not callable(zero_block_ids): - return False - - for n_blocks in _ZERO_KV_N_BLOCKS: - zero_block_ids(list(range(n_blocks))) - return True - - -def _warm_zero_kv_blocks_kernel( - device: torch.device, config: _ZeroKvWarmupConfig -) -> None: - from vllm.v1.worker.utils import _zero_kv_blocks_kernel - - max_n_blocks = max(_ZERO_KV_N_BLOCKS) - max_page_size = int(config.seg_page_sizes.max().item()) - scratch = torch.empty( - max_n_blocks * max_page_size, - dtype=torch.int32, - device=device, - ) - seg_addrs = torch.tensor( - [scratch.data_ptr()] * config.n_segs, - dtype=torch.uint64, - device=device, - ) - - for n_blocks in _ZERO_KV_N_BLOCKS: - block_ids = torch.arange(n_blocks, dtype=torch.int64, device=device) - grid = (n_blocks * config.n_segs * config.max_chunks,) - _zero_kv_blocks_kernel[grid]( - seg_addrs, - config.seg_page_sizes, - block_ids, - n_blocks, - N_SEGS=config.n_segs, - MAX_CHUNKS=config.max_chunks, - BLOCK_SIZE=config.block_size, - ) - - -def _warm_compute_slot_mapping_kernel(device: torch.device) -> None: - from vllm.v1.worker.block_table import BlockTable - - # num_tokens/max_num_tokens are do_not_specialize; keep the launch tiny. - num_tokens = 1 - query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) - positions = torch.arange(num_tokens, dtype=torch.int64, device=device) - - for block_table_stride in _SLOT_MAPPING_BLOCK_TABLE_STRIDES: - # Use BlockTable so the JIT key matches the production slot-mapping call. - block_table = BlockTable( - block_size=_SLOT_MAPPING_KV_BLOCK_SIZE, - max_num_reqs=1, - max_num_blocks_per_req=block_table_stride, - max_num_batched_tokens=num_tokens, - pin_memory=False, - device=device, - kernel_block_size=_SLOT_MAPPING_KV_BLOCK_SIZE, - cp_kv_cache_interleave_size=_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE, - ) - block_table.add_row(list(range(block_table_stride)), 0) - block_table.commit_block_table(num_reqs=1) - block_table.compute_slot_mapping(1, query_start_loc, positions) - - def _warm_causal_conv1d_fwd_kernel( device: torch.device, config: _QwenGDNWarmupConfig ) -> None: @@ -373,16 +269,6 @@ def qwen_triton_warmup( device = getattr(runner, "device", torch.device("cuda")) logger.info("Warming up Qwen Triton kernels for model_type=%s.", model_type) - zero_config = _zero_kv_warmup_config(runner) - warmed_zeroer = _warm_zero_kv_blocks_with_runner_zeroer(runner) - if zero_config is not None: - _warm_zero_kv_blocks_kernel(device, zero_config) - elif not warmed_zeroer: - logger.info("Skipping Qwen zero-kv warmup: no KVBlockZeroer metadata.") - - _warm_compute_slot_mapping_kernel(device) - _synchronize_device(device) - compilation_config = getattr(runner, "compilation_config", None) static_forward_context = getattr(compilation_config, "static_forward_context", None) gdn_config = _qwen_gdn_warmup_config(static_forward_context) diff --git a/vllm/model_executor/warmup/v1_block_table_warmup.py b/vllm/model_executor/warmup/v1_block_table_warmup.py index 8d2328432e..d49e1ba7cc 100644 --- a/vllm/model_executor/warmup/v1_block_table_warmup.py +++ b/vllm/model_executor/warmup/v1_block_table_warmup.py @@ -2,42 +2,28 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Warm up v1 block-table Triton kernels.""" +from typing import TYPE_CHECKING + import torch +if TYPE_CHECKING: + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + _SLOT_MAPPING_WARMUP_TOKENS = 8 -_SLOT_MAPPING_WARMUP_BLOCK_SIZES = (3, 16) -_SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE = 1 -def warm_v1_block_table_kernels( - device: torch.device, - max_tokens: int, -) -> None: - from vllm.v1.worker.block_table import BlockTable +def warm_v1_block_table_kernels(runner: "GPUModelRunner") -> None: + """JIT-compile ``_compute_slot_mapping_kernel`` for the real block tables.""" - num_tokens = max(0, min(_SLOT_MAPPING_WARMUP_TOKENS, max_tokens)) + device = runner.device + block_table = runner.input_batch.block_table + num_tokens = min( + _SLOT_MAPPING_WARMUP_TOKENS, + runner.scheduler_config.max_num_batched_tokens, + ) if num_tokens <= 0: return query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) positions = torch.arange(num_tokens, dtype=torch.int64, device=device) - for block_size in _SLOT_MAPPING_WARMUP_BLOCK_SIZES: - max_num_blocks_per_req = max( - 1, (max(num_tokens, max_tokens) + block_size - 1) // block_size - ) - max_num_blocks_per_req = ((max_num_blocks_per_req + 15) // 16) * 16 - block_table = BlockTable( - block_size=block_size, - max_num_reqs=1, - max_num_blocks_per_req=max_num_blocks_per_req, - max_num_batched_tokens=max(num_tokens, max_tokens), - pin_memory=False, - device=device, - kernel_block_size=block_size, - cp_kv_cache_interleave_size=( - _SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE - ), - ) - block_table.add_row(list(range(max_num_blocks_per_req)), 0) - block_table.commit_block_table(1) - block_table.compute_slot_mapping(1, query_start_loc, positions) + block_table.compute_slot_mapping(1, query_start_loc, positions) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 785188dc3c..260ca7a5f1 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -157,13 +157,10 @@ def warmup_kernels( worker_execute_model: Callable[[SchedulerOutput], Any], worker_sample_tokens: Callable[[GrammarOutput | None], Any], ) -> None: - """Run two execute_model + sample_tokens iterations to JIT compile - triton kernels. We must call the provided worker's execute_model for - pipeline parallel coordination. + """Run scheduler-realistic prefill and decode steps to JIT compile kernels. - The first iteration simulates a prefill with requests of - decode_query_len + 1 prompt tokens each. The second iteration simulates - a decode step with all requests generating decode_query_len tokens. + We must call the provided worker's execute_model for pipeline parallel + coordination. """ num_spec_steps = model_runner.num_speculative_steps decode_query_len = model_runner.decode_query_len @@ -172,8 +169,13 @@ def warmup_kernels( # a uniform decode batch. prompt_len = decode_query_len + 1 prompt_token_ids = list(range(prompt_len)) - # After prefill, decode generates decode_query_len tokens. - decode_len = prompt_len + decode_query_len + # Upper bound on the decode steps built in `decode_steps` below. + num_decode_steps = 1 + if not model_runner.is_pooling_model: + num_decode_steps = 5 if num_spec_steps > 0 else 3 + # Size the block allocation for the worst case: every request advancing + # decode_query_len tokens on every decode step. + decode_len = prompt_len + num_decode_steps * decode_query_len kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups num_kv_cache_groups = len(kv_cache_groups) @@ -208,9 +210,6 @@ def warmup_kernels( kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups] prefill_block_counts = [_warmup_block_count(prompt_len, s) for s in kv_cache_specs] decode_block_counts = [_warmup_block_count(decode_len, s) for s in kv_cache_specs] - decode_block_deltas = [ - d - p for d, p in zip(decode_block_counts, prefill_block_counts) - ] max_blocks_per_req = sum(decode_block_counts) num_reqs = min( @@ -243,6 +242,11 @@ def warmup_kernels( nonlocal next_block_id return list(range(next_block_id, next_block_id := next_block_id + num_blocks)) + # The KV-block zeroing kernel is driven by the scheduler's + # new_block_ids_to_zero, so none of the steps below reach it. + if model_runner.kv_block_zeroer is not None: + model_runner.kv_block_zeroer.warmup(model_runner.kv_cache_config.num_blocks) + # Step 1: Prefill all requests with 1 + decode_query_len prompt tokens each. new_reqs = [ NewRequestData.from_request( @@ -287,33 +291,78 @@ def warmup_kernels( worker_sample_tokens(grammar_output) - # Step 2: Decode all requests with decode_query_len tokens each. - cached_req_data = CachedRequestData.make_empty() - cached_req_data.req_ids = list(req_ids) - cached_req_data.num_computed_tokens = [prompt_len] * num_reqs - cached_req_data.num_output_tokens = [1] * num_reqs - new_block = any(decode_block_deltas) - cached_req_data.new_block_ids = [ - tuple(_alloc_blocks(n) for n in decode_block_deltas) if new_block else None - for _ in range(num_reqs) + # Per-request state carried across the decode steps. + req_computed = [prompt_len] * num_reqs + req_blocks = [list(prefill_block_counts) for _ in range(num_reqs)] + + def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: + """Decode `indices`, spec-decoding the ones flagged in `spec_flags`.""" + cached_req_data = CachedRequestData.make_empty() + cached_req_data.req_ids = [req_ids[i] for i in indices] + cached_req_data.num_computed_tokens = [req_computed[i] for i in indices] + cached_req_data.num_output_tokens = [1] * len(indices) + cached_req_data.new_block_ids = [] + + step_num_scheduled_tokens: dict[str, int] = {} + step_spec_tokens: dict[str, list[int]] = {} + for i, use_spec in zip(indices, spec_flags): + num_tokens = decode_query_len if use_spec else 1 + after = req_computed[i] + num_tokens + deltas = [ + _warmup_block_count(after, spec) - held + for spec, held in zip(kv_cache_specs, req_blocks[i]) + ] + cached_req_data.new_block_ids.append( + tuple(_alloc_blocks(n) for n in deltas) if any(deltas) else None + ) + req_blocks[i] = [ + held + delta for held, delta in zip(req_blocks[i], deltas) + ] + step_num_scheduled_tokens[req_ids[i]] = num_tokens + if use_spec: + step_spec_tokens[req_ids[i]] = [0] * num_spec_steps + + decode_output = SchedulerOutput.make_empty() + decode_output.scheduled_cached_reqs = cached_req_data + decode_output.num_scheduled_tokens = step_num_scheduled_tokens + decode_output.scheduled_spec_decode_tokens = step_spec_tokens + decode_output.total_num_scheduled_tokens = sum( + step_num_scheduled_tokens.values() + ) + decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups + + worker_execute_model(decode_output) + worker_sample_tokens(None) + + for i, use_spec in zip(indices, spec_flags): + req_computed[i] += decode_query_len if use_spec else 1 + + all_indices = list(range(num_reqs)) + use_spec_decode = num_spec_steps > 0 + + # Decode steps to warm, as (request indices, per-request spec flag). + # Under spec decoding the scheduler drops requests the drafter proposed + # nothing for, so warm each batch shape with and without draft tokens. + decode_steps: list[tuple[list[int], list[bool]]] = [ + (all_indices, [use_spec_decode] * num_reqs), ] + if num_reqs >= 2: + # Mixed spec / non-spec: GDN and KDA reclassify the non-spec decode + # as a prefill and split the batch into spec/non-spec token indices. + decode_steps.append(([0, 1], [use_spec_decode, False])) + if use_spec_decode: + # Exercise the model paths that split a batch by whether each + # request received draft tokens. + decode_steps.append(([0, 1], [False, False])) + if num_reqs > 1: + decode_steps.append(([0], [use_spec_decode])) + if use_spec_decode: + decode_steps.append(([0], [False])) + elif use_spec_decode: + decode_steps.append(([0], [False])) - decode_output = SchedulerOutput.make_empty() - decode_output.scheduled_cached_reqs = cached_req_data - decode_output.num_scheduled_tokens = { - req_id: decode_query_len for req_id in req_ids - } - if num_spec_steps > 0: - decode_output.scheduled_spec_decode_tokens = { - req_id: [0] * num_spec_steps for req_id in req_ids - } - decode_output.total_num_scheduled_tokens = sum( - decode_output.num_scheduled_tokens.values() - ) - decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups - - worker_execute_model(decode_output) - worker_sample_tokens(None) + for step_indices, step_spec_flags in decode_steps: + _run_decode_step(step_indices, step_spec_flags) # Clean up - process finish_req_ids. cleanup_output = SchedulerOutput.make_empty() diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 7f611a2eef..36b72efe35 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -150,7 +150,7 @@ def _copy_mamba_state_block( tl.store(tail_dst + tail_off, tail_data, mask=tail_mask) -@triton.jit +@triton.jit(do_not_specialize=["num_reqs"]) def postprocess_mamba_fused_kernel( # Decision inputs (per-request) num_accepted_tokens_ptr, @@ -280,7 +280,7 @@ def postprocess_mamba_fused_kernel( ) -@triton.jit +@triton.jit(do_not_specialize=["num_reqs"]) def preprocess_mamba_align_fused_kernel( idx_mapping_ptr, state_idx_ptr, @@ -326,7 +326,7 @@ def preprocess_mamba_align_fused_kernel( tl.store(num_accepted_tokens_ptr + req_indices, 1, mask=mask & should_reset) -@triton.jit +@triton.jit(do_not_specialize=["num_reqs"]) def precopy_mamba_align_fused_kernel( # Per-request-slot inputs (indexed by req_idx via idx_mapping), produced by # the V2 fused align preprocess kernel for the current step: diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index d00974046f..4afb941c12 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -40,7 +40,7 @@ from vllm.v1.kv_cache_interface import ( logger = init_logger(__name__) -@triton.jit +@triton.jit(do_not_specialize=["n_blocks"]) def _zero_kv_blocks_kernel( seg_addrs_ptr, seg_page_sizes_ptr, @@ -206,6 +206,11 @@ class KVBlockZeroer: BLOCK_SIZE=blk_size, ) + def warmup(self, num_kv_blocks: int) -> None: + """JIT-compile the zeroing kernel before the first real request.""" + if num_kv_blocks > 0: + self.zero_block_ids([0]) + @dataclass class AttentionGroup: