Skip to content

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

Open
PerryZhang01 wants to merge 63 commits into
mainfrom
perzhang/dsv41-lmcache-031
Open

PerryZhang01 wants to merge 63 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 Registers recurrent_page_groups as recurrent engine groups, 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; it also fail-stops a transfer past its deadline that no round has polled since, as a submit that never returns would leave it. A polled future holds the server's imported IPC event: it is handed back and dropped on the step thread once LMCache no longer holds it, while the stream is idle. An adapter that records its own producer events is refused. A chunk with no image names -1. 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 enforce_transfer_deadlines is public. 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). 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

  • 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

  • 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

@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 at 3593a34, core offload / MP only this time; the plugin side I'll leave to its own reviewers. Thanks for the round: the slot-owner cleanup, V4's starts=, the double store on descriptor slot 1, _load_poll_lock, PAGE_LAYOUT_VERSION back at 3 and the dsv41 refusal for other models all read right. From reading the code; nothing was run on a GPU.

1. Native V4.1's namespace changes, and the description doesn't say so (offload/config.py:246-247). Native V4.1 used to fall through compress_ratios to hybrid; it now resolves to dsv41, so the page mode becomes dsv41-page-state and dsv41_geometry enters the namespace hash. A native lmcache_mp V4.1 deployment therefore misses on every object stored before the upgrade (the tier is cold after rollout), and an offload_layout override that used to be accepted is now refused. The description says "Native V4.1: same geometry". That may be the right call, but it's a deployment-visible change and needs to be stated in the description and the offload docs.

2. close() destroys the step events before the server has drained (mp/worker.py:357-362). self._step_events.clear() runs before adapter.shutdown(), and adapter.shutdown() is the step that drains in-flight operations. StepEventRing.clear assumes no transfer can still wait on an event, but TransferDrain.close doesn't wait for transfers to reach a terminal state, and InlineTransfers.close returns True at once. A save or load still pending at shutdown leaves the server waiting on an interprocess event this side already destroyed: a HIP invalid-argument error on the server, or a hang. Clearing the ring after adapter.shutdown() restores the precondition.

3. The last look costs up to 1 ms a step while any load waits (mp/worker.py:820-826). With the drain thread, get_finished calls wait_round(_LAST_LOOK_WAIT_S) (1 ms) whenever _wants_last_look_locked(), which is bool(self._pending_loads). The round has to submit what's queued first, and drain.py:17 measures a submit at 18.7 ms, so the 1 ms window almost never catches a completion. For as long as any request sits in WAITING_FOR_REMOTE_KVS, each step spends up to 1 ms of step-thread time for nothing. That's the hot-path cost the drain thread was added to remove. Waiting only when nothing is queued to submit (a pure poll round), or not at all, would fix it.

4. Still open from 10-29 (item 8): tp_rank_collapse=auto never collapses V4.1 (mp/deployment.py:161-167). The shared rule returns False for every state-owning layout "until their backend declares otherwise", but the V4.1 backend already declares it: tp_replication_factor=tp_size and native_state_tp_replication_factor=tp_size (deepseek_v41/backend.py:501-502). So native V4.1 at TP4 with auto has every rank write, and the tier holds 4 identical copies. The recipe sets false explicitly, so auto is never exercised. Having the rule read the factor the backend declares (Kimi-K3 declares less than TP and stays uncollapsed) would replace the layout list with the real criterion.

5. Dead drain-era scaffolding in NativeState (mp/native_state_worker.py:53, :221-233). 3593a34 put NativeState back on InlineTransfers, but the drain-thread machinery stayed:

  • queued is set and cleared, but nothing observes it, so the not entry.queued filter never excludes anything;
  • the _transfers_in_flight_locked override (:376) only feeds wake(), which is a no-op inline;
  • _outstanding_locked (:383) is only reached through fence, which only the plugin's drain calls;
  • InlineTransfers.fence and _FENCE_POLL_S (drain.py:277-296) run only in a unit test;
  • worker.py:820 still branches on self._transfers.inline instead of the two strategies sharing one wait_round contract.

That's ~100 lines no production path exercises. Please either return NativeState to main's polling in get_finished, or make it actually share the drain.

6. Unrelated or out-of-date bits:

  • metadata.py:51 only lengthens a chunk-size error and adds a comment about Kimi-K3's 128 → 1536 block size. That's a K3 change, so please move it to its own PR.
  • tests/test_lmcache_mp.py:732/752/883 rewrite lambda: next(ticks) as iter(...).__next__, which changes nothing; please revert.
  • chunked_scheduler._retire_finished_save (:674) hand-writes the release-lease sequence a third time, next to _source_group_finished (:1239) and _release_operation_lease (:1365). One helper would keep them from drifting.
  • The description still lists RecurrentPageGroup / RecurrentStateTransfer / install_recurrent_state_hook, disaggregation/types.py and mooncake_store/scheduler.py, none of which are in the diff. Its mp/scheduler.py and chunked_scheduler.py rows describe changes that aren't there, while it omits items 1 and 4.

- NativeState submits on the step thread and polls in `get_finished`,
  as on main, holding the step's event from the ring until the transfer
  is terminal. The drain-era `queued`/`terminal` state, the
  `_perform_submit`/`_poll_round` overrides and `_outstanding_locked`
  are gone.
- `get_finished`'s last look waits only for a pure poll round (nothing
  queued to submit, no round running): a round that has to submit first
  outlasts any wait short enough for a step. Dropping the look entirely
  raised the plugin's c16 TTFT p50 from 302 to 324 ms. Both executors
  share one `report_round(wait_s)` (inline: the step's poll; drain: a
  wake and that bounded look), so nothing branches on the strategy.
  `InlineTransfers` keeps only what the worker calls; its `fence`
  refuses, since only the plugin fences.
- `close` destroys the step events only after `adapter.shutdown()` has
  drained the server, and keeps them when the shutdown fails.
- The three lease releases in `chunked_scheduler.py` share
  `_return_leased_blocks`.
- Reverted: the Kimi-K3 chunk-size message in `metadata.py` and the
  `iter(...).__next__` test rewrites.
- The offload README states native V4.1's `dsv41` layout, its one-time
  cold tier after upgrading, and that `tp_rank_collapse: auto` keeps
  state-owning layouts one copy per rank.

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

Copy link
Copy Markdown
Contributor Author

Round 15 pushed 7081137 for the second pass of the 11-02 review: items 1, 2, 3, 5 and 6 are fixed; 4 stays as on main.

  • branch @ 7081137: unit tests fail only what 3593a34 fails; ruff, black clean.
  • Native V4.1 two-pass, interleaved with main: pass 1 in 62-63 s against 69 s, more stored per second, same accuracy.
  • Plugin DSpark two-pass: 0.9219 / 0.9272, 0 faults. Perf c16: 9.05 req/s, TTFT p50 307 ms, as 3593a34.
Per item
  1. Kept, and stated: the description's shared-code table and the offload README say native V4.1 resolves to dsv41, its tier is cold once after upgrading, and an override it accepted before is refused.
  2. close destroys the step events only after adapter.shutdown() has drained the server, and keeps them when the shutdown fails. A test covers both orders and both executors.
  3. The last look waits only for a pure poll round: nothing queued to submit and no round running. Dropping it entirely raised TTFT p50 from 302 to 324 ms (table below), so it stays in that form. Both executors share report_round(wait_s) (inline: the step's poll), so worker.py no longer branches on the strategy.
  4. Unchanged, as on main. auto is decided before any backend exists, and the backend modules that could declare the factor import AITER, which the CPU CI cannot load. The README and Known limitations say auto keeps native V4.1 per rank and that true stores one copy; the plugin resolves auto to true for V4.1.
  5. NativeState submits on the step thread and polls in get_finished, as on main; queued, the _perform_submit/_poll_round/_transfers_in_flight_locked/_outstanding_locked overrides and InlineTransfers.fence's poll loop are gone. Inline fence refuses: only the plugin fences.
  6. metadata.py and the three iter(...).__next__ rewrites are reverted; the three lease releases share _return_leased_blocks; the description's tables now match the diff.
Runs (cards 4-7, same container)

Native and plugin two-pass ran at 28593d185; 7081137 changes only the plugin's last look.

Arm Code Pass 1 Stored by end of pass 1 (per second) Pass-2 loads p1 / p2
Native main e281d58 69 s 897,024 (13.0k) 953,600 0.9303 / 0.9257
Native 28593d185 62 s 848,896 (13.7k) 919,552 0.9272 / 0.9295
Native main e281d58 69 s 878,080 (12.7k) 943,360 0.9272 / 0.9234
Native 28593d185 63 s 857,600 (13.6k) 926,976 0.9318 / 0.9303

Plugin, DSpark, two-pass at 28593d185: p1 0.9219, p2 0.9272; pass-2 tier hits 5.84M / 7.18M; 0 stale state.

Plugin perf, DSpark, c16, 900 s, OSL 64, HBM pool 4000 blocks, one run per row:

Code Last look req/s TTFT avg / p50 / p90 (ms) ITL avg (ms)
7081137 pure poll only 9.05 318 / 307 / 389 22.96
28593d185 none 8.94 338 / 324 / 418 23.01
28593d185 none 8.68 376 / 330 / 433 23.24
3593a34 always, 1 ms 9.04 313 / 302 / 380 23.07
3593a34 always, 1 ms 8.93 324 / 307 / 399 23.24

Conflicts were with #2906 (lmcache_mp lookup thread, transfer thread,
same-step saves) in `mp/worker.py` and `mp/native_state_worker.py`.

Native ATOM's MP path follows main: its producer event is
`_producer_fence` (a fresh event per step, or `_ProducerFences` under
OFFLOAD_MP_TRANSFER_THREAD), recorded on the forward thread and passed
as a `RecordedStepEvent`; submissions go through `_dispatch` and the
optional `_TransferThread`; same-step saves, `_submit_kwargs` and
`_finished_housekeeping` are main's. The step-event ring now serves the
plugin's drain thread only, which does not start main's transfer
thread. The base worker's `transfer_thread` argument is renamed
`drain_thread` so it does not read as main's transfer thread.

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

Copy link
Copy Markdown
Contributor Author

Round 16 pushed 9e70cbc, merging main 498133e: native ATOM keeps #2906's MP path; the step-event ring serves only the plugin's drain (drain_thread).

  • branch @ 9e70cbc: unit tests fail a subset of what main fails; ruff, black clean.
  • Native V4.1 two-pass, interleaved with main: same accuracy, pass 1 in 65 s against 69-70 s.
  • Plugin DSpark: two-pass 0.9189 / 0.9272, 0 faults; perf c16 9.06 req/s, as before the merge.
Runs (cards 4-7, same container)
Arm Code Pass 1 Stored by end of pass 1 (per second) Pass-2 loads p1 / p2
Native main 498133e 69 s 895,744 (13.0k) 956,416 0.9257 / 0.9242
Native 9e70cbc 65 s 814,848 (12.5k) 884,032 0.9234 / 0.9265
Native main 498133e 70 s 923,392 (13.2k) 975,872 0.9257 / 0.9287
Native 9e70cbc 65 s 836,096 (12.9k) 905,216 0.9257 / 0.9295

Plugin, DSpark, two-pass at 9e70cbc: p1 0.9189, p2 0.9272; 0 stale state, 0 invalid-argument waits.

Plugin perf, DSpark, c16, 900 s, OSL 64, HBM pool 4000 blocks, one run each:

Code req/s TTFT avg / p50 / p90 (ms) ITL avg (ms)
9e70cbc 9.06 323 / 311 / 397 22.86
7081137 9.05 318 / 307 / 389 22.96

`OwnedFutures` now holds futures only for the plugin's drain thread
(`hold=drain_thread`). Native ATOM, all on the step thread, drops a
future where it finished, as on main. Held until the stream was idle,
native's steps ran 6% faster but its saves fell behind: a GSM8K
two-pass on V4.1 loaded 7% fewer tier tokens in pass 2. Dropped, pass
time, stores per second and pass-2 loads sit within main's spread.

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

Copy link
Copy Markdown
Contributor Author

Round 17 pushed 97c8dde: native ATOM drops a finished future where it finished, as main does; holding futures is now the plugin's drain thread alone.

  • branch @ 97c8dde: unit tests fail only what 9e70cbc fails; ruff, black clean.
  • Native V4.1 two-pass with futures dropped, interleaved with main: pass time, stores per second and pass-2 loads within main's spread.
Native V4.1, LMCache MP, GSM8K two-pass (cards 4-7, same container)

Held futures cost no accuracy but made pass 1 end sooner while saves kept their pace, so pass 2 found fewer of them. A counter showed the stream busy at most step starts, holding up to 422 futures.

Code Futures Pass 1 Stored by end of pass 1 (per second) Pass-2 loads p1 / p2
main 498133e dropped 68 s 913,152 (13.4k) 968,192 0.9295 / 0.9340
9e70cbc + drop dropped 70 s 889,600 (12.7k) 955,904 0.9280 / 0.9257
9e70cbc held until idle 65 s 828,672 (12.8k) 898,304 0.9272 / 0.9242
main 498133e dropped 69 s 884,480 (12.8k) 948,224 0.9227 / 0.9287
9e70cbc + drop dropped 70 s 907,520 (13.0k) 966,656 0.9325 / 0.9242

97c8dde changes only native's future lifetime, which the "+ drop" rows ran; the plugin keeps its drain path unchanged.

`dsv41` changed native ATOM's V4.1 namespace (page mode
`dsv41-page-state`, plus `sliding_window` and `candidate_block_size`),
so a tier an earlier build filled missed on every object after the
upgrade. Nothing native needed from it: its native-state key already
carries the checkpoint layout.

`build_page_namespace` now gives `dsv41` the `dsv4-page-regions` page
mode `hybrid` gave V4.1 and hashes no V4.1 fields, so native V4.1's
namespace is byte-identical to main's. The vLLM plugin, which had no
V4.1 tiers before, adds the two fields to the HF geometry its config
view snapshots, in the worker and the scheduler alike.

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

Copy link
Copy Markdown
Contributor Author

Round 18 pushed 85d55d8: native V4.1's namespace is main's again, correcting my reply to item 1 of the 11-02 review; its tiers stay reachable.

  • dsv41 keeps the namespace hybrid gave V4.1; only the plugin's config view hashes the two V4.1 geometry fields.
  • Same V4.1 config: main ...b06c652d19c0fb36, 85d55d8 the same, 97c8dde ...ec795cfd458be88f.
  • branch @ 85d55d8: unit tests fail only what 97c8dde fails; ruff, black clean.
  • Plugin DSpark two-pass: 0.9249 / 0.9310, tier hits unchanged.

@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 offload / MP side at 85d55d8, by reading the code and running 7 MP worker / scheduler / config test files (329 passed); no GPU. The plugin side I'm leaving to its own reviewers. Thanks for rounds 16–18: native now keeps #2906's MP path, drops finished futures as main does, and V4.1's namespace is main's again, which settles my earlier namespace and drain-cost points for native.

What's left is in the core files that this PR adds for the drain (used today only with drain_thread=True), plus two scheduler fixes. Each item ends with the change I'm asking for.

A. Correctness

1. StepEventRing has no re-record bound (mp/step_events.py:39-49). record_locked pops a free index and records the same interprocess event again, with no count. mp/transfer_thread.py:35-37 documents that on ROCm 7.2.4 a server stream wait on an event its exporter has recorded more than 32 times fails, which is why _ProducerFences swaps in a fresh event after _FENCE_MAX_RECORDS = 31. In steady state the drain's transfers finish before the next step, so the ring re-records the same index every step and passes 32 records after ~32 steps. The two modules also disagree on the semantics: the ring says re-recording an event a transfer still waits on errors, while _ProducerFences says the wait just becomes longer. (The invalid argument waits mentioned earlier in the thread fit this.)
→ Change: settle which statement is true and write it once. Give the ring the same per-event record count and replacement as _ProducerFences, or better, item 3.

