Repository navigation
DeepSeek-V4.1-Flash: LMCache tier offload through the MP server (vLLM plugin) - #2598
Conversation
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
70733ee to
d3c4ec5
Compare
There was a problem hiding this comment.
Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit
ruff
Import block is un-sorted or un-formatted
ATOM/tests/plugin/test_vllm_v41_shutdown.py
Lines 4 to 10 in d3c4ec5
Import block is un-sorted or un-formatted
ATOM/tests/test_state_byte_codec_allocation.py
Lines 4 to 10 in d3c4ec5
pytest imported but unused
ATOM/tests/test_tree_is_importable.py
Line 17 in d3c4ec5
Call startswith once with a tuple
ATOM/tests/test_tree_is_importable.py
Line 46 in d3c4ec5
inspect imported but unused
ATOM/tests/test_tree_is_importable.py
Line 66 in d3c4ec5
Import block is un-sorted or un-formatted
ATOM/tests/test_v41_admission_needs_a_load.py
Lines 17 to 20 in d3c4ec5
Import block is un-sorted or un-formatted
ATOM/tests/test_v41_state_slot_reserve.py
Lines 13 to 16 in d3c4ec5
Do not catch blind exception: Exception
ATOM/tools/kv_offload_throughput_arms.py
Line 92 in d3c4ec5
Do not catch blind exception: Exception
ATOM/tools/kv_offload_throughput_arms.py
Line 154 in d3c4ec5
Do not catch blind exception: Exception
ATOM/tools/v41_shape_hit_probe.py
Line 48 in d3c4ec5
Do not catch blind exception: Exception
ATOM/tools/v41_shape_hit_probe.py
Line 103 in d3c4ec5
Use a context manager for opening files
ATOM/tools/v41_shape_hit_probe.py
Line 107 in d3c4ec5
Use a context manager for opening files
ATOM/tools/v41_shape_hit_probe.py
Line 110 in d3c4ec5
Do not catch blind exception: Exception
ATOM/tools/v41_shape_logprob_probe.py
Line 127 in d3c4ec5
Use a context manager for opening files
ATOM/tools/v41_shape_logprob_probe.py
Line 132 in d3c4ec5
5e4bdbb to
b1eaef0
Compare
6cc5c05 to
c8cf505
Compare
Review: V4.1 plugin LMCache MP and ATOM native impactReviewed Recommendation: address the P1 correctness/lifetime issues before merging. This is not limited to the new V4.1 plugin path: the shared PAGE-only MP worker also introduces a completion-notification latency regression in native ATOM. NativeState MP and traditional in-process LMCache need to be distinguished from that worker. This review used source/call-chain inspection, CPU tests, and targeted execution of original functions extracted through AST. It did not run GPU inference, DMA stress tests, GSM8K, or a throughput benchmark. CPU reproductions establish the control-flow defects described below, not an observed GPU accuracy regression. P1 — 1. The preemption fence misses queued/submitting transfers and fails open on timeoutLocation: mp_worker_adapter.py:135–164
A running request can therefore enqueue a PAGE save, be preempted, and have its source blocks reused before the background save reads them. The resulting bytes can be stored under the preempted request's prefix key. The PAGE leases acquired on normal request finish do not protect this path: running-save dispatch records the source mapping but does not acquire those leases, and preemption does not take the normal There is a second fail-open exit in the same fence: after 30 seconds it logs that a transfer is still in flight, then returns normally and allows the next forward to overwrite its blocks. CPU reproduction: the original fence returns with queued-only work, with submission in progress but no published future, and with a nonterminal future after the deadline. Suggested fix: cover queued → submitting → pending → terminal as one lifetime. A timeout must fail-stop, or first establish cancellation/quiescence; logging and continuing is not a safe source fence. Scope: plugin MP, not only V4.1. This finding is not a claim that native uses this plugin preemption path. P1 — 2. Removing the Request at finish prevents late suffix saves and strands their block leasesLocations: connector.py:2512–2537, connector.py:887–899
The saved frontier does not advance, the save tracker does not retire, and the suffix PAGE leases remain held. There is no dispatched operation for the transfer deadline to expire, and MP CPU reproduction: after finish, the Request registry is empty, the late save returns Suggested fix: retain the Request/hash information needed by late saves until their retirement. Alternatively, abandoning the suffix must explicitly retire the tracker/session and release its leases. Scope: V4.1 plugin MP. Already-dispatched saves carrying complete metadata are a different case. P1 — 3. A failed load's old destination is reused as a save source without a pinLocation: page_ride.py:160–200 The save branch prefers A later save can consequently select a freed/reused old destination instead of the current checkpoint. Because CPU positive/negative control: with recorded destination Suggested fix: keep load-destination bookkeeping separate from save-source resolution. Resolve saves from the current hash-backed source and pin it for the transfer lifetime; invalidate load records on recovery/reallocation. Scope: V4.1 plugin MP load-failure/recompute recovery, not every successful load. P1 — 4. Cursor equality is not sufficient to decide that a recycled slot needs no restoreLocation: cache.py:453–458 The new hook restores only rows where CPU reproduction: on the V2 host-history path, a same-position/different-prefix case makes zero restore calls, retains the old history, and passes the stale check. A concrete configuration for this path is eager execution with Suggested fix: track slot/request identity or generation and whether the newly admitted request has restored its checkpoint. Token-position equality alone cannot establish state ownership. Scope: plugin checkpoint restoration; native does not install this restore hook. P2 — 5. Native PAGE-only MP can wait an extra ~50 ms to report an idle load completionLocation: mp/worker.py:606–626 The new drain thread polls futures after Native's 1 ms idle progress polling can therefore repeatedly call Controlled CPU comparison, five trials per version:
The future was made ready immediately after the drain entered its next wait; these are controlled notification-latency measurements, not production average latency or GPU throughput measurements. Suggested fix: provide a completion wakeup, poll promptly while transfers are pending, or let the idle-progress path drive future polling. Scope: generic PAGE-only native MP. NativeState's P2 — 6. Current-step Engram staging runs before the previous speculative checkpoint is finalizedLocations: deepseek_v41_bridge.py:808–825, staging call:614–628
If the previous DSpark step committed at checkpoint boundary N, and the next step clears drafts to coexist with a waiting prefill, that next non-tentative staging can advance cursor/history to N+1 before the pending checkpoint at N is copied. The checkpoint can then combine the window at N with cursor/history at N+1. This ordering problem is independent of async scheduling; Suggested fix: finalize the previous committed checkpoint before current-step staging mutates its state. The relevant checkpoint image here contains window and cursor/history; this is not a claim about copying compressor rings into the image. Scope: V4.1 plugin DSpark checkpoint ordering; established by tracing the staging/store order, not a GPU reproduction. P2 — 7. The V2 facade exposes an optimistic CPU count as the exact committed startLocation: req_id_passthrough_patch.py:137–144
For example, a step starting at 255 can actually commit to 256 while the already-scheduled next step carries CPU count 261. A checkpoint candidate exists at 256, but CPU positive/negative control: start 261 produces zero stores and discards the candidate; start 256 stores it once. Suggested fix: use the actual committed frontier, or explicitly reject async scheduling with this checkpoint path until supported. The recipe's Scope: plugin V2 + async speculative scheduling. P2 — 8. Local checkpoint hits bypass DSpark's single-token-tail protectionLocation: v41_state.py:229–243 The external-hit path applies With Suggested fix: enforce the DSpark batch/tail constraint on local admission too, without inventing a checkpoint boundary that does not exist. Scope: opt-in local checkpoint hits + DSpark, not the recipe's default local-hits-off path. P2 — 9. Full-prompt hits disagree on the PAGE transfer end and STATE destination boundaryLocation: chunked_scheduler.py:900–918 For chunk/checkpoint interval 256, a 1024-token prompt fully present in the tier, HBM frontier 0, and
The request has already been parked based on the external-hit promise. It eventually falls back to recompute through the promise-withdrawal checks instead of using the available tier hit. Those checks are scheduler iterations, not a fixed wall-clock timeout; this is not a claim of a permanent hang. Suggested fix: make the physical PAGE transfer endpoint, recurrent destination, and resumed logical frontier consistent. Scope: V4.1 plugin MP. The old MP override becomes problematic in combination with the new recurrent-destination lookup; ordinary native MP without this hook does not inherit this failure. Native impact summary
Validation and coverage limits
Priority regressions to add alongside the fixes: preemption with queued/submitting transfers, finish with an undispatched suffix save, failed-load recomputation into different blocks, same-position slot reuse across different prefixes, speculative checkpoint ordering, and native idle load completion. Automated review prepared with Claude Code; posted at the user's request. |
|
IPC event lifetime on ROCm 7.2.x — two concerns with the per-step event + drain-thread release
Suggestions:
Repro scripts, live-service probe and raw data: https://gist.github.com/yhl-amd/6f8d9b4a0b7dd9c30d4ee4170b98f7a2 Caveat: the live numbers are from MiniMax-M3 on the native ATOM path, not the V4.1 vLLM-plugin path. |
… plugin) V4.1's per-request state (window ring, compressor rings, Engram cursor) is checkpointed at PAGE boundaries into vLLM-managed MambaSpec groups; the MP worker registers those groups as recurrent engine groups so a tier object carries the PAGE KV and the state checkpoint of the same boundary. Every hit, local or external, is capped to a boundary whose state is stored. - connector: V4.1 planner/ride, MP-only state transport, request-level multi-group load failures (vLLM 0.31 failed_recving), async tier lookup with detach once a local hit rules the load out - chunked scheduler: recurrent hand-off hooks for the base scheduler - MP worker: one fresh interprocess event per step, recurrent groups registered one chunk wide, LMCache 0.5.5 submit_request lookups - scheduler patches: prefill chunks end on the checkpoint grid; local hits capped to a stored checkpoint - recipe: recipes/atom_vllm/DeepSeek-V4.1-Flash-LMCache-Byte-Offload.md Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…entation Local hits: remove ATOM_V41_CHECKPOINT_LOCAL_HITS. With checkpoint groups registered, vLLM's local hit already ends on a boundary the groups hold a checkpoint for, so it stands and the tier is asked only for the rest -- vLLM's normal HBM-then-tier order. Without checkpoint groups the hit is declined as before. Debug instrumentation removed: the planner's ride/cap counters and the periodic "V4.1 state leg" stats line, last_load_skip_reason, the PAGE namespace and MP adapter namespace logs, and the startup lines that only reported which hooks were installed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…sources, slot ownership) - MP preemption fence: also wait for transfers still queued for the drain thread (they reach the server after the fence), and fail-stop on timeout instead of letting the forward overwrite blocks a transfer still reads. - MP worker: get_finished wakes the drain thread while transfers are pending, so an idle engine sees a completed load without the 50 ms backstop. - Tail saves of finished requests: the ride hook reads vLLM's request off the SeqView (which lives as long as the save is owed); a finished request whose save still cannot carry its state retires its tracker and gives back the leased suffix not read by an earlier operation. - Ride sources: a save rides only on the block committed under its hash and is always pinned; a recorded load destination (stale after a failed load) is never a save source. - Slot ownership: a real request that enters another request's STATE slot mid-prompt has its cursor invalidated before staging, so it is restored or refused even when the old cursor matches its position. Synthetic (capture/warmup) batches are left alone. - Speculative decode keeps are stored at commit time, before the next step's staging advances the cursor. - Local hits under DSpark keep a prefill tail (with or without a connector); checkpoints refuse async scheduling with speculative decoding. - MP: no full-prompt transfer_end override when recurrent state rides the load; the state is restored at the logical hit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…s measured Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A V4.1 tier hit is shortened after the lookup: `_loadable_hit` steps a whole-prompt hit back one chunk (prompt length % 256 == 0), and under DSpark `cap_hit` does the same to keep a two-token tail (length % 256 == 1). The MP server prefetched the recurrent state of the lookup's endpoint, not of the shortened hit, so the load failed and the request recomputed from its HBM frontier, usually token 0. Truncate the query instead, as the native state scheduler does, and place the prompt-end checkpoint where such a hit can end. Signed-off-by: linsun12 <linsun12@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…tream is idle A finished transfer's future holds its step event, so the drain thread dropped the last reference to an interprocess event. On ROCm 7.2.x that destructor synchronises every stream of the device, and while it runs, creating an interprocess event on any other thread (the step thread, every step) blocks as long (review by yhl-amd on #2598). Terminal futures, and the event of a submission that never produced one, are parked for the step thread. start_load_kv drops them while its stream is idle (under synchronous scheduling, the start of a step), or past 256 parked regardless. Every event is now created and destroyed on the step thread. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
6fb1f34 to
807d20b
Compare
|
Round 2 pushed 807d20b, rebased on main d6681e1. All nine items are fixed in d10a698; item 9 is completed by 95c687d (#2642). Each fix has a regression test. Item to fix
GPU at 51a1ef5 (the fixes, before the rebase): GSM8K two-pass on the tier-only and HBM plus tier paths, DSpark and no spec, pass 2 within 0.7 pp of pass 1, zero faults. Throughput is in the recipe. |
|
Round 2 pushed 807d20b with suggestion (b). Finished futures, which hold the step event, are parked for the step thread; it drops them at |
yhl-amd
left a comment
There was a problem hiding this comment.
Reviewed current head 807d20be0ddfad33fcae407eefe1c48ddc10dcd9. Three inline findings remain: a P1 checkpoint-restore defect and two P2 checkpoint-store/event-lifetime regressions. I recommend addressing them before merging.
Validation: 367 focused regression tests passed, 1 skipped; 8 controlled CPU probes passed; git diff --check passed. The probes execute selected unchanged source functions and the pinned LMCache adapter's pruning method with controlled stand-ins. Full plugin integration and GPU inference/accuracy/performance were not run in this review because the local environment has no vLLM, LMCache or AITER installation.
Current-head native accuracy CI passed, but vLLM plugin accuracy was skipped. The non-GPU CI job excludes tests/plugin, and the PR body says its plugin GPU runs predate the rebase/event-release change. Those results do not validate these latest plugin paths.
| cursors = ( | ||
| torch.index_select(self.cursor, 0, step.slots[:count]).cpu().numpy() | ||
| ) | ||
| cursors = self._restore_stale_rows(cursors, wanted, ids, fresh, step, count) |
There was a problem hiding this comment.
[P1] Restore checkpoint state on the Engram device-history path
DeepseekV41MetadataBuilder.prepare_model_inputs() calls prepare_state(histories=False) when Engram uses UVA or overlap (both default to enabled), so this restore hook is skipped on that path.
For a local or tier prefix hit entering a new/reassigned STATE slot, _v41_invalidate_reowned_slots() sets cursor -1. The device-history branch then records that cursor for the delayed stale-state check without restoring the checkpoint. The later planner does not repair it: _v41_stage_outside_forward() returns no reset_rows, but _v41_state_checkpoints_around_forward() uses getattr(prepared, "reset_rows", ()) as fresh_rows, leaving its resumed-row list empty. The next prepare_state() can therefore reject the recorded state instead of resuming the prefix.
A controlled CPU execution of the unchanged functions gives zero restore calls and needs state at 256, found -1 on the next probe. The same setup with histories=True calls the installed hook and accepts cursor 256. This establishes the missed control-flow branch, not a GPU inference reproduction.
Please restore the state before device Engram staging and the delayed probe, and carry the row ownership information needed by the planner. The host-history branch alone does not cover the advertised UVA path.
There was a problem hiding this comment.
Fixed in 13bae6f. With a restore hook installed, prepare_state now reads the cursors blocking whatever the caller asks, so the hook runs before the reset and the Engram staging. The plugin already forced that read unless ATOM_V41_BLOCKING_STATE_PROBE=0, which is how the deferred path was reachable. The bridge now passes the planner fresh_rows=() explicitly, since prepare_state has restored by then.
Test
test_an_installed_restore_hook_keeps_the_read_blocking: on 807d20b prepare_state(histories=False) returns None and the hook is never offered the row; after the fix it restores (0, 7, 27) and the next step is accepted.
| get_current_req_ids, | ||
| ) | ||
|
|
||
| batch = get_current_input_batch() |
There was a problem hiding this comment.
[P2] Preserve the batch through the speculative commit store
On the V2 runner these getters expose thread-local data only for the duration of the wrapped ModelState.prepare_attn(). Its finally restores the previous batch/request IDs. Pinned vLLM db9527a46873454610df6dbedf79a36d6bf1a7f6 calls postprocess_sampled() later from sample_tokens(), after attention preparation and the forward have returned. At this new commit-time store site the getters consequently return None and this helper silently returns.
The decode candidate remains for the old next-build fallback. If the following step is non-tentative (for example, drafts are cleared to admit a prefill), Engram staging advances the cursor before store_decode_keeps() runs, so a candidate filed at boundary 256 can copy cursor/history 257. That leaves an inconsistent checkpoint which a later prefix restore can reject.
The CPU probe executes the original snapshot wrapper and helper: it stores while inside prepare_attn, skips after the wrapper returns, and the original next-build fallback subsequently copies a cursor marker 257 for the 256 candidate. The existing regression test mocks both getters, so it misses their real lifetime. No GPU accuracy failure is claimed here.
Please pass the current batch/request data explicitly from the runner (for example, self.execute_model_state.input_batch) or retain a per-step snapshot through sampling, and store the checkpoint immediately after commit, before the next staging operation.
There was a problem hiding this comment.
Fixed in 6fdb823. sample_tokens also clears execute_model_state before postprocess_sampled, so the build keeps its own rows: record_decode_candidates stores the request ids and scheduled ends, and store_committed_decode_keeps stores at those ends minus num_rejected. It runs from postprocess_sampled after a speculative step, and from execute_model after a step without drafts, which had the same next-build problem.
Test
test_decode_keeps_are_stored_after_prepare_attn_returns runs the real snapshot wrapper and stores after it returns. On 807d20b the same sequence stores nothing and leaves both candidates for the next build.
| with self._lock: | ||
| for operation_id, entry, result in terminal_ops: | ||
| self._pending_loads.pop(operation_id, None) | ||
| self._retired.append(entry.future) |
There was a problem hiding this comment.
[P2] Transfer ownership of the IPC-event resources, not only ATOM's future reference
Parking the future here does not guarantee that the step thread drops its last event reference. The pinned LMCache wheel source f22dec28 also retains the future in AtomMPWorkerAdapter._context_operation_futures, and the future retains the step event. _release_retired() drops only ATOM's reference; LMCache prunes terminal futures in _track_operation_future() on the next submission, which this PR executes on the drain thread. That can still destroy the event on the drain thread after step-thread retirement completed. The save twin at line 870 has the same issue.
Two additional paths remain: the drain's pending, terminal_ops and entry locals hold references after publication to _retired, allowing the step to empty the list before the drain drops its locals; and a valid adapter return of None creates no future retaining the event but does not park the event explicitly.
Controlled CPU probes reproduce drain-thread destruction for all three paths, including the unchanged _track_operation_future() from the pinned LMCache source, with a step-thread positive control when no other owner remains. These are reference-ownership reproductions, not new GPU timing measurements. On ROCm 7.2.x they leave the device-synchronization and cross-thread event-creation stalls from the earlier review possible, including for native PAGE MP users.
Please arrange explicit step-thread ownership/release of the event-bearing resources, accounting for LMCache's retained references, and publish retired resources only after the drain has relinquished its references.
There was a problem hiding this comment.
Fixed in a7d323e. The step thread owns each step event from creation and each finished future from its terminal poll. It destroys one only while its stream is idle and once sys.getrefcount shows no other holder. That covers the three paths you list and one more: DefaultEventIPCBackend._exported_events, a 2048-entry deque that evicts on the drain thread at every export.
Test
test_ipc_resources_are_destroyed_on_the_step_thread models those holders: futures pruned at the next submission, a two-entry export ring, a submit returning None and one raising. On 807d20b, future4-7 and event2/4/5/6 are destroyed on the drain thread; after the fix every destroy is on the step thread.
|
No blocking issues from (c) on the V4.1 path: the change cuts step-thread event creation from 7.1 to 5.6 ms mean (p99 15.9 to 11.6 ms), and the release it adds costs 0.017 ms. End to end it is flat (8.16 vs 8.12 req/s, one run each). Creating the interprocess event still costs about 5 ms on each step with transfers, so the remaining win is reuse, your (a). Setup and numbersMI355X TP4, DeepSeek-V4.1-Flash DSpark on the vLLM plugin, FULL_AND_PIECEWISE, LMCache MP. aiperf concurrency 16 for 180 s, 7097-token prompts, HBM pool 4000 blocks. Image
|
`prepare_state(histories=False)` ships the cursors to pinned memory and judges them on the next call, after the reset and the Engram staging. The checkpoint restore hook runs only on the blocking branch, so a prefix hit entering a re-owned slot was never restored there and the next call rejected it (review by yhl-amd on #2598). The plugin's default already forces the blocking read; `ATOM_V41_BLOCKING_STATE_PROBE=0` reached the deferred one. With a restore hook installed the read now blocks whatever the caller asks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… rows The commit-time store read the batch through the pass-through getters, but the V2 runner commits from `sample_tokens`, after `prepare_attn` returned and took them along, so it returned without storing. The candidates fell through to the next build, where a step that is not speculative has already advanced the cursor while staging, and the keep paired one step's window with the next step's cursor (review by yhl-amd on #2598). `record_decode_candidates` keeps the batch-ordered request ids and scheduled ends it was built from, and `store_committed_decode_keeps` stores at those ends less the rejected drafts: from `postprocess_sampled` after a speculative step commits, and from `execute_model` after one that did not speculate. The bridge no longer hands the planner restore rows: `prepare_state` has restored them already. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…holds them Parking finished futures for the step thread was not enough: LMCache keeps each future until its next submission and each exported step event in a 2048-entry ring, both released on the drain thread, and the drain's locals or a submit that returned no future could still drop an event last there (review by yhl-amd on #2598). The step thread now owns every step event from creation and every finished future from its terminal poll, and destroys one only while its stream is idle and once `sys.getrefcount` shows nothing else holds it. Past 8192 owned entries the idle check is skipped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 3 pushed a7d323e: the three findings of the 10-26 review are fixed, one commit each, with replies in their threads. GPU reruns are pending because the test cards are reserved until 2026-10-27T03:00Z; I will post decode-keep, accuracy and IPC timing results then. Unit tests
|
|
Round 23 pushed ba157dd: all nine items are done, 1, 3 and 4 included. Please approve.
Evidence
Each new test fails with its fix reverted. Unit, at ba157dd in the PR's container: the CI command and GPU, MI355X TP4, DeepSeek-V4.1-Flash DSpark:
|
valarLip
left a comment
There was a problem hiding this comment.
Re-read the core files at ba157dd, by reading the code and running the related tests (offload / MP / native / config / shell 812 passed, block manager / V4.1 1001 passed, test_lmcache_mp* ten times in a row) and two probes; no GPU. Thanks for round 23. Items 1 and 3 from last time are fixed.
Item 4 (deadlines during a stuck submit) still has a gap, reproduced below as item 1. So I can't approve yet. Items 1, 4 and 9 continue my last items 4, 1 and 5. Each item ends with the change I'm asking for.
1. A stuck submit can still hide a missed deadline (mp/worker.py:532-554). _enforce_stalled_deadlines stops a transfer only while last_round_at < due. But last_round_at is when the last round ended (transfer_thread.py:358), and a round checks each entry when it reaches it. Suppose a round starts at due − 2 ms, polls transfer X (not yet due, no raise), spends a few ms on a terminal import, and ends at due + 3 ms. Then the next submit hangs. From then on last_round_at < due is false on every step, and X is never stopped. A probe confirmed it: 90 s past the deadline, no raise. The method also restates _enforce_transfer_deadline with a magic _POLL_STALL_GRACE_S = 1.0, as a second deadline path.
→ Change: on the forward thread, call _enforce_transfer_deadline(op, entry.started_at, deadline) on every pending entry each step (a CPU comparison, which is what I meant last time). Then delete _enforce_stalled_deadlines, _POLL_STALL_GRACE_S, last_round_at and _last_round_at.
2. Native V4.1's namespace leaves out two byte-geometry fields (offload/config.py:44-56, :380-387). Native dsv41 maps to dsv4-page-regions, and _HF_PAGE_FIELDS doesn't include sliding_window or candidate_block_size. The plugin hashes both, because, in its own words, either one moves the bytes a stored object holds. So under native ATOM, two V4.1 checkpoints with the same name and revision but a different window or candidate block share a namespace. A hit then restores a window ring or index plane of the wrong length.
→ Change: add sliding_window and candidate_block_size to the hashed fields for dsv41 in core build_page_namespace, and accept the one-time native namespace change explicitly, instead of keeping two rules for one model.
3. The new adapter refusal runs after registration and outside its cleanup (mp/worker.py:320, :430-433). _start_transfer_thread refuses an off-step worker whose adapter records its own producer events. It runs after adapter.register_kv_caches(...) and outside the try / except: adapter.shutdown(). When it raises, the server registration and heartbeat are already live, and nothing shuts them down until process exit or the server's reaper.
→ Change: do the check before registration (it depends only on the adapter type), or move _start_transfer_thread inside the try.
4. close() ignores a timed-out join (mp/worker.py:345, mp/transfer_thread.py:291-294). transfer.close() drains and joins with a timeout and returns nothing, and close() goes on to adapter.shutdown(). In polling mode the thread runs rounds on its own. If a submit or terminal poll is stuck past 60 s, the join times out and the adapter drains and unregisters, while _poll_round may still be inside _terminal_future_result on that adapter's futures.
→ Change: have close() report whether the thread stopped (is_alive() after the join). If it didn't, log it and skip adapter.shutdown() rather than tearing the adapter down under a live thread.
5. A frontend's speculative form is hashed raw (offload/config.py:375-377). When offload_speculative_config is set it goes into the document as is; only the fallback goes through _stable_config_value. A torch.dtype in it makes json.dumps raise, and a NaN is hashed instead of refused.
→ Change: run both through _stable_config_value.
6. An unneeded change to rank collapse (mp/deployment.py:165). _config_has_fully_replicated_tp_pages now asks select_offload_layout(config) in offcfg._STATE_OWNING_LAYOUTS instead of main's model_type == "kimi_linear". For V4.1 and GDN it changes nothing, since they have no kv_lora_rank and returned False on main already. It reaches into another module's private set and can raise on an override. The only real behaviour change is that a dense MLA model with an explicit kimi_k3 / qwen override loses collapse (8 copies instead of 1). test_auto_rank_collapse_keeps_kimi_k3_per_rank had to grow six hf fields so the layout lookup could run.
→ Change: revert to main's check, and drop the extra fixture fields.
7. The polling thread spins while idle (mp/transfer_thread.py:313-316). In polling mode _queue.get(timeout=_IDLE_POLL_S) turns every timeout into a round, even with nothing pending or queued. A probe on an idle rank counted 19 empty rounds a second, each taking _lock twice and copying both pending dicts, indefinitely. While steps run, get_finished already asks for a round each step via look(). The timer only matters when a transfer is in flight and the rank has stopped stepping.
→ Change: block without a timeout when nothing is pending or queued, and use the timeout only while a transfer is outstanding.
8. Four fields and three entry points for "wait for a round" (mp/transfer_thread.py:125-129, :217-280). _rounds_started / _rounds_done / _round_asked / _last_round_at back look, _wait_round and _wait_for_round. In fence, or self._poll_failed (:248) is dead, because raise_error() at the top of the loop already raises when _poll_failed is set.
→ Change: one generation counter under _idle, with request_round() -> target and wait_round(target, timeout). Delete the dead check. With item 1, _last_round_at goes too.
9. A lookup that didn't answer reopens saves (chunked_scheduler.py:825-832, hybrid/dsv4/connector.py:2321). Every False from _ensure_lookup_pin now calls _reopen_saves, including hit is None, which a lookup returns on timeout or RPC error. mp/lookup.py:516 calls that "a non-answer, not an empty answer". The tier most likely still holds the prefix, yet [hbm, lmc) is stored again: extra D2H and store traffic, not corruption.
→ Change: have _ensure_lookup_pin tell the two cases apart. Reopen only on a real answer below the promise; on None, just drop the load.
10. Three special cases for family-bound layouts (offload/config.py:160-185). There's natural == "m3" and mapped != "m3", natural in _STATE_OWNING_LAYOUTS and mapped != natural, and the reverse mapped == "dsv41" and natural != "dsv41". The reverse direction is checked only for dsv41, so a dense model overridden to m3 / kimi_k3 / qwen is still accepted, and the next family-bound layout will need another special case.
→ Change: one table of family-bound layouts, checked in both directions, with an explicit allowance for the test fixtures that use the kimi_k3 override.
11. The README module table is stale (offload/README.md:116). It says mp/native_state_{layout,scheduler,worker}.py own native-state registration. Registration (_register_native_state_layout) and block-id validation now live in mp/worker.py, and require_native_state_server in mp/deployment.py.
→ Change: update the row to point at the files that now own these.
12. A sleep-timed test (tests/test_lmcache_mp_async_threads.py:339-343). test_a_failed_poll_round_stops_the_polling relies on time.sleep(3 * _IDLE_POLL_S) and time.sleep(0.02) for the round to have run. drain() waits only on _outstanding, and the _WAKE sentinel isn't counted. On a loaded CI runner the round may not run within ~170 ms, rounds == [], and the test fails intermittently.
→ Change: wait with the existing _wait_for_round(target, timeout) (or item 8's wait_round) instead of sleeping.
Addresses valarLip's review of ba157dd. - With the polling thread, the forward thread enforces every pending transfer's deadline each step (`_enforce_transfer_deadline`, a clock comparison), so a submit that never returns cannot hide one. The stall heuristic, its grace and `last_round_at` are gone (1). - An off-step worker refuses an adapter that records its own producer events before registering it, inside the cleanup that shuts the adapter down (3). - `_TransferThread.close` reports whether the thread stopped; when it did not, `close` leaves the adapter registered for the server's reaper rather than shut it down under a live thread (4). - A frontend's speculative form is hashed through `_stable_config_value` like the fallback (5). - Auto rank collapse is main's `kimi_linear` check again; the test fixture is main's (6). - The polling thread polls on a timer only while transfers are in flight (its poll returns whether they are); idle, it waits for a submission or a request (7). - One request/wait pair: `request_round()` returns the round to wait for, `wait_round(target, timeout)` waits for it. Three round counters replace the flag and the timestamp, and the fence's dead check is gone (8). - `_ensure_lookup_pin` answers None when the lookup did not answer, and the caller reopens saves only on False, an answer below the promise (9). - README: native-state registration lives in `mp/worker.py`, the server check in `mp/deployment.py` (11). - The failed-poll test waits on rounds instead of sleeping (12). Not changed: the native V4.1 namespace (2) stays main's, and the override table (10) keeps main's accepted combinations. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 24 pushed ce7dd45: items 1, 3-9, 11 and 12 are done; 2 and 10 are not changed. Item 1 was the open blocker; please approve.
Evidence
Each new test fails with its fix reverted. Unit, at ce7dd45 in the PR's container: the CI command and GPU, MI355X TP4, DeepSeek-V4.1-Flash DSpark:
|
valarLip
left a comment
There was a problem hiding this comment.
Re-read the core files at ce7dd45, by reading the code and running the related core tests (816 passed, 1 skipped; the MP thread tests three times over); no GPU. Thanks for round 24 and the GPU runs. I accept your reasons for leaving items 2 and 10 to their own PRs.
Item 1 below is a false fail-stop introduced by the deadline fix I proposed last round. That's on me, and it's the one thing I'd like fixed before approving. Items 2–9 can go in a follow-up if you prefer. Items 2, 3 and 5 also come from my earlier asks (7, 9 and 4). Each item ends with the change I'm asking for.
1. The forward-thread deadline check fires on transfers that already finished (mp/worker.py:543-547). In polling mode, _settle_transfer_thread calls _enforce_transfer_deadline on every entry in _pending_saves / _pending_loads using only started_at, without first asking whether its future is terminal. Terminal entries leave _pending_* only when the transfer thread runs a round. Take a save or load that finished before its deadline while the transfer thread is held in a slow submit (or in _wait_ready for an after-forward save). The next start_load_kv / get_finished raises LMCacheTransferUnprovable ("still not terminal") and the engine fail-stops over a transfer that completed. _poll_entries and the inline get_finished check terminal first and deadline second; this copy reverses that. My suggestion last round left out the terminal check.
→ Change: check _terminal_future_result(entry.future) before the deadline, reusing _poll_entries' per-entry step rather than a third copy. Add a test where a save finishes, the thread is held in a submit past the deadline, and _settle_transfer_thread must not raise.
2. req_id_of / raw_req_id still miss the isinstance spellings (Fix-then-sweep). The last sweep covered the getattr / hasattr form, but the str(req_id.req_id if isinstance(req_id, LoadOperationId / SaveOperationId) else req_id) form remains at:
chunked_scheduler.py:1146,:1378,:1402and:1424, including this PR's ownload_failed;hybrid/dsv4/connector.py:2624and:2673, plus the paths around:2588/:2677-2679.
→ Change: use the helper everywhere the id is unwrapped. Keep the isinstance checks only where the code branches on the operation type itself.
3. The tier-lost block is duplicated (chunked_scheduler.py:820-835, hybrid/dsv4/connector.py:2312-2328). Both dispatch loops repeat the same twelve lines:
- write the floor into
_load_save_floors; - call
_ensure_lookup_pin; - branch on its three-way
True / False / Noneresult; _mark_load_skip,_reopen_saveson False,_clear_pending_load.
The floor is written only to be read back four lines later by _reopen_saves.
→ Change: one OffloadSchedulerMixin helper, e.g. _confirm_load_pin(seq, sid, ls, hbm, ...) -> bool, that does the skip / reopen / clear and takes the floor as an argument. Both loops call it, and _ensure_lookup_pin's tri-state stays private to that helper.
4. The README describes the collapse rule you reverted (offload/README.md:212-216). It says tp_rank_collapse: auto "leaves every state-owning layout (kimi_k3, qwen, dsv41) one copy per rank". After item 6, _config_has_fully_replicated_tp_pages again uses main's kv_lora_rank + kimi_linear test and doesn't look at the layout. dsv41 and qwen come out per-rank only because they have no kv_lora_rank. A future MLA model with per-request state would be collapsed under auto, contrary to the README.
→ Change: state the actual rule (MLA pages without kimi_linear collapse; everything else stays per rank).
5. close() can leave the connector half-closed (mp/worker.py:352-367). When the thread doesn't stop, close() sets _transfer = None, logs and returns, but keeps _adapter. A second close() (a repeated plugin shutdown) then finds no transfer and calls adapter.shutdown() under the still-running thread, which is what the first call avoided. Any later get_finished sees _adapter set and no transfer, so it polls the same futures inline on the forward thread. Separately, the log says "within 60s", but drain and join each wait _TRANSFER_DRAIN_TIMEOUT_S, up to 120 s.
→ Change: on that path, drop _adapter too (leave the registration to the server's reaper) or set a closed flag every entry point checks, and log the real bound.
6. The first recorded error can mask a poll failure (mp/transfer_thread.py:178-186, :313, :348). _record_error keeps only the first error. If a submit error (e.g. duplicate operation) was recorded and not yet raised, the poll failure that stops the polling is never reported: every step re-raises the stale submit error. When _poll raises, _in_flight also keeps its last value, so a True leaves the thread waking at 20 Hz until exit.
→ Change: when a poll fails, prefer that error (or chain the two), and set _in_flight = False on the failure path.
7. A test whose assertions hold by construction (tests/test_lmcache_mp_native_worker.py:781). test_native_transfers_are_submitted_and_polled_on_the_step_thread builds its worker with _make_worker(), which never calls register_kv_caches / _start_transfer_thread. So _transfer is always None, and the thread count is read after construction. If the shared _register_native_state_layout ever started a thread for native, saves would leave the step thread and this test would stay green. Its docstring line "No transfer thread starts" also isn't checked on the real registration path with OFFLOAD_MP_TRANSFER_THREAD=1.
→ Change: register through the real path and compare the thread count before and after.
8. A shared-fixtures module with one user (tests/lmcache_mp_support.py). Its docstring says it holds the fixtures "every LMCache MP worker test file shares", but only test_lmcache_mp.py imports it (via # noqa: F401 so the autouse fixtures apply). step_event(worker) ignores its argument, and start_transfers only wraps _start_transfer_thread.
→ Change: move the fixtures into test_lmcache_mp.py (or a conftest.py if other files should use them), and drop the two wrappers.
9. Small:
_poll_roundnames two different things: the worker callback (mp/worker.py:904, returns bool) and the thread method (mp/transfer_thread.py:343, counts and records errors), which calls the first throughself._poll.LMCacheMPConnector's class docstring (mp/worker.py:104) still says only "for backend-published PAGE views".offload/config.py:83-84says "qwenis the same tier over a GDN PAGE image", which contradicts the new refusal that the two tiers can't restore each other's state.
→ Change:
- Rename the worker callback, e.g.
_poll_transfers. - Update the class docstring to what it does now.
- Reword the layout comment to match the refusal.
Addresses item 1 of valarLip's review of ce7dd45. The forward thread enforced every pending transfer's deadline from its start time alone. A transfer that had finished, but waited for a round held back by a slow submit, was failed as "still not terminal". The forward thread now polls only transfers already past their deadline, through the same per-entry step as a round (`_poll_entry`: terminal first, deadline second), so a finished one is left for the next round to settle. Transfers within their deadline are not polled there, so terminal polls (the IPC event import) stay off the forward thread. The check runs only between rounds (`between_rounds`): the first query of a finished future imports its event, and a round must not be polling the same future at the same time. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 25 pushed 4b0dc75: item 1 is fixed; items 2-9 go to a follow-up PR, as you offered. Please approve. The forward thread now polls only a transfer past its deadline, through the round's per-entry step (terminal first), and only between rounds. A finished one is left for the next round to settle. EvidenceWhy not a terminal check on every entry each step: a terminal poll imports the peer's IPC event (
Unit, at 4b0dc75 in the PR's container: the CI command and GPU, MI355X TP4, DeepSeek-V4.1-Flash DSpark:
|
valarLip
left a comment
There was a problem hiding this comment.
Re-read the core files at 4b0dc75, by reading the code, the LMCache adapter it drives (lmcache/integration/atom/multi_process_adapter.py, lmcache/v1/multiprocess/futures.py at f22dec28), the core tests (557 passed) and three probes; no GPU. Thanks for round 25 and the A/B runs. The terminal-first check is right, and its tests fail when the fix is reverted.
I said I'd approve once item 1 was fixed. Reading the fix, I think the forward-thread deadline path I asked for two rounds ago shouldn't exist. That ask was mine, and I'm withdrawing it. Removing the path is the one change I'd like before approving. Everything else below can go to the follow-up PR with items 2–9.
Before approval: delete the forward-thread deadline path
1. Remove _enforce_overdue_deadlines, between_rounds and _polling (mp/worker.py:543-567, mp/transfer_thread.py:136, :239-253, :366). Four reasons:
- It races a submit, and only when it's actually used.
between_roundsexcludes only a poll round. A submit runs (_run,fn(*args)) without_polling. Every LMCache submit goes through_track_operation_future(multi_process_adapter.py:722-738), which callsquery()on every future it tracks for the context, on the transfer thread.query()on a finished raw future runs_on_raw_future_complete(futures.py:181-195): check_raw_response_processed,import_event, then set the flag, not atomically. So when the forward thread polls an overdue, finished, not-yet-settled entry while a submit is running, both threads can make the first query of the same future and import its IPC event twice, which is what the lock was meant to prevent. That window ("a slow submit is holding rounds back") is exactly when this path does anything. A probe confirmedbetween_roundsruns its function during a submit. I didn't verify what a double import does on ROCm. - It can't see the transfer it exists for. A submit that never returns has its operation only in
_submitting_*._PendingLoad/_PendingSave, and sostarted_at, are created after the submit returns (worker.py:705-714). A probe holding a submit and moving the clock an hour past the deadline:_settle_transfer_threaddoesn't raise, withpending_loads {} submitting {'load:16:1'}. - The scheduler already covers it.
LMCacheMPConnectorScheduler._enforce_transfer_deadlines(mp/scheduler.py:341-355) bounds every dispatched operation, the stuck one included, at the deadline plus_SCHEDULER_DEADLINE_MARGIN_S(10 s), natively and under the plugin. What the forward path adds is about 10 s earlier detection for transfers submitted before the stuck one. - Half of its locking is untested. With
with self._polling:removed from_poll_round,test_lmcache_mp.py+test_lmcache_mp_async_threads.pystill pass (182). The one related test (test_lmcache_mp.py:2729) holds the lock by hand on the test thread instead of running a round.
→ Change: delete _enforce_overdue_deadlines, the between_rounds call in _settle_transfer_thread, between_rounds, _polling, and the tests built on them (test_a_stuck_submit_cannot_hide_a_transfer_past_its_deadline, test_the_forward_thread_polls_no_transfer_before_its_deadline, test_the_forward_thread_leaves_deadlines_to_a_running_round). Keep test_a_transfer_that_finished_is_not_failed_while_rounds_are_held, which then checks that nothing on the forward thread fails a finished transfer. Rounds and the scheduler deadline remain the two bounds. If the submit-side wait should be shorter, lower the scheduler margin rather than add a third path.
Follow-up PR (with items 2–9)
2. A request during a submit runs two rounds (mp/transfer_thread.py:219-224, :355-359). request_round sets _rounds_requested = started + 1 and queues a _WAKE. After the submit, asked runs one round. The _WAKE is then taken with asked False but queue.empty() True, so a second, empty round runs. A probe gives rounds after one request + one submit: 2 target 1. Each round copies both pending dicts under _lock and queries every future. The comment above _IDLE_POLL_S ("rounds are otherwise asked for by the forward thread") also contradicts the module docstring's "whenever its queue runs dry".
→ Change: on a _WAKE, run a round only if _rounds_requested > _rounds_started, and make the _IDLE_POLL_S comment match the docstring.
3. The core retire change is tested only under tests/plugin (chunked_scheduler.py:668). _retire_finished_save now keeps the blocks another operation of the request still reads instead of returning the whole lease. The only test is in tests/plugin/test_vllm_offload_completions.py, which imports vLLM through the plugin and isn't collected by the core job. With _retire_finished_save reverted to main's whole-lease return, the related core offload / LMCache tests still all pass.
→ Change: add a core test (e.g. tests/test_offload_early_block_release.py) where one request has two save operations and the first retires while the second still reads.
4. Three scans for "this seq's operations" (chunked_scheduler.py:668, :977, :999). _retire_finished_save, should_defer_free and protected_block_ids each scan all of _save_operation_owner for one seq's operations. Every retire is O(all in-flight saves), and the rule "which of this seq's blocks are still read" exists three times.
→ Change: keep an owner → operations index and one _blocks_still_read_by(seq) that the three call.
5. Two registration skeletons in one class (mp/worker.py:169, :269). _register_native_state_layout repeats register_kv_caches' sequence:
- the rank, config and replication checks;
_make_worker_adapter;try { _refuse_unfenceable_adapter; register } except { shutdown; raise };- assigning
_adapter/chunk_size,_start_transfer_thread, and the log line.
The two have drifted: chunk validation is _strict_integer + % block_size in one and int() in the other, and shutdown goes through getattr in one only. Last round's "refuse inside the try" fix had to be made twice.
→ Change: one registration skeleton that both paths call, each supplying only its views and groups.
6. The "no answer reopens nothing" branch is tested on one side only (chunked_scheduler.py:828). The pinned is False / None distinction exists in both chunked_scheduler.py and hybrid/dsv4/connector.py, but only DSV4 has a test. With the chunked side's None turned into False, the related tests still all pass; the same change on the DSV4 side fails only test_dsv4_a_lookup_that_did_not_answer_reopens_nothing. native_state_scheduler.py:312-328 also returns None, also untested. This folds into last round's item 3: one helper, one test.
7. Docs (offload/README.md:24-30, :146, docs/environment_variables.md:430):
- The README's layout rule ("
compress_ratiospresent meanshybrid") predates CSA2 being checked first, so it says V4.1 resolves tohybridwhile:34-39saysdsv41. :146namesrecurrent_page_groups, which doesn't exist; the contract ispaged_state_checkpoint_spec.- The
OFFLOAD_MP_TRANSFER_THREADrow doesn't say that atransfers_off_stepworker ignores the env and always starts a polling thread, or that an adapter recording its own producer events is refused at registration rather than falling back to inline submission. The warn branch in_start_transfer_threadis now unreachable under off-step.
→ Change: bring the three in line with the code, and drop the unreachable warning.
Addresses item 1 of valarLip's review of 4b0dc75, which withdraws the earlier ask for this path. The check could not see the transfer it was for: a submit that never returns has no pending entry, since one is recorded only after the submit returns. While a submit ran it could poll a future that the submit's own operation tracking polls on the transfer thread, which is the concurrent first query its lock was meant to rule out. Polling rounds already enforce every pending transfer's deadline, and the scheduler's transfer deadline bounds a submit that never returns. Removed: `_enforce_overdue_deadlines`, `_TransferThread.between_rounds` and its lock, and the three tests built on them. `test_a_transfer_that_finished_is_not_failed_while_rounds_are_held` stays: nothing on the forward thread may fail a finished transfer. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 26 pushed ce1187d: the forward-thread deadline path is removed. Items 2-7 join last review's 2-9 in the follow-up PR. Please approve. EvidenceRemoved: The two remaining bounds, as checked: rounds enforce each pending transfer's deadline in Unit, at ce1187d in the PR's container: the CI command and GPU, MI355X TP4, DeepSeek-V4.1-Flash DSpark, plugin GSM8K two-pass: 0.9227 / 0.928, 5.85M tier hit tokens, 0 faults, 0 poll failures. |
Picks up #2809 (OPUS FP4 index scoring), whose aiter API change made the branch's V4.1 accuracy jobs fail to import. Conflicts: - deepseek_v41_bridge.py: _build_v41_pool_geometry keeps its required keywords and gains main's index_layout, also required; the proxy geometry passes None, what Config resolves for its FP8 plane. - tests/test_v41_geometry_parity.py stays deleted: both sides build through build_v41_pool_geometry. build_v41_pool_geometry takes index_layout and forwards it, so the native builder's call (auto-merged from main) builds the same geometry as main's inline construction. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 27 pushed f259bbb: merged main To review:
Evidence
|
Conflict: the kimi_k3 connector import list; keep both StateTouchId (this PR) and raw_req_id (#2598).
Conflicts after #2598 moved the shared native-state helpers out of mp/native_state_worker.py; both sides kept: - mp/deployment.py: TransferSpec next to require_native_state_server. - mp/native_state_scheduler.py: restore_slot_count beside the store channel. - mp/native_state_worker.py: TransferSpec and _mp_session_id only; restore_slot_count stays here. - mp/worker.py: TransferSpec with main's _transfer_block_ids. - tests/test_lmcache_mp.py: both imports. Co-authored-by: Cursor <cursoragent@cursor.com>
…qwen38-lmcache Brings in #3027's merge of main a30463d (after #2598, which moved the shared native-state helpers out of mp/native_state_worker.py). One conflict, in atom/kv_transfer/offload/config.py build_page_namespace: main's index_cache_layout (#2809) and this branch's QSA_GDN page segments are both kept, index_cache_layout first. Co-authored-by: Cursor <cursoragent@cursor.com>
… object (#3117) * Kimi-K3 plugin: ride the KDA state on the PAGE object through the LMCache MP server Kimi-K3 on the MP backend was refused ("recurrent (mamba) state offloads in-process only"). It now takes the path DeepSeek-V4.1's checkpoints take (#2598): the KDA state at a boundary rides the PAGE object as a native-state image, so KV and state commit and restore as one LMCache object. No copy is needed. vLLM lays every KV cache group out layer-outer from byte 0 of one allocation, so mamba group layer j aliases MLA plane j. On K3 the three mamba groups of 23 KDA layers sit on the first 23 of the 29 MLA planes, and the state a group holds in block b is PAGE unit b of those planes. A boundary's image is one unit per mamba group (3 x 23 x 884,736 B). - kda_ride.py: KdaRidePlanner (caps hits to the chunk grid, names the load destination row (n - 1) // block in every mamba group or none), kda_ride_geometry (the checkpoint spec both halves derive from vLLM's KV cache config), publish_kda_ride_state (checks on the registered tensors that every KDA layer starts at, and strides with, one of the leading PAGE planes before declaring the image), and the startup refusals (align mode, prefix caching, kv_load_failure_policy=recompute). - V41StateMPScheduler had nothing V4.1-specific; renamed StateRideMPScheduler (ride_mp_scheduler.py) and shared by both models. _v41_request_of is now _request_of. - The KDA state is per rank (TP-sharded heads); the in-process K3 leg is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Kimi-K3 plugin: size the KDA image by attention layer, and ride only K3 Review of #3117: - The image covered the leading N PAGE planes with N = mamba layers per group, which counts layers, not planes. An attention layer with a DSA index cache publishes two planes, so the count fell short. The spec now sums each leading attention layer's movable bytes from vLLM's specs (an index cache folded into the layer that owns it, as build_kv_cache_tensors publishes it), and the worker counts the region planes off the published regions' layer names before it checks the aliasing. - Only Kimi-K3 rides on MP. Its KDA layers report GDN_ATTN, as Qwen's GDN layers do, so the gate is the kimi_k3 offload layout; any other mamba model on MP is refused as before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Kimi-K3 plugin: put the layers the KDA state aliases first in the PAGE order The first K3 TP8 boot over mp refused at registration: mamba group 0 layer 18 started on PAGE plane 23, outside the image's leading 23 planes. ATOM publishes PAGE planes by layer name, while vLLM aliases mamba layer j onto the j-th attention layer it allocates, and with DSpark the two orders differ (29 attention layers, five sort ahead of the first aliased one). kda_ride_page_order moves the attention layers the mamba layers alias to the front, in mamba-layer order, matched on addresses, and the scheduler sizes the image from vLLM's allocation order. The end-to-end test gains a layer that sorts first without being aliased; without the reorder it raises the same refusal the boot did. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Kimi-K3 plugin: default spec_state_indices_in_tensor so DSpark starts #2365 made the shared K3 KDA forward read `kda_metadata.spec_state_indices_in_tensor` on its speculative-decode branch. vLLM's KimiK3KDAMetadata has no such field, and the plugin only defaulted its non-spec twin, so K3 on the plugin dies in profiling under DSpark (AttributeError on the first verify-shaped forward). None keeps the forward on the write slot, as the non-spec default does. The same hunk as #3137's, so the two apply cleanly in either order. The K3 TP8 LMCache MP runs on this PR ran with it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: perzhang <perzhang@amd.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Summary
LMCache prefix offload for DeepSeek-V4.1-Flash on the vLLM plugin, through
the LMCache MP server (no in-process LMCache). Works with
FULL_AND_PIECEWISE cudagraph and DSpark.
A V4.1 request cannot resume from PAGE KV alone. CSA2 attention also reads
per-request state: the sliding-window ring, the compressor rings and the
Engram cursor. This PR checkpoints that state at every 256-token PAGE boundary.
The MP worker stores each checkpoint in the same LMCache object as the PAGE
KV that ends at that boundary. Prefix hits follow vLLM's usual order: HBM
first, at a boundary whose checkpoint is cached, and then the tier for the
rest. Every hit is capped to a boundary whose checkpoint is stored.
How it fits together:
--enable-prefix-caching --mamba-cache-mode alignregisters 8 single-layer
MambaSpeccheckpoint groups. vLLM allocates and hashes them like any othergroup. They alias the V4.1 PAGE proxy pool, so no second pool is needed.
Prefill chunks are clipped to end on the 256-token grid, so each chunk end
has a checkpoint.
B/block/rank), and then each checkpoint group as a recurrent engine group one
chunk wide. A save carries the checkpoint of its last chunk only. The other
chunks hold the server's null block, -1, which the server drops. A load restores the PAGE
KV and the checkpoint at its end in one operation.
hold a checkpoint. The connector caps the external hit to a stored boundary.
Under DSpark both keep a prefill tail of at least two tokens.
Shared ATOM code changes
kv_transfer/offload/mp/worker.py,transfer_thread.pypaged_state_checkpoint_spec, registers each checkpoint image ordinal as a recurrent engine group, one chunk wide. A worker built withtransfers_off_step=True(the plugin's) runs #2906's transfer thread with a poll round, whateverOFFLOAD_MP_TRANSFER_THREADsays: it submits and polls off the step thread, with fences from_ProducerFences. A roundget_finishedorfence()asks for runs after the submission in progress, whatever is still queued. While a load is in flightget_finishedwaits up to 1 ms for that round when it is a pure poll (nothing queued, no round running).fence()waits on its rounds up to the transfer deadline. A failed round stops the polling, and the forward thread raises it on every step; rounds enforce every pending transfer's deadline, and the scheduler's transfer deadline bounds a submit that never returns. The thread polls on a timer only while transfers are in flight. A polled future holds the server's imported IPC event: it is handed back and dropped on the step thread once LMCache no longer holds it, while the stream is idle. An adapter that records its own producer events is refused. A chunk with no image names -1.STORE block ID underflow). On the plugin's step thread a submit cost 18.7 ms and a terminal poll 13.8 ms under launch pressure (the IPC handle export and import), so a wait short enough for a step never sees a round that has to submit first; a fixed 2 ms poll cadence with loads outstanding cost more end to end (ITL p50 23.1 vs 28.8 ms) than the work it moved off; without the 1 ms wait, c16 ITL p50 was 23.4 vs 22.8 ms and TTFT p50 326 vs 319 ms. Destroying an imported IPC event synchronises the device on ROCm 7.2.x; dropped mid-forward, the plugin's c16 ITL p50 rose from 23.0 to 23.8 ms. One submission thread and one event pool serve both frontends.get_finished, producer fence, future lifetime and submit-onlyOFFLOAD_MP_TRANSFER_THREADthread.kv_transfer/offload/mp/native_state_worker.py,native_state_layout.py,native_state_scheduler.pyget_finished, with main's producer fence, as on main. Registration moved into the base worker;stores_state_imagesmoved tonative_state_layout.pyandrequire_native_state_servertodeployment.py.kv_transfer/offload/chunked_scheduler.py,hybrid/dsv4/connector.py_retire_finished_saveretires a finished request's save that will never be emitted and releases its leased blocks, except those another of its operations still reads. Every lease release goes through_return_leased_blocks. The save floor is recorded once, when a load is decided; a load refused astier_lost_prefixor failing reopens[HBM, hit)for saving throughOffloadSchedulerMixin._reopen_saves, which reads it (both schedulers).kv_transfer/offload/mp/scheduler.py_enforce_transfer_deadlinesthrough theirenforce_transfer_deadlines. A promise withdrawn before dispatch no longer callscomplete_retrieve, so the lookup cleanup frees its read locks.complete_retrieveforgot the lookup first, so the server held the locks until session or TTL cleanup.kv_transfer/offload/config.py,connector.py,mp/deployment.pyattn_family) resolves to its owndsv41layout instead of DSV4'shybrid, and the in-processlmcache_offloadshell refuses it. Its namespace is the onehybridgave it; the plugin's config view addssliding_windowandcandidate_block_sizeto the geometry it hashes. A state-owning model accepts only its own layout as an override (qwenandkimi_k3are refused for each other too), and no other model may namedsv41. The namespace takes a frontend's stable speculative config when it supplies one (Config.offload_speculative_config).tp_rank_collapse: autokeeps every state-owning layout one copy per rank.hybrid, V4.1 accepted adenseoverride, which restores KV over stale state. Either geometry field moves the bytes a plugin object holds. A draft changes the cached bytes, and a per-process speculative field split one deployment's namespace in two.PAGE_LAYOUT_VERSIONstays 3). Native V4.1's namespace is byte-identical to main's, so its tiers stay reachable; anoffload_layoutoverride it accepted before is refused.models/deepseek_v41/pool_geometry.py,deepseek_v41/backend.pybuild_v41_pool_geometry(), one source of the pool geometry for native and plugin, with every argument that moves it required; it takes main'skv_dtype(#2725) andindex_layout(#2809).pool_layoutstill imports nothing ofatom.plugin/vllm/paged_state_checkpoint.pyprefill_tail, ATOM's checkpoint defaults, andsizes_pool_without_checkpoints.seq_lens - query_lens. A keep sharing the boundary store's buffer waited on the CPU behind the whole forward. One owner of the DSpark prefill tail keeps the scheduler, the local cap and the connector from disagreeing.plugin/config.pyplugin/vllm/platform.pygates the connector by vLLM's names.Plugin changes
plugin/vllm/kv_transfer/connector.pyAtomLMCacheOffloadConnector. It refusesinproc, refuses to start without checkpoint groups, and refuses retired options. It runs the boundary planner and a non-blocking tier lookup (a request is deferred while its lookup is in flight), and detaches the lookup once a local hit rules a load out. It handles multi-group load failures throughfailed_recvingand withdraws promised loads. A save reads vLLM's request off theSeqView, so a finished request's tail save still finds it.plugin/vllm/kv_transfer/v41_state.py,page_ride.pyV41BoundaryPlanner(the grid, hit caps, the DSpark prefill tail, which blocks a load names) and the ride bookkeeping. A load rides on its recorded destination. A save rides only on the block committed under its hash, and is always pinned.plugin/vllm/kv_transfer/mp_page_layout.py,mp_worker_adapter.pyfence(). A failed load reports only the blocks it writes.plugin/vllm/kv_transfer/boundary_planner.py,kda_state.pyKdaBoundaryPlannerand V4.1's planner.plugin/vllm/kv_transfer/offload_config.py,seq_view.pytp_rank_collapse: autoresolved to true for V4.1.SeqViewgainsnum_tokensand_deferred_save_at, which the shared scheduler reads and writes.plugin/vllm/paged_state_checkpoint_scheduler.pyplugin/vllm/deepseek_v4_bridge.py,deepseek_v4_prefix_patch.py,deepseek_v4_state_checkpoint.pyplugin/vllm/kv_transfer/plugin_schedulers.py,v41_mp_scheduler.pyhas_pending_block_frees), and refusing a save whose source blocks vLLM took back. V4.1's MP scheduler retires a finished request's save it declines to build.plugin/vllm/deepseek_v41_state_checkpoint.pyV41PoolGeometry. The checkpoints are on with--enable-prefix-cachingand--mamba-cache-mode align.plugin/vllm/v41_core_overrides.pyVllmBackedPagedAttentionCacherestores checkpoints into the rows core's contract would reject, then calls core's blocking read, so the contract judges what the restore produced. It defers the cursor read as native ATOM does while vLLM schedules synchronously, and blocks under async scheduling (where this frontend's deferred read raised false stale-state errors), during vLLM's warmup and on a step that can restore; it waits for a probe that has not landed rather than drop its verdict. It leaves vLLM's_warmup_N_rows, run through slots nothing wrote, out of the contract.plugin/vllm/deepseek_v41_bridge.pyregister_kv_caches. A real request that enters another request's STATE slot mid-prompt has its cursor invalidated, so it is restored or refused even when the old cursor matches its position.plugin/vllm/scheduler.py,register.pyplugin/vllm/spec_decode_patch.py,v41_speculative.pykv_holding_waitingrequest that is about to prefill as one: a resumed request must not share a step with a DSpark verify. At the speculative commit, flush the step's checkpoint stores, which the copier checks against the committed cursor.plugin/vllm/req_id_passthrough_patch.pyserving_synchronouslyopens afterWorker.compile_or_warm_up_model, on V2 under synchronous scheduling.plugin/vllm/platform.py,model_wrapper.py,deepseek_v41_state_reserve_patch.pyTests
All runs are on MI355X (cards 4–7), TP4, FULL_AND_PIECEWISE with
max_cudagraph_capture_size512, at51a1ef540, before the rebase onto maind6681e153and the review fixes since. The accuracy and timing reruns ata7d323eebare in this comment. Each run gets a fresh model and MP server. The full commands are inrecipes/atom_vllm/DeepSeek-V4.1-Flash-LMCache-Byte-Offload.md.Accuracy. GSM8K (1319 questions, 5-shot) runs twice on one server, with
an HBM pool of 950 blocks. Pass 1 fills the tier. Pass 2 repeats the prompts.
"HBM + tier" puts a shared 3842-token preamble in front of every prompt: the
preamble is served from HBM and the rest of each prompt from the tier.
Zero engine faults and zero
STORE block ID underflow. Pass-to-passdifferences range from −0.5 to +0.3 pp; one standard error is about 0.7 pp.
¹ Run at
c1ae98269. With checkpoints off, that commit runs the same code onthis path.
Performance (DSpark). aiperf at concurrency 16 for 180 s. Each prompt is
one of four 4096-token prefixes plus a 3000-token tail (7097 tokens on
average), with 64 output tokens and an HBM pool of 4000 blocks.
Under DSpark, HBM alone keeps only the shared prefixes; the tier supplies the
tails. Without speculative decoding the same workload fits in HBM: no reuse
4.78, HBM-only 14.55, HBM + LMCache 14.00 req/s, with the tier idle.
Theoretical vs measured (DSpark).
grid). Measured: 94.3%.
t(c) = 17.6 us × c + 91.7 ms,calibrated on the two arms that move no bytes. At LMCache's 408
recomputed tokens it predicts 10.11 req/s. Measured: 7.80 (77%).
50 GB/s. The measured server-side retrieve is 40 ms p50, so the gap is
per-load latency (TTFT +101 ms against HBM-only), not bandwidth.
Unit tests.
tests/pluginplus the offload, LMCache-MP and V4.1 suites,and the CI command (
tests/, withouttests/plugin): no failure that maindoes not also have in the same container. This PR's tests are in
tests/plugin/test_vllm_dsv41_lmcache.py; the other test changes adaptexisting suites.
Since the previous revision
10b99a085. It brings feat(indexer): migrate V4 and V4.1 FP4 scoring to the OPUS config API #2809, whose aiter API change failedthis branch's V4.1 accuracy jobs at import.
build_v41_pool_geometrytakesfeat(indexer): migrate V4 and V4.1 FP4 scoring to the OPUS config API #2809's
index_layout: native passesindex_cache_layoutas on main, andthe plugin passes None for its FP8 index plane.
removed: polling rounds and the scheduler's transfer deadline bound every
transfer.
longer failed by the forward thread's deadline check while a slow submit
holds rounds back.
thread every step. A self-recording adapter is refused before registration,
and
closeleaves the adapter to a thread that did not stop. The idlepolling thread no longer spins. A lookup that did not answer reopens no saves.
Auto rank collapse is main's check again.
to the thread it runs on, so shutdown with a transfer in flight is clean. The
foreground and background polls share one pass. A failed round is raised on
every step, and a stuck submit cannot hold a transfer past its deadline.
Every core unwrap of a completion id goes through
raw_req_idorreq_id_of.error is swapped under its lock, and a failed poll round stops the polling.
Native
get_finishedis main's single locked pass again, and native'ssubmit-only transfer thread keeps no fence bookkeeping. The save floor is
recorded once and read by every load that does not go out; the tier-lost
path no longer loses it to
_ensure_lookup_pin's cleanup.worker with
transfers_off_stepwhen it builds it, so the env flag no longerbreaks its startup. A round
get_finishedor a fence asks for runs after thesubmission in progress. The in-process shell refuses
dsv41. NativeState ismain's apart from the moved registration, and the MP worker tests run their
failure paths on both the step-thread and the polling worker.
run on [kv_transfer] lmcache_mp: lookup thread, transfer thread, same-step saves #2906's transfer thread, with its fence pool and its refusal of an
adapter that records its own events; the drain, step-event and IPC-lifetime
modules are gone. A lost or failed tier hit is stored again, and a withdrawn
promise frees its lookup locks; it still spends the request's tier attempt,
which stops a re-promise loop. The
_unbuilt_save_is_finalhook andlmcache.mp.null_block_idare gone, and V4.1 is detected byattn_family.dsv41keeps the namespacehybridgave V4.1, so native ATOM's tiers stay reachable; only the plugin hashes
V4.1's window and candidate-block fields.
finished. Held until the stream was idle, its steps ran 6% faster but its
saves fell behind: a GSM8K two-pass loaded 7% fewer tier tokens.
498133ed9. The conflicts were with [kv_transfer] lmcache_mp: lookup thread, transfer thread, same-step saves #2906 (lookup thread,transfer thread, same-step saves) in the MP worker; native ATOM keeps [kv_transfer] lmcache_mp: lookup thread, transfer thread, same-step saves #2906's
path.
polls as main does, without the drain-era state; the plugin's last look
waits only for a pure poll round; the three lease releases share one helper. The Kimi-K3
chunk-size message and two test rewrites are reverted, and native V4.1's
namespace change and
autocollapse are documented.polls loads. Plugin-only scheduler interfaces
moved to the plugin, and V4.1's namespace fields no longer move other
models'. The lookup shim, the per-request failure channel,
_canonical_head_dim, the core restore hook and both V4.1 env switches aregone.
next build from the step's own lengths, so the sampler's rejection counts
are no longer copied to the host. The V4.1 build reads PAGE rows, lengths
and cursors from the host as native ATOM does, and from the device during
vLLM's warmup and under async scheduling, where the cursor read still
blocks. Host rows are armed only by a runner that keeps them. A step
records its IPC event only when a transfer exports it, which fixes failed
server waits under TP rank collapse.
rank reports and the plugin drives MP's transfer deadline every step; a
state-owning model accepts only its own
offload_layout; the fence waits upto the transfer deadline. V4.1 offload requires
kv_load_failure_policy=recompute, and the V4.1 adapter matches only withits checkpoints on.
STATE tail withheld twice without checkpoints, save retirement and declined
loads, the preemption fence, preemption behind a stuck save, failed-load
blocks, the V2 table copy and a rebind. The Engram cursor change, which was
a no-op, is reverted.
and the vLLM warmup exemption live in the plugin's cache.
481c2bfd4. The conflicts came from feat(moe): route before the rows move; per-pass real row counts; V4.1 FP8 KV #2725 replacing thepool geometry's
packedflag withkv_dtype; both paths still build thegeometry through
build_v41_pool_geometry.read is deferred. The step thread destroys finished futures only once LMCache has let
go. Each fix has a CPU test that fails on
807d20be0; GPU results are inthis comment.
d6681e153(fix(mla): migrate Triton decode to maintained aiter API #2669 replaces theTriton MLA module, feat(v41): support P/D disaggregation (Mooncake) #2666's Mooncake P/D admitted), the lookup extent of dsv41 offload: look up only the extent a load can end on #2642,
local hits served by default, debug instrumentation removed, the fixes for the
review of 10-19, squashed onto main and slimmed from 11.6k to 6.5k added lines.
Known limitations
sliding_windoworcandidate_block_size(the plugin's does): two checkpoints of one name thatdiffer only there, through
--hf-overridessay, share a tier. Closing itmoves every native V4.1 namespace, so it is left to its own PR.
tp_rank_collapse: autokeeps native V4.1 one copy per rank, as on main, soTP4 stores four identical copies. Set
trueto store one; the pluginresolves
autoto true for V4.1.costs about 4% (no spec: 14.00 vs 14.55 req/s).
prefix's checkpoint is evicted, requests load past it: 1125 of 1311 loads in
the DSpark run started at token 0.
because V4.1 cannot mix a prefill with a speculative verify.
ends every prefix hit inside the prompt.
🤖 Generated with Claude Code