Trimmable ArraysCache alternative

[zcbenz]

Refs #990

The ArraysCache only stores the last hidden states so it is not trimmable and does not work with speculative decoding like MTP.

This PR experiments with a new RecurrentCache implementation that, stores the hidden states in a temporal manner so we can roll back in speculative decoding. The downside is that the hidden states per token is much larger than traditional KV cache and the sequence length would be very limited, so it would not help much in the case of prefix prompt cache, but for speculative decoding we usually only draft 4~6 tokens and additional RAM usage would be quite small.

To try this branch:

mlx_lm.generate --model mlx-community/Qwen3.8-27B-4bit --max-tokens 200 --prompt "Write a story about George Washington" --chat-template-config "{\"enable_thinking\": false}" --draft-model mlx-community/Qwen3.5-0.8B-4bit --num-draft-tokens 4

[pierre427]

Thanks for working on this. Making the recurrent Gated DeltaNet state safely rewindable is the right problem to solve, and it unlocks useful target/draft combinations for Qwen3.5/Qwen3.8. I compared the current PR head (c3081c333099f7958f4d5a5833f6a77a3f1023e5) against an exact-rollback implementation we have been exercising locally, ran the new TestRecurrentCache tests (9/9 pass), and then added a few focused integration probes.

@zcbenz I'm dropping a series of suggestions based not only on the needed functionality that is missing here, but also from some hard won lessons here. If you'd rather that I upstream my own implementation here, I can open a new PR. But I also thought this would be a good channel here to explain what is needed;

I found several correctness issues that I think need addressing before this is used by speculative generation (MTP, PLD, Dflash etc). I also wanted to share an alternative design that has worked well for us and avoids changing the ordinary GDN execution path. In these comments is a bare minimum to make it work. Not included:

does not include ragged/continuous batching, dynamic batch membership, persistent prefix checkpoints, arbitrary prefix-cache truncation, or our Qwen4/self-MTP serving machinery. It is intentionally limited to transaction-scoped exact rollback for single-sequence speculative decoding with Qwen3.5/Qwen3-Next targets and drafts.

If you're interested in seeing the pieces missing, will include links to PRs covering the topic.

1. A fresh RecurrentCache is rejected before the advertised path reaches prefill

speculative_generate_step() creates fresh caches and immediately calls:

if not can_trim_prompt_cache(model_cache, num_draft_tokens):
    ...

At this point a fresh RecurrentCache has conv_offset == 0. With four draft tokens, is_trimmable(4) returns false because 0 - 4 < conv_kernel_size - 1. Consequently, the example command in the PR raises the existing “requires a trimmable prompt cache” error before _prefill() runs.

I reproduced the precondition directly:

empty=True
is_trimmable(4)=False
can_trim_prompt_cache(..., 4)=False

I would distinguish two questions in the cache API:

1. Can this cache trim n tokens from its current populated state?
2. Can this cache support a future speculative transaction of width n?

For example:

class RecurrentCache(_BaseCache):
    def supports_speculative_trim(self, n: int) -> bool:
        # One boundary state plus n speculative state transitions are needed.
        return n >= 0 and n + 1 <= self.state_capacity


def _can_speculate(cache, model, n):
    return cache.is_trimmable(n) or (
        getattr(model, "supports_speculative_rollback", False)
        and getattr(cache, "supports_speculative_trim", lambda _: False)(n)
    )

That avoids making is_trimmable(n) claim that an empty cache can actually be trimmed. Another option is to perform the populated-state check after the first target forward, but that still needs careful handling for a one-token prompt.

The draft cache should be capability-checked as well as the target cache. A Qwen3.5 draft is recurrent too.

2. The maximum advertised rollback depth is missing its boundary state

The cache keeps max_size post-token SSM states. To rewind n transitions, however, it needs the state immediately before those n transitions as well. In other words, a six-token rollback needs seven state boundaries.

With max_size=6, after retaining six states and calling trim(6):

before: ssm_offset=6, current state is slot 5
after:  ssm_offset=0

__getitem__(-1) then reads ssm_states[ssm_offset - 1], which is ssm_states[-1]: the newest stale state. More importantly, the required pre-window state was never retained.

This can be fixed by giving the ring explicit boundary-state semantics. For example, if max_rollback means the maximum number of token transitions:

self.max_rollback = max_rollback
self.state_capacity = max_rollback + 1

The cache must retain the pre-forward state plus up to max_rollback post-token states. It should also reject reading the current SSM state when ssm_offset == 0 rather than allowing negative-index wraparound:

def current_ssm_state(self):
    if self.ssm_states is None:
        return None
    if self.ssm_offset <= 0:
        raise RuntimeError("No current recurrent-state boundary is available")
    return self.ssm_states[self.ssm_offset - 1]

For the current fixed-size implementation, the smaller immediate correction would be to advertise at most max_size - 1 rollback tokens. Supporting the stated 4–6-token range requires sizing from num_draft_tokens or storing at least seven state boundaries.

Please add explicit tests for trimming zero, one, max_size - 1, max_size, and max_size + 1; the test should validate the restored state against a prefix-only reference, not only the cursor value.

3. Right-padded batches advance the recurrence through padding

BatchGenerator passes lengths and right_padding into prepare(), but RecurrentCache.prepare() retains only right_padding and make_mask() always returns None.

In qwen3_5.py, qkv is zeroed only when the mask is non-None, and the same mask is passed into gated_delta_update(). Therefore a shorter row in a right-padded prompt slab advances its convolution and SSM recurrence through padding tokens as though they were real tokens.

A minimal mask implementation would look approximately like:

def prepare(self, lengths=None, right_padding=None, **kwargs):
    self._lengths = None if lengths is None else mx.array(lengths)
    self._right_padding = (
        None if right_padding is None else mx.array(right_padding)
    )

def make_mask(self, n: int):
    if self._lengths is None:
        return None
    return mx.arange(n)[None, :] < self._lengths[:, None]

There is a second part: finalize() currently rotates only conv_states. The temporal SSM history must also be aligned independently for every batch row, or later trim() calls walk repeated padding states rather than logical token boundaries. Conceptually:

def finalize(self):
    if self._right_padding is not None:
        self.conv_states = [
            None if state is None
            else dynamic_roll(state, self._right_padding, axis=1)
            for state in self.conv_states
        ]

        if self.ssm_states is not None:
            # dynamic_roll expects batch first.
            states = mx.swapaxes(self.ssm_states, 0, 1)
            states = dynamic_roll(states, self._right_padding, axis=1)
            self.ssm_states = mx.swapaxes(states, 0, 1)

    self._lengths = None
    self._right_padding = None

This deserves an end-to-end test comparing each row of a mixed-length batch with the same prompt run independently. The current test manually zeroes qkv before invoking the cache, which bypasses the model/cache mask integration and therefore does not expose this bug.

For speculative batches, a single shared ssm_offset also becomes problematic once rows accept different numbers of draft tokens. Either the cache needs per-row logical offsets/rollback depths, or ragged speculative trimming must be rejected explicitly.

4. Merging populated and empty caches drops the empty lanes

RecurrentCache.merge() constructs valid_states using only caches whose state is non-None, then derives B from those entries. If one prompt has a populated cache and another is fresh, the fresh row disappears from both the convolution and SSM tensors.

I reproduced:

input: one populated cache + one empty cache
expected batch size: 2
merged conv batch:  1
merged SSM batch:   1

The merge should advance the output row cursor for every input cache, inserting zero state for empty lanes:

def cache_batch_size(cache):
    for state in cache.conv_states:
        if state is not None:
            return state.shape[0]
    if cache.ssm_states is not None:
        return cache.ssm_states.shape[1]
    # An empty cache supplied for one request still represents one row.
    return 1

batch_sizes = [cache_batch_size(c) for c in caches]
total_batch = sum(batch_sizes)

j = 0
for source, source_batch in zip(caches, batch_sizes):
    if source.ssm_states is not None:
        # Copy/right-align the valid source history into [j:j+source_batch].
        ...
    # Always advance, including for an empty source.
    j += source_batch

The existing “merge with empty” test should additionally assert the merged batch dimension and validate both output rows. Checking only that the populated state survived does not detect a dropped empty row.

The same empty-state audit should cover filter(), extract(), and extend(), all of which currently have paths that assume ssm_states or every convolution state is populated.

5. The speculative generator needs exception-safe round initialization

This is existing generator behavior rather than a new cache-only issue, but recurrent draft support makes it important.

At the beginning of a later speculative round, num_draft is updated but n still contains the preceding round's accepted-token count. If _draft_generate() or the target _step() raises, finally calls _rewind_cache(num_draft, n) with stale values. That can trim target and draft caches even though the current target verification never completed, and a rollback exception can mask the original error.

