[Kimi-K3][LMCache] LMCache offload for Kimi-K3 by zejunchen-zejun · Pull Request #2053 · ROCm/ATOM

zejunchen-zejun · 2026-08-27T08:20:47Z

[Kimi-K3][LMCache] LMCache offload for Kimi-K3

Adds a CPU offload tier for K3's KDA recurrent state, alongside the KV tier
lmcache_offload already provides. A prefix whose KV came back from LMCache is
useless without the recurrent state that goes with it, so the two are fetched
together or not at all.

Builds on #2045, which turned a K3 checkpoint from a whole Active Slot into a
PAGE image living in ordinary KV blocks. That is what makes this tier cheap:
the source of a save is a set of reserved blocks that nobody can take away,
so nothing has to be rescued anywhere before the D2H.

Design

A third layout, inheriting dense rather than written beside it.
K3's paged KV is ordinary dense MLA — block-addressed, sliceable, moved by
chunk — so KimiK3OffloadConnector / KimiK3OffloadScheduler extend
DenseOffloadConnector / DenseOffloadScheduler and add one extra leg for the
69 layers of KDA state. (DSV4 is written in parallel instead because its PAGE
leg carries an indexer and has a different shape.) Every override wraps a dense
call rather than replacing it; the KV leg's behaviour is unchanged.

Layout selection is config-only (select_offload_layout), because the
scheduler process has no model.

Both legs save when the bytes land in HBM, into one pool.

	KV	state
trigger	a prefill chunk completes	a checkpoint reaches `READY`
source	the request's blocks, held by `should_defer_free`	the checkpoint's PAGE units, held by a pin
pool	one LMCache engine, one allocator, one LRU	same pool, same LRU

No second engine and no OFFLOAD_STATE_CPU_SIZE. A request writes its KV
chunks and its one state object inside the same prefill window, so both enter
the LRU together and cool at the same rate — state is retired alongside its own
KV, which is what we want, since a boundary whose KV is gone is worthless. Two
pools would buy two eviction policies that must drift, while a joint boundary
needs both legs to survive together.

Loads pick one boundary for both legs. The ceiling comes from KV, the
landing point from state; the KV leg fetches into blocks the request already
holds, and the state leg's H2D writes the request's Active Slot directly — no
units are reserved for it.

Everything is asynchronous. submit() returns immediately; the Triton
gather runs on its own stream, the D2H on another, and the engine advances on
the worker's report. Nothing waits on the compute stream.

Why the boundary has to be joint

KDA state at position B is the compressed history of exactly [0, B). Raise
the KV-loaded length past B and the linear layers never see the tokens in
between: wrong output, no exception. So the two legs cannot be decided
separately.

Example

A request whose prompt is 12 hash blocks, A B C D E F G H I J K L
(--block-size 128, LMCACHE_CHUNK_SIZE 256). Block C has been evicted from
HBM; LMCache still holds the KV for A..K and the state checkpoint filed at
H.

