Gemma 4 at every MLX affine width, uniform or per module: oQ packs on the native Metal engine (stacked on #975)

THIRD_PARTY_NOTICES.md

@@ -69,6 +69,45 @@ In the Python engine, frozen at 0.6.6, the MLX engine of `glm5_next`
   row, so each row keeps MLX's one-row bits. The native `zig/kernels/metal/glm/qmv*.metal`, `gemv*.metal`,
   `moe_*.metal` and `router.metal` are generated from these.
 
+## Gemma 4 on Apple Silicon
+
+The native Gemma 4 engine (`zig/src/families/gemma4/`) ports the Python engine's dense Gemma 4 decode, which @gilby
+wrote in [#360](https://github.com/ashhart/TensorFold/pull/360) on @cshintov's Gemma kernels from
+[#10](https://github.com/ashhart/TensorFold/pull/10) ([credits](https://github.com/ashhart/TensorFold/blob/main/docs/credits.md)).
+Its reference is gilby's `gemma4-dense` branch at `a6ff795e`: `RowDecode` in `src/tensorfold/kernels/gemma/v1/decode.py`
+and the `_EMIT`, `qkv_prep`, `attn_tail` and `dense_tail` kernels and the attention partials and merge in `glue.py` and
+`attention.py`. `zig/kernels/metal/gemma4/glue.metal` repeats that arithmetic, and each kernel names its source.
+
+Where that reference calls MLX or mlx-lm, `glue.metal` repeats their kernels' arithmetic (MLX and mlx-lm, MIT License,
+Copyright © 2023 Apple Inc.; see [the license text](LICENSES/MIT.txt)): `gemma_rms_looped` is MLX's `rms_looped` from
+`mlx/backend/metal/kernels/rms_norm.metal`, the kernel `mx.fast.rms_norm` runs for rows past 4,096, with its thread
+counts, reads and rounding points; `gemma_embed` is MLX's affine dequantize followed by a bf16 multiply;
+`gemma_geglu` and `gemma_softcap` are `geglu` and `logit_softcap` from mlx-lm's `models/gemma4_text.py` as MLX's
+bf16 ops; and `gemma_rope_freqs` is mlx-lm's `ProportionalRoPE` frequencies, computed as MLX's `pow` computes them.
+
+The image tower (`zig/src/families/gemma4/vision.zig`, `zig/kernels/metal/gemma4/vision.metal`) follows
+`Gemma4VisionModel` and `Gemma4MultimodalEmbedder` in transformers' `models/gemma4/modeling_gemma4.py`, and
+`zig/src/families/gemma4/image.zig` follows `Gemma4ImageProcessorPil` in `image_processing_pil_gemma4.py`
+(`get_aspect_ratio_preserving_size`, the patch layout and the rescale). Copyright 2026 the HuggingFace Team, Apache
+License 2.0; see [the license text](LICENSES/Apache-2.0.txt). The engine runs none of their code.
+
+The drafter (`zig/src/families/gemma4/dflash.zig`) implements z-lab's DFlash for Gemma 4 31B from the MLX model listed
+under [vendored code](#vendored-code-and-weights) (z-lab/dflash, MIT License, Copyright © 2026 Z Lab); the native
+engine vendors no copy.
+
+The 12B (`gemma4_unified`) and its MTP assistant (`zig/src/families/gemma4/assistant.zig`, `assistant_weights.zig`)
+follow @gilby's `gemma4_unified` family at `6b0ead7a` ([gilby/TensorFold](https://github.com/gilby/TensorFold),
+`src/tensorfold/families/gemma4_unified/model.py` and `assistant.py`), which follows `Gemma4UnifiedAssistantForCausalLM`
+in transformers. The 12B's image embedder (`zig/src/families/gemma4/vision_unified.zig`) follows
+`Gemma4UnifiedVisionEmbedder` in `models/gemma4_unified/modeling_gemma4_unified.py`, and `image.zig`'s merged patches
+`patches_merge` in `image_processing_gemma4_unified.py` (transformers, Copyright 2026 the HuggingFace Team, Apache
+License 2.0; see [the license text](LICENSES/Apache-2.0.txt)). The engine runs none of their code.
+
+Other widths (`zig/src/families/gemma4/widths.zig` and the stacks in `weights.zig`) follow @gilby's oQ per-layer widths
+at `76c83c37` on the `gemma4-dense` branch: each module's MLX affine width from the config, and a q|k|v or gate|up stack
+whose members differ split into runs of equal widths (`src/tensorfold/kernels/gemma/v1/matmul.py`). The engine runs none of
+that code.
+
 ## CUDA
 
 The Python engine's CUDA backends, frozen at 0.6.6, use [PyTorch](https://github.com/pytorch/pytorch), BSD-3-Clause, and

app/electron/agent/engine/cli.test.ts

@@ -46,14 +46,14 @@ async function main() {
     assert.equal(parseVersion((await runEngineCli(h.binary, ['--version'], h.env)).stdout), '1.0.7')
     assert.deepEqual(parseModelsOutput((await runEngineCli(h.binary, ['models'], h.env)).stdout), [])
     await seedModel(h.hub, 'mlx-community/Qwen3.5-2B-MLX-4bit', { bytes: 4096 })
-    await seedModel(h.hub, 'Org/Gemma', { modelType: 'gemma4' })
+    await seedModel(h.hub, 'Org/Gemma', { modelType: 'gemma4_unified' })
     assert.deepEqual(parseModelsOutput((await runEngineCli(h.binary, ['models'], h.env)).stdout).map((m) => m.repo), ['mlx-community/Qwen3.5-2B-MLX-4bit'])
     const ok = await runEngineCli(h.binary, ['info', 'mlx-community/Qwen3.5-2B-MLX-4bit'], h.env)
     assert.equal(ok.code, 0)
     assert.equal(parseInfoOutput(ok.stdout).format, '4-bit, group 64')
     const refused = await runEngineCli(h.binary, ['info', 'Org/Gemma'], h.env)
     assert.equal(refused.code, 1)
-    assert.match(refused.stdout, /no registered Zig family serves model_type gemma4/)
+    assert.match(refused.stdout, /no registered Zig family serves model_type gemma4_unified/)
     await assert.rejects(() => runEngineCli({ command: '/nonexistent/tensorfold-native', args: [] }, ['--version'], h.env), /Could not run tensorfold-native/)
   } finally {
     await h.cleanup()

app/electron/agent/engine/hub.test.ts

@@ -108,10 +108,10 @@ async function main() {
   assert.equal(installed.verdict, 'installed')
 
   // Unsupported family, and a drafter on its own.
-  const gemma = mockHub(repoRoutes('google/gemma-9', { modelType: 'gemma4' }))
+  const gemma = mockHub(repoRoutes('google/gemma-9', { modelType: 'gemma4_unified' }))
   const unsupported = await preflightRepo('google/gemma-9', null, context(gemma.fetchImpl))
   assert.equal(unsupported.verdict, 'unsupported')
-  assert.match(unsupported.reasons[0], /serves no gemma4 family/)
+  assert.match(unsupported.reasons[0], /serves no gemma4_unified family/)
   const drafterOnly = mockHub(repoRoutes('z-lab/Qwen3.8-27B-DFlash2', { dflash: true }))
   assert.equal((await preflightRepo('z-lab/Qwen3.8-27B-DFlash2', null, context(drafterOnly.fetchImpl))).verdict, 'unsupported')
 

app/electron/agent/engine/model-revisions.test.ts

@@ -28,7 +28,7 @@ function reviewedFetch(options: { invalidSha?: boolean; missingDrafter?: boolean
     if (route === `/api/models/${REPO}/revision/unsupported`) body = { sha: OTHER }
     if ([`/api/models/${REPO}/tree/${SHA}`, `/api/models/${REPO}/tree/${OTHER}`].includes(route)) body = [{ type: 'file', path: 'model.safetensors', size: 8192 }]
     if (route === `/${REPO}/resolve/${SHA}/config.json`) body = { model_type: 'qwen3_5', quantization: { bits: 4, group_size: 64 } }
-    if (route === `/${REPO}/resolve/${OTHER}/config.json`) body = { model_type: 'gemma4' }
+    if (route === `/${REPO}/resolve/${OTHER}/config.json`) body = { model_type: 'gemma4_unified' }
     return body ? new Response(JSON.stringify(body), { headers: { 'content-type': 'application/json' } }) : new Response('missing', { status: 404 })
   }) as typeof fetch
 }
@@ -45,7 +45,7 @@ async function main() {
     const pull = (await h.records()).find((r) => r.kind === 'pull')
     assert.equal(pull.revision, SHA, 'the actual child must pull the immutable checkpoint whose format and memory were reviewed')
     await service.preflight(REPO) // Re-populate the main-review cache after the completed pull invalidates it.
-    await assert.rejects(() => service.pull(REPO, 'unsupported'), /serves no gemma4/, 'a cached main preflight cannot authorize another revision')
+    await assert.rejects(() => service.pull(REPO, 'unsupported'), /serves no gemma4_unified/, 'a cached main preflight cannot authorize another revision')
 
     const registry = path.join(h.root, 'engine', 'model-revisions.json')
     const persisted = JSON.parse(await fs.readFile(registry, 'utf8'))
@@ -71,7 +71,7 @@ async function main() {
     assert.equal(serve.args[serve.args.indexOf('--name') + 1], REPO, 'logical model identity is preserved for the API')
     await assert.rejects(() => restarted.pull(REPO), /Unload/, 'a pull cannot change a currently serving selection')
     await restarted.unload()
-    await seedModel(h.hub, REPO, { modelType: 'gemma4' })
+    await seedModel(h.hub, REPO, { modelType: 'gemma4_unified' })
     const beforeRef = await fs.readFile(ref, 'utf8')
     restarted = fresh()
     assert.equal((await restarted.installed())[0].modelType, 'qwen3_5')

app/electron/agent/engine/service.test.ts

@@ -85,8 +85,8 @@ async function main() {
     await assert.rejects(() => offline.search('qwen'), /HF_HUB_OFFLINE/)
 
     // Unsupported repos are refused before any pull starts.
-    const refusing = new EngineService({ binaryPath: '/fake/tensorfold-native', binary: h.binary, appHome: h.root, childEnv: h.env, fetch: hubFetch('google/gemma-9', 1024, 'gemma4'), totalMemoryBytes: 32 * GIB, emit: () => {} })
-    await assert.rejects(() => refusing.pull('google/gemma-9'), /serves no gemma4 family/)
+    const refusing = new EngineService({ binaryPath: '/fake/tensorfold-native', binary: h.binary, appHome: h.root, childEnv: h.env, fetch: hubFetch('google/gemma-9', 1024, 'gemma4_unified'), totalMemoryBytes: 32 * GIB, emit: () => {} })
+    await assert.rejects(() => refusing.pull('google/gemma-9'), /serves no gemma4_unified family/)
     await assert.rejects(() => refusing.pull('../../etc'), /owner\/name/)
 
     // Load, then delete is refused until unloaded; delete removes only the cache folder.

build.zig

@@ -440,6 +440,12 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b
         .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" },
         .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" },
         .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" },
+        .{ .name = "tf-gemma4-kernels", .path = "zig/tests/gemma4/kernels.zig", .about = "Gemma 4 dense decode kernels against gemma4_dump.py's fixtures, bit for bit by window width" },
+        .{ .name = "tf-gemma4-glue", .path = "zig/tests/gemma4/glue.zig", .about = "Gemma 4 dense glue kernels against gemma4_dump.py's fixtures: embed, norm, q/k/v, attention, tails, softcap" },
+        .{ .name = "tf-gemma4-run", .path = "zig/tests/gemma4/run.zig", .about = "Gemma 4 dense greedy decode, tokens and top-two margins for gemma4_gate.py check" },
+        .{ .name = "tf-gemma4-lanes", .path = "zig/tests/gemma4/lanes.zig", .about = "Gemma 4 dense on the lane core: plain, copy-drafted and shared-round streams against the oracle's tokens" },
+        .{ .name = "tf-gemma4-vision", .path = "zig/tests/gemma4/vision.zig", .about = "Gemma 4's image tower against gemma4_vision_fixtures.py's dumps: patches bit for bit, each stage's error beside torch bf16's" },
+        .{ .name = "tf-gemma4-bench", .path = "zig/tests/gemma4/bench.zig", .about = "Gemma 4 dense step costs: projections alone by rows against whole forwards" },
         .{ .name = "tf-qwen36-row-seam", .path = "zig/tests/qwen36_row_seam.zig", .about = "Model-free fixture: the row projection's bias sums stay fp32 at every width" },
     };
     for (engine_programs) |p| {

docs/api.md

@@ -73,7 +73,7 @@ chat fields that shape the prompt (`tools`, `reasoning_effort`, `chat_template_k
 request, images and videos expanded. `/detokenize` takes `tokens` and returns `prompt`.
 In the Python 0.6.6 engine, `--vision` lets supported Qwen3.5/3.8 dense checkpoints accept `image_url` content parts
 alongside text in user messages and in tool results (`role: "tool"`), such as an agent's screenshots. The native
-server reads images for GLM-5.3-Flash and Qwen3.8-27B on Metal, with data URLs, and public HTTPS URLs under
+server reads images for GLM-5.3-Flash, Qwen3.8-27B and Gemma 4 31B on Metal, with data URLs, and public HTTPS URLs under
 `--vision-urls` ([image input](vision.md#the-native-server)); for GLM-5.3-Flash it also takes `video_url` parts in
 user messages ([video](vision.md#video-on-the-native-server)).
 See [image input](vision.md) for data URLs, public image URLs, limits and cache behavior.
@@ -264,16 +264,18 @@ macOS arm64, Linux x86_64 and Linux aarch64; from source, `zig build grammar -Dx
 v0.2.8 checkout (`git clone --branch v0.2.8 --recurse-submodules https://github.com/mlc-ai/xgrammar`) into
 `zig-out/native/lib` on macOS and `zig-out/lib` on Linux. The server looks in `../lib` beside its binary, then beside
 it, then on the loader's path, and loads it at the first structured request; without it, everything else serves and a
-structured request gets HTTP 400. Nemotron (Metal and CUDA) and Qwen3.5-2B (Metal) enforce it; other native engines
-answer a structured request with HTTP 400. Drafting stays on: the drafts the head holds on the GPU come back to the host
-before each verify, the grammar cuts the ones it rejects and drafts the text it forces itself (a key, closing markup),
-and each row's logits are masked before its draw, so drafted, `"draft": false` and concurrent replies are the same.
-
-On the native server, `tool_choice: "required"` and a named function are a grammar too, on the engines above: the answer
-is one or more calls in the chat template's own markup (`<tool_call>`, with Qwen3-Coder XML parameters as Nemotron and
-Qwen3.5 write them, or a JSON object), each call to an offered tool (the named one only) with its arguments inside the
-tool's parameter schema; blank lines may come between calls, prose may not. `parallel_tool_calls: false` allows exactly
-one call. With thinking on, the grammar starts after the think end.
+structured request gets HTTP 400. Nemotron (Metal and CUDA), Qwen3.5-2B and Gemma 4 31B (Metal) enforce it; other
+native engines answer a structured request with HTTP 400. Drafting stays on: the drafts the head holds on the GPU come
+back to the host before each verify, the grammar cuts the ones it rejects and drafts the text it forces itself (a key,
+closing markup), and each row's logits are masked before its draw, so drafted, `"draft": false` and concurrent replies
+are the same.
+
+On the native server, `tool_choice: "required"` and a named function are a grammar too, on Nemotron and Qwen3.5-2B
+(Gemma 4's template writes `<|tool_call>call:NAME{...}<tool_call|>`, so there they get HTTP 400 asking for `auto`): the
+answer is one or more calls in the chat template's own markup (`<tool_call>`, with Qwen3-Coder XML parameters as
+Nemotron and Qwen3.5 write them, or a JSON object), each call to an offered tool (the named one only) with its arguments
+inside the tool's parameter schema; blank lines may come between calls, prose may not. `parallel_tool_calls: false`
+allows exactly one call. With thinking on, the grammar starts after the think end.
 
 HTTP 400 comes before any token for a malformed field, a grammar xgrammar cannot compile (with xgrammar's reason), a
 server without xgrammar (with the install command; on the native server, a server that cannot load `libtfgrammar`),

docs/credits.md

@@ -107,6 +107,12 @@ contributions are in it.
   [prefill.py](https://github.com/ashhart/TensorFold/blob/v0.6.6/src/tensorfold/kernels/qwen/flash_next/v1/prefill.py)).
   The native hyper-connection prompt kernels in `zig/kernels/metal/prefill/fn_prompt.metal` are the #29 kernels, and
   `fn_attn.metal` reads only each row's selected blocks and tail, as the #29 attention did.
+- @gilby wrote the dense Gemma 4 31B decode (`attn_tail`, `dense_tail`, the ring cache) in
+  [#360](https://github.com/ashhart/TensorFold/pull/360), the native port's reference and source of its glue.
+- z-lab trained and published the Gemma 4 31B DFlash drafter
+  ([`z-lab/gemma-4-31B-it-DFlash`](https://huggingface.co/z-lab/gemma-4-31B-it-DFlash)); their MLX model, vendored in
+  the Python engine as `drafters/vendor/z_lab_dflash/model_mlx.py`, is the reference for the native drafter's block,
+  masks and positions (`zig/src/families/gemma4/dflash.zig`).
 - @gilby applied an n-gram table's `weight_scale` at lookup in [#148](https://github.com/ashhart/TensorFold/pull/148)
   ([embed.py](https://github.com/ashhart/TensorFold/blob/v0.6.6/src/tensorfold/kernels/qwen/flash_next/v1/embed.py)),
   so oQ checkpoints of Flash Next load. The native `zig/src/families/flashnext/ple_scale.zig` applies the same rule.
@@ -169,7 +175,8 @@ credited here.
   Flash Next vision patch adapted in #146), @mgoldwasser ([#101](https://github.com/ashhart/TensorFold/pull/101)),
   @shantanugoel ([#146](https://github.com/ashhart/TensorFold/pull/146)), @BobClawblaw
   ([#473](https://github.com/ashhart/TensorFold/pull/473)).
-- Model families: @cshintov (Gemma 4 in [#10](https://github.com/ashhart/TensorFold/pull/10), Qwen3.6 MoE on Macs in
+- Model families: @cshintov (Gemma 4 in [#10](https://github.com/ashhart/TensorFold/pull/10), whose kernels the native
+  dense 31B carries, Qwen3.6 MoE on Macs in
   [#138](https://github.com/ashhart/TensorFold/pull/138)), @di37 ([#224](https://github.com/ashhart/TensorFold/pull/224)),
   @philip-pentatonic (Qwen3.6 on CUDA in [#78](https://github.com/ashhart/TensorFold/pull/78) and
   [#84](https://github.com/ashhart/TensorFold/pull/84)), @BHCC2025 ([#270](https://github.com/ashhart/TensorFold/pull/270)).

docs/kernel-contract.md

@@ -94,8 +94,8 @@ A replacement kernel names the kernel it replaces and gives `replaced`, which ru
 
 A gate failure on a kernel already shipped is reported, not hidden. Until its fix lands, its spec can carry a
 `known` entry naming the gate and what fails. That gate must keep failing (the run prints `KNOWN`); once a fix makes
-it pass, the run fails until the entry is removed. Today: affine_mm's 4-bit form gives NaN for an infinite input where
-MLX's formula keeps the infinity, because its codes enter the MMAs as q + 128 and its bias as b - 128 s.
+it pass, the run fails until the entry is removed. Today there is none: affine_mm's 4-bit form gave NaN for an
+infinite input until its tile took each weight as MLX's loader does, bf16(fma(s, q, b)), instead of q + 128.
 
 ## Coverage
 

docs/recipes/README.md

@@ -13,6 +13,8 @@ the native engine, or the Python engine frozen at 0.6.6.
 | Ternary Bonsai 2 27B | [MLX](ternary-bonsai-2.md) |
 | GLM-5.3-Flash | [MLX on a 256 GB Mac, two-rank CUDA](glm-5.3-flash.md) |
 | Gemma 4 26B-A4B | [MLX, fused one-row decode](gemma-4.md) |
+| Gemma 4 31B | [Native Metal, affine 4-bit](gemma-4-native.md) |
+| Gemma 4 12B | [Native Metal, affine 4-bit, MTP assistant](gemma-4-12b-native.md) |
 | DeepSeek-V4-Flash | [MLX on a 256 GB Mac, DSpark and MTP drafts](deepseek-v4-flash.md) |
 | Qwen3.6-35B-A3B | [MLX with MTP drafts, one-GPU CUDA](qwen3.6-moe.md) |