2. The drain bypasses the pooled-producer-event refusal (mp/worker.py:482-490). if not envs.OFFLOAD_MP_TRANSFER_THREAD or self._drain_thread: return skips the check that an adapter with its own _producer_events must keep transfers on the forward thread. In drain mode _fences is None, so producer_event_recorded=True is never passed either. An adapter that records its pooled event on the drain thread's current stream records it after nothing in particular, so a save can read KV the forward hasn't finished writing, and a load can overwrite blocks the forward is reading. This depends on the adapter (I don't have LMCache here), but the guard exists for exactly that adapter.
→ Change: apply the same refusal in drain mode (_adapter_takes_recorded_events() or refuse), and pass producer_event_recorded from the drain's submit path as the transfer thread does.

3. A refused or withdrawn load doesn't lower the save floor (chunked_scheduler.py:1419, :854). load_failed lowers entry[1] only when _load_save_floors[sid] is set, and that's written only at dispatch (:854). The tracker's floor is set to the whole tier hit at tracking time (chunked_scheduler.py:575, _lmcache_hit_save_floor). So a load refused before dispatch (e.g. tier_lost_prefix, the tier evicted the prefix) leaves the floor at the hit: the request recomputes [HBM, hit) but never stores that span back, the exact span the tier just lost.
→ Change: record the floor when tracking starts, not at dispatch, so load_failed can always restore it, or have load_failed fall back to the HBM hit when no dispatch floor exists.

4. A withdrawn promise leaks the lookup's read locks (mp/scheduler.py:283-292). With active is None, a raw-id load_failed counts as current and calls complete_retrieve(False) for a retrieve that was never submitted. That pops _lookups, so the following clear_lookup_status sees state is None and skips free_lookup_locks. The server then holds the read locks until session or TTL cleanup, blocking eviction and writes for that prefix. The same withdrawal also spends the request's one tier attempt through _record_failed_load_attempt. (From reading; forcing a promise past its 50 steps and checking the server lock count would confirm it.)
→ Change: for an operation that was never submitted, release the lookup's locks directly (free_lookup_locks) instead of complete_retrieve, and don't count it as a failed tier attempt.

B. Structure

5. The drain is a second implementation of the transfer thread, living in core (mp/drain.py:43, mp/step_events.py, mp/ipc_lifetime.py). TransferDrain, StepEventRing / PendingStepEvent / RecordedStepEvent, OwnedFutures, and InlineTransfers (an empty object that exists to match a constructor, drain.py:238) re-implement _TransferThread, _ProducerFences and take_spent from transfer_thread.py. The only caller is drain_thread=True from the plugin connector. The worker now holds both self._transfer and self._transfers, one letter apart, and worker.py is 1028 lines. Items 1 and 2 are the two copies already disagreeing.
→ Change: add what the drain needs (a poll round and a fence) to transfer_thread.py's thread and fences, so there's one submission thread and one event pool. Then delete drain.py, step_events.py, ipc_lifetime.py, InlineTransfers and the drain_thread switch. If the plugin needs something native doesn't, it belongs in the plugin's adapter, not in the core worker.

6. Single-model hooks and knobs in shared code (chunked_scheduler.py:666, mp/worker.py:136, config.py:98).

  • _unbuilt_save_is_final() is a new hook on the shared base that only V41StateMPScheduler overrides.
  • lmcache.mp.null_block_id adds a layout field, a parameter, a check, a worker attribute and a README paragraph for a value only tests set.
  • _DSV41_MODEL_TYPES is a hand copy of selector._V41_TYPES.

→ Change:

  • In v41_mp_scheduler._build_save_request, call _retire_finished_save where it declines, and delete the hook and its branch (:938-946).
  • Keep -1 as the constant in layout.transfer_block_ids and delete the knob.
  • Use attn_family(hf) is Family.CSA2 instead of the copied set.

zgplvyou and others added 2 commits November 3, 2026 08:51
Addresses valarLip's core review of 85d55d8.

The vLLM plugin's off-step submit and poll run on #2906's transfer thread
instead of a second implementation of it (items 1, 2, 5):

- `_TransferThread` takes an optional `poll`. With it, a round of polls
  runs whenever the queue is empty, `look` wakes it from `get_finished`,
  and `fence` waits on its rounds up to the transfer deadline.
  `hand_back`/`take_returned` drop a polled future on the forward thread
  once LMCache no longer holds it.
- `LMCacheMPConnector.run_transfers_off_step` starts it with
  `_ProducerFences`, so the plugin's events have the env-gated thread's
  re-record bound, and refuses an adapter that records its own pooled
  event: neither thread could fence it.
- `drain.py`, `step_events.py`, `ipc_lifetime.py` and the `drain_thread`
  switch are removed. Native ATOM keeps main's inline path and its
  env-gated transfer thread.

Scheduler fixes (items 3, 4):

- The save floor a failed load restores is recorded at allocation, and a
  load refused as `tier_lost_prefix` reopens `[HBM, hit)` for saving, in
  the chunked and DSV4 schedulers: the request recomputes the span the
  tier lost, and it was never stored back.
- A promise withdrawn before dispatch no longer calls `complete_retrieve`,
  which forgot the lookup first, so the lookup cleanup frees its read
  locks.
- The withdrawal still counts as the request's failed tier attempt; that
  is what stops a re-promise/withdraw loop
  (test_dsv4_failed_load_is_not_retried_for_the_same_request,
  test_a_failed_load_leaves_no_remembered_hit_behind).

Single-model hooks (item 6): V4.1's MP scheduler calls
`_retire_finished_save` where it declines, replacing the shared
`_unbuilt_save_is_final` hook; `lmcache.mp.null_block_id` is dropped for
the `-1` constant; V4.1 is detected with `attn_family(...) is
Family.CSA2`.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A future the plugin's polling transfer thread finished holds the
server's imported IPC event, and on ROCm 7.2.x destroying one
synchronises every stream of the device. 0132de1 dropped them in
`get_finished`, mid-forward: three interleaved c16 pairs against
85d55d8 ran 3.8% fewer requests per second (8.65 vs 8.99) with ITL
p50 23.85 vs 23.05 ms, every new run below every old one.

The forward thread now drops them only while its stream is idle, as
85d55d8's drain did, and a stream that is never idle (async
scheduling) drops a batch per call past 8192 held. Native ATOM hands
no futures back and is unchanged.

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

Copy link
Copy Markdown
Contributor Author

Round 19 pushed 4b589d8. The six items of the 11-03 core review are addressed, except 4b.

  • 1, 2, 5: the plugin runs on _TransferThread with a poll round and _ProducerFences, and refuses a self-recording adapter. The drain modules are deleted.
  • 3: the floor is set at allocation; tier_lost_prefix reopens [HBM, hit), DSV4 too.
  • 4a: a withdrawn promise skips complete_retrieve; cleanup frees the locks.
  • 4b: the attempt stays spent, or the request re-promises.
  • 6: as suggested.
Evidence

4b: without the spent attempt, test_dsv4_failed_load_is_not_retried_for_the_same_request and test_a_failed_load_leaves_no_remembered_hit_behind fail: the withdrawn request is promised and withdrawn again.

Unit, in the PR's container at 4b589d8: the CI command and tests/plugin fail exactly what they fail at 85d55d8. The 7 new tests for items 1-4 each fail with their fix reverted.

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

Run Result
Plugin GSM8K two-pass, 4b589d8 0.9204 / 0.9219, 5.85M tier hit tokens, 0 faults, 0 invalid argument
Native V4.1 two-pass, main 498133e x2 0.928 / 0.9295, 0.9265 / 0.9249
Native V4.1 two-pass, 0132de1 x2 (interleaved) 0.9257 / 0.9265, 0.934 / 0.9356; no transfer thread, as on main
Plugin c16 900 s, 85d55d8 x5 8.92-9.08 req/s, ITL p50 22.70-23.16 ms
Plugin c16 900 s, 4b589d8 x2 (ABBA) 8.82 / 8.90 req/s, ITL p50 22.84 / 22.92 ms

