Gemma 1–4 support and Gemma performance (1.2.0)

CHANGELOG.md

@@ -1,5 +1,29 @@
 # Changelog
 
+## 1.2.0
+
+Gemma support (`gemma`, `gemma2`, `gemma3` / `gemma3_text`, `gemma4` / `gemma4_text`, `gemma4_unified`):
+
+- Gemma 1–3: `(1 + w)` RMSNorm weights, scaled embeddings, GeGLU (exact or tanh GELU), sandwich norms fused with the residual add, attention and final-logit soft-capping (Gemma 2), sliding-window layers with their own RoPE base (Gemma 3). Multimodal `gemma3` exports run their language model.
+- Gemma 4 (dense models): 256/512 head sizes in one model, proportional RoPE, KV cache shared with an earlier layer, K = V global attention (12B), per-layer input embeddings (E2B/E4B), `layer_scalar`, mixed 4/8-bit (OptiQ) checkpoints. Mixture-of-experts Gemma 4 (26B-A4B) is not supported. Thinking mode is not wired up; models run with it off.
+- Chat formats: `<start_of_turn>` (Gemma 1–3) and `<|turn>` (Gemma 4), checked against transformers' `apply_chat_template`. No tool calling for Gemma.
+- Tokenizer: Gemma's SentencePiece tokenizer adds no dummy prefix, and runs of spaces (added tokens) decode correctly.
+- Verified against mlx-lm, teacher-forced over 40 positions after a 664-token prompt: Gemma 1.1 2B, Gemma 2 2B, Gemma 3 270M/1B/4B, Gemma 4 E2B and 12B produce the same argmax at every position (one near-tie aside), logits within 0.035 of a float32 mlx-lm run.
+- Attention kernels take a sliding window and soft-cap; new kernel tests compare decode, split decode, prefill and flash prefill against a CPU reference.
+
+Gemma performance (M2 Pro, 3 alternating pairs against the first Gemma build; mlx-lm 0.32.2 on the same files, 200 greedy tokens):
+
+- The 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. The tail lives in its own instantiations: inside the one kernel it cost every model 2–11%.
+- Flash prefill for head_dim 256 (8-key blocks fit threadgroup memory): prefill +7% on Gemma 3 1B, +6% on Gemma 4 E2B (attention 66 → 43 ms and 177 → 121 ms per 512 tokens).
+- Split decode attention for head sizes 256 and 512 (query heads in sub-groups of 2, each K/V row read once per pair): decode at 1K context +3.5% on Gemma 3 1B, +3.7% on Gemma 3 4B, +6.1% on Gemma 4 E2B.
+- Gemma 4's per-layer input gate folds its GELU into the gate matvec's epilogue (one dispatch fewer per layer).
+- Against mlx-lm (mlx-lm → MLX Runner), decode after a short prompt: 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 (raw forward pass of 664 tokens vs `mlxr bench` at 512): 1B 1949 vs 1453, 2B 768 vs 538, 4B 378 vs 345, E2B 627 vs 545 tok/s; the remaining prefill gap is the engine's fp32 GEMM, not Gemma-specific.
+- Rejected after measuring (kept out): a gate/up matvec with the activation in its epilogue (−4% to −11% on every model), split-K for few-row matrices (−3% to −11% on Gemma), 8-byte lane loads for a 512 remainder.
+- Existing models: greedy output is token-identical and decode/prefill throughput is unchanged (3 alternating pairs on Qwen3 0.6B, Llama 3.2 1B, Mistral 7B, all within noise).
+- `mlxr bench` adds a per-kernel-class prefill breakdown; the profile sweep no longer crashes on 8-bit layers.
+
+Review fixes (from an adversarial review of the Gemma branch): Gemma replies keep their leading spaces and indentation (the dummy-prefix rule now applies only to tokenizers that have one); raw completions add `<bos>` when the tokenizer's post-processor does (Gemma 2, Llama 3); Gemma 3 encoding of long HTML/code chats is 100× faster (added tokens are matched with a trie instead of a 6,000-way regex); non-special added tokens such as Gemma's HTML tags are kept inside message content, as the reference tokenizer does; Gemma 4's per-layer-input norm uses the reference's exact eps; Gemma 4 mixture-of-experts repos are marked not runnable in Discover and My Models; Gemma 4 OptiQ exports no longer map or budget their 1 GB vision sidecar; flat `gemma3` / transformers 5 `rope_parameters` configs, the outer config's stop tokens, Gemma 2's default soft-caps and an out-of-range `sliding_window` are handled; later system messages render as `<|turn>system`; `layer_scalar` is read only for Gemma 4; model tests cover every installed Gemma checkpoint (Gemma 1.1 2B, 3 4B and 4 12B included) instead of passing silently when one is missing.
+
 ## 1.1.1
 
 Performance (M2 Pro, measured against 1.1.0 by alternating the two builds, 4 pairs; greedy output is token-identical on Qwen3 0.6B, Qwen2.5 0.5B, Llama 3.2 1B, Phi-3.5 and Mistral 7B):

README.md

@@ -90,13 +90,19 @@ gh attestation verify MLX-Runner-1.0.0-macos-arm64.dmg --repo BackendStack21/mlx
 | `llama` | Llama-3.2-1B-Instruct-4bit | Llama-3 RoPE scaling |
 | `mistral` | Mistral-7B-Instruct-v0.3-4bit | SentencePiece BPE + byte fallback |
 | `phi3` | Phi-3.5-mini-instruct-4bit | fused qkv / gate_up, LongRoPE |
+| `gemma` | gemma-1.1-2b-it-4bit | (1 + w) RMSNorm, GeGLU, scaled embeddings |
+| `gemma2` | gemma-2-2b-it-4bit | sandwich norms, attention and final-logit soft-capping, sliding window |
+| `gemma3_text`, `gemma3` | gemma-3-1b-it-qat-4bit, gemma-3-270m-it-4bit | q/k norms, 5:1 sliding/global layers with their own RoPE bases; `gemma3` multimodal exports run their language model |
+| `gemma4`, `gemma4_text` | gemma-4-e2b-it-OptiQ-4bit | head dims 256/512, proportional RoPE, KV sharing across layers, per-layer input embeddings, mixed 4/8-bit layers |
+| `gemma4_unified` | gemma-4-12B-it-OptiQ-4bit | as above, K = V in the global layers; the vision part of the export is ignored |
 
-Each checkpoint above passes end-to-end generation and the prefill-versus-decode consistency test. Other models of the same type, such as larger Qwen3, Qwen2.5-Coder or other Llama 3 fine-tunes, should work too, unless they need a feature that isn't implemented, such as sliding-window attention.
+Each checkpoint above passes end-to-end generation and the prefill-versus-decode consistency test; the Gemma checkpoints also match mlx-lm's logits position by position (teacher-forced: the same argmax at every position, within 0.05 of a float32 mlx-lm run). Other models of the same type, such as larger Qwen3, Qwen2.5-Coder or other Llama 3 fine-tunes, should work too, unless they need a feature that isn't implemented.
 
 ## Limitations
 
