gemma3n: KV-shared layers only reuse cached KV when a cache is present, so no-cache forward (LoRA training, evaluation) diverges from generation

[binaydhakal]

Summary

In gemma3n, the KV-shared decoder layers only reuse the source layer's
keys/values when a cache is present. Without a cache they fall through and
compute their own k_proj/v_proj, so a no-cache forward pass does not apply KV
sharing and diverges from the cached (generation) path.

In Gemma3nAttention.__call__ (mlx_lm/models/gemma3n.py):

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
    offset = cache.offset
else:
    # shared layers ALSO land here when cache is None, and compute their own KV
    if cache is not None:
        offset = cache.offset
    keys = self.k_proj(x)...
    values = self.v_proj(x)...

Generation always runs with a cache, so generation is unaffected. But any
no-cache forward is wrong for the shared layers — this affects
mlx_lm.evaluate / perplexity and, notably, LoRA fine-tuning, whose forward
passes have no KV cache. The model being trained therefore does not match the
model used at inference.

Reproduction

A fresh, empty cache should be a no-op for a plain prefill, but for gemma3n it
changes the output. Building the model from the gemma3n entry in
tests/test_models.py::test_all_models (random weights):

import mlx.core as mx
from mlx_lm.models.cache import make_prompt_cache

toks = mx.array([[1, 2, 3, 4, 5, 6]])
no_cache = model(toks)
cached   = model(toks, cache=make_prompt_cache(model))
print(float(mx.max(mx.abs(no_cache - cached))))   # ~1.6, expected ~0

config	max|no_cache − cached|
num_kv_shared_layers=2 (default)	~1.6
num_kv_shared_layers=0 (sharing off)	0.0

Setting num_kv_shared_layers=0 removes the discrepancy entirely, confirming the
KV-shared path is the cause. A dense control model (e.g. helium) reports 0.0
for the same check. Weights are random, so the magnitude is arbitrary — the point
is that no_cache != cached for a model that shares KV.

Suggested direction

The no-cache path needs to apply the same KV sharing as the cached path. The
source (concrete) layers run before the shared layers, so their K/V could be made
available to the shared layers without going through a cache. (Internally
allocating a cache when cache is None would also make the two paths agree, but
at the cost of extra memory on every training/eval forward and possible
interaction with the rotating sliding-window cache, so a targeted reuse seems
preferable.) Happy to open a PR if a maintainer can confirm the intended approach.

Environment

• mlx-lm main @ 254d153
• mlx 0.32.0