0132de1 dropped the polled futures mid-forward and ran 8.41-8.77 req/s, ITL p50 23.71-24.07 ms; 4b589d8 drops them on an idle stream. 4b589d8 changes nothing native runs.

@valar-cop valar-cop left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

a few small things..

ci green.

Comment thread atom/kv_transfer/offload/mp/worker.py Outdated
plugin, where a submit cost 18.7 ms and a terminal poll 13.8 ms under
its launch pressure. Called once the caches are registered.
"""
if self._transfer is not None or self._rank is None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

plugin + this env hard-fails startup here, misleading msg. replace the non-polling thread instead, or skip the env-gated start when frontend will poll?

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 45355d2. The plugin builds the worker with transfers_off_step=True, and registration starts the polling thread whatever OFFLOAD_MP_TRANSFER_THREAD says; run_transfers_off_step is gone. Covered by test_off_step_transfers_poll_whatever_the_transfer_thread_flag_says.

# from a remembered hit has none yet.
if not self._ensure_lookup_pin(seq, sid, ls):
self._mark_load_skip(seq, "tier_lost_prefix", hbm, lmc, need, chunk)
# The tier no longer holds the prefix this request recomputes.

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

same 3 lines inline twice, and chunked scheduler has _reopen_saves already. move it into OffloadSchedulerMixin and call from both?

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 45355d2. _reopen_saves is on OffloadSchedulerMixin; the DSV4 scheduler calls it for tier_lost_prefix and in load_failed, and the chunked scheduler's copy is gone.

@valarLip 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 offload / MP side at 4b589d8, by reading the code; no GPU. Thanks for round 19. The drain modules are gone, the plugin now runs on _TransferThread with _ProducerFences, a self-recording adapter is refused, the floor is set at allocation, and a withdrawn promise frees its locks. I accept your reason for keeping the spent attempt (4b). I'm not repeating the DSV4 inline reopen (valar-cop's comment). Each item ends with the change I'm asking for.

A. Correctness

1. The polling thread polls only when its queue is empty (mp/transfer_thread.py:280-310). The _WAKE branch is if self._queue.empty(): self._poll_round(), the post-item poll has the same condition, and the _IDLE_POLL_S backstop fires only on queue.Empty. So while the queue never drains, every wake from look() and from a fence is dropped, which contradicts the module docstring ("polls when woken"). Under heavy prefill, a submit takes several to ~20 ms and an after-forward save also waits in _wait_ready. If each step's submits take at least a step's time, the queue never empties: loads that have finished are never polled and requests sit in WAITING_FOR_REMOTE_KVS behind background saves, and a fence's round never comes, so it waits out the transfer deadline and fail-stops. (From reading. Watching _rounds_done under concurrent long prefills would confirm it.)
→ Change: make a wake from look() / a fence force a poll round on the next loop iteration, whatever the queue holds (e.g. a _poll_requested flag checked before taking the next item), or put loads and saves on separate queues so a load's completion isn't behind saves.

2. The in-process offload shell doesn't refuse dsv41 (offload/connector.py:40-108). V4.1 now resolves to the dsv41 layout (offload/config.py:234-235), but _build_worker / _build_scheduler have no branch for it and fall through to DenseOffloadConnector / DenseOffloadScheduler. On main it went to DSV4 hybrid. The selector's own comment says the dense codec must never serve V4.1 (it would restore KV over a window ring that doesn't exist). Nothing breaks today only because two gates in other files stop it first (deepseek_v41/config.py:382 and the plugin's refusal). The README says dsv41 is lmcache_mp-only.
→ Change: in the shell, raise for dsv41 with a message pointing at lmcache_mp, so the one place that maps layouts to implementations can't route V4.1 to the dense codec.

B. Performance

3. Up to 1 ms of step-thread wait per step while any load is pending (mp/worker.py:960-967). wait_s = _LAST_LOOK_WAIT_S if self._pending_loads else 0.0, then transfer.look(wait_s), on every get_finished, decode-only steps included. In your DSpark run (82.9% of tokens from the tier) a load is in flight almost every step, so that's up to 1 ms on a ~23 ms ITL. With item 1 the poll round runs when woken, so the step thread needn't wait for it.
→ Change: look() should wake the poller and return. Report what the previous round polled, and let a load that completes in between be reported next step, as main does. If you keep a wait, show the ITL with and without it.

C. Boundaries

4. A plugin-only transfer mode lives in the core worker (mp/worker.py:422, :932, :952, :960). run_transfers_off_step() is called only by the plugin. fence() raises NotImplementedError in the other modes, and get_finished / _retire_future branch on transfer.polls, so what the core worker does depends on who drives it. Future retirement is inferred from sys.getrefcount (transfer_thread.py:62-64, :184). Also:

  • native_state_layout.py:24 imports deployment._storage_kv_transfer_config (private) for a check that belongs in deployment.py;
  • the offload README still describes mp/worker.py as the "Generic PAGE-only" half (:110) and doesn't mention the plugin constructing LMCacheMPConnector directly.

→ Change: give _TransferThread the poll round as a constructor option the front end passes when it builds the worker, so the worker has one mode with a flag, not a mode switched on afterwards; then delete run_transfers_off_step. Retire futures by completion state (future.done() plus the stream-idle check), not refcounts. Move the native-state server check into deployment.py as a public function, and update the README table and the paths section.

5. Native's inline path lost its test coverage, and native_state_worker carries no-op rewrites (tests/test_lmcache_mp.py:1121, mp/native_state_worker.py). The shared _worker fixture now starts the plugin's polling thread (_run_off_step), so the ~23 failure-path tests no longer exercise native's inline get_finished; 3 test_inline_* tests remain. Yet that path was rewritten: load polled before save, lock-free snapshot, results through _polled_*. The query-exception test also lost main's "recovers and reports finished_*" assertion. Meanwhile native_state_worker.py (+73 / −187) has changes with no behaviour change:

  • _submit_native's check-and-insert is split into two lock sections;
  • _poll_native_load is split up;
  • _emit_page_source_safe is renamed;
  • restore_unprovable replaces main's pending.future = _UnprovableSubmission().

→ Change: parametrize the worker fixture over inline and polling, so every failure-path test runs on both, and restore the recovery assertion. Revert the native_state_worker.py restructuring that doesn't change behaviour, keeping only what the poll round needs.

…w fixes

Addresses valarLip's and valar-cop's reviews of 4b589d8.

- `LMCacheMPConnector(config, transfers_off_step=True)` replaces
  `run_transfers_off_step()`: the frontend asks for the polling transfer
  thread when it builds the worker, and registration starts it whatever
  OFFLOAD_MP_TRANSFER_THREAD says. With that flag set, plugin startup no
  longer fails after the env had started a submit-only thread.
- A round asked for by `look` or a fence runs after the submission in
  progress, whatever is still queued, so a backlog of submissions holds
  back neither a finished load's report nor a fence.
- The in-process offload shell refuses the `dsv41` layout instead of
  routing DeepSeek-V4.1 to the dense codec.
- `_reopen_saves` lives on `OffloadSchedulerMixin`; the DSV4 scheduler
  calls it in place of its two inline copies.
- `require_native_state_server` is public in `mp/deployment.py`.
- `native_state_worker.py` is main's apart from the registration that
  moved to the base worker; the rewrites that changed no behaviour are
  reverted.
- The worker tests' failure paths run on both the step-thread worker and
  the polling one, and the query-exception test checks recovery again.
- README: the module table and the plugin's direct use of the worker.

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

Copy link
Copy Markdown
Contributor Author

Round 20 pushed 45355d2. All five items are addressed, except retiring futures by done().

  • 1: an asked-for round runs after the submission in progress.
  • 2: the in-process shell raises for dsv41.
  • 3: the wait stays: without it, c16 ITL p50 rose from 22.8 to 23.4 ms.
  • 4: done() holds for every handed-back future; LMCache still holds it until its next submission (10-26).
  • 5: 22 failure-path tests run in both modes.
Evidence

Item 4: the worker takes transfers_off_step at construction and run_transfers_off_step is gone; require_native_state_server is public in mp/deployment.py; the README's module table and lmcache_mp section describe the plugin's use of the worker. A future is handed back only once terminal, and LMCache's AtomMPWorkerAdapter keeps it in _context_operation_futures until _track_operation_future prunes it at its next submission, on the transfer thread. Dropping ours on done() would leave the last reference, and the imported event's destroy, to that submission.

Item 5: native_state_worker.py is main's apart from the registration that moved to the base worker (+3 / -143 against main); the query-exception test checks recovery again.

Unit, in the PR's container at 45355d2: the CI command and tests/plugin fail exactly what they fail at 85d55d8. The new tests for items 1 and 2 and for the env flag fail with their fix reverted.

GPU, MI355X TP4, DeepSeek-V4.1-Flash DSpark, FULL_AND_PIECEWISE. "No wait" is the same tree with _LAST_LOOK_WAIT_S = 0.0, run interleaved (wait, no wait, no wait, wait):

Run Result
Plugin GSM8K two-pass 0.9249 / 0.9257, 5.85M tier hit tokens, 0 faults, 0 invalid argument
Native V4.1 GSM8K two-pass 0.9295 / 0.9265, no transfer thread
Plugin c16 900 s, 1 ms wait x2 8.77 / 9.00 req/s, ITL p50 22.97 / 22.64 ms, TTFT p50 318 / 319 ms
Plugin c16 900 s, no wait x2 8.59 / 8.55 req/s, ITL p50 23.47 / 23.32 ms, TTFT p50 327 / 325 ms

@valar-cop valar-cop left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

a few small things..

ci green.

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

Follow-up on the current head 45355d23a. My previous review (on f5d97882e) covered only the non-plugin files; this one covers the whole PR, plugin included. Several earlier points are no longer reported as fixed or no longer applicable. The relevant CPU tests pass locally (test_v41_index_plane_parity, test_v41_plugin_refusals, test_lmcache_mp*, test_offload_shell: 240); no GPU runs. Statuses are noted inline where a point is plausible rather than confirmed.

Action items for the author

Affects other models — please fix first

  • E1 Restore exact decode-keep starts for DeepSeek-V4 (deepseek_v4_bridge.py:1372).
  • E3 Report block frees only when they actually return blocks; don't stall preemption on pins/stores (boundary_planner.py:208).
  • E4 Revert the V4.1 state-probe default flip and keep the env escape hatch; move it to its own PR with a test (v41_core_overrides.py:337).
  • E8 Scope has_pending_push_work's scheduler check to MP (connector.py:2075).

V4.1 correctness

  • E2 One STATE-tail owner that counts every KV group (deepseek_v41_state_reserve_patch.py:107).
  • E5 Let checkpoint-restore failures raise (v41_core_overrides.py:404).
  • E6 Decide restore by slot owner, not request-id text (v41_core_overrides.py:352).
  • E9 Remove the cache.pending = None bypass; run the checkpoint step next to the keep (deepseek_v41_bridge.py:2133).
  • E10 Return (empty, 0) when the local hit caps to 0 (plugin/vllm/scheduler.py:262).
  • E12 Translate the connector name to lmcache_mp instead of clearing kv_transfer_config; close the module-path bypass (platform.py:254).
  • E13 Swap the transfer error under the lock (transfer_thread.py:161).

Performance

  • E7 Release returned futures on a recorded event, not stream idleness (mp/worker.py:517).
  • E11 Coalesce WAKEs; don't wait on decode steps or interleave polls into the submit backlog (transfer_thread.py:200).

Remove transitional design / duplication

  • E14 One TP-rank-collapse rule; drop the config rewrite and connector V4.1 branch; explain or fix the recipe's false (offload_config.py:172).
  • E15 Move plugin-only transfer modes out of the core worker; stop monkeypatching paged_state_checkpoint.py; remove the copied get_kv_transfer_tensors, the second MP scheduler, duplicate rules/helpers, forwarders and dead checks (mp/worker.py:425).

name="DeepSeek-V4",
# The first-token positions `_build_dsv4_metadata` derives, exact under
# MTP; the input batch's V2 count is an optimistic upper bound.
starts=np.asarray(md.chunk_start_per_seq_cpu[:rows], dtype=np.int64),

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.

[E1] DeepSeek-V4 (not V4.1) decode keep can restore wrong state. The decode-keep start is now md.chunk_start_per_seq_cpu, which on decode rows comes from seq_lens_cpu_upper_bound (all drafts accepted), instead of main's _exact_starts reading vLLM's corrected num_computed_tokens_cpu. With V4 + MTP + async scheduling (vLLM's default) + prefix caching on the V1 runner — nothing refuses async+spec for V4 — on a step whose drafts are rejected, by_end.get(starts[i]) looks up the optimistic end: it either stores stale state as boundary B's checkpoint or drops an already-hashed B. A later hit on that prefix restores wrong state and output is silently wrong. The comment's "exact under MTP" doesn't hold for decode rows (bridge:1724 is seq_np - lens). Please keep exact starts for V4.

v41_proxy_state_reserve_blocks,
)

if _state_checkpoints_own_the_tail(vllm_config):

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.

[E2] One STATE-tail owner, counting every group. With checkpoints on, _state_checkpoints_own_the_tail makes the reserve patch return 0, so its own "one request must fit" check — which counts all KV groups — no longer runs; the remaining check in paged_state_checkpoint.py:470/503 computes per_req over the proxy group only. With DSpark + --enable-prefix-caching + align mode and a pool sized between the proxy-only and all-groups requirement, startup passes but a full-max_model_len request can never be admitted — exactly the livelock this patch's own comment describes. Two wrappers own the same thing; keep one, and count every group.

"""Whether only a further engine step can finish this leg's work."""
return False

