[jkyberneees]

Summary

Gemma support for the engine, app, CLI and server: gemma (1), gemma2, gemma3 / gemma3_text, and the dense Gemma 4 models (gemma4 / gemma4_text, gemma4_unified). Verified checkpoints: gemma-1.1-2b, gemma-2-2b, gemma-3-270m / 1b / text-4b, gemma-4-e2b and gemma-4-12B (all mlx-community 4-bit / OptiQ).

• Model code: (1 + w) RMSNorm, scaled embeddings, GeGLU, sandwich norms fused into the residual-add kernel, attention and final-logit soft-capping, sliding-window layers with their own RoPE base, mixed 256/512 head sizes, proportional RoPE, KV cache shared across layers, K = V global attention, per-layer input embeddings, layer scalars, mixed 4/8-bit (OptiQ) checkpoints. Multimodal exports run their language model.
• Chat formats: <start_of_turn> and <|turn> renderers checked against transformers' apply_chat_template; Gemma's tokenizer adds no dummy prefix. No tool calling and no Gemma 4 thinking mode (runs with thinking off); Gemma 4 MoE (26B-A4B) and Gemma 3n are not supported and are marked not runnable.
• Correctness: teacher-forced logits match mlx-lm on all seven checkpoints over 40 positions after a 664-token prompt — the same argmax at every position, within 0.05 of a float32 mlx-lm run. Existing models: greedy output token-identical on Qwen3, Qwen2.5, Llama 3.2, Phi-3.5 and Mistral 7B.

Performance (M2 Pro, 3 alternating pairs; mlx-lm 0.32.2 on the same files)

• 4-bit decode matvec gets a kernel per remainder of hidden_size mod 1024 (Gemma's 1152 / 2304 / 2560 / 1536 left a pass with 4–16 of 32 lanes busy): decode +12% on Gemma 3 1B, +5% on Gemma 2 2B, +4% on Gemma 3 4B.
• Flash prefill for head_dim 256: prefill +6–7%. Split decode attention for head sizes 256/512 in head pairs: +3.5–6% at 1K context. Gemma 4's per-layer gate folds its GELU into the matvec epilogue.
• Decode vs mlx-lm (mlx-lm → this): Gemma 3 270M 396 → 481 tok/s, Gemma 3 1B 187 → 193, Gemma 2 2B 97 → 94, Gemma 3 4B 71 → 67, Gemma 4 E2B 78 → 73; after a 664-token prompt 1B 177 → 179, 2B 93 → 87, 4B 67 → 64, E2B 75 → 70. Prefill trails mlx-lm's raw forward by 9–30%: the engine's fp32 GEMM, not Gemma-specific.
• Rejected after measuring: a GLU epilogue in the gate/up matvec, split-K, tail code inside the one matvec kernel, 8-byte tail loads, 4 heads per split-attention simdgroup. Existing models' decode/prefill throughput is unchanged (within noise).

Review fixes

From a /code-review max pass on the first commit: Gemma replies keep leading spaces and indentation; raw completions add <bos> when the tokenizer's post-processor does; Gemma 3 encoding of long HTML/code chats is ~100× faster (trie instead of a 6,000-way regex); non-special added tokens are kept inside message content; Gemma 4's per-layer-input norm uses the reference's exact eps; MoE repos marked not runnable; OptiQ vision sidecars no longer mapped or budgeted; flat gemma3 / transformers 5 rope_parameters configs, outer eos union, Gemma 2 soft-cap defaults, bounded sliding_window; later system messages render as <|turn>system; layer_scalar read only for Gemma 4; model tests run for every installed Gemma checkpoint.

Not fixed (pre-existing, engine-wide, reported): MLX safetensors tensors are never 16-byte aligned, so SafeTensorFile copies every tensor at load instead of using the mmap zero-copy.

Test plan

• swift test — 131 tests pass (kernel tests for window/soft-cap/head sizes 128–512, split attention, 4-bit tails, Gemma configs/templates/tokenizers, prefill-vs-decode on every installed Gemma checkpoint)
• Teacher-forced parity with mlx-lm (float32) on all seven Gemma checkpoints
• Greedy output token-identical on the five existing architectures, short and 664-token prompts
• mlxr bench A/B against main, 3 alternating pairs, Qwen3 0.6B / Llama 3.2 1B / Mistral 7B within noise

🤖 Generated with Claude Code

[cloudflare-workers-and-pages]

🚀 Deploying Preview to Cloudflare 🚀

Preview Deployments by commit

Status	Deployment URL	Commit	Updated (UTC)	See this deployment's details
Build: Failed ❌ View logs ↗		bde1e20	2026-11-01T16:19:16.224Z	View logs ↗
Build: Failed ❌ View logs ↗		b5498d4	2026-11-01T16:05:53.405Z	View logs ↗