-- **Unsupported model types:** Gemma, mixture-of-experts (Mixtral, Qwen3-MoE) and vision-language models. Discover marks them as not runnable.
-- **Chat templates:** the prompt format is detected (ChatML, Llama 3, Mistral, Phi-3) rather than executed as Jinja, so a model with an unusual template may get a near-match format.
+- **Unsupported model types:** mixture-of-experts (Mixtral, Qwen3-MoE, Gemma 4 26B-A4B), Gemma 3n, and the vision and audio parts of multimodal models (Gemma 3 and Gemma 4 exports run as text-only). Discover marks unsupported types as not runnable.
+- **Gemma:** no tool calling and no Gemma 4 thinking mode (Gemma 4 runs with thinking off, its template's default). Sliding-window layers keep a full-length KV cache, so long contexts cost as much memory as on a model without sliding windows.
+- **Chat templates:** the prompt format is detected (ChatML, Llama 3, Mistral, Phi-3, Gemma, Gemma 4) rather than executed as Jinja, so a model with an unusual template may get a near-match format.
 - **Prefill speed:** about 165 tok/s on a 7B model on an M2 Pro, because the prefill GEMM reaches about 35% of fp32 peak.
 - **Speculative decoding is opt-in.** Prompt-lookup drafting with exact verification is implemented and tested (greedy output is identical), but checking 4 drafted tokens on a 7B model still costs about 3 decode steps. So it only helps when most drafts are accepted (copying or editing text), and it is off by default (`SamplingParameters.speculative`, or `mlxr run --spec`).
 - **Glass effects:** set `MLXRUNNER_NO_GLASS=1` to use plain materials instead of Liquid Glass.

Sources/MLXEngine/ChatTemplate.swift

@@ -19,18 +19,22 @@ public struct ChatMessage: Codable, Hashable, Sendable {
 /// Prompt formats for the supported model families. The format is detected
 /// from the model's Jinja `chat_template` instead of executing it.
 public enum ChatFormat: String, Sendable {
-    case chatML, llama3, mistral, phi3, plain
+    case chatML, llama3, mistral, phi3, gemma, gemma4, plain
 
     public static func detect(template: String?, modelType: String) -> ChatFormat {
         let t = template ?? ""
         if t.contains("<|im_start|>") { return .chatML }
         if t.contains("<|start_header_id|>") { return .llama3 }
         if t.contains("[INST]") { return .mistral }
         if t.contains("<|user|>") && t.contains("<|end|>") { return .phi3 }
+        if t.contains("<start_of_turn>") { return .gemma }
+        if t.contains("<|turn>") { return .gemma4 }
         switch modelType {
         case "qwen2", "qwen3": return .chatML
         case "mistral": return .mistral
         case "phi3": return .phi3
+        case "gemma", "gemma2", "gemma3", "gemma3_text": return .gemma
+        case "gemma4", "gemma4_text", "gemma4_unified", "gemma4_unified_text": return .gemma4
         case "llama": return t.isEmpty ? .plain : .llama3
         default: return .chatML
         }
@@ -52,22 +56,28 @@ public struct ChatTemplate: Sendable {
     /// Qwen2.5's template prints its tool-call example with doubled braces (`{{"name": …}}`, a
     /// Jinja quirk the model was trained with); Qwen3's prints single braces.
     let hermesDoubleBraces: Bool
+    /// Gemma 4 templates that close an empty thought channel after the model turn opener when
+    /// thinking is off (the 12B), so the model answers directly.
+    let emptyThoughtPrefill: Bool
 
     /// `toolTemplate`: a separate tool-use template, when the model ships one (checked for tool support).
     public init(template: String?, modelType: String, bosToken: String?, toolTemplate: String? = nil) {
         format = ChatFormat.detect(template: template, modelType: modelType)
-        supportsThinkingToggle = template?.contains("enable_thinking") ?? false
+        // Gemma 4's thinking (`<|think|>`, `<|channel>thought…<channel|>`) is not wired up: it runs
+        // with thinking off, its template's default.
+        supportsThinkingToggle = format != .gemma4 && (template?.contains("enable_thinking") ?? false)
         bos = bosToken ?? ""
         let qwen = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
         defaultSystem = format == .chatML && (template?.contains(qwen) ?? false) ? qwen : nil
         llamaDatePreamble = format == .llama3 && (template?.contains("Cutting Knowledge Date") ?? false)
         let t = (template ?? "") + (toolTemplate ?? "")
         hermesDoubleBraces = t.contains("{{\\\"name\\\": <function-name>")
+        emptyThoughtPrefill = format == .gemma4 && (template ?? "").contains("'<|channel>thought\\n<channel|>'")
         switch format {
         case .chatML: supportsTools = t.contains("<tool_call>")
         case .llama3: supportsTools = llamaDatePreamble && t.contains("ipython")
         case .mistral: supportsTools = t.contains("[AVAILABLE_TOOLS]")
-        case .phi3, .plain: supportsTools = false
+        case .phi3, .gemma, .gemma4, .plain: supportsTools = false
         }
     }
 
@@ -228,6 +238,39 @@ public struct ChatTemplate: Sendable {
                 lit("<|\(m.role.rawValue)|>\n"); content(m.content); lit("<|end|>\n")
             }
             lit("<|assistant|>\n")
+        case .gemma:
+            // Gemma's templates fold a system prompt into the first turn, call the assistant
+            // "model" and trim message content. There is no tool role: results go in as user turns.
+            lit(bos.isEmpty ? "<bos>" : bos)
+            var rest = messages[...]
+            var prefix: String?
+            if rest.first?.role == .system { prefix = rest.removeFirst().content + "\n\n" }
+            for m in rest {
+                let role = m.role == .assistant ? "model" : "user"
+                lit("<start_of_turn>\(role)\n")
+                if let p = prefix { content(p); prefix = nil }
+                content(m.content.trimmingCharacters(in: .whitespacesAndNewlines)); lit("<end_of_turn>\n")
+            }
+            if let p = prefix { lit("<start_of_turn>user\n"); content(p); lit("<end_of_turn>\n") }
+            lit("<start_of_turn>model\n")
+        case .gemma4:
+            // `<|turn>role … <turn|>`; a system prompt gets its own (trimmed) turn; consecutive
+            // assistant messages share one model turn. No tools: results go in as user turns.
+            lit(bos.isEmpty ? "<bos>" : bos)
+            var rest = messages[...]
+            if rest.first?.role == .system {
+                lit("<|turn>system\n"); content(rest.removeFirst().content.trimmingCharacters(in: .whitespacesAndNewlines)); lit("<turn|>\n")
+            }
+            let list = Array(rest)
+            for (i, m) in list.enumerated() {
+                let model = m.role == .assistant
+                let role = model ? "model" : m.role == .system ? "system" : "user"
+                if !(model && i > 0 && list[i - 1].role == .assistant) { lit("<|turn>\(role)\n") }
+                content(m.content.trimmingCharacters(in: .whitespacesAndNewlines))
+                if !(model && i + 1 < list.count && list[i + 1].role == .assistant) { lit("<turn|>\n") }
+            }
+            lit("<|turn>model\n")
+            if emptyThoughtPrefill { lit("<|channel>thought\n<channel|>") }
         case .plain:
             lit(bos)
             for m in messages {
@@ -252,6 +295,8 @@ public struct ChatTemplate: Sendable {
         case .llama3: return ["<|eot_id|>", "<|end_of_text|>", "<|eom_id|>"]
         case .mistral: return ["</s>"]
         case .phi3: return ["<|end|>", "<|endoftext|>"]
+        case .gemma: return ["<end_of_turn>", "<eos>"]
+        case .gemma4: return ["<turn|>", "<eos>"]
         case .plain: return []
         }
     }

Sources/MLXEngine/Engine.swift

@@ -454,7 +454,7 @@ public final class InferenceEngine: @unchecked Sendable {
             stats.generatedTokens += 1
             var text = decoder.append(tok.bytes(for: t, includeSpecial: reveal.contains(t)))
             // SentencePiece models start the reply with a word-boundary space.
-            if atStart, tok.mode == .sentencePiece { text = String(text.drop(while: { $0 == " " })) }
+            if atStart, tok.stripsLeadingSpace { text = String(text.drop(while: { $0 == " " })) }
             if !text.isEmpty { atStart = false; emit(.text(text)) }
             history.append(t)
             if nameOpen, let texts {
@@ -585,13 +585,13 @@ extension InferenceEngine {
     static func checkMemory(directory: URL, contextLength: Int) throws {
         let fm = FileManager.default
         var weights = 0
-        for f in (try? fm.contentsOfDirectory(atPath: directory.path)) ?? [] where f.hasSuffix(".safetensors") {
+        for f in (try? WeightStore.shardFiles(in: directory)) ?? [] {
             let p = directory.appendingPathComponent(f).resolvingSymlinksInPath().path
             weights += (try? fm.attributesOfItem(atPath: p)[.size] as? Int) ?? 0
         }
         let cfg = try ModelConfig(directory: directory)
         let ctx = min(contextLength, max(cfg.maxPositionEmbeddings, 256))
-        let kv = cfg.numLayers * 2 * cfg.numKVHeads * cfg.headDim * 2 * ctx
+        let kv = cfg.kvBytesPerToken * ctx
         let need = UInt64(weights + kv)
         let budget = GPU.shared.device.recommendedMaxWorkingSetSize
         guard need < budget else {