Conversation
…ff checkpointing The interval is a guess about where reuse will resume; a demand rung is a position a request was actually refused at. Measured 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. Every rung also costs the prompt that keeps it an extra prefill chunk, which is where the ladder's measured 17.5% throughput regression on Qwen3.5-27B tp2 comes from. So the ladder needs an off switch that is not the off switch for the feature: >0 a rung every N tokens (unchanged, still the default) 0 state checkpointing off entirely (unchanged) -1 no interval rungs; the demand rung still places checkpoints -1 rather than reusing 0, even though "no interval" reads like zero: 0 is the documented contract and it is also 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. Giving 0 a second meaning would make that typo silently enable a caching policy instead of disabling one. Four sites, three of which are not the arithmetic you would guess: - the clamp at __init__ stops flattening the negative - the grid snap tests `> 0` rather than truthiness, or -1 snaps to -4 and warns about a flag the user set deliberately - `checkpointers_at` cannot ask `pos % interval` under -1: -1 divides every integer, so that test admits *every* position as on-grid rather than none - the decode path returns early under -1 for the same reason -- `pos - last < -1` is true for every pos, which would checkpoint on every decode step Swept over 76 prompt lengths x 3 requests: -1 costs 68 cuts against the ladder's 108 (-37%), and every cut is still kept in all three modes. Co-Authored-By: Claude <noreply@anthropic.com>
The ladder guesses where reuse will resume; the demand rung learns it one refusal late. Neither covers the case that dominates agentic traffic: a conversation whose next turn resends this whole prompt and continues past it. That resumes at *this* prompt's end, and the only request in a position to leave a checkpoint there is this one. Measured on the SemiAnalysis cc-traces (4,808 resumes with a nonzero KV hit): 93.5% land on a previous prompt's end, 0.0% on the 8192 ladder. `_record_checkpoint_end` reserves that position at admission, off the grid on the same terms as a demand. It steps back to the rightmost grid position that still leaves `successor_room` behind it — the exact floored end leaves at most `hash_block_size - 1`, so for any class whose room reaches a block or more (V4's 131 against 256) it is never keepable, and an anchor placed there would be cut for and then refused. `checkpoint_cut` now takes the *earliest* candidate rather than the latest. With an anchor at 36 and a rung at 32, `max` cuts at 36 and the forward never ends at 32, so the rung is not deferred but lost — a class the anchor is out of reach for then loses its resume point permanently while the demand counter climbs. Within the grid a later rung still dominates, so the grid collapses to one candidate first. `chunks_cut_for_end` is counted apart from `chunks_cut_for_demand`: the anchor fires on nearly every prompt and would drown the convergence signal the demand counter exists to expose. Co-Authored-By: Claude <noreply@anthropic.com>
…s that can A checkpoint is kept at the position a forward ended at, because that is the only position a backend can be assumed to expose. So the interval ladder is paid for in forwards: `checkpoint_cut` shortens a prefill chunk onto every rung, and a prompt with N rungs on it is N shortened chunks. Measured at 17.5% throughput on Qwen3.5-27B tp2. A chunk kernel already materializes the recurrent state at each of its interior chunk boundaries. A backend whose kernel copies those out does not need the forward to end anywhere in particular — it can snapshot every rung the chunk covers and still run the chunk full length. This is the engine half of letting it: one checkpoint per forward at its end becomes N checkpoints at interior boundaries, each filed under its own block hash. `StateTransfer.readable_midstep` carries the capability, defaulting False and crossing the process boundary as its own scalar. A separate axis from `kind`, and deliberately not derived from it: `kind` says how a group is handed to another request, this says where within one forward a snapshot can be taken at all. GDN is fork(1) and will be readable; DeepSeek-V4 is copy() and is not. Keeping one splits into two moments the end-of-forward path runs as one. `plan_midstep` takes a destination group per position before the batch is built — the free list has to be committed against while admission for the pass is finished, or a checkpoint's destination competes with a request still to be let in. `commit_midstep` files the hashes after, from inside `hash_blocks`, which is the point all three prefill call sites converge on and precisely when the reserved bytes exist. Publishing at reservation time would index a group over bytes the forward had not written; `cancel_midstep` hands back a reservation whose forward never ran, released vacant rather than indexed. Naming a position the forward has not reached needs a hash the admission scan never computed — `block_hashes` stops at the first miss. `_extend_hash_chain` runs the chain over the whole prompt onto `seq.block_hashes`, only for a backend that reserves midstep. The two gates are a pair: `checkpoint_cut` returns 0 when every applicable class is readable, and `checkpointers_at` drops those classes on the aimed path so the same boundary is not kept twice. Moving either alone is silent — suppressing the cut alone leaves `checkpointers_at` refusing every off-grid position it is then handed, so zero checkpoints are kept with no error; dropping the second gate alone spends two groups on one hash and, under fork, makes the seq hand its live group away for a replacement it never needed. Inert until a backend declares the flag, which none does yet: the kernel-side copy-out and the runner's read of `batch.state_save_all` are still to land. 19 tests, and 17 mutations run against the two files — 16 killed, the survivor an unreachable `interval == 0` check now removed with the reason recorded. Co-Authored-By: Claude <noreply@anthropic.com>
… forward
The engine layer landed in 2c64d69 reserves N checkpoints inside one
forward for backends that declare `readable_midstep`. Nothing declared
it, so nothing was written. This is the GDN backend making the claim
true.
A checkpoint at a *step end* was already nearly free: the chunk kernel
leaves the final state in the runtime slot and the existing scatter puts
it there. A checkpoint *inside* the step was not reachable at all —
which is why the engine had to cut the prefill chunk short so the
position it wanted became a step end. That cut is what cost 17.5%
throughput on Qwen3.5-27B tp2.
The state at an interior position is not gone, though: the chunked
recurrence materializes it at every 64-token boundary as `h[:, j]` and
then drops it on the floor. So:
chunk.py returns `h` regardless of SUPPRESS_LEVEL and parks
it on the Function for one pop, behind an explicit
`keep_intermediate_states` flag (~33 MB at T=8192,
so callers that never pop must not ask).
state_checkpoint.py one Triton launch over a split grid copies both
halves of every target — recurrent state and conv
window — for all targets at once.
gdn_attn.py `_checkpoint_targets` turns the engine's
reservations into device index tensors, tagging
each as interior (source: `h`) or end (source: the
runtime slot, since `h` holds only boundaries
strictly *before* the end).
attention_gdn.py invokes the write after the scatter, so an `is_end`
target reads this step's final state.
Kimi/KDA is excluded, deliberately. It inherits `GDNStateMixin` but runs
`fla.ops.kda.chunk_kda`, which never exposes per-chunk states, and it
overrides the `prepare_prefill` that builds the targets. Inheriting the
capability would have suppressed the chunk cut AND made `checkpointers_at`
drop every off-grid position — zero checkpoints, silently, with nothing
on that path able to tell a skipped checkpoint from an unwanted one. So
`_KimiMLAGDNCommon` overrides `state_transfer()` back to a plain fork.
Porting the branch's `chunk_kda_paged` would let that override go away.
Two things this is NOT:
- Bit-exact. `h` is bf16 while the recurrence carries fp32, so resuming
from an interior checkpoint is ~1.2e-4 relative (0.19%) off a cold run.
Measured on the source branch, T=512, resume at 256, output absmax
0.065. Positions at a step end read the runtime slot and stay exact.
- Executed. No kernel launch in this commit has ever run. The 130 CPU
tests mock AITER and stay green whether or not the Triton kernel is
correct; the full suite diff is 120 failures before and after, i.e.
all pre-existing. GPU parity is the next step and is not covered here.
…ca2 claimed about it
c6f1f4ca2 added the interior-checkpoint copy-out and admitted its kernel had
never been launched. It has now, on an MI355, and the run disproved the
accuracy claim that commit reasoned its way into.
That claim was: `h` is bf16 while the recurrence carries fp32, so a state
sliced out of `h` is pre-rounded and resuming from it must be ~1.2e-4 off a
cold run. The reasoning is sound and the conclusion is wrong, because it stops
one step early. The state does not stay in `h` — it goes into the GDN state
pool, and `GDNStateMixin._state_dtypes` allocates that pool at
`config.torch_dtype`, the same bf16. Ending the forward at that position
instead would round the same fp32 value on the way in. Both paths round; they
round identically; the difference is nil.
Measured, not argued: across 56 (seed, length, boundary) combinations an
interior checkpoint equals a shortened re-run's stored state exactly, and
resuming from one reproduces the remaining tokens' outputs bit for bit —
rtol=0, atol=0.
This matters because "interior checkpoints are approximate" is precisely the
argument for keeping the prefill cut that costs 17.5% throughput: pay the cut,
get the exact state. There was nothing to buy, and the old docstrings were
talking anyone who found them into buying it.
The exactness rests on the two dtypes agreeing, not on the copy, so it does not
survive a pool wider than `h`. `_state_dtypes` builds exactly one such pool —
kimi_linear's fp32 v side — and that model is already off this path.
`test_midstep_backends_are_the_ones_with_a_matching_pool` states that as an
invariant over the backends declaring `readable_midstep`, so a future GDN-like
backend has to satisfy it rather than merely resemble one. Widening a pool
without widening `h` would otherwise break nothing loudly: cached requests
would quietly resume from a bf16-rounded state uncached ones never see, which
surfaces only as an eval delta between cache-hit and cache-miss runs.
tests/test_gdn_state_checkpoint_gpu.py 16 tests, GPU-only, skips without one
tests/test_state_checkpoint.py +3 pinning the dtype premise on CPU,
since CI has no GPU to run the above
The GPU file also covers what the copy-out can get wrong without looking wrong:
varlen bases (ragged lengths, so a dropped `chunk_offsets[row]` cannot land
right by luck), the `is_end` branch (row 0's end is row 1's first chunk in
`h` — a real state, wrong sequence), several targets per row, and the conv half
checked alongside every SSM assertion because the two share `slots`. Ten
mutants of the kernel were introduced and all ten were killed.
Two test-isolation bugs surfaced on the way and are fixed here, because the new
tests are red without the first and the suite cannot be read without both:
- `test_rtpllm_forward_context_semantics.py` installed stubs over real modules
in `sys.modules` at import time and never restored them, so
`atom.model_ops.attentions.gdn_attn` and `atom.utils.forward_context` stayed
unreachable for the rest of the session. That was costing 85 tests across
four files, as collection errors reading "unknown location" that named the
imported module and never the file responsible.
- `test_vllm_kimi_k3.py::test_kimi_k3_plugin_registries_are_synchronized` now
runs out-of-process, like the two tests beside it already did. Several plugin
tests import torch inside a `patch.dict(sys.modules, ...)` whose exit deletes
the torch tree while torch's C-level state stays initialized; the next
in-process `import torch` dies in an unrelated file. Unmasked by the fix
above, not caused by it — it reproduces identically on an unmodified tree.
Full suite: 120 failures before, 35 after, no new ones.
Co-Authored-By: Claude <noreply@anthropic.com>
The three premise tests imported `gdn_attn`/`kimi_mla_gdn_attn` directly,
and those do `from aiter import ...` at module scope. Importing aiter runs
its arch probe, which on a CPU-only box raises `0 active drivers` and then
falls back to a `jax` import that is not installed -- so the tests passed
only where a GPU happened to be visible, which is not the property a premise
check should have.
Stubbed via a meta-path finder scoped to the `aiter` namespace, installed
and removed per-test. A finder rather than a `sys.modules` list because the
transitive set is 19 submodules and is aiter's business, not this file's.
Two things the obvious version gets wrong, both found by running it:
- Six `atom` modules are already sitting in `sys.modules` as bare
ModuleTypes that other test files installed and never removed. Any on
the import path re-raises as "cannot import name X (unknown location)",
so they are evicted too and restored after.
- Modules imported *under* the stub hold MagicMocks but have a real
`__file__`, so they cannot be identified the same way. Snapshotted
before the import and removed after instead, leaving `sys.modules` as
found -- the earlier stubs included, since removing them is not this
fixture's call.
Mutation-tested rather than assumed: widening the GDN pool to fp32, letting
KDA declare `readable_midstep`, and dropping the kimi_linear special case
each fail exactly the tests that should catch them, in full-suite context
and not merely in isolation. Without that check a stub that imports cleanly
and asserts nothing looks identical to a pass.
Full suite vs merge-base: no new failures with a GPU visible (120 -> 35),
and CPU collection errors return to the pre-existing 22.
merge_attn_states_kernel declared prefill_tokens_with_context as
tl.constexpr, but it is fed a runtime per-batch token count (defaulting
to num_tokens). Every tl.constexpr is baked into the compiled artifact
and forms part of the cache key, so each distinct batch token count
minted a fresh kernel.
Measured on a Kimi-K3 agentic run: 184 variants in a single 8-minute
run, never converging. First compile ~42ms against 0.022ms of warm
execution. TTIR confirms the bake -- one variant carries
`arith.constant 16384`, others `arith.constant 292`.
The parameter feeds exactly one expression:
prefix_mask = token_idx < prefill_tokens_with_context
whose left side is already runtime (tl.program_id(0)), so the branch was
never resolved at compile time and nothing is lost by demoting it. Moved
it above the constexpr block in the signature and correspondingly in the
launch's positional args; parameter order verified to still align.
The public merge_attn_states() signature is unchanged, and all three
in-tree call sites pass it by keyword.
Introduced by the MLA chunked cached-prefix prefill path (758c276).
Co-Authored-By: Claude <noreply@anthropic.com>
…ttributed The hit rate cannot say why reuse was missing. A prompt with no shared prefix and a prefix evicted an hour ago produce the same number, and they want opposite fixes. Every counter needed to tell them apart already existed and none of them was ever printed: `StateGroupPool` has kept four checkpoint fates since #1771, `BlockManager.checkpoint_funnel` assembles them, and nothing calls it. The paged pool counted nothing at all. Adds the paged side and logs both next to the existing Cache Stats line: [Pool Pressure] paged: N/M used, R reusable-free, V vacant, I indexed | evicted: E, retired: T | state: ... [Checkpoint Fates] kept/dropped/evicted/orphaned Two distinctions the counters have to keep, because collapsing either gives an answer that points at the wrong pool: evicted vs retired the pool ran out, vs the boundary moved over cached content. First says the pool is too small, second says the split is wrong. evicted vs orphaned a checkpoint spent for room, vs one whose prefix left the KV index first. The second is paged pressure surfacing in a state counter. `_unindex` now reports whether an entry actually went, because `on_evict` also fires for relocation (`_adopt`), where the hash survives on the block that adopted it — counting that would report a move as a loss. Vacant is called out because it is the leading indicator: evictions cannot begin until it hits 0, so a run ending with vacant blocks never had paged pressure whatever its hit rate. Pressure is read through a callable at log time, not passed per update: the free-list scan behind `num_reusable_free` is O(free blocks) (~10k here) for a line printed once per 100 prefills. `debug_helper/cache_pressure.py` reads a server log and prints the verdict. Written for the open question on this branch — whether the paged pool is actually the binding constraint at conc 4 — which so far rests on arithmetic from the startup sizing dump, with no eviction event ever observed. These counters make that falsifiable either way. Tests: 7 new, each verified to fail without the change. A/B over the cache/scheduler/block suites shows no new failures (132 -> 125, the delta being exactly these). The 125 remaining are pre-existing CPU-suite failures unrelated to this change. Co-Authored-By: Claude <noreply@anthropic.com>
Two kinds of position get checkpointed and LRU could not tell them
apart. The prompt-end anchor is where the next turn actually resumes;
the demand rung and the ladder rung guess. Measured on the cc-traces at
conc 4: anchors are read back 85.2% of the time, demands 2.8% -- and
demands are 47% of all writes (1,370 of 2,919). Under plain LRU the
placement almost never read was evicting the one almost always read.
So file a guess at the LRU head rather than the tail, and promote it to
the tail if a request does resume from it. Demote, not drop: the demand
exists to fill a gap once, and one spent before it is read cost nothing
that never taking it would have saved.
A group mark rather than an argument to `release`, because the two
checkpoint paths hand the group back differently -- `checkpoint` under
`fork` *pins* it and `release_pins` two passes later knows only the
index. Marking is what both paths can reach. Promotion sits in
`_attach_state_group`, not `claim`, because `_set_hash` also calls
`claim` to re-file a group nobody read.
Replayed through the real BlockManager on 1,560 requests of cc-traces,
at today's pool sizes and zero extra bytes:
max_num_seqs=4 86.9% -> 88.2% retention 89.6% -> 96.5%
max_num_seqs=8 85.1% -> 86.8% retention 88.1% -> 96.1%
Co-Authored-By: Claude <noreply@anthropic.com>
The STATE pool was sized at exactly `max_num_seqs * entries_per_req`, so
the room to retain a checkpoint was whatever concurrency happened to
leave spare. That makes a *lower* max_num_seqs measure a *worse* prefix
hit rate on traffic that reuses prefixes -- the opposite of what the
knob is for, and a colleague hit it at concurrency 1.
`--state-checkpoint-groups N` adds N per-request entry sets on top of
the in-flight floor, paid for out of the paged KV pool. Multiplied by
`1 + num_spec` in `gdn_attn.state_spec` because BlockManager derives its
group count as `entries // entries_per_req`, so headroom in anything but
whole widths is rounded away.
Replayed through the real BlockManager on cc-traces at conc 8, holding
total bytes fixed (each group bought costs paged blocks):
32 groups (today) 88.8%
48 groups 90.6%
64 groups 92.0% <- peak
96 groups 91.9%
128 groups 88.2% KV ceiling falls faster than retention rises
Default 0 keeps the old coupling.
Co-Authored-By: Claude <noreply@anthropic.com>
`cevicted` counts checkpoints displaced by later checkpoints, which any run longer than the pool does by construction -- it tracks `kept` minus the pool size, not shortage, so it declared "the state pool is the wall" on every run. `cdropped` is the one that means a keep was refused for want of a group; the verdict now turns on that. Also surfaces the lost-to-checkpoint percentage as a ceiling line, since it bounds what growing the state pool can possibly buy back. The scheduler's prefill log called a cursor a hit count: `num_cached_tokens` advances by each finished chunk, so a chunked prompt logged the same req_id repeatedly with the number climbing, reading as a growing cache hit. Labelled "done". Co-Authored-By: Claude <noreply@anthropic.com>
A demand rung is 47% of all checkpoint writes on the cc-traces and reads back 2.8% of the time, against 85.2% for a prompt-end anchor. Every one of those writes evicts something, so the rung may cost more than its reuse is worth — but the CPU replay scores "delete the rung" and "keep it and demote it in LRU" identically (86.9% -> 88.2% both), which is the signature of a harness that models eviction order and not the cost of the write itself. Only hardware separates them. Gated independently of --state-checkpoint-interval-tokens, because the demand is not part of the interval grid: it is the one placement a refused hit makes for itself. Default unchanged. The refusal is still measured when the placement is off — `num_wanted_hit_blocks` is what CacheStats splits declined reuse by, and switching off a rung must not blind the diagnostic that justifies it. Co-Authored-By: Claude <noreply@anthropic.com>
`--no-state-checkpoint-demand` already drops the rung, but changing it means editing a launch script. ATOM_STATE_CHECKPOINT_DEMAND overrides the flag when exported: =0 leaves the prompt-end anchor as the only checkpoint placement, =1 forces the rung on over a script that disables it. Unset changes nothing, so an exported variable cannot pin the policy for every server on the box by accident. Worth being able to flip per-run because the CPU replay cannot settle it: it scores "delete every demand rung" (2,919 -> 1,550 writes) and "keep them but demote them in LRU" (2,925 writes) identically, which is the signature of a harness that models eviction order but not the cost of the write. On the cc-traces a demand is 47% of all checkpoint writes and reads back 2.8% of the time against the anchor's 85.2%. Co-Authored-By: Claude <noreply@anthropic.com>
`[Cache Stats]` already splits declined reuse between the two pools, but in tokens, and the traffic this runs on makes that ambiguous: prompts span 481 to 688k tokens, so a handful of deep conversations decide every ratio on the line. 9.65% lost tokens reads identically whether every request lost a tenth of its prefix or a tenth of requests lost all of theirs, and those want opposite fixes. `[Cache Freq]` counts requests instead. The bucket that matters for sizing is `paged-match+state-miss`: the paged pool still had the prefix and only the missing checkpoint stood between the request and reuse, so state-pool capacity binds exactly to the extent it is large. When `paged-miss` dominates instead, no checkpoint tuning helps. The buckets deliberately overlap and do not sum to the request count. A prompt whose paged prefix ran out early AND whose remaining reuse needed a checkpoint has two independent problems; charging it to one pool would have to pick a winner and would undercount the other, which is how a measurement points sizing at the wrong pool. Co-Authored-By: Claude <noreply@anthropic.com>
`[Cache Freq]` shipped two counters that could not report anything but
their extreme. `full_reuse` tested `cached >= full` and `no_paged`
tested `compressed < full`, but `full` is unreachable by construction:
`can_allocate` matches over `range(n_hash_blocks - 1)`, because prefill
must forward at least one block to produce sampler logits, so a
request's trailing block is never a reuse candidate. `compressed < full`
therefore holds for every request that can exist. The line read
`full-reuse: 0.0%, paged-miss: 100.0%` on the first run that used it,
and would have read that on a flawless one.
`full` was the denominator for the token rates too, which is the same
error costing accuracy rather than meaning: it charges both pools for a
block neither was offered, and the unreachable part is a fixed
`hash_block_size`, so it is ~13% of a 1k prompt and ~0.05% of a 275k
one. Hit rate then drifts with the prompt-length mix even when both
pools behave identically -- enough to make two runs incomparable.
Add `reusable` as the ceiling, computed at the call site from the same
rule the matcher uses, and divide by it everywhere.
That fixes the counters but still reports one blended rate, which cannot
say which pool to fix: the same 85% is a KV pool that lost the prefix or
a state cache that refused to resume from it, and those want opposite
changes. The pools run in series, so scoring them apart is a matter of
picking the right denominator for each:
paged = compressed / reusable was the prefix still in KV?
state = cached / compressed given it was, could we resume?
paged * state = cached / reusable, the end-to-end rate
`state` is scored against what the paged pool actually handed over, not
against `reusable`. Otherwise a KV eviction lands in the state cache's
score and points tuning at the wrong pool -- a property pinned by test,
since it is the whole reason for the split. The product being exact
means the smaller factor is the bottleneck, which `[Cache Pools]` names
outright.
`state-hit+ckpt` is where the state cache would land with a dense
ladder: the ceiling on what checkpoint placement or more groups can buy,
and the number that says when to stop spending bytes on them.
The nesting the rates depend on is now asserted rather than assumed. It
was documented as holding "by construction", but a violation surfaces as
a negative percentage in a log line rather than a failure, which is how
the tautology above survived being read several times.
Fixed in `/debug/cache_stats` as well, which divided by `full_tokens` in
its own copy of the arithmetic.
`checkpoints_dropped` is not the "pool too small" signal both docstrings claimed. `_commit_pending` only counts a drop when `has_free()` is false, but a finished request hands its group back and `pop` never refuses -- with nothing vacant it spends the LRU checkpoint and counts an eviction. The pool overwrites checkpoints rather than turning them away, so `dropped` reads 0 at every pool size, including one far too small. The real reading is `checkpoints_evicted` against `kept - num_groups`. On a Kimi-K3 agentic run that identity held exactly -- kept 198, evicted 166, 32 groups -- at the same flush where state-hit fell from 95.69% to 89.50%. 84% of checkpoints destroyed for space, and the 0 in `dropped` was read as proof the pool had room to spare. `occupancy` pointed at the same wrong counter, and its `groups_vacant` invites a second misread: `5/32 used, 22 checkpointed, 5 vacant` is a pool 84% full, not 84% idle. Adds a test that thrashes a 4-group pool with 16 checkpoints and pins `dropped == 0` alongside `evicted == kept - num_groups`, so the counter cannot quietly read as slack again. Co-Authored-By: Claude <noreply@anthropic.com>
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 and reads back into slot 0 — so filing one cost a full group and wasted `num_spec/(1 + num_spec)` of its bytes. At the two speculative tokens this model runs with, that is two thirds. The slot is now the unit. A live request takes `1 + num_spec` of them because it speculates; a checkpoint takes exactly one. `extra_entries` is passed through rather than multiplied by the width, so `--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, which is sized out of what is left after this. 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`, and `Sequence.per_req_cache_group` becomes `state_slots`, a list whose element 0 is the committed state. The setter re-points [0] and preserves [1:], because speculation scratch persists across forwards — step N's accepted slot is step N+1's initial state — so it belongs to the request, not to whichever slot it currently commits into. `--state-checkpoint-groups` still parses, as an alias. DeepSeek-V4 is unaffected: it declares `entries_per_req=1`, so a slot and a group were always the same thing there. Tests: new `TestPerNeedWidth` runs at `entries_per_req=3`, where a request and a checkpoint differ in width and the old code could not tell them apart — every pre-existing case ran at 1, where nothing distinguishes them. Full suite failure set is unchanged from before this series (192, all pre-existing). Co-Authored-By: Claude <noreply@anthropic.com>
`prepare_state_indices` is the one place a slot number crosses from the scheduler into a kernel, and nothing downstream can catch a mistake in it: every value is a valid index into the pool, so a wrong one reads another request's state and produces plausible tokens. No shape check, no NaN, no assert fires. It was also the site that manufactured the contiguity the previous commit removed, and it had no direct coverage — the only test naming it was asserting the old `group * (1 + num_spec)` arithmetic from a docstring. Pins the slot list going through verbatim, column 0 as the committed state on both paths, and a fork reading its source while writing its fresh slot. The spec-path cases use a scattered `[4, 1, 6]` on purpose: a `base + i` implementation writes `[4, 5, 6]` there and passes everything else in the file. Verified by reintroducing that exact bug — two cases fail, and only those two. Runs on CPU: `.np` is the host staging array, and the copy to device is a separate step that is not what could be wrong here. Co-Authored-By: Claude <noreply@anthropic.com>
Comments still saying "group" where the code now hands out slots. Stale names mislead future readers, and these were the ones that would have made someone believe there is still a fixed-width unit. The one remaining "groups" in state_pool.py is deliberate: it names the concept being contrasted against. Co-Authored-By: Claude <noreply@anthropic.com>
…code path The persistent aiter MLA decode asks the metadata kernel for a global KV-split budget of min(cu_num, max_split_per_batch * batch_size), and we hardcoded max_split_per_batch=16. At batch_size=1 that caps a decode at 16 of the 256 gfx950 CUs. A batch-1 decode has no parallelism other than the KV walk, and this kernel streams the whole context per layer per step, so the cap was an occupancy ceiling rather than a tuning choice. ATOM_MLA_MAX_SPLIT_PER_BATCH exposes it, defaulting to 16 so an unset environment behaves exactly as before. The raised cap is also threaded into get_mla_metadata_info_v1: mla_decode_fwd sizes its fp32 logits from reduce_partial_map.size(0), so a runtime cap above the sized one would write out of bounds -- sizing and runtime have to quote the same number. Measured on Kimi-K3 MXFP4, 8x MI355X, tp8, DSpark N=2, the InferenceX AgentX trace replay at concurrency 1, with ATOM_MLA_MAX_SPLIT_PER_BATCH=256: mla_a8w8_qh16_qseqlen4_gqaratio16_v3_ps 416.7us -> 48.0us per call paired kn_mla_reduce_v1 7.8us -> 41.0us per call aggregate ITL 12.68ms -> 9.76ms (1.30x) replay of a fixed 168-request set 3,358s -> 2,813s (1.19x) Output sequence lengths are bit-identical to the unpatched run on every request. The reduce gets more partials to combine, so the win narrows as batch_size grows and the budget stops binding; 16 stays the default for that reason. ATOM_MLA_KV_GRANULARITY comes along as an escape hatch only. Lowering granularity bought no ITL on its own, and stacked on the raised budget it added 2.5% ITL while costing 34% on TTFT, because prefill already has query-dimension parallelism and splitting its KV further only pays for a bigger reduce. It defaults to off.
…ecode
Speculative decoding made the GDN/KDA state pool scale with the verify
window: every request reserved `1 + num_speculative_tokens` copies of the
recurrent state so a rejected draft could roll back by resuming from a
different slot. On Kimi-K3 + DSpark(7) at tp=8 that is 8 slots x 57.81 MiB
= 28.91 GiB per rank, more than half the KV budget, leaving the paged
cache with less memory than the rollback machinery.
ReplaySSM keeps ONE checkpoint per request plus a small ring of the SSM
inputs `(k, u, g)`, and reconstructs any intermediate state on demand:
S_h = exp(G_h)*S_0 + sum_j exp(G_h - G_j) * k_j (x) u_j, G_j = sum_{i<=j} g_i
Rollback stops being a state copy and becomes a cursor move: the records a
rejected draft wrote are simply never read, and the next step overwrites
them. `slots_per_req()` drops from `1 + num_spec` to 1.
The buffer stores `u` (post-delta-correction), not raw `v`. Storing `v`
would leave `u_j = beta_j (v_j - S_{j-1}^T k_j)` depending on the previous
state, making the replay irreducibly serial; with `u` the expansion above
is a plain weighted outer-product sum.
Divergence from the reference ReplaySSM: the checkpoint absorbs committed
records at the START of a step rather than the end. Folding at the end
strands the step's speculative records at a non-zero offset and forces a
true ring with modular indexing; folding at the start means the buffer
always refills from offset 0 -- a linear buffer with reset, same state
traffic and flush cadence, no wrap-around arithmetic.
Flush fires at `h + 2T > L` rather than `h + T > L`, one window early.
That keeps a full window free (a step landing at h = L-T followed by a
full accept would otherwise truncate the next window to one draft) and
guarantees `h + T <= L`, so appends never run off the end. Hence
`cache_len >= 2*(mtp_k+1)`, raised automatically with a warning if the
env var asks for less.
Off by default; enable with ATOM_ENABLE_REPLAYSSM=1.
Measured on Kimi-K3 + DSpark(--num-speculative-tokens 7), tp=8, MI355X:
state pool 28.91 -> 4.42 GiB per rank (-24.49 GiB, ~196 GiB per node)
KV blocks 8325 -> 18813 (2.26x)
GSM8K 0.9515 -> 0.9530 (3-shot, full 1319, z=0.18)
acceptance 47.64% -> 47.37% over 217k draft tokens (z=1.73)
throughput parity -- interleaved multi-instance A/B at conc=64 gives
-0.69% at ISL/OSL 512/512 and +1.22% at 1024/1024 (t=0.50),
both smaller than the spread BETWEEN base instances
Also covers the GDN backbone (Qwen3.5 / Qwen3-Next): validated at mtp_k=3
with -22.72 GiB at max_num_seqs=256 and GSM8K within 0.61 sigma.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`prepare_state_indices` grew a `self.replayssm` fork when ReplaySSM landed, but this file builds its subject with `object.__new__` and hands it only the three staging arrays — so every spec-path case died on AttributeError rather than testing anything. Production sets the attribute in `__init__`; the stub is what was stale. Adds the ReplaySSM spec branch to the cases, which had no coverage at all: that a request holding one slot does not fan out across the row, that the 1-D tensor is filled here too (`_attach_replayssm` reads `slot_idx` off it even on the spec path), and the contrast case that the baseline still fans out. Co-Authored-By: Claude <noreply@anthropic.com>
The anchor seatbelt added in #1861 read `self.model.args.vocab_size`, but `args` is set only by the DeepSeek-V4 draft (deepseek_v4_dspark.py:869). KimiK3DSpark carries `hf_config` instead, so every TP rank died at warmup: AttributeError: 'KimiK3DSpark' object has no attribute 'args' dspark_proposer.py:427 -> model_runner.py:1238 warmup_model The proposer is shared by both flavors, so the bound has to be too. What the clamp protects is `markov_w1`, a raw nn.Embedding inside DSparkMarkovHead -- and both drafts build that head with its own `vocab_size`. Reading it there is both flavor-independent and the table actually being indexed. Co-Authored-By: Claude <noreply@anthropic.com>
…base
Brings `zejun/state-cache-lmcache-load`'s state-checkpoint offload tier onto
this branch: an evicted GDN/KDA checkpoint is spilled to LMCache and fetched
back when a later prefix hit lands on that boundary. Off unless `OFFLOAD_STATE`
is set, and it additionally requires `--kv-transfer-config` to name the
`lmcache_offload` connector.
Only the offload work is taken. The source branch also carries ~57 commits'
worth of unrelated main history; none of it is here, so `OFFLOAD_STATE` unset
should leave this branch's behaviour as it was.
The one real adaptation: a checkpoint is a SLOT here, not a GROUP
-----------------------------------------------------------------
The source branch addresses per-request state as a group of `1 + num_spec`
contiguous rows and indexes the staging ring as `num_groups + slot`. This
branch allocates slots individually from a free list, so slot indices are
arbitrary and a checkpoint occupies exactly one of them -- the committed state,
without the speculation-rollback rows beside it.
So `state_entry_views` takes a slot and returns one row per (cache, layer),
which is what the Triton staging packer needs (it refuses a strided view). Three
consequences, all in this branch's favour:
* a resume moves 57.4MB rather than ~172MB,
* the staging ring is counted in slots, one row each, so the "ring sized in
bare entries runs off the end of the arena" hazard cannot arise,
* `--state-checkpoint-slots` stays the single reader of checkpoint capacity;
the source branch's `STATE_CKPT_EXTRA_ENTRIES` override is deliberately not
carried, because two readers of one number is the bug its own docstring
warns about.
Spills ride `StateMaintenanceOps.spills` rather than a second list beside it,
which is what this branch's batch struct wants and also makes the ordering
guarantee a property of one consumer: `CommonAttentionBuilder.build` issues
spill copies BEFORE the batch's other state copies. That order is load-bearing.
`pop()` spills the slot it is about to hand out, so a slot returned as a
checkpoint destination is frequently the same step's spill source; issue the
copies first and the tier stores the new occupant's bytes under the evicted
checkpoint's hash -- present, valid-looking, someone else's, and undetectable on
any later load.
`pop_many` grew `spill_first`, which reaches element 0 only: that is the
committed slot and the only one an offload load ever writes into. The rollback
slots are scratch, so their evictions spill normally.
`_commit_pending` now releases its destinations after the loop rather than
inside it. Both branches touched these lines; this half has to win because
`pop()` stages what it evicts, so a `dst` released and indexed mid-loop can be
popped again by a later iteration, evicted, and spilled -- while its own
`src -> dst` copy is still queued, leaving the previous occupant's state under
the new hash.
Verification
------------
* CPU suite: 1619 passed. The 4 failures are `test_fused_compress_ragged`, which
fails identically on an unmodified `k3-dev-rebase` worktree.
* On MI355X (Kimi-K3, TP8, fp8 KV, DSpark N=2): the tier builds, refuses
correctly when the shared GPU staging buffer is smaller than one state entry,
and spills under load -- 90 spills requested, 89 indexed, 0 dropped, 0
forgone. Output verified correct against a known answer.
* The load direction has been exercised only as far as the first attempt; a
full-length run with loads completing is still outstanding.
Three bugs found on hardware, all in code that merged cleanly and so was never
read during conflict resolution:
* `state_spills_for_batch` addressed the ring as `cache.num_groups`; this
branch's pool is a `StateSlotPool` with `num_slots`. Fires on the first
spill only.
* `_request_state_load` read `seq.per_req_cache_group`; the field here is
`seq.state_slot`. Fires the first time the tier tries to serve a load, which
is why `loads_attempted` was still 0 when it crashed.
* `state_pool.py` ended up with two definitions each of `_resumable_from`,
`_spill`, `has_pending_spill` and `take_spill_copies` -- the branch's
group-based ones and this port's slot-based ones. Python kept the last, so
the slot versions were live and spilling worked, but the dead copies would
have taken over silently on any reordering. Removed.
Note for operators: a crash with the tier on wedges the process. The executors
are joined at interpreter exit and nothing bounds that wait, so the ranks sit in
`threading._shutdown()` holding their VRAM, and only SIGKILL on the rank
processes themselves (`ATOM::TP*`) releases it.
…d them
All three are the same shape: the ported code reads an attribute this branch
does not define, in a file that merged CLEANLY and so was never opened during
conflict resolution.
* `_request_state_load` read `seq.per_req_cache_group`; the field here is
`seq.state_slot`. Fires the first time the tier tries to SERVE a load, which
is why `loads_attempted` was still 0 when the engine died.
* `settle_state_load` and the abandon path read `self._orphan_load_groups`;
the resolution had renamed it `_orphan_load_slots`. Fires when a load report
lands for a request deallocated mid-flight. The local variable is renamed
with it -- swapping only the attribute would leave `group` undefined on the
next line.
* `self._warned_no_checkpoint_fates` lost its initialiser: it shared a
conflict block with the tier construction and the resolution kept this
branch's `state_caches` line while dropping it.
`state_pool.py` also ended up with two definitions each of `_resumable_from`,
`_spill`, `has_pending_spill` and `take_spill_copies` -- the source branch's
group-based ones, which merged cleanly, and this port's slot-based ones. Python
keeps the last, so the slot versions were live and spilling worked, but the dead
copies read as the truth to anyone opening the file and would have taken over
silently on any reordering. Removed.
The last of these had not crashed and would not have for a while, which is the
point: a keyword list cannot find a name the port itself invented. They came out
of comparing, per file, the `self.X` names that are READ against those that are
ASSIGNED or defined, and reading whatever is only ever read. That audit is now
empty across the six offload-touched files, and it is the check to re-run after
any further merge from the source branch.
On hardware (MI355X, Kimi-K3, TP8, fp8 KV, DSpark N=2, concurrency 10) the tier
now runs end to end: 141 spills requested / 140 indexed / 5 dropped, and 5 loads
attempted / 5 completed / 0 failed -- checkpoints leaving GPU state slots and
coming back from the CPU tier. CPU suite unchanged at 1619 passed, with the
same 4 `test_fused_compress_ragged` failures an unmodified worktree has.
…equest state
`_decide_load_after_alloc` refuses the offload KV load for every sequence with
`has_per_req_cache`, unconditionally and for a correctness reason: the KV leg's
boundary comes from LMCache's chunk-floored `lookup()`, the state leg's from
`_gated_hit` snapped to a checkpoint rung, and `L > P` is silent wrong output.
So on a hybrid every byte this connector wrote to the CPU tier was a byte it
would never read. Measured on Kimi-K3: 63.7 MB stored per 4096-token chunk, per
rank, for no reader at all.
Gated at REGISTRATION, not at emission, and that distinction is the whole bug
this went through first. Skipping a tracked sequence in the loop that emits
saves leaves it in `_save_tracker`, and
_has_pending_save(seq) -> _save_frontier(seq) > entry[1]
compares the frontier the prompt keeps advancing against a saved offset that now
never moves. It is true forever, `should_defer_free` is built on it, so the
request's blocks are never freed: the paged pool leaks and the engine crawls to
a halt (observed as a warmup stuck at 21/109 with 13 in flight). Never entering
the tracker makes `_save_tracker.get(sid)` return None, which
`_has_pending_save` already reads as "nothing pending".
`OFFLOAD_SAVE_PER_REQ_CACHE=1` restores the old behaviour. It is only useful
when a separate stateless consumer shares this LMCache instance -- there is no
such reader in tree.
The state tier is untouched: its spills ride `StateMaintenanceOps.spills` into
`StateOffloadTier.submit_spill` and never pass through `_save_tracker`. A state
load needs its KV prefix resident in HBM, not in LMCache, so nothing here
narrows what the tier can recover.
Measured, Kimi-K3 TP8 fp8 KV DSpark N=2, agentic replay at concurrency 10,
3600s, same node, same flags, this switch the only difference:
tok/s/GPU hit 1/ITL_p90 ITL p50 TTFT p90
tier on, KV save on 4131.5 92.28% 28.65 16.27ms 9646ms
tier on, KV save off 4489.9 92.98% 37.48 15.13ms 5371ms
and against the best configuration without the tier at all (96 checkpoint
slots, which is where that curve peaks):
no tier, slots 96 4236.2 89.24% 34.00 15.22ms 8407ms
+6.0% throughput and +10.2% interactivity over the no-tier optimum. The ITL p50
is the tell: it returns from 16.27ms to 15.13ms, i.e. the saves were on the
per-step critical path, which is also why TTFT p90 falls by a third.
The tier's own configuration matters and is the opposite of the no-tier one:
the no-tier curve peaks at 96 checkpoint slots, the tier's at 16, because slots
it does not need are bytes the paged pool does. Spilling only pays once the
paged pool is large enough that a prefix outlives its checkpoint --
`checkpoints_orphaned` falls 433 -> 62 across that range, and completed loads
rise 3 -> 72.
CPU suite: 1619 passed, with the same 4 `test_fused_compress_ragged` failures an
unmodified worktree has.
gbyu-amd
marked this pull request as draft
August 18, 2026 04:57
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
- atom_kv_byte_codec.py: iterate kv_caches.values() (key unused) (PERF102) - block_manager.py: collapse nested if into single condition (SIM102) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…e to build A state entry is one entry; the GPU staging buffer is sized in LMCache chunks, and at the shipped `OFFLOAD_GPU_STAGING_CHUNKS=2` that is ~8MiB against a 54.78MiB entry on Kimi-K3. The worker then refused to build the tier -- one WARNING, after which the engine-side index carries on handing out staging slots and counting spills nobody performs. Every state number reads as zero and nothing says why. That is not a hypothetical: a c8 and a c10 measurement were taken this way before the warning was spotted, and both were reported as "lmcache brings no benefit at c8" and "the accuracy fix costs 7% throughput". Neither was true. Growing the shared buffer is the obvious alternative and it charges every KV worker thread for a size only the state thread needs. So the tier takes its own `StagedTransfer` of exactly one entry, on the `lmc-state` thread that packs into it, and the log says how to make the two share instead. Co-authored-by: Cursor <cursoragent@cursor.com>
The tier could bring a state checkpoint back but the connector refused the KV leg for every `has_per_req_cache` sequence, and refusing was right: the two legs are chosen by different matchers -- the KV leg from LMCache's chunk- floored `lookup()`, the state leg from `_gated_hit` snapped to a checkpoint rung -- so raising the KV-loaded length L past the state boundary P has the scheduler forward only `[L, num_prompt)` while the linear layers never see `[P, L)`. Wrong output, and silent. So one matcher picks a boundary for both. `BlockManager._joint_kv_boundary` takes the rightmost rung that is state-resumable, sits inside the LMCache KV prefix, and whose covering chunk LMCache holds; the connector clamps the KV leg to it. The state has to come from the tier rather than HBM, which is the ordinary case and not a lucky one: `unindex` spills a checkpoint whose KV blocks left HBM, so the state follows its KV down. Asymmetric by design where the grids disagree. A rung is a hash-block boundary and the KV leg moves whole chunks, so the KV leg is aimed at the chunk *covering* B while the request claims only B (`_claim_after_load`). Overshooting costs one chunk of transfer into blocks the forward rewrites; undershooting is the silent failure above. Both legs report through one set -- `get_finished` unions the tier's report into the KV channel -- so an id in both would collapse into a single wake and resume the suffix prefill while the other transfer is still writing. `_JointPark` (unused until now) holds those until both land; either leg failing fails the pair, and `num_cached_tokens` does not move until then, so any failure degrades to a recompute from the HBM prefix. Off by default. Measured at c10 on the agentic replay, it fires almost never: `joint kv: boundaries=1 | lmcache_within_hbm=1119`, i.e. in 1119 of 1130 admissions the HBM prefix cache already reached at least as far as LMCache's copy. The workload that needs this is one whose KV working set does not fit HBM; this one's does. Enabling it also requires `OFFLOAD_SAVE_PER_REQ_CACHE=1`, which costs ~8% throughput, so the pair stays opt-in. Every gate that declines counts itself (`joint kv:` in the periodic stats). The first attempt at this feature could not fire at all -- the chunk size came off the connector object, was 0 there, and the early return it took was the one path with no counter, so the log showed neither a boundary nor a decline. The chunk size now comes from the LMCache config where the config is. Co-authored-by: Cursor <cursoragent@cursor.com>
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
|
|
||
| sys.path.insert(0, str(Path(__file__).resolve().parent)) | ||
|
|
||
| from atom.model_engine.block_manager import BlockManager # noqa: E402 |
Contributor
| sys.path.insert(0, str(Path(__file__).resolve().parent)) | ||
|
|
||
| from atom.model_engine.block_manager import BlockManager # noqa: E402 | ||
| from atom.model_engine.sequence import Sequence # noqa: E402 |
Contributor
|
|
||
| from atom.model_engine.block_manager import BlockManager # noqa: E402 | ||
| from atom.model_engine.sequence import Sequence # noqa: E402 | ||
| from atom.model_engine.state_runtime import ( # noqa: E402 |
Contributor
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist