Repository navigation
DeepSeek-V4.1-Flash: LMCache tier offload through the MP server (vLLM plugin) - #2598
PerryZhang01 wants to merge 63 commits into
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
|
valarLip
left a comment
There was a problem hiding this comment.
Re-read at 3593a34, core offload / MP only this time; the plugin side I'll leave to its own reviewers. Thanks for the round: the slot-owner cleanup, V4's starts=, the double store on descriptor slot 1, _load_poll_lock, PAGE_LAYOUT_VERSION back at 3 and the dsv41 refusal for other models all read right. From reading the code; nothing was run on a GPU.
1. Native V4.1's namespace changes, and the description doesn't say so (offload/config.py:246-247). Native V4.1 used to fall through compress_ratios to hybrid; it now resolves to dsv41, so the page mode becomes dsv41-page-state and dsv41_geometry enters the namespace hash. A native lmcache_mp V4.1 deployment therefore misses on every object stored before the upgrade (the tier is cold after rollout), and an offload_layout override that used to be accepted is now refused. The description says "Native V4.1: same geometry". That may be the right call, but it's a deployment-visible change and needs to be stated in the description and the offload docs.
2. close() destroys the step events before the server has drained (mp/worker.py:357-362). self._step_events.clear() runs before adapter.shutdown(), and adapter.shutdown() is the step that drains in-flight operations. StepEventRing.clear assumes no transfer can still wait on an event, but TransferDrain.close doesn't wait for transfers to reach a terminal state, and InlineTransfers.close returns True at once. A save or load still pending at shutdown leaves the server waiting on an interprocess event this side already destroyed: a HIP invalid-argument error on the server, or a hang. Clearing the ring after adapter.shutdown() restores the precondition.
3. The last look costs up to 1 ms a step while any load waits (mp/worker.py:820-826). With the drain thread, get_finished calls wait_round(_LAST_LOOK_WAIT_S) (1 ms) whenever _wants_last_look_locked(), which is bool(self._pending_loads). The round has to submit what's queued first, and drain.py:17 measures a submit at 18.7 ms, so the 1 ms window almost never catches a completion. For as long as any request sits in WAITING_FOR_REMOTE_KVS, each step spends up to 1 ms of step-thread time for nothing. That's the hot-path cost the drain thread was added to remove. Waiting only when nothing is queued to submit (a pure poll round), or not at all, would fix it.
4. Still open from 10-29 (item 8): tp_rank_collapse=auto never collapses V4.1 (mp/deployment.py:161-167). The shared rule returns False for every state-owning layout "until their backend declares otherwise", but the V4.1 backend already declares it: tp_replication_factor=tp_size and native_state_tp_replication_factor=tp_size (deepseek_v41/backend.py:501-502). So native V4.1 at TP4 with auto has every rank write, and the tier holds 4 identical copies. The recipe sets false explicitly, so auto is never exercised. Having the rule read the factor the backend declares (Kimi-K3 declares less than TP and stays uncollapsed) would replace the layout list with the real criterion.
5. Dead drain-era scaffolding in NativeState (mp/native_state_worker.py:53, :221-233). 3593a34 put NativeState back on InlineTransfers, but the drain-thread machinery stayed:
queuedis set and cleared, but nothing observes it, so thenot entry.queuedfilter never excludes anything;- the
_transfers_in_flight_lockedoverride (:376) only feedswake(), which is a no-op inline; _outstanding_locked(:383) is only reached throughfence, which only the plugin's drain calls;InlineTransfers.fenceand_FENCE_POLL_S(drain.py:277-296) run only in a unit test;worker.py:820still branches onself._transfers.inlineinstead of the two strategies sharing onewait_roundcontract.
That's ~100 lines no production path exercises. Please either return NativeState to main's polling in get_finished, or make it actually share the drain.
6. Unrelated or out-of-date bits:
metadata.py:51only lengthens a chunk-size error and adds a comment about Kimi-K3's 128 → 1536 block size. That's a K3 change, so please move it to its own PR.tests/test_lmcache_mp.py:732/752/883rewritelambda: next(ticks)asiter(...).__next__, which changes nothing; please revert.chunked_scheduler._retire_finished_save(:674) hand-writes the release-lease sequence a third time, next to_source_group_finished(:1239) and_release_operation_lease(:1365). One helper would keep them from drifting.- The description still lists
RecurrentPageGroup/RecurrentStateTransfer/install_recurrent_state_hook,disaggregation/types.pyandmooncake_store/scheduler.py, none of which are in the diff. Itsmp/scheduler.pyandchunked_scheduler.pyrows describe changes that aren't there, while it omits items 1 and 4.
- NativeState submits on the step thread and polls in `get_finished`, as on main, holding the step's event from the ring until the transfer is terminal. The drain-era `queued`/`terminal` state, the `_perform_submit`/`_poll_round` overrides and `_outstanding_locked` are gone. - `get_finished`'s last look waits only for a pure poll round (nothing queued to submit, no round running): a round that has to submit first outlasts any wait short enough for a step. Dropping the look entirely raised the plugin's c16 TTFT p50 from 302 to 324 ms. Both executors share one `report_round(wait_s)` (inline: the step's poll; drain: a wake and that bounded look), so nothing branches on the strategy. `InlineTransfers` keeps only what the worker calls; its `fence` refuses, since only the plugin fences. - `close` destroys the step events only after `adapter.shutdown()` has drained the server, and keeps them when the shutdown fails. - The three lease releases in `chunked_scheduler.py` share `_return_leased_blocks`. - Reverted: the Kimi-K3 chunk-size message in `metadata.py` and the `iter(...).__next__` test rewrites. - The offload README states native V4.1's `dsv41` layout, its one-time cold tier after upgrading, and that `tp_rank_collapse: auto` keeps state-owning layouts one copy per rank. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 15 pushed 7081137 for the second pass of the 11-02 review: items 1, 2, 3, 5 and 6 are fixed; 4 stays as on main.
Per item
Runs (cards 4-7, same container)Native and plugin two-pass ran at 28593d185; 7081137 changes only the plugin's last look.
Plugin, DSpark, two-pass at 28593d185: p1 0.9219, p2 0.9272; pass-2 tier hits 5.84M / 7.18M; 0 stale state. Plugin perf, DSpark, c16, 900 s, OSL 64, HBM pool 4000 blocks, one run per row:
|
Conflicts were with #2906 (lmcache_mp lookup thread, transfer thread, same-step saves) in `mp/worker.py` and `mp/native_state_worker.py`. Native ATOM's MP path follows main: its producer event is `_producer_fence` (a fresh event per step, or `_ProducerFences` under OFFLOAD_MP_TRANSFER_THREAD), recorded on the forward thread and passed as a `RecordedStepEvent`; submissions go through `_dispatch` and the optional `_TransferThread`; same-step saves, `_submit_kwargs` and `_finished_housekeeping` are main's. The step-event ring now serves the plugin's drain thread only, which does not start main's transfer thread. The base worker's `transfer_thread` argument is renamed `drain_thread` so it does not read as main's transfer thread. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 16 pushed 9e70cbc, merging main 498133e: native ATOM keeps #2906's MP path; the step-event ring serves only the plugin's drain (
Runs (cards 4-7, same container)
Plugin, DSpark, two-pass at 9e70cbc: p1 0.9189, p2 0.9272; 0 stale state, 0 invalid-argument waits. Plugin perf, DSpark, c16, 900 s, OSL 64, HBM pool 4000 blocks, one run each:
|
`OwnedFutures` now holds futures only for the plugin's drain thread (`hold=drain_thread`). Native ATOM, all on the step thread, drops a future where it finished, as on main. Held until the stream was idle, native's steps ran 6% faster but its saves fell behind: a GSM8K two-pass on V4.1 loaded 7% fewer tier tokens in pass 2. Dropped, pass time, stores per second and pass-2 loads sit within main's spread. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 17 pushed 97c8dde: native ATOM drops a finished future where it finished, as main does; holding futures is now the plugin's drain thread alone.
Native V4.1, LMCache MP, GSM8K two-pass (cards 4-7, same container)Held futures cost no accuracy but made pass 1 end sooner while saves kept their pace, so pass 2 found fewer of them. A counter showed the stream busy at most step starts, holding up to 422 futures.
97c8dde changes only native's future lifetime, which the "+ drop" rows ran; the plugin keeps its drain path unchanged. |
`dsv41` changed native ATOM's V4.1 namespace (page mode `dsv41-page-state`, plus `sliding_window` and `candidate_block_size`), so a tier an earlier build filled missed on every object after the upgrade. Nothing native needed from it: its native-state key already carries the checkpoint layout. `build_page_namespace` now gives `dsv41` the `dsv4-page-regions` page mode `hybrid` gave V4.1 and hashes no V4.1 fields, so native V4.1's namespace is byte-identical to main's. The vLLM plugin, which had no V4.1 tiers before, adds the two fields to the HF geometry its config view snapshots, in the worker and the scheduler alike. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 18 pushed 85d55d8: native V4.1's namespace is main's again, correcting my reply to item 1 of the 11-02 review; its tiers stay reachable.
|
valarLip
left a comment
There was a problem hiding this comment.
Re-read the core offload / MP side at 85d55d8, by reading the code and running 7 MP worker / scheduler / config test files (329 passed); no GPU. The plugin side I'm leaving to its own reviewers. Thanks for rounds 16–18: native now keeps #2906's MP path, drops finished futures as main does, and V4.1's namespace is main's again, which settles my earlier namespace and drain-cost points for native.
What's left is in the core files that this PR adds for the drain (used today only with drain_thread=True), plus two scheduler fixes. Each item ends with the change I'm asking for.
A. Correctness
1. StepEventRing has no re-record bound (mp/step_events.py:39-49). record_locked pops a free index and records the same interprocess event again, with no count. mp/transfer_thread.py:35-37 documents that on ROCm 7.2.4 a server stream wait on an event its exporter has recorded more than 32 times fails, which is why _ProducerFences swaps in a fresh event after _FENCE_MAX_RECORDS = 31. In steady state the drain's transfers finish before the next step, so the ring re-records the same index every step and passes 32 records after ~32 steps. The two modules also disagree on the semantics: the ring says re-recording an event a transfer still waits on errors, while _ProducerFences says the wait just becomes longer. (The invalid argument waits mentioned earlier in the thread fit this.)
→ Change: settle which statement is true and write it once. Give the ring the same per-event record count and replacement as _ProducerFences, or better, item 3.
2. The drain bypasses the pooled-producer-event refusal (mp/worker.py:482-490). if not envs.OFFLOAD_MP_TRANSFER_THREAD or self._drain_thread: return skips the check that an adapter with its own _producer_events must keep transfers on the forward thread. In drain mode _fences is None, so producer_event_recorded=True is never passed either. An adapter that records its pooled event on the drain thread's current stream records it after nothing in particular, so a save can read KV the forward hasn't finished writing, and a load can overwrite blocks the forward is reading. This depends on the adapter (I don't have LMCache here), but the guard exists for exactly that adapter.
→ Change: apply the same refusal in drain mode (_adapter_takes_recorded_events() or refuse), and pass producer_event_recorded from the drain's submit path as the transfer thread does.
3. A refused or withdrawn load doesn't lower the save floor (chunked_scheduler.py:1419, :854). load_failed lowers entry[1] only when _load_save_floors[sid] is set, and that's written only at dispatch (:854). The tracker's floor is set to the whole tier hit at tracking time (chunked_scheduler.py:575, _lmcache_hit_save_floor). So a load refused before dispatch (e.g. tier_lost_prefix, the tier evicted the prefix) leaves the floor at the hit: the request recomputes [HBM, hit) but never stores that span back, the exact span the tier just lost.
→ Change: record the floor when tracking starts, not at dispatch, so load_failed can always restore it, or have load_failed fall back to the HBM hit when no dispatch floor exists.
4. A withdrawn promise leaks the lookup's read locks (mp/scheduler.py:283-292). With active is None, a raw-id load_failed counts as current and calls complete_retrieve(False) for a retrieve that was never submitted. That pops _lookups, so the following clear_lookup_status sees state is None and skips free_lookup_locks. The server then holds the read locks until session or TTL cleanup, blocking eviction and writes for that prefix. The same withdrawal also spends the request's one tier attempt through _record_failed_load_attempt. (From reading; forcing a promise past its 50 steps and checking the server lock count would confirm it.)
→ Change: for an operation that was never submitted, release the lookup's locks directly (free_lookup_locks) instead of complete_retrieve, and don't count it as a failed tier attempt.
B. Structure
5. The drain is a second implementation of the transfer thread, living in core (mp/drain.py:43, mp/step_events.py, mp/ipc_lifetime.py). TransferDrain, StepEventRing / PendingStepEvent / RecordedStepEvent, OwnedFutures, and InlineTransfers (an empty object that exists to match a constructor, drain.py:238) re-implement _TransferThread, _ProducerFences and take_spent from transfer_thread.py. The only caller is drain_thread=True from the plugin connector. The worker now holds both self._transfer and self._transfers, one letter apart, and worker.py is 1028 lines. Items 1 and 2 are the two copies already disagreeing.
→ Change: add what the drain needs (a poll round and a fence) to transfer_thread.py's thread and fences, so there's one submission thread and one event pool. Then delete drain.py, step_events.py, ipc_lifetime.py, InlineTransfers and the drain_thread switch. If the plugin needs something native doesn't, it belongs in the plugin's adapter, not in the core worker.
6. Single-model hooks and knobs in shared code (chunked_scheduler.py:666, mp/worker.py:136, config.py:98).
_unbuilt_save_is_final()is a new hook on the shared base that onlyV41StateMPScheduleroverrides.lmcache.mp.null_block_idadds a layout field, a parameter, a check, a worker attribute and a README paragraph for a value only tests set._DSV41_MODEL_TYPESis a hand copy ofselector._V41_TYPES.
→ Change:
- In
v41_mp_scheduler._build_save_request, call_retire_finished_savewhere it declines, and delete the hook and its branch (:938-946). - Keep
-1as the constant inlayout.transfer_block_idsand delete the knob. - Use
attn_family(hf) is Family.CSA2instead of the copied set.
Addresses valarLip's core review of 85d55d8. The vLLM plugin's off-step submit and poll run on #2906's transfer thread instead of a second implementation of it (items 1, 2, 5): - `_TransferThread` takes an optional `poll`. With it, a round of polls runs whenever the queue is empty, `look` wakes it from `get_finished`, and `fence` waits on its rounds up to the transfer deadline. `hand_back`/`take_returned` drop a polled future on the forward thread once LMCache no longer holds it. - `LMCacheMPConnector.run_transfers_off_step` starts it with `_ProducerFences`, so the plugin's events have the env-gated thread's re-record bound, and refuses an adapter that records its own pooled event: neither thread could fence it. - `drain.py`, `step_events.py`, `ipc_lifetime.py` and the `drain_thread` switch are removed. Native ATOM keeps main's inline path and its env-gated transfer thread. Scheduler fixes (items 3, 4): - The save floor a failed load restores is recorded at allocation, and a load refused as `tier_lost_prefix` reopens `[HBM, hit)` for saving, in the chunked and DSV4 schedulers: the request recomputes the span the tier lost, and it was never stored back. - A promise withdrawn before dispatch no longer calls `complete_retrieve`, which forgot the lookup first, so the lookup cleanup frees its read locks. - The withdrawal still counts as the request's failed tier attempt; that is what stops a re-promise/withdraw loop (test_dsv4_failed_load_is_not_retried_for_the_same_request, test_a_failed_load_leaves_no_remembered_hit_behind). Single-model hooks (item 6): V4.1's MP scheduler calls `_retire_finished_save` where it declines, replacing the shared `_unbuilt_save_is_final` hook; `lmcache.mp.null_block_id` is dropped for the `-1` constant; V4.1 is detected with `attn_family(...) is Family.CSA2`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A future the plugin's polling transfer thread finished holds the server's imported IPC event, and on ROCm 7.2.x destroying one synchronises every stream of the device. 0132de1 dropped them in `get_finished`, mid-forward: three interleaved c16 pairs against 85d55d8 ran 3.8% fewer requests per second (8.65 vs 8.99) with ITL p50 23.85 vs 23.05 ms, every new run below every old one. The forward thread now drops them only while its stream is idle, as 85d55d8's drain did, and a stream that is never idle (async scheduling) drops a batch per call past 8192 held. Native ATOM hands no futures back and is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 19 pushed 4b589d8. The six items of the 11-03 core review are addressed, except 4b.
Evidence4b: without the spent attempt, Unit, in the PR's container at 4b589d8: the CI command and GPU, MI355X TP4, DeepSeek-V4.1-Flash DSpark, FULL_AND_PIECEWISE:
0132de1 dropped the polled futures mid-forward and ran 8.41-8.77 req/s, ITL p50 23.71-24.07 ms; 4b589d8 drops them on an idle stream. 4b589d8 changes nothing native runs. |
| plugin, where a submit cost 18.7 ms and a terminal poll 13.8 ms under | ||
| its launch pressure. Called once the caches are registered. | ||
| """ | ||
| if self._transfer is not None or self._rank is None: |
There was a problem hiding this comment.
plugin + this env hard-fails startup here, misleading msg. replace the non-polling thread instead, or skip the env-gated start when frontend will poll?
There was a problem hiding this comment.
Fixed in 45355d2. The plugin builds the worker with transfers_off_step=True, and registration starts the polling thread whatever OFFLOAD_MP_TRANSFER_THREAD says; run_transfers_off_step is gone. Covered by test_off_step_transfers_poll_whatever_the_transfer_thread_flag_says.
| # from a remembered hit has none yet. | ||
| if not self._ensure_lookup_pin(seq, sid, ls): | ||
| self._mark_load_skip(seq, "tier_lost_prefix", hbm, lmc, need, chunk) | ||
| # The tier no longer holds the prefix this request recomputes. |
There was a problem hiding this comment.
same 3 lines inline twice, and chunked scheduler has _reopen_saves already. move it into OffloadSchedulerMixin and call from both?
There was a problem hiding this comment.
Fixed in 45355d2. _reopen_saves is on OffloadSchedulerMixin; the DSV4 scheduler calls it for tier_lost_prefix and in load_failed, and the chunked scheduler's copy is gone.
valarLip
left a comment
There was a problem hiding this comment.
Re-read the core offload / MP side at 4b589d8, by reading the code; no GPU. Thanks for round 19. The drain modules are gone, the plugin now runs on _TransferThread with _ProducerFences, a self-recording adapter is refused, the floor is set at allocation, and a withdrawn promise frees its locks. I accept your reason for keeping the spent attempt (4b). I'm not repeating the DSV4 inline reopen (valar-cop's comment). Each item ends with the change I'm asking for.
A. Correctness
1. The polling thread polls only when its queue is empty (mp/transfer_thread.py:280-310). The _WAKE branch is if self._queue.empty(): self._poll_round(), the post-item poll has the same condition, and the _IDLE_POLL_S backstop fires only on queue.Empty. So while the queue never drains, every wake from look() and from a fence is dropped, which contradicts the module docstring ("polls when woken"). Under heavy prefill, a submit takes several to ~20 ms and an after-forward save also waits in _wait_ready. If each step's submits take at least a step's time, the queue never empties: loads that have finished are never polled and requests sit in WAITING_FOR_REMOTE_KVS behind background saves, and a fence's round never comes, so it waits out the transfer deadline and fail-stops. (From reading. Watching _rounds_done under concurrent long prefills would confirm it.)
→ Change: make a wake from look() / a fence force a poll round on the next loop iteration, whatever the queue holds (e.g. a _poll_requested flag checked before taking the next item), or put loads and saves on separate queues so a load's completion isn't behind saves.
2. The in-process offload shell doesn't refuse dsv41 (offload/connector.py:40-108). V4.1 now resolves to the dsv41 layout (offload/config.py:234-235), but _build_worker / _build_scheduler have no branch for it and fall through to DenseOffloadConnector / DenseOffloadScheduler. On main it went to DSV4 hybrid. The selector's own comment says the dense codec must never serve V4.1 (it would restore KV over a window ring that doesn't exist). Nothing breaks today only because two gates in other files stop it first (deepseek_v41/config.py:382 and the plugin's refusal). The README says dsv41 is lmcache_mp-only.
→ Change: in the shell, raise for dsv41 with a message pointing at lmcache_mp, so the one place that maps layouts to implementations can't route V4.1 to the dense codec.
B. Performance
3. Up to 1 ms of step-thread wait per step while any load is pending (mp/worker.py:960-967). wait_s = _LAST_LOOK_WAIT_S if self._pending_loads else 0.0, then transfer.look(wait_s), on every get_finished, decode-only steps included. In your DSpark run (82.9% of tokens from the tier) a load is in flight almost every step, so that's up to 1 ms on a ~23 ms ITL. With item 1 the poll round runs when woken, so the step thread needn't wait for it.
→ Change: look() should wake the poller and return. Report what the previous round polled, and let a load that completes in between be reported next step, as main does. If you keep a wait, show the ITL with and without it.
C. Boundaries
4. A plugin-only transfer mode lives in the core worker (mp/worker.py:422, :932, :952, :960). run_transfers_off_step() is called only by the plugin. fence() raises NotImplementedError in the other modes, and get_finished / _retire_future branch on transfer.polls, so what the core worker does depends on who drives it. Future retirement is inferred from sys.getrefcount (transfer_thread.py:62-64, :184). Also:
native_state_layout.py:24importsdeployment._storage_kv_transfer_config(private) for a check that belongs indeployment.py;- the offload README still describes
mp/worker.pyas the "Generic PAGE-only" half (:110) and doesn't mention the plugin constructingLMCacheMPConnectordirectly.
→ Change: give _TransferThread the poll round as a constructor option the front end passes when it builds the worker, so the worker has one mode with a flag, not a mode switched on afterwards; then delete run_transfers_off_step. Retire futures by completion state (future.done() plus the stream-idle check), not refcounts. Move the native-state server check into deployment.py as a public function, and update the README table and the paths section.
5. Native's inline path lost its test coverage, and native_state_worker carries no-op rewrites (tests/test_lmcache_mp.py:1121, mp/native_state_worker.py). The shared _worker fixture now starts the plugin's polling thread (_run_off_step), so the ~23 failure-path tests no longer exercise native's inline get_finished; 3 test_inline_* tests remain. Yet that path was rewritten: load polled before save, lock-free snapshot, results through _polled_*. The query-exception test also lost main's "recovers and reports finished_*" assertion. Meanwhile native_state_worker.py (+73 / −187) has changes with no behaviour change:
_submit_native's check-and-insert is split into two lock sections;_poll_native_loadis split up;_emit_page_source_safeis renamed;restore_unprovablereplaces main'spending.future = _UnprovableSubmission().
→ Change: parametrize the worker fixture over inline and polling, so every failure-path test runs on both, and restore the recovery assertion. Revert the native_state_worker.py restructuring that doesn't change behaviour, keeping only what the poll round needs.
…w fixes Addresses valarLip's and valar-cop's reviews of 4b589d8. - `LMCacheMPConnector(config, transfers_off_step=True)` replaces `run_transfers_off_step()`: the frontend asks for the polling transfer thread when it builds the worker, and registration starts it whatever OFFLOAD_MP_TRANSFER_THREAD says. With that flag set, plugin startup no longer fails after the env had started a submit-only thread. - A round asked for by `look` or a fence runs after the submission in progress, whatever is still queued, so a backlog of submissions holds back neither a finished load's report nor a fence. - The in-process offload shell refuses the `dsv41` layout instead of routing DeepSeek-V4.1 to the dense codec. - `_reopen_saves` lives on `OffloadSchedulerMixin`; the DSV4 scheduler calls it in place of its two inline copies. - `require_native_state_server` is public in `mp/deployment.py`. - `native_state_worker.py` is main's apart from the registration that moved to the base worker; the rewrites that changed no behaviour are reverted. - The worker tests' failure paths run on both the step-thread worker and the polling one, and the query-exception test checks recovery again. - README: the module table and the plugin's direct use of the worker. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 20 pushed 45355d2. All five items are addressed, except retiring futures by
EvidenceItem 4: the worker takes Item 5: Unit, in the PR's container at 45355d2: the CI command and GPU, MI355X TP4, DeepSeek-V4.1-Flash DSpark, FULL_AND_PIECEWISE. "No wait" is the same tree with
|
valarLip
left a comment
There was a problem hiding this comment.
Follow-up on the current head 45355d23a. My previous review (on f5d97882e) covered only the non-plugin files; this one covers the whole PR, plugin included. Several earlier points are no longer reported as fixed or no longer applicable. The relevant CPU tests pass locally (test_v41_index_plane_parity, test_v41_plugin_refusals, test_lmcache_mp*, test_offload_shell: 240); no GPU runs. Statuses are noted inline where a point is plausible rather than confirmed.
Action items for the author
Affects other models — please fix first
- E1 Restore exact decode-keep starts for DeepSeek-V4 (
deepseek_v4_bridge.py:1372). - E3 Report block frees only when they actually return blocks; don't stall preemption on pins/stores (
boundary_planner.py:208). - E4 Revert the V4.1 state-probe default flip and keep the env escape hatch; move it to its own PR with a test (
v41_core_overrides.py:337). - E8 Scope
has_pending_push_work's scheduler check to MP (connector.py:2075).
V4.1 correctness
- E2 One STATE-tail owner that counts every KV group (
deepseek_v41_state_reserve_patch.py:107). - E5 Let checkpoint-restore failures raise (
v41_core_overrides.py:404). - E6 Decide restore by slot owner, not request-id text (
v41_core_overrides.py:352). - E9 Remove the
cache.pending = Nonebypass; run the checkpoint step next to the keep (deepseek_v41_bridge.py:2133). - E10 Return
(empty, 0)when the local hit caps to 0 (plugin/vllm/scheduler.py:262). - E12 Translate the connector name to
lmcache_mpinstead of clearingkv_transfer_config; close the module-path bypass (platform.py:254). - E13 Swap the transfer error under the lock (
transfer_thread.py:161).
Performance
- E7 Release returned futures on a recorded event, not stream idleness (
mp/worker.py:517). - E11 Coalesce WAKEs; don't wait on decode steps or interleave polls into the submit backlog (
transfer_thread.py:200).
Remove transitional design / duplication
- E14 One TP-rank-collapse rule; drop the config rewrite and connector V4.1 branch; explain or fix the recipe's
false(offload_config.py:172). - E15 Move plugin-only transfer modes out of the core worker; stop monkeypatching
paged_state_checkpoint.py; remove the copiedget_kv_transfer_tensors, the second MP scheduler, duplicate rules/helpers, forwarders and dead checks (mp/worker.py:425).
| name="DeepSeek-V4", | ||
| # The first-token positions `_build_dsv4_metadata` derives, exact under | ||
| # MTP; the input batch's V2 count is an optimistic upper bound. | ||
| starts=np.asarray(md.chunk_start_per_seq_cpu[:rows], dtype=np.int64), |
There was a problem hiding this comment.
[E1] DeepSeek-V4 (not V4.1) decode keep can restore wrong state. The decode-keep start is now md.chunk_start_per_seq_cpu, which on decode rows comes from seq_lens_cpu_upper_bound (all drafts accepted), instead of main's _exact_starts reading vLLM's corrected num_computed_tokens_cpu. With V4 + MTP + async scheduling (vLLM's default) + prefix caching on the V1 runner — nothing refuses async+spec for V4 — on a step whose drafts are rejected, by_end.get(starts[i]) looks up the optimistic end: it either stores stale state as boundary B's checkpoint or drops an already-hashed B. A later hit on that prefix restores wrong state and output is silently wrong. The comment's "exact under MTP" doesn't hold for decode rows (bridge:1724 is seq_np - lens). Please keep exact starts for V4.
| v41_proxy_state_reserve_blocks, | ||
| ) | ||
|
|
||
| if _state_checkpoints_own_the_tail(vllm_config): |
There was a problem hiding this comment.
[E2] One STATE-tail owner, counting every group. With checkpoints on, _state_checkpoints_own_the_tail makes the reserve patch return 0, so its own "one request must fit" check — which counts all KV groups — no longer runs; the remaining check in paged_state_checkpoint.py:470/503 computes per_req over the proxy group only. With DSpark + --enable-prefix-caching + align mode and a pool sized between the proxy-only and all-groups requirement, startup passes but a full-max_model_len request can never be admitted — exactly the livelock this patch's own comment describes. Two wrappers own the same thing; keep one, and count every group.
| """Whether only a further engine step can finish this leg's work.""" | ||
| return False | ||
|
|
||
| def has_pending_frees(self, horizon_s: float) -> bool: |
There was a problem hiding this comment.
[E3] Don't block preemption on frees that can't free anything. The new has_pending_block_frees makes vLLM break instead of preempting on an allocation failure, and this KDA leg returns True whenever _pending_stores is non-empty — ignoring the horizon and whether releasing the pin actually frees a block; the lease and deferred-free legs have no sole-holder check either. Kimi-K3 + LMCache under KV pressure: the pinned boundary block is still in a running request's block table, so releasing it frees nothing, but with KDA stores in flight (nearly every step at steady state) vLLM's scheduler.py:925 breaks, nothing is preempted and later running requests aren't scheduled that step — indefinitely while stores keep coming (lease/deferred legs stall up to ~10 s each). Only report frees that will actually return blocks to the pool.
| if override is not None and override != "": | ||
| blocking = override not in ("0", "false", "False") | ||
| else: | ||
| blocking = not serving_synchronously() |
There was a problem hiding this comment.
[E4] Revert this default flip (unrelated to offload). main's default blocking state probe (with the comment "keep this workaround until a separate test does") and the ATOM_V41_BLOCKING_STATE_PROBE escape hatch are removed in favour of blocking = not serving_synchronously(). This affects every V4.1 plugin run, connector and checkpoints or not: on the V2 runner with synchronous scheduling and FULL_AND_PIECEWISE, the blocking read added for stale-state reporting becomes a deferred read, with no test driving prepare_state under serving_synchronously() == True to show it's safe, and users who set the env silently lose it. Please move it to its own PR with that test.
| ) | ||
| for i in wrong | ||
| ] | ||
| try: |
There was a problem hiding this comment.
[E5] Don't swallow restore failures. _restore_stale_rows wraps the checkpoint restore in except Exception: logger.warning, defeating its fail-stop: if one row's checkpoint unit is null, state_checkpoint_unit_ids raises "refusing to run from an unrestored slot", no row in the batch is restored, the real cause is a WARNING, and core then reports "needs state at N, found -1" — a -1 the bridge just wrote — pointing at the stale-state contract instead of the failed restore. Let it raise.
| settles is reported on a later step.""" | ||
| with self._idle: | ||
| pure = self._outstanding == 0 and self._rounds_done == self._rounds_started | ||
| target = self._ask_round() |
There was a problem hiding this comment.
[E11] Don't wait on decode steps or interleave polls into the submit backlog. With any load in flight, get_finished makes the forward thread wait up to 1 ms every step, decode included; look() calls _ask_round and enqueues a _WAKE every step, forcing poll rounds into the middle of a submit backlog, and stale _WAKEs run extra rounds when dequeued. At c16 with prefix-hit loads that's up to ~1 ms per decode step (~4% of a 23 ms ITL), and one or two full poll rounds per step inserted into the save queue (this PR prices a terminal poll at 13.8 ms). Coalesce WAKEs and only wait when the step has no runnable work. (Plausible — worth counting rounds per step to confirm.)
| name = getattr(kv, "kv_connector", None) | ||
| if name is None and isinstance(kv, dict): | ||
| name = kv.get("kv_connector") | ||
| return name == "AtomLMCacheOffloadConnector" |
There was a problem hiding this comment.
[E12] The connector whitelist can be bypassed. plugin/config.py:483 clears V4.1's kv_transfer_config to {}, so the model's own validate_runtime_config never sees the connector, and this whitelist compares only the class name, ignoring kv_connector_module_path (which vLLM's factory prefers). {"kv_connector": "AtomLMCacheOffloadConnector", "kv_connector_module_path": "other.mod"} passes and loads other.mod's class, skipping every V4.1 refusal. Root fix: translate the vLLM connector name to ATOM's lmcache_mp instead of clearing the config.
| @@ -118,6 +161,106 @@ def raise_error(self) -> None: | |||
| if error is not None: | |||
There was a problem hiding this comment.
[E13] Swap the error under the lock. raise_error does error, self._error = self._error, None unlocked (LOAD_ATTR and STORE_ATTR are separate bytecodes), and the transfer thread check-then-sets unlocked too; this round adds the poll round as another writer and the fence as another reader. A GIL switch between the two bytecodes lets the forward thread overwrite a just-written exception with None. Deadline errors recur next round, but a submit exception that escaped the _UnprovableSubmission catch is lost for good. Do the swap under the _idle lock. (Plausible.)
There was a problem hiding this comment.
Fixed in a278334. raise_error swaps _error under the _idle lock, and the thread records a failure through _record_error under the same lock, keeping the first.
| role = getattr(kv_transfer_config, "kv_role", None) or "kv_both" | ||
| extra = getattr(kv_transfer_config, "kv_connector_extra_config", None) | ||
| extra = dict(extra) if isinstance(extra, dict) else {} | ||
| collapse = extra.get("lmcache.mp.tp_rank_collapse", "auto") |
There was a problem hiding this comment.
[E14] One rule for TP rank collapse. It currently has three contradictory answers: deployment.py:161-166 makes auto False for every state-owning layout (dsv41 included); this rewrites V4.1's auto to True; and connector.py:642-700's _page_tp_replication_factor has a V4.1-only branch whose comment claims a test binds the two sides, but no test in tests/ calls it. The README says auto is per-rank, while the recipe sets tp_rank_collapse: false without saying why. So native and plugin store the same model with different collapse factors, the recipe's published numbers use 4× tier bytes at TP4 (not the plugin's default path), and if either side's "not TP-sharded" claim drifts, other ranks silently read rank 0's bytes. Keep one rule — auto = the replication factor the backend publishes, resolved once in the worker — and remove the config-layer rewrite and the connector's V4.1 branch.
| @@ -296,23 +423,39 @@ def _flush_after_forward(self) -> None: | |||
| # -- submission thread --------------------------------------------------- | |||
|
|
|||
| def _start_transfer_thread(self, rank: int) -> None: | |||
There was a problem hiding this comment.
[E15] Please remove the transitional design and duplication. (a) Plugin-only machinery — transfers_off_step, the poll thread, fence, getrefcount-style take_returned — lives in the shared core worker, giving three transfer modes; the native subclass keeps a second pending mechanism (_native_saves / _native_loads) that fence meets only with NotImplementedError; and the _adapter_records_its_own_events branch is unreachable with released LMCache (dead). (b) The plugin monkeypatches ATOM's own paged_state_checkpoint.py in four layers, with a docstring claiming "the vendored installer only hooks V1" — wrong, base _apply_runner_patches already walks V1 and V2. (c) v41_page_planes and _publish_v41_state_checkpoints copy backend.get_kv_transfer_tensors line by line (a consequence of clearing kv_transfer_config, E12). (d) V41StateMPScheduler is a second implementation parallel to NativeStateLMCacheMPConnectorScheduler and calls past its parent to the grandparent; the prefill_tail rule is written in 5 places; checkpoint-group discovery 4–5 times (all() vs any() disagree); _req_id_of is defined 3 times; and the _planners tuple is forced to at most one element, with 6 no-op hooks and a 303-line test of an unreachable configuration. (e) Pure forwarders _build_v41_pool_geometry and mp_worker_adapter.register_kv_caches (whose docstring claims it "does not extend the MP worker"); _RETIRED_STATE_OPTIONS refuses options that never existed on main; _save_source_is_backed's is not None is a dead check (vLLM's null block is 0). Each concept should have one mechanism, with plugin-only modes kept in the plugin.
Addresses valarLip's E13 and valar-cop's note on 45355d2. - `raise_error` swapped `_error` out unlocked while the transfer thread checked and set it unlocked too, so a failure recorded between the two bytecodes of the swap could be overwritten with None and lost. Both sides now take the thread's `_idle` lock; the thread records through `_record_error`. - The tests' `settle()` drains the queue before it waits for a round: an asked-for round runs after the submission in progress, not after everything queued, so the round alone no longer proves every queued submission went out. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 21 pushed a278334: E13 and the E7, E11 and E15(a) ask to undo what your 11-03 and 11-04 reviews asked for or proposed; two are measured (details). They stay unless you withdraw those requests. The plugin items (E1-E6, E8-E10, E12, E14, E15 b-e) are not changed in this round. Mark which block merge. Evidence
Unit, at a278334 in the PR's container: the CI command and |
valarLip
left a comment
There was a problem hiding this comment.
Re-read the core files at a278334, by reading the code and running the affected core tests (offload / MP / native / config / shell: 798 passed, 1 skipped; MP worker: 228 passed); no GPU. I found nothing in core that produces wrong output. This review covers only new items. It doesn't restate E1–E15, and it doesn't ask again for anything you implemented at my 11-03 / 11-04 request. Each item ends with the change I'm asking for.
1. A failed poll round keeps polling and logging (mp/transfer_thread.py:336-346). When _poll() raises (e.g. _enforce_transfer_deadline times out), _record_error keeps the first error and the thread carries on. Every _IDLE_POLL_S (50 ms) it runs another round, raises again, and logs logger.exception with a full traceback. Fail-stop happens only when the forward thread next calls _settle_transfer_thread / fence. If the front end stops stepping (no new requests), the engine doesn't stop and the log grows by about 20 tracebacks a second. I haven't triggered this.
→ Change: once an error is recorded, stop scheduling poll rounds (or exit the loop) and log it once; leave the raise to the forward thread as now.
2. Native pays for _queued_requests, which only the polling mode reads (mp/worker.py:479-505). Each _dispatch onto the thread now wraps the submit in _submit_queued: it takes _lock twice and updates _queued_requests. Only fence and get_finished's polls branch read it (:881, :890), and both are polling-mode only. With native OFFLOAD_MP_TRANSFER_THREAD=1, every transfer takes two extra lock acquisitions and contends with the forward thread's get_finished, to maintain state nobody reads.
→ Change: count queued requests only when transfer.polls, and submit directly otherwise.
3. Native get_finished was restructured (mp/worker.py:786-878, :157-160). It used to be one locked pass. It's now _poll_round → _poll_pending_loads / _poll_pending_saves, each snapshotting under the lock and committing under it again. Results go into four instance sets (_polled_done_save, _polled_completions, _polled_done_load, _polled_failed_load), which get_finished then copies and clears under the lock. So native goes from 1 lock acquisition per step to up to 5, plus copies of both pending dicts, and the four sets are written and read within the same call. The description says native keeps main's inline submit and poll.
→ Change: keep main's single-pass poll for native, and replace the four sets with one pending KVConnectorOutput that the polling mode fills and get_finished swaps out.
4. "Reopen the save floor when a load doesn't go out" is written three ways (chunked_scheduler.py:569, :854, :831):
update_state_after_allocrecords_load_save_floors(new);- dispatch records it again at
:854; - the tier-lost path ignores the floor and calls
_reopen_saves(sid, hbm)directly at:831.
The comment at :566 gives "the tier lost the prefix" as a reason for the alloc-time record, which the code doesn't follow. In core the alloc-time record has no observable reader: its one consumer (:1413) is the raw-id load_failed, which core never issues. hybrid/dsv4/connector.py:2079 copies the same thing, where it's dead.
→ Change: record the floor once, have every load that doesn't go out (tier lost, withdrawn) go through one helper that reads it, and delete the DSV4 copy.
5. Measurements in code comments: mp/transfer_thread.py:25 ("18.7 ms per submit and 13.8 ms per terminal poll"), :56 ("ITL p50 23.1 -> 28.8 ms"), mp/worker.py:80 ("c16 TTFT p50 rose from 302 to 324 ms"), :85 ("c16 ITL p50 rose from 23.0 to 23.8 ms"), and :106, which repeats 18.7 / 13.8 ms. These date with ROCm / LMCache versions and are duplicated.
→ Change: reduce these comments to the constraint each one protects, and keep the numbers in the description.
6. Names that no longer match (mp/worker.py:261, :310, :595-601). In the base worker, _register_native_state_layout, the native_state.{i} registration keys, require_native_state_server and _native_layout / _native_block_ids now serve the path that restores no Active SLOT; the docstring itself says "This worker restores no Active SLOT". The base _check_state_destination(state, *, loading) begins with del loading, so two meanings share one signature.
→ Change: rename them for what the base does (e.g. checkpoint_image_layout), and express the native / base difference as data rather than an override whose base ignores its argument.
7. Duplicate helpers and checks (mp/worker.py:91, mp/transfer.py:135, mp/worker.py:905, mp/transfer_thread.py:212):
_req_id_ofis identical to the one in the plugin'smp_worker_adapter.py._transfer_operation_idspells the samegetattr(completion, "req_id", completion).- "only the polling mode can fence" is checked in both
worker.fence(transfer is None) and_TransferThread.fence(self._poll is None), with different wording.
→ Change: one req_id_of helper in core that the others import, and one fence-precondition check.
8. Core reads a plugin-only attribute (offload/config.py:409). build_page_namespace now uses getattr(config, "offload_speculative_config", None) or .... No core code defines that attribute; only the plugin's offload_config.py sets it. (Its precedent, offload_page_hf_geometry, is set by atom/config.py itself.) The or also treats a falsy stable form (e.g. {}) as absent and falls back to hashing the raw object, so scheduler and worker could compute different namespaces. In the same vein, enforce_transfer_deadlines (mp/scheduler.py:340) became public only for the plugin's getattr call.
→ Change: pass the frontend's stable speculative form through an explicit config field (declared in core, like offload_page_hf_geometry), test it with is not None, and keep the deadline method private behind a declared interface.
9. Docstrings that overclaim (chunked_scheduler.py:677, :700, :601):
_return_leased_blockssays it is "the one sequence every lease release runs", and_retire_finished_savesays it is "the one rule for every path that retires a finished request's save". Butabandon_save(:1386) andreclaim_stale_leases(:1083) still pop_save_lease_blocksdirectly._release_lookup's comment at:601("clear_lookup_status only clears the client's memo; it does not release worker pins") is the opposite of what this PR relies on inmp/scheduler.py:286, where MP'sclear_lookup_statusfrees the read locks.
Someone who trusts the first will add bookkeeping that the abort paths skip. Someone who trusts the second will add back the complete_retrieve whose lock leak this PR removed.
→ Change: route the two direct pops through _return_leased_blocks (or drop "the one"), and correct the :601 comment for the MP client.
10. README and docstring out of date (offload/README.md:34, config.py:188). The README still says "A GDN or K3 model cannot be overridden onto a layout that owns no state tier (dense)", but config.py:167 now refuses any mapped != natural for state-owning layouts, so qwen ↔ kimi_k3 is refused too. An operator who sets offload_layout: kimi_k3 on Qwen3-Next per the README gets a ValueError at startup. _layout_from_model's docstring still lists "dense/hybrid/kimi_k3/qwen".
→ Change: state the actual rule in the README and update the docstring's layout list.
Addresses valarLip's review of a278334. - A failed poll round stops the polling thread's rounds and is logged once; the forward thread raises it, as before (1). - A transfer thread that does not poll submits as main does: only the polling thread counts queued requests for its fence (2). - Native `get_finished` is main's single locked pass again. The polling thread fills one pending `KVConnectorOutput` that `get_finished` swaps out, in place of four sets (3). - The save floor a load reopens to is recorded once, when the load is decided. `_reopen_saves` reads it for a load the tier lost, now also in `_ensure_lookup_pin`'s refusal, which used to clear the floor first, and for a failed load. The alloc-time copies are gone; a promise withdrawn before dispatch reopens nothing, the tier still holds the span (4). - Code comments state the constraint, not the measurement (5). - Whether a checkpoint load restores into an Active SLOT is the class attribute `restores_into_slot`, not an override whose base ignored its argument (6). - One `req_id_of` in `disaggregation/types.py`; the fence's polling precondition is checked once, in the worker (7). - `Config.offload_speculative_config` is declared in core and tested with `is not None`; MP's deadline method is private again, called by the plugin's schedulers through `enforce_transfer_deadlines` (8). - Docstrings that claimed one lease release path, and the in-process description of `clear_lookup_status`, are corrected; the README and `_layout_from_model` state the actual override rule (9, 10). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The plugin builds PluginDenseScheduler, whose enforce_transfer_deadlines the connector now calls directly; these two tests had wired the core DenseOffloadScheduler in its place. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Round 22 pushed d7619fe: all ten items are done except the renames in 6, which are refused. Your review found nothing in core that produces wrong output; please approve next review, or name what blocks merge.
Evidence
New tests for 1, 2, 4 and 8 fail with their fix reverted. Unit, at d7619fe in the PR's container: the CI command and |
valarLip
left a comment
There was a problem hiding this comment.
Re-read the core files at d7619fe, by reading the code, running the worker / scheduler / config tests (817 passed) and a shutdown probe; no GPU. Thanks for round 22. The nine fixes read right, and I accept your reason on the item-6 rename: main names the format native_state. Items 2, 3, 5 and 7 below follow from my last round's items 3, 1, 4 and 7, so they're partly on me.
On approval: I'd like items 1, 3 and 4 fixed before I approve. One is a reproduced bug, and the other two are ways the polling thread can stop without anyone noticing. The rest can go in a follow-up if you prefer. Each item ends with the change I'm asking for.
1. close() races the polling thread (mp/worker.py:338-343). close() swaps self._transfer to None first and then spends up to _TRANSFER_DRAIN_TIMEOUT_S draining and joining a thread that is still polling. A round that meets a terminal future in that window calls _retire_future, which calls hand_back on None. Reproduced with a probe that shuts down with a save in flight: AttributeError: 'NoneType' object has no attribute 'hand_back'. The consequences:
- polling stops;
- that completion is popped from
_pending_*and never reported; - the future holding the imported IPC event is destroyed on the transfer thread, which the module docstring forbids;
- every such shutdown logs a "polling stops" exception.
→ Change: clear self._transfer only after transfer.close(...) returns, or have _poll_round capture the transfer object it was started with. Add the shutdown-with-in-flight test.
2. The background poll copies the foreground poll (mp/worker.py:788-860 vs get_finished :936-1000). _poll_pending_loads / _poll_pending_saves repeat get_finished's body line for line: take_completed_ranges, _terminal_future_result, _enforce_transfer_deadline, the tombstone, _source_safe_completions, building ConnectorCompletion. That's about 70 lines twice. They differ only in lock scope, _retire_future, and where the result goes. A future change to terminal or deadline handling has to be made in both.
→ Change: one _poll_pending(into, retire) that both the foreground pass and the polling thread call.
3. A failed poll is reported once, then the thread stays dead silently (mp/transfer_thread.py:156-165, :342). This comes from my item 1. After a failure, _poll_failed stays set and no more rounds run. But raise_error swaps _error to None, so the forward thread raises exactly once. If that one exception is caught above, the engine keeps running with no poller: every later load sits in WAITING_FOR_REMOTE_KVS forever and nothing is reported again. The inline path, by contrast, re-raises the deadline error on every get_finished.
→ Change: in _settle_transfer_thread (mp/worker.py:506), raise on every step while _poll_failed is set (keep the first error, don't clear it), so a dead poller can't go unnoticed.
4. A stuck submit stops every deadline (mp/transfer_thread.py:305). Submissions and poll rounds share one thread, and a round runs only between submissions. If one submit RPC hangs (the LMCache server stalls), no round runs, so _enforce_transfer_deadline never fires for any in-flight transfer. The failure then surfaces only through the fence timeout or the scheduler's own deadline, later and with less information. Before this PR, polling stayed on the forward thread even with OFFLOAD_MP_TRANSFER_THREAD, so it wasn't exposed to this.
→ Change: don't make deadline enforcement depend on the submit thread getting a turn. Check deadlines from the forward thread's _settle_transfer_thread (comparing against submission time is cheap), or bound each submit with its own timeout.
5. The tier-lost cleanup runs twice (chunked_scheduler.py:825-829, _offload_common.py:851-852). This follows from your item-4 fix. The mixin's _ensure_lookup_pin already calls _reopen_saves(sid) and _clear_pending_load(sid) before returning False, and the caller then calls both again. On the base path the second pair is dead (the floor is already gone). hybrid/dsv4/connector.py:2319 repeats it. The caller's pair only matters for native_state_scheduler's override, which returns False without cleaning up.
→ Change: do the cleanup in one layer. Either _ensure_lookup_pin only answers and every caller cleans up, or it cleans up and the native override does too. Delete the other copy.
6. A stale module docstring (mp/native_state_layout.py:10-14). It says every worker and scheduler "asks the same two questions here: whether saves carry images, and whether the server's chunk is the configured one". The second check, require_native_state_server, now lives in mp/deployment.py:367.
→ Change: either describe only stores_state_images here, or move the chunk check back.
7. req_id_of replaced one copy of the pattern (Fix-then-sweep). This follows from my item 7. The new helper is used in the worker and transfer.py, but the same hand-written unwrap is still at:
mp/scheduler.py:269and:284, right beside the changedload_failed;_offload_common.py:519,:524,:543and:546;hybrid/kimi_k3/connector.py:872.
→ Change: use req_id_of at every one of them.
8. The new override rule refuses combinations that used to start (offload/config.py:167). mapped != natural for state-owning layouts now refuses offload_layout: qwen on a Kimi-linear model and kimi_k3 on a GDN model, besides the V4.1 cases. The README now states the rule, but the description mentions only V4.1, and no test covers the newly refused pairs (test_lmcache_offload_config:231 tests dense only).
→ Change: confirm in the description that the kimi_k3 ↔ qwen refusal is intended, and add a test for that pair.
9. Comments still carry version findings (rest of my item 5). mp/worker.py:83 ("on ROCm 7.2.x destroying one synchronises every stream") and mp/transfer_thread.py:55-58 ("a fixed short cadence with loads outstanding cost more end to end than the forward-thread work it moved off") are findings, not constraints.
→ Change: state the constraint (e.g. "destroy imported events only on the forward thread at a step boundary"), and keep the version and the comparison in the description.
Addresses valarLip's review of d7619fe. - A polling round gets the thread it runs on and hands finished futures back to it; `close` forgets the thread only once it is closed. Shutdown with a save in flight no longer meets `hand_back` on None (1). - One poll pass, two callers: `_poll_entries` polls without touching worker state and `_settle_terminal_locked` reports what finished. Native `get_finished` runs both under one lock, as main's single pass; the polling thread polls unlocked and settles under the lock (2). - A failed poll round is raised on every `raise_error`, not once: no round runs after it. The failure is recorded and marked in one locked step (3). - The forward thread fail-stops a transfer past its deadline when no polling round has run since: a submit that never returns stops the rounds that enforce deadlines (4). - `_ensure_lookup_pin` only answers; the caller drops the load (5). - `native_state_layout`'s docstring names only `stores_state_images` (6). - `raw_req_id` unwraps a completion id in the type the engine compares, and `req_id_of` is its string form; every hand-written unwrap in core uses one of them. The sets returned to the engine keep raw ids: stringifying them fails test_offload_completion_processing_calls_load_ terminal_once (7). - A test covers the refused `kimi_k3` / `qwen` swap (8). - Comments state the constraint, not the finding (9). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
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:
|
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.pyrecurrent_page_groupsas recurrent engine groups, 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; it also fail-stops a transfer past its deadline that no round has polled since, as a submit that never returns would leave it. 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.pyenforce_transfer_deadlinesis public. 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).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
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
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