[Yu-Haikuo]

Summary

This PR fixes #1674, where Gemma 3n’s KV-shared layers previously computed their own keys and values when no cache was provided, causing paths such as LoRA training, perplexity, and quantization calibration to differ from cached generation.

The fix passes the source layer’s keys and values directly to shared layers, making KV sharing consistent with or without a cache.

Also crediting #1702 for earlier work.

How to Reproduce

Scripts

Both scripts use small models with fixed-seed random weights, so no model download is required.

reproduce_gemma3n_kv_sharing_bug.py

# gemma3n: compare a forward pass without a cache to one with a fresh cache.
# Both must give the same logits. Random weights, so no download is needed.
import mlx.core as mx
from mlx_lm.models import gemma3n
from mlx_lm.models.cache import make_prompt_cache


def build(num_kv_shared_layers):
    args = gemma3n.ModelArgs(
        model_type="gemma3n",
        text_config={
            "model_type": "gemma3n",
            "hidden_size": 32,
            "num_hidden_layers": 4,
            "intermediate_size": 64,
            "num_attention_heads": 2,
            "head_dim": 16,
            "rms_norm_eps": 1e-5,
            "vocab_size": 64,
            "num_key_value_heads": 1,
            "num_kv_shared_layers": num_kv_shared_layers,
            "vocab_size_per_layer_input": 64,
            "sliding_window": 8,
            "max_position_embeddings": 128,
            "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.0, 0.0, 0.0, 0.0],
            "hidden_size_per_layer_input": 8,
            "altup_num_inputs": 1,
            "altup_coef_clip": 1.0,
            "altup_correct_scale": True,
            "altup_active_idx": 0,
            "laurel_rank": 4,
        },
    )
    mx.random.seed(0)
    model = gemma3n.Model(args)
    mx.eval(model.parameters())
    return model


inputs = mx.array([[1, 2, 3, 4, 5, 6]])
# 0 shared layers turns KV sharing off and is the control.
for num_kv_shared_layers in (2, 0):
    model = build(num_kv_shared_layers)
    no_cache = model(inputs)
    cached = model(inputs, cache=make_prompt_cache(model))
    diff = mx.abs(no_cache - cached).max().item()
    print(f"num_kv_shared_layers={num_kv_shared_layers}: max |no_cache - cached| = {diff:.4g}")

check_gemma3n_cached_output_unchanged.py

# gemma3n: print a SHA-256 prefix for logits from the cached generation path.
# Compare main and the fix in the same environment.
import hashlib

import mlx.core as mx
import numpy as np
from mlx_lm.models import gemma3n
from mlx_lm.models.cache import make_prompt_cache

# 6 layers, 3 shared. The full and sliding source layers are different layers.
layer_types = ["sliding_attention", "full_attention", "sliding_attention"] * 2
args = gemma3n.ModelArgs(
    model_type="gemma3n",
    text_config={
        "model_type": "gemma3n",
        "hidden_size": 32,
        "num_hidden_layers": 6,
        "intermediate_size": 64,
        "num_attention_heads": 2,
        "head_dim": 16,
        "rms_norm_eps": 1e-5,
        "vocab_size": 64,
        "num_key_value_heads": 1,
        "num_kv_shared_layers": 3,
        "vocab_size_per_layer_input": 64,
        "sliding_window": 8,
        "max_position_embeddings": 128,
        "rope_local_base_freq": 1.0,
        "rope_theta": 1000.0,
        "final_logit_softcapping": 1.0,
        "layer_types": layer_types,
        "activation_sparsity_pattern": [0.0] * 6,
        "hidden_size_per_layer_input": 8,
        "altup_num_inputs": 1,
        "altup_coef_clip": 1.0,
        "altup_correct_scale": True,
        "altup_active_idx": 0,
        "laurel_rank": 4,
    },
)
mx.random.seed(0)
model = gemma3n.Model(args)
mx.eval(model.parameters())

logits = []

# Prefill 10 tokens with an 8-token window, then decode 12 tokens.
cache = make_prompt_cache(model)
logits.append(model(mx.array([list(range(1, 11))]), cache=cache))
for token in range(20, 32):
    logits.append(model(mx.array([[token]]), cache=cache))

# Prefill in chunks across the window boundary.
cache = make_prompt_cache(model)
for chunk in ([1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12], [13, 14, 15]):
    logits.append(model(mx.array([chunk]), cache=cache))

digest = hashlib.sha256()
for out in logits:
    digest.update(np.array(out, dtype=np.float32).tobytes())
print(f"cached logits: {len(logits)} arrays, sha256 prefix {digest.hexdigest()[:16]}")

Script	Purpose
reproduce_gemma3n_kv_sharing_bug.py	Reproduces the bug and verifies the fix. With sharing enabled, the no-cache/fresh-cache difference should become 0. The sharing-disabled control is already 0.
check_gemma3n_cached_output_unchanged.py	Checks cached generation for regressions. Covers a 10-token prefill with an 8-token sliding window, 12 decode steps, and chunked prefill.

How to Verify

Please run the below commands on this fix branch and main branch:

python reproduce_gemma3n_kv_sharing_bug.py
python check_gemma3n_cached_output_unchanged.py

Results

Before the Fix

$ python reproduce_gemma3n_kv_sharing_bug.py
num_kv_shared_layers=2: max |no_cache - cached| = 1.445
num_kv_shared_layers=0: max |no_cache - cached| = 0

$ python check_gemma3n_cached_output_unchanged.py
cached logits: 16 arrays, sha256 prefix 99f9d3f4d6f5bf4b

After the Fix

$ python reproduce_gemma3n_kv_sharing_bug.py
num_kv_shared_layers=2: max |no_cache - cached| = 0
num_kv_shared_layers=0: max |no_cache - cached| = 0

$ python check_gemma3n_cached_output_unchanged.py
cached logits: 16 arrays, sha256 prefix 99f9d3f4d6f5bf4b

Environment

Component	Version / Details
Machine	Apple M5 Pro (Mac17,9), 64 GB
OS	macOS 26.7.1 (25G241), arm64
Python	3.14.6
mlx	0.32.4
numpy	2.5.4

AI Disclose

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: I used Claude Code and Codex to help with problem exploration, environment preparation, coding, and testing.