[Bugfix][MRV2] Reserve spec-decode lookahead blocks in V2 warmup (#51438)
Signed-off-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Co-authored-by: Janelle Cai <janelle.cai@modal.com>
Nick Hill committed
Aug 8, 2026 at 16:58 UTC
d608dfabfdba7e1496e5fef7c66d27c131f59432
4 files changed
+449
-37
tests/v1/worker/test_gpu_warmup_blocks.py
new
+361
@@ -0,0 +1,361 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""The V2 warmup must reserve the KV blocks the scheduler would reserve.
4
+
5
+`KVCacheManager.allocate_slots` sizes every allocation for
6
+`num_computed + num_scheduled + num_lookahead_tokens`, because the speculator
7
+writes KV past the token range the target model was scheduled for. The warmup
8
+hand-builds its `SchedulerOutput`s, so it has to reserve the same blocks.
9
+"""
10
+
11
+from types import SimpleNamespace
12
+
13
+import pytest
14
+import torch
15
+
16
+from tests.v1.core.test_prefix_caching import make_kv_cache_manager
17
+from tests.v1.core.utils import create_requests
18
+from vllm.config.speculative import SpeculativeConfig
19
+from vllm.config.vllm import VllmConfig
20
+from vllm.platforms import current_platform
21
+from vllm.utils.math_utils import cdiv
22
+from vllm.v1.kv_cache_interface import (
23
+ FullAttentionSpec,
24
+ KVCacheConfig,
25
+ KVCacheGroupSpec,
26
+ MambaSpec,
27
+)
28
+from vllm.v1.worker.gpu.warmup import (
29
+ _reserved_block_count,
30
+ run_mixed_prefill_decode_warmup,
31
+ warmup_kernels,
32
+)
33
+
34
+BLOCK_SIZE = 16
35
+MAX_MODEL_LEN = 1024
36
+NUM_SPEC_STEPS = 3
37
+
38
+# `warmup_kernels` ends on `torch.accelerator.synchronize()`.
39
+pytestmark = pytest.mark.skipif(
40
+ not current_platform.is_cuda(), reason="warmup synchronizes on the accelerator"
41
+)
42
+
43
+
44
+def _attention_group() -> KVCacheGroupSpec:
45
+ return KVCacheGroupSpec(
46
+ ["layer"],
47
+ FullAttentionSpec(
48
+ block_size=BLOCK_SIZE,
49
+ num_kv_heads=1,
50
+ head_size=1,
51
+ dtype=torch.float32,
52
+ ),
53
+ )
54
+
55
+
56
+def _mamba_group(mamba_cache_mode: str) -> KVCacheGroupSpec:
57
+ # Name carries the mode so several groups can coexist in one config.
58
+ return KVCacheGroupSpec(
59
+ [f"mamba_{mamba_cache_mode}"],
60
+ MambaSpec(
61
+ block_size=BLOCK_SIZE,
62
+ shapes=((1,),),
63
+ dtypes=(torch.float32,),
64
+ mamba_cache_mode=mamba_cache_mode,
65
+ num_speculative_blocks=NUM_SPEC_STEPS,
66
+ ),
67
+ )
68
+
69
+
70
+def _make_runner(
71
+ kv_cache_groups: list[KVCacheGroupSpec],
72
+ num_lookahead_tokens: int,
73
+ num_spec_steps: int = NUM_SPEC_STEPS,
74
+) -> SimpleNamespace:
75
+ """Stub model runner exposing only what the warmup entry points read."""
76
+ return SimpleNamespace(
77
+ num_speculative_steps=num_spec_steps,
78
+ decode_query_len=num_spec_steps + 1,
79
+ is_pooling_model=False,
80
+ is_encoder_only=False,
81
+ is_encoder_decoder=False,
82
+ is_last_pp_rank=True,
83
+ max_num_reqs=4,
84
+ max_model_len=MAX_MODEL_LEN,
85
+ model_config=SimpleNamespace(get_vocab_size=lambda: 64),
86
+ model_state=SimpleNamespace(max_encoder_len=0),
87
+ scheduler_config=SimpleNamespace(max_num_seqs=4, max_num_batched_tokens=2048),
88
+ kv_cache_config=SimpleNamespace(
89
+ kv_cache_groups=kv_cache_groups, num_blocks=1024
90
+ ),
91
+ vllm_config=SimpleNamespace(num_lookahead_tokens=num_lookahead_tokens),
92
+ kv_block_zeroer=None,
93
+ kv_connector=SimpleNamespace(set_disabled=lambda disabled: None),
94
+ )
95
+
96
+
97
+class _StepRecorder:
98
+ """Rebuilds each request's per-group block holdings from the warmup steps."""
99
+
100
+ def __init__(self) -> None:
101
+ # (blocks held per group, num_computed_tokens, num_scheduled_tokens)
102
+ self.steps: list[tuple[list[int], int, int]] = []
103
+ self._held: dict[str, list[int]] = {}
104
+
105
+ def execute_model(self, scheduler_output) -> None:
106
+ for new_req in scheduler_output.scheduled_new_reqs:
107
+ self._held[new_req.req_id] = [len(ids) for ids in new_req.block_ids]
108
+ self._record(new_req.req_id, new_req.num_computed_tokens, scheduler_output)
109
+ cached = scheduler_output.scheduled_cached_reqs
110
+ for i, req_id in enumerate(cached.req_ids):
111
+ new_block_ids = cached.new_block_ids[i]
112
+ if new_block_ids is not None:
113
+ self._held[req_id] = [
114
+ held + len(ids)
115
+ for held, ids in zip(self._held[req_id], new_block_ids)
116
+ ]
117
+ self._record(req_id, cached.num_computed_tokens[i], scheduler_output)
118
+
119
+ def _record(self, req_id: str, num_computed: int, scheduler_output) -> None:
120
+ self.steps.append(
121
+ (
122
+ list(self._held[req_id]),
123
+ num_computed,
124
+ scheduler_output.num_scheduled_tokens[req_id],
125
+ )
126
+ )
127
+
128
+ def sample_tokens(self, grammar_output=None) -> None:
129
+ return None
130
+
131
+
132
+def _assert_covers_lookahead(
133
+ steps: list[tuple[list[int], int, int]], num_lookahead_tokens: int
134
+) -> None:
135
+ assert steps, "warmup ran no steps"
136
+ for num_blocks, num_computed, num_scheduled in steps:
137
+ num_tokens = min(
138
+ num_computed + num_scheduled + num_lookahead_tokens, MAX_MODEL_LEN
139
+ )
140
+ assert num_blocks[0] >= cdiv(num_tokens, BLOCK_SIZE), (
141
+ f"{num_blocks[0]} blocks for {num_computed}+{num_scheduled} tokens "
142
+ f"and {num_lookahead_tokens} lookahead tokens"
143
+ )
144
+
145
+
146
+# 0 covers eagle / MTP / draft models, 1 covers DFlash's extra in-fill query.
147
+@pytest.mark.parametrize("extra_lookahead", [0, 1])
148
+@pytest.mark.parametrize("num_spec_steps", [2, 3, 5, 7])
149
+def test_warmup_kernels_reserves_lookahead_blocks(num_spec_steps, extra_lookahead):
150
+ num_lookahead_tokens = num_spec_steps + extra_lookahead
151
+ recorder = _StepRecorder()
152
+
153
+ warmup_kernels(
154
+ _make_runner([_attention_group()], num_lookahead_tokens, num_spec_steps),
155
+ recorder.execute_model,
156
+ recorder.sample_tokens,
157
+ )
158
+
159
+ _assert_covers_lookahead(recorder.steps, num_lookahead_tokens)
160
+
161
+
162
+def test_mixed_warmup_reserves_lookahead_blocks():
163
+ num_lookahead_tokens = NUM_SPEC_STEPS + 1
164
+ recorder = _StepRecorder()
165
+
166
+ assert run_mixed_prefill_decode_warmup(
167
+ _make_runner([_attention_group()], num_lookahead_tokens),
168
+ worker_execute_model=recorder.execute_model,
169
+ worker_sample_tokens=recorder.sample_tokens,
170
+ num_tokens=128,
171
+ )
172
+
173
+ _assert_covers_lookahead(recorder.steps, num_lookahead_tokens)
174
+
175
+
176
+@pytest.mark.parametrize("mamba_cache_mode", ["none", "all", "align"])
177
+def test_warmup_reserves_mamba_speculative_blocks(mamba_cache_mode):
178
+ """Mamba groups hold the running-state block plus the speculative tail.
179
+
180
+ `MambaManager` reserves `num_speculative_blocks` blocks past the token
181
+ range in every cache mode, and the mamba kernels read all
182
+ `1 + num_speculative_blocks` of those block-table columns (see
183
+ `mamba_get_block_table_tensor`).
184
+ """
185
+ recorder = _StepRecorder()
186
+
187
+ warmup_kernels(
188
+ _make_runner(
189
+ [_attention_group(), _mamba_group(mamba_cache_mode)], NUM_SPEC_STEPS
190
+ ),
191
+ recorder.execute_model,
192
+ recorder.sample_tokens,
193
+ )
194
+
195
+ assert recorder.steps, "warmup ran no steps"
196
+ for num_blocks, num_computed, num_scheduled in recorder.steps:
197
+ # `MambaManager` drops the lookahead tokens in align mode to keep the
198
+ # allocation block-aligned, and keeps them otherwise. Pin the token
199
+ # range and the speculative tail separately, so an implementation that
200
+ # returned a flat `1 + num_speculative_blocks` would fail.
201
+ if mamba_cache_mode == "align":
202
+ # Align mode sizes from the uncapped main-model range.
203
+ lookahead, num_tokens = 0, num_computed + num_scheduled
204
+ else:
205
+ lookahead = NUM_SPEC_STEPS
206
+ num_tokens = min(num_computed + num_scheduled + lookahead, MAX_MODEL_LEN)
207
+ assert num_blocks[1] == cdiv(num_tokens, BLOCK_SIZE) + NUM_SPEC_STEPS, (
208
+ f"{mamba_cache_mode}: {num_blocks[1]} mamba blocks for "
209
+ f"{num_computed}+{num_scheduled} tokens and {lookahead} lookahead"
210
+ )
211
+
212
+
213
+def _hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig:
214
+ """Full attention plus one Mamba group per cache mode, so a single manager
215
+ exercises every branch `_reserved_block_count` has.
216
+
217
+ "none" and "all" reach the same branch of both `_reserved_block_count` and
218
+ `MambaManager`, which tests only for "align". They are still both listed:
219
+ the mode is a spec-level input, and having the real manager confirm the
220
+ prediction for each is what keeps a future divergence between them from
221
+ landing unnoticed.
222
+ """
223
+ return KVCacheConfig(
224
+ num_blocks=num_blocks,
225
+ kv_cache_tensors=[],
226
+ kv_cache_groups=[
227
+ _attention_group(),
228
+ _mamba_group("none"),
229
+ _mamba_group("all"),
230
+ _mamba_group("align"),
231
+ ],
232
+ )
233
+
234
+
235
+def test_reserved_block_count_matches_real_kv_cache_manager():
236
+ """`_reserved_block_count` must predict exactly what the real
237
+ `KVCacheManager.allocate_slots` consumes, for every group, at every step
238
+ of a warmup-realistic trajectory: a prefill followed by decode steps that
239
+ mix spec and non-spec shapes, all at the fixed `num_lookahead_tokens`
240
+ `_warmup_block_counter` binds once per model runner.
241
+
242
+ The existing tests above pin this arithmetic against a hand-built
243
+ `_StepRecorder` fake, which only proves the fake and the production
244
+ function agree with each other. Driving `allocate_slots` itself is the
245
+ only way to catch a real divergence from the allocator warmup has to
246
+ match.
247
+
248
+ Since `num_computed_tokens` only grows and `num_lookahead_tokens` is held
249
+ fixed across the trajectory, the real manager's held-block count is
250
+ non-decreasing in lockstep with the prediction, so equality (not just
251
+ the reservation lower bound) holds at every step -- exactly the shape of
252
+ trajectory `warmup_kernels` itself drives the allocator through.
253
+ """
254
+ num_lookahead_tokens = NUM_SPEC_STEPS + 1
255
+ decode_query_len = NUM_SPEC_STEPS + 1
256
+
257
+ kv_cache_config = _hybrid_kv_cache_config(num_blocks=256)
258
+ specs = [g.kv_cache_spec for g in kv_cache_config.kv_cache_groups]
259
+ manager = make_kv_cache_manager(
260
+ kv_cache_config,
261
+ max_model_len=MAX_MODEL_LEN,
262
+ hash_block_size=BLOCK_SIZE,
263
+ enable_caching=False,
264
+ )
265
+ (request,) = create_requests(
266
+ num_requests=1, num_tokens=decode_query_len + 1, block_size=BLOCK_SIZE
267
+ )
268
+
269
+ def _step(num_new_tokens: int) -> None:
270
+ computed_blocks, num_new_computed, _ = manager.get_computed_blocks(request)
271
+ blocks = manager.allocate_slots(
272
+ request,
273
+ num_new_tokens,
274
+ num_new_computed,
275
+ computed_blocks,
276
+ num_lookahead_tokens=num_lookahead_tokens,
277
+ )
278
+ assert blocks is not None, "block pool exhausted"
279
+ held = manager.get_block_ids(request.request_id)
280
+ num_tokens = request.num_computed_tokens + num_new_tokens
281
+ for spec, group_held in zip(specs, held):
282
+ expected = _reserved_block_count(
283
+ num_tokens,
284
+ spec,
285
+ num_lookahead_tokens=num_lookahead_tokens,
286
+ max_model_len=MAX_MODEL_LEN,
287
+ max_encoder_len=0,
288
+ )
289
+ assert len(group_held) == expected, (
290
+ f"{type(spec).__name__} "
291
+ f"mode={getattr(spec, 'mamba_cache_mode', None)}: real manager "
292
+ f"holds {len(group_held)} blocks for {num_tokens} tokens, "
293
+ f"_reserved_block_count predicts {expected}"
294
+ )
295
+ request.num_computed_tokens = num_tokens
296
+
297
+ # Prefill, a spec decode step, a no-draft decode step, and another spec
298
+ # decode step: the same shapes `warmup_kernels` drives its decode steps
299
+ # through (see `decode_steps` in `warmup.py`).
300
+ _step(decode_query_len + 1)
301
+ _step(decode_query_len)
302
+ _step(1)
303
+ _step(decode_query_len)
304
+
305
+ manager.free(request)
306
+
307
+
308
+@pytest.mark.parametrize(
309
+ ("method", "expected"),
310
+ [
311
+ ("eagle", NUM_SPEC_STEPS),
312
+ ("eagle3", NUM_SPEC_STEPS),
313
+ ("mtp", NUM_SPEC_STEPS),
314
+ ("dspark", NUM_SPEC_STEPS),
315
+ ("draft_model", NUM_SPEC_STEPS),
316
+ # DFlash's in-fill decoding adds a query for the last sampled token.
317
+ ("dflash", NUM_SPEC_STEPS + 1),
318
+ ("ngram", 0),
319
+ ("ngram_gpu", 0),
320
+ ("medusa", 0),
321
+ ("mlp_speculator", 0),
322
+ ("suffix", 0),
323
+ ("extract_hidden_states", 0),
324
+ ],
325
+)
326
+def test_num_lookahead_tokens_per_method(method: str, expected: int):
327
+ """`VllmConfig.num_lookahead_tokens` is the single source of the reservation.
328
+
329
+ Both the scheduler and the warmup read it, so a wrong answer here silently
330
+ changes how many blocks every request holds. Exercises the real property
331
+ and the real `SpeculativeConfig` predicates; the config is built without
332
+ `__post_init__` because the speculative methods otherwise require a draft
333
+ model to be resolvable.
334
+ """
335
+
336
+ class _Config:
337
+ speculative_config: SpeculativeConfig | None = None
338
+ diffusion_config = None
339
+ num_speculative_tokens = VllmConfig.num_speculative_tokens
340
+ num_lookahead_tokens = VllmConfig.num_lookahead_tokens
341
+
342
+ speculative_config = object.__new__(SpeculativeConfig)
343
+ object.__setattr__(speculative_config, "method", method)
344
+ object.__setattr__(speculative_config, "num_speculative_tokens", NUM_SPEC_STEPS)
345
+
346
+ config = _Config()
347
+ config.speculative_config = speculative_config
348
+
349
+ assert config.num_lookahead_tokens == expected
350
+
351
+
352
+def test_num_lookahead_tokens_without_speculation():
353
+ class _Config:
354
+ speculative_config: SpeculativeConfig | None = None
355
+ diffusion_config = None
356
+ num_speculative_tokens = VllmConfig.num_speculative_tokens
357
+ num_lookahead_tokens = VllmConfig.num_lookahead_tokens
358
+
359
+ config = _Config()
360
+
361
+ assert config.num_lookahead_tokens == 0
vllm/config/vllm.py
+27
@@ -592,6 +592,33 @@ class VllmConfig:
592
return self.diffusion_config.canvas_length
593
return 0
594
595
+ @property
596
+ def num_lookahead_tokens(self) -> int:
597
+ """KV slots to reserve past the tokens the target model is scheduled for.
598
+
599
+ The drafter writes KV for positions beyond the target model's query
600
+ range, so every component that reserves blocks must add this margin:
601
+ the scheduler through `allocate_slots`, and the worker warmup, which
602
+ builds its own `SchedulerOutput`s. Consumers must read this property
603
+ rather than re-deriving their own per-method lookahead, so the
604
+ scheduler and warmup cannot drift apart.
605
+ """
606
+ speculative_config = self.speculative_config
607
+ if speculative_config is None:
608
+ return 0
609
+ if speculative_config.use_dflash():
610
+ # DFlash requires an extra lookahead slot since it uses in-fill-style
611
+ # decoding instead of standard next-token sampling, so it has a query
612
+ # for the last sampled token plus queries for each draft token.
613
+ return self.num_speculative_tokens + 1
614
+ if speculative_config.use_eagle() or speculative_config.uses_draft_model():
615
+ # DSpark (covered by use_eagle) drafts a block of num_speculative_tokens
616
+ # query tokens in which the anchor itself is the first prediction
617
+ # position (no separate bonus query), so it needs exactly
618
+ # num_speculative_tokens lookahead slots.
619
+ return self.num_speculative_tokens
620
+ return 0
621
+
622
@property
623
def use_v2_model_runner(self) -> bool:
624
use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
vllm/v1/core/sched/scheduler.py
+2
-16
@@ -245,7 +245,7 @@ class Scheduler(SchedulerInterface):
245
speculative_config = vllm_config.speculative_config
246
self.use_eagle = False
247
self.num_spec_tokens = vllm_config.num_speculative_tokens
248
- self.num_lookahead_tokens = 0
248
+ self.num_lookahead_tokens = vllm_config.num_lookahead_tokens
249
self.dynamic_sd_lookup: list[int] | None = None
250
if speculative_config is not None:
251
if speculative_config.num_speculative_tokens_per_batch_size:
@@ -254,21 +254,7 @@ class Scheduler(SchedulerInterface):
254
vllm_max_batch_size=self.scheduler_config.max_num_seqs,
255
vllm_num_speculative_tokens=self.num_spec_tokens,
256
)
257
- if speculative_config.use_eagle():
258
- self.use_eagle = True
259
- self.num_lookahead_tokens = self.num_spec_tokens
260
- if speculative_config.uses_draft_model():
261
- self.num_lookahead_tokens = self.num_spec_tokens
262
- if speculative_config.use_dflash():
263
- # DFlash requires an extra lookahead slot since it uses in-fill-style
264
- # decoding instead of standard next-token sampling, so it has a query
265
- # for the last sampled token plus queries for each draft token.
266
- self.num_lookahead_tokens = self.num_spec_tokens + 1
267
- if speculative_config.use_dspark():
268
- # DSpark drafts a block of num_spec_tokens query tokens in which the
269
- # anchor itself is the first prediction position (no separate bonus
270
- # query), so it needs exactly num_spec_tokens lookahead slots.
271
- self.num_lookahead_tokens = self.num_spec_tokens
257
+ self.use_eagle = speculative_config.use_eagle()
258
259
# Create the KV cache manager.
260
if hash_block_size is None:
vllm/v1/worker/gpu/warmup.py
+59
-21
@@ -18,13 +18,61 @@ from vllm.v1.core.sched.output import (
18
NewRequestData,
19
SchedulerOutput,
20
)
21
-from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec
21
+from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec, MambaSpec
22
from vllm.v1.request import Request
23
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
24
25
logger = init_logger(__name__)
26
27
28
+def _reserved_block_count(
29
+ num_tokens: int,
30
+ kvcache_spec: KVCacheSpec,
31
+ *,
32
+ num_lookahead_tokens: int,
33
+ max_model_len: int,
34
+ max_encoder_len: int,
35
+) -> int:
36
+ """Number of blocks the scheduler would hold for a request of `num_tokens`.
37
+
38
+ Warmup hand-builds its `SchedulerOutput`s, so it must reserve what
39
+ `KVCacheManager.allocate_slots` reserves: the token range plus
40
+ `num_lookahead_tokens`, where the speculator writes the KV of its drafts.
41
+ """
42
+ if isinstance(kvcache_spec, CrossAttentionSpec):
43
+ # Cross-attention blocks cover the encoder sequence only.
44
+ return cdiv(max_encoder_len, kvcache_spec.block_size)
45
+ num_speculative_blocks = 0
46
+ if isinstance(kvcache_spec, MambaSpec):
47
+ # MambaManager appends speculative running-state blocks in every cache
48
+ # mode; align mode sizes from the uncapped, lookahead-free token range.
49
+ num_speculative_blocks = kvcache_spec.num_speculative_blocks
50
+ if kvcache_spec.mamba_cache_mode == "align":
51
+ return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks
52
+ num_tokens = min(num_tokens + num_lookahead_tokens, max_model_len)
53
+ return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks
54
+
55
+
56
+def _warmup_block_counter(
57
+ model_runner: GPUModelRunner,
58
+) -> Callable[[int, KVCacheSpec], int]:
59
+ """Bind `_reserved_block_count` to `model_runner`'s reservation policy."""
60
+ num_lookahead_tokens = model_runner.vllm_config.num_lookahead_tokens
61
+ max_model_len = model_runner.max_model_len
62
+ max_encoder_len = getattr(model_runner.model_state, "max_encoder_len", 0)
63
+
64
+ def block_count(num_tokens: int, kvcache_spec: KVCacheSpec) -> int:
65
+ return _reserved_block_count(
66
+ num_tokens,
67
+ kvcache_spec,
68
+ num_lookahead_tokens=num_lookahead_tokens,
69
+ max_model_len=max_model_len,
70
+ max_encoder_len=max_encoder_len,
71
+ )
72
+
73
+ return block_count
74
+
75
+
76
def run_mixed_prefill_decode_warmup(
77
model_runner: GPUModelRunner,
78
worker_execute_model: Callable[[SchedulerOutput], Any],
@@ -48,21 +96,20 @@ def run_mixed_prefill_decode_warmup(
96
97
kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
98
num_kv_cache_groups = len(kv_cache_groups)
51
- group_block_sizes = [g.kv_cache_spec.block_size for g in kv_cache_groups]
99
+ block_count = _warmup_block_counter(model_runner)
100
+ kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups]
101
decode_prefill_block_counts = [
53
- cdiv(decode_prompt_len, block_size) for block_size in group_block_sizes
102
+ block_count(decode_prompt_len, s) for s in kv_cache_specs
103
]
104
decode_block_counts = [
56
- cdiv(decode_prompt_len + decode_scheduled_tokens, block_size)
57
- for block_size in group_block_sizes
105
+ block_count(decode_prompt_len + decode_scheduled_tokens, s)
106
+ for s in kv_cache_specs
107
]
108
decode_block_deltas = [
109
decode - prefill
110
for decode, prefill in zip(decode_block_counts, decode_prefill_block_counts)
111
]
63
- prefill_block_counts = [
64
- cdiv(prefill_len, block_size) for block_size in group_block_sizes
65
- ]
112
+ prefill_block_counts = [block_count(prefill_len, s) for s in kv_cache_specs]
113
required_blocks = sum(decode_block_counts) + sum(prefill_block_counts)
114
if model_runner.kv_cache_config.num_blocks <= required_blocks:
115
logger.warning(
@@ -200,19 +247,10 @@ def warmup_kernels(
247
]
248
249
# Compute per-request block counts for each KV cache group.
203
- def _warmup_block_count(num_tokens: int, spec: Any) -> int:
204
- if isinstance(spec, CrossAttentionSpec):
205
- num_tokens = max_encoder_len
206
- num_blocks = cdiv(num_tokens, spec.block_size)
207
- if isinstance(spec, MambaSpec) and spec.mamba_cache_mode == "align":
208
- # Align mode reserves extra blocks beyond the token range for the
209
- # speculative-decode running-state snapshots.
210
- num_blocks += spec.num_speculative_blocks
211
- return num_blocks
212
-
250
+ block_count = _warmup_block_counter(model_runner)
251
kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups]
214
- prefill_block_counts = [_warmup_block_count(prompt_len, s) for s in kv_cache_specs]
215
- decode_block_counts = [_warmup_block_count(decode_len, s) for s in kv_cache_specs]
252
+ prefill_block_counts = [block_count(prompt_len, s) for s in kv_cache_specs]
253
+ decode_block_counts = [block_count(decode_len, s) for s in kv_cache_specs]
254
max_blocks_per_req = sum(decode_block_counts)
255
256
num_reqs = min(
@@ -316,7 +354,7 @@ def warmup_kernels(
354
num_tokens = decode_query_len if use_spec else 1
355
after = req_computed[i] + num_tokens
356
deltas = [
319
- _warmup_block_count(after, spec) - held
357
+ block_count(after, spec) - held
358
for spec, held in zip(kv_cache_specs, req_blocks[i])
359
]
360
cached_req_data.new_block_ids.append(