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

mlx_lm/models/gemma3n.py

@@ -120,22 +120,21 @@ def __call__(
         x: mx.array,
         mask: Optional[mx.array] = None,
         cache: Optional[Any] = None,
-    ) -> mx.array:
+        shared_kv: Optional[tuple] = None,
+        offset: Optional[Any] = None,
+    ):
         B, L, _ = x.shape
 
         queries = self.q_proj(x)
         queries = queries.reshape(B, L, -1, self.head_dim)
         queries = self.q_norm(queries)
 
-        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
-
+        if shared_kv is not None:
+            # KV-shared layers reuse the keys/values and the RoPE offset of the
+            # designated source layer, whether or not a cache is present.
+            keys, values = shared_kv
         else:
-            if cache is not None:
-                offset = cache.offset
+            offset = cache.offset if cache is not None else 0
             keys = self.k_proj(x).reshape(B, L, -1, self.head_dim)
             keys = self.k_norm(keys)
             keys = keys.transpose(0, 2, 1, 3)
@@ -157,7 +156,7 @@ def __call__(
 
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
 
-        return self.o_proj(output)
+        return self.o_proj(output), (keys, values), offset
 
 
 @partial(mx.compile, shapeless=True)
@@ -328,17 +327,21 @@ def __call__(
         mask: Optional[mx.array] = None,
         cache: Optional[Any] = None,
         per_layer_input: Optional[mx.array] = None,
+        shared_kv: Optional[tuple] = None,
+        offset: Optional[Any] = None,
     ):
         predictions = self.altup.predict(x)
         active_prediction = predictions[self.config.altup_active_idx]
 
         active_prediction_normed = self.input_layernorm(active_prediction)
         laurel_output = self.laurel(active_prediction_normed)
 
-        attn = self.self_attn(
+        attn, kvs, offset = self.self_attn(
             active_prediction_normed,
             mask,
             cache,
+            shared_kv=shared_kv,
+            offset=offset,
         )
 
         attn = self.post_attention_layernorm(attn)
@@ -366,7 +369,7 @@ def __call__(
 
         corrected_predictions[1:] = corrected_predictions[1:] + first_prediction
 
-        return corrected_predictions
+        return corrected_predictions, kvs, offset
 
 
 @partial(mx.compile, shapeless=True)
@@ -491,6 +494,9 @@ def __call__(
         h = mx.stack(h_list, axis=0)
         mags = mx.mean(h[1:] ** 2, axis=-1, keepdims=True) ** 0.5
         h[1:] = h[1:] * (target_magnitude / mx.maximum(mags, mx.finfo(h0.dtype).min))
+        # Save each layer's keys/values and RoPE offset so the KV-shared layers
+        # can pick up the ones from their source layer.
+        intermediates = [(None, None)] * len(self.layers)
         for i, layer in enumerate(self.layers):
             per_layer_input = per_layer_inputs[:, :, i, :]
 
@@ -501,13 +507,19 @@ def __call__(
             else:
                 mask = sliding_window_mask
 
-            h = layer(
+            shared_kv, offset = intermediates[self.layer_idx_to_cache_idx[i]]
+
+            h, kvs, offset = layer(
                 h,
                 mask,
                 cache[self.layer_idx_to_cache_idx[i]],
                 per_layer_input,
+                shared_kv=shared_kv,
+                offset=offset,
             )
 
+            intermediates[i] = (kvs, offset)
+
         # Per-layer inputs to single output
         target_magnitude = mx.mean(h[0] ** 2, axis=-1, keepdims=True) ** 0.5
         for i, proj in enumerate(self.altup_unembed_projections):

tests/test_models.py

@@ -1727,6 +1727,72 @@ def test_gemma4_input_embeddings_reconstruct_per_layer_inputs(self):
             mx.allclose(direct.astype(mx.float32), explicit.astype(mx.float32))
         )
 
+    def test_gemma3n_kv_sharing_is_cache_independent(self):
+        """KV-shared layers must reuse the source layer's keys/values and RoPE
+        offset whether or not a cache is present, so a no-cache forward (LoRA
+        training, evaluation) agrees with generation."""
+        from mlx_lm.models import gemma3n
+
+        def build(num_kv_shared_layers):
+            args = gemma3n.ModelArgs(
+                model_type="gemma3n",
+                text_config={
+                    "model_type": "gemma3n",
+                    "hidden_size": 128,
+                    "num_hidden_layers": 4,
+                    "intermediate_size": 128,
+                    "num_attention_heads": 4,
+                    "head_dim": 32,
+                    "rms_norm_eps": 1e-5,
+                    "vocab_size": 1000,
+                    "num_key_value_heads": 2,
+                    "num_kv_shared_layers": num_kv_shared_layers,
+                    "vocab_size_per_layer_input": 1000,
+                    "sliding_window": 8,
+                    "max_position_embeddings": 1000,
+                    "rope_local_base_freq": 1.0,
+                    "rope_theta": 1000.0,
+                    "final_logit_softcapping": 1.0,
+                    "layer_types": [
+                        "sliding_attention",
+                        "full_attention",
+                        "sliding_attention",
+                        "full_attention",
+                    ],
+                    "activation_sparsity_pattern": [0.5, 0.5, 0.5, 0.5],
+                    "hidden_size_per_layer_input": 256,
+                    "altup_num_inputs": 4,
+                    "altup_coef_clip": 1.0,
+                    "altup_correct_scale": True,
+                    "altup_active_idx": 0,
+                    "laurel_rank": 8,
+                },
+            )
+            return gemma3n.Model(args)
+
+        tokens = [1, 2, 3, 4, 5, 6]
+        inputs = mx.array([tokens])
+
+        # num_kv_shared_layers=0 turns sharing off and is the control: it agreed
+        # on both checks before this behavior was fixed.
+        for num_kv_shared_layers in (2, 0):
+            model = build(num_kv_shared_layers)
+
+            # A forward with no cache must match a forward with a fresh cache.
+            no_cache = model(inputs)
+            cached = model(inputs, cache=make_prompt_cache(model))
+            self.assertTrue(mx.allclose(no_cache, cached, atol=1e-4, rtol=1e-4))
+
+            # Prefilling the prompt in one step must match decoding it token by
+            # token, which pins the RoPE offset the shared layers query with.
+            cache = make_prompt_cache(model)
+            prefill = model(inputs, cache=cache)[:, -1, :]
+
+            cache = make_prompt_cache(model)
+            for token in tokens:
+                incremental = model(mx.array([[token]]), cache=cache)[:, -1, :]
+            self.assertTrue(mx.allclose(prefill, incremental, atol=1e-4, rtol=1e-4))
+
     def test_gpt_bigcode(self):
         from mlx_lm.models import gpt_bigcode