Add Gemma 4 assistant (MTP drafter) model class

[broomva]

What

Adds the gemma4_assistant model class — the lightweight drafter companion released by Google alongside Gemma 4 for multi-token-prediction (MTP) speculative decoding. Loading mlx-community/gemma-4-*-it-assistant-bf16 via mlx_lm.load(...) now succeeds and the forward pass produces correctly-shaped output.

Architecture

• Q-only cross-attention. Each of the 4 layers has only q_proj/q_norm/o_proj. K and V come from the target model via a shared_kv_states dict keyed by layer_type (full_attention or sliding_attention).
• Pre/post projection (5120 → 256 → 2560). Caller feeds concat(target_embed(last_token), target_last_hidden_state); the drafter returns its own last_hidden_state_2560 for the next draft step.
• Centroid logit head. 2048-cluster top-K=32 gating reduces logit computation ~64× over the 262144-token vocab. Optional via use_ordered_embeddings config flag.
• Per-layer global_head_dim dispatch. Full-attention layers use global_head_dim (512 in E4B), sliding-attention layers use head_dim (256) — mirrors gemma4_text.Attention.
• layer_scalar (per-layer 11th tensor) applied after the FFN residual.

Scope

This PR lands ONLY the model class. It loads cleanly via mlx_lm.load(...) and produces correctly-shaped output given synthetic shared_kv_states. Integration into the existing stream_generate speculative-decoding loop (target must extract per-layer-type K/V and forward the target's last-hidden-state to the drafter each step) is intended as a follow-up PR; happy to take guidance from maintainers on whether to land that as a separate PR or roll into this one.

Tests

• test_gemma4_assistant: synthetic forward in float32 + float16, covers all submodules, GQA expansion, centroid scatter, copy.deepcopy, make_cache raise, quant_predicate exclusion list
• test_gemma4_assistant_no_ordered_embeddings: non-clustered logit path (tied embedding matmul) — verifies the alternate code path
• test_gemma4_assistant_published_checkpoint_forward_shapes: opt-in via MLX_LM_RUN_NETWORK_TESTS=1, default-skipped in CI. Loads the real 159MB mlx-community/gemma-4-E4B-it-assistant-bf16 and runs one forward pass with synthetic K/V matching the real target's shapes

All existing test_gemma4_* tests still pass (10 default + 1 opt-in network test).

Verified locally

$ python -c "from mlx_lm import load; m, t = load('mlx-community/gemma-4-E4B-it-assistant-bf16'); print(type(m).__name__, len(m.layers))"
Model 4

$ pytest tests/test_models.py -k gemma4 -q
10 passed, 1 skipped

$ MLX_LM_RUN_NETWORK_TESTS=1 pytest tests/test_models.py::TestModels::test_gemma4_assistant_published_checkpoint_forward_shapes
1 passed

Notes on the design

• Why make_cache raises NotImplementedError — the drafter owns no KV cache of its own (it cross-attends to the target's shared_kv_states each forward). Returning [] would silently no-op in zip-based cache iteration paths, leading to attention without cached context. Raising loudly is the safer contract.
• Why _token_ordering (leading underscore) — the centroid lookup is an int32 buffer, not a parameter. Without the underscore prefix it shows up in Module.parameters() and gets corrupted by model.update(tree_map(astype, ...)). The leading underscore excludes it from parameters(). To make load_weights (which is strict and rejects underscored keys) happy, the buffer is installed directly on the submodule inside Model.sanitize and then stripped from the returned weight dict.
• No .item() calls in the forward path — MaskedEmbedder's mask_value is constructed on-device; AssistantAttention's RoPE offset is passed as mx.array. Both are in the MTP per-token hot loop where any GPU→host sync would defeat the speedup.

References

• Google blog (2026-05-05): https://blog.google/innovation-and-ai/technology/developers-tools/multi-token-prediction-gemma-4/
• HF reference implementation: transformers/models/gemma4_assistant/modeling_gemma4_assistant.py
• Drafter checkpoint: https://huggingface.co/mlx-community/gemma-4-E4B-it-assistant-bf16