[Bugfix] Defer adaptive verification until after kernel warmup (#55455)
Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: Codex <noreply@openai.com>
Kevin H. Luu committed
Sep 6, 2026 at 11:31 UTC
6865e67f0be02d53694517f6f71d7fb96492792d
3 files changed
+38
tests/v1/worker/test_gpu_warmup_blocks.py
+1
@@ -88,6 +88,7 @@ def _make_runner(
88
"""Stub model runner exposing only what the warmup entry points read."""
89
return SimpleNamespace(
90
num_speculative_steps=num_spec_steps,
91
+ adaptive_verification=None,
92
decode_query_len=num_spec_steps + 1,
93
is_pooling_model=False,
94
is_encoder_decoder=False,
tests/v1/worker/test_mixed_warmup_gate.py
+22
@@ -6,6 +6,7 @@ from types import SimpleNamespace
6
7
import pytest
8
9
+from vllm.v1.worker.gpu import warmup
10
from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup
11
12
@@ -28,3 +29,24 @@ def test_mixed_warmup_skipped_for_single_seq(max_num_reqs):
29
)
30
is False
31
)
32
+
33
+
34
+@pytest.mark.parametrize("fail_warmup", [False, True])
35
+def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_warmup):
36
+ """Startup must warm fixed drafts before calibration and retain its manager."""
37
+ manager = SimpleNamespace(cost_tables=None)
38
+ runner = SimpleNamespace(adaptive_verification=manager)
39
+
40
+ def run_steps(model_runner, execute, sample):
41
+ assert model_runner.adaptive_verification is None
42
+ if fail_warmup:
43
+ raise RuntimeError("warmup failed")
44
+
45
+ monkeypatch.setattr(warmup, "_warmup_kernels", run_steps)
46
+ if fail_warmup:
47
+ with pytest.raises(RuntimeError, match="warmup failed"):
48
+ warmup.warmup_kernels(runner, _fail, _fail)
49
+ else:
50
+ warmup.warmup_kernels(runner, _fail, _fail)
51
+ assert runner.adaptive_verification is manager
52
+ assert manager.cost_tables is None
vllm/v1/worker/gpu/warmup.py
+15
@@ -220,6 +220,21 @@ def warmup_kernels(
220
We must call the provided worker's execute_model for pipeline parallel
221
coordination.
222
"""
223
+ # Adaptive costs are calibrated during capture, after this warmup. Exercise
224
+ # fixed draft counts here, then restore the manager for capture and serving.
225
+ adaptive_verification = model_runner.adaptive_verification
226
+ model_runner.adaptive_verification = None
227
+ try:
228
+ _warmup_kernels(model_runner, worker_execute_model, worker_sample_tokens)
229
+ finally:
230
+ model_runner.adaptive_verification = adaptive_verification
231
+
232
+
233
+def _warmup_kernels(
234
+ model_runner: GPUModelRunner,
235
+ worker_execute_model: Callable[[SchedulerOutput], Any],
236
+ worker_sample_tokens: Callable[[GrammarOutput | None], Any],
237
+) -> None:
238
if model_runner.vllm_config.is_mm_encoder_only:
239
return
240