[fix][model_engine][rollout]: blank a capture's cache write targets, and keep the KV pool size across a wake - #2267
Conversation
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
df5d8dd to
6b0f89c
Compare
xysheng-AMD
changed the title
|
Two real defects, and the diagnosis of 2 is a good catch. But both fixes land on the symptom rather than the shared root cause, and one of them has a hole. Details below. 1. The stated difference between startup and wake does not holdThe comment says "on startup the bytes land in a pool no request has reached yet; here the pool is one step from serving." Both paths go through the same The actual difference is the address written, not the pool: 2.
|
Review — sizing the KV budget against our own shareReviewed at The thesis is right and the fix is correct as far as it goes. Provenance: [verified] means I read the deciding code myself on this SHA. [reported] means the shape matches the code but I did not trace the whole chain. 1. The fix is at the wrong altitude1.1 A wake-time shrink never reaches the engine process, so the corruption path stays open [verified]The startup path deliberately carries the number across the process boundary by hand, and says why: # atom/model_engine/engine_core.py:128-144
block_info = self.runner_mgr.call_func("get_num_blocks", wait_out=True)
num_blocks = block_info["num_kvcache_blocks"]
# Sizing happens in the runner subprocess, so nothing it wrote to
# its own `config` is visible here. ...
config.num_kvcache_blocks = num_blocks
The wake path does none of it: # atom/rollout/memory_manager.py:408-415 -- runs in the RUNNER subprocess
available_blocks = self.get_num_blocks()["num_kvcache_blocks"]
num_blocks = min(saved_blocks, available_blocks)
if num_blocks < saved_blocks:
logger.warning(f"... KV cache blocks reduced from {saved_blocks} to {num_blocks} ...")
self.allocate_kv_cache(num_blocks)
This PR removes the shrink for one configuration. It does not close the path, and the one signal an operator gets is that 1.2 On wake, each TP rank sizes itself independently [reported]At startup exactly one reply is taken from The PR reduces how often ranks disagree; it does not make them agree. 1.3 One of four live inputs was latched [reported]For the captured decode graphs to stay valid the wake pool must equal the sizing pool. Allocating 2. The principle is established but not sweptCLAUDE.md: "Fix-then-sweep: after fixing a bug, immediately grep for the same pattern across the codebase and fix all occurrences in one pass." Two occurrences of the exact pattern this PR names as a bug are untouched. 2.1
|
6ddd752 to
802af4c
Compare
|
@yitingw1 @valarLip — thank you both. You arrived from different sides at the Root cause. Two fixes.
Evidence. Qwen3-8B smoke, 12 steps, releasing sleep. Same base, image and
The middle column is the point: with only the blanking applied the pool still On the 30B config the fault path is closed by fix 2: the wake-time count never Review closure. @yitingw1 #1–#5 are addressed (#3 by construction — nothing Behaviour change worth naming: if the memory is genuinely gone, Gate. All four Pre Checkin jobs green on One disclosure: the GPU arms ran on @valarLip this should close 1.1. If you both agree, an approval would let the |
Second pass — reviewed at
|
32282bc to
d87233e
Compare
|
@valarLip thank you — #3 is the one that mattered. #0 — the body did not describe the diff. Rewritten. Process failure on my #3 — correct, and the interface is rebuilt in the shape you named. The There is a layer past name-matching, too: But GDN is not covered, and that is after trying. I added the override filling That leaves GDN's recurrent state exactly where it was on One correction to your mechanism: On the three you marked [reported]: I checked two, and the new interface does #1 — correct, fixed. #2 — correct, the comment now says so instead of claiming the failure is clean. #4 — fixed, to §6, taken: Not in this PR, and I will take any of them: the GDN and V4 / Qwen4-Exp state write targets, which are one VerificationNon-GPU: 6 failed / 5938 passed, the 6 being GPU, on this head: Qwen3-8B 12 steps — 0/12 NaN steps, 0 wake-time shrinks, 112 CI: Pre Checkin all green, Qwen3-Next back to success. The remaining reds sit in |
xysheng-AMD
changed the title
…batch's `capture_cudagraph` runs the model for real -- a warmup forward and then the captured one -- so both write cache wherever the backend's index buffers point. The capture builders synthesize the *read* side (`kv_indptr`, `kv_indices`, `kv_last_page_lens`) precisely because a stale value there walks the kernel off the KV list, and leave the write side at whatever the last real batch left. At startup those buffers are freshly zeroed, so a capture writes row 0 and nothing else. `forward_vars` is not released on sleep, so on a wake they still hold a served batch's rows and the capture scatters writes across live entries -- or past the end of the pool, if it came back smaller. Both were measured. On Qwen3-30B-A3B, where a wake also resized the pool from 33562 blocks to 8316, the stale ids fell outside it and all eight replicas died with `Memory access fault by GPU node-N` during `Capturing bs=256` -- `expandable_segments` was not set and `_warn_if_recapture_will_fault` never fired, so the alloc-config path the warning exists for is not this. On Qwen3-8B at 4k, where the pool kept its size, the same writes landed on live slots instead: one decode row per batch came back NaN across the whole 151669-entry vocabulary while its `context_lens` incremented normally, the sampler drew token 0 out of that row, and that single logprob turned the step's entire rollout -- every aggregate, and `grad_norm` with it -- into NaN, so all eight actors skipped the update. Ten of twelve steps were lost that way; with this, none of twelve. Applying only this commit and leaving the pool free to shrink on every wake is also 0 of 12, which is what places the NaN on the write rather than on the sizing. So the capture is pointed at nothing first. `PAD_SLOT_ID` is the established sentinel: every step already puts it on the tail of a CUDAGraph-padded batch, the paged cache kernels skip it (`slot_idx < 0`), and so do GDN's `fused_recurrent_gated_delta_rule` and replayssm kernels. Replay overwrites the buffers with real rows, so nothing about the captured graph changes -- the same reason the builders can pin the read side to block 0. Which buffers those are is the backend's knowledge, so the enumeration is `AttentionMetadataBuilder.cache_write_targets` rather than a scan of `forward_vars` key names. Paged attention aims writes with `slot_mapping`, and the default covers it by suffix so TBO's per-ubatch mirrors come along. GDN aims recurrent-state writes with `spec_state_indices` / `non_spec_state_indices`, which are not spelled that way, and with `non_spec_state_indices_in_tensor`, which is not in `forward_vars` at all -- no naming rule could have reached the second kind, which is why the enumeration moved rather than got a better regex. GDN's own state indices stay uncovered: blanking them reads as correct, but it was tried and Qwen3-Next stopped coming up inside `LLMEngine.__init__`, where the startup capture runs. Uncovered is the state that backend was already in, so it waits for a fix that can be run against a checkpoint of one. Both live on the builder, next to the capture metadata they belong with. Co-authored-by: Cursor <cursoragent@cursor.com>
`_resume_kv_cache` re-derived the block count and `min()`-ed it with the saved one, so anything holding device memory at wake time silently handed the engine back a smaller pool. Nothing downstream is told. `BlockManager` sizes its `BlockPool` in the ENGINE process, from the `config.num_kvcache_blocks` that `EngineCore.__init__` copies out of the startup reply by hand -- with a comment saying why: sizing runs in the runner subprocess, so nothing it writes to its own `config` is visible there. The wake path has no such step. `allocate_kv_cache` sets the subprocess's copy, and neither `_handle_resume_memory`, `broadcast_utility_command_sync` nor `AsyncLLMEngine.wake_up` carries the count back, so the scheduler goes on issuing block ids past the end of the reallocated pool. It was not even one count. `resume_memory` broadcasts the command rather than a number, so every TP rank sized itself against its own `mem_get_info()`; the `all_reduce` in `get_num_blocks` only fires under pipeline parallelism. A peer allocating unevenly across cards could therefore leave a block id valid on rank 0 and out of bounds on rank 3. Measured on Qwen3-30B-A3B at utilization 0.45: 33562 blocks at startup, 8316 on the first wake while 176GB was still free, and eight `Memory access fault by GPU node-N`. The pool's size is settled when the engine is sized -- the scheduler and the captured decode graphs are both built against it -- so the wake path has no business re-deriving it. `reset_peak_memory_stats()` goes with the re-derivation, `get_num_blocks` having been its only reader here; the headroom it used to consult is logged instead, so wake-time pressure still leaves a record. That makes `allocate_kv_cache` the call that can fail, which changes two things. `_kv_cache_num_blocks` is now cleared only once it has returned: clearing it first loses the only record of the pool size on the way out, and the next wake then takes the "nothing to resume from" guard and lets `resume_memory` report success for an engine with no pool at all. And the raise is not a clean failure -- `resume_memory` is not a barrier func and `AsyncIOProc` does not guard the dispatch, so one rank raising kills that worker while the others block on a collective it will never reach. The comment says so rather than claiming otherwise; what makes this preferable to clamping is that clamping was silent corruption, and a cross-rank preflight is the thing that would make it clean. Co-authored-by: Cursor <cursoragent@cursor.com>
d87233e to
d38aedc
Compare
3 participants
Description
After a sleep/wake cycle the CUDA graph recapture writes cache through the rows
of the last real batch, and the wake re-sizes the KV pool to something the
scheduler does not know about. Independent defects, fixed separately.
1. A capture writes cache through live write targets
capture_cudagraphruns the model for real — a warmup forward and then thecaptured one — so both write cache through the backend's index buffers. The
capture builders synthesize the read side (
kv_indptr,kv_indices,kv_last_page_lens) because a stale value there walks the kernel off the KV list;the write side is left alone. And
forward_varsis not released on sleep.So at startup those buffers are still zeroed and a capture writes row 0 — safe by
luck, not design. After a wake they hold a served batch's rows. One cause, two
symptoms, depending on whether the pool also changed size:
eight replicas die with
Memory access faultduringCapturing bs=256.all-NaN, the sampler draws token 0 out of it, and that single logprob turns the
step's whole rollout —
grad_normincluded — into NaN. Generation quality looksunchanged, which is why this was hard to spot.
Fix: point the write targets at nothing before capturing.
PAD_SLOT_IDis theestablished convention — every step already puts it on the tail of a
CUDAGraph-padded batch, and the cache kernels skip it. Replay overwrites the
buffers with real rows, so nothing about the captured graph changes; that is the
same reason the builders can pin the read side to block 0.
Which buffers are write targets is the backend's knowledge, so the enumeration
is
AttentionMetadataBuilder.cache_write_targets()rather than a scan offorward_varskey names. Paged attention aims writes withslot_mapping, and thedefault covers it by suffix so TBO's per-ubatch mirrors come along. This matters
beyond tidiness: GDN aims recurrent-state writes with buffers that are not spelled
that way, one of which (
non_spec_state_indices_in_tensor) is not inforward_varsat all, so no naming rule could have reached it.GDN's own state indices are left uncovered here. Blanking them reads as
correct, but it was tried and Qwen3-Next stopped coming up, dying inside
LLMEngine.__init__where the startup capture runs; Kimi-K3, which inherits thesame mixin, was unaffected. That is the state GDN was already in on
main— thisPR does not change it, and covering it wrongly breaks a model that works today —
so it waits for a fix that can be run against a checkpoint of one. The PR is still
a net improvement for that family: their attention layers are aimed by
slot_mapping, and that half is fixed.2. A wake re-sizes the pool, and nothing downstream is told
_resume_kv_cachere-derived the block count andmin()-ed it with the saved one,so anything holding device memory at wake time handed the engine back a smaller
pool.
BlockManagersizes itsBlockPoolin the engine process, from the countEngineCore.__init__copies out of the startup reply by hand — with a commentsaying why: sizing runs in the runner subprocess, so nothing it writes to its own
configis visible there. The wake path has no such step, so the scheduler goes onissuing block ids past the end of the pool. It was not even one count:
resume_memorybroadcasts the command rather than a number, so every TP rank sizeditself and they could disagree.
Fix: allocate
saved_blocks. The pool's size is settled when the engine issized — the scheduler and the captured decode graphs are both built against it — so
the wake path has no business re-deriving it. Net deletion; the headroom the clamp
used to consult is logged instead, so wake-time pressure still leaves a record.
Two behaviour changes.
allocate_kv_cachecan now raise, so_kv_cache_num_blocksis cleared only once it has returned — clearing it firstloses the only record of the pool size, and the next wake then reports success for
an engine with no pool. And that raise is not a clean failure:
resume_memoryis not a barrier func, so one rank raising kills that worker whilethe others block on a collective it will never reach. It beats clamping only
because clamping was silent corruption; a cross-rank preflight is what would make
it clean, and is not in this change.
Measurements
Qwen3-8B DAPO smoke, 12 steps, releasing sleep. Same base, image and config, same
112 sleep/wake cycles; the only variable is which commit is applied:
rollout_pplgrad_normThe middle column is what makes this one root cause: with only the blanking
applied the pool still shrinks on every wake and the NaN is already gone, so the
NaN is the capture's write and not the sizing.
This head was re-run on that config: 0/12 NaN, 0 shrinks, 112 recaptures completed,
0 faults. That run is also what covers the call site, which the non-GPU suite
cannot reach.
Qwen3-30B-A3B, 8× MI355X, resumed from
global_step_110for 31 steps: CUDA graphrecapture from 0/8 all faulted to 272 completed,
Memory access faultfrom 8 to 0,wake-time shrinks 0.
All GPU arms ran on
6ced6ce51plus these two commits rather than the rebasedhead, because current
mainneeds a neweraiterthan our validation imagecarries; both apply there with zero conflicts and identical code.
Checks
tests/test_heavy_ci_gate.py, which shells out tojq(absent from thecontainer, green in real CI) — identical to base's 6 / 5924. Each commit green on
its own: 5924 → 5935 → 5938.
black --check .870 files, 0 would-reformat; ruff with the diff-context filteragainst the merge-base, 0 findings would block the job.