MLX: _GDNStateCapture.rollback assumes "non-trimmable cache count == GDN layer count" without an assertion

[shaun0927]

Summary

_GDNStateCapture.rollback in dflash/model_mlx.py:338-355 indexes self._gdn_inputs[j] and self.conv_data[j] with j advancing only for not c.is_trimmable() cache entries:

def rollback(self, cache, accepted, trim):
    j = 0
    for c in cache:
        if c.is_trimmable():
            c.trim(trim)
        else:
            q, k, v, a, b, A_log, dt_bias, init_state, mask = self._gdn_inputs[j]
            ...
            conv_input, K = self.conv_data[j]
            c.cache[0] = conv_input[:, accepted + 1 : accepted + K]
            j += 1

This silently relies on the invariant "every non-trimmable cache entry corresponds to exactly one GatedDeltaNet layer". The capture lists are populated only by the monkey-patched GatedDeltaNet.__call__, so if mlx_lm ever introduces another non-trimmable cache type (a different SSM/recurrent layer, a custom DDT cache, etc.), j would still advance for it while _gdn_inputs[j] holds data from a different layer — silently producing rollback corruption rather than an exception.

Why it matters

Today this is true for the Qwen3.5 family on mlx_lm 0.31.x, but the invariant is undocumented and unguarded:

• It would silently break the day someone adds a new architecture with mixed cache types (a likely scenario given the existing model-request issues Qwen3.5 27b and 122b #40, HF backend support for Qwen3.5 model family #44, Dflash support for gemma4 #47, Training Dataset #48).
• mlx-lm could introduce another non-trimmable cache between releases without it being obvious that DFlash needs to update.

Suggested fix

Add a single assertion at the top of rollback:

def rollback(self, cache, accepted, trim):
    n_non_trimmable = sum(1 for c in cache if not c.is_trimmable())
    assert n_non_trimmable == len(self._gdn_inputs), (
        f"non-trimmable cache count ({n_non_trimmable}) != "
        f"captured GDN inputs ({len(self._gdn_inputs)}); "
        "DFlash MLX rollback assumes every non-trimmable cache is a GatedDeltaNet"
    )
    j = 0
    for c in cache:
        ...

Cheap to add; converts a silent corruption into a loud failure that future mlx-lm bumps catch immediately.

Happy to send a PR if the team agrees.

Environment

• repo at HEAD 6e0c951
• mlx-lm 0.31.x (per pyproject.toml's mlx-lm>=0.31.2)
• finding from a static review of PR Add MLX backend for Apple Silicon #59