HBM   KV:  A B  ✗C evicted✗
HBM ckpt:  (gone with C's neighbourhood)
CPU   KV:  A B C D E F G H I J K
CPU ckpt:                  H

can_allocate walks the HBM prefix and stops at the first miss, C.
Two blocks matched.
No checkpoint at or below B, so the request may claim 0 tokens as
cached.
_joint_kv_boundary:
- ceiling from KV — LMCache covers 11 blocks, floored to the chunk grid
  gives 10; capped by n_hash_blocks - 1 (prefill must forward one block to
  produce logits) → 10;
- the chained hash is a function of the prompt alone, so the chain is
  continued past where HBM stopped, which is what makes a boundary that
  exists only in LMCache addressable at all;
- the rightmost reachable checkpoint on that chain is H → B = 1024;
- KV leg target = the chunk covering B = 1024;
- block table may claim what the walk matched, floored to the grid = 256.
allocate points the block table at A B — free, since those blocks sit
below where the forward will start — and leaves num_cached_tokens at 0.
State: HBM misses, the tier votes, a load is issued for H.
KV: transfer [256, 1024) = C D E F G H. The already-resident A B is
not re-fetched; doing so would move bytes the GPU holds and leave an
unindexed second copy of them in HBM.
_claim_after_load → 1024, the state boundary, not the transfer end.
Both legs report; the request unparks with num_cached_tokens = 1024 and
the forward computes I J K L.

Three positions, all different, and each is needed: 0 is what may be called
cached, 256 is what the block table may point at, 1024 is where both
legs are aimed. Claiming past the first is safe only because a joint boundary
exists — those blocks are below the forward's start, so nobody writes them.
Transferring past B costs one chunk the forward immediately rewrites;
claiming past B is silent wrong output.

Notes for reviewers

state_pool.py and sub_pool_spec.py have an empty diff against main.
The tier does not reach into the slot pool at all.
The state key folds in layout_id. CacheEngineKey has no field saying what
an entry is, and KV and state keys come from different hash functions into one
integer space — now sharing one pool. Folding the layout in makes the two
spaces disjoint by construction, and separately stops a build that changed the
state geometry from reading another's image back as valid.
A checkpoint is nominated at READY but pinned only when handed to the
worker, bounded by OFFLOAD_MAX_PENDING_SAVES. Pinning at READY would make
every checkpoint un-evictable for a window that spans several scheduler
passes, breaking the admission rule that a READY unpinned checkpoint counts
as space available to live KV.
Both the KV and the state save gather through the same Triton kernel
(fused_pack_chunk_major), unmodified.
Removed env vars: OFFLOAD_STATE, OFFLOAD_KV_FOR_HYBRID,
OFFLOAD_SAVE_MAX_INFLIGHT, OFFLOAD_SAVE_STALL_S,
OFFLOAD_STATE_MIN_LOAD_TOKENS, OFFLOAD_STATE_CPU_SIZE,
OFFLOAD_STATE_STAGING_GROUPS. A resumable prefix needs both legs, so the
connector being configured is the only switch.

Verification

CON16

CON32

github-actions · 2026-08-27T08:21:37Z

🏷️ 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 2053 --add-label <label>

CI's reviewdog run on #2053 flagged five findings, all introduced by this branch. Every one is real and still present. atom/kv_transfer/offload/hybrid/dsv4/connector.py:27 I001 atom/kv_transfer/offload/hybrid/kimi_k3/connector.py:11 I001 atom/kv_transfer/offload/hybrid/kimi_k3/connector.py:425 SIM102 tests/test_scheduler.py:1466 I001 tests/test_scheduler.py:1498 I001 plus one more the bot has not reported yet because reviewdog filters to diff context and that line has not changed since: atom/model_engine/page_unit_checkpoint.py:6 I001 --- why the local checks said clean `ruff check` with no `--select` runs the default rule set (E4/E7/E9/F), and neither `I` (isort) nor `SIM` is in it. There is no `[tool.ruff]` in `pyproject.toml` and no `ruff.toml`, so every local run in this branch's history measured the default set and reported "36, unchanged" against the base. That number was true and I reported it as if it meant CI would be clean, which it does not. Checked from now on with `--select I,SIM` as well. --- scope `--fix` over the whole tree repairs 100 findings in files this branch never touched; that was reverted. Fixed here are only the findings this branch introduced, established per file by running the same rules against the same file on the base: for f in $(git diff --name-only <base> HEAD -- '*.py'); do compare ruff --select I,SIM on <base>:$f against $f done Six files come back with more findings than their base; the six above are them. After this commit that comparison returns nothing. The 14 findings that remain across files this branch touches are all pre-existing on untouched lines, which is also why reviewdog does not report them. SIM102 is `should_defer_free`'s stalled-save escape: two nested `if`s collapsed into one condition, same short-circuit order, no behaviour change. Verified: pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors) -> 3936 passed, zero failures. black clean; ruff default set 36, unchanged; ruff `--select I,SIM` shows zero findings introduced by this branch. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Copilot

Pull request overview

Adds Kimi-K3 support to the lmcache_offload connector by introducing a CPU offload tier for K3’s per-request recurrent (KDA) state and making KV+state resumes joint (both legs must agree on a single safe boundary). This extends the existing dense offload path with K3-specific state spill/load plumbing, new engine-side bookkeeping, and substantial test coverage.

Changes:

Introduce kimi_k3 offload layout: worker-side state tier (store/load), scheduler-side joint-boundary policy, and shared-pool keying with build-safety (layout_id) separation.
Add engine-side StateOffloadIndex + BlockManager/Scheduler wiring to enqueue state loads/stores, settle reports, and expose new cache stats counters.
Register per-request state tensors via KVCacheTensor.per_request_state and exclude them from block-addressed KV movers; add backend hooks (state_entry_views, page_unit_views) to name state bytes for packing.

Reviewed changes

Copilot reviewed 39 out of 39 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/test_state_checkpoint.py	Adds extensive tests for tier load behavior and joint boundary semantics.
tests/test_scheduler.py	Tests offload-resume accounting and stalled deferred-save reclamation behavior.
tests/test_page_unit_checkpoint.py	Tests PAGE checkpoint nomination/pinning/offload vote behavior.
tests/test_multi_connector.py	Ensures multi-connector routes state loads and exposes `_state_tier`.
tests/test_lmcache_offload_v4_page_slot.py	Updates tests to use shared `max_pending_saves` helper.
tests/test_lmcache_offload_connector.py	Adds K3-specific scheduler contract tests (joint boundary, state channels, shell forwarding).
tests/test_kda_checkpoint_slot_copy.py	Adds coverage for `page_unit_views` correctness and image trimming.
tests/test_block_manager.py	Tests widened claim behavior + state-tier index installation gating.
docs/environment_variables.md	Documents offload connector behavior and layout selection; introduces a new section.
atom/plugin/sglang/attention_backend/attention_gdn.py	Marks recurrent per-request state tensors as `per_request_state=True`.
atom/plugin/rtpllm/utils/forward_context.py	Same `per_request_state` marking for RTP plugin path.
atom/model_ops/attentions/kimi_mla_gdn_attn.py	Implements `page_unit_views` + marks per-request state tensors.
atom/model_ops/attentions/gdn_attn.py	Adds `state_entry_views` and marks recurrent state tensors as per-request state.
atom/model_ops/attentions/deepseek_v4_attn.py	Adds `state_entry_views` and clarifies state-copy docstring/comments.
atom/model_ops/attentions/backends.py	Defines new `state_entry_views` abstract hook for backends owning per-request state.
atom/model_engine/state_offload.py	New engine-side state offload index + connector-host gating helper.
atom/model_engine/sequence.py	Adds fields to track state load intent and joint-boundary extents.
atom/model_engine/page_unit_checkpoint.py	Adds offload nomination/pinning lifecycle + vote through CPU-tier reachability.
atom/model_engine/model_runner.py	Publishes attention backend to transfer tensors so tier can access `state_entry_views`.
atom/model_engine/llm_engine.py	Exposes new joint/tier counters in cache statistics.
atom/model_engine/engine_core.py	Adds worker termination hardening + periodic stalled-save reconciliation + dispatch of state load/store publishing.
atom/model_engine/block_manager.py	Implements joint boundary selection + widened claim + state-tier load/store APIs and funnel aggregation changes.
atom/kv_transfer/offload/metadata.py	Adds `state_loads` / `state_stores` to offload metadata and work-field declaration.
atom/kv_transfer/offload/hybrid/kimi_k3/state_tier.py	New worker-side tier executor for state store/load plus joint-park helper.
atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py	New codec for state objects (stable keying + packing/unpacking).
atom/kv_transfer/offload/hybrid/kimi_k3/staging.py	New bounded staging helper for whole-entry pack/unpack.
atom/kv_transfer/offload/hybrid/kimi_k3/connector.py	New K3-specific connector/scheduler implementation (state tier + joint KV clamp + stall guards).
atom/kv_transfer/offload/hybrid/kimi_k3/init.py	New package marker for K3 offload layout.
atom/kv_transfer/offload/hybrid/dsv4/connector.py	Uses shared `max_pending_saves` + routes claim computation through `_claim_after_load`.
atom/kv_transfer/offload/dense/kv_byte_codec.py	Skips `per_request_state` tensors in dense KV packing.
atom/kv_transfer/offload/dense/connector.py	Adds `_may_emit_save` seam + uses `_claim_after_load` for post-load claim.
atom/kv_transfer/offload/connector.py	Adds `kimi_k3` variant selection and strengthens shell docs/forwarding.
atom/kv_transfer/offload/config.py	Adds `kimi_k3` alias and selection rule based on HF `model_type`.
atom/kv_transfer/offload/atom_lmcache_staging.py	Makes `_env_flag` robust to whitespace/empty values.
atom/kv_transfer/offload/_offload_common.py	Adds `_claim_after_load` seam + factors `max_pending_saves` helper.
atom/kv_transfer/disaggregation/types.py	Adds `ConnectorMetadata.WORK_FIELDS` + `has_work()` and uses it in work detection.
atom/kv_transfer/disaggregation/multi/multi_connector.py	Delegates `has_work()` to subs, re-exports `_state_tier`, and routes state loads.
atom/config.py	Adds `KVCacheTensor.per_request_state` flag to separate per-request state from block-addressed KV.

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

+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self._state_tier = None
+        # Inert until a request has both legs, which only a joint boundary
+        # produces; costs one dict lookup per report otherwise.
+        self._joint_park = _JointPark()
+


+    def get_finished(self):
+        out = super().get_finished()
+        if self._state_tier is None:
+            return out
+        indexed, index_failed = self._state_tier.take_store_reports()


+| Variable | Type | Default | Description |
+|----------|------|---------|-------------|


Split out of PR #1960 so they can land on their own. None of them is about state offload, and none needs #2045; blocking them behind either is pure loss, and pulling them out narrows what is left to review by a few hundred lines. --- a worker that outlives the join keeps the node `EngineCore.exit` joined each worker with a 5s timeout and did nothing after it. A worker that outlives that join keeps its VRAM slice and its all-reduce IPC handles, and `multiprocessing`'s atexit handler then joins the same process again with NO timeout -- so an engine already on its way out parks in `_exit_function -> join` forever while /metrics keeps answering 200. Seen twice on the k3-dev line; only a manual kill recovered the node. `enable_orphan_reaping` cannot help: `PR_SET_PDEATHSIG` fires when the parent *dies*, and this parent never does. Now escalates terminate -> kill. --- a stalled offload save hung the engine with every GPU idle A deferred save's blocks are freed only on `finished_saving`, but LMCache's pin monitor force-unpins a stalled transfer after `pin_timeout_sec` WITHOUT emitting that report. `has_pending_kv_work()` then never clears and the engine busy-loops with every GPU idle -- a hard hang under a tight pool, a slow block leak under a large one. `Scheduler._reconcile_stalled_deferred_saves`, wired into the KV progress poll, reclaims a save deferred past an abandon window and self-throttles to 5s so a 1ms poll costs nothing. The window reads LMCache's own `LMCACHE_EC_PIN_TIMEOUT_SEC` and adds 30s rather than taking an ATOM knob of its own, because that ordering IS the safety argument: two independent env vars would let ours sit below the timeout it must exceed, and nothing would say so. Unset gives 330s. Non-positive disables reclamation and restores the wait-forever behaviour. What makes reclaiming safe is not the clock -- our clock starts when the free is deferred, LMCache's when it pins, and nothing aligns them. It is that `OffloadWorkerMixin._guard` reports `finished_saving` on both the success and the exception path, so a report is lost only when `store()` neither returns nor raises. Two cases follow and they are exhaustive: the parked save is not copying (LMCache force-unpinned its source and stopped reading), and a save queued behind a parked one never reached `store()` at all. `finished_sending`'s `assert seq is not None` becomes a `continue`: after a reclaim, a late completion report has nothing left to free. --- a state-only step under `multi` parked its request forever `connector_metadata_has_work` decides whether a step's metadata reaches the worker, and it decided from a hardcoded field list shared by three connector families. `MultiConnectorMetadata` holds none of the work itself -- everything is in `metas` -- so it mirrored the subs' fields as aggregating properties, and missed one. Every step whose only work was in an unmirrored field was discarded, and the request parked on it waited for a report nobody was asked to produce. Fixed by asking the metadata instead of inspecting it: each class declares its own `WORK_FIELDS` next to the fields it defines, `has_work()` accumulates down the hierarchy, and `MultiConnectorMetadata` delegates to its subs rather than mirroring them -- so a field added to a sub can no longer be missed here. The old union list stays only as the fallback for duck-typed doubles, marked as not the place to register anything new. --- two behaviour-neutral seams and one marker `OffloadSchedulerMixin._claim_after_load` returns `max(hbm, lmc)`, which is what all four inlined call sites did. A seam rather than four `max`es because a layout whose claim and transfer aim at different boundaries has to override it, and getting that wrong is silent: the request starts further along than its state supports and the output is wrong with no exception. `DenseOffloadScheduler._may_emit_save` returns True. It matters for a layout whose `should_defer_free` pins a finished request's blocks until its save drains -- an unbounded queue then lets a slow backend hold an unbounded slice of the pool. `max_pending_saves` moves to `_offload_common` so every layout reads the same `OFFLOAD_MAX_PENDING_SAVES`. `KVCacheTensor.per_request_state` marks a hybrid's slot-addressed recurrent state. It must be registered in `kv_cache_data` -- the linear-attention forward reads it from there -- but no block-addressed mover may touch it, and `DenseKVByteCodec` now skips it: including it either fails the divisibility check or, if the slot count happens to divide `num_blocks`, inflates `bytes_per_block` past what `block_regions` describes. Verified on top of #2045 (`4df58434`): pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors that fail on #2045 too) -> 3871 passed, 326 skipped, 3 xfailed, zero failures. black clean; ruff 36, the same count as #2045. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Phase 1 of remounting PR #1960 on #2045. Deliberately a RENAME, not a redesign: the tier still hangs off `StateSlotPool` and still spills on eviction through the staging ring. Phase 2 deletes the ring; Phase 3 moves the mount point to `PageUnitCheckpointStore`. Keeping those apart is what lets a reviewer tell rebase noise from design decisions. Squashed rather than replayed for the reason #2045 gives for its own squash: other's resolutions. #1960 keeps the original history. --- what #2045 renamed underneath the branch StateGroupPool -> StateSlotPool lookup_group(h) -> lookup(h) seq.per_req_cache_group -> seq.state_slots (list) + seq.state_slot num_per_req_cache_groups -> num_state_slots + state_slots_per_req pop() / release() -> pop_many(width) / release_many() has_free() -> has_free(n) _attach_state_group() -> _attach_state_slots() The semantic change hides in that table: **the "group" concept is gone.** Slots are allocated one at a time, `pop_many` states outright that the set is not adjacent and nothing may assume it is, and a checkpoint occupies exactly ONE slot -- speculation scratch is not state anybody resumes into. Every `group * span` address computation in #1960 lost its basis, so the tier's `_spill` / `_resumable_from` / `has_pending_spill` / `take_spill_copies` are re-expressed over slots, and `(group, staging)` becomes `(slot, staging)`. --- the three places the two PRs actually collide `_attach_state_slots` now returns bool. #2045 rewrote its body for per-slot allocation; #1960 needs it to be able to say "the state behind this boundary is not really here", because with the tier voting a hash can be accepted by `_resumable_from` and still not be in HBM. The two compose: #2045's width logic runs, and only the committed slot (element 0) is the load's destination, so only it takes `pop(spill=False)` -- the rollback scratch is this request's own and has nothing to forgo. `allocate` keeps #1960's two extents (`num_cached_blocks` vs `claim_blocks`) on top of #2045's per-slot attach, and the fresh-block loop starts from `len(seq.block_table)` rather than `num_cached_blocks`, since the widened claim already put blocks there. `CacheStats` gains #1960's `total_offload_tokens` alongside #2045's `total_reusable_tokens`. They are not the same series and neither subsumes the other: `reusable` is the honest denominator for the HBM walk, while tier tokens sit ABOVE that walk by design, so folding them in would make `cached > wanted`. `update()` therefore takes five required arguments and one defaulted. `deallocate` keeps #2045's midstep cancellation and `release_many`, and #1960's in-flight-load hold: the committed slot is parked in `_orphan_load_slots` until `settle_state_load`, because a worker is writing it; the scratch goes back immediately. `checkpoint_funnel` reports through `state_checkpoint_fates()` (a sum across `state_caches`) rather than the single `_state_checkpoint_cache` call. Today `state_caches` is a 1-tuple and the two are equivalent -- kept so Phase 3's tier counters, which arrive from a second class, are not silently dropped. --- functionality is NOT verified by this commit Three gates all take the else branch for K3 under #2045, so the tier is inert: 1. `isinstance(cache, StateSlotPool)` -- K3's `state_caches` is `(PagedStateCheckpointCoordinator,)`, so `cache.offload` is never set. 2. `_joint_kv_boundary` returns `_no_joint("not_hybrid")` whenever `paged_state_checkpoints is not None`, which for K3 is always. 3. `_attach_state_slots` returns from its PAGE branch before reaching `_tier_can_serve`. That is deliberate and is what Phase 3/5 undo. Do not read the green suite below as evidence the feature works -- the exit criterion for the tier is a nonzero `state_tier` counter in the logs, which no unit test can produce. Verified: pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors that fail on #2045 too) -> 3915 passed, 326 skipped, 3 xfailed, zero failures. black clean; ruff 36, the same count as #2045. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…to rescue Phase 2. About 1000 lines go, and none of them for tidiness: the ring existed to solve exactly one problem, and #2045 removed the problem. --- why the ring existed at all A checkpoint used to live in a state slot, and `StateSlotPool.pop` handed that slot to the next request the instant it was evicted. The bytes therefore had to be rescued somewhere before they were overwritten, which is what the ring was: K entries allocated inside the state arena, never leased to a request, and addressed as `num_slots + staging` so `state_entry_views` could reach them with no second scheme. Everything else -- `staging_entries`, `admission_entries`, the drain protocol, the starvation detector, the reclaimer, `pop(spill=)`, `spills_forgone`, the spill-before-relocate ordering in `CommonAttentionBuilder.build` -- was a consequence. Under #2045 a K3 checkpoint is a PAGE image in the KV pool whose 127 units are *reserved*, and reserved units are structurally un-takeable: `_take_free` draws only from `_free` and they sit in `_used`, `BlockPool.free` raises on one, and `retire_top` refuses outright (`if top in self._raw_unit_owner: return None`). The source holds still on its own, so a transfer needs no rescue -- exactly the property dense KV has always had, where `should_defer_free` alone carries a save through its D2H and no ring was ever needed. So this is not a simplification, it is the removal of a workaround whose premise is gone. `state_pool.py` and `sub_pool_spec.py` now have an EMPTY diff against --- what stays, and why `state_entry_views` stays (re-indexed from group to slot, since a checkpoint is one slot wide now). Phase 4's load writes the committed slot directly, so this is its destination -- the plan's earlier draft had it deleted, which was wrong. `StagedTransfer` / `_StagingBuffer` stays. It shares a word with the ring and is a different thing: a per-thread pinned host bounce buffer that D2H and H2D both need, sized in bytes and unrelated to slot count. `StateOffloadIndex` keeps its load half (`request_load` / `complete_load` / `fail_load` / `abandon_load`) and its optimistic hash set. `confirm_spill` becomes `note_stored`, which is what it always did. `StateOffloadTier.submit_spill` becomes `submit_store(h, slot)` and loses its `ready_event`: there is no compute-stream producer to fence against once the packer gathers straight from a source that is not moving. --- what this leaves inert, deliberately **Nothing calls `submit_store` after this commit.** Phase 2 removed the ring that used to drive it and Phase 3 gives it its new caller, whose source is the checkpoint's PAGE units rather than a slot. Nothing populates `StateOffloadIndex.hashes` in between, so `_tier_can_serve` is false and no load is ever offered -- the tier is fully inert, on top of the three gates Phase 1 already documented. That is a WIP state, not a shippable one. --- one dead-code find, pre-existing on #1960 A method between `_tier_can_serve` and `_request_state_load` had lost its `def` line in an earlier rebase, leaving its docstring and `return` as unreachable statements after `_tier_can_serve`'s own return. It reads like a live invariant ("Whether `hit_hash`'s state can be produced at all, from either tier") and has no caller: the tier's vote actually reached `_gated_hit` through `StateSlotPool._resumable_from`. Removed rather than restored, since Phase 4 puts the vote on the coordinator side. Verified: pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors that fail on #2045 too) -> 3906 passed, 326 skipped, 3 xfailed, zero failures. black clean; ruff 36, the same count as #2045. `git diff` against Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Phase 3a. Deletes the tier's separate `atom-state-{rank}` LMCache engine, its `OFFLOAD_STATE_CPU_SIZE` knob, and its `cache_policy = "FIFO"` override. The codec now binds the same `storage_manager` the KV leg uses. --- what changed - `_state_storage_manager()` is gone (~50 lines). `bind_storage_manager` is handed `self._engine.storage_manager` directly. The `gib <= 0` branch that already did this was the whole implementation; what goes is the `> 0` half. - `OFFLOAD_STATE_CPU_SIZE` removed, from the code and from `docs/environment_variables.md`. - `cache_policy = "FIFO"` removed; state inherits the KV pool's LRU. --- why the two-pool argument does not hold It rested on "the KV write stream is several times the state volume, so it evicts the checkpoints". Measured against the code rather than assumed: `page_unit_bytes = num_MLA_layers * block_size * entry * itemsize = 442,368` with `--block-size 128` pinned by the recipe, so KV costs 3,456 B/token and a 53.6 MiB state image is worth ~16k tokens of it -- 13% of a 117k-token prompt's KV. (Read the MLA layer count the other way, as a quarter of 93 layers, and it is 3.3%. The conclusion is the same either way, which is why it did not need settling first.) And the byte ratio is not what decides an LRU eviction anyway. A request writes its ~457 KV chunks and its one state object inside the same prefill window, so both enter the LRU at the same position and cool at the same rate. State is retired ALONGSIDE its own KV -- which is the behaviour we want, because a boundary whose KV is gone is worthless. Two pools cost what one does not: two eviction policies that must drift, while a joint boundary needs both legs to survive together. #2045 made this same argument one tier up -- state moved out of its reserved slots into the KV pool, 81.85% -> 93.60% -- and `OFFLOAD_STATE_CPU_SIZE` is the same reserved capacity one tier down, so keeping it repeats the thing #2045 disproved. `LMCACHE_MAX_LOCAL_CPU_SIZE` is now the one size to tune. --- the cost, stated rather than hidden `cache_policy` is per-engine, so one pool means one policy for both legs. That is a real loss of freedom and it is currently free, because LRU is right for both. The FIFO override's premise was "a state entry is written once and read once"; #2045's own numbers refute it (~4,808 resumes over ~1,508 checkpoints is ~3 reads each), and re-access is bursty -- branches and retries on one boundary, then superseded -- which is LRU's home ground. The genuine want, if it shows up, is "stickier than LRU gives it", which FIFO does not provide either; the lever then is touching the state key on a KV hit, not splitting the pool back apart. Nothing yet measures whether that is needed. The counter to add with the store path (Phase 3d) is `joint_boundary_state_miss`: a boundary whose KV leg landed and whose state leg missed. Zero means the two pools never needed separating. Verified: pytest tests/test_lmcache_offload_connector.py -> 186 passed. black clean; ruff unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Phase 3b. `StateByteCodec.key()` stops being the bare ATOM hash; it becomes `xxh64(hash || layout_id)`. `get()` gains a size check on what comes back. --- what changed - `StateByteCodec.__init__` takes a required `layout_id` and refuses an empty one. `_build_state_tier` reads it off `state_runtime.checkpoint_spec` and declines to build the tier at all when it is absent, rather than falling back to an unqualified key. - `key(h)` hashes `h` and the layout id together with `xxh64`. - `get(h, ...)` measures the returned `MemoryObj` and, on a size other than `entry_bytes`, discharges its reference, counts `_misfit_reads`, warns, and returns a miss. - `_object_bytes()` reads `get_size()` with the tensor as fallback, and returns None when neither answers. --- why the key needed both *Build safety.* The same prefix hash names a completely different image under a different `num_spec`, TP size, or conv/ssm dtype. A size mismatch throws at `entry_bytes`, but **the same size with a different order or meaning reads back silently wrong state** -- and a request resumed onto wrong state produces wrong output with nothing raised anywhere. #2045 already encodes all of it in `layout_id` (layers / conv shape+dtype / ssm shape+dtype / order / tp / spec / carry) and enforces it HBM-side in `_validate_paged_state_op`. This is the CPU side of the same check, and it was missing. *Namespace separation.* `CacheEngineKey` is `(model_name, world_size, worker_id, chunk_hash, dtype)` -- no field says what an entry IS. KV chunk keys carry `ChunkedTokenDatabase`'s chunk hash and state keys carry ATOM's chained block hash: two different hash functions writing plain integers into one space. Until Phase 3a they at least lived in different pools; now they share one, and the only thing keeping them apart was accidental -- this side hard-codes `torch.uint8` while KV carries the KV dtype. `MemoryFormat` is not available as a discriminator either: the tier's own comment records that its `KV_2LTD` is inert, chosen only to pass the allocator's `raise ValueError`. Folding the layout in makes the two spaces disjoint by construction, since no KV key can carry a K3 layout id. `xxh64` rather than Python's `hash((h, layout_id))`: `hash` of a str is salted per process, so a restart would silently orphan every entry the previous run wrote -- a cache that quietly starts cold and never says so. `xxh64` is what the block hashes themselves use. --- why the size check, if the key already prevents it It is unreachable, and that is the point: reaching it means two things collided in the shared pool, and unpacking a colliding object writes another entry's bytes over a request's live state. Degrading to a miss costs one recompute; not checking costs silent wrong output. `_misfit_reads` is the counter that says the key stopped doing its job. None from `_object_bytes` means "cannot measure", never "size 0" -- refusing what cannot be measured would turn an unknown into a guaranteed miss. --- tests Five, in `tests/test_lmcache_offload_connector.py`, on a `CacheEngineKey` fixture rather than the full lmcache stub (these are about two keys being equal or not, which the real dataclass answers by value): a layout change makes a different key; the layout does not collapse distinct hashes; the key survives a restart; a state key cannot equal a raw-hash KV key; a wrong-sized hit is a miss with its reference discharged; an object that will not report a size is still loaded. Verified: pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors that fail on #2045 too) -> 3911 passed, zero failures. black clean; ruff 36, unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Phase 3c. Adds `_KimiMLAGDNCommon.page_unit_views(unit_ids)` -- the tensor-view counterpart of #2045's `_page_unit_regions`. No behaviour changes; nothing calls it until Phase 3d. --- what changed - `page_unit_views(unit_ids) -> list[torch.Tensor]`, one view per (unit, MLA row), unit major and row minor. - Seven tests in `tests/test_kda_checkpoint_slot_copy.py`, on the harness #2045 already built for `_page_unit_bases`. --- why a second way to name the same bytes `_page_unit_regions` hands the Triton descriptor raw int64 addresses, which is what a D2D copy the GPU issues wants. The LMCache staging packer takes `list[torch.Tensor]` -- `fused_pack_chunk_major` is already a fully parameterised gather over `segment_ptrs[] + segment_block_bytes[]`, so it needs no kernel change, only its segments named the way it names them. This is the ONLY new seam the tier needs. A load writes the Active Slot directly and reuses `state_entry_views`, so the asymmetry is deliberate: the save gathers from units, the load scatters to a slot, and the two compose because #2045's `plan_segmented_copy` intersects two *ordered byte streams* -- the slot's ranges and the units' regions are two cuts of the same logical image, so a blob gathered in one order is byte-identical to one gathered in the other. --- the two traps, both asserted rather than assumed *Order is the contract.* Unit major, row minor -- the same ravel as `_page_unit_bases`, which is the order `_checkpoint_copy_plan` builds the destination stream in. The packer indexes segments positionally, so this order IS the blob layout. Self-consistency would be enough for one build; what makes it a contract across builds is `layout_id` in the key (Phase 3b), so a build that reordered cannot read another's blob. *`block_ratio`.* `unit_ids` carries **logical** block ids while `kv_cache` is shaped in **physical** ones, and K3's ratio is 128. Flattening the two block axes together and indexing by `unit * block_size` is the conversion -- the same arithmetic `_page_unit_bases` does in addresses, so the two cannot drift into disagreeing about which bytes a unit owns. The range check is against the *logical* count for the same reason: the physical count is `ratio` times larger in rows, so checking the wrong axis admits ids that read past the end. Its test computes the bound from the tensor rather than from a constant, because taking the harness's label would let the very confusion under test pass unnoticed. `_page_unit_regions` is called first for its side effect: the contiguity and granularity checks live there, and they have to run before any view is handed out rather than after. --- verification The unit tests skip on a non-GPU runner (the module imports aiter at load), so the geometry was checked directly by extracting the four methods and comparing `page_unit_views` against `_page_unit_bases` byte for byte: 9 segments for 3 units x 3 rows, identical `data_ptr()`s, lengths equal to `_page_unit_stream_sizes`, every view contiguous, no two units sharing a byte, an out-of-range id refused, and a granularity mismatch refused before any view is produced. pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors that fail on #2045 too) -> 3911 passed, zero failures. black clean; ruff 36, unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

… evicted Phase 3d, the last of Phase 3. This is what turns the tier back on for saves: a checkpoint that reaches READY is offered to the CPU tier, and the whole engine->worker->engine round trip is wired. Loads are still off (Phase 4). --- what changed `page_unit_checkpoint.py` - `_offload_ready`: checkpoints nominated at the READY transition. NOT pinned. - `take_offload_stores(max_inflight)`: hands `(hash, unit_ids)` over and takes the pin at that moment, up to the cap. - `settle_offload_store(hash)` / `reclaim_stale_offload_pins(timeout_s)` / `offload_pins_reclaimed`, exposed on the coordinator and reported in `checkpoint_fates`. - `PageUnitCheckpointStore(..., offload_sink=)` -- False nominates nothing. `block_manager.py` - `take_state_stores` / `settle_state_store(hash, ok)` / `reclaim_stale_state_store_pins`. The switch (`kv_connector_hosts_state_tier`) moves above the coordinator's construction, which now needs it. `scheduler.py` - `_publish_state_stores()`, beside `_publish_state_loads()` and on the same two call sites plus the idle path. - `_offload_max_pending_saves()`, reading the KV leg's own `OFFLOAD_MAX_PENDING_SAVES`. - The report drain settles the pin on both outcomes and runs the reconciler. `kimi_k3/connector.py` -- `enqueue_state_stores`, `meta.state_stores`, `_start_state_stores`, and `_state_store_failed_locally` for stores a worker with no tier cannot even attempt. `kimi_k3/state_tier.py` -- `submit_store(h, unit_ids)`, no `ready_event`. `kimi_k3/state_object.py` -- `put(h, unit_ids)` gathers `page_unit_views`; `get(h, slot)` still scatters into `state_entry_views`; `puts_refused`. `metadata.py` / `types.py` -- `state_stores` declared as work. --- why READY, and why the pin is not taken there Three candidate trigger points and only one works. `begin_store` is too early: the scatter has not ridden a batch, so the units do not hold the image yet. `_evict` / `unindex` are too late -- they fire when the pool wants those units *now*, so a multi-millisecond D2H started there holds 127 of the most-wanted blocks at the worst possible moment. READY is when the bytes first exist and when the record is least wanted: it has just entered the LRU at the cold end. The pin is a separate question, and getting it wrong broke a real invariant -- `test_a_ready_unpinned_checkpoint_is_available_to_live_kv` caught the first draft. A pin lives in this process while the copy runs in the worker, so it spans several scheduler passes. Pinning at READY would make EVERY checkpoint un-evictable for that window, which is exactly what #2045's admission argument forbids. So READY only *nominates*: a nominee stays unpinned and spendable, and `take_offload_stores` pins only the few actually in flight. A checkpoint the pool needed more than the tier did simply loses its copy, which is the right price for never making the pool wait on the CPU. The cap is `OFFLOAD_MAX_PENDING_SAVES`, shared with the KV leg rather than given a knob of its own: a KV save and a state store both hold bytes out of the same pool while they run, so the depth is one question about one resource, and two numbers would let an operator raise one and unknowingly double the other. --- write volume, which is the objection this trigger invites Storing every checkpoint rather than only evicted ones is more traffic, and it is affordable: per request state adds 53.6 MiB on top of ~404 MB of KV, **13%** on a path already in flight (~12 MB/s against the KV leg's ~90 MB/s over a 3600s run). What it buys is the trigger being at the cheapest possible moment and dedup for free -- `begin_store` already refuses a hash it holds or has pending, so one hash is offered once. --- unindex during a copy does not cancel it `unindex` means the boundary's KV block was spent, so the image is unreachable in HBM from here on. That is precisely when the CPU copy is worth having, so the pin holds the units to the end of the D2H and `_release_record` runs after. The CPU copy outliving the HBM one is the entire point of the tier. --- the lost report, and why 20 lines is enough `reclaim_stale_offload_pins` takes the same window as `_reconcile_stalled_deferred_saves`, derived from LMCache's own pin timeout for the same reason. Recovery is total: a leaked pin breaks no `BlockPool` invariant and leaves no half-released state, so zeroing the count restores the record exactly -- which is why this needs a counter and a timer rather than any repair logic. `offload_pins_reclaimed` reaches `checkpoint_fates`, because otherwise the symptom is a pool that has quietly shrunk. --- still off Loads. Nothing calls `_request_state_load` for K3 yet: `_attach_state_slots` returns from its PAGE branch before reaching `_tier_can_serve`, and the `begin_restore` raise above it has to become an issue-a-load branch first (Phase 4). `StateOffloadIndex.hashes` now fills, so the tier is half live: it writes, and nothing reads yet. --- tests 11 new in `tests/test_page_unit_checkpoint.py` -- nomination does not pin, the hand-off does, a nominee stays spendable and may be lost, no sink nominates nothing, offered once, the cap makes the rest wait rather than drop, a waiting nominee stays evictable, both outcomes release, a report for a hash never sent is a no-op, `unindex` mid-copy holds the units to the end, and a lost report is reclaimed with the record left unbroken. Two more in `tests/test_lmcache_offload_connector.py`: stores drain exactly once, and a step whose only work is a store still reaches the worker. Verified: pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors that fail on #2045 too) -> 3924 passed, zero failures. black clean; ruff 36, unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…raised Phase 4a, and it has to be first: with nothing else changed this commit is behaviour-preserving, but without it the first request that actually uses the CPU tier takes the engine down. --- what changed `BlockManager._attach_state_slots`, PAGE branch. Was: if hit_hash != -1 and not begin_restore(hit_hash, seq.state_slot): release_many(seq.state_slots); seq.state_slots = [] raise RuntimeError("gated PAGE checkpoint disappeared before state attach") Now: cold start returns early; `begin_restore` success returns as before; a miss falls through to `_request_state_load`, and only if that also declines is the boundary disowned (`return False`, which `allocate` turns into `num_cached_tokens = 0`). New counter `state_gate_lost_boundary`, reported in `checkpoint_funnel`. --- why the raise could not stay It asserted that the gate and the HBM store agree. That held while `can_allocate` only ever accepted boundaries the HBM index carried. Phase 4b gives the CPU tier a vote, so a hash the gate accepted may live only there -- and every such request would hit this raise. Turning it into a load branch has to land before the vote does, not after, or the two commits are only safe as a pair. --- why disowning is reachable rather than defensive Two ways to arrive with neither tier able to produce it: * the tier's index is optimistic by construction -- `hashes` means "was stored once", never "is still there", because LMCache's own LRU can drop bytes at any time; * an HBM checkpoint can be unindexed between `can_allocate` and `allocate` by another seq's `_fresh_block` in the same pass. So this is a normal path with a counter, not an assertion. A large fraction of `joint_boundaries` landing here would mean the gate is accepting boundaries that do not survive to attach. --- the slots are kept, unlike the old abort The old code released the slots because it was aborting. Disowning is not an abort: the request is about to recompute its whole prefix, and it writes that state into these very slots. Releasing them would hand the next request a buffer this one is still filling. `test_missing_gated_checkpoint_releases_the_new_slot_and_raises` becomes `test_a_gated_boundary_neither_tier_has_is_disowned_not_raised`, and now pins all three: returns False, keeps its slots, counts one. Verified: pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors that fail on #2045 too) -> zero failures. black clean; ruff unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…m it Phase 4b. This is the commit that makes the tier reachable: until now K3 wrote 53.6 MiB per checkpoint to LMCache and nothing ever read it. --- what changed `page_unit_checkpoint.py` - `PagedStateCheckpointCoordinator.offload` + `attach_offload(index)`. - `resumable_hit` now asks `_reachable(h)` instead of `store.contains(h)`: HBM first, then `offload.hashes`. - The constructor's `offload_sink: bool` is gone -- `attach_offload` sets the vote and the sink together (see below). `block_manager.py` - Attaches the index to the coordinator once both exist. The coordinator is built from `state_runtime.checkpoint_spec` and the index from the connector config, so they are constructed in the wrong order for a ctor argument. `kimi_k3/connector.py` - Refuses to build the tier when `image_bytes != entry_bytes`. --- why the vote is the whole feature `can_allocate` walks the HBM prefix and asks every `Pool.STATE` class "the rightmost boundary <= X that you accept". With the coordinator answering only from `store.contains`, a hash whose image went to LMCache is invisible: the walk stops short, nothing ever asks for it back, and the store path is pure cost. `_reachable` is where the tier gets to say "I have that one". It stops at the first boundary it accepts, scanning right to left, so a hash accepted here that nothing can deliver does not cost a wasted lookup -- it costs the whole walk-back, hiding every shorter checkpoint still resident in HBM. That is why the disown path (Phase 4a) had to land first. No preference rule between the tiers is needed: both are keyed by the same content hash, so the scan takes the rightmost boundary wherever it lives, and `_attach_state_slots` tries `begin_restore` before `_request_state_load` regardless. A resident image never pays a park. --- why the vote is optimistic, deliberately `hashes` means "was stored once", never "is still there" -- LMCache's own LRU can drop bytes under it at any moment. A false positive costs one park plus a recompute, and retracts itself: `fail_load` calls `forget(h)`, so the next request over that prefix is not sent down the same hole. Being certain instead would mean a synchronous cross-process lookup on the admission path. --- why attaching is one call, not two fields Half-attached is worse than off in both directions. A vote with no sink accepts hashes nothing will ever store, so every accepting request parks and recomputes. A sink with no vote pins 127 units per checkpoint for a transfer nobody will read. `attach_offload` sets both, so they cannot diverge. --- the load's destination is the Active Slot Not 127 freshly reserved units. The two are interchangeable because #2045's copy plan intersects two *ordered byte streams* -- the slot's ranges and the units' regions are two cuts of one image -- so a blob gathered in unit order unpacks correctly in slot order. Writing the slot directly means no `ensure_free_units` on the admission path, no refusal branch, no units pinned for the length of an H2D, and no cap needed on concurrent loads. It does mean the store reads `image_bytes` and the load writes a whole slot, so the connector now refuses to build a tier where those differ, at the one point both numbers are in scope. They are equal for K3; a model where they are not would silently truncate the store or over-read the load. --- tests Six in `tests/test_page_unit_checkpoint.py`: a tier-only hash is accepted, no tier means no acceptance, the scan still takes the rightmost, attaching turns on both halves together, and the optimistic half retracts on `forget`. Three in `tests/test_state_checkpoint.py`, end to end through `_attach_state_slots`: a tier-only boundary becomes a queued load with no restore op and the slot as its destination; a boundary in both is served from HBM with no park; a tier that declines disowns and counts. Verified: pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors that fail on #2045 too) -> zero failures. black clean; ruff unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…ason to exist Phase 5, the last one. `_joint_kv_boundary` refused every PAGE sequence; it now serves only PAGE. With this the tier is reachable end to end. --- what changed `block_manager.py` - The gate inverts: - if not seq.has_per_req_cache or self.paged_state_checkpoints is not None: - return self._no_joint("not_hybrid") + if not seq.has_per_req_cache or self.paged_state_checkpoints is None: + return self._no_joint("no_paged_checkpoints") - The hbm/tier split reads the coordinator instead of the slot pool. - `joint_boundaries_hbm` / `_tier` become `state_hbm` / `state_tier`. `page_unit_checkpoint.py` -- `contains(h)` delegation, so the split does not reach into `.store`. `scheduler.py` / `llm_engine.py` -- the log line and the Prometheus tuple follow the rename. --- why the old refusal was right, and why it stopped being right It read `is not None` and refused, because a K3 checkpoint was an Active Slot that the tier spilled out of `StateSlotPool` -- a PAGE image lived in the KV pool, which the KV connector already offloads on its own, so there was no second boundary to agree on. #2045 moved K3's image into the KV pool and with it the entire reason the joint path exists. HBM now enforces `state ⊆ KV`: `_record_evicted` unindexes a checkpoint the instant its boundary block is spent, so a checkpoint can no longer outlive its KV there. **When LMCache hands the KV back, nothing hands the state back unless the two are fetched together.** Leaving this refusal in place makes the tier dead weight -- it writes 53.6 MiB per checkpoint and nothing ever reads it -- and no other counter in the system would say so. --- why the counters were renamed rather than left alone `joint_boundaries_hbm/_tier` described the boundary as a whole, which was true when the KV leg had no tier of its own to be served from. It has one now (`_decide_load_after_alloc` decides it separately), so the old names read as if they covered both legs. `state_hbm` / `state_tier` say what they measure. The split also means something new. It used to say "the state pool is too small for this concurrency". Under #2045 checkpoints share the paged pool with the KV write stream, so a ratio dominated by `state_tier` says **the paged pool** is too small -- checkpoints are being squeezed out by the KV they now live beside. That diagnosis did not exist before. `state_tier` is also the only counter here that cannot be non-zero with the CPU tier switched off, which makes it the honest test of "did this feature run". No passing unit test can produce it in production. --- tests Three in `tests/test_state_checkpoint.py`, on a harness that had to reproduce the shape the joint path is for: publish a prompt, place ONE rung at an interior boundary (a checkpoint under the last block is one no scan can look up, since `can_allocate` matches over `range(n_hash_blocks - 1)`), then evict a KV block below it so the walk stops short. A PAGE seq gets a joint boundary at the rung with real work for the KV leg; the split reports `hbm` while the checkpoint is resident and `tier` once it is not; no tier means no boundary. Verified: pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors that fail on #2045 too) -> 3935 passed, zero failures. black clean; ruff unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…res` Found by the first end-to-end run of the K3 CPU state tier on this branch: the tier reported healthy on both sides and stored exactly nothing. 94 checkpoints were nominated, handed to `_publish_state_stores`, refused, and had their PAGE units released again. --- the symptom Server log, 94 times over a 6-minute agentic replay at concurrency 8: state offload: LMCacheOffloadConnectorScheduler did not carry 1 state store(s); releasing their units now. with everything else looking correct: lmcache_offload: worker family=kimi_k3 kimi_k3 offload: state tier up, entry=55.39 MiB rank=0, sharing the paged-KV CPU pool, layout=kda-paged-state-v1:layers=69:... state checkpoints: checkpoints_kept=112 ... offload_pins_reclaimed=0 So the worker half built the tier, the coordinator nominated checkpoints at READY, and `take_offload_stores` pinned and handed them over -- and the scheduler half dropped every one. --- the cause `LMCacheOffloadConnectorScheduler` is a shell that owns no behaviour: it holds `self._impl = _build_scheduler(config)` and forwards each method by hand. `7b6b58e0` added `enqueue_state_stores` to `KimiK3OffloadScheduler` (kimi_k3/connector.py:352) but did not add the matching forwarder to the shell, so the method existed on the implementation and was invisible from outside: $ git grep -c enqueue_state_stores atom/kv_transfer/offload/connector.py 0 `Scheduler._publish_state_stores` probes with `getattr(self.kv_connector, "enqueue_state_stores", None)`, and `self.kv_connector` is the shell, so the probe returned None unconditionally and the store path took its "nothing will carry these" branch on every pass. This is not a race or a config error -- it could never have worked. The load leg was unaffected only by luck: `enqueue_state_loads` and `take_state_reports` were both given forwarders when they were added. Loads still moved no bytes, because with nothing ever stored `offload.hashes` stayed empty and the tier had nothing to vote for. --- the fix Add the missing forwarder, in the same shape as its `enqueue_state_loads` neighbour so the two cannot drift again. --- the second change, and why it is in this commit The warning that reported the failure named a cause that was not the one: The tier's index should not have been installed against this connector. That reads as a configuration mistake, and it sent the first pass of this diagnosis at `kv_connector_hosts_state_tier` and the layout selector -- both of which were behaving correctly. The class name printed is the *shell*, which tells you nothing about whether `_impl` supports the call. Both warnings now name the two possible causes and say which class is being reported. --- verification Static: the probe is `getattr` on a method the shell does not define, so the refusal is unconditional; adding the forwarder is the whole of the repair, and it is additive -- no existing path changes. Runtime: re-running the same workload with this commit applied, the `did not carry` warning no longer appears. Lines conform to black's 88-column limit; black and ruff are not installed in the benchmark image, so formatting was checked by inspection against the surrounding style rather than by running them. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…that holds none Three observability defects, found while trying to judge whether the K3 CPU state tier was working. Together they meant the tier could store nothing at all and no number anywhere would say so -- which is exactly what happened in 8a5364a, where 94 refused stores produced one warning line and four counters that all read zero for unrelated reasons. --- 1. the store leg had no counters at all `StateOffloadIndex.stats()` reported `loads_attempted / completed / failed` and `indexed`. On a K3 run the load counters stay zero until something has been stored, so the only store-side signal was `indexed` growing -- and `indexed` is a set size that also grows for reasons unrelated to this pass. Added `stores_attempted` (handed to the connector, counted in `take_state_stores`), `stores_completed` / `stores_failed` (split in `settle_state_store` by the worker's report), and `stores_refused`. `stores_refused` is deliberately apart from `stores_failed`, and it is the one that matters most: failed means the worker tried and could not, refused means nobody tried. The shell-forwarding bug fixed in 8a5364a was a pure `refused` condition, and with this counter it would have been one glance at the periodic line instead of a hunt through the connector factory. The worker's own `puts_refused` needs no separate channel: a refused `put` already returns False and settles as `ok=False`, so it lands in `stores_failed`. --- 2. `[Pool Pressure]` read the pool that holds no checkpoints `BlockManager.pool_pressure()` folded in `self.state.checkpoint_fates()`. Under #2045 a K3 checkpoint is a PAGE image the coordinator owns and the slot pool holds none, so the line printed [Checkpoint Fates] kept: 0, dropped: 0, evicted: 0, orphaned: 0 on the same run whose other line said `checkpoints_kept=112`. Two counters for one fact, disagreeing in one log, with the wrong one printed in the section a reader goes to for pool pressure. Now reads `self._state_checkpoint_cache`, which is already the coordinator when there is one and the slot pool otherwise -- the same object every other checkpoint question is asked of. `occupancy()` stays on the slot pool, which is the thing that has slots. --- 3. the tier's counters vanished when K3's checkpoints moved to PAGE `StateGroupPool.checkpoint_fates` used to fold the tier's stats in under a `state_offload_` prefix. `PagedStateCheckpointCoordinator.checkpoint_fates` was written without that, so once #2045 moved K3's checkpoints into the coordinator the whole `state_offload_*` family disappeared from the logs -- a silent regression, since nothing errors when a key stops being emitted. Restored on the coordinator. The fold is `getattr`-guarded because `attach_offload` accepts anything answering `hashes`, and the tests attach a double with no counters; a missing `stats` means "nothing to fold", not an error. (`test_a_tier_that_declines_disowns_rather_than_parking` caught this on the first draft.) --- also added: `state_offload_store_backlog` A gauge rather than a counter: `len(_offload_ready)`, the nominations waiting for a slot under `OFFLOAD_MAX_PENDING_SAVES`. That queue is unbounded and drained oldest-first, so checkpoints reaching READY faster than the tier drains them shows up as a value that climbs and stays high, with fresh nominations queueing behind stale ones. There is no other way to see it -- the stale entries are skipped cheaply on the way out, so nothing else registers the backlog. --- how to read the new line state_offload_stores_refused > 0 -- wiring fault, the connector never took it stores_attempted - completed - failed -- in flight stores_completed vs checkpoints_kept -- how much of what HBM keeps the tier got loads_failed / loads_attempted -- the index's false-positive rate store_backlog climbing -- READY outruns the drain Verified: pytest tests/test_page_unit_checkpoint.py tests/test_state_checkpoint.py tests/test_lmcache_offload_connector.py tests/test_block_manager.py -> 466 passed. Lines conform to black's 88-column limit; black and ruff are absent from the benchmark image, so formatting was checked by inspection. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Every state store on this branch has failed since the store path was wired. Found by the first run where stores were actually carried (8a5364a): 3,904 tracebacks, all identical, and not one checkpoint reached LMCache. --- the symptom Traceback (most recent call last): File "kimi_k3/state_tier.py", line 109, in _do_store stored = bool(self.codec.put(h, unit_ids)) File "kimi_k3/state_object.py", line 141, in put self._staged.pack(self._backend.page_unit_views(unit_ids), obj) File "kimi_k3/staging.py", line 175, in pack from atom.kv_transfer.offload.triton_kv_staging import fused_pack_chunk_major ModuleNotFoundError: No module named 'atom.kv_transfer.offload.triton_kv_staging' state offload: store of hash 10785477046942485201 failed 3,904 of them over a 50-minute concurrency-16 agentic replay -- one per store attempt, on every rank. `stores_failed` would have counted all of them, and `indexed` never moved. --- cause 1: the module is one package deeper `fused_pack_chunk_major` lives in `atom.kv_transfer.offload.dense.triton_kv_staging`. The dense codec, which has used it all along, imports it from there (`dense/kv_byte_codec.py:127`). The K3 stager's import dropped the `.dense`. Nothing caught it because the import is function-local, inside `pack`, so it is only attempted when a store actually runs -- and until 8a5364a no store ever reached the worker. The unit tests do not reach it either: they exercise the codec against doubles that never call the real packer. --- cause 2: a failed `put` leaked its allocation, once per failure obj = self._allocate(self.entry_bytes) # ref_count = 1 self._staged.pack(...) # raises self._storage.batched_put([key], [obj]) # never reached `batched_put` discharges the reference it is handed (`StorageManager.batched_put` ends in `memory_obj.ref_count_down()`), so the success path correctly owes nothing -- but only if it is reached. With `pack` throwing, the allocation was stranded at ref_count=1 and LMCache reported it much later: MemoryObj at ... is being garbage collected with ref_count=1, pin_count=0. This indicates ref_count_down()/unpin() was not called 3,768 such warnings against 3,904 failed stores: one leak per failure. Left alone this shrinks the shared CPU pool by one 53.6 MiB entry per failed store until nothing can be allocated at all -- so the second-order effect of the import bug was a slow starvation of the pool the KV leg shares. `get` already guarded its own reference with `finally` for exactly this reason; `put` did not. It does now, on the exception path only, matching what the DSV4 codec does (`hybrid/dsv4/codec.py:1293`). --- why one commit The leak is only reachable through a throwing `pack`, and the import was the only thing throwing. Fixing the import alone would hide the leak rather than remove it, and the next exception from that line -- an out-of-memory in the staging buffer, say -- would strand allocations again with no traceback to explain the pool shrinking. --- verification python -c "from atom.kv_transfer.offload.dense.triton_kv_staging import fused_pack_chunk_major" -> OK pytest tests/test_lmcache_offload_connector.py tests/test_state_checkpoint.py tests/test_page_unit_checkpoint.py -> 412 passed The import is still function-local, so this remains untested by anything that does not run a real store. A unit test that calls `StagedTransfer.pack` against real CUDA tensors would have caught it; that is worth adding and is not in this commit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…ot reach Fix-then-sweep over `dfab0fa0` and `8a5364ab`. Both of those repaired one instance of a bug that had two, and neither instance the fix reached is covered by any test. --- 1. `unpack` imports the same module that does not exist `dfab0fa0` corrected `StagedTransfer.pack`'s import of `fused_pack_chunk_major` (the module is one package deeper, under `.dense`). `unpack` imports `fused_unpack_chunk_major` from the same wrong path, eleven lines further down, and was left alone. It has never fired because no state load has ever run -- with nothing stored, `offload.hashes` stayed empty and the tier had nothing to vote for. The first successful store makes the first load reachable, and it would have raised `ModuleNotFoundError` on that load, on every rank. Both kernels live in `atom.kv_transfer.offload.dense.triton_kv_staging`; verified by importing both. --- 2. `chunk_size` is not on the delegating shell either `8a5364ab` added the missing `enqueue_state_stores` forwarder. Sweeping every member `Scheduler` reads off `self.kv_connector` against what `LMCacheOffloadConnectorScheduler` defines turns up one more: `chunk_size`. scheduler.py:2284 seq.offload_kv_chunk_tokens = int( getattr(self.kv_connector, "chunk_size", 0) or 0) `self.kv_connector` is the shell, the shell has no `chunk_size`, so this has always stamped 0. It is latent rather than live: `_joint_kv_boundary` reads `self._joint_chunk_tokens` first and only falls back to this, and that primary comes from `BlockManager.__init__` reading the LMCache config directly. But the fallback exists precisely for the case where that read failed -- and it was dead, so a build that could not read the config at startup would refuse every joint boundary with `no_chunk_size` and never say why. Forwarded as a property, with the reason in its docstring. --- 3. the class of bug, closed rather than fixed twice Both are the same shape: the shell forwards by hand, so a member added to an implementation is invisible from outside, and every reader is a `getattr(..., default)` that does not raise -- it takes the default forever. Two instances in two days says the next one is a matter of time. `test_every_member_the_scheduler_reads_is_reachable_through_the_shell` scans `scheduler.py`'s source for `getattr(self.kv_connector, "x")`, `self.kv_connector.x`, and `_connector_flag("x")`, and asserts every name is reachable on the shell class. It reads the source rather than a hand-kept list, because a hand-kept list is the same failure one level up. The scan asserts it found something, so a regex that stops matching fails loudly instead of passing vacuously. Falsified: renaming the new `chunk_size` property makes it fail and name the missing member. Verified: pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors) -> 3936 passed, zero failures. black clean; ruff 36, unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…it were image Fourth and last blocker in the chain that kept the K3 CPU tier from storing anything. With the previous three fixed, a 25-minute concurrency-16 replay still ended at state_offload_stores_attempted=462 state_offload_stores_completed=0 state_offload_stores_failed=450 state_offload_indexed=0 and 3,696 identical tracebacks: File "kimi_k3/state_object.py", line 149, in put self._staged.pack(self._backend.page_unit_views(unit_ids), obj) File "kimi_k3/staging.py", line 182, in pack dst_tensor = self.memory_tensor(dst, nbytes) ValueError: ATOM LMCache connector: MemoryObj tensor is too small for 59867136 bytes; got 58079232 --- the geometry From the run's own startup line: unit_bytes=2138112 image_bytes=58079232 units_per_checkpoint=28 ceil(58079232 / 2138112) = ceil(27.163) = 28 28 x 2138112 = 59867136 <- what page_unit_views gathered 58079232 <- what put() allocated difference = 1787904 <- padding in the 28th unit An image occupies whole units, and the last one is 83.6% padding here. Nothing owns those bytes. --- why only the CPU side hit it The HBM path never had the problem because it already passes the image size down: `_checkpoint_copy_plan` calls plan_segmented_copy(src_sizes, self._page_unit_stream_sizes(units), spec.image_bytes) so the D2D copy writes exactly `image_bytes` into the unit stream and leaves the tail untouched. The image is, by construction, the LEADING `image_bytes` of that stream. `page_unit_views` was written as the tensor-view counterpart of the same addressing but did not carry the third argument across, so it named whole units to a consumer that sizes its destination at `entry_bytes` -- which for K3 is exactly `image_bytes`. The connector's `image_bytes != entry_bytes` guard cannot catch this: both are 58079232 and agree. The quantity that disagrees is a third one, `sum(page_unit_views)`, which no check compared against either. --- the fix Trim the gathered stream to `image_bytes`, dropping whole views past the budget and slicing the one that straddles it. The slice is byte-exact rather than view-aligned because the image does not end on a view boundary: 58079232 leaves 350208 B of the last unit in use, which is 20.96 of its rows. That view is reinterpreted as `uint8` and sliced; the packer sizes segments as `numel * element_size`, so a uint8 segment costs it nothing. The load side needs no counterpart: it scatters into `state_entry_views`, which sums to `entry_bytes` exactly, so both ends of the round trip are now the same 58079232 bytes. A budget left unspent raises rather than storing a short blob -- that would be a real geometry disagreement, and a truncated image read back as valid is the one outcome worth crashing over. --- two notes on shape Inline rather than a helper method: the existing tests drive `page_unit_views` with a `types.SimpleNamespace` stand-in for `self`, and a second attribute would have to be stubbed in every one of them. For the same reason the spec is read with `getattr(spec, "image_bytes", 0)` -- a fork build carries no spec at all, and there the whole-unit stream is the right answer. --- verification pytest tests/test_kda_checkpoint_slot_copy.py tests/test_lmcache_offload_connector.py tests/test_state_checkpoint.py tests/test_page_unit_checkpoint.py -> 441 passed Also confirmed from the same run that the previous commit's leak guard works: `MemoryObj ... garbage collected with ref_count=1` went from 3,768 occurrences to 0, even though every store still failed. None of the four bugs in this chain was reachable by the unit tests, because all of them live past the first line that only a real store executes. A test that packs real CUDA tensors through `StagedTransfer` would have caught three of the four; it is still worth adding. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

About half of every K3 state store failed. Measured over a concurrency-16 agentic replay with the four earlier blockers fixed: state_offload_stores_attempted=46 state_offload_stores_completed=19 state_offload_stores_failed=27 <- 59% with 296 identical tracebacks: File "kimi_k3/state_object.py", line 153, in put self._storage.batched_put([self.key(h)], [obj]) ^^^^^^^^^^^ File "kimi_k3/state_object.py", line 107, in key digest.update(int(h).to_bytes(8, "little", signed=True)) OverflowError: int too big to convert state offload: store of hash 12983297547834104159 failed --- cause 1: the hash is unsigned and the key said otherwise 12983297547834104159 > 2**63-1, so `to_bytes(8, signed=True)` cannot hold it. An ATOM block hash is `xxhash.xxh64().intdigest()`, which spans the whole 0..2**64-1, and `BlockManager.compute_hash` chains one into the next with `prefix.to_bytes(8, "little")` -- unsigned, no keyword. The state key was the only place in the chain that read the same value as signed, so it rejected every hash with the top bit set. That is half the key space, and the observed 59% is that with sampling noise. Nothing downstream noticed because the failure is inside a `try` in `_do_store`, which counts a failed store and moves on. Before the counters added in a804561 it produced no number at all. --- cause 2: the raise stranded a MemoryObj, in a window the last guard missed `dfab0fa0` wrapped `pack` so a throwing gather would discharge the allocation. `key(h)` is evaluated in `batched_put`'s argument list, which is after the allocation and before `batched_put` can take ownership -- outside that guard. So each of the 296 overflows also leaked a 53.6 MiB entry, and the LMCache warning reappeared next to the traceback: MemoryObj at 136839168 is being garbage collected with ref_count=1 Fixed by computing the key *before* allocating. `key` is pure, so hoisting it removes the window rather than guarding it, and the `except` now covers `batched_put` itself as well -- matching what the DSV4 codec does. --- what this unblocks The same `key` is used by `get`, so roughly half of all loads would have missed in a way indistinguishable from an LRU eviction -- `fail_load` would then `forget` a hash whose bytes were never stored under a name it could find, and the index would decay for a reason nothing reported. Verified: pytest over the state and offload suites -> 442 passed. The overflow itself is confirmed directly: >>> (12983297547834104159).to_bytes(8, "little", signed=True) OverflowError: int too big to convert >>> (12983297547834104159).to_bytes(8, "little", signed=False).hex() '5f19dacf54f02db4' Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

`961f80d6` fixed `page_unit_views` gathering a checkpoint's padding, and the suite it ran stayed green before and after — because no test reaches the trim at all. The addressing tests build their stub spec as checkpoint_spec=SimpleNamespace(page_unit_bytes=...) with no `image_bytes`, so `budget` is 0 and every one of them takes the `return views` early path. The trim shipped untested, which is how a bug that failed 3,696 stores in a row got there in the first place. --- what is pinned Seven cases in `TestPageUnitViewsStopAtTheImage`, on the harness the addressing tests already use (`build` gains optional `image_bytes` and `dtype`; both stay off by default so the existing tests keep exercising whole units): * the gathered stream is exactly `image_bytes` — the invariant `put`'s allocation depends on, and the one that was broken; * a whole view past the budget is dropped, not truncated to zero length; * an image that lands on a view boundary is not sliced at all; * **the kept bytes are the leading bytes and nothing else** — trimming truncates, it does not reorder. The blob is read back by scattering into the slot in the same order, so a byte that moved here would land in the wrong layer there with nothing raised; * a multi-byte dtype is sliced by BYTES, not elements — the straddling view is reinterpreted as `uint8` first, and on bf16 an element-wise slice would keep twice what it should. The production cache is 1 byte per element, so only a test with a wider dtype can hold this; * a spec with no image size keeps whole units (the fork build, which has no spec to trim against); * units that cannot cover the image raise rather than storing a short blob. --- verification `tests/test_kda_checkpoint_slot_copy.py` skips wholesale on a non-GPU runner (the module imports aiter at load), so the seven cases were additionally run by extracting `page_unit_views` and its three neighbours out of the source and driving them directly. All seven hold, including the bf16 one. pytest tests/ (less tests/plugin and the two msgpack/aiter collection errors) -> 3936 passed, zero failures. black clean; ruff 36, unchanged. --- one note on `put`'s exception shape, checked not changed `2c5c3457` moved `batched_put` inside the `try` whose `except` calls `ref_count_down`. That is safe only if `batched_put` takes ownership just on success — and it is the pattern the DSV4 codec already uses (`hybrid/dsv4/codec.py:1284`, "StorageManager owns the MemoryObj after a successful batched_put"). Consistent with the repo rather than a new assumption, so left alone. The window between `_allocate` and `batched_put` is now fully covered: `page_unit_views` can raise (the geometry check above) and is inside the try; `key(h)` is computed before the allocation. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Copilot

Pull request overview

Copilot reviewed 39 out of 39 changed files in this pull request and generated no new comments.

zejunchen-zejun · 2026-08-31T11:12:42Z

Thanks for the extremely thorough pass, @valarLip — the ordering by consequence made this easy to work through. Status below, ordered to match your list. Everything is pushed to zejun/lmcache_with_2045 in 0d704202..7252da66 (5 commits) plus two observability follow-ups (c531ddb4, c88f46b3); all py_compile-clean and added lines kept ≤88 cols.

Fixed (10)

#	Finding	Commit	Fix
1	`_state_store_failed_locally` written from the worker but only created on the scheduler → `AttributeError` before `super().start_load_kv`, wedging the whole step's KV	`951dab4c`	No-tier store-failure is now routed through the `STATE_INDEX`/`SOURCE` channels instead of a cross-process attribute write
2	`unpack`'s `except` missing the `_drain_device()` that `pack` has → slot recycled while the copy kernel is still queued	`18a9869a`	Both `pack` and `unpack` failure fast-paths now drain the device first. This was the real silent state-corruption path
4	Joint park armed with no tier → `get_finished` returns at the tier-`None` guard before `_settle_joint`, leaking one `_need`/`_alias` entry per load	`951dab4c`	Guard the arm so the park is only armed when a tier is present
6	`MultiConnectorScheduler` never forwards `chunk_size` → joint KV refused with `no_chunk_size` under `kv_connector: multi`	`ccddb8fe`	`chunk_size` is now propagated on `MultiConnectorScheduler`
7	`_joint_chunk_tokens` reads the wrong config level under `multi`	`73eccab1`	Unwrap the offload sub-config before reading `_joint_chunk_tokens`
8	`_adopt_state_tier` finds nothing under `multi`	`ccddb8fe`	Added a `_state_tier` property on the worker shell so the `_adopt_state_tier` dup-guard revives
9	`has_pending_work()` doesn't count state work	`951dab4c`	Override now ORs `_pending_state_loads` / `_pending_state_stores`
10	`_orphan_load_slots` has no reclaimer → stranded slots wedge the `can_allocate` state gate	`73eccab1`	Added `reconcile_orphan_load_slots`, the load-side twin of `reclaim_stale_state_store_pins`
11	`take_offload_stores` drops an in-flight-hash nomination instead of re-queueing	`7252da66`	Defer + requeue in-flight-hash store nominations rather than dropping them
12	`_build_state_tier` validates only one of two backend APIs	`951dab4c`	Probe that `page_unit_views` is callable at tier-build time

Not fixed — assessed, no code change (rationale)

3 — dropped free_event wait / free_event_valid is write-only. No active correctness bug: ordering on that path is already guaranteed by producer.synchronize(), which is why the wait was dropped. free_event_valid is now vestigial. Re-adding the wait would be redundant and, in testing, risks regressions on the fast path; the cleaner move is to delete the dead field rather than restore a second barrier. Happy to do that removal if you'd prefer the dead-code cleanup in this PR.

5 — state-only load reported on neither finished_loading nor failed_loading. Not reachable in the current wiring: the connector only arms a load when load_spec is not None and hardcodes needs_kv=True, so a load always carries a KV id and the state tier reports under that id — _settle/_release fires correctly. Even keyed on a bare req_id (kv_id=None) the same id is reported. The alias seam only engages when there is a KV id, where it is already correct. A "fix" here is either a no-op or risks inverting the joint case, so I left it and documented the invariant.

Observability

You're right that the state-tier observability was mostly dead; one clarification and two follow-up commits:

Clarification: the engine-side StateOffloadIndex counters are actually live — stats() → checkpoint_fates → checkpoint_funnel (engine_utility.py:317/:358). So "every hook is dead" is a slight overstatement; the dead ones were the worker-side accumulators.
c531ddb4 — folded my own dead _orphan_load_slots_reclaimed (introduced in 73eccab1, no reader) into checkpoint_funnel(), symmetric with the store-side offload_pins_reclaimed. Agreed this funnel wiring is the highest-leverage bit.
c88f46b3 — removed the genuinely-dead worker-side StateOffloadTier.stats() and its three write-only accumulators (load_queue_wait_ms_last/max, loads_started) — only tests called them, there's no worker→engine drain. Kept the one live signal: the _do_load load-wait logger.warning (≥ _LOAD_WAIT_WARN_MS, with oldest_store_age_s()).

Two "smaller items" I looked at but deliberately left (not bugs):

StateOffloadIndex.hashes forget-on-evict would be wrong — hashes tracks the CPU/LMCache tier (note_stored on the worker report), whereas checkpoints_orphaned is the paged-HBM pool; the design intentionally uses fail_load as the only engine-visible LMCache-eviction signal.
Capability gate by shared name vs resolved variant — stores_refused already catches the harmless mis-nomination, so this is a behavior change, not a correctness fix; non-blocking.

valarLip · 2026-08-31T13:40:15Z

Thanks for the fast and precise turnaround — the commit-per-finding table made this easy to re-check. Re-reviewed at 7e08f5994.

You're right about the observability, and I overstated it. The engine-side StateOffloadIndex counters are live via stats() → checkpoint_fates → checkpoint_funnel (engine_utility.py:317/:358); what was dead was the worker-side accumulators. c531ddb4 and c88f46b3 are the right calls — folding your own _orphan_load_slots_reclaimed into the funnel rather than leaving a second write-only counter is exactly the discipline that was missing.

Two things below: one rebuttal I think has to be reopened, and a new batch.

The #5 rebuttal needs reopening — and it is your #4 fix that made it certain

Your rationale was:

the connector only arms a load when load_spec is not None and hardcodes needs_kv=True, so a load always carries a KV id

That described the pre-fix wiring. On 7e08f5994, _arm_joint_loads now returns early — this is the guard you added for #4:

if self._state_tier is None:
    # No tier means the state leg is unsettleable: `_start_state_loads`
    # fails these loads for recompute, and `get_finished` returns at the
    # `self._state_tier is None` guard BEFORE it drains the park ...
    return

while _fail_state_loads is unchanged:

def _fail_state_loads(self, loads) -> None:
    for req_id, _h, _group in loads:
        self._joint_park.settle_state(req_id, False)

So in the no-tier case the park is now deliberately not armed, and then settled — _settle finds need is None and returns. Nothing produces a failed_loading id, even though the comment you wrote one function away says _start_state_loads "fails these loads for recompute". The two fixes are individually right and jointly leave the load unreported: a state-only load parks until reconcile_orphan_load_slots' abandon window, and a joint load is woken by the KV leg alone so _settle_state_load(ok=True) counts a success that never happened.

tests/test_lmcache_offload_connector.py:5187's docstring asserts the opposite of what the code does, so the suite will not catch it either.

New findings at `7e08f5994`

1. `connector_completion` falls through to a `super()` that does not exist

connector.py:601 ends with return super().connector_completion(completion). git grep -n "def connector_completion" atom/ returns exactly two hits — hybrid/kimi_k3/connector.py:589 and hybrid/dsv4/connector.py:2402 — and dsv4 is a sibling, not an ancestor: the MRO is KimiK3OffloadScheduler → DenseOffloadScheduler → OffloadSchedulerMixin → KVConnectorSchedulerBase.

Any completion on a third channel raises AttributeError instead of returning the documented False. That is trivially reachable under kv_connector: multi, where MultiConnector.get_finished unions every sub's connector_completions. It raises inside Scheduler._update_from_kv_xfer_finished, and _offload_common.process_completions is written around callback(completion) is False — it expects a bool, never an exception — so the engine step dies and every TP rank wedges on the next collective.

2. Every joint-boundary disown clears `num_cached_tokens` but leaves the widened claim in the block table

allocate claims claim_blocks = max(num_cached_blocks, state_joint_claim_tokens // hbs) canonical hash-indexed blocks, justified by "those blocks are below the joint boundary the forward will start from, so nobody writes them". Four paths invalidate that premise without trimming seq.block_table:

allocate's if not state_holds: seq.num_cached_tokens = 0 (block_manager.py:946)
_schedule_prefill's [JOINT-DISOWN] branch (scheduler.py:1878)
BlockManager.cancel_state_load (block_manager.py:1121)
_consume_failed_remote_kv (scheduler.py:2237)

A tier miss — which the index itself calls "optimistic", so a normal event — then makes the prefill re-run from token 0 and write KV in place into hot shared prefix blocks another sequence is decoding out of. Not bitwise reproducible on this engine, so a concurrent reader can observe a torn value rather than a clean stale one.

3. Skipping `per_request_state` in the codec removed the fail-fast that kept non-K3 GDN models off the dense path

kv_byte_codec.py:74 now continues past per_request_state tensors. Before this PR, registering a GDN model (Qwen3-Next, MiniMax-M3, Qwen3.5) with lmcache_offload raised ValueError: ... not divisible by num_blocks at register time, because the slot-indexed mamba caches were in the segment list. Registration now succeeds, so the dense path saves and reloads KV for those models.

The replacement rule — "a hybrid may only load KV up to the boundary its recurrent state covers" — exists only in KimiK3OffloadScheduler._decide_load_after_alloc, keyed on has_per_req_cache. git grep -n has_per_req_cache atom/kv_transfer/offload/dense/ has no hits. So a GDN model on the dense layout now gets a KV prefix restored while its linear-attention state is stale, and the forward skips [hbm, lmc) — silent wrong output, on exactly the models the old ValueError protected. Worth flagging because this one reaches beyond K3.

4. `_save_stalled` latches permanently on one lost save report

_save_inflight is popped only by DenseOffloadScheduler.save_finished (dense/connector.py:705). Scheduler._reconcile_stalled_deferred_saves frees the blocks but never touches _save_inflight, so _refresh_save_stall keeps stamping that sid, oldest never advances, and after SAVE_STALL_SECONDS _save_stalled is True for the life of the process.

_may_emit_save() then returns False on every build_connector_meta, whose save loop breaks — no KV chunk is ever written to LMCache again, with the one warning line suppressed by _warned_save_stalled. The same omission keeps has_pending_work() True forever, so EngineCore busy-loops on the 1 ms idle-drain with every GPU idle — which is the failure _reconcile_stalled_deferred_saves's own docstring says it fixes.

5. The save-stall escape in `should_defer_free` bypasses the base's load check

DenseOffloadScheduler.should_defer_free (dense/connector.py:674-676) checks _has_active_load(seq) first and returns True unconditionally, because the deferred free is what keeps the load's destination blocks alive. The K3 override at connector.py:511 evaluates its stall escape before ever calling super(): with _save_stalled true, the sid not in _save_inflight and _has_pending_save(seq) true, it pops _save_tracker[sid] and returns False without asking about the load.

A request aborted or finished while the backend is stalled but with a live _active_load_operations entry has its blocks deallocated while the worker is still writing into them; the pool hands them to another request whose KV is overwritten mid-transfer and indexed under the wrong prefix. The escape's own docstring reasons only about saves.

6. Two `multi` problems the `_state_tier` property does not cover

:441 — _first_with selects on attribute presence, but LMCacheOffloadConnectorScheduler defines all the state methods unconditionally for every layout. With connectors: [lmcache_offload(dense), lmcache_offload(kimi_k3)], every state call goes to the dense shell, whose _impl has none of them: enqueue_state_stores returns False so _publish_state_stores settles every store failed; enqueue_state_loads returns False so every load is abandoned and recomputed; take_state_reports returns two empty sets so the K3 sub's real reports are never drained; and chunk_size returns the dense sub's grid, which _joint_kv_boundary aligns against. _adopt_state_tier's dup-guard cannot catch this — it runs in the worker, keys on _state_tier, and only raises when two tiers were built. Here one was.
:401 — get_num_new_matched_tokens is first-hit-wins in its return value only. It queries every sub, never withdraws the losers, then fans update_state_after_alloc (:415) out to all of them — and the offload lookup is not pure (dense/connector.py:440-465 records _load_specs[sid], appends to _lookup_in_step, takes an LMCache pin). With [moriio, lmcache_offload(kimi_k3)], moriio answers first, but the offload sub still sets ls.can_load = True, enters _reqs_need_recv, and emits a full KV load into the same block table moriio is receiving into: two writers, plus a finished_loading the scheduler never accounted for. The mirror case drops a larger hit — sub A (5, False), sub B (10, True) gives (5, False) while B believes it owns a load. vLLM's MultiConnector breaks on first hit; there is no cancel_pending_load forwarder here.

7. Smaller, but worth a line each

scheduler.py:2628 — _publish_state_loads dereferences self.block_manager bare, while both siblings on the same new path guard it (_publish_state_stores uses getattr(getattr(self, "block_manager", None), ...), and so does _settle_state_load). EngineCore._dispatch_idle_offload_work (:532-538) has a comment saying "this path is reached with scheduler doubles that implement only the connector surface" — so this raises out of the 1 ms idle drain.
block_manager.py:873 — a load-bearing assert in the widened-claim loop. Under python -O, assert i >= num_cached_blocks vanishes; an eviction inside [0, num_cached_blocks) then breaks with hit_hash never assigned, _attach_state_slots(-1) takes the cold-start exit, and the next line still sets seq.num_cached_tokens = num_cached_blocks * hbs. Separately, the intended break writes back only state_joint_claim_tokens — state_joint_boundary_tokens and state_joint_kv_tokens stay stale.
_offload_common.py:291 — StateStoreOperationId lands on finished_saving. It is (prefix_hash, generation) with no req_id, so two raw ops per store (INDEX + SOURCE) get folded in. This is only absorbed because the PR simultaneously replaced assert seq is not None with continue in that loop. MultiConnector.get_finished keys self._saved by completion_req_key(r) with the only pop gated on a matching self._sent[key], which a state-store op never has — so _saved grows unbounded for the life of the process.
block_manager.py:792 — can_allocate runs _joint_kv_boundary on the admission probe. _can_admit_head_prefill (scheduler.py:1256) calls it purely for a bool, and can_allocate can still return -1 at the _has_page_units check afterwards. _chain_to is an xxhash plus an int64 numpy tobytes() per block (~2-3 µs), so a 32k prompt at hash_block_size=64 is ~511 blocks ≈ 1.5 ms of host hashing per probe, paid every tick a long request sits behind a full pool — and allocate re-hashes the same range afterwards. The same call also increments joint_boundaries / state_hbm / state_tier / joint_skips, i.e. the counters this PR exports through checkpoint_funnel, so one queued request can add dozens of phantom counts to the numbers an operator uses to size the pool.
page_unit_checkpoint.py:455 — the requeue in the fix gpt_oss accuracy drop #11 fix goes to the tail. take_offload_stores now defers rather than drops (good), but re-inserts with self._offload_ready.extend(deferred) while the drain uses popleft(). The deferred, older checkpoint now sits behind every nomination queued since, and under sustained prefill it is pushed back until its record is spent and dropped at the record.state != READY check — the exact loss the deferral was added to prevent. checkpoint_fates also documents this queue as "drained oldest-first". extendleft(reversed(deferred)) restores it. Same shape in build_connector_meta: the new if not self._may_emit_save(): break sits before the continue guards and iterates a never-rotated _save_tracker, so at the default cap of 2 the same early requests win every pass while should_defer_free pins the rest.

8. Still open from the last round

_offload_max_pending_saves (scheduler.py:71) is still a third, divergent reader of OFFLOAD_MAX_PENDING_SAVES; the two inspect.getsource(...) assertions are still there; and no .md changed, so the README still states the two-way dense/hybrid selector rule, has no rows for the four new hybrid/kimi_k3/ modules, and documents OFFLOAD_MAX_PENDING_SAVES with a default the new reader contradicts.

On the dead-function references: state_tier._do_spill joins AttentionBackend._submit_state_spills — StagedTransfer's class docstring names it in bold as the producer fence for the state leg and warns "Do not delete either believing this class covers it", and it does not exist anywhere in the tree. Given your #3 rationale (the ordering is producer.synchronize()'s), the docstring is now the thing to fix rather than the code — and yes, deleting the vestigial free_event_valid in this PR would be the cleaner close.

A few below the cut, one line each: CacheStats.update clamps the five nested counters but adds num_offload_tokens unclamped, so lmcache_hit_rate and the [Cache Tiers] combined line can exceed 100%; stores_refused and stores_failed double-count the same refusal (_publish_state_stores bumps refused, then settle_state_store(ok=False) bumps failed); and pool_pressure() was not swept — it still calls _state_checkpoint_cache.checkpoint_fates() directly, which is the per-class omission the new state_checkpoint_fates() aggregator exists to prevent, so two log lines print the same metric names with different values.

`connector_completion` fell through to `super().connector_completion()` for any channel other than STATE_SOURCE / STATE_INDEX. No class in the MRO (`DenseOffloadConnector` and up) defines that method -- only the two sibling hybrid connectors (kimi_k3, dsv4) do -- so the super() call raised AttributeError the moment a completion on an unowned channel arrived, wedging the offload apply path. `_offload_common._apply_connector_completions` already treats a `False` return as "unhandled channel: log and skip", which is exactly the intent here and what the DSV4 sibling returns. Return False. Addresses review finding #1.

`_publish_state_loads` dereferenced `self.block_manager` directly, while its two siblings on the same idle-drain path -- `_settle_state_load` and `_publish_state_stores` -- both go through `getattr(getattr(self, "block_manager", None), <method>, None)` precisely because a scheduler built without a block manager (the connector test doubles) reaches this code. Mirror the sibling guard so the bare deref can no longer AttributeError. Addresses review finding #7 (block_manager deref).

The claim loop guarded the "cached block vanished during allocate" case with a bare `assert i >= num_cached_blocks`. Under `python -O` the assert is stripped and the fall-through was wrong two ways: * `num_cached_tokens` was set to `num_cached_blocks * hbs` (line below the loop) even though only `i < num_cached_blocks` blocks had been claimed -- the forward then resumed over a prefix that was not actually in the pool, and `hit_hash` was never assigned so the state leg took its cold-start exit. * only `state_joint_claim_tokens` was clamped; `state_joint_boundary_*` and `state_joint_kv_tokens` stayed above the claimed region, so the KV leg would load to a boundary the request could not claim and `_claim_after_load` would raise `num_cached_tokens` past it. Both are silent wrong output. Convert to real control flow: clamp `num_cached_blocks` to the blocks actually claimed, and drop the joint boundary (all three joint vars + hash) whenever it sits above the claimed prefix, so the request recomputes rather than resuming over a gap. Correctness no longer depends on assertions being enabled. Addresses review finding #7 (allocate assert).

`get_finished`'s send/save pairing path parks every save completion in `self._saved` keyed by `completion_req_key`, and pops it only when a matching send with the same key drains from `self._sent`. State-tier store completions are `StateStoreOperationId` -- a (prefix_hash, generation) pair with no `req_id` and, crucially, no send counterpart -- so their key never enters `self._sent` and their `_saved` entry was never popped. Over a long run `self._saved` grew without bound. State stores are terminal on their own; they do not pair with a send. Detect them by exact type (`isinstance(r, StateStoreOperationId)`) and release them straight into `finished_saving`, pairing only the request-scoped KV saves. The check must be `isinstance`, not `hasattr(r, "req_id")`: a bare `ReqId` save completion (a plain str/int) also lacks a `req_id` attribute and must still go through send/save pairing on a producer node. Addresses review finding #7 (StateStoreOperationId unbounded _saved).

…cache `pool_pressure()` read `self._state_checkpoint_cache.checkpoint_fates()` directly, while `checkpoint_funnel()` (block_manager.py:1921) reports the same four fate counters through `self.state_checkpoint_fates()` -- the aggregator that sums every state class. The two paths therefore printed different numbers under the same metric names: `pool_pressure` undercounts any second state class the aggregator folds in (and prints four zeros under PAGE, where `self.state` is a `StateTransfer.none()` that never checkpoints), while the funnel line shows the real total. Two outputs disagreeing on "state checkpoints" is exactly how a tuning session concludes checkpointing never fires. Route `pool_pressure` through the aggregator too; occupancy stays on `self.state`, the slot pool, which is unchanged. Ref: review #8-4g.

The deferred-nomination requeue used `self._offload_ready.extend(deferred)`, appending the older, already-waiting nominations to the *right* of the deque. But the drain pops from the left (`popleft()`, oldest-first), so newer nominations queued after them win every drain and the deferred ones starve at the back until they age out at the `state != READY` check. Put them back at the front with `extendleft(reversed(deferred))`, which restores their original oldest-first order ahead of newer arrivals. Ref: review #8-4h (page_unit requeue).

`test_a_late_report_on_a_reclaimed_save_is_not_an_assertion` (test_scheduler.py) and `test_the_runner_warms_the_builder_once_the_pools_are_reachable` (test_v4_checkpoint_slot_copy.py) asserted on `inspect.getsource(...)` substrings and statement ordering rather than on behaviour. They break on any whitespace/refactor that leaves the behaviour intact and pass even when the behaviour regresses, so they test the text, not the code. Remove them. The dynamic scan in test_lmcache_offload_connector.py:5042 is a legitimate capability probe and is kept. Ref: review #8-4b.

…ills Three docstrings referenced methods that no longer exist, and one of them described an architecture that has since changed: * staging.py `StagedTransfer` / `pack` claimed the state producer fence is `state_tier._do_spill`'s `ready_event.synchronize()`. `_do_spill` is gone, and the current state path needs no event at all: `state_tier.submit_store` reserves the PAGE units out of the KV pool and engine-pins them for the whole transfer, so nothing on the compute stream writes them. The KV path is still fenced by the caller's `save_ready_event`. Rewrite to describe both callers accurately. * multi_connector.py (module docstring + __init__) claimed `AttentionBackend._submit_state_spills` reads/probes `_state_tier`. That symbol is dead; `_state_tier` is populated by `_adopt_state_tier` at `register_kv_caches` time (which also refuses a config listing two offload subs). Point the comments at the real mechanism. Docstring-only; no behaviour change. Ref: review #8-4d (docstrings).

…dge the engine `Scheduler._reconcile_stalled_deferred_saves` frees the blocks of a save the backend never reported (LMCache force-unpins a stalled transfer without a completion), but it only touched the scheduler's `deferred_free_blocks` -- the connector still held the save in `_save_inflight`. `save_finished` cannot clear it: handed the raw request id while an exact `SaveOperationId` generation is parked, it deliberately refuses (a delayed TP notification must not complete a newer lifecycle). So `_save_inflight` never drained, `should_defer_free` stayed True and `has_pending_kv_work()` never went False, and the EngineCore busy-looped with every GPU idle -- reproduced under a tight pool on the k3-dev line. Add an `abandon_save(req_id)` hook that drops the entry unconditionally: - `_offload_common`: `_cancel_save_statistics`, mirror of the existing `_cancel_load_statistics` -- the bytes were never persisted, so cancel the inflight-tokens gauge rather than counting a success. - `DenseOffloadConnector.abandon_save`: pop `_save_inflight` (+cancel stats) and `_save_tracker` so the save loop cannot re-emit against freed blocks. - `KimiK3OffloadConnector.abandon_save`: also re-run `_refresh_save_stall` so the stall latch sheds the abandoned sid instead of staying stuck True. - offload wrapper + MultiConnector: forward the call to the owning sub. - `Scheduler._reconcile_stalled_deferred_saves`: notify the connector via a getattr-guarded `_connector_abandon_save` before deallocating. Tests: make the reclaim stub realistic (production always sets `kv_connector`) and add a test asserting the connector is notified with the string request id. Addresses review finding #4 (_save_inflight never cleared on reclaim → busy-loop).

…cape `KimiK3OffloadConnector.should_defer_free` short-circuits to `return False` (release the blocks) when the backend has stalled and the save was never handed out. But it evaluated that escape *before* delegating to the base, whose first act is `if self._has_active_load(seq): return True`. A request whose save has stalled can still have a live load into the same block table, so the escape freed blocks while a load was mid-transfer -- the next request writes into them and the load's result is indexed under the wrong prefix (free-while-writing corruption, matching the offload-only, concurrency-gated garbage-output fingerprint). Hoist `if self._has_active_load(seq): return True` ahead of the stall escape. The predicate is identical to the base's first check, so requests without a live load are unaffected; only a stalled-save-with-active-load request changes -- its blocks now stay deferred until the load finishes. Addresses review finding #5 (save-stall escape bypasses the active-load guard).

…om success When the worker state tier never built, a joint load (KV + recurrent state) was left unreported. `_arm_joint_loads` skipped the park on `_state_tier is None`, so `get_finished`'s `_state_tier is None` early-return let the KV leg pass straight through as `finished_loading`. `Scheduler._settle_state_load` (scheduler.py:3503) then counted that as `ok=True` -- a state restore that never happened -- and the request resumed decode reading an Active Slot the state H2D never wrote: silent wrong output. The maintainer reopened the #5 rebuttal on exactly this: the skip was the guard I added *for* #4 (avoid arming under a `kv_id` the state leg could never settle, which would leak a park entry). The two fixes were individually right and jointly left the load unreported. Fix, in cooperation with #4 (not a revert of it): - `_arm_joint_loads` arms a joint load with or without a tier. `_start_state_loads` already fails the state leg on the no-tier path (`_fail_state_loads` -> `settle_state(ok=False)`); with the arm in place that failure lands on a real park entry, so the pair is marked failed and owes only the KV leg. - `get_finished` drains the park via `_settle_joint` (empty state reports) on the no-tier path too, before its early return. The landed KV completion settles the pair into `failed_loading` -> the request recomputes. This does not reintroduce #4's leak: the arm is safe precisely because the drain runs -- the KV completion always arrives (finished or failed) and releases the park entry. A state-only load (no KV leg) is still not armed and is reclaimed by `reconcile_orphan_load_slots`' abandon window, as before. Updates the stale unit test (no-tier is now a fail-for-recompute, not a KV-only passthrough) and its helper docstring.

…he dense path Review finding #3. The dense byte codec's per-request-state branch was changed from a hard ValueError to a silent `continue`. That skip removes the only guard that keeps a GDN / linear-attention model (Qwen3-Next, MiniMax-M3, Qwen3.5) off the plain dense offload path. Those models register their slot-indexed recurrent (mamba) state in the same kv_caches dict the codec walks. The dense path has no rule aligning a restored KV prefix with that recurrent state -- nothing like KimiK3OffloadScheduler's per-request-state load decision exists for it. So a `continue` lets registration succeed, and at load time the KV prefix `[hbm, lmc)` is restored while the linear state stays stale and the forward skips recomputing it: silent wrong output. Pre -PR, those models were kept off the path by the divisibility ValueError the skip now also swallows. Fix: gate the skip behind a keyword-only `permit_per_request_state` flag on DenseKVByteCodec, default False -> restore the fail-fast ValueError. The dense connector passes its class attr `_permit_per_request_state = False`; KimiK3OffloadConnector overrides it to True because it owns a state tier that moves the recurrent state separately, so for it the skip is correct. DSV4 has its own register_kv_caches / codec and is unaffected. tests: add test_dense_codec_rejects_per_request_state_by_default (raises by default, skips + excludes the state segment when permitted). 212 passed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…ared fairly Review finding #8-4h. DenseOffloadScheduler.build_connector_meta scans _save_tracker in fixed insertion order and breaks as soon as _may_emit_save says the outstanding-save budget is full. The base _may_emit_save is unbounded, but KimiK3OffloadScheduler overrides it to cap saves at OFFLOAD_MAX_PENDING_SAVES (default max(2, 2*save_workers)). Under that cap the scan restarts at the same head every step. A long, multi-chunk request at the head produces a fresh chunk each step and re-wins the freed slot every time, so later requests in the tracker never get a save dispatched -- and their finished blocks stay pinned by should_defer_free until a save drains, which for them never happens. Starvation, worse the more concurrent prefills share the cap. Fix: keep a round-robin cursor `_save_rr_last` (the last sid that emitted a save) and resume the scan just after it, wrapping around. When the budget is unbounded (dense/dsv4 base) every eligible save still emits in the same step, so behaviour there is unchanged; only the bounded kimi_k3 path changes, and it now serves requests in rotation. tests: add test_bounded_saves_are_shared_round_robin_not_head_first (cap=1, three mid-prefill requests -> emits [100,101,102] in rotation, not [100,100,100]). 213 passed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

… decode corruption Finding #2 (review 5479206239). A widened joint claim shares the matched prefix by `BlockPool.claim` (ref_count += 1, not a copy), so several live sequences point their block_table at the same canonical blocks. When a boundary later turns out to have no state behind it, the request is "disowned": num_cached_tokens is set to 0 so the forward recomputes from token 0 -- but the block_table still points at those shared canonical blocks. The recompute writes KV in place into blocks another sequence is actively decoding out of, tearing its values. This is non-reproducible, grows with N and concurrency, and is the main suspect behind the pass2 precision regression (offload-specific, byte-level codecs already cleared). Four disown sites only zeroed num_cached_tokens and never re-pointed the block_table: - block_manager.allocate (state-less joint boundary) - block_manager.cancel_state_load (withdrawn load, same boundary) - scheduler _schedule_prefills (needs-remote-load / joint-disown) - scheduler _consume_failed_remote_kv (offload-resume reuses the table) Fix: add BlockManager.disown_claimed_prefix(seq, keep_blocks=0). For each claimed block above keep_blocks that is still shared (hash != -1), kv.free it and drop a fresh private _fresh_block into the same slot -- table length and every other mapping unchanged, the already-private fresh tail (hash==-1) untouched. Only ref_count>1 blocks cost a net PAGE unit (free releases nothing while a peer holds them, and can_allocate never reserved them), so reserve n_shared units up front via _ensure_page_units; return False if the pool cannot back them. Thread that False through: allocate and cancel_state_load now return bool, and all callers (both allocate sites, both cancel_state_load sites, and _consume_failed_remote_kv) deallocate + requeue the seq for a clean recompute rather than run a forward over shared blocks. keep_blocks is 0 at every site because hybrid disown forces recompute from 0 (has_initial_state = num_cached_tokens > 0). Tests: TestDisownClaimedPrefix in test_block_manager.py -- privatises shared blocks in place (length/tail unchanged, slots now hash==-1 and distinct from the peer's, peer back to ref_count==1), no leak after deallocate, refuses (returns False, table untouched) when the pool cannot back the copies, and is a no-op without prefix caching. Full suite green (356 passed).

CacheStats.update() clamps the ordered cache-hit counters (cached <= wanted <= compressed <= reusable <= full) inside the `if not ordered:` block, but num_offload_tokens is assigned afterwards and never bounded. The class docstring (line 463) states the invariant `cached + offload <= reusable`: offload reuse and HBM-cache reuse partition the same served-reuse pool, so their sum cannot exceed it. When a connector reports more offload tokens than the residual (reusable - cached) -- e.g. a stale lookup or an over-counted load -- lmcache_hit_rate = offload/reusable and the combined hit_rate can exceed 100%, corrupting operator-facing dashboards. Clamp num_offload_tokens to [0, num_reusable_tokens - num_cached_tokens] so it always fits the residual reuse budget the docstring promises.

_publish_state_stores throttled the publish batch with the scheduler's own _offload_max_pending_saves() (scheduler.py:71), which only reads the OFFLOAD_MAX_PENDING_SAVES env default "2". That ignores both the canonical default max(2, 2*save_workers) computed in _offload_common.max_pending_saves and any per-connector kv_connector_extra_config["max_pending_saves"] override the connector already stored on self._max_pending_saves. The result is a third, inconsistent reader of the same knob: with more save workers the connector self-throttles at 2*save_workers while the scheduler publishes only 2 per drain, needlessly serializing state stores. Add _state_store_pending_cap() that prefers the live connector's _max_pending_saves (through _impl for wrapped connectors), and falls back to _offload_max_pending_saves() only when no connector value is available. Route the _publish_state_stores take() through it so the publish cap tracks the same value the connector enforces.

_publish_state_stores bumps state_offload.stores_refused by the whole batch length when a refusal short-circuits publishing, then still walks each op calling settle_state_store(op, ok=False). settle_state_store's else-branch bumps state_offload.stores_failed += 1 per op, so every refused store lands in BOTH counters -- refused and failed -- inflating the failure rate by exactly the refused count. A refusal is not an attempted-then-failed store: the unit must be released, but the op never reached the offload path. Add an attempted=True parameter to settle_state_store and gate the stores_failed bump behind `elif attempted:`; the refusal caller passes attempted=False so it releases the reserved unit without recording a phantom failure.

can_allocate() runs _joint_kv_boundary() purely to test whether a sequence fits, but that helper bumps the checkpoint funnel counters (joint_boundaries, state_hbm, state_tier, joint_skips) on every call. The scheduler probes can_allocate() speculatively (scheduler.py:1266 schedules-ahead without committing), so each probe pollutes the funnel that checkpoint_funnel exports to operators -- the counts no longer correspond to real joint-claim decisions. Thread a record flag through _joint_kv_boundary()/_no_joint() (default True, so allocate() and every existing caller keep recording). Guard the four counter bumps behind `if record:`. can_allocate() gains a record parameter it forwards, and the speculative probe call passes record=False so only the committing allocate() path advances the funnel.

The offload README still described a two-way dense/hybrid selector and had no module-map rows for the four new hybrid/kimi_k3/ files, even though select_offload_layout() (config.py:100) resolves a third layout: kimi_k3 when hf_config.model_type == "kimi_linear". - Rewrite the selector paragraph to name all three layouts (kimi_k3 / hybrid / dense) and the condition that picks each. - Update the connector.py module-map row to dense/hybrid/kimi_k3. - Add rows for hybrid/kimi_k3/{connector,staging,state_object,state_tier}.py. The OFFLOAD_MAX_PENDING_SAVES row (line 856) already documents the canonical max(2, 2 x OFFLOAD_COPY_WORKERS) default; with the #8-4a reader unification the scheduler now honors that value, so doc and code agree.

StagedTransfer._finish records free_event and sets free_event_valid=True before calling producer.synchronize(). The synchronize() is the whole fence -- by the time _finish returns, the producer stream has drained and the buffer is free -- and StagedTransfer never reads free_event or free_event_valid anywhere (its only wait_event is on ready_event, in _handoff). So the free-fence it publishes has no consumer: the record, the True write, and the four defensive `= False` resets in ensure_buffer / release_buffer_if_requested / the pack and unpack except paths are all dead within this class. The maintainer flagged free_event_valid as vestigial. It is NOT globally dead, though: _StagingBuffer.free_event_valid is a live stream-reuse fence in run_staged_pipeline (atom_lmcache_staging.py:123), which waits on free_event before refilling the bounded buffer. StagedTransfer imports that shared _StagingBuffer class but drives its own synchronous pipeline and uses a disjoint per-thread buffer set, so it never reaches that reader. Remove only StagedTransfer's writes; keep the field and its live reader. _finish's docstring now states it is a synchronous fence that deliberately leaves the shared buffer's free_event/free_event_valid untouched. Tests that assert free_event_valid resets all exercise the BlockGPUConnector / run_staged_pipeline path and are unaffected.

…resence Review #6a. Under `kv_connector: multi` the composite selected the sub-connector for every state call with `_first_with`, which picks the first sub that *has* the attribute. That is the wrong test for the state tier: every `LMCacheOffloadConnectorScheduler` shell -- dense, hybrid, and kimi_k3 alike -- defines the whole state face (`enqueue_state_stores`, `enqueue_state_loads`, `take_state_reports`, `take_state_source_releases`) unconditionally, delegating through `getattr` to its `_impl` and returning the no-tier default when the impl has none. So with a `[dense_offload, kimi_k3_offload]` (or moriio-first) layout every sub answers `hasattr`, and `_first_with` routes state stores, loads and reports to a dense shell whose `_impl` silently refuses the stores (recording them failed), discards the loads (forcing recompute), drains no reports, and reports the wrong `chunk_size` -- while the sub that actually built the tier never sees them. The CPU/HBM state tier cannot fill or be found. Give the shell a `has_state_tier` property that probes `_impl` for the store entry point (`enqueue_state_stores` is defined only on `KimiK3OffloadScheduler`, so its presence on the impl is the true tier discriminator), and select the state sub in the composite with a new `_state_tier_sub()` that keys on that property. Route `chunk_size` and the four state forwarders through it; `chunk_size` keeps its `_first_with` fallback for non-state offload layouts. `_first_with` still drives the KV-only members (park/prefill-chunk hooks). tests/test_multi_connector.py: teach the mock about `has_state_tier` and add a regression that puts a tier-less offload shell (methods present) first and the tier-carrying shell second, asserting every state call lands on the second.

…cellation Review #6b. The composite's `get_num_new_matched_tokens` is first-hit-wins in its *return value*, but it queries every sub and an offload sub's lookup is not side-effect-free: matching a prefix arms a KV load (LMCache lookup pin plus a `_load_specs`/`_lookup_in_step` entry). Because `update_state_after_alloc` fans to every sub, a losing offload sub's armed load is later flipped to `can_load=True` and recv-queued -- firing into the same block table the winner already owns. Under `[moriio, lmcache_offload]` moriio answers first while the offload tier still arms a second write over moriio's KV, plus a `finished_loading` the scheduler never accounted; under two offload subs the dense shell can win while kimi_k3 arms behind it. Once a winner is chosen, cancel every other sub's pending load. `cancel_pending_load` is idempotent (guarded by `_load_lifecycles`) and only offload subs define it, so a sub that armed nothing -- a miss, or moriio -- is a no-op. The querying pass over all subs is unchanged, so every sub still sees every request; only the harmful armed-load side effect is undone. Also add the `cancel_pending_load` forwarder the composite was missing entirely: the scheduler abandons a parked load by calling it on the connector it holds, which under `multi` is the composite. With no forwarder the offload sub never heard the cancel and leaked its `_load_specs`/`_reqs_need_recv`/lookup pin for the abandoned request. Fan to every sub that arms loads. tests/test_multi_connector.py: mock now arms a load on a matching offload lookup and clears it on cancel; add regressions for loser-cancelled, winner-kept, and the standalone cancel forwarder.

Brings the branch up to origin/main (tip 171de45) ahead of PR #2053 re-review. Only three conflict hunks, all in the model_engine stats path: - block_manager.py: kept HEAD's `_joint_kv_boundary(record=...)` probe call (#7 fix), adopted main's EngineStats comment wording. - scheduler.py: main's #1770 refactor moved SpecStats+CacheStats out of scheduler.py into the new engine_stats.py as a unified `EngineStats` (self.cache_stats -> self.engine_stats, update -> update_cache, hit_rate -> cache_hit_rate, etc.). Dropped the branch's now-relocated stats classes, kept the branch's offload helpers and the stats caller (which passes num_offload_tokens). Ported the branch's offload-token tracking feature (absent from main's EngineStats) onto engine_stats.py EngineStats so the relocated class keeps it: num_offload_tokens param on update_cache with the cached+offload<=reusable clamp, total/interval offload counters, lmcache_hit_rate property, the [Cache Stats offload] / [Cache Tiers] log lines, and the cache_statistics entry. Moved the branch's TierSplit tests onto EngineStats.update_cache. Verified: offload suite 662 passed; full suite 4913 passed (pre-existing env-only failures unrelated to the merge: 4 fused_compress_ragged GPU-kernel numerics, byte-identical kernel+test across both sides; 4 sglang-plugin collection errors from a missing optional dep). black + ruff clean on all changed files. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

github-actions · 2026-08-31T16:15:00Z

+        # 6 blocks: holder takes 4, seq's fresh tail takes 1, leaving 1 free --
+        # fewer than the 4 private copies the disown needs, so it must refuse
+        # rather than silently reuse the shared blocks.
+        bm, holder, seq, shared = self._shared_prefix(seq_factory, num_kvcache_blocks=6)


⚠️ [ruff] <RUF059> _{reported by reviewdog 🐶}
Unpacked variable holder is never used

Suggested change

bm, holder, seq, shared = self._shared_prefix(seq_factory, num_kvcache_blocks=6)

bm, _holder, seq, shared = self._shared_prefix(seq_factory, num_kvcache_blocks=6)

Copilot

Pull request overview

Copilot reviewed 42 out of 42 changed files in this pull request and generated 1 comment.

 class TestWarmup:
    """That the checkpoint copy path is paid for before a request can pay it."""

    def test_a_backend_without_paged_checkpoints_warms_nothing(self):
        """The hook is on the base builder, so every backend answers it."""
        from atom.model_ops.attentions.backends import AttentionMetadataBuilder

        # No-op by contract: a backend that never copies has nothing to warm,
        # and ModelRunner calls this unconditionally.
        assert AttentionMetadataBuilder.warmup_per_req_cache(object()) is None


zejunchen-zejun mentioned this pull request Aug 27, 2026

[Kimi-K3][LMCache] LMCache offload for Kimi-K3 #1960

Closed

zejunchen-zejun force-pushed the zejun/lmcache_with_2045 branch from 1fe7ce5 to 76a10fa Compare August 28, 2026 07:53

zejunchen-zejun marked this pull request as ready for review August 28, 2026 09:14

Copilot AI lite review requested due to automatic review settings August 28, 2026 09:14

Copilot started reviewing on behalf of zejunchen-zejun August 28, 2026 09:15 View session

zejunchen-zejun requested review from ganyi1996ppo, gbyu-amd, valarLip and yhl-amd August 28, 2026 09:19

Copilot AI reviewed Aug 28, 2026

View reviewed changes

zejunchen-zejun and others added 17 commits August 28, 2026 18:46

Copilot AI reviewed Aug 31, 2026

View reviewed changes

zejchen and others added 23 commits August 31, 2026 14:26

Copilot AI review requested due to automatic review settings August 31, 2026 16:14

github-actions Bot reviewed Aug 31, 2026

View reviewed changes

Copilot started reviewing on behalf of zejunchen-zejun August 31, 2026 16:15 View session

Copilot AI reviewed Aug 31, 2026

View reviewed changes

Navigation Menu

Uh oh!

[Kimi-K3][LMCache] LMCache offload for Kimi-K3 - #2053

[Kimi-K3][LMCache] LMCache offload for Kimi-K3#2053
zejunchen-zejun wants to merge 65 commits into
mainfrom
zejun/lmcache_with_2045

zejunchen-zejun commented Aug 27, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Aug 27, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

zejunchen-zejun commented Aug 31, 2026

Uh oh!

valarLip commented Aug 31, 2026

Uh oh!

github-actions Bot Aug 31, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		\| Variable \| Type \| Default \| Description \|
		\|----------\|------\|---------\|-------------\|

	bm, holder, seq, shared = self._shared_prefix(seq_factory, num_kvcache_blocks=6)
	bm, _holder, seq, shared = self._shared_prefix(seq_factory, num_kvcache_blocks=6)

Uh oh!

Conversation

zejunchen-zejun commented Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

[Kimi-K3][LMCache] LMCache offload for Kimi-K3

Design

Why the boundary has to be joint

Example

Notes for reviewers

Verification

Uh oh!

github-actions Bot commented Aug 27, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

zejunchen-zejun commented Aug 31, 2026

Fixed (10)

Not fixed — assessed, no code change (rationale)

Observability

Uh oh!

valarLip commented Aug 31, 2026

The #5 rebuttal needs reopening — and it is your #4 fix that made it certain

New findings at 7e08f5994

1. connector_completion falls through to a super() that does not exist

2. Every joint-boundary disown clears num_cached_tokens but leaves the widened claim in the block table

3. Skipping per_request_state in the codec removed the fail-fast that kept non-K3 GDN models off the dense path

4. _save_stalled latches permanently on one lost save report

5. The save-stall escape in should_defer_free bypasses the base's load check

6. Two multi problems the _state_tier property does not cover

7. Smaller, but worth a line each

8. Still open from the last round

Uh oh!

github-actions Bot Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

zejunchen-zejun commented Aug 27, 2026 •

edited

Loading

New findings at `7e08f5994`

1. `connector_completion` falls through to a `super()` that does not exist

2. Every joint-boundary disown clears `num_cached_tokens` but leaves the widened claim in the block table

3. Skipping `per_request_state` in the codec removed the fail-fast that kept non-K3 GDN models off the dense path

4. `_save_stalled` latches permanently on one lost save report

5. The save-stall escape in `should_defer_free` bypasses the base's load check

6. Two `multi` problems the `_state_tier` property does not cover