Set n to the no-rewind value before any fallible work in every round:

while True:
    num_draft = min(max_tokens - ntoks, num_draft_tokens)

    # Nothing from this round is rewindable until target verification advances.
    n = num_draft

    draft_tokens = _draft_generate(draft_y, num_draft)
    y = mx.concatenate([y, draft_tokens])
    tokens, logprobs = _step(model, model_cache, y, num_draft + 1)

    # Verification completed; n now tracks this round's accepted prefix.
    n = 0
    ...

Recording/rollback must be enabled on both model_cache and draft_cache, and cleanup should always happen in a finally block after the final rewind.

Tests worth adding here are: exception during the second draft round, exception during target verification, closing the generator while suspended at a yield, and reusing the same prompt cache after early close.

6. Retaining temporal states currently disables the optimized GDN path

The largest performance concern is that qwen3_5.py passes:

return_num_states=cache.max_size

for every cached GDN call. In gated_delta_update(), any non-None return_num_states selects gated_delta_ops() instead of the optimized Metal kernel. This affects normal cached prefill and decode, not only speculative rejection.

That means the tradeoff is not just a small amount of extra memory for 4–6 states. It structurally changes the hot path and may create a large prefill/decode regression. I would add a test or mechanism receipt confirming which implementation executes, plus an A/B benchmark for:

• ordinary cached prefill;
• ordinary single-token decode;
• target verification at widths 2, 4, and 6;
• recurrent draft generation at width 1;
• total speculative throughput and peak memory.

Alternative that worked for us: transaction-scoped exact replay

We landed on retaining the exact inputs needed to replay only the open speculative window, rather than retaining temporal SSM states during every forward.

The lifecycle is:

1. Use the ordinary ArraysCache and ordinary optimized GDN path for prefill.
2. After prefill, call start_speculation() on target and draft caches.
3. Before a speculative GDN forward, capture its pre-forward convolution/SSM state and the exact normalized q, k, v, a, and b tensors consumed by the recurrence.
4. Record a closure that can replay the first m tokens from the pre-forward state.
5. If only m proposals survive, replay that prefix and replace the cache with the exact state after m tokens.
6. Clear all retained inputs at the end of the speculative transaction.

The model-side shape is roughly:

state_before = cache[1] if cache is not None else None
conv_before = cache[0] if cache is not None else None
q, k = normalize_qk(q, k)

if cache is not None and cache.speculating:
    n_keep = self.conv_kernel_size - 1

    def rollback(m, q=q, k=k, v=v, a=a, b=b,
                 state_before=state_before, conv_input=conv_input):
        _, state_m = gated_delta_update(
            q[:, :m],
            k[:, :m],
            v[:, :m],
            a[:, :m],
            b[:, :m],
            self.A_log,
            self.dt_bias,
            state_before,
            mask=None,
            use_kernel=not self.training,
        )
        conv_m = mx.contiguous(conv_input[:, m : m + n_keep, :])
        return [conv_m, state_m]

    cache.record_rollback(
        num_tokens=S,
        fn=rollback,
        snapshot=[conv_before, state_before],
    )

# The live path remains the existing optimized implementation.
out, state = gated_delta_update(
    q, k, v, a, b,
    self.A_log, self.dt_bias,
    state_before, mask,
    use_kernel=not self.training,
)

A few details were important in making this reliable:

• Capture the post-normalization inputs actually consumed by GDN. Re-running projections would read weights again and can change rounding or transformed inputs.
• Use a stack of rollback records. The target usually records one wider verification forward, while a recurrent draft records several T=1 forwards.
• Make rollback strict. Never clamp to the available history; partially rewinding one cache type silently desynchronizes it from the KV layers.
• Bound retained history by logical tokens, not record count. We use a small fixed transaction window; nothing prompt-sized is retained.
• Start recording after prefill, otherwise the closure holds prompt-sized arrays.
• The convolution state after accepting m tokens is a slice of the already-computed convolution input; it does not require another convolution.
• A full rollback should restore the captured pre-forward state directly. A partial rollback replays only the accepted prefix.
• After a partial rollback, preserve/re-record whatever prefix remains valid if another trim may occur in the same transaction.
• Batch membership changes invalidate closures that capture whole-batch tensors. Clear them or transform them explicitly; never replay them against a differently shaped/reordered batch.
• For a right-padded speculative slab, record the valid depth per row. Crediting every row with the slab width is silently wrong.

This approach leaves ordinary prefill and decode unchanged and pays the replay cost only when a speculative suffix is rejected. In our local real-weight checks, dense-draft and recurrent-draft combinations remained token-identical to vanilla greedy generation and showed approximately 1.37x and 1.33x end-to-end speedups respectively at the tested operating points. Those figures are workload-specific, but the important result for this discussion is that exact rollback did not require replacing the normal GDN cache or disabling its optimized forward kernel.

Suggested scope

For an initial upstreamable change, I would keep the scope narrow:

• transaction-scoped rollback support in ArraysCache;
• exact replay closures for Qwen3.5 and Qwen3-Next GDN;
• capability checks for both target and draft caches;
• start/stop lifecycle after prefill;
• strict trimming and exception-safe cleanup;
• single-sequence correctness tests.

I would leave ragged batching, dynamic membership, long-lived prefix checkpoints, and arbitrary prefix-cache truncation to follow-ups. Speculative rollback needs a tiny temporary window; general prefix-cache rewind is a different storage/lifetime problem and does not need to burden the normal decode path.

Tests I would use as the acceptance gate

• Same accepted prefix with different rejected tails produces bit-identical restored cache state and next logits.
• Greedy speculative output is token-identical to vanilla greedy output.
• Rejection at zero, first, middle, and final draft positions.
• All draft tokens accepted.
• Dense draft with recurrent target.
• Recurrent draft with recurrent target.
• Multiple T=1 draft records in one round.
• Multiple consecutive speculative rounds.
• Maximum supported rollback depth and one beyond it.
• Generator early close and mid-round exceptions.
• Prompt-cache reuse after early close.
• Mixed empty/populated cache merge.
• Batched rows with different prompt lengths compared independently.
• Filter/extract/merge during or immediately after rollback recording.
• A receipt or benchmark proving ordinary cached prefill/decode still selects the optimized GDN kernel.

I hope this is useful. I think the PR is attacking an important gap, and the temporal-cache approach may still be useful for bounded prefix checkpoints. For speculative decoding specifically, though, transaction-scoped exact replay has a much smaller steady-state cost and isolates the additional machinery to the operation that actually needs it.

--

Related PRs with a more fleshed out implementation:

raullenchai/Rapid-MLX#2842

Semi related and different implementation:

raullenchai/Rapid-MLX#3141

Shows the serving integration, fail-closed eligibility, request-private ownership and APC publish-back. Its public https://github.com/pierre427/rapid-mlx-decode-lane-deps contains the compiled-replay modules (part of megakernel tests). It does not contain the complete RingKVCache class because that wiring lives in heavily diverged core files. That cache was used to avoid regenerating dispatch graphs. Megakernel was interesting to work on but optimizing that path takes away idle times that can be used for verify on speculative lanes. Eager mode and self MTP when available are the better optimizations, but including my thoughts here to avoid duplication.

[AirRunner]

@pierre427 Since you are comparing designs, here is what I have on #990: ArraysCache carries a single rollback_state snapshot (conv + SSM state) rather than a bounded history. The GatedDeltaNet forward is split into a "confirmed" chunk and a "draft" chunk, and the snapshot is taken right after the confirmed chunk, before running the draft chunk. On rejection I just restore that one snapshot, on acceptance I drop it. It's depth-1 only right now but I have a feat/mtp-depth branch generalizing it to a list of snapshots per confirmed position.

Do you think the "split the forward, snapshot the boundary, replay nothing" shape could avoid problems your closure-replay approach has to handle? Or maybe does it just move the same complexity somewhere else?

[pierre427]

https://github.com/pierre427/mlx-lm-unified/blob/a7ecfae3194538df8f02ec399a88ede302b75a46/results/qwen4-gdn-replay-20260917/README.md

Here's where I landed, reach out to me directly if you'd like to chat pierre @ userid.org

[pierre427]

So I finally got around to publishing some updated work about caches.

https://github.com/pierre427/mlx2/blob/main/docs/mlx2_deep_dive_feature_architecture_review.pdf - section 3+

Some related code: https://github.com/pierre427/mlx2/blob/main/src/mlx2/runtime/apc_v2.py

Why? Because legacy cache can't really scale and support advanced features & functionality. Splitting it into layers + segments allows for cow-branching, rollbacks, APC suspend to disk etc.

No performance regressions, I had some early on in the dev process but it's at or slightly above parity with legacy.