fix(gemma3n): share KV independently of the cache (#1674)

[Anai-Guo]

Closes #1674.

Problem

In Gemma3nAttention.__call__, KV-shared layers only reused the source layer's keys/values when a cache was present:

offset = 0
if self.is_kv_shared_layer and cache is not None:
    keys, values = cache.state
    offset = cache.offset
else:
    ...   # shared layers ALSO land here when cache is None

Coupling KV sharing to the cache produces two distinct defects.

1. No-cache forwards don't share KV (the reported issue). Generation always runs with a cache, so it is unaffected, but any cache-free forward — mlx_lm.evaluate/perplexity and, notably, LoRA fine-tuning — falls into the else branch and computes the shared layers' own k_proj/v_proj. The model being trained is not the model used at inference. For reference, transformers threads a shared_kv_states dict through the forward and never consults a cache for this, and it does not even build k_proj/v_proj on shared layers.

2. With a cache, the shared layers rotate their queries at the wrong offset. cache here is the source layer's cache, and by the time a shared layer runs, that layer has already called update_and_fetch, so cache.offset is offset + L, not offset. The concrete layers rotate their queries at offset. So prefilling a prompt in one step does not agree with decoding it token by token.

Both are visible on the gemma3n config already in tests/test_models.py (random weights, so magnitudes are arbitrary — the point is that quantities which must agree don't):

check (num_kv_shared_layers=2)	before	after
max abs(model(x) - model(x, cache=fresh))	1.736722	0.0
max abs(prefill_6_tokens - decode_6x1_tokens)	0.868520	0.0

With num_kv_shared_layers=0 (sharing off) both are 0.0 before and after, which is the control that pins the KV-shared path as the cause.

Change

Thread the source layer's (keys, values) and its RoPE offset through the decoder loop instead of recovering them from a cache — the same shape gemma4_text already uses (Attention.__call__(..., shared_kv=..., offset=...) returning (out, (keys, values), offset), with Gemma4TextModel keeping an intermediates list). Gemma3nModel already computes layer_idx_to_cache_idx, which is exactly the "source layer of the same type" mapping, so it doubles as the index into intermediates and no new bookkeeping is needed.

Sharing now behaves identically with and without a cache, and the shared layers rotate their queries at the same offset as the layer whose K/V they reuse.

Deliberately not changed: shared layers keep their (now unused in the shared path) k_proj/v_proj/k_norm/v_norm parameters, so weight loading is untouched. Dropping them the way gemma4_text does would change what strict=True accepts and belongs in a separate change.

Tests

New test_gemma3n_kv_sharing_is_cache_independent asserts both invariants, for num_kv_shared_layers of 2 and 0.

Verified with mlx-cuda 0.30.0 on the CPU backend (no Apple hardware available here):

• new test: fails on main (assertTrue(mx.allclose(no_cache, cached, ...)) → array(False)), passes with this change.
• full tests/test_models.py: 78 tests, the same 3 errors before and after (test_bitnet, test_gated_delta, test_gated_delta_masked — custom-Metal-kernel tests that don't run on this backend). No new failures; test_all_models covers gemma3n and passes.
• black 25.1.0 and isort 6.0.0 (--profile=black) report both files unchanged.

🤖 Generated with Claude Code