feat(lmcache): add DSV4 page and slot offload - #1880
Conversation
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit
ruff
Import block is un-sorted or un-formatted
ATOM/tests/test_deepseek_v4_transfer_regions.py
Lines 4 to 20 in e4c43e6
Do not assert blind exception: Exception
ATOM/tests/test_kv_aggregator.py
Line 396 in e4c43e6
Import block is un-sorted or un-formatted
ATOM/tests/test_lmcache_offload_config.py
Lines 4 to 14 in e4c43e6
Import block is un-sorted or un-formatted
ATOM/tests/test_lmcache_offload_connector.py
Lines 21 to 55 in e4c43e6
Prefer dict over useless lambda
ATOM/tests/test_lmcache_offload_connector.py
Line 306 in e4c43e6
Unused noqa directive (non-enabled: E402)
Unused noqa directive (non-enabled: E402)
Import block is un-sorted or un-formatted
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Lines 4 to 49 in e4c43e6
Unnecessary dict() call (rewrite as a literal)
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Lines 622 to 632 in e4c43e6
Do not explicitly return None in function if it is the only possible return value
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Line 1644 in e4c43e6
Useless return statement at end of function
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Line 1644 in e4c43e6
Import block is un-sorted or un-formatted
ATOM/tests/test_multi_connector.py
Lines 12 to 30 in e4c43e6
Do not catch blind exception: Exception
Line 174 in e4c43e6
Import block is un-sorted or un-formatted
ATOM/tests/test_slot_sidecar_codec.py
Lines 4 to 16 in e4c43e6
Import block is un-sorted or un-formatted
ATOM/tests/test_slot_sidecar_format.py
Lines 4 to 25 in e4c43e6
Import block is un-sorted or un-formatted
ATOM/tests/test_slot_sidecar_store.py
Lines 4 to 11 in e4c43e6
Unused noqa directive (non-enabled: E402)
ATOM/tests/test_slot_sidecar_store.py
Line 15 in e4c43e6
Unused noqa directive (non-enabled: E402)
ATOM/tests/test_slot_sidecar_store.py
Line 20 in e4c43e6
Prefer TypeError exception for invalid type
ATOM/tests/test_slot_sidecar_store.py
Line 49 in e4c43e6
There was a problem hiding this comment.
Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit
ruff
try-except-pass detected, consider logging the exception
ATOM/atom/kv_transfer/offload/hybrid/connector.py
Lines 624 to 625 in b5f5a06
Do not catch blind exception: Exception
Do not catch blind exception: Exception
Do not catch blind exception: Exception
Do not catch blind exception: Exception
Do not catch blind exception: Exception
ATOM/atom/kv_transfer/offload/hybrid/connector.py
Line 1420 in b5f5a06
Do not catch blind exception: Exception
ATOM/atom/kv_transfer/offload/hybrid/connector.py
Line 1597 in b5f5a06
Do not catch blind exception: Exception
ATOM/atom/kv_transfer/offload/hybrid/connector.py
Line 1610 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/atom/kv_transfer/offload/hybrid/page_region_codec.py
Lines 6 to 16 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/atom/kv_transfer/offload/hybrid/policy.py
Lines 6 to 12 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/atom/kv_transfer/offload/hybrid/sidecar_format.py
Lines 6 to 14 in b5f5a06
Prefer TypeError exception for invalid type
Import block is un-sorted or un-formatted
ATOM/atom/kv_transfer/offload/hybrid/slot_codec.py
Lines 6 to 16 in b5f5a06
Prefer TypeError exception for invalid type
Import block is un-sorted or un-formatted
ATOM/atom/kv_transfer/offload/hybrid/store.py
Lines 6 to 14 in b5f5a06
Prefer TypeError exception for invalid type
Unused noqa directive (unused: BLE001)
ATOM/atom/kv_transfer/offload/hybrid/store.py
Line 176 in b5f5a06
Unused noqa directive (unused: BLE001)
ATOM/atom/kv_transfer/offload/hybrid/store.py
Line 252 in b5f5a06
Unused noqa directive (unused: BLE001)
ATOM/atom/kv_transfer/offload/hybrid/store.py
Line 317 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/tests/test_deepseek_v4_transfer_regions.py
Lines 4 to 20 in b5f5a06
Do not assert blind exception: Exception
ATOM/tests/test_kv_aggregator.py
Line 440 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/tests/test_lmcache_offload_config.py
Lines 4 to 14 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/tests/test_lmcache_offload_connector.py
Lines 21 to 60 in b5f5a06
Prefer dict over useless lambda
ATOM/tests/test_lmcache_offload_connector.py
Line 329 in b5f5a06
Prefer list over useless lambda
ATOM/tests/test_lmcache_offload_connector.py
Line 1626 in b5f5a06
Prefer list over useless lambda
ATOM/tests/test_lmcache_offload_connector.py
Line 1664 in b5f5a06
Unused noqa directive (non-enabled: E402)
Import block is un-sorted or un-formatted
ATOM/tests/test_lmcache_offload_disk_integration.py
Lines 32 to 39 in b5f5a06
Unused noqa directive (non-enabled: E402)
Import block is un-sorted or un-formatted
ATOM/tests/test_lmcache_offload_gpu_disk_e2e.py
Lines 56 to 72 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Lines 4 to 54 in b5f5a06
Unnecessary dict() call (rewrite as a literal)
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Lines 638 to 650 in b5f5a06
Do not explicitly return None in function if it is the only possible return value
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Line 1666 in b5f5a06
Useless return statement at end of function
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Line 1666 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/tests/test_multi_connector.py
Lines 12 to 31 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/tests/test_offload_profiles.py
Lines 4 to 12 in b5f5a06
Do not catch blind exception: Exception
Line 249 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/tests/test_slot_sidecar_codec.py
Lines 4 to 16 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/tests/test_slot_sidecar_format.py
Lines 4 to 25 in b5f5a06
Import block is un-sorted or un-formatted
ATOM/tests/test_slot_sidecar_store.py
Lines 4 to 11 in b5f5a06
Unused noqa directive (non-enabled: E402)
ATOM/tests/test_slot_sidecar_store.py
Line 15 in b5f5a06
Unused noqa directive (non-enabled: E402)
ATOM/tests/test_slot_sidecar_store.py
Line 20 in b5f5a06
Prefer TypeError exception for invalid type
ATOM/tests/test_slot_sidecar_store.py
Line 50 in b5f5a06
Use a single with statement with multiple contexts instead of nested with statements
ATOM/tests/test_slot_sidecar_store.py
Lines 521 to 522 in b5f5a06
There was a problem hiding this comment.
Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit
ruff
Prefer list over useless lambda
ATOM/tests/test_lmcache_offload_connector.py
Line 1723 in 4abf370
Unused noqa directive (non-enabled: E402)
Import block is un-sorted or un-formatted
ATOM/tests/test_lmcache_offload_disk_integration.py
Lines 32 to 39 in 4abf370
Unused noqa directive (non-enabled: E402)
Import block is un-sorted or un-formatted
ATOM/tests/test_lmcache_offload_gpu_disk_e2e.py
Lines 56 to 72 in 4abf370
Import block is un-sorted or un-formatted
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Lines 4 to 55 in 4abf370
Unnecessary dict() call (rewrite as a literal)
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Lines 656 to 668 in 4abf370
Do not explicitly return None in function if it is the only possible return value
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Line 1901 in 4abf370
Useless return statement at end of function
ATOM/tests/test_lmcache_offload_v4_page_slot.py
Line 1901 in 4abf370
Import block is un-sorted or un-formatted
ATOM/tests/test_multi_connector.py
Lines 12 to 31 in 4abf370
Import block is un-sorted or un-formatted
ATOM/tests/test_offload_profiles.py
Lines 4 to 12 in 4abf370
Do not catch blind exception: Exception
Line 285 in 4abf370
Import block is un-sorted or un-formatted
ATOM/tests/test_slot_sidecar_codec.py
Lines 4 to 16 in 4abf370
Import block is un-sorted or un-formatted
ATOM/tests/test_slot_sidecar_format.py
Lines 4 to 25 in 4abf370
Import block is un-sorted or un-formatted
ATOM/tests/test_slot_sidecar_store.py
Lines 4 to 11 in 4abf370
Unused noqa directive (non-enabled: E402)
ATOM/tests/test_slot_sidecar_store.py
Line 15 in 4abf370
Unused noqa directive (non-enabled: E402)
ATOM/tests/test_slot_sidecar_store.py
Line 20 in 4abf370
Prefer TypeError exception for invalid type
ATOM/tests/test_slot_sidecar_store.py
Line 50 in 4abf370
Use a single with statement with multiple contexts instead of nested with statements
ATOM/tests/test_slot_sidecar_store.py
Lines 521 to 522 in 4abf370
Route page-major compressed KV through LMCache and persist full per-request state as a fail-closed sidecar so hybrid prefixes resume safely. Co-authored-by: Cursor <cursoragent@cursor.com>
Restructure lmcache_offload around a thin public shell, shared block/staging primitives, and dense/hybrid family implementations while preserving DSV4 PAGE+SLOT semantics. Harden exact transfer generations, deferred idle progress, namespace isolation, and corrupt sidecar recovery; add focused regression coverage.
|
Real DSV4-Pro MTP + LMCache CPU E2E update Configuration:
Results:
Caveat: functional correctness passes, but strict token parity does not. The cold response and offload-hit response use slightly different wording while both extract to 18. This was reproduced after a full-shape warmup and with the experimental deferred MTP proposal path both enabled and disabled, so it is not a first-JIT or deferred-proposal-only effect. If acceptance requires byte-for-byte generation equality, that item remains open. The real MTP startup also exposed a PAGE namespace serialization failure because speculative_config contains torch.dtype. I fixed it in validation/pr1880 at 507b692 (133 targeted tests passed). That fix is not yet on this PR head. |
e31c121 to
507b692
Compare
|
Update: the rebased validation commit chain has now been pushed to this PR head. feature/dsv4-lmcache-page-slot now points to 507b692 and includes both the unified PAGE/SLOT adaptation and the torch.dtype PAGE-namespace startup fix described above. CI has been retriggered. |
… a -1 interval Rebased onto main's PAGE-backed checkpoint work (#1874, #1894, #1943, #1880), which rewrote this subsystem underneath the branch. Squashed to one commit because the three original commits each re-conflicted against the new base and against each other's resolutions; the reasoning from all three is kept below. --- allocate state slots per need, not by group A "state cache group" was `1 + num_spec` slots wide and was the unit of everything: allocation, admission, sizing, and the checkpoint index. But a checkpoint has no speculation to roll back -- it holds a committed state -- so filing one cost a full group and wasted `num_spec/(1 + num_spec)` of its bytes. At two speculative tokens that is two thirds. The slot is now the unit. `--state-checkpoint-slots 64` buys 64 checkpoints for 64 slots instead of 192, and the slots it no longer takes stay in the paged KV pool. This is only possible because a request's slots need not be adjacent, which the kernels never required: the ssm kernel gathers each index out of the indices tensor and the conv path is handed column 0 alone. Contiguity was manufactured by `prepare_state_indices` writing `arange(base, base + width)`; it now writes the seq's own slot list straight in. `StateGroupPool` -> `StateSlotPool`, `Sequence.per_req_cache_group` -> `state_slots`, whose element 0 is the committed state. The setter re-points [0] and preserves [1:], because speculation scratch persists across forwards. `--state-checkpoint-groups` still parses, as an alias. --- -1 turns off the interval ladder without turning off checkpointing The interval is a guess about where reuse will resume; a demand rung is a position a request was actually refused at. On the SemiAnalysis cc-traces the 8192 ladder placed ~30x the writes of the demand rung alone and caught reuse the demand already reaches -- 0.0% of resumes landed on a ladder rung -- while every rung costs the prompt that keeps it an extra prefill chunk. >0 a rung every N tokens (unchanged, still the default) 0 state checkpointing off entirely (unchanged) -1 no interval rungs; the demand rung and prompt-end anchor still place them -1 rather than reusing 0 because 0 is the documented contract and is reachable by accident: the grid snap rounds an off-grid interval down and can land on 0, so a --block-size typo currently fails safe. Three of the four sites are not the arithmetic you would guess -- `pos % interval` under -1 admits *every* position rather than none, and `pos - last < -1` is true for every pos. --- make the demand rung switchable A demand rung is 47% of checkpoint writes on the cc-traces and reads back 2.8% of the time, against 85.2% for a prompt-end anchor. Gated independently of --state-checkpoint-interval-tokens, because the demand is not part of the interval grid. Default unchanged. The refusal is still measured when the placement is off -- switching off a rung must not blind the diagnostic that justifies it. --- reconciliation with main Dropped: the copy/pending-checkpoint path (`_commit_pending`, `record_copy`, `take_copies`, `Sequence.pending_checkpoint`). DeepSeek-V4 checkpointing moved to `PagedStateCheckpointCoordinator`, and `StateSlotPool` now rejects `transfer.copies` outright. The fork path (GDN), where the measured wins are, is kept in full. `readable_midstep` is carried on main's `StateTransfer` in `state_runtime.py` (wire format included) rather than on the branch's copy. Two bugs this rebase exposed, both fixed here: - `PagedStateCheckpointCoordinator` did not implement the midstep half of the `StateCache` protocol, so `checkpoint_cut` raised on every V4 batch. A PAGE image is not readable midstep, so the three methods are the no-ops the protocol documents. - `_record_checkpoint_end` could place the anchor past the last matchable block. `can_allocate` stops one block short of the prompt, so a checkpoint filed under the final block's hash is one no scan looks up -- and being stored, it evicted the ladder rung that would have served the resume, taking an identical re-request from 8 hit blocks to 0. Capped at `(n_hash_blocks - 1) * hbs`. Tests: `tests/test_state_checkpoint.py` 170 passed. Full suite 45 failed / 2300 passed, a strict subset of origin/main's own 114 pre-existing failures -- zero regressions, verified by set difference against a clean origin/main worktree. black clean; ruff no new findings. Co-Authored-By: Claude <noreply@anthropic.com>
… a -1 interval Rebased onto main's PAGE-backed checkpoint work (#1874, #1894, #1943, #1880), which rewrote this subsystem underneath the branch. Squashed to one commit because the three original commits each re-conflicted against the new base and against each other's resolutions; the reasoning from all three is kept below. --- allocate state slots per need, not by group A "state cache group" was `1 + num_spec` slots wide and was the unit of everything: allocation, admission, sizing, and the checkpoint index. But a checkpoint has no speculation to roll back -- it holds a committed state -- so filing one cost a full group and wasted `num_spec/(1 + num_spec)` of its bytes. At two speculative tokens that is two thirds. The slot is now the unit. `--state-checkpoint-slots 64` buys 64 checkpoints for 64 slots instead of 192, and the slots it no longer takes stay in the paged KV pool. This is only possible because a request's slots need not be adjacent, which the kernels never required: the ssm kernel gathers each index out of the indices tensor and the conv path is handed column 0 alone. Contiguity was manufactured by `prepare_state_indices` writing `arange(base, base + width)`; it now writes the seq's own slot list straight in. `StateGroupPool` -> `StateSlotPool`, `Sequence.per_req_cache_group` -> `state_slots`, whose element 0 is the committed state. The setter re-points [0] and preserves [1:], because speculation scratch persists across forwards. `--state-checkpoint-groups` still parses, as an alias. --- -1 turns off the interval ladder without turning off checkpointing The interval is a guess about where reuse will resume; a demand rung is a position a request was actually refused at. On the SemiAnalysis cc-traces the 8192 ladder placed ~30x the writes of the demand rung alone and caught reuse the demand already reaches -- 0.0% of resumes landed on a ladder rung -- while every rung costs the prompt that keeps it an extra prefill chunk. >0 a rung every N tokens (unchanged, still the default) 0 state checkpointing off entirely (unchanged) -1 no interval rungs; the demand rung and prompt-end anchor still place them -1 rather than reusing 0 because 0 is the documented contract and is reachable by accident: the grid snap rounds an off-grid interval down and can land on 0, so a --block-size typo currently fails safe. Three of the four sites are not the arithmetic you would guess -- `pos % interval` under -1 admits *every* position rather than none, and `pos - last < -1` is true for every pos. --- make the demand rung switchable A demand rung is 47% of checkpoint writes on the cc-traces and reads back 2.8% of the time, against 85.2% for a prompt-end anchor. Gated independently of --state-checkpoint-interval-tokens, because the demand is not part of the interval grid. Default unchanged. The refusal is still measured when the placement is off -- switching off a rung must not blind the diagnostic that justifies it. --- DeepSeek-V4 Unaffected by the slot-vs-group change, and that claim is now checked rather than asserted: DSV4 declares `entries_per_req=1` unconditionally (the MTP/DSpark lookahead widens the slot via `win_with_spec`, it never multiplies the count), so `state_slots_per_req == 1`, `pop_many(1)` pops the same index `pop()` did, and slot == group exactly as before. `v4_pool_geometry.py` and `sub_pool_spec.py` are untouched, so the `_physical_slots` reversal and DSV4's pool size are both byte-identical to main. DSV4 passes no `extra_entries`, so `--state-checkpoint- slots` is inert for it. The *anchor*, though, did reach DSV4 -- and cost it. `PagedStateCheckpointCoord- inator.applies()` is true for any V4 seq, so `_record_checkpoint_end` reserved a prompt-end anchor and `checkpoint_cut` shortened a prefill chunk onto it. But the coordinator files one pending checkpoint per seq (`_pending[id(seq)]`, and it `del`s `boundary_blocks`), so the prompt-end checkpoint landing a chunk later overwrote the anchor before either was stored. Measured: one extra prefill chunk per prompt for a hit rate that did not move (identical re-send 0 blocks either way, continuation 11 either way). So the anchor is now gated on `keeps_interior_boundaries`, which `StateSlotPool` answers True (each boundary is its own slot in the index, so both survive) and the PAGE coordinator answers False. Asked as a capability rather than by naming the backend, so a future multi-boundary copy class opts in by answering yes. DSV4 is back to main's one-cut prefill; GDN keeps the anchor. Pinned by `test_a_last_boundary_only_class_is_not_anchored_for`. --- reconciliation with main Dropped: the copy/pending-checkpoint path (`_commit_pending`, `record_copy`, `take_copies`, `Sequence.pending_checkpoint`). DeepSeek-V4 checkpointing moved to `PagedStateCheckpointCoordinator`, and `StateSlotPool` now rejects `transfer.copies` outright. The fork path (GDN), where the measured wins are, is kept in full. `readable_midstep` is carried on main's `StateTransfer` in `state_runtime.py` (wire format included) rather than on the branch's copy. Two bugs this rebase exposed, both fixed here: - `PagedStateCheckpointCoordinator` did not implement the midstep half of the `StateCache` protocol, so `checkpoint_cut` raised on every V4 batch. This was introduced *by this branch*, not latent in main: `readable_midstep` does not exist on main at all, and main's `checkpoint_cut` never consults it. A PAGE image is not readable midstep, so the three methods are the no-ops the protocol documents. - `_record_checkpoint_end` could place the anchor past the last matchable block. `can_allocate` stops one block short of the prompt, so a checkpoint filed under the final block's hash is one no scan looks up -- and being stored, it evicted the ladder rung that would have served the resume, taking an identical re-request from 8 hit blocks to 0. Capped at `(n_hash_blocks - 1) * hbs`. Tests: `tests/test_state_checkpoint.py` 171 passed; the state/checkpoint and DSV4/LMCache suites together 632 passed. Full suite 45 failed / 2300 passed, a strict subset of origin/main's own 114 pre-existing failures -- zero regressions, verified by set difference against a clean origin/main worktree. black clean; ruff no new findings. Co-Authored-By: Claude <noreply@anthropic.com>
… a -1 interval Rebased onto main's PAGE-backed checkpoint work (#1874, #1894, #1943, #1880), which rewrote this subsystem underneath the branch. Squashed to one commit because the three original commits each re-conflicted against the new base and against each other's resolutions; the reasoning from all three is kept below. --- allocate state slots per need, not by group A "state cache group" was `1 + num_spec` slots wide and was the unit of everything: allocation, admission, sizing, and the checkpoint index. But a checkpoint has no speculation to roll back -- it holds a committed state -- so filing one cost a full group and wasted `num_spec/(1 + num_spec)` of its bytes. At two speculative tokens that is two thirds. The slot is now the unit. `--state-checkpoint-slots 64` buys 64 checkpoints for 64 slots instead of 192, and the slots it no longer takes stay in the paged KV pool. This is only possible because a request's slots need not be adjacent, which the kernels never required: the ssm kernel gathers each index out of the indices tensor and the conv path is handed column 0 alone. Contiguity was manufactured by `prepare_state_indices` writing `arange(base, base + width)`; it now writes the seq's own slot list straight in. `StateGroupPool` -> `StateSlotPool`, `Sequence.per_req_cache_group` -> `state_slots`, whose element 0 is the committed state. The setter re-points [0] and preserves [1:], because speculation scratch persists across forwards. `--state-checkpoint-groups` still parses, as an alias. --- -1 turns off the interval ladder without turning off checkpointing The interval is a guess about where reuse will resume; a demand rung is a position a request was actually refused at. On the SemiAnalysis cc-traces the 8192 ladder placed ~30x the writes of the demand rung alone and caught reuse the demand already reaches -- 0.0% of resumes landed on a ladder rung -- while every rung costs the prompt that keeps it an extra prefill chunk. >0 a rung every N tokens (unchanged, still the default) 0 state checkpointing off entirely (unchanged) -1 no interval rungs; the demand rung and prompt-end anchor still place them -1 rather than reusing 0 because 0 is the documented contract and is reachable by accident: the grid snap rounds an off-grid interval down and can land on 0, so a --block-size typo currently fails safe. Three of the four sites are not the arithmetic you would guess -- `pos % interval` under -1 admits *every* position rather than none, and `pos - last < -1` is true for every pos. --- make the demand rung switchable A demand rung is 47% of checkpoint writes on the cc-traces and reads back 2.8% of the time, against 85.2% for a prompt-end anchor. Gated independently of --state-checkpoint-interval-tokens, because the demand is not part of the interval grid. Default unchanged. The refusal is still measured when the placement is off -- switching off a rung must not blind the diagnostic that justifies it. --- DeepSeek-V4 Unaffected by the slot-vs-group change, and that claim is now checked rather than asserted: DSV4 declares `entries_per_req=1` unconditionally (the MTP/DSpark lookahead widens the slot via `win_with_spec`, it never multiplies the count), so `state_slots_per_req == 1`, `pop_many(1)` pops the same index `pop()` did, and slot == group exactly as before. `v4_pool_geometry.py` and `sub_pool_spec.py` are untouched, so the `_physical_slots` reversal and DSV4's pool size are both byte-identical to main. DSV4 passes no `extra_entries`, so `--state-checkpoint- slots` is inert for it. The *anchor*, though, did reach DSV4 -- and cost it. `PagedStateCheckpointCoord- inator.applies()` is true for any V4 seq, so `_record_checkpoint_end` reserved a prompt-end anchor and `checkpoint_cut` shortened a prefill chunk onto it. But the coordinator files one pending checkpoint per seq (`_pending[id(seq)]`, and it `del`s `boundary_blocks`), so the prompt-end checkpoint landing a chunk later overwrote the anchor before either was stored. Measured: one extra prefill chunk per prompt for a hit rate that did not move (identical re-send 0 blocks either way, continuation 11 either way). So the anchor is now gated on `keeps_interior_boundaries`, which `StateSlotPool` answers True (each boundary is its own slot in the index, so both survive) and the PAGE coordinator answers False. Asked as a capability rather than by naming the backend, so a future multi-boundary copy class opts in by answering yes. DSV4 is back to main's one-cut prefill; GDN keeps the anchor. Pinned by `test_a_last_boundary_only_class_is_not_anchored_for`. --- reconciliation with main Dropped: the copy/pending-checkpoint path (`_commit_pending`, `record_copy`, `take_copies`, `Sequence.pending_checkpoint`). DeepSeek-V4 checkpointing moved to `PagedStateCheckpointCoordinator`, and `StateSlotPool` now rejects `transfer.copies` outright. The fork path (GDN), where the measured wins are, is kept in full. `readable_midstep` is carried on main's `StateTransfer` in `state_runtime.py` (wire format included) rather than on the branch's copy. Two bugs this rebase exposed, both fixed here: - `PagedStateCheckpointCoordinator` did not implement the midstep half of the `StateCache` protocol, so `checkpoint_cut` raised on every V4 batch. This was introduced *by this branch*, not latent in main: `readable_midstep` does not exist on main at all, and main's `checkpoint_cut` never consults it. A PAGE image is not readable midstep, so the three methods are the no-ops the protocol documents. - `_record_checkpoint_end` could place the anchor past the last matchable block. `can_allocate` stops one block short of the prompt, so a checkpoint filed under the final block's hash is one no scan looks up -- and being stored, it evicted the ladder rung that would have served the resume, taking an identical re-request from 8 hit blocks to 0. Capped at `(n_hash_blocks - 1) * hbs`. Tests: `tests/test_state_checkpoint.py` 171 passed; the state/checkpoint and DSV4/LMCache suites together 632 passed. Full suite 45 failed / 2300 passed, a strict subset of origin/main's own 114 pre-existing failures -- zero regressions, verified by set difference against a clean origin/main worktree. black clean; ruff no new findings. Co-Authored-By: Claude <noreply@anthropic.com>
… a -1 interval Rebased onto main's PAGE-backed checkpoint work (#1874, #1894, #1943, #1880), which rewrote this subsystem underneath the branch. Squashed to one commit because the three original commits each re-conflicted against the new base and against each other's resolutions; the reasoning from all three is kept below. --- allocate state slots per need, not by group A "state cache group" was `1 + num_spec` slots wide and was the unit of everything: allocation, admission, sizing, and the checkpoint index. But a checkpoint has no speculation to roll back -- it holds a committed state -- so filing one cost a full group and wasted `num_spec/(1 + num_spec)` of its bytes. At two speculative tokens that is two thirds. The slot is now the unit. `--state-checkpoint-slots 64` buys 64 checkpoints for 64 slots instead of 192, and the slots it no longer takes stay in the paged KV pool. This is only possible because a request's slots need not be adjacent, which the kernels never required: the ssm kernel gathers each index out of the indices tensor and the conv path is handed column 0 alone. Contiguity was manufactured by `prepare_state_indices` writing `arange(base, base + width)`; it now writes the seq's own slot list straight in. `StateGroupPool` -> `StateSlotPool`, `Sequence.per_req_cache_group` -> `state_slots`, whose element 0 is the committed state. The setter re-points [0] and preserves [1:], because speculation scratch persists across forwards. `--state-checkpoint-groups` still parses, as an alias. --- -1 turns off the interval ladder without turning off checkpointing The interval is a guess about where reuse will resume; a demand rung is a position a request was actually refused at. On the SemiAnalysis cc-traces the 8192 ladder placed ~30x the writes of the demand rung alone and caught reuse the demand already reaches -- 0.0% of resumes landed on a ladder rung -- while every rung costs the prompt that keeps it an extra prefill chunk. >0 a rung every N tokens (unchanged, still the default) 0 state checkpointing off entirely (unchanged) -1 no interval rungs; the demand rung and prompt-end anchor still place them -1 rather than reusing 0 because 0 is the documented contract and is reachable by accident: the grid snap rounds an off-grid interval down and can land on 0, so a --block-size typo currently fails safe. Three of the four sites are not the arithmetic you would guess -- `pos % interval` under -1 admits *every* position rather than none, and `pos - last < -1` is true for every pos. --- make the demand rung switchable A demand rung is 47% of checkpoint writes on the cc-traces and reads back 2.8% of the time, against 85.2% for a prompt-end anchor. Gated independently of --state-checkpoint-interval-tokens, because the demand is not part of the interval grid. Default unchanged. The refusal is still measured when the placement is off -- switching off a rung must not blind the diagnostic that justifies it. --- DeepSeek-V4 Unaffected by the slot-vs-group change, and that claim is now checked rather than asserted: DSV4 declares `entries_per_req=1` unconditionally (the MTP/DSpark lookahead widens the slot via `win_with_spec`, it never multiplies the count), so `state_slots_per_req == 1`, `pop_many(1)` pops the same index `pop()` did, and slot == group exactly as before. `v4_pool_geometry.py` and `sub_pool_spec.py` are untouched, so the `_physical_slots` reversal and DSV4's pool size are both byte-identical to main. DSV4 passes no `extra_entries`, so `--state-checkpoint- slots` is inert for it. The *anchor*, though, did reach DSV4 -- and cost it. `PagedStateCheckpointCoord- inator.applies()` is true for any V4 seq, so `_record_checkpoint_end` reserved a prompt-end anchor and `checkpoint_cut` shortened a prefill chunk onto it. But the coordinator files one pending checkpoint per seq (`_pending[id(seq)]`, and it `del`s `boundary_blocks`), so the prompt-end checkpoint landing a chunk later overwrote the anchor before either was stored. Measured: one extra prefill chunk per prompt for a hit rate that did not move (identical re-send 0 blocks either way, continuation 11 either way). So the anchor is now gated on `keeps_interior_boundaries`, which `StateSlotPool` answers True (each boundary is its own slot in the index, so both survive) and the PAGE coordinator answers False. Asked as a capability rather than by naming the backend, so a future multi-boundary copy class opts in by answering yes. DSV4 is back to main's one-cut prefill; GDN keeps the anchor. Pinned by `test_a_last_boundary_only_class_is_not_anchored_for`. --- reconciliation with main Dropped: the copy/pending-checkpoint path (`_commit_pending`, `record_copy`, `take_copies`, `Sequence.pending_checkpoint`). DeepSeek-V4 checkpointing moved to `PagedStateCheckpointCoordinator`, and `StateSlotPool` now rejects `transfer.copies` outright. The fork path (GDN), where the measured wins are, is kept in full. `readable_midstep` is carried on main's `StateTransfer` in `state_runtime.py` (wire format included) rather than on the branch's copy. Two bugs this rebase exposed, both fixed here: - `PagedStateCheckpointCoordinator` did not implement the midstep half of the `StateCache` protocol, so `checkpoint_cut` raised on every V4 batch. This was introduced *by this branch*, not latent in main: `readable_midstep` does not exist on main at all, and main's `checkpoint_cut` never consults it. A PAGE image is not readable midstep, so the three methods are the no-ops the protocol documents. - `_record_checkpoint_end` could place the anchor past the last matchable block. `can_allocate` stops one block short of the prompt, so a checkpoint filed under the final block's hash is one no scan looks up -- and being stored, it evicted the ladder rung that would have served the resume, taking an identical re-request from 8 hit blocks to 0. Capped at `(n_hash_blocks - 1) * hbs`. Tests: `tests/test_state_checkpoint.py` 171 passed; the state/checkpoint and DSV4/LMCache suites together 632 passed. Full suite 45 failed / 2300 passed, a strict subset of origin/main's own 114 pre-existing failures -- zero regressions, verified by set difference against a clean origin/main worktree. black clean; ruff no new findings. Co-Authored-By: Claude <noreply@anthropic.com>
* feat(state-cache): per-slot allocation, a switchable demand rung, and a -1 interval Rebased onto main's PAGE-backed checkpoint work (#1874, #1894, #1943, #1880), which rewrote this subsystem underneath the branch. Squashed to one commit because the three original commits each re-conflicted against the new base and against each other's resolutions; the reasoning from all three is kept below. --- allocate state slots per need, not by group A "state cache group" was `1 + num_spec` slots wide and was the unit of everything: allocation, admission, sizing, and the checkpoint index. But a checkpoint has no speculation to roll back -- it holds a committed state -- so filing one cost a full group and wasted `num_spec/(1 + num_spec)` of its bytes. At two speculative tokens that is two thirds. The slot is now the unit. `--state-checkpoint-slots 64` buys 64 checkpoints for 64 slots instead of 192, and the slots it no longer takes stay in the paged KV pool. This is only possible because a request's slots need not be adjacent, which the kernels never required: the ssm kernel gathers each index out of the indices tensor and the conv path is handed column 0 alone. Contiguity was manufactured by `prepare_state_indices` writing `arange(base, base + width)`; it now writes the seq's own slot list straight in. `StateGroupPool` -> `StateSlotPool`, `Sequence.per_req_cache_group` -> `state_slots`, whose element 0 is the committed state. The setter re-points [0] and preserves [1:], because speculation scratch persists across forwards. `--state-checkpoint-groups` still parses, as an alias. --- -1 turns off the interval ladder without turning off checkpointing The interval is a guess about where reuse will resume; a demand rung is a position a request was actually refused at. On the SemiAnalysis cc-traces the 8192 ladder placed ~30x the writes of the demand rung alone and caught reuse the demand already reaches -- 0.0% of resumes landed on a ladder rung -- while every rung costs the prompt that keeps it an extra prefill chunk. >0 a rung every N tokens (unchanged, still the default) 0 state checkpointing off entirely (unchanged) -1 no interval rungs; the demand rung and prompt-end anchor still place them -1 rather than reusing 0 because 0 is the documented contract and is reachable by accident: the grid snap rounds an off-grid interval down and can land on 0, so a --block-size typo currently fails safe. Three of the four sites are not the arithmetic you would guess -- `pos % interval` under -1 admits *every* position rather than none, and `pos - last < -1` is true for every pos. --- make the demand rung switchable A demand rung is 47% of checkpoint writes on the cc-traces and reads back 2.8% of the time, against 85.2% for a prompt-end anchor. Gated independently of --state-checkpoint-interval-tokens, because the demand is not part of the interval grid. Default unchanged. The refusal is still measured when the placement is off -- switching off a rung must not blind the diagnostic that justifies it. --- DeepSeek-V4 Unaffected by the slot-vs-group change, and that claim is now checked rather than asserted: DSV4 declares `entries_per_req=1` unconditionally (the MTP/DSpark lookahead widens the slot via `win_with_spec`, it never multiplies the count), so `state_slots_per_req == 1`, `pop_many(1)` pops the same index `pop()` did, and slot == group exactly as before. `v4_pool_geometry.py` and `sub_pool_spec.py` are untouched, so the `_physical_slots` reversal and DSV4's pool size are both byte-identical to main. DSV4 passes no `extra_entries`, so `--state-checkpoint- slots` is inert for it. The *anchor*, though, did reach DSV4 -- and cost it. `PagedStateCheckpointCoord- inator.applies()` is true for any V4 seq, so `_record_checkpoint_end` reserved a prompt-end anchor and `checkpoint_cut` shortened a prefill chunk onto it. But the coordinator files one pending checkpoint per seq (`_pending[id(seq)]`, and it `del`s `boundary_blocks`), so the prompt-end checkpoint landing a chunk later overwrote the anchor before either was stored. Measured: one extra prefill chunk per prompt for a hit rate that did not move (identical re-send 0 blocks either way, continuation 11 either way). So the anchor is now gated on `keeps_interior_boundaries`, which `StateSlotPool` answers True (each boundary is its own slot in the index, so both survive) and the PAGE coordinator answers False. Asked as a capability rather than by naming the backend, so a future multi-boundary copy class opts in by answering yes. DSV4 is back to main's one-cut prefill; GDN keeps the anchor. Pinned by `test_a_last_boundary_only_class_is_not_anchored_for`. --- reconciliation with main Dropped: the copy/pending-checkpoint path (`_commit_pending`, `record_copy`, `take_copies`, `Sequence.pending_checkpoint`). DeepSeek-V4 checkpointing moved to `PagedStateCheckpointCoordinator`, and `StateSlotPool` now rejects `transfer.copies` outright. The fork path (GDN), where the measured wins are, is kept in full. `readable_midstep` is carried on main's `StateTransfer` in `state_runtime.py` (wire format included) rather than on the branch's copy. Two bugs this rebase exposed, both fixed here: - `PagedStateCheckpointCoordinator` did not implement the midstep half of the `StateCache` protocol, so `checkpoint_cut` raised on every V4 batch. This was introduced *by this branch*, not latent in main: `readable_midstep` does not exist on main at all, and main's `checkpoint_cut` never consults it. A PAGE image is not readable midstep, so the three methods are the no-ops the protocol documents. - `_record_checkpoint_end` could place the anchor past the last matchable block. `can_allocate` stops one block short of the prompt, so a checkpoint filed under the final block's hash is one no scan looks up -- and being stored, it evicted the ladder rung that would have served the resume, taking an identical re-request from 8 hit blocks to 0. Capped at `(n_hash_blocks - 1) * hbs`. Tests: `tests/test_state_checkpoint.py` 171 passed; the state/checkpoint and DSV4/LMCache suites together 632 passed. Full suite 45 failed / 2300 passed, a strict subset of origin/main's own 114 pre-existing failures -- zero regressions, verified by set difference against a clean origin/main worktree. black clean; ruff no new findings. Co-Authored-By: Claude <noreply@anthropic.com> * remove gpu unit test Signed-off-by: ganyi <ygan@amd.com> * remove the triton import part Signed-off-by: ganyi <ygan@amd.com> * match main's array('i') token_ids contract in a relocation test `#1990` added an assertion that `Block.token_ids` is an `array('i')`, not a list -- a list never compares equal to what the production publish paths store, so every hit on the block would read as a hash collision. This test was written before that landed and still passed a bare list. The file already has `toks()` for exactly this; the test just did not use it. * feat(state-cache): keep Kimi-K3's KDA checkpoints as PAGE images A KDA Active Slot is 53.6 MiB. Held as a checkpoint it competed with live requests for the pool that admits them, so retaining one cost the workload the concurrency it was retained for. Held as PAGE units it is 127 ordinary KV blocks -- 0.112% of the paged pool -- drawn from the same free list as everything else and evicted by the same LRU. This is the mechanism `main` already ships and DeepSeek-V4 already uses (`PagedStateCheckpointCoordinator`). Nothing about the coordinator changes; what is added is the source side of the copy for a state that is two strided tensors rather than one contiguous slab. `plan_segmented_copy` intersects two ordered byte streams and needs neither block alignment nor equal segments, so the state tensors keep their layout: a slot is 138 ranges (69 conv + 69 ssm) and the planner cuts them against 127 units. `_checkpoint_layer_ranges` is the sole owner of that order -- both the sizes and the addresses read it, because a plan cut against one order and addressed through another lands whole layers in the wrong unit. Two things the port had to get right, both now asserted rather than assumed: - A PAGE unit is a *logical* block, but `kv_cache` is shaped in physical ones and K3's `block_ratio` is 128. `_page_unit_regions` derives its stride from `runner.block_size` and checks `num_rows * region == page_unit_bytes`, so a granularity mix-up is a startup error instead of 127 blocks of scrambled state. Unit ids are range-checked against the logical count for the same reason. - K3's slots are strided by `num_slots`, so an off-by-one in `(layer * num_slots + slot)` lands inside a neighbouring request's live state rather than off the end of the tensor. V4 cannot fail this way and its tests do not look for it; `test_no_bystander_slot_is_touched` does. `state_spec` now asks for no spare checkpoint slots under PAGE. `--state-checkpoint-slots` buys Active Slots for checkpoints to sit in, which a copy does not need -- 1.7 GiB reserved for nothing, and it is the same memory the paged pool wants in order to absorb the images. Both fall back to `fork` under pipeline parallelism and RapidServe, where `get_num_blocks` raises on a copying transfer: answering `copy` there would turn "K3 keeps no state cache" into "K3 does not start". The dtype objection in the old `state_transfer` docstring is retired, not ignored. It was that `_state_dtypes` gives kimi_linear an fp32 v side while the chunked states are bf16, so a checkpoint cut from the kernel's `h` would hand cached requests a rounded state. A PAGE image is copied out of the slot and back into a slot -- both fp32, no kernel output in between, no conversion anywhere. Both dtypes are named in the layout id, so a build that changed either cannot read another's images. Not yet flipped on in anger: `execute_paged_state_copies` is reachable only from `build()`, and the GPU verification (probe at conc 1 and 8 against the known-good 0/1 and 0/8, then GSM8K, then a matched-N hit-rate A/B) is the next step. Known follow-up, measured before it is fixed: the coordinator keeps one checkpoint per sequence (`_pending` is last-writer-wins), where the fork path indexed every boundary. A 24k prompt files at 8192/16384/24576 today and would keep only the last. If the A/B shows the drop, the lever is to make `_pending` hold a list -- deliberately not bundled here, because a mechanism swap plus a policy change is a regression nobody can attribute. * feat(state-cache): keep every boundary a PAGE seq reaches, not just its last `_pending` was keyed by sequence, so a prompt's second checkpoint overwrote its first before either was stored. That made the prompt-end anchor worthless -- it sits under a block from the prompt's end, lands in the same or the adjacent prefill chunk, and was reliably the loser. `_record_checkpoint_end` reads `keeps_interior_boundaries` and duly declined to reserve one. Keyed by `(sequence, prefix hash)` both survive. Reaching the *same* hash twice still collapses, which is what the hash in the key is for: that is one boundary reached again, not two boundaries. This matters because the anchor is the placement that pays. The measurement is already in `_record_checkpoint_end`'s docstring: of 4,808 cc-trace resumes with a nonzero KV hit, 93.5% land on a previous prompt end and 0.0% on the 8192 ladder. The ladder was cutting a prefill chunk every 8192 tokens to store something nothing ever resumed from -- and on this workload a prompt averages 117k tokens, so that is ~14 rungs per request, each one a shortened forward and an image in the paged pool. What makes keeping both affordable is the price a PAGE image pays: 127 blocks, 0.112% of the paged pool, against a whole 53.6 MiB Active Slot under `fork`. The measured run that preceded this kept 1,508 checkpoints with `checkpoints_evicted: 0` -- capacity was never the binding constraint. Run with `--state-checkpoint-interval-tokens -1` to drop the ladder entirely and leave the anchor and the demand as the only two placements. Three tests changed rather than deleted, because each pinned the old behaviour deliberately and each now pins its replacement: - `test_latest_pending_checkpoint_replaces_the_previous_intent` becomes `test_two_boundaries_of_one_seq_are_both_stored`, plus a new sibling for the same-hash-twice case. - `test_a_last_boundary_only_class_is_not_anchored_for` becomes `test_both_classes_are_anchored_for`. - Two demand tests rested a tightened pool on "exactly one image"; a prompt now stores two, so they spend down to the deepest -- which is both the resume target and what `_next_victim` would keep longest. Not yet measured. The preceding PAGE run at conc 8 reached 91.79% at N=791 against a 96.9% trace ceiling; this is the change aimed at that gap, and the A/B is the next step. * refactor(state-cache): drop the parts of the PR nothing reads Three removals, none of which change behaviour. Verified against the same 4610-passed baseline, and `ruff` on the touched files goes 16 -> 14 findings. `cache_pressure.py` had no importer anywhere in the tree, and the log field its regex parses (`Cached/Total:`) was renamed to `Cached/Reusable:` by this same PR -- so it could not have matched a line this branch produces. `keeps_interior_boundaries` was a capability hook with one reader and no implementor that answered `False`: the `getattr` default was `True`, both classes set `True`, and the case it existed for -- the PAGE coordinator overwriting its own anchor -- was fixed earlier in this branch by re-keying `_pending` on `(seq, hash)`. The measurement that justified it (of 4,808 cc-trace resumes with a nonzero KV hit, 93.5% land on a previous prompt end, 0.0% on the 8192 ladder) moves onto `checkpoint`, which is where the key it argues for lives. `_log_frequency` and its four `reqs_*` counters cost four `__slots__` entries and four per-request branches to render one log line, and are read by nothing else -- not `metrics.py`, not either aggregation tuple in `llm_engine.py`. `_log_pools` stays: its three rates are pure ratios of totals already kept, and the paged/state split is this PR's central claim. `_log_pressure` stays because `checkpoints_*` and `demands_recorded` do reach Prometheus. Left alone deliberately: the `record_relocation` / `take_relocations` / `relocate_state_slots` chain is equally unreachable, but it is that way on `main` too. Deleting main's debt from this branch would widen the diff it is meant to narrow. * docs(state-cache): tighten the comments this PR added No code changes; 375 tests pass and `ruff` on the touched files stays at 14 findings against main's 16. The bf16/fp32 accuracy argument was written out in full three times -- `GDNStateMixin.state_transfer`, `pop_last_intermediate_states`, and inverted again in `_KimiMLAGDNCommon.state_transfer` -- each time as a rebuttal to an objection nobody raised, and two of the three cited `tests/test_gdn_state_checkpoint_gpu.py`, deleted in a04ce7f. It now lives once, in the present tense, where the dtypes are chosen; the other two point at it. That alone is ~30 lines and both dead citations. Two measurements had spread to four and five sites. The prompt-end anchor's read-back rate stays in `_record_checkpoint_end`, which exists because of it; the demand rung's stays in `mark_speculative`, the only place it decides behaviour, and in the `--state-checkpoint-demand` help text, where a CLI user cannot follow a code reference. `config.py`, `envs.py`, `sequence.py`, `page_unit_checkpoint.py` and `checkpointers_at` now reference rather than restate, so there is one copy to update when the number moves. The rest is history that git already holds: what an earlier Python-loop version got wrong, what the upstream branch does with `state_cache_base`, what this pool "used to allocate", which objection "kept this on fork". Each is restated as the invariant it was arguing for. Also two stragglers of the group->slot rename in `attention_gdn.py`, and a call-site comment in `gdn_attn.py` that restated `_checkpoint_targets`' own docstring. Left long on purpose: `_page_unit_regions`' logical-vs-physical block-id trap (K3's block_ratio is 128, and getting it wrong scrambles 127 blocks silently), `_assert_checkpoint_geometry_still_holds`, the conv-window claim in `state_transfer`, and `CacheStats`' argument for `reusable` over `full` as the denominator -- that last reads like a rebuttal but the objection is one a reader will actually raise. * docs: describe the two model-agnostic features and the instrumentation The description covered only the K3 PAGE port, which is 1,221 of the 4,694 added lines. Three things it shipped were undocumented: Per-slot allocation. `StateGroupPool` -> `StateSlotPool`, and a request's state goes from one fixed-width group of `1 + num_spec` adjacent slots to a list of ids that need not be adjacent. The point is that a checkpoint takes one slot rather than a whole group, since a resumed prefix has no speculation to roll back. Documents the one consumer that reads past element 0 -- the spec-decode path, which stopped deriving the set from `base = group * slots_per_group` -- and states why DeepSeek-V4 is a rename rather than a behaviour change. Midstep checkpoints. A mamba-like backend can now take every boundary a forward covers out of the chunk kernel's own `h`, instead of having its prefill cut so the forward *ends* on each one. Covers the reserve/publish/cancel split (the bytes do not exist when the destination must be chosen), the `is_end` targets that read the runtime slot because `h` does not hold the final state, and the paired gate in `checkpoint_cut`/`checkpointers_at` -- suppressing one alone keeps zero checkpoints with no error. Names Qwen3-Next and Qwen3.5 as the models on this path and K3 as the one that cannot be, and adds the latter to the follow-ups. Hit-rate instrumentation. Every measurement in this PR was read off these lines. The `[Cache Stats]` denominator was `full`, which includes the trailing block `can_allocate` never matches -- so it charged both pools for a block neither was offered and reported an unreachable ceiling; it is now `reusable`. `[Cache Pools]` splits the series into `paged * state = combined`, which is what showed the paged index matching 99.4% while the state gate discarded it. `[Checkpoint Fates]` separates four fates that argue for different fixes, and `kept: 1508, dropped: 0, evicted: 0` is the evidence behind the "capacity stopped being the binding constraint" claim. Also refreshes the numbers the rebase and the two cleanup commits invalidated: 33 files / +4694, the current commit hashes, the per-file table, and the test baseline (4610 passed / 50 pre-existing failures, 40 of them sglang files that score identically on origin/main). * docs(state-cache): KDA's interior h exists; aiter just does not return it The follow-up said K3 cannot be `readable_midstep` because `chunk_kimi_delta_attn` "exposes only `output_final_state`". True of the API, misleading about the cause: in aiter's `_triton_kernels/chunk_delta_attn/chunk_fwd.py` the per-chunk `h` is computed at line 170 -- by `chunk_gated_delta_rule_fwd_h`, the same function the GDN path uses -- consumed by `chunk_gla_fwd_o`, then set to None at line 202 and left out of the returned tuple. So the two backends differ in plumbing, not in what their kernels produce. ATOM vendors GDN's chunk entry under `model_ops/fla_ops/`, which is why `keep_intermediate_states` could be added there; KDA goes out to aiter, which has no equivalent. Whoever picks this up is adding a return value, not an algorithm -- worth stating, because the old wording invites the conclusion that the kernel would have to be rewritten. Behaviour is unchanged: K3 stays `readable_midstep = False` and keeps cutting a chunk per placement. Under the shipped anchor-only policy that is 0% of prompts cut at 1.00 checkpoints per request, since the anchor lands where the last prefill chunk was going to end anyway. * test(gdn): pin the claim readable_midstep rests on, on real hardware `readable_midstep` asserts that `h[:, j]` is the recurrent state after `j * 64` tokens, and `BlockManager` acts on it by suppressing `checkpoint_cut` outright -- the prefill runs full length and the boundaries are harvested from `h` afterwards. If that assertion is false, every checkpoint the readable path stores is subtly wrong: a resuming request inherits a state its prefix never produced, silently. `TestMidstepCheckpoints` pins everything *around* the claim (which positions are chosen, reserve/publish/cancel, that the cut is suppressed) but stubs the kernel, so it cannot see the claim itself fail. This asks the kernel. Measured on MI355, 8 chunks of 64: all 7 interior boundaries are **bit-exact** against a forward stopped at that position -- `torch.equal`, not a tolerance, which is the right bar because both arms round the same fp32 value into the same dtype (`h` is `k.new_empty`; `_state_dtypes` returns `config.torch_dtype`). Two smaller guards alongside it: popping consumes the reference, so a later layer cannot read the previous one's `h` and file it under its own slot; and a forward that was not asked to keep retains nothing, so the plugins that never pop do not pin a large tensor past their last forward. Needs one GPU and a few hundred MB -- no server, no TP, no weights -- and skips at module level otherwise, following `test_compress_chunk_equivalence`. * docs: record the midstep hardware result, and narrow what is still unmeasured "Qwen3.5 is not measured" was true when written and is now too blunt. The claim `readable_midstep` rests on -- that `h[:, j]` equals the state a forward stopped at `j * 64` would leave -- has been asked of the kernel directly: 7 of 7 interior boundaries bit-exact on MI355. That belongs in Verification, because it is the one part of the midstep path a CPU test cannot reach and a failure there would be silent. What remains unmeasured is narrower and worth saying precisely: no server has been stood up on a readable backend, so there is no hit rate, accuracy, or TTFT for it. Named the three things a single sequence through one kernel cannot show -- the per-sequence `chunk_offsets[row]` base with two prefills in a batch, the ordering an `is_end` target depends on, and a resume landing on a stored midstep boundary -- so the gap is actionable rather than a blanket disclaimer. * fix(state-cache): address review findings 1, 7, 10 and the instrumentation Findings from @valarLip on #2045, each re-verified against the code rather than taken on report -- two of the sixteen did not survive that check (`_rehome_checkpoint` does not exist; `chunk_gated_delta_rule` carries `@torch.compiler.disable`, so the CUDA-graph half of #12 cannot happen). **#1, a regression this branch introduced.** `eb058321d` re-keyed `_pending` to `(seq, hash)` so two boundaries of one prompt could coexist, and did not touch the drain, which still resolves a single `seq.state_slot` for all of them. Both images are then copied out of whatever the last forward left there, filing the earlier hash over the later state -- a request resuming on it continues from ahead of its own prefix, and `_validate_paged_state_op` passes because layout, size and unit count are all still correct. `_supersede` keeps one pending boundary per sequence. Ordinarily a drain follows every forward and both boundaries are stored correctly from their own slots; the exception is a pass that schedules nothing, where `state_maintenance_ops=None` carries `_pending` into the next drain. The newer boundary wins because it is the one the slot holds, and the older is counted `dropped` -- it is reuse the placement asked for and did not get. The two tests that pinned the old behaviour asserted coexistence without asserting each was stored from its own slot, which the drain cannot do; they now pin the fix and a sibling covers the ordinary drain-between-forwards case. **#7** was the same invariant read from the other end: the descriptor buffer is sized `2 * max_num_seqs` on "one store per sequence", which the re-key removed and `_supersede` restores. No resize -- the docstrings here and in `deepseek_v4_attn.py` now name what holds the bound instead of asserting it. **#2/#3/#4/#5/#8 are all on the midstep write path, so `readable_midstep` goes back to False.** The write path declines on six conditions `commit_midstep` cannot see and publishes the hash regardless; `_checkpoint_targets` indexes three differently scoped sequence lists with one `i`; the SSM read floors to a 64 grid `midstep_positions` does not enforce (`hash_block_size` defaults to 16); the conv window is `conv_kernel-1+num_spec` in the kernel and `conv_kernel-1` in the guard. Each stores a findable image holding the wrong state. None of it has run under a server -- K3 takes the PAGE path and cannot reach it -- so the honest state is off. The machinery and its bit-exactness test stay; `test_midstep_is_off_in_production.py` pins the decision, and is deliberately not behind `importorskip` so the non-GPU runner actually runs it. **#10** `pool_pressure` read `self.state`, which under PAGE is a different object built with `StateTransfer.none()` that never sees a `checkpoint()`. It printed four zeros for the life of the server while `checkpoint_funnel`, the next method, reported the real numbers from the coordinator. Smaller: `chunks_cut_for_end` and `checkpoints_orphaned` reach both aggregation whitelists (a cut counter without its sibling is unreadable; `orphaned` argues for a bigger paged pool where `evicted` argues for a bigger state pool); `paged_hit` is dropped as a second name for `compressed_hit`; `_warn_if_unschedulable` compares against `state_slots_per_req` again, so a pool too narrow for one request warns instead of waiting forever in silence. `clear_index` still moves no counter, now stated as a decision: each fate argues for a different fix and an operator emptying the cache argues for none. * fix(state-cache): address review findings 6, 9, 11, 13 and 14 **#9 inverts the policy it implements, on the majority of prompts.** `mark_speculative` exists so a guessed resume point is spent before a known one -- anchors are read back 85.2% of the time against a demand rung's 2.8%. Both call sites gated on `if anchor and pos != anchor`, so a seq whose `checkpoint_end_pos` is 0 demoted *nothing* and filed its guesses at the LRU tail beside real anchors. `_record_checkpoint_end` leaves it at 0 on four paths, one being every prompt too short for a keepable end -- the common shape of an agentic first turn. `_anchor_of` answers None rather than 0 there, which compares unequal to every position, so those seqs demote everything. Shared by both sites because two spellings of one rule is how they drift apart. `publish_midstep(seq=None)` still demotes nothing, now as the stated other end of the rule: a caller with no sequence cannot tell a guess from knowledge, and over-keeping costs one eviction where over-demoting spends an anchor. **#6 could take the engine down over a log line.** The two asserts run for every prefill seq, over counters with four independent writers -- the CPU-offload wake sets `num_cached_tokens` without touching the hit-block counters the rest derive from, so an LMCache resume that loads more prefix than the GPU index held produces `cached > wanted` legitimately. Now a warning and a clamp, which also removes a behaviour difference between `-O` and not. **#11 had two consumers disagreeing about what -1 means.** `BlockManager` clamps to `max(-1, ...)` and reads -1 as "grid off, anchor and demand still placing"; the DSV4 offload policy clamped to `max(0, ...)`, folding it into 0, which for that consumer means no sidecar checkpoints at all -- so the engine kept checkpointing while offload resume silently degraded to zero reuse. With no grid to align to the sidecar now takes `resume_alignment` alone. **#13b/#13c.** `_extend_hash_chain` sat one line above the `_has_page_units` refusal, so a 128k prompt queued behind a full pool paid ~2000 xxhash rounds per waiting request per pass for a list that was then discarded. Moved below it; verified nothing between consumes it, and that its one reader is `midstep_positions`. The comment claiming it "reads `checkpoint_end_pos`" was false and is replaced with what actually orders the call. **#14.** `chunk_gated_delta_rule_fwd` returned `h` unconditionally, so the caller's frame pinned ~33 MB for the rest of that layer's forward even with `keep_intermediate_states=False` -- every GDN prefill with checkpointing off, which is the default, and every vLLM/SGLang/rtpllm caller. The flag now reaches the producer, so the reference dies with the fwd frame. No compute changes; the kernel computed it either way. `test_gdn_midstep_state_gpu.py` covers both values and still passes on hardware. Suite: 4622 passed against 4618 before, with the same 50 pre-existing failures. * docs: carry the slot rename into the guides, and document the new knobs The rename landed in code and left five guides describing a `group` model that no longer exists. One of them was actively dangerous: the `deallocate` snippet in the scheduling guide released `seq.per_req_cache_group` — a single slot — where the real function calls `release_many(seq.state_slots)`. Copied as written it leaks `num_spec` slots per request, and admission cannot see the loss because it gates on the free list this never returns them to. Corrected across `scheduling_kv_cache_guide.md` (the pool construction snippet, the allocation and deallocation prose, the pool-field list, the Sequence table, and the fork-checkpoint capacity paragraph), plus the Sequence rows in `architecture_guide.md` and the GDN state paragraph in `model_support_guide.md`. `state_slots` is documented as a list with `[0]` committed and `[1:]` rollback, explicitly not adjacent, with `state_slot` as the property over element 0 — which is the contract a backend has to know before it indexes anything. Newly documented rather than merely renamed: * `--state-checkpoint-interval-tokens -1`. The guides described `0` as the only off switch, so the ladder-off-but-anchor-on regime this PR added was reachable and undocumented. * `--state-checkpoint-slots` (and its `--state-checkpoint-groups` alias), with the note that a PAGE backend zeroes it out. * `--state-checkpoint-demand` / `--no-state-checkpoint-demand`. * `ATOM_STATE_CHECKPOINT_DEMAND`, under a new "State checkpoints" section in `environment_variables.md` — it had no entry at all. Every symbol the guides now name was checked to exist in `atom/`. The two `*_plan.md` files still say `group`; they are dated design notes rather than reference docs, and rewriting them would misrepresent what was planned. * revert: drop two changes that belong to other PRs Neither touches per-request state, checkpoints, or the pools. They rode along on this branch and widen its review surface for no reason. `triton_merge_attn_states.py` moves `prefill_tokens_with_context` off `tl.constexpr`. That is a real fix — a per-batch token count as a constexpr mints a fresh kernel per distinct batch size, 184 of them in one 8-minute agentic run — but it is an attention-kernel compile-time bug, not a state-cache one, and belongs in a PR that says so. `tests/plugin/test_vllm_kimi_k3.py` moved its registry check out-of-process to survive `sys.modules` damage other plugin tests do. Also genuine, also unrelated; verified it does not pollute the session on its own. `test_rtpllm_forward_context_semantics.py` is NOT reverted, though it looked like the same category. Its change makes the stubs it installs restore what they displaced, and without it `atom.model_ops.attention_gdn` and `atom.utils.forward_context` stay shadowed for the rest of the session: reverting it turned 3 collection errors into 6, taking `test_cudagraph_capture_bounds.py` (9 passed alone) and three sglang plugin modules down with `cannot import name ... (unknown location)`. That is load-bearing for whether this branch's own suite can be run at all. * remove --state-checkpoint-slots, which never took effect The flag sized a flat cushion of spare Active Slots for checkpoints to sit in. It defaults to 0, DeepSeek-V4 never declared it, and Kimi-K3 overrode it back to 0 — so on every shipped path it added nothing, and the only configuration where it did anything was GDN's `fork` with someone passing a value by hand, which no measurement in this PR or before it covers. What it was for is real: a checkpoint held as a slot competes with live requests, so how many can be retained is set by concurrency rather than by how much reuse the traffic has. The PAGE path solves that properly, by keeping the image in KV blocks instead of a slot. A cushion would buy the same decoupling for `fork` at the price of a knob nobody can size without measuring first. Removing it collapses two things it had propped up. `_KimiMLAGDNCommon.state_spec` existed only to zero the field and is deleted — with the flag gone the base spec is already right, and `super()` needs no correction. And `TestTheSpareSlotsGoBackToTheKvPool` went with it: it monkeypatched `GDNStateMixin.state_spec` to a lambda returning a literal `extra_entries=32`, so it asserted against its own stub and would have passed unchanged if production stopped reading the field altogether. That is the shape @valarLip flagged, and deleting the feature removes the test's subject rather than its substitute. `SubPoolSpec.extra_entries` stays. It is the sizing layer's general capability, no backend passes a nonzero value today, and the test that pins its arithmetic now says so — a future cushion should get a flat one, not `width x` what it asked for. 4620 passed against 4622 before, the difference being the two deleted tests; same 50 pre-existing failures. `ruff` on the touched files matches origin/main exactly. * test: keep the midstep watch on the CPU-only side of the aiter line `test_gdn_does_not_declare_itself_midstep_readable` reached `GDNStateMixin.state_transfer`, which means importing `gdn_attn`, which imports aiter at module level. The non-GPU CI runner installs CPU torch and neither aiter nor triton, so that is a collection-time ModuleNotFoundError, not a skip -- it failed the job. Split the flag's two halves by what a CPU runner can actually see. The PAGE coordinator's `readable_midstep`, the three-call midstep protocol on `StateCache`, and `StateTransfer`'s field are all pure Python and stay watched here. GDN's declaration is the half that needs aiter; it belongs with the kernel tests, and the docstring now says so rather than leaving the next reader to rediscover it by breaking CI. Coverage lost is one assertion, not the mechanism: `TestMidstepCheckpoints` builds its own `StateTransfer(readable_midstep=True)` and exercises the write path regardless of what production declares. --------- Signed-off-by: ganyi <ygan@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Guanbao Yu <Guanbao.Yu@amd.com>
2 participants
Summary
Test plan
pytest -q tests/test_deepseek_v4_transfer_regions.py tests/test_kv_aggregator.py tests/test_lmcache_offload_config.py tests/test_lmcache_offload_connector.py tests/test_lmcache_offload_v4_page_slot.py tests/test_lmcache_offload_disk_integration.py tests/test_lmcache_offload_gpu_disk_e2e.py tests/test_multi_connector.py tests/test_page_region_codec.py tests/test_per_req_cache_decoupling.py tests/test_scheduler.py tests/test_slot_sidecar_codec.py tests/test_slot_sidecar_format.py tests/test_slot_sidecar_store.py tests/test_state_checkpoint.py tests/test_v4_pool_geometry.py(781 passed, 7 prerequisite skips)ruff checkandblack --checkon all changed Python filespython3 -m compileallon changed packages andgit diff --checkROCm regression results
Validated at commit
4cf68deawith DeepSeek-V4-Pro, TP8, FP8 KV/index cache, and standalone DSV4 PAGE/SLOT LMCache LocalCPU offload.GSM8K two-pass offload regression
Workload: GSM8K 20-shot, 1,319 questions, concurrency 32.
1,316 x 8 = 10,528successful rank-local completions; the remaining 3 requests were served by the native HBM prefix cache.LMBenchmark CxS recompute vs offload
Workload: C=2, S=4, eight 262,144-token sources, one seed plus three follow-up turns, answer length 20,
max_num_seqs=2.The HBM capacity gate was
524,288 < 1,498,624 < 2,097,152tokens (active pair < HBM capacity < full working set), forcing older prefixes out of HBM while keeping the active pair safe.LocalCPUBackendon all 8/8 ranks.load_requests=18,loaded_tokens=4,128,768,load_failures=0,save_requests=110,saved_tokens=1,884,928,loads_pending=0, andsaves_pending=0.18 x 8 = 144), all withstatus=ok.Measurement caveat: concurrent MoE generation is numerically non-deterministic, and later turns include earlier generated text. The two variants therefore differed by 6 cumulative prompt tokens out of approximately 7.54M (
7,544,640vs7,544,646) and by 3 completion tokens (660vs657). Request keys, source data, questions, and benchmark configuration were identical.Made with Cursor