A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 04:00 UTC
[Model Runner V2] Fix warmup for very small kvcache and/or blocksizes (#36176)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Mar 5, 2026 at 14:48 UTC
a73af584fe6d4c1c2781d537c35e3cc85f58480b
2 files changed
+34
-4
vllm/v1/worker/gpu/warmup.py
+30
-4
@@ -5,6 +5,7 @@ import numpy as np
5
import torch
6
7
from vllm import PoolingParams, SamplingParams
8
+from vllm.utils.math_utils import cdiv
9
from vllm.v1.core.sched.output import (
10
CachedRequestData,
11
GrammarOutput,
@@ -26,12 +27,27 @@ def warmup_kernels(model_runner: GPUModelRunner) -> None:
27
"""
28
prompt_token_ids = [0, 1]
29
prompt_len = len(prompt_token_ids)
30
+ decode_len = prompt_len + 1 # After prefill, one decode token is added.
31
+
32
+ kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
33
+ num_kv_cache_groups = len(kv_cache_groups)
34
+
35
+ # Compute per-request block counts for each KV cache group.
36
+ group_block_sizes = [g.kv_cache_spec.block_size for g in kv_cache_groups]
37
+ prefill_block_counts = [cdiv(prompt_len, bs) for bs in group_block_sizes]
38
+ decode_block_counts = [cdiv(decode_len, bs) for bs in group_block_sizes]
39
+ decode_block_deltas = [
40
+ d - p for d, p in zip(decode_block_counts, prefill_block_counts)
41
+ ]
42
+ max_blocks_per_req = sum(decode_block_counts)
43
+
44
num_reqs = min(
45
model_runner.scheduler_config.max_num_seqs,
46
model_runner.scheduler_config.max_num_batched_tokens // prompt_len,
47
+ # Reserve block 0 (null block) and ensure we have enough blocks.
48
+ max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req),
49
)
50
34
- num_kv_cache_groups = len(model_runner.kv_cache_config.kv_cache_groups)
51
req_ids = [f"_warmup_{i}_" for i in range(num_reqs)]
52
53
# SamplingParams exercising all sampling features.
@@ -42,12 +58,18 @@ def warmup_kernels(model_runner: GPUModelRunner) -> None:
58
sampling_params = SamplingParams.for_sampler_warmup()
59
pooling_params = None
60
61
+ # Assign distinct block IDs per request per group. 0 null block, start from 1.
62
+ next_block_id = 1
63
+
64
+ def _alloc_blocks(num_blocks: int) -> list[int]:
65
+ nonlocal next_block_id
66
+ return list(range(next_block_id, next_block_id := next_block_id + num_blocks))
67
+
68
# Step 1: Prefill all requests with 2 prompt tokens each.
69
new_reqs = [
70
NewRequestData.from_request(
71
Request(req_ids[i], prompt_token_ids, sampling_params, pooling_params),
49
- # Each request uses a distinct block per KV cache group.
50
- block_ids=tuple([i] for _ in range(num_kv_cache_groups)),
72
+ block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts),
73
prefill_token_ids=prompt_token_ids,
74
)
75
for i in range(num_reqs)
@@ -84,9 +106,13 @@ def warmup_kernels(model_runner: GPUModelRunner) -> None:
106
# Step 2: Decode all requests with 1 token each.
107
cached_req_data = CachedRequestData.make_empty()
108
cached_req_data.req_ids = list(req_ids)
87
- cached_req_data.new_block_ids = [None] * num_reqs
109
cached_req_data.num_computed_tokens = [prompt_len] * num_reqs
110
cached_req_data.num_output_tokens = [1] * num_reqs
111
+ new_block = any(decode_block_deltas)
112
+ cached_req_data.new_block_ids = [
113
+ tuple(_alloc_blocks(n) for n in decode_block_deltas) if new_block else None
114
+ for _ in range(num_reqs)
115
+ ]
116
117
decode_output = SchedulerOutput.make_empty()
118
decode_output.scheduled_cached_reqs = cached_req_data
vllm/v1/worker/gpu_worker.py
+4
@@ -464,6 +464,10 @@ class Worker(WorkerBase):
464
def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
465
"""Allocate GPU KV cache with the specified kv_cache_config."""
466
467
+ # Update local config with adjusted num blocks after profiling,
468
+ # so that it's available to the warmup stage.
469
+ self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks
470
+
471
# Init kv cache connector here, because it requires
472
# `kv_cache_config`.
473
# NOTE(Kuntai): This need to be done before `initialize_kv_cache`,