gemma3n broken with batch KV caches: cache.state unpack crash + RoPE offset mutated in place across update_and_fetch

[Cyb3rb1ade]

Summary

Running a gemma3n text model through the batched generation path (BatchKVCache / BatchRotatingKVCache, i.e. BatchGenerator and any server built on it) fails in two independent ways:

1. Crash — ValueError: too many values to unpack (expected 2, got 4) in Gemma3nAttention.__call__ at keys, values = cache.state, because batch caches return a 4-tuple (k, v, offset, left_padding) while plain KVCache/RotatingKVCache return 2-tuples.
2. Degenerate output — after working around (1), generation produces repetition loops (e.g. HierHierHierHier...) because the RoPE offset used for the queries is silently mutated by the cache update (details below). The same model + prompt generates correctly with plain caches.

Environment

• mlx-lm 0.31.3 (latest release), mlx 0.32.0
• macOS, Apple Silicon
• Model: any gemma3n text conversion, e.g. mlx-community/Huihui-gemma-3n-E4B-it-abliterated-lm-8bit
• Code in main (mlx_lm/models/gemma3n.py) is identical as of 2026-08-09.

Reproduction

from mlx_lm import load
from mlx_lm.generate import generate
from mlx_lm.models.cache import (
    BatchKVCache, BatchRotatingKVCache, KVCache, RotatingKVCache, CacheList,
)

model, tok = load("mlx-community/Huihui-gemma-3n-E4B-it-abliterated-lm-8bit")
prompt = tok.apply_chat_template(
    [{"role": "user", "content": "What is the capital of France?"}],
    tokenize=False, add_generation_prompt=True)

def to_batch(c):
    if type(c) is KVCache:
        return BatchKVCache([0])
    if isinstance(c, RotatingKVCache):
        return BatchRotatingKVCache(c.max_size, [0])
    if isinstance(c, CacheList):
        return CacheList(*(to_batch(x) for x in c.caches))
    return c

# Plain caches: works -> "The capital of France is Paris."
print(generate(model, tok, prompt=prompt, max_tokens=48,
               verbose=False, prompt_cache=model.make_cache()))

# Batch caches (what _make_cache() produces for the batched engine):
batch = [to_batch(c) for c in model.make_cache()]
print(generate(model, tok, prompt=prompt, max_tokens=48,
               verbose=False, prompt_cache=batch))
# stock: ValueError: too many values to unpack (expected 2, got 4)
# with `cache.state[:2]` workaround: "HierHierHierHier..." (garbage)

Root cause of bug 2 (the interesting one)

In Gemma3nAttention.__call__ (non-shared branch), the offset for the queries RoPE is bound before update_and_fetch but applied after it:

if cache is not None:
    offset = cache.offset          # bound BEFORE the update
...
if cache is not None:
    keys, values = cache.update_and_fetch(keys, values)  # batch: self.offset += S

queries = queries.transpose(0, 2, 1, 3)
queries = self.rope(queries, offset=offset)              # applied AFTER the update

With plain caches offset is a Python int — no problem. The batch caches, however, keep self.offset as an mx.array, and self.offset += S mutates the buffer in place, so the previously bound offset silently becomes the post-update value. Every query position is then shifted by the prompt length and attention degenerates.

Minimal demonstration of the aliasing trap:

import mlx.core as mx
a = mx.array([0])
b = a
a += 24
mx.eval(a)
print(b)   # array([24], dtype=int32) — b changed under our feet

Bisection evidence (same prompt, tiny sizes, window never reached — so this is unrelated to the rotation logic):

• KVCache → BatchKVCache (full-attention layers): output still correct
• RotatingKVCache → BatchRotatingKVCache (sliding layers): garbage
• Inside one sliding layer: SDPA inputs k, v, mask are bit-identical, but q differs (max|diff| ≈ 2.4) — the queries' RoPE is the only divergent step, consistent with the offset being shifted after the update.

Suggested fix

Apply the queries RoPE before update_and_fetch (same offset semantics as the keys), and read the shared-layer state tolerantly:

offset = 0
if self.is_kv_shared_layer and cache is not None:
    # For shared layers, retrieve KV from the designated cache layer
    keys, values = cache.state[:2]          # batch caches return 4-tuples
    offset = cache.offset
    queries = queries.transpose(0, 2, 1, 3)
    queries = self.rope(queries, offset=offset)
else:
    if cache is not None:
        offset = cache.offset
    keys = self.k_proj(x).reshape(B, L, -1, self.head_dim)
    keys = self.k_norm(keys)
    keys = keys.transpose(0, 2, 1, 3)
    keys = self.rope(keys, offset=offset)

    values = self.v_proj(x).reshape(B, L, -1, self.head_dim)
    values = self.v_norm(values)
    values = values.transpose(0, 2, 1, 3)

    # RoPE queries BEFORE update_and_fetch: batch cache offsets are
    # mlx arrays mutated in place by the update, so an offset variable
    # bound earlier would silently shift afterwards.
    queries = queries.transpose(0, 2, 1, 3)
    queries = self.rope(queries, offset=offset)

    if cache is not None:
        keys, values = cache.update_and_fetch(keys, values)

With this patch, plain-cache and batch-cache generation produce identical, correct output.

Notes

• Other models that read cache.offset or cache.state across an update_and_fetch call may have the same latent bug when used with batch caches — worth a sweep (grep -n "cache.state" / cache.offset in mlx_lm/models/).
• Found while serving gemma3n through Rapid-MLX (which drives the mlx-lm batched engine); reproduced in pure mlx-lm as shown above.