fix(engine): a fourth consumer of token ids, and a guard for the next one by valarLip · Pull Request #1990 · ROCm/ATOM

valarLip · 2026-08-23T10:35:27Z

What

Sequence.token_ids became an array("i") in #1989. That PR named three
places that compare token ids and fixed each. Sweeping the rest of the chain
turned up a fourth consumer of a different kind — one that serializes
rather than compares — and it fails silently too.

The bug

BlockManager.publish_loaded_prefix hands _hash_block_tokens' output
straight into a BlockStored event, which is msgpack-encoded:

list         encoded 33 bytes
array('i')   FAILS: TypeError: Encoding objects of type array.array is unsupported

KVEventPublisher.publish counts encode failures rather than raising
(kv_events.py:195-206), so with KV events enabled the whole event stream goes
dark after one log line. Latent today — KVEventsConfig.enable defaults to
false — but it is a real path.

The other two BlockStored sites accumulate into a list via extend and were
never affected.

The fix: assert at the boundary, don't convert at each site

boundary	takes	why
`Block.update`	`array`	a list never compares equal to the arrays the other publish paths store, so every hit on that block reads as a collision — and it hands the collector one traversal slot per token
`_make_block_stored`	`list`	the msgpack encoder above

Both refuse the wrong type loudly instead of failing quietly downstream.

The rest of the axis

Everything that grows once per token is an array now: Sequence.output_tokens
(the completion half of token_ids, kept in step with it by every writer),
Sequence.logprobs, the API server's three per-request accumulators, the
streaming detokenizer's buffer, and atomesh's two.

tokenizer.decode takes an array as is, so only the JSON edges convert back —
and only LLMEngine's, because the HTTP response builders read text and the
counters and never token_ids. A test pins that, so a builder that starts
forwarding the key fails there rather than in production.

Four annotations said list[int] where an array now flows.

Measurements

Live V4-Flash-DSpark tp1, 94,681 blocks, gc.callbacks probe on the EngineCore:

storage	gen-2 pause
`list[int]` (before)	288.6 ms
`array("i")` (after)	242.8 ms

With arrays the pause also stops growing as the prefix cache fills — the cost
becomes independent of pool occupancy.

Throughput is unchanged at this scale (32098 vs 32012 tok/s): gen-2 fires
about once per seven minutes in steady state at c=32/tp1, so this is a
tail-latency and memory change, not a throughput one.

Test plan

5 new tests in tests/test_token_ids_storage.py, each with a positive
control asserting the failure shape it guards
Mutation-tested: reverting the conversion is caught by the integration
test; additionally removing the boundary assert is caught by the encoder
check underneath it
tests/test_block_pool.py updated to hold the production type
Full CI suite: 9 failed / 2429 passed — the 9 are pre-existing
(eplb ×4, WoA ×3, dspark ×2)
ruff/black clean; import-without-aiter probe passes on all touched tests

… one `Sequence.token_ids` became an `array("i")` in #1989. That commit named three places that compare token ids and fixed each. Sweeping the rest of the chain turned up a fourth consumer of a different kind -- one that serializes rather than compares -- and it fails silently too. `publish_loaded_prefix` hands `_hash_block_tokens`' output straight into a `BlockStored` event, which is msgpack-encoded. msgspec has no encoding for `array.array`, and `KVEventPublisher.publish` counts encode failures rather than raising, so with KV events enabled the whole event stream goes dark after one log line. The other two BlockStored sites accumulate into a list via `extend` and were never affected. Rather than convert at each site, the two boundaries where the type is load-bearing now assert it: - `Block.update` takes an array and refuses a list. A list there costs twice and neither cost announces itself: it never compares equal to the arrays the other publish paths store, so every hit on that block reads as a collision; and it hands the collector one traversal slot per token. - `_make_block_stored` takes a list and refuses an array, for the encoder. The rest of the chain follows the same axis: everything that grows once per token is an array now. `Sequence.output_tokens` (the completion half of `token_ids`, kept in step with it by every writer), `Sequence.logprobs`, the API server's three per-request accumulators, the streaming detokenizer's buffer, and atomesh's two. `tokenizer.decode` takes an array as is, so only the JSON edges convert back -- and only `LLMEngine`'s, because the HTTP response builders read `text` and the counters and never `token_ids`. A test pins that, so a builder that starts forwarding the key fails there rather than in production. Four annotations said `list[int]` where an array now flows. Measured on a live V4-Flash-DSpark tp1 (94,681 blocks) with a `gc.callbacks` probe: the EngineCore's stop-the-world gen-2 pause is 288.6ms with list storage and 242.8ms with arrays, and with arrays it stops growing as the prefix cache fills. Throughput is unchanged at this scale -- gen-2 fires about once per seven minutes in steady state here, so the pause matters at the concurrency #1980 measured at, not at this one.

github-actions · 2026-08-23T10:35:56Z

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

`compute_hash` was annotated `list[int] | array.array` on the reasoning that one caller still passed a list. It does not. All six production call sites pass an `array("i")`: - `block_manager.py` x5, from `_hash_block_tokens` / `seq.block(i)`, both slices of `Sequence.token_ids` - `policy.py`, from `_chained_prefix_hashes(seq.token_ids, ...)` The union came from reading `_chained_prefix_hashes(token_ids: list[int], ...)` and taking the signature for the fact -- but that annotation is stale in the same way, since `connector.py` hands it `seq.token_ids`. Both are narrowed to what actually arrives. The int64 pin stays. Its reason changes rather than disappearing: nothing passes a list today, so it is no longer reconciling two live callers, but it is what stops a caller who does pass one from silently computing a different digest for the same tokens. Only tests pass lists now, and `test_compute_hash_does_not_depend_on_its_ argument_type` keeps asserting the two agree -- the annotation says what callers should pass, that test says the function is robust if they do not.

) * perf(engine): stop the collector walking a heap it never reclaims At steady state the garbage collector in these processes finds nothing. Measured on DeepSeek-V4-Flash-DSpark tp1 with a `gc.callbacks` probe: over 1754 gen-0, 158 gen-1 and 50 gen-2 passes after startup, it reclaimed **zero** objects. Everything the hot loop allocates is acyclic and reference counting takes it. What each pass does cost, being stop-the-world and proportional to the live heap: EngineCore 268 ms of 819,923 tracked TP0 worker 979 ms of 959,694 api_server 265 ms of 632,424 Almost all of that heap is startup state -- model, compiled graph, tokenizer, KV block pool -- so `gc.freeze()` after warmup removes the work rather than merely spacing it out, which is all raising thresholds (#1980) can do. With it on, gen-2 stops firing entirely in all three processes for the rest of the run, while gen-0/1 keep going at under 2 ms: this narrows the collector, it does not disable it, so a cycle written by later code is still caught. `unfreeze_gc_heap` on shutdown is not optional. A frozen object is invisible to the collector, so an engine destroyed inside a live interpreter -- a test, an RL loop that rebuilds it -- would leave its weights and KV cache unreachable *and* uncollectable, which presents as a GPU memory leak with nothing about GC in the symptom. vLLM hit this and handles it the same way. Placement is one call per process at the boundary where startup ends: - api_server, in the lifespan just before it yields - EngineCore, at the end of `__init__` -- KV cache allocated, graphs captured, BlockPool built (94,763 `Block` objects on a V4-Flash tp1) - each ModelRunner worker, over an RPC the EngineCore sends, because only the caller knows warmup is over: weights, compile and capture all arrive as RPCs it issues, and `--enforce-eager` has no capture step to hook Disaggregated decode is the exception and is left partly covered. Its block pool is built in `_init_disagg`, after `super().__init__()` has already sent READY, so a freeze there could make queued requests permanent. It keeps the base freeze and forgoes the pool. `tune_gc` moves into the new module and stays as the fallback for `ATOM_GC_FREEZE=0`. `ATOM_GC_DEBUG` is added alongside: it logs every collection's generation, duration and reclaimed count, which is the only way to see these pauses -- a stall in the EngineCore idles the workers, and an idle worker emits no trace event at all. It is expensive (~90s of extra startup) and off by default. Two gaps found next door and closed, since the fix needed a shared per-process setup anyway: `PrefillEngineCore.run_engine` and `DecodeEngineCore.run_engine` override the base without calling it, so a disaggregated deployment got neither `enable_orphan_reaping` nor `tune_gc`, and neither ever set a process title -- both showed as bare `python` in `ps`. `_setup_engine_process` now owns identity, reaping and GC policy for all three, which also collapses the base's three `set_process_title` branches into one and makes the process title and the GC log lines agree by construction. Throughput is unchanged at this scale (32098 vs 32122 tok/s, and 32098 with freezing off), because gen-2 fires roughly once per seven minutes in steady state here. This is a tail-latency change; it pays at the concurrency #1980 measured at, which is not reproduced here. * fix(engine): one name per worker process, not three The worker side had the drift the EngineCore side just lost. Three places produced its name and only one knew about data parallelism: - `set_process_title` in `AsyncIOProc.__init__`, dp-aware (`DP0TP0`) - the GC debug callback, `f"TP{rank}"` - `ModelRunner.freeze_gc_heap`, `f"TP{rank}"` So under dp>1 every rank's worker logged as `TP0`, which is exactly the case worth telling apart -- and `ps` disagreed with both log lines. `engine_process_name` and `worker_process_name` now live next to `set_process_title`, since naming a process is what that function is for, and all five call sites take their string from them. `AsyncIOProc` names itself once, before anything logs, which also retires a `try/except Exception` that existed only to repeat the fallback the helper now expresses directly. `EngineCore._process_name_for` moves out to join them rather than staying a method, so a reader asking "what are these processes called" finds one answer. * fix(entrypoints): the atomesh frontend was the third process to miss this Enumerating which processes get the GC policy has now gone stale twice. #1980 reached the API server but not the disaggregated EngineCores, whose `run_engine` does not call the base one. The first version of this branch reached those but not atomesh, which builds its own engine in `launch_atom_standalone` and never runs the FastAPI lifespan the API server applies the policy from -- so that frontend had neither `tune_gc` (since #1980) nor the freeze, while the EngineCore and workers it spawns had both. It is the API server's counterpart, with the same profile: the tokenizer, the per-request accumulators, ~632k tracked objects and a 265 ms gen-2 pause. Rather than fix the third instance and wait for a fourth, the rule is now tested: `test_every_serving_frontend_applies_the_gc_policy` walks every module that builds an engine and serves, and fails naming the ones that do not apply the policy. Written the obvious way it was vacuous -- an unused import satisfies a search for the name -- so it requires the call. The prose enumerations in `tune_gc` and the env-var docs are replaced by the rule they kept failing to track. * fix(engine): restore the decode freeze, dropped on a misreading An earlier commit on this branch removed `_freeze_after_startup` from `DecodeEngineCore` on the grounds that READY had already gone out, so freezing could catch requests the input thread had queued and make them permanent. That is not what the code does. `DecodeEngineCore.__init__` sets `_ready_deferred = True` before calling `super().__init__()`, precisely so the base's ready signal is suppressed; the real one is sent nine lines *after* the block pool is built, once the kvcache IPC import and graph capture are done. Nothing can have been admitted at that point, so there is no window and the freeze is safe -- and it is what covers decode's 94,763 `Block` objects. The comment left in its place asserted the false premise in so many words, which would have taught the next reader the same wrong thing.

CI runs ruff through reviewdog with `-filter-mode=diff_context`, so it reports anything landing on a changed line or in the context around it -- pre-existing or not. Adding an import inside this file's already-unsorted block pulled a long-standing `I001` into review and failed the check. Comparing whole-file findings against `main` is the wrong local check for that: identical counts still fail if one of them is near your edit. Cheaper to leave the file with none. Sorting the imports is the fix CI asked for. The other two go with it: `initialize_standalone_service` declared `global engine, tokenizer` but only reads them, which `global` is not for; and the blind except around the version banner gets the reason it always had.

CI's non-GPU runner has no aiter, and `atom.model_engine.model_runner` imports it. Importing `ModelRunner` inside the test to read `freeze_gc_heap` off it therefore failed there with `ImportError: cannot import name destroy_dist_env`. The import-without-aiter probe missed it because that probe models collection: it imports each test module and stops. An import inside a test body only runs when the test does. The contract is about what the source says -- an RPC target must return something or `call_func(..., wait_out=True)` hangs -- so it is now read with `ast` straight from the file. No import, no aiter, no GPU, and the whole module drops from 5.7s to 1.6s. Re-verified by running every test file this branch touches with `aiter` and `triton` blocked at the import hook: 142 passed.

Rebasing onto main brought #1989/#1990, which turned `Sequence.token_ids` and `Sequence.block_table` into `array("i")` for the GC pause it buys at long context. An `array("i")` never compares equal to a list, so the two assertions this branch added in `TestJointClaimReusesResidentBlocks` -- correct when written against a `list` block table -- started failing. Same fix main applied to its own test in #1990 (`assert list(seq.output_tokens) == [...]`): convert the array side. assert list(seq.block_table[:4]) == resident `seq.block_table = [0]` in `test_scheduler.py` is deliberately left alone: main has two of its own at lines 469 and 489 of that file, so a lone converted third would read as the odd one out, and it passes. The formatting half is `black .` over the nine source files and three tests this branch had left non-compliant; main is black-clean, so every one of them was ours. Three of the reformats join a wrapped expression onto one line and the rest are blank lines -- `git diff -w --ignore-blank-lines` is empty outside those three. Verified: `pytest tests/ --ignore=tests/plugin` -> 3780 passed, 192 skipped, 3 xfailed. The two remaining collection errors (msgpack, aiter) are on main as well. `tests/plugin` is excluded because `test_rtpllm_forward_context_semantics.py` and `test_rtpllm_glm5_sparse_backend_contract.py` replace `sys.modules["atom.utils.forward_context"]` without restoring it, which breaks every later test that builds a Scheduler -- 66 of main's own included. Pre-existing and not this branch's to fix. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

`#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.

@valarLip

* 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>

valarLip mentioned this pull request Aug 23, 2026

perf(engine): stop the collector walking a heap it never reclaims #1991

Merged

5 of 7 tasks

valarLip added 2 commits August 23, 2026 10:45

valarLip added the ci:atom label Aug 23, 2026

valarLip added 2 commits August 23, 2026 11:31

valarLip merged commit da7698b into main Aug 23, 2026
31 of 32 checks passed

valarLip deleted the valarLip/token-ids-fourth-consumer branch August 23, 2026 12:21

valarLip mentioned this pull request Aug 25, 2026

perf(frontend): merge_chunk rebuilt the whole backlog on every merge #2020

Merged

1 task done

valarLip mentioned this pull request Sep 14, 2026

Enhance metrics with scheduler, cache, and decoding improvements #2186

Merged

Navigation Menu

fix(engine): a fourth consumer of token ids, and a guard for the next one - #1990

fix(engine): a fourth consumer of token ids, and a guard for the next one#1990
valarLip merged 5 commits into
mainfrom
valarLip/token-ids-fourth-consumer

valarLip commented Aug 23, 2026

Uh oh!

github-actions Bot commented Aug 23, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

valarLip commented Aug 23, 2026

What

The bug

The fix: assert at the boundary, don't convert at each site

The rest of the axis

Measurements

Test plan

Uh oh!

github-actions Bot commented Aug 23, 2026

🏷️ CI Guide

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant