[Bugfix][V2] Warm up kernels before capturing CUDA graphs - #55341
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
warmup_kernels runs a scheduler-realistic prefill and a max_num_seqs decode step. Running it after capture_model() means those shapes reach the shared workspace arena only once graphs already hold pointers into it, and a resize swaps in a fresh buffer and frees the one they captured, so the next replay writes into freed memory. Hoist it next to kernel_warmup(), which already runs before capture for the same reason. The V1 sampler warmup stays where it is; its comment documents why it has to follow capture. Fixes vllm-project#55336 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
9f85583 to
c587b02
Compare
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (6)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughThe V2 runner now warms up kernels before CUDA graph capture and locks workspace storage after non-profile captures. Profiling captures remain unlocked, and related capture stubs and workspace tests support the new behavior. ChangesV2 workspace safety and warmup
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to V2 kernel warmup now occurs before CUDA graph capture, and execution captures lock the workspace afterward to prevent later reallocation from invalidating captured buffers. Profiling remains able to grow the workspace, with matching test coverage; the change is ready to merge. Sequence Diagram(s)sequenceDiagram
participant GPUWorker
participant GPUModelRunner
participant WorkspaceManager
GPUWorker->>GPUModelRunner: warmup_kernels()
GPUWorker->>GPUModelRunner: capture_model()
GPUModelRunner->>WorkspaceManager: lock_workspace()
WorkspaceManager-->>GPUModelRunner: reject post-capture growth
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
/ci run |
he-yufeng
left a comment
he-yufeng
left a comment
There was a problem hiding this comment.
Read this against the V1 path since #55336 quoted the missing lock. The hoist looks right as the primary fix: sizing the arena to scheduler-realistic shapes before capture means the graphs bake in a buffer that never needs to grow for warmup-bounded demand, which is what corrupts today.
One gap worth considering. This covers demand bounded by the warmup shapes. If anything later exceeds them (a request shape outside what warmup_kernels exercised, or a path the warmup does not hit), _ensure_workspace_size still replaces the arena after capture and the corruption is back, silently. V1 handles that class by locking after capture (lock_workspace() at gpu_model_runner.py:6998), so residual growth fails with the manager's precise error instead of writing into freed memory. Would it make sense to also lock on the V2 path once capture_model() returns? It is a one-liner that turns every remaining unknown growth path from silent corruption into a loud failure, and it mirrors the invariant V1 already enforces.
Side note: the format/pre-commit lanes are red on the current head, looks like trivial formatting.
sylvesterkaczmarek
left a comment
sylvesterkaczmarek
left a comment
There was a problem hiding this comment.
The ordering invariant is correct: any warmup that can resize or replace shared workspace has to finish before CUDA graph capture freezes addresses. Hoisting the existing V2 warmup rather than adding a second allocator path keeps the fix narrow, and leaving the V1 sampler warmup in its documented post-capture position avoids conflating two different lifecycle requirements.
Signed-off-by: Nick Hill <nickhill123@gmail.com>
|
Thanks @aoshen02, good catch! I added some additional changes to lock the workspace as additional protection, along with some tests. |
|
/ci run |
|
✅ Triggered Buildkite CI #87320 for commit |
…ct#55341) Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: Jyotirmoy Roy <jyotirmoyroy649@gmail.com>
4 participants
Purpose
Fixes #55336.
warmup_kernelsruns a scheduler-realistic prefill plus a decode step atmax_num_seqs. On the V2 path it is called aftercapture_model(), so thoseshapes reach the shared workspace arena only once CUDA graphs already hold
pointers into it.
WorkspaceManager._ensure_workspace_sizegrows the arena byreplacing the tensor, which frees the buffer the captured graphs baked in, and
the next replay writes into freed memory.
This hoists the V2 warmup next to
kernel_warmup(), which already runs beforecapture for the same class of reason. The V1 sampler warmup is untouched — its
comment documents why it has to follow capture.
Warming up before capture is also what TensorRT-LLM and SGLang do:
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pywarms up ahead of eachcapture with the comment "This also lets us initialize states in the
attn_metadata and resize the shared attention workspace before any graph is
captured", and SGLang runs two warmup iterations per shape before capturing.
Test Plan
pre-commit run --all-files # ruff, ruff-format, mypyServing DeepSeek-V4-Flash on GB200, which aborts during startup on main with
under
VLLM_BATCH_INVARIANT=1,--max-num-seqs 256,cudagraph_mode=FULL_AND_PIECEWISE.Two setups: 8x GB200 (DP8/EP8/TP1, 43 layers) and a single GB200 with a 4-layer
proxy of the same model. Measured four things: whether this patch alone fixes
the crash, cold-start time, KV cache sizing, and throughput on a healthy
configuration that never hit the bug.
Test Result
pre-commit: passed (ruff, ruff-format, mypy).Fix, isolated. This patch is the only change in the third row; a separate
one-line model-side fix I am upstreaming as #55299 is not sufficient on its own,
which rules it out as the cause of the repair:
max-num-seqs 256, FULL_AND_PIECEWISE8x GB200, 43 layers: 37 FULL + 53 PIECEWISE graphs captured, zero illegal
memory accesses, batch-invariance probes intact, server serving.
Cold start (first log line to health, 8x GB200, 43 layers): 8m19s with this
patch vs 8m10s without, i.e. +1.8%, within run-to-run noise. The workload is
unchanged, only its position moves.
Memory accounting is byte-identical. The workspace this bug is about is
allocated either way — the patch only moves when — and the KV cache is sized
before
compile_or_warm_up_modelruns, so nothing shifts. Every startup figureon the 8x GB200 run matches:
No throughput regression on a configuration that was already healthy —
DeepSeek-V4-Flash,
VLLM_BATCH_INVARIANT=0, FULL_AND_PIECEWISE,--max-num-batched-tokens 8192, 256 concurrency, 150/3000 in/out:Notes
AI assistance was used for this change: the fault was located with
instrumentation written by Claude Code, and the patch was drafted with it. The
diff moves one call and I have reviewed it end to end.
🤖 Generated with Claude Code