Conversation
yhl-amd
pushed a commit
that referenced
this pull request
Aug 31, 2026
* feat(state-cache): per-slot allocation, a switchable demand rung, and a -1 interval Rebased onto main's PAGE-backed checkpoint work (ROCm#1874, ROCm#1894, ROCm#1943, ROCm#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 `ROCm#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 ROCm#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. **ROCm#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>
yhl-amd
pushed a commit
that referenced
this pull request
Sep 2, 2026
…and full attention (ROCm#2089) * [feat](k3): fuse activation+fp8 per-token quant for dense/shared MLP and full attention Two ptpc_fp8 fusions for Kimi-K3, each collapsing an activation kernel and a separate per-token FP8 quant into a single pass so the bf16 activation never round-trips through HBM. Both gate on the consuming Linear's runtime scheme (_effective_layer_quant == per_Token/fp8) and fall back to the unfused path otherwise. 1. Dense MLP + shared-expert down_proj: SiTUv2-and-mul + per-token fp8 quant via the aiter `situv2_and_mul_quant` kernel. New compile-safe wrapper `situ_and_mul_quant` (torch_compile_guard, returns (fp8, scale)); SituAndMul grows a `fused_quant` flag and KimiMLP forwards the scale into down_proj. Routed experts are excluded by the online-quant globs; only dense + shared (which resolve to per_Token/fp8) take the fused path. 2. KimiFullAttention o_proj input: sigmoid(g_proj) * attn_out + per-token fp8 quant. Adds a `per_token=True` mode to the existing triton_fused_sigmoid_mul_quant kernel (previously per-1x128 only, used by qwen3_next), returning a single [M,1] scale; KimiFullAttention.forward feeds (fp8, scale) straight to o_proj. Validation: - Kernel UTs: scale err ~1e-9, fp8 dequant rel err ~3%, zero/empty inputs clean. - Device time (torch.profiler per-kernel, reconciled against a kineto trace): ~1.7x (decode) to ~2.6x (prefill) for #1; ~1.5x-2x for #2. - GSM8K: 0.9674 strict / 0.9666 flexible (tp8, ptpc_fp8) -- no regression. Depends on aiter `situv2_and_mul_quant` (branch situv2-and-mul-quant); no-ops gracefully where the scheme doesn't match if that kernel is absent at build. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * [refactor](k3): unify sigmoid-mul quant path behind one QuantType-selected helper Address review feedback (valarLip): collapse the two-branch KimiFullAttention.forward into a single `o_proj(x, x_scale=scale)` call site. - fused_sigmoid_mul_maybe_quant(attn, gate, quant, quant_type): returns (fp8, scale) when quant=True and (bf16, None) otherwise, so o_proj always takes the same x_scale= path (None -> o_proj self-quantizes). - quant_type (QuantType.per_Token or QuantType.per_1x128) names the scheme explicitly instead of a per_token bool; per-1x128 delegates to the existing block-quant path, unsupported types raise. K3's o_proj fusion gate now also engages for per-1x128 block fp8, not just per-token ptpc_fp8. - Import hoisted to the module top-level (no circular dep). No behavior change for K3's per-token path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [refactor](k3): unify SiTUv2 quant/non-quant path behind situ_and_mul_maybe_quant Mirror the fusion #2 unification onto fusion #1: SituAndMul.forward now always returns (activation, scale) — (fp8, scale) when it fuses SiTUv2 + per-token FP8 quant, else (bf16, None) — so KimiMLP takes a single down_proj(x, x_scale=scale) call site with no isinstance branch. Unlike fused_sigmoid_mul_maybe_quant there is no scheme selector: the aiter situv2_and_mul_quant kernel implements only the per-token scheme. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [fix](k3): guard fused sigmoid-mul-quant triton launchers for torch.compile The two triton launchers in triton_fused_sigmoid_mul_quant.py launched with a raw grid=(M, ...) inside the @support_torch_compile Kimi-K3 forward. With no compile guard, Dynamo traced into them and specialized the token dim M to a constant (16384 = max_num_batched_tokens). After main ROCm#2088 ("Settle a step's shape once") newly marks that dim dynamic, warmup compile raised ConstraintViolationError on hidden_states.size()[0], failing K3 accuracy CI. Wrap both launchers as torch_compile_guard custom ops with matching fakes, mirroring situ_and_mul_quant in kimi_k3/activations.py: - _fused_sigmoid_mul_fp8_per_token_quant (the path K3 actually uses) - _fused_sigmoid_mul_fp8_group_quant (per-1x128; extracted from the public fused_sigmoid_mul_fp8_quant, which now resolves the env transpose_scale to a concrete bool then dispatches) Guarded ops are opaque to Dynamo, so the raw triton launch never forces M to a constant. Non-quant path was already safe (pure torch). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Revert "[fix](k3): guard fused sigmoid-mul-quant triton launchers for torch.compile" This reverts commit fccfcc8. The guards did not fix the K3 accuracy CI failure: CI reproduced the same ConstraintViolationError with them in place. The triton launchers were never what specialized the token dim; the real cause is fixed in the next commit. Unguarded triton launchers inside a @support_torch_compile forward are the norm here (rmsnorm_nw in deepseek_v4, triton_mrope in qwen3_next). Co-Authored-By: Claude <noreply@anthropic.com> * [fix](k3): pass o_proj quant scheme into forward as QuantType.value KimiFullAttention stored the raw QuantType and passed it into fused_sigmoid_mul_maybe_quant, which compared it against QuantType.per_Token inside the traced forward. QuantType is a pybind11 type, so Dynamo cannot compare two of them: Failed to compare UserDefinedObjectVariable(QuantType) with UserDefinedObjectVariable(QuantType), because ... is not a Python constant That graph break split KimiModel.forward, so AttentionResidual.forward compiled as its own frame, where block_residual re-entered as a fresh input with a concrete token dim. Its torch.cat then pinned the dim to max_num_batched_tokens, raising ConstraintViolationError on hidden_states.size()[0] at warmup and failing K3 accuracy CI. Store o_type.value and compare .value on both sides, matching qknorm_quant_type_value in the same __init__ (and deepseek_v2, llama, activation.py). Verified with my_script/run_atom.sh: no ConstraintViolationError, and the graph-break log is clean -- compile now traces past the attention layer. Co-Authored-By: Claude <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: ganyi <ygan@amd.com>
This branch has not been deployed
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
1 participant
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.
Summary
fastpath=fused_chunkwhenOFFLOAD_NATIVE_KV_STAGING=1fastpath=chunkrequest-level segment-major path as fallbackValidation
python3 -m py_compile atom/kv_transfer/offload/*.pypython3 -m pytest tests/test_lmcache_offload_connector.py -q-> 19 passed, 22 skippedyhl_kvoff_009:python3 -m py_compile atom/kv_transfer/offload/*.py && python3 -m pytest tests/test_lmcache_offload_connector.py -q-> 41 passedOFFLOAD_NATIVE_KV_STAGING=1 OFFLOAD_CODEC_LAYOUT=segment_indexed->native_kv_staging smoke okgit diff --checkBench
Config: MiniMax-M2.5, TP=2, prefix cache on, chunked prefill on,
MAXBATCH=16384,OFFLOAD_REQUEST_FASTPATH=0,OFFLOAD_NATIVE_STITCH=1,OFFLOAD_NATIVE_KV_STAGING=1,OFFLOAD_MIN_LOAD_TOKENS=1024.fastpath=fused_chunk=80, no failuresfastpath=fused_chunk=86, no failuresDetailed results are recorded in
/shared/amdgpu/home/hyi_qle/yhl/project/009-kv-off-llmcache/18_SUMMARY_009_kv_offload.mdandcode-review/lmcache-compatible-fused-staging-plan.md.