def has_pending_frees(self, horizon_s: float) -> bool:

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.

[E3] Don't block preemption on frees that can't free anything. The new has_pending_block_frees makes vLLM break instead of preempting on an allocation failure, and this KDA leg returns True whenever _pending_stores is non-empty — ignoring the horizon and whether releasing the pin actually frees a block; the lease and deferred-free legs have no sole-holder check either. Kimi-K3 + LMCache under KV pressure: the pinned boundary block is still in a running request's block table, so releasing it frees nothing, but with KDA stores in flight (nearly every step at steady state) vLLM's scheduler.py:925 breaks, nothing is preempted and later running requests aren't scheduled that step — indefinitely while stores keep coming (lease/deferred legs stall up to ~10 s each). Only report frees that will actually return blocks to the pool.

if override is not None and override != "":
blocking = override not in ("0", "false", "False")
else:
blocking = not serving_synchronously()

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.

[E4] Revert this default flip (unrelated to offload). main's default blocking state probe (with the comment "keep this workaround until a separate test does") and the ATOM_V41_BLOCKING_STATE_PROBE escape hatch are removed in favour of blocking = not serving_synchronously(). This affects every V4.1 plugin run, connector and checkpoints or not: on the V2 runner with synchronous scheduling and FULL_AND_PIECEWISE, the blocking read added for stale-state reporting becomes a deferred read, with no test driving prepare_state under serving_synchronously() == True to show it's safe, and users who set the env silently lose it. Please move it to its own PR with that test.

)
for i in wrong
]
try:

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.

[E5] Don't swallow restore failures. _restore_stale_rows wraps the checkpoint restore in except Exception: logger.warning, defeating its fail-stop: if one row's checkpoint unit is null, state_checkpoint_unit_ids raises "refusing to run from an unrestored slot", no row in the batch is restored, the real cause is a WARNING, and core then reports "needs state at N, found -1" — a -1 the bridge just wrote — pointing at the stale-state contract instead of the failed restore. Let it raise.

settles is reported on a later step."""
with self._idle:
pure = self._outstanding == 0 and self._rounds_done == self._rounds_started
target = self._ask_round()

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.

[E11] Don't wait on decode steps or interleave polls into the submit backlog. With any load in flight, get_finished makes the forward thread wait up to 1 ms every step, decode included; look() calls _ask_round and enqueues a _WAKE every step, forcing poll rounds into the middle of a submit backlog, and stale _WAKEs run extra rounds when dequeued. At c16 with prefix-hit loads that's up to ~1 ms per decode step (~4% of a 23 ms ITL), and one or two full poll rounds per step inserted into the save queue (this PR prices a terminal poll at 13.8 ms). Coalesce WAKEs and only wait when the step has no runnable work. (Plausible — worth counting rounds per step to confirm.)

name = getattr(kv, "kv_connector", None)
if name is None and isinstance(kv, dict):
name = kv.get("kv_connector")
return name == "AtomLMCacheOffloadConnector"

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.

[E12] The connector whitelist can be bypassed. plugin/config.py:483 clears V4.1's kv_transfer_config to {}, so the model's own validate_runtime_config never sees the connector, and this whitelist compares only the class name, ignoring kv_connector_module_path (which vLLM's factory prefers). {"kv_connector": "AtomLMCacheOffloadConnector", "kv_connector_module_path": "other.mod"} passes and loads other.mod's class, skipping every V4.1 refusal. Root fix: translate the vLLM connector name to ATOM's lmcache_mp instead of clearing the config.

@@ -118,6 +161,106 @@ def raise_error(self) -> None:
if error is not None:

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.

[E13] Swap the error under the lock. raise_error does error, self._error = self._error, None unlocked (LOAD_ATTR and STORE_ATTR are separate bytecodes), and the transfer thread check-then-sets unlocked too; this round adds the poll round as another writer and the fence as another reader. A GIL switch between the two bytecodes lets the forward thread overwrite a just-written exception with None. Deadline errors recur next round, but a submit exception that escaped the _UnprovableSubmission catch is lost for good. Do the swap under the _idle lock. (Plausible.)

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 a278334. raise_error swaps _error under the _idle lock, and the thread records a failure through _record_error under the same lock, keeping the first.

role = getattr(kv_transfer_config, "kv_role", None) or "kv_both"
extra = getattr(kv_transfer_config, "kv_connector_extra_config", None)
extra = dict(extra) if isinstance(extra, dict) else {}
collapse = extra.get("lmcache.mp.tp_rank_collapse", "auto")

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.

[E14] One rule for TP rank collapse. It currently has three contradictory answers: deployment.py:161-166 makes auto False for every state-owning layout (dsv41 included); this rewrites V4.1's auto to True; and connector.py:642-700's _page_tp_replication_factor has a V4.1-only branch whose comment claims a test binds the two sides, but no test in tests/ calls it. The README says auto is per-rank, while the recipe sets tp_rank_collapse: false without saying why. So native and plugin store the same model with different collapse factors, the recipe's published numbers use 4× tier bytes at TP4 (not the plugin's default path), and if either side's "not TP-sharded" claim drifts, other ranks silently read rank 0's bytes. Keep one rule — auto = the replication factor the backend publishes, resolved once in the worker — and remove the config-layer rewrite and the connector's V4.1 branch.

@@ -296,23 +423,39 @@ def _flush_after_forward(self) -> None:
# -- submission thread ---------------------------------------------------

def _start_transfer_thread(self, rank: int) -> None:

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.

[E15] Please remove the transitional design and duplication. (a) Plugin-only machinery — transfers_off_step, the poll thread, fence, getrefcount-style take_returned — lives in the shared core worker, giving three transfer modes; the native subclass keeps a second pending mechanism (_native_saves / _native_loads) that fence meets only with NotImplementedError; and the _adapter_records_its_own_events branch is unreachable with released LMCache (dead). (b) The plugin monkeypatches ATOM's own paged_state_checkpoint.py in four layers, with a docstring claiming "the vendored installer only hooks V1" — wrong, base _apply_runner_patches already walks V1 and V2. (c) v41_page_planes and _publish_v41_state_checkpoints copy backend.get_kv_transfer_tensors line by line (a consequence of clearing kv_transfer_config, E12). (d) V41StateMPScheduler is a second implementation parallel to NativeStateLMCacheMPConnectorScheduler and calls past its parent to the grandparent; the prefill_tail rule is written in 5 places; checkpoint-group discovery 4–5 times (all() vs any() disagree); _req_id_of is defined 3 times; and the _planners tuple is forced to at most one element, with 6 no-op hooks and a 303-line test of an unreachable configuration. (e) Pure forwarders _build_v41_pool_geometry and mp_worker_adapter.register_kv_caches (whose docstring claims it "does not extend the MP worker"); _RETIRED_STATE_OPTIONS refuses options that never existed on main; _save_source_is_backed's is not None is a dead check (vLLM's null block is 0). Each concept should have one mechanism, with plugin-only modes kept in the plugin.

Addresses valarLip's E13 and valar-cop's note on 45355d2.

- `raise_error` swapped `_error` out unlocked while the transfer thread
  checked and set it unlocked too, so a failure recorded between the two
  bytecodes of the swap could be overwritten with None and lost. Both
  sides now take the thread's `_idle` lock; the thread records through
  `_record_error`.
- The tests' `settle()` drains the queue before it waits for a round:
  an asked-for round runs after the submission in progress, not after
  everything queued, so the round alone no longer proves every queued
  submission went out.

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

Copy link
Copy Markdown
Contributor Author

Round 21 pushed a278334: E13 and the settle() note are fixed; CI is green, 11/11 accuracy cells pass. Please decide on merge next review.

E7, E11 and E15(a) ask to undo what your 11-03 and 11-04 reviews asked for or proposed; two are measured (details). They stay unless you withdraw those requests.

The plugin items (E1-E6, E8-E10, E12, E14, E15 b-e) are not changed in this round. Mark which block merge.

Evidence
Item What it asks What an earlier review asked, and the measurement
E7 release returned futures on an event, not stream idleness 11-04 item 4 proposed "future.done() plus the stream-idle check". Dropped mid-forward, c16 ITL p50 was 23.71-24.07 ms; with the idle check 22.84 / 22.92 ms, against 22.70-23.16 ms at 85d55d8. Under synchronous scheduling take_returned also runs in start_load_kv, after the previous step's sampled tokens reached the host.
E11, the wait don't wait up to 1 ms on decode steps 11-04 item 3: "If you keep a wait, show the ITL with and without it". Posted, interleaved: with the wait ITL p50 22.97 / 22.64 ms, without it 23.47 / 23.32 ms; TTFT p50 318 / 319 vs 327 / 325 ms.
E11, interleaving don't put poll rounds into the submit backlog 11-04 item 1: "force a poll round on the next loop iteration, whatever the queue holds".
E15(a), the mode move transfers_off_step out of the core worker 11-04 item 4: "give _TransferThread the poll round as a constructor option the front end passes when it builds the worker".
E15(a), the refusal and NativeState _adapter_records_its_own_events is dead; _native_saves / _native_loads is a second mechanism 11-03 item 2: "apply the same refusal in drain mode". 11-04 item 5: revert native_state_worker.py to main, whose pending maps these are.

Unit, at a278334 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.

@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 a278334, by reading the code and running the affected core tests (offload / MP / native / config / shell: 798 passed, 1 skipped; MP worker: 228 passed); no GPU. I found nothing in core that produces wrong output. This review covers only new items. It doesn't restate E1–E15, and it doesn't ask again for anything you implemented at my 11-03 / 11-04 request. Each item ends with the change I'm asking for.

1. A failed poll round keeps polling and logging (mp/transfer_thread.py:336-346). When _poll() raises (e.g. _enforce_transfer_deadline times out), _record_error keeps the first error and the thread carries on. Every _IDLE_POLL_S (50 ms) it runs another round, raises again, and logs logger.exception with a full traceback. Fail-stop happens only when the forward thread next calls _settle_transfer_thread / fence. If the front end stops stepping (no new requests), the engine doesn't stop and the log grows by about 20 tracebacks a second. I haven't triggered this.
→ Change: once an error is recorded, stop scheduling poll rounds (or exit the loop) and log it once; leave the raise to the forward thread as now.

2. Native pays for _queued_requests, which only the polling mode reads (mp/worker.py:479-505). Each _dispatch onto the thread now wraps the submit in _submit_queued: it takes _lock twice and updates _queued_requests. Only fence and get_finished's polls branch read it (:881, :890), and both are polling-mode only. With native OFFLOAD_MP_TRANSFER_THREAD=1, every transfer takes two extra lock acquisitions and contends with the forward thread's get_finished, to maintain state nobody reads.
→ Change: count queued requests only when transfer.polls, and submit directly otherwise.

3. Native get_finished was restructured (mp/worker.py:786-878, :157-160). It used to be one locked pass. It's now _poll_round → _poll_pending_loads / _poll_pending_saves, each snapshotting under the lock and committing under it again. Results go into four instance sets (_polled_done_save, _polled_completions, _polled_done_load, _polled_failed_load), which get_finished then copies and clears under the lock. So native goes from 1 lock acquisition per step to up to 5, plus copies of both pending dicts, and the four sets are written and read within the same call. The description says native keeps main's inline submit and poll.
→ Change: keep main's single-pass poll for native, and replace the four sets with one pending KVConnectorOutput that the polling mode fills and get_finished swaps out.

4. "Reopen the save floor when a load doesn't go out" is written three ways (chunked_scheduler.py:569, :854, :831):

  • update_state_after_alloc records _load_save_floors (new);
  • dispatch records it again at :854;
  • the tier-lost path ignores the floor and calls _reopen_saves(sid, hbm) directly at :831.

The comment at :566 gives "the tier lost the prefix" as a reason for the alloc-time record, which the code doesn't follow. In core the alloc-time record has no observable reader: its one consumer (:1413) is the raw-id load_failed, which core never issues. hybrid/dsv4/connector.py:2079 copies the same thing, where it's dead.
→ Change: record the floor once, have every load that doesn't go out (tier lost, withdrawn) go through one helper that reads it, and delete the DSV4 copy.

5. Measurements in code comments: mp/transfer_thread.py:25 ("18.7 ms per submit and 13.8 ms per terminal poll"), :56 ("ITL p50 23.1 -> 28.8 ms"), mp/worker.py:80 ("c16 TTFT p50 rose from 302 to 324 ms"), :85 ("c16 ITL p50 rose from 23.0 to 23.8 ms"), and :106, which repeats 18.7 / 13.8 ms. These date with ROCm / LMCache versions and are duplicated.
→ Change: reduce these comments to the constraint each one protects, and keep the numbers in the description.

6. Names that no longer match (mp/worker.py:261, :310, :595-601). In the base worker, _register_native_state_layout, the native_state.{i} registration keys, require_native_state_server and _native_layout / _native_block_ids now serve the path that restores no Active SLOT; the docstring itself says "This worker restores no Active SLOT". The base _check_state_destination(state, *, loading) begins with del loading, so two meanings share one signature.
→ Change: rename them for what the base does (e.g. checkpoint_image_layout), and express the native / base difference as data rather than an override whose base ignores its argument.

7. Duplicate helpers and checks (mp/worker.py:91, mp/transfer.py:135, mp/worker.py:905, mp/transfer_thread.py:212):

  • _req_id_of is identical to the one in the plugin's mp_worker_adapter.py.
  • _transfer_operation_id spells the same getattr(completion, "req_id", completion).
  • "only the polling mode can fence" is checked in both worker.fence (transfer is None) and _TransferThread.fence (self._poll is None), with different wording.

→ Change: one req_id_of helper in core that the others import, and one fence-precondition check.

8. Core reads a plugin-only attribute (offload/config.py:409). build_page_namespace now uses getattr(config, "offload_speculative_config", None) or .... No core code defines that attribute; only the plugin's offload_config.py sets it. (Its precedent, offload_page_hf_geometry, is set by atom/config.py itself.) The or also treats a falsy stable form (e.g. {}) as absent and falls back to hashing the raw object, so scheduler and worker could compute different namespaces. In the same vein, enforce_transfer_deadlines (mp/scheduler.py:340) became public only for the plugin's getattr call.
→ Change: pass the frontend's stable speculative form through an explicit config field (declared in core, like offload_page_hf_geometry), test it with is not None, and keep the deadline method private behind a declared interface.

9. Docstrings that overclaim (chunked_scheduler.py:677, :700, :601):

  • _return_leased_blocks says it is "the one sequence every lease release runs", and _retire_finished_save says it is "the one rule for every path that retires a finished request's save". But abandon_save (:1386) and reclaim_stale_leases (:1083) still pop _save_lease_blocks directly.
  • _release_lookup's comment at :601 ("clear_lookup_status only clears the client's memo; it does not release worker pins") is the opposite of what this PR relies on in mp/scheduler.py:286, where MP's clear_lookup_status frees the read locks.

Someone who trusts the first will add bookkeeping that the abort paths skip. Someone who trusts the second will add back the complete_retrieve whose lock leak this PR removed.
→ Change: route the two direct pops through _return_leased_blocks (or drop "the one"), and correct the :601 comment for the MP client.

10. README and docstring out of date (offload/README.md:34, config.py:188). The README still says "A GDN or K3 model cannot be overridden onto a layout that owns no state tier (dense)", but config.py:167 now refuses any mapped != natural for state-owning layouts, so qwen ↔ kimi_k3 is refused too. An operator who sets offload_layout: kimi_k3 on Qwen3-Next per the README gets a ValueError at startup. _layout_from_model's docstring still lists "dense/hybrid/kimi_k3/qwen".
→ Change: state the actual rule in the README and update the docstring's layout list.

zgplvyou and others added 2 commits November 5, 2026 00:34
Addresses valarLip's review of a278334.

- A failed poll round stops the polling thread's rounds and is logged
  once; the forward thread raises it, as before (1).
- A transfer thread that does not poll submits as main does: only the
  polling thread counts queued requests for its fence (2).
- Native `get_finished` is main's single locked pass again. The polling
  thread fills one pending `KVConnectorOutput` that `get_finished` swaps
  out, in place of four sets (3).
- The save floor a load reopens to is recorded once, when the load is
  decided. `_reopen_saves` reads it for a load the tier lost, now also
  in `_ensure_lookup_pin`'s refusal, which used to clear the floor first,
  and for a failed load. The alloc-time copies are gone; a promise
  withdrawn before dispatch reopens nothing, the tier still holds the
  span (4).
- Code comments state the constraint, not the measurement (5).
- Whether a checkpoint load restores into an Active SLOT is the class
  attribute `restores_into_slot`, not an override whose base ignored its
  argument (6).
- One `req_id_of` in `disaggregation/types.py`; the fence's polling
  precondition is checked once, in the worker (7).
- `Config.offload_speculative_config` is declared in core and tested
  with `is not None`; MP's deadline method is private again, called by
  the plugin's schedulers through `enforce_transfer_deadlines` (8).
- Docstrings that claimed one lease release path, and the in-process
  description of `clear_lookup_status`, are corrected; the README and
  `_layout_from_model` state the actual override rule (9, 10).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The plugin builds PluginDenseScheduler, whose enforce_transfer_deadlines
the connector now calls directly; these two tests had wired the core
DenseOffloadScheduler in its place.

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

Copy link
Copy Markdown
Contributor Author

Round 22 pushed d7619fe: all ten items are done except the renames in 6, which are refused. Your review found nothing in core that produces wrong output; please approve next review, or name what blocks merge.

  • 4: _ensure_lookup_pin's refusal cleared the floor before the tier-lost reopen; it now reopens first.
  • 6: native_state is main's name for the checkpoint-image format, not the frontend. restores_into_slot replaces the override.
  • 9: "the one" is dropped; main's two abort paths stay.
Evidence
Item Change
1 A failed round sets _poll_failed: no further rounds, one log line, raise_error unchanged.
2 _dispatch submits directly when the thread does not poll; only the polling thread counts _queued_requests.
3 Native get_finished is main's single locked pass; the polling thread fills one KVConnectorOutput that get_finished swaps out.
4 The floor is recorded once, when the load is decided; _reopen_saves(sid) reads it for a tier-lost load (in _ensure_lookup_pin and the dispatch loop) and in load_failed. The alloc-time copies are gone. A promise withdrawn before dispatch reopens nothing: the tier still holds what the request recomputes. The DSV4 test runs the real _ensure_lookup_pin and fails without the fix.
5 Comments state the constraint; the numbers are in the description.
6 Main names the format NativeStateTransfer, native_state_layout.py, build_native_state_mp_layout and stores_state_images, and native registers its planes as native_state.{i}; renaming the base alone would give one format two names.
7 One req_id_of in disaggregation/types.py, imported by the worker, transfer.py and the plugin; the fence's polling check is in the worker only.
8 Config.offload_speculative_config is declared in core, tested with is not None; _enforce_transfer_deadlines is private again, called through the plugin schedulers' enforce_transfer_deadlines.
9 abandon_save and reclaim_stale_leases are main's and drop whole leases; routing them through _return_leased_blocks would change native's abort paths for a docstring. The _release_lookup comment now describes both clients.
10 README and _layout_from_model state the rule config.py enforces.

New tests for 1, 2, 4 and 8 fail with their fix reverted. Unit, at d7619fe in the PR's container: the CI command and tests/plugin fail what they fail at 85d55d8.

@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 d7619fe, by reading the code, running the worker / scheduler / config tests (817 passed) and a shutdown probe; no GPU. Thanks for round 22. The nine fixes read right, and I accept your reason on the item-6 rename: main names the format native_state. Items 2, 3, 5 and 7 below follow from my last round's items 3, 1, 4 and 7, so they're partly on me.

On approval: I'd like items 1, 3 and 4 fixed before I approve. One is a reproduced bug, and the other two are ways the polling thread can stop without anyone noticing. The rest can go in a follow-up if you prefer. Each item ends with the change I'm asking for.

1. close() races the polling thread (mp/worker.py:338-343). close() swaps self._transfer to None first and then spends up to _TRANSFER_DRAIN_TIMEOUT_S draining and joining a thread that is still polling. A round that meets a terminal future in that window calls _retire_future, which calls hand_back on None. Reproduced with a probe that shuts down with a save in flight: AttributeError: 'NoneType' object has no attribute 'hand_back'. The consequences:

  • polling stops;
  • that completion is popped from _pending_* and never reported;
  • the future holding the imported IPC event is destroyed on the transfer thread, which the module docstring forbids;
  • every such shutdown logs a "polling stops" exception.

→ Change: clear self._transfer only after transfer.close(...) returns, or have _poll_round capture the transfer object it was started with. Add the shutdown-with-in-flight test.

2. The background poll copies the foreground poll (mp/worker.py:788-860 vs get_finished :936-1000). _poll_pending_loads / _poll_pending_saves repeat get_finished's body line for line: take_completed_ranges, _terminal_future_result, _enforce_transfer_deadline, the tombstone, _source_safe_completions, building ConnectorCompletion. That's about 70 lines twice. They differ only in lock scope, _retire_future, and where the result goes. A future change to terminal or deadline handling has to be made in both.
→ Change: one _poll_pending(into, retire) that both the foreground pass and the polling thread call.

3. A failed poll is reported once, then the thread stays dead silently (mp/transfer_thread.py:156-165, :342). This comes from my item 1. After a failure, _poll_failed stays set and no more rounds run. But raise_error swaps _error to None, so the forward thread raises exactly once. If that one exception is caught above, the engine keeps running with no poller: every later load sits in WAITING_FOR_REMOTE_KVS forever and nothing is reported again. The inline path, by contrast, re-raises the deadline error on every get_finished.
→ Change: in _settle_transfer_thread (mp/worker.py:506), raise on every step while _poll_failed is set (keep the first error, don't clear it), so a dead poller can't go unnoticed.

4. A stuck submit stops every deadline (mp/transfer_thread.py:305). Submissions and poll rounds share one thread, and a round runs only between submissions. If one submit RPC hangs (the LMCache server stalls), no round runs, so _enforce_transfer_deadline never fires for any in-flight transfer. The failure then surfaces only through the fence timeout or the scheduler's own deadline, later and with less information. Before this PR, polling stayed on the forward thread even with OFFLOAD_MP_TRANSFER_THREAD, so it wasn't exposed to this.
→ Change: don't make deadline enforcement depend on the submit thread getting a turn. Check deadlines from the forward thread's _settle_transfer_thread (comparing against submission time is cheap), or bound each submit with its own timeout.

5. The tier-lost cleanup runs twice (chunked_scheduler.py:825-829, _offload_common.py:851-852). This follows from your item-4 fix. The mixin's _ensure_lookup_pin already calls _reopen_saves(sid) and _clear_pending_load(sid) before returning False, and the caller then calls both again. On the base path the second pair is dead (the floor is already gone). hybrid/dsv4/connector.py:2319 repeats it. The caller's pair only matters for native_state_scheduler's override, which returns False without cleaning up.
→ Change: do the cleanup in one layer. Either _ensure_lookup_pin only answers and every caller cleans up, or it cleans up and the native override does too. Delete the other copy.

6. A stale module docstring (mp/native_state_layout.py:10-14). It says every worker and scheduler "asks the same two questions here: whether saves carry images, and whether the server's chunk is the configured one". The second check, require_native_state_server, now lives in mp/deployment.py:367.
→ Change: either describe only stores_state_images here, or move the chunk check back.

7. req_id_of replaced one copy of the pattern (Fix-then-sweep). This follows from my item 7. The new helper is used in the worker and transfer.py, but the same hand-written unwrap is still at:

  • mp/scheduler.py:269 and :284, right beside the changed load_failed;
  • _offload_common.py:519, :524, :543 and :546;
  • hybrid/kimi_k3/connector.py:872.

→ Change: use req_id_of at every one of them.

8. The new override rule refuses combinations that used to start (offload/config.py:167). mapped != natural for state-owning layouts now refuses offload_layout: qwen on a Kimi-linear model and kimi_k3 on a GDN model, besides the V4.1 cases. The README now states the rule, but the description mentions only V4.1, and no test covers the newly refused pairs (test_lmcache_offload_config:231 tests dense only).
→ Change: confirm in the description that the kimi_k3 ↔ qwen refusal is intended, and add a test for that pair.

9. Comments still carry version findings (rest of my item 5). mp/worker.py:83 ("on ROCm 7.2.x destroying one synchronises every stream") and mp/transfer_thread.py:55-58 ("a fixed short cadence with loads outstanding cost more end to end than the forward-thread work it moved off") are findings, not constraints.
→ Change: state the constraint (e.g. "destroy imported events only on the forward thread at a step boundary"), and keep the version and the comparison in the description.

Addresses valarLip's review of d7619fe.

- A polling round gets the thread it runs on and hands finished futures
  back to it; `close` forgets the thread only once it is closed. Shutdown
  with a save in flight no longer meets `hand_back` on None (1).
- One poll pass, two callers: `_poll_entries` polls without touching
  worker state and `_settle_terminal_locked` reports what finished. Native
  `get_finished` runs both under one lock, as main's single pass; the
  polling thread polls unlocked and settles under the lock (2).
- A failed poll round is raised on every `raise_error`, not once: no round
  runs after it. The failure is recorded and marked in one locked step
  (3).
- The forward thread fail-stops a transfer past its deadline when no
  polling round has run since: a submit that never returns stops the
  rounds that enforce deadlines (4).
- `_ensure_lookup_pin` only answers; the caller drops the load (5).
- `native_state_layout`'s docstring names only `stores_state_images` (6).
- `raw_req_id` unwraps a completion id in the type the engine compares,
  and `req_id_of` is its string form; every hand-written unwrap in core
  uses one of them. The sets returned to the engine keep raw ids:
  stringifying them fails test_offload_completion_processing_calls_load_
  terminal_once (7).
- A test covers the refused `kimi_k3` / `qwen` swap (8).
- Comments state the constraint, not the finding (9).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@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

This branch has not been deployed

No deployments
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