Description
BatchKVCache and BatchRotatingKVCache (mlx_lm/models/cache.py) store offset as an mx.array and advance/adjust it in place (self.offset += S / -= n / -= roll). Because mx.array += mutates in place, any model that captures offset = cache.offset and reuses that reference after update_and_fetch sees the captured value silently changed to the post-update offset.
Impact — a real bug on main
gemma3n.py captures offset = cache.offset, calls cache.update_and_fetch(...), then RoPEs the queries with the captured offset. Under the batched caches (what mlx_lm.server continuous batching uses) offset is an mx.array, so the in-place advance corrupts the captured reference and the queries are rotated at the post-update offset — gemma3n batched / continuous-batching decode diverges from sequential from the first token.
Two models already carry a defensive consumer-side snapshot that this footgun makes necessary:
gpt2.py: offset = mx.array(offset)gemma4_text.py: offset = mx.array(cache.offset)
Models that RoPE before update_and_fetch (the common pattern) are unaffected — a derived graph freezes the offset's value at construction, so a later in-place mutation cannot change it.
Fix
Rebind (self.offset = self.offset +/- X) at every site that touches offset in both batched caches (update_and_fetch, prepare, finalize, trim, merge, and the rotating cache's roll), so offset is never mutated in place and a captured reference keeps value semantics like the int-offset caches. left_padding is left in place (no consumer aliases it and reuses it after a mutation). Regression test asserts a captured offset survives update_and_fetch and trim for both caches.
Description
BatchKVCacheandBatchRotatingKVCache(mlx_lm/models/cache.py) storeoffsetas anmx.arrayand advance/adjust it in place (self.offset += S/-= n/-= roll). Becausemx.array+=mutates in place, any model that capturesoffset = cache.offsetand reuses that reference afterupdate_and_fetchsees the captured value silently changed to the post-update offset.Impact — a real bug on main
gemma3n.pycapturesoffset = cache.offset, callscache.update_and_fetch(...), then RoPEs the queries with the capturedoffset. Under the batched caches (whatmlx_lm.servercontinuous batching uses)offsetis anmx.array, so the in-place advance corrupts the captured reference and the queries are rotated at the post-update offset — gemma3n batched / continuous-batching decode diverges from sequential from the first token.Two models already carry a defensive consumer-side snapshot that this footgun makes necessary:
gpt2.py:offset = mx.array(offset)gemma4_text.py:offset = mx.array(cache.offset)Models that RoPE before
update_and_fetch(the common pattern) are unaffected — a derived graph freezes the offset's value at construction, so a later in-place mutation cannot change it.Fix
Rebind (
self.offset = self.offset +/- X) at every site that touches offset in both batched caches (update_and_fetch, prepare, finalize, trim, merge, and the rotating cache's roll), so offset is never mutated in place and a captured reference keeps value semantics like the int-offset caches.left_paddingis left in place (no consumer aliases it and reuses it after a mutation). Regression test asserts a captured offset survives update_and_fetch and trim for both caches.