Fix MLXVLM Gemma4 loader to honor num_kv_shared_layers (E-series targets)

[GoodOlClint]

Proposed changes

Fixes loading of the Gemma 4 E-series targets (E2B/E4B), which currently fail
with keyNotFound … layers.{15,24}.self_attn.k_proj.weight. This is independent
of MTP — plain loading (speculative=false) fails identically.

Root cause. The E-series checkpoints use num_kv_shared_layers > 0: the last
N decoder layers reuse an earlier layer's K/V and ship no k_proj/v_proj. They
carry vision_config, so they load through VLMModelFactory →
MLXVLM/Models/Gemma4.swift. That backbone builds every decoder layer with the
default kvSharedOnly: false, and Gemma4TextAttention.init gated K/V projection
construction on kvSharedOnly alone — even though it already computes
isKVSharedLayer. So the shared-tail layers declared a k_proj the checkpoint
doesn't have, and loadWeights threw. The text-only backbone
(MLXLLM/Models/Gemma4Text.swift) already gates on the shared-layer predicate;
this was drift between the two copies.

Fix. Gate K/V projection construction on !kvSharedOnly && !isKVSharedLayer,
matching the text backbone. The runtime already passes sharedKV to these layers
and sanitize already drops redundant K/V weights for PTQ checkpoints, so both
QAT (omit) and PTQ (ship+drop) layouts load. Dense / 26B-A4B MoE
(num_kv_shared_layers = 0) are unaffected.

Relates to #282 (Gemma 4 family loader gaps). Complements the E-series MTP
centroid-embedder PR — that fixes the drafter; this fixes the target it pairs with.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code
• I have added tests that prove my fix is effective
  • testGemma4TextSharedLayersDeclareNoKProj (unit, no checkpoint): a 4-layer
    model with 2 KV-shared layers declares k_proj only on the KV-owning layers.
  • Gemma4EseriesLoadIntegrationTests (cache-gated): the real E2B/E4B 4bit
    targets load through VLMModelFactory and generate. Verified locally — both
    load and produce output (the path that threw pre-fix).
• Documentation updates — n/a

[mandrael]

Independently hit this exact bug and can confirm the fix.

Repro: Loading mlx-community/gemma-4-E4B-it-qat-4bit through VLMModelFactory fails with UpdateError.keyNotFound at language_model.model.layers.24.self_attn.k_proj.weight. E4B has num_hidden_layers=42 / num_kv_shared_layers=18, so layers 24–41 are KV-shared and the checkpoint ships no k_proj/v_proj/k_norm for them, yet Gemma4TextAttention.init builds those modules for every layer. The PTQ ...-4bit build fails the same way one step earlier, at layers.24.self_attn.k_norm.weight: sanitize() (merged #330) correctly drops the redundant tensors, which leaves the wrongly-declared modules unfilled. So #330 alone doesn't cover this — the module construction still needs the gate this PR adds.

Confirmation: With this PR applied, gemma-4-E4B-it-qat-4bit loads and generates coherent output through VLMModelFactory. No regression for non-E-series configs: isKVSharedLayer is false whenever num_kv_shared_layers == 0, so those models build every layer exactly as before.

As a bonus this also unblocks the E-series gemma4_assistant MTP drafter end-to-end (paired with #383): with both applied, the E4B target + assistant runs speculative decoding losslessly (byte-identical greedy output) at ~2× tokens/s on an M-series Mac.

Thanks for the fix!

[davidkoski]

Changes look good but it is failing swift-format -- please rerun that. Thank you!

[GoodOlClint]

Done — rerun and pushed (785f963).

The failing file was actually an unrelated one: Tests/MLXLMTests/MTPSpeculativeTokenIteratorTests.swift, not the KV-sharing change. Root cause is a swift-format version skew: the CI lint job builds swift-format from releases/latest (currently 603.0.0), whose OrderedImports rule groups attributed imports (@_spi/@testable) into a trailing block, whereas the repo's pre-commit hook is language: system and picks up whatever's on the contributor's PATH (603 vs my local 602.x disagree on exactly this file). So it had drifted on main and pre-commit run --all (whole-repo) tripped it on this PR. Reformatted to 603's ordering here.

Heads-up for the maintainers: because CI tracks releases/latest while pre-commit is unpinned, every swift-format release that changes a rule can red-flag otherwise-untouched PRs until someone reformats. Pinning the CI swift-format tag to match a pinned pre-commit rev would stop the recurring churn — happy to send that as a separate PR if useful.

swift-format has been run

[davidkoski]

Ugh, trying to get CI to rerun on this.

[GoodOlClint]

Claude thinks it sees the issue with the CI, PR incoming.
#404 should hopefully unblock CI, looks like it's been failing across the board for a week.

[davidkoski]

Changes look good, thank you!