Exact speculative rollback for hybrid caches (model-agnostic cache API)#1486
Open
lBroth wants to merge 4 commits into
Open
Exact speculative rollback for hybrid caches (model-agnostic cache API)#1486lBroth wants to merge 4 commits into
lBroth wants to merge 4 commits into
Conversation
Speculative decoding currently raises for models with recurrent
ArraysCache layers ("requires a trimmable prompt cache", ml-explore#980): the
recurrent state summarizes the whole past and cannot be trimmed.
This adds an exact, cheap rollback instead of forbidding the trim:
- ArraysCache: start/stop_speculation + record_rollback(T, fn, snapshot);
while speculating the cache is trimmable, and trim(n) restores the
entries via the recorded fn (or the snapshot for a full rollback).
- GDN layers (qwen3_5, qwen3_next): while speculating, record a rollback
closure over the exact per-token tensors the kernel consumed. Replaying
the recurrence from the pre-forward state over the first m stashed
inputs reproduces the post-m state bit-for-bit (causality: rows of a
chunk do not depend on later rows), and the conv state is a slice of
conv_input. No deepcopy, no re-forward, no extra weight streaming —
the rollback costs one tiny kernel launch per GDN layer.
- speculative_generate_step: allow models declaring
supports_speculative_rollback, and start/stop recording around the
decode loop (recording is off during prefill so no prompt-sized
tensors are retained).
Tests: tail-independence bit-exactness (same m-prefix, different tails
=> identical caches and logits after rollback), greedy speculative ==
greedy vanilla end-to-end with a mismatched draft (rollback every
round), and the unsupported-model error path.
Addresses findings from an adversarial audit of the first version: - ArraysCache keeps a small stack of rollback records (bounded by a 64-token window) instead of a single record, so multi-forward windows can be rewound — enabling HYBRID DRAFT models: a Qwen3.5 draft records one entry per T=1 draft step and its cache is now rewound exactly (the setup reported in ml-explore#1446 uses a hybrid draft). Partial rollbacks re-record the still-valid prefix. - trim() raises when asked to trim beyond the recorded window instead of silently clamping (a clamp would desync ArraysCache layers from KVCache layers, which trim the full amount). - speculative_generate_step: per-cache capability check (trimmable OR rollback-capable with model support) instead of a model-wide bypass; start/stop speculation on the draft cache too; n is initialized to num_draft at the top of each round so a mid-round exception or generator close makes the finally-block rewind a no-op rather than trimming with stale values — previously the rollback RuntimeError could mask the original exception. Tests: hybrid draft == vanilla greedy (rollback across stacked T=1 records), mid-round draft exception surfaces unmasked, trim beyond the window raises. 8/8 in this file; full suite failure set identical to main (pre-existing failures only). Real-weights e2e: dense draft 1.37x identical; hybrid draft pair 1.33x identical.
Extends the exact-rollback protocol to the OTHER non-trimmable cache: sliding-window RotatingKVCache. Once it wraps (offset >= max_size) the tokens a verify step pushes out of the window are gone, so stock trim can't rewind. Same protocol as ArraysCache: start/stop_speculation + record_rollback; update_and_fetch stashes the pre-forward window + new K/V while speculating, trim restores + re-appends the accepted prefix. No model-file hook (KV is not recurrent). is_trimmable() is True while speculating; gpt_oss declares supports_speculative_rollback=True, so speculative_generate_step accepts it unchanged -> --draft-model spec now works on gpt-oss-family sliding-window models, not just GDN hybrids. 4 tests incl. bit-exact tail-independence (two verify chunks sharing an m-prefix with different tails leave identical windows after rollback-to-m) + stacked draft records. 37/37 across prompt_cache + all rollback + MTP. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pierre Lamy <pierre@userid.org>
|
Thanks for cutting this out into its own PR, @lBroth — clean split. Confirming from my side: the |
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.
Up front: I'm not a professional ML engineer and this was done with heavy AI
assistance — I've validated everything I could on-device (the matrix below is
all real measured runs), and I'd welcome expert scrutiny on anything I've missed.
Split of #1456 as @pierre427 proposed: this is the rollback cache-API alone
(the shared dependency), landed standalone. The qwen3.5 MTP-head loading follows
as a separate PR on top. Co-developed with @pierre427, who contributed the
RotatingKVCachesupport and reviews.What & why
Speculative decoding needs to rewind the KV/state cache when a draft is
rejected. mlx-lm's existing
--draft-modelpath does this bytrim()ing thecache — but that silently corrupts any cache that isn't trimmable:
state (
ArraysCache) has no validtrim()(Prefix cache reuse is broken for all hybrid-architecture models (sliding window, SSM/Mamba) #980, mlx_lm.server speculative decoding fails on Qwen 35B target: ArraysCache is not trimmable #1446) — a multi-tokendraft advances the recurrence and there is no way to undo it by truncation;
RotatingKVCache:trim()only adjustsoffset/_idx, it does notshrink the key buffer, so a multi-token block feed that crosses the sliding
window leaves stale keys.
This PR adds an exact rollback protocol so speculative decoding is
lossless on these caches.
The API (model-agnostic — proven on two cache families)
A small protocol on the cache objects:
start/stop_speculation(),record_rollback(num_tokens, fn, snapshot),is_trimmable(), stricttrim().Two independent implementations demonstrate the abstraction isn't cache-specific:
the pre-block
conv_state+ recurrentS, capture (via hooks) the exacttensors the GatedDeltaNet kernel consumed, and reconstruct
S_mby re-runninggated_delta_updateon the committed prefix. Exact by construction (noapproximation), because the recurrence can't be truncated.
at snapshot time, since
trim()can't shrink them.Same protocol, opposite mechanisms — that's the evidence the API is right.
Correctness (the first thing you'll ask)
identity between speculative and non-speculative greedy decode.
argmax at that position (the accept rule always commits the target's own
token) — not bit-identity with sequential
generate_step, which differs byinherent FP noise regardless of speculation.
tests/test_speculative_rollback.py(8) +tests/test_rotating_rollback.py(4) — 12 total, all green. Cover cache-reuse across turns, early-close, and
offset-exactness after each step.
No overhead when unused, no regressions
start_speculation()is onlycalled on the speculative branch; normal generation is byte-identical.
M2 Pro 16GB / M3 16GB / CPU(ops) — unit tests 15/15, cache regression 22/22
(failure set ⊆ main's), real-weights e2e identical (80B Coder-Next + 0.8B
draft, token-for-token vs non-speculative). No regressions found anywhere.
Scope
Rollback cache-API + speculative wiring + GDN & RotatingKVCache support + tests.
No MTP head (separate follow-up). Diff is +639/-6 across 7 files.