fix(plugin/dsv4): correct DeepSeek-V4 prefix hits and cudagraph replay on the vLLM plugin

[linsun12]

This PR fixes three correctness problems of DeepSeek-V4 on the vLLM plugin (DeepseekV4ForCausalLM: V4-Flash, V4-Flash-0731, V4-Pro, V4-Pro-0813), all on main's default settings with prefix caching on:

1. A prefill step whose rows are all one token long crashes the engine (Memory access fault by GPU, also with --enforce-eager). This happens on the 1-token tail of a prefix hit, a 1-token last chunk, and a 1-token prompt.
2. Wrong output with cudagraphs on. vLLM's V1 runner replays the FULL decode graph for steps where a row is still prefilling. Separately, whenever vLLM's breakable cudagraph is on, prefill and mixed batches are wrong.
3. Prefix hits resume from a wrong state. The sliding-window ring and compressor tails at the hit boundary are not cached, so every hit continues from an all-zero state.

Only atom/plugin/vllm/ and tests/plugin/ change, plus the registration and documentation of the plugin's new environment variables (atom/utils/envs.py, docs/environment_variables.md); ATOM core is untouched. Base: main 661ef7661 (vLLM 0.31.0 pin db9527a46, #2236).

part	commits
1. 1-token prefill rows (same change as #2501)	0b8ea573d fix(plugin/dsv4): score 1-token prefill rows on the dense indexer path
2. cudagraph	135f5f7df fix(plugin/dsv4): keep ATOM's V4 off vLLM's breakable cudagraph 3cf6a1d5c fix(plugin/dsv4): a still-prefilling row never replays the FULL graph
3. prefix-hit state (was #2503)	4601b0207 feat(plugin): keep per-request state checkpoint images in vLLM blocks e7b91038a feat(plugin/dsv4): restore prefix hits from checkpoint images
review fixes (since 53eff4452)	27556888f fix(plugin): do not publish a pending image hash on abort 847f038ac fix(plugin/dsv4): leave models that are not ATOM's V4 to vLLM's allocator ba56222e6 fix(plugin): count the null block and image working blocks in the pool check cf54c1a88 fix(plugin): re-fit an auto-fitted max_model_len after withholding the slot tail 021201b50 fix(plugin/dsv4): write the cudagraph counters on every step 83c8b63e9 fix(plugin): register and document the image and cudagraph diagnostics env vars
follow-up review round	0b91463b4 test(plugin): pin the decode keep at a request's end against native placement 9a4b01b90 fix(plugin): refuse a leftover image store instead of copying a stale slot 087ca8e53 test(plugin): pin that a deferred decode keep is published while the request runs

The parts are stacked in this order because each one relies on the one before. Part 2 sends 1-token prefill steps down the prefill path, which faults without part 1. Part 3's stores run after FULL replays and rely on part 2's dispatch rule.

Changes since the review of 53eff4452

Each review finding is fixed in its own commit on top of the five commits above, with a regression test that fails on the tree before the commit and passes after it.

finding	fix	test
Aborting an async decode could crash EngineCore (PerryZhang01 #1). A decode keep whose block is closed by the token still in flight is hashed later. An abort is handled before the step's output lands, so on cancellation on_free published a boundary hash that did not exist yet and BlockPool.cache_full_blocks raised IndexError, which took EngineCore down.	27556888f: a pending keep is published only if its boundary hash exists, read at the image block size through vLLM's resolve_block_hashes (vLLM may hash at a finer hash_block_size); otherwise it is dropped and the request is freed normally. The two other places that read deferred hashes use the same view.	test_paged_state_image_scheduler.py on vLLM's real AsyncScheduler: abort with the keep deferred (hash block 256 and 128), natural completion (published), preemption (dropped, resumes)
The V4 adapter matched vLLM's own V4.1 and V4 (PerryZhang01 #2, Copilot). It matched any architecture containing DeepseekV4, and the KV-config hook sized V4 images before looking for an ATOM proxy: vLLM's V4.1 failed with unknown V4 compress ratios [1, 2], and vLLM's own V4 got align forced or MTP refused.	847f038ac: DeepseekV4ForCausalLM exactly and only when vLLM's registry maps it to an ATOM class (the check part 2's opt-out uses); the hook passes through unless the ATOM proxy is in the KV-cache specs.	vLLM's V4.1 with prefix caching on and off, V4 resolved to vLLM's class, ATOM's V4, a matching model without the proxy
The minimum-pool check missed two blocks' worth of accounting (PerryZhang01 #4). vLLM's null block was counted as allocatable, and with ATOM_STATE_CHECKPOINT_INTERVAL_TOKENS=0 the image groups' working blocks were not counted.	ba56222e6: the check needs a running request's peak: its KV blocks, two blocks per image group (the running one and the previous one, published or not) and the null block.	real scheduler at the minimum ±1, with and without preempt / resume from the prompt-end anchor; interval 0
Slot reservation came after vLLM's auto-fit (PerryZhang01 #3). With --max-model-len -1, vLLM fitted the length to every physical block and the hook then refused it.	cf54c1a88: vLLM's own auto-fit runs again against the blocks left after the slot tail (minus the null block); the allocation is unchanged, so the slot bytes stay backed and nothing is reserved twice.	vLLM's real get_kv_cache_configs: the length drops from 23,808 to 20,992 tokens and a request of that length is served
Diagnostic counters lost their tail (PerryZhang01, Copilot).	021201b50: the opt-in counters file is rewritten on every step.	a subprocess with 127 steps records 127 (64 before)
New environment variables were not registered or documented (Copilot).	83c8b63e9: ATOM_STATE_CHECKPOINT_INTERVAL_TOKENS, ATOM_STATE_IMAGE_CHECK_BLOCKS and ATOM_V4_CG_INSTR_DIR are registered in atom/utils/envs.py, documented, and read through envs.	registration and documentation test; tests/test_envs.py
A leftover image store could cache a stale slot (Copilot 4255518933). A store registered by a metadata build runs right after that step's forward; if the forward never ran, the next build flushed it anyway, copying a slot that still holds an earlier position's state into image blocks already hashed for the planned boundary.	9a4b01b90: such a store is refused with a RuntimeError naming the broken contract, and nothing is copied. Failing loudly is the right response: if the forward were skipped, that step's paged KV would not have been written either, while vLLM had already hashed those KV blocks at scheduling time, so vLLM's own prefix cache would already be wrong and no image-side recovery could make serving correct. It cannot be cancelled cleanly from the worker: the blocks are hashed in EngineCore's block pool, and with async scheduling the next step, possibly already granted a prefix hit on them, is scheduled before the worker could report it. The path is not reachable with the pinned vLLM (execute_model always runs _model_forward after the build; the early returns come before it), and every GPU run counted 0 such stores.	a build without its forward, then the next build: no copy, an error (the stale copy happened silently before)
A deferred decode keep might only be published when the request ends (Copilot 4261779415): the reviewer read cache_image_blocks as advancing the image groups' cached-block count past a decode keep's block before the keep is deferred, which would strand the keep until on_free.	Not reproducible with the pinned vLLM, so no code change; 087ca8e53 adds the regression test. Both of vLLM's cache_blocks call sites count only finalized tokens, and a block's hash exists as soon as its tokens are known, so the count reaches the keep's block only together with its hash: the keep is then hashed at once, or, if deferred, published on the next schedule while the request keeps generating. On GPU, follow-ups sent while the first request was still generating hit its decode keep (below).	the keep is cached while the request runs, a later keep does not replace it, and a second request hits it mid-generation (passes on the code as it was); the CPU equivalence and the GPU run below

Not changed: the suggestion to replace _attach_native() with an explicit shared layout / copy component is kept as a follow-up (see below).

1. 1-token prefill rows: GPU memory access fault

Problem. With prefix caching on, also with --enforce-eager, the engine dies with Memory access fault by GPU node-… on all TP ranks, followed by EngineDeadError, whenever the last prompt chunk of a request is a single token:

• a prompt of 128·m + 1 tokens sent again: the block-aligned prefix hit leaves one token to compute (129, 897, 1025, 1153, 2049, 8193, … tokens);
• a fresh prompt of 8192·k + 1 tokens with --max-num-batched-tokens 8192: its last chunk is one token;
• a fresh 1-token prompt.

Repro on main: send a 1025-token prompt, then send it again. The second request returns HTTP 500.

Root cause. The bridge builds such a step as a prefill, but build_atom_v4_attention_metadata split the indexer rows by length only. The leading run of rows with the first row's length, if <= decode_q, went to the paged decode indexer. When every row is one token long, that is every row, so the dense prefill group was empty. _populate_indexer's cu[-1] = max(cu[-1], 1) still claimed one committed row, so IndexerVllm.indexer_score_topk ran _score_topk_prefill on the empty slice and gathered one row through an empty block table: an out-of-bounds read.

Fix.

• _indexer_decode_group(lens, decode_q, is_prefilling): the paged decode group now also ends at the first row that is still prefilling. Such a row takes the dense prefill path, as in native ATOM, so the dense group is not empty on these steps. Batches without a prefilling row split exactly as before. is_prefilling is vLLM's CPU num_computed < num_prompt mask, so this adds no sync.
• IndexerVllm.indexer_score_topk: if no prefill rows are left, it returns the decode result instead of running the dense path on an empty slice.

2. Wrong output with cudagraphs on

Problem.

• Default settings (no cudagraph flags). Steps whose rows all have the decode query length (1, or 1 + k with MTP) are wrong while some row is still in its prompt. These are the 1-token tail after a prefix hit, a 1-token last chunk (8193), and a 1-token prompt. Under load, other requests in the same batches also get wrong output, and requests start failing: with part 1 alone, the mixed load below gives 31 % garbage and 12 % HTTP 400. On vLLM 0.28 the engine crashed with Memory access fault by GPU.
• With the breakable cudagraph (VLLM_USE_BREAKABLE_CUDAGRAPH=1). Every prefill or mixed batch of at most 128 tokens (the largest capture size) is wrong: fresh prompts of ≤ 128 tokens, a short last chunk of a long prompt (8194 … 8320), prefix hits with a 2 … 128-token tail, and any decode batch a short prefill joins. vLLM sets the variable itself for V4 when it is unset on vLLM ≤ 0.28 (all platforms) and on 0.31 off ROCm. VLLM_BATCH_INVARIANT=1 can also set it, and vLLM's V4 recipe sets it explicitly. On vLLM 0.31 / ROCm, vLLM no longer enables it automatically for V4, so on main this only happens when the variable is set.

Root cause: V1 replays the FULL decode graph for a still-prefilling row. On ROCm, V4 runs on Model Runner V1 (see below). GPUModelRunner._is_uniform_decode decides uniform decode by shape only (max_num_scheduled_tokens == uniform_decode_query_len and num_tokens == max * num_reqs). The V4 bridge builds such a step as PREFILL (_is_pure_uniform_decode rejects rows with is_prefilling) with fresh tensors. It does not refresh the persistent decode buffers the FULL graph reads, so the replay reads the previous step's metadata (positions, slot mapping, block tables).

Root cause: vLLM's breakable cudagraph replaces ATOM's PIECEWISE graphs. With VLLM_USE_BREAKABLE_CUDAGRAPH=1, VllmConfig.__post_init__ switches the compilation mode to NONE. ATOM's V4 is compiled with ATOM's own torch.compile and splits the graph at the V4 attention op (v4_attention_with_output). Attention therefore runs eagerly between the PIECEWISE graphs and reads the current batch's metadata. With mode NONE, BreakableCUDAGraphWrapper captures the whole forward instead. The V4 attention has no @eager_break_during_capture (in the plugin only the MLA path, Kimi-K3 and MiniMax-M3 have one), so it is recorded with the capture batch's metadata and replayed with it.

Fix: prefilling rows are never uniform decode (3cf6a1d5c). For a runner whose model is ATOM's V4 (decided once per runner from model_config and the registry), the PR wraps V1's GPUModelRunner._determine_batch_execution_and_padding. On real steps (force_uniform_decode unset), if input_batch.num_computed_tokens_cpu[:num_reqs] < input_batch.num_prompt_tokens[:num_reqs] for any row, the wrapper passes force_uniform_decode=False. _dummy_run and capture pass a bool and are left untouched. The arrays are the NumPy views of the tensors _build_attention_metadata uses, and the predicate is the same as is_prefilling, which _is_pure_uniform_decode uses to reject DECODE. So dispatch and the bridge agree on every step. The rule does not depend on the decode length, so MTP (1 + k) is covered. Pure decode is unchanged.

Guard (always on). The wrapper records the step's dispatched cudagraph mode. A _build_attention_metadata wrapper exposes it only while that step's target metadata is built (drafter builds and dummy runs see nothing). AtomDeepseekV4ProxyMetadataBuilder raises RuntimeError if vLLM dispatched FULL for a step it did not build as DECODE. This uses CPU scalars only, with no sync. Per-process counters (dispatched modes, overridden steps, guard trips) are written only when ATOM_V4_CG_INSTR_DIR is set; it is off by default and was used for the numbers below.

Fix: no breakable cudagraph for ATOM's V4 (135f5f7df). register_model (the general plugin) wraps VllmConfig.__post_init__. It runs in the API server before VllmConfig is built, and EngineCore and the workers inherit the environment. The wrapper acts before vLLM's code runs, when the plugin is enabled, model_config.architectures has DeepseekV4ForCausalLM or DeepSeekV4MTPModel, and vLLM's ModelRegistry resolves it to an ATOM class:

• VLLM_USE_BREAKABLE_CUDAGRAPH unset: set it to 0 and log [atom-v4-cg] DeepseekV4ForCausalLM is served by ATOM: setting VLLM_USE_BREAKABLE_CUDAGRAPH=0 …;
• explicitly 1: startup fails with ValueError: VLLM_USE_BREAKABLE_CUDAGRAPH=1 is not supported for DeepseekV4ForCausalLM served by ATOM: …, before any engine process starts. Pydantic reports it as ValidationError … Value error, …;
• explicitly 0: kept.

Every other architecture is untouched. Kimi-K3 and MiniMax-M3, whose plugins add eager breaks, keep vLLM's behaviour, including an explicit 1. On vLLM 0.31 / ROCm the unset path is a no-op: vLLM's ROCm default list holds only DeepseekV41ForCausalLM, and the startup logs of main and this branch differ only by the [atom-v4-cg] line. The commit is kept for:

• an explicit =1 (vLLM's V4 recipe);
• VLLM_BATCH_INVARIANT=1;
• non-ROCm platforms, where DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES still lists V4;
• pins that put V4 back on the ROCm list.

Model Runner V1 only (and vllm#51865)

3cf6a1d5c wraps V1's runner, so it only acts where V4 runs on Model Runner V1. On vLLM 0.31 that is the default on ROCm: DeepseekV4ForCausalLM is in ROCM_DEFAULT_MRV1_ARCHITECTURES, and every server log shows Defaulting to V1 model runner on ROCm for model architectures: DeepseekV4ForCausalLM.

Upstream, vllm#51865 fixed this rule in Model Runner V2 only (get_uniform_decode_token_count(..., has_prefill)). V1's _is_uniform_decode is unchanged, and the V1 fixes (vllm#53059, vllm#47123) are open without review while V1 is deprecated. If V1 gets the fix, this commit becomes a no-op and can be dropped.

For the next vLLM upgrade: vllm#56360 (not in 0.31) removes the ROCm exception, so V4 will default to Model Runner V2 and this wrapper will not attach. V2's rule (vllm#51865 + vllm#58400) still allows FULL for a row that has context and a single prompt token left, assuming the backend builds it as a decode. ATOM's bridge builds it as PREFILL. That upgrade therefore has to either make the bridge build these rows as decode or pin VLLM_USE_V2_MODEL_RUNNER=0.

3. Prefix hits resume from a wrong state

Problem. To resume a DeepSeek-V4 request at token b, two things are needed: the paged history (compressed KV and indexer keys, in 128-token pages) and the request's state slot. The state slot holds the sliding-window ring (KV of the last 128 tokens) and the compressor tails (partially filled compression windows). Native ATOM stores both: at chosen boundaries it takes a slot checkpoint (PagedStateCheckpointSpec), and on a hit it restores that checkpoint into the new request's slot.

On the vLLM plugin only the pages are cached. The slot lives in a per-request area of the proxy layer that no vLLM block covers, so a hit starts from a freshly reset slot with the ring and compressor tails all zero. Every token after the boundary attends to zero ring rows instead of the previous 128 tokens, and the compressors restart mid-window from zeros. The pages these tokens write are then hashed and reused by later hits.

deepseek_v4_prefix_patch (#1454) rolls back 128 tokens and recomputes them. That is also wrong, because the recompute starts from the same reset slot. It is not active on main anyway: it is installed from ATOMPlatform.check_and_update_config, and on vLLM 0.31 register_platform() runs but the returned platform is not adopted (current_platform stays RocmPlatform).

The effect is visible directly. In part 2's shape probe on part 2's head (zero slot), the prefix-hit re-send of a 128·m + 1-token prompt returns a different first token than the first send of the same prompt in 9 / 14 cells; miss vs first send differs in 0 / 14. For example, a repeated sentence continues with Ir. The quick brown fox … (log-prob −3.07) instead of quick brown fox … (−8e-6).

End-task accuracy does not show this. GSM8K with zero-state hits, with prefix caching off, and with this fix are all within run-to-run noise. The error is confined to the sliding window right after each boundary, while the compressed history, which carries most of the context, is cached correctly. Teacher-forced log-probabilities after the boundary do separate them: the zero state deviates from a miss by 2.3–3.1× the miss-vs-miss noise, this fix by 1.00×. GSM8K only reacts to a grossly wrong state: resuming from the state 128 tokens behind the boundary drops it to 0.35.

Approach: native checkpoints, stored in vLLM blocks. The plugin reuses native ATOM's PAGE-backed state checkpoints (PagedStateCheckpointSpec, execute_paged_state_copies, placement as in BlockManager):

• The active state stays in ATOM's fixed slot, so kernel addressing is unchanged.
• At the boundaries native ATOM checkpoints at, the slot's checkpoint image is copied into k vLLM blocks. Those boundaries are the 8,192-token interval rungs, the prompt-end anchor, the demand rung, and decode keeps under native's spacing rule; at most two images per request and one per boundary. vLLM manages the k blocks as k single-layer MambaSpec(mamba_cache_mode="align") groups next to the proxy. For V4-Flash with fp8 KV the image is 5,324,800 bytes, k = 11 at block 128. vLLM files the blocks under the prefix hash, evicts them block by block (its hybrid-Mamba semantics), and reports a hit only when the proxy and image blocks are all cached.
• On a hit, the k image blocks come back with the proxy blocks and are restored into the new request's slot before its forward, by native ATOM's own copy (execute_paged_state_copies). The image is therefore cut at the same byte ranges, through the same descriptor kernel, as a native checkpoint.
• vLLM-side byte operations on these blocks are off: the align pre-copy copies zero bytes, new-block zeroing is off, and CoW copies are refused. A block id names the same unit to vLLM and ATOM, but not the same bytes.
• Slots are released only on vLLM's finished, preempted and resumed events (_V4StateSlotAllocator(evict_lru=False) + release), so a slot with a pending store is never handed out.

Generic layer + per-family adapter.

• 4601b0207: paged_state_image.py and paged_state_image_scheduler.py hold the mechanism and import no model code. They cover the image layers and groups, the get_kv_cache_configs hook and its self-check, withholding the slot area from the pool, the runner patches, the store / restore order, and native placement on vLLM's scheduler (hash, withdraw, dedupe, two-image cap, decode-keep hash deferred under async scheduling). No adapter is registered in this commit, so every hook is a pass-through.
• PagedStateImageAdapter is what a model family supplies: model match, proxy layer name, block size, sizing from the native PagedStateCheckpointSpec (k = units_per_checkpoint), image layer names, placement defaults, and the copier.
• e7b91038a: deepseek_v4_image.py is the V4 adapter, plus the bridge, model_wrapper, register, platform and req_id_passthrough_patch wiring.

The interface was also checked against DeepSeek-V4.1 (#2494: block 256, image = the whole STATE slot including the Engram cursor, native StateCopies as the copier). It fits without interface changes; the V4.1 adapter will follow once #2494 lands.

Under cudagraph. On a FULL replay vLLM never enters the model's Python forward. Stores therefore run from a GPUModelRunner._model_forward wrapper, right after the forward or the replay is enqueued, on the same stream and outside any capture, for NONE, PIECEWISE and FULL steps. Restores and slot resets run in the metadata build, outside any graph. The image groups leave vLLM's modes unchanged (VLLM_COMPILE, FULL_AND_PIECEWISE, capture sizes 1 … 128).

vLLM 0.31 specifics this relies on.

• A KVCacheTensor is a placement inside one shared allocation, and the layer views stop at num_blocks. The group self-check therefore verifies that the proxy and every image group start at byte 0 with block stride = page, and the V4 carve takes the slot tail right after the view.
• The image groups' cache_blocks accepts and ignores 0.31's replay_boundaries, because the image path decides its own boundaries.
• The model provides get_mamba_state_copy_funcs(mamba_types), which 0.31's align path asks for (a zero-byte copy).
• The image backend implements get_supported_kernel_block_sizes(kv_cache_spec=None), and the proxy is read through _proxy_block_major_view.

The MTP draft proxy carves its slots from its own tail. The slot tail is withheld from the whole KV-cache config, and vLLM 0.31 places the target and the MTP draft proxies one after the other, each with its own tail after its view. Only the target used to carve its slots there; the draft carved them inside its view, which no longer has room for them now that the slots are not amortized into every block, so MTP with --no-enable-prefix-caching failed at startup (DeepSeek V4 proxy cache too small). Every proxy now carves its view's blocks plus its own slots from its own tail, sized from its own layers, and refuses to start if they would exceed what the pool withheld. The target's carve is unchanged.

Behaviour changes

• Images are on by default for DeepseekV4ForCausalLM when prefix caching is on (vLLM's default), and only when vLLM's registry serves that architecture with ATOM's model and its proxy layer is in the KV-cache specs. Other models, including vLLM's own V4 and V4.1, are unchanged. V4 with --no-enable-prefix-caching keeps no images but uses the same KV capacity accounting (next item).
• KV capacity. The proxy is now priced at its real per-block bytes, and the per-request slot area is withheld from the pool as a tail instead of being amortized into every block. V4-Flash, TP4, fp8 KV, same memory: 45,513 blocks on main (GPU KV cache size: 5,825,664 tokens) → 402,305 schedulable blocks.
• --speculative-config together with prefix caching is refused at startup, with a message pointing at --no-enable-prefix-caching (see follow-ups). MTP with --no-enable-prefix-caching runs, without images.
• VLLM_USE_BREAKABLE_CUDAGRAPH=1 with ATOM's V4 is refused at startup.
• Minimum pool. Startup requires one max-length request's peak: ceil(max_model_len / 128) proxy blocks, two blocks per image group and vLLM's null block (279 blocks for V4-Flash at 32,768 tokens).
• --max-model-len -1. The automatically selected length leaves room for the withheld slot area: vLLM's auto-fit runs again on the schedulable blocks, and the lower length reaches the workers and the API server as usual. Under tight memory the length is therefore shorter than vLLM alone would pick.
• ATOM_STATE_CHECKPOINT_INTERVAL_TOKENS=0 publishes no image. vLLM reports a hit only when every image group is cached, so there is no prefix hit at all (never a proxy-only, zero-state resume); the image groups still allocate their working blocks. -1 keeps the prompt-end anchor and the demand rung without interval rungs.
• A decode keep is published as soon as its boundary hash exists. During generation that is the step after it is stored (under async scheduling the token closing its block is still in flight when the keep is planned), so a follow-up sent while the request is still generating can resume from it, as with native ATOM. At a request's end it is published when its boundary hash exists and dropped otherwise: a natural end keeps it (so a follow-up turn can resume there), and an abort whose closing token never reached the CPU drops it. That matches native ATOM, which keeps a decode checkpoint on a request's last step too (a PAGE checkpoint's successor_room is 0) and skips an aborted request's hashing; the one difference left is engine scheduling, not placement: when max_tokens ends the answer one token short of a boundary, native's deferred output forwards the last token once more and keeps that boundary, while vLLM schedules no forward past max_tokens, so the state never exists.
• A metadata build whose forward never ran (not reachable with the pinned vLLM) stops the engine with an error instead of caching a stale image.

Tests

CPU (new):

• test_vllm_deepseek_v4_indexer_split.py: a prefilling row of decode length ends the decode group; no dense call when no prefill rows are left; mixed batches keep their split.
• test_vllm_deepseek_v4_breakable_opt_out.py (18): unset → 0 before vLLM decides; explicit 1 → ValueError for both V4 architectures; other architectures, vLLM's own V4 class and a disabled plugin are untouched; explicit 0 is kept; idempotent; the real pydantic VllmConfig honours the class-level wrap.
• test_vllm_deepseek_v4_prefill_not_uniform.py (18): covers the 1-token tail, mixed batches with decodes, and the MTP 2-token tail (not uniform); pure decode, the first decode after the prompt, MTP verify, non-V4 runners, capture / dummy calls and stale padded rows (untouched); the guard; the counters, including a subprocess that checks no step is lost at exit.
• test_paged_state_image.py (generic layer, fake adapter with a 256-token block, no model imports): registry routing, group self-check, kv-config hook, placement at another block size, step order, stores right after the runner forward and never in a capture, a store whose forward never ran is refused rather than copied, store / restore round trip, refusal to resume from an unrestored slot.
• test_paged_state_image_scheduler.py (vLLM's real AsyncScheduler, BlockPool and get_kv_cache_configs, fake adapter, no checkpoint): abort / completion / preemption with a deferred decode keep, a deferred keep published while the request keeps running, the request's end around a decode keep checked against native's BlockManager on the same request, minimum pool ±1 with preempt / resume, interval 0, --max-model-len -1 auto-fit.
• test_vllm_deepseek_v4_image.py: event-driven slots, sizing, architecture matching (vLLM's V4.1 and V4 left alone), the MTP draft proxy carving its slots from its own withheld tail, GPU slot ↔ PAGE round trip through the plugin carve, native placement ladder, partial eviction, and the V4 adapter spec.

Suites (vLLM 0.31 image):

• Plugin (tests/plugin without sglang), on main 661ef7661: main 222 passed / 62 failed → this branch 308 passed / 61 failed (plus, on both, one collection error: main's test_rtpllm_qwen4_exp_config.py needs rtp_llm, which the vLLM image does not have). All 61 also fail on main (offload connector, proxy-state arena layout, rtpllm, MHA backend, kv-cache layout, Kimi-K3). The 62nd main failure (test_rocm_skips_temporary_cudagraph_capture) is test-order dependent and passed in this run. The tests this PR adds or changes all pass, together with tests/test_envs.py and tests/test_sparse_indexer_fp4.py (233).
• Native: main 11,558 passed / 27 failed → this branch 11,559 / 26; no new failures (the extra main failure, test_gc_utils.py::test_freezing_twice_is_additive_and_harmless, is in code this PR does not touch and passed in this run).

CPU equivalence of part 3's placement against native ATOM's BlockManager: 33 scenarios with sync scheduling and the same 33 with async scheduling, 0 mismatches, including a request ending at, before and after a decode keep, aborts around it, and a follow-up arriving while the first request is still generating past its decode keep (both hit the decode keep). One async case differs for an engine-scheduling reason (max_tokens one token short of the boundary; see Behaviour changes) and is listed as such. Splice simulations on vLLM's real scheduler: 2,469 stores, 739 restores, 0 wrong. The off-by-one negative control is judged wrong 239 / 239. Outputs are identical to the same code on vLLM 0.28.

GPU validation

The numbers below were measured on the same changes on main a0527c9db (before the rebases, the MTP draft-proxy fix and the review fixes). The head 9a4b01b90, on main 661ef7661, was checked again (the head 087ca8e53 adds only a test; its state-level smoke repeated the numbers below: 300 / 300, decode-keep 48 / 48, guard 0). That main adds #2556 (FP4 decode bounds) and #2376 (V4.1 DPA / prefill TBO); the plugin runs the FP8 indexer and vLLM's own V4.1 does not run ATOM's V4.1 code, so neither reaches it. The checks on the head were part 3's state-level smoke (300 / 300 restores right in all 11 groups, decode-keep images 48 / 48, 1,382 stores run right after the forward, 0 deferred or left over, 0 in a capture, guard 0, no errors in the log); the small pool (656 / 656 restores right, including 104 / 104 cross-group splices and 160 / 160 preemption resumes; 2,007 stores, 0 deferred, 0 in a capture; guard 0; hot-prefix GSM8K 0.9583 / 0.9483); part 2's cudagraph checks (shape probe, 108 cells × 4 samples: 0 cells differ between default and E; prefix hits vs misses 0 flags; 1-token-tail prefix hits match the first send in 12 / 12 cells on both; mixed load: 37 / 360 garbage vs E's 35 / 360, Fisher p = 0.90, 38 steps re-routed, guard 0); one GSM8K pass (0.9431 / 0.9371, against 0.9426 / 0.9355 pooled over 4 passes before); the abort stress below (64 / 64 cancelled, 0 errors, served afterwards); MTP with --no-enable-prefix-caching (probe subset 48 / 48, acceptance 0.884, guard 0), with MTP plus prefix caching still refused at startup; and a live follow-up (three 1,000-token prompts each streaming 9,000 tokens; at 8,300 generated tokens, while they are still running, a request over the prompt, those tokens and 200 new ones is sent): each follow-up hits the decode keep at 9,088 (cached_tokens 9,088 of 9,500) and its 3 / 3 restores are judged right, with the first requests at 8,308–8,321 tokens when their follow-ups finish. The first round of review fixes was also checked end to end, on the heads before and after it (4ccd6bd84 and ced271c74, before the rebase onto 661ef7661); the abort stress was repeated on 9a4b01b90 with the same result:

check	before the review fixes	after them
64 streaming 1-token-prompt decodes cancelled at 8,180–8,200 tokens (around the 8,192 decode keep), with background traffic	EngineCore died: IndexError in BlockPool.cache_full_blocks, raised from EngineCore._process_aborts_queue	64 / 64 cancelled, 0 errors; afterwards healthy, a new request answered; 0 Traceback / EngineDeadError
--max-model-len -1 --kv-cache-memory-bytes 1206986240 (2,270 physical blocks)	refused at startup: vLLM picked 287,616 tokens, then plugin pool too small … 230 left	starts; vLLM's log: auto-fit 287,616 → 26,496; /v1/models reports 26,496; a 26,336-token prompt is answered
vLLM's own DeepseekV41ForCausalLM (V4.1-Flash) with the plugin installed, prefix caching on, --enforce-eager	ValueError: unknown V4 compress ratios [1, 2] in this PR's KV-config hook	the hook passes it through; startup then stops in main's spec_decode_patch.py (see Known issues)

DeepSeek-V4-Flash, TP4, fp8 KV, prefix caching on, --max-num-seqs 64 --max-model-len 32768 --max-num-batched-tokens 8192, async scheduling, rocm/atom-dev:vllm-v0.31.0-nightly_20261012. Servers:

• main = a0527c9db;
• "part 1 only" = 92a9ab671 (now 0b8ea573d);
• "parts 1+2" = 2d9519f2c (now 3cf6a1d5c);
• this branch = c1c950dd7, the five commits on a0527c9db (rebased as 0b8ea573d … e7b91038a; the rebase changed no plugin line, and the later fix only touches the MTP draft proxy's carve);
• E = --enforce-eager.

Startup on every default server: VLLM_COMPILE, FULL_AND_PIECEWISE, capture sizes 1 … 128, [atom] Dynamo bytecode transform on all 4 workers, and no breakable cudagraph. Output is nondeterministic run to run even on eager, so each server was sampled 4× and compared statistically: first-token log-probability per cell against eager's 4 samples, and the mixed-load garbage rate against eager standalone references (Fisher's exact test).

Part 1 (V4-Pro-0813, TP8, from #2501):

	main	this fix
1025-token prompt sent twice (--enforce-eager)	2nd request: GPU memory access fault, EngineDeadError	answered (1024 tokens from cache)
shape probe, 120 requests (20 prompt lengths 1 … 8321 × first send / re-send / fresh × 2 texts)	–	eager and default cudagraph: 120 / 120 answered, no fault
GSM8K 20-shot, default cudagraph	0.956 ± 0.0056	0.9545 ± 0.0057

Part 2 (V4-Flash; measured on parts 1+2):

check	main default	part 1 only, default	parts 1+2, default	E
shape probe, 120 cells: cells whose first-token log-prob differs from E	18 (1-token-tail hits 10/14, 1-token prompts 6/6, 8193 2/2)	18 (same cells)	0	independent eager sample: 1
same probe with VLLM_USE_BREAKABLE_CUDAGRAPH=1	70	–	refuses to start	–
mixed batches under 32 running decodes (360 requests: hit tails 1/2/8/64/104/127, fresh 1–128-token prompts): garbage	–	112 / 360 + 44 × HTTP 400	38 / 360	41 / 360 (Fisher p = 0.81)
step counters under that load	–	–	FULL 2,009 / PIECEWISE 65 / NONE 102; 37 steps re-routed; guard 0	–
GSM8K gsm8k_cot, 64 concurrent, flexible / strict, 4 passes pooled	–	–	0.9278 / 0.9270	0.9299 / 0.9295
MTP k=1 (with a test-only shim for the main issue below)	–	–	probe subset 48 / 48; acceptance 0.885; 57 steps re-routed; guard 0	–
decode TPOT, 1 request / 16 concurrent (2 rounds, GPU groups swapped)	8.69 / 10.52 ms	–	8.99 / 10.71 ms (1.035× / 1.018×)	–

Part 3 (V4-Flash; this branch). Restores were checked at the state level on an instrumented server, for every image group of every restore: block position, bytes unchanged, slot == unit after the copy, hash.

check	E	default cudagraph
restores right in all 11 groups (L1 probe + multi-turn workloads)	300 / 300	300 / 300
of which decode-keep images (stored after FULL replays on default)	48 / 48	48 / 48
zero ring rows read / resumes without a restore	0 / 0	0 / 0
stores planned = run right after the forward / deferred	1,382 = 1,382 / 0	1,382 = 1,382 (FULL 35, PIECEWISE 10, NONE 1,337) / 0
copies inside a capture / align pre-copy bytes	0 / 0 B	0 / 0 B
part 2's guard	–	0 (FULL 26,180 / PIECEWISE 520 / NONE 3,252 steps)
teacher-forced log-prob deviation after the boundary vs miss, over miss-vs-miss noise	1.00×	0.996×

check	result
small pool (--num-gpu-blocks-override 2850 → 830 schedulable blocks), multi-turn + hot-prefix GSM8K, default cudagraph	651 / 651 restores right, incl. 177 / 177 cross-group splices and 124 / 124 preemption resumes (605 preemptions, 108,885 evicted blocks); 1,915 stores, 0 deferred, 0 in a capture; guard 0; hot-prefix GSM8K 0.9617 / 0.9567
shape probe, 108 cells × 4 samples: default vs E	0 / 108 cells differ; prefix hit vs miss: 0 (E), 0 (default)
1-token-tail prefix hit vs first send (first token; part 2's 20 prompt lengths, 14 cells × 4 samples)	0 / 14 (E), 0 / 14 (default); parts 1+2 without images: 9 / 14 (miss vs first send: 0 / 14 everywhere)
2–128-token-tail prefix hit vs first send	0 / 16 (E), 0 / 16 (default); parts 1+2: 3 / 16
mixed load (as in part 2) vs E	34 / 360 garbage vs E under the same load 35 / 360 (Fisher p = 1.0; a second default run: 38 / 360, p = 0.81); FULL 2,024 / PIECEWISE 94 / NONE 122 steps; 44 steps re-routed; guard 0
GSM8K gsm8k_cot vs prefix caching off, 4 + 4 passes pooled, flexible / strict	images 0.9426 / 0.9355 vs off 0.9410 / 0.9335: +0.15 / +0.21 pt (on a second node: images 0.9450 / 0.9386 over 4 passes, off 0.9431 / 0.9375 over 2)
decode TPOT vs prefix caching off, 1 request / 16 concurrent (2 rounds, GPU groups swapped)	8.957 vs 8.839 ms (1.013×) / 10.825 vs 10.792 ms (1.003×)
MTP k=1 with --no-enable-prefix-caching (rebased head, with the main xdrope fix below applied for the test)	starts; probe subset 48 / 48 answered; acceptance 0.885 overall (0.855 in the probe window); mixed load 0 errors, 50 steps re-routed, guard 0
MTP k=1 with prefix caching on	refused at startup: ValueError: DeepSeek-V4 prefix caching through checkpoint images does not drive the MTP draft proxy; drop --speculative-config or pass --no-enable-prefix-caching.
errors in the server logs (memory access fault, EngineDeadError, Traceback, RuntimeError, guard, HTTP 500)	0 in all 13 servers that served traffic (eager, default, small pool, prefix caching off, MTP). The only errors are the intended MTP startup refusal above, and the MTP startup failure of the build before the draft-proxy fix.

Discrimination check, run earlier on vLLM 0.28 with test-only negative-control arms (not part of this PR):

• zero-state arm: 96 / 96 (eager) and 338 / 338 (cudagraph) restores judged wrong; log-prob deviation 3.1× (eager) and 2.3–2.8× (cudagraph);
• 128-tokens-behind arm: 93 / 93 and 743 / 743 restores judged wrong; log-prob deviation 14.3× and 8.2–15.1×; hot-prefix GSM8K 0.347 vs 0.963 with images.

V4-Flash, V4-Flash-0731, V4-Pro and V4-Pro-0813 pass sizing (equal to native's own checkpoint spec: k = 11, 11, 10, 11) and binding on CPU.

Known issues and follow-ups

• Speculative decoding with images. --speculative-config with prefix caching is refused at startup; with --no-enable-prefix-caching it runs without images. Support for images under MTP is a follow-up.
• MTP is broken on main with vLLM 0.31, independent of this PR: atom/plugin/vllm/spec_decode_patch.py reads uses_xdrope_dim / draft_uses_xdrope_dim, which vLLM 0.31's proposer no longer has. The first drafted step raises AttributeError. The MTP row in part 2 used a test-only shim that sets both to 0, and the MTP rows in part 3 applied the upcoming fix to the test build; the fix will be a separate PR.
• Known cost. Part 2: decode TPOT is 1.035× main for a single request and 1.018× at 16 concurrent decodes, over two rounds with the GPU groups swapped. This is the per-step cost of the code it adds: a ≤ 64-element NumPy comparison and two Python wrappers on the runner, plus a CPU-scalar check in the builder. Part 3 vs prefix caching off: 1.013× for a single request and 1.003× at 16 concurrent decodes, over two rounds with the GPU groups swapped (1.02× / 1.00× on vLLM 0.28).
• vLLM's own V4.1 does not start with the plugin installed, independent of this PR: main's atom/plugin/vllm/spec_decode_patch.py (patched_get_kv_cache_groups) takes V4.1's KV specs for a heterogeneous MLA / MHA (EAGLE3) layout and builds groups that fail vLLM's All attention layers in the same KV cache group must be SlidingWindowMLASpec assertion (during cudagraph memory profiling, or in get_kv_cache_configs with --enforce-eager). It fails the same way before and after this PR's changes.
• Copier interface. _PluginV4StateCopier borrows native DeepseekV4AttentionMetadataBuilder methods (_attach_native()); an explicit shared layout / copy component would make native changes less likely to break it. Planned as a follow-up, together with the V4.1 adapter, which needs the same component.
• After merge: coordinate with feat(plugin): DeepSeek-V4.1-Flash on the vLLM plugin backend #2494, which moves _V4StateSlotAllocator into a shared state_slot_allocator.py with LRU eviction. This PR needs evict_lru=False and event-driven release there.

Not in this PR

The DeepSeek-V4.1 adapter (after #2494), images under speculative decoding, Model Runner V2.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2502 --add-label <label>

[Copilot]

🔵 Needs a closer look

It changes high-risk CUDA-graph dispatch through runtime patches, and the diagnostic counter persistence can report stale totals.

1 open finding

• Flush diagnostic counters at shutdown to prevent undercounting · New

What changed in this PR

Fixes DeepSeek-V4 output corruption in the vLLM plugin under CUDA-graph execution.

Changes:

• Disables unsupported breakable CUDA graphs for ATOM V4.
• Routes still-prefilling rows away from FULL decode graphs.
• Corrects indexer splitting and adds regression coverage.

File	Description
atom/​plugin/​vllm/​deepseek_v4_cudagraph_patch.py	Adds CUDA-graph dispatch safeguards and instrumentation.
atom/​plugin/​vllm/​deepseek_v4_bridge.py	Aligns metadata and indexer classification with prefill state.
atom/​plugin/​vllm/​models/​deepseek_v4.py	Avoids dense scoring on empty prefill slices.
atom/​plugin/​vllm/​register.py	Installs the V4 compatibility patches.
tests/​plugin/​test_vllm_deepseek_v4_breakable_opt_out.py	Tests breakable-graph opt-out behavior.
tests/​plugin/​test_vllm_deepseek_v4_indexer_split.py	Tests safe indexer splitting.
tests/​plugin/​test_vllm_deepseek_v4_prefill_not_uniform.py	Tests prefill dispatch and consistency guards.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🔵 Needs a closer look

The diff includes an undocumented scheduler/cache redesign and conflicts with the stated MTP validation setup.

3 open findings

• Prefix caching with MTP is rejected despite claimed MTP coverage · New
• Checkpoint-image redesign is out of scope and insufficiently validated · New
• Flush diagnostic counters at shutdown to prevent undercounting

🧠 Review effort: Balanced

[Copilot]

🟡 Changes recommended

The image adapter currently alters non-ATOM DeepSeek-V4 configurations, and new environment variables are not registered or documented.

6 open findings

• Restrict adapter matching to configurations resolving to ATOM models · New
• Checkpoint-image redesign is out of scope and insufficiently validated
• Prefix caching with MTP is rejected despite claimed MTP coverage
• Flush diagnostic counters at shutdown to prevent undercounting
• Register and document ATOM_V4_CG_INSTR_DIR · New
• Register and document ATOM_STATE_CHECKPOINT_INTERVAL_TOKENS · New

🧠 Review effort: Balanced

[PerryZhang01]

Review of 53eff4452b5d

The prefill/indexer and CUDA-graph dispatch fixes make sense, and restoring the window/compressor state is necessary for correct prefix hits. However, I would hold off merging the checkpoint-image path until the lifecycle and capacity issues below are addressed.

This review was initially performed against 356df49878c1 and refreshed against 53eff4452b5d27d6172452cce28d092f08b3bc4e, using the pinned vLLM commit db9527a46873454610df6dbedf79a36d6bf1a7f6. The code underlying the four findings below is unchanged. I also checked the new draft-proxy carve fix and the updated validation section: I am not carrying forward the earlier draft-tail concern, or claiming that Part 3 mixed-load/GSM8K/TPOT results are still pending. The MTP-plus-images restriction and the separate baseline xdrope issue are now clearly documented.

1. Blocking: aborting an async decode can crash EngineCore while publishing a pending image hash

Source: paged_state_image_scheduler.py:361–373

if request.is_finished() and pend is not None:
    self._hash_block(request, pend)

is_finished() includes FINISHED_ABORTED; it does not establish that the last token/hash of the pending image has reached the CPU.

A concrete reachable state with the default checkpoint interval is:

• a one-token prompt followed by async decode;
• num_computed_tokens == 8192, but the CPU currently knows only num_tokens == 8191;
• 63 complete prefix-block hashes exist, and pending_hash[rid] == 63;
• the client disconnects or cancels the request.

on_free() calls _hash_block() for block index 63, which calls upstream BlockPool.cache_full_blocks(). Its new_block_hashes[i] access raises IndexError because that hash does not exist yet.

The timing is supported by the actual upstream ordering, not just an assumed race: EngineCore processes aborts before updating the request from the model output. Without a KV connector requiring deferred frees, the normal free path reaches this hook immediately. Under the default fault-handling path, the exception propagates to EngineCore's fatal-error handler and shuts the engine down; this is not merely an error returned to the cancelled request.

Verification: an isolated CPU harness using the actual pinned scheduler/free/block-pool methods and the PR's image methods reproduced:

Scheduler.finish_requests(FINISHED_ABORTED)
  -> _free_request -> _free_blocks -> _free_request_blocks
  -> on_free -> _hash_block -> BlockPool.cache_full_blocks
  -> IndexError: list index out of range

Controls passed for normal completion with the hash available, and preemption with the pending hash discarded.

Please validate that the complete boundary hash is available before publishing it, accounting for the upstream hash granularity. If it is not available on cancellation, discard the pending publication and release the request normally. Add regression coverage for cancellation, natural completion, and preemption.

2. Blocking: the V4 adapter matches native V4.1 and fails before the original allocator runs

Matcher: deepseek_v4_image.py:62–66 · Hook: paged_state_image.py:304–313

This is a concrete additional consequence of the adapter-scope issue already raised in the inline review. The adapter reuses the platform predicate:

any("DeepseekV4" in str(a) for a in arches)

Consequently, DeepseekV41ForCausalLM also matches. The new globally installed KV-config hook calls adapter.sizing(vllm_config) before checking whether any actual ATOM proxy spec/group exists.

The pinned vLLM has a native V4.1 implementation, whose valid compression ratios include 1 and 2. Passing those through V4 sizing reaches UnifiedPoolGeometry, which accepts only -1/0/4/128, and raises:

ValueError: unknown V4 compress ratios [1, 2]

This happens with prefix caching both enabled and disabled. ATOM does not override the V4.1 registry entry in this revision, so excluding the V4.1 adapter from this PR's scope is not sufficient: the generic hook must leave that native model alone.

The substring helper predates this PR, but its use in this unconditional image-sizing hook is new. The baseline platform constraint did not construct V4 pool geometry for these configurations.

Verification: source-extracted CPU probes using the actual matcher, sizing entry point, and geometry reproduced the exception for both prefix-cache settings, with original_called=False.

Please require an exact supported architecture and an actual ATOM registry resolution, as the graph patch already does. Also check for the expected ATOM proxy in the input specs before sizing or mutating configuration. Add a native V4.1 no-op test.

3. Capacity: slot reservation happens after upstream auto-fit, so memory-limited --max-model-len -1 can fail to start

Source: paged_state_image.py:323–364

The hook first calls the original allocator, which completes automatic max-model-length selection, and only afterward reduces cfg.num_blocks using sizing.usable_blocks().

The upstream auto-fit budget includes the proxy/image specs and the null-block reservation, but not the fixed slot tail that this PR removed from the per-page price. When memory limits the selected length, the hook then rejects the length the original allocator just selected instead of fitting a shorter one.

For the default two-image peak, let N be the physical block count, R the withheld slot reserve, and k the image group count. Upstream can select:

ceil(L / block_size) + 2*k = N - 1

The hook subsequently compares that demand against N - R. With R > 1, this selection can fail immediately.

Verification: an isolated CPU probe used the pinned auto-fit/binary-search/group-memory methods and the PR's generic fake-adapter scale: 100 physical blocks, 11 reserved blocks, block size 256, and three image groups. Upstream selected 23,808 tokens requiring 99 blocks; the hook then reduced capacity to 89 and raised plugin pool too small. These are toy-adapter values, not a V4 GPU measurement; the ordering defect does not depend on k == 3.

Please include slot reservation in the effective auto-fit budget, or re-fit after determining the schedulable count and propagate the resulting length through the normal EngineCore path. Preserve the physical backing bytes needed for the slots; do not simply shrink the allocation and reserve the same bytes again. Add a test in which the real upstream auto-fit actually reduces the length.

4. Capacity: the minimum-pool check omits the null block, and interval zero does not eliminate image working blocks

Source: paged_state_image.py:351–354

need = per_req + images_per_req * k
if usable < need:
    ...

The value assigned to cfg.num_blocks still includes the null block permanently held by vLLM. The new check counts it as allocatable.

A concrete boundary case using the actual pinned allocation methods:

• max_model_len=32768, block size 128, k=11, default interval 8192;
• post-slot-reservation cfg.num_blocks=278, which passes 256 + 2*11;
• a 32,767-token prompt reaches 32,640 computed tokens;
• the last 127-token chunk needs one proxy block plus eleven image blocks, but only eleven blocks are free;
• allocation returns None.

Changing only the pool count to 279 makes the same path succeed. Freeing/preempting and trying to resume the cached 32,640-token anchor still fails with 278 total blocks, so this is not just transient contention.

There is a larger accounting discrepancy when ATOM_STATE_CHECKPOINT_INTERVAL_TOKENS=0: images_per_request() returns zero, but the Mamba image groups remain registered and still allocate running/working blocks. A CPU probe of the first 1,000-token forward produced no store plan but held eleven image blocks; enabling ATOM_STATE_IMAGE_CHECK_BLOCKS then immediately raises because the configured limit is zero.

Please base startup capacity and diagnostics on the actual allocation peak, including the null block and working blocks even when no checkpoint is published. If interval zero is intended to remove image allocation entirely, that must be reflected in the allocation path without permitting unsafe proxy-only prefix hits. Add minimum-capacity ±1, interval-zero, and preempt/resume tests.

Smaller follow-ups

• Diagnostic persistence: I reproduced the existing counter-flush finding. A normally exiting subprocess with 127 non-flipped steps leaves a file recording only 64. This affects total steps/mode counts; flipped steps and guard trips flush immediately, so this does not establish that those counters lose the same tail. Please add a normal worker-teardown flush or write each step while the opt-in diagnostic is enabled. A shutdown hook cannot guarantee persistence on forced termination.
• Copier interface: _attach_native() transfers a list of private methods from the native builder and manually recreates their required state. I did not confirm a failure here, but a shared explicit layout/copy component would make future native changes less likely to break the adapter implicitly. Please consider this as a follow-up rather than a prerequisite rewrite.

Verification scope

• No new GPU server or end-to-end model run was performed for this review.
• The failure probes are isolated CPU harnesses using source-extracted methods; heavyweight runtime/config dependencies are replaced with minimal fixtures. They do not constitute a full vLLM deployment test.
• The changed-test subset on the originally reviewed head returned 38 passed, 5 skipped. Missing vLLM/AITER dependencies caused the skips, including the image suites; those are not reported as passing integration coverage.
• The current image scheduler, generic allocator hook, adapter, native geometry, and vLLM pin were checked against the original review and are unchanged. The latest draft carve and regression test were inspected separately; I am not reporting that earlier concern as open.
• Current CI shows passing style/non-GPU/smoke checks, while the heavy accuracy jobs are skipped. I acknowledge the expanded GPU results in the updated PR description, including the rebased-head checks; those are author-reported results, not independent measurements from this review.

The focused missing regression coverage is lifecycle transitions and exact capacity boundaries, not another broad steady-state accuracy run. Please address the four issues above before merging the default-enabled image mechanism.

---

🤖 Generated with Claude Code

[Copilot]

🟡 Changes recommended

A leftover image store can cache stale slot state when its corresponding forward did not execute.

7 open findings

• Prevent stale pending stores from corrupting future hashed blocks · New
• Restrict adapter matching to configurations resolving to ATOM models
• Checkpoint-image redesign is out of scope and insufficiently validated
• Prefix caching with MTP is rejected despite claimed MTP coverage
• Flush diagnostic counters at shutdown to prevent undercounting
• Register and document ATOM_STATE_CHECKPOINT_INTERVAL_TOKENS
• Register and document ATOM_V4_CG_INSTR_DIR

🧠 Review effort: Balanced

[Copilot]

🔵 Needs a closer look

The extensive scheduler monkeypatching, cache lifecycle changes, and GPU state-copy path warrant final human review despite strong test coverage.

5 open findings

• Prevent stale pending stores from corrupting future hashed blocks
• Checkpoint-image redesign is out of scope and insufficiently validated
• Prefix caching with MTP is rejected despite claimed MTP coverage
• Register and document ATOM_STATE_CHECKPOINT_INTERVAL_TOKENS
• Register and document ATOM_V4_CG_INSTR_DIR

2 resolved since last review

• Restrict adapter matching to configurations resolving to ATOM models
• Flush diagnostic counters at shutdown to prevent undercounting

🧠 Review effort: Balanced

[Copilot]

🟡 Changes recommended

Deferred decode checkpoint hashes are not published while requests remain active.

5 open findings

• Deferred decode keep is lost after cache boundary advances · New
• Checkpoint-image redesign is out of scope and insufficiently validated
• Prefix caching with MTP is rejected despite claimed MTP coverage
• Register and document ATOM_STATE_CHECKPOINT_INTERVAL_TOKENS
• Register and document ATOM_V4_CG_INSTR_DIR

1 resolved since last review

• Prevent stale pending stores from corrupting future hashed blocks

🧠 Review effort: Balanced

[linsun12]

@PerryZhang01 Thanks for the careful review and the harnesses. All four findings reproduced on my side and are fixed; details per item below.

All fixes are separate commits on top of the five original commits, one per finding, each with a regression test that fails
on the tree before it and passes after it. The branch has since been rebased onto main 661ef7661 (vLLM 0.31 pin
db9527a46); every SHA below is the rebased one. Head: 9a4b01b90.

1. Aborting an async decode can crash EngineCore while publishing a pending image hash — confirmed, fixed in 27556888f

Confirmed. on_free published the deferred decode keep for every finished request, including FINISHED_ABORTED, and
EngineCore handles aborts before it applies the step's output, so the boundary hash can still be missing.

• Reproduced on CPU with vLLM's real AsyncScheduler / BlockPool and a fake adapter: a 1-token prompt stops right
  after the step that defers the keep (num_computed_tokens == 512, num_tokens == 511), then
  finish_requests(FINISHED_ABORTED) → _free_request → _free_blocks → _free_request_blocks → on_free → _hash_block → BlockPool.cache_full_blocks → IndexError, your path.
• Reproduced on GPU with the previous head e7b91038a (V4-Flash TP4, default cudagraph, async scheduling): 64 streaming
  1-token-prompt decodes cancelled at 8,180–8,200 generated tokens, with background traffic. EngineCore died with the same
  IndexError, raised from EngineCore._process_aborts_queue → abort_requests (36 EngineDeadError; the server stopped
  answering).
• Fix: a pending keep is published only if its boundary hash exists; otherwise the publication is dropped (counter
  pending_hashes_dropped) and the request is freed normally. The hash is read at the image block size through vLLM's
  resolve_block_hashes(request.block_hashes, hash_block_size, block), since hash_block_size can be finer. The two other
  places that read deferred hashes (_hash_or_defer, _image_cached) indexed request.block_hashes directly and now use
  the same view. (At hash_block_size=128 the request already has 3 fine-grained hashes but only 1 at block 256, so a plain
  length check would still have crashed.)
• Tests (tests/plugin/test_paged_state_image_scheduler.py, real scheduler): test_abort_with_a_deferred_decode_keep_frees_the_request
  (hash_block_size 256 and 128; IndexError before), test_natural_completion_publishes_the_deferred_decode_keep,
  test_preemption_drops_the_deferred_decode_keep_and_resumes.
• At a request's end the plugin now matches native exactly (checked after the fix, 0b91463b4): native also keeps a
  decode checkpoint on a request's last step (a PAGE checkpoint's successor_room is 0 and the scheduler passes
  next_forward_tokens=0 there), and an aborted seq skips the hashing postprocess. The plugin publishes a pending keep
  exactly when its boundary hash exists. test_request_end_around_the_decode_keep_matches_native runs the same request on
  vLLM's AsyncScheduler and native's BlockManager (natural end on / after the keep step, aborts after 509–512
  delivered tokens), and the CPU equivalence harness has the same shapes at V4 scale (all match; one documented
  engine-scheduling difference: a max_tokens end one token short of the boundary, where only native runs the extra forward).
• GPU after the fix (same stress; head of that round, before the rebase): 64 / 64 decodes cancelled at 8,180–8,200 tokens, 0 errors, background requests all answered (29), the server healthy afterwards and a new request answered; 0 Traceback / EngineDeadError / IndexError in the log.

2. The V4 adapter matches native V4.1 and fails before the original allocator runs — confirmed, fixed in 847f038ac

Confirmed, including Copilot's variant (V4 resolved to vLLM's own class: align forced, MTP + prefix caching refused).

• Fix: the adapter matches DeepseekV4ForCausalLM exactly and only when vLLM's registry maps it to an ATOM class
  (atom_v4_arch, the check the cudagraph opt-out uses). The KV-config hook now passes through untouched unless the
  adapter's proxy layer is in the input specs, before any sizing or config change.
• Tests: test_vllms_own_v41_is_left_to_vllm (prefix on / off; ValueError: unknown V4 compress ratios [1, 2] before),
  test_v4_served_by_vllm_itself_is_left_to_vllm (real registry), test_atoms_v4_is_matched_only_by_exact_arch_and_registry,
  test_a_matching_model_without_the_proxy_layer_passes_through.
• GPU: vLLM's own V4.1-Flash (DeepseekV41ForCausalLM) with the plugin installed, prefix caching on, --enforce-eager: before the fix, ValueError: unknown V4 compress ratios [1, 2] from this hook; after it, the hook passes the model through. Startup then stops in main's spec_decode_patch.py: patched_get_kv_cache_groups takes V4.1's specs for a heterogeneous MLA / MHA (EAGLE3) layout and vLLM's SlidingWindowMLASpec assertion fails. That happens identically before and after this PR's changes (with the default cudagraph mode even earlier, during cudagraph memory profiling), so I listed it in the description as a separate known issue rather than fixing it here.
• platform.py's _is_deepseek_v4 substring helper is on main and unchanged: it is only read by
  ATOMPlatform.check_and_update_config (not adopted on vLLM 0.31 here), where it now reaches only the pass-through install
  and main's single-forward guard.

3. Slot reservation happens after upstream auto-fit — confirmed, fixed in cf54c1a88

Confirmed with the real upstream get_kv_cache_configs and your toy scale (100 physical blocks, 11 withheld): vLLM picked
23,808 tokens and the hook refused it.

• Fix: with max_model_len=-1 the hook runs vLLM's own _auto_fit_max_model_len again after the allocator, against the
  schedulable blocks (physical minus the slot tail, minus the null block). The allocation stays as vLLM built it, so the
  slots keep their backing bytes in the withheld tail and nothing is reserved twice; EngineCore propagates a length lowered
  inside get_kv_cache_configs to the workers and the frontend.
• Test: test_auto_fit_max_model_len_leaves_room_for_the_slot_tail (real auto-fit: 23,808 → 20,992 tokens, 89 schedulable
  blocks, tensors still cover all 100 physical blocks, a request of that length is served).
• GPU (V4-Flash, --max-model-len -1 --kv-cache-memory-bytes 1206986240, i.e. 2,270 physical blocks): before (e7b91038a): refused at startup, vLLM picked 287,616 tokens and then plugin pool too small: 2270 blocks of 531712 B, 2040 withheld …, 230 left. After: starts; vLLM logs auto-fit 287,616 → 26,496, /v1/models reports max_model_len 26,496, and a 26,336-token prompt (+32 generated) is answered.

4. The minimum-pool check omits the null block, and interval zero does not eliminate image working blocks — confirmed, fixed in ba56222e6

• Fix: the startup check needs a running request's actual peak: its KV blocks, 2 blocks per image group (the running one
  and the previous one, published or not) and the null block. At V4-Flash scale that is 279 blocks instead of 278. The block
  diagnostic uses the same 2-per-group limit.
• Interval 0: I kept it as the smallest safe change. The image groups stay registered and allocate their working blocks,
  no image is published, and since vLLM reports a hybrid hit only when every group is cached there is no prefix hit at all,
  so there is never a proxy-only (zero-state) resume. Documented in docs/environment_variables.md.
• Tests (real scheduler / allocator): test_min_pool_check_counts_the_null_block (one block short is refused; was
  accepted), test_smallest_accepted_pool_serves_a_max_length_request (with and without preempt / resume from the prompt-end
  anchor: the minimum serves it, one block less does not), test_interval_zero_still_counts_working_blocks_and_never_hits
  (before: accepted at startup, then holds 3 image blocks, more than 0 on the first step).

Smaller follow-ups

• Diagnostic persistence — fixed in 021201b50: the counters file is rewritten on every step while
  ATOM_V4_CG_INSTR_DIR is set (no exit hook, as you note it cannot survive a killed worker).
  test_counters_hold_every_step_when_the_process_exits: 127 steps → 127 (was 64).
• Copier interface (_attach_native()): agreed that an explicit shared layout / copy component is the better long-term
  shape; I'd like to keep it as a follow-up PR together with the V4.1 adapter, which needs the same component.

Re-validation (head of the review-fix round, before the rebase onto 661ef7661; repeated on 9a4b01b90, see the description)

V4-Flash TP4, fp8 KV, vLLM 0.31; the previous numbers are in the description.

• State-level smoke (instrumented, default cudagraph, L1 + L2 + L3): 300 / 300 restores right in all 11 groups, decode-keep images 48 / 48, 1,382 stores run right after the forward, 0 deferred, 0 in a capture, align pre-copy 0 B, guard 0, no errors. Sizing line identical to the previous head (pool check now 279 blocks instead of 278).
• Small pool (830 schedulable blocks) with hot-prefix GSM8K: 635 / 635 restores right, including 144 / 144 cross-group splices and 111 / 111 preemption resumes (670 preemptions); guard 0.
• Shape probe, 108 cells × 4: 1 / 108 cells differ between default and eager (a 1-token fresh prompt); hit vs miss 0 / 0; 1-token-tail hit vs first send 0 / 12 (default), 1 / 12 (eager, where miss vs first send is also 1 / 12).
• Mixed load: 38 / 360 garbage vs eager's 35 / 360, Fisher p = 0.81, guard 0.
• GSM8K gsm8k_cot: 0.9416 / 0.9356 (vs 0.9426 / 0.9355 pooled over 4 passes before).
• MTP k=1 with --no-enable-prefix-caching: 48 / 48, acceptance 0.885, guard 0; MTP with prefix caching still refused at startup.
• Unit suites vs main d12e06405: plugin 292 passed / 62 failed (main 215 / 63, no new failures), native 11,035 / 26 on both.

Re-validation on 9a4b01b90 (on main 661ef7661)

• State-level smoke: 300 / 300 restores right, decode-keep 48 / 48, 0 stores left over, 0 in a capture, guard 0, no errors.
• Small pool: 656 / 656 restores right, including 104 / 104 splices and 160 / 160 preemption resumes; guard 0; hot-prefix GSM8K 0.9583 / 0.9483.
• Shape probe 0 / 108 vs eager; hit vs miss 0 / 0; 1-token-tail hit vs first send 0 / 12; mixed load 37 / 360 vs 35 / 360 (p = 0.90), guard 0; GSM8K 0.9431 / 0.9371.
• Abort stress (64 decodes cancelled at 8,180–8,200 tokens): 64 / 64, 0 errors, served afterwards.
• MTP with --no-enable-prefix-caching: 48 / 48, acceptance 0.884, guard 0; with prefix caching still refused.
• Units vs main 661ef7661: no new failures (plugin 306 / 61, native 11,559 / 26); CPU equivalence vs native placement: 30 + 30 scenarios match (one documented engine-scheduling case).

Copilot's latest comment on the deferred decode-keep hash (#2502 (comment)) came in after this round; I am looking at it separately and will follow up there.

[Copilot]

🟡 Changes recommended

Same-schedule prefix readers can restore an image before the producing forward has written it.

2 open findings

• Image prefix cache publishes state before producing forward completes · New
• Deferred decode keep is lost after cache boundary advances

4 resolved since last review

• Checkpoint-image redesign is out of scope and insufficiently validated
• Prefix caching with MTP is rejected despite claimed MTP coverage
• Register and document ATOM_STATE_CHECKPOINT_INTERVAL_TOKENS
• Register and document ATOM_V4_CG_INSTR_DIR

Previously missed (1)

In code that hasn't changed since last review

Test comment mislabels decode/verify metadata as one-token prefill

tests/​plugin/​test_vllm_deepseek_v4_indexer_split.py:84

This comment describes a prefilling one-token extend, but such a row is precisely what the preceding test verifies must leave the paged group and take the dense path. This test's metadata instead models a prefill-classified step containing only decode/verify rows. Please name that actual scenario so the regression test does not document the opposite of the fix.

🧠 Review effort: Balanced