GDN ingredient-replay: cut MTP rollback VRAM - #331
Conversation
Extends ggml_gated_delta_net with an opt-in emit_mode param (0 = today's full K-snapshot output, default; 1 = small per-token (k,v,g,beta) ingredients plus a fixed-cost trailing final-state block). Replaying the ingredients through a K=1 call reconstructs the same state as a full-snapshot run at O(S_v) storage per retained step instead of O(S_v^2) -- verified bit-exact on CPU (tests/test-gdn-ingredient-replay.cpp) and numerically matching on CUDA via test-backend-ops' existing CPU-vs-CUDA diff (all 4 new emit_mode=1 cases pass on first build). emit_mode=0 is unchanged for every existing caller; this is purely additive. ggml_cuda_try_gdn_cache_fusion gets an explicit emit_mode!=0 guard (previously safe only by shape coincidence). Vulkan's debug-clone call site updated for the new constructor signature; Vulkan itself does not yet implement emit_mode=1 (its own supports_op path is unaffected since that op-clone path is compile-time-gated debug tooling, not the normal dispatch). Also adds llama_hparams::n_embd_s_ingredient(), sized off ssm_d_inner to match the mamba-branch derivation qwen35/qwen35moe actually use for GDN dimensions (S_v = ssm_d_inner / ssm_dt_rank), not the unused KDA branch. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Wires the new gated_delta_net emit_mode=1 into the real MTP speculative-
decoding rollback path for GDN/delta-net models (qwen35, qwen35moe),
opt-in via LLAMA_GDN_REPLAY=1 (env var for now; defaults off, zero
behavior change otherwise).
llama_memory_recurrent: replaces the (1+n_rs_seq)-wide state snapshot
tensor with one authoritative state row (s_l) plus a small ingredient
ring (ingr_l) when gdn_replay is on. seq_rm's bounded-rollback branch
marks a pending replay length instead of the old free index rewind.
Also persists a second checkpoint, s_ckpt_l ("state as of n_rs_seq
tokens before the end of the last decode"). This was not in the
original design and is required for correctness: at graph-build time
it isn't yet known which trailing tokens will be accepted, so the
"optimistic" final state in s_l cannot serve as a replay base once a
rollback is discovered -- delta-net's rank-1 update has no inverse, so
replay must start from a checkpoint that predates the whole uncertain
window and replay forward through only the accepted ingredients.
llm_graph_input_rs::can_reuse now invalidates the graph when the
pending replay length changes, since a replay reconstruction is a
topology change (a new ggml_gated_delta_net(K=1) node sized by however
many steps need replaying), not swappable input data the way the old
snapshot-index rollback was.
build_recurrent_attn's gdn_replay branch: on a pending replay, chains
single-token K=1 calls from s_ckpt_l through the accepted ingredients
to reconstruct the correct state before the main call; the main call
itself switches to emit_mode=1 and updates ingr_l/s_l/s_ckpt_l for the
next round.
Verified end-to-end through real llama_decode calls (not just the ggml
op level): test-recurrent-state-rollback passes with the flag on
(including with -ngl 999, genuine GPU execution of the GDN math),
comparing logits between a rolled-back-and-replayed context and an
independently checkpoint-restored one to 1e-5, and passes unchanged
with the flag off. Full test-backend-ops GATED_DELTA_NET suite (CPU
and CUDA) unaffected.
Measured against the real Qwen3.8-27B model (--spec-chain 2/4/6/8):
VRAM stays flat under gdn_replay while the old design grows linearly
with depth (939 MiB saved at depth 8) -- the predicted win, confirmed.
Latency is currently a small regression (+0.6 to +2.0 ms/round): the
checkpoint-update logic's "batch longer than retained window" branch
was assumed to be the rare case, but the verify batch is always
exactly one token longer than the window by construction, so it fires
on every decode rather than occasionally. Root cause understood, fix
pending in a follow-up.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Both found via the real-model --spec-chain 2/4/6/8 benchmark, not synthetic tests -- neither changes output values (bit-exact before and after on the CPU reference, and draft_n/n_rounds identical on the real model, confirming these are pure perf fixes). 1. build_recurrent_attn's checkpoint-update logic assumed "batch longer than the retained window" was the rare case. It isn't: the verify batch is always exactly one token longer than the window by construction (n_draft + 1 vs n_rs_seq = n_draft), so the old code fired an extra full ggml_gated_delta_net(K=1) call to recompute the prefix on every single decode, not occasionally. Fixed at the op level: emit_mode=1 now also captures the state immediately before the K-token window starts as a further fixed-cost trailing block, whenever n_tokens > K (ggml.h). The recurrence already passes through that exact value on its way to the final state, so exposing it is a free byproduct of the existing loop, not a second kernel launch. build_recurrent_attn now reads it via a view instead of an extra op call. 2. The CUDA ingredient-write path had every column-owning warp (up to S_v of them) redundantly write the column-independent k/g/beta values, turning what should have been O(S_v) traffic per component into O(S_v^2) -- the same order as a full snapshot write this design exists to avoid, done three times over. Fixed by gating those writes to a single warp (col == 0); v's write (genuinely per-column) is unaffected. Measured impact against the real Qwen3.8-27B model, ms/round vs the --spec-chain-off baseline (was / now): depth 2: +0.65 -> +0.26 depth 4: +0.70 -> +0.28 depth 6: +0.61 -> +0.11 (near parity) depth 8: +1.95 -> +1.18 The regression is substantially smaller at every depth; depth 6 is close to parity. Depth 8 still has a real gap, most likely the replay chain's own per-launch kernel overhead (which scales with how many ingredient steps get replayed) -- a further fix would need batching that chain into fewer launches, which was deliberately deferred during the original design to avoid a negative-stride view over the ingredient ring; not attempted here. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Threads a real experimental flag through the full chain, mirroring
exactly how n_rs_seq itself is already threaded (same call sites, same
llama_memory_hybrid[_iswa] -> llama_memory_recurrent forwarding):
common_params.gdn_replay (--gdn-replay / LLAMA_ARG_GDN_REPLAY)
-> llama_context_params.gdn_replay [EXPERIMENTAL] (public API)
-> llama_cparams.gdn_replay
-> llama_memory_recurrent's new gdn_replay_req constructor param
(via llama_memory_hybrid / llama_memory_hybrid_iswa for the
real qwen35/qwen35moe hybrid-arch path)
llama_memory_recurrent keeps the original LLAMA_GDN_REPLAY env var too,
OR'd with the new flag, as a fallback for quick testing without
touching CLI args -- this is what every test and benchmark in the
earlier DRC commits already used, and changing that mechanism now
would risk invalidating already-verified results for no benefit.
Verified end-to-end with the real flag (no env var) against the real
Qwen3.8-27B model via llama-server: loads, serves a completion
correctly with --gdn-replay --spec-chain 4. Full existing regression
suite (test-recurrent-state-rollback on/off/GPU, test-backend-ops
CPU+CUDA, test-gdn-ingredient-replay) still passes unchanged, and the
full project (all 290 build targets) builds clean.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Chases the remaining latency gap at --spec-chain 8 after the previous perf-fix commit. Root cause: emit_mode=1 ingredients were stored most-recent-first (matching emit_mode=0's convention), so replaying the accepted prefix -- which is the *oldest* end of the retained window -- required either a negative-stride view or, as built previously, a chain of m separate single-token K=1 kernel launches. That chain's per-launch overhead scaled with how many steps needed replaying, hence the largest remaining gap at the deepest tested depth (where rollbacks tend to be longest). Fix: emit_mode=1 now writes ingredients in chronological order (slot 0 = oldest of the K retained tokens) instead of most-recent-first. This convention is private to this op's own replay call site -- unlike emit_mode=0's snapshots, nothing else (s_copy/rs_idx row selection) depends on emit_mode=1's slot ordering -- so changing it carries no compatibility cost. With chronological storage, "the accepted (K - replay_len) ingredients from the oldest end" is exactly slots [0, K - replay_len): a single contiguous, forward-order view, replayed in ONE multi-token ggml_gated_delta_net(K=1) call instead of a chain. Verified: tests/test-gdn-ingredient-replay.cpp's replay check (already using a single batched call at the pure-op level, now updated to the new slot-order convention) stays bit-exact. test-recurrent-state-rollback still passes on/off/GPU -- the real end-to-end logit-comparison check that already caught one design flaw earlier in this effort. Full test-backend-ops GATED_DELTA_NET suite (CPU 40/40, CUDA 40/40) unaffected. Measured impact against the real Qwen3.8-27B model, ms/round vs. the --spec-chain-off baseline (previous perf-fix commit -> now): depth 2: +0.26 -> +0.31 (noise-level) depth 4: +0.28 -> +0.22 depth 6: +0.11 -> +0.32 (noise-level; see below) depth 8: +1.18 -> +0.77 -- confirmed reproducible (re-ran: +0.81) Depth 8 (this fix's actual target, since it has the longest average replay chains) improved consistently and reproducibly; depths 2/4/6's small movements are within the ~0.3-0.4ms run-to-run variance already observed on unchanged configurations earlier in this benchmark series, not attributed to this change specifically. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…rify Hand-ports the CPU/CUDA ingredient-replay layout (k,v,g,beta ingredients, trailing final-state block, conditional checkpoint block) to the Metal kernel, including the same column-invariant-write fix CUDA needed (k/g/beta don't vary per row, so only the row-0 threadgroup writes them). Unbuilt and unverified: this dev machine is Linux with no Metal toolchain, and Metal only compiles on macOS/iOS. supports_op keeps an explicit `emit_mode != 0 -> false` guard so the untested path is unreachable until someone with real Apple hardware verifies it and flips the guard.
Same ingredient-replay port as Metal/CUDA (k,v,g,beta ingredients, trailing final-state block, conditional checkpoint block), threaded in as a new push-constant field since it doesn't affect the S_v/K/KDA-keyed pipeline specialization. Unlike Metal, this was actually built and run: a separate -DGGML_VULKAN=ON tree against the real RTX 5090 passes test-backend-ops GATED_DELTA_NET 40/40 (all 4 emit_mode=1 cases, including the K<n_seq_tokens checkpoint case), and test-recurrent-state-rollback passes end-to-end with LLAMA_GDN_REPLAY=1 both on and off. No supports_op guard needed -- it already accepted emit_mode=1 and now actually computes it correctly.
|
Thanks for putting this together, Giveen. I tested the previously gated Metal ingredient-replay path on an Apple M5 Max and pushed Local validation:
I could not run the model-level recurrent rollback test because it requires a compatible local model, but the kernel and focused replay coverage are clean. Fresh CI should start on the updated head. |
|
I also updated this onto the current base after today’s merges and pushed Post-merge validation on the M5 Max is clean:
Fresh CI is running from the resolved head. |
26d920e
into
TheTom:feature/turboquant-kv-cache
i was just in the process of fixing that conflict, lol |
GDN ingredient-replay: cut MTP rollback VRAM for Qwen3.5 / Qwen3.5-MoE (opt-in)
Branch:
DRC→feature/turboquant-kv-cacheStatus: opt-in via
--gdn-replay(env varLLAMA_GDN_REPLAY=1also works), default off — zero behavior change unless explicitly enabled.What this is
MTP speculative decoding needs to roll back a model's recurrent state when a drafted token gets rejected. For Gated DeltaNet (GDN) layers — the recurrent-attention layers in Qwen3.5 / Qwen3.5-MoE's hybrid architecture — the existing design makes that rollback free by eagerly computing and storing a full state snapshot for every one of the last
n_rs_seq(= max speculative depth) draft positions, every single decode, whether or not a rollback ever happens. Snapshot storage and per-decode write cost both scale linearly with speculative depth:O(depth × S_v²)whereS_vis the recurrent head width (128 for this model).This change replaces that with ingredient-replay: instead of storing a full
S_v × S_vstate matrix per retained draft position, it records the small per-token(k, v, g, beta)values that produced that step (O(S_v)instead ofO(S_v²)— a ~32x reduction per retained step at this model's dimensions) and reconstructs the correct state on demand by replaying those values through the same recurrence, only when a rollback actually occurs. The common case (no rollback) pays only the cheapO(S_v)recording cost; the previous design paid the fullO(S_v²)cost on every decode regardless.This required extending
ggml_gated_delta_netwith a new opt-in output mode (emit_mode), plus a real architectural fix found partway through: a single "last known state" isn't enough to serve as a replay base, because at the point a decode is computed it isn't yet known which of its trailing tokens will be accepted — the design ended up needing two persisted checkpoints per sequence (the optimistic final state, and a second checkpoint from before the uncertain window), not one.What tests we did
Every piece was verified at multiple levels before moving to the next, and every fix in this series was found by testing, not by inspection:
tests/test-gdn-ingredient-replay.cpp) proves the core numerics: recording ingredients and replaying them through the op reconstructs the exact same state (max_abs = 0.0, not just "close") as the original full-snapshot path, across several sub-cases (the trailing final-state block, the "checkpoint before the retained window" block, and full multi-step replay).test-backend-opsGATED_DELTA_NET suite (which already covers K>1 snapshot cases) was extended withemit_mode=1cases across shapes, KDA and non-KDA gating, and multi-token batches. Both CPU (40/40) and CUDA (40/40) pass, with CUDA's output diffed against the CPU reference automatically by the harness.tests/test-recurrent-state-rollback.cppdrives an actual accept/reject cycle through realllama_decodecalls and compares logits between a rolled-back-and-replayed context and an independently checkpoint-restored one, to 1e-5 tolerance. This ran on CPU, on CPU with the flag on, and with-ngl 999(genuine GPU execution of the real kernel) — and it's the test that caught the two-checkpoint design flaw before it ever reached the real model.llama-serveragainst the actual Qwen3.8-27B checkpoint (general.architecture=qwen35), sweeping--spec-chain 2/4/6/8, flag on vs. off, measuring both VRAM and per-round latency — the numbers below are from this, not a synthetic model.emit_mode=0paths, the pre-existing non-speculative rollback test) re-run after every change; the full project (all 290 build targets, includingllama-server/llama-cli) builds clean.Bugs this process actually caught, in order: a shape-check gap that made a naive negative-stride replay design unsafe (never shipped); a
ggml_new_objectmetadata-arena overflow that crashed model load on the real 27B model (the synthetic test model had too few layers to expose it); an always-firing "redundant recompute" that was supposed to be a rare-case fallback; a CUDA kernel writing 3x more memory traffic than intended due to redundant per-warp writes; and a replay design that needed a full kernel-launch chain until the ingredient storage order was changed to avoid it.The benefit: lower VRAM, and it doesn't grow with depth
The old design's VRAM use grows linearly with speculative depth (900 MiB of growth from depth 2→8 alone); this change's VRAM use is essentially flat (57 MiB of growth over the same range) — because ingredient storage doesn't scale with the state matrix size the way full snapshots do.
Why this matters beyond the raw number: this saving is fixed and independent of context length — recurrent state size doesn't grow with how much context is loaded, so freeing this VRAM doesn't compete with context at all; it's pure headroom. Concretely, that means the ~939 MiB saved at a typical speculative depth can be redirected entirely toward more usable context — a bigger KV-cache/compression budget, a longer context window, or more parallel sequences in the same VRAM envelope — rather than being spent on bookkeeping that scales the wrong way as you push depth up for better acceptance rates. For a deployment that's already VRAM-constrained (the whole reason this fork's KV-cache compression work exists), that's a direct, reusable budget increase, not just a number that looks better in isolation.
The honest tradeoff: this is currently a VRAM-only win, not also a speed win. Per-round latency regresses slightly (+0.2 to +0.8ms/round depending on depth, down from an initial +0.6 to +2.0ms before three rounds of profiling-driven fixes) — root-caused to the replay mechanism's own overhead (a handful of
ggml_cont()materializations and a graph rebuild whenever a rollback occurs, both real costs the original snapshot design avoided by paying upfront every decode instead). If VRAM is the binding constraint, this is a clear net positive. If raw tokens/sec is what matters most, the current regression is real, if small.Metal port (written, gated off — not yet verified on real hardware)
This branch also adds a hand-ported Metal kernel for
emit_mode=1(ggml/src/ggml-metal/ggml-metal.metal,ggml-metal-ops.cpp,ggml-metal-impl.h,ggml-metal-device.m), mirroring the CPU/CUDA ingredient/final-state/checkpoint layout exactly, including the same "column-invariant values written once, not redundantly" fix the CUDA kernel needed. It has not been built or run — this development machine is Linux with no Metal toolchain (xcrun/metaldon't exist here), and Metal only compiles on macOS/iOS. It's mechanically correct as far as line-by-line translation from the verified CPU reference can guarantee, but it carries none of the verification the CPU/CUDA paths got (notest-backend-opsdiff, no real build).Because of that,
ggml-metal-device.m'ssupports_opstill explicitly declinesemit_mode=1(ggml_get_op_params_i32(op, 1) == 0guard), so on a Metal-offloaded model this path is unreachable today — GDN layers would just run the ordinary CPU fallback ifgdn_replaywere ever turned on. This is intentional: the code exists in the tree for someone with real Apple hardware to pick up, but is not exposed until it's actually verified.To test it on real Apple hardware, two separate things need to happen — one flag, one code edit:
gdn_replayitself — either the CLI flag--gdn-replay, or the environment variableLLAMA_GDN_REPLAY=1(either works; the flag was added later in this branch specifically so the env var wouldn't be the only way). Setting this alone does not exercise the new Metal kernel — it only turns on the feature, and every backend'ssupports_opdecides whether that backend actually runs it.ggml-metal-device.m, theGGML_OP_GATED_DELTA_NETcase currently ends with&& ggml_get_op_params_i32(op, 1) == 0. That clause has to be deleted (or changed to acceptemit_mode==1) before Metal will ever be asked to run the new kernel instead of quietly falling back to CPU.With both of those done, the fastest real check is the existing
tests/test-backend-opsGATED_DELTA_NET suite (it already hasemit_mode=1cases; they'll simply start running on the Metal backend instead of being skipped) plustests/test-recurrent-state-rollbackwith the flag on, on a Metal-offloaded context — the same two tests that caught every real CUDA-port bug in this branch.Vulkan port (written AND verified on real hardware)
Unlike Metal, Vulkan runs on this Linux dev machine too — this box has a real Vulkan-capable NVIDIA driver, so this port got the same treatment as CPU/CUDA, not a blind one. Added
emit_modeas a new push-constant field (ggml-vulkan.cpp,gated_delta_net.comp), ported the same ingredient/final-state/checkpoint logic (same "column-invariant values written once" fix the CUDA and Metal kernels needed), then built a separate-DGGML_VULKAN=ONtree and ran it for real:tests/test-backend-opsGATED_DELTA_NET: 40/40 pass on the real RTX 5090 via Vulkan, including all 4emit_mode=1cases (K>1, KDA, and the case withn_seq_tokens=8 > K=3that exercises the checkpoint block) — the harness diffs Vulkan's output against the CPU reference automatically.tests/test-recurrent-state-rollback(the real accept/reject/replay integration test, not just the op in isolation): passes on Vulkan both withLLAMA_GDN_REPLAY=1and with it off.Because this is actually verified, Vulkan's
supports_opneeded no defensive guard — it already acceptedemit_mode=1correctly (it just didn't compute the right thing before this port existed; now it does, and the tests confirm it).Scope — what this does and doesn't apply to
This is not a general MTP/speculative-decoding VRAM improvement. It's specific to one thing: the Gated DeltaNet recurrent-state rollback mechanism, in the two architectures where this fork enables rollback for it at all — Qwen3.5 and Qwen3.5-MoE (keyed on GGUF architecture tag, so any model using that arch benefits, not just this one checkpoint). It does not apply to Qwen3-Next, Kimi-Linear, Mamba/Mamba2, or other recurrent architectures (rollback isn't enabled for them in this fork regardless of this change), nor to DeepSeek-V4 (which has its own, separate, untouched rollback implementation), nor to ordinary transformer models doing MTP (their rollback is a cheap KV-cache truncation with no equivalent snapshot cost to begin with).
Commits in this branch
b7dce42a2— ggml: addgated_delta_netemit_mode=1(ingredient replay) for CPU+CUDAbc14836e0— llama: wire ingredient-replay into MTP recurrent-state rollback (opt-in)ae5124f38— fix two real perf regressions (redundant prefix-recompute, redundant CUDA writes)ac375f654— add--gdn-replayCLI flag, replacing the raw env-var-only gate99ba79a49— replay rollback with one batched call instead of a kernel-launch chain4f99bc4b5— ggml/metal: portemit_mode=1, gated off pending hardware verificationa7f8d290f— ggml/vulkan: portemit_mode=1, verified on the real RTX 5090Not done / open follow-ups
emit_mode=1to SYCL, OpenCL, WebGPU — not attempted, not needed for this deployment's single-GPU CUDA target. (Hexagon needs no work: itssupports_opshape check already happens to rejectemit_mode=1's output shape and falls back to CPU safely, by coincidence rather than an explicit guard.)AI Usage: yes