Skip to content

DeepSeek-V4.1-Flash: LMCache tier offload through the MP server (vLLM plugin) - #2598

Merged
valarLip merged 67 commits into
mainfrom
perzhang/dsv41-lmcache-031
Nov 9, 2026
Merged

valarLip merged 67 commits into
mainfrom
perzhang/dsv41-lmcache-031

Conversation

@PerryZhang01

@PerryZhang01 PerryZhang01 commented Oct 15, 2026 •

Copy link
Copy Markdown
Contributor

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:

  • Checkpoint groups. --enable-prefix-caching --mamba-cache-mode align
    registers 8 single-layer MambaSpec checkpoint groups. vLLM allocates and hashes them like any other
    group. 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.
  • MP transport. The worker registers the PAGE planes (5 planes, 739,840
    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.
  • Hit capping. vLLM already ends a local hit where the checkpoint groups
    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

File Change Why Effect on other models / paths
kv_transfer/offload/mp/worker.py, transfer_thread.py Given a published paged_state_checkpoint_spec, registers each checkpoint image ordinal as a recurrent engine group, one chunk wide. A worker built with transfers_off_step=True (the plugin's) runs #2906's transfer thread with a poll round, whatever OFFLOAD_MP_TRANSFER_THREAD says: it submits and polls off the step thread, with fences from _ProducerFences. A round get_finished or fence() asks for runs after the submission in progress, whatever is still queued. While a load is in flight get_finished waits up to 1 ms for that round when it is a pure poll (nothing queued, no round running). fence() waits on its rounds up to the transfer deadline. A failed round stops the polling, and the forward thread raises it on every step; rounds enforce every pending transfer's deadline, and the scheduler's transfer deadline bounds a submit that never returns. The thread polls on a timer only while transfers are in flight. A polled future holds the server's imported IPC event: it is handed back and dropped on the step thread once LMCache no longer holds it, while the stream is idle. An adapter that records its own producer events is refused. A chunk with no image names -1. Carries state with PAGE KV in one object. At chunk > block, registering at block width skipped every store (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. Without recurrent groups, registration is unchanged. Native ATOM keeps main's single-pass get_finished, producer fence, future lifetime and submit-only OFFLOAD_MP_TRANSFER_THREAD thread.
kv_transfer/offload/mp/native_state_worker.py, native_state_layout.py, native_state_scheduler.py NativeState submits on the step thread and polls in get_finished, with main's producer fence, as on main. Registration moved into the base worker; stores_state_images moved to native_state_layout.py and require_native_state_server to deployment.py. One registration path for native and the plugin's recurrent groups. Native V4/V4.1 submit and report when main does.
kv_transfer/offload/chunked_scheduler.py, hybrid/dsv4/connector.py _retire_finished_save retires 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 as tier_lost_prefix or failing reopens [HBM, hit) for saving through OffloadSchedulerMixin._reopen_saves, which reads it (both schedulers). A tracker left behind holds its lease forever, and a block still read must stay leased. A request whose tier hit is lost recomputes that span, which was never stored back. Any model: a failed save of a finished request keeps the blocks another of its operations still reads, and a lost or failed tier hit is stored again.
kv_transfer/offload/mp/scheduler.py The plugin's schedulers call the private _enforce_transfer_deadlines through their enforce_transfer_deadlines. A promise withdrawn before dispatch no longer calls complete_retrieve, so the lookup cleanup frees its read locks. The plugin drives the deadlines every step. complete_retrieve forgot the lookup first, so the server held the locks until session or TTL cleanup. Any MP model: a withdrawn promise frees its lookup locks.
kv_transfer/offload/config.py, connector.py, mp/deployment.py DeepSeek-V4.1 (CSA2 attention, by attn_family) resolves to its own dsv41 layout instead of DSV4's hybrid, and the in-process lmcache_offload shell refuses it. Its namespace is the one hybrid gave it; the plugin's config view adds sliding_window and candidate_block_size to the geometry it hashes. A state-owning model accepts only its own layout as an override (qwen and kimi_k3 are refused for each other too), and no other model may name dsv41. The namespace takes a frontend's stable speculative config when it supplies one (Config.offload_speculative_config). tp_rank_collapse: auto keeps every state-owning layout one copy per rank. Under hybrid, V4.1 accepted a dense override, 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. Other models' PAGE namespace is unchanged (PAGE_LAYOUT_VERSION stays 3). Native V4.1's namespace is byte-identical to main's, so its tiers stay reachable; an offload_layout override it accepted before is refused.
models/deepseek_v41/pool_geometry.py, deepseek_v41/backend.py build_v41_pool_geometry(), one source of the pool geometry for native and plugin, with every argument that moves it required; it takes main's kv_dtype (#2725) and index_layout (#2809). Removes a duplicated geometry (and its parity test). Beside the model's config, so a frontend that sizes the pool imports neither the attention backend nor AITER, and pool_layout still imports nothing of atom. Native V4.1: same geometry.
plugin/vllm/paged_state_checkpoint.py Decode keeps take the caller's exact starts and upload their descriptors on a buffer of their own. The adapter base gains prefill_tail, ATOM's checkpoint defaults, and sizes_pool_without_checkpoints. V2's input-batch count runs ahead of the step under speculation, so callers pass the step's own 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. V4 decode keeps take their starts from V4's metadata; V4 is otherwise unchanged.
plugin/config.py For V4.1, vLLM's transport is not carried into the ATOM Config. The model's validator knows only ATOM's connector registry, where every vLLM name is unknown. plugin/vllm/platform.py gates the connector by vLLM's names. Other models still see the transport.

Plugin changes

File Why
plugin/vllm/kv_transfer/connector.py The V4.1 path in AtomLMCacheOffloadConnector. It refuses inproc, 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 through failed_recving and withdraws promised loads. A save reads vLLM's request off the SeqView, so a finished request's tail save still finds it.
plugin/vllm/kv_transfer/v41_state.py, page_ride.py V41BoundaryPlanner (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.py Publish V4.1's PAGE planes and checkpoint groups to the MP worker. vLLM's preemption fence is the worker's fence(). A failed load reports only the blocks it writes.
plugin/vllm/kv_transfer/boundary_planner.py, kda_state.py The rules every recurrent-state planner shares (hit caps, lookup arming), now the base of Kimi-K3's KdaBoundaryPlanner and V4.1's planner.
plugin/vllm/kv_transfer/offload_config.py, seq_view.py The plugin's config view: stable speculative fields, the HF geometry snapshotted before a worker builds the model, and tp_rank_collapse: auto resolved to true for V4.1. SeqView gains num_tokens and _deferred_save_at, which the shared scheduler reads and writes.
plugin/vllm/paged_state_checkpoint_scheduler.py An anchor never leaves fewer tokens to compute than the prefill tail (one, two under DSpark).
plugin/vllm/deepseek_v4_bridge.py, deepseek_v4_prefix_patch.py, deepseek_v4_state_checkpoint.py V4's decode keeps pass the step's own starts; the profiling mark is set before the original KV init runs; V4's checkpoint defaults come from the shared adapter base.
plugin/vllm/kv_transfer/plugin_schedulers.py, v41_mp_scheduler.py ATOM's MP and in-process schedulers as this frontend drives them: which loads a build declined, whether a recent save or lease will free blocks (vLLM's has_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.py The V4.1 adapter for the PAGE-backed checkpoint mechanism on main: the layer count and spec come from V41PoolGeometry. The checkpoints are on with --enable-prefix-caching and --mamba-cache-mode align.
plugin/vllm/v41_core_overrides.py VllmBackedPagedAttentionCache restores 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.py Builds each step from the host where native ATOM does: PAGE rows from the host rows of the proxy's KV cache group, and lengths from vLLM's host upper bound once vLLM's warmup has returned under synchronous scheduling (that warmup schedules every draft as accepted, so its host count runs ahead); otherwise from the device. Stores the decode keeps the last step committed to at the next build, before it stages. Stores checkpoints around the forward and sets the cache's restore on every bind. Binds the proxy at register_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.py Clip prefill chunks to the checkpoint grid. Under DSpark, decline a local hit that would leave a one-token tail (vLLM would pad it into a verify row).
plugin/vllm/spec_decode_patch.py, v41_speculative.py Keep checkpoint layers out of the proxy and draft groups. Count a kv_holding_waiting request 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.py Keeps the V2 runner's block-table rows on the host, rebuilt from the scheduler output on V4.1's runner only, for the checkpoint copier and the V4.1 build. The first runner that keeps the rows arms them for readers. serving_synchronously opens after Worker.compile_or_warm_up_model, on V2 under synchronous scheduling.
plugin/vllm/platform.py, model_wrapper.py, deepseek_v41_state_reserve_patch.py Admit prefix caching only together with the connector or the checkpoints. Register the checkpoint layers, and refuse them with async speculative scheduling (its optimistic CPU counts would misplace checkpoints). Do not withhold the state tail twice. Skip vLLM V2's profiling capture for V4.1, which runs before the KV cache exists.

Tests

All runs are on MI355X (cards 4–7), TP4, FULL_AND_PIECEWISE with
max_cudagraph_capture_size 512, at 51a1ef540, before the rebase onto main
d6681e153 and the review fixes since. The accuracy and timing reruns at
a7d323eeb are in this comment. Each run gets a fresh model and MP server. The full commands are in
recipes/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.

Configuration Pass 1 Pass 2 Pass 2 from HBM / tier
DSpark, tier only 0.9272 0.9303 0 / 84.3%
No spec, tier only 0.9295 0.9265 0 / 84.8%
DSpark, HBM + tier 0.9151 0.9158 21.8 / 74.7%
No spec, HBM + tier 0.9181 0.9128 19.5 / 77.0%
DSpark, preamble, LMCache off (control)¹ 0.9227 0.9204 –

Zero engine faults and zero STORE block ID underflow. Pass-to-pass
differences 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 on
this 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.

Configuration req/s TTFT / ITL avg (ms) HBM / tier / recomputed
No prefix reuse (as on main) 4.62 361 / 49.1 –
HBM prefix caching, LMCache off 6.90 286 / 32.2 57.3 / 0 / 42.7%
HBM prefix caching + LMCache 7.80 (+13%) 387 / 26.3 11.4 / 82.9 / 5.7%

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).

  • Hit rate. The ceiling is 97.4% (6912 of 7097 tokens fall on the 256
    grid). Measured: 94.3%.
  • Throughput. Per-request engine time is 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%).
  • Transfer. One load is 21.3 MiB per rank, which would take 0.45 ms at
    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/plugin plus the offload, LMCache-MP and V4.1 suites,
and the CI command (tests/, without tests/plugin): no failure that main
does not also have in the same container. This PR's tests are in
tests/plugin/test_vllm_dsv41_lmcache.py; the other test changes adapt
existing suites.

Since the previous revision

  • Merged main 10b99a085. It brings feat(indexer): migrate V4 and V4.1 FP4 scoring to the OPUS config API #2809, whose aiter API change failed
    this branch's V4.1 accuracy jobs at import. build_v41_pool_geometry takes
    feat(indexer): migrate V4 and V4.1 FP4 scoring to the OPUS config API #2809's index_layout: native passes index_cache_layout as on main, and
    the plugin passes None for its FP8 index plane.
  • Review of 11-07 (valarLip). The forward thread's deadline check is
    removed: polling rounds and the scheduler's transfer deadline bound every
    transfer.
  • Review of 11-06 (valarLip, second). A transfer that finished is no
    longer failed by the forward thread's deadline check while a slow submit
    holds rounds back.
  • Review of 11-06 (valarLip). Deadlines are enforced from the forward
    thread every step. A self-recording adapter is refused before registration,
    and close leaves the adapter to a thread that did not stop. The idle
    polling thread no longer spins. A lookup that did not answer reopens no saves.
    Auto rank collapse is main's check again.
  • Review of 11-05 (valarLip, second). A polling round hands futures back
    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_id or
    req_id_of.
  • Reviews of 11-04 and 11-05 (valarLip, valar-cop). The transfer thread's
    error is swapped under its lock, and a failed poll round stops the polling.
    Native get_finished is main's single locked pass again, and native's
    submit-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.
  • Reviews of 11-04 (valarLip, valar-cop). The plugin asks for the polling
    worker with transfers_off_step when it builds it, so the env flag no longer
    breaks its startup. A round get_finished or a fence asks for runs after the
    submission in progress. The in-process shell refuses dsv41. NativeState is
    main's apart from the moved registration, and the MP worker tests run their
    failure paths on both the step-thread and the polling worker.
  • Core review of 11-03 (valarLip). The plugin's off-step submit and poll
    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_final hook and
    lmcache.mp.null_block_id are gone, and V4.1 is detected by attn_family.
  • Native V4.1's namespace as on main. dsv41 keeps the namespace hybrid
    gave V4.1, so native ATOM's tiers stay reachable; only the plugin hashes
    V4.1's window and candidate-block fields.
  • Native futures as on main. Native ATOM drops a finished future where it
    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.
  • Merged main 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.
  • Second pass of the 11-02 core review (valarLip). NativeState submits and
    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 auto collapse are documented.
  • Core review of 11-02 (valarLip). The plugin's step thread no longer
    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 are
    gone.
  • Reviews of 10-31 and 11-01 (valar-cop). Decode keeps are stored at the
    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.
  • Follow-up review of 10-29 (valarLip). On MP a stalled save keeps its
    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 up
    to the transfer deadline. V4.1 offload requires
    kv_load_failure_policy=recompute, and the V4.1 adapter matches only with
    its checkpoints on.
  • Review of 10-28 (valarLip). Fixes for vLLM 0.33's hit-cap call, the
    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.
  • The plugin's restore stays out of core. The restore, its blocking read
    and the vLLM warmup exemption live in the plugin's cache.
  • Merged main 481c2bfd4. The conflicts came from feat(moe): route before the rows move; per-pass real row counts; V4.1 FP8 KV #2725 replacing the
    pool geometry's packed flag with kv_dtype; both paths still build the
    geometry through build_v41_pool_geometry.
  • Review of 10-26 (yhl-amd). Checkpoint restores also run when the cursor
    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 in
    this comment.
  • Earlier in this revision: rebased onto main d6681e153 (fix(mla): migrate Triton decode to maintained aiter API #2669 replaces the
    Triton 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

  • Native V4.1's namespace, as on main, does not hash sliding_window or
    candidate_block_size (the plugin's does): two checkpoints of one name that
    differ only there, through --hf-overrides say, share a tier. Closing it
    moves every native V4.1 namespace, so it is left to its own PR.
  • tp_rank_collapse: auto keeps native V4.1 one copy per rank, as on main, so
    TP4 stores four identical copies. Set true to store one; the plugin
    resolves auto to true for V4.1.
  • When the working set fits in HBM, the tier is idle, and attaching LMCache
    costs about 4% (no spec: 14.00 vs 14.55 req/s).
  • A tier load restores only the checkpoint at its own end. Once a shared
    prefix's checkpoint is evicted, requests load past it: 1125 of 1311 loads in
    the DSpark run started at token 0.
  • DSpark: a request resumed from the tier prefills in a step of its own,
    because V4.1 cannot mix a prefill with a speculative verify.
  • DSpark: a hit never reaches a decode keep. The draft group's cached extent
    ends every prefix hit inside the prompt.

🤖 Generated with Claude Code

@github-actions

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

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

Heavy model tests:

  • ✅ Run after the PR is approved and Pre Checkin passes
  • ✅ Run immediately when an approval review is submitted
  • ✅ Can be requested before approval with labels
Label Tests
ci:full Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom Run native ATOM model accuracy tests
ci:vllm Run ATOM vLLM OOT model accuracy tests
ci:sglang Run ATOM SGLang model accuracy tests

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

@PerryZhang01 PerryZhang01 changed the title DeepSeek-V4.1: run checkpoint images on the V2 model runner, with LMCache tier offload DeepSeek-V4.1-Flash on the vLLM plugin: state checkpoints on the V2 runner, with LMCache tier offload Oct 15, 2026
@PerryZhang01
PerryZhang01 force-pushed the perzhang/dsv41-lmcache-031 branch from 70733ee to d3c4ec5 Compare October 15, 2026 15:40

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

ruff

⚠️ [ruff] <I001> reported by reviewdog 🐶
Import block is un-sorted or un-formatted

import importlib
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
from atom.plugin.vllm import v41_shutdown as patch


⚠️ [ruff] <I001> reported by reviewdog 🐶
Import block is un-sorted or un-formatted

import threading
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import StateByteCodec


⚠️ [ruff] <F401> reported by reviewdog 🐶
pytest imported but unused


⚠️ [ruff] <PIE810> reported by reviewdog 🐶
Call startswith once with a tuple

if line.startswith("<<<<<<< ") or line.startswith(">>>>>>> "):


⚠️ [ruff] <F401> reported by reviewdog 🐶
inspect imported but unused


⚠️ [ruff] <I001> reported by reviewdog 🐶
Import block is un-sorted or un-formatted

from types import SimpleNamespace
import pytest
from atom.plugin.vllm.kv_transfer.v41_state import V41BoundaryPlanner


⚠️ [ruff] <I001> reported by reviewdog 🐶
Import block is un-sorted or un-formatted

from copy import deepcopy
import pytest
from atom.plugin.vllm.state_slot_allocator import StateSlotAllocator


⚠️ [ruff] <BLE001> reported by reviewdog 🐶
Do not catch blind exception: Exception

except Exception:


⚠️ [ruff] <BLE001> reported by reviewdog 🐶
Do not catch blind exception: Exception

except Exception as exc:


⚠️ [ruff] <BLE001> reported by reviewdog 🐶
Do not catch blind exception: Exception

except Exception:


⚠️ [ruff] <BLE001> reported by reviewdog 🐶
Do not catch blind exception: Exception

except Exception as exc:


⚠️ [ruff] <SIM115> reported by reviewdog 🐶
Use a context manager for opening files

json.dump(rows, open(OUT, "w"))


⚠️ [ruff] <SIM115> reported by reviewdog 🐶
Use a context manager for opening files

json.dump(rows, open(OUT, "w"))


⚠️ [ruff] <BLE001> reported by reviewdog 🐶
Do not catch blind exception: Exception

except Exception as exc:


⚠️ [ruff] <SIM115> reported by reviewdog 🐶
Use a context manager for opening files

json.dump(out, open(OUT, "w"))

@zufayu
zufayu requested a review from yhl-amd October 16, 2026 01:08
@PerryZhang01
PerryZhang01 force-pushed the perzhang/dsv41-lmcache-031 branch from 5e4bdbb to b1eaef0 Compare October 16, 2026 08:06
@PerryZhang01 PerryZhang01 changed the title DeepSeek-V4.1-Flash on the vLLM plugin: state checkpoints on the V2 runner, with LMCache tier offload DeepSeek-V4.1-Flash: LMCache tier offload through the MP server (vLLM plugin) Oct 17, 2026
@PerryZhang01
PerryZhang01 force-pushed the perzhang/dsv41-lmcache-031 branch from 6cc5c05 to c8cf505 Compare October 17, 2026 16:48
@PerryZhang01

Copy link
Copy Markdown
Contributor Author

Review: V4.1 plugin LMCache MP and ATOM native impact

Reviewed 8b330dcbfbb92a2c8b5cb19e84105226b2447acc against df86c48298f926c98ee71972ace8bda237bc9958. The findings and permalinks below refer to that revision.

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 timeout

Location: mp_worker_adapter.py:135–164

wait_for_requests() only inspects _pending_saves and _pending_loads. A transfer still in _transfer_queue, or inside synchronous submission before its future is inserted into a pending map, is invisible to the fence.

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 request_finished path. Pinning recurrent-state blocks separately does not protect PAGE blocks either.

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 leases

Locations: connector.py:2512–2537, connector.py:887–899

_request_finished() removes _requests[req_id] before processing unfinished saves and acquiring their source leases. A final suffix save that has not yet been dispatched subsequently calls _v41_state_for_transfer(), which requires that Request to resolve checkpoint hashes. With the Request gone, the hook always returns None and the scheduler keeps skipping the save.

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 reclaim_stale_leases() returns no releases.

CPU reproduction: after finish, the Request registry is empty, the late save returns None, three suffix-block leases remain, has_pending_work is true, and the live-transfer set is empty.

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 pin

Location: page_ride.py:160–200

The save branch prefers recorded over source_blocks_for(request, end), but recorded is the destination from an earlier load. After a multi-group load failure, vLLM's recompute recovery can free the old block table and recompute the same boundary into different checkpoint blocks. The failure path does not clear the recorded destination.

A later save can consequently select a freed/reused old destination instead of the current checkpoint. Because recorded is present, not recorded is false and the source-pin branch is also skipped.

CPU positive/negative control: with recorded destination (7,) and current canonical source (42,), the original function saves from (7,) and acquires no pin. Clearing the stale record selects and pins (42,).

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 restore

Location: cache.py:453–458

The new hook restores only rows where cursors[:, 0] != starts. Two different prefixes can reuse the same state slot at the same token position: the old occupant leaves cursor 256, and the new request resumes from its own checkpoint at 256. The predicate then skips restoration, and the subsequent stale-position check also passes, despite the history/window belonging to the old request.

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 ATOM_ENGRAM_UVA=0. This is deliberately not a claim that every default-UVA run fails.

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 completion

Location: mp/worker.py:606–626

The new drain thread polls futures after _drain_wake.wait(timeout=0.05), while get_finished() only consumes results already collected by that thread. When native ATOM has no runnable forward and is waiting for a load, no new transfer metadata is dispatched to wake the drain. A future becoming ready does not wake it either.

Native's 1 ms idle progress polling can therefore repeatedly call get_finished() without observing the completion until the next 50 ms drain tick.

Controlled CPU comparison, five trials per version:

Version Future-ready → completion observed by get_finished()
PR head 50.109–51.034 ms
Base 0.009–0.013 ms

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 get_finished() overrides this path and directly polls its native futures, so this finding must not be generalized to that worker or to in-process LMCache.

P2 — 6. Current-step Engram staging runs before the previous speculative checkpoint is finalized

Locations: deepseek_v41_bridge.py:808–825, staging call:614–628

_v41_stage_outside_forward() calls Engram stage() before _v41_state_checkpoints_around_forward() processes the previous step's decode candidates.

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; --no-async-scheduling does not remove it.

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 start

Location: req_id_passthrough_patch.py:137–144

num_computed_tokens_cpu is sourced from V2's num_computed_tokens_np, then consumed by checkpoint _exact_starts(). Under async speculative scheduling, that CPU count can include optimistic advancement assuming all drafts are accepted; it is not necessarily the committed GPU frontier.

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 by_end.get(261) misses it and the candidate collection is cleared.

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 --no-async-scheduling avoids this case, but the implementation does not reject the other configuration.

Scope: plugin V2 + async speculative scheduling.

P2 — 8. Local checkpoint hits bypass DSpark's single-token-tail protection

Location: v41_state.py:229–243

The external-hit path applies keep_a_prefill_tail, but the enabled local-checkpoint-hit branch returns the hit unchanged.

With ATOM_V41_CHECKPOINT_LOCAL_HITS=1 and DSpark5, a 1025-token prompt hitting at 1024 has one token left. Waiting-request admission can expand that into a speculative verify row. If another waiting request contributes a normal prefill in the same batch, the V4.1 bridge rejects the verify/prefill mixture with NotImplementedError. The earlier running-draft clearing does not protect placeholders added later by waiting-request admission.

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 boundary

Location: chunked_scheduler.py:900–918

For chunk/checkpoint interval 256, a 1024-token prompt fully present in the tier, HBM frontier 0, and OFFLOAD_MIN_LOAD_TOKENS=256:

  1. Leaving a prefill tail and flooring the logical hit produces 768.
  2. V4.1 admission records a STATE destination only for boundary 768.
  3. MP's existing full-prompt override sets physical transfer_end back to 1024.
  4. The new recurrent hook looks for destination 1024, finds none, and refuses the load.

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

Path Result
Traditional in-process LMCache Does not use the new MP drain thread. Inspection of shared scheduler/config/metadata and the selected CPU regressions did not identify an additional normal-operation read/write regression. Namespace canonicalization still changes behavior for affected configurations and can change reuse of pre-upgrade cache entries; this is not a blanket compatibility guarantee.
Generic PAGE-only native MP Confirmed control-flow regression in finding 5.
NativeState/checkpoint-aware native MP Overrides submission and completion polling, so finding 5 does not apply. It still inherits other shared code; this does not establish that the entire PR has zero effect on it.
Native V4.1 model-side refactors Geometry-construction arguments matched base in 16 tested combinations; cursor-advance predicates matched in 8 valid-state combinations. No native behavior change was found in those two refactors.

Validation and coverage limits

  • Ran 14 native LMCache-related CPU test files against a separate git archive snapshot of the reviewed commit: 789 passed, 3 skipped.
  • V4.1 runtime/refusal CPU tests: 50 passed.
  • Targeted original-function/AST CPU reproductions and positive/negative controls for the lifetime, cursor, committed-start, and idle-completion issues described above.
  • tests/plugin/test_vllm_dsv41_lmcache.py failed during collection because vllm is not installed in the review environment; no pass is claimed for it.
  • No GPU inference/accuracy/performance rerun. The recipe's GSM8K/performance numbers are author-reported results, not measurements from this review.
  • When inspected, GPU accuracy CI jobs were skipped. The non-GPU CI script also excludes tests/plugin, so its green status does not cover the new plugin tests.

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.

@yhl-amd

yhl-amd commented Oct 22, 2026

Copy link
Copy Markdown
Collaborator

IPC event lifetime on ROCm 7.2.x — two concerns with the per-step event + drain-thread release

  1. Destroying an IPC event blocks on all queued GPU work. On 7.2.x, ~IPCEvent() → ihipHostUnregister → SyncAllStreams() (ROCm issue #7520, rocm-systems#11529). We measured in a live MiniMax-M3 TP4 c48 service that at start_load_kv / get_finished the GPU is busy >95% of the time, and destroying an IPC event there takes p50 ~8–10 ms, max 2.4 s (long prefill steps). In this PR the last reference usually drops on the drain thread (item/rest locals and the future), so the drain thread stalls that long every step.
  2. While one thread destroys an IPC event, creating an IPC event on another thread blocks for the same duration. A single-GPU repro shows 40 ms vs 40 ms; kernel launch, record, query and ipc_handle() are not affected. _record_step_event() creates a fresh IPC event on the step thread every step, so the step thread can stall behind the drain thread's destroy.

Suggestions:

  • (a) Reuse a pool of IPC events, rotating each one after ≤31 records. The invalid argument on re-record matches the 7.2.x wait-index overflow past 32 records (rocm-systems#12230/#11749) or the importer-unlink bug (#11511). This needs the LMCache side to cache imported handles.
  • (b) Hand spent events back to the step thread and release them only when the stream is idle; otherwise defer.
  • (c) Add timing around event creation and release to confirm on the V4.1 path.

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.

zgplvyou and others added 8 commits October 25, 2026 21:35
… 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>
@PerryZhang01
PerryZhang01 force-pushed the perzhang/dsv41-lmcache-031 branch from 6fb1f34 to 807d20b Compare October 26, 2026 03:17
@PerryZhang01

Copy link
Copy Markdown
Contributor Author

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
Item Fix
P1-1 fence misses queued transfers The fence also waits for queued and submitting transfers; a timeout is fail-stop.
P1-2 tail save loses its Request The ride hook reads the request off the SeqView; a finished save that cannot carry its state retires and frees its lease.
P1-3 stale load destination as save source A save rides only on the block committed under its hash, always pinned.
P1-4 same-position slot reuse A request entering another request's slot mid-prompt has its cursor invalidated, so it is restored or refused.
P2-5 native idle completion latency get_finished wakes the drain thread while transfers are pending.
P2-6 decode keeps after staging Decode keeps are stored at commit, at the committed position.
P2-7 async speculative counts Checkpoints refuse async scheduling with speculative decoding.
P2-8 local hit without the DSpark tail Local hits keep the prefill tail, with or without a connector.
P2-9 transfer end vs state destination The lookup stops at the last boundary a load can end on.

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.

Re #2598 (comment)

@PerryZhang01

Copy link
Copy Markdown
Contributor Author

Round 2 pushed 807d20b with suggestion (b). Finished futures, which hold the step event, are parked for the step thread; it drops them at start_load_kv while its stream is idle, or past 256 parked. Every event is now created and destroyed on that thread. (a) needs LMCache to cache imported handles, so it is not in this PR. (c) on the V4.1 path waits for GPU time; I will post the numbers.

@yhl-amd re #2598 (comment)

@yhl-amd yhl-amd left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread atom/plugin/vllm/v41_speculative.py Outdated
get_current_req_ids,
)

batch = get_current_input_batch()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread atom/kv_transfer/offload/mp/worker.py Outdated
with self._lock:
for operation_id, entry, result in terminal_ops:
self._pending_loads.pop(operation_id, None)
self._retired.append(entry.future)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@PerryZhang01

Copy link
Copy Markdown
Contributor Author

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 numbers

MI355X 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 rocm/atom-dev:vllm-v0.31.0-nightly_20261025 with ATOM_DSV41_MHC_BACKEND=triton. A probe outside the PR wraps the worker methods.

per call, ms before (08d6753) after (807d20b)
create step event, mean / p50 / p99 7.10 / 5.75 / 15.92 5.57 / 4.77 / 11.59
release on the step thread, mean / max none 0.017 / 0.03
drain poll of loads, mean / p90 2.20 / 7.24 1.85 / 6.21
start_load_kv, mean / p99 1.49 / 11.10 1.36 / 12.51
req/s, ITL avg, TTFT avg 8.16, 25.05, 375 8.12, 25.32, 370

@yhl-amd re #2598 (comment)

zgplvyou and others added 3 commits October 26, 2026 06:56
`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>
@PerryZhang01

Copy link
Copy Markdown
Contributor Author

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
  • Non-plugin suite at a7d323e: 48 failed, 13160 passed. The same 48 fail on 807d20b, which also fails a 49th (test_every_step_records_a_fresh_event, only on a GPU host, fixed here).
  • tests/plugin (not run by CI): 65 failed, 676 passed, the same set as 807d20b (rtpllm and sglang not installed, fakes behind main).
  • ruff 0.17.2 and black clean on the changed files.

@yhl-amd

@PerryZhang01

Copy link
Copy Markdown
Contributor Author

Round 23 pushed ba157dd: all nine items are done, 1, 3 and 4 included. Please approve.

  • 1, 3, 4: a round hands futures back to its own thread; a failed round raises every step; the forward thread fail-stops a transfer no round reached past its deadline.
  • 7: the engine-facing unwraps use raw_req_id: stringified, test_offload_completion_processing_calls_load_terminal_once fails.
  • 8: the kimi_k3 / qwen refusal is intended; neither tier restores the other's state.
Evidence
Item Change
1 _poll_round(transfer) gets the thread it runs on and hands futures back to it; close clears _transfer after transfer.close returns. The new test runs a round inside close with a finished save: no poll error, and the completion is settled.
2 _poll_entries polls without touching worker state; _settle_terminal_locked reports what finished. Native get_finished runs both under one lock, saves then loads, as main; the polling thread polls unlocked and settles under the lock.
3 raise_error keeps a poll failure and raises it on every call; the failure is recorded and marked in one locked step.
4 _settle_transfer_thread fail-stops a pending transfer whose deadline passed, plus a 1 s grace, with no round since. The test blocks the thread in a submit and goes through _settle_transfer_thread.
5 _ensure_lookup_pin only answers; its two callers drop the load.
6 The docstring names only stores_state_images.
7 raw_req_id returns the id in the type the engine compares (seq.id), req_id_of its string; used at mp/scheduler.py, _offload_common.py, hybrid/kimi_k3, hybrid/dsv4 and pp_kv_aggregator.py.
8 test_offload_layout_override_cannot_swap_the_two_state_tiers; the description states the rule.
9 Comments state the constraint; the ROCm and cadence findings are in the description.

Each new test fails with its fix reverted. Unit, at ba157dd in the PR's container: the CI command and tests/plugin fail what they fail at 85d55d8.

GPU, MI355X TP4, DeepSeek-V4.1-Flash DSpark:

Run Result
Plugin GSM8K two-pass 0.9227 / 0.9234, 5.84M tier hit tokens, 0 faults, 0 poll failures
Native V4.1 two-pass, ba157dd x2 0.9219 / 0.9189, 0.9249 / 0.931; pass-2 tier loads 950k, 967k tokens
Native V4.1 two-pass, main x2 (interleaved) 0.9265 / 0.9333, 0.9272 / 0.9204; pass-2 tier loads 968k, 968k tokens
Plugin c16 900 s, ba157dd x2 (ABBA) 8.63 / 8.89 req/s, ITL p50 23.24 / 23.01 ms, TTFT p50 316 / 314 ms
Plugin c16 900 s, d7619fe x2 8.76 / 8.73 req/s, ITL p50 23.22 / 23.51 ms, TTFT p50 310 / 315 ms

@valarLip valarLip left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-read the core files at ba157dd, by reading the code and running the related tests (offload / MP / native / config / shell 812 passed, block manager / V4.1 1001 passed, test_lmcache_mp* ten times in a row) and two probes; no GPU. Thanks for round 23. Items 1 and 3 from last time are fixed.

Item 4 (deadlines during a stuck submit) still has a gap, reproduced below as item 1. So I can't approve yet. Items 1, 4 and 9 continue my last items 4, 1 and 5. Each item ends with the change I'm asking for.

1. A stuck submit can still hide a missed deadline (mp/worker.py:532-554). _enforce_stalled_deadlines stops a transfer only while last_round_at < due. But last_round_at is when the last round ended (transfer_thread.py:358), and a round checks each entry when it reaches it. Suppose a round starts at due − 2 ms, polls transfer X (not yet due, no raise), spends a few ms on a terminal import, and ends at due + 3 ms. Then the next submit hangs. From then on last_round_at < due is false on every step, and X is never stopped. A probe confirmed it: 90 s past the deadline, no raise. The method also restates _enforce_transfer_deadline with a magic _POLL_STALL_GRACE_S = 1.0, as a second deadline path.
→ Change: on the forward thread, call _enforce_transfer_deadline(op, entry.started_at, deadline) on every pending entry each step (a CPU comparison, which is what I meant last time). Then delete _enforce_stalled_deadlines, _POLL_STALL_GRACE_S, last_round_at and _last_round_at.

2. Native V4.1's namespace leaves out two byte-geometry fields (offload/config.py:44-56, :380-387). Native dsv41 maps to dsv4-page-regions, and _HF_PAGE_FIELDS doesn't include sliding_window or candidate_block_size. The plugin hashes both, because, in its own words, either one moves the bytes a stored object holds. So under native ATOM, two V4.1 checkpoints with the same name and revision but a different window or candidate block share a namespace. A hit then restores a window ring or index plane of the wrong length.
→ Change: add sliding_window and candidate_block_size to the hashed fields for dsv41 in core build_page_namespace, and accept the one-time native namespace change explicitly, instead of keeping two rules for one model.

3. The new adapter refusal runs after registration and outside its cleanup (mp/worker.py:320, :430-433). _start_transfer_thread refuses an off-step worker whose adapter records its own producer events. It runs after adapter.register_kv_caches(...) and outside the try / except: adapter.shutdown(). When it raises, the server registration and heartbeat are already live, and nothing shuts them down until process exit or the server's reaper.
→ Change: do the check before registration (it depends only on the adapter type), or move _start_transfer_thread inside the try.

4. close() ignores a timed-out join (mp/worker.py:345, mp/transfer_thread.py:291-294). transfer.close() drains and joins with a timeout and returns nothing, and close() goes on to adapter.shutdown(). In polling mode the thread runs rounds on its own. If a submit or terminal poll is stuck past 60 s, the join times out and the adapter drains and unregisters, while _poll_round may still be inside _terminal_future_result on that adapter's futures.
→ Change: have close() report whether the thread stopped (is_alive() after the join). If it didn't, log it and skip adapter.shutdown() rather than tearing the adapter down under a live thread.

5. A frontend's speculative form is hashed raw (offload/config.py:375-377). When offload_speculative_config is set it goes into the document as is; only the fallback goes through _stable_config_value. A torch.dtype in it makes json.dumps raise, and a NaN is hashed instead of refused.
→ Change: run both through _stable_config_value.

6. An unneeded change to rank collapse (mp/deployment.py:165). _config_has_fully_replicated_tp_pages now asks select_offload_layout(config) in offcfg._STATE_OWNING_LAYOUTS instead of main's model_type == "kimi_linear". For V4.1 and GDN it changes nothing, since they have no kv_lora_rank and returned False on main already. It reaches into another module's private set and can raise on an override. The only real behaviour change is that a dense MLA model with an explicit kimi_k3 / qwen override loses collapse (8 copies instead of 1). test_auto_rank_collapse_keeps_kimi_k3_per_rank had to grow six hf fields so the layout lookup could run.
→ Change: revert to main's check, and drop the extra fixture fields.

7. The polling thread spins while idle (mp/transfer_thread.py:313-316). In polling mode _queue.get(timeout=_IDLE_POLL_S) turns every timeout into a round, even with nothing pending or queued. A probe on an idle rank counted 19 empty rounds a second, each taking _lock twice and copying both pending dicts, indefinitely. While steps run, get_finished already asks for a round each step via look(). The timer only matters when a transfer is in flight and the rank has stopped stepping.
→ Change: block without a timeout when nothing is pending or queued, and use the timeout only while a transfer is outstanding.

8. Four fields and three entry points for "wait for a round" (mp/transfer_thread.py:125-129, :217-280). _rounds_started / _rounds_done / _round_asked / _last_round_at back look, _wait_round and _wait_for_round. In fence, or self._poll_failed (:248) is dead, because raise_error() at the top of the loop already raises when _poll_failed is set.
→ Change: one generation counter under _idle, with request_round() -> target and wait_round(target, timeout). Delete the dead check. With item 1, _last_round_at goes too.

9. A lookup that didn't answer reopens saves (chunked_scheduler.py:825-832, hybrid/dsv4/connector.py:2321). Every False from _ensure_lookup_pin now calls _reopen_saves, including hit is None, which a lookup returns on timeout or RPC error. mp/lookup.py:516 calls that "a non-answer, not an empty answer". The tier most likely still holds the prefix, yet [hbm, lmc) is stored again: extra D2H and store traffic, not corruption.
→ Change: have _ensure_lookup_pin tell the two cases apart. Reopen only on a real answer below the promise; on None, just drop the load.

10. Three special cases for family-bound layouts (offload/config.py:160-185). There's natural == "m3" and mapped != "m3", natural in _STATE_OWNING_LAYOUTS and mapped != natural, and the reverse mapped == "dsv41" and natural != "dsv41". The reverse direction is checked only for dsv41, so a dense model overridden to m3 / kimi_k3 / qwen is still accepted, and the next family-bound layout will need another special case.
→ Change: one table of family-bound layouts, checked in both directions, with an explicit allowance for the test fixtures that use the kimi_k3 override.

11. The README module table is stale (offload/README.md:116). It says mp/native_state_{layout,scheduler,worker}.py own native-state registration. Registration (_register_native_state_layout) and block-id validation now live in mp/worker.py, and require_native_state_server in mp/deployment.py.
→ Change: update the row to point at the files that now own these.

12. A sleep-timed test (tests/test_lmcache_mp_async_threads.py:339-343). test_a_failed_poll_round_stops_the_polling relies on time.sleep(3 * _IDLE_POLL_S) and time.sleep(0.02) for the round to have run. drain() waits only on _outstanding, and the _WAKE sentinel isn't counted. On a loaded CI runner the round may not run within ~170 ms, rounds == [], and the test fails intermittently.
→ Change: wait with the existing _wait_for_round(target, timeout) (or item 8's wait_round) instead of sleeping.

Addresses valarLip's review of ba157dd.

- With the polling thread, the forward thread enforces every pending
  transfer's deadline each step (`_enforce_transfer_deadline`, a clock
  comparison), so a submit that never returns cannot hide one. The
  stall heuristic, its grace and `last_round_at` are gone (1).
- An off-step worker refuses an adapter that records its own producer
  events before registering it, inside the cleanup that shuts the
  adapter down (3).
- `_TransferThread.close` reports whether the thread stopped; when it did
  not, `close` leaves the adapter registered for the server's reaper
  rather than shut it down under a live thread (4).
- A frontend's speculative form is hashed through `_stable_config_value`
  like the fallback (5).
- Auto rank collapse is main's `kimi_linear` check again; the test
  fixture is main's (6).
- The polling thread polls on a timer only while transfers are in flight
  (its poll returns whether they are); idle, it waits for a submission or
  a request (7).
- One request/wait pair: `request_round()` returns the round to wait for,
  `wait_round(target, timeout)` waits for it. Three round counters replace
  the flag and the timestamp, and the fence's dead check is gone (8).
- `_ensure_lookup_pin` answers None when the lookup did not answer, and
  the caller reopens saves only on False, an answer below the promise (9).
- README: native-state registration lives in `mp/worker.py`, the server
  check in `mp/deployment.py` (11).
- The failed-poll test waits on rounds instead of sleeping (12).

Not changed: the native V4.1 namespace (2) stays main's, and the
override table (10) keeps main's accepted combinations.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@PerryZhang01

Copy link
Copy Markdown
Contributor Author

Round 24 pushed ce7dd45: items 1, 3-9, 11 and 12 are done; 2 and 10 are not changed. Item 1 was the open blocker; please approve.

  • 1: the forward thread enforces every pending deadline each step; the stall heuristic is gone.
  • 2: main has the same gap, and closing it moves every native V4.1 namespace: its own PR.
  • 10: main accepts those overrides and test_block_manager builds on one; its own PR.
Evidence
Item Change
1 _settle_transfer_thread calls _enforce_transfer_deadline on every pending entry each step in polling mode; _enforce_stalled_deadlines, _POLL_STALL_GRACE_S and last_round_at are deleted. The test hangs a submit and calls _settle_transfer_thread across the deadline.
2 Main's _HF_PAGE_FIELDS is this PR's: V4.1 under main's hybrid hashed neither field either.
3 _refuse_unfenceable_adapter(adapter) runs inside the registration try, before register_kv_caches; the test checks nothing was registered and the adapter was shut down.
4 _TransferThread.close returns whether the thread stopped; if not, close logs and skips adapter.shutdown().
5 Both forms go through _stable_config_value; a dtype hashes as its name and NaN is refused. The plugin's stable form is a dict of plain values, so its namespace is unchanged.
6 Main's kimi_linear check and main's test fixture.
7 poll returns whether transfers are in flight; only then does queue.get time out.
8 request_round() -> target and wait_round(target, timeout); _rounds_started, _rounds_done and _rounds_requested replace the flag and the timestamp; the fence's dead check is gone.
9 _ensure_lookup_pin returns None for a lookup that did not answer; the dispatch loops reopen saves only on False.
10 tests/test_block_manager.py builds its state tier with offload_layout: kimi_k3 on a dense model.
11 README row points at mp/worker.py and mp/deployment.py.
12 The test waits on wait_round; no sleep.

Each new test fails with its fix reverted. Unit, at ce7dd45 in the PR's container: the CI command and tests/plugin fail what they fail at 85d55d8, plus test_atomesh_spur_dispatch.py timeouts that main's tree also hits in that container under load.

GPU, MI355X TP4, DeepSeek-V4.1-Flash DSpark:

Run Result
Plugin GSM8K two-pass 0.9234 / 0.9265, 0 faults, 0 poll failures
Native V4.1 two-pass, ce7dd45 0.9249 / 0.9204; pass-2 tier loads 967k tokens in 1301 requests
Native V4.1 two-pass, main (back to back) 0.931 / 0.9227; pass-2 tier loads 971k tokens in 1306 requests
Plugin c16 900 s, ce7dd45 x2 (ABBA) 8.88 / 8.86 req/s, ITL p50 23.13 / 23.22 ms, TTFT p50 310 / 313 ms
Plugin c16 900 s, ba157dd x2 8.79 / 8.59 req/s, ITL p50 22.87 / 23.27 ms, TTFT p50 309 / 316 ms

@valarLip valarLip left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-read the core files at ce7dd45, by reading the code and running the related core tests (816 passed, 1 skipped; the MP thread tests three times over); no GPU. Thanks for round 24 and the GPU runs. I accept your reasons for leaving items 2 and 10 to their own PRs.

Item 1 below is a false fail-stop introduced by the deadline fix I proposed last round. That's on me, and it's the one thing I'd like fixed before approving. Items 2–9 can go in a follow-up if you prefer. Items 2, 3 and 5 also come from my earlier asks (7, 9 and 4). Each item ends with the change I'm asking for.

1. The forward-thread deadline check fires on transfers that already finished (mp/worker.py:543-547). In polling mode, _settle_transfer_thread calls _enforce_transfer_deadline on every entry in _pending_saves / _pending_loads using only started_at, without first asking whether its future is terminal. Terminal entries leave _pending_* only when the transfer thread runs a round. Take a save or load that finished before its deadline while the transfer thread is held in a slow submit (or in _wait_ready for an after-forward save). The next start_load_kv / get_finished raises LMCacheTransferUnprovable ("still not terminal") and the engine fail-stops over a transfer that completed. _poll_entries and the inline get_finished check terminal first and deadline second; this copy reverses that. My suggestion last round left out the terminal check.
→ Change: check _terminal_future_result(entry.future) before the deadline, reusing _poll_entries' per-entry step rather than a third copy. Add a test where a save finishes, the thread is held in a submit past the deadline, and _settle_transfer_thread must not raise.

2. req_id_of / raw_req_id still miss the isinstance spellings (Fix-then-sweep). The last sweep covered the getattr / hasattr form, but the str(req_id.req_id if isinstance(req_id, LoadOperationId / SaveOperationId) else req_id) form remains at:

  • chunked_scheduler.py:1146, :1378, :1402 and :1424, including this PR's own load_failed;
  • hybrid/dsv4/connector.py:2624 and :2673, plus the paths around :2588 / :2677-2679.

→ Change: use the helper everywhere the id is unwrapped. Keep the isinstance checks only where the code branches on the operation type itself.

3. The tier-lost block is duplicated (chunked_scheduler.py:820-835, hybrid/dsv4/connector.py:2312-2328). Both dispatch loops repeat the same twelve lines:

  • write the floor into _load_save_floors;
  • call _ensure_lookup_pin;
  • branch on its three-way True / False / None result;
  • _mark_load_skip, _reopen_saves on False, _clear_pending_load.

The floor is written only to be read back four lines later by _reopen_saves.
→ Change: one OffloadSchedulerMixin helper, e.g. _confirm_load_pin(seq, sid, ls, hbm, ...) -> bool, that does the skip / reopen / clear and takes the floor as an argument. Both loops call it, and _ensure_lookup_pin's tri-state stays private to that helper.

4. The README describes the collapse rule you reverted (offload/README.md:212-216). It says tp_rank_collapse: auto "leaves every state-owning layout (kimi_k3, qwen, dsv41) one copy per rank". After item 6, _config_has_fully_replicated_tp_pages again uses main's kv_lora_rank + kimi_linear test and doesn't look at the layout. dsv41 and qwen come out per-rank only because they have no kv_lora_rank. A future MLA model with per-request state would be collapsed under auto, contrary to the README.
→ Change: state the actual rule (MLA pages without kimi_linear collapse; everything else stays per rank).

5. close() can leave the connector half-closed (mp/worker.py:352-367). When the thread doesn't stop, close() sets _transfer = None, logs and returns, but keeps _adapter. A second close() (a repeated plugin shutdown) then finds no transfer and calls adapter.shutdown() under the still-running thread, which is what the first call avoided. Any later get_finished sees _adapter set and no transfer, so it polls the same futures inline on the forward thread. Separately, the log says "within 60s", but drain and join each wait _TRANSFER_DRAIN_TIMEOUT_S, up to 120 s.
→ Change: on that path, drop _adapter too (leave the registration to the server's reaper) or set a closed flag every entry point checks, and log the real bound.

6. The first recorded error can mask a poll failure (mp/transfer_thread.py:178-186, :313, :348). _record_error keeps only the first error. If a submit error (e.g. duplicate operation) was recorded and not yet raised, the poll failure that stops the polling is never reported: every step re-raises the stale submit error. When _poll raises, _in_flight also keeps its last value, so a True leaves the thread waking at 20 Hz until exit.
→ Change: when a poll fails, prefer that error (or chain the two), and set _in_flight = False on the failure path.

7. A test whose assertions hold by construction (tests/test_lmcache_mp_native_worker.py:781). test_native_transfers_are_submitted_and_polled_on_the_step_thread builds its worker with _make_worker(), which never calls register_kv_caches / _start_transfer_thread. So _transfer is always None, and the thread count is read after construction. If the shared _register_native_state_layout ever started a thread for native, saves would leave the step thread and this test would stay green. Its docstring line "No transfer thread starts" also isn't checked on the real registration path with OFFLOAD_MP_TRANSFER_THREAD=1.
→ Change: register through the real path and compare the thread count before and after.

8. A shared-fixtures module with one user (tests/lmcache_mp_support.py). Its docstring says it holds the fixtures "every LMCache MP worker test file shares", but only test_lmcache_mp.py imports it (via # noqa: F401 so the autouse fixtures apply). step_event(worker) ignores its argument, and start_transfers only wraps _start_transfer_thread.
→ Change: move the fixtures into test_lmcache_mp.py (or a conftest.py if other files should use them), and drop the two wrappers.

9. Small:

  • _poll_round names two different things: the worker callback (mp/worker.py:904, returns bool) and the thread method (mp/transfer_thread.py:343, counts and records errors), which calls the first through self._poll.
  • LMCacheMPConnector's class docstring (mp/worker.py:104) still says only "for backend-published PAGE views".
  • offload/config.py:83-84 says "qwen is the same tier over a GDN PAGE image", which contradicts the new refusal that the two tiers can't restore each other's state.

→ Change:

  • Rename the worker callback, e.g. _poll_transfers.
  • Update the class docstring to what it does now.
  • Reword the layout comment to match the refusal.

Addresses item 1 of valarLip's review of ce7dd45.

The forward thread enforced every pending transfer's deadline from its
start time alone. A transfer that had finished, but waited for a round
held back by a slow submit, was failed as "still not terminal".

The forward thread now polls only transfers already past their
deadline, through the same per-entry step as a round (`_poll_entry`:
terminal first, deadline second), so a finished one is left for the
next round to settle. Transfers within their deadline are not polled
there, so terminal polls (the IPC event import) stay off the forward
thread. The check runs only between rounds (`between_rounds`): the first
query of a finished future imports its event, and a round must not be
polling the same future at the same time.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@PerryZhang01

Copy link
Copy Markdown
Contributor Author

Round 25 pushed 4b0dc75: item 1 is fixed; items 2-9 go to a follow-up PR, as you offered. Please approve.

The forward thread now polls only a transfer past its deadline, through the round's per-entry step (terminal first), and only between rounds. A finished one is left for the next round to settle.

Evidence

Why not a terminal check on every entry each step: a terminal poll imports the peer's IPC event (DeviceMessagingFuture._on_raw_future_complete), which the polling thread exists to keep off the forward thread, and query() takes no lock, so two threads polling a future at once could import its event twice. Hence overdue entries only, and _TransferThread.between_rounds, a lock a round holds while it polls; the forward thread skips the check while a round runs, since that round polls the same transfers.

Test Fails with
test_a_transfer_that_finished_is_not_failed_while_rounds_are_held[load/save]: the transfer finishes, a submit holds the thread past its deadline, _settle_transfer_thread does not raise, and the next round reports it ce7dd45's worker
test_a_stuck_submit_cannot_hide_a_transfer_past_its_deadline: still raises for one that has not finished removing the forward thread's check
test_the_forward_thread_polls_no_transfer_before_its_deadline polling every entry each step
test_the_forward_thread_leaves_deadlines_to_a_running_round ignoring the round's lock

Unit, at 4b0dc75 in the PR's container: the CI command and tests/plugin fail what they fail at 85d55d8, plus test_atomesh_spur_dispatch.py timeouts that main's tree also hits in that container under load.

GPU, MI355X TP4, DeepSeek-V4.1-Flash DSpark:

Run Result
Plugin GSM8K two-pass 0.9219 / 0.9219, 0 faults, 0 poll failures
Native V4.1 two-pass, 4b0dc75 0.9242 / 0.9272; pass-2 tier loads 962,304 tokens in 1295 requests
Native V4.1 two-pass, main (back to back) 0.9257 / 0.9272; pass-2 tier loads 962,048 tokens in 1294 requests
Plugin c16 900 s, 4b0dc75 x2 (ABBA) 8.76 / 8.86 req/s, ITL p50 23.27 / 23.14 ms, TTFT p50 314 / 315 ms
Plugin c16 900 s, ce7dd45 x2 8.74 / 8.82 req/s, ITL p50 23.54 / 23.24 ms, TTFT p50 313 / 317 ms

@valarLip valarLip left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-read the core files at 4b0dc75, by reading the code, the LMCache adapter it drives (lmcache/integration/atom/multi_process_adapter.py, lmcache/v1/multiprocess/futures.py at f22dec28), the core tests (557 passed) and three probes; no GPU. Thanks for round 25 and the A/B runs. The terminal-first check is right, and its tests fail when the fix is reverted.

I said I'd approve once item 1 was fixed. Reading the fix, I think the forward-thread deadline path I asked for two rounds ago shouldn't exist. That ask was mine, and I'm withdrawing it. Removing the path is the one change I'd like before approving. Everything else below can go to the follow-up PR with items 2–9.

Before approval: delete the forward-thread deadline path

1. Remove _enforce_overdue_deadlines, between_rounds and _polling (mp/worker.py:543-567, mp/transfer_thread.py:136, :239-253, :366). Four reasons:

  • It races a submit, and only when it's actually used. between_rounds excludes only a poll round. A submit runs (_run, fn(*args)) without _polling. Every LMCache submit goes through _track_operation_future (multi_process_adapter.py:722-738), which calls query() on every future it tracks for the context, on the transfer thread. query() on a finished raw future runs _on_raw_future_complete (futures.py:181-195): check _raw_response_processed, import_event, then set the flag, not atomically. So when the forward thread polls an overdue, finished, not-yet-settled entry while a submit is running, both threads can make the first query of the same future and import its IPC event twice, which is what the lock was meant to prevent. That window ("a slow submit is holding rounds back") is exactly when this path does anything. A probe confirmed between_rounds runs its function during a submit. I didn't verify what a double import does on ROCm.
  • It can't see the transfer it exists for. A submit that never returns has its operation only in _submitting_*. _PendingLoad / _PendingSave, and so started_at, are created after the submit returns (worker.py:705-714). A probe holding a submit and moving the clock an hour past the deadline: _settle_transfer_thread doesn't raise, with pending_loads {} submitting {'load:16:1'}.
  • The scheduler already covers it. LMCacheMPConnectorScheduler._enforce_transfer_deadlines (mp/scheduler.py:341-355) bounds every dispatched operation, the stuck one included, at the deadline plus _SCHEDULER_DEADLINE_MARGIN_S (10 s), natively and under the plugin. What the forward path adds is about 10 s earlier detection for transfers submitted before the stuck one.
  • Half of its locking is untested. With with self._polling: removed from _poll_round, test_lmcache_mp.py + test_lmcache_mp_async_threads.py still pass (182). The one related test (test_lmcache_mp.py:2729) holds the lock by hand on the test thread instead of running a round.

→ Change: delete _enforce_overdue_deadlines, the between_rounds call in _settle_transfer_thread, between_rounds, _polling, and the tests built on them (test_a_stuck_submit_cannot_hide_a_transfer_past_its_deadline, test_the_forward_thread_polls_no_transfer_before_its_deadline, test_the_forward_thread_leaves_deadlines_to_a_running_round). Keep test_a_transfer_that_finished_is_not_failed_while_rounds_are_held, which then checks that nothing on the forward thread fails a finished transfer. Rounds and the scheduler deadline remain the two bounds. If the submit-side wait should be shorter, lower the scheduler margin rather than add a third path.

Follow-up PR (with items 2–9)

2. A request during a submit runs two rounds (mp/transfer_thread.py:219-224, :355-359). request_round sets _rounds_requested = started + 1 and queues a _WAKE. After the submit, asked runs one round. The _WAKE is then taken with asked False but queue.empty() True, so a second, empty round runs. A probe gives rounds after one request + one submit: 2 target 1. Each round copies both pending dicts under _lock and queries every future. The comment above _IDLE_POLL_S ("rounds are otherwise asked for by the forward thread") also contradicts the module docstring's "whenever its queue runs dry".
→ Change: on a _WAKE, run a round only if _rounds_requested > _rounds_started, and make the _IDLE_POLL_S comment match the docstring.

3. The core retire change is tested only under tests/plugin (chunked_scheduler.py:668). _retire_finished_save now keeps the blocks another operation of the request still reads instead of returning the whole lease. The only test is in tests/plugin/test_vllm_offload_completions.py, which imports vLLM through the plugin and isn't collected by the core job. With _retire_finished_save reverted to main's whole-lease return, the related core offload / LMCache tests still all pass.
→ Change: add a core test (e.g. tests/test_offload_early_block_release.py) where one request has two save operations and the first retires while the second still reads.

4. Three scans for "this seq's operations" (chunked_scheduler.py:668, :977, :999). _retire_finished_save, should_defer_free and protected_block_ids each scan all of _save_operation_owner for one seq's operations. Every retire is O(all in-flight saves), and the rule "which of this seq's blocks are still read" exists three times.
→ Change: keep an owner → operations index and one _blocks_still_read_by(seq) that the three call.

5. Two registration skeletons in one class (mp/worker.py:169, :269). _register_native_state_layout repeats register_kv_caches' sequence:

  • the rank, config and replication checks;
  • _make_worker_adapter;
  • try { _refuse_unfenceable_adapter; register } except { shutdown; raise };
  • assigning _adapter / chunk_size, _start_transfer_thread, and the log line.

The two have drifted: chunk validation is _strict_integer + % block_size in one and int() in the other, and shutdown goes through getattr in one only. Last round's "refuse inside the try" fix had to be made twice.
→ Change: one registration skeleton that both paths call, each supplying only its views and groups.

6. The "no answer reopens nothing" branch is tested on one side only (chunked_scheduler.py:828). The pinned is False / None distinction exists in both chunked_scheduler.py and hybrid/dsv4/connector.py, but only DSV4 has a test. With the chunked side's None turned into False, the related tests still all pass; the same change on the DSV4 side fails only test_dsv4_a_lookup_that_did_not_answer_reopens_nothing. native_state_scheduler.py:312-328 also returns None, also untested. This folds into last round's item 3: one helper, one test.

7. Docs (offload/README.md:24-30, :146, docs/environment_variables.md:430):

  • The README's layout rule ("compress_ratios present means hybrid") predates CSA2 being checked first, so it says V4.1 resolves to hybrid while :34-39 says dsv41.
  • :146 names recurrent_page_groups, which doesn't exist; the contract is paged_state_checkpoint_spec.
  • The OFFLOAD_MP_TRANSFER_THREAD row doesn't say that a transfers_off_step worker ignores the env and always starts a polling thread, or that an adapter recording its own producer events is refused at registration rather than falling back to inline submission. The warn branch in _start_transfer_thread is now unreachable under off-step.

→ Change: bring the three in line with the code, and drop the unreachable warning.

@zufayu
zufayu requested a review from valarLip November 9, 2026 01:58
@zufayu zufayu assigned valarLip and unassigned yhl-amd Nov 9, 2026
Addresses item 1 of valarLip's review of 4b0dc75, which withdraws the
earlier ask for this path.

The check could not see the transfer it was for: a submit that never
returns has no pending entry, since one is recorded only after the
submit returns. While a submit ran it could poll a future that the
submit's own operation tracking polls on the transfer thread, which is
the concurrent first query its lock was meant to rule out. Polling
rounds already enforce every pending transfer's deadline, and the
scheduler's transfer deadline bounds a submit that never returns.

Removed: `_enforce_overdue_deadlines`, `_TransferThread.between_rounds`
and its lock, and the three tests built on them.
`test_a_transfer_that_finished_is_not_failed_while_rounds_are_held`
stays: nothing on the forward thread may fail a finished transfer.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@PerryZhang01

PerryZhang01 commented Nov 9, 2026 •

Copy link
Copy Markdown
Contributor Author

Round 26 pushed ce1187d: the forward-thread deadline path is removed. Items 2-7 join last review's 2-9 in the follow-up PR. Please approve.

Evidence

Removed: _enforce_overdue_deadlines and its call in _settle_transfer_thread, _TransferThread.between_rounds, _polling, and the three tests built on them. transfer_thread.py is byte-identical to ce7dd45's. test_a_transfer_that_finished_is_not_failed_while_rounds_are_held stays.

The two remaining bounds, as checked: rounds enforce each pending transfer's deadline in _poll_entries; LMCacheMPConnectorScheduler._enforce_transfer_deadlines bounds every dispatched operation at the deadline plus _SCHEDULER_DEADLINE_MARGIN_S, from process_completions natively and from the plugin's build_connector_meta each step.

Unit, at ce1187d in the PR's container: the CI command and tests/plugin fail what they fail at 85d55d8, plus test_atomesh_spur_dispatch.py timeouts that main's tree also hits in that container under load.

GPU, MI355X TP4, DeepSeek-V4.1-Flash DSpark, plugin GSM8K two-pass: 0.9227 / 0.928, 5.85M tier hit tokens, 0 faults, 0 poll failures.

Picks up #2809 (OPUS FP4 index scoring), whose aiter API change made
the branch's V4.1 accuracy jobs fail to import.

Conflicts:
- deepseek_v41_bridge.py: _build_v41_pool_geometry keeps its required
  keywords and gains main's index_layout, also required; the proxy
  geometry passes None, what Config resolves for its FP8 plane.
- tests/test_v41_geometry_parity.py stays deleted: both sides build
  through build_v41_pool_geometry.

build_v41_pool_geometry takes index_layout and forwards it, so the
native builder's call (auto-merged from main) builds the same geometry
as main's inline construction.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@PerryZhang01

Copy link
Copy Markdown
Contributor Author

Round 27 pushed f259bbb: merged main 10b99a085 for #2809, whose aiter change failed the V4.1 accuracy jobs at import.

To review: build_v41_pool_geometry and the bridge's _build_v41_pool_geometry take index_layout, required. Native passes index_cache_layout; the plugin passes None for its FP8 plane.

  • Native geometry and PAGE namespace: identical to main.
  • Unit tests: no failure main lacks.
  • GSM8K two-pass: plugin 0.9310 / 0.9257; native merge 0.9303 / 0.9295, 0.9356 / 0.9272; main 0.9340 / 0.9333, 0.9234 / 0.9257.
Evidence
  • tests/test_v41_geometry_parity.py stays deleted: both sides build through build_v41_pool_geometry.
  • Geometry equals main's inline construction, layout_id included, for 45 combinations: kv bf16/fp4/fp8 x drafts 0/1/3 x index (fp8, fp4 with None/preshuffle/mfma/natural). Both refuse a layout the format lacks.
  • Namespace identical for 24 configurations: gfx950 opus (mfma), gfx950 flydsl (preshuffle), gfx1250 opus (natural) x FP4/FP8 index x bf16/fp8 KV x chunk 256/8192. That gives 16 distinct namespaces, so the layout reaches the hash.
  • Unit tests: the CI command and tests/plugin, against main 10b99a085 in the same container.
  • GPU: CI's aiter wheel 5c003a5e3, TP4 on MI355X. Plugin two-pass: 0 engine faults, 0 STORE block ID underflow.
  • Native ran at --level 0 on both trees. At level 3, main itself does not finish Inductor compilation in this image: formatting a lowering error (torch/_inductor/exc.py operator_str) stringifies a deeply nested IR and does not return.
  • Native two-pass, main and merge interleaved: MP offload, prefix caching, DSpark, FULL graph. Pass 2 loaded 908k / 888k tier tokens on the merge and 911k / 901k on main. 0 faults on every arm.

@valarLip
valarLip merged commit b986b65 into main Nov 9, 2026
33 checks passed
@valarLip
valarLip deleted the perzhang/dsv41-lmcache-031 branch November 9, 2026 10:09
ZLkanyo009 added a commit that referenced this pull request Nov 9, 2026
Conflict: the kimi_k3 connector import list; keep both StateTouchId
(this PR) and raw_req_id (#2598).
ZLkanyo009 added a commit that referenced this pull request Nov 9, 2026
Conflicts after #2598 moved the shared native-state helpers out of
mp/native_state_worker.py; both sides kept:
- mp/deployment.py: TransferSpec next to require_native_state_server.
- mp/native_state_scheduler.py: restore_slot_count beside the store channel.
- mp/native_state_worker.py: TransferSpec and _mp_session_id only;
  restore_slot_count stays here.
- mp/worker.py: TransferSpec with main's _transfer_block_ids.
- tests/test_lmcache_mp.py: both imports.

Co-authored-by: Cursor <cursoragent@cursor.com>
ZLkanyo009 added a commit that referenced this pull request Nov 9, 2026
…qwen38-lmcache

Brings in #3027's merge of main a30463d (after #2598, which moved the
shared native-state helpers out of mp/native_state_worker.py).

One conflict, in atom/kv_transfer/offload/config.py build_page_namespace:
main's index_cache_layout (#2809) and this branch's QSA_GDN page segments
are both kept, index_cache_layout first.

Co-authored-by: Cursor <cursoragent@cursor.com>
valarLip pushed a commit that referenced this pull request Nov 11, 2026
… object (#3117)

* Kimi-K3 plugin: ride the KDA state on the PAGE object through the LMCache MP server

Kimi-K3 on the MP backend was refused ("recurrent (mamba) state offloads
in-process only"). It now takes the path DeepSeek-V4.1's checkpoints take
(#2598): the KDA state at a boundary rides the PAGE object as a native-state
image, so KV and state commit and restore as one LMCache object.

No copy is needed. vLLM lays every KV cache group out layer-outer from byte 0
of one allocation, so mamba group layer j aliases MLA plane j. On K3 the three
mamba groups of 23 KDA layers sit on the first 23 of the 29 MLA planes, and the
state a group holds in block b is PAGE unit b of those planes. A boundary's
image is one unit per mamba group (3 x 23 x 884,736 B).

- kda_ride.py: KdaRidePlanner (caps hits to the chunk grid, names the load
  destination row (n - 1) // block in every mamba group or none),
  kda_ride_geometry (the checkpoint spec both halves derive from vLLM's
  KV cache config), publish_kda_ride_state (checks on the registered tensors
  that every KDA layer starts at, and strides with, one of the leading PAGE
  planes before declaring the image), and the startup refusals (align mode,
  prefix caching, kv_load_failure_policy=recompute).
- V41StateMPScheduler had nothing V4.1-specific; renamed StateRideMPScheduler
  (ride_mp_scheduler.py) and shared by both models. _v41_request_of is now
  _request_of.
- The KDA state is per rank (TP-sharded heads); the in-process K3 leg is
  unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* Kimi-K3 plugin: size the KDA image by attention layer, and ride only K3

Review of #3117:

- The image covered the leading N PAGE planes with N = mamba layers per
  group, which counts layers, not planes. An attention layer with a DSA
  index cache publishes two planes, so the count fell short. The spec now
  sums each leading attention layer's movable bytes from vLLM's specs (an
  index cache folded into the layer that owns it, as build_kv_cache_tensors
  publishes it), and the worker counts the region planes off the published
  regions' layer names before it checks the aliasing.
- Only Kimi-K3 rides on MP. Its KDA layers report GDN_ATTN, as Qwen's GDN
  layers do, so the gate is the kimi_k3 offload layout; any other mamba model
  on MP is refused as before.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* Kimi-K3 plugin: put the layers the KDA state aliases first in the PAGE order

The first K3 TP8 boot over mp refused at registration: mamba group 0
layer 18 started on PAGE plane 23, outside the image's leading 23 planes.
ATOM publishes PAGE planes by layer name, while vLLM aliases mamba layer j
onto the j-th attention layer it allocates, and with DSpark the two orders
differ (29 attention layers, five sort ahead of the first aliased one).

kda_ride_page_order moves the attention layers the mamba layers alias to the
front, in mamba-layer order, matched on addresses, and the scheduler sizes
the image from vLLM's allocation order. The end-to-end test gains a layer
that sorts first without being aliased; without the reorder it raises the
same refusal the boot did.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* Kimi-K3 plugin: default spec_state_indices_in_tensor so DSpark starts

#2365 made the shared K3 KDA forward read
`kda_metadata.spec_state_indices_in_tensor` on its speculative-decode
branch. vLLM's KimiK3KDAMetadata has no such field, and the plugin only
defaulted its non-spec twin, so K3 on the plugin dies in profiling under
DSpark (AttributeError on the first verify-shaped forward). None keeps the
forward on the write slot, as the non-spec default does.

The same hunk as #3137's, so the two apply cleanly in either order. The
K3 TP8 LMCache MP runs on this PR ran with it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: perzhang <perzhang@amd.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants