llama : add opt-in boundary aligned deep rollback for DSV4 recurrent state - #26528
Closed
TacoTakumi wants to merge 13 commits into
Closed
llama : add opt-in boundary aligned deep rollback for DSV4 recurrent state#26528TacoTakumi wants to merge 13 commits into
TacoTakumi wants to merge 13 commits into
Conversation
…state rollback test - allow DEEPSEEK4 in test-llama-archs: supply the arch-specific metadata (compress ratios 0/4/128, MQA dims, sqrtsoftplus gating, hyper-connection count 4) and drop the refusal - model saver: write the DSV4 hparams and hc head tensors, and trim the per-layer swiglu clamp arrays to block_count so the saved model reloads - run test-recurrent-state-rollback against the generated DSV4 fixture (CPU only: with rollback state enabled the per-plane inputs overflow GGML_SCHED_MAX_SPLIT_INPUTS when one device takes the whole graph)
…e fix) Four cases pinning the deep-rewind contract, appended after the existing per-token cases which stay untouched and green: - aligned deep rollback beyond n_rs_seq replays reference-equal - a rollback the raw SWA window cannot cover is refused or falls back, never silently applied with wrong output - a second seq_rm before any decode is rejected or cumulative-correct - a failed decode does not lose or corrupt an armed rollback All four fail on the current base: the deep-rollback capability is absent, stacked rollbacks are accepted beyond plane coverage with silently wrong replay, and nothing guards window coverage. The DEEPSEEK4 fixture now generates with weight sigma 0.1: at the default 0.01 its attention output is context-inert (history perturbation moves logits ~9e-8, under the comparison eps), which made context-content comparisons vacuous. The backend NMSE comparison keeps 0.01 for its own models. generate-models is seeded so logits-vs-reference outcomes are deterministic.
… rollback Adds the context param and its plumbing only: the value is validated (requires n_rs_seq > 0, else clamped to 0), stored in cparams, and exposed via llama_n_rs_aligned(). Nothing consumes it yet - the boundary-aligned rollback tier that allocates and uses the slots lands in the next commit. Split out so the DSV4 deep-rollback regression tests, which request aligned slots, can be applied to the merge base and run there to demonstrate the failures the tier fixes.
…ate rollback Extends the rs_rollback framework with a second tier of rollback planes holding compressor-ring snapshots taken at HCA (128-token) boundaries, so seq_rm can accept targets far deeper than the per-token tier: - the n_rs_aligned context param (added in the previous commit, requires n_rs_seq > 0) takes effect: CSA and LID states allocate that many extra boundary-aligned planes, HCA allocates none - an aligned replay rewrites the whole HCA ring before its next commit, so it has nothing to restore - at an aligned boundary the only state that does not self-heal on replay is the overlap compressor's previous-window rows; the boundary snapshot is one in-graph plane copy per 128 tokens, added to the existing snapshot gather - the per-ubatch plane shift is untouched and per-step cost does not grow with aligned depth - seq_rm accepts p0 at a 128 boundary when the boundary's slot still holds its snapshot; the restore rides the existing marked in-graph restore path (rs_idx values past n_rs_seq select aligned planes) - rollback soundness guards: a pending un-consumed rollback rejects further rollback requests (stacking against shifted planes was silently wrong), and any rollback the raw SWA window can no longer cover is refused instead of silently replaying with reclaimed cells
…k cases The deep-case contexts now request n_rs_aligned slots, turning the four red cases green: aligned deep rollback replays reference-equal, and the stacked / uncovered / failed-decode cases observe the clean-refusal contracts.
…ubatch Plans for every ubatch of a batch were built from the same pre-reset rollback markers, so a replay batch split into multiple ubatches re-applied the restore in each of them, clobbering ring rows earlier ubatches had already persisted. Unreachable with the per-token tier (replay depth is at most n_rs_seq tokens) but real for aligned deep replays spanning hundreds of tokens.
The batch form must match the reference bit-for-bit; before the restore-once fix it diverged by 1.3e-3.
…q_rm_align Memories with a boundary-aligned deep-rollback tier accept partial sequence removal below the tip at aligned positions only; a consumer has to know the quantum to compute a target. Default is 1 (no aligned tier); the DSV4 cache reports its HCA block size when aligned slots are enabled.
…emories When the target context supports bounded partial removal and reports an alignment quantum above 1, the divergence path aligns the reuse point down to the quantum boundary and asks the memory for the rollback directly, instead of restoring a context checkpoint or re-processing the whole prompt. The memory stays the authority on coverage: a refused rollback falls through to the unchanged checkpoint machinery. The rollback is requested once and the removed suffix is re-decoded in the same pass; targets are kept strictly below the task size so the final-prompt re-eval never needs a second removal below an armed rollback. Contexts without the capability take the exact pre-existing path. The tier is enabled with --rs-aligned N (default 0 = off). Enabling it without speculative decoding raises n_rs_seq to 1, since the aligned tier rides the per-token rollback machinery and single-token re-eval removals below the tip need per-token depth 1. Draft contexts keep both tiers disabled.
…nding (red until the fix) The server prompt cache sizes and saves sequence state on task start, which can land while an aligned rollback marker is armed. States without aligned planes made the save throw. The case arms an aligned rollback, requires the size query to succeed, and requires the saved state to restore and replay within a round-trip tolerance: the save/load repack changes attention reduction order, which measures ~5e-5 here and hits the pre-existing per-token save path equally (~2.4e-5 on the same prompt), while real state corruption measures ~1e-3.
state_write selected the plane to save by the raw rollback marker and threw when the marker exceeded the state's own plane count - which every aligned marker does on a state without aligned slots (HCA), so any sequence-state save with an aligned rollback pending failed. Such a state has nothing snapshot-held to save: its rows are rewritten by the replay, exactly like the restore graph's self-copy rule for the same markers. Degrade the selection to the live plane instead of throwing.
…emoval memories The context-checkpoint restore path resumed at max(pos_min + 1, pos_max), making the post-restore memory_seq_rm a removal AT the restored tip whenever pos_min < pos_max. Aligned-removal memories (seq_rm_align > 1, i.e. the DSV4 compressed cache) cannot reliably serve that: restoring a sequence state invalidates the rollback history (the restored blob may describe a different branch than the snapshots, so state_read resets the rollback markers), leaving a dead tail strictly past pos_max as the only removal the memory is guaranteed to accept. Resume at pos_max + 1 instead - the same shape the speculative rollback already uses (restore, then seq_rm [pos_max + 1, end)) - and require checkpoint pos_max strictly below pos_next in the search so the resume point stays inside the verified common prefix and at least one token remains to decode [TAG_PROMPT_LOGITS]. Memories with seq_rm_align == 1 keep byte-identical behavior.
The DeepSeek-V4 graph with the recurrent-state rollback tiers active exceeds the scheduler's input cap. Each of the three compressed-cache paths (csa, hca, lid) carries a set of index/position input tensors for the rollback state machinery (state_pos, persist/restore/snapshot src+dst idxs, read/write idxs, write_pos) in addition to the per-path kq_mask and k_rot and the generic and raw-SWA inputs. Measured on a 5-GPU LAYER split with pipeline parallelism (all layers offloaded), at the worst-case shapes probed by graph_reserve during model load: n_graph_inputs max = 39 per-split n_inputs max = 44 A divergence-heavy workload exercising the aligned deep rollback and a sequence-state restore afterwards raised neither maximum - the reserve shapes are the binding case. Both counters assert with a strict less-than against the same macro, so 44 consumed inputs need a cap of at least 45. At the compiled default the load aborts: GGML_ASSERT(n_graph_inputs < GGML_SCHED_MAX_SPLIT_INPUTS) failed (ggml-backend.cpp:1356, via graph_reserve at model load) 48 is chosen as the floor of 45 plus one more consumed input per compressed path, should a shape not exercised here light one up. It is not a round number picked for headroom. The macro sizes two fixed pointer arrays, the split-boundary heuristic and the scheduler's reserved allocation expressions, so reserved scheduler memory scales linearly with it; the arrays are malloc/calloc backed and pages fault in on demand. There is no ABI exposure: ggml_backend_sched is defined in this translation unit. An alternative, if a global increase is unwelcome, is to expose this as a CMake cache variable the way GGML_SCHED_MAX_COPIES already is and leave the compiled default at 30.
TacoTakumi
requested review from
a team,
CISC,
JohannesGaessler and
ggerganov
as code owners
August 3, 2026 17:15
github-actions
Bot
added
testing
Author
|
@am17an why was this closed with no explanation? If you could offer me direction that would be great. |
Contributor
|
Because it's slop, please read the contributing guidelines. |
Author
|
@am17an Slop means lazy low effort AI output. This is weeks of work with a regression test, memory accounting, a 1200 iteration reproduction of the exact issue it fixes. You closed it immediately without evaluating it. I'm shocked at the hostility towards AI in an AI project. The AI disclosure is your policy and I followed it honestly. I know you're drowning in PRs but come on. |
Contributor
|
Consider this out of scope then. Thanks |
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
2 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.
Overview
DSV4's recurrent state can only roll back through the per token tier (llama_n_rs_seq), which is shallow. Any deeper chat turn divergence forces a full re-prefill which is the stall in #25452.
This PR snapshots the state every 128 tokens (the compressor boundary) into a fixed ring of n_rs_aligned slots. On divergence the server rolls back to the nearest boundary at or before the divergence point and delta decodes the rest.
Past the ring's reach it refuses and the existing checkpoint path takes over.
It is opt-in: --rs-aligned N (default 0 = disabled everywhere).
Non-DSV4 and default config behavior is unchanged. Every kernel path gates on n_rs_aligned, every server path on seq_rm_type == RS && seq_rm_align > 1. Draft contexts force 0.
The reach at --rs-aligned 64 is 8192 tokens (64 x 128).
The cost at depth 64 on DSV4-Flash is only 116.64 MiB.
Memory accounting: two identical server loads, with and without --rs-aligned 64 (DSV4-Flash, 5 GPUs, ctx 131072)
Overhead: 128.28 - 11.64 = 116.64 MiB of device-memory state buffers at depth 64.
Corroboration: total VRAM at server READY 85219 MiB vs 84979 MiB (+240 MiB, of
which +116.6 MiB is the state planes, the rest fit-placement variation).
This PR fixes #25452. The crash documented in the issue no longer reproduces on the merge base but the stall does. On the merge base 397 of 400 iterations recovered by re-prefilling from a checkpoint (the stall). On this branch the same workload run to 1200 iterations delta-decoded 1193 divergences with 4 checkpoint fallbacks.
This PR targets 128-aligned boundaries, not arbitrary positions.
The shipped test fails on base plus the test commits alone and passes on the tip.
This also passes on long context quality. A scripted 61-turn agentic session at ctx 65536 had no coherence failures, 4/4 content probes, zero HTTP errors.
Re speculative decoding coexistence, I ran a 150 iteration churn with the ngram drafter active (peak 56 drafted tokens) with no problems.
Additional information
One of my commits (ggml: raise GGML_SCHED_MAX_SPLIT_INPUTS 30 -> 48) is a global ggml constant bump. 30 aborts, 48 leaves margin over the measured 44. #25832 helped but this graph still exceeds the cap after it. I see #22789 merged after this branch's base. This commit drops out on a rebase.
Decode throughput with the tier armed (llama-bench + served-config smoke)
llama-bench cannot enable the tier from flags, so the candidate runs used a
local default-params patch arming exactly the served shape (planes verified in
the run logs). Baseline is the previously served build of the predecessor fix
branch, same model, same config, r=2:
Same-binary control (tier armed vs stock defaults, uniform fitt 2048): tg128
9.47 vs 9.77, a 3.1 percent cost at that placement, consistent with the
~1 GiB expert spill the state planes displace at that margin.
Served-config smoke with --rs-aligned 64: 913-token prefill at 82.32 t/s,
105-token decode at 9.48 t/s, inside the baseline band (9.25-9.69).
Reproduction workload from #25452 against this branch: 1200 churned chat-reuse iterations, ctx 16384
network/parse errors.
find_slot "failed to find free space" 0, CUDA error / ggml_abort /
GGML_ASSERT / out of memory 0, send_error 0, failed to remove sequence 0.
the suffix; 4 by the existing checkpoint fallback (the workload hard-cuts
~11k tokens there, beyond the 8192-token reach at depth 64 - the refusal
path, by design); 2 full re-prefills, both in the first two minutes after
launch, before any cache or checkpoint existed.
Fails-before / passes-after: how to reproduce the red state
The test commits are separable from the fix commits by construction. On the
merge base, cherry-pick the five tests: commits plus the inert plumbing commit
"llama: accept an n_rs_aligned context param for boundary-aligned deep
rollback" (the param is accepted but consumed by nothing, so the tree builds
and behavior is unchanged). The DSV4 rollback test then fails on exactly the
deep-tier cases, with the per-token checkpoint cases staying green:
The identical test at the branch tip passes.
Requirements