diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 74546d20cd..ae9d1bd81a 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/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 diff --git a/app/electron/agent/engine/cli.test.ts b/app/electron/agent/engine/cli.test.ts index bd92709797..f71aec3cdb 100644 --- a/app/electron/agent/engine/cli.test.ts +++ b/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() diff --git a/app/electron/agent/engine/hub.test.ts b/app/electron/agent/engine/hub.test.ts index 75f899d5df..57a5ab2d8d 100644 --- a/app/electron/agent/engine/hub.test.ts +++ b/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') diff --git a/app/electron/agent/engine/model-revisions.test.ts b/app/electron/agent/engine/model-revisions.test.ts index db652451a8..bad59dcd27 100644 --- a/app/electron/agent/engine/model-revisions.test.ts +++ b/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') diff --git a/app/electron/agent/engine/service.test.ts b/app/electron/agent/engine/service.test.ts index 129d1c95f2..71ceaff0d8 100644 --- a/app/electron/agent/engine/service.test.ts +++ b/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. diff --git a/build.zig b/build.zig index 380a38b3de..17feddc7d6 100644 --- a/build.zig +++ b/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| { diff --git a/docs/api.md b/docs/api.md index 85194d2150..16f39de73b 100644 --- a/docs/api.md +++ b/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 (``, 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{...}`, so there they get HTTP 400 asking for `auto`): the +answer is one or more calls in the chat template's own markup (``, 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`), diff --git a/docs/credits.md b/docs/credits.md index d9e3fcad1b..bf56fb06cb 100644 --- a/docs/credits.md +++ b/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)). diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index e625bcb2fd..c07b1abc32 100644 --- a/docs/kernel-contract.md +++ b/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 diff --git a/docs/recipes/README.md b/docs/recipes/README.md index f3e82f73b4..2fe6ccce32 100644 --- a/docs/recipes/README.md +++ b/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) | diff --git a/docs/recipes/gemma-4-12b-native.md b/docs/recipes/gemma-4-12b-native.md new file mode 100644 index 0000000000..5fa5f61be5 --- /dev/null +++ b/docs/recipes/gemma-4-12b-native.md @@ -0,0 +1,147 @@ +# Gemma 4 12B, native Metal + +This recipe covers text and image input with Google's +[`google/gemma-4-12B-it`](https://huggingface.co/google/gemma-4-12B-it/tree/707f0a3b8a3c7ad586ed01e27eafbad8a27dd0f7), +revision `707f0a3b8a3c7ad586ed01e27eafbad8a27dd0f7` (`gemma4_unified`), its language model converted to MLX affine +4-bit in groups of 64 and its image embedder kept in bf16, with Google's MTP assistant +[`google/gemma-4-12B-it-assistant`](https://huggingface.co/google/gemma-4-12B-it-assistant/tree/46d4c6f13f0ac0ad827b915669b8df9b81c64c51) +(revision `46d4c6f13f0ac0ad827b915669b8df9b81c64c51`) as the drafter. It runs on the native Gemma 4 family that serves +[the 31B](gemma-4-native.md), on M5 tensor-unit GPUs. Build with Zig **0.17.0** and Xcode's Metal toolchain, as +described in [the native preview](../../ZIG-PREVIEW.md#build). + +```bash +hf download google/gemma-4-12B-it --revision 707f0a3b8a3c7ad586ed01e27eafbad8a27dd0f7 --local-dir "$HOME/models/gemma-4-12b-hf" +hf download google/gemma-4-12B-it-assistant --revision 46d4c6f13f0ac0ad827b915669b8df9b81c64c51 \ + --local-dir "$HOME/models/gemma-4-12b-assistant" +python tools/zig/gemma4_12b_convert.py "$HOME/models/gemma-4-12b-hf" "$HOME/models/gemma-4-12b" # mlx and mlx-lm installed +zig build native -Dcpu=apple_m1 +zig-out/native/bin/tensorfold-native serve "$HOME/models/gemma-4-12b" --port 8888 --parallel 4 \ + --drafter "$HOME/models/gemma-4-12b-assistant" --vision +``` + +`gemma4_12b_convert.py` quantizes the language model's tensors with mlx-lm's converter to affine 4-bit in groups of +64 (6.3 GB, 4.5 bits a weight with the scales), writes the image embedder's ten tensors as stored (bf16, 100 MB) to +`model-vision.safetensors`, drops the audio input's, and writes the checkpoint's own `gemma4_unified` config back with +that quantization. Run twice, it writes the same tensors bit for bit. + +## The model + +The 12B is Gemma 4's dense decoder at another size: 48 layers, five sliding-window layers (window 1,024, head dim 256, +16 query heads over 8 key heads) to each global one (head dim 512 over one key head, keys reused as values, +proportional RoPE), hidden 3,840, GeGLU at 15,360, per-layer scalars, and the tied 262,144-token head with its logit +soft-cap at 30. The family reads both sizes from one geometry the config picks (`zig/src/families/gemma4/config.zig`): +every dimension must match the 31B's or the 12B's exactly, and other layouts (per-layer inputs, shared KV layers, MoE) +are refused by name. The kernels take each size as a define, so the 12B runs the 31B's kernels. + +Without `--drafter` the 12B decodes plain at 88 tokens a second for one stream. + +## The MTP assistant + +`--drafter` reads the drafter's `config.json`: a Gemma 4 assistant type goes to the MTP assistant +(`zig/src/families/gemma4/assistant.zig`), which admits `gemma4_unified_assistant` and refuses others by name, and a +DFlash checkpoint to DFlash (the 31B's), each refused by name on the wrong target. The assistant is four layers 1,024 +wide (three sliding, one global) with no keys or values of its own: each layer's queries attend over the target's own +caches, the last sliding layer's and the last global layer's, from the last kept row's position. Its input is the +pending token's scaled embedding beside the target's final-normed row, then its own projected output; its head is tied +to its own 262,144-row embedding. Its bf16 weights are quantized at load to 4-bit groups of 64 by MLX's affine rule. A +round drafts up to six tokens (0.97 ms a draft), every stream's chain in one command buffer, the next step's token +straight from the GPU's picks. Ordered (centroid) embeddings are refused by name. + +Served, one stream decodes at 140-248 tokens a second with it (prose 140, code 223, structured 205, list 248; +3.2 accepted tokens a round), 183 against 88 undrafted on the quality suite's draft check, with the same tokens. + +## Serving + +| Checkpoint / platform | Qualification | +| --- | --- | +| The converted 12B, M5 Ultra, macOS 27 | Greedy and sampled decode, prompts to 252k, shared rounds, MTP drafts, structured output, prompt reuse, the server, below. | +| `google/gemma-4-12B-it-assistant` `46d4c6f1` | The one MTP assistant admitted for the 12B. | +| M1 to M4 | Refused: the lane projection needs M5 tensor units. | +| Images | `--vision`: PNG and JPEG, 256 to 280 tokens an image, [below](#images). | +| Audio | Not read: the converter drops the audio input, and the server takes no audio. | +| Other widths | MLX affine 2-6 or 8 bits in groups of 64, uniform or per module: [other widths](#other-widths). | + +Structured output (`response_format`, `guided_*`) works with drafts on: the harness's structured check obeys 18 of 18 +(each field greedy, sampled and with thinking), and a bad regex gets HTTP 400. Thinking is on with +`chat_template_kwargs: {"enable_thinking": true}`; replies carry it in `reasoning_content`. + +Prompts are reused between requests as for the 31B: an 11,481-token conversation took 3.0 s to its first token fresh +and 0.09 s asked again (11,264 tokens resumed). A 1-token request after 2 to 120 idle seconds took 57-74 ms, where a +busy server takes 37 ms. + +Thinking replies fill the window: at `--context 32768`, three greedy prompts (a proof with a checked list, a +tested red-black tree, a question on a 25,801-token document) thought until the window was full (32,705, 32,714 and +6,959 tokens), each drafted reply equal to the same request with `"draft": false`. + +## Images + +`serve --vision` loads the image embedder (`vision_embedder.`, `embed_vision.`) and reads images as for +[the 31B](gemma-4-native.md#images): PNG and JPEG data URLs in user messages and tool results (public HTTPS URLs with +`--vision-urls`), decoded to Pillow's pixels by core's picture decoder (JPEG needs libjpeg-turbo 3.1.4.1, as +[image input](../vision.md#the-native-server) says). An image is sized as for the 31B (the largest aspect-preserving +size whose sides are multiples of 48 and whose 16-pixel patches fit 280 tokens, PIL's bicubic resize, then 1/255), +and `zig/src/families/gemma4/image.zig` cuts it into 48-pixel patches, the order transformers' +`Gemma4UnifiedImageProcessor` gives with `patches_merge`, one token each: 256 to 280 tokens (70 or 140 with +`detail: low` or a smaller `--vision-image-tokens`). + +The 12B has no image encoder: `zig/src/families/gemma4/vision_unified.zig` follows `Gemma4UnifiedVisionEmbedder`, +LayerNorm, the 6,912 to 3,840 Linear and LayerNorm on each patch, the x and y position rows and LayerNorm, then +the unscaled RMS norm and the 3,840 x 3,840 projection, in bf16 at torch's rounding on gemm_nax's tensor-unit +GEMMs (about 2 ms an image). The language model does the rest: the image's rows replace its tokens' embeddings, and on +the 40 sliding layers they see one another both ways (`use_bidirectional_attention: "vision"`, as for the 31B). + +Checked with `tools/zig/gemma4_vision_fixtures.py` (transformers 5.17 and torch, bf16 and fp32 runs of the embedder) +and `tf-gemma4-vision` on the 31B's five drawn images: the patches equal those of the PIL processor's resize through +`patches_merge` bit for bit, and each stage is as close to the fp32 run as torch's own bf16 run (the projected rows +0.30% relative error averaged over the images, torch's 0.31%). The checkpoint's own processor resizes with +torchvision, 0 to 2 levels of 255 from PIL's. Served with the assistant, `tools/zig/vision_served.py` passes its 19 +checks (drafted equal to `"draft": false`, concurrent image requests equal to their runs alone, a later turn resuming +past its image, the refusals), and questions about the images answer right (the circle red and the rectangle blue, +the circle on the left, the text "GEMMA 4", the tall image taller than wide). mlx-vlm 0.7.2 on the same checkpoint +gives the same pixels and the same five short answers; an 80-token description parts from it at its third token, on +a near-tie in mlx-vlm's own logits ("design" -1.625, "illustration" -1.75). + +## Other widths + +The family reads each linear's MLX affine width from the checkpoint's `quantization` table +(`zig/src/families/gemma4/widths.zig`, on core's `core/quant/mlx.zig`): 2, 3, 4, 5, 6 or 8 bits in groups of 64, +uniform or per module, the embedding and the image projection included. A tensor stored at another width than its +entry, groups of 32 or 128, 7 bits, other modes (mxfp4, nvfp4) and a module declared unquantized are refused by name, +never guessed from the shapes. A q|k|v or gate|up stack whose members differ runs one lane projection for each run of +equal widths, each writing its own columns of the stack's rows, as gilby's runs do. 4-bit groups of 64 keep the reg +kernel and, when every linear has that width, the fused one-row pipelines; other widths take the lane kernel's general +unpack, unfused, so they decode slower than their bytes alone would say. + +| Checkpoint | Checked | +| --- | --- | +| [`unigilby/gemma-4-12B-it-oQ8e`](https://huggingface.co/unigilby/gemma-4-12B-it-oQ8e/tree/a47ebf1a53598bb9ebe474724b20b2c0cd2e3b73) `a47ebf1a`, oMLX oQ uniform 8-bit | Decode equal to gilby's oracle on the four prompts x 128 tokens; long prompts against mlx-lm: 32k and 128k equal, 8k parts at a near-tie (23.375 / 23.125); `tf-gemma4-lanes` plain, copy and assistant drafts equal at one, four and eight streams, seeded at one and four; images: the embedder's rows 0.31% from fp32 (torch bf16 0.31%), `vision_served.py` passes. | +| `gemma4_12b_convert.py --recipe mixed_4_6` (v_proj 6-bit in 16 layers, down_proj in 24) | Decode: two of four prompts equal, two part at near-ties (22.875 / 22.75, 19.375 / 19.25); long prompts: 8k and 32k equal, 128k parts at a near-tie (27.375 / 27.25); lanes as above, all equal. | + +Served (harness target `gemma4-12b-oq8e-m5`), the oQ8e pack with the assistant: one stream 93, 142, 146 and 158 +tokens a second (prose, code, structured, list; 4-bit: 140, 223, 205, 248), four streams 262-340 together; the draft +check 137 against 46 undrafted (4-bit: 183 against 88); structured output 18 of 18, concurrent 12 of 12. + +## Quality + +The recipe harness (`gemma4-12b-m5`), thinking off: tools 58/60, JSON 30/30, GSM8K 239/250, MMLU 264/342; with +thinking on, the behaviour check's word problems 8 of 10 (every reply with its reasoning, none leaking tags). The +harness's thinking plan (`--plan thinking`) was not run for this release. + +## Verification + +The references are gilby's `gemma4_unified` Python engine (`6b0ead7a`, MLX 0.32.3, mlx-lm 0.32.0, lane kernels) for +decode and mlx-lm for long prompts (`tools/zig/gemma4_gate.py ... --family gemma4_unified`). + +- Decode, four chat prompts x 128 greedy tokens: three equal the reference's tokens; the fourth parts at its second + token on a near-tie (18.25 against 17.875 there, 18.125 for both here). As for the 31B since its speed work + (decode attention and the projections' K splits), the logits are not the reference's bit for bit, and the last + prompt rows sit as close to mlx-lm's as the reference's do (RMS 0.10-0.18 apart for both). +- Long prompts, 49 tokens after each, against mlx-lm in 2,048-token steps: 8k and 128k equal; 32k, 64k and 252k part + at near-ties (23.125 / 23.0, 25.625 / 25.5, and an exact tie 24.5 / 24.5 at 252k, where mlx-lm's own reply runs + into a loop). +- `tf-gemma4-lanes`: plain and copy drafts at one, four and eight streams, and the MTP assistant greedy at one, four + and eight and seeded sampled at one and four streams, all equal to their references. +- The 31B's decode, long-prompt, lanes, DFlash and image receipts are byte-identical to `21e8040e`'s. + +Native prompts in 512-row chunks (`tf-gemma4-run ... 49 512`): 3,961 tokens a second at 8k, 3,407 at 32k, 2,907 at +64k, 2,130 at 128k and 1,429 at 252k. diff --git a/docs/recipes/gemma-4-native.md b/docs/recipes/gemma-4-native.md new file mode 100644 index 0000000000..a333a823ca --- /dev/null +++ b/docs/recipes/gemma-4-native.md @@ -0,0 +1,204 @@ +# Gemma 4 31B, native Metal + +This recipe covers text generation with +[`mlx-community/gemma-4-31B-it-4bit`](https://huggingface.co/mlx-community/gemma-4-31B-it-4bit/tree/696d436c404745a59f30e4939a658162b0a9e57f), +revision `696d436c404745a59f30e4939a658162b0a9e57f`, the dense 31B in MLX affine 4-bit, groups of 64. +Build from `main` with Zig **0.17.0** and Xcode's Metal toolchain, as described in +[the native preview](../../ZIG-PREVIEW.md#build). + +```bash +hf download mlx-community/gemma-4-31B-it-4bit --revision 696d436c404745a59f30e4939a658162b0a9e57f \ + --local-dir "$HOME/models/gemma-4-31b" +hf download z-lab/gemma-4-31B-it-DFlash --revision eabd648301ce28583cc14757912e5e0f84e152e1 \ + --local-dir "$HOME/models/gemma-4-31b-dflash" +zig build native -Dcpu=apple_m1 +zig-out/native/bin/tensorfold-native serve "$HOME/models/gemma-4-31b" --port 8888 --parallel 4 \ + --drafter "$HOME/models/gemma-4-31b-dflash" +``` + +Without `--drafter` the model decodes plain, at about 51 tokens a second for one stream at short context (mlx-lm 46) +and 40 at 32k. With z-lab's DFlash drafter +a stream decodes at 150-190 tokens a second (6-7 accepted tokens a round on chat and code, 10 on copying replies, +2 on open-ended prose), with the same tokens. + +For the 31B the loader admits its dense geometry exactly: 60 layers, five sliding-window layers (window 1,024, head dim +256) to each global one (head dim 512, keys reused as values, proportional RoPE), GeGLU, the tied 262,144-token head +with its logit soft-cap, MLX affine weights at the widths [other widths](#other-widths) lists. It refuses the 26B-A4B +MoE, per-layer inputs and shared KV layers, each by name. With `--vision` it also loads the image tower +([images](#images)). The 12B (`gemma4_unified`) runs on the same family with its MTP assistant: [Gemma 4 12B, native +Metal](gemma-4-12b-native.md). + +Every linear, the tied head included, runs on the core lane projection (`zig/src/core/lane_projection.zig`) on M5 +tensor units; the rest of a layer is `zig/kernels/metal/gemma4/glue.metal`. The glue ports the Python reference's +decode kernels (gilby's `gemma4-dense` branch, #360, built on cshintov's Gemma kernels from #10) with their rounding +points, and MLX's own ops where that reference calls them: the looped RMS norm, the 4-bit embedding times its bf16 +scale, `geglu` and the soft-cap as bf16 ops, and the global layers' RoPE frequencies from MLX's precise `pow`. +Each output equals the reference's bit for bit at 1, 2 and 16 rows (`tf-gemma4-kernels`, `tf-gemma4-glue`, on +fixtures from `tools/zig/gemma4_dump.py`) at the reference's K split. Decode takes each shape's fastest split instead +(`weights.splitK`: 4 for o and down, 8 for the gate/up stack and the head), another fp32 summation order: at one row a +handful of outputs a linear move by one bf16 step (5 of 43,008 on gate_up, 99 of 262,144 on the head, none on qkv). + +A stream keeps a ring of window + 512 slots for each sliding layer and a buffer for each global one; rolling back is +moving the cache's length. A decode forward takes up to 16 rows of one stream and 128 of several, and every row keeps +its bits. A one-row window runs five or six dispatches a sliding layer rather than ten, and seven a global one: the +previous layer's dense tail is the qkv projection's prologue, the attention tail the gate/up projection's with geglu as +its epilogue (the stack's rows interleaved in 16-row blocks, so a lane holds a gate column and its up column), and the +down projection sums its own input (`zig/kernels/metal/gemma4/fused.metal` on the lane core's `Fused` hook); a sliding +layer's attention kernel preps its own query heads and writes the row's key and value, and with one key chunk writes +the merged output itself. The same bits as the kernels apart. +Prompts run in chunks of 512 rows on `core/affine_mm` and a tensor-unit attention kernel, at about 1,360 +tokens a second at 8k, 1,200 at 32k and 560 at 252k. + +The drafter is z-lab's DFlash (`zig/src/families/gemma4/dflash.zig`): five Qwen3-style layers reading six of the +target's layer outputs (1, 12, 23, 35, 46, 57) through `fc` and a norm, drafting a block of the pending token and up +to 15 mask tokens with the target's own embedding and tied head. Its bf16 weights are quantized at load to 4-bit +groups of 64 by MLX's affine rule and run on the same lane projection; the context keys and values of each stream +live in a 2,047-slot ring for the four sliding layers and a growing buffer for the full one, as in z-lab's MLX model. +It sits in front of the engine as the lane core's `Drafted` wrapper (chains, no trees): the target hands it each +prompt chunk's and each verify's tapped rows, and the depth rule prices its drafts against the forward's cost by rows +timed at load. A round's streams draft in one block forward (3.4 ms for one stream, 9 for eight) and absorb in one +pass, whose command buffer the block settles. + +The context defaults to 32,768 tokens; `--context` raises it to the checkpoint's 262,144, at the cost of attention +scratch that grows with it. + +Prompts are reused between requests (`zig/src/families/gemma4/snapshot.zig`). A kept state is each sliding layer's +last 1,024 keys and values, every global layer's keys and values so far, and the DFlash drafter's own context for the +same stream (4 KB a token in its full layer, 34 MB of sliding rings), so a resumed stream drafts as a fresh one does. A +prompt chunk's rows are their own arithmetic, so states are kept and resumed only where the request's planned chunks +start, as Step-5's are; a drafting stream does not resume a state kept without the drafter's context, and none is +resumed inside an image. The budget is what 70% of RAM leaves beside every stream's caches, 16 GiB by default +(`--prompt-cache-gib`, 0 for none). A later turn of a 3,336-token conversation resumed 3,329 tokens and answered as the +first pass did; one with an image after 2,300 tokens of system prompt resumed 3,584 of 3,607 tokens, alike. Served +with the drafter and `--vision`, an 11,481-token conversation took 7.8 s to its first token fresh and 0.23 s asked +again (11,264 tokens resumed); with an image in it, 8.1 s and 0.13 s. + +Structured output (`response_format`, `guided_json`, `guided_regex`, `guided_choice`, `guided_grammar`) is enforced +through `libtfgrammar`, as [structured output](../api.md#structured-output) describes, with drafts on: before each +verify the `Drafted` wrapper hands the drafter's held drafts to the host, where the grammar cuts the ones it rejects. +The recipe harness's structured check obeys 18 of 18 (each field greedy, sampled and with thinking), and a bad regex +gets HTTP 400. Tools are called with `tool_choice: "auto"`; `"required"` and a named function get HTTP 400, as the +grammar for them reads `` markup and Gemma's template writes `<|tool_call>call:NAME{...}`. The +weights, both scratches and the drafter's and tower's buffers stay in one residency set, which the idle keepalive's +tick keeps wired: a 1-token request after 2 to 120 idle seconds took 0.10-0.12 s to its first token, where a busy +server took 0.07 s (ten idle gaps, none slower). + +| Checkpoint / platform | Qualification | +| --- | --- | +| The pinned 31B 4-bit, M5 Ultra, macOS 27 | Greedy decode, prompts, shared rounds, copy and DFlash drafts and the server, below. | +| `z-lab/gemma-4-31B-it-DFlash` `eabd6483` | The one drafter admitted: DFlash v1, block 16, six taps, head dim 128. | +| M1 to M4 | Refused: the lane projection needs M5 tensor units. | +| Images (`--vision`) | PNG and JPEG data URLs, and public HTTPS URLs with `--vision-urls`, below. | +| Structured output | `response_format` and the `guided_*` fields, drafts on and off, below. | +| Other widths | MLX affine 2-6 or 8 bits in groups of 64, uniform or per module: [other widths](#other-widths). | + +## Images + +`serve --vision` loads the checkpoint's image tower (`vision_tower.`, `embed_vision.`; 1.1 GB, bf16 but for the 4-bit +projection) and reads PNG and JPEG data URLs in user messages and tool results (with `--vision-urls`, public HTTPS URLs +too), decoded to Pillow's pixels by core's picture decoder, as [image input](../vision.md#the-native-server) describes. +`zig/src/families/gemma4/image.zig` prepares an image as transformers' `Gemma4ImageProcessorPil` does: the +largest aspect-preserving size whose sides are multiples of 48 and whose 16-pixel patches fit 280 tokens (PIL's bicubic +resize, then 1/255), so an image becomes 256 to 280 tokens (70 or 140 with `detail: low` or a smaller +`--vision-image-tokens`). The template's `<|image|>` expands to `boi`, the image's rows, `eoi`. + +`zig/src/families/gemma4/vision.zig` runs the 27-layer encoder on gemm_nax's tensor-unit GEMMs at torch's per-op bf16 +rounding: the patch projection and the x and y position rows, RMS norms in fp32 times their weights, per-head q, k and +v norms, the axial 2D RoPE (theta 100, x on dims 0-35 and y on 36-71), attention with scale 1 as two batched GEMMs and +a softmax, GELU-tanh gating, the 3 x 3 pooler, the standardization and the projection on core affine_mm. An image takes +about 90 ms, and adds about 270 ms to a request's first token (a 290-token prompt with an image 365 ms, the text +alone 92 ms). In the language model the image's rows replace its tokens' embeddings, and on the 50 sliding layers they +see one another both ways, as `use_bidirectional_attention: "vision"` asks (the prompt attention kernel takes each +row's last key; global layers stay causal); a prompt chunk never cuts an image. + +Checked with `tools/zig/gemma4_vision_fixtures.py` (transformers 5.17 and torch, bf16 and fp32 runs of +`Gemma4VisionModel`) and `tf-gemma4-vision` on five drawn images, from 20 x 15 to 1200 x 800 and 180 x 1400: the patches +equal the PIL processor's bit for bit, and the tower's rows are as close to the fp32 run as torch's own bf16 run (7.6% +relative error averaged over the images, torch's 7.8%). AutoProcessor with torch installed picks the torchvision +processor instead, whose resize differs from PIL's by a level at a few pixels in a thousand. Served with the drafter, +a described image answers token for token as mlx-vlm 0.7.2 does (80 greedy tokens, the same 287-token prompt), with +drafts on and off; `tools/zig/vision_served.py` passes its 19 checks, concurrent image requests equal their runs +alone, and text-only prompts keep their bits (the receipts above, unchanged). + +## Other widths + +The family reads each linear's width from the checkpoint's `quantization` table, as [the 12B +recipe](gemma-4-12b-native.md#other-widths) describes: 2, 3, 4, 5, 6 or 8 bits in groups of 64, uniform or per module, +stacks of mixed widths in runs, every other width or mode refused by name. 4-bit groups of 64 keep the reg kernel and +the fused one-row pipelines: the 4-bit decode, 512-row prompt and 8k and 32k receipts are the same bytes as before +other widths were read. Other widths run the lane kernel's general unpack, unfused. + +[`unigilby/gemma-4-31B-it-qat-oQ8e`](https://huggingface.co/unigilby/gemma-4-31B-it-qat-oQ8e/tree/53dc2af091c166c8cefe5f670e744864f964ad47) +(`53dc2af0`, oMLX oQ uniform 8-bit from Google's QAT checkpoint, 33.8 GB): decode against gilby's oracle, the four +prompts x 128 tokens, two equal and two parting at near-ties (24.625 / 24.5, 21.25 / 21.125); `tf-gemma4-lanes` +plain, copy and DFlash drafts equal at one, four and eight streams, seeded at one and four; images: the tower's rows +7.6% from fp32 (torch bf16 7.8%), the 8-bit projection included, and `vision_served.py` passes. + +Served with DFlash (harness target `gemma4-31b-oq8e-m5`): one stream 40, 63, 52 and 55 tokens a second (prose, code, +structured, list), four streams 110-180 together; the draft check 50 against 22 undrafted; structured output 18 of 18, +concurrent 12 of 12. On the four prompts DFlash accepts 1.6-4.1 tokens a round on this pack against 2.0-5.8 on the 4-bit +one (`tf-gemma4-lanes`), whose unfused path accepts as many as its fused one. + +## Verification + +The oracle is gilby's `gemma4-dense` Python engine (`a6ff795e`, MLX 0.32.3, mlx-lm 0.32.0) with lane kernels. + +```bash +python tools/zig/gemma4_gate.py oracle --model M --out O --tokens 128 --prompt decode # in the oracle's venv +zig-out/bin/tf-gemma4-run M O N 128 16 && python tools/zig/gemma4_gate.py check --oracle O --native N +zig-out/bin/tf-gemma4-lanes M O && GEMMA_DRAFTER=D zig-out/bin/tf-gemma4-lanes M O +zig-out/bin/tf-gemma4-run M O R 128 512 && python tools/zig/gemma4_gate.py served --oracle O --native R +python tools/zig/gemma4_gate.py prompts --model M --out IDS # 8k, 32k, 64k, 128k and 252k prompts from the repo's text +python tools/zig/gemma4_gate.py oracle --model M --out L --ids IDS --tokens 49 --prompt mlx +zig-out/bin/tf-gemma4-run M L NL 49 512 && python tools/zig/gemma4_gate.py check --oracle L --native NL +``` + +- Four chat prompts, 128 greedy tokens, prompts fed one row at a time or in windows of 16: three equal the oracle + token for token; the story prompt parts at token 56, where the native run ties the oracle's top two (24.5 and 24.5, + the oracle 24.5 and 24.375) and takes the lower id. Every row keeps its bits at any window width and beside other + streams' rows. +- Long prompts against the oracle's mlx-lm prompt forward (fed in 2048-token chunks, as `mlx_lm.generate` feeds a + prompt), 49 greedy tokens each: equal at 32k, 64k and 252k. At 8k and 128k the two part at tokens 25 and 30, each + where mlx-lm's top two logits are equal (28.25 and 26.25); the native run puts them 0.125 and 0.75 apart. +- The lane core: copy-drafted and DFlash-drafted streams, alone and 4 and 8 in shared rounds, equal the engine's + plain replies token for token, on the four prompts and on two copy-heavy ones. +- Served on an M5 Ultra with the drafter: greedy token SHAs equal the engine's receipts alone, with drafts off and on, + and four at once; thinking replies keep their reasoning apart, and tool calls parse. + +With the drafter, one stream decodes at 150-190 tokens a second (6.4, 6.4, 2.1 and 6.7 tokens a round on the four +prompts; 10.2 and 7.8 on copying replies) and eight streams at about 210 together. Plain decode is 51 tokens a second +for one stream. The drafts match z-lab's own MLX drafter's acceptance. Shared rounds draft shallow on purpose: past 16 +rows the forward is bound by tensor-op compute at about 0.8 ms a row (40 ms at 32 rows, 71 at 64, 130 at 128, the +same rate as the prompt path and mlx-lm's), so the depth rule, given those costs, finds deeper shared drafts do not +pay; the drafter's own cost is 9 ms a round for eight streams. + +Long prompts, one after the other on the same M5 Ultra: native in 512-row chunks (`tf-gemma4-run ... 49 512`), mlx-lm +0.32.0 in 2048-token steps into a fresh cache. Native decode is the run's 49 tokens after the prompt, timed to 0.1 s. + +| Prompt | Native prompt tok/s | mlx-lm prompt tok/s | Native decode tok/s | mlx-lm decode tok/s | +|---|---|---|---|---| +| 8k | 1385 | 1305 | about 38 | 39 | +| 32k | 1241 | 1139 | about 35 | 32 | +| 64k | 1072 | 944 | about 33 | 32 | +| 128k | 882 | 651 | about 27 | 25 | +| 252k | 570 | 440 | about 20 | 16 | + +Prompt attention takes 32 query rows of a head a threadgroup: each simdgroup scores 32 keys of a block over the whole +head and accumulates its own 64 output columns, the row maxes and fp32 probabilities crossing threadgroup memory 128 +keys at a time. A 512-row chunk's 60 layers spend 138 ms there at 32k and 527 ms at 128k, where 16 rows a threadgroup +spent 194 and 810 (mlx-lm's fused attention about 161 at 32k): the kernel streams keys and values, so what each load +serves decides its rate, and vector loads, fetching a step ahead, and sparing interior blocks the mask and the rescale +moved nothing. The prompt GEMMs are core affine_mm's, which dequantizes each step's weights to bf16 in its tile (scale +times code plus bias in fp32, rounded once, as mlx-lm's quantized matmul does) and accumulates on the tensor unit over +the whole depth: a 512-row chunk's 60 layers take 284 ms there (mlx-lm's quantized matmuls 302), where the exact-codes +design, its per-group scale and bias epilogue and its row-sums pass took 344. So the prompt path rounds its weights as +mlx-lm does while the decode path's lane kernels keep the codes exact: over 40 tokens their last-row logits part by up +to 0.8. + +A plain token at short context is 19.7 ms on the M5 Ultra (mlx-lm 21.6), where it was 23.7: the 241 projections read +17.3 GB in about 18 ms at the decode splits, 20.3 at the reference's (the 168-tile o and down projections ran two +waves of 256-thread groups at a split of 8 and one wave at 4; the gate/up stack and the head streamed 84-group chains +at 1, eight times the simdgroups at 8); the attention's 120 dispatches took 0.9 ms and the glue's 0.7, which running +as the projections' prologues shortens by 0.2 (a prologue costs what the tail kernel cost, and only the dispatch gap +goes) and the attention's prep and merge folded into its kernel by 0.6. `tf-gemma4-bench` prints the stages +(`GEMMA_DSTAGES`), each shape plain and fused (`GEMMA_FUSED`) and the glue kernels alone (`GEMMA_GLUE`). diff --git a/docs/recipes/gemma-4.md b/docs/recipes/gemma-4.md index 3d53e29500..197afa7558 100644 --- a/docs/recipes/gemma-4.md +++ b/docs/recipes/gemma-4.md @@ -4,7 +4,8 @@ This page describes the Python engine's `gemma4` family, frozen at 0.6.6 and ava [`v0.6.6` tag](https://github.com/ashhart/TensorFold/tree/v0.6.6); the packages, tools and tests it names are at that tag. @cshintov wrote the family and its fused one-row kernels in [#10](https://github.com/ashhart/TensorFold/pull/10), and @di37 made its `qkv_rows` kernel launch on M1 and M2 in [#224](https://github.com/ashhart/TensorFold/pull/224). -The native binary has no Gemma 4 family yet, so `tensorfold pull` and `tensorfold info` refuse the checkpoint. +The native binary serves the dense 31B ([Gemma 4 31B, native Metal](gemma-4-native.md)) and the 12B +([Gemma 4 12B, native Metal](gemma-4-12b-native.md)); it refuses the 26B-A4B. Measured on an M4 Pro with 64 GB (273 GB/s) and MLX 0.31.2, with `mlx-community/gemma-4-26b-a4b-it-4bit`. Packages: [`src/tensorfold/families/gemma4/`](https://github.com/ashhart/TensorFold/tree/v0.6.6/src/tensorfold/families/gemma4) diff --git a/docs/vision.md b/docs/vision.md index 81c7b3610c..b674ce98fc 100644 --- a/docs/vision.md +++ b/docs/vision.md @@ -146,8 +146,9 @@ For Qwen, each image request carries its own multimodal rotary positions and con ## The native server `tensorfold-native serve` takes `--vision`, `--vision-max-images` (default 4) and `--vision-image-tokens` (default -4,096; 13,312 for Nemotron 3.5 Super VL) for engines that read images: GLM-5.3-Flash and Qwen3.8-27B on one Mac, -and Nemotron 3.5 Super VL on CUDA. Other native engines refuse the flag at startup. +4,096; 13,312 for Nemotron 3.5 Super VL) for engines that read images: GLM-5.3-Flash, Qwen3.8-27B, Gemma 4 +[31B](recipes/gemma-4-native.md#images) and [12B](recipes/gemma-4-12b-native.md#images) on one Mac, and Nemotron 3.5 +Super VL on CUDA. Other native engines refuse the flag at startup. - Images arrive as `image_url` data URLs in user messages and tool results. With `--vision-urls` (it needs `--vision`), a part may name a public HTTPS URL instead, fetched as 0.6.6 fetched it: port 443 without credentials @@ -166,8 +167,8 @@ and Nemotron 3.5 Super VL on CUDA. Other native engines refuse the flag at start folder, beside the server, Homebrew's `jpeg-turbo` on macOS, then the system's library path. It checks the version libjpeg-turbo writes into the library and refuses JPEG images by name when the library is missing or another release; PNG images still serve. WebP is refused by name until a decoder exists; other formats are refused. -- The limits and the processing are 0.6.6's: EXIF and XMP orientation, alpha composited over white, the GLM5-Next - processor's sizing and Pillow's bicubic resampler. The decoder gives 0.6.6's pixels for every picture in +- The limits and the processing are 0.6.6's: EXIF and XMP orientation, alpha composited over white, the model's + processor's sizing (GLM5-Next's, or Gemma 4's PIL processor) and Pillow's bicubic resampler. The decoder gives 0.6.6's pixels for every picture in `zig/tests/fixtures/lfm2/images.json` and `zig/tests/fixtures/picture/edges.json`, except where it departs on purpose: a CMYK JPEG is refused, and a JPEG holding several pictures (MPO) gives its first, which 0.6.6 refused. - Images are counted from the request's image parts. A user or tool message whose text spells the model's image diff --git a/tools/zig/gemma4_12b_convert.py b/tools/zig/gemma4_12b_convert.py new file mode 100644 index 0000000000..164a7cf9bd --- /dev/null +++ b/tools/zig/gemma4_12b_convert.py @@ -0,0 +1,74 @@ +"""Gemma 4 12B (gemma4_unified) to the native checkpoint: text weights MLX affine 4-bit g64, the image embedder bf16.""" + +from __future__ import annotations + +import argparse +import json +import shutil +import tempfile +from pathlib import Path + +import mlx.core as mx +from mlx_lm import convert + +# the image embedder's tensors, kept in bf16 in their own shard; the audio input's are dropped +VISION = ("model.vision_embedder.", "model.embed_vision.") +DROPPED = VISION + ("model.embed_audio.",) +VISION_SHARD = "model-vision.safetensors" + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("source", type=Path, help="google/gemma-4-12B-it as downloaded (config.json, model*.safetensors)") + p.add_argument("output", type=Path) + p.add_argument("--recipe", help="an mlx-lm mixed recipe (mixed_4_6: v_proj and down_proj at 6 bits in some layers)") + args = p.parse_args() + if args.output.exists(): + raise SystemExit(f"{args.output} exists") + config = json.loads((args.source / "config.json").read_text()) + if config.get("model_type") != "gemma4_unified": + raise SystemExit("expected a gemma4_unified checkpoint") + with tempfile.TemporaryDirectory() as tmp: + staged = Path(tmp) + for f in args.source.iterdir(): + weights_file = f.name.endswith(".safetensors") or f.name == "model.safetensors.index.json" + if f.name != "config.json" and not weights_file: + (staged / f.name).symlink_to(f.resolve()) + weights = {} + for f in sorted(args.source.glob("*.safetensors")): + weights.update(mx.load(str(f))) + kept = {k: v for k, v in weights.items() if not k.startswith(DROPPED)} + if any(not k.startswith("model.language_model.") for k in kept): + raise SystemExit("unexpected tensors beside the language model") + print(f"{len(weights)} tensors, {len(kept)} kept (the language model)") + mx.save_safetensors(str(staged / "model.safetensors"), kept, metadata={"format": "mlx"}) + vision = {k.removeprefix("model."): v for k, v in weights.items() if k.startswith(VISION)} + if len(vision) != 10 or any(v.dtype != mx.bfloat16 for v in vision.values()): + raise SystemExit("expected the 10 bf16 tensors of the image embedder") + # mlx-lm's gemma4 module reads the unified text model as it reads the dense 31B's + remapped = dict(config, model_type="gemma4", architectures=["Gemma4ForConditionalGeneration"]) + remapped["text_config"] = dict(config["text_config"], model_type="gemma4_text") + (staged / "config.json").write_text(json.dumps(remapped, indent=2)) + convert(str(staged), str(args.output), quantize=True, q_bits=4, q_group_size=64, quant_predicate=args.recipe) + out = json.loads((args.output / "config.json").read_text()) + for key in ("model_type", "architectures"): + out[key] = config[key] + out["text_config"]["model_type"] = config["text_config"]["model_type"] + out["vision_config"] = config["vision_config"] + out.pop("audio_config", None) + (args.output / "config.json").write_text(json.dumps(out, indent=2) + "\n") + mx.save_safetensors(str(args.output / VISION_SHARD), vision, metadata={"format": "pt"}) + index_path = args.output / "model.safetensors.index.json" + index = json.loads(index_path.read_text()) + index["weight_map"].update({k: VISION_SHARD for k in sorted(vision)}) + index_path.write_text(json.dumps(index, indent=4) + "\n") + for name in ("chat_template.jinja", "processor_config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json"): + if (args.source / name).exists() and not (args.output / name).exists(): + shutil.copy(args.source / name, args.output / name) + q = out["quantization"] + widths = sorted({(v["bits"], v["group_size"]) for v in q.values() if isinstance(v, dict)} | {(q["bits"], q["group_size"])}) + print(f"wrote {args.output}: {out['model_type']}, affine (bits, group) {widths}") + + +if __name__ == "__main__": + main() diff --git a/tools/zig/gemma4_dump.py b/tools/zig/gemma4_dump.py new file mode 100644 index 0000000000..7b3cb62994 --- /dev/null +++ b/tools/zig/gemma4_dump.py @@ -0,0 +1,98 @@ +"""Gemma 4 dense kernel fixtures from gilby's RowDecode: each decode kernel's inputs and outputs as npy, by rows.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--layers", default="0,5", help="layers to dump (a sliding one and a global one)") + parser.add_argument("--rows", default="1,2,16", help="window widths, each after the same prompt") + parser.add_argument("--prompt", type=int, default=1100, help="tokens decoded before each window (past the window)") + args = parser.parse_args() + + import mlx.core as mx + import numpy as np + from mlx_lm.models.gemma4_text import geglu + from tensorfold.families.gemma4.model import load + from tensorfold.kernels.gemma.v1.attention import Rows, attend + from tensorfold.kernels.gemma.v1.glue import attn_tail, dense_tail, qkv_prep + + family, _ = load(args.model, backend="lane", check=False) + d = family.decode + out = args.out + out.mkdir(parents=True, exist_ok=False) + rng = np.random.default_rng(7) + + def save(name: str, a: mx.array) -> None: + a = a.astype(mx.float32) if a.dtype == mx.bfloat16 else a + np.save(out / f"{name}.npy", np.asarray(a)) + + for i in (int(x) for x in args.layers.split(",")): + save(f"L{i}.inv_freq", d.inv_freq[i]) + meta = {"layers": args.layers, "prompt": args.prompt, "eps": d.eps_value, "windows": {}} + for rows in (int(x) for x in args.rows.split(",")): + cache = family.make_cache() + layers = family._layers(cache) + prompt = rng.integers(1000, 200000, args.prompt).tolist() + for t in prompt: + mx.eval(family.hidden(mx.array([t]), cache)) + tokens = mx.array(rng.integers(1000, 200000, rows).tolist()) + p0 = int(layers[0].offset) + positions = list(range(p0, p0 + rows)) + at = mx.array(positions, dtype=mx.int32) + tag = f"r{rows}" + h = d.backbone.embed_tokens(tokens) * d.backbone.embed_scale + normed = mx.fast.rms_norm(h, d.layers[0].input_layernorm.weight, d.eps_value) + save(f"{tag}.tokens", tokens.astype(mx.int32)) + save(f"{tag}.embed", h) + save(f"{tag}.normed0", normed) + wanted = {int(x) for x in args.layers.split(",")} + for i in range(max(wanted) + 1): + layer, attn = d.layers[i], d.layers[i].self_attn + qkv = d.qkv[i](normed) + q, k, v = qkv_prep(qkv, attn.q_norm.weight, attn.k_norm.weight, d.inv_freq[i], at, d.eps, + heads=attn.n_heads, kv_heads=attn.n_kv_heads, head_dim=attn.head_dim, + values_are_keys=attn.use_k_eq_v) + sliding = d.sliding(i) + keys, values = layers[i].write(k[None], v[None]) + kinds = Rows(positions, d.window if sliding else 0, layers[i].ring, attn.head_dim) + o_in = attend(q, keys, values, kinds, k, v, float(attn.scale)) + o = d.o[i](o_in.reshape(rows, -1)) + pre = layer.pre_feedforward_layernorm.weight + hn, n_mlp, _, _ = attn_tail(h, o, layer.post_attention_layernorm.weight, pre, pre, pre, d.eps) + gu = d.gate_up[i](n_mlp) + gate, up = d.gate_up[i].split(gu) + act = geglu(gate, up) + y = d.down[i](act) + nxt = d.layers[i + 1].input_layernorm.weight if i + 1 < len(d.layers) else d.backbone.norm.weight + h2, normed2 = dense_tail(hn, y, layer.post_feedforward_layernorm.weight, layer.layer_scalar, nxt, d.eps) + if i in wanted: + p = f"{tag}.L{i}" + for name, a in (("x", normed), ("qkv", qkv), ("q", q), ("k", k), ("v", v), ("o_in", o_in), ("o", o), + ("h", h), ("hn", hn), ("n_mlp", n_mlp), ("gate_up", gu), ("act", act), ("y", y), + ("h2", h2), ("normed2", normed2)): + save(f"{p}.{name}", a) + # the cache rows the window's attention read: every earlier key and value in position order + lo = max(0, p0 - d.window + 1) if sliding else 0 + idx = mx.array([(x % layers[i].ring) if sliding else x for x in range(lo, p0)], dtype=mx.int32) + save(f"{p}.keys_before", keys[0][:, idx]) + save(f"{p}.values_before", values[0][:, idx]) + meta["windows"][f"{p}"] = {"positions": positions, "first_key": lo, "scale": float(attn.scale)} + h, normed = h2, normed2 + final = mx.random.normal((rows, d.layers[0].input_layernorm.weight.shape[0])).astype(mx.bfloat16) + logits = d.logits(final) + save(f"{tag}.head_in", final) + save(f"{tag}.head_raw", d.head(final)) + save(f"{tag}.logits", logits) + mx.eval(logits) + print(f"rows {rows}: dumped layers {sorted(wanted)} at positions {p0}..{p0 + rows - 1}", flush=True) + (out / "meta.json").write_text(json.dumps(meta, indent=1) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/zig/gemma4_gate.py b/tools/zig/gemma4_gate.py new file mode 100644 index 0000000000..939443f83c --- /dev/null +++ b/tools/zig/gemma4_gate.py @@ -0,0 +1,236 @@ +"""Gemma 4 dense token oracle (gilby's gemma4-dense Python engine, lane kernels) and the native engine's checks.""" +from __future__ import annotations + +import argparse +import json +import platform +import time +from pathlib import Path + +PROMPTS = { + "copy": "Reply with exactly: The quick brown fox jumps over the lazy dog.", + "arithmetic": "Compute 37 times 19 and explain the calculation briefly.", + "story": "Write a short story about a lighthouse keeper finding a message in a bottle.", + "code": "Write a Python function that returns the first n Fibonacci numbers.", +} + + +def first_difference(a: list[int], b: list[int]) -> int | None: + for i, (x, y) in enumerate(zip(a, b)): + if x != y: + return i + return min(len(a), len(b)) if len(a) != len(b) else None + + +def environment() -> dict: + import mlx.core as mx + import mlx_lm + import tensorfold + + return {"mlx": mx.__version__, "mlx_lm": mlx_lm.__version__, "tensorfold": tensorfold.__version__, + "device": mx.metal.device_info().get("device_name", ""), "macos": platform.mac_ver()[0]} + + +def chat_ids(tokenizer, text: str) -> list[int]: + ids = tokenizer.apply_chat_template([{"role": "user", "content": text}], tokenize=True, + add_generation_prompt=True, enable_thinking=False) + return list(ids["input_ids"] if hasattr(ids, "keys") else ids) # transformers 5 returns an encoding + + +def greedy(family, ids: list[int], count: int, prompt: str) -> tuple[list[int], list[dict]]: + """`count` greedy tokens after `ids`; the prompt through mlx_lm's forward ("mlx") or one decode row a token.""" + import mlx.core as mx + + from qwen27_margins import capture + + cache = family.make_cache() + if prompt == "mlx": + # 2048-token chunks, as mlx_lm's generate feeds a prompt: one 252k-row forward builds 252k x 252k masks + for at in range(0, len(ids), 2048): + hidden = family.prefill(mx.array([ids[at:at + 2048]]), cache) + mx.eval(hidden) + mx.clear_cache() + hidden = hidden[:, -1:] + else: + for token in ids[:-1]: + mx.eval(family.hidden(mx.array([token]), cache)) + hidden = family.hidden(mx.array([ids[-1]]), cache) + tokens, top2 = [], [] + for step in range(count): + logits = family.head(hidden) + mx.eval(logits) + if step == 0: + family.last_row = logits.reshape(-1).astype(mx.float32) + top2.extend(capture(logits, [len(ids) + step])) + token = top2[-1]["ids"][0] + tokens.append(token) + hidden = family.hidden(mx.array([token]), cache) + return tokens, top2 + + +def oracle(args: argparse.Namespace) -> None: + import tensorfold + + # the 31B's oracle is gilby's gemma4-dense branch (0.6.5), the 12B's his gemma4_unified commit (0.6.6) + want = "0.6.6" if args.family == "gemma4_unified" else "0.6.5" + if tensorfold.__version__ != want: + raise SystemExit(f"Need gilby's {args.family} oracle (tensorfold {want}), got {tensorfold.__version__}") + if args.family == "gemma4_unified": + from tensorfold.families.gemma4_unified.model import load + else: + from tensorfold.families.gemma4.model import load + + family, tokenizer = load(args.model, backend="lane", check=False) + out = args.out + out.mkdir(parents=True, exist_ok=False) + prompts = json.loads(args.ids.read_text()) if args.ids else {n: chat_ids(tokenizer, t) for n, t in PROMPTS.items()} + results = {} + for name, ids in prompts.items(): + (out / f"{name}.ids.json").write_text(json.dumps(ids) + "\n") + started = time.perf_counter() + tokens, top2 = greedy(family, ids, args.tokens, args.prompt) + results[name] = {"tokens": tokens, "elapsed_s": round(time.perf_counter() - started, 3)} + (out / f"{name}.reference.json").write_text(json.dumps(tokens) + "\n") + (out / f"{name}.top2.json").write_text(json.dumps(top2) + "\n") + import numpy as np + np.asarray(family.last_row).astype(np.float32).tofile(out / f"{name}.last.f32") + print(f"{name}: {len(ids)} prompt tokens, {tokenizer.decode(tokens)[:80]!r}", flush=True) + receipt = {"environment": environment(), "model": str(args.model), "prompt_path": args.prompt, + "decode": family.decode.backend, + "kinds": sorted({k for p in family.decode.qkv for k in getattr(p, "kinds", ())}), + "tokens": args.tokens, "results": results} + (out / "oracle.json").write_text(json.dumps(receipt, indent=1) + "\n") + + +def prompts(args: argparse.Namespace) -> None: + """Long chat prompts of about the requested sizes from the repo's own text, written as {name: ids} JSON.""" + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(str(args.model)) + root = Path(__file__).resolve().parents[2] + # the repo's own text, docs first then source, until the largest prompt fits (no passage repeats) + files = [root / f for f in ("CHANGELOG.md", "README.md", "CONTRIBUTING.md", "docs/api.md")] + files += sorted(p for p in (root / "docs").rglob("*.md") if p not in files) + sorted((root / "zig" / "src").rglob("*.zig")) + words: list[int] = [] + for f in files: + if len(words) >= max(args.sizes): + break + words += tokenizer.encode(f.read_text(errors="replace") + "\n\n", add_special_tokens=False) + if len(words) < max(args.sizes): + raise SystemExit(f"only {len(words)} tokens of repo text for a {max(args.sizes)}-token prompt") + out = {} + for size in args.sizes: + body = tokenizer.decode(words[:size]) + ask = "\n\nIn one paragraph, what does the text above say about prompt caching?" + out[f"long{round(size / 1024)}k"] = chat_ids(tokenizer, body + ask) + args.out.write_text(json.dumps(out) + "\n") + print({k: len(v) for k, v in out.items()}) + + +def check(args: argparse.Namespace) -> None: + """Compare a native run's tokens and top-two margins with the oracle's, prompt by prompt.""" + from qwen27_margins import verdict + + oracle_dir, native_dir = args.oracle, args.native + report, failed = {}, False + for ref in sorted(oracle_dir.glob("*.reference.json")): + name = ref.name.removesuffix(".reference.json") + want = json.loads(ref.read_text()) + got = json.loads((native_dir / f"{name}.tokens.json").read_text()) + at = first_difference(want, got) + entry = {"equal": at is None, "first_difference": at} + if at is not None and at < min(len(want), len(got)): + python = json.loads((oracle_dir / f"{name}.top2.json").read_text())[at] + native = json.loads((native_dir / f"{name}.top2.json").read_text())[at] + entry["margins"] = verdict(native, python, want[at], got[at]) + failed |= at is not None + report[name] = entry + print(f"{name}: {'equal' if at is None else f'differs at {at}'}", flush=True) + (native_dir / "token-gate.json").write_text(json.dumps(report, indent=1) + "\n") + if failed: + raise SystemExit(1) + + +def post(url: str, path: str, body: dict) -> dict: + from urllib.request import Request, urlopen + + request = Request(url.rstrip("/") + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}) + with urlopen(request, timeout=900) as response: + return json.load(response) + + +def served(args: argparse.Namespace) -> None: + """The served model's greedy replies equal the oracle's, alone and four at once; thinking and tool calls parse.""" + import hashlib + from concurrent.futures import ThreadPoolExecutor + + names = sorted(p.name.removesuffix(".reference.json") for p in args.oracle.glob("*.reference.json")) + # the native CLI's tokens when given (its prompt path), else the oracle's + src = args.native or args.oracle + want = {n: json.loads((src / f"{n}.{'tokens' if args.native else 'reference'}.json").read_text()) for n in names} + ids = {n: json.loads((args.oracle / f"{n}.ids.json").read_text()) for n in names} + + def one(name: str, draft: bool) -> tuple[str, bool]: + got = post(args.url, "/v1/completions", {"model": "gemma", "prompt": ids[name], "max_tokens": len(want[name]), + "temperature": 0, "draft": draft, "ignore_eos": True}) + digest = hashlib.sha256(",".join(str(t) for t in want[name]).encode()).hexdigest()[:12] + return got.get("tensorfold", {}).get("token_sha", ""), got.get("tensorfold", {}).get("token_sha") == digest + + failed = 0 + for draft in (False, True): + for name in names: + sha, ok = one(name, draft) + failed += not ok + print(f"{'ok ' if ok else 'FAIL'} alone, drafts {'on ' if draft else 'off'}: {name} {sha}", flush=True) + with ThreadPoolExecutor(len(names)) as pool: + for name, (sha, ok) in zip(names, pool.map(lambda n: one(n, True), names)): + failed += not ok + print(f"{'ok ' if ok else 'FAIL'} four at once: {name} {sha}", flush=True) + think = post(args.url, "/v1/chat/completions", {"model": "gemma", "max_tokens": 800, "temperature": 0, + "chat_template_kwargs": {"enable_thinking": True}, + "messages": [{"role": "user", "content": "What is 17 * 23? Think first."}]}) + msg = think["choices"][0]["message"] + ok = bool(msg.get("reasoning_content")) and "391" in (msg.get("content") or "") and "<|channel>" not in (msg.get("content") or "") + failed += not ok + print(f"{'ok ' if ok else 'FAIL'} thinking apart from the answer: {len(msg.get('reasoning_content') or '')} reasoning chars, answer {(msg.get('content') or '')[:60]!r}") + tools = [{"type": "function", "function": {"name": "get_weather", "description": "Weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, + "required": ["city"]}}}] + call = post(args.url, "/v1/chat/completions", {"model": "gemma", "max_tokens": 200, "temperature": 0, "tools": tools, + "messages": [{"role": "user", "content": "What's the weather in Lisbon?"}]}) + calls = call["choices"][0]["message"].get("tool_calls") or [] + ok = len(calls) == 1 and calls[0]["function"]["name"] == "get_weather" and "Lisbon" in calls[0]["function"]["arguments"] + failed += not ok + print(f"{'ok ' if ok else 'FAIL'} tool call: {[c['function'] for c in calls]}") + print("PASS" if not failed else f"FAIL: {failed}") + if failed: + raise SystemExit(1) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="mode", required=True) + o = sub.add_parser("oracle") + o.add_argument("--model", type=Path, required=True) + o.add_argument("--out", type=Path, required=True) + o.add_argument("--ids", type=Path, help="JSON {name: [ids]} in place of the four chat prompts") + o.add_argument("--tokens", type=int, default=128) + o.add_argument("--prompt", choices=("decode", "mlx"), default="decode") + o.add_argument("--family", choices=("gemma4", "gemma4_unified"), default="gemma4", help="the 31B's, or the 12B's") + p = sub.add_parser("prompts") + p.add_argument("--model", type=Path, required=True) + p.add_argument("--out", type=Path, required=True) + p.add_argument("--sizes", type=int, nargs="+", default=[8192, 32768, 65536, 131072, 258048]) + v = sub.add_parser("served") + v.add_argument("--url", default="http://127.0.0.1:8888") + v.add_argument("--oracle", type=Path, required=True) + v.add_argument("--native", type=Path, help="tf-gemma4-run's output (prompts in 512-row chunks): the reference") + c = sub.add_parser("check") + c.add_argument("--oracle", type=Path, required=True) + c.add_argument("--native", type=Path, required=True) + args = parser.parse_args() + {"oracle": oracle, "prompts": prompts, "served": served, "check": check}[args.mode](args) + + +if __name__ == "__main__": + main() diff --git a/tools/zig/gemma4_vision_fixtures.py b/tools/zig/gemma4_vision_fixtures.py new file mode 100644 index 0000000000..0ef505cbd7 --- /dev/null +++ b/tools/zig/gemma4_vision_fixtures.py @@ -0,0 +1,236 @@ +"""Gemma 4 images (the 31B's tower, the 12B's embedder) against transformers' on drawn images (tf-gemma4-vision).""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch +from PIL import Image, ImageDraw +from safetensors import safe_open + + +def draw_shapes(size: tuple[int, int]) -> Image.Image: + w, h = size + im = Image.new("RGB", size, (245, 240, 230)) + d = ImageDraw.Draw(im) + d.ellipse((w * 0.08, h * 0.25, w * 0.42, h * 0.7), fill=(210, 40, 40)) + d.rectangle((w * 0.55, h * 0.25, w * 0.9, h * 0.7), fill=(40, 70, 200)) + for x in range(int(w * 0.1), int(w * 0.9), 6): + d.line((x, h * 0.8, x, h * 0.92), fill=(20, 20, 20), width=2) + d.text((w * 0.4, h * 0.06), "GEMMA 4", fill=(0, 0, 0)) + return im + + +def fixtures(folder: Path) -> list[Path]: + folder.mkdir(parents=True, exist_ok=True) + out = [] + + def save(name: str, im: Image.Image, **kw) -> None: + p = folder / name + im.save(p, **kw) + out.append(p) + + save("shapes.png", draw_shapes((517, 389))) # a mild downscale + big = draw_shapes((1200, 800)) + px = big.load() + for y in range(0, 800, 3): # fine detail for the bicubic filter + for x in range(0, 1200, 7): + px[x, y] = ((x * 3) % 256, (y * 5) % 256, (x + y) % 256) + save("wide.png", big) # 3:2, the full 280-token budget + save("tall.png", draw_shapes((180, 1400))) # one 48-px column wide + save("tiny.png", draw_shapes((20, 15))) # upscaled + save("square.png", draw_shapes((672, 672))) # already a multiple of 48 + return out + + +def tensors(model: Path, prefix: str) -> dict[str, torch.Tensor]: + index = json.loads((model / "model.safetensors.index.json").read_text())["weight_map"] + out = {} + for name, shard in index.items(): + if name.startswith(prefix): + with safe_open(model / shard, "pt") as f: + out[name[len(prefix):]] = f.get_tensor(name) + return out + + +def dequantize(w: torch.Tensor, scales: torch.Tensor, biases: torch.Tensor, group: int, bits: int) -> torch.Tensor: + """MLX affine: `bits`-wide codes packed little-endian through each row's uint32 words; a scale and bias a group.""" + words = (w.to(torch.int64) & 0xFFFFFFFF).numpy().astype(" tuple[int, int]: + """embed_vision.embedding_projection's (bits, group): its own config entry, else the checkpoint's.""" + q = json.loads((model / "config.json").read_text())["quantization"] + entry = q.get("embed_vision.embedding_projection", q) + return entry["bits"], entry["group_size"] + + +def build(model: Path, dtype: torch.dtype): + from transformers.models.gemma4.configuration_gemma4 import Gemma4Config + from transformers.models.gemma4.modeling_gemma4 import Gemma4MultimodalEmbedder, Gemma4VisionModel + + cfg = Gemma4Config.from_pretrained(model) + cfg.vision_config._attn_implementation = "eager" + tower = Gemma4VisionModel(cfg.vision_config) + missing, unexpected = tower.load_state_dict(tensors(model, "vision_tower."), strict=False) + if [m for m in missing if "inv_freq" not in m] or unexpected: + raise SystemExit(f"vision tower tensors: missing {missing}, unexpected {unexpected}") + embed = Gemma4MultimodalEmbedder(cfg.vision_config, cfg.text_config) + q = tensors(model, "embed_vision.embedding_projection.") + bits, group = projection_width(model) + embed.embedding_projection.weight.data = dequantize(q["weight"], q["scales"], q["biases"], group, bits) + return tower.to(dtype).eval(), embed.to(dtype).eval(), cfg + + +def build_unified(model: Path, dtype: torch.dtype): + from transformers.models.gemma4_unified.configuration_gemma4_unified import Gemma4UnifiedConfig + from transformers.models.gemma4_unified.modeling_gemma4_unified import Gemma4UnifiedVisionEmbedder + + cfg = Gemma4UnifiedConfig.from_pretrained(model) + embedder = Gemma4UnifiedVisionEmbedder(cfg.vision_config, cfg.text_config) + state = tensors(model, "vision_embedder.") + q = tensors(model, "embed_vision.embedding_projection.") + projection = q["weight"] + if "scales" in q: # an oQ pack's quantized projection + bits, group = projection_width(model) + projection = dequantize(q["weight"], q["scales"], q["biases"], group, bits) + state["multimodal_embedder.embedding_projection.weight"] = projection + embedder.load_state_dict(state, strict=True) + return embedder.to(dtype).eval() + + +@torch.no_grad() +def unified_stages(embedder, pixels: torch.Tensor, positions: torch.Tensor) -> dict[str, torch.Tensor]: + """Gemma4UnifiedVisionEmbedder.forward's steps: the patch norms and Linear, positions and pos_norm, projection.""" + dtype = next(embedder.parameters()).dtype + x = embedder.patch_ln2(embedder.patch_dense(embedder.patch_ln1(pixels.to(dtype)))) + out = {"patch": x[0]} + axes = torch.arange(2) + x = embedder.pos_norm(x + embedder.pos_embedding[positions.long(), axes].sum(-2)) + out["posed"] = x[0] + out["features"] = embedder(pixels.to(dtype), positions).pooler_output[0] + return out + + +def merged(pil, image: Image.Image, max_tokens: int): + """The PIL processor's 16-px patches (its 8-bit resize, the engine's) through the unified processor's merge.""" + from transformers.models.gemma4_unified.image_processing_gemma4_unified import patches_merge + + done = pil([image], return_tensors="pt", max_soft_tokens=max_tokens) + real = ~(done["image_position_ids"] == -1).all(dim=-1)[0] + teacher, positions = done["pixel_values"][0][real], done["image_position_ids"][0][real] + count = int(done["num_soft_tokens_per_image"][0]) + pixels, merged_positions = patches_merge(teacher.unsqueeze(0), positions.unsqueeze(0), count) + return pixels, merged_positions, (positions.max(dim=0).values + 1).tolist() + + +def unified_reference(model: Path, image_path: Path, out: Path, runs, processor, pil, max_tokens: int) -> dict: + out.mkdir(parents=True, exist_ok=True) + image = Image.open(image_path).convert("RGB") + pixels, positions, grid = merged(pil, image, max_tokens) + np.asarray(pixels[0], dtype=np.float32).tofile(out / "patches.f32") + np.asarray(positions[0], dtype=np.int32).tofile(out / "positions.i32") + for suffix, embedder in runs.items(): + for name, a in unified_stages(embedder, pixels, positions).items(): + np.asarray(a.float(), dtype=np.float32).tofile(out / f"{name}{suffix}.f32") + # the checkpoint's own processor (torchvision's resize) beside the PIL one, as levels of 1/255 + own = processor.image_processor([image], return_tensors="pt", max_soft_tokens=max_tokens)["pixel_values"] + own = own[0][: pixels.shape[1]] + levels = int(((own - pixels[0]).abs().max() * 255).round()) + messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Describe this image."}]}] + text = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + ids = processor(text=[text], images=[image], return_tensors="np", images_kwargs={"max_soft_tokens": max_tokens}) + (out / "prompt.ids.json").write_text(json.dumps(ids["input_ids"][0].tolist()) + "\n") + meta = {"image": str(image_path), "size": list(image.size), "patches": grid[0] * grid[1], "grid_wh": grid, + "tokens": int(pixels.shape[1]), "max_tokens": max_tokens, "prompt": text, "torchvision_levels": levels} + (out / "meta.json").write_text(json.dumps(meta, indent=1) + "\n") + return meta + + +@torch.no_grad() +def stages(tower, embed, pixels: torch.Tensor, positions: torch.Tensor) -> dict[str, torch.Tensor]: + dtype = next(tower.parameters()).dtype + pad = (positions == -1).all(dim=-1) + x = tower.patch_embedder(pixels.to(dtype), positions, pad) + out = {"embed": x[~pad]} + mask = torch.zeros(1, 1, x.shape[1], x.shape[1], dtype=dtype) + mask[..., pad[0]] = torch.finfo(dtype).min + cos_sin = tower.encoder.rotary_emb(x, positions) + for i, layer in enumerate(tower.encoder.layers): + x = layer(x, attention_mask=mask, position_embeddings=cos_sin, position_ids=positions) + if i == 0: + out["block0"] = x[~pad] + out["blocks"] = x[~pad] + pooled = tower(pixel_values=pixels.to(dtype), pixel_position_ids=positions).last_hidden_state + out["pooled"] = pooled + out["features"] = embed(pooled) + return out + + +def reference(model: Path, image_path: Path, out: Path, runs, processor, pil, max_tokens: int) -> dict: + out.mkdir(parents=True, exist_ok=True) + image = Image.open(image_path).convert("RGB") + done = pil([image], return_tensors="pt", max_soft_tokens=max_tokens) + pixels, positions = done["pixel_values"], done["image_position_ids"] + real = ~(positions == -1).all(dim=-1)[0] + np.asarray(pixels[0][real], dtype=np.float32).tofile(out / "patches.f32") + np.asarray(positions[0][real], dtype=np.int32).tofile(out / "positions.i32") + for suffix, (tower, embed) in runs.items(): + for name, a in stages(tower, embed, pixels, positions).items(): + np.asarray(a.float().reshape(-1, a.shape[-1]), dtype=np.float32).tofile(out / f"{name}{suffix}.f32") + messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Describe this image."}]}] + text = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + ids = processor(text=[text], images=[image], return_tensors="np", images_kwargs={"max_soft_tokens": max_tokens}) + (out / "prompt.ids.json").write_text(json.dumps(ids["input_ids"][0].tolist()) + "\n") + if "mm_token_type_ids" in ids: + (out / "prompt.types.json").write_text(json.dumps(ids["mm_token_type_ids"][0].tolist()) + "\n") + grid = positions[0][real].max(dim=0).values + 1 + meta = {"image": str(image_path), "size": list(image.size), "patches": int(real.sum()), "grid_wh": grid.tolist(), + "tokens": int(done["num_soft_tokens_per_image"][0]), "max_tokens": max_tokens, "prompt": text} + (out / "meta.json").write_text(json.dumps(meta, indent=1) + "\n") + return meta + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--model", required=True, help="the 31B (mlx-community/gemma-4-31B-it-4bit) or the converted 12B") + ap.add_argument("--out", required=True, help="where the images and their references go") + ap.add_argument("--max-tokens", type=int, default=280, help="soft tokens an image is sized to: 70, 140 or 280") + args = ap.parse_args() + from transformers import AutoProcessor + + model, out = Path(args.model), Path(args.out) + from transformers.models.gemma4.image_processing_pil_gemma4 import Gemma4ImageProcessorPil + + processor = AutoProcessor.from_pretrained(model) + # the PIL processor, whose 8-bit resize the engine matches (AutoProcessor's torchvision one differs by a level) + config = json.loads((model / "processor_config.json").read_text())["image_processor"] + pil = Gemma4ImageProcessorPil(**{k: v for k, v in config.items() if k != "image_processor_type"}) + if json.loads((model / "config.json").read_text())["model_type"] == "gemma4_unified": + runs = {"": build_unified(model, torch.bfloat16), "_fp32": build_unified(model, torch.float32)} + for path in fixtures(out / "images"): + meta = unified_reference(model, path, out / path.stem, runs, processor, pil, args.max_tokens) + size, grid = meta["size"], meta["grid_wh"] + print(f"{path.name}: {size[0]}x{size[1]} -> grid {grid[0]}x{grid[1]}, {meta['tokens']} tokens, " + f"torchvision's resize {meta['torchvision_levels']} levels from PIL's") + return 0 + tower, embed, _ = build(model, torch.bfloat16) + tower32, embed32, _ = build(model, torch.float32) + runs = {"": (tower, embed), "_fp32": (tower32, embed32)} + for path in fixtures(out / "images"): + meta = reference(model, path, out / path.stem, runs, processor, pil, args.max_tokens) + print(f"{path.name}: {meta['size'][0]}x{meta['size'][1]} -> grid {meta['grid_wh'][0]}x{meta['grid_wh'][1]}, {meta['patches']} patches, {meta['tokens']} tokens") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/kernels/metal/core/affine_mm.metal b/zig/kernels/metal/core/affine_mm.metal index f705be08c9..1591b308bd 100644 --- a/zig/kernels/metal/core/affine_mm.metal +++ b/zig/kernels/metal/core/affine_mm.metal @@ -1,4 +1,4 @@ -// x W^T for MLX affine W on the tensor units: exact codes (4-bit ones plus 128) in the MMAs, fp32 group sums scaled, bias times fp32 row sums. +// x W^T for MLX affine W on the tensor units: a step's weights dequantized to bf16 in the tile, one rounding each. #include using namespace metal; #include "../nax.h" @@ -11,61 +11,17 @@ using namespace tfp; constant constexpr int STEP_BYTES = 8 * TF_BITS; // a weight row's bytes in one 64-deep step constant constexpr int HALF_BYTES = 4 * TF_BITS; // a thread's half of them (32 codes) constant constexpr int NG = 64 / TF_GROUP; // groups a step -constant constexpr int PAD = 64 + 8; // the code tile's row pitch (bf16) -constant constexpr float OFF = TF_BITS == 4 ? 128.0f : 0.0f; // bf16(128 + q) is q's bits ORed with 0x4300; bias - 128 scale undoes it +constant constexpr int PAD = 64 + 8; // the weight tile's row pitch (bf16) struct MmArgs { int rows, n, k, experts; // rows (sorted by expert for a gather), outputs a row, depth, experts (gather) }; -// A dense matmul's pitches, and each batch's start in x, y (elements), the sums and the weights (rows). +// A dense matmul's pitches, and each batch's start in x, y (elements) and the weights (rows). struct MmStrides { - int x_row, y_row, sums_row, x_batch, y_batch, sums_batch, w_batch, pad; + int x_row, y_row, x_batch, y_batch, w_batch, pad; }; -// Each row's fp32 sum over each group of x (row pitch a.n, or a.k when 0), in order: the bias's operand. -[[kernel]] void tf_affine_row_sums(const device bfloat* X [[buffer(0)]], constant MmArgs& a [[buffer(5)]], - device float* XS [[buffer(6)]], uint2 gid [[thread_position_in_grid]]) { - const int g = int(gid.x), r = int(gid.y), groups = a.k / TF_GROUP; - if (g >= groups || r >= a.rows) return; - const device bfloat* x = X + long(r) * (a.n > 0 ? a.n : a.k) + g * TF_GROUP; - float sum = 0.0f; - for (int i = 0; i < TF_GROUP; i++) { - sum += float(x[i]); - } - XS[long(r) * groups + g] = sum; -} - -// (lo | hi) = a * (b0 | b1): a group's first 16-deep slice, its destination overwritten. -inline void am_first(thread frag& lo, thread frag& hi, thread const frag& a, - thread const frag& b0, thread const frag& b1) { -#ifndef TF_SIMD_FRAGS - using namespace mpp::tensor_ops; - constexpr auto shape = matmul2d_descriptor(16, 32, 16, false, true, true, matmul2d_descriptor::mode::multiply); - matmul2d op; - auto left = op.template get_left_input_cooperative_tensor(); - auto right = op.template get_right_input_cooperative_tensor(); - auto acc = op.template get_destination_cooperative_tensor, - metal::remove_addrspace_t, float>(); - TF_UNROLL - for (short e = 0; e < 8; e++) { - left[e] = a[e]; - right[e] = b0[e]; - right[8 + e] = b1[e]; - } - op.run(left, right, acc); - TF_UNROLL - for (short e = 0; e < 8; e++) { - lo[e] = acc[e]; - hi[e] = acc[8 + e]; - } -#else - lo = frag(0); - hi = frag(0); - mma_16x32(lo, hi, a, b0, b1); -#endif -} - // A thread's 32 codes of one weight row in a step (16-byte loads; 8 bytes at 2 bits; words at 3, 5 and 6 bits). #if TF_BITS == 2 struct Codes { @@ -103,67 +59,43 @@ inline Codes am_load(const device uchar* w) { } #endif -// The codes plus OFF, exactly, as bf16 into the tile row (16-byte stores). -inline void am_codes(thread const Codes& c, threadgroup bfloat* out) { - threadgroup uint4* o = (threadgroup uint4*)out; +// Code i (0..31) of a thread's step: MLX's little-endian bit stream at every width. +inline float am_code(thread const Codes& c, short i) { #if TF_BITS == 2 - TF_UNROLL - for (short i = 0; i < 2; i++) { // a word's 16 codes, as themselves (no offset): byte j holds codes 4j..4j+3 - const uint w = c.v[i]; - const float4 f0 = float4(as_type(w & 0x03030303u)), f1 = float4(as_type((w >> 2) & 0x03030303u)); - const float4 f2 = float4(as_type((w >> 4) & 0x03030303u)), f3 = float4(as_type((w >> 6) & 0x03030303u)); - TF_UNROLL - for (short j = 0; j < 4; j++) { - o[2 * i + j / 2][2 * (j % 2)] = as_type(ushort2(as_type(bfloat(f0[j])), as_type(bfloat(f1[j])))); - o[2 * i + j / 2][2 * (j % 2) + 1] = as_type(ushort2(as_type(bfloat(f2[j])), as_type(bfloat(f3[j])))); - } - } + return float((c.v[i / 16] >> (2 * (i % 16))) & 3u); #elif TF_BITS % 4 != 0 - TF_UNROLL - for (short i = 0; i < 8; i++) { // codes 4i .. 4i + 3 of the little-endian bit stream, as themselves - float f[4]; - TF_UNROLL - for (short j = 0; j < 4; j++) { - const short at = TF_BITS * (4 * i + j), word = at / 32, shift = at % 32; - uint v = c.v[word] >> shift; - if (shift + TF_BITS > 32) v |= c.v[word + 1] << (32 - shift); - f[j] = float(v & ((1u << TF_BITS) - 1u)); - } - o[i / 2][2 * (i % 2)] = as_type(ushort2(as_type(bfloat(f[0])), as_type(bfloat(f[1])))); - o[i / 2][2 * (i % 2) + 1] = as_type(ushort2(as_type(bfloat(f[2])), as_type(bfloat(f[3])))); - } + const short at = TF_BITS * i, word = at / 32, shift = at % 32; + uint v = c.v[word] >> shift; + if (shift + TF_BITS > 32) v |= c.v[word + 1] << (32 - shift); + return float(v & ((1u << TF_BITS) - 1u)); #elif TF_BITS == 4 - TF_UNROLL - for (short i = 0; i < 4; i++) { - const uint w = c.v[0][i]; - const uint lo = w & 0x0f0f0f0fu, hi = (w >> 4) & 0x0f0f0f0fu; - o[i] = uint4(0x43004300u | (lo & 0xffu) | ((hi & 0xffu) << 16), - 0x43004300u | ((lo >> 8) & 0xffu) | (((hi >> 8) & 0xffu) << 16), - 0x43004300u | ((lo >> 16) & 0xffu) | (((hi >> 16) & 0xffu) << 16), - 0x43004300u | (lo >> 24) | ((hi >> 24) << 16)); - } + return float((c.v[0][i / 8] >> (4 * (i % 8))) & 15u); #else + return float((c.v[i / 16][(i / 4) % 4] >> (8 * (i % 4))) & 255u); +#endif +} + +// The thread's 32 weights, scale * code + bias in fp32 rounded once to bf16 (as MLX dequantizes), into the tile row. +inline void am_weights(thread const Codes& c, float s, float b, threadgroup bfloat* out) { + threadgroup uint4* o = (threadgroup uint4*)out; TF_UNROLL - for (short i = 0; i < 8; i++) { - const float4 f = float4(as_type(c.v[i / 4][i % 4])); - o[i / 2][2 * (i % 2)] = as_type(ushort2(as_type(bfloat(f.x)), as_type(bfloat(f.y)))); - o[i / 2][2 * (i % 2) + 1] = as_type(ushort2(as_type(bfloat(f.z)), as_type(bfloat(f.w)))); + for (short j = 0; j < 4; j++) { // 16-byte stores: codes 8j .. 8j + 7 + uint4 v; + TF_UNROLL + for (short m = 0; m < 4; m++) { + const bfloat w0 = bfloat(fma(s, am_code(c, 8 * j + 2 * m), b)); + const bfloat w1 = bfloat(fma(s, am_code(c, 8 * j + 2 * m + 1), b)); + v[m] = uint(as_type(w0)) | (uint(as_type(w1)) << 16); + } + o[j] = v; } -#endif } -// A block of the code tile as the transposed right operand: 8-byte loads on M5 (PAD rows, 16-column blocks). -#ifndef TF_SIMD_FRAGS -#define am_tile_frag frag_get_tg -#else -#define am_tile_frag frag_get_t -#endif -// A lane's two rows (home.y and home.y + 8) of each of its TM fragments: their starts in x and in the row sums. +// A lane's two rows (home.y and home.y + 8) of each of its TM fragments: their starts in x. template struct AmRows { const device bfloat* x[TM][2]; - const device float* s[TM][2]; }; // The fragment of x at column c from a lane's two row starts (8-byte loads on M5). @@ -180,45 +112,34 @@ inline frag am_frag(const device bfloat* r0, const device bfloat* r1, in #endif } -// One step's MMAs and group epilogues for the first TL of TM row fragments (TL < TM: a tile's tail, the rest dead). +// The tile's 16 x 16 block at (r, c) as a transposed right operand (on M5 a lane's four columns of a row in one load). +inline frag am_tile_frag(const threadgroup bfloat* tile, int r, int c, short2 home) { +#ifndef TF_SIMD_FRAGS + const threadgroup bfloat* p = tile + (r + home.y) * PAD + c + home.x; + return frag(*(const threadgroup bfloat4*)p, *(const threadgroup bfloat4*)(p + 8 * PAD)); +#else + frag f; + frag_get_t(f, tile, PAD, r, c, home); + return f; +#endif +} + +// One step's MMAs for the first TL of TM row fragments (TL < TM: a tile's tail, the rest dead). template inline void am_step(thread frag (&acc)[TM][2], thread const AmRows& rows, int k, - const threadgroup bfloat* tile, const threadgroup float* sb, int tn, short2 home) { - frag part[TL][2]; + const threadgroup bfloat* tile, int tn, short2 home) { TF_UNROLL for (short q = 0; q < 4; q++) { // the step's four 16-deep slices, every row fragment's chain interleaved frag b0, b1, a[TL]; - am_tile_frag(b0, tile, PAD, tn, 16 * q, home); - am_tile_frag(b1, tile, PAD, tn + 16, 16 * q, home); + b0 = am_tile_frag(tile, tn, 16 * q, home); + b1 = am_tile_frag(tile, tn + 16, 16 * q, home); TF_UNROLL for (short m = 0; m < TL; m++) { a[m] = am_frag(rows.x[m][0], rows.x[m][1], k + 16 * q, home); } TF_UNROLL for (short m = 0; m < TL; m++) { - if (q % (TF_GROUP / 16) == 0) { - am_first(part[m][0], part[m][1], a[m], b0, b1); - } else { - mma_16x32(part[m][0], part[m][1], a[m], b0, b1); - } - } - if ((q + 1) % (TF_GROUP / 16) == 0) { // a group's sums done: its scale, its bias times the row sums - const int g = q / (TF_GROUP / 16); - const threadgroup float* sg = sb + g * 128; - TF_UNROLL - for (short m = 0; m < TL; m++) { - const float xr0 = rows.s[m][0][k / TF_GROUP + g], xr1 = rows.s[m][1][k / TF_GROUP + g]; - TF_UNROLL - for (short j = 0; j < 2; j++) { - TF_UNROLL - for (short c = 0; c < 4; c++) { - const int col = tn + 16 * j + home.x + TF_COL(c); - const float s = sg[col], b = sg[64 + col]; - acc[m][j][c] += s * part[m][j][c] + b * xr0; - acc[m][j][4 + c] += s * part[m][j][4 + c] + b * xr1; - } - } - } + mma_16x32(acc[m][0], acc[m][1], a[m], b0, b1); } } } @@ -227,27 +148,31 @@ inline void am_step(thread frag (&acc)[TM][2], thread const AmRows& r template inline void am_k_loop(thread frag (&acc)[TM][2], thread const AmRows& rows, int K, int live, const device uchar* wq, const device bfloat* scales, const device bfloat* biases, - threadgroup bfloat* tile, threadgroup float* sb, int tn, uint t, short2 home) { + threadgroup bfloat* tile, int tn, uint t, short2 home) { threadgroup bfloat* mine = tile + (t / 2) * PAD + 32 * (t % 2); - const int g_mine = TF_GROUP == 32 ? int(t % 2) : 0; // at g32 each thread of a pair takes one group's scale and bias + const int g_mine = TF_GROUP == 32 ? int(t % 2) : 0; // at g32 each thread of a pair dequantizes its own group TF_UNROLL for (short i = 0; i < TM; i++) { acc[i][0] = frag(0); acc[i][1] = frag(0); } + Codes next = am_load(wq); // a step's codes, scale and bias fetched a step ahead: the loads' latency off the barriers + float s_next = float(scales[g_mine]), b_next = float(biases[g_mine]); for (int k = 0; k < K; k += 64) { - threadgroup_barrier(mem_flags::mem_threadgroup); - am_codes(am_load(wq), mine); - if (TF_GROUP == 32 || t % 2 == 0) { - const float s = float(scales[g_mine]); - sb[g_mine * 128 + t / 2] = s; - sb[g_mine * 128 + 64 + t / 2] = float(biases[g_mine]) - OFF * s; + const Codes codes = next; + const float s_now = s_next, b_now = b_next; + if (k + 64 < K) { + next = am_load(wq + STEP_BYTES); + s_next = float(scales[NG + g_mine]); + b_next = float(biases[NG + g_mine]); } threadgroup_barrier(mem_flags::mem_threadgroup); + am_weights(codes, s_now, b_now, mine); + threadgroup_barrier(mem_flags::mem_threadgroup); if (live > 16 * (TM - 1)) { - am_step(acc, rows, k, tile, sb, tn, home); + am_step(acc, rows, k, tile, tn, home); } else if (TM > 1 && live > 0) { // a tail with only the first fragment live - am_step(acc, rows, k, tile, sb, tn, home); + am_step(acc, rows, k, tile, tn, home); } wq += STEP_BYTES; scales += NG; @@ -311,20 +236,19 @@ inline bool am_tile(const device int32_t* offsets, int experts, int total, int t return false; } -// A simdgroup's row starts: rows [row, row + rows) of x and its row sums, or with `order` sorted rows' token rows (order / topk). +// A simdgroup's row starts: rows [row, row + rows) of x, or with `order` the sorted rows' token rows (order / topk). template -inline AmRows am_rows(const device bfloat* x, int ldx, const device float* xs, int ldxs, int row, int rows, int tm, - short2 home, const device int32_t* order, int topk) { +inline AmRows am_rows(const device bfloat* x, int ldx, int row, int rows, int tm, short2 home, + const device int32_t* order, int topk) { AmRows r; TF_UNROLL for (short m = 0; m < TM; m++) { TF_UNROLL for (short h = 0; h < 2; h++) { const int i = tm + 16 * m + home.y + 8 * h; - const int at = row + (i < rows ? i : 0); // a dead row reads a live one; its sums never reach y + const int at = row + (i < rows ? i : 0); // a dead row reads a live one; its products never reach y const long src = order != nullptr ? long(order[at] / topk) : long(at); r.x[m][h] = x + src * ldx; - r.s[m][h] = xs + src * ldxs; } } return r; @@ -332,10 +256,9 @@ inline AmRows am_rows(const device bfloat* x, int ldx, const device float* x // Rows [row, row + rows) of x times weight rows [wrow0, wrow0 + 64) of a [*, K] matrix into y (ld N) at column col. template -inline void am_block(const device bfloat* x, int ldx, const device float* xs, int ldxs, int row, int rows, int K, - int ldy, int col, long wrow0, const device uint32_t* w, const device bfloat* scales, - const device bfloat* biases, device TF_OUT_T* y, threadgroup bfloat* tile, threadgroup float* sb, - uint sg, uint lane, const device int32_t* dst = nullptr, const device int32_t* order = nullptr, +inline void am_block(const device bfloat* x, int ldx, int row, int rows, int K, int ldy, int col, long wrow0, const device uint32_t* w, const device bfloat* scales, + const device bfloat* biases, device TF_OUT_T* y, threadgroup bfloat* tile, uint sg, uint lane, + const device int32_t* dst = nullptr, const device int32_t* order = nullptr, int topk = 1) { constexpr int SM = BM / 2; constexpr int TM = SM / 16; @@ -343,10 +266,10 @@ inline void am_block(const device bfloat* x, int ldx, const device float* xs, in const int tm = SM * int(sg / 2), tn = 32 * int(sg % 2), live = clamp(rows - tm, 0, SM); const long wrow = wrow0 + t / 2; const short2 home = frag_home(ushort(lane)); - const AmRows starts = am_rows(x, ldx, xs, ldxs, row, rows, tm, home, order, topk); + const AmRows starts = am_rows(x, ldx, row, rows, tm, home, order, topk); frag acc[TM][2]; am_k_loop(acc, starts, K, live, (const device uchar*)w + wrow * (K * TF_BITS / 8) + HALF_BYTES * (t % 2), - scales + wrow * (K / TF_GROUP), biases + wrow * (K / TF_GROUP), tile, sb, tn, uint(t), home); + scales + wrow * (K / TF_GROUP), biases + wrow * (K / TF_GROUP), tile, tn, uint(t), home); if (dst != nullptr) { if (live > 0) am_scatter(acc, y, dst + row + tm, ldy, col + tn, live, home); } else { @@ -357,32 +280,29 @@ inline void am_block(const device bfloat* x, int ldx, const device float* xs, in // y [rows, n] = x [rows, k] W^T, n a multiple of 64: 64x64 tiles, a batch along tg.z (MmStrides). [[kernel]] void tf_affine_mm(const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], - const device float* XS [[buffer(4)]], constant MmArgs& a [[buffer(5)]], - device TF_OUT_T* Y [[buffer(6)]], constant MmStrides& st [[buffer(7)]], + constant MmArgs& a [[buffer(5)]], device TF_OUT_T* Y [[buffer(6)]], + constant MmStrides& st [[buffer(7)]], uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { threadgroup bfloat tile[64 * PAD]; - threadgroup float sb[NG * 128]; const int row = int(tg.y) * 64, col = int(tg.x) * 64, b = int(tg.z); - am_block<64>(X + long(b) * st.x_batch, st.x_row, XS + long(b) * st.sums_batch * (a.k / TF_GROUP), - st.sums_row * (a.k / TF_GROUP), row, min(64, a.rows - row), a.k, st.y_row, col, - long(b) * st.w_batch + col, W, S, B, Y + long(b) * st.y_batch, tile, sb, sg, lane); + am_block<64>(X + long(b) * st.x_batch, st.x_row, row, min(64, a.rows - row), a.k, st.y_row, col, + long(b) * st.w_batch + col, W, S, B, Y + long(b) * st.y_batch, tile, sg, lane); } #define TF_GATHER(BM) \ [[kernel]] void tf_affine_gather_##BM( \ const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], \ const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], \ - const device float* XS [[buffer(4)]], constant MmArgs& a [[buffer(5)]], device TF_OUT_T* Y [[buffer(6)]], \ + constant MmArgs& a [[buffer(5)]], device TF_OUT_T* Y [[buffer(6)]], \ const device int32_t* OFFS [[buffer(7)]], uint sg [[simdgroup_index_in_threadgroup]], \ uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { \ threadgroup bfloat tile[64 * PAD]; \ - threadgroup float sb[NG * 128]; \ int expert, row, rows; \ if (!am_tile(OFFS, a.experts, a.rows, int(tg.y), lane, expert, row, rows)) return; \ const int col = int(tg.x) * 64; \ - am_block(X, a.k, XS, a.k / TF_GROUP, row, rows, a.k, a.n, col, long(expert) * a.n + col, W, S, B, Y, \ - tile, sb, sg, lane); \ + am_block(X, a.k, row, rows, a.k, a.n, col, long(expert) * a.n + col, W, S, B, Y, \ + tile, sg, lane); \ } TF_GATHER(32) TF_GATHER(64) @@ -392,17 +312,16 @@ TF_GATHER(64) [[kernel]] void tf_affine_gather_scatter_##BM( \ const device bfloat* X [[buffer(0)]], const device uint32_t* W [[buffer(1)]], \ const device bfloat* S [[buffer(2)]], const device bfloat* B [[buffer(3)]], \ - const device float* XS [[buffer(4)]], constant MmArgs& a [[buffer(5)]], device TF_OUT_T* Y [[buffer(6)]], \ + constant MmArgs& a [[buffer(5)]], device TF_OUT_T* Y [[buffer(6)]], \ const device int32_t* OFFS [[buffer(7)]], const device int32_t* DST [[buffer(8)]], \ uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], \ uint3 tg [[threadgroup_position_in_grid]]) { \ threadgroup bfloat tile[64 * PAD]; \ - threadgroup float sb[NG * 128]; \ int expert, row, rows; \ if (!am_tile(OFFS, a.experts, a.rows, int(tg.y), lane, expert, row, rows)) return; \ const int col = int(tg.x) * 64; \ - am_block(X, a.k, XS, a.k / TF_GROUP, row, rows, a.k, a.n, col, long(expert) * a.n + col, W, S, B, Y, \ - tile, sb, sg, lane, DST); \ + am_block(X, a.k, row, rows, a.k, a.n, col, long(expert) * a.n + col, W, S, B, Y, \ + tile, sg, lane, DST); \ } TF_GATHER_SCATTER(32) TF_GATHER_SCATTER(64) @@ -451,7 +370,7 @@ inline void am_store_glu(thread const frag (&acc)[TM][2], device bfloat* [[kernel]] void tf_affine_gather_glu_##BM( \ const device bfloat* X [[buffer(0)]], const device uint32_t* WG [[buffer(1)]], \ const device bfloat* SG [[buffer(2)]], const device bfloat* BG [[buffer(3)]], \ - const device float* XS [[buffer(4)]], constant MmArgs& a [[buffer(5)]], device bfloat* Y [[buffer(6)]], \ + constant MmArgs& a [[buffer(5)]], device bfloat* Y [[buffer(6)]], \ const device int32_t* OFFS [[buffer(7)]], const device int32_t* ORDER [[buffer(8)]], \ const device uint32_t* WU [[buffer(9)]], const device bfloat* SU [[buffer(10)]], \ const device bfloat* BU [[buffer(11)]], constant float2& glu [[buffer(12)]], \ @@ -460,7 +379,6 @@ inline void am_store_glu(thread const frag (&acc)[TM][2], device bfloat* constexpr int SM = BM / 2; \ constexpr int TM = SM / 16; \ threadgroup bfloat tile[64 * PAD]; \ - threadgroup float sb[NG * 128]; \ int expert, row, rows; \ if (!am_tile(OFFS, a.experts, a.rows, int(tg.y), lane, expert, row, rows)) return; \ const int c0 = int(tg.x) * 32, t = int(sg) * 32 + int(lane), tr = t / 2; \ @@ -468,12 +386,12 @@ inline void am_store_glu(thread const frag (&acc)[TM][2], device bfloat* const long wrow = long(expert) * a.n + c0 + (tr >> 5) * 16 + (tr & 15); \ const int tm = SM * int(sg / 2), tn = 32 * int(sg % 2), live = clamp(rows - tm, 0, SM); \ const short2 home = frag_home(ushort(lane)); \ - const AmRows starts = am_rows(X, a.k, XS, a.k / TF_GROUP, row, rows, tm, home, ORDER, int(glu.x)); \ + const AmRows starts = am_rows(X, a.k, row, rows, tm, home, ORDER, int(glu.x)); \ frag acc[TM][2]; \ am_k_loop(acc, starts, a.k, live, \ (const device uchar*)(up ? WU : WG) + wrow * (a.k * TF_BITS / 8) + HALF_BYTES * (t % 2), \ - (up ? SU : SG) + wrow * (a.k / TF_GROUP), (up ? BU : BG) + wrow * (a.k / TF_GROUP), tile, sb, tn, \ - uint(t), home); \ + (up ? SU : SG) + wrow * (a.k / TF_GROUP), (up ? BU : BG) + wrow * (a.k / TF_GROUP), tile, tn, uint(t), \ + home); \ am_store_glu(acc, Y + long(row + tm) * a.n + c0 + tn / 2, a.n, live, home, glu.y); \ } TF_GATHER_GLU(32) diff --git a/zig/kernels/metal/core/lane_projection.metal b/zig/kernels/metal/core/lane_projection.metal index 81cfd5d9d4..301d604c57 100644 --- a/zig/kernels/metal/core/lane_projection.metal +++ b/zig/kernels/metal/core/lane_projection.metal @@ -276,36 +276,87 @@ uint group_id [[simdgroup_index_in_threadgroup]],uint lane [[thread_index_in_sim #if TF_BITS == 4 && defined(TF_REG) struct tf_reg_load { uint4 w0, w1, q0, q1; float x0, x1; }; // a lane's codes, its two metadata quads and its two row sums for one group -inline tf_reg_load tf_reg_fetch(const device uint4* wl, const device uint4* sbv, const device float* XS, int g, int n, int m, int stride) { - return {wl[(size_t)g * 64], wl[(size_t)g * 64 + 32], sbv[((size_t)g * TF_N + n) / 4], sbv[((size_t)g * TF_N + n + 16) / 4], XS[g * stride + m], XS[g * stride + m + 8]}; +inline tf_reg_load tf_reg_fetch_w(const device uint4* wl, const device uint4* sbv, int g, int n) { + return {wl[(size_t)g * 64], wl[(size_t)g * 64 + 32], sbv[((size_t)g * TF_N + n) / 4], sbv[((size_t)g * TF_N + n + 16) / 4], 0.0f, 0.0f}; } +template +inline void tf_reg_sums(thread tf_reg_load& l, XSP XS, int g, int m, int stride) { + l.x0 = XS[g * stride + m]; + l.x1 = XS[g * stride + m + 8]; +} + +template +inline tf_reg_load tf_reg_fetch(const device uint4* wl, const device uint4* sbv, XSP XS, int g, int n, int m, int stride) { + tf_reg_load l = tf_reg_fetch_w(wl, sbv, g, n); + tf_reg_sums(l, XS, g, m, stride); + return l; +} + +#ifndef TF_TPT +#define TF_TPT 1 +#endif +// TF_FUSED (Fused in lane_projection.zig): a family's prologue makes the row and its sums under the first weight fetch. + // tf_lane's arithmetic with codes unpacked straight into tensor-op registers: every word equals tf_lane_pair's. [[kernel]] void tf_lane_reg(const device bfloat* X [[buffer(0)]], const device uint4* W [[buffer(1)]], const device bfloat* SB [[buffer(2)]], const device int* mdims [[buffer(3)]], device TF_OUT* Y [[buffer(4)]], const device float* XS [[buffer(5)]], +#ifdef TF_FUSED + const device bfloat* A0 [[buffer(6)]], const device bfloat* A1 [[buffer(7)]], const device bfloat* A2 [[buffer(8)]], + const device bfloat* A3 [[buffer(9)]], const device bfloat* A4 [[buffer(10)]], const device bfloat* A5 [[buffer(11)]], + device bfloat* O0 [[buffer(12)]], device bfloat* O1 [[buffer(13)]], constant float4& F [[buffer(14)]], constant int4& I [[buffer(15)]], +#endif uint sgi [[simdgroup_index_in_threadgroup]], uint lanei [[thread_index_in_simdgroup]], uint2 tg [[threadgroup_position_in_grid]]) { const int r0 = int(tg.y) * 16, rows = mdims[0] - r0; if (rows <= 0) return; X += (size_t)r0 * TF_K; Y += (size_t)r0 * TF_YS; XS += r0; - const int lane = int(lanei), ks = int(sgi), M = min(rows, 16), stride = mdims[1], tile = tf_tile(int(tg.x)), n0 = tile * NT; +#ifdef TF_FUSED + const int lane = int(lanei), ks = int(sgi) % TF_SK, M = min(rows, 16), stride = mdims[1]; + const int tile = tf_tile(int(tg.x) * TF_TPT + int(sgi) / TF_SK), n0 = tile * NT, tp = int(sgi) / TF_SK; +#else + const int lane = int(lanei), ks = int(sgi), M = min(rows, 16), stride = mdims[1], tile = tf_tile(int(tg.x)), n0 = tile * NT, tp = 0; +#endif const int nb = 4 * (lane & 1) + 8 * ((lane >> 3) & 1), mb = ((lane >> 1) & 3) + 4 * ((lane >> 4) & 1); const int first = TF_G0 + (ks * TF_GN) / TF_SK, last = TF_G0 + ((ks + 1) * TF_GN) / TF_SK; constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); matmul2d op; + const device uint4* wl = W + (size_t)tile * KG * 64 + lane; + const device uint4* sbv = (const device uint4*)SB; + tf_reg_load next = {}; +#ifdef TF_FUSED + if (first < last) next = tf_reg_fetch_w(wl, sbv, first, n0 + nb); // the first weights fly while the prologue runs + const tf_fused fa = {A0, A1, A2, A3, A4, A5, O0, O1, F, I}; + threadgroup bfloat xs[TF_XROWS > 0 ? TF_XROWS * TF_K : 1]; + threadgroup float xsum[KG + 16]; // one row's sums; the 16 past the end feed the discarded rows' lanes + threadgroup float pscratch[16]; // the prologue's reduction partials + const uint tid = sgi * 32 + lanei; + if (tid < 16) xsum[KG + tid] = 0.0f; + tf_prologue(fa, X, xs, xsum, pscratch, tg.x == 0, tid); // args, X, row, sums, scratch, first group, thread + threadgroup_barrier(mem_flags::mem_threadgroup); + const threadgroup float* xsp = xsum; +#if TF_XROWS > 0 + tensor, tensor_inline> a(xs, dextents(TF_K, M)); +#else tensor, tensor_inline> a((device bfloat*)X, dextents(TF_K, M)); +#endif +#else + const device float* xsp = XS; + tensor, tensor_inline> a((device bfloat*)X, dextents(TF_K, M)); +#endif auto a0 = a.slice(0, 0); auto cb = op.template get_right_input_cooperative_tensor(); auto P = op.template get_destination_cooperative_tensor, metal::remove_addrspace_t, float>(); thread half2* b2 = (thread half2*)&cb[0]; - const device uint4* wl = W + (size_t)tile * KG * 64 + lane; - const device uint4* sbv = (const device uint4*)SB; const half2 bias = half2(1024.0h); float C[16]; for (int i = 0; i < 16; i++) C[i] = 0.0f; - tf_reg_load next = {}; - if (first < last) next = tf_reg_fetch(wl, sbv, XS, first, n0 + nb, mb, stride); +#ifdef TF_FUSED + if (first < last) tf_reg_sums(next, xsp, first, mb, stride); +#else + if (first < last) next = tf_reg_fetch(wl, sbv, xsp, first, n0 + nb, mb, stride); // main's order, unfused +#endif for (int g = first; g < last; g++) { const tf_reg_load cur = next; for (int w = 0; w < 8; w++) { @@ -317,7 +368,7 @@ inline tf_reg_load tf_reg_fetch(const device uint4* wl, const device uint4* sbv, b2[e + 8] = as_type(((lo >> 8) & 0x00FF00FFu) | 0x64006400u) - bias; b2[e + 9] = as_type(((hi >> 8) & 0x00FF00FFu) | 0x64006400u) - bias; } - if (g + 1 < last) next = tf_reg_fetch(wl, sbv, XS, g + 1, n0 + nb, mb, stride); + if (g + 1 < last) next = tf_reg_fetch(wl, sbv, xsp, g + 1, n0 + nb, mb, stride); auto x = a.slice(g * GS, 0); op.run(x, cb, P); for (int j = 0; j < 4; j++) { @@ -329,16 +380,20 @@ inline tf_reg_load tf_reg_fetch(const device uint4* wl, const device uint4* sbv, C[12 + j] = fma(s1, P[12 + j], fma(b1, cur.x1, C[12 + j])); } } - threadgroup float part[(TF_SK > 1 ? TF_SK - 1 : 1) * 16 * 32]; + threadgroup float part[TF_TPT * (TF_SK > 1 ? TF_SK - 1 : 1) * 16 * 32]; if (TF_SK > 1) { - if (ks > 0) for (int i = 0; i < 16; i++) part[((ks - 1) * 16 + i) * 32 + lane] = C[i]; + if (ks > 0) for (int i = 0; i < 16; i++) part[((tp * (TF_SK - 1) + ks - 1) * 16 + i) * 32 + lane] = C[i]; threadgroup_barrier(mem_flags::mem_threadgroup); - if (ks == 0) for (int s2 = 1; s2 < TF_SK; s2++) for (int i = 0; i < 16; i++) C[i] += part[((s2 - 1) * 16 + i) * 32 + lane]; + if (ks == 0) for (int s2 = 1; s2 < TF_SK; s2++) for (int i = 0; i < 16; i++) C[i] += part[((tp * (TF_SK - 1) + s2 - 1) * 16 + i) * 32 + lane]; } - if (ks == 0) - for (int i = 0; i < 16; i++) { - const int m = mb + 8 * ((i >> 2) & 1), n = n0 + nb + 16 * (i >> 3) + (i & 3); - if (m < M && n < TF_N) Y[(size_t)m * TF_YS + n] = static_cast(C[i]); - } + if (ks != 0) return; +#ifdef TF_EPILOGUE + tf_epilogue(fa, C, M, mb, n0 + nb, tile, Y); // the family's store of the lane's 16 accumulators +#else + for (int i = 0; i < 16; i++) { + const int m = mb + 8 * ((i >> 2) & 1), n = n0 + nb + 16 * (i >> 3) + (i & 3); + if (m < M && n < TF_N) Y[(size_t)m * TF_YS + n] = static_cast(C[i]); + } +#endif } #endif diff --git a/zig/kernels/metal/gemma4/fused.metal b/zig/kernels/metal/gemma4/fused.metal new file mode 100644 index 0000000000..5074fad736 --- /dev/null +++ b/zig/kernels/metal/gemma4/fused.metal @@ -0,0 +1,125 @@ +// Gemma 4's decode glue inside the core reg lane kernel (TF_FUSED), one row, f.x = eps: glue.metal's arithmetic. +#include +using namespace metal; +constant constexpr int FT = 256; // the tails' threads, the glue kernels' T (their reductions' order) +constant constexpr int FKG = TF_K / TF_GROUP; + +inline float gf_reduce(float acc, threadgroup float* p, uint lane, uint sg) { + const float s = simd_sum(acc); + if (lane == 0) p[sg] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + float out = 0.0f; + for (int q = 0; q < FT / 32; q++) out += p[q]; + return out; +} + +// gemma_row_sums over the row in threadgroup memory: each group's values added in order from zero. +inline void gf_sums(const threadgroup bfloat* row, threadgroup float* xsum, uint tid) { + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int g = int(tid); g < FKG; g += FT) { + float acc = 0.0f; + for (int i = 0; i < TF_GROUP; i++) acc += float(row[g * TF_GROUP + i]); + xsum[g] = acc; + } +} + +#if TF_FUSE == 1 +static_assert(TF_THREADS == FT, "the dense tail reduces over 256 threads"); +// TF_FUSE 1, gemma_dense_tail (the qkv input): a0 H, a1 Y, a2 post-MLP norm, a3 scalar, a4 next norm; o0 HN, o1 tap +inline void tf_prologue(tf_fused fa, const device bfloat* X, threadgroup bfloat* xs, threadgroup float* xsum, threadgroup float* p, bool first, uint tid) { + constexpr int PER = TF_K / FT; + threadgroup float* p1 = p; + threadgroup float* p2 = p + FT / 32; + const uint lane = tid & 31, sg = tid >> 5; + float yv[PER], hin[PER], wp[PER], wn[PER]; + for (int i = 0; i < PER; i++) { + const int c = int(tid) + i * FT; + yv[i] = float(fa.a1[c]); hin[i] = float(fa.a0[c]); + wp[i] = float(fa.a2[c]); wn[i] = float(fa.a4[c]); + } + const float sc = float(fa.a3[0]); + float ss1 = 0.0f; + for (int i = 0; i < PER; i++) ss1 = fma(yv[i], yv[i], ss1); + const float inv1 = metal::precise::rsqrt(gf_reduce(ss1, p1, lane, sg) / float(TF_K) + fa.f.x); + float hv[PER]; + float ss2 = 0.0f; + for (int i = 0; i < PER; i++) { + const int c = int(tid) + i * FT; + const float b = float(bfloat(wp[i] * float(bfloat(yv[i] * inv1)))); + const bfloat hs = bfloat(hin[i] + b); + const bfloat hn = bfloat(float(hs) * sc); + if (first) fa.o0[c] = hn; + if (first && fa.i.x) fa.o1[c] = hn; + hv[i] = float(hn); + ss2 = fma(hv[i], hv[i], ss2); + } + const float inv2 = metal::precise::rsqrt(gf_reduce(ss2, p2, lane, sg) / float(TF_K) + fa.f.x); + for (int i = 0; i < PER; i++) { + const int c = int(tid) + i * FT; + xs[c] = bfloat(wn[i] * float(bfloat(hv[i] * inv2))); + } + gf_sums(xs, xsum, tid); +} +#endif + +#if TF_FUSE == 2 +static_assert(TF_THREADS == FT, "the attention tail reduces over 256 threads"); +// TF_FUSE 2, gemma_attn_tail (the gate/up input): a0 H, a1 O, a2 post-attention norm, a3 pre-MLP norm; o0 HN. +inline void tf_prologue(tf_fused fa, const device bfloat* X, threadgroup bfloat* xs, threadgroup float* xsum, threadgroup float* p, bool first, uint tid) { + constexpr int PER = TF_K / FT; + threadgroup float* p1 = p; + threadgroup float* p2 = p + FT / 32; + const uint lane = tid & 31, sg = tid >> 5; + float ov[PER], hin[PER], wa[PER], w1[PER]; + for (int i = 0; i < PER; i++) { + const int c = int(tid) + i * FT; + ov[i] = float(fa.a1[c]); hin[i] = float(fa.a0[c]); + wa[i] = float(fa.a2[c]); w1[i] = float(fa.a3[c]); + } + float ss = 0.0f; + for (int i = 0; i < PER; i++) ss = fma(ov[i], ov[i], ss); + const float inv1 = metal::precise::rsqrt(gf_reduce(ss, p1, lane, sg) / float(TF_K) + fa.f.x); + float hv[PER]; + float ss2 = 0.0f; + for (int i = 0; i < PER; i++) { + const int c = int(tid) + i * FT; + const float a = float(bfloat(wa[i] * float(bfloat(ov[i] * inv1)))); + const bfloat hn = bfloat(hin[i] + a); + if (first) fa.o0[c] = hn; + hv[i] = float(hn); + ss2 = fma(hv[i], hv[i], ss2); + } + const float inv2 = metal::precise::rsqrt(gf_reduce(ss2, p2, lane, sg) / float(TF_K) + fa.f.x); + for (int i = 0; i < PER; i++) { + const int c = int(tid) + i * FT; + xs[c] = bfloat(w1[i] * float(bfloat(hv[i] * inv2))); + } + gf_sums(xs, xsum, tid); +} + +// gemma_geglu over the gate/up stack's 16-row interleaved blocks: C[i] a gate column, C[i + 8] its up column; Y = act. +inline void tf_epilogue(tf_fused fa, thread const float* C, int M, int mb, int n, int tile, device bfloat* Y) { + for (int i = 0; i < 8; i++) { + const int m = mb + 8 * ((i >> 2) & 1), c = tile * 16 + (n - tile * 32) + (i & 3); + if (m >= M) continue; + const bfloat x = bfloat(C[i]), up = bfloat(C[i + 8]); + const bfloat cube = bfloat(metal::pow(float(x), 3.0f)); + const bfloat inner = bfloat(float(x) + float(bfloat(float(bfloat(0.044715f)) * float(cube)))); + const bfloat th = bfloat(metal::precise::tanh(float(bfloat(float(bfloat(0.7978845608028654f)) * float(inner))))); + const bfloat half_x = bfloat(0.5f * float(x)); + const bfloat gl = bfloat(float(half_x) * float(bfloat(1.0f + float(th)))); + Y[(size_t)m * (TF_N / 2) + c] = bfloat(float(gl) * float(up)); + } +} +#endif + +#if TF_FUSE == 3 +// TF_FUSE 3 (the down input): gemma_geglu's group sums over the act row X, each group's values added in order. +inline void tf_prologue(tf_fused fa, const device bfloat* X, threadgroup bfloat* xs, threadgroup float* xsum, threadgroup float* p, bool first, uint tid) { + for (int g = int(tid); g < FKG; g += TF_THREADS) { + float acc = 0.0f; + for (int i = 0; i < TF_GROUP; i++) acc += float(X[g * TF_GROUP + i]); + xsum[g] = acc; + } +} +#endif diff --git a/zig/kernels/metal/gemma4/glue.metal b/zig/kernels/metal/gemma4/glue.metal new file mode 100644 index 0000000000..ea914874e5 --- /dev/null +++ b/zig/kernels/metal/gemma4/glue.metal @@ -0,0 +1,492 @@ +// Gemma 4's decode glue around the core lane projections: gilby's RowDecode kernels (#360) and MLX's ops it calls. +#include +using namespace metal; + +// Each pipeline's constants: D, T, DH, NQ, NK, VK, W, HK, G, CK, S, BLK, SCALE, GELU_C0, GELU_C1, EB (kernels.zig). + +// A row's group sums for the reg lane kernel in tf_lane_sums' order: XS[g * stride + r], none when stride is 0. +inline void gemma_row_sums(const device bfloat* row, int len, device float* XS, int stride, int r, uint t, uint nt) { + if (stride == 0) return; + for (int g = int(t); g < len / 64; g += int(nt)) { + float acc = 0.0f; + for (int i = 0; i < 64; i++) acc += float(row[g * 64 + i]); + XS[g * stride + r] = acc; + } +} + +// The table's codes 2b and 2b + 1 of row `id`, EB bits each in MLX's little-endian packing. +inline uint2 gemma_codes(const device uchar* weight, size_t id, uint b) { +#if EB == 4 + const uchar v = weight[id * (D / 2) + b]; + return uint2(v & 0xf, v >> 4); +#else + const device uchar* row = weight + id * (D * EB / 8); + uint2 out; + for (uint j = 0; j < 2; j++) { + const uint bit = (2 * b + j) * EB, at = bit / 8, sh = bit % 8; + uint v = uint(row[at]) >> sh; + if (sh + EB > 8) v |= uint(row[at + 1]) << (8 - sh); + out[j] = v & ((1u << EB) - 1u); + } + return out; +#endif +} + +// Embedding rows of the EB-bit table times the bf16 embedding scale: MLX's dequantize, then a bf16 multiply. +kernel void gemma_embed(const device uint* ids [[buffer(0)]], const device uchar* weight [[buffer(1)]], + const device bfloat* scales [[buffer(2)]], const device bfloat* biases [[buffer(3)]], + device bfloat* out [[buffer(4)]], constant float& scale [[buffer(5)]], + uint2 pos [[thread_position_in_grid]]) { + const uint b = pos.x, r = pos.y; + const size_t id = ids[r]; + const uint2 v = gemma_codes(weight, id, b); + const size_t g = id * (D / 64) + (2 * b) / 64; + const float s = float(scales[g]), z = float(biases[g]); + out[size_t(r) * D + 2 * b] = bfloat(float(bfloat(s * float(v.x) + z)) * scale); + out[size_t(r) * D + 2 * b + 1] = bfloat(float(bfloat(s * float(v.y) + z)) * scale); +} + +// MLX 0.32's rms_looped (rms_norm.metal, MIT) for bf16, stride-1 w and D wide, its two loops as one; row sums after. +kernel void gemma_rms_looped(const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], + device bfloat* out [[buffer(2)]], constant float& eps [[buffer(3)]], + device float* XS [[buffer(4)]], constant int& xst [[buffer(5)]], + uint gid [[threadgroup_position_in_grid]], uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float inv[1]; + threadgroup float sums[32]; + const device bfloat* xr = x + gid * size_t(D) + lid * 4; + const device bfloat* wr = w + lid * 4; + float acc = 0; + for (uint r = 0; r < D; r += lsize * 4) { + for (int i = 0; i < 4; i++) { + if (r + lid * 4 + i < D) { + const float xi = xr[i + r]; + acc += xi * xi; + } + } + } + acc = simd_sum(acc); + if (sg == 0) sums[lane] = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (lane == 0) sums[sg] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + acc = simd_sum(sums[lane]); + if (lane == 0) inv[0] = metal::precise::rsqrt(acc / D + eps); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + device bfloat* o = out + gid * size_t(D) + lid * 4; + for (uint r = 0; r < D; r += lsize * 4) + for (int i = 0; i < 4; i++) + if (r + lid * 4 + i < D) o[r + i] = bfloat(float(wr[i + r]) * float(bfloat(float(xr[r + i]) * inv[0]))); + threadgroup_barrier(mem_flags::mem_device); + gemma_row_sums(out + gid * size_t(D), D, XS, xst, int(gid), lid, lsize); +} + +// gilby's _EMIT's arithmetic: a head's raw values xv (lane + 32 i) normed (v unweighted), q and k roped at pos, in y. +inline void gemma_norm_rope(int kind, uint lane, thread const float* xv, const device bfloat* QW, const device bfloat* KW, + const device float* INVF, float pos, float eps, thread float* y) { + constexpr int PER = DH / 32; + float ss = 0.0f; + for (int i = 0; i < PER; i++) ss = fma(xv[i], xv[i], ss); + ss = simd_sum(ss); + const float inv = metal::precise::rsqrt(ss / float(DH) + eps); + if (kind == 2) { + for (int i = 0; i < PER; i++) y[i] = float(bfloat(xv[i] * inv)); + return; + } + const device bfloat* w = kind == 0 ? QW : KW; + for (int i = 0; i < PER; i++) y[i] = float(bfloat(float(w[int(lane) + 32 * i]) * float(bfloat(xv[i] * inv)))); + for (int i = 0; i < PER / 2; i++) { + const float theta = pos * INVF[int(lane) + 32 * i]; + const float c = metal::fast::cos(theta), s = metal::fast::sin(theta); + const float x1 = y[i], x2 = y[i + PER / 2]; + y[i] = x1 * c - x2 * s; + y[i + PER / 2] = x1 * s + x2 * c; + } +} + +// The stacked q|k|v row's head `slot` as raw values in a simdgroup's lanes; VK: v is k's raw values. +inline void gemma_head_raw(const device bfloat* QKV, int r, int slot, uint lane, thread float* xv) { + constexpr int PER = DH / 32; + const int kind = slot < NQ ? 0 : (slot < NQ + NK ? 1 : 2); + const int head = kind == 0 ? slot : (kind == 1 ? slot - NQ : slot - NQ - NK); + const int src = kind == 2 ? (VK ? NQ * DH : (NQ + NK) * DH) + head * DH : slot * DH; + for (int i = 0; i < PER; i++) xv[i] = float(QKV[r * W + src + int(lane) + 32 * i]); +} + +// gilby's _EMIT: the head's values to Q, K or V, and with cache slots (cm: ring or 0, slots) k and v to the row's slot. +inline void gemma_emit(int kind, int head, int r, int R, uint lane, thread float* xv, const device bfloat* QW, + const device bfloat* KW, const device float* INVF, const device int* POS, float eps, + device bfloat* Q, device bfloat* K, device bfloat* V, device bfloat* KC, device bfloat* VC, + int3 cm) { + const size_t at = cm.y ? (size_t(head) * cm.y + (cm.x ? POS[r] % cm.x : POS[r])) * DH : 0; + constexpr int PER = DH / 32; + float y[PER]; + gemma_norm_rope(kind, lane, xv, QW, KW, INVF, float(POS[r]), eps, y); + if (kind == 0) { + for (int i = 0; i < PER; i++) Q[(r * NQ + head) * DH + int(lane) + 32 * i] = bfloat(y[i]); + return; + } + device bfloat* out = (kind == 1 ? K : V) + (head * R + r) * DH; + device bfloat* cache = kind == 1 ? KC : VC; + for (int i = 0; i < PER; i++) { + out[int(lane) + 32 * i] = bfloat(y[i]); + if (cm.y) cache[at + int(lane) + 32 * i] = bfloat(y[i]); + } +} + +// gilby's qkv_prep: a simdgroup per (head slot, row) of the stacked q|k|v rows; VK: v is k's raw values. +kernel void gemma_qkv_prep(const device bfloat* QKV [[buffer(0)]], const device bfloat* QW [[buffer(1)]], + const device bfloat* KW [[buffer(2)]], const device float* INVF [[buffer(3)]], + const device int* POS [[buffer(4)]], constant float& eps [[buffer(5)]], + device bfloat* Q [[buffer(6)]], device bfloat* K [[buffer(7)]], device bfloat* V [[buffer(8)]], + device bfloat* KC [[buffer(9)]], device bfloat* VC [[buffer(10)]], constant int3& cm [[buffer(11)]], + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]], + uint3 tgs [[threadgroups_per_grid]]) { + const int slot = int(tg.x), r = int(tg.y), R = int(tgs.y); + constexpr int PER = DH / 32; + const int kind = slot < NQ ? 0 : (slot < NQ + NK ? 1 : 2); + const int head = kind == 0 ? slot : (kind == 1 ? slot - NQ : slot - NQ - NK); + float xv[PER]; + gemma_head_raw(QKV, r, slot, lane, xv); + gemma_emit(kind, head, r, R, lane, xv, QW, KW, INVF, POS, eps, Q, K, V, KC, VC, cm); +} + +// gilby's attend partials, sums fixed by key position; META: first chunk, chunks, ring, rows, first new key. +kernel void gemma_attention_partial(const device bfloat* Q [[buffer(0)]], const device bfloat* K [[buffer(1)]], + const device bfloat* V [[buffer(2)]], const device bfloat* KN [[buffer(3)]], + const device bfloat* VN [[buffer(4)]], const device int* POS [[buffer(5)]], + const device int* LO [[buffer(6)]], const device int* META [[buffer(7)]], + constant int& CAP [[buffer(8)]], device float* PM [[buffer(9)]], + device float* PL [[buffer(10)]], device float* PO [[buffer(11)]], + const device bfloat* QKV [[buffer(12)]], const device bfloat* QW [[buffer(13)]], + const device bfloat* KW [[buffer(14)]], const device float* INVF [[buffer(15)]], + constant float& eps [[buffer(16)]], constant int2& FUSE [[buffer(17)]], + device bfloat* OUT [[buffer(18)]], device float* XS [[buffer(19)]], + constant int& xst [[buffer(20)]], device bfloat* KC [[buffer(21)]], + device bfloat* VC [[buffer(22)]], + uint lane [[thread_index_in_simdgroup]], uint sgu [[simdgroup_index_in_threadgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int sgi = int(sgu); + const int g = sgi / S, s = sgi % S; + const int C0 = META[0], NCH = META[1], RING = META[2], R = META[3]; + const int c = C0 + int(tg.x); + const int h = int(tg.y); + const int r = int(tg.z); + const int qh = h * G + g; + const int P0 = META[4]; + constexpr int DPL = DH / 32; + const int k0 = metal::max(c * CK, LO[r]), k1 = metal::min((c + 1) * CK, POS[r] + 1); + float q[DPL], o[DPL]; + threadgroup bfloat qs[G][DH]; + // FUSE.x (one row, every key cached): gemma_qkv_prep's work here; the last chunk's group writes k and v to the cache + if (FUSE.x) { // simdgroups 0..G-1 prep the query heads; G and G + 1 the key and value, into the cache at the row's slot + float xv[DPL], y[DPL]; + if (sgi < G) { + gemma_head_raw(QKV, r, h * G + sgi, lane, xv); + gemma_norm_rope(0, lane, xv, QW, KW, INVF, float(POS[r]), eps, y); + for (int i = 0; i < DPL; i++) qs[sgi][int(lane) + 32 * i] = bfloat(y[i]); + } else if (sgi < G + 2 && int(tg.x) + 1 == NCH) { + const int kind = sgi - G + 1; + gemma_head_raw(QKV, r, NQ + (kind - 1) * NK + h, lane, xv); + gemma_norm_rope(kind, lane, xv, QW, KW, INVF, float(POS[r]), eps, y); + device bfloat* cache = (kind == 1 ? KC : VC) + (size_t(h) * CAP + (RING ? POS[r] % RING : POS[r])) * DH; + for (int i = 0; i < DPL; i++) cache[int(lane) + 32 * i] = bfloat(y[i]); + } + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + for (int i = 0; i < DPL; i++) { q[i] = float(qs[g][lane * DPL + i]) * SCALE; o[i] = 0.0f; } + } else { + const device bfloat* qp = Q + (size_t(r) * (HK * G) + qh) * DH + lane * DPL; + for (int i = 0; i < DPL; i++) { q[i] = float(qp[i]) * SCALE; o[i] = 0.0f; } + } + float m = -INFINITY, l = 0.0f; + const device bfloat* kb = K + size_t(h) * CAP * DH + lane * DPL; + const device bfloat* vb = V + size_t(h) * CAP * DH + lane * DPL; + const device bfloat* kn = KN + size_t(h) * R * DH + lane * DPL; + const device bfloat* vn = VN + size_t(h) * R * DH + lane * DPL; + for (int base = k0 + s; base < k1; base += S * BLK) { + float sc[BLK]; + size_t at[BLK]; + bool fresh[BLK], live[BLK]; + float bm = -INFINITY; + for (int j = 0; j < BLK; j++) { + const int p = base + j * S; + live[j] = p < k1; + fresh[j] = p >= P0; + at[j] = size_t(fresh[j] ? p - P0 : (RING ? p % RING : p)) * DH; + float d = 0.0f; + if (live[j]) { + const device bfloat* kr = (fresh[j] ? kn : kb) + at[j]; + for (int i = 0; i < DPL; i++) d = fma(q[i], float(kr[i]), d); + } + sc[j] = simd_sum(d); + if (live[j]) bm = metal::max(bm, sc[j]); + } + const float mn = metal::max(m, bm); + const float a = metal::exp(m - mn); + l *= a; + for (int i = 0; i < DPL; i++) o[i] *= a; + for (int j = 0; j < BLK; j++) { + if (!live[j]) continue; + const float b = metal::exp(sc[j] - mn); + l += b; + const device bfloat* vr = (fresh[j] ? vn : vb) + at[j]; + for (int i = 0; i < DPL; i++) o[i] = fma(b, float(vr[i]), o[i]); + } + m = mn; + } + const size_t slot = (size_t(qh) * R + r) * NCH + (c - C0); + if (S == 1) { + if (FUSE.y) { + for (int i = 0; i < DPL; i++) OUT[(size_t(r) * NQ + qh) * DH + lane * DPL + i] = bfloat(o[i] / l); + simdgroup_barrier(mem_flags::mem_device); + gemma_row_sums(OUT + (size_t(r) * NQ + qh) * DH, DH, XS + qh * (DH / 64) * xst, xst, r, lane, 32); + return; + } + if (lane == 0) { PM[slot] = m; PL[slot] = l; } + for (int i = 0; i < DPL; i++) PO[slot * DH + lane * DPL + i] = o[i]; + return; + } + threadgroup float sm[G * S], sl[G * S]; + threadgroup float so[S > 1 ? G * S : 1][S > 1 ? DH : 1]; + if (lane == 0) { sm[sgi] = m; sl[sgi] = l; } + for (int i = 0; i < DPL; i++) so[sgi][lane * DPL + i] = o[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (s != 0) return; + float top = -INFINITY; + for (int t = 0; t < S; t++) top = metal::max(top, sm[g * S + t]); + float lsum = 0.0f, acc[DPL]; + for (int i = 0; i < DPL; i++) acc[i] = 0.0f; + for (int t = 0; t < S; t++) { + const float e = sl[g * S + t] > 0.0f ? metal::exp(sm[g * S + t] - top) : 0.0f; + lsum = fma(sl[g * S + t], e, lsum); + for (int i = 0; i < DPL; i++) acc[i] = fma(so[g * S + t][lane * DPL + i], e, acc[i]); + } + if (FUSE.y) { // FUSE.y (one chunk): gemma_attention_merge's output and sums, exp(top - top) being 1, to OUT and XS + for (int i = 0; i < DPL; i++) OUT[(size_t(r) * NQ + qh) * DH + lane * DPL + i] = bfloat(acc[i] / lsum); + simdgroup_barrier(mem_flags::mem_device); + gemma_row_sums(OUT + (size_t(r) * NQ + qh) * DH, DH, XS + qh * (DH / 64) * xst, xst, r, lane, 32); + return; + } + if (lane == 0) { PM[slot] = top; PL[slot] = lsum; } + for (int i = 0; i < DPL; i++) PO[slot * DH + lane * DPL + i] = acc[i]; +} + +// gilby's attend, merge: a simdgroup per (query head, row) adds the chunks' partials in chunk order. +kernel void gemma_attention_merge(const device float* PM [[buffer(0)]], const device float* PL [[buffer(1)]], + const device float* PO [[buffer(2)]], const device int* META [[buffer(3)]], + device bfloat* OUT [[buffer(4)]], device float* XS [[buffer(5)]], + constant int& xst [[buffer(6)]], uint lane [[thread_index_in_simdgroup]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int qh = int(tg.y), r = int(tg.z); + const int NCH = META[1], R = META[3]; + constexpr int DPL = DH / 32; + const size_t base = (size_t(qh) * R + r) * NCH; + float top = -INFINITY; + for (int c = 0; c < NCH; c++) top = metal::max(top, PM[base + c]); + float lsum = 0.0f, acc[DPL]; + for (int i = 0; i < DPL; i++) acc[i] = 0.0f; + for (int c = 0; c < NCH; c++) { + const float e = PL[base + c] > 0.0f ? metal::exp(PM[base + c] - top) : 0.0f; + lsum = fma(PL[base + c], e, lsum); + for (int i = 0; i < DPL; i++) acc[i] = fma(PO[(base + c) * DH + lane * DPL + i], e, acc[i]); + } + for (int i = 0; i < DPL; i++) OUT[(size_t(r) * NQ + qh) * DH + lane * DPL + i] = bfloat(acc[i] / lsum); + if (xst == 0) return; + simdgroup_barrier(mem_flags::mem_device); + gemma_row_sums(OUT + (size_t(r) * NQ + qh) * DH, DH, XS + qh * (DH / 64) * xst, xst, r, lane, 32); +} + +// Sum of squares over T threads into `out` (p: a threadgroup array of T / 32 floats). +inline float gemma_reduce(float acc, threadgroup float* p, uint lane, uint sg) { + const float s = simd_sum(acc); + if (lane == 0) p[sg] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + float out = 0.0f; + for (int q = 0; q < T / 32; q++) out += p[q]; + return out; +} + +// gilby's attn_tail: hn = h + RMSNorm(o) w_attn, then RMSNorm(hn) times the pre-MLP norm (a dense layer's one input). +kernel void gemma_attn_tail(const device bfloat* H [[buffer(0)]], const device bfloat* O [[buffer(1)]], + const device bfloat* WA [[buffer(2)]], const device bfloat* W1 [[buffer(3)]], + constant float& eps [[buffer(4)]], device bfloat* HN [[buffer(5)]], device bfloat* N1 [[buffer(6)]], + device float* XS [[buffer(7)]], constant int& xst [[buffer(8)]], + uint t [[thread_position_in_threadgroup]], uint r [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + constexpr int PER = D / T; + threadgroup float p1[T / 32], p2[T / 32]; + float ov[PER], hin[PER], wa[PER], w1[PER]; + for (int i = 0; i < PER; i++) { + const int c = int(t) + i * T; + ov[i] = float(O[int(r) * D + c]); hin[i] = float(H[int(r) * D + c]); + wa[i] = float(WA[c]); w1[i] = float(W1[c]); + } + float ss = 0.0f; + for (int i = 0; i < PER; i++) ss = fma(ov[i], ov[i], ss); + const float inv1 = metal::precise::rsqrt(gemma_reduce(ss, p1, lane, sg) / float(D) + eps); + float hv[PER]; + float ss2 = 0.0f; + for (int i = 0; i < PER; i++) { + const int c = int(t) + i * T; + const float a = float(bfloat(wa[i] * float(bfloat(ov[i] * inv1)))); + const bfloat hn = bfloat(hin[i] + a); + HN[int(r) * D + c] = hn; + hv[i] = float(hn); + ss2 = fma(hv[i], hv[i], ss2); + } + const float inv2 = metal::precise::rsqrt(gemma_reduce(ss2, p2, lane, sg) / float(D) + eps); + for (int i = 0; i < PER; i++) { + const int c = int(t) + i * T; + N1[int(r) * D + c] = bfloat(w1[i] * float(bfloat(hv[i] * inv2))); + } + threadgroup_barrier(mem_flags::mem_device); + gemma_row_sums(N1 + int(r) * D, D, XS, xst, int(r), t, T); +} + +// gilby's dense_tail: hn = (h + RMSNorm(y) w_post) * scalar, and RMSNorm(hn) w_next for the next layer. +kernel void gemma_dense_tail(const device bfloat* H [[buffer(0)]], const device bfloat* Y [[buffer(1)]], + const device bfloat* WP [[buffer(2)]], const device bfloat* SC [[buffer(3)]], + const device bfloat* WN [[buffer(4)]], constant float& eps [[buffer(5)]], + device bfloat* HN [[buffer(6)]], device bfloat* NEXT [[buffer(7)]], + device float* XS [[buffer(8)]], constant int& xst [[buffer(9)]], + device bfloat* TAP [[buffer(10)]], constant int& tst [[buffer(11)]], + uint t [[thread_position_in_threadgroup]], uint r [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + constexpr int PER = D / T; + threadgroup float p1[T / 32], p2[T / 32]; + float yv[PER], hin[PER], wp[PER], wn[PER]; + for (int i = 0; i < PER; i++) { + const int c = int(t) + i * T; + yv[i] = float(Y[int(r) * D + c]); hin[i] = float(H[int(r) * D + c]); + wp[i] = float(WP[c]); wn[i] = float(WN[c]); + } + const float sc = float(SC[0]); + float ss1 = 0.0f; + for (int i = 0; i < PER; i++) ss1 = fma(yv[i], yv[i], ss1); + const float inv1 = metal::precise::rsqrt(gemma_reduce(ss1, p1, lane, sg) / float(D) + eps); + float hv[PER]; + float ss2 = 0.0f; + for (int i = 0; i < PER; i++) { + const int c = int(t) + i * T; + const float b = float(bfloat(wp[i] * float(bfloat(yv[i] * inv1)))); + const bfloat hs = bfloat(hin[i] + b); + const bfloat hn = bfloat(float(hs) * sc); + HN[int(r) * D + c] = hn; + if (tst) TAP[size_t(r) * tst + c] = hn; // a drafter's tapped layer: its row at the tap's columns + hv[i] = float(hn); + ss2 = fma(hv[i], hv[i], ss2); + } + const float inv2 = metal::precise::rsqrt(gemma_reduce(ss2, p2, lane, sg) / float(D) + eps); + for (int i = 0; i < PER; i++) { + const int c = int(t) + i * T; + NEXT[int(r) * D + c] = bfloat(wn[i] * float(bfloat(hv[i] * inv2))); + } + threadgroup_barrier(mem_flags::mem_device); + gemma_row_sums(NEXT + int(r) * D, D, XS, xst, int(r), t, T); +} + +// mlx_lm's geglu as MLX's bf16 ops; IL: gate and up columns in interleaved 16-wide blocks, else two halves. +kernel void gemma_geglu(const device bfloat* GU [[buffer(0)]], device bfloat* ACT [[buffer(1)]], + constant uint& n [[buffer(2)]], device float* XS [[buffer(3)]], constant int& xst [[buffer(4)]], + constant int& IL [[buffer(5)]], uint2 tg [[threadgroup_position_in_grid]], uint2 tp [[thread_position_in_threadgroup]]) { + const uint g = tg.x, r = tg.y, i = tp.x, c = g * 64 + i; + const uint gi = IL ? 32 * (c / 16) + c % 16 : c, ui = IL ? gi + 16 : n + c; + const bfloat x = GU[size_t(r) * 2 * n + gi], up = GU[size_t(r) * 2 * n + ui]; + const bfloat cube = bfloat(metal::pow(float(x), 3.0f)); + const bfloat inner = bfloat(float(x) + float(bfloat(float(bfloat(GELU_C1)) * float(cube)))); + const bfloat th = bfloat(metal::precise::tanh(float(bfloat(float(bfloat(GELU_C0)) * float(inner))))); + const bfloat half_x = bfloat(0.5f * float(x)); + const bfloat gl = bfloat(float(half_x) * float(bfloat(1.0f + float(th)))); + const bfloat out = bfloat(float(gl) * float(up)); + ACT[size_t(r) * n + c] = out; + if (xst == 0) return; + threadgroup float vals[64]; + vals[i] = float(out); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (i != 0) return; + float sum = 0.0f; + for (int j = 0; j < 64; j++) sum += vals[j]; + XS[g * xst + r] = sum; +} + +// mlx_lm's logit_softcap: tanh(x / cap) * cap as bf16 ops. +kernel void gemma_softcap(device bfloat* X [[buffer(0)]], constant float& cap [[buffer(1)]], + uint i [[thread_position_in_grid]]) { + X[i] = bfloat(float(bfloat(metal::precise::tanh(float(bfloat(float(X[i]) / cap))))) * cap); +} + +// mlx_lm's ProportionalRoPE frequencies: base ** (arange(0, rotated, 2) / dims) in float32, as MLX's pow computes it. +kernel void gemma_rope_freqs(device float* F [[buffer(0)]], constant float& base [[buffer(1)]], + constant uint& dims [[buffer(2)]], uint i [[thread_position_in_grid]]) { + F[i] = metal::precise::pow(base, float(2 * i) / float(dims)); +} + +// A window's new keys and values [HK, R, DH] into a layer's cache [HK, cap, DH] at slot p % ring (0: p). +kernel void gemma_kv_write(const device bfloat* KN [[buffer(0)]], const device bfloat* VN [[buffer(1)]], + device bfloat* K [[buffer(2)]], device bfloat* V [[buffer(3)]], + const device int* POS [[buffer(4)]], constant int3& m [[buffer(5)]], + uint3 pos [[thread_position_in_grid]]) { + const int d = int(pos.x), r = int(pos.y), h = int(pos.z), R = m.x, ring = m.y, cap = m.z; + const int p = POS[r]; + const size_t dst = (size_t(h) * cap + (ring ? p % ring : p)) * DH + d; + const size_t src = (size_t(h) * R + r) * DH + d; + K[dst] = KN[src]; + V[dst] = VN[src]; +} + +// DFlash drafter rows (a: x stride, width, first position, ring or 0) into cache slots, position % ring or position. +kernel void gemma_draft_store(const device bfloat* X [[buffer(0)]], device bfloat* C [[buffer(1)]], + constant uint4& a [[buffer(2)]], uint2 pos [[thread_position_in_grid]]) { + if (pos.x >= a.y) return; + const uint p = a.z + pos.y, slot = a.w ? p % a.w : p; + C[size_t(slot) * a.y + pos.x] = X[size_t(pos.y) * a.x + pos.x]; +} + +// tf_draft_swiglu's bf16 steps on GU rows of gate then up (pitch 2n) into ACT (pitch n), 64 threads a group of sums. +kernel void gemma_draft_swiglu(const device bfloat* GU [[buffer(0)]], device bfloat* ACT [[buffer(1)]], + constant uint& n [[buffer(2)]], device float* XS [[buffer(3)]], constant int& xst [[buffer(4)]], + uint2 tg [[threadgroup_position_in_grid]], uint2 tp [[thread_position_in_threadgroup]]) { + const uint g = tg.x, r = tg.y, i = tp.x, c = g * 64 + i; + const float gate = float(GU[size_t(r) * 2 * n + c]), up = float(GU[size_t(r) * 2 * n + n + c]); + const bfloat exponential = bfloat(metal::exp(metal::abs(gate))); + const bfloat denominator = bfloat(1.0f + float(exponential)); + const bfloat inverse = bfloat(1.0f / float(denominator)); + const bfloat sigmoid = gate < 0.0f ? inverse : bfloat(1.0f - float(inverse)); + const bfloat activation = bfloat(float(bfloat(gate)) * float(sigmoid)); + const bfloat out = bfloat(float(activation) * float(bfloat(up))); + ACT[size_t(r) * n + c] = out; + threadgroup float vals[64]; + vals[i] = float(out); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (i != 0) return; + float sum = 0.0f; + for (int j = 0; j < 64; j++) sum += vals[j]; + XS[g * xst + r] = sum; +} + +// The drafter's residual: H += Y, each sum rounded to bf16. +kernel void gemma_draft_add(device bfloat* H [[buffer(0)]], const device bfloat* Y [[buffer(1)]], + uint i [[thread_position_in_grid]]) { + H[i] = bfloat(float(H[i]) + float(Y[i])); +} + +// The MTP assistant's rows [embedding | H's row], 2 D wide: gemma_embed's on ids[r * ist.x] (a pick), kept in vocab. +kernel void gemma_assist_input(const device uint* ids [[buffer(0)]], constant uint2& ist [[buffer(1)]], + const device uchar* weight [[buffer(2)]], const device bfloat* scales [[buffer(3)]], + const device bfloat* biases [[buffer(4)]], const device bfloat* H [[buffer(5)]], + device bfloat* out [[buffer(6)]], constant float& scale [[buffer(7)]], + uint2 pos [[thread_position_in_grid]]) { + const uint b = pos.x, r = pos.y; + const size_t id = metal::min(ids[r * ist.x], ist.y - 1); + const uint2 v = gemma_codes(weight, id, b); + const size_t g = id * (D / 64) + (2 * b) / 64; + const float s = float(scales[g]), z = float(biases[g]); + device bfloat* row = out + size_t(r) * 2 * D; + row[2 * b] = bfloat(float(bfloat(s * float(v.x) + z)) * scale); + row[2 * b + 1] = bfloat(float(bfloat(s * float(v.y) + z)) * scale); + row[D + 2 * b] = H[size_t(r) * D + 2 * b]; + row[D + 2 * b + 1] = H[size_t(r) * D + 2 * b + 1]; +} diff --git a/zig/kernels/metal/gemma4/prompt_attention.metal b/zig/kernels/metal/gemma4/prompt_attention.metal new file mode 100644 index 0000000000..ad84840705 --- /dev/null +++ b/zig/kernels/metal/gemma4/prompt_attention.metal @@ -0,0 +1,329 @@ +// Gemma 4 prompt attention on the tensor units: 32 query rows of a head a threadgroup, 32 keys a simdgroup a block. +#include "../nax.h" +using namespace tfp; + +// Compiled with DH (head dim), NQ, G (query heads a key head), SL (= DH / 128), WINDOW (0: global), RING (0: linear). + +constant constexpr float kMasked = -3.402823466e+38f; +constant constexpr float kLog2e = 1.44269504089f; +constant constexpr int BK = 32; + +// Row max (MAX) or sum over a fragment's 32 columns into r[0], r[1] (rows home.y, home.y + 8), as attention_nax folds. +template +inline void fold(thread const frag& f, thread float* r) { + TF_UNROLL + for (short h = 0; h < 2; h++) { + const short b = 4 * h; + float t = MAX ? max(max(f[b], f[b + 1]), max(f[b + 2], f[b + 3])) : (f[b] + f[b + 1]) + (f[b + 2] + f[b + 3]); + const float u = simd_shuffle_xor(t, ushort(1)); + t = MAX ? max(t, u) : t + u; + const float w = simd_shuffle_xor(t, ushort(8)); + t = MAX ? max(t, w) : t + w; + r[h] = MAX ? max(r[h], t) : r[h] + t; + } +} + +// A tile is RT query rows of one head: NS simdgroups each score KS keys of a block over the whole head (Q hot in L1). +constant constexpr int RT = 32; // query rows a tile +// Each simdgroup accumulates its own 64 output columns; once a block the tile's row maxes cross threadgroup memory. +constant constexpr int NS = DH / 64; // simdgroups a tile +constant constexpr int KS = 32; // keys a simdgroup a block +// The fragment-step loops stay rolled: fully unrolled, the code outgrew the instruction cache and ran 5x slower. +constant constexpr int BKT = KS * NS; // keys a block +// The probabilities cross threadgroup memory PH keys at a time, as fp32 (bf16 moved an 8k prompt's logits off mlx-lm). +constant constexpr int PH = 128; // keys whose probabilities cross at once +constant constexpr int PP = PH + 8; // the probabilities' row pitch +using PT = float; + +// A bf16 matrix's 16x16 block at (r, c): a left or (TR) transposed right operand (on M5 a lane's row in one load). +template +inline frag blk(const device bfloat* p, int ld, int r, int c, short2 home) { +#ifndef TF_SIMD_FRAGS + const device bfloat* q = p + (r + home.y) * ld + c + home.x; + return frag(*(const device bfloat4*)q, *(const device bfloat4*)(q + 8 * ld)); +#else + frag f; + if (TR) { + frag_get_t(f, p, ld, r, c, home); + } else { + frag_get(f, p, ld, r, c, home); + } + return f; +#endif +} + +// The probabilities' 16 x 16 block at (r, c) of the tile, and a scored fragment's store into it. +inline frag prob_frag(const threadgroup PT* pt, int r, int c, short2 home) { +#ifndef TF_SIMD_FRAGS + const threadgroup PT* q = pt + (r + home.y) * PP + c + home.x; + return frag(*(const threadgroup vec*)q, *(const threadgroup vec*)(q + 8 * PP)); +#else + frag f; + frag_get(f, pt, PP, r, c, home); + return f; +#endif +} + +inline void prob_put(threadgroup PT* pt, int r, int c, short2 home, thread const frag& s) { +#ifndef TF_SIMD_FRAGS + threadgroup PT* q = pt + (r + home.y) * PP + c + home.x; + *(threadgroup vec*)q = vec(s[0], s[1], s[2], s[3]); + *(threadgroup vec*)(q + 8 * PP) = vec(s[4], s[5], s[6], s[7]); +#else + TF_UNROLL + for (short e = 0; e < 8; e++) pt[(r + home.y + (e >> 2) * 8) * PP + c + home.x + TF_COL(e)] = PT(s[e]); +#endif +} + +// Q [R, NQ, DH] at positions p0..; K, V the cache [HK, cap, DH]; row j of top and total is the tile's home.y + 8j. +[[max_total_threads_per_threadgroup(256)]] +// Fragment i of s: tile rows 16 (i >> 1).., keys 16 (i & 1)..; of acc: rows 16 (i >> 2).., own columns 16 (i & 3)... +[[kernel]] void gemma_prompt_attention(const device bfloat* Q [[buffer(0)]], const device bfloat* K [[buffer(1)]], + const device bfloat* V [[buffer(2)]], device bfloat* O [[buffer(3)]], + constant int3& P [[buffer(4)]], + const device int* HI [[buffer(5)]], // sliding layers: a row's last key (image blocks) + uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { + const int R = P.x, p0 = P.y, cap = P.z; + const int qh = int(tg.y), kv = qh / G; + const int r0 = int(tg.x) * RT; + const int q_rows = min(RT, R - r0); + const int qs = NQ * DH; + const device bfloat* q = Q + r0 * qs + qh * DH; + const device bfloat* kc = K + long(kv) * cap * DH; + const device bfloat* vc = V + long(kv) * cap * DH + int(sg) * 64; + const short2 home = frag_home(ushort(lane)); + threadgroup PT pt[RT * PP]; + threadgroup float rowmax[NS][RT]; + frag acc[8]; + TF_UNROLL + for (short i = 0; i < 8; i++) acc[i] = frag(0); + float top[4] = {kMasked, kMasked, kMasked, kMasked}, total[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + const int first_pos = p0 + r0, last_pos = WINDOW ? HI[r0 + q_rows - 1] : p0 + r0 + q_rows - 1; // the tile's last key + const int lo = WINDOW ? max(0, first_pos - WINDOW + 1) : 0; + for (int kb = lo / BKT; kb <= last_pos / BKT; kb++) { + const int kpos = kb * BKT + int(sg) * KS; + const int slot = RING ? kpos % RING : kpos; + const device bfloat* kblk = kc + long(slot) * DH; + frag s[4]; + TF_UNROLL + for (short i = 0; i < 4; i++) s[i] = frag(0); + for (int d = 0; d < DH / 16; d++) { + const frag k0 = blk(kblk, DH, 0, 16 * d, home); + const frag k1 = blk(kblk, DH, 16, 16 * d, home); + TF_UNROLL + for (short h = 0; h < 2; h++) { // rows past q_rows sit inside the 512-row buffer; their outputs stay unwritten + const frag qf = blk(q, qs, 16 * h, 16 * d, home); + mma_16x32(s[2 * h], s[2 * h + 1], qf, k0, k1); + } + } + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) { + const int local = 16 * (i >> 1) + home.y + (e >> 2) * 8, row = first_pos + local; + const int col = kpos + 16 * (i & 1) + home.x + TF_COL(e); + const int hi = WINDOW && local < q_rows ? HI[r0 + local] : row; + const bool out = col > hi || (WINDOW && col < row - WINDOW + 1); + s[i][e] = out ? kMasked : s[i][e] * kLog2e; + } + } + float top_new[4] = {top[0], top[1], top[2], top[3]}; + TF_UNROLL + for (short i = 0; i < 4; i++) fold(s[i], top_new + 2 * (i >> 1)); + if (home.x == 0) { + TF_UNROLL + for (short j = 0; j < 4; j++) rowmax[sg][home.y + 8 * j] = top_new[j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + TF_UNROLL + for (short j = 0; j < 4; j++) { + for (int t = 0; t < NS; t++) top_new[j] = max(top_new[j], rowmax[t][home.y + 8 * j]); + } + TF_UNROLL + for (short i = 0; i < 4; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) { + s[i][e] = s[i][e] == kMasked ? 0.0f : fast::exp2(s[i][e] - top_new[2 * (i >> 1) + (e >> 2)]); + } + } + float keep[4]; + TF_UNROLL + for (short j = 0; j < 4; j++) { + keep[j] = top[j] == kMasked ? 0.0f : fast::exp2(top[j] - top_new[j]); + top[j] = top_new[j]; + total[j] = total[j] * keep[j]; + } + TF_UNROLL + for (short i = 0; i < 4; i++) fold(s[i], total + 2 * (i >> 1)); + TF_UNROLL + for (short i = 0; i < 8; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * keep[2 * (i >> 2) + (e >> 2)]; + } + for (int ph = 0; ph < BKT / PH; ph++) { + if (ph > 0) threadgroup_barrier(mem_flags::mem_threadgroup); // the phase before has read the tile + if (int(sg) / (PH / KS) == ph) { + TF_UNROLL + for (short i = 0; i < 4; i++) prob_put(pt, 16 * (i >> 1), int(sg) * KS - ph * PH + 16 * (i & 1), home, s[i]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int kk = 0; kk < PH / 16; kk++) { + const int kp = kb * BKT + ph * PH + 16 * kk; + if (kp > last_pos) break; // keys past the tile's rows: masked, and their cache rows may hold anything + const int valid = min(16, last_pos + 1 - kp); // a value row never written would poison the tile as 0 x NaN + const device bfloat* vblk = vc + long(RING ? kp % RING : kp) * DH; + frag v0, v1; + TF_UNROLL + for (short d = 0; d < 4; d += 2) { + if (valid == 16) { + v0 = blk(vblk, DH, 0, 16 * d, home); + v1 = blk(vblk, DH, 0, 16 * d + 16, home); + } else { + frag_get_in(v0, vblk, DH, 0, 16 * d, home, valid, DH); + frag_get_in(v1, vblk, DH, 0, 16 * d + 16, home, valid, DH); + } + TF_UNROLL + for (short h = 0; h < 2; h++) { + const frag pf = prob_frag(pt, 16 * h, 16 * kk, home); + mma_16x32(acc[4 * h + d], acc[4 * h + d + 1], pf, v0, v1); + } + } + } + } + } + // the tile's totals: every simdgroup's partial sums, scaled alike by the shared maxes + threadgroup_barrier(mem_flags::mem_threadgroup); + if (home.x == 0) { + TF_UNROLL + for (short j = 0; j < 4; j++) rowmax[sg][home.y + 8 * j] = total[j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float inv[4]; + TF_UNROLL + for (short j = 0; j < 4; j++) { + float t = 0.0f; + for (int u = 0; u < NS; u++) t += rowmax[u][home.y + 8 * j]; + inv[j] = t > 0.0f ? 1.0f / t : 0.0f; + } + device bfloat* o = O + r0 * qs + qh * DH + int(sg) * 64; + TF_UNROLL + for (short i = 0; i < 8; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * inv[2 * (i >> 2) + (e >> 2)]; + frag_put_in(acc[i], o, qs, 16 * (i >> 2), 16 * (i & 3), home, q_rows, 64); + } +} + + +// Decode rows, a key head's G query heads as a tile's rows; partials by CD-key chunk. P: rows, p0, cap, chunks. +constant constexpr int CD = 256; + +[[max_total_threads_per_threadgroup(128)]] +[[kernel]] void gemma_decode_attention(const device bfloat* Q [[buffer(0)]], const device bfloat* K [[buffer(1)]], + const device bfloat* V [[buffer(2)]], constant int4& P [[buffer(3)]], + device float* PM [[buffer(4)]], device float* PL [[buffer(5)]], + device float* PO [[buffer(6)]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { + constexpr int RPT = 16 / G; // decode rows a tile + const int R = P.x, p0 = P.y, cap = P.z, NCH = P.w; + const int c = int(tg.x), kv = int(tg.y), r0 = int(tg.z) * RPT; + const int qs = NQ * DH; + const short2 home = frag_home(ushort(lane)); + // tile row t: decode row r0 + t / G, query head kv * G + t % G; the fragment's two row halves are two decode rows + const int ra = r0 + home.y / G, rb = r0 + (home.y + 8) / G; + const bool la = ra < R, lb = rb < R; + const device bfloat* qa = Q + ra * qs + (kv * G + home.y % G) * DH + int(sg) * 128; + const device bfloat* qb = Q + rb * qs + (kv * G + (home.y + 8) % G) * DH + int(sg) * 128; + frag qf[8]; + TF_UNROLL + for (short d = 0; d < 8; d++) frag_get_rows(qf[d], qa, qb, 16 * d, home, la, lb); + const device bfloat* kc = K + long(kv) * cap * DH + int(sg) * 128; + const device bfloat* vc = V + long(kv) * cap * DH + int(sg) * 128; + frag acc[8]; + TF_UNROLL + for (short i = 0; i < 8; i++) acc[i] = frag(0); + float top[2] = {kMasked, kMasked}, total[2] = {0.0f, 0.0f}; + const int pos_a = p0 + ra, pos_b = p0 + rb; + const int last = min(p0 + R - 1, (c + 1) * CD - 1); + threadgroup float part[SL][32][16]; + for (int kpos = c * CD; kpos <= last; kpos += BK) { + const device bfloat* kblk = kc + long(kpos) * DH; + const device bfloat* vblk = vc + long(kpos) * DH; + frag s[2] = {frag(0), frag(0)}; + TF_UNROLL + for (short d = 0; d < 8; d++) { + frag k0, k1; + frag_get_t(k0, kblk, DH, 0, 16 * d, home); + frag_get_t(k1, kblk, DH, 16, 16 * d, home); + mma_16x32(s[0], s[1], qf[d], k0, k1); + } + if (SL > 1) { + TF_UNROLL + for (short e = 0; e < 8; e++) { part[sg][lane][e] = s[0][e]; part[sg][lane][8 + e] = s[1][e]; } + threadgroup_barrier(mem_flags::mem_threadgroup); + TF_UNROLL + for (short e = 0; e < 8; e++) { + float a = part[0][lane][e], b = part[0][lane][8 + e]; + for (int t = 1; t < SL; t++) { a += part[t][lane][e]; b += part[t][lane][8 + e]; } + s[0][e] = a; + s[1][e] = b; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) { + const int col = kpos + 16 * f + home.x + TF_COL(e); + s[f][e] = col > ((e >> 2) ? pos_b : pos_a) ? kMasked : s[f][e] * 1.44269504089f; + } + } + float top_new[2] = {top[0], top[1]}; + fold(s[0], top_new); + fold(s[1], top_new); + TF_UNROLL + for (short f = 0; f < 2; f++) { + TF_UNROLL + for (short e = 0; e < 8; e++) s[f][e] = s[f][e] == kMasked ? 0.0f : fast::exp2(s[f][e] - top_new[e >> 2]); + } + float keep[2]; + TF_UNROLL + for (short h = 0; h < 2; h++) { + keep[h] = top[h] == kMasked ? 0.0f : fast::exp2(top[h] - top_new[h]); + top[h] = top_new[h]; + total[h] = total[h] * keep[h]; + } + fold(s[0], total); + fold(s[1], total); + TF_UNROLL + for (short i = 0; i < 8; i++) { + TF_UNROLL + for (short e = 0; e < 8; e++) acc[i][e] = acc[i][e] * keep[e >> 2]; + } + TF_UNROLL + for (short d = 0; d < 8; d += 2) { + TF_UNROLL + for (short k = 0; k < 2; k++) { + frag v0, v1; + frag_get(v0, vblk, DH, 16 * k, 16 * d, home); + frag_get(v1, vblk, DH, 16 * k, 16 * d + 16, home); + mma_16x32(acc[d], acc[d + 1], s[k], v0, v1); + } + } + } + // partials in gemma_attention_merge's layout: max and sum in log2 units back to natural (merge uses exp) + TF_UNROLL + for (short h = 0; h < 2; h++) { + const int r = h ? rb : ra; + if (!(h ? lb : la)) continue; + const int qh = kv * G + (home.y + 8 * h) % G; + const size_t slot = (size_t(qh) * R + r) * NCH + c; + if (home.x == 0 && sg == 0) { PM[slot] = top[h] == kMasked ? -INFINITY : top[h] * 0.69314718056f; PL[slot] = total[h]; } + TF_UNROLL + for (short i = 0; i < 8; i++) { + TF_UNROLL + for (short e = 0; e < 4; e++) PO[slot * DH + int(sg) * 128 + 16 * i + home.x + TF_COL(e)] = acc[i][4 * h + e]; + } + } +} diff --git a/zig/kernels/metal/gemma4/vision.metal b/zig/kernels/metal/gemma4/vision.metal new file mode 100644 index 0000000000..472b8074c7 --- /dev/null +++ b/zig/kernels/metal/gemma4/vision.metal @@ -0,0 +1,211 @@ +// Gemma 4's image tower at torch's per-op bf16 rounding (transformers' modeling_gemma4.py); after gemm_nax.metal. + +#define G4V_GEMM(NAME, T, TA, TB, AM, AN) \ + [[max_total_threads_per_threadgroup(256)]] [[kernel]] void NAME( \ + const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], const device int32_t* P [[buffer(2)]], \ + device T* D [[buffer(3)]], uint sg [[simdgroup_index_in_threadgroup]], \ + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { \ + tfp::gemm_nax(A, B, D, P, tg, sg, lane); \ + } +G4V_GEMM(g4v_gemm_bf16_nt_00, bfloat16_t, false, true, false, false) +G4V_GEMM(g4v_gemm_bf16_nt_01, bfloat16_t, false, true, false, true) +G4V_GEMM(g4v_gemm_bf16_nt_10, bfloat16_t, false, true, true, false) +G4V_GEMM(g4v_gemm_bf16_nt_11, bfloat16_t, false, true, true, true) +G4V_GEMM(g4v_gemm_f32_nt_00, float, false, true, false, false) +G4V_GEMM(g4v_gemm_f32_nt_01, float, false, true, false, true) +G4V_GEMM(g4v_gemm_f32_nt_10, float, false, true, true, false) +G4V_GEMM(g4v_gemm_f32_nt_11, float, false, true, true, true) +G4V_GEMM(g4v_gemm_f32_nn_00, float, false, false, false, false) +G4V_GEMM(g4v_gemm_f32_nn_01, float, false, false, false, true) +G4V_GEMM(g4v_gemm_f32_nn_10, float, false, false, true, false) +G4V_GEMM(g4v_gemm_f32_nn_11, float, false, false, true, true) + +constant constexpr uint G4V_W = 1152, G4V_HEADS = 16, G4V_DH = 72, G4V_HALF = 36, G4V_QKV = 3 * G4V_W; + +inline float g4v_bf(float x) { return float(bfloat(x)); } + +inline float g4v_group_sum(float v, threadgroup float* part, uint n, uint lane, uint sg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float w = lane < (n + 31) / 32 ? part[lane] : 0.0f; + w = simd_sum(w); + if (lane == 0) part[32] = w; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float r = part[32]; + threadgroup_barrier(mem_flags::mem_threadgroup); + return r; +} + +struct G4vRows { uint width; uint weighted; float eps; }; + +// The patches' x and y position rows into the projected patches: x = bf16(x + bf16(tx[px] + ty[py])), grid gw wide. +[[kernel]] void g4v_pos(device bfloat* x [[buffer(0)]], const device bfloat* table [[buffer(1)]], + constant uint& gw [[buffer(2)]], uint2 p [[thread_position_in_grid]]) { + if (p.x >= G4V_W) return; + const uint px = p.y % gw, py = p.y / gw; + const float e = g4v_bf(float(table[ulong(px) * G4V_W + p.x]) + float(table[(10240ul + py) * G4V_W + p.x])); + const ulong i = ulong(p.y) * G4V_W + p.x; + x[i] = bfloat(float(x[i]) + e); +} + +// Gemma4RMSNorm: x * (mean(x^2) + eps)^-0.5 in fp32, times the weight when `weighted`, rounded to bf16 once. +[[kernel]] void g4v_rms(const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], device bfloat* y [[buffer(2)]], + constant G4vRows& p [[buffer(3)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint n [[threads_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + const device bfloat* xr = x + ulong(row) * p.width; + device bfloat* yr = y + ulong(row) * p.width; + float s = 0.0f; + for (uint i = t; i < p.width; i += n) s += float(xr[i]) * float(xr[i]); + const float r = metal::precise::rsqrt(g4v_group_sum(s, part, n, lane, sg) / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) yr[i] = bfloat(p.weighted ? float(xr[i]) * r * float(w[i]) : float(xr[i]) * r); +} + +// A branch's post norm into the residual: x = bf16(x + Gemma4RMSNorm(y) * w). +[[kernel]] void g4v_rms_add(device bfloat* x [[buffer(0)]], const device bfloat* y [[buffer(1)]], const device bfloat* w [[buffer(2)]], + constant float& eps [[buffer(3)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint n [[threads_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + device bfloat* xr = x + ulong(row) * G4V_W; + const device bfloat* yr = y + ulong(row) * G4V_W; + float s = 0.0f; + for (uint i = t; i < G4V_W; i += n) s += float(yr[i]) * float(yr[i]); + const float r = metal::precise::rsqrt(g4v_group_sum(s, part, n, lane, sg) / float(G4V_W) + eps); + for (uint i = t; i < G4V_W; i += n) xr[i] = bfloat(float(xr[i]) + g4v_bf(float(yr[i]) * r * float(w[i]))); +} + +// qkv [P, 3456] to head-major q, k, v fp32: per-head norms (v unweighted), q and k by the 2D RoPE (x, then y). +[[kernel]] void g4v_qkv(const device bfloat* qkv [[buffer(0)]], const device bfloat* qw [[buffer(1)]], + const device bfloat* kw [[buffer(2)]], const device bfloat* cs [[buffer(3)]], + const device bfloat* sn [[buffer(4)]], device float* Q [[buffer(5)]], device float* K [[buffer(6)]], + device float* V [[buffer(7)]], constant uint& P [[buffer(8)]], constant float& eps [[buffer(9)]], + uint g [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) { + threadgroup float nrm[G4V_DH]; + const uint part = g % 3, head = (g / 3) % G4V_HEADS, row = g / (3 * G4V_HEADS); + if (row >= P) return; + const device bfloat* x = qkv + ulong(row) * G4V_QKV + part * G4V_W + head * G4V_DH; + const device bfloat* w = part == 0 ? qw : kw; + float s = 0.0f; + for (uint d = lane; d < G4V_DH; d += 32) s += float(x[d]) * float(x[d]); + const float r = metal::precise::rsqrt(simd_sum(s) / float(G4V_DH) + eps); + for (uint d = lane; d < G4V_DH; d += 32) nrm[d] = g4v_bf(part == 2 ? float(x[d]) * r : float(x[d]) * r * float(w[d])); + threadgroup_barrier(mem_flags::mem_threadgroup); + device float* out = (part == 0 ? Q : part == 1 ? K : V) + (ulong(head) * P + row) * G4V_DH; + for (uint d = lane; d < G4V_DH; d += 32) { + if (part == 2) { out[d] = nrm[d]; continue; } + const uint base = d / G4V_HALF * G4V_HALF, j = d % G4V_HALF; + const float other = j < G4V_HALF / 2 ? -nrm[base + j + G4V_HALF / 2] : nrm[base + j - G4V_HALF / 2]; + const float c = float(cs[row * G4V_DH + d]), sv = float(sn[row * G4V_DH + d]); + out[d] = g4v_bf(g4v_bf(nrm[d] * c) + g4v_bf(other * sv)); + } +} + +// Each row of S (n scores) to its softmax, in place (the tower's attention scale is 1). +[[kernel]] void g4v_softmax(device float* S [[buffer(0)]], constant uint& n [[buffer(1)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint tn [[threads_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + device float* s = S + ulong(row) * n; + float m = -INFINITY; + for (uint i = t; i < n; i += tn) m = max(m, s[i]); + m = simd_max(m); + if (lane == 0) part[sg] = m; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float v = lane < (tn + 31) / 32 ? part[lane] : -INFINITY; + v = simd_max(v); + if (lane == 0) part[32] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + m = part[32]; + threadgroup_barrier(mem_flags::mem_threadgroup); + float z = 0.0f; + for (uint i = t; i < n; i += tn) { + const float e = metal::precise::exp(s[i] - m); + s[i] = e; + z += e; + } + const float inv = 1.0f / g4v_group_sum(z, part, tn, lane, sg); + for (uint i = t; i < n; i += tn) s[i] *= inv; +} + +// O [16, rows, 72] fp32 to rows r0.. of out [P, 1152] bf16. +[[kernel]] void g4v_heads(const device float* O [[buffer(0)]], device bfloat* out [[buffer(1)]], + constant uint2& p [[buffer(2)]], uint2 at [[thread_position_in_grid]]) { + const uint rows = p.x, r0 = p.y, h = at.x / G4V_DH, d = at.x % G4V_DH, r = at.y; + if (at.x >= G4V_W || r >= rows) return; + out[ulong(r0 + r) * G4V_W + at.x] = bfloat(O[(ulong(h) * rows + r) * G4V_DH + d]); +} + +// gate | up [P, 2 n] to bf16(bf16(gelu_tanh(gate)) * up) [P, n] (torch's gelu, approximate "tanh", on bf16). +[[kernel]] void g4v_geglu(const device bfloat* gu [[buffer(0)]], device bfloat* act [[buffer(1)]], + constant uint& n [[buffer(2)]], uint2 p [[thread_position_in_grid]]) { + if (p.x >= n) return; + const float g = float(gu[ulong(p.y) * 2 * n + p.x]), u = float(gu[ulong(p.y) * 2 * n + n + p.x]); + const float a = g4v_bf(0.5f * g * (1.0f + metal::precise::tanh(0.7978845608028654f * (g + 0.044715f * g * g * g)))); + act[ulong(p.y) * n + p.x] = bfloat(a * u); +} + +// Gemma4VisionPooler and standardize: 3 x 3 block means to bf16, then (x sqrt(1152) - bias) scale in fp32, to bf16. +[[kernel]] void g4v_pool(const device bfloat* x [[buffer(0)]], const device bfloat* bias [[buffer(1)]], + const device bfloat* scale [[buffer(2)]], device bfloat* out [[buffer(3)]], + constant uint& gw [[buffer(4)]], uint2 p [[thread_position_in_grid]]) { + if (p.x >= G4V_W) return; + const uint bw = gw / 3, ty = p.y / bw, tx = p.y % bw; + float s = 0.0f; + for (uint dy = 0; dy < 3; dy++) + for (uint dx = 0; dx < 3; dx++) s += float(x[(ulong(3 * ty + dy) * gw + 3 * tx + dx) * G4V_W + p.x]) * (1.0f / 9.0f); + const float h = g4v_bf(s) * 33.941125496954285f; + out[ulong(p.y) * G4V_W + p.x] = bfloat((h - float(bias[p.x])) * float(scale[p.x])); +} + +// An image's rows over the prompt's embedded rows, G4V_OUT (the language model's width) bf16 a row, unscaled. +[[kernel]] void g4v_rows(const device uint4* src [[buffer(0)]], device uint4* dst [[buffer(1)]], + uint2 p [[thread_position_in_grid]]) { + if (p.x >= G4V_OUT / 8) return; + dst[ulong(p.y) * (G4V_OUT / 8) + p.x] = src[ulong(p.y) * (G4V_OUT / 8) + p.x]; +} + +// The 12B's image embedder (Gemma4UnifiedVisionEmbedder): a bf16 nn.LayerNorm a row, in fp32 and rounded once. +struct G4uNorm { + uint width; + uint add; // 1: plus a Linear's bias in fp32 (addmm's one rounding); 2: bf16(x + bf16(pos[px][0] + pos[py][1])) + uint gw; // the merged grid's width (add 2) + float eps; +}; + +inline float g4u_in(const device bfloat* x, const device bfloat* extra, constant G4uNorm& p, uint row, uint i) { + const float v = float(x[ulong(row) * p.width + i]); + if (p.add == 1) return v + float(extra[i]); + if (p.add == 2) { + const uint px = row % p.gw, py = row / p.gw; + const float e = g4v_bf(float(extra[ulong(2 * px) * p.width + i]) + float(extra[ulong(2 * py + 1) * p.width + i])); + return g4v_bf(v + e); + } + return v; +} + +[[kernel]] void g4u_norm(const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], + const device bfloat* b [[buffer(2)]], const device bfloat* extra [[buffer(3)]], + device bfloat* y [[buffer(4)]], constant G4uNorm& p [[buffer(5)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], + uint n [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + float s = 0.0f; + for (uint i = t; i < p.width; i += n) s += g4u_in(x, extra, p, row, i); + const float mean = g4v_group_sum(s, part, n, lane, sg) / float(p.width); + float d2 = 0.0f; + for (uint i = t; i < p.width; i += n) { + const float d = g4u_in(x, extra, p, row, i) - mean; + d2 += d * d; + } + const float r = metal::precise::rsqrt(g4v_group_sum(d2, part, n, lane, sg) / float(p.width) + p.eps); + device bfloat* yr = y + ulong(row) * p.width; + for (uint i = t; i < p.width; i += n) yr[i] = bfloat((g4u_in(x, extra, p, row, i) - mean) * r * float(w[i]) + float(b[i])); +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 0573c8832b..379e4ec288 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -5,6 +5,14 @@ pub const nemotron = @import("nemotron/kernels.zig"); pub const qwen35 = @import("qwen3_5/kernels.zig"); pub const qwen35_layout = @embedFile("qwen3_5/layout.metal"); +/// Gemma 4 dense decode glue around the core lane projections (families/gemma4/kernels.zig). +pub const gemma4_glue = @embedFile("gemma4/glue.metal"); +/// Gemma 4 decode glue as the reg lane kernel's prologues and epilogue (families/gemma4/weights.zig). +pub const gemma4_fused = @embedFile("gemma4/fused.metal"); +/// Gemma 4 prompt attention on the tensor units (families/gemma4/kernels.zig). +pub const gemma4_prompt_attention = @embedFile("gemma4/prompt_attention.metal"); +/// Gemma 4's image tower: gemm_nax and its glue (families/gemma4/vision.zig). +pub const gemma4_vision = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("gemma4/vision.metal"); /// LFM2-MoE: embedding, norms, the short conv, attention, the route and the experts' glue. pub const lfm2 = @embedFile("lfm2/lfm2.metal"); diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 99d6835a49..a68e6c56a4 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -14,6 +14,8 @@ const titles = [_]struct { model_type: []const u8, title: []const u8 }{ .{ .model_type = "qwen4_exp", .title = "Qwen3.8 Flash Next" }, .{ .model_type = "glm5_next", .title = "GLM-5.3-Flash" }, .{ .model_type = "qwen3_5", .title = "Qwen3.5-2B or Qwen3.8-27B" }, + .{ .model_type = "gemma4", .title = "Gemma 4 31B" }, + .{ .model_type = "gemma4_unified", .title = "Gemma 4 12B" }, .{ .model_type = "lfm2_moe", .title = "LFM2.5-24B-A2B" }, }; @@ -174,7 +176,8 @@ fn readSmall(a: Allocator, io: std.Io, path: []const u8) ?[]u8 { test "family detection matches the registry and refuses unknown types" { try std.testing.expectEqualStrings("Nemotron 3.5 Lightning", family("nemotron_h").?.title); try std.testing.expectEqualStrings("Qwen3.8 Flash Next", family("qwen4_exp").?.title); - try std.testing.expect(family("gemma4") == null); + try std.testing.expectEqualStrings("Gemma 4 12B", family("gemma4_unified").?.title); + try std.testing.expect(family("gemma3n") == null); try std.testing.expect(family("unknown") == null); } diff --git a/zig/src/cli/info.zig b/zig/src/cli/info.zig index c0d36b1abc..9303bef424 100644 --- a/zig/src/cli/info.zig +++ b/zig/src/cli/info.zig @@ -72,9 +72,9 @@ test "info prints family, format, context, weights and floor for a cached checkp var refused: std.Io.Writer.Allocating = .init(a); try w.createDirPath(io, try std.fs.path.join(a, &.{ root, "models--Org--Draft/snapshots/rev2" })); - try w.writeFile(io, .{ .sub_path = try std.fs.path.join(a, &.{ root, "models--Org--Draft/snapshots/rev2/config.json" }), .data = "{\"model_type\": \"gemma4\"}" }); + try w.writeFile(io, .{ .sub_path = try std.fs.path.join(a, &.{ root, "models--Org--Draft/snapshots/rev2/config.json" }), .data = "{\"model_type\": \"gemma3n\"}" }); try std.testing.expectEqual(@as(u8, 1), try run(a, io, &refused.writer, null, root, "Org/Draft")); - try std.testing.expect(has(refused.written(), "no registered Zig family serves model_type gemma4")); + try std.testing.expect(has(refused.written(), "no registered Zig family serves model_type gemma3n")); } fn has(text: []const u8, needle: []const u8) bool { diff --git a/zig/src/cli/models.zig b/zig/src/cli/models.zig index d04a773545..8d0f1040f7 100644 --- a/zig/src/cli/models.zig +++ b/zig/src/cli/models.zig @@ -49,7 +49,7 @@ test "models lists serve-able checkpoints and skips families Zig cannot serve" { try w.writeFile(io, .{ .sub_path = try std.fs.path.join(a, &.{ root, "models--Org--Flash/refs/main" }), .data = "rev1" }); try w.writeFile(io, .{ .sub_path = try std.fs.path.join(a, &.{ root, "models--Org--Flash/snapshots/rev1/config.json" }), .data = "{\"model_type\": \"qwen4_exp\", \"quantization\": {\"bits\": 6, \"group_size\": 32}}" }); try w.createDirPath(io, try std.fs.path.join(a, &.{ root, "models--Org--Draft/snapshots/rev2" })); - try w.writeFile(io, .{ .sub_path = try std.fs.path.join(a, &.{ root, "models--Org--Draft/snapshots/rev2/config.json" }), .data = "{\"model_type\": \"gemma4\"}" }); + try w.writeFile(io, .{ .sub_path = try std.fs.path.join(a, &.{ root, "models--Org--Draft/snapshots/rev2/config.json" }), .data = "{\"model_type\": \"gemma3n\"}" }); var out: std.Io.Writer.Allocating = .init(a); const code = try run(a, io, &out.writer, null, root); diff --git a/zig/src/cli/pull_test.zig b/zig/src/cli/pull_test.zig index cc330426c9..cf09b361df 100644 --- a/zig/src/cli/pull_test.zig +++ b/zig/src/cli/pull_test.zig @@ -153,7 +153,7 @@ const FakeHub = struct { } else if (std.mem.eql(u8, path, "/api/models/Org/Draft/tree/ccccccccccccccccccccccccccccccccccccccc3")) { respond(fd, "[{\"type\": \"file\", \"path\": \"config.json\", \"size\": 27}]"); } else if (std.mem.eql(u8, path, "/Org/Draft/resolve/ccccccccccccccccccccccccccccccccccccccc3/config.json")) { - respond(fd, "{\"model_type\": \"gemma4\"}"); + respond(fd, "{\"model_type\": \"gemma3n\"}"); } else if (std.mem.eql(u8, path, "/api/models/Org/DFlash/revision/main")) { respond(fd, "{\"sha\": \"ddddddddddddddddddddddddddddddddddddddd4\"}"); } else if (std.mem.eql(u8, path, "/api/models/Org/DFlash/tree/ddddddddddddddddddddddddddddddddddddddd4")) { @@ -306,7 +306,7 @@ test "pull downloads, verifies, retries a cut body, resumes and refuses a family var refused_out: std.Io.Writer.Allocating = .init(a); var refused_err: std.Io.Writer.Allocating = .init(a); try std.testing.expectEqual(@as(u8, 1), try pull.run(a, io, &refused_out.writer, &refused_err.writer, &env, root, "Org/Draft")); - try std.testing.expect(std.mem.indexOf(u8, refused_err.written(), "model_type gemma4") != null); + try std.testing.expect(std.mem.indexOf(u8, refused_err.written(), "model_type gemma3n") != null); try std.testing.expect(std.mem.indexOf(u8, refused_err.written(), "tensorfold@0.6") != null); // A DFlash2 drafter repo comes down for serve --drafter although no family serves its model_type. diff --git a/zig/src/core/affine_mm.zig b/zig/src/core/affine_mm.zig index 25d23a3ba9..912b352bd0 100644 --- a/zig/src/core/affine_mm.zig +++ b/zig/src/core/affine_mm.zig @@ -1,4 +1,4 @@ -//! x W^T for affine-quantized W on the tensor units, dense or by expert: exact codes, fp32 group sums and fp32 row sums. +//! x W^T for affine-quantized W on the tensor units, dense or by expert: each step's weights dequantized to bf16 in the tile. const std = @import("std"); const mtl = @import("metal"); const ks = @import("kernel_sources"); @@ -21,50 +21,35 @@ pub fn source(device: mtl.Device, a: std.mem.Allocator, f: Format) ![]u8 { return std.fmt.allocPrint(a, head, .{ f.bits, f.group, out, glu, body }); } -/// Entry points: row sums, the dense matmul, gathers by tile height, scattering gathers, the experts' gate and up through a pair order into SwiGLU. -pub const names = [8][:0]const u8{ "tf_affine_row_sums", "tf_affine_mm", "tf_affine_gather_32", "tf_affine_gather_64", "tf_affine_gather_scatter_32", "tf_affine_gather_scatter_64", "tf_affine_gather_glu_32", "tf_affine_gather_glu_64" }; +/// Entry points: the dense matmul, gathers by tile height, scattering gathers, the experts' gate and up through a pair order into SwiGLU. +pub const names = [7][:0]const u8{ "tf_affine_mm", "tf_affine_gather_32", "tf_affine_gather_64", "tf_affine_gather_scatter_32", "tf_affine_gather_scatter_64", "tf_affine_gather_glu_32", "tf_affine_gather_glu_64" }; /// A format's pipelines, in `names` order. -pub const Pipes = [8]mtl.Pipeline; +pub const Pipes = [7]mtl.Pipeline; pub const Args = extern struct { rows: i32, n: i32, k: i32, experts: i32 }; -/// A dense matmul's pitches and batch offsets, in elements (the sums' in rows); the default is one plain [rows, k] x [k, n] product. -pub const Strides = extern struct { x_row: i32 = 0, y_row: i32 = 0, sums_row: i32 = 1, x_batch: i32 = 0, y_batch: i32 = 0, sums_batch: i32 = 0, w_batch: i32 = 0, pad: i32 = 0 }; +/// A dense matmul's pitches and batch offsets, in elements (the weights' in rows); the default is one plain [rows, k] x [k, n] product. +pub const Strides = extern struct { x_row: i32 = 0, y_row: i32 = 0, x_batch: i32 = 0, y_batch: i32 = 0, w_batch: i32 = 0, pad: i32 = 0 }; fn bind(e: mtl.ComputeEncoder, i: usize, r: anytype) void { e.setBuffer(r.buf, r.off, i); } -/// Each row's sum over each group of x [rows, k] into `sums` (fp32 [rows, k / group]): the bias's operand. -pub fn rowSums(e: mtl.ComputeEncoder, p: Pipes, group: u32, x: anytype, sums: anytype, rows: u32, k: u32) void { - rowSumsAt(e, p, group, x, sums, rows, k, k); -} - -/// `rowSums` over rows `ld` elements apart (k of them each). -pub fn rowSumsAt(e: mtl.ComputeEncoder, p: Pipes, group: u32, x: anytype, sums: anytype, rows: u32, k: u32, ld: u32) void { - e.setPipeline(p[0]); - bind(e, 0, x); - e.setValue(Args{ .rows = @intCast(rows), .n = @intCast(ld), .k = @intCast(k), .experts = 0 }, 5); - bind(e, 6, sums); - e.dispatchThreads(mtl.Size.of(k / group, rows, 1), mtl.Size.of(@min(k / group, 64), 1, 1)); -} - -/// y [rows, n] = x [rows, q.k] W^T with x's group sums `sums`; n is q.n rounded up to 64 (the padded weight rows, y's pitch). -pub fn dense(e: mtl.ComputeEncoder, p: Pipes, x: anytype, sums: anytype, q: anytype, y: anytype, rows: u32) void { +/// y [rows, n] = x [rows, q.k] W^T; n is q.n rounded up to 64 (the padded weight rows, y's pitch). +pub fn dense(e: mtl.ComputeEncoder, p: Pipes, x: anytype, q: anytype, y: anytype, rows: u32) void { const n = std.mem.alignForward(u32, q.n, 64); - denseBatch(e, p, x, sums, q, y, rows, 1, .{ .x_row = @intCast(q.k), .y_row = @intCast(n) }); + denseBatch(e, p, x, q, y, rows, 1, .{ .x_row = @intCast(q.k), .y_row = @intCast(n) }); } /// `batch` dense products at once, each over its own rows of the weights (`s`; q.n a multiple of 64), e.g. a per-head projection. -pub fn denseBatch(e: mtl.ComputeEncoder, p: Pipes, x: anytype, sums: anytype, q: anytype, y: anytype, rows: u32, batch: u32, s: Strides) void { +pub fn denseBatch(e: mtl.ComputeEncoder, p: Pipes, x: anytype, q: anytype, y: anytype, rows: u32, batch: u32, s: Strides) void { const n = std.mem.alignForward(u32, q.n, 64); - e.setPipeline(p[1]); + e.setPipeline(p[0]); bind(e, 0, x); bind(e, 1, q.w); bind(e, 2, q.s); bind(e, 3, q.b); - bind(e, 4, sums); e.setValue(Args{ .rows = @intCast(rows), .n = @intCast(n), .k = @intCast(q.k), .experts = 0 }, 5); bind(e, 6, y); e.setValue(s, 7); @@ -72,22 +57,21 @@ pub fn denseBatch(e: mtl.ComputeEncoder, p: Pipes, x: anytype, sums: anytype, q: } /// y [rows, q.n] = x times each row's expert's W^T, rows sorted by expert with `offsets` (each expert's first row); q is [experts, n, k]. -pub fn gather(e: mtl.ComputeEncoder, p: Pipes, x: anytype, sums: anytype, q: anytype, offsets: anytype, y: anytype, rows: u32, experts: u32) void { - gatherTo(e, p, x, sums, q, offsets, y, null, rows, experts); +pub fn gather(e: mtl.ComputeEncoder, p: Pipes, x: anytype, q: anytype, offsets: anytype, y: anytype, rows: u32, experts: u32) void { + gatherTo(e, p, x, q, offsets, y, null, rows, experts); } /// `gather` with sorted row i's output at row dst[i] of y (i32 [rows]; null: in sorted order). -pub fn gatherTo(e: mtl.ComputeEncoder, p: Pipes, x: anytype, sums: anytype, q: anytype, offsets: anytype, y: anytype, dst: anytype, rows: u32, experts: u32) void { +pub fn gatherTo(e: mtl.ComputeEncoder, p: Pipes, x: anytype, q: anytype, offsets: anytype, y: anytype, dst: anytype, rows: u32, experts: u32) void { std.debug.assert(q.n % 64 == 0 and q.k % 64 == 0); const bm: u32 = if (rows / experts < 64) 32 else 64; const scatter = @TypeOf(dst) != @TypeOf(null); - e.setPipeline(p[@as(usize, if (bm == 64) 3 else 2) + if (scatter) 2 else 0]); + e.setPipeline(p[@as(usize, if (bm == 64) 2 else 1) + if (scatter) 2 else 0]); if (scatter) bind(e, 8, dst); bind(e, 0, x); bind(e, 1, q.w); bind(e, 2, q.s); bind(e, 3, q.b); - bind(e, 4, sums); e.setValue(Args{ .rows = @intCast(rows), .n = @intCast(q.n), .k = @intCast(q.k), .experts = @intCast(experts) }, 5); bind(e, 6, y); bind(e, 7, offsets); @@ -96,15 +80,14 @@ pub fn gatherTo(e: mtl.ComputeEncoder, p: Pipes, x: anytype, sums: anytype, q: a } /// act = SwiGLU(x W_gate^T, x W_up^T) clipped at `limit` for expert-sorted pairs (`offsets`), pair i on x's row order[i] / topk. -pub fn gatherGlu(e: mtl.ComputeEncoder, p: Pipes, x: anytype, sums: anytype, gate: anytype, up: anytype, offsets: anytype, order: anytype, act: anytype, rows: u32, experts: u32, topk: u32, limit: f32) void { +pub fn gatherGlu(e: mtl.ComputeEncoder, p: Pipes, x: anytype, gate: anytype, up: anytype, offsets: anytype, order: anytype, act: anytype, rows: u32, experts: u32, topk: u32, limit: f32) void { std.debug.assert(gate.n % 32 == 0 and gate.k % 64 == 0 and gate.n == up.n and gate.k == up.k); const bm: u32 = if (rows / experts < 200) 32 else 64; // 32-row tiles waste less of each expert's last tile until ~200 rows - e.setPipeline(p[if (bm == 64) 7 else 6]); + e.setPipeline(p[if (bm == 64) 6 else 5]); bind(e, 0, x); bind(e, 1, gate.w); bind(e, 2, gate.s); bind(e, 3, gate.b); - bind(e, 4, sums); e.setValue(Args{ .rows = @intCast(rows), .n = @intCast(gate.n), .k = @intCast(gate.k), .experts = @intCast(experts) }, 5); bind(e, 6, act); bind(e, 7, offsets); @@ -168,8 +151,8 @@ test "dense and gathered, each format, against exact sums on the host" { const w_bytes = E * N * K * bits / 8; const s_n = E * N * K / group; const opts = mtl.ResourceOptions.shared; - const bufs = [_]mtl.Buffer{ try device.buffer(R * K * 2, opts), try device.buffer(w_bytes, opts), try device.buffer(s_n * 2, opts), try device.buffer(s_n * 2, opts), try device.buffer(E * 4, opts), try device.buffer(5 * R * N * 4, opts), try device.buffer(R * (K / group) * 4, opts), try device.buffer(R * 4, opts) }; - const dst = bufs[7].slice(i32, R); // the scatter's destinations: rows reversed + const bufs = [_]mtl.Buffer{ try device.buffer(R * K * 2, opts), try device.buffer(w_bytes, opts), try device.buffer(s_n * 2, opts), try device.buffer(s_n * 2, opts), try device.buffer(E * 4, opts), try device.buffer(5 * R * N * 4, opts), try device.buffer(R * 4, opts) }; + const dst = bufs[6].slice(i32, R); // the scatter's destinations: rows reversed for (dst, 0..) |*d, i| d.* = @intCast(R - 1 - i); defer for (bufs) |b| b.deinit(); const x = bufs[0].slice(u16, R * K); @@ -198,17 +181,15 @@ test "dense and gathered, each format, against exact sums on the host" { } defer for (pipes) |pp| for (pp) |p| p.deinit(); const xr = .{ .buf = bufs[0], .off = @as(usize, 0) }; - const sums = .{ .buf = bufs[6], .off = @as(usize, 0) }; const q = .{ .w = .{ .buf = bufs[1], .off = @as(usize, 0) }, .s = .{ .buf = bufs[2], .off = @as(usize, 0) }, .b = .{ .buf = bufs[3], .off = @as(usize, 0) }, .n = @as(u32, N), .k = @as(u32, K) }; const out = [5]struct { buf: mtl.Buffer, off: usize }{ .{ .buf = bufs[5], .off = 0 }, .{ .buf = bufs[5], .off = R * N * 4 }, .{ .buf = bufs[5], .off = 2 * R * N * 4 }, .{ .buf = bufs[5], .off = 3 * R * N * 4 }, .{ .buf = bufs[5], .off = 4 * R * N * 4 } }; const cb = queue.commandBuffer(); const enc = cb.compute(.serial); - rowSums(enc, pipes[0], @intCast(group), xr, sums, R, K); - gather(enc, pipes[0], xr, sums, q, .{ .buf = bufs[4], .off = @as(usize, 0) }, out[0], R, E); // bf16 - gather(enc, pipes[1], xr, sums, q, .{ .buf = bufs[4], .off = @as(usize, 0) }, out[1], R, E); // fp32 - dense(enc, pipes[0], xr, sums, q, out[2], R); // expert 0's weights for every row, bf16 - dense(enc, pipes[1], xr, sums, q, out[3], R); // fp32 - gatherTo(enc, pipes[1], xr, sums, q, .{ .buf = bufs[4], .off = @as(usize, 0) }, out[4], .{ .buf = bufs[7], .off = @as(usize, 0) }, R, E); // fp32, rows reversed + gather(enc, pipes[0], xr, q, .{ .buf = bufs[4], .off = @as(usize, 0) }, out[0], R, E); // bf16 + gather(enc, pipes[1], xr, q, .{ .buf = bufs[4], .off = @as(usize, 0) }, out[1], R, E); // fp32 + dense(enc, pipes[0], xr, q, out[2], R); // expert 0's weights for every row, bf16 + dense(enc, pipes[1], xr, q, out[3], R); // fp32 + gatherTo(enc, pipes[1], xr, q, .{ .buf = bufs[4], .off = @as(usize, 0) }, out[4], .{ .buf = bufs[6], .off = @as(usize, 0) }, R, E); // fp32, rows reversed enc.end(); cb.commit(); cb.wait(); @@ -225,27 +206,14 @@ test "dense and gathered, each format, against exact sums on the host" { for ([2]usize{ ex, 0 }, 0..) |use, pass| for (0..N) |n| { const row = use * N + n; var sum: f64 = 0; - var mag: f64 = 0; - for (0..K / group) |g| { - var dot: f64 = 0; - var dmag: f64 = 0; // the sizes the fp32 sums carry: codes (4-bit ones plus 128) times x - var xs: f64 = 0; - var xmag: f64 = 0; - const off: f64 = if (bits == 4) 128 else 0; - for (0..group) |j| { - const k = g * group + j; - const qv: u32 = code(wb, (row * K + k) * bits, bits); - const xv: f64 = f32of(x[r * K + k]); - dot += @as(f64, @floatFromInt(qv)) * xv; - dmag += (@as(f64, @floatFromInt(qv)) + off) * @abs(xv); - xs += xv; - xmag += @abs(xv); - } - const sg = row * (K / group) + g; - const s: f64 = f32of(sc[sg]); - const b: f64 = f32of(bi[sg]); - sum += s * dot + b * xs; - mag += s * dmag + @abs(b - off * s) * xmag; + var mag: f64 = 0; // the sizes the fp32 sums carry: the weights as the tile rounds them, times x + for (0..K) |k| { + const sg = row * (K / group) + k / group; + const qv: f32 = @floatFromInt(code(wb, (row * K + k) * bits, bits)); + const w = f32of(bf16(@mulAdd(f32, @floatCast(f32of(sc[sg])), qv, @floatCast(f32of(bi[sg]))))); + const xv: f64 = f32of(x[r * K + k]); + sum += w * xv; + mag += @abs(w * xv); } const got = if (pass == 0) g32[r * N + n] else d32[r * N + n]; const got16 = if (pass == 0) g16[r * N + n] else d16[r * N + n]; @@ -272,7 +240,7 @@ test "a batch of dense products with pitches: each batch's rows of x, the weight var prng = std.Random.DefaultPrng.init(5); const rnd = prng.random(); const opts = mtl.ResourceOptions.shared; - const bufs = [_]mtl.Buffer{ try device.buffer(R * B * K * 2, opts), try device.buffer(B * N * K / 2, opts), try device.buffer(B * N * K / 64 * 2, opts), try device.buffer(B * N * K / 64 * 2, opts), try device.buffer(R * B * N * 2, opts), try device.buffer(R * B * (K / 64) * 4, opts) }; + const bufs = [_]mtl.Buffer{ try device.buffer(R * B * K * 2, opts), try device.buffer(B * N * K / 2, opts), try device.buffer(B * N * K / 64 * 2, opts), try device.buffer(B * N * K / 64 * 2, opts), try device.buffer(R * B * N * 2, opts) }; defer for (bufs) |b| b.deinit(); const x = bufs[0].slice(u16, R * B * K); for (x) |*v| v.* = bf16(rnd.float(f32) * 2 - 1); @@ -292,12 +260,11 @@ test "a batch of dense products with pitches: each batch's rows of x, the weight for (names, 0..) |n, i| pipes[i] = try mtl.Pipeline.init(device, lib, n, false); defer for (pipes) |p| p.deinit(); const xr = .{ .buf = bufs[0], .off = @as(usize, 0) }; - const sums = .{ .buf = bufs[5], .off = @as(usize, 0) }; const q = .{ .w = .{ .buf = bufs[1], .off = @as(usize, 0) }, .s = .{ .buf = bufs[2], .off = @as(usize, 0) }, .b = .{ .buf = bufs[3], .off = @as(usize, 0) }, .n = @as(u32, N), .k = @as(u32, K) }; const cb = queue.commandBuffer(); const enc = cb.compute(.serial); - rowSums(enc, pipes, 64, xr, sums, R * B, K); // x as [R * B, K]: batch b's row r is row r * B + b - denseBatch(enc, pipes, xr, sums, q, .{ .buf = bufs[4], .off = @as(usize, 0) }, R, B, .{ .x_row = B * K, .y_row = B * N, .sums_row = B, .x_batch = K, .y_batch = N, .sums_batch = 1, .w_batch = N }); + // x as [R * B, K]: batch b's row r is row r * B + b + denseBatch(enc, pipes, xr, q, .{ .buf = bufs[4], .off = @as(usize, 0) }, R, B, .{ .x_row = B * K, .y_row = B * N, .x_batch = K, .y_batch = N, .w_batch = N }); enc.end(); cb.commit(); cb.wait(); @@ -306,24 +273,15 @@ test "a batch of dense products with pitches: each batch's rows of x, the weight for (0..R) |r| for (0..B) |b| for (0..N) |n| { const row = b * N + n; var sum: f64 = 0; - var mag: f64 = 0; - for (0..K / 64) |g| { - var dot: f64 = 0; - var xs: f64 = 0; - var xmag: f64 = 0; - const xrow = (r * B + b) * K; - for (0..64) |j| { - const k = g * 64 + j; - const qv: u32 = (wb[(row * K + k) / 2] >> @intCast(4 * (k % 2))) & 15; - const xv: f64 = f32of(x[xrow + k]); - dot += @as(f64, @floatFromInt(qv)) * xv; - xs += xv; - xmag += @abs(xv); - } - const sg = row * (K / 64) + g; - const s: f64 = f32of(sc[sg]); - sum += s * dot + @as(f64, f32of(bi[sg])) * xs; - mag += s * 143 * xmag + @abs(f32of(bi[sg]) - 128 * s) * xmag; // the sizes the fp32 sums carry (codes plus 128) + var mag: f64 = 0; // the sizes the fp32 sums carry: the weights as the tile rounds them, times x + const xrow = (r * B + b) * K; + for (0..K) |k| { + const sg = row * (K / 64) + k / 64; + const qv: f32 = @floatFromInt((wb[(row * K + k) / 2] >> @intCast(4 * (k % 2))) & 15); + const w = f32of(bf16(@mulAdd(f32, @floatCast(f32of(sc[sg])), qv, @floatCast(f32of(bi[sg]))))); + const xv: f64 = f32of(x[xrow + k]); + sum += w * xv; + mag += @abs(w * xv); } const got: f64 = f32of(y[(r * B + b) * N + n]); try std.testing.expect(@abs(got - sum) <= @abs(sum) / 128.0 + 4e-6 * mag); // one bf16 rounding of the fp32 sum diff --git a/zig/src/core/contract/affine.zig b/zig/src/core/contract/affine.zig index 41e6ba0359..44874f035e 100644 --- a/zig/src/core/contract/affine.zig +++ b/zig/src/core/contract/affine.zig @@ -3,8 +3,8 @@ const std = @import("std"); const c = @import("contract.zig"); const rows = @import("rows.zig"); -/// The formula a kernel computes: MLX's per-group s*sum(q x) + b*sum(x), or each weight s*q + b first. -pub const Form = enum { factored, dequantized }; +/// The formula a kernel computes: MLX's s*sum(q x) + b*sum(x), each weight s*q + b first, or bf16(fma(s, q, b)) first. +pub const Form = enum { factored, dequantized, rounded }; /// Host inputs of y = x W^T: x [rows, k] bf16, codes [n, k], scales and biases [n, k / group] bf16. pub const Affine = struct { @@ -102,6 +102,11 @@ pub const Affine = struct { return c.fromBf16(f.x[r * f.shape.k + i]); } + /// A weight as `rounded` loads it: bf16(fma(s, q, b)) in fp32. + fn roundedWeight(p: [2]f64, q: f64) f64 { + return c.fromBf16(c.toBf16(@mulAdd(f32, @floatCast(p[0]), @floatCast(q), @floatCast(p[1])))); + } + fn sb(f: Affine, j: usize, g: usize) [2]f64 { const at = j * (f.shape.k / f.group) + g; return .{ c.fromBf16(f.scales[at]), c.fromBf16(f.biases[at]) }; @@ -127,12 +132,13 @@ pub const Affine = struct { xs += x; }, .dequantized => qx += (p[0] * q + p[1]) * x, + .rounded => qx += roundedWeight(p, q) * x, } - qa += @abs((q + f.off) * x); + qa += if (f.form == .rounded) @abs(roundedWeight(p, q) * x) else @abs((q + f.off) * x); xa += @abs(x); } y += if (f.form == .factored) p[0] * qx + p[1] * xs else qx; - mag += @abs(p[0]) * qa + @abs(p[1] - f.off * p[0]) * xa; + mag += if (f.form == .rounded) qa else @abs(p[0]) * qa + @abs(p[1] - f.off * p[0]) * xa; } return .{ y, mag }; } @@ -166,7 +172,10 @@ pub const Affine = struct { y.* = rows.sumIn(acc, .sequential, t[g * f.group ..][0..f.group]) * pg[0] + rows.sumIn(acc, .sequential, t[k + g * f.group ..][0..f.group]) * (pg[1] - f.off * pg[0]); } const grouped = rows.sumIn(acc, order, gs); - for (0..k) |i| t[i] = (p[0] * @as(f64, @floatFromInt(f.codes[j * k + i])) + p[1]) * f.xv(r, i); + for (0..k) |i| { + const q: f64 = @floatFromInt(f.codes[j * k + i]); + t[i] = (if (f.form == .rounded) roundedWeight(p, q) else p[0] * q + p[1]) * f.xv(r, i); + } const direct = rows.sumIn(acc, order, t[0..k]); var best = whole; for ([_]f64{ grouped, direct }) |v| if (!std.math.isFinite(best) or @abs(v - want) < @abs(best - want)) { diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b79337b5c5..f8fe6993e5 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -80,9 +80,8 @@ pub const metal = [_]M{ .{ .name = "row projection q6 g32 tf_row_projection_q6_f32", .threads = 64, .tg_memory = 0, .m1 = 512, .m2 = 512 }, .{ .name = "row projection q8 g64 tf_row_projection_q8_f32", .threads = 64, .tg_memory = 0, .m1 = 512, .m2 = 512 }, .{ .name = "row projection q3 g32 tf_row_projection_q3_f32", .threads = 64, .tg_memory = 0, .m1 = 512, .m2 = 512 }, - .{ .name = "affine_mm q4 g64 tf_affine_row_sums", .threads = 64, .tg_memory = 0, .m1 = 1024, .m2 = 1024 }, - .{ .name = "affine_mm q4 g64 tf_affine_mm", .threads = 128, .tg_memory = 9728, .offline = false }, - .{ .name = "affine_mm q8 g32 tf_affine_mm", .threads = 128, .tg_memory = 10240, .offline = false }, + .{ .name = "affine_mm q4 g64 tf_affine_mm", .threads = 128, .tg_memory = 9216, .offline = false }, + .{ .name = "affine_mm q8 g32 tf_affine_mm", .threads = 128, .tg_memory = 9216, .offline = false }, .{ .name = "simd_qmm tf_sq_mma_16_4_1", .threads = 512, .tg_memory = 16384, .m1 = 768, .m2 = 768, .adapts = true }, .{ .name = "simd_qmm tf_sq_mma_16_2_2", .threads = 512, .tg_memory = 16384, .m1 = 704, .m2 = 704, .adapts = true }, .{ .name = "simd_qmm tf_sq_mma_8_4_1", .threads = 256, .tg_memory = 8192, .m1 = 768, .m2 = 768, .adapts = true }, diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 243b3386a4..74f614e92c 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -121,10 +121,10 @@ pub const Vision = struct { image_token: u32, /// The tokens the template wraps an image in (the image token among them): text that spells one stays text. markers: []const u32 = &.{}, + part_type: []const u8 = "image_url", // the content part type the template reads: OpenAI's, or "image" (most HF ones) /// Tokens the server writes around each image's rows (Nemotron Omni's and ); null: the template's own. wrap: ?[2]u32 = null, - /// Image tokens one request may hold in all, and images. - image_tokens: u32, + image_tokens: u32, // the image tokens one request may hold in all, and its images max_images: u32, /// Video input, when the engine's family reads it. video: ?VideoOffer = null, diff --git a/zig/src/core/lane_projection.zig b/zig/src/core/lane_projection.zig index 001b684892..441367a9ea 100644 --- a/zig/src/core/lane_projection.zig +++ b/zig/src/core/lane_projection.zig @@ -9,6 +9,17 @@ pub const reg_layout = "u32[N/32][K/64][2][32 lanes][4] 4-bit codes, regBlock or pub const metadata_layout = "bf16[K/group][N][scale,bias]"; pub const Output = enum { bf16, f32 }; +/// A family's Metal text around the reg kernel: its tf_prologue, and with `epilogue` its tf_epilogue (the .metal). +pub const Fused = struct { + source: []const u8, + defines: []const u8 = "", + /// Rows the prologue makes in threadgroup memory (0: X stays in device memory, the prologue writes the sums only). + x_rows: usize = 1, + epilogue: bool = false, + /// Threads a threadgroup takes (a multiple of 32 * sk, its tiles sharing the prologue): a reduction's fixed count. + threads: usize = 256, +}; + pub const Format = struct { bits: u8, group: usize, @@ -36,6 +47,8 @@ pub const Layout = struct { reg: bool = false, /// Elements between output rows (default N): a member of a fused stack writes its columns of the fused rows. y_stride: ?usize = null, + /// The reg kernel with a family's prologue and epilogue around it (every tensor-op word the same). + fused: ?Fused = null, pub fn validate(l: Layout) !void { try l.format.validate(); @@ -46,6 +59,7 @@ pub const Layout = struct { if (l.n == 0 or l.n % 16 != 0 or l.k == 0 or l.k % l.format.group != 0 or l.n > std.math.maxInt(i32) / 16 or l.k > std.math.maxInt(i32) / 16) return error.InvalidShape; if (l.sk == 0 or l.sk > 32 or (l.pf != 1 and l.pf != 2)) return error.InvalidSchedule; + if (l.fused) |f| if (!l.reg or f.threads == 0 or f.threads > 1024 or f.threads % (32 * l.sk) != 0 or l.tileCount() % l.tilesPerGroup() != 0 or f.x_rows > 16) return error.InvalidFusedShape; const g = l.groupCut(); if (g[1] == 0 or g[0] >= l.k / l.format.group or g[1] > l.k / l.format.group - g[0]) return error.InvalidGroupCut; if (l.groups != null and l.output != .f32) return error.PartialOutputNeedsFloat; @@ -76,7 +90,14 @@ pub const Layout = struct { } pub fn threadgroupBytes(l: Layout) usize { - return l.sk * 32 * l.format.group + @max(l.sk -| 1, 1) * 16 * 32 * @sizeOf(f32); + const part = l.tilesPerGroup() * @max(l.sk -| 1, 1) * 16 * 32 * @sizeOf(f32); + if (l.fused) |f| return part + f.x_rows * l.k * 2 + (l.k / l.format.group + 16) * 4 + 16 * 4; + return if (l.reg) part else l.sk * 32 * l.format.group + part; // the reg kernel holds no code tile + } + + /// Tiles a threadgroup takes: one, or for a fused pipeline enough for its threads to share the prologue. + pub fn tilesPerGroup(l: Layout) usize { + return if (l.fused) |f| f.threads / (32 * l.sk) else 1; } pub fn weightWords(l: Layout) usize { @@ -115,6 +136,11 @@ pub fn source(a: std.mem.Allocator, l: Layout) ![]u8 { try text.print(a, "#define TF_PRECOMPUTE_SUMS {d}\n", .{@intFromBool(l.precompute_sums)}); if (l.reg) try text.appendSlice(a, "#define TF_REG 1\n"); if (l.y_stride) |s| try text.print(a, "#define TF_YS {d}\n", .{s}); + if (l.fused) |f| { + try text.print(a, "#define TF_FUSED 1\n#define TF_TPT {d}\n#define TF_XROWS {d}\n#define TF_THREADS {d}\n{s}{s}", .{ l.tilesPerGroup(), f.x_rows, f.threads, if (f.epilogue) "#define TF_EPILOGUE 1\n" else "", f.defines }); + try text.appendSlice(a, "struct tf_fused { const device bfloat* a0; const device bfloat* a1; const device bfloat* a2; const device bfloat* a3; const device bfloat* a4; const device bfloat* a5; device bfloat* o0; device bfloat* o1; float4 f; int4 i; };\n"); + try text.appendSlice(a, f.source); + } try text.appendSlice(a, ks.core_lane_projection); return text.toOwnedSlice(a); } @@ -232,6 +258,10 @@ pub const Projection = struct { sums_buffer: ?mtl.Buffer = null, threads_per_slice: usize = 32, reg: bool = false, + /// Tiles a threadgroup takes, its threads, and the rows a fused encode takes (0: not fused). + tpt: usize = 1, + threads: usize = 32, + fused_rows: usize = 0, pub fn init(a: std.mem.Allocator, device: mtl.Device, l: Layout) !Projection { try l.validate(); @@ -252,7 +282,7 @@ pub const Projection = struct { errdefer sums_pipe.?.deinit(); sums_buffer = try device.buffer(@as(usize, if (l.reg) reg_rows else 16) * (l.k / l.format.group) * 4, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); } - return .{ .pipe = p, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .y_stride = l.y_stride orelse l.n, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg }; + return .{ .pipe = p, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .y_stride = l.y_stride orelse l.n, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg, .tpt = l.tilesPerGroup(), .threads = if (l.fused) |f| f.threads else 32 * l.sk, .fused_rows = if (l.fused) |f| (if (f.x_rows > 0) f.x_rows else 16) else 0 }; } pub fn deinit(p: Projection) void { @@ -274,7 +304,30 @@ pub const Projection = struct { return p.encodeImpl(e, x, w, sb, sums, y, rows); } + /// A fused pipeline's one encode: the prologue's inputs, the outputs it writes, and its constants. + pub fn encodeFused(p: Projection, e: mtl.ComputeEncoder, x: Ref, w: Ref, sb: Ref, y: Ref, args: FusedArgs, rows: u32) !void { + if (p.fused_rows == 0) return error.NotFused; + if (rows == 0 or rows > p.fused_rows) return error.InvalidRows; + try checkRef(x, rows * p.k * 2, 2); + try checkRef(w, p.weight_bytes, 16); + try checkRef(sb, p.metadata_bytes, 16); + try checkRef(y, ((rows - 1) * p.y_stride + p.n) * p.output_bytes, p.output_bytes); + e.setPipeline(p.pipe); + e.setBuffer(x.buf, x.off, 0); + e.setBuffer(w.buf, w.off, 1); + e.setBuffer(sb.buf, sb.off, 2); + e.setValue([2]i32{ @intCast(rows), 1 }, 3); + e.setBuffer(y.buf, y.off, 4); + e.setBuffer(x.buf, x.off, 5); // the sums come from the prologue + for (args.a, 6..) |r, i| e.setBuffer(r.buf, r.off, i); + for (args.o, 12..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(args.f, 14); + e.setValue(args.i, 15); + e.dispatchGroups(mtl.Size.of(p.tile_count / p.tpt, 1, 1), mtl.Size.of(p.threads, 1, 1)); + } + fn encodeImpl(p: Projection, e: mtl.ComputeEncoder, x: Ref, w: Ref, sb: Ref, given: ?Sums, y: Ref, rows: u32) !void { + if (p.fused_rows != 0) return error.FusedPipeline; if (rows == 0 or rows > @as(u32, if (p.reg) reg_rows else 16)) return error.InvalidRows; const blocks = (rows + 15) / 16; const stride: u32 = if (given) |g| g.stride else 16 * blocks; @@ -309,6 +362,9 @@ pub const reg_rows = 128; /// Group sums of a projection's input: XS[group * stride + row], fp32, each group's values added in order from zero. pub const Sums = struct { ref: Ref, stride: u32 }; +/// What a fused encode binds for its prologue and epilogue: six inputs, two outputs, four floats and four ints. +pub const FusedArgs = struct { a: [6]Ref, o: [2]Ref, f: [4]f32 = .{ 0, 0, 0, 0 }, i: [4]i32 = .{ 0, 0, 0, 0 } }; + fn checkRef(r: Ref, need: usize, alignment: usize) !void { if (r.off % alignment != 0) return error.MisalignedBuffer; const len = r.buf.length(); diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 8378c44900..ee826d79e4 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -45,6 +45,12 @@ pub const Features = struct { ready: Fence = .none, // .none: written; else the fence to wait on before reading, kind naming the backend }; +/// Takes a prompt chunk's tapped rows, the first at position `start`, while they are still held. +pub const FeatureSink = struct { + ptr: *anyopaque, + take: *const fn (ptr: *anyopaque, s: *Stream, f: Features, start: u64) anyerror!void, +}; + /// Absorb a stream's kept rows into its draft head and hold `depth` drafts for its next round. pub const DraftRequest = struct { stream: *Stream, @@ -92,6 +98,8 @@ pub const Backend = struct { prepare_features: ?*const fn (ptr: *anyopaque, taps: []const u32) anyerror!void = null, /// Cache rows [start, start + count) of the prepared taps, within what `Features` says is held (null: none). features: ?*const fn (ptr: *anyopaque, s: *Stream, taps: []const u32, start: u64, count: u32) anyerror!Features = null, + /// The prompt pass handing `sink` each chunk's tapped rows as it makes them (null: features() after the pass). + prefill_features: ?*const fn (ptr: *anyopaque, s: *Stream, taps: []const u32, sink: FeatureSink) anyerror!void = null, /// The first token's logprob row (the prompt's last row) for a stream with `logprobs`; null: not given. first_row: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!Row = null, /// `first` with the first token drawn among `mask`'s allowed tokens (structured output). diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig index 3bdffe7c7a..27d45d6ce0 100644 --- a/zig/src/core/lanes/drafted.zig +++ b/zig/src/core/lanes/drafted.zig @@ -44,7 +44,7 @@ pub const Drafted = struct { } pub fn backend(x: *Drafted) be.Backend { - return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .first_masked = firstMasked, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release } }; + return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .first_masked = firstMasked, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .held = held, .release = release } }; } /// The target's facts with the drafter's depth, step cost, prior, plain guard and batching; chains drafted late. @@ -76,6 +76,13 @@ pub const Drafted = struct { /// The target's prompt pass, then the drafter absorbs its rows but the last; a later failure releases both. fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void { const x = self(ptr); + if (s.drafts) if (x.target.vtable.prefill_features) |streamed| { // chunk by chunk: no whole prompt's taps held + try x.lanes.ensureUnusedCapacity(x.gpa, 1); + try x.drafter.open(s); + x.lanes.putAssumeCapacity(s, .{}); + errdefer x.dropStream(s); + return streamed(x.target.ptr, s, x.drafter.taps(), .{ .ptr = x, .take = take }); + }; try x.target.prefill(s); if (!s.drafts) return; errdefer x.dropStream(s); @@ -90,6 +97,17 @@ pub const Drafted = struct { try x.drafter.absorb(&.{.{ .stream = s, .features = f, .start = from, .follow = ids[from + 1 ..] }}); } + /// A prompt chunk's rows but the prompt's last (absorbed with the first token, as `prefill` leaves it). + fn take(ptr: *anyopaque, s: *Stream, f: be.Features, start: u64) anyerror!void { + const x = self(ptr); + const ids = s.prompt(); + const last = ids.len - 1; + if (start >= last) return; + var rows = f; + rows.rows = @intCast(@min(f.rows, last - start)); + try x.drafter.absorb(&.{.{ .stream = s, .features = rows, .start = start, .follow = ids[@intCast(start + 1)..] }}); + } + fn first(ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64 { return self(ptr).target.first(s, position); } @@ -142,6 +160,17 @@ pub const Drafted = struct { return x.target.verify(x.windows.items, out); } + /// One stream's held drafts as host tokens for a grammar to walk (fillers: the pending token, kept for the verify). + fn held(ptr: *anyopaque, s: *Stream, out: []u32) anyerror!void { + const x = self(ptr); + const lane = x.lanes.getPtr(s) orelse return error.NotDrafting; + if (lane.filler > 0) { + @memset(out, s.pending.?); + return; + } + try x.drafter.held(&.{s}, &.{out}); + } + /// The target keeps from the windows it verified (host tokens): the round loop's, position for position. fn keep(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { const x = self(ptr); diff --git a/zig/src/core/lanes/drafted_grammar_test.zig b/zig/src/core/lanes/drafted_grammar_test.zig new file mode 100644 index 0000000000..468aed6bcf --- /dev/null +++ b/zig/src/core/lanes/drafted_grammar_test.zig @@ -0,0 +1,98 @@ +//! Structured output through the drafts wrapper: the grammar walks held drafts, and the verify sees what it walked. +const std = @import("std"); +const fake = @import("fake.zig"); +const be = @import("backend.zig"); +const grammar = @import("grammar.zig"); +const Sampling = @import("sampling.zig").Sampling; +const dt = @import("drafted_test.zig"); + +test "a grammar's read of fillers leaves them for the verify: both see the pending token" { + var r: dt.Restored = undefined; + const b = try r.init(); + defer r.deinit(); + const s = &r.s; + const n: u64 = dt.long.len; + const t = try b.read(try b.first(s, n)); + s.pending = t; + try b.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = t }, .rows = null, .start = n, .position = n + 1, .depth = 3 }}); + var walked: [3]u32 = undefined; + try b.vtable.held.?(b.ptr, s, &walked); + try std.testing.expectEqualSlices(u32, &.{ t, t, t }, &walked); + // the verify of the same held window reads the same fillers, not the head's empty chain + var drawn: [4]u32 = undefined; + var echo: [3]u32 = undefined; + const w = [_]be.Window{.{ .stream = s, .pending = t, .held = 3, .tokens = &.{}, .parents = null, .positions = &.{ n + 1, n + 2, n + 3, n + 4 } }}; + var o = [_]be.Verified{.{ .sampled = &drawn, .drafts = &echo }}; + try b.verify(&w, &o); + try std.testing.expectEqual(@as(u32, 0), r.head.helds); + try std.testing.expectEqualSlices(u32, &walked, &echo); +} + +/// A test grammar: `length` odd tokens, then 96. +const Odd = struct { + length: u32 = 12, + at: u32 = 0, + done: bool = false, + + fn rules(o: *Odd) grammar.Rules { + return .{ .ptr = o, .vtable = &.{ .accept = accept, .rollback = rollback, .terminated = terminated, .fill = fill, .drafts = drafts, .free = free } }; + } + + fn cast(ptr: *anyopaque) *Odd { + return @ptrCast(@alignCast(ptr)); + } + + fn allows(o: *const Odd, t: u32) bool { + if (o.done) return false; + return if (o.at == o.length) t == 96 else t % 2 == 1; + } + + fn accept(ptr: *anyopaque, t: u32) anyerror!bool { + const o = cast(ptr); + if (!o.allows(t)) return false; + if (o.at == o.length) o.done = true else o.at += 1; + return true; + } + + fn rollback(ptr: *anyopaque, n: usize) anyerror!void { + const o = cast(ptr); + for (0..n) |_| { + if (o.done) o.done = false else o.at -= 1; + } + } + + fn terminated(ptr: *anyopaque) bool { + return cast(ptr).done; + } + + fn fill(ptr: *anyopaque, words: []u32) anyerror!void { + const o = cast(ptr); + @memset(words, 0); + for (0..fake.vocab) |t| if (o.allows(@intCast(t))) { + words[t / 32] |= @as(u32, 1) << @intCast(t % 32); + }; + } + + fn drafts(_: *anyopaque, _: []u32) anyerror!usize { + return 0; + } + + fn free(_: *anyopaque) void {} +}; + +test "structured output through the wrapper: the grammar walks the drafter's held drafts, and drafted == plain" { + for ([_]?Sampling{ null, dt.sampled }) |s| { + var outs: [2]dt.Out = undefined; + for (&outs, [_]dt.Mode{ .wrapped, .plain }) |*o, mode| { + var g: Odd = .{}; + var k: grammar.Constraint = .init(g.rules(), (fake.vocab + 31) / 32, null); + o.* = try dt.run(&.{.{ .prompt = &dt.p1, .sampling = s, .drafts = mode == .wrapped, .grammar = &k }}, mode, false); + } + defer for (outs) |o| dt.freeOut(o); + try std.testing.expectEqualSlices(u32, outs[1].tokens[0], outs[0].tokens[0]); + try std.testing.expect(outs[0].head.helds > 0 and outs[0].drafted > 0); + var g: Odd = .{}; + for (outs[0].tokens[0]) |t| try std.testing.expect(try Odd.accept(&g, t)); + try std.testing.expect(g.done); + } +} diff --git a/zig/src/core/lanes/drafted_test.zig b/zig/src/core/lanes/drafted_test.zig index ec4003297c..ff37253c3b 100644 --- a/zig/src/core/lanes/drafted_test.zig +++ b/zig/src/core/lanes/drafted_test.zig @@ -11,11 +11,12 @@ const dr = @import("drafter.zig"); const Drafted = @import("drafted.zig").Drafted; const SuffixLookup = @import("proposer.zig").SuffixLookup; const Sampling = @import("sampling.zig").Sampling; +const grammar = @import("grammar.zig"); const gpa = std.testing.allocator; /// Rebuilds a stream's history from the target's features and drafts its next tokens, now and then wrong. -const FakeDrafter = struct { +pub const FakeDrafter = struct { lanes: std.AutoHashMapUnmanaged(*sm.Stream, Lane) = .empty, fail: enum { none, open, absorb } = .none, facts_: dr.Facts = .{ .depth = 6, .step_ms = 0.5 }, @@ -132,7 +133,7 @@ const FakeDrafter = struct { }; /// The fake target holding features as `Features` promises (a keep overwrites what it drops), refusing anything else. -const Strict = struct { +pub const Strict = struct { inner: be.Backend, snaps: std.AutoHashMapUnmanaged(*sm.Stream, Snap) = .empty, refused: u32 = 0, @@ -149,7 +150,7 @@ const Strict = struct { } fn backend(x: *Strict) be.Backend { - return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draftNone, .prepare_features = prepare, .features = features, .release = release } }; + return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .first_masked = firstMasked, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draftNone, .prepare_features = prepare, .features = features, .release = release } }; } fn self(ptr: *anyopaque) *Strict { @@ -176,6 +177,11 @@ const Strict = struct { return self(ptr).inner.first(s, position); } + fn firstMasked(ptr: *anyopaque, s: *sm.Stream, position: u64, mask: []const u32) anyerror!u64 { + const in = self(ptr).inner; + return in.vtable.first_masked.?(in.ptr, s, position, mask); + } + fn queue(ptr: *anyopaque, s: *sm.Stream, feed: be.Feed, position: u64) anyerror!u64 { const x = self(ptr); (try x.snap(s)).valid = 0; @@ -242,9 +248,9 @@ const Strict = struct { } }; -const Mode = enum { plain, wrapped }; -const Case = struct { prompt: []const u32, max_new: u32 = 40, sampling: ?Sampling = null, drafts: bool = true, restored: u32 = 0 }; -const Out = struct { +pub const Mode = enum { plain, wrapped }; +const Case = struct { prompt: []const u32, max_new: u32 = 40, sampling: ?Sampling = null, drafts: bool = true, restored: u32 = 0, grammar: ?*grammar.Constraint = null }; +pub const Out = struct { tokens: [][]u32, drafted: u64 = 0, accepted: u64 = 0, @@ -260,7 +266,7 @@ fn base() Model { return .{ .exact_width = 16, .gpu_tokens = true, .window_costs = &costs, .hidden_rows = true, .batch_rows = 32, .max_streams = 8 }; } -fn run(cases: []const Case, mode: Mode, batched: bool) !Out { +pub fn run(cases: []const Case, mode: Mode, batched: bool) !Out { var target: fake.Fake = .{ .gpa = gpa }; defer target.deinit(); var head: FakeDrafter = .{ .facts_ = .{ .depth = 6, .step_ms = 0.5, .batched = batched } }; @@ -291,7 +297,7 @@ fn run(cases: []const Case, mode: Mode, batched: bool) !Out { b.release(&early); } p.* = try SuffixLookup.init(gpa, .{ .min_match = 4 }); - s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = c.prompt, .max_new = c.max_new, .eos = &.{96}, .sampling = c.sampling, .drafts = c.drafts, .proposer = p.proposer(), .reuse = .{ .saved = k.*, .at = c.restored } }); + s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = c.prompt, .max_new = c.max_new, .eos = &.{96}, .sampling = c.sampling, .drafts = c.drafts, .proposer = p.proposer(), .reuse = .{ .saved = k.*, .at = c.restored }, .grammar = c.grammar }); } defer for (streams, proposers) |*s, *p| { s.deinit(gpa); @@ -310,15 +316,15 @@ fn run(cases: []const Case, mode: Mode, batched: bool) !Out { return out; } -fn freeOut(o: Out) void { +pub fn freeOut(o: Out) void { for (o.tokens) |t| gpa.free(t); gpa.free(o.tokens); } -const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5 }; +pub const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5 }; const p2 = [_]u32{ 2, 7, 1, 8, 2, 8, 1, 8, 2, 8, 4, 5, 9 }; -const long = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2, 3, 8, 4 }; -const sampled: Sampling = .{ .seed = 5, .temperature = 0.7, .top_k = 0, .top_p = 0.95 }; +pub const long = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2, 3, 8, 4 }; +pub const sampled: Sampling = .{ .seed = 5, .temperature = 0.7, .top_k = 0, .top_p = 0.95 }; test "an external drafter's rounds commit the one-token decode, greedy and sampled, and its drafts land" { for ([_]?Sampling{ null, sampled }) |s| { @@ -459,48 +465,68 @@ test "the strict target holds only what Features promises: a keep drops the reje try std.testing.expectError(error.RowsNotHeld, b.features(&s, &.{49}, p1.len, 3)); } +/// The wrapper over the strict target, streaming `long` restored whole from a kept state: the head has no row yet. +pub const Restored = struct { + target: fake.Fake = .{ .gpa = gpa }, + strict: Strict = undefined, + head: FakeDrafter = .{}, + wrapped: Drafted = undefined, + saved: *anyopaque = undefined, + s: sm.Stream = undefined, + + pub fn init(r: *Restored) !be.Backend { + r.* = .{}; + r.strict = .{ .inner = r.target.backend() }; + r.wrapped = try Drafted.init(gpa, r.strict.backend(), r.head.drafter()); + var early = try sm.Stream.init(gpa, .{ .id = "early", .prompt = &long, .max_new = 1, .eos = &.{96} }); + defer early.deinit(gpa); + try r.target.backend().prefill(&early); + r.saved = try r.target.save(&early, long.len); + r.target.backend().release(&early); + r.s = try sm.Stream.init(gpa, .{ .id = "s", .prompt = &long, .max_new = 32, .eos = &.{96}, .reuse = .{ .saved = r.saved, .at = long.len } }); + const b = r.wrapped.backend(); + try b.prefill(&r.s); + return b; + } + + pub fn deinit(r: *Restored) void { + r.wrapped.backend().release(&r.s); + r.s.deinit(gpa); + r.target.drop(r.saved); + r.wrapped.deinit(); + r.head.deinit(); + r.strict.deinit(); + r.target.deinit(); + } +}; + test "a copy round after a whole restore does not leave fillers in place of the drafter's first real drafts" { - var target: fake.Fake = .{ .gpa = gpa }; - defer target.deinit(); - var strict: Strict = .{ .inner = target.backend() }; - defer strict.deinit(); - var head: FakeDrafter = .{}; - defer head.deinit(); - var wrapped = try Drafted.init(gpa, strict.backend(), head.drafter()); - defer wrapped.deinit(); - const b = wrapped.backend(); - // the whole prompt as a kept state - var early = try sm.Stream.init(gpa, .{ .id = "early", .prompt = &long, .max_new = 1, .eos = &.{96} }); - defer early.deinit(gpa); - try target.backend().prefill(&early); - const saved = try target.save(&early, long.len); - defer target.drop(saved); - target.backend().release(&early); - var s = try sm.Stream.init(gpa, .{ .id = "s", .prompt = &long, .max_new = 32, .eos = &.{96}, .reuse = .{ .saved = saved, .at = long.len } }); - defer s.deinit(gpa); - try b.prefill(&s); - defer b.release(&s); + var r: Restored = undefined; + const b = try r.init(); + defer r.deinit(); + const s = &r.s; + const head = &r.head; const n: u64 = long.len; - const t = try b.read(try b.first(&s, n)); + const t = try b.read(try b.first(s, n)); // the head is asked for 3 drafts but has no row: fillers - try b.draft(&.{.{ .stream = &s, .follow = &.{}, .first = .{ .value = t }, .rows = null, .start = n, .position = n + 1, .depth = 3 }}); + try b.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = t }, .rows = null, .start = n, .position = n + 1, .depth = 3 }}); // the round loop takes a copy proposal instead: the fillers are never asked for const copy = [_]u32{ 7, 8 }; var drawn: [3]u32 = undefined; var echo: [2]u32 = undefined; - const w1 = [_]be.Window{.{ .stream = &s, .pending = t, .held = 0, .tokens = ©, .parents = null, .positions = &.{ n + 1, n + 2, n + 3 } }}; + const w1 = [_]be.Window{.{ .stream = s, .pending = t, .held = 0, .tokens = ©, .parents = null, .positions = &.{ n + 1, n + 2, n + 3 } }}; var o1 = [_]be.Verified{.{ .sampled = &drawn, .drafts = &echo }}; try b.verify(&w1, &o1); try b.keep(&w1, &.{&.{ 0, 1 }}); // the next request has real rows: the drafter holds real drafts - try b.draft(&.{.{ .stream = &s, .follow = &.{ copy[0], drawn[1] }, .rows = &.{ 0, 1 }, .start = n, .position = n + 3, .depth = 3 }}); + try b.draft(&.{.{ .stream = s, .follow = &.{ copy[0], drawn[1] }, .rows = &.{ 0, 1 }, .start = n, .position = n + 3, .depth = 3 }}); try std.testing.expectEqual(@as(u32, 1), head.holds); var drawn2: [4]u32 = undefined; var echo2: [3]u32 = undefined; - const w2 = [_]be.Window{.{ .stream = &s, .pending = drawn[1], .held = 3, .tokens = &.{}, .parents = null, .positions = &.{ n + 3, n + 4, n + 5, n + 6 } }}; + const w2 = [_]be.Window{.{ .stream = s, .pending = drawn[1], .held = 3, .tokens = &.{}, .parents = null, .positions = &.{ n + 3, n + 4, n + 5, n + 6 } }}; var o2 = [_]be.Verified{.{ .sampled = &drawn2, .drafts = &echo2 }}; try b.verify(&w2, &o2); // the verify read the drafter (not stale fillers) and the target saw its drafts try std.testing.expectEqual(@as(u32, 1), head.helds); - try std.testing.expectEqualSlices(u32, head.lanes.getPtr(&s).?.held.items[0..3], &echo2); + try std.testing.expectEqualSlices(u32, head.lanes.getPtr(s).?.held.items[0..3], &echo2); } diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index e943bf9d46..a159257baf 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -44,6 +44,7 @@ test { _ = @import("recorder_test.zig"); _ = @import("gpu_full_test.zig"); _ = @import("drafted_test.zig"); + _ = @import("drafted_grammar_test.zig"); _ = @import("pricing_test.zig"); } diff --git a/zig/src/families/gemma4/assistant.zig b/zig/src/families/gemma4/assistant.zig new file mode 100644 index 0000000000..fa75acbf1d --- /dev/null +++ b/zig/src/families/gemma4/assistant.zig @@ -0,0 +1,535 @@ +//! Gemma 4's MTP assistant as a drafter: a stream's last kept row, then chains over the target's own last caches. +const std = @import("std"); +const mtl = @import("metal"); +const lanes = @import("lanes"); +const cfg = @import("config.zig"); +const aw = @import("assistant_weights.zig"); +const kernels = @import("kernels.zig"); +const weights = @import("weights.zig"); +const st = @import("state.zig"); +const fwd = @import("forward.zig"); +const backend = @import("backend.zig"); +const round = @import("tree_round_gpu"); +const lane = @import("../../core/lane_projection.zig"); +const Model = @import("model.zig").Model; + +const dr = lanes.drafter; +const Stream = lanes.Stream; +const Ref = lane.Ref; +const Pick = round.contract.Pick; +const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const private = mtl.ResourceOptions.private | mtl.ResourceOptions.untracked; + +/// Drafts a round at most (the reference's default; the head was trained for 3). +pub const depth = 6; +/// Streams one hold or absorb call takes. +const max_items = 32; +/// Keys a gemma_decode_attention threadgroup takes (its CD), and a sliding partial's (the sliding kind's CK). +const global_chunk = 256; +const sliding_chunk = 128; + +/// One stream's assistant state: the target's final-normed last kept row, its position, the token after it. +const Lane = struct { + h: mtl.Buffer, + position: u64 = 0, + token: u32 = 0, + rows: bool = false, // a row was absorbed + held: [depth]u32 = @splat(0), + held_len: u32 = 0, +}; + +/// A stream's chain in a hold: its lane and the target cache it reads, from the row's position. +const Item = struct { lane: *Lane, cache: *const st.Cache, pending: u32, position: u64, depth: u32 }; + +pub const Drafter = struct { + gpa: std.mem.Allocator, + io: std.Io, + target: *Model, + metal: *backend.Metal, + config: aw.Config, + w: aw.Weights, + /// The assistant's glue at its own width: sliding and global kinds (q-only qkv_prep, the target's key heads). + kinds: [2]kernels.Kind, + projections: weights.Projections, + lanes: std.AutoHashMapUnmanaged(*Stream, *Lane) = .empty, + /// The target's layers each kind reads: its last sliding and its last global layer. + sources: [2]usize, + tap: [1]u32, + step_ms: f64 = 0, + chain_ms: [2]f64 = .{ 0, 0 }, // a 1-draft and a 6-draft chain's wall time at load + trace: bool = false, // GEMMA_DRAFT_TRACE: each hold's rows and drafts on stderr + /// An absorb's command buffer, committed without a wait: the hold that follows (the same queue) runs after it. + pending: ?mtl.CommandBuffer = null, + // one chain's scratch, a row a stream + ids: mtl.Buffer, + pos: mtl.Buffer, + lo: mtl.Buffer, + meta: [2]mtl.Buffer, + picks: mtl.Buffer, + hin: mtl.Buffer, + x: mtl.Buffer, + h: [2]mtl.Buffer, + normed: mtl.Buffer, + hn: mtl.Buffer, + n_mlp: mtl.Buffer, + q_raw: mtl.Buffer, + q: mtl.Buffer, + attn: mtl.Buffer, + o: mtl.Buffer, + gate_up: mtl.Buffer, + act: mtl.Buffer, + y: mtl.Buffer, + pm: mtl.Buffer, + pl: mtl.Buffer, + po: mtl.Buffer, + sums: mtl.Buffer, + logits: mtl.Buffer, + + const scratch = .{ "ids", "pos", "lo", "picks", "hin", "x", "normed", "hn", "n_mlp", "q_raw", "q", "attn", "o", "gate_up", "act", "y", "pm", "pl", "po", "sums", "logits" }; + + /// `metal`: the target's backend, whose stream caches the assistant reads; its streams' capacity sizes the scratch. + pub fn init(gpa: std.mem.Allocator, io: std.Io, metal: *backend.Metal, dir: []const u8) !*Drafter { + const target = metal.model; + const geo = target.config.geo; + const d = try gpa.create(Drafter); + errdefer gpa.destroy(d); + d.* = undefined; + d.gpa = gpa; + d.io = io; + d.target = target; + d.metal = metal; + d.lanes = .empty; + d.step_ms = 0; + d.chain_ms = .{ 0, 0 }; + d.pending = null; + d.trace = std.c.getenv("GEMMA_DRAFT_TRACE") != null; + d.projections = .{ .geo = geo }; + errdefer d.projections.deinit(gpa); // warm's chains fill it + d.config = try aw.Config.read(gpa, io, dir, geo); + const c = &d.config; + d.sources = .{ 0, 0 }; + for (0..geo.layers) |i| d.sources[@intFromBool(cfg.global(i))] = i; + d.tap = .{@intCast(geo.layers - 1)}; + d.w = try aw.load(gpa, target.device, io, dir, c, geo); + errdefer d.w.deinit(gpa); + d.kinds[0] = try kernels.Kind.init(gpa, target.device, c.glue, false, true, target.weights.embed_bits); + errdefer d.kinds[0].deinit(); + d.kinds[1] = try kernels.Kind.init(gpa, target.device, c.glue, true, true, target.weights.embed_bits); + errdefer d.kinds[1].deinit(); + const n = max_items; + const hid: usize = c.hidden; + const qw = c.heads * cfg.global_dim; + const chunks = @max(metal.options.capacity / global_chunk + 2, (cfg.window + sliding_chunk) / sliding_chunk + 2); + const slots = c.heads * chunks; + const sizes = .{ n * 4, n * 4, n * 4, depth * n * @sizeOf(Pick), n * geo.hidden * 2, n * 2 * geo.hidden * 2, n * hid * 2, n * hid * 2, n * hid * 2, n * qw * 2, n * qw * 2, n * qw * 2, n * hid * 2, n * 2 * c.intermediate * 2, n * c.intermediate * 2, n * hid * 2, slots * 4, slots * 4, slots * cfg.global_dim * 4, n * (@max(c.intermediate, qw) / 64) * 4, n * cfg.vocab * 2 }; + const host = .{ true, true, true, true, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false }; + var made: usize = 0; + errdefer d.free(made); + inline for (scratch, sizes, host) |f, bytes, on_host| { + @field(d, f) = try target.device.buffer(bytes, if (on_host) shared else private); + made += 1; + } + for (&d.meta) |*b| { + b.* = try target.device.buffer(n * st.meta_ints * 4, shared); + made += 1; + } + for (&d.h) |*b| { + b.* = try target.device.buffer(n * hid * 2, private); + made += 1; + } + try d.warm(); + metal.resume_needs_draft = false; + return d; + } + + pub fn deinit(d: *Drafter) void { + d.settle() catch {}; + var it = d.lanes.valueIterator(); + while (it.next()) |l| d.dropLane(l.*); + d.lanes.deinit(d.gpa); + d.free(std.math.maxInt(usize)); + d.projections.deinit(d.gpa); + for (d.kinds) |k| k.deinit(); + d.w.deinit(d.gpa); + d.gpa.destroy(d); + } + + /// Releases the first `made` scratch buffers in init's order, so a failed init frees what it made. + fn free(d: *Drafter, made: usize) void { + var i: usize = 0; + inline for (scratch) |f| { + if (i < made) @field(d, f).deinit(); + i += 1; + } + for (d.meta ++ d.h) |b| { + if (i < made) b.deinit(); + i += 1; + } + } + + pub fn drafter(d: *Drafter) dr.Drafter { + return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn } }; + } + + fn cast(ptr: *anyopaque) *Drafter { + return @ptrCast(@alignCast(ptr)); + } + + /// The target's last layer: its output before the final norm, which absorb applies. + fn tapsFn(ptr: *anyopaque) []const u32 { + return &cast(ptr).tap; + } + + /// Chains of up to 6, one assistant forward a draft: the rule prices each draft and may choose none. + fn factsFn(ptr: *anyopaque) dr.Facts { + const d = cast(ptr); + return .{ .depth = depth, .step_ms = d.step_ms, .batched = true, .stream_bytes = d.target.config.geo.hidden * 2 }; + } + + fn openFn(ptr: *anyopaque, s: *Stream) anyerror!void { + const d = cast(ptr); + if (d.lanes.contains(s)) return error.DraftStreamOpen; + const l = try d.newLane(); + errdefer d.dropLane(l); + try d.lanes.put(d.gpa, s, l); + } + + fn newLane(d: *Drafter) !*Lane { + const l = try d.gpa.create(Lane); + errdefer d.gpa.destroy(l); + l.* = .{ .h = try d.target.device.buffer(d.target.config.geo.hidden * 2, private) }; + return l; + } + + fn dropLane(d: *Drafter, l: *Lane) void { + l.h.deinit(); + d.gpa.destroy(l); + } + + fn releaseFn(ptr: *anyopaque, s: *Stream) void { + const d = cast(ptr); + if (d.lanes.fetchRemove(s)) |e| d.dropLane(e.value); + } + + fn heldFn(ptr: *anyopaque, streams: []const *Stream, out: []const []u32) anyerror!void { + const d = cast(ptr); + for (streams, out) |s, o| { + const l = d.lanes.get(s) orelse return error.DraftStreamUnknown; + if (o.len > l.held_len) return error.DraftHeldFewer; + @memcpy(o, l.held[0..o.len]); + } + } + + fn failed(msg: [*:0]const u8) anyerror { + std.log.err("gemma4 assistant: command buffer failed: {s}", .{msg}); + return error.GpuFailed; + } + + /// Waits for a pending absorb's command buffer (a hold ordered after it on the queue has run it already). + fn settle(d: *Drafter) !void { + const cb = d.pending orelse return; + d.pending = null; + defer mtl.objc.release(cb.id); // retained when it was left pending (an autoreleased object otherwise) + cb.wait(); + if (cb.failure()) |msg| return failed(msg); + } + + /// Each stream's last absorbed row through the target's final norm into its lane; committed, not awaited. + fn absorbFn(ptr: *anyopaque, items: []const dr.Absorb) anyerror!void { + const d = cast(ptr); + const t = d.target; + const width = t.config.geo.hidden * 2; + try d.settle(); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + for (items) |it| { + const l = d.lanes.get(it.stream) orelse return error.DraftStreamUnknown; + const f = it.features; + if (f.space != .device or f.dtype != .bf16 or f.row_bytes != width or f.rows == 0 or it.follow.len < f.rows) return error.DraftFeaturesUnsupported; + const last = f.rows - 1; + e.setPipeline(t.kernels.sliding.rms); + e.setBuffer(mtl.Buffer{ .id = @ptrFromInt(f.buffer) }, f.offset + last * f.row_bytes, 0); + e.setBuffer(t.weights.final_norm.buf, t.weights.final_norm.off, 1); + e.setBuffer(l.h, 0, 2); + e.setValue(cfg.eps, 3); + e.setBuffer(d.sums, 0, 4); + e.setValue(@as(i32, 0), 5); // no group sums: the row is copied into a chain's input before it is read + e.dispatchGroups(.{ .width = 1 }, .{ .width = 1024 }); + l.position = it.start + last; + l.token = it.follow[last]; + l.rows = true; + l.held_len = 0; + if (d.trace) std.debug.print("absorb {s}: rows [{d}, {d}), kept row {d}, token {d}\n", .{ it.stream.id, it.start, it.start + f.rows, l.position, l.token }); + } + e.end(); + cb.commit(); + _ = mtl.objc.retain(cb.id); + d.pending = cb; + } + + /// A projection; `summed`: its input's group sums are in d.sums (stride rows), written by the input's own kernel. + fn project(d: *Drafter, e: mtl.ComputeEncoder, l: aw.Linear, x: mtl.Buffer, y: mtl.Buffer, rows: usize, summed: bool) !void { + const p = try d.projections.get(d.gpa, d.target.device, .{ .w = l.w, .sb = l.sb, .n = l.n, .k = l.k }); + if (summed) return p.encodeSums(e, .{ .buf = x }, .{ .buf = l.w }, .{ .buf = l.sb }, .{ .ref = .{ .buf = d.sums }, .stride = @intCast(rows) }, .{ .buf = y }, @intCast(rows)); + try p.encode(e, .{ .buf = x }, .{ .buf = l.w }, .{ .buf = l.sb }, .{ .buf = y }, @intCast(rows)); + } + + fn sumsOut(d: *Drafter, e: mtl.ComputeEncoder, rows: usize, index: usize) void { + e.setBuffer(d.sums, 0, index); + e.setValue(@as(i32, @intCast(rows)), index + 1); + } + + fn bindRef(e: mtl.ComputeEncoder, r: Ref, index: usize) void { + e.setBuffer(r.buf, r.off, index); + } + + /// The host inputs of a chain: each row's pending token, position, first sliding key and both kinds' meta. + fn prepare(d: *Drafter, items: []const Item) void { + const ids = d.ids.slice(u32, max_items); + const pos = d.pos.slice(i32, max_items); + const lo = d.lo.slice(i32, max_items); + for (items, 0..) |it, r| { + const p = it.position; + ids[r] = it.pending; + pos[r] = @intCast(p); + lo[r] = @intCast(if (p + 1 <= cfg.window) 0 else p + 1 - cfg.window); + const c0: usize = @as(usize, @intCast(lo[r])) / sliding_chunk; + const metas = [2][5]i32{ .{ @intCast(c0), @intCast(p / sliding_chunk - c0 + 1), st.ring, 1, std.math.maxInt(i32) }, .{ 0, @intCast(p / global_chunk + 1), 0, 1, std.math.maxInt(i32) } }; + for (d.meta, metas) |m, v| m.slice(i32, max_items * st.meta_ints)[r * st.meta_ints ..][0..5].* = v; + } + } + + /// Item r's queries over its target cache's layer of kind `global` (keys up to its position), merged into d.attn. + fn attend(d: *Drafter, e: mtl.ComputeEncoder, it: Item, r: usize, rows: usize, global: bool) void { + const kd = &d.kinds[@intFromBool(global)]; + const src = d.sources[@intFromBool(global)]; + const c = it.cache; + const nk = if (global) d.config.glue.global_kv_heads else d.config.glue.sliding_kv_heads; + const dh: usize = if (global) cfg.global_dim else cfg.sliding_dim; + const q_off = r * d.config.heads * dh * 2; + const meta = d.meta[@intFromBool(global)]; + const nch: usize = @intCast(meta.slice(i32, max_items * st.meta_ints)[r * st.meta_ints + 1]); + if (global) { + e.setPipeline(kd.decode_attn); + e.setBuffer(d.q, q_off, 0); + e.setBuffer(c.k[src], 0, 1); + e.setBuffer(c.v[src], 0, 2); + e.setValue([4]i32{ 1, @intCast(it.position), @intCast(st.slotsOf(src, c.capacity)), @intCast(nch) }, 3); + for ([_]mtl.Buffer{ d.pm, d.pl, d.po }, 4..) |b, i| e.setBuffer(b, 0, i); + e.dispatchGroups(.{ .width = nch, .height = nk, .depth = (kd.heads_per_kv + 15) / 16 }, .{ .width = 32 * dh / 128 }); + } else { + e.setPipeline(kd.partial); + e.setBuffer(d.q, q_off, 0); + for ([_]mtl.Buffer{ c.k[src], c.v[src], c.k[src], c.v[src] }, 1..) |b, i| e.setBuffer(b, 0, i); // none new: P0 past all + e.setBuffer(d.pos, r * 4, 5); + e.setBuffer(d.lo, r * 4, 6); + e.setBuffer(meta, r * st.meta_ints * 4, 7); + e.setValue(@as(i32, @intCast(st.slotsOf(src, c.capacity))), 8); + for ([_]mtl.Buffer{ d.pm, d.pl, d.po, d.q_raw }, 9..) |b, i| e.setBuffer(b, 0, i); + const qn = d.w.layers[0].q_norm; + bindRef(e, qn, 13); + bindRef(e, qn, 14); + e.setBuffer(d.target.inv_sliding, 0, 15); + e.setValue(cfg.eps, 16); + e.setValue([2]i32{ 0, 0 }, 17); // q prepped already, partials merged below + e.setBuffer(d.attn, q_off, 18); + e.setBuffer(d.sums, r * 4, 19); + e.setValue(@as(i32, @intCast(rows)), 20); + e.setBuffer(c.k[src], 0, 21); + e.setBuffer(c.v[src], 0, 22); + e.dispatchGroups(.{ .width = nch, .height = nk }, .{ .width = 32 * kd.heads_per_kv * kd.split }); + } + e.setPipeline(kd.merge); + for ([_]mtl.Buffer{ d.pm, d.pl, d.po }, 0..) |b, i| e.setBuffer(b, 0, i); + e.setBuffer(meta, r * st.meta_ints * 4, 3); + e.setBuffer(d.attn, q_off, 4); + e.setBuffer(d.sums, r * 4, 5); // this row's o projection sums + e.setValue(@as(i32, @intCast(rows)), 6); + e.dispatchGroups(.{ .width = 1, .height = d.config.heads, .depth = 1 }, .{ .width = 32 }); + } + + /// Assistant layer i on the rows: input normed in d.normed (sums written), residual d.h[cur] into d.h[1 - cur]. + fn layer(d: *Drafter, e: mtl.ComputeEncoder, items: []const Item, i: usize, cur: usize) !void { + const c = &d.config; + const rows = items.len; + const global = !c.sliding[i]; + const kd = &d.kinds[@intFromBool(global)]; + const glue = &d.kinds[0]; + const l = &d.w.layers[i]; + try d.project(e, l.q, d.normed, d.q_raw, rows, true); + e.setPipeline(kd.qkv); + e.setBuffer(d.q_raw, 0, 0); + bindRef(e, l.q_norm, 1); + bindRef(e, l.q_norm, 2); + e.setBuffer(if (global) d.target.inv_global else d.target.inv_sliding, 0, 3); + e.setBuffer(d.pos, 0, 4); + e.setValue(cfg.eps, 5); + for ([_]mtl.Buffer{ d.q, d.o, d.o, d.o, d.o }, 6..) |b, j| e.setBuffer(b, 0, j); // k, v and the caches unwritten + e.setValue([3]i32{ 0, 0, 0 }, 11); + e.dispatchGroups(.{ .width = c.heads, .height = rows }, .{ .width = 32 }); + for (items, 0..) |it, r| d.attend(e, it, r, rows, global); + try d.project(e, l.o, d.attn, d.o, rows, true); + e.setPipeline(glue.attn_tail); + e.setBuffer(d.h[cur], 0, 0); + e.setBuffer(d.o, 0, 1); + bindRef(e, l.post_attention, 2); + bindRef(e, l.pre_feedforward, 3); + e.setValue(cfg.eps, 4); + e.setBuffer(d.hn, 0, 5); + e.setBuffer(d.n_mlp, 0, 6); + d.sumsOut(e, rows, 7); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + try d.project(e, l.gate_up, d.n_mlp, d.gate_up, rows, true); + e.setPipeline(glue.geglu); + e.setBuffer(d.gate_up, 0, 0); + e.setBuffer(d.act, 0, 1); + e.setValue(@as(u32, c.intermediate), 2); + d.sumsOut(e, rows, 3); + e.setValue(@as(i32, 0), 5); // gate rows then up rows, as quantized + e.dispatchGroups(.{ .width = c.intermediate / 64, .height = rows }, .{ .width = 64 }); + try d.project(e, l.down, d.act, d.y, rows, true); + e.setPipeline(glue.dense_tail); + e.setBuffer(d.hn, 0, 0); + e.setBuffer(d.y, 0, 1); + bindRef(e, l.post_feedforward, 2); + bindRef(e, l.scalar, 3); + bindRef(e, if (i + 1 < c.layers) d.w.layers[i + 1].input_norm else d.w.norm, 4); + e.setValue(cfg.eps, 5); + e.setBuffer(d.h[1 - cur], 0, 6); + e.setBuffer(d.normed, 0, 7); + d.sumsOut(e, rows, 8); + e.setBuffer(d.hn, 0, 10); + e.setValue(@as(i32, 0), 11); // no tap + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + + /// Every item's chain of `steps` drafts in one pass: step s's picks at d.picks[s * max_items + row]. + fn encodeChain(d: *Drafter, e: mtl.ComputeEncoder, items: []const Item, steps: u32) !void { + const c = &d.config; + const t = d.target; + const geo = t.config.geo; + const rows = items.len; + d.prepare(items); + for (items, 0..) |it, r| { // step 0's hidden rows: each lane's normed row + e.setPipeline(t.kernels.sliding.draft_store); + e.setBuffer(it.lane.h, 0, 0); + e.setBuffer(d.hin, 0, 1); + e.setValue([4]u32{ @intCast(geo.hidden), @intCast(geo.hidden), @intCast(r), 0 }, 2); + e.dispatchThreads(.{ .width = geo.hidden, .height = 1 }, .{ .width = 256 }); + } + for (0..steps) |s| { + e.setPipeline(t.kernels.sliding.assist_input); + if (s == 0) e.setBuffer(d.ids, 0, 0) else e.setBuffer(d.picks, (s - 1) * max_items * @sizeOf(Pick), 0); + e.setValue([2]u32{ if (s == 0) 1 else @sizeOf(Pick) / 4, cfg.vocab }, 1); + bindRef(e, t.weights.embed_w, 2); + bindRef(e, t.weights.embed_s, 3); + bindRef(e, t.weights.embed_b, 4); + e.setBuffer(d.hin, 0, 5); + e.setBuffer(d.x, 0, 6); + e.setValue(fwd.embedScale(geo.hidden), 7); + e.dispatchThreads(.{ .width = geo.hidden / 2, .height = rows }, .{ .width = 256 }); + try d.project(e, d.w.pre, d.x, d.h[0], rows, false); + e.setPipeline(d.kinds[0].rms); + e.setBuffer(d.h[0], 0, 0); + bindRef(e, d.w.layers[0].input_norm, 1); + e.setBuffer(d.normed, 0, 2); + e.setValue(cfg.eps, 3); + d.sumsOut(e, rows, 4); + e.dispatchGroups(.{ .width = rows }, .{ .width = 1024 }); + var cur: usize = 0; + for (0..c.layers) |i| { + try d.layer(e, items, i, cur); + cur = 1 - cur; + } + try d.project(e, d.w.head, d.normed, d.logits, rows, true); + var first: usize = 0; + while (first < rows) : (first += 16) { // tree_round's argmax in 16-row pieces, as the target's forward + const out = (s * max_items + first) * @sizeOf(Pick); + try t.picks.argmax(e, .{ .buf = d.logits, .off = first * cfg.vocab * 2 }, .{ .buf = d.picks, .off = out }, .{ .rows = @intCast(@min(16, rows - first)), .vocab = cfg.vocab, .stride = cfg.vocab }); + } + if (s + 1 < steps) try d.project(e, d.w.post, d.normed, d.hin, rows, true); // the next step's hidden rows + } + } + + /// Each stream's chain of `depth` drafts after `pending`, all streams' rows through each step together. + fn holdFn(ptr: *anyopaque, holds: []const dr.Hold) anyerror!void { + const d = cast(ptr); + if (holds.len > max_items) return error.DraftTooManyStreams; + const t0 = std.Io.Clock.awake.now(d.io); + var items: [max_items]Item = undefined; + var steps: u32 = 0; + for (holds, items[0..holds.len]) |h, *it| { + const l = d.lanes.get(h.stream) orelse return error.DraftStreamUnknown; + const c = d.metal.caches.get(h.stream) orelse return error.DraftStreamUnknown; + if (h.depth == 0 or h.depth > depth) return error.DraftDepth; + if (!l.rows or l.position + 2 != h.position or c.len < h.position - 1) return error.DraftContextGap; + it.* = .{ .lane = l, .cache = c, .pending = h.pending, .position = l.position, .depth = h.depth }; + steps = @max(steps, h.depth); + } + try d.run(items[0..holds.len], steps); + const picks = d.picks.slice(Pick, depth * max_items); + for (items[0..holds.len], holds, 0..) |it, h, r| { + for (it.lane.held[0..it.depth], 0..) |*tok, s| { + const p = picks[s * max_items + r]; + if (p.nonfinite != 0 or p.token >= cfg.vocab) return error.DraftNonfinite; + tok.* = p.token; + } + it.lane.held_len = it.depth; + if (d.trace) std.debug.print("draft {s} row {d} pending {d} depth {d}: {any}\n", .{ h.stream.id, it.position, it.pending, it.depth, it.lane.held[0..it.depth] }); + } + if (d.trace) std.debug.print("hold of {d} streams: {d:.2} ms\n", .{ holds.len, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(d.io)).toNanoseconds())) / 1e6 }); + } + + /// One command buffer of chains, waited on; the absorb before it (the same queue) has run by then. + fn run(d: *Drafter, items: []const Item, steps: u32) !void { + const cb = d.target.queue.commandBuffer(); + const e = cb.compute(.serial); + try d.encodeChain(e, items, steps); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| return failed(msg); + try d.settle(); + } + + /// Chains over a throwaway cache (the two layers it reads): the pipelines compiled, then 1 and 6 drafts timed. + fn warm(d: *Drafter) !void { + const geo = d.target.config.geo; + var cache: st.Cache = .{ .k = undefined, .v = undefined, .geo = geo, .capacity = 256, .logits = undefined }; + var made: usize = 0; + defer for (d.sources[0..made]) |i| { + cache.k[i].deinit(); + cache.v[i].deinit(); + }; + for (d.sources) |i| { + const bytes = geo.kvHeads(i) * st.slotsOf(i, cache.capacity) * cfg.headDim(i) * 2; + cache.k[i] = try d.target.device.buffer(bytes, private); + cache.v[i] = d.target.device.buffer(bytes, private) catch |e| { + cache.k[i].deinit(); + return e; + }; + made += 1; + } + const l = try d.newLane(); + defer d.dropLane(l); + var ms: [2]f64 = undefined; + for ([_]u32{ depth, 1, depth, 1, depth }, 0..) |steps, i| { + const t0 = std.Io.Clock.awake.now(d.io); + try d.run(&.{.{ .lane = l, .cache = &cache, .pending = 2, .position = 128, .depth = steps }}, steps); + if (i >= 3) ms[i - 3] = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(d.io)).toNanoseconds())) / 1e6; + } + d.chain_ms = ms; + d.step_ms = @max(ms[1] - ms[0], 0) / @as(f64, depth - 1); + } +}; + +test "a hold's key range and chunk counts follow the kept row's position" { + const cases = [_][3]u64{ .{ 0, 0, 1 }, .{ 1023, 0, 8 }, .{ 1024, 1, 9 }, .{ 5000, 3977, 9 } }; + for (cases) |k| { + const p = k[0]; + const lo: u64 = if (p + 1 <= cfg.window) 0 else p + 1 - cfg.window; + try std.testing.expectEqual(k[1], lo); + try std.testing.expectEqual(k[2], p / sliding_chunk - lo / sliding_chunk + 1); + try std.testing.expect(p + 1 - lo <= cfg.window); + } +} diff --git a/zig/src/families/gemma4/assistant_weights.zig b/zig/src/families/gemma4/assistant_weights.zig new file mode 100644 index 0000000000..445640f014 --- /dev/null +++ b/zig/src/families/gemma4/assistant_weights.zig @@ -0,0 +1,301 @@ +//! Gemma 4's MTP assistant (gemma4_unified_assistant): config admitted, bf16 linears and the tied head made 4-bit. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); +const ckpt = @import("../../core/checkpoint_metal.zig"); +const shards = @import("../../core/checkpoint.zig"); +const dw = @import("dflash_weights.zig"); +const weights = @import("weights.zig"); + +pub const Linear = dw.Linear; +pub const Ref = dw.Ref; + +/// Layers an admitted assistant has at most: the size of per-layer arrays. +pub const max_layers = 8; +/// The stock assistant's sizes (google/gemma-4-12B-it-assistant); another config is refused with its reason. +const stock = .{ .hidden = 1024, .layers = 4, .intermediate = 8192, .heads = 16 }; + +pub const Config = struct { + layers: u32, + hidden: u32, + intermediate: u32, + heads: u32, + /// Each layer's kind: true for a sliding layer (the target's last sliding cache), false for a full one. + sliding: [max_layers]bool = @splat(true), + /// The assistant's sizes as the glue's geometry: its width, its query heads over the target's key heads. + glue: cfg.Geometry, + + /// `geo`: the target's sizes, whose width the assistant was trained on and whose caches it reads. + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, geo: cfg.Geometry) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(bytes); + return parse(gpa, bytes, geo); + } + + pub fn parse(gpa: std.mem.Allocator, bytes: []const u8, geo: cfg.Geometry) !Config { + var parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const o = if (parsed.value == .object) parsed.value.object else return error.BadDraftConfig; + if (!isString(o.get("model_type"), "gemma4_unified_assistant")) return error.UnsupportedAssistantType; + if (isTrue(o, "use_ordered_embeddings")) return error.AssistantOrderedEmbeddingsNotServed; + if (try integer(o, "backbone_hidden_size") != geo.hidden) return error.AssistantForAnotherTarget; + if (o.get("quantization") != null or o.get("quantization_config") != null) return error.UnsupportedAssistantQuantization; + const t = if (o.get("text_config")) |v| (if (v == .object) v.object else return error.BadDraftConfig) else return error.BadDraftConfig; + if (!isString(t.get("model_type"), "gemma4_unified_text")) return error.BadDraftConfig; + var c: Config = .{ .layers = @intCast(try integer(t, "num_hidden_layers")), .hidden = @intCast(try integer(t, "hidden_size")), .intermediate = @intCast(try integer(t, "intermediate_size")), .heads = @intCast(try integer(t, "num_attention_heads")), .glue = undefined }; + if (c.hidden != stock.hidden or c.layers != stock.layers or c.intermediate != stock.intermediate or c.heads != stock.heads) return error.UnsupportedAssistantGeometry; + if (try integer(t, "num_key_value_heads") != geo.sliding_kv_heads or try integer(t, "num_global_key_value_heads") != geo.global_kv_heads) return error.AssistantForAnotherTarget; + if (try integer(t, "num_kv_shared_layers") != c.layers) return error.UnsupportedAssistantGeometry; // every layer reads the target's + const dims = .{ .{ "head_dim", cfg.sliding_dim }, .{ "global_head_dim", cfg.global_dim }, .{ "sliding_window", cfg.window }, .{ "vocab_size", cfg.vocab } }; + inline for (dims) |d| if (try integer(t, d[0]) != d[1]) return error.UnsupportedAssistantGeometry; + if (isTrue(t, "enable_moe_block") or isTrue(t, "use_double_wide_mlp") or isTrue(t, "attention_bias")) return error.UnsupportedAssistantGeometry; + if (t.get("hidden_size_per_layer_input")) |v| if (v != .integer or v.integer != 0) return error.UnsupportedAssistantGeometry; + if (!isTrue(t, "attention_k_eq_v") or !isTrue(t, "tie_word_embeddings")) return error.UnsupportedAssistantGeometry; + if (!isString(t.get("hidden_activation"), "gelu_pytorch_tanh")) return error.UnsupportedAssistantGeometry; + if (t.get("final_logit_softcapping")) |v| if (v != .null) return error.UnsupportedAssistantGeometry; + if (try number(t, "rms_norm_eps") != 1e-6) return error.UnsupportedAssistantGeometry; + const ropes = if (t.get("rope_parameters")) |v| (if (v == .object) v.object else return error.BadDraftConfig) else return error.BadDraftConfig; + try rope(ropes, "sliding_attention", "default", cfg.sliding_theta, null); + try rope(ropes, "full_attention", "proportional", cfg.global_theta, 0.25); + const kinds = t.get("layer_types") orelse return error.BadDraftConfig; + if (kinds != .array or kinds.array.items.len != c.layers) return error.BadDraftConfig; + for (kinds.array.items, 0..) |k, i| { + c.sliding[i] = isString(k, "sliding_attention"); + if (!c.sliding[i] and !isString(k, "full_attention")) return error.UnsupportedAssistantGeometry; + } + c.glue = .{ .name = "assistant", .hidden = c.hidden, .layers = c.layers, .intermediate = c.intermediate, .query_heads = c.heads, .sliding_kv_heads = geo.sliding_kv_heads, .global_kv_heads = geo.global_kv_heads }; + return c; + } + + /// Layer `i`'s head dim: the target's sliding or global one. + pub fn headDim(c: *const Config, i: usize) usize { + return if (c.sliding[i]) cfg.sliding_dim else cfg.global_dim; + } + + /// The shape tensor `name` must have, or null for a tensor the assistant does not read. + pub fn expected(c: *const Config, geo: cfg.Geometry, name: []const u8) ?[2]usize { + const hid: usize = c.hidden; + if (std.mem.eql(u8, name, "pre_projection.weight")) return .{ hid, 2 * geo.hidden }; + if (std.mem.eql(u8, name, "post_projection.weight")) return .{ geo.hidden, hid }; + if (std.mem.eql(u8, name, "model.embed_tokens.weight")) return .{ cfg.vocab, hid }; + if (std.mem.eql(u8, name, "model.norm.weight")) return .{ hid, 1 }; + const rest = if (std.mem.startsWith(u8, name, "model.layers.")) name["model.layers.".len..] else return null; + const dot = std.mem.indexOfScalar(u8, rest, '.') orelse return null; + const i = std.fmt.parseInt(usize, rest[0..dot], 10) catch return null; + if (i >= c.layers) return null; + const q = c.heads * c.headDim(i); + const shapes = .{ + .{ "self_attn.q_proj.weight", [2]usize{ q, hid } }, .{ "self_attn.o_proj.weight", [2]usize{ hid, q } }, + .{ "self_attn.q_norm.weight", [2]usize{ c.headDim(i), 1 } }, .{ "mlp.gate_proj.weight", [2]usize{ c.intermediate, hid } }, + .{ "mlp.up_proj.weight", [2]usize{ c.intermediate, hid } }, .{ "mlp.down_proj.weight", [2]usize{ hid, c.intermediate } }, + .{ "input_layernorm.weight", [2]usize{ hid, 1 } }, .{ "post_attention_layernorm.weight", [2]usize{ hid, 1 } }, + .{ "pre_feedforward_layernorm.weight", [2]usize{ hid, 1 } }, .{ "post_feedforward_layernorm.weight", [2]usize{ hid, 1 } }, + .{ "layer_scalar", [2]usize{ 1, 1 } }, + }; + inline for (shapes) |s| if (std.mem.eql(u8, rest[dot + 1 ..], s[0])) return s[1]; + return null; + } +}; + +/// Whether `dir` holds a Gemma 4 assistant (its config's model_type names one), not a DFlash drafter. +pub fn isAssistant(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !bool { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(bytes); + return typeIsAssistant(gpa, bytes); +} + +fn typeIsAssistant(gpa: std.mem.Allocator, bytes: []const u8) !bool { + var parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const t = if (parsed.value == .object) parsed.value.object.get("model_type") else null; + return t != null and t.? == .string and std.mem.startsWith(u8, t.?.string, "gemma4") and std.mem.endsWith(u8, t.?.string, "assistant"); +} + +fn isString(v: ?std.json.Value, expected: []const u8) bool { + const s = v orelse return false; + return s == .string and std.mem.eql(u8, s.string, expected); +} + +fn isTrue(o: std.json.ObjectMap, key: []const u8) bool { + const v = o.get(key) orelse return false; + return v == .bool and v.bool; +} + +fn integer(o: std.json.ObjectMap, key: []const u8) !i64 { + const v = o.get(key) orelse return error.BadDraftConfig; + return if (v == .integer) v.integer else error.BadDraftConfig; +} + +fn number(o: std.json.ObjectMap, key: []const u8) !f64 { + const v = o.get(key) orelse return error.BadDraftConfig; + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadDraftConfig, + }; +} + +fn rope(o: std.json.ObjectMap, kind: []const u8, kind_type: []const u8, theta: f64, partial: ?f64) !void { + const v = o.get(kind) orelse return error.BadDraftConfig; + const r = if (v == .object) v.object else return error.BadDraftConfig; + if (!isString(r.get("rope_type"), kind_type) or try number(r, "rope_theta") != theta) return error.UnsupportedAssistantRotary; + if (partial) |p| if (try number(r, "partial_rotary_factor") != p) return error.UnsupportedAssistantRotary; +} + +pub const Layer = struct { + q: Linear, + o: Linear, + gate_up: Linear, + down: Linear, + q_norm: Ref, + input_norm: Ref, + post_attention: Ref, + pre_feedforward: Ref, + post_feedforward: Ref, + scalar: Ref, +}; + +pub const Weights = struct { + /// [embedding | hidden] to the assistant's width, and its output back to the target's. + pre: Linear, + post: Linear, + /// The tied head: embed_tokens as a 4-bit linear over the vocabulary. + head: Linear, + norm: Ref, + layers: [max_layers]Layer, + owned: std.ArrayList(mtl.Buffer) = .empty, + + pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void { + for (w.owned.items) |b| b.deinit(); + w.owned.deinit(gpa); + } +}; + +/// Every tensor the directory holds has the shape `c` expects (an unknown one is refused), and each linear is 4-bit. +pub fn load(gpa: std.mem.Allocator, device: mtl.Device, io: std.Io, dir: []const u8, c: *const Config, geo: cfg.Geometry) !Weights { + var check = ckpt.Checkpoint.init(gpa); + defer check.deinit(); + const files = try shards.shardFiles(gpa, io, dir); + defer shards.freeShardFiles(gpa, files); + for (files) |path| try check.addFile(device, path, ""); + var it = check.tensors.iterator(); + while (it.next()) |e| { + const want = c.expected(geo, e.key_ptr.*) orelse { + std.log.err("gemma4 assistant: {s} is not a tensor the assistant reads", .{e.key_ptr.*}); + return error.BadDraftWeight; + }; + const t = e.value_ptr.*; + const got: [2]usize = if (t.rank == 1) .{ t.shape[0], 1 } else .{ t.shape[0], t.shape[1] }; + if (t.dtype != .bf16 or t.rank == 0 or t.rank > 2 or !std.mem.eql(usize, &got, &want)) return error.BadDraftWeight; + } + var w: Weights = undefined; + w.owned = .empty; + errdefer w.deinit(gpa); + var p: [64]u8 = undefined; + var name: [160]u8 = undefined; + w.pre = try dw.quantized(gpa, device, &check, &w.owned, geo, "", &.{"pre_projection"}); + w.post = try dw.quantized(gpa, device, &check, &w.owned, geo, "", &.{"post_projection"}); + w.head = try dw.quantized(gpa, device, &check, &w.owned, geo, "model.", &.{"embed_tokens"}); + w.norm = try dw.vector(gpa, device, &check, &w.owned, "model.norm.weight", c.hidden); + for (w.layers[0..c.layers], 0..) |*l, i| { + const pre = try std.fmt.bufPrint(&p, "model.layers.{d}.", .{i}); + l.q = try dw.quantized(gpa, device, &check, &w.owned, geo, pre, &.{"self_attn.q_proj"}); + l.o = try dw.quantized(gpa, device, &check, &w.owned, geo, pre, &.{"self_attn.o_proj"}); + l.gate_up = try dw.quantized(gpa, device, &check, &w.owned, geo, pre, &.{ "mlp.gate_proj", "mlp.up_proj" }); + l.down = try dw.quantized(gpa, device, &check, &w.owned, geo, pre, &.{"mlp.down_proj"}); + const fields = .{ "q_norm", "input_norm", "post_attention", "pre_feedforward", "post_feedforward", "scalar" }; + const names = .{ "self_attn.q_norm.weight", "input_layernorm.weight", "post_attention_layernorm.weight", "pre_feedforward_layernorm.weight", "post_feedforward_layernorm.weight", "layer_scalar" }; + inline for (fields, names) |field, t| { + const full = try std.fmt.bufPrint(&name, "{s}{s}", .{ pre, t }); + @field(l, field) = try dw.vector(gpa, device, &check, &w.owned, full, c.expected(geo, full).?[0]); + } + } + return w; +} + +const fixture = @embedFile("fixtures/assistant_config.json"); + +fn edited(gpa: std.mem.Allocator, from: []const u8, to: []const u8) ![]u8 { + const out = try std.mem.replaceOwned(u8, gpa, fixture, from, to); + if (std.mem.eql(u8, out, fixture)) return error.FixtureUnchanged; + return out; +} + +test "the 12B's assistant is admitted against the 12B, with the target's key heads" { + const c = try Config.parse(std.testing.allocator, fixture, cfg.unified12b); + try std.testing.expectEqual(@as(u32, 4), c.layers); + try std.testing.expect(c.sliding[0] and c.sliding[2] and !c.sliding[3]); + try std.testing.expectEqual(@as(usize, 8), c.glue.sliding_kv_heads); + try std.testing.expectEqual(@as(usize, 1), c.glue.global_kv_heads); + try std.testing.expectEqual(@as(usize, 1024), c.glue.hidden); + try std.testing.expectEqual(@as(usize, 512), c.headDim(3)); +} + +test "an assistant for another target or layout is refused by name" { + const gpa = std.testing.allocator; + try std.testing.expectError(error.AssistantForAnotherTarget, Config.parse(gpa, fixture, cfg.dense31b)); + const cases = .{ + .{ "\"use_ordered_embeddings\": false", "\"use_ordered_embeddings\": true", error.AssistantOrderedEmbeddingsNotServed }, + .{ "\"backbone_hidden_size\": 3840", "\"backbone_hidden_size\": 5376", error.AssistantForAnotherTarget }, + .{ "\"model_type\": \"gemma4_unified_assistant\"", "\"model_type\": \"gemma4_assistant\"", error.UnsupportedAssistantType }, + .{ "\"num_key_value_heads\": 8", "\"num_key_value_heads\": 16", error.AssistantForAnotherTarget }, + .{ "\"final_logit_softcapping\": null", "\"final_logit_softcapping\": 30.0", error.UnsupportedAssistantGeometry }, + .{ "\"hidden_size\": 1024", "\"hidden_size\": 2048", error.UnsupportedAssistantGeometry }, + .{ "\"num_kv_shared_layers\": 4", "\"num_kv_shared_layers\": 2", error.UnsupportedAssistantGeometry }, + .{ "\"rope_theta\": 10000.0", "\"rope_theta\": 1000.0", error.UnsupportedAssistantRotary }, + .{ "\"tie_word_embeddings\": true,\n \"transformers", "\"quantization\": {\"bits\": 4, \"group_size\": 64},\n \"transformers", error.UnsupportedAssistantQuantization }, + }; + inline for (cases) |k| { + const bytes = try edited(gpa, k[0], k[1]); + defer gpa.free(bytes); + try std.testing.expectError(k[2], Config.parse(gpa, bytes, cfg.unified12b)); + } +} + +test "a drafter directory is the assistant's by its model type" { + const gpa = std.testing.allocator; + try std.testing.expect(try typeIsAssistant(gpa, fixture)); + try std.testing.expect(try typeIsAssistant(gpa, "{\"model_type\":\"gemma4_assistant\"}")); + try std.testing.expect(!try typeIsAssistant(gpa, "{\"architectures\":[\"DFlashDraftModel\"],\"model_type\":\"qwen3\"}")); + try std.testing.expect(!try typeIsAssistant(gpa, "{\"architectures\":[\"DFlashDraftModel\"]}")); +} + +test "the stock checkpoint's tensors are the shapes the assistant reads, and others are unknown" { + const c = try Config.parse(std.testing.allocator, fixture, cfg.unified12b); + const g = cfg.unified12b; + const cases = .{ + .{ "model.embed_tokens.weight", [2]usize{ 262144, 1024 } }, .{ "pre_projection.weight", [2]usize{ 1024, 7680 } }, + .{ "post_projection.weight", [2]usize{ 3840, 1024 } }, .{ "model.norm.weight", [2]usize{ 1024, 1 } }, + .{ "model.layers.0.self_attn.q_proj.weight", [2]usize{ 4096, 1024 } }, .{ "model.layers.3.self_attn.q_proj.weight", [2]usize{ 8192, 1024 } }, + .{ "model.layers.0.self_attn.o_proj.weight", [2]usize{ 1024, 4096 } }, .{ "model.layers.3.self_attn.o_proj.weight", [2]usize{ 1024, 8192 } }, + .{ "model.layers.2.self_attn.q_norm.weight", [2]usize{ 256, 1 } }, .{ "model.layers.3.self_attn.q_norm.weight", [2]usize{ 512, 1 } }, + .{ "model.layers.1.mlp.gate_proj.weight", [2]usize{ 8192, 1024 } }, .{ "model.layers.1.mlp.down_proj.weight", [2]usize{ 1024, 8192 } }, + .{ "model.layers.3.post_feedforward_layernorm.weight", [2]usize{ 1024, 1 } }, .{ "model.layers.3.layer_scalar", [2]usize{ 1, 1 } }, + }; + inline for (cases) |k| try std.testing.expectEqual(@as(?[2]usize, k[1]), c.expected(g, k[0])); + for ([_][]const u8{ "model.layers.0.self_attn.k_proj.weight", "model.layers.4.layer_scalar", "lm_head.weight", "model.token_ordering" }) |name| + try std.testing.expectEqual(@as(?[2]usize, null), c.expected(g, name)); +} + +test "every assistant linear fits the reg lane layout at the target's splits" { + const c = try Config.parse(std.testing.allocator, fixture, cfg.unified12b); + const g = cfg.unified12b; + var shapes: [3 + 4 * max_layers][2]usize = undefined; + var count: usize = 0; + for ([_][]const u8{ "pre_projection.weight", "post_projection.weight", "model.embed_tokens.weight" }) |name| { + shapes[count] = c.expected(g, name).?; + count += 1; + } + for (0..c.layers) |i| { + const q = c.heads * c.headDim(i); + for ([_][2]usize{ .{ q, c.hidden }, .{ c.hidden, q }, .{ 2 * c.intermediate, c.hidden }, .{ c.hidden, c.intermediate } }) |s| { + shapes[count] = s; + count += 1; + } + } + for (shapes[0..count]) |s| try weights.layout(g, .{ .w = undefined, .sb = undefined, .n = s[0], .k = s[1] }).validate(); +} diff --git a/zig/src/families/gemma4/backend.zig b/zig/src/families/gemma4/backend.zig new file mode 100644 index 0000000000..96084d96ef --- /dev/null +++ b/zig/src/families/gemma4/backend.zig @@ -0,0 +1,381 @@ +//! Gemma 4 dense behind the lane engine's backend: windows of up to 16 rows, 128 rows of streams a forward, taps. +const std = @import("std"); +const lanes = @import("lanes"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); +const st = @import("state.zig"); +const fwd = @import("forward.zig"); +const prompt = @import("prompt.zig"); +const snapshot = @import("snapshot.zig"); +const Model = @import("model.zig").Model; +const draw = @import("../qwen3_5/backend.zig").draw; +const be = lanes.backend; + +pub const Options = struct { capacity: usize, streams: u32 = 8 }; + +pub const Metal = struct { + gpa: std.mem.Allocator, + model: *Model, + options: Options, + scratch: st.Scratch, + /// Prompt chunks' activations: 512 rows on core affine_mm and the tensor-unit attention. + prompt_scratch: prompt.Scratch, + caches: std.AutoHashMapUnmanaged(*lanes.Stream, *st.Cache) = .empty, + /// Forwards run on `scratch`: a stream's tapped rows there are current while its tap's gen matches. + gen: u64 = 0, + /// A window's forward ms by rows, timed by measure() (the depth rule prices drafts with them); none until then. + costs: [st.max_rows]lanes.config.Cost = undefined, + timed: usize = 0, + /// Wider forwards' ms by total rows (several streams' windows), timed at load; the rule interpolates between. + shared: [shared_widths.len]lanes.config.Cost = undefined, + shared_timed: usize = 0, + io: std.Io, + trace: bool = false, // GEMMA_DRAFT_TRACE: each verify's rows and wall time on stderr + /// Whether the drafter keeps a context with a kept state (DFlash does); the MTP assistant reads the target's alone. + resume_needs_draft: bool = true, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, model: *Model, options: Options) !*Metal { + const out = try gpa.create(Metal); + errdefer gpa.destroy(out); + var scratch = try st.Scratch.init(model.device, model.config.geo, options.capacity); + errdefer scratch.deinit(); + out.* = .{ .gpa = gpa, .io = io, .trace = std.c.getenv("GEMMA_DRAFT_TRACE") != null, .model = model, .options = options, .scratch = scratch, .prompt_scratch = try prompt.Scratch.init(model.device, model.config.geo, options.capacity) }; + return out; + } + + pub fn deinit(self: *Metal) void { + var it = self.caches.valueIterator(); + while (it.next()) |cache| { + cache.*.deinit(); + self.gpa.destroy(cache.*); + } + self.caches.deinit(self.gpa); + self.scratch.deinit(); + self.prompt_scratch.deinit(); + self.gpa.destroy(self); + } + + /// Every kernel keeps a row's bits at any window width and beside any other stream's rows (tf-gemma4-glue). + pub fn facts(self: *const Metal) lanes.Model { + return .{ .exact_width = st.max_rows, .hidden_rows = true, .batch_rows = st.prompt_rows, .max_streams = self.options.streams, .streams_exact = true, .window_costs = self.costs[0..self.timed], .shared_costs = self.shared[0..self.shared_timed] }; + } + + /// A window of each width 1..16 at position 0 on a throwaway cache, the second of two runs timed. + pub fn measure(self: *Metal) !void { + const io = self.io; + var cache = try st.Cache.init(self.model.device, self.model.config.geo, 256); + defer cache.deinit(); + var ids: [st.max_rows]u32 = @splat(2); + for (1..st.max_rows + 1) |rows| { + var ms: f64 = 0; + for (0..2) |run| { + cache.len = 0; + const t0 = std.Io.Clock.awake.now(io); + try fwd.run(self.model, &self.scratch, &.{.{ .cache = &cache, .rows = rows }}, ids[0..rows]); + if (run == 1) ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds())) / 1e6; + } + self.costs[rows - 1] = .{ .width = @intCast(rows), .ms = ms }; + } + self.timed = st.max_rows; + for (&self.shared, shared_widths) |*cost, rows| cost.* = .{ .width = @intCast(rows), .ms = try self.timeRows(rows) }; + self.shared_timed = shared_widths.len; + self.gen += 1; + } + + /// The forward's ms at `rows` rows (up to 128) of one stream, the second of two runs. + pub fn timeRows(self: *Metal, rows: usize) !f64 { + const io = self.io; + var cache = try st.Cache.init(self.model.device, self.model.config.geo, 256); + defer cache.deinit(); + var ids: [st.prompt_rows]u32 = @splat(2); + var ms: f64 = 0; + for (0..2) |run| { + cache.len = 0; + const t0 = std.Io.Clock.awake.now(io); + try fwd.run(self.model, &self.scratch, &.{.{ .cache = &cache, .rows = rows }}, ids[0..rows]); + if (run == 1) ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds())) / 1e6; + } + self.gen += 1; + return ms; + } + + pub fn backend(self: *Metal) be.Backend { + return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .prepare_features = prepareFn, .features = featuresFn, .prefill_features = prefillFeaturesFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .release = releaseFn } }; + } + + fn cast(ptr: *anyopaque) *Metal { + return @ptrCast(@alignCast(ptr)); + } + + fn cacheOf(self: *Metal, stream: *lanes.Stream) !*st.Cache { + return self.caches.get(stream) orelse error.UnknownGemmaStream; + } + + /// The last row of the scratch logits into the stream's cache. + fn keepLast(self: *Metal, cache: *st.Cache, row: usize) void { + @memcpy(cache.logits.slice(u16, cfg.vocab), self.scratch.logits.slice(u16, st.prompt_rows * cfg.vocab)[row * cfg.vocab ..][0..cfg.vocab]); + } + + fn prefillFn(ptr: *anyopaque, stream: *lanes.Stream) !void { + return cast(ptr).prefillWith(stream, null); + } + + fn prefillFeaturesFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, sink: be.FeatureSink) !void { + const self = cast(ptr); + if (taps.len != self.model.taps) return error.UnpreparedGemmaTaps; + return self.prefillWith(stream, sink); + } + + /// A tapped row's bytes: every tap's hidden row, bf16. + fn tapBytes(self: *const Metal) usize { + return self.model.taps * self.model.config.geo.hidden * 2; + } + + /// The drafter's layers: both scratches keep their rows from now on. + fn prepareFn(ptr: *anyopaque, taps: []const u32) !void { + const self = cast(ptr); + if (self.scratch.taps != null) return if (taps.len == self.model.taps) self.model.setTaps(taps) else error.BadGemmaTaps; + try self.model.setTaps(taps); + const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + self.scratch.taps = try self.model.device.buffer(st.prompt_rows * self.tapBytes(), shared); + self.prompt_scratch.taps = try self.model.device.buffer(st.chunk_rows * self.tapBytes(), shared); + } + + /// Rows [start, start + count) of the stream's latest tapped rows, on the device. + fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, taps: []const u32, start: u64, count: u32) !be.Features { + const self = cast(ptr); + const t = (try self.cacheOf(stream)).tap; + if (taps.len != self.model.taps or t.buf == null) return error.UnpreparedGemmaTaps; + if (t.gen) |g| if (g != self.gen) return error.StaleGemmaTaps; + if (count == 0 or start < t.start or start + count > t.start + t.rows) return error.GemmaTapsNotHeld; + const width = self.tapBytes(); + return .{ .buffer = @intFromPtr(t.buf.?.id), .offset = t.off + (start - t.start) * width, .rows = count, .row_bytes = @intCast(width), .space = .device, .dtype = .bf16 }; + } + + fn prefillWith(self: *Metal, stream: *lanes.Stream, sink: ?be.FeatureSink) !void { + const ids = stream.prompt(); + if (ids.len == 0 or ids.len + stream.max_new > self.options.capacity or self.caches.contains(stream)) return error.InvalidGemmaPrompt; + const cache = try self.gpa.create(st.Cache); + errdefer self.gpa.destroy(cache); + const capacity = @min(self.options.capacity, (ids.len + stream.max_new + st.max_rows + 63) / 64 * 64); + cache.* = try st.Cache.init(self.model.device, self.model.config.geo, capacity); + errdefer cache.deinit(); + try self.caches.put(self.gpa, stream, cache); // a mark's save finds it there + errdefer _ = self.caches.remove(stream); + stream.cached = 0; + if (stream.reuse.saved) |saved| { // a kept state of this prompt's prefix: the pass starts there + const snap: *const snapshot.Snap = @ptrCast(@alignCast(saved)); + var inside = false; // an image's rows see each other: a state inside one is not resumed + for (stream.images) |im| inside = inside or (im.at < snap.at and snap.at < im.at + im.tokens); + // a drafting stream resumes only a state kept with its drafter's context, if any (dflash.zig keepLane) + const drafter_less = stream.drafts and self.prompt_scratch.taps != null and snap.draft == null and self.resume_needs_draft; + if (snap.at < ids.len and !inside and !drafter_less) { + if (snapshot.restore(self.gpa, self.model.queue, snap, cache)) stream.cached = snap.at else |_| stream.reuse_failed = true; + } else stream.reuse_failed = true; + } + var spans: std.ArrayList(prompt.Span) = .empty; + defer spans.deinit(self.gpa); + const rows_buf = try self.encodeImages(stream, cache.len, &spans); + defer if (rows_buf) |b| b.deinit(); + var at: usize = cache.len; + while (at < ids.len) { + if (stream.isCancelled()) return error.Cancelled; + var end = @min(ids.len, at + st.chunk_rows); + // planned chunk starts and kept marks cut a chunk (resumed and fresh passes cut alike past them) + for ([_][]const u32{ stream.chunks, stream.reuse.marks }) |cuts| for (cuts) |x| if (x > at and x < end) { + end = x; + }; + // an image's rows see each other on sliding layers: no cut inside one (before it, or past it when it leads) + for (spans.items) |span| if (span.at < end and span.at + span.n > end) { + end = if (span.at > at) span.at else span.at + span.n; + }; + try prompt.run(self.model, &self.prompt_scratch, cache, ids[at..end], spans.items); + if (sink) |k| { + const width = self.tapBytes(); + const rows = self.prompt_scratch.taps.?; + try k.take(k.ptr, stream, .{ .buffer = @intFromPtr(rows.id), .rows = @intCast(end - at), .row_bytes = @intCast(width), .space = .device, .dtype = .bf16 }, at); + if (end == ids.len) { // the prompt's last row, kept for the first draft (another prompt may reuse the scratch) + const last = try self.model.device.buffer(width, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + cache.last_tap = last; + @memcpy(last.slice(u8, width), rows.slice(u8, st.chunk_rows * width)[(end - 1 - at) * width ..][0..width]); + cache.tap = .{ .buf = last, .start = ids.len - 1, .rows = 1 }; + } + } + at = end; + if (std.mem.indexOfScalar(u32, stream.reuse.marks, @intCast(cache.len)) != null) if (stream.reuse.hook) |k| k.at(k.ptr, stream, @intCast(cache.len)); + } + @memcpy(cache.logits.slice(u16, cfg.vocab), self.prompt_scratch.logits.slice(u16, cfg.vocab)); + } + + /// The stream's images from `from` on through the tower into one buffer of rows (null: none), a span each. + fn encodeImages(self: *Metal, stream: *const lanes.Stream, from: usize, spans: *std.ArrayList(prompt.Span)) !?mtl.Buffer { + var total: usize = 0; + for (stream.images) |im| { + if (im.at + im.tokens > stream.prompt().len or im.tokens * 9 != im.gh * im.gw) return error.InvalidGemmaImage; + if (im.at >= from) total += im.tokens; // one before `from` is in the restored state + } + if (total == 0) return null; + const v = self.model.vision orelse return error.NoGemmaVision; + const hidden = self.model.config.geo.hidden; + const buf = try self.model.device.buffer(total * hidden * 2, mtl.ResourceOptions.private | mtl.ResourceOptions.untracked); + errdefer buf.deinit(); + var row: usize = 0; + for (stream.images) |im| { // one command buffer an image: the tower's host-written inputs are one image's + if (im.at < from) continue; + const rows: prompt.Span = .{ .at = im.at, .n = im.tokens, .rows = .{ .buf = buf, .off = row * hidden * 2 } }; + const cb = self.model.queue.commandBuffer(); + const e = cb.compute(.serial); + try v.encode(e, im.pixels, im.gh, im.gw, rows.rows); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("gemma4: image tower failed: {s}", .{msg}); + return error.GpuFailed; + } + try spans.append(self.gpa, rows); + row += im.tokens; + } + return buf; + } + + fn firstFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len) return error.InvalidGemmaPosition; + return try draw(self.gpa, cache.logits.slice(u16, cfg.vocab), stream.sampling, position, &.{}); + } + + fn firstMaskedFn(ptr: *anyopaque, stream: *lanes.Stream, position: u64, mask: []const u32) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len) return error.InvalidGemmaPosition; + return try draw(self.gpa, cache.logits.slice(u16, cfg.vocab), stream.sampling, position, mask); + } + + fn queueFn(ptr: *anyopaque, stream: *lanes.Stream, feed: be.Feed, position: u64) !u64 { + const self = cast(ptr); + const cache = try self.cacheOf(stream); + if (position != cache.len + 1) return error.InvalidGemmaPosition; + const token: u32 = switch (feed) { + .value => |v| v, + .handle => |h| @intCast(h), + }; + cache.from = cache.len; + try fwd.run(self.model, &self.scratch, &.{.{ .cache = cache, .rows = 1 }}, &.{token}); + self.gen += 1; + self.keepLast(cache, 0); + return try draw(self.gpa, cache.logits.slice(u16, cfg.vocab), stream.sampling, position, &.{}); + } + + fn readFn(_: *anyopaque, handle: u64) !u32 { + if (handle >= cfg.vocab) return error.InvalidGemmaToken; + return @intCast(handle); + } + + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) !void { + const self = cast(ptr); + var segments: [st.prompt_rows]fwd.Segment = undefined; + var ids: [st.prompt_rows]u32 = undefined; + var total: usize = 0; + if (windows.len > segments.len or windows.len != out.len) return error.InvalidGemmaWindows; + for (windows, 0..) |w, i| { + const rows = w.rows(); + if (w.parents != null or w.held != 0 or rows > st.max_rows or total + rows > st.prompt_rows or w.positions.len != rows) return error.UnsupportedGemmaWindow; + if (w.masks.len != 0 and w.masks.len != rows * mask_words) return error.InvalidGemmaMasks; + const cache = try self.cacheOf(w.stream); + for (w.positions, 0..) |p, r| if (p != cache.len + r + 1) return error.InvalidGemmaPosition; + cache.from = cache.len; + segments[i] = .{ .cache = cache, .rows = rows }; + ids[total] = w.pending; + @memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens); + total += rows; + } + const t0 = std.Io.Clock.awake.now(self.io); + try fwd.run(self.model, &self.scratch, segments[0..windows.len], ids[0..total]); + if (self.trace) std.debug.print("verify of {d} streams, {d} rows: {d:.2} ms\n", .{ windows.len, total, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(self.io)).toNanoseconds())) / 1e6 }); + self.gen += 1; + if (self.scratch.taps) |taps| { + var row: usize = 0; + for (segments[0..windows.len]) |seg| { + seg.cache.tap = .{ .buf = taps, .off = row * self.tapBytes(), .start = seg.cache.from, .rows = seg.rows, .gen = self.gen }; + row += seg.rows; + } + } + const logits = self.scratch.logits.slice(u16, total * cfg.vocab); + const picks = self.scratch.picks.slice(Pick, total); + var base: usize = 0; + for (windows, out) |w, o| { + if (o.sampled.len != w.rows() or o.drafts.len != w.tokens.len) return error.InvalidGemmaOutputs; + const greedy = w.masks.len == 0 and (w.stream.sampling == null or w.stream.sampling.?.temperature <= 0); + for (o.sampled, w.positions, 0..) |*t, p, r| { + if (greedy) { // the GPU's argmax, the same token the host draw finds + const pick = picks[base + r]; + if (pick.nonfinite != 0 or pick.token >= cfg.vocab) return error.NonfiniteGemmaLogits; + t.* = pick.token; + continue; + } + const mask = if (w.masks.len > 0) w.masks[r * mask_words ..][0..mask_words] else &.{}; + t.* = try draw(self.gpa, logits[(base + r) * cfg.vocab ..][0..cfg.vocab], w.stream.sampling, p, mask); + } + @memcpy(o.drafts, w.tokens); + base += w.rows(); + } + } + + /// A chain's kept rows: the cache rolls back to where the verify began plus them; later rows are rewritten. + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) !void { + const self = cast(ptr); + for (windows, paths) |w, path| { + const cache = try self.cacheOf(w.stream); + for (path, 0..) |r, j| if (r != j) return error.UnsupportedGemmaTree; + if (cache.from + path.len > cache.len) return error.InvalidGemmaKeep; + cache.len = cache.from + path.len; + } + } + + fn draftFn(_: *anyopaque, requests: []const be.DraftRequest) !void { + for (requests) |r| if (r.depth != 0) return error.GemmaCheckpointHasNoDraftHead; + } + + /// The prompt cache's Snapshots functions: a stream's caches at a mark, copied on the GPU. + pub fn snapBytes(ptr: *anyopaque, at: u32) u64 { + return snapshot.bytes(cast(ptr).model.config.geo, at); + } + + pub fn snapSave(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const self = cast(ptr); + const cache = try self.cacheOf(@ptrCast(@alignCast(owner orelse return error.NoStream))); + if (cache.len != at) return error.SnapshotOutOfStep; + const snap = try self.gpa.create(snapshot.Snap); + errdefer self.gpa.destroy(snap); + snap.* = try snapshot.save(self.gpa, self.model.device, self.model.queue, cache); + return snap; + } + + pub fn snapRestore(_: *anyopaque, _: ?*anyopaque, _: *anyopaque) anyerror!void { + return error.BackendRestores; // the prompt pass restores, before its first chunk + } + + pub fn snapDrop(ptr: *anyopaque, saved: *anyopaque) void { + const snap: *snapshot.Snap = @ptrCast(@alignCast(saved)); + snap.free(); + cast(ptr).gpa.destroy(snap); + } + + fn releaseFn(ptr: *anyopaque, stream: *lanes.Stream) void { + const self = cast(ptr); + if (self.caches.fetchRemove(stream)) |entry| { + entry.value.deinit(); + self.gpa.destroy(entry.value); + } + } +}; + +/// Total rows a shared round is timed at beyond one window (the forward's cost grows past 16 rows: tensor-op compute). +const shared_widths = [_]usize{ 32, 48, 64, 96, 128 }; + +/// A row's allowed-token bits (structured output): one bit a logit. +const mask_words = (cfg.vocab + 31) / 32; + +const Pick = @import("tree_round_gpu").contract.Pick; diff --git a/zig/src/families/gemma4/config.zig b/zig/src/families/gemma4/config.zig new file mode 100644 index 0000000000..70f5b04163 --- /dev/null +++ b/zig/src/families/gemma4/config.zig @@ -0,0 +1,335 @@ +//! Admission for dense Gemma 4 MLX affine checkpoints (31B and 12B): pinned geometries and widths, each refusal named. +const std = @import("std"); +pub const widths = @import("widths.zig"); + +pub const vocab = 262144; +pub const sliding_dim = 256; +pub const global_dim = 512; +pub const window = 1024; +pub const group = widths.group; +pub const eps: f32 = 1e-6; +pub const softcap: f32 = 30.0; +pub const sliding_theta: f64 = 10000.0; +pub const global_theta: f64 = 1000000.0; +/// Pairs the global layers' proportional RoPE rotates: partial_rotary_factor 0.25 of 512 dims, in pairs. +pub const global_rotated = 64; + +/// A model's sizes; everything else (head dims, window, vocab, RoPE) is the same in each admitted one. +pub const Geometry = struct { + name: []const u8, + hidden: usize, + layers: usize, + intermediate: usize, + query_heads: usize, + sliding_kv_heads: usize, + global_kv_heads: usize, + + /// The 12B's unified checkpoint: its image input is the patch embedder, not the 31B's tower. + pub fn unified(g: Geometry) bool { + return std.mem.eql(u8, g.name, unified12b.name); + } + + /// Layer `index`'s key heads. + pub fn kvHeads(g: Geometry, index: usize) usize { + return if (global(index)) g.global_kv_heads else g.sliding_kv_heads; + } + + /// The widest q a layer writes: every query head at the global head dim. + pub fn queryWidth(g: Geometry) usize { + return g.query_heads * global_dim; + } + + /// The widest k (or v) a layer writes, of a sliding and a global layer. + pub fn kvWidth(g: Geometry) usize { + return @max(g.sliding_kv_heads * sliding_dim, g.global_kv_heads * global_dim); + } +}; + +/// Gemma 4 31B (model_type gemma4). +pub const dense31b: Geometry = .{ .name = "31B", .hidden = 5376, .layers = 60, .intermediate = 21504, .query_heads = 32, .sliding_kv_heads = 16, .global_kv_heads = 4 }; +/// Gemma 4 12B's text model (model_type gemma4_unified, text_config gemma4_unified_text). +pub const unified12b: Geometry = .{ .name = "12B", .hidden = 3840, .layers = 48, .intermediate = 15360, .query_heads = 16, .sliding_kv_heads = 8, .global_kv_heads = 1 }; +/// Layers an admitted geometry has at most: the size of per-layer arrays. +pub const max_layers = 60; + +comptime { + for ([_]Geometry{ dense31b, unified12b }) |g| std.debug.assert(g.layers <= max_layers); +} + +pub const Error = error{ + BadGemmaConfig, + GemmaMoeNotServed, + GemmaUnifiedNotServed, + GemmaPerLayerInputsNotServed, + GemmaSharedKvNotServed, + GemmaKeysNotValues, + UnsupportedGemmaGeometry, + UnsupportedGemmaRotary, + UnsupportedGemmaQuantization, +} || widths.Error || std.json.ParseError(std.json.Scanner); + +pub const Config = struct { + geo: Geometry, + /// Each linear's MLX affine width (the config's quantization table). + widths: widths.Widths = .{}, + context: usize, + eos: [4]u32 = .{ 0, 0, 0, 0 }, + eos_count: usize = 0, + + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } +}; + +/// Whether layer `index` is a global (full attention) layer: every sixth, and the last. +pub fn global(index: usize) bool { + return index % 6 == 5; +} + +/// A layer's head dim and key heads. +pub fn headDim(index: usize) usize { + return if (global(index)) global_dim else sliding_dim; +} + +/// RoPE inverse frequencies for a layer kind, as float32 [dim / 2]: zero for pairs proportional RoPE leaves alone. +pub fn inverseFrequencies(is_global: bool, out: []f32) void { + const dim: f64 = if (is_global) global_dim else sliding_dim; + const theta = if (is_global) global_theta else sliding_theta; + for (out, 0..) |*f, i| { + const rotated = !is_global or i < global_rotated; + f.* = if (rotated) @floatCast(1.0 / std.math.pow(f64, theta, 2.0 * @as(f64, @floatFromInt(i)) / dim)) else 0.0; + } +} + +fn object(v: std.json.Value) Error!std.json.ObjectMap { + return if (v == .object) v.object else error.BadGemmaConfig; +} + +fn field(o: std.json.ObjectMap, key: []const u8) Error!std.json.Value { + return o.get(key) orelse error.BadGemmaConfig; +} + +fn number(v: std.json.Value) Error!f64 { + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadGemmaConfig, + }; +} + +fn integer(o: std.json.ObjectMap, key: []const u8) Error!i64 { + const v = try field(o, key); + return if (v == .integer) v.integer else error.BadGemmaConfig; +} + +fn optionalInt(o: std.json.ObjectMap, key: []const u8) i64 { + const v = o.get(key) orelse return 0; + return if (v == .integer) v.integer else 0; +} + +fn isString(v: ?std.json.Value, expected: []const u8) bool { + const s = v orelse return false; + return s == .string and std.mem.eql(u8, s.string, expected); +} + +fn isTrue(o: std.json.ObjectMap, key: []const u8) bool { + const v = o.get(key) orelse return false; + return v == .bool and v.bool; +} + +/// The widths of `quantization`, and of `quantization_config` when present, which must name the same ones. +fn quantization(gpa: std.mem.Allocator, root: std.json.ObjectMap, layers: usize) Error!widths.Widths { + const w = try widths.read(gpa, try object(root.get("quantization") orelse return error.UnsupportedGemmaQuantization), layers); + if (root.get("quantization_config")) |q| { + if (!std.meta.eql(w, try widths.read(gpa, try object(q), layers))) return error.UnsupportedGemmaQuantization; + } + return w; +} + +fn rope(o: std.json.ObjectMap, kind: []const u8, kind_type: []const u8, theta: f64, partial: ?f64) Error!void { + const r = try object(try field(o, kind)); + if (!isString(r.get("rope_type"), kind_type)) return error.UnsupportedGemmaRotary; + if (try number(try field(r, "rope_theta")) != theta) return error.UnsupportedGemmaRotary; + if (partial) |p| if (try number(try field(r, "partial_rotary_factor")) != p) return error.UnsupportedGemmaRotary; +} + +fn eosIds(v: std.json.Value, c: *Config) void { + switch (v) { + .integer => |i| if (c.eos_count < c.eos.len) { + c.eos[c.eos_count] = @intCast(i); + c.eos_count += 1; + }, + .array => |a| for (a.items) |x| eosIds(x, c), + else => {}, + } +} + +pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) Error!Config { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const root = try object(parsed.value); + const unified = isString(root.get("model_type"), "gemma4_unified"); + if (!unified and !isString(root.get("model_type"), "gemma4")) return error.BadGemmaConfig; + const text = try object(try field(root, "text_config")); + if (unified and !isString(text.get("model_type"), "gemma4_unified_text")) return error.GemmaUnifiedNotServed; + const geo = if (unified) unified12b else dense31b; + if (isTrue(text, "enable_moe_block")) return error.GemmaMoeNotServed; + if (optionalInt(text, "hidden_size_per_layer_input") != 0) return error.GemmaPerLayerInputsNotServed; + if (optionalInt(text, "num_kv_shared_layers") != 0) return error.GemmaSharedKvNotServed; + if (!isTrue(text, "attention_k_eq_v")) return error.GemmaKeysNotValues; + if (isTrue(text, "use_double_wide_mlp") or isTrue(text, "attention_bias")) return error.UnsupportedGemmaGeometry; + if (!isString(text.get("hidden_activation"), "gelu_pytorch_tanh")) return error.UnsupportedGemmaGeometry; + if (!isTrue(text, "tie_word_embeddings")) return error.UnsupportedGemmaGeometry; + const dims = .{ + .{ "hidden_size", geo.hidden }, .{ "vocab_size", vocab }, .{ "num_hidden_layers", geo.layers }, + .{ "intermediate_size", geo.intermediate }, .{ "num_attention_heads", geo.query_heads }, .{ "num_key_value_heads", geo.sliding_kv_heads }, + .{ "head_dim", sliding_dim }, .{ "global_head_dim", global_dim }, .{ "num_global_key_value_heads", geo.global_kv_heads }, + .{ "sliding_window", window }, + }; + inline for (dims) |d| if (try integer(text, d[0]) != @as(i64, @intCast(d[1]))) return error.UnsupportedGemmaGeometry; + if (try number(try field(text, "rms_norm_eps")) != 1e-6) return error.UnsupportedGemmaGeometry; + if (try number(try field(text, "final_logit_softcapping")) != softcap) return error.UnsupportedGemmaGeometry; + const types = try field(text, "layer_types"); + if (types != .array or types.array.items.len != geo.layers) return error.UnsupportedGemmaGeometry; + for (types.array.items, 0..) |kind, i| { + if (!isString(kind, if (global(i)) "full_attention" else "sliding_attention")) return error.UnsupportedGemmaGeometry; + } + const ropes = try object(try field(text, "rope_parameters")); + try rope(ropes, "sliding_attention", "default", sliding_theta, null); + try rope(ropes, "full_attention", "proportional", global_theta, 0.25); + const ws = try quantization(gpa, root, geo.layers); + const context: usize = @intCast(try integer(text, "max_position_embeddings")); + if (context < 1 or context > 262144) return error.UnsupportedGemmaGeometry; + var c: Config = .{ .geo = geo, .widths = ws, .context = context }; + eosIds(root.get("eos_token_id") orelse try field(text, "eos_token_id"), &c); + if (c.eos_count == 0) return error.BadGemmaConfig; + return c; +} + +const fixture = @embedFile("fixtures/config.json"); +/// The 4-bit 12B's config as tools/zig/gemma4_12b_convert.py writes it: unified model types (the text model's part). +const fixture_12b = @embedFile("fixtures/config_12b.json"); + +fn edited(gpa: std.mem.Allocator, from: []const u8, to: []const u8) ![]u8 { + return std.mem.replaceOwned(u8, gpa, fixture, from, to); +} + +fn edited12b(gpa: std.mem.Allocator, edits: []const [2][]const u8) ![]u8 { + var bytes = try gpa.dupe(u8, fixture_12b); + for (edits) |e| { + const next = try std.mem.replaceOwned(u8, gpa, bytes, e[0], e[1]); + gpa.free(bytes); + bytes = next; + } + return bytes; +} + +test "the 31B 4-bit checkpoint's config is admitted, with its context and end tokens" { + const c = try parse(std.testing.allocator, fixture); + try std.testing.expectEqualStrings("31B", c.geo.name); + try std.testing.expectEqual(dense31b, c.geo); + try std.testing.expectEqual(@as(usize, 262144), c.context); + try std.testing.expectEqualSlices(u32, &.{ 1, 106, 50 }, c.eos[0..c.eos_count]); +} + +test "the 12B's unified text model is admitted with its own geometry" { + const c = try parse(std.testing.allocator, fixture_12b); + try std.testing.expectEqualStrings("12B", c.geo.name); + try std.testing.expectEqual(unified12b, c.geo); + try std.testing.expectEqual(@as(usize, 48), c.geo.layers); + try std.testing.expectEqual(@as(usize, 262144), c.context); + try std.testing.expectEqualSlices(u32, &.{ 1, 106, 50 }, c.eos[0..c.eos_count]); +} + +test "a model type with the other geometry's sizes is refused" { + const gpa = std.testing.allocator; + const dense12b = try edited12b(gpa, &.{.{ "\"model_type\": \"gemma4_unified\",", "\"model_type\": \"gemma4\"," }}); + defer gpa.free(dense12b); + try std.testing.expectError(error.UnsupportedGemmaGeometry, parse(gpa, dense12b)); + const unified31b = try edited(gpa, "\"model_type\": \"gemma4\"", "\"model_type\": \"gemma4_unified\""); + defer gpa.free(unified31b); + const unified31b_text = try std.mem.replaceOwned(u8, gpa, unified31b, "\"gemma4_text\"", "\"gemma4_unified_text\""); + defer gpa.free(unified31b_text); + try std.testing.expectError(error.UnsupportedGemmaGeometry, parse(gpa, unified31b_text)); + const sixty = try edited12b(gpa, &.{.{ "\"num_hidden_layers\": 48", "\"num_hidden_layers\": 60" }}); + defer gpa.free(sixty); + try std.testing.expectError(error.UnsupportedGemmaGeometry, parse(gpa, sixty)); + const wide = try edited12b(gpa, &.{.{ "\"num_key_value_heads\": 8", "\"num_key_value_heads\": 16" }}); + defer gpa.free(wide); + try std.testing.expectError(error.UnsupportedGemmaGeometry, parse(gpa, wide)); +} + +test "each other unified layout is refused by name" { + const gpa = std.testing.allocator; + const cases = .{ + .{ "\"model_type\": \"gemma4_unified_text\"", "\"model_type\": \"gemma4_text\"", error.GemmaUnifiedNotServed }, + .{ "\"hidden_size_per_layer_input\": 0", "\"hidden_size_per_layer_input\": 256", error.GemmaPerLayerInputsNotServed }, + .{ "\"num_kv_shared_layers\": 0", "\"num_kv_shared_layers\": 18", error.GemmaSharedKvNotServed }, + .{ "\"enable_moe_block\": false", "\"enable_moe_block\": true", error.GemmaMoeNotServed }, + .{ "\"bits\": 4", "\"bits\": 7", error.UnsupportedGemmaQuantization }, + }; + inline for (cases) |c| { + const bytes = try edited12b(gpa, &.{.{ c[0], c[1] }}); + defer gpa.free(bytes); + try std.testing.expect(!std.mem.eql(u8, bytes, fixture_12b)); + try std.testing.expectError(c[2], parse(gpa, bytes)); + } +} + +test "each other Gemma 4 layout is refused by name" { + const gpa = std.testing.allocator; + const cases = .{ + .{ "\"enable_moe_block\": false", "\"enable_moe_block\": true", error.GemmaMoeNotServed }, + .{ "\"model_type\": \"gemma4\"", "\"model_type\": \"gemma4_unified\"", error.GemmaUnifiedNotServed }, + .{ "\"attention_k_eq_v\": true", "\"attention_k_eq_v\": false", error.GemmaKeysNotValues }, + .{ "\"num_kv_shared_layers\": 0", "\"num_kv_shared_layers\": 18", error.GemmaSharedKvNotServed }, + .{ "\"hidden_size_per_layer_input\": 0", "\"hidden_size_per_layer_input\": 256", error.GemmaPerLayerInputsNotServed }, + .{ "\"sliding_window\": 1024", "\"sliding_window\": 512", error.UnsupportedGemmaGeometry }, + }; + inline for (cases) |c| { + const bytes = try edited(gpa, c[0], c[1]); + defer gpa.free(bytes); + try std.testing.expect(!std.mem.eql(u8, bytes, fixture)); + try std.testing.expectError(c[2], parse(gpa, bytes)); + } +} + +test "other widths are admitted per module, and quantization_config must agree" { + const gpa = std.testing.allocator; + const eight = try edited12b(gpa, &.{.{ "\"bits\": 4", "\"bits\": 8" }}); // both tables + defer gpa.free(eight); + const c8 = try parse(gpa, eight); + try std.testing.expectEqual(@as(u8, 8), c8.widths.of(47, .down).bits); + try std.testing.expect(!c8.widths.allFast(c8.geo.layers)); + try std.testing.expect((try parse(gpa, fixture_12b)).widths.allFast(48)); + const entry = "\"group_size\": 64, \"language_model.model.layers.0.mlp.down_proj\": {\"bits\": 8, \"group_size\": 64},"; + const one = try edited(gpa, "\"quantization\": {\n \"group_size\": 64,", "\"quantization\": {" ++ entry); + defer gpa.free(one); + try std.testing.expect(!std.mem.eql(u8, one, fixture)); + try std.testing.expectError(error.UnsupportedGemmaQuantization, parse(gpa, one)); // quantization_config disagrees + const both = try edited(gpa, "\"group_size\": 64,", entry); + defer gpa.free(both); + const c = try parse(gpa, both); + try std.testing.expectEqual(@as(u8, 8), c.widths.of(0, .down).bits); + try std.testing.expectEqual(@as(u8, 4), c.widths.of(1, .down).bits); +} + +test "layer kinds, head dims and the proportional RoPE table" { + try std.testing.expect(!global(0) and global(5) and global(59) and !global(58)); + try std.testing.expectEqual(@as(usize, 512), headDim(11)); + try std.testing.expectEqual(@as(usize, 16), dense31b.kvHeads(1)); + try std.testing.expectEqual(@as(usize, 1), unified12b.kvHeads(47)); + try std.testing.expectEqual(@as(usize, 4096), dense31b.kvWidth()); + try std.testing.expectEqual(@as(usize, 2048), unified12b.kvWidth()); + var g: [global_dim / 2]f32 = undefined; + inverseFrequencies(true, &g); + try std.testing.expectEqual(@as(f32, 1.0), g[0]); + try std.testing.expect(g[63] > 0 and g[64] == 0 and g[255] == 0); + var s: [sliding_dim / 2]f32 = undefined; + inverseFrequencies(false, &s); + try std.testing.expect(s[127] > 0 and s[127] < s[1]); +} diff --git a/zig/src/families/gemma4/dflash.zig b/zig/src/families/gemma4/dflash.zig new file mode 100644 index 0000000000..8ec994a756 --- /dev/null +++ b/zig/src/families/gemma4/dflash.zig @@ -0,0 +1,582 @@ +//! z-lab's DFlash drafter on the target's taps: a stream's context keys, a block forward of up to 16 rows, held chains. +const std = @import("std"); +const mtl = @import("metal"); +const lanes = @import("lanes"); +const cfg = @import("config.zig"); +const dw = @import("dflash_weights.zig"); +const lane = @import("../../core/lane_projection.zig"); +const ops = @import("draft_ops"); +const attention = @import("../../core/shared_attention.zig"); +const round = @import("tree_round_gpu"); +const weights = @import("weights.zig"); +const fwd = @import("forward.zig"); +const Model = @import("model.zig").Model; +const snapshot = @import("snapshot.zig"); + +const dr = lanes.drafter; +const Stream = lanes.Stream; +const Ref = lane.Ref; +const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const private = mtl.ResourceOptions.private | mtl.ResourceOptions.untracked; + +/// A block's rows at most (the pending token and 15 drafts); the rows one batched hold or absorb pass takes. +pub const block_rows = 16; +const absorb_rows = lane.reg_rows; +const hold_rows = lane.reg_rows; +/// Streams one hold or absorb call takes, and host-written rows one absorb call takes before it flushes. +const max_items = 32; +const pos_rows = 1024; + +pub const Options = struct { capacity: usize }; + +/// One stream's drafter state: context keys and values a layer, and the chain it holds. +const Lane = struct { + k: [8]mtl.Buffer, + v: [8]mtl.Buffer, + capacity: usize, // positions a full layer's cache holds + end: u64 = 0, // context rows absorbed: positions [0, end) + held: [block_rows]u32 = @splat(0), + held_len: u32 = 0, +}; + +pub const Drafter = struct { + gpa: std.mem.Allocator, + io: std.Io, + target: *Model, + config: dw.Config, + w: dw.Weights, + options: Options, + ops: ops.Ops, + attention: attention.Attention, + picks: round.Ops, + projections: weights.Projections, + lanes: std.AutoHashMapUnmanaged(*Stream, *Lane) = .empty, + step_ms: f64 = 0, + block_ms: [2]f64 = .{ 0, 0 }, // a 1-draft and a 15-draft block's wall time at load + trace: bool = false, // GEMMA_DRAFT_TRACE: each hold's anchor, depth and drafts on stderr + /// An absorb's command buffer, committed without a wait: the hold that follows (the same queue) runs after it. + pending: ?mtl.CommandBuffer = null, + // the block forward's scratch (one block at a time; items of a call run in order on one command buffer) + ids: mtl.Buffer, + pos: mtl.Buffer, + apos: mtl.Buffer, // an absorb's positions, apart from a block's (its command buffer may still be running) + xs: mtl.Buffer, // the next projection's group sums, written by the kernel that made its input + mask: mtl.Buffer, + ones: mtl.Buffer, + pick_buf: mtl.Buffer, + h: mtl.Buffer, + x: mtl.Buffer, + qkv: mtl.Buffer, + q: mtl.Buffer, + q2: mtl.Buffer, + k16: mtl.Buffer, + k16r: mtl.Buffer, + attn: mtl.Buffer, + o: mtl.Buffer, + gate_up: mtl.Buffer, + act: mtl.Buffer, + logits: mtl.Buffer, + // the absorb scratch: 128 context rows at a time + fcx: mtl.Buffer, + hctx: mtl.Buffer, + kv: mtl.Buffer, + kn: mtl.Buffer, + knr: mtl.Buffer, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, target: *Model, dir: []const u8, options: Options) !*Drafter { + const d = try gpa.create(Drafter); + errdefer gpa.destroy(d); + d.* = undefined; + d.gpa = gpa; + d.io = io; + d.target = target; + d.options = options; + d.projections = .{ .geo = target.config.geo }; + d.lanes = .empty; + d.step_ms = 0; + d.block_ms = .{ 0, 0 }; + d.pending = null; + d.trace = std.c.getenv("GEMMA_DRAFT_TRACE") != null; + errdefer d.projections.deinit(gpa); // warm's blocks fill it + d.config = try dw.Config.read(gpa, io, dir, target.config.geo); + const c = &d.config; + d.w = try dw.load(gpa, target.device, io, dir, c, target.config.geo); + errdefer d.w.deinit(gpa); + d.ops = try ops.Ops.init(target.device); + errdefer d.ops.deinit(); + const keys_max = @max(options.capacity + block_rows, c.window - 1 + block_rows); + d.attention = try attention.Attention.init(gpa, target.device, .{ .query_heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.head_dim, .rotary_dim = c.head_dim, .capacity = keys_max, .output = .bf16 }); + errdefer d.attention.deinit(); + d.picks = try round.Ops.init(target.device); + errdefer d.picks.deinit(); + const dev = target.device; + const hid = c.hidden; + const qw = c.heads * c.head_dim; + const kw = c.kv_heads * c.head_dim; + const r = hold_rows; + var made: usize = 0; + errdefer d.release(made); + inline for (.{ .{ "ids", max_items * block_rows * 4, true }, .{ "pos", max_items * block_rows * 8, true }, .{ "apos", pos_rows * 8, true }, .{ "xs", (c.tap_count * hid / 64) * absorb_rows * 4, false }, .{ "mask", max_items * block_rows * (c.window - 1 + block_rows), true }, .{ "ones", block_rows * keys_max, true }, .{ "pick_buf", max_items * block_rows * @sizeOf(round.contract.Pick), true }, .{ "h", r * hid * 2, false }, .{ "x", r * hid * 2, false }, .{ "qkv", r * (qw + 2 * kw) * 2, false }, .{ "q", r * qw * 2, false }, .{ "q2", r * qw * 2, false }, .{ "k16", r * kw * 2, false }, .{ "k16r", r * kw * 2, false }, .{ "attn", r * qw * 2, false }, .{ "o", r * hid * 2, false }, .{ "gate_up", r * 2 * c.intermediate * 2, false }, .{ "act", r * c.intermediate * 2, false }, .{ "logits", r * cfg.vocab * 2, false }, .{ "fcx", absorb_rows * hid * 2, false }, .{ "hctx", absorb_rows * hid * 2, d.trace }, .{ "kv", absorb_rows * 2 * kw * 2, false }, .{ "kn", absorb_rows * kw * 2, false }, .{ "knr", absorb_rows * kw * 2, false } }) |f| { + @field(d, f[0]) = try dev.buffer(f[1], if (f[2]) shared else private); + made += 1; + } + @memset(d.ones.slice(u8, block_rows * keys_max), 0xff); + try d.warm(io); + return d; + } + + pub fn deinit(d: *Drafter) void { + d.settle() catch {}; + var it = d.lanes.valueIterator(); + while (it.next()) |l| d.dropLane(l.*); + d.lanes.deinit(d.gpa); + d.release(std.math.maxInt(usize)); + d.projections.deinit(d.gpa); + d.picks.deinit(); + d.attention.deinit(); + d.ops.deinit(); + d.w.deinit(d.gpa); + d.gpa.destroy(d); + } + + /// Releases the first `made` buffers in init's order, so a failed init frees what it made. + fn release(d: *Drafter, made: usize) void { + var i: usize = 0; + inline for (.{ "ids", "pos", "apos", "xs", "mask", "ones", "pick_buf", "h", "x", "qkv", "q", "q2", "k16", "k16r", "attn", "o", "gate_up", "act", "logits", "fcx", "hctx", "kv", "kn", "knr" }) |f| { + if (i < made) @field(d, f).deinit(); + i += 1; + } + } + + pub fn drafter(d: *Drafter) dr.Drafter { + return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn } }; + } + + fn cast(ptr: *anyopaque) *Drafter { + return @ptrCast(@alignCast(ptr)); + } + + fn tapsFn(ptr: *anyopaque) []const u32 { + const d = cast(ptr); + return d.config.taps[0..d.config.tap_count]; + } + + /// Acceptance by depth, measured on the Python reference (114 rounds, 3.5 a round), until a stream has its own. + const prior = [_]f64{ 0.77, 0.64, 0.71, 0.82, 0.84, 0.81, 0.9, 0.9, 0.9, 0.9, 0.83, 0.8, 0.83, 0.8, 0.8 }; + + /// Chains of up to 15, always drafted, no plain guard: a block costs the same at any depth, a guard bills per row. + fn factsFn(ptr: *anyopaque) dr.Facts { + const d = cast(ptr); + return .{ .depth = block_rows - 1, .step_ms = d.step_ms, .prior = &prior, .plain_guard = false, .batched = true, .stream_bytes = d.laneBytes(d.options.capacity) }; + } + + fn laneBytes(d: *const Drafter, capacity: usize) usize { + const c = &d.config; + var bytes: usize = 0; + for (c.sliding[0..c.layers]) |sliding| bytes += 2 * (if (sliding) c.window - 1 + block_rows else capacity + block_rows) * c.kv_heads * c.head_dim * 2; + return bytes; + } + + fn openFn(ptr: *anyopaque, s: *Stream) anyerror!void { + const d = cast(ptr); + if (d.lanes.contains(s)) return error.DraftStreamOpen; + const capacity = @min(d.options.capacity, (s.prompt_len + s.max_new + 2 * block_rows + 63) / 64 * 64); + const l = try d.newLane(capacity); + errdefer d.dropLane(l); + try d.lanes.put(d.gpa, s, l); + } + + fn newLane(d: *Drafter, capacity: usize) !*Lane { + const c = &d.config; + const l = try d.gpa.create(Lane); + errdefer d.gpa.destroy(l); + l.* = .{ .k = undefined, .v = undefined, .capacity = capacity }; + var made: usize = 0; + errdefer for (0..made) |i| { + l.k[i].deinit(); + l.v[i].deinit(); + }; + for (0..c.layers) |i| { + const bytes = (if (c.sliding[i]) c.window - 1 + block_rows else capacity + block_rows) * c.kv_heads * c.head_dim * 2; + l.k[i] = try d.target.device.buffer(bytes, private); + l.v[i] = d.target.device.buffer(bytes, private) catch |e| { + l.k[i].deinit(); + return e; + }; + made += 1; + } + return l; + } + + fn dropLane(d: *Drafter, l: *Lane) void { + for (0..d.config.layers) |i| { + l.k[i].deinit(); + l.v[i].deinit(); + } + d.gpa.destroy(l); + } + + fn releaseFn(ptr: *anyopaque, s: *Stream) void { + const d = cast(ptr); + if (d.lanes.fetchRemove(s)) |e| d.dropLane(e.value); + } + + fn heldFn(ptr: *anyopaque, streams: []const *Stream, out: []const []u32) anyerror!void { + const d = cast(ptr); + for (streams, out) |s, o| { + const l = d.lanes.get(s) orelse return error.DraftStreamUnknown; + if (o.len > l.held_len) return error.DraftHeldFewer; + @memcpy(o, l.held[0..o.len]); + } + } + + /// A projection; `summed`: its input's group sums are in d.xs (stride rows), written by the input's own kernel. + fn project(d: *Drafter, e: mtl.ComputeEncoder, l: dw.Linear, x: Ref, y: Ref, rows: usize, summed: bool) !void { + if (l.bf16) return d.ops.linear(e, .{ .buf = x.buf, .off = x.off }, .{ .buf = l.w }, .{ .buf = y.buf, .off = y.off }, .{ .rows = @intCast(rows), .n = @intCast(l.n), .k = @intCast(l.k), .x_stride = @intCast(l.k), .w_stride = @intCast(l.k), .y_stride = @intCast(l.n) }); + const p = try d.projections.get(d.gpa, d.target.device, .{ .w = l.w, .sb = l.sb, .n = l.n, .k = l.k }); + if (summed) return p.encodeSums(e, x, .{ .buf = l.w }, .{ .buf = l.sb }, .{ .ref = .{ .buf = d.xs }, .stride = @intCast(rows) }, y, @intCast(rows)); + try p.encode(e, x, .{ .buf = l.w }, .{ .buf = l.sb }, y, @intCast(rows)); + } + + /// A hidden-wide RMS norm (w times the normalized row, the target's kernel) writing the next projection's sums. + fn normSums(d: *Drafter, e: mtl.ComputeEncoder, x: Ref, gain: Ref, y: Ref, rows: usize) void { + e.setPipeline(d.target.kernels.sliding.rms); + e.setBuffer(x.buf, x.off, 0); + e.setBuffer(gain.buf, gain.off, 1); + e.setBuffer(y.buf, y.off, 2); + e.setValue(d.config.eps, 3); + e.setBuffer(d.xs, 0, 4); + e.setValue(@as(i32, @intCast(rows)), 5); + e.dispatchGroups(.{ .width = rows }, .{ .width = 1024 }); + } + + /// Bytes a lane's context at `at` takes beside a kept prompt state: the sliding rings whole, the full layer's rows. + pub fn keptBytes(d: *const Drafter, at: usize) u64 { + const c = &d.config; + const row: u64 = c.kv_heads * c.head_dim * 2; + var total: u64 = 0; + for (c.sliding[0..c.layers]) |sliding| total += 2 * row * (if (sliding) c.window - 1 + block_rows else at); + return total; + } + + /// Stream `s`'s context copied into `snap` when its lane stands at the snapshot's position (a prompt pass's mark). + pub fn keepLane(d: *Drafter, s: *Stream, snap: *snapshot.Snap) !void { + const l = d.lanes.get(s) orelse return; + try d.settle(); // the chunk's absorb, left pending + if (l.end != snap.at) return; + const buf = try d.target.device.buffer(@intCast(d.keptBytes(snap.at)), private); + errdefer buf.deinit(); + try d.copyLane(l, snap.at, buf, true); + snap.draft = buf; + } + + /// The lane takes the context kept with the state its stream's prompt pass resumed at `start`. + fn resumeLane(d: *Drafter, s: *Stream, l: *Lane, start: u64) !void { + const saved = s.reuse.saved orelse return error.DraftContextGap; + const snap: *const snapshot.Snap = @ptrCast(@alignCast(saved)); + if (s.cached != start or snap.at != start or snap.draft == null or l.end != 0) return error.DraftContextGap; + try d.copyLane(l, start, snap.draft.?, false); + l.end = start; + } + + /// A lane's keys and values and a kept buffer, one blit pass, waited on (`out`: lane to buffer). + fn copyLane(d: *Drafter, l: *const Lane, at: usize, buf: mtl.Buffer, out: bool) !void { + const c = &d.config; + const row: usize = c.kv_heads * c.head_dim * 2; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = d.target.queue.commandBuffer(); + const e = cb.blit(); + var to: usize = 0; + for (0..c.layers) |i| for ([_]mtl.Buffer{ l.k[i], l.v[i] }) |live| { + const n = row * (if (c.sliding[i]) c.window - 1 + block_rows else at); + if (n == 0) continue; + if (out) e.copy(live, 0, buf, to, n) else e.copy(buf, to, live, 0, n); + to += n; + }; + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| return failed(msg); + } + + /// Waits for a pending absorb's command buffer (a hold ordered after it on the queue has run it already). + fn settle(d: *Drafter) !void { + const cb = d.pending orelse return; + d.pending = null; + defer mtl.objc.release(cb.id); // retained when it was left pending (an autoreleased object otherwise) + cb.wait(); + if (cb.failure()) |msg| return failed(msg); + } + + fn norm(d: *Drafter, e: mtl.ComputeEncoder, x: Ref, gain: Ref, y: Ref, rows: usize, heads: u32, dim: u32, x_stride: u32, y_stride: u32) !void { + try d.ops.norm(e, .{ .buf = x.buf, .off = x.off }, .{ .buf = gain.buf, .off = gain.off }, .{ .buf = y.buf, .off = y.off }, .{ .rows = @intCast(rows), .heads = heads, .dim = dim, .x_stride = x_stride, .y_stride = y_stride, .eps = d.config.eps }); + } + + /// Rows of `x` (pitch `stride` values) into cache slots from `first`: position % ring on a ring, else the position. + fn store(d: *Drafter, e: mtl.ComputeEncoder, x: Ref, width: u32, stride: u32, cache: mtl.Buffer, rows: usize, first: u64, ring: u32) void { + e.setPipeline(d.target.kernels.sliding.draft_store); + e.setBuffer(x.buf, x.off, 0); + e.setBuffer(cache, 0, 1); + e.setValue([4]u32{ stride, width, @intCast(first), ring }, 2); + e.dispatchThreads(.{ .width = width, .height = rows }, .{ .width = 256 }); + } + + fn add(d: *Drafter, e: mtl.ComputeEncoder, h: mtl.Buffer, y: mtl.Buffer, count: usize) void { + e.setPipeline(d.target.kernels.sliding.draft_add); + e.setBuffer(h, 0, 0); + e.setBuffer(y, 0, 1); + e.dispatchThreads(.{ .width = count }, .{ .width = 256 }); + } + + /// Up to 128 rows of one absorb pass: each piece's rows (one stream, one source) sit at `row` of the pass. + const Piece = struct { lane: *Lane, src: Ref, rows: usize, p0: u64, row: usize }; + + /// Rows [start, start + rows) replace the context from `start`, in passes of pieces; committed, not awaited. + fn absorbFn(ptr: *anyopaque, items: []const dr.Absorb) anyerror!void { + const d = cast(ptr); + const c = &d.config; + const t0 = std.Io.Clock.awake.now(d.io); + defer if (d.trace) std.debug.print("absorb of {d} streams: {d:.2} ms\n", .{ items.len, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(d.io)).toNanoseconds())) / 1e6 }); + try d.settle(); // the positions buffer is rewritten below + var cb = d.target.queue.commandBuffer(); + var e = cb.compute(.serial); + var used: usize = 0; // host position rows written this command buffer + var pieces: [max_items]Piece = undefined; + var count: usize = 0; + var rows: usize = 0; + for (items) |it| { + const l = d.lanes.get(it.stream) orelse return error.DraftStreamUnknown; + const f = it.features; + if (f.space != .device or f.dtype != .bf16 or f.row_bytes != c.tap_count * c.hidden * 2) return error.DraftFeaturesUnsupported; + if (it.start > l.end) try d.resumeLane(it.stream, l, it.start); // a prompt pass that resumed a kept state + if (it.start + f.rows > l.capacity) return error.DraftContextGap; + const src = mtl.Buffer{ .id = @ptrFromInt(f.buffer) }; + var at: usize = 0; + while (at < f.rows) : (at += absorb_rows) { + const n = @min(absorb_rows, f.rows - at); + if (rows + n > absorb_rows or count == pieces.len) { // the pass is full: run it + if (used + rows > pos_rows) { // and the positions buffer with it: a new command buffer + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| return failed(msg); + cb = d.target.queue.commandBuffer(); + e = cb.compute(.serial); + used = 0; + } + try d.absorbPass(e, pieces[0..count], rows, used); + used += rows; + count = 0; + rows = 0; + } + pieces[count] = .{ .lane = l, .src = .{ .buf = src, .off = f.offset + at * f.row_bytes }, .rows = n, .p0 = it.start + at, .row = rows }; + count += 1; + rows += n; + } + if (d.trace) { + const first = src.slice(u16, (f.offset + f.row_bytes) / 2)[f.offset / 2 ..][0..4]; + std.debug.print("absorb {s}: rows [{d}, {d}) replace from end {d}; first tap bf16 {any}\n", .{ it.stream.id, it.start, it.start + f.rows, l.end, first.* }); + } + l.end = it.start + f.rows; + l.held_len = 0; + } + if (count > 0) { + if (used + rows > pos_rows) { + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| return failed(msg); + cb = d.target.queue.commandBuffer(); + e = cb.compute(.serial); + used = 0; + } + try d.absorbPass(e, pieces[0..count], rows, used); + } + e.end(); + cb.commit(); + _ = mtl.objc.retain(cb.id); + d.pending = cb; + } + + /// One pass: fc of each piece into its rows, then hidden_norm, kv, k_norm and RoPE over all rows, stores a piece. + fn absorbPass(d: *Drafter, e: mtl.ComputeEncoder, pieces: []const Piece, rows: usize, pos_at: usize) !void { + const c = &d.config; + const kw = c.kv_heads * c.head_dim; + const positions = d.apos.slice(u64, pos_rows)[pos_at..][0..rows]; + for (pieces) |pc| { + for (positions[pc.row..][0..pc.rows], 0..) |*p, i| p.* = pc.p0 + i; + try d.project(e, d.w.fc, pc.src, .{ .buf = d.fcx, .off = pc.row * c.hidden * 2 }, pc.rows, false); + } + const pos = Ref{ .buf = d.apos, .off = pos_at * 8 }; + d.normSums(e, .{ .buf = d.fcx }, d.w.hidden_norm, .{ .buf = d.hctx }, rows); + for (0..c.layers) |i| { + const lw = &d.w.layers[i]; + try d.project(e, lw.kv, .{ .buf = d.hctx }, .{ .buf = d.kv }, rows, true); + try d.norm(e, .{ .buf = d.kv }, lw.k_norm, .{ .buf = d.kn }, rows, c.kv_heads, c.head_dim, 2 * kw, kw); + try d.ops.rope(e, .{ .buf = d.kn }, .{ .buf = pos.buf, .off = pos.off }, .{ .buf = d.knr }, .{ .rows = @intCast(rows), .heads = c.kv_heads, .dim = c.head_dim, .x_stride = kw, .y_stride = kw, .theta = c.theta }); + const ring: u32 = if (c.sliding[i]) c.window - 1 else 0; + for (pieces) |pc| { + d.store(e, .{ .buf = d.knr, .off = pc.row * kw * 2 }, kw, kw, pc.lane.k[i], pc.rows, pc.p0, ring); + d.store(e, .{ .buf = d.kv, .off = (pc.row * 2 * kw + kw) * 2 }, kw, 2 * kw, pc.lane.v[i], pc.rows, pc.p0, ring); + } + } + } + + fn failed(msg: [*:0]const u8) anyerror { + std.log.err("gemma4 drafter: command buffer failed: {s}", .{msg}); + return error.GpuFailed; + } + + /// The sliding layers' mask for item `slot`: ring slots by the position each holds, the block causal. + fn slidingMask(d: *Drafter, slot: usize, anchor: u64, rows: usize) Ref { + const ring: u64 = d.config.window - 1; + const keys = ring + rows; + const base = slot * block_rows * (ring + block_rows); + const m = d.mask.slice(u8, max_items * block_rows * (ring + block_rows))[base..][0 .. rows * keys]; + for (0..rows) |i| { + const row = m[i * keys ..][0..keys]; + for (0..ring) |s| { + const held = anchor > s; // a slot past the anchor was never written + const p = if (held) anchor - 1 - ((anchor - 1 - s) % ring) else 0; + row[s] = @intFromBool(held and p + ring >= anchor + i); + } + for (0..rows) |j| row[ring + j] = @intFromBool(j <= i); + } + return .{ .buf = d.mask, .off = base }; + } + + /// A stream's block in a batch: its rows [row, row + depth + 1) of the pass, its slot's host regions. + const Block = struct { lane: *Lane, pending: u32, anchor: u64, depth: u32, slot: usize, row: usize }; + + /// A batch of blocks [pending, mask x depth] through the drafter together, attention a block; picks a row. + fn encodeBlocks(d: *Drafter, e: mtl.ComputeEncoder, blocks: []const Block, rows: usize) !void { + const c = &d.config; + const t = d.target; + const hid = c.hidden; + const qw = c.heads * c.head_dim; + const kw = c.kv_heads * c.head_dim; + const ids = d.ids.slice(u32, max_items * block_rows)[0..rows]; + const positions = d.pos.slice(u64, max_items * block_rows)[0..rows]; + var masks: [max_items]Ref = undefined; + for (blocks, 0..) |b, i| { + ids[b.row] = b.pending; + @memset(ids[b.row + 1 ..][0..b.depth], c.mask); + for (positions[b.row..][0 .. b.depth + 1], 0..) |*p, j| p.* = b.anchor + j; + masks[i] = d.slidingMask(b.slot, b.anchor, b.depth + 1); + } + e.setPipeline(t.kernels.sliding.embed); + e.setBuffer(d.ids, 0, 0); + e.setBuffer(t.weights.embed_w.buf, t.weights.embed_w.off, 1); + e.setBuffer(t.weights.embed_s.buf, t.weights.embed_s.off, 2); + e.setBuffer(t.weights.embed_b.buf, t.weights.embed_b.off, 3); + e.setBuffer(d.h, 0, 4); + e.setValue(fwd.embedScale(t.config.geo.hidden), 5); + e.dispatchThreads(.{ .width = hid / 2, .height = rows }, .{ .width = 256 }); + const scale = 1.0 / @sqrt(@as(f32, @floatFromInt(c.head_dim))); + const pos = Ref{ .buf = d.pos }; + for (0..c.layers) |i| { + const lw = &d.w.layers[i]; + const sliding = c.sliding[i]; + d.normSums(e, .{ .buf = d.h }, lw.input_norm, .{ .buf = d.x }, rows); + try d.project(e, lw.qkv, .{ .buf = d.x }, .{ .buf = d.qkv }, rows, true); + try d.norm(e, .{ .buf = d.qkv }, lw.q_norm, .{ .buf = d.q }, rows, c.heads, c.head_dim, qw + 2 * kw, qw); + try d.norm(e, .{ .buf = d.qkv, .off = qw * 2 }, lw.k_norm, .{ .buf = d.k16 }, rows, c.kv_heads, c.head_dim, qw + 2 * kw, kw); + try d.ops.rope(e, .{ .buf = d.q }, .{ .buf = pos.buf, .off = pos.off }, .{ .buf = d.q2 }, .{ .rows = @intCast(rows), .heads = c.heads, .dim = c.head_dim, .x_stride = qw, .y_stride = qw, .theta = c.theta }); + try d.ops.rope(e, .{ .buf = d.k16 }, .{ .buf = pos.buf, .off = pos.off }, .{ .buf = d.k16r }, .{ .rows = @intCast(rows), .heads = c.kv_heads, .dim = c.head_dim, .x_stride = kw, .y_stride = kw, .theta = c.theta }); + for (blocks, 0..) |b, bi| { + const n = b.depth + 1; + const first: u64 = if (sliding) c.window - 1 else b.anchor; + d.store(e, .{ .buf = d.k16r, .off = b.row * kw * 2 }, kw, kw, b.lane.k[i], n, first, 0); + d.store(e, .{ .buf = d.qkv, .off = (b.row * (qw + 2 * kw) + qw + kw) * 2 }, kw, qw + 2 * kw, b.lane.v[i], n, first, 0); + const keys: usize = @intCast(first + n); + const mask = if (sliding) masks[bi] else Ref{ .buf = d.ones }; + try d.attention.encode(e, n, keys, keys, scale, .{ .buf = d.q2, .off = b.row * qw * 2 }, .{ .buf = b.lane.k[i] }, .{ .buf = b.lane.v[i] }, .{ .buf = mask.buf, .off = mask.off }, .{ .buf = d.attn, .off = b.row * qw * 2 }); + } + try d.project(e, lw.o, .{ .buf = d.attn }, .{ .buf = d.o }, rows, false); + d.add(e, d.h, d.o, rows * hid); + d.normSums(e, .{ .buf = d.h }, lw.post_attention, .{ .buf = d.x }, rows); + try d.project(e, lw.gate_up, .{ .buf = d.x }, .{ .buf = d.gate_up }, rows, true); + e.setPipeline(t.kernels.sliding.draft_swiglu); + e.setBuffer(d.gate_up, 0, 0); + e.setBuffer(d.act, 0, 1); + e.setValue(@as(u32, c.intermediate), 2); + e.setBuffer(d.xs, 0, 3); + e.setValue(@as(i32, @intCast(rows)), 4); + e.dispatchGroups(.{ .width = c.intermediate / 64, .height = rows }, .{ .width = 64 }); + try d.project(e, lw.down, .{ .buf = d.act }, .{ .buf = d.o }, rows, true); + d.add(e, d.h, d.o, rows * hid); + } + d.normSums(e, .{ .buf = d.h }, d.w.norm, .{ .buf = d.x }, rows); + const tw = t.weights.head; + const head = try t.projections.get(t.gpa, t.device, tw); + if (tw.reg()) try head.encodeSums(e, .{ .buf = d.x }, .{ .buf = tw.w }, .{ .buf = tw.sb }, .{ .ref = .{ .buf = d.xs }, .stride = @intCast(rows) }, .{ .buf = d.logits }, @intCast(rows)) else try tw.encodeRows(head, e, .{ .buf = d.x }, .{ .buf = d.logits }, rows); + var first: usize = 0; + while (first < rows) : (first += 16) { + const n = @min(16, rows - first); + try d.picks.argmax(e, .{ .buf = d.logits, .off = first * cfg.vocab * 2 }, .{ .buf = d.pick_buf, .off = first * @sizeOf(round.contract.Pick) }, .{ .rows = @intCast(n), .vocab = cfg.vocab, .stride = cfg.vocab }); + } + } + + /// Each stream's chain of `depth` drafts after `pending` at `position`, the streams in batches of up to 128 rows. + fn holdFn(ptr: *anyopaque, items: []const dr.Hold) anyerror!void { + const d = cast(ptr); + if (items.len > max_items) return error.DraftTooManyStreams; + const t0 = std.Io.Clock.awake.now(d.io); + defer if (d.trace) std.debug.print("hold of {d} streams: {d:.2} ms\n", .{ items.len, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(d.io)).toNanoseconds())) / 1e6 }); + var at: usize = 0; + while (at < items.len) { + var blocks: [max_items]Block = undefined; + var count: usize = 0; + var rows: usize = 0; + while (at < items.len) : (at += 1) { + const it = items[at]; + const l = d.lanes.get(it.stream) orelse return error.DraftStreamUnknown; + if (it.depth == 0 or it.depth >= block_rows or it.position == 0) return error.DraftDepth; + const anchor = it.position - 1; + if (l.end != anchor or anchor + block_rows > l.capacity + block_rows) return error.DraftContextGap; + if (rows + it.depth + 1 > hold_rows) break; + blocks[count] = .{ .lane = l, .pending = it.pending, .anchor = anchor, .depth = it.depth, .slot = count, .row = rows }; + count += 1; + rows += it.depth + 1; + } + const cb = d.target.queue.commandBuffer(); + const e = cb.compute(.serial); + try d.encodeBlocks(e, blocks[0..count], rows); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| return failed(msg); + try d.settle(); + const picks = d.pick_buf.slice(round.contract.Pick, max_items * block_rows); + for (blocks[0..count], items[at - count .. at]) |b, it| { + for (picks[b.row + 1 ..][0..b.depth], 0..) |p, i| { + if (p.nonfinite != 0 or p.token >= cfg.vocab) return error.DraftNonfinite; + b.lane.held[i] = p.token; + } + b.lane.held_len = b.depth; + if (d.trace) std.debug.print("draft {s} row {d} anchor {d} pending {d} depth {d}: {any}\n", .{ it.stream.id, b.row, b.anchor, b.pending, b.depth, b.lane.held[0..b.depth] }); + } + } + } + + /// Blocks on an empty context: the pipelines compiled, then a 1-draft and a 15-draft block timed (the second runs). + fn warm(d: *Drafter, io: std.Io) !void { + const l = try d.newLane(64); + defer d.dropLane(l); + var ms: [2]f64 = undefined; + for ([_]u32{ block_rows - 1, 1, block_rows - 1, 1, block_rows - 1 }, 0..) |depth, i| { + const t0 = std.Io.Clock.awake.now(io); + const cb = d.target.queue.commandBuffer(); + const e = cb.compute(.serial); + try d.encodeBlocks(e, &.{.{ .lane = l, .pending = 2, .anchor = 0, .depth = depth, .slot = 0, .row = 0 }}, depth + 1); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| return failed(msg); + if (i >= 3) ms[i - 3] = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds())) / 1e6; + } + d.block_ms = ms; + d.step_ms = @max(ms[1] - ms[0], 0) / @as(f64, block_rows - 2); + } +}; diff --git a/zig/src/families/gemma4/dflash_weights.zig b/zig/src/families/gemma4/dflash_weights.zig new file mode 100644 index 0000000000..e208682c5c --- /dev/null +++ b/zig/src/families/gemma4/dflash_weights.zig @@ -0,0 +1,322 @@ +//! z-lab's DFlash drafter for Gemma 4 31B: config admitted, bf16 linears quantized to the 4-bit lane layout at load. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); +const ckpt = @import("../../core/checkpoint_metal.zig"); +const shards = @import("../../core/checkpoint.zig"); +const weights = @import("weights.zig"); + +const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + +/// A packed 4-bit lane linear, or (GEMMA_DRAFT_BF16, a check of the quantizer) the bf16 rows as stored, `sb` unused. +pub const Linear = struct { w: mtl.Buffer, sb: mtl.Buffer, n: usize, k: usize, bf16: bool = false }; +pub const Ref = weights.Ref; + +/// The stock drafter's geometry (z-lab/gemma-4-31B-it-DFlash); any other config is refused with its reason. +pub const Config = struct { + taps: [8]u32 = @splat(0), + tap_count: u32, + layers: u32, + hidden: u32, + intermediate: u32, + heads: u32, + kv_heads: u32, + head_dim: u32, + block: u32, + mask: u32, + window: u32, + eps: f32, + theta: f32, + /// Each layer's kind: true for a sliding layer (causal within the block), false for a full one. + sliding: [8]bool = @splat(true), + + /// `geo`: the target's sizes, which the drafter's width and taps must fit. + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, geo: cfg.Geometry) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(bytes); + return parse(gpa, bytes, geo); + } + + pub fn parse(gpa: std.mem.Allocator, bytes: []const u8, geo: cfg.Geometry) !Config { + var parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const o = if (parsed.value == .object) parsed.value.object else return error.BadDraftConfig; + const archs = o.get("architectures") orelse return error.BadDraftConfig; + if (archs != .array or archs.array.items.len != 1 or archs.array.items[0] != .string or !std.mem.eql(u8, archs.array.items[0].string, "DFlashDraftModel")) return error.NotDFlashV1; + const d = if (o.get("dflash_config")) |v| (if (v == .object) v.object else return error.BadDraftConfig) else return error.BadDraftConfig; + var c: Config = .{ .tap_count = 0, .layers = @intCast(try integer(o, "num_hidden_layers")), .hidden = @intCast(try integer(o, "hidden_size")), .intermediate = @intCast(try integer(o, "intermediate_size")), .heads = @intCast(try integer(o, "num_attention_heads")), .kv_heads = @intCast(try integer(o, "num_key_value_heads")), .head_dim = @intCast(try integer(o, "head_dim")), .block = @intCast(if (d.get("block_size") != null) try integer(d, "block_size") else try integer(o, "block_size")), .mask = @intCast(try integer(d, "mask_token_id")), .window = @intCast(try integer(o, "sliding_window")), .eps = @floatCast(try number(o, "rms_norm_eps")), .theta = @floatCast(try number(o, "rope_theta")) }; + if (c.hidden != geo.hidden) return error.DraftForAnotherTarget; + const taps = d.get("target_layer_ids") orelse return error.BadDraftConfig; + if (taps != .array or taps.array.items.len == 0 or taps.array.items.len > 8) return error.BadDraftConfig; + for (taps.array.items, 0..) |t, i| { + if (t != .integer or t.integer < 0) return error.BadDraftConfig; + if (t.integer >= geo.layers) return error.DraftForAnotherTarget; + c.taps[i] = @intCast(t.integer); + } + c.tap_count = @intCast(taps.array.items.len); + if (c.layers == 0 or c.layers > 8) return error.BadDraftConfig; + if (o.get("layer_types")) |kinds| { + if (kinds != .array or kinds.array.items.len != c.layers) return error.BadDraftConfig; + for (kinds.array.items, 0..) |k, i| c.sliding[i] = k == .string and std.mem.eql(u8, k.string, "sliding_attention"); + } + if (o.get("is_causal") != null) return error.UnsupportedDraftMask; // the vendored default decides the block's mask + if (c.head_dim != 128 or c.heads % c.kv_heads != 0 or c.heads > 128 or c.block != 16 or c.window < 32) return error.UnsupportedDraftShape; + if (c.hidden % 64 != 0 or c.intermediate % 64 != 0 or (c.heads * c.head_dim) % 32 != 0 or (c.kv_heads * c.head_dim) % 64 != 0) return error.UnsupportedDraftShape; + if (c.mask >= cfg.vocab) return error.BadDraftConfig; + return c; + } +}; + +fn integer(o: std.json.ObjectMap, key: []const u8) !i64 { + const v = o.get(key) orelse return error.BadDraftConfig; + return if (v == .integer) v.integer else error.BadDraftConfig; +} + +fn number(o: std.json.ObjectMap, key: []const u8) !f64 { + const v = o.get(key) orelse return error.BadDraftConfig; + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadDraftConfig, + }; +} + +pub const Layer = struct { + /// q, k and v stacked (the block's rows), and k and v alone (context rows at absorb). + qkv: Linear, + kv: Linear, + o: Linear, + gate_up: Linear, + down: Linear, + input_norm: Ref, + post_attention: Ref, + q_norm: Ref, + k_norm: Ref, +}; + +pub const Weights = struct { + fc: Linear, + hidden_norm: Ref, + norm: Ref, + layers: [8]Layer, + owned: std.ArrayList(mtl.Buffer) = .empty, + + pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void { + for (w.owned.items) |b| b.deinit(); + w.owned.deinit(gpa); + } +}; + +/// MLX's affine 4-bit codes of 64 bf16 values: bias the larger edge, scale |edge| / round(|edge| / ((max - min) / 15)). +fn quantizeGroup(x: []const u16, words: []u32, scale: *u16, bias: *u16) void { + var lo: f32 = std.math.inf(f32); + var hi: f32 = -std.math.inf(f32); + for (x) |v| { + const f = bf16ToF32(v); + lo = @min(lo, f); + hi = @max(hi, f); + } + const bins: f32 = 15; + const scale0 = @max(hi - lo, 1e-7) / bins; + const from_min = @abs(lo) > @abs(hi); + const edge = if (from_min) lo else hi; + const q0 = @round(@abs(edge) / scale0); + const s_f: f32 = if (q0 == 0) scale0 else (if (from_min) @abs(edge) else -@abs(edge)) / q0; + scale.* = f32ToBf16(s_f); + bias.* = f32ToBf16(if (q0 == 0) lo else edge); + const s = bf16ToF32(scale.*); + const b = bf16ToF32(bias.*); + @memset(words, 0); + for (x, 0..) |v, i| { + const q = std.math.clamp(@round((bf16ToF32(v) - b) / s), 0, bins); + words[i / 8] |= @as(u32, @intFromFloat(q)) << @intCast((i % 8) * 4); + } +} + +fn bf16ToF32(v: u16) f32 { + return @bitCast(@as(u32, v) << 16); +} + +/// Round to nearest even, as MLX casts. +fn f32ToBf16(f: f32) u16 { + const bits: u32 = @bitCast(f); + if ((bits & 0x7f80_0000) == 0x7f80_0000) return @truncate(bits >> 16); + const round = 0x7fff + ((bits >> 16) & 1); + return @truncate((bits + round) >> 16); +} + +/// A thread's share of a tensor's rows quantized. +const Job = struct { + src: []const u16, + k: usize, + words: []u32, + scales: []u16, + biases: []u16, + fn run(j: Job, from: usize, to: usize) void { + const groups = j.k / 64; + for (from..to) |r| for (0..groups) |g| { + quantizeGroup(j.src[r * j.k + g * 64 ..][0..64], j.words[(r * groups + g) * 8 ..][0..8], &j.scales[r * groups + g], &j.biases[r * groups + g]); + }; + } +}; + +/// The members' bf16 rows stacked along N, quantized to 4-bit groups of 64 on every core, packed for the reg kernel. +fn quantizeStack(gpa: std.mem.Allocator, device: mtl.Device, c: *const ckpt.Checkpoint, owned: *std.ArrayList(mtl.Buffer), geo: cfg.Geometry, prefix: []const u8, members: []const []const u8) !Linear { + _ = try stackShape(c, prefix, members); + const bf16_env = std.c.getenv("GEMMA_DRAFT_BF16"); + if (bf16_env != null and (std.mem.eql(u8, std.mem.span(bf16_env.?), "all") or std.mem.eql(u8, std.mem.span(bf16_env.?), members[0]))) return bf16Stack(gpa, device, c, owned, prefix, members); + return quantized(gpa, device, c, owned, geo, prefix, members); +} + +/// The stack's rows and columns: bf16 matrices of one width, n % 32 and k % 64 for the lane layout. +fn stackShape(c: *const ckpt.Checkpoint, prefix: []const u8, members: []const []const u8) ![2]usize { + var name: [160]u8 = undefined; + var n: usize = 0; + var k: usize = 0; + for (members) |m| { + const t = try c.get(try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, m })); + if (t.dtype != .bf16 or t.rank != 2) return error.BadDraftWeight; + if (k != 0 and k != t.shape[1]) return error.BadDraftWeight; + n += t.shape[0]; + k = t.shape[1]; + } + if (k % 64 != 0 or n % 32 != 0) return error.BadDraftWeight; + return .{ n, k }; +} + +/// The members' bf16 rows stacked along N, always 4-bit: DFlash's by default, the MTP assistant's linears and head. +pub fn quantized(gpa: std.mem.Allocator, device: mtl.Device, c: *const ckpt.Checkpoint, owned: *std.ArrayList(mtl.Buffer), geo: cfg.Geometry, prefix: []const u8, members: []const []const u8) !Linear { + var name: [160]u8 = undefined; + const n, const k = try stackShape(c, prefix, members); + const words = try gpa.alloc(u32, n * k / 8); + defer gpa.free(words); + const scales = try gpa.alloc(u16, n * k / 64); + defer gpa.free(scales); + const biases = try gpa.alloc(u16, n * k / 64); + defer gpa.free(biases); + var row: usize = 0; + for (members) |m| { + const t = try c.get(try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, m })); + const rows = t.shape[0]; + const job = Job{ .src = t.host(u16)[0 .. rows * k], .k = k, .words = words[row * k / 8 ..][0 .. rows * k / 8], .scales = scales[row * k / 64 ..][0 .. rows * k / 64], .biases = biases[row * k / 64 ..][0 .. rows * k / 64] }; + try parallel(gpa, job, rows); + row += rows; + } + const lin = try weights.packArrays(gpa, device, owned, geo, words, scales, biases, .{ .w = undefined, .sb = undefined, .n = n, .k = k }); + return .{ .w = lin.w, .sb = lin.sb, .n = n, .k = k }; +} + +/// The members' bf16 rows stacked along N as stored (the slow per-row kernel reads them). +fn bf16Stack(gpa: std.mem.Allocator, device: mtl.Device, c: *const ckpt.Checkpoint, owned: *std.ArrayList(mtl.Buffer), prefix: []const u8, members: []const []const u8) !Linear { + var name: [160]u8 = undefined; + var n: usize = 0; + var k: usize = 0; + for (members) |m| { + const t = try c.get(try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, m })); + n += t.shape[0]; + k = t.shape[1]; + } + const buf = try device.buffer(n * k * 2, shared); + errdefer buf.deinit(); + var at: usize = 0; + for (members) |m| { + const t = try c.get(try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, m })); + @memcpy(buf.slice(u16, n * k)[at..][0 .. t.shape[0] * k], t.host(u16)[0 .. t.shape[0] * k]); + at += t.shape[0] * k; + } + try owned.append(gpa, buf); + return .{ .w = buf, .sb = buf, .n = n, .k = k, .bf16 = true }; +} + +/// `job` over its rows on every CPU. +fn parallel(gpa: std.mem.Allocator, job: Job, rows: usize) !void { + const threads = @min(std.Thread.getCpuCount() catch 1, 16); + const per = (rows + threads - 1) / threads; + const handles = try gpa.alloc(std.Thread, threads); + defer gpa.free(handles); + var spawned: usize = 0; + var at: usize = 0; + while (at < rows) : (at += per) { + handles[spawned] = try std.Thread.spawn(.{}, Job.run, .{ job, at, @min(rows, at + per) }); + spawned += 1; + } + for (handles[0..spawned]) |h| h.join(); +} + +/// A bf16 vector copied into its own buffer (the checkpoint's buffers are released after load). +pub fn vector(gpa: std.mem.Allocator, device: mtl.Device, c: *const ckpt.Checkpoint, owned: *std.ArrayList(mtl.Buffer), name: []const u8, len: usize) !Ref { + const t = try c.get(name); + if (t.dtype != .bf16 or t.count() != len) return error.BadDraftWeight; + const buf = try device.buffer(len * 2, shared); + errdefer buf.deinit(); + @memcpy(buf.slice(u16, len), t.host(u16)[0..len]); + try owned.append(gpa, buf); + return .{ .buf = buf }; +} + +pub fn load(gpa: std.mem.Allocator, device: mtl.Device, io: std.Io, dir: []const u8, c: *const Config, geo: cfg.Geometry) !Weights { + var check = ckpt.Checkpoint.init(gpa); + defer check.deinit(); + const files = try shards.shardFiles(gpa, io, dir); + defer shards.freeShardFiles(gpa, files); + for (files) |path| try check.addFile(device, path, ""); + var w: Weights = undefined; + w.owned = .empty; + errdefer w.deinit(gpa); + var p: [64]u8 = undefined; + var name: [160]u8 = undefined; + w.fc = try quantizeStack(gpa, device, &check, &w.owned, geo, "", &.{"fc"}); + if (w.fc.k != c.tap_count * c.hidden or w.fc.n != c.hidden) return error.BadDraftWeight; + w.hidden_norm = try vector(gpa, device, &check, &w.owned, "hidden_norm.weight", c.hidden); + w.norm = try vector(gpa, device, &check, &w.owned, "norm.weight", c.hidden); + for (w.layers[0..c.layers], 0..) |*l, i| { + const pre = try std.fmt.bufPrint(&p, "layers.{d}.", .{i}); + l.qkv = try quantizeStack(gpa, device, &check, &w.owned, geo, pre, &.{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj" }); + l.kv = try quantizeStack(gpa, device, &check, &w.owned, geo, pre, &.{ "self_attn.k_proj", "self_attn.v_proj" }); + l.o = try quantizeStack(gpa, device, &check, &w.owned, geo, pre, &.{"self_attn.o_proj"}); + l.gate_up = try quantizeStack(gpa, device, &check, &w.owned, geo, pre, &.{ "mlp.gate_proj", "mlp.up_proj" }); + l.down = try quantizeStack(gpa, device, &check, &w.owned, geo, pre, &.{"mlp.down_proj"}); + if (l.qkv.n != (c.heads + 2 * c.kv_heads) * c.head_dim or l.o.k != c.heads * c.head_dim or l.gate_up.n != 2 * c.intermediate or l.down.k != c.intermediate) return error.BadDraftWeight; + inline for (.{ "input_norm", "post_attention", "q_norm", "k_norm" }, .{ "input_layernorm.weight", "post_attention_layernorm.weight", "self_attn.q_norm.weight", "self_attn.k_norm.weight" }, .{ c.hidden, c.hidden, c.head_dim, c.head_dim }) |field, t, len| + @field(l, field) = try vector(gpa, device, &check, &w.owned, try std.fmt.bufPrint(&name, "{s}{s}", .{ pre, t }), len); + } + return w; +} + +test "the stock drafter config is admitted and others refused with a reason" { + const stock = + \\{"architectures":["DFlashDraftModel"],"dflash_config":{"mask_token_id":4,"target_layer_ids":[1,12,23,35,46,57]},"block_size":16, + \\"hidden_size":5376,"intermediate_size":10752,"num_hidden_layers":5,"num_attention_heads":64,"num_key_value_heads":8,"head_dim":128, + \\"sliding_window":2048,"rms_norm_eps":1e-6,"rope_theta":1000000,"layer_types":["sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention"]} + ; + const c = try Config.parse(std.testing.allocator, stock, cfg.dense31b); + try std.testing.expectEqual(@as(u32, 6), c.tap_count); + try std.testing.expectEqual(@as(u32, 57), c.taps[5]); + try std.testing.expect(c.sliding[3] and !c.sliding[4]); + const two = try std.mem.replaceOwned(u8, std.testing.allocator, stock, "DFlashDraftModel", "DFlash2DraftModel"); + defer std.testing.allocator.free(two); + try std.testing.expectError(error.NotDFlashV1, Config.parse(std.testing.allocator, two, cfg.dense31b)); + const causal = try std.mem.replaceOwned(u8, std.testing.allocator, stock, "\"hidden_size\"", "\"is_causal\":false,\"hidden_size\""); + defer std.testing.allocator.free(causal); + try std.testing.expectError(error.UnsupportedDraftMask, Config.parse(std.testing.allocator, causal, cfg.dense31b)); + try std.testing.expectError(error.DraftForAnotherTarget, Config.parse(std.testing.allocator, stock, cfg.unified12b)); +} + +test "4-bit quantization keeps the larger edge exact and zero within the scale's bf16 rounding, as MLX does" { + var x: [64]u16 = undefined; + for (&x, 0..) |*v, i| v.* = f32ToBf16(@as(f32, @floatFromInt(i)) * 0.125 - 2.0); // -2.0 .. 5.875: the max is the edge + var words: [8]u32 = undefined; + var s: u16 = undefined; + var b: u16 = undefined; + quantizeGroup(&x, &words, &s, &b); + const scale = bf16ToF32(s); + const bias = bf16ToF32(b); + try std.testing.expect(scale < 0 and bias == 5.875); // codes descend from the max + const q63: u32 = (words[7] >> 28) & 0xf; + try std.testing.expectEqual(@as(u32, 0), q63); + const q16: u32 = words[2] & 0xf; // x[16] = 0.0 lands on a code, off by the bf16 rounding of the scale times the code + try std.testing.expectApproxEqAbs(@as(f32, 0.0), scale * @as(f32, @floatFromInt(q16)) + bias, 0.05); + try std.testing.expect(@abs(scale * @as(f32, @floatFromInt(q16)) + bias) < @abs(scale) / 2); +} diff --git a/zig/src/families/gemma4/fixtures/assistant_config.json b/zig/src/families/gemma4/fixtures/assistant_config.json new file mode 100644 index 0000000000..9889070ccf --- /dev/null +++ b/zig/src/families/gemma4/fixtures/assistant_config.json @@ -0,0 +1,88 @@ +{ + "architectures": [ + "Gemma4UnifiedAssistantForCausalLM" + ], + "audio_token_id": 258881, + "backbone_hidden_size": 3840, + "boa_token_id": 256000, + "boi_token_id": 255999, + "centroid_intermediate_top_k": 32, + "dtype": "bfloat16", + "eoa_token_index": 258883, + "eoi_token_id": 258882, + "image_token_id": 258880, + "model_type": "gemma4_unified_assistant", + "num_centroids": 2048, + "text_config": { + "_name_or_path": "", + "architectures": null, + "attention_bias": false, + "attention_dropout": 0.0, + "attention_k_eq_v": true, + "bos_token_id": 2, + "chunk_size_feed_forward": 0, + "dtype": "bfloat16", + "enable_moe_block": false, + "eos_token_id": 1, + "final_logit_softcapping": null, + "global_head_dim": 512, + "head_dim": 256, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 1024, + "hidden_size_per_layer_input": 0, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "initializer_range": 0.02, + "intermediate_size": 8192, + "is_encoder_decoder": false, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "layer_types": [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention" + ], + "max_position_embeddings": 262144, + "model_type": "gemma4_unified_text", + "moe_intermediate_size": null, + "num_attention_heads": 16, + "num_experts": null, + "num_global_key_value_heads": 1, + "num_hidden_layers": 4, + "num_key_value_heads": 8, + "num_kv_shared_layers": 4, + "output_attentions": false, + "output_hidden_states": false, + "pad_token_id": 0, + "problem_type": null, + "return_dict": true, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "full_attention": { + "partial_rotary_factor": 0.25, + "rope_theta": 1000000.0, + "rope_type": "proportional" + }, + "sliding_attention": { + "rope_theta": 10000.0, + "rope_type": "default" + } + }, + "sliding_window": 1024, + "tie_word_embeddings": true, + "top_k_experts": null, + "use_bidirectional_attention": "vision", + "use_cache": true, + "use_double_wide_mlp": false, + "vocab_size": 262144, + "vocab_size_per_layer_input": 0 + }, + "tie_word_embeddings": true, + "transformers_version": "5.10.0.dev0", + "use_ordered_embeddings": false +} diff --git a/zig/src/families/gemma4/fixtures/config.json b/zig/src/families/gemma4/fixtures/config.json new file mode 100644 index 0000000000..566583018d --- /dev/null +++ b/zig/src/families/gemma4/fixtures/config.json @@ -0,0 +1,205 @@ +{ + "architectures": [ + "Gemma4ForConditionalGeneration" + ], + "audio_config": null, + "audio_token_id": 258881, + "boa_token_id": 256000, + "boi_token_id": 255999, + "do_sample": true, + "dtype": "bfloat16", + "eoa_token_id": 258883, + "eoa_token_index": 258883, + "eoi_token_id": 258882, + "eos_token_id": [ + 1, + 106, + 50 + ], + "generation_config": { + "bos_token_id": 2, + "do_sample": true, + "eos_token_id": [ + 1, + 106, + 50 + ], + "pad_token_id": 0, + "temperature": 1.0, + "top_k": 64, + "top_p": 0.95, + "transformers_version": "5.5.0.dev0" + }, + "image_token_id": 258880, + "initializer_range": 0.02, + "model_type": "gemma4", + "quantization": { + "group_size": 64, + "bits": 4, + "mode": "affine" + }, + "quantization_config": { + "group_size": 64, + "bits": 4, + "mode": "affine" + }, + "temperature": 1.0, + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "attention_k_eq_v": true, + "bos_token_id": 2, + "dtype": "bfloat16", + "enable_moe_block": false, + "eos_token_id": 1, + "expert_intermediate_size": null, + "final_logit_softcapping": 30.0, + "global_head_dim": 512, + "head_dim": 256, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 5376, + "hidden_size_per_layer_input": 0, + "initializer_range": 0.02, + "intermediate_size": 21504, + "layer_types": [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention" + ], + "max_position_embeddings": 262144, + "model_type": "gemma4_text", + "num_attention_heads": 32, + "num_experts": null, + "num_global_key_value_heads": 4, + "num_hidden_layers": 60, + "num_key_value_heads": 16, + "num_kv_shared_layers": 0, + "pad_token_id": 0, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "full_attention": { + "partial_rotary_factor": 0.25, + "rope_theta": 1000000.0, + "rope_type": "proportional" + }, + "sliding_attention": { + "rope_theta": 10000.0, + "rope_type": "default" + } + }, + "sliding_window": 1024, + "tie_word_embeddings": true, + "top_k_experts": null, + "use_bidirectional_attention": "vision", + "use_cache": true, + "use_double_wide_mlp": false, + "vocab_size": 262144, + "vocab_size_per_layer_input": 262144 + }, + "tie_word_embeddings": true, + "top_k": 64, + "top_p": 0.95, + "transformers_version": "5.5.0.dev0", + "video_token_id": 258884, + "vision_config": { + "_name_or_path": "", + "architectures": null, + "attention_bias": false, + "attention_dropout": 0.0, + "chunk_size_feed_forward": 0, + "default_output_length": 280, + "dtype": "bfloat16", + "global_head_dim": 72, + "head_dim": 72, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 1152, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "initializer_range": 0.02, + "intermediate_size": 4304, + "is_encoder_decoder": false, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "max_position_embeddings": 131072, + "model_type": "gemma4_vision", + "num_attention_heads": 16, + "num_hidden_layers": 27, + "num_key_value_heads": 16, + "output_attentions": false, + "output_hidden_states": false, + "patch_size": 16, + "pooling_kernel_size": 3, + "position_embedding_size": 10240, + "problem_type": null, + "return_dict": true, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "rope_theta": 100.0, + "rope_type": "default" + }, + "standardize": true, + "use_clipped_linears": false + }, + "vision_soft_tokens_per_image": 280 +} \ No newline at end of file diff --git a/zig/src/families/gemma4/fixtures/config_12b.json b/zig/src/families/gemma4/fixtures/config_12b.json new file mode 100644 index 0000000000..2375fb2bed --- /dev/null +++ b/zig/src/families/gemma4/fixtures/config_12b.json @@ -0,0 +1,128 @@ +{ + "architectures": [ + "Gemma4UnifiedForConditionalGeneration" + ], + "audio_token_id": 258881, + "boa_token_id": 256000, + "boi_token_id": 255999, + "dtype": "bfloat16", + "eoa_token_index": 258883, + "eoi_token_id": 258882, + "eos_token_id": [ + 1, + 106, + 50 + ], + "image_token_id": 258880, + "initializer_range": 0.02, + "model_type": "gemma4_unified", + "quantization": { + "group_size": 64, + "bits": 4, + "mode": "affine" + }, + "quantization_config": { + "group_size": 64, + "bits": 4, + "mode": "affine" + }, + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "attention_k_eq_v": true, + "bos_token_id": 2, + "enable_moe_block": false, + "eos_token_id": 1, + "final_logit_softcapping": 30.0, + "global_head_dim": 512, + "head_dim": 256, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 3840, + "hidden_size_per_layer_input": 0, + "initializer_range": 0.02, + "intermediate_size": 15360, + "layer_types": [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention" + ], + "max_position_embeddings": 262144, + "model_type": "gemma4_unified_text", + "moe_intermediate_size": null, + "num_attention_heads": 16, + "num_experts": null, + "num_global_key_value_heads": 1, + "num_hidden_layers": 48, + "num_key_value_heads": 8, + "num_kv_shared_layers": 0, + "pad_token_id": 0, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "full_attention": { + "partial_rotary_factor": 0.25, + "rope_theta": 1000000.0, + "rope_type": "proportional" + }, + "sliding_attention": { + "rope_theta": 10000.0, + "rope_type": "default" + } + }, + "sliding_window": 1024, + "tie_word_embeddings": true, + "top_k_experts": null, + "use_bidirectional_attention": "vision", + "use_cache": true, + "use_double_wide_mlp": false, + "vocab_size": 262144, + "vocab_size_per_layer_input": 262144 + }, + "tie_word_embeddings": true, + "transformers_version": "5.10.0.dev0", + "video_token_id": 258884 +} diff --git a/zig/src/families/gemma4/forward.zig b/zig/src/families/gemma4/forward.zig new file mode 100644 index 0000000000..218b6e9126 --- /dev/null +++ b/zig/src/families/gemma4/forward.zig @@ -0,0 +1,316 @@ +//! One stream's decode window through Gemma 4 dense: embed, each layer's glue and lane projections, the head. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); +const state = @import("state.zig"); +const kernels = @import("kernels.zig"); +const weights = @import("weights.zig"); +const lane = @import("../../core/lane_projection.zig"); +const Model = @import("model.zig").Model; + +const Kind = kernels.Kind; + +/// One stream's rows in a forward: its cache, and how many rows it adds at the cache's length. +pub const Segment = struct { cache: *state.Cache, rows: usize }; + +/// Which of a layer's stages a call encodes (the bench times each alone); every forward takes all three. +pub const Stages = struct { mm: bool = true, attn: bool = true, glue: bool = true }; + +/// Rows a call of the attention kernels takes: its partials' scratch is sized for this many. +const attend_rows = state.max_rows; +/// Keys a gemma_decode_attention threadgroup takes (its CD). +const decode_chunk = 256; + +/// Each sub-batch's chunk count, by its first row; the kernels' meta sits at that row too. +pub const Chunks = [state.prompt_rows][2]usize; + +/// Host-side inputs: ids, each row's position and first key, and each attention sub-batch's meta at its first row. +pub fn prepare(s: *state.Scratch, segments: []const Segment, ids: []const u32, chunks: *Chunks) void { + @memcpy(s.ids.slice(u32, ids.len), ids); + var at: usize = 0; + for (segments) |seg| { + const p0 = seg.cache.len; + for (s.pos.slice(i32, ids.len)[at..][0..seg.rows], 0..) |*p, i| p.* = @intCast(p0 + i); + for (0..2) |kind| { + const global = kind == 1; + const ck: usize = if (global) decode_chunk else 128; + const lo = s.lo[kind].slice(i32, ids.len); + for (lo[at..][0..seg.rows], 0..) |*l, i| { + const p = p0 + i; + l.* = @intCast(if (global or p + 1 <= cfg.window) 0 else p + 1 - cfg.window); + } + var sub: usize = 0; + while (sub < seg.rows) : (sub += attend_rows) { + const n = @min(attend_rows, seg.rows - sub); + const first_chunk = @as(usize, @intCast(lo[at + sub])) / ck; + chunks[at + sub][kind] = (p0 + sub + n - 1) / ck - first_chunk + 1; + const ring: i32 = if (global) 0 else state.ring; + // every key the sub-batch reads is in the cache already: no row counts as new (first new key past all) + s.meta[kind].slice(i32, state.prompt_rows * state.meta_ints)[(at + sub) * state.meta_ints ..][0..5].* = .{ @intCast(first_chunk), @intCast(chunks[at + sub][kind]), ring, @intCast(n), std.math.maxInt(i32) }; + } + } + at += seg.rows; + } +} + +/// A projection: at 4/64 on the group sums in s.sums its input's kernel wrote; other widths sum their input themselves. +fn project(m: *Model, s: *state.Scratch, e: mtl.ComputeEncoder, l: weights.Linear, x: mtl.Buffer, y: mtl.Buffer, rows: usize) !void { + const p = try m.projections.get(m.gpa, m.device, l); + const out: lane.Ref = .{ .buf = y, .off = l.col * 2 }; + if (!l.reg()) return l.encodeRows(p, e, .{ .buf = x }, out, rows); + try p.encodeSums(e, .{ .buf = x }, .{ .buf = l.w }, .{ .buf = l.sb }, .{ .ref = .{ .buf = s.sums }, .stride = @intCast(rows) }, out, @intCast(rows)); +} + +/// A stack's runs, each writing its columns of y's rows. +fn projectStack(m: *Model, s: *state.Scratch, e: mtl.ComputeEncoder, st: *const weights.Stack, x: mtl.Buffer, y: mtl.Buffer, rows: usize) !void { + for (st.list()) |l| try project(m, s, e, l, x, y, rows); +} + +/// The sums a producing kernel writes: s.sums with the window's rows as stride. +fn sumsOut(e: mtl.ComputeEncoder, s: *state.Scratch, rows: usize, index: usize) void { + e.setBuffer(s.sums, 0, index); + e.setValue(@as(i32, @intCast(rows)), index + 1); +} + +fn bind(e: mtl.ComputeEncoder, bufs: []const mtl.Buffer, first: usize) void { + for (bufs, first..) |b, i| e.setBuffer(b, 0, i); +} + +fn bindRef(e: mtl.ComputeEncoder, r: weights.Ref, i: usize) void { + e.setBuffer(r.buf, r.off, i); +} + +/// Every segment's rows in one forward: soft-capped logits in s.logits by segment; each cache grows by its rows. +pub fn run(m: *Model, s: *state.Scratch, segments: []const Segment, ids: []const u32) !void { + const geo = m.config.geo; + var rows: usize = 0; + for (segments) |seg| { + if (seg.rows == 0 or seg.cache.len + seg.rows > seg.cache.capacity) return error.ContextFull; + rows += seg.rows; + } + if (rows == 0 or rows > state.prompt_rows or rows != ids.len) return error.BadWindow; + var chunks: Chunks = undefined; + prepare(s, segments, ids, &chunks); + var at: usize = 0; + for (segments) |seg| { + var sub: usize = 0; + while (sub < seg.rows) : (sub += attend_rows) for (chunks[at + sub]) |n| if (geo.query_heads * attend_rows * n > s.slots) return error.ContextFull; + at += seg.rows; + } + const w = &m.weights; + const ks = &m.kernels.sliding; + const cb = m.queue.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(ks.embed); + bind(e, &.{s.ids}, 0); + bindRef(e, w.embed_w, 1); + bindRef(e, w.embed_s, 2); + bindRef(e, w.embed_b, 3); + bind(e, &.{s.h[0]}, 4); + e.setValue(embedScale(geo.hidden), 5); + e.dispatchThreads(.{ .width = geo.hidden / 2, .height = rows }, .{ .width = 256 }); + e.setPipeline(ks.rms); + bind(e, &.{s.h[0]}, 0); + bindRef(e, w.layers[0].input_norm, 1); + bind(e, &.{s.normed[0]}, 2); + e.setValue(cfg.eps, 3); + sumsOut(e, s, rows, 4); + e.dispatchGroups(.{ .width = rows }, .{ .width = 1024 }); + var cur: usize = 0; + const fused = rows == 1 and m.fuse; + for (0..geo.layers) |i| { + try layer(m, e, s, segments, &chunks, i, cur, rows, .{}, fused); + cur = 1 - cur; + } + if (fused) denseTail(m, e, s, geo.layers - 1, 1 - cur, rows); // the last layer's tail, no projection to carry it + try project(m, s, e, w.head, s.normed[cur], s.logits, rows); + e.setPipeline(ks.softcap); + bind(e, &.{s.logits}, 0); + e.setValue(cfg.softcap, 1); + e.dispatchThreads(.{ .width = rows * cfg.vocab }, .{ .width = 256 }); + // each row's greedy token (the softcap keeps the order): tree_round's argmax in 16-row pieces + var first: usize = 0; + while (first < rows) : (first += 16) { + const n = @min(16, rows - first); + try m.picks.argmax(e, .{ .buf = s.logits, .off = first * cfg.vocab * 2 }, .{ .buf = s.picks, .off = first * 8 }, .{ .rows = @intCast(n), .vocab = cfg.vocab, .stride = cfg.vocab }); + } + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("gemma4: command buffer failed: {s}", .{msg}); + return error.GpuFailed; + } + for (segments) |seg| seg.cache.len += seg.rows; +} + +/// One stream's window at the cache's length (`p0` must equal it). +pub fn window(m: *Model, s: *state.Scratch, c: *state.Cache, ids: []const u32, p0: usize) !void { + if (p0 != c.len) return error.BadPosition; + return run(m, s, &.{.{ .cache = c, .rows = ids.len }}, ids); +} + +/// A dense tail's tap output: layer i's rows at its tap's columns and their stride, or `spare` with stride 0. +pub fn tapRef(m: *const Model, taps: ?mtl.Buffer, spare: mtl.Buffer, i: usize) struct { ref: lane.Ref, stride: i32 } { + const t = m.tap_of[i]; + if (taps == null or t < 0) return .{ .ref = .{ .buf = spare }, .stride = 0 }; + const hidden = m.config.geo.hidden; + return .{ .ref = .{ .buf = taps.?, .off = @as(usize, @intCast(t)) * hidden * 2 }, .stride = @intCast(m.taps * hidden) }; +} + +pub fn tapOut(e: mtl.ComputeEncoder, m: *const Model, taps: ?mtl.Buffer, spare: mtl.Buffer, i: usize, index: usize) void { + const t = tapRef(m, taps, spare, i); + e.setBuffer(t.ref.buf, t.ref.off, index); + e.setValue(t.stride, index + 1); +} + +/// Layer i's dense tail on the window's rows: the next residual and normed input in s.h and s.normed[1 - cur]. +fn denseTail(m: *Model, e: mtl.ComputeEncoder, s: *state.Scratch, i: usize, cur: usize, rows: usize) void { + const l = &m.weights.layers[i]; + e.setPipeline(m.kernels.sliding.dense_tail); + bind(e, &.{ s.hn, s.y }, 0); + bindRef(e, l.post_feedforward, 2); + bindRef(e, l.scalar, 3); + bindRef(e, if (i + 1 < m.config.geo.layers) m.weights.layers[i + 1].input_norm else m.weights.final_norm, 4); + e.setValue(cfg.eps, 5); + bind(e, &.{ s.h[1 - cur], s.normed[1 - cur] }, 6); + sumsOut(e, s, rows, 8); + tapOut(e, m, s.taps, s.hn, i, 10); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); +} + +/// A fused projection on one row: its prologue's inputs and outputs (gemma4/fused.metal); eps in f.x, a tap in i.x. +fn projectFused(m: *Model, e: mtl.ComputeEncoder, l: weights.Linear, fuse: u8, x: mtl.Buffer, y: mtl.Buffer, a: [6]lane.Ref, o: [2]lane.Ref, tap: i32) !void { + const p = try m.projections.getFused(m.gpa, m.device, l, fuse); + try p.encodeFused(e, .{ .buf = x }, .{ .buf = l.w }, .{ .buf = l.sb }, .{ .buf = y }, .{ .a = a, .o = o, .f = .{ cfg.eps, 0, 0, 0 }, .i = .{ tap, 0, 0, 0 } }, 1); +} + +/// sqrt(hidden) as MLX multiplies by it: a weak-typed scalar, rounded to bf16 first. +pub fn embedScale(hidden: usize) f32 { + const x: f32 = @sqrt(@as(f32, @floatFromInt(hidden))); + const b: u32 = @bitCast(x); + return @bitCast((b + 0x7FFF + ((b >> 16) & 1)) & 0xFFFF0000); +} + +/// Layer i on a window's rows; `fused` (one row): dense tails as the next qkv's prologue, attention tails as gate/up's. +pub fn layer(m: *Model, e: mtl.ComputeEncoder, s: *state.Scratch, segments: []const Segment, chunks: *const Chunks, i: usize, cur: usize, rows: usize, st: Stages, fused: bool) !void { + const geo = m.config.geo; + const global = cfg.global(i); + const kd: *const Kind = m.kernels.kind(global); + const l = &m.weights.layers[i]; + const nk = geo.kvHeads(i); + const dh = cfg.headDim(i); + const kind: usize = @intFromBool(global); + const width = l.qkv.n; + if (st.mm and fused and i > 0) { + const pl = &m.weights.layers[i - 1]; + const tap = tapRef(m, s.taps, s.hn, i - 1); + try projectFused(m, e, l.qkv.runs[0], 1, s.y, s.qkv, .{ .{ .buf = s.hn }, .{ .buf = s.y }, pl.post_feedforward, pl.scalar, l.input_norm, l.input_norm }, .{ .{ .buf = s.h[cur] }, tap.ref }, tap.stride); + } else if (st.mm) try projectStack(m, s, e, &l.qkv, s.normed[cur], s.qkv, rows); + var at: usize = 0; + for (segments) |seg| { + const n = seg.rows; + const c = seg.cache; + if (!st.attn) { + at += n; + continue; + } + const kv_off = at * nk * dh * 2; // this stream's keys and values [nk, n, dh] in s.k, s.v + const q_off = at * geo.query_heads * dh * 2; + const prep = !(fused and !global); // a fused sliding layer's attention kernel preps its own heads + if (prep) { + e.setPipeline(kd.qkv); + e.setBuffer(s.qkv, at * width * 2, 0); + bindRef(e, l.q_norm, 1); + bindRef(e, l.k_norm, 2); + bind(e, &.{if (global) m.inv_global else m.inv_sliding}, 3); + e.setBuffer(s.pos, at * 4, 4); + e.setValue(cfg.eps, 5); + e.setBuffer(s.q, q_off, 6); + e.setBuffer(s.k, kv_off, 7); + e.setBuffer(s.v, kv_off, 8); + bind(e, &.{ c.k[i], c.v[i] }, 9); // keys and values straight into the cache at each row's slot + e.setValue([3]i32{ if (global) 0 else state.ring, @intCast(state.slotsOf(i, c.capacity)), 0 }, 11); + e.dispatchGroups(.{ .width = geo.query_heads + 2 * nk, .height = n }, .{ .width = 32 }); + } + var sub: usize = 0; + while (sub < n) : (sub += attend_rows) { + const r = at + sub; + const rn = @min(attend_rows, n - sub); + if (global) { // the tensor units, a key head's eight query heads as one tile: each key read once a token + const nch = chunks[r][kind]; // from the first key, in 256-key chunks (prepare) + e.setPipeline(kd.decode_attn); + e.setBuffer(s.q, r * geo.query_heads * dh * 2, 0); + bind(e, &.{ c.k[i], c.v[i] }, 1); + e.setValue([4]i32{ @intCast(rn), @intCast(c.len + sub), @intCast(state.slotsOf(i, c.capacity)), @intCast(nch) }, 3); + bind(e, &.{ s.pm, s.pl, s.po }, 4); + e.dispatchGroups(.{ .width = nch, .height = nk, .depth = (rn * kd.heads_per_kv + 15) / 16 }, .{ .width = 32 * dh / 128 }); + e.setPipeline(kd.merge); + bind(e, &.{ s.pm, s.pl, s.po }, 0); + e.setBuffer(s.meta[kind], r * state.meta_ints * 4, 3); + e.setBuffer(s.attn, r * geo.query_heads * dh * 2, 4); + e.setBuffer(s.sums, r * 4, 5); // this sub-batch's rows of the o projection's sums + e.setValue(@as(i32, @intCast(rows)), 6); + e.dispatchGroups(.{ .width = 1, .height = geo.query_heads, .depth = rn }, .{ .width = 32 }); + continue; + } + const direct = !prep and chunks[r][kind] == 1; // one chunk: the kernel writes the merged output itself + e.setPipeline(kd.partial); + e.setBuffer(s.q, r * geo.query_heads * dh * 2, 0); + bind(e, &.{ c.k[i], c.v[i], c.k[i], c.v[i] }, 1); // the new rows were written first: every key is a cached one + e.setBuffer(s.pos, r * 4, 5); + e.setBuffer(s.lo[kind], r * 4, 6); + e.setBuffer(s.meta[kind], r * state.meta_ints * 4, 7); + e.setValue(@as(i32, @intCast(state.slotsOf(i, c.capacity))), 8); + bind(e, &.{ s.pm, s.pl, s.po }, 9); + e.setBuffer(s.qkv, r * width * 2, 12); + bindRef(e, l.q_norm, 13); + bindRef(e, l.k_norm, 14); + bind(e, &.{m.inv_sliding}, 15); + e.setValue(cfg.eps, 16); + e.setValue([2]i32{ @intFromBool(!prep), @intFromBool(direct) }, 17); + e.setBuffer(s.attn, r * geo.query_heads * dh * 2, 18); + e.setBuffer(s.sums, r * 4, 19); + e.setValue(@as(i32, @intCast(rows)), 20); + bind(e, &.{ c.k[i], c.v[i] }, 21); + e.dispatchGroups(.{ .width = chunks[r][kind], .height = nk, .depth = rn }, .{ .width = 32 * kd.heads_per_kv * kd.split }); + if (direct) continue; + e.setPipeline(kd.merge); + bind(e, &.{ s.pm, s.pl, s.po }, 0); + e.setBuffer(s.meta[kind], r * state.meta_ints * 4, 3); + e.setBuffer(s.attn, r * geo.query_heads * dh * 2, 4); + e.setBuffer(s.sums, r * 4, 5); + e.setValue(@as(i32, @intCast(rows)), 6); + e.dispatchGroups(.{ .width = 1, .height = geo.query_heads, .depth = rn }, .{ .width = 32 }); + } + at += n; + } + if (st.mm) try project(m, s, e, l.o, s.attn, s.o, rows); + if (!st.glue) { + if (st.mm) try projectStack(m, s, e, &l.gate_up, s.n_mlp, s.gate_up, rows); + if (st.mm) try project(m, s, e, l.down, s.act, s.y, rows); + return; + } + if (fused and st.mm) { + try projectFused(m, e, l.gate_up.runs[0], 2, s.o, s.act, .{ .{ .buf = s.h[cur] }, .{ .buf = s.o }, l.post_attention, l.pre_feedforward, l.pre_feedforward, l.pre_feedforward }, .{ .{ .buf = s.hn }, .{ .buf = s.hn } }, 0); + try projectFused(m, e, l.down, 3, s.act, s.y, .{ .{ .buf = s.act }, .{ .buf = s.act }, .{ .buf = s.act }, .{ .buf = s.act }, .{ .buf = s.act }, .{ .buf = s.act } }, .{ .{ .buf = s.y }, .{ .buf = s.y } }, 0); + return; + } + e.setPipeline(kd.attn_tail); + bind(e, &.{ s.h[cur], s.o }, 0); + bindRef(e, l.post_attention, 2); + bindRef(e, l.pre_feedforward, 3); + e.setValue(cfg.eps, 4); + bind(e, &.{ s.hn, s.n_mlp }, 5); + sumsOut(e, s, rows, 7); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + if (st.mm) try projectStack(m, s, e, &l.gate_up, s.n_mlp, s.gate_up, rows); + e.setPipeline(kd.geglu); + bind(e, &.{ s.gate_up, s.act }, 0); + e.setValue(@as(u32, @intCast(geo.intermediate)), 2); + sumsOut(e, s, rows, 3); + e.setValue(@as(i32, @intFromBool(l.gate_up.interleaved)), 5); // one run: its rows interleaved in 16-row blocks + e.dispatchGroups(.{ .width = geo.intermediate / 64, .height = rows }, .{ .width = 64 }); + if (st.mm) try project(m, s, e, l.down, s.act, s.y, rows); + denseTail(m, e, s, i, cur, rows); +} diff --git a/zig/src/families/gemma4/gemma4.zig b/zig/src/families/gemma4/gemma4.zig new file mode 100644 index 0000000000..2714bf4771 --- /dev/null +++ b/zig/src/families/gemma4/gemma4.zig @@ -0,0 +1,27 @@ +//! Gemma 4 dense on Metal: config admission and the decode glue around the core lane projections. +pub const config = @import("config.zig"); +pub const kernels = @import("kernels.zig"); +pub const weights = @import("weights.zig"); +pub const state = @import("state.zig"); +pub const forward = @import("forward.zig"); +pub const prompt = @import("prompt.zig"); +pub const Model = @import("model.zig").Model; +pub const backend = @import("backend.zig"); +pub const dflash = @import("dflash.zig"); +pub const dflash_weights = @import("dflash_weights.zig"); +pub const assistant = @import("assistant.zig"); +pub const assistant_weights = @import("assistant_weights.zig"); +pub const image = @import("image.zig"); +pub const vision = @import("vision.zig"); +pub const vision_unified = @import("vision_unified.zig"); +pub const snapshot = @import("snapshot.zig"); + +test { + _ = assistant; + _ = assistant_weights; + _ = config; + _ = dflash_weights; + _ = image; + _ = snapshot; + _ = vision; +} diff --git a/zig/src/families/gemma4/image.zig b/zig/src/families/gemma4/image.zig new file mode 100644 index 0000000000..7a7482ac38 --- /dev/null +++ b/zig/src/families/gemma4/image.zig @@ -0,0 +1,110 @@ +//! An image's bytes to Gemma 4's tower inputs, as transformers' Gemma4ImageProcessorPil makes them (the 12B: merged). +const std = @import("std"); +const picture = @import("../../core/picture/picture.zig"); + +/// The shared decoder's refusals, and sizes or budgets the processor cannot take. +pub const Error = picture.Error || error{ ImageTooSmall, ImageTooLarge, TooFewTokens }; +pub const decoderProblem = picture.decoderProblem; + +/// The processor's constants: patch side, the pooling kernel, and the soft-token budgets it accepts. +pub const patch = 16; +pub const pool = 3; +pub const patch_width = 3 * patch * patch; // a patch's values: row, column, channel (input_proj's weight order) +const side_mult = pool * patch; +pub const budgets = [_]u32{ 70, 140, 280, 560, 1120 }; +/// The 31B's processor_config.json max_soft_tokens: the most an image takes. +pub const max_budget = 280; +pub const Ids = struct { boi: u32 = 255999, image: u32 = 258880, eoi: u32 = 258882 }; + +/// The 12B's patches: Gemma4UnifiedImageProcessor's patches_merge of each 3 x 3 group of 16-px patches, a 48-px patch. +pub const merged_width = pool * pool * patch_width; + +/// The tower's layout: 16-px patches in grid order (the 31B), or 48-px ones, each a 3 x 3 group's (the 12B). +pub const Layout = enum { patches, merged }; + +/// An image as tower inputs (engine_api.PreparedImage's fields): patches in [0, 1]; gh and gw count 16-px patches. +pub const Prepared = struct { pixels: []f32, gh: u32, gw: u32, tokens: u32, hash: u64 }; + +/// get_aspect_ratio_preserving_size: the largest sides, multiples of 48, whose patches fit `max_patches`. +pub fn targetSize(height: u32, width: u32, max_patches: u32) Error![2]u32 { + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + const factor = @sqrt(@as(f64, @floatFromInt(max_patches * patch * patch)) / (h * w)); + var th: u32 = @intFromFloat(@floor(factor * h / side_mult) * side_mult); + var tw: u32 = @intFromFloat(@floor(factor * w / side_mult) * side_mult); + if (th == 0 and tw == 0) return error.ImageTooSmall; + const longest = (max_patches / (pool * pool)) * side_mult; + if (th == 0) { + th = side_mult; + tw = @min(@as(u32, @intFromFloat(@floor(w / h))) * side_mult, longest); + } else if (tw == 0) { + tw = side_mult; + th = @min(@as(u32, @intFromFloat(@floor(h / w))) * side_mult, longest); + } + if (@as(u64, th) * tw > @as(u64, max_patches) * patch * patch) return error.ImageTooLarge; + return .{ th, tw }; +} + +/// The largest budget the processor takes within `max_tokens` (the server's share an image), up to the checkpoint's. +pub fn budget(max_tokens: u32) Error!u32 { + var best: ?u32 = null; + for (budgets) |b| if (b <= @min(max_tokens, max_budget)) { + best = b; + }; + return best orelse error.TooFewTokens; +} + +/// Bytes to patches (core's picture decoder, Pillow's bicubic resize, [0, 1]); the hash covers the bytes and the grid. +pub fn prepare(io: std.Io, gpa: std.mem.Allocator, bytes: []const u8, max_tokens: u32, layout: Layout) Error!Prepared { + const max_patches = try budget(max_tokens) * pool * pool; + var src = try picture.decode(io, gpa, bytes); + defer src.deinit(gpa); + const size = try targetSize(@intCast(src.height), @intCast(src.width), max_patches); + const resized = try picture.resample.resize(gpa, .pillow, .of(src.pixels, src.width, src.height), size[1], size[0]); + defer gpa.free(resized); + const gh = size[0] / patch; + const gw = size[1] / patch; + const grid = [2]u32{ gh, gw }; + const hash = picture.identity.withGrid(picture.identity.ofBytes(if (layout == .patches) 0x6e34 else 0x6e35, bytes), grid); + const side: usize = if (layout == .patches) patch else side_mult; + const rgb: picture.Rgb = .{ .width = size[1], .height = size[0], .pixels = resized }; + return .{ .pixels = try patches(gpa, rgb, side), .gh = gh, .gw = gw, .tokens = gh * gw / (pool * pool), .hash = hash }; +} + +/// convert_image_to_patches after np_rescale at `side` px (48: patches_merge's order): pixels in row order, RGB. +fn patches(gpa: std.mem.Allocator, src: picture.Rgb, side: usize) Error![]f32 { + const gh = src.height / side; + const gw = src.width / side; + const width = 3 * side * side; + const out = try gpa.alloc(f32, gh * gw * width); + for (0..gh) |py| for (0..gw) |px| for (0..side) |y| for (0..side) |x| for (0..3) |c| { + const v = src.pixels[((py * side + y) * src.width + px * side + x) * 3 + c]; + out[(py * gw + px) * width + (y * side + x) * 3 + c] = @floatCast(@as(f64, @floatFromInt(v)) * (1.0 / 255.0)); + }; + return out; +} + +test "target sizes as the processor sizes the fixtures" { + const cases = [_][4]u32{ // height, width, then the processor's patch grid (rows, columns) at 280 tokens + .{ 389, 517, 42, 57 }, + .{ 800, 1200, 39, 60 }, + .{ 1400, 180, 138, 18 }, + .{ 15, 20, 42, 57 }, + .{ 672, 672, 48, 48 }, + }; + for (cases) |c| { + const s = try targetSize(c[0], c[1], 280 * 9); + try std.testing.expectEqual(c[2], s[0] / patch); + try std.testing.expectEqual(c[3], s[1] / patch); + } + // one side under 48 once scaled: that side is 48 and the other its whole ratio, capped + try std.testing.expectEqual([2]u32{ 48, 48 * 280 }, try targetSize(10, 30000, 280 * 9)); + try std.testing.expectEqual([2]u32{ 48 * 280, 48 }, try targetSize(60000, 2, 280 * 9)); +} + +test "budgets within the server's share" { + try std.testing.expectEqual(@as(u32, 280), try budget(4096)); + try std.testing.expectEqual(@as(u32, 140), try budget(256)); + try std.testing.expectEqual(@as(u32, 70), try budget(70)); + try std.testing.expectError(error.TooFewTokens, budget(69)); +} diff --git a/zig/src/families/gemma4/kernels.zig b/zig/src/families/gemma4/kernels.zig new file mode 100644 index 0000000000..f3cf5edd54 --- /dev/null +++ b/zig/src/families/gemma4/kernels.zig @@ -0,0 +1,128 @@ +//! Gemma 4 dense pipelines: the glue compiled once for sliding layers and once for global ones, MLX's compile options. +const std = @import("std"); +const mtl = @import("metal"); +const glue = @import("kernel_sources").gemma4_glue; +const prompt_attention = @import("kernel_sources").gemma4_prompt_attention; +const frags = @import("../../core/frags.zig"); +const cfg = @import("config.zig"); + +/// One layer kind's pipelines; the kind-free ones (embed, norms, tails, geglu, softcap) are the same in both. +pub const Kind = struct { + embed: mtl.Pipeline, + rms: mtl.Pipeline, + qkv: mtl.Pipeline, + partial: mtl.Pipeline, + merge: mtl.Pipeline, + attn_tail: mtl.Pipeline, + dense_tail: mtl.Pipeline, + geglu: mtl.Pipeline, + softcap: mtl.Pipeline, + freqs: mtl.Pipeline, + kv_write: mtl.Pipeline, + draft_store: mtl.Pipeline, + draft_add: mtl.Pipeline, + draft_swiglu: mtl.Pipeline, + /// gemma_assist_input: the MTP assistant's rows, the pending token's embedding beside a hidden row. + assist_input: mtl.Pipeline, + /// gemma_prompt_attention: a prompt chunk's rows over the cache, on the tensor units. + prompt_attn: mtl.Pipeline, + /// gemma_decode_attention: decode rows' partials on the tensor units, a key head's query heads as one tile. + decode_attn: mtl.Pipeline, + /// gilby's attend shape by head dim: keys a chunk, simdgroups a query head, keys scored before an update. + chunk: usize, + split: usize, + block: usize, + heads_per_kv: usize, + + /// `q_only`: gemma_qkv_prep reads q alone (the MTP assistant's layers); `embed_bits`: the embedding table's width. + pub fn init(a: std.mem.Allocator, device: mtl.Device, geo: cfg.Geometry, global: bool, q_only: bool, embed_bits: u8) !Kind { + const dh: usize = if (global) cfg.global_dim else cfg.sliding_dim; + const nk: usize = if (global) geo.global_kv_heads else geo.sliding_kv_heads; + const shape: [3]usize = if (global) .{ 64, 1, 4 } else .{ 128, 4, 4 }; + const width = geo.query_heads * dh + (if (q_only) 0 else nk * dh * @as(usize, if (global) 1 else 2)); + const text = try std.fmt.allocPrint(a, + \\#define D {d} + \\#define T 256 + \\#define DH {d} + \\#define NQ {d} + \\#define NK {d} + \\#define VK {d} + \\#define W {d} + \\#define HK {d} + \\#define G {d} + \\#define CK {d} + \\#define S {d} + \\#define BLK {d} + \\#define SCALE 1.0f + \\#define GELU_C0 0.7978845608028654f + \\#define GELU_C1 0.044715f + \\#define EB {d} + \\{s} + , .{ geo.hidden, dh, geo.query_heads, nk, @intFromBool(global), width, nk, geo.query_heads / nk, shape[0], shape[1], shape[2], embed_bits, glue }); + defer a.free(text); + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var k: Kind = undefined; + inline for (.{ "embed", "rms", "qkv", "partial", "merge", "attn_tail", "dense_tail", "geglu", "softcap", "freqs", "kv_write", "draft_store", "draft_add", "draft_swiglu", "assist_input" }, .{ "gemma_embed", "gemma_rms_looped", "gemma_qkv_prep", "gemma_attention_partial", "gemma_attention_merge", "gemma_attn_tail", "gemma_dense_tail", "gemma_geglu", "gemma_softcap", "gemma_rope_freqs", "gemma_kv_write", "gemma_draft_store", "gemma_draft_add", "gemma_draft_swiglu", "gemma_assist_input" }) |field, name| + @field(k, field) = try mtl.Pipeline.init(device, lib, name, false); + { + const body = try frags.source(device, a, prompt_attention); + defer a.free(body); + const src = try std.fmt.allocPrint(a, "#define DH {d}\n#define NQ {d}\n#define G {d}\n#define SL {d}\n#define WINDOW {d}\n#define RING {d}\n{s}", .{ dh, geo.query_heads, geo.query_heads / nk, dh / 128, @as(usize, if (global) 0 else cfg.window), @as(usize, if (global) 0 else @import("state.zig").ring), body }); + defer a.free(src); + const plib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + defer plib.deinit(); + k.prompt_attn = try mtl.Pipeline.init(device, plib, "gemma_prompt_attention", false); + k.decode_attn = try mtl.Pipeline.init(device, plib, "gemma_decode_attention", false); + } + k.chunk = shape[0]; + k.split = shape[1]; + k.block = shape[2]; + k.heads_per_kv = geo.query_heads / nk; + return k; + } + + pub fn deinit(k: Kind) void { + inline for (.{ "embed", "rms", "qkv", "partial", "merge", "attn_tail", "dense_tail", "geglu", "softcap", "freqs", "kv_write", "draft_store", "draft_add", "draft_swiglu", "assist_input", "prompt_attn", "decode_attn" }) |field| @field(k, field).deinit(); + } +}; + +pub const Kernels = struct { + sliding: Kind, + global: Kind, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, geo: cfg.Geometry, embed_bits: u8) !Kernels { + const s = try Kind.init(a, device, geo, false, false, embed_bits); + errdefer s.deinit(); + return .{ .sliding = s, .global = try Kind.init(a, device, geo, true, false, embed_bits) }; + } + + pub fn kind(k: *const Kernels, global: bool) *const Kind { + return if (global) &k.global else &k.sliding; + } + + pub fn deinit(k: Kernels) void { + k.sliding.deinit(); + k.global.deinit(); + } +}; + +/// RoPE inverse frequencies as gilby takes them from mlx_lm: sliding ones in f64; global ones from MLX's GPU pow. +pub fn inverseFrequencies(k: *const Kernels, queue: mtl.Queue, device: mtl.Device, global: bool, out: []f32) !void { + if (!global) return cfg.inverseFrequencies(false, out); + const b = try device.buffer(cfg.global_rotated * 4, mtl.ResourceOptions.shared); + defer b.deinit(); + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + e.setPipeline(k.global.freqs); + e.setBuffer(b, 0, 0); + e.setValue(@as(f32, @floatCast(cfg.global_theta)), 1); + e.setValue(@as(u32, cfg.global_dim), 2); + e.dispatchThreads(.{ .width = cfg.global_rotated }, .{ .width = cfg.global_rotated }); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.GpuFailed; + @memset(out, 0); + for (b.slice(f32, cfg.global_rotated), out[0..cfg.global_rotated]) |f, *o| o.* = @floatCast(1.0 / @as(f64, f)); +} diff --git a/zig/src/families/gemma4/model.zig b/zig/src/families/gemma4/model.zig new file mode 100644 index 0000000000..d54aad4ff1 --- /dev/null +++ b/zig/src/families/gemma4/model.zig @@ -0,0 +1,181 @@ +//! A native Gemma 4 dense model (31B or 12B): admitted config, packed weights, the glue pipelines, lane projections. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); +const ckpt = @import("../../core/checkpoint_metal.zig"); +const shards = @import("../../core/checkpoint.zig"); +const weights = @import("weights.zig"); +const kernels = @import("kernels.zig"); +const affine_mm = @import("../../core/affine_mm.zig"); +const round = @import("tree_round_gpu"); +const vision = @import("vision.zig"); +const vision_unified = @import("vision_unified.zig"); +const image = @import("image.zig"); +const lane = @import("../../core/lane_projection.zig"); + +/// The image input: the 31B's tower (27 encoder layers, pooling) or the 12B's embedder (48-px patches, no layers). +pub const Tower = union(enum) { + encoder: *vision.Vision, + embedder: *vision_unified.Embedder, + + /// An image's patches (image.prepare's, in `layout`'s order) to `out` [gh * gw / 9, the model's hidden]. + pub fn encode(t: Tower, e: mtl.ComputeEncoder, pixels: []const f32, gh: u32, gw: u32, out: lane.Ref) !void { + switch (t) { + inline else => |v| try v.encode(e, pixels, gh, gw, out, null), + } + } + + pub fn inject(t: Tower, e: mtl.ComputeEncoder, src: lane.Ref, dst: lane.Ref, n: u32) void { + switch (t) { + inline else => |v| v.inject(e, src, dst, n), + } + } + + /// The patches' order the processor writes for this input. + pub fn layout(t: Tower) image.Layout { + return if (t == .encoder) .patches else .merged; + } + + /// The working buffers a residency set holds, appended to `list`. + pub fn appendBuffers(t: Tower, gpa: std.mem.Allocator, list: *std.ArrayList(mtl.Buffer)) !void { + switch (t) { + .encoder => |v| try list.appendSlice(gpa, &.{ v.scratch, v.patches, v.cos, v.sin }), + .embedder => |v| try list.appendSlice(gpa, &.{ v.scratch, v.patches }), + } + } + + pub fn deinit(t: Tower) void { + switch (t) { + inline else => |v| v.deinit(), + } + } +}; + +pub const Model = struct { + gpa: std.mem.Allocator, + device: mtl.Device, + queue: mtl.Queue, + config: cfg.Config, + checkpoint: ckpt.Checkpoint, + weights: weights.Weights, + kernels: kernels.Kernels, + projections: weights.Projections, + /// core/affine_mm's pipelines for prompt chunks, one set per width the checkpoint's linears have (by bit_widths). + mm: [cfg.widths.bit_widths.len]?affine_mm.Pipes = @splat(null), + /// RoPE inverse frequencies, float32 [dim / 2], for sliding and for global layers. + inv_sliding: mtl.Buffer, + inv_global: mtl.Buffer, + /// The core's bf16 argmax: a window's greedy tokens on the GPU (the lowest index among ties, as mx.argmax). + picks: round.Ops, + /// A drafter's tapped layers: each layer's index among the taps (-1: untapped), and how many there are. + tap_of: [cfg.max_layers]i8 = @splat(-1), + taps: u32 = 0, + /// One-row windows take the fused pipelines (glue as prologues) when every linear is 4/64; the bench may not. + fuse: bool = true, + /// The image input, once loadVision read it (serve --vision). + vision: ?Tower = null, + + pub fn load(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !*Model { + const m = try gpa.create(Model); + errdefer gpa.destroy(m); + m.* = undefined; + m.gpa = gpa; + m.tap_of = @splat(-1); + m.taps = 0; + m.vision = null; + m.mm = @splat(null); + m.config = try cfg.Config.read(gpa, io, dir); + m.fuse = m.config.widths.allFast(m.config.geo.layers); + m.projections = .{ .geo = m.config.geo }; + m.device = try mtl.Device.init(); + errdefer m.device.deinit(); + if (!m.device.tensorUnits()) return error.GemmaNeedsTensorUnits; + m.queue = try m.device.queue(); + errdefer m.queue.deinit(); + m.checkpoint = ckpt.Checkpoint.init(gpa); + errdefer m.checkpoint.deinit(); + const files = try shards.shardFiles(gpa, io, dir); + defer shards.freeShardFiles(gpa, files); + for (files) |path| try m.checkpoint.addFileSelected(m.device, path, "", "language_model."); + m.weights = try weights.load(gpa, m.device, &m.checkpoint, m.config.geo, &m.config.widths); + errdefer m.weights.deinit(gpa); + m.kernels = try kernels.Kernels.init(gpa, m.device, m.config.geo, m.weights.embed_bits); + errdefer m.kernels.deinit(); + m.picks = try round.Ops.init(m.device); + errdefer m.picks.deinit(); + errdefer for (m.mm) |o| if (o) |ps| for (ps) |p| p.deinit(); + for (m.config.widths.layers[0..m.config.geo.layers]) |l| for (l) |w| try m.pipesFor(w.bits); + try m.pipesFor(m.config.widths.vision.bits); + const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + m.inv_sliding = try m.device.buffer(cfg.sliding_dim / 2 * 4, opts); + m.inv_global = try m.device.buffer(cfg.global_dim / 2 * 4, opts); + try kernels.inverseFrequencies(&m.kernels, m.queue, m.device, false, m.inv_sliding.slice(f32, cfg.sliding_dim / 2)); + try kernels.inverseFrequencies(&m.kernels, m.queue, m.device, true, m.inv_global.slice(f32, cfg.global_dim / 2)); + return m; + } + + /// affine_mm's pipelines at `bits` (groups of 64), compiled the first time a width needs them. + fn pipesFor(m: *Model, bits: u8) !void { + const i = std.mem.indexOfScalar(u8, &cfg.widths.bit_widths, bits).?; + if (m.mm[i] != null) return; + const text = try affine_mm.source(m.device, m.gpa, .{ .bits = bits, .group = cfg.group }); + defer m.gpa.free(text); + const lib = try mtl.Library.fromSource(m.device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var ps: affine_mm.Pipes = undefined; + var made: usize = 0; + errdefer for (ps[0..made]) |p| p.deinit(); + for (&ps, affine_mm.names) |*p, name| { + p.* = try mtl.Pipeline.init(m.device, lib, name, false); + made += 1; + } + m.mm[i] = ps; + } + + /// affine_mm's pipelines for a linear of `bits`. + pub fn pipes(m: *const Model, bits: u8) affine_mm.Pipes { + return m.mm[std.mem.indexOfScalar(u8, &cfg.widths.bit_widths, bits).?].?; + } + + /// The checkpoint's image input (the 12B's embedder, the 31B's tower); error.NoGemmaVision when it has none. + pub fn loadVision(m: *Model, io: std.Io, dir: []const u8) !void { + if (m.vision != null) return; + const files = try shards.shardFiles(m.gpa, io, dir); + defer shards.freeShardFiles(m.gpa, files); + if (m.config.geo.unified()) { + try vision_unified.addTensors(&m.checkpoint, m.device, files); + if (!vision_unified.present(&m.checkpoint)) return error.NoGemmaVision; + m.vision = .{ .embedder = try vision_unified.Embedder.load(m.gpa, m.device, &m.checkpoint, m.config.geo.hidden, m.pipes(m.config.widths.vision.bits), m.config.widths.vision) }; + return; + } + try vision.addTensors(&m.checkpoint, m.device, files); + if (!vision.present(&m.checkpoint)) return error.NoGemmaVision; + m.vision = .{ .encoder = try vision.Vision.load(m.gpa, m.device, &m.checkpoint, m.pipes(m.config.widths.vision.bits), m.config.widths.vision, m.config.geo.hidden) }; + } + + /// The drafter's taps, in its feature order: each a layer whose output (layer scalar applied) it reads. + pub fn setTaps(m: *Model, layers: []const u32) !void { + if (layers.len == 0 or layers.len > 8) return error.BadGemmaTaps; + m.tap_of = @splat(-1); + for (layers, 0..) |l, t| { + if (l >= m.config.geo.layers or m.tap_of[l] >= 0) return error.BadGemmaTaps; + m.tap_of[l] = @intCast(t); + } + m.taps = @intCast(layers.len); + } + + pub fn deinit(m: *Model) void { + if (m.vision) |v| v.deinit(); + m.inv_sliding.deinit(); + m.inv_global.deinit(); + m.projections.deinit(m.gpa); + m.picks.deinit(); + for (m.mm) |o| if (o) |ps| for (ps) |p| p.deinit(); + m.kernels.deinit(); + m.weights.deinit(m.gpa); + m.checkpoint.deinit(); + m.queue.deinit(); + m.device.deinit(); + m.gpa.destroy(m); + } +}; diff --git a/zig/src/families/gemma4/prompt.zig b/zig/src/families/gemma4/prompt.zig new file mode 100644 index 0000000000..c5c7b864bc --- /dev/null +++ b/zig/src/families/gemma4/prompt.zig @@ -0,0 +1,237 @@ +//! A prompt chunk of up to 512 rows of one stream: core/affine_mm projections, the decode glue, the last row's logits. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); +const state = @import("state.zig"); +const kernels = @import("kernels.zig"); +const affine_mm = @import("../../core/affine_mm.zig"); +const lane = @import("../../core/lane_projection.zig"); +const Model = @import("model.zig").Model; +const fwd = @import("forward.zig"); + +const Ref = lane.Ref; +const R = state.chunk_rows; +const private = mtl.ResourceOptions.private | mtl.ResourceOptions.untracked; +const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + +const Q = @import("weights.zig").Mlx; + +pub const Scratch = struct { + ids: mtl.Buffer, + pos: mtl.Buffer, + /// Each row's last key on sliding layers: its position, or its image's last row (bidirectional vision blocks). + hi: mtl.Buffer, + h: [2]mtl.Buffer, + normed: [2]mtl.Buffer, + qkv: mtl.Buffer, + q: mtl.Buffer, + k: mtl.Buffer, + v: mtl.Buffer, + attn: mtl.Buffer, + o: mtl.Buffer, + hn: mtl.Buffer, + n_mlp: mtl.Buffer, + gate_up: mtl.Buffer, + act: mtl.Buffer, + y: mtl.Buffer, + logits: mtl.Buffer, + /// The drafter's tapped rows [rows, taps x hidden] bf16, once a drafter asked for them. + taps: ?mtl.Buffer = null, + + pub fn init(device: mtl.Device, geo: cfg.Geometry, capacity: usize) !Scratch { + const d = geo.hidden; + const qd = geo.queryWidth(); + const kv = geo.kvWidth(); + _ = capacity; + var s: Scratch = undefined; + s.taps = null; + var made: usize = 0; + errdefer s.release(made); + inline for (.{ .{ "ids", R * 4, true }, .{ "pos", R * 4, true }, .{ "hi", R * 4, true }, .{ "qkv", R * (qd + 2 * kv) * 2, false }, .{ "q", R * qd * 2, false }, .{ "k", R * kv * 2, false }, .{ "v", R * kv * 2, false }, .{ "attn", R * qd * 2, false }, .{ "o", R * d * 2, false }, .{ "hn", R * d * 2, false }, .{ "n_mlp", R * d * 2, false }, .{ "gate_up", R * 2 * geo.intermediate * 2, false }, .{ "act", R * geo.intermediate * 2, false }, .{ "y", R * d * 2, false }, .{ "logits", cfg.vocab * 2, true } }) |f| { + @field(s, f[0]) = try device.buffer(f[1], if (f[2]) shared else private); + made += 1; + } + for ([_][]mtl.Buffer{ &s.h, &s.normed }) |pair| for (pair) |*b| { + b.* = try device.buffer(R * d * 2, private); + made += 1; + }; + return s; + } + + pub fn deinit(s: *Scratch) void { + s.release(std.math.maxInt(usize)); + if (s.taps) |b| b.deinit(); + } + + /// Releases the first `made` buffers in init's order, so a failed init frees what it made. + fn release(s: *Scratch, made: usize) void { + var i: usize = 0; + inline for (.{ "ids", "pos", "hi", "qkv", "q", "k", "v", "attn", "o", "hn", "n_mlp", "gate_up", "act", "y", "logits" }) |f| { + if (i < made) @field(s, f).deinit(); + i += 1; + } + for (s.h ++ s.normed) |b| { + if (i < made) b.deinit(); + i += 1; + } + } +}; + +/// y[:, col..col + q.n] (row pitch `pitch`) = x q^T. +fn mm(m: *Model, e: mtl.ComputeEncoder, q: Q, x: mtl.Buffer, y: mtl.Buffer, col: usize, pitch: usize, rows: u32) void { + affine_mm.denseBatch(e, m.pipes(q.bits), Ref{ .buf = x }, q, Ref{ .buf = y, .off = col * 2 }, rows, 1, .{ .x_row = @intCast(q.k), .y_row = @intCast(pitch) }); +} + +/// The glue's optional group sums off (the lane kernels' operand; affine_mm dequantizes its weights instead). +fn noSums(e: mtl.ComputeEncoder, s: *Scratch, index: usize) void { + e.setBuffer(s.pos, 0, index); + e.setValue(@as(i32, 0), index + 1); +} + +fn bind(e: mtl.ComputeEncoder, bufs: []const mtl.Buffer, first: usize) void { + for (bufs, first..) |b, i| e.setBuffer(b, 0, i); +} + +/// An encoded image's rows in the prompt: positions [at, at + n), its rows from `rows` (hidden bf16 each). +pub const Span = struct { at: usize, n: u32, rows: Ref }; + +/// `ids` (at most 512) at the cache's length, `images`' rows over theirs (a span wholly in or out); s.logits: the last. +pub fn run(m: *Model, s: *Scratch, c: *state.Cache, ids: []const u32, images: []const Span) !void { + const rows = ids.len; + if (rows == 0 or rows > R) return error.BadChunk; + if (c.len + rows > c.capacity) return error.ContextFull; + const p0 = c.len; + @memcpy(s.ids.slice(u32, rows), ids); + for (s.pos.slice(i32, rows), s.hi.slice(i32, rows), 0..) |*p, *h, i| { + p.* = @intCast(p0 + i); + h.* = p.*; + } + for (images) |span| { + if (span.at + span.n <= p0 or span.at >= p0 + rows) continue; + if (span.at < p0 or span.at + span.n > p0 + rows) return error.ImageAcrossChunks; + @memset(s.hi.slice(i32, rows)[span.at - p0 ..][0..span.n], @intCast(span.at + span.n - 1)); + } + const geo = m.config.geo; + const ks = &m.kernels.sliding; + const w = &m.weights; + const cb = m.queue.commandBuffer(); + const e = cb.compute(.serial); + const r32: u32 = @intCast(rows); + e.setPipeline(ks.embed); + bind(e, &.{s.ids}, 0); + e.setBuffer(w.embed_w.buf, w.embed_w.off, 1); + e.setBuffer(w.embed_s.buf, w.embed_s.off, 2); + e.setBuffer(w.embed_b.buf, w.embed_b.off, 3); + bind(e, &.{s.h[0]}, 4); + e.setValue(fwd.embedScale(geo.hidden), 5); + e.dispatchThreads(.{ .width = geo.hidden / 2, .height = rows }, .{ .width = 256 }); + for (images) |span| if (span.at >= p0 and span.at < p0 + rows) { + const v = m.vision orelse return error.NoGemmaVision; + v.inject(e, span.rows, .{ .buf = s.h[0], .off = (span.at - p0) * geo.hidden * 2 }, span.n); + }; + e.setPipeline(ks.rms); + bind(e, &.{s.h[0]}, 0); + e.setBuffer(w.layers[0].input_norm.buf, w.layers[0].input_norm.off, 1); + bind(e, &.{s.normed[0]}, 2); + e.setValue(cfg.eps, 3); + noSums(e, s, 4); + e.dispatchGroups(.{ .width = rows }, .{ .width = 1024 }); + var cur: usize = 0; + for (0..geo.layers) |i| { + try layer(m, e, s, c, i, cur, r32, .{}); + cur = 1 - cur; + } + // the head on the last row only (the lane kernel's one-row bits), soft-capped + const head = try m.projections.get(m.gpa, m.device, w.head); + try head.encode(e, .{ .buf = s.normed[cur], .off = (rows - 1) * geo.hidden * 2 }, .{ .buf = w.head.w }, .{ .buf = w.head.sb }, .{ .buf = s.logits }, 1); + e.setPipeline(ks.softcap); + bind(e, &.{s.logits}, 0); + e.setValue(cfg.softcap, 1); + e.dispatchThreads(.{ .width = cfg.vocab }, .{ .width = 256 }); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("gemma4: prompt chunk failed: {s}", .{msg}); + return error.GpuFailed; + } + c.len += rows; +} + +/// Which of a layer's dispatches to encode: a chunk takes them all; the bench times the stages apart. +pub const Stages = struct { mm: bool = true, attn: bool = true, glue: bool = true }; + +/// Layer `i` over `rows` rows of s.normed[cur] (keys at the cache's length) into s.h[1 - cur] and s.normed[1 - cur]. +pub fn layer(m: *Model, e: mtl.ComputeEncoder, s: *Scratch, c: *state.Cache, i: usize, cur: usize, rows: u32, st: Stages) !void { + const geo = m.config.geo; + const global = cfg.global(i); + const kd: *const kernels.Kind = m.kernels.kind(global); + const l = &m.weights.layers[i]; + const nk = geo.kvHeads(i); + const dh = cfg.headDim(i); + const qn = geo.query_heads * dh; + const width = qn + nk * dh * @as(usize, if (global) 1 else 2); + if (st.mm) { + mm(m, e, l.mlx[0], s.normed[cur], s.qkv, 0, width, rows); + mm(m, e, l.mlx[1], s.normed[cur], s.qkv, qn, width, rows); + if (!global) mm(m, e, l.mlx[2], s.normed[cur], s.qkv, qn + nk * dh, width, rows); + } + if (st.glue) { + e.setPipeline(kd.qkv); + bind(e, &.{s.qkv}, 0); + e.setBuffer(l.q_norm.buf, l.q_norm.off, 1); + e.setBuffer(l.k_norm.buf, l.k_norm.off, 2); + bind(e, &.{ if (global) m.inv_global else m.inv_sliding, s.pos }, 3); + e.setValue(cfg.eps, 5); + bind(e, &.{ s.q, s.k, s.v, c.k[i], c.v[i] }, 6); // keys and values straight into the cache too + e.setValue([3]i32{ if (global) 0 else state.ring, @intCast(state.slotsOf(i, c.capacity)), 0 }, 11); + e.dispatchGroups(.{ .width = geo.query_heads + 2 * nk, .height = rows }, .{ .width = 32 }); + } + if (st.attn) { + e.setPipeline(kd.prompt_attn); + bind(e, &.{ s.q, c.k[i], c.v[i], s.attn }, 0); + e.setValue([3]i32{ @intCast(rows), @intCast(c.len), @intCast(state.slotsOf(i, c.capacity)) }, 4); + bind(e, &.{s.hi}, 5); + e.dispatchGroups(.{ .width = (rows + 31) / 32, .height = geo.query_heads }, .{ .width = 32 * dh / 64 }); + } + if (st.mm) { + mm(m, e, l.mlx[3], s.attn, s.o, 0, geo.hidden, rows); + } + if (st.glue) { + e.setPipeline(kd.attn_tail); + bind(e, &.{ s.h[cur], s.o }, 0); + e.setBuffer(l.post_attention.buf, l.post_attention.off, 2); + e.setBuffer(l.pre_feedforward.buf, l.pre_feedforward.off, 3); + e.setValue(cfg.eps, 4); + bind(e, &.{ s.hn, s.n_mlp }, 5); + noSums(e, s, 7); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); + } + if (st.mm) { + mm(m, e, l.mlx[4], s.n_mlp, s.gate_up, 0, 2 * geo.intermediate, rows); + mm(m, e, l.mlx[5], s.n_mlp, s.gate_up, geo.intermediate, 2 * geo.intermediate, rows); + } + if (st.glue) { + e.setPipeline(kd.geglu); + bind(e, &.{ s.gate_up, s.act }, 0); + e.setValue(@as(u32, @intCast(geo.intermediate)), 2); + noSums(e, s, 3); + e.setValue(@as(i32, 0), 5); + e.dispatchGroups(.{ .width = geo.intermediate / 64, .height = rows }, .{ .width = 64 }); + } + if (st.mm) { + mm(m, e, l.mlx[6], s.act, s.y, 0, geo.hidden, rows); + } + if (!st.glue) return; + e.setPipeline(kd.dense_tail); + bind(e, &.{ s.hn, s.y }, 0); + e.setBuffer(l.post_feedforward.buf, l.post_feedforward.off, 2); + e.setBuffer(l.scalar.buf, l.scalar.off, 3); + const nx = if (i + 1 < geo.layers) m.weights.layers[i + 1].input_norm else m.weights.final_norm; + e.setBuffer(nx.buf, nx.off, 4); + e.setValue(cfg.eps, 5); + bind(e, &.{ s.h[1 - cur], s.normed[1 - cur] }, 6); + noSums(e, s, 8); + fwd.tapOut(e, m, s.taps, s.hn, i, 10); + e.dispatchGroups(.{ .width = rows }, .{ .width = 256 }); +} diff --git a/zig/src/families/gemma4/snapshot.zig b/zig/src/families/gemma4/snapshot.zig new file mode 100644 index 0000000000..97d1c3a617 --- /dev/null +++ b/zig/src/families/gemma4/snapshot.zig @@ -0,0 +1,112 @@ +//! A stream's prompt state kept for reuse: each sliding ring's last window of keys and values, the global layers' all. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); +const st = @import("state.zig"); + +/// One kept state: its bytes on the GPU (private) and the prompt tokens it holds. +pub const Snap = struct { + buf: mtl.Buffer, + at: u32, + bytes: u64, + /// The DFlash drafter's context for the same stream at `at` (dflash.zig keepLane), when it had one there. + draft: ?mtl.Buffer = null, + + pub fn free(self: *Snap) void { + self.buf.deinit(); + if (self.draft) |b| b.deinit(); + } +}; + +/// One contiguous piece of a layer's key (or value) buffer: its offset and length in bytes. +const Piece = struct { layer: usize, value: bool, off: usize, len: usize }; + +/// Positions [lo, hi) of one head's ring of `width`-byte slots from `base`: one piece, or two where it wraps. +fn ring(layer: usize, value: bool, base: usize, width: usize, lo: usize, hi: usize, out: *std.ArrayList(Piece), gpa: std.mem.Allocator) !void { + if (hi <= lo) return; + const first = lo % st.ring; + const n = @min(hi - lo, st.ring - first); + try out.append(gpa, .{ .layer = layer, .value = value, .off = base + first * width, .len = n * width }); + if (n < hi - lo) try out.append(gpa, .{ .layer = layer, .value = value, .off = base, .len = (hi - lo - n) * width }); +} + +/// The pieces of `cache` (its capacity places the global heads) a state at `at` holds: what the next rows read. +fn pieces(cache: *const st.Cache, at: usize, out: *std.ArrayList(Piece), gpa: std.mem.Allocator) !void { + for (0..cache.geo.layers) |i| { + const width = cfg.headDim(i) * 2; + const slots = st.slotsOf(i, cache.capacity); + for (0..cache.geo.kvHeads(i)) |h| for ([_]bool{ false, true }) |value| { + const base = h * slots * width; + if (cfg.global(i)) { + try out.append(gpa, .{ .layer = i, .value = value, .off = base, .len = at * width }); + } else try ring(i, value, base, width, at -| cfg.window, at, out, gpa); + }; + } +} + +/// Bytes a state after `at` prompt tokens takes. +pub fn bytes(geo: cfg.Geometry, at: usize) u64 { + var total: u64 = 0; + for (0..geo.layers) |i| { + const rows = if (cfg.global(i)) at else @min(at, cfg.window); + total += 2 * geo.kvHeads(i) * rows * cfg.headDim(i) * 2; + } + return total; +} + +/// Copy `cache`'s state at its length into a new snapshot. +pub fn save(gpa: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, cache: *const st.Cache) !Snap { + const n = bytes(cache.geo, cache.len); + const buf = try device.buffer(@intCast(@max(n, 16)), mtl.ResourceOptions.private); + errdefer buf.deinit(); + try copy(gpa, queue, cache, cache.len, buf, true); + return .{ .buf = buf, .at = @intCast(cache.len), .bytes = n }; +} + +/// Make `cache` (empty, its capacity past the state) hold `snap`'s state. +pub fn restore(gpa: std.mem.Allocator, queue: mtl.Queue, snap: *const Snap, cache: *st.Cache) !void { + if (snap.at >= cache.capacity) return error.GemmaSnapshotTooLong; + try copy(gpa, queue, cache, snap.at, snap.buf, false); + cache.len = snap.at; +} + +/// The pieces between the cache and a snapshot buffer, one blit pass, waited on. +fn copy(gpa: std.mem.Allocator, queue: mtl.Queue, cache: *const st.Cache, at: usize, buf: mtl.Buffer, out: bool) !void { + var list: std.ArrayList(Piece) = .empty; + defer list.deinit(gpa); + try pieces(cache, at, &list, gpa); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = queue.commandBuffer(); + const e = cb.blit(); + var to: usize = 0; + for (list.items) |p| { + if (p.len == 0) continue; + const live = if (p.value) cache.v[p.layer] else cache.k[p.layer]; + if (out) e.copy(live, p.off, buf, to, p.len) else e.copy(buf, to, live, p.off, p.len); + to += p.len; + } + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.log.err("gemma4 prompt state copy: {s}", .{msg}); + return error.GpuFailed; + } +} + +test "a state's bytes match its pieces, rings wrapped or not" { + for ([_]cfg.Geometry{ cfg.dense31b, cfg.unified12b }) |geo| for ([_]usize{ 1, 5, 1000, 1024, 1500, 1536, 2600, 100_003 }) |at| { + var cache = st.Cache{ .k = undefined, .v = undefined, .geo = geo, .logits = undefined, .capacity = 1 << 18 }; + var list: std.ArrayList(Piece) = .empty; + defer list.deinit(std.testing.allocator); + try pieces(&cache, at, &list, std.testing.allocator); + var sum: u64 = 0; + for (list.items) |p| { + sum += p.len; + const slots = st.slotsOf(p.layer, cache.capacity); + try std.testing.expect(p.off + p.len <= geo.kvHeads(p.layer) * slots * cfg.headDim(p.layer) * 2); + } + try std.testing.expectEqual(bytes(geo, at), sum); + }; +} diff --git a/zig/src/families/gemma4/state.zig b/zig/src/families/gemma4/state.zig new file mode 100644 index 0000000000..9b377bfcea --- /dev/null +++ b/zig/src/families/gemma4/state.zig @@ -0,0 +1,159 @@ +//! Gemma 4 dense decode state: each stream's caches (a ring a sliding layer, a buffer a global one) and scratch. +const std = @import("std"); +const mtl = @import("metal"); +const cfg = @import("config.zig"); + +const opts = mtl.ResourceOptions.private | mtl.ResourceOptions.untracked; +const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + +/// A verify window's rows at most (the exact width), and a prompt window's (the reg projection's widest). +pub const max_rows = 16; +pub const prompt_rows = 128; +/// Ints a row's attention meta takes: first chunk, chunks, ring, rows, first new key, padding. +pub const meta_ints = 8; +/// Rows a prompt chunk takes (prompt.zig), and a sliding layer's ring: the window plus a chunk written past it. +pub const chunk_rows = 512; +pub const ring = cfg.window + chunk_rows; + +/// Where a stream's latest tapped rows sit: `rows` rows from cache position `start`, current while `gen` matches. +pub const Tap = struct { buf: ?mtl.Buffer = null, off: usize = 0, start: u64 = 0, rows: usize = 0, gen: ?u64 = null }; + +/// One stream's keys and values, per layer [kv heads, slots, head dim] bf16, and how many positions it holds. +pub const Cache = struct { + k: [cfg.max_layers]mtl.Buffer, + v: [cfg.max_layers]mtl.Buffer, + geo: cfg.Geometry, + capacity: usize, + len: usize = 0, + /// The cache length the last verify started from (keep rolls back to it plus the kept rows). + from: usize = 0, + /// The last row's logits, bf16 [vocab]: the prompt's (first) or a one-token round's (queue). + logits: mtl.Buffer, + /// The drafter's rows: the prompt's last (in `last_tap`), then each verify's in the scratch. + tap: Tap = .{}, + last_tap: ?mtl.Buffer = null, + + pub fn init(device: mtl.Device, geo: cfg.Geometry, capacity: usize) !Cache { + var c: Cache = .{ .k = undefined, .v = undefined, .geo = geo, .capacity = capacity, .logits = try device.buffer(cfg.vocab * 2, shared) }; + errdefer c.logits.deinit(); + var made: usize = 0; + errdefer for (0..made) |i| { + c.k[i].deinit(); + c.v[i].deinit(); + }; + for (0..geo.layers) |i| { + const slots = slotsOf(i, capacity); + const bytes = geo.kvHeads(i) * slots * cfg.headDim(i) * 2; + c.k[i] = try device.buffer(bytes, opts); + c.v[i] = device.buffer(bytes, opts) catch |e| { + c.k[i].deinit(); + return e; + }; + made += 1; + } + return c; + } + + pub fn deinit(c: *Cache) void { + c.logits.deinit(); + if (c.last_tap) |b| b.deinit(); + for (0..c.geo.layers) |i| { + c.k[i].deinit(); + c.v[i].deinit(); + } + } +}; + +/// Bytes a stream's caches take at `capacity` (Cache.init's buffers). +pub fn cacheBytes(geo: cfg.Geometry, capacity: usize) u64 { + var total: u64 = cfg.vocab * 2; + for (0..geo.layers) |i| total += 2 * geo.kvHeads(i) * slotsOf(i, capacity) * cfg.headDim(i) * 2; + return total; +} + +/// Cache slots for layer `i`: the ring on a sliding layer, the whole capacity on a global one. +pub fn slotsOf(i: usize, capacity: usize) usize { + return if (cfg.global(i)) capacity else ring; +} + +/// Activations for one window of up to `max_rows` rows, shared by every layer. +pub const Scratch = struct { + ids: mtl.Buffer, + pos: mtl.Buffer, + /// Per layer kind (0 sliding, 1 global): each row's first key, and at a sub-batch's first row its meta. + lo: [2]mtl.Buffer, + meta: [2]mtl.Buffer, + h: [2]mtl.Buffer, + normed: [2]mtl.Buffer, + qkv: mtl.Buffer, + q: mtl.Buffer, + k: mtl.Buffer, + v: mtl.Buffer, + attn: mtl.Buffer, + o: mtl.Buffer, + hn: mtl.Buffer, + n_mlp: mtl.Buffer, + gate_up: mtl.Buffer, + act: mtl.Buffer, + y: mtl.Buffer, + pm: mtl.Buffer, + pl: mtl.Buffer, + po: mtl.Buffer, + logits: mtl.Buffer, + /// The next projection's group sums, written by the kernel that made its input (reg lane kernel order). + sums: mtl.Buffer, + /// The drafter's tapped rows [rows, taps x hidden] bf16, once a drafter asked for them. + taps: ?mtl.Buffer = null, + /// Each row's greedy token by the GPU argmax (tree_round Pick: token, nonfinite flag). + picks: mtl.Buffer, + /// Attention partial slots: heads x rows x chunks (256 keys a chunk on global layers; sliding ones in the window). + slots: usize, + + pub fn init(device: mtl.Device, geo: cfg.Geometry, capacity: usize) !Scratch { + const r = prompt_rows; + const d = geo.hidden; + const qd = geo.queryWidth(); + const kv = geo.kvWidth(); + const slots = geo.query_heads * max_rows * @max(capacity / 256 + 2, (cfg.window + max_rows) / 128 + 2); + var s: Scratch = undefined; + s.slots = slots; + s.taps = null; + var made: usize = 0; + errdefer s.release(made); + inline for (.{ .{ "ids", r * 4, true }, .{ "pos", r * 4, true }, .{ "qkv", r * (qd + 2 * kv) * 2, false }, .{ "q", r * qd * 2, false }, .{ "k", r * kv * 2, false }, .{ "v", r * kv * 2, false }, .{ "attn", r * qd * 2, false }, .{ "o", r * d * 2, false }, .{ "hn", r * d * 2, false }, .{ "n_mlp", r * d * 2, false }, .{ "gate_up", r * 2 * geo.intermediate * 2, false }, .{ "act", r * geo.intermediate * 2, false }, .{ "y", r * d * 2, false }, .{ "pm", slots * 4, false }, .{ "pl", slots * 4, false }, .{ "po", slots * cfg.global_dim * 4, false }, .{ "logits", r * cfg.vocab * 2, true }, .{ "sums", r * (geo.intermediate / 64) * 4, false }, .{ "picks", r * 8, true } }) |f| { + @field(s, f[0]) = try device.buffer(f[1], if (f[2]) shared else opts); + made += 1; + } + for (&s.lo) |*b| { + b.* = try device.buffer(r * 4, shared); + made += 1; + } + for (&s.meta) |*b| { + b.* = try device.buffer(r * meta_ints * 4, shared); + made += 1; + } + for ([_][]mtl.Buffer{ &s.h, &s.normed }) |pair| for (pair) |*b| { + b.* = try device.buffer(r * d * 2, opts); + made += 1; + }; + return s; + } + + pub fn deinit(s: *Scratch) void { + s.release(std.math.maxInt(usize)); + if (s.taps) |b| b.deinit(); + } + + /// Releases the first `made` buffers in init's order, so a failed init frees what it made. + fn release(s: *Scratch, made: usize) void { + var i: usize = 0; + inline for (.{ "ids", "pos", "qkv", "q", "k", "v", "attn", "o", "hn", "n_mlp", "gate_up", "act", "y", "pm", "pl", "po", "logits", "sums", "picks" }) |f| { + if (i < made) @field(s, f).deinit(); + i += 1; + } + for (s.lo ++ s.meta ++ s.h ++ s.normed) |b| { + if (i < made) b.deinit(); + i += 1; + } + } +}; diff --git a/zig/src/families/gemma4/vision.zig b/zig/src/families/gemma4/vision.zig new file mode 100644 index 0000000000..d5b5db259a --- /dev/null +++ b/zig/src/families/gemma4/vision.zig @@ -0,0 +1,352 @@ +//! Gemma 4's image tower on Metal: patch projection and 2D position rows, 27 encoder layers, 3 x 3 pooling, projection. +const std = @import("std"); +const mtl = @import("metal"); +const ks = @import("kernel_sources"); +const frags = @import("../../core/frags.zig"); +const ck = @import("../../core/checkpoint_metal.zig"); +const affine_mm = @import("../../core/affine_mm.zig"); +const lane = @import("../../core/lane_projection.zig"); +const Mlx = @import("weights.zig").Mlx; +const image = @import("image.zig"); + +pub const width = 1152; +pub const heads = 16; +pub const head_dim = 72; +pub const depth = 27; +pub const mlp_inter = 4304; +pub const max_patches = image.max_budget * image.pool * image.pool; +const table_rows = 10240; // position_embedding_size: each axis's table +const chunk = 256; // query rows an attention pass scores at once +const eps: f32 = 1e-6; +const rope_theta: f32 = 100; + +const Width = @import("widths.zig").Width; + +pub const Ref = lane.Ref; +fn ref(t: ck.Tensor) Ref { + return .{ .buf = t.buffer, .off = t.offset }; +} + +/// The projection as affine_mm reads it at its config width: words [out_hidden, 1152 * bits / 32], bf16 tables. +fn projectionFits(w: ck.Tensor, s: ck.Tensor, b: ck.Tensor, q: Width, out_hidden: usize) bool { + const words = @as(usize, width) * q.bits / 32; + if (w.dtype != .u32 or w.rank != 2 or w.shape[0] != out_hidden or w.shape[1] != words or w.count() != out_hidden * words) return false; + for ([_]ck.Tensor{ s, b }) |t| { + if (t.dtype != .bf16 or t.rank != 2 or t.shape[0] != out_hidden or t.shape[1] != width / q.group or t.count() != out_hidden * width / q.group) return false; + } + return true; +} + +const glue_names = [_][:0]const u8{ "g4v_pos", "g4v_rms", "g4v_rms_add", "g4v_qkv", "g4v_softmax", "g4v_heads", "g4v_geglu", "g4v_pool", "g4v_rows" }; +const G = enum { pos, rms, rms_add, qkv, softmax, heads_, geglu, pool, rows }; + +const Layer = struct { in_norm: Ref, post_attn: Ref, pre_ff: Ref, post_ff: Ref, q: Ref, k: Ref, v: Ref, o: Ref, q_norm: Ref, k_norm: Ref, gate: Ref, up: Ref, down: Ref }; + +/// No tower tensor stays in its file: vision_tower. and embed_vision. pack, 256-byte aligned, into their own buffers. +fn never(_: []const u8, _: usize) bool { + return false; +} + +/// The checkpoint's tower tensors, packed 256-byte aligned beside the language model's (the checkpoint frees them). +pub fn addTensors(c: *ck.Checkpoint, device: mtl.Device, files: []const [:0]u8) !void { + for (files) |path| for ([_][]const u8{ "vision_tower.", "embed_vision." }) |p| try c.addFileSplit(device, path, "", p, never); +} + +/// Whether the checkpoint has the tower (a text-only conversion does not). +pub fn present(c: *const ck.Checkpoint) bool { + _ = c.get("vision_tower.patch_embedder.input_proj.weight") catch return false; + return true; +} + +pub const Vision = struct { + gpa: std.mem.Allocator, + lib: mtl.Library, + gemm: [3][2][2]mtl.Pipeline, // [bf16 nt, f32 nt, f32 nn][M % 64 == 0][N % 128 == 0] + glue: [glue_names.len]mtl.Pipeline, + mm: affine_mm.Pipes, // the model's at the projection's width + out_hidden: usize, // the language model's width: the projection's rows + input_proj: Ref, + table: Ref, + std_bias: Ref, + std_scale: Ref, + layers: [depth]Layer, + proj: Mlx, + scratch: mtl.Buffer, + x: Ref, + xn: Ref, + qkv: Ref, + q: Ref, + k: Ref, + v: Ref, + s: Ref, + o: Ref, + att: Ref, + gu: Ref, + act: Ref, + pooled: Ref, + pn: Ref, + /// Host-written per image: the patches as bf16 2 (p - 0.5), and the RoPE's cos and sin rows. + patches: mtl.Buffer, + cos: mtl.Buffer, + sin: mtl.Buffer, + + pub fn load(gpa: std.mem.Allocator, device: mtl.Device, c: *const ck.Checkpoint, mm: affine_mm.Pipes, q: Width, out_hidden: usize) !*Vision { + const v = try gpa.create(Vision); + errdefer gpa.destroy(v); + v.gpa = gpa; + v.mm = mm; + v.out_hidden = out_hidden; + const body = try frags.source(device, gpa, ks.gemma4_vision); + defer gpa.free(body); + const src = try std.fmt.allocPrint(gpa, "#define G4V_OUT {d}\n{s}", .{ out_hidden, body }); + defer gpa.free(src); + v.lib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + errdefer v.lib.deinit(); + var made: usize = 0; + errdefer for (0..made) |i| v.pipeline(i).deinit(); + const kinds = [_][]const u8{ "bf16_nt", "f32_nt", "f32_nn" }; + for (kinds, 0..) |kind, ki| for (0..2) |m| for (0..2) |n| { + var name: [64]u8 = undefined; + const full = try std.fmt.bufPrintSentinel(&name, "g4v_gemm_{s}_{d}{d}", .{ kind, m, n }, 0); + v.gemm[ki][m][n] = try mtl.Pipeline.init(device, v.lib, full, false); + made += 1; + }; + for (glue_names, 0..) |n, i| { + v.glue[i] = try mtl.Pipeline.init(device, v.lib, n, false); + made += 1; + } + const t = struct { + fn get(cp: *const ck.Checkpoint, comptime fmt: []const u8, args: anytype, count: usize) !ck.Tensor { + var b: [160]u8 = undefined; + const name = try std.fmt.bufPrint(&b, fmt, args); + const x = try cp.get(name); + if (x.dtype != .bf16 or x.count() != count) { + std.log.err("gemma4 vision: {s} is not bf16[{d}]", .{ name, count }); + return error.UnexpectedGemmaVisionTensor; + } + return x; + } + }; + const pe = "vision_tower.patch_embedder."; + v.input_proj = ref(try t.get(c, pe ++ "input_proj.weight", .{}, width * image.patch_width)); + v.table = ref(try t.get(c, pe ++ "position_embedding_table", .{}, 2 * table_rows * width)); + v.std_bias = ref(try t.get(c, "vision_tower.std_bias", .{}, width)); + v.std_scale = ref(try t.get(c, "vision_tower.std_scale", .{}, width)); + for (&v.layers, 0..) |*l, i| { + const p = "vision_tower.encoder.layers.{d}."; + l.* = .{ + .in_norm = ref(try t.get(c, p ++ "input_layernorm.weight", .{i}, width)), + .post_attn = ref(try t.get(c, p ++ "post_attention_layernorm.weight", .{i}, width)), + .pre_ff = ref(try t.get(c, p ++ "pre_feedforward_layernorm.weight", .{i}, width)), + .post_ff = ref(try t.get(c, p ++ "post_feedforward_layernorm.weight", .{i}, width)), + .q = ref(try t.get(c, p ++ "self_attn.q_proj.linear.weight", .{i}, width * width)), + .k = ref(try t.get(c, p ++ "self_attn.k_proj.linear.weight", .{i}, width * width)), + .v = ref(try t.get(c, p ++ "self_attn.v_proj.linear.weight", .{i}, width * width)), + .o = ref(try t.get(c, p ++ "self_attn.o_proj.linear.weight", .{i}, width * width)), + .q_norm = ref(try t.get(c, p ++ "self_attn.q_norm.weight", .{i}, head_dim)), + .k_norm = ref(try t.get(c, p ++ "self_attn.k_norm.weight", .{i}, head_dim)), + .gate = ref(try t.get(c, p ++ "mlp.gate_proj.linear.weight", .{i}, mlp_inter * width)), + .up = ref(try t.get(c, p ++ "mlp.up_proj.linear.weight", .{i}, mlp_inter * width)), + .down = ref(try t.get(c, p ++ "mlp.down_proj.linear.weight", .{i}, width * mlp_inter)), + }; + } + const pp = "embed_vision.embedding_projection."; + const pw = try c.get(pp ++ "weight"); + const ps = try c.get(pp ++ "scales"); + const pb = try c.get(pp ++ "biases"); + if (!projectionFits(pw, ps, pb, q, out_hidden)) { + std.log.err("gemma4 vision: {s}* is not u32[{d}, {d}] with bf16[{d}, {d}] scales and biases", .{ pp, out_hidden, @as(usize, width) * q.bits / 32, out_hidden, width / q.group }); + return error.UnexpectedGemmaVisionTensor; + } + v.proj = .{ .w = ref(pw), .s = ref(ps), .b = ref(pb), .n = @intCast(out_hidden), .k = width, .bits = q.bits }; + const P: usize = max_patches; + const T: usize = image.max_budget; + const sizes = [_]usize{ P * width * 2, P * width * 2, P * 3 * width * 2, heads * P * head_dim * 4, heads * P * head_dim * 4, heads * P * head_dim * 4, heads * chunk * P * 4, heads * chunk * head_dim * 4, P * width * 2, P * 2 * mlp_inter * 2, P * mlp_inter * 2, T * width * 2, T * width * 2 }; + var total: usize = 0; + for (sizes) |s| total += std.mem.alignForward(usize, s, 256); + v.scratch = try device.buffer(total, mtl.ResourceOptions.private); + errdefer v.scratch.deinit(); + var off: usize = 0; + const fields = [_]*Ref{ &v.x, &v.xn, &v.qkv, &v.q, &v.k, &v.v, &v.s, &v.o, &v.att, &v.gu, &v.act, &v.pooled, &v.pn }; + for (fields, sizes) |f, s| { + f.* = .{ .buf = v.scratch, .off = off }; + off += std.mem.alignForward(usize, s, 256); + } + const shared = mtl.ResourceOptions.shared; + v.patches = try device.buffer(P * image.patch_width * 2, shared); + errdefer v.patches.deinit(); + v.cos = try device.buffer(P * head_dim * 2, shared); + errdefer v.cos.deinit(); + v.sin = try device.buffer(P * head_dim * 2, shared); + return v; + } + + fn pipeline(v: *Vision, i: usize) mtl.Pipeline { + const g = 3 * 2 * 2; + return if (i < g) v.gemm[i / 4][(i / 2) % 2][i % 2] else v.glue[i - g]; + } + + pub fn deinit(v: *Vision) void { + for (0..3 * 2 * 2 + glue_names.len) |i| v.pipeline(i).deinit(); + v.lib.deinit(); + v.scratch.deinit(); + v.patches.deinit(); + v.cos.deinit(); + v.sin.deinit(); + v.gpa.destroy(v); + } + + /// An image's patches (image.prepare's) to `out` [gh * gw / 9, out_hidden]; `stop`: leave `x` after that many. + pub fn encode(v: *Vision, e: mtl.ComputeEncoder, pixels: []const f32, gh: u32, gw: u32, out: Ref, stop: ?u32) !void { + const P: u32 = gh * gw; + if (gh == 0 or gw == 0 or gh % image.pool != 0 or gw % image.pool != 0 or P > max_patches) return error.GemmaImageGrid; + if (pixels.len != @as(usize, P) * image.patch_width) return error.GemmaImageGrid; + for (pixels, v.patches.slice(u16, pixels.len)) |p, *d| d.* = bf16(2 * (p - 0.5)); + rope(v.cos.slice(u16, @as(usize, P) * head_dim), v.sin.slice(u16, @as(usize, P) * head_dim), gw, P); + v.linear(e, .{ .buf = v.patches }, v.input_proj, v.x, P, width, image.patch_width, width); + v.run(e, .pos, &.{ v.x, v.table }, .{ width, P, 1 }, .{ 256, 1, 1 }, false, gw); + if (stop) |s| if (s == 0) return; + for (v.layers, 0..) |l, i| { + v.layer(e, l, P); + if (stop) |s| if (i + 1 == s) return; + } + const T = P / (image.pool * image.pool); + v.run(e, .pool, &.{ v.x, v.std_bias, v.std_scale, v.pooled }, .{ width, T, 1 }, .{ 256, 1, 1 }, false, gw); + v.rms(e, v.pooled, null, v.pn, T); + affine_mm.dense(e, v.mm, v.pn, v.proj, out, T); + } + + /// Rows of an encoded image over `n` embedded prompt rows (`dst`, out_hidden bf16 a row), in the prompt's encoder. + pub fn inject(v: *Vision, e: mtl.ComputeEncoder, src: Ref, dst: Ref, n: u32) void { + v.run(e, .rows, &.{ src, dst }, .{ v.out_hidden / 8, n, 1 }, .{ 64, 1, 1 }, false, null); + } + + fn layer(v: *Vision, e: mtl.ComputeEncoder, l: Layer, P: u32) void { + v.rms(e, v.x, l.in_norm, v.xn, P); + for ([_]Ref{ l.q, l.k, l.v }, 0..) |w, i| v.linear(e, v.xn, w, .{ .buf = v.qkv.buf, .off = v.qkv.off + i * width * 2 }, P, width, width, 3 * width); + e.setPipeline(v.glue[@backingInt(G.qkv)]); + for ([_]Ref{ v.qkv, l.q_norm, l.k_norm, .{ .buf = v.cos }, .{ .buf = v.sin }, v.q, v.k, v.v }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(P, 8); + e.setValue(eps, 9); + e.dispatchGroups(mtl.Size.of(@as(usize, P) * heads * 3, 1, 1), mtl.Size.of(32, 1, 1)); + var q0: u32 = 0; + while (q0 < P) : (q0 += chunk) { + const lc = @min(chunk, P - q0); + v.gemmAt(e, 1, .{ .buf = v.q.buf, .off = v.q.off + @as(usize, q0) * head_dim * 4 }, v.k, v.s, lc, P, head_dim, head_dim, head_dim, P, heads, @as(usize, P) * head_dim, @as(usize, P) * head_dim, @as(usize, lc) * P); + e.setPipeline(v.glue[@backingInt(G.softmax)]); + e.setBuffer(v.s.buf, v.s.off, 0); + e.setValue(P, 1); + e.dispatchGroups(mtl.Size.of(@as(usize, heads) * lc, 1, 1), mtl.Size.of(256, 1, 1)); + v.gemmAt(e, 2, v.s, v.v, v.o, lc, head_dim, P, P, head_dim, head_dim, heads, @as(usize, lc) * P, @as(usize, P) * head_dim, @as(usize, lc) * head_dim); + e.setPipeline(v.glue[@backingInt(G.heads_)]); + e.setBuffer(v.o.buf, v.o.off, 0); + e.setBuffer(v.att.buf, v.att.off, 1); + e.setValue([2]u32{ lc, q0 }, 2); + e.dispatchThreads(mtl.Size.of(width, lc, 1), mtl.Size.of(256, 1, 1)); + } + v.linear(e, v.att, l.o, v.xn, P, width, width, width); + v.rmsAdd(e, v.x, v.xn, l.post_attn, P); + v.rms(e, v.x, l.pre_ff, v.xn, P); + v.linear(e, v.xn, l.gate, v.gu, P, mlp_inter, width, 2 * mlp_inter); + v.linear(e, v.xn, l.up, .{ .buf = v.gu.buf, .off = v.gu.off + mlp_inter * 2 }, P, mlp_inter, width, 2 * mlp_inter); + v.run(e, .geglu, &.{ v.gu, v.act }, .{ mlp_inter, P, 1 }, .{ 256, 1, 1 }, false, mlp_inter); + v.linear(e, v.act, l.down, v.xn, P, width, mlp_inter, width); + v.rmsAdd(e, v.x, v.xn, l.post_ff, P); + } + + fn run(v: *Vision, e: mtl.ComputeEncoder, g: G, bufs: []const Ref, grid: [3]usize, group: [3]usize, groups: bool, value: ?u32) void { + e.setPipeline(v.glue[@backingInt(g)]); + for (bufs, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + if (value) |x| e.setValue(x, bufs.len); + const a = mtl.Size.of(grid[0], grid[1], grid[2]); + const b = mtl.Size.of(group[0], group[1], group[2]); + if (groups) e.dispatchGroups(a, b) else e.dispatchThreads(a, b); + } + + /// y = Gemma4RMSNorm(x) [rows, width], by `w` when there is one. + fn rms(v: *Vision, e: mtl.ComputeEncoder, x: Ref, w: ?Ref, y: Ref, rows: u32) void { + e.setPipeline(v.glue[@backingInt(G.rms)]); + for ([_]Ref{ x, w orelse x, y }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(extern struct { width: u32, weighted: u32, eps: f32 }{ .width = width, .weighted = @intFromBool(w != null), .eps = eps }, 3); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// x += Gemma4RMSNorm(y) by w: a branch's post norm into the residual. + fn rmsAdd(v: *Vision, e: mtl.ComputeEncoder, x: Ref, y: Ref, w: Ref, rows: u32) void { + e.setPipeline(v.glue[@backingInt(G.rms_add)]); + for ([_]Ref{ x, y, w }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(eps, 3); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// y [M, N] (rows `ldy` apart) = x [M, K] w^T in bf16 (w as stored, [N, K]). + fn linear(v: *Vision, e: mtl.ComputeEncoder, x: Ref, w: Ref, y: Ref, M: u32, N: u32, K: u32, ldy: u32) void { + v.gemmAt(e, 0, x, w, y, M, N, K, K, K, ldy, 1, 0, 0, 0); + } + + /// tfp::gemm_nax: d [batch, M, N] = a x b (kind 0: bf16 a b^T, 1: fp32 a b^T, 2: fp32 a b), batches `s*` apart. + fn gemmAt(v: *Vision, e: mtl.ComputeEncoder, kind: usize, a: Ref, b: Ref, d: Ref, M: u32, N: u32, K: u32, lda: u32, ldb: u32, ldd: u32, batch: u32, sa: usize, sb: usize, sd: usize) void { + const tn = (N + 127) / 128; + const tm = (M + 63) / 64; + e.setPipeline(v.gemm[kind][@intFromBool(M % 64 == 0)][@intFromBool(N % 128 == 0)]); + e.setBuffer(a.buf, a.off, 0); + e.setBuffer(b.buf, b.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ M, N, K, lda, ldb, ldd, tn, tm, 2, K / 256, sa, sb, sd, 0, 0 }; + for (vals, 0..) |x, i| p[i] = @intCast(x); + e.setBytes(std.mem.asBytes(&p), 2); + e.setBuffer(d.buf, d.off, 3); + e.dispatchThreads(mtl.Size.of(@as(usize, tn << 2) * 32, (tm + 3) / 4 * 4, @as(usize, batch) * 2), mtl.Size.of(32, 4, 2)); + } +}; + +/// fp32 to bf16, round to nearest even. +pub fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +/// Gemma4VisionRotaryEmbedding, gw wide: dims 0-35 by x and 36-71 by y, 18 frequencies twice each, bf16 cos and sin. +fn rope(cs: []u16, sn: []u16, gw: u32, P: u32) void { + const half = head_dim / 2; + var inv: [half / 2]f32 = undefined; + for (&inv, 0..) |*f, i| f.* = 1.0 / std.math.pow(f32, rope_theta, @as(f32, @floatFromInt(2 * i)) / @as(f32, half)); + for (0..P) |p| for (0..head_dim) |d| { + const axis: f32 = @floatFromInt(if (d < half) p % gw else p / gw); + const angle = axis * inv[(d % half) % (half / 2)]; + cs[p * head_dim + d] = bf16(@cos(angle)); + sn[p * head_dim + d] = bf16(@sin(angle)); + }; +} + +test "the projection is admitted only as packed u32 words at its width with bf16 scales and biases per group" { + const T = struct { + fn of(dtype: ck.DType, rows: usize, cols: usize) ck.Tensor { + var t: ck.Tensor = .{ .buffer = undefined, .offset = 0, .bytes = rows * cols * dtype.size(), .dtype = dtype, .rank = 2 }; + t.shape[0] = rows; + t.shape[1] = cols; + return t; + } + }; + const out_hidden = @import("config.zig").dense31b.hidden; + const four: Width = .{ .bits = 4, .group = 64 }; + const words = T.of(.u32, out_hidden, width / 8); + const scales = T.of(.bf16, out_hidden, width / 64); + try std.testing.expect(projectionFits(words, scales, scales, four, out_hidden)); + const f32s = T.of(.f32, out_hidden, width / 64); + try std.testing.expect(!projectionFits(words, f32s, scales, four, out_hidden)); + try std.testing.expect(!projectionFits(words, scales, f32s, four, out_hidden)); + const group32 = T.of(.bf16, out_hidden, width / 32); + try std.testing.expect(!projectionFits(words, group32, group32, four, out_hidden)); + var short = scales; // the header's shape, the file's bytes cut short + short.bytes -= 2; + try std.testing.expect(!projectionFits(words, scales, short, four, out_hidden)); + var flat = scales; + flat.rank = 1; + try std.testing.expect(!projectionFits(words, flat, scales, four, out_hidden)); + try std.testing.expect(!projectionFits(T.of(.u32, out_hidden + 1, width / 8), scales, scales, four, out_hidden)); + try std.testing.expect(!projectionFits(T.of(.bf16, out_hidden, width / 8), scales, scales, four, out_hidden)); + const eight: Width = .{ .bits = 8, .group = 64 }; + try std.testing.expect(projectionFits(T.of(.u32, out_hidden, width / 4), scales, scales, eight, out_hidden)); + try std.testing.expect(!projectionFits(words, scales, scales, eight, out_hidden)); // 4-bit words, 8-bit entry +} diff --git a/zig/src/families/gemma4/vision_unified.zig b/zig/src/families/gemma4/vision_unified.zig new file mode 100644 index 0000000000..283745ccc7 --- /dev/null +++ b/zig/src/families/gemma4/vision_unified.zig @@ -0,0 +1,218 @@ +//! Gemma 4 12B's image embedder on Metal: 48-px patches through LayerNorm, Linear and LayerNorm, positions, projection. +const std = @import("std"); +const mtl = @import("metal"); +const ks = @import("kernel_sources"); +const frags = @import("../../core/frags.zig"); +const ck = @import("../../core/checkpoint_metal.zig"); +const lane = @import("../../core/lane_projection.zig"); +const affine_mm = @import("../../core/affine_mm.zig"); +const Mlx = @import("weights.zig").Mlx; +const Width = @import("widths.zig").Width; +const image = @import("image.zig"); +const vision = @import("vision.zig"); + +pub const width = 3840; // mm_embed_dim and output_proj_dims +const table_rows = 1120; // mm_posemb_size: each axis's rows +const ln_eps: f32 = 1e-5; // nn.LayerNorm's default +const rms_eps: f32 = 1e-6; + +pub const Ref = lane.Ref; +fn ref(t: ck.Tensor) Ref { + return .{ .buf = t.buffer, .off = t.offset }; +} + +fn never(_: []const u8, _: usize) bool { + return false; +} + +/// The checkpoint's embedder tensors, packed 256-byte aligned beside the language model's (the checkpoint frees them). +pub fn addTensors(c: *ck.Checkpoint, device: mtl.Device, files: []const [:0]u8) !void { + for (files) |path| for ([_][]const u8{ "vision_embedder.", "embed_vision." }) |p| try c.addFileSplit(device, path, "", p, never); +} + +/// Whether the checkpoint has the embedder (tools/zig/gemma4_12b_convert.py keeps it). +pub fn present(c: *const ck.Checkpoint) bool { + return c.has("vision_embedder.patch_dense.weight"); +} + +pub const Embedder = struct { + gpa: std.mem.Allocator, + lib: mtl.Library, + gemm: [2][2]mtl.Pipeline, // bf16 a b^T: [M % 64 == 0][N % 128 == 0] + norm: mtl.Pipeline, + rms: mtl.Pipeline, + rows: mtl.Pipeline, + out_hidden: usize, + ln1_w: Ref, + ln1_b: Ref, + dense_w: Ref, + dense_b: Ref, + ln2_w: Ref, + ln2_b: Ref, + table: Ref, + pos_w: Ref, + pos_b: Ref, + proj: Ref, + /// The projection quantized (an oQ pack's, at its config width) with the model's affine_mm pipelines for it. + proj_q: ?struct { w: Mlx, mm: affine_mm.Pipes } = null, + scratch: mtl.Buffer, + a: Ref, // [T, merged_width]: patch_ln1's rows + b: Ref, // [T, width]: patch_dense's, before its bias + c: Ref, // [T, width]: patch_ln2's, then the RMSNorm's + d: Ref, // [T, width]: pos_norm's + /// Host-written per image: the merged patches as bf16. + patches: mtl.Buffer, + + /// `mm` and `q`: the model's affine_mm pipelines and the config's width, for a quantized projection. + pub fn load(gpa: std.mem.Allocator, device: mtl.Device, ckpt: *const ck.Checkpoint, out_hidden: usize, mm: affine_mm.Pipes, q: Width) !*Embedder { + if (out_hidden != width) return error.UnexpectedGemmaVisionTensor; + const v = try gpa.create(Embedder); + errdefer gpa.destroy(v); + v.* = undefined; + v.proj_q = null; + v.gpa = gpa; + v.out_hidden = out_hidden; + const body = try frags.source(device, gpa, ks.gemma4_vision); + defer gpa.free(body); + const src = try std.fmt.allocPrint(gpa, "#define G4V_OUT {d}\n{s}", .{ out_hidden, body }); + defer gpa.free(src); + v.lib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + errdefer v.lib.deinit(); + var made: usize = 0; + errdefer for (0..made) |i| v.pipeline(i).deinit(); + for (0..2) |m| for (0..2) |n| { + var name: [64]u8 = undefined; + v.gemm[m][n] = try mtl.Pipeline.init(device, v.lib, try std.fmt.bufPrintSentinel(&name, "g4v_gemm_bf16_nt_{d}{d}", .{ m, n }, 0), false); + made += 1; + }; + for ([_]*mtl.Pipeline{ &v.norm, &v.rms, &v.rows }, [_][:0]const u8{ "g4u_norm", "g4v_rms", "g4v_rows" }) |p, n| { + p.* = try mtl.Pipeline.init(device, v.lib, n, false); + made += 1; + } + const t = struct { + fn get(cp: *const ck.Checkpoint, name: []const u8, count: usize) !Ref { + const x = try cp.get(name); + if (x.dtype != .bf16 or x.count() != count) { + std.log.err("gemma4 vision: {s} is not bf16[{d}]", .{ name, count }); + return error.UnexpectedGemmaVisionTensor; + } + return ref(x); + } + }; + const e = "vision_embedder."; + const mw = image.merged_width; + v.ln1_w = try t.get(ckpt, e ++ "patch_ln1.weight", mw); + v.ln1_b = try t.get(ckpt, e ++ "patch_ln1.bias", mw); + v.dense_w = try t.get(ckpt, e ++ "patch_dense.weight", width * mw); + v.dense_b = try t.get(ckpt, e ++ "patch_dense.bias", width); + v.ln2_w = try t.get(ckpt, e ++ "patch_ln2.weight", width); + v.ln2_b = try t.get(ckpt, e ++ "patch_ln2.bias", width); + v.table = try t.get(ckpt, e ++ "pos_embedding", table_rows * 2 * width); + v.pos_w = try t.get(ckpt, e ++ "pos_norm.weight", width); + v.pos_b = try t.get(ckpt, e ++ "pos_norm.bias", width); + const pp = "embed_vision.embedding_projection."; + if (ckpt.has(pp ++ "scales")) { + const s = try ckpt.get(pp ++ "scales"); + const w = try ckpt.get(pp ++ "weight"); + const b = try ckpt.get(pp ++ "biases"); + const words = @as(usize, width) * q.bits / 32; + const groups = width / q.group; + const fits = w.dtype == .u32 and w.rank == 2 and w.shape[0] == out_hidden and w.shape[1] == words and + w.count() == out_hidden * words and s.dtype == .bf16 and b.dtype == .bf16 and s.rank == 2 and b.rank == 2 and + s.shape[0] == out_hidden and s.shape[1] == groups and b.shape[0] == out_hidden and b.shape[1] == groups and + s.count() == out_hidden * groups and b.count() == s.count(); + if (!fits) { + std.log.err("gemma4 vision: {s}* is not {d}-bit in groups of {d} [{d}, {d}]", .{ pp, q.bits, q.group, out_hidden, width }); + return error.UnexpectedGemmaVisionTensor; + } + v.proj_q = .{ .w = .{ .w = ref(w), .s = ref(s), .b = ref(b), .n = @intCast(out_hidden), .k = width, .bits = q.bits }, .mm = mm }; + } else v.proj = try t.get(ckpt, pp ++ "weight", out_hidden * width); + const T: usize = image.max_budget; + const sizes = [_]usize{ T * mw * 2, T * width * 2, T * width * 2, T * width * 2 }; + var total: usize = 0; + for (sizes) |s| total += std.mem.alignForward(usize, s, 256); + v.scratch = try device.buffer(total, mtl.ResourceOptions.private); + errdefer v.scratch.deinit(); + var off: usize = 0; + for ([_]*Ref{ &v.a, &v.b, &v.c, &v.d }, sizes) |f, s| { + f.* = .{ .buf = v.scratch, .off = off }; + off += std.mem.alignForward(usize, s, 256); + } + v.patches = try device.buffer(T * mw * 2, mtl.ResourceOptions.shared); + return v; + } + + fn pipeline(v: *Embedder, i: usize) mtl.Pipeline { + return switch (i) { + 0...3 => v.gemm[i / 2][i % 2], + 4 => v.norm, + 5 => v.rms, + else => v.rows, + }; + } + + pub fn deinit(v: *Embedder) void { + for (0..7) |i| v.pipeline(i).deinit(); + v.lib.deinit(); + v.scratch.deinit(); + v.patches.deinit(); + v.gpa.destroy(v); + } + + /// Merged patches (image.prepare's .merged) to `out` [gh * gw / 9, out_hidden]; `stop` 0: patch_ln2's in c, 1: d's. + pub fn encode(v: *Embedder, e: mtl.ComputeEncoder, pixels: []const f32, gh: u32, gw: u32, out: Ref, stop: ?u32) !void { + if (gh == 0 or gw == 0 or gh % image.pool != 0 or gw % image.pool != 0) return error.GemmaImageGrid; + const T: u32 = gh * gw / (image.pool * image.pool); + if (T > image.max_budget or pixels.len != @as(usize, T) * image.merged_width) return error.GemmaImageGrid; + if (gw / image.pool > table_rows or gh / image.pool > table_rows) return error.GemmaImageGrid; + for (pixels, v.patches.slice(u16, pixels.len)) |p, *d| d.* = vision.bf16(p); + const mw: u32 = image.merged_width; + v.layerNorm(e, .{ .buf = v.patches }, v.ln1_w, v.ln1_b, v.ln1_b, v.a, T, mw, 0, 0); + v.linear(e, v.a, v.dense_w, v.b, T, width, mw); + v.layerNorm(e, v.b, v.ln2_w, v.ln2_b, v.dense_b, v.c, T, width, 1, 0); + if (stop) |s| if (s == 0) return; + v.layerNorm(e, v.c, v.pos_w, v.pos_b, v.table, v.d, T, width, 2, gw / image.pool); + if (stop) |s| if (s == 1) return; + e.setPipeline(v.rms); + for ([_]Ref{ v.d, v.d, v.c }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(extern struct { width: u32, weighted: u32, eps: f32 }{ .width = width, .weighted = 0, .eps = rms_eps }, 3); + e.dispatchGroups(mtl.Size.of(T, 1, 1), mtl.Size.of(256, 1, 1)); + if (v.proj_q) |p| return affine_mm.dense(e, p.mm, v.c, p.w, out, T); + v.linear(e, v.c, v.proj, out, T, @intCast(v.out_hidden), width); + } + + /// Rows of an encoded image over `n` embedded prompt rows (`dst`, out_hidden bf16 a row), in the prompt's encoder. + pub fn inject(v: *Embedder, e: mtl.ComputeEncoder, src: Ref, dst: Ref, n: u32) void { + e.setPipeline(v.rows); + e.setBuffer(src.buf, src.off, 0); + e.setBuffer(dst.buf, dst.off, 1); + e.dispatchThreads(mtl.Size.of(v.out_hidden / 8, n, 1), mtl.Size.of(64, 1, 1)); + } + + fn layerNorm(v: *Embedder, e: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, extra: Ref, y: Ref, rows: u32, cols: u32, add: u32, gw: u32) void { + e.setPipeline(v.norm); + for ([_]Ref{ x, w, b, extra, y }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); + e.setValue(extern struct { width: u32, add: u32, gw: u32, eps: f32 }{ .width = cols, .add = add, .gw = gw, .eps = ln_eps }, 5); + e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + /// y [M, N] = x [M, K] w^T in bf16 (w as stored, [N, K]): tfp::gemm_nax as the 31B's tower runs it. + fn linear(v: *Embedder, e: mtl.ComputeEncoder, x: Ref, w: Ref, y: Ref, M: u32, N: u32, K: u32) void { + const tn = (N + 127) / 128; + const tm = (M + 63) / 64; + e.setPipeline(v.gemm[@intFromBool(M % 64 == 0)][@intFromBool(N % 128 == 0)]); + e.setBuffer(x.buf, x.off, 0); + e.setBuffer(w.buf, w.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ M, N, K, K, K, N, tn, tm, 2, K / 256, 0, 0, 0, 0, 0 }; + for (vals, 0..) |x_, i| p[i] = @intCast(x_); + e.setBytes(std.mem.asBytes(&p), 2); + e.setBuffer(y.buf, y.off, 3); + e.dispatchThreads(mtl.Size.of(@as(usize, tn << 2) * 32, (tm + 3) / 4 * 4, 2), mtl.Size.of(32, 4, 2)); + } +}; + +test "the embedder's sizes fit its kernels" { + try std.testing.expectEqual(@as(usize, 6912), image.merged_width); + try std.testing.expect(image.merged_width % 256 == 0 and width % 256 == 0 and width % 128 == 0); +} diff --git a/zig/src/families/gemma4/weights.zig b/zig/src/families/gemma4/weights.zig new file mode 100644 index 0000000000..4aeca6b1f0 --- /dev/null +++ b/zig/src/families/gemma4/weights.zig @@ -0,0 +1,293 @@ +//! Gemma 4 dense weights: each linear stack packed once for the core lane projection, norms read in place. +const std = @import("std"); +const mtl = @import("metal"); +const ckpt = @import("../../core/checkpoint_metal.zig"); +const lane = @import("../../core/lane_projection.zig"); +const ks = @import("kernel_sources"); +const cfg = @import("config.zig"); + +const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + +/// A buffer and an offset into it. +pub const Ref = lane.Ref; + +const fast: lane.Format = .{ .bits = cfg.widths.fast.bits, .group = cfg.group }; + +/// One packed linear (members of one width along N); with `y_stride`, it writes columns col.. of a stack's rows. +pub const Linear = struct { + w: mtl.Buffer, + sb: mtl.Buffer, + n: usize, + k: usize, + format: lane.Format = fast, + col: usize = 0, + y_stride: ?usize = null, + + /// The reg kernel's width, which reads the group sums its input's kernel wrote; others sum their input themselves. + pub fn reg(x: Linear) bool { + return std.meta.eql(x.format, fast); + } + + /// `rows` of x [rows, k] through a general-unpack pipeline, which takes 16 rows an encode: y's rows y_stride apart. + pub fn encodeRows(x: Linear, p: lane.Projection, e: mtl.ComputeEncoder, in: Ref, y: Ref, rows: usize) !void { + const pitch = x.y_stride orelse x.n; + var r: usize = 0; + while (r < rows) : (r += 16) { + const n = @min(16, rows - r); + try p.encode(e, .{ .buf = in.buf, .off = in.off + r * x.k * 2 }, .{ .buf = x.w }, .{ .buf = x.sb }, .{ .buf = y.buf, .off = y.off + r * pitch * 2 }, @intCast(n)); + } + } +}; + +/// A q|k|v or gate|up stack: one Linear per run of equal widths (gilby's runs), together `n` wide. +pub const Stack = struct { + runs: [3]Linear = undefined, + count: usize = 0, + n: usize = 0, + /// gate|up in one run, its rows interleaved in 16-row blocks (gemma_geglu's mode 1); split, gate then up. + interleaved: bool = false, + + pub fn list(s: *const Stack) []const Linear { + return s.runs[0..s.count]; + } +}; + +/// A linear as core/affine_mm reads it (prompt chunks): MLX's words, scales and biases, copied to aligned storage. +pub const Mlx = struct { w: Ref, s: Ref, b: Ref, n: u32, k: u32, bits: u8 = 4 }; + +pub const Layer = struct { + /// q, k, v (k again on global layers, which have none), o, gate, up, down for prompt chunks. + mlx: [7]Mlx, + qkv: Stack, + o: Linear, + gate_up: Stack, + down: Linear, + q_norm: Ref, + k_norm: Ref, + input_norm: Ref, + post_attention: Ref, + pre_feedforward: Ref, + post_feedforward: Ref, + scalar: Ref, +}; + +/// Python's lane_qmm.split_k: K slices fixed by the weight's shape, so every window width keeps one summation order. +pub fn referenceSplit(n: usize, k: usize) usize { + const tiles = (n + 31) / 32; + var sk: usize = 1; + while (sk < 8 and tiles * sk < 1024 and (k / 64) / (sk * 2) >= 8) sk *= 2; + return sk; +} + +/// The split each shape streams fastest at on one row (tf-gemma4-bench GEMMA_FUSED), not the reference's order. +pub fn splitK(geo: cfg.Geometry, n: usize, k: usize) usize { + if (k == geo.hidden and n >= 2 * geo.intermediate) return 8; // the gate/up stack and the head: long chains at 1 + if (n == geo.hidden) return 4; // o and down: two waves of 256-thread groups at 8 + return referenceSplit(n, k); +} + +/// 4/64: the reg kernel (codes in tensor-op registers, lane_qmm's words); other widths, the general unpack. +pub fn layout(geo: cfg.Geometry, x: Linear) lane.Layout { + return .{ .n = x.n, .k = x.k, .format = x.format, .sk = splitK(geo, x.n, x.k), .precompute_sums = x.reg(), .reg = x.reg(), .y_stride = x.y_stride }; +} + +/// The same kernel with a layer's glue as its prologue (gemma4/fused.metal's TF_FUSE kinds) for one-row windows. +pub fn fusedLayout(geo: cfg.Geometry, x: Linear, fuse: u8) lane.Layout { + var l = layout(geo, x); + l.fused = .{ .source = ks.gemma4_fused, .defines = switch (fuse) { + 1 => "#define TF_FUSE 1\n", + 2 => "#define TF_FUSE 2\n", + else => "#define TF_FUSE 3\n", + }, .x_rows = if (fuse == 3) 0 else 1, .epilogue = fuse == 2, .threads = if (fuse == 3) 32 * l.sk else 256 }; + return l; +} + +pub const Weights = struct { + /// The geometry's layers first; the rest unset. + layers: [cfg.max_layers]Layer, + head: Linear, + /// The tied table's bits (gemma_embed's EB). + embed_bits: u8, + embed_w: Ref, + embed_s: Ref, + embed_b: Ref, + final_norm: Ref, + owned: std.ArrayList(mtl.Buffer) = .empty, + + pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void { + for (w.owned.items) |b| b.deinit(); + w.owned.deinit(gpa); + } +}; + +/// A linear's three tensors copied into one aligned buffer: checkpoint tensors can sit at any byte offset in a shard. +fn mlxCopy(gpa: std.mem.Allocator, device: mtl.Device, c: *const ckpt.Checkpoint, owned: *std.ArrayList(mtl.Buffer), prefix: []const u8, member: []const u8, width: cfg.widths.Width) !Mlx { + var name: [160]u8 = undefined; + const w = try c.get(try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, member })); + const k = try storedK(c, prefix, member, width); + const s = try c.get(try std.fmt.bufPrint(&name, "{s}{s}.scales", .{ prefix, member })); + const b = try c.get(try std.fmt.bufPrint(&name, "{s}{s}.biases", .{ prefix, member })); + const sb_at = std.mem.alignForward(usize, w.bytes, 256); + const bb_at = sb_at + std.mem.alignForward(usize, s.bytes, 256); + const buf = try device.buffer(bb_at + b.bytes, shared); + errdefer buf.deinit(); + const dst = buf.slice(u8, bb_at + b.bytes); + @memcpy(dst[0..w.bytes], w.host(u8)[0..w.bytes]); + @memcpy(dst[sb_at..][0..s.bytes], s.host(u8)[0..s.bytes]); + @memcpy(dst[bb_at..][0..b.bytes], b.host(u8)[0..b.bytes]); + try owned.append(gpa, buf); + return .{ .w = .{ .buf = buf }, .s = .{ .buf = buf, .off = sb_at }, .b = .{ .buf = buf, .off = bb_at }, .n = @intCast(w.shape[0]), .k = @intCast(k), .bits = width.bits }; +} + +/// A member's K at its config width (words [N, K * bits / 32], tables [N, K / group]); other storage refused. +fn storedK(c: *const ckpt.Checkpoint, prefix: []const u8, member: []const u8, width: cfg.widths.Width) !usize { + var name: [160]u8 = undefined; + const w = try c.get(try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, member })); + const s = try c.get(try std.fmt.bufPrint(&name, "{s}{s}.scales", .{ prefix, member })); + const b = try c.get(try std.fmt.bufPrint(&name, "{s}{s}.biases", .{ prefix, member })); + const ok = w.dtype == .u32 and w.rank == 2 and s.rank == 2 and b.rank == 2 and s.dtype == .bf16 and b.dtype == .bf16 and + s.shape[0] == w.shape[0] and b.shape[0] == w.shape[0] and s.shape[1] == b.shape[1] and + w.count() == w.shape[0] * w.shape[1] and s.count() == s.shape[0] * s.shape[1] and b.count() == s.count(); + const k = s.shape[1] * width.group; + if (!ok or w.shape[1] * 32 != k * width.bits) { + std.log.warn("gemma4: {s}{s} is not stored at its config width, {d}-bit in groups of {d}", .{ prefix, member, width.bits, width.group }); + return error.BadGemmaWeight; + } + return k; +} + +fn ref(c: *const ckpt.Checkpoint, name: []const u8) !Ref { + const t = try c.get(name); + return .{ .buf = t.buffer, .off = t.offset }; +} + +/// Members of one width stacked along N into the lane layout; `interleave`: two members in alternate 16-row blocks. +fn pack(gpa: std.mem.Allocator, device: mtl.Device, c: *const ckpt.Checkpoint, owned: *std.ArrayList(mtl.Buffer), geo: cfg.Geometry, prefix: []const u8, members: []const []const u8, width: cfg.widths.Width, interleave: bool) !Linear { + var name: [160]u8 = undefined; + var n: usize = 0; + var k: usize = 0; + for (members) |m| { + const mk = try storedK(c, prefix, m, width); + if (k != 0 and k != mk) return error.BadGemmaWeight; + n += (try c.get(try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, m }))).shape[0]; + k = mk; + } + const row_words = k * width.bits / 32; + const groups = k / width.group; + const words = try gpa.alloc(u32, n * row_words); + defer gpa.free(words); + const scales = try gpa.alloc(u16, n * groups); + defer gpa.free(scales); + const biases = try gpa.alloc(u16, n * groups); + defer gpa.free(biases); + if (interleave and (members.len != 2 or n % 32 != 0)) return error.BadGemmaWeight; + var row: usize = 0; + for (members, 0..) |m, mi| { + const w = (try c.get(try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, m }))).host(u32); + const s = (try c.get(try std.fmt.bufPrint(&name, "{s}{s}.scales", .{ prefix, m }))).host(u16); + const b = (try c.get(try std.fmt.bufPrint(&name, "{s}{s}.biases", .{ prefix, m }))).host(u16); + const rows = w.len / row_words; + if (interleave and rows != n / 2) return error.BadGemmaWeight; + for (0..rows) |r| { + const dst = if (interleave) 32 * (r / 16) + 16 * mi + r % 16 else row + r; + @memcpy(words[dst * row_words ..][0..row_words], w[r * row_words ..][0..row_words]); + @memcpy(scales[dst * groups ..][0..groups], s[r * groups ..][0..groups]); + @memcpy(biases[dst * groups ..][0..groups], b[r * groups ..][0..groups]); + } + row += rows; + } + return packArrays(gpa, device, owned, geo, words, scales, biases, .{ .w = undefined, .sb = undefined, .n = n, .k = k, .format = .{ .bits = width.bits, .group = width.group } }); +} + +/// A q|k|v or gate|up stack: each run of members of one width packed as one linear writing its columns of the rows. +fn stack(gpa: std.mem.Allocator, device: mtl.Device, c: *const ckpt.Checkpoint, owned: *std.ArrayList(mtl.Buffer), geo: cfg.Geometry, prefix: []const u8, members: []const []const u8, ws: []const cfg.widths.Width, glu: bool) !Stack { + var s: Stack = .{}; + var first: usize = 0; + while (first < members.len) { + var end = first + 1; + while (end < members.len and std.meta.eql(ws[end], ws[first])) end += 1; + const whole = first == 0 and end == members.len; + var x = try pack(gpa, device, c, owned, geo, prefix, members[first..end], ws[first], glu and whole); + x.col = s.n; + s.n += x.n; + s.runs[s.count] = x; + s.count += 1; + first = end; + } + s.interleaved = glu and s.count == 1; + if (s.count > 1) for (s.runs[0..s.count]) |*x| { + x.y_stride = s.n; // each run writes its columns of the stack's rows (the lane words do not depend on it) + }; + return s; +} + +/// MLX affine words, scales and biases [x.n, x.k] at x's width in the lane layout (reg order at 4/64). +pub fn packArrays(gpa: std.mem.Allocator, device: mtl.Device, owned: *std.ArrayList(mtl.Buffer), geo: cfg.Geometry, words: []const u32, scales: []const u16, biases: []const u16, x: Linear) !Linear { + const l = layout(geo, x); + const wb = try device.buffer(l.weightWords() * 4, shared); + errdefer wb.deinit(); + const sb = try device.buffer(l.metadataElements() * 2, shared); + errdefer sb.deinit(); + try lane.pack(l, words, scales, biases, wb.slice(u32, l.weightWords()), sb.slice(u16, l.metadataElements())); + if (l.reg) try lane.regOrder(device, &.{.{ .w = .{ .buf = wb }, .n = x.n, .k = x.k }}); + try owned.append(gpa, wb); + try owned.append(gpa, sb); + var out = x; + out.w = wb; + out.sb = sb; + return out; +} + +pub fn load(gpa: std.mem.Allocator, device: mtl.Device, c: *const ckpt.Checkpoint, geo: cfg.Geometry, ws: *const cfg.widths.Widths) !Weights { + var w: Weights = undefined; + w.owned = .empty; + errdefer w.deinit(gpa); + var p: [96]u8 = undefined; + var name: [160]u8 = undefined; + for (w.layers[0..geo.layers], 0..) |*l, i| { + const pre = try std.fmt.bufPrint(&p, "language_model.model.layers.{d}.", .{i}); + const lw = ws.layers[i]; + const qkv: []const []const u8 = if (cfg.global(i)) &.{ "self_attn.q_proj", "self_attn.k_proj" } else &.{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj" }; + l.qkv = try stack(gpa, device, c, &w.owned, geo, pre, qkv, lw[0..qkv.len], false); + l.o = try pack(gpa, device, c, &w.owned, geo, pre, &.{"self_attn.o_proj"}, lw[3], false); + l.gate_up = try stack(gpa, device, c, &w.owned, geo, pre, &.{ "mlp.gate_proj", "mlp.up_proj" }, lw[4..6], true); + l.down = try pack(gpa, device, c, &w.owned, geo, pre, &.{"mlp.down_proj"}, lw[6], false); + for (&l.mlx, cfg.widths.module_names, lw) |*x, member, width| { + // global layers have no v_proj: k stands in, never read + x.* = if (cfg.global(i) and std.mem.eql(u8, member, "self_attn.v_proj")) l.mlx[1] else try mlxCopy(gpa, device, c, &w.owned, pre, member, width); + } + inline for (.{ "q_norm", "k_norm", "input_norm", "post_attention", "pre_feedforward", "post_feedforward", "scalar" }, .{ "self_attn.q_norm.weight", "self_attn.k_norm.weight", "input_layernorm.weight", "post_attention_layernorm.weight", "pre_feedforward_layernorm.weight", "post_feedforward_layernorm.weight", "layer_scalar" }) |field, t| + @field(l, field) = try ref(c, try std.fmt.bufPrint(&name, "{s}{s}", .{ pre, t })); + } + w.head = try pack(gpa, device, c, &w.owned, geo, "language_model.model.", &.{"embed_tokens"}, ws.embed, false); + w.embed_bits = ws.embed.bits; + w.embed_w = try ref(c, "language_model.model.embed_tokens.weight"); + w.embed_s = try ref(c, "language_model.model.embed_tokens.scales"); + w.embed_b = try ref(c, "language_model.model.embed_tokens.biases"); + w.final_norm = try ref(c, "language_model.model.norm.weight"); + return w; +} + +/// The lane pipelines, one per distinct (n, k, width, stride, prologue). +pub const Projections = struct { + /// The model's sizes, which pick each shape's split. + geo: cfg.Geometry, + list: std.ArrayList(struct { n: usize, k: usize, format: lane.Format, y_stride: ?usize, fuse: u8, p: lane.Projection }) = .empty, + + pub fn get(ps: *Projections, gpa: std.mem.Allocator, device: mtl.Device, x: Linear) !lane.Projection { + return ps.getFused(gpa, device, x, 0); + } + + /// `fuse`: a gemma4/fused.metal prologue kind, or 0 for the plain kernel. + pub fn getFused(ps: *Projections, gpa: std.mem.Allocator, device: mtl.Device, x: Linear, fuse: u8) !lane.Projection { + for (ps.list.items) |e| if (e.n == x.n and e.k == x.k and std.meta.eql(e.format, x.format) and e.y_stride == x.y_stride and e.fuse == fuse) return e.p; + const p = try lane.Projection.init(gpa, device, if (fuse == 0) layout(ps.geo, x) else fusedLayout(ps.geo, x, fuse)); + try ps.list.append(gpa, .{ .n = x.n, .k = x.k, .format = x.format, .y_stride = x.y_stride, .fuse = fuse, .p = p }); + return p; + } + + pub fn deinit(ps: *Projections, gpa: std.mem.Allocator) void { + for (ps.list.items) |e| e.p.deinit(); + ps.list.deinit(gpa); + } +}; diff --git a/zig/src/families/gemma4/widths.zig b/zig/src/families/gemma4/widths.zig new file mode 100644 index 0000000000..d643f3a91e --- /dev/null +++ b/zig/src/families/gemma4/widths.zig @@ -0,0 +1,133 @@ +//! Gemma 4's MLX affine width per module, read once from the config's quantization table (oQ packs among them). +const std = @import("std"); +const mlx = @import("../../core/quant/mlx.zig"); +const cfg = @import("config.zig"); + +pub const Width = mlx.Width; + +/// The bits the lane projection and core affine_mm both read; the glue's group sums are 64 wide. +pub const bit_widths = [_]u8{ 2, 3, 4, 5, 6, 8 }; +pub const group = 64; +/// The width the reg lane kernel and the fused one-row prologues take. +pub const fast: Width = .{ .bits = 4, .group = group }; + +/// A layer's linears, in Layer.mlx's order. +pub const Module = enum { q, k, v, o, gate, up, down }; +pub const module_names = [_][]const u8{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj" }; + +pub const Widths = struct { + layers: [cfg.max_layers][module_names.len]Width = @splat(@splat(fast)), + /// The tied embedding and head. + embed: Width = fast, + /// The 31B tower's (and a quantized 12B's) embed_vision projection. + vision: Width = fast, + + pub fn of(w: *const Widths, layer: usize, m: Module) Width { + return w.layers[layer][@backingInt(m)]; + } + + /// Every linear at the fast width: the one-row windows may take the fused pipelines. + pub fn allFast(w: *const Widths, layers: usize) bool { + if (!std.meta.eql(w.embed, fast)) return false; + for (w.layers[0..layers]) |l| for (l) |x| if (!std.meta.eql(x, fast)) return false; + return true; + } +}; + +pub const Error = error{ UnsupportedGemmaQuantization, GemmaWidthNotRead, GemmaModuleUnquantized } || std.mem.Allocator.Error; + +/// Whether the family's kernels read `w`. +pub fn readable(w: Width) bool { + return w.group == group and std.mem.indexOfScalar(u8, &bit_widths, w.bits) != null; +} + +/// Each module's width from `quant` (`quantization` or `quantization_config`); global layers have no v_proj. +pub fn read(gpa: std.mem.Allocator, quant: std.json.ObjectMap, layers: usize) Error!Widths { + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var it = quant.iterator(); + while (it.next()) |e| { + const k = e.key_ptr.*; + if (std.mem.eql(u8, k, "bits") or std.mem.eql(u8, k, "group_size") or std.mem.eql(u8, k, "mode")) continue; + if (e.value_ptr.* != .object or e.value_ptr.object.count() == 0) { + std.log.warn("gemma4: {s} is not a width table (false or {{}} marks a module unquantized); every linear is affine", .{k}); + return error.GemmaModuleUnquantized; + } + } + const q = (mlx.detect(a, .{ .root = quant, .quantization = quant }) catch |err| return switch (err) { + error.UnsupportedQuantization => error.UnsupportedGemmaQuantization, // two spellings of one module disagree + else => |e| e, + }) orelse return error.UnsupportedGemmaQuantization; // not MLX affine + var w: Widths = .{}; + var key: [96]u8 = undefined; + for (0..layers) |i| for (module_names, 0..) |name, mi| { + if (cfg.global(i) and mi == @backingInt(Module.v)) continue; + w.layers[i][mi] = try named(q, std.fmt.bufPrint(&key, "layers.{d}.{s}", .{ i, name }) catch unreachable); + }; + w.embed = try named(q, "embed_tokens"); + w.vision = try named(q, "embed_vision.embedding_projection"); + return w; +} + +fn named(q: mlx.Config, key: []const u8) Error!Width { + const w = q.width(key) catch { + std.log.warn("gemma4: {s}'s width is not MLX affine bits 2-6 or 8", .{key}); + return error.GemmaWidthNotRead; + }; + if (!readable(w)) { + std.log.warn("gemma4: {s} is {d}-bit in groups of {d}; the kernels read bits 2-6 and 8 in groups of 64", .{ key, w.bits, w.group }); + return error.GemmaWidthNotRead; + } + return w; +} + +fn table(text: []const u8) !Widths { + const parsed = try std.json.parseFromSlice(std.json.Value, std.testing.allocator, text, .{}); + defer parsed.deinit(); + return read(std.testing.allocator, parsed.value.object, 12); +} + +test "a uniform table gives every module the global width" { + const w4 = try table("{\"bits\": 4, \"group_size\": 64, \"mode\": \"affine\"}"); + try std.testing.expect(w4.allFast(12)); + const w8 = try table("{\"bits\": 8, \"group_size\": 64}"); + try std.testing.expect(!w8.allFast(12)); + try std.testing.expectEqual(Width{ .bits = 8, .group = 64 }, w8.of(11, .down)); + try std.testing.expectEqual(Width{ .bits = 8, .group = 64 }, w8.embed); +} + +test "per-module entries under any spelling reach their module only" { + const w = try table( + \\{"bits": 4, "group_size": 64, "mode": "affine", + \\ "language_model.model.layers.1.self_attn.v_proj": {"bits": 6, "group_size": 64}, + \\ "model.layers.2.mlp.down_proj": {"bits": 6, "group_size": 64, "mode": "affine"}, + \\ "language_model.model.embed_tokens": {"bits": 8, "group_size": 64}} + ); + try std.testing.expectEqual(@as(u8, 6), w.of(1, .v).bits); + try std.testing.expectEqual(@as(u8, 4), w.of(1, .k).bits); + try std.testing.expectEqual(@as(u8, 4), w.of(0, .v).bits); + try std.testing.expectEqual(@as(u8, 6), w.of(2, .down).bits); + try std.testing.expectEqual(@as(u8, 8), w.embed.bits); + try std.testing.expect(!w.allFast(12)); +} + +test "widths no kernel reads are refused by name, never guessed" { + const cases = .{ + .{ "{\"bits\": 7, \"group_size\": 64}", error.UnsupportedGemmaQuantization }, + .{ "{\"bits\": 4, \"group_size\": 32}", error.GemmaWidthNotRead }, + .{ "{\"bits\": 4, \"group_size\": 128}", error.GemmaWidthNotRead }, + .{ "{\"bits\": 4, \"group_size\": 64, \"mode\": \"mxfp4\"}", error.UnsupportedGemmaQuantization }, + .{ "{\"bits\": 4, \"group_size\": 64, \"layers.3.mlp.up_proj\": {\"bits\": 7, \"group_size\": 64}}", error.GemmaWidthNotRead }, + .{ "{\"bits\": 4, \"group_size\": 64, \"layers.3.mlp.up_proj\": {\"bits\": 4, \"group_size\": 64, \"mode\": \"nvfp4\"}}", error.GemmaWidthNotRead }, + .{ "{\"bits\": 4, \"group_size\": 64, \"layers.3.mlp.up_proj\": false}", error.GemmaModuleUnquantized }, + .{ "{\"bits\": 4, \"group_size\": 64, \"layers.3.mlp.up_proj\": {}}", error.GemmaModuleUnquantized }, + .{ "{\"bits\": 4, \"group_size\": 64, \"layers.3.mlp.up_proj\": {\"bits\": 6, \"group_size\": 64}, \"model.layers.3.mlp.up_proj\": {\"bits\": 8, \"group_size\": 64}}", error.UnsupportedGemmaQuantization }, + }; + inline for (cases) |c| try std.testing.expectError(c[1], table(c[0])); +} + +test "a global layer's missing v_proj takes no entry" { + const w = try table("{\"bits\": 4, \"group_size\": 64, \"layers.5.self_attn.v_proj\": {\"bits\": 4, \"group_size\": 32}}"); + try std.testing.expectEqual(@as(u8, 4), w.of(5, .k).bits); +} diff --git a/zig/src/families/glm/checks.zig b/zig/src/families/glm/checks.zig index 50c0285c40..d553ce5eef 100644 --- a/zig/src/families/glm/checks.zig +++ b/zig/src/families/glm/checks.zig @@ -171,8 +171,7 @@ pub fn checkMatmul(e: *Engine, prompt: []const u32) !void { fwd.embed(&x, b.enc, e.prompt_ids, rows, 0); fwd.boundary(&x, b.enc, rows, false, L.hc.?[0], L.in_norm); fwd.qmv(&x, b.enc, e.k.qmv_kda_in, e.sc.normed, a.in_proj, out, rows); - affine_mm.rowSums(b.enc, e.k.mm_bf16, 64, e.sc.normed, pr.sums, rows, a.in_proj.k); - affine_mm.dense(b.enc, e.k.mm_bf16, e.sc.normed, pr.sums, a.in_proj, pr.proj, rows); + affine_mm.dense(b.enc, e.k.mm_bf16, e.sc.normed, a.in_proj, pr.proj, rows); try e.finish(b.cb, b.enc); const want: [*]const u16 = @ptrCast(@alignCast(out.addr())); const got: [*]const u16 = @ptrCast(@alignCast(pr.proj.addr())); diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index b325424805..92a30a86c1 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -61,7 +61,6 @@ pub const Prompt = struct { m_x: Ref, m_xn: Ref, m_out: Ref, - sums: Ref, // fp32 group sums of a matmul's rows (affine_mm's bias operand) sparse_nax: bool, // the sparse MLA attention on the tensor units (GLM_SPARSE_NAX=0: the row kernel) rows: u32, // a chunk's rows at most (the buffers' height) kda_rows: u32, // the rows a KDA block's planes hold (half a chunk) @@ -129,7 +128,6 @@ fn buffers(p: *Prompt, arena: *st.Arena, c: *const cfg.Config, decode: *const st p.offsets = try big.of(arena, c.experts * 4); p.order = try big.of(arena, n * 4); p.sorted = try big.of(arena, n * 4); - p.sums = try big.of(arena, @max(R * 16384, n * D) / 64 * 4); // the widest dense K (MLA's out-projection), or a gather's // the rest share three regions by layer; proj: KDA's projection, MLA's queries and key lists then values, the experts' planes, the dense gate and up const pitch = std.mem.alignForward(usize, c.kdaProj(), 64); // the matmul's padded pitch const qp = up(R * c.qrProj() * 2); @@ -189,8 +187,7 @@ fn run(e: mtl.ComputeEncoder, grid: [3]usize, group: [3]usize) void { /// y [M, N] = x [M, K] W^T for a 4-bit matrix (64x64 tiles; N rounded up to the weights' padded rows, y's row pitch). fn qmm(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Affine, y: Ref, M: u32) void { if (!wide.native(q)) return qmmAt(p, e, x, q.k, q, y, M); - affine_mm.rowSums(e, p.mm.mm_bf16, 64, x, p.sums, M, q.k); - affine_mm.dense(e, p.mm.mm_bf16, x, p.sums, q, y, M); + affine_mm.dense(e, p.mm.mm_bf16, x, q, y, M); } /// `qmm` over x's rows `ld` elements apart (a slice of a wider row). @@ -201,24 +198,18 @@ fn qmmAt(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, ld: u32, q: wts.Affine e, q, .{ .in = x, .ldx = ld, .out = y, .ldo = pitch, .rows = M }, - p.sums, ); - const mm = p.mm.mm_bf16; - affine_mm.rowSumsAt(e, mm, 64, x, p.sums, M, q.k, ld); - affine_mm.denseBatch(e, mm, x, p.sums, q, y, M, 1, .{ .x_row = @intCast(ld), .y_row = @intCast(pitch) }); + affine_mm.denseBatch(e, p.mm.mm_bf16, x, q, y, M, 1, .{ .x_row = @intCast(ld), .y_row = @intCast(pitch) }); } /// `qmm` with fp32 out (TP2: one Mac's partial of a row-split projection). fn qmmF32(p: *const Prompt, e: mtl.ComputeEncoder, x: Ref, q: wts.Affine, y: Ref, M: u32) void { - affine_mm.rowSums(e, p.mm.mm_bf16, 64, x, p.sums, M, q.k); - affine_mm.dense(e, p.mm.mm_f32, x, p.sums, q, y, M); + affine_mm.dense(e, p.mm.mm_f32, x, q, y, M); } /// The routed experts' gate, up and SwiGLU on the M rows of `x_in` its pairs read (in expert order) into `act` ([n, gate.n]). -fn gateUp(p: *const Prompt, mm: affine_mm.Pipes, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, M: u32, n: u32) void { - affine_mm.rowSums(e, mm, 64, x_in, p.sums, M, w.gate.k); - affine_mm.gatherGlu(e, mm, x_in, p.sums, w.gate, w.up, p.offsets, p.order, p.act, n, x.c.experts, x.c.topk, x.c.swiglu_limit); - affine_mm.rowSums(e, mm, 64, p.act, p.sums, n, w.down.k); +fn gateUp(p: *const Prompt, mm: affine_mm.Pipes, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts.Moe, x_in: Ref, n: u32) void { + affine_mm.gatherGlu(e, mm, x_in, w.gate, w.up, p.offsets, p.order, p.act, n, x.c.experts, x.c.topk, x.c.swiglu_limit); } fn swiglu(x: *const fwd.Ctx, e: mtl.ComputeEncoder, gu: Ref, act: Ref, rows: u32, width: u32) void { @@ -258,8 +249,8 @@ fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts if (s & Class.route == 0) route(p, x, e, w, x_in, M); if (ep) |t| { if (s & Class.routed == 0) { // down's fp32 partials put back in (row, slot) order for the decode's combine - gateUp(p, p.mm.mm_bf16, x, e, w, x_in, M, n); - affine_mm.gatherTo(e, p.mm.mm_f32, p.act, p.sums, w.down, p.offsets, p.yf, p.order, n, E); + gateUp(p, p.mm.mm_bf16, x, e, w, x_in, n); + affine_mm.gatherTo(e, p.mm.mm_f32, p.act, w.down, p.offsets, p.yf, p.order, n, E); } if (s & Class.exchange == 0) t.sendRows(e, p.yf, p.wts, M); if (s & Class.shared == 0) shared(p, x, e, w, x_in, M); @@ -292,8 +283,8 @@ fn moe(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *const wts ); } else if (s & Class.routed == 0) { // down's rows put back in (row, slot) order for the combine const mm = if (path == .mm4) p.mm.mm_bf16 else p.mm.wide.?.mm2; - gateUp(p, mm, x, e, w, x_in, M, n); - affine_mm.gatherTo(e, mm, p.act, p.sums, w.down, p.offsets, p.yp, p.order, n, E); + gateUp(p, mm, x, e, w, x_in, n); + affine_mm.gatherTo(e, mm, p.act, w.down, p.offsets, p.yp, p.order, n, E); } if (s & Class.combine != 0) return; e.setPipeline(k.moe_combine); @@ -398,15 +389,13 @@ fn unabsorb(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, w: *cons H, .{ .in = p.att, .ldx = H * K, .out = p.vals, .ldo = H * c.v_dim, .rows = M }, K, - p.sums, ); var q = w.kv_b; q.w = q.w.at(@as(usize, c.nope) * K / 2); q.s = q.s.at(@as(usize, c.nope) * (K / 64) * 2); q.b = q.b.at(@as(usize, c.nope) * (K / 64) * 2); q.n = c.v_dim; - affine_mm.rowSums(e, p.mm.mm_bf16, 64, p.att, p.sums, M * H, K); - affine_mm.denseBatch(e, p.mm.mm_bf16, p.att, p.sums, q, p.vals, M, H, .{ .x_row = @intCast(H * K), .y_row = @intCast(H * c.v_dim), .sums_row = @intCast(H), .x_batch = @intCast(K), .y_batch = @intCast(c.v_dim), .sums_batch = 1, .w_batch = @intCast(per) }); + affine_mm.denseBatch(e, p.mm.mm_bf16, p.att, q, p.vals, M, H, .{ .x_row = @intCast(H * K), .y_row = @intCast(H * c.v_dim), .x_batch = @intCast(K), .y_batch = @intCast(c.v_dim), .w_batch = @intCast(per) }); } /// `n` rows' key lists from their index scores over every pooled block they read (positions p0..), the scores on the tensor units. @@ -443,7 +432,6 @@ fn mla(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, mi: usize, w: c.mla_heads, .{ .in = p.qp, .ldx = c.qrProj(), .out = p.ql, .ldo = c.mla_heads * c.kv_lora, .rows = M }, c.nope, - p.sums, ); // mlx-lm's embed_q } else if (fwd.on(x, "mla_absorb")) { e.setPipeline(if (c.tp > 1) x.k.absorb_nax_tp else x.k.absorb_nax); // the row kernel's qvm, weights in three bf16 parts diff --git a/zig/src/families/glm/wide.zig b/zig/src/families/glm/wide.zig index 54e857a7b8..883a107779 100644 --- a/zig/src/families/glm/wide.zig +++ b/zig/src/families/glm/wide.zig @@ -241,19 +241,16 @@ pub fn perHead(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, heads: u32, io } /// A chunk's x slices (`kx` apart) through q's heads into out [rows, heads, n / heads]: one affine_mm batch, else rows. -pub fn headsBatch(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, h: u32, io: Io, kx: u32, sums: Ref) void { +pub fn headsBatch(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, h: u32, io: Io, kx: u32) void { const ny = q.n / h; const pipes = mm(p, q) orelse return perHead(p, e, q, h, io, kx, ny); var head = q; head.n = ny; - affine_mm.rowSumsAt(e, pipes, 64, io.in, sums, io.rows, h * kx, io.ldx); // a head's sums one row apart - affine_mm.denseBatch(e, pipes, io.in, sums, head, io.out, io.rows, h, .{ + affine_mm.denseBatch(e, pipes, io.in, head, io.out, io.rows, h, .{ .x_row = @intCast(io.ldx), .y_row = @intCast(io.ldo), - .sums_row = @intCast(h), .x_batch = @intCast(kx), .y_batch = @intCast(ny), - .sums_batch = 1, .w_batch = @intCast(ny), }); } @@ -408,8 +405,7 @@ pub fn mm(p: *const Pipes, q: Affine) ?affine_mm.Pipes { } /// A chunk's product into io.out: each member on affine_mm at its width if it fills 64-row tiles, else the row kernels. -pub fn denseMm(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, io: Io, sums: Ref) void { - var summed = false; +pub fn denseMm(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, io: Io) void { for (0..@as(usize, q.more) + 1) |i| { const r = q.run(i); const one: Affine = .{ .w = r.w, .s = r.s, .b = r.b, .n = r.n, .k = q.k, .bits = r.bits, .group = r.group }; @@ -419,13 +415,10 @@ pub fn denseMm(p: *const Pipes, e: mtl.ComputeEncoder, q: Affine, io: Io, sums: dense(p, e, one, part); continue; }; - if (!summed) affine_mm.rowSumsAt(e, pipes, 64, io.in, sums, io.rows, q.k, io.ldx); // every member's: one k, g64 - summed = true; affine_mm.denseBatch( e, pipes, io.in, - sums, one, part.out, io.rows, diff --git a/zig/src/families/lfm2_moe/forward.zig b/zig/src/families/lfm2_moe/forward.zig index 0bbfc3ef02..96cedbeedd 100644 --- a/zig/src/families/lfm2_moe/forward.zig +++ b/zig/src/families/lfm2_moe/forward.zig @@ -52,14 +52,11 @@ const Encoder = struct { if (!self.prompt) return null; return self.m.kernels.naxPipes(bits); } - /// y[rows] = x[rows] W^T: the prompt tiles (on x's group sums), or their twin at decode and window rows. + /// y[rows] = x[rows] W^T: the tensor-unit prompt GEMM, else the tiles (on x's group sums) or their twin. fn projection(self: Encoder, w: row.Weights, x: Buffer, x_off: usize, y: Buffer, y_off: usize, rows: usize) void { const t = &self.m.kernels.tiles; if (self.nax(w.bits)) |p| { - const xr = Ref{ .buf = x, .off = x_off }; - const sums = Ref{ .buf = self.s.sums }; - am.rowSums(self.e, p, c.group, xr, sums, @intCast(rows), @intCast(w.k)); - am.dense(self.e, p, xr, sums, affine(w), Ref{ .buf = y, .off = y_off }, @intCast(rows)); + am.dense(self.e, p, Ref{ .buf = x, .off = x_off }, affine(w), Ref{ .buf = y, .off = y_off }, @intCast(rows)); return; } if (self.prompt or !self.m.kernels.twin) { @@ -289,16 +286,13 @@ fn moe(enc: Encoder, s: *st.Scratch, x: wts.Moe, scale: f32, total: usize) bool enc.groups(.{ 1, 1, 1 }, .{ c.experts * 16, 1, 1 }); if (enc.nax(x.gate.one.bits)) |p| { // Sorted slot j reads token row tokens[j]; gate and up SwiGLU'd into act, then down over act's sorted rows. - const sums = Ref{ .buf = s.sums }; const starts = Ref{ .buf = s.starts }; const act = Ref{ .buf = s.act }; - am.rowSums(e, p, c.group, Ref{ .buf = s.x }, sums, @intCast(total), c.hidden); const gate = affine(x.gate.one); const up = affine(x.up.one); const tokens = Ref{ .buf = s.tokens }; - am.gatherGlu(e, p, Ref{ .buf = s.x }, sums, gate, up, starts, tokens, act, @intCast(slots), c.experts, 1, 0); - am.rowSums(e, p, c.group, act, sums, @intCast(slots), c.moe_intermediate); - am.gather(e, p, act, sums, affine(x.down.one), starts, Ref{ .buf = s.experts }, @intCast(slots), c.experts); + am.gatherGlu(e, p, Ref{ .buf = s.x }, gate, up, starts, tokens, act, @intCast(slots), c.experts, 1, 0); + am.gather(e, p, act, affine(x.down.one), starts, Ref{ .buf = s.experts }, @intCast(slots), c.experts); } else { tile.encodeWork(e, &enc.m.kernels.tiles, s.starts, c.experts, s.work); tile.encodeSums(e, t, c.hidden, s.x, 0, s.sums, total); diff --git a/zig/src/native/gemma4.zig b/zig/src/native/gemma4.zig new file mode 100644 index 0000000000..29a2eebe9d --- /dev/null +++ b/zig/src/native/gemma4.zig @@ -0,0 +1,227 @@ +//! Gemma 4 dense (31B, 12B) on the lane host: windows of up to 16 rows, several streams a forward, external drafts. +const std = @import("std"); +const mtl = @import("metal"); +const api = @import("engine_api"); +const tf = @import("tensorfold"); +const g = tf.gemma4; +const lanes = tf.lanes; +const cache_fit = @import("cache_fit.zig"); +const Resident = @import("resident.zig").Resident; + +/// The context a stream gets without --context: attention's per-call scratch grows with it (pass --context for more). +const default_context = 32768; + +/// --drafter loads z-lab's DFlash drafter (31B) or Google's MTP assistant (12B) in front of the engine. +pub const takes_drafter = true; + +const Host = struct { + gpa: std.mem.Allocator, + io: std.Io, // the image decoder's, on the HTTP thread + model: *g.Model, + metal: *g.backend.Metal, + draft: ?*g.dflash.Drafter = null, + assist: ?*g.assistant.Drafter = null, + wrap: ?lanes.drafted.Drafted = null, + warm: mtl.keepalive.Target = undefined, // the residency set's queue, for the lane host's idle ticker + resident: ?Resident = null, // weights, scratches, the drafter's and the tower's buffers, wired between requests + config: lanes.Config, + clock: lanes.backend.WallClock, + core: lanes.Engine, + host: api.LaneHost, + cache: ?api.prompt_cache.Store = null, // kept prompt states (--prompt-cache-gib 0: none) + /// The image wrapper's tokens (--vision): begin, image, end, and each image's rows between the first and last. + markers: [3]u32 = .{ 255999, 258880, 258882 }, + + fn close(ctx: *anyopaque) void { + const h: *Host = @ptrCast(@alignCast(ctx)); + h.host.stop(); + if (h.cache) |*store| store.deinit(); + h.core.deinit(); + if (h.wrap) |*w| w.deinit(); + if (h.draft) |d| d.deinit(); + if (h.assist) |d| d.deinit(); + h.config.deinit(h.gpa); + if (h.resident) |*r| r.deinit(); + h.metal.deinit(); + h.model.deinit(); + h.gpa.destroy(h); + } +}; + +pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const h = start(a, gpa, io, o, problem) catch |err| switch (err) { + error.Refused => return null, + else => return err, + }; + return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; +} + +/// The host, or error.Refused once `problem` says why; either error path releases what was made before it. +fn start(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !*Host { + const model = g.Model.load(gpa, io, o.dir) catch |err| { + problem.* = try std.fmt.allocPrint(a, "the native Gemma 4 engine cannot load this checkpoint ({s}): it serves Gemma 4 31B and 12B as MLX affine 2-6 or 8 bits in groups of 64, uniform or per module", .{@errorName(err)}); + return error.Refused; + }; + errdefer model.deinit(); + const window: i64 = o.context orelse @min(default_context, @as(i64, @intCast(model.config.context))); + if (window < 1 or window > model.config.context) { + problem.* = "--context must fit the Gemma 4 checkpoint's positive context window"; + return error.Refused; + } + if (o.vision != null) model.loadVision(io, o.dir) catch |err| switch (err) { + error.NoGemmaVision => { + const names = if (model.config.geo.unified()) "vision_embedder. and embed_vision. (tools/zig/gemma4_12b_convert.py keeps them)" else "vision_tower. and embed_vision."; + problem.* = try std.fmt.allocPrint(a, "--vision: this Gemma 4 {s} checkpoint has no image input ({s} tensors)", .{ model.config.geo.name, names }); + return error.Refused; + }, + else => return err, + }; + const h = try gpa.create(Host); + errdefer gpa.destroy(h); + h.gpa = gpa; + h.io = io; + h.model = model; + h.markers = .{ 255999, 258880, 258882 }; + h.metal = try g.backend.Metal.init(gpa, io, model, .{ .capacity = @intCast(window), .streams = @max(o.lanes, 2) }); + errdefer h.metal.deinit(); + try h.metal.measure(); // the depth rule prices drafts against the window forward's cost by rows + h.draft = null; + h.assist = null; + h.wrap = null; + h.resident = null; + const draft_dir = if (o.drafts) o.drafter else null; + const assistant = if (draft_dir) |dir| g.assistant_weights.isAssistant(a, io, dir) catch false else false; + if (draft_dir) |dir| if (assistant) { + h.assist = g.assistant.Drafter.init(gpa, io, h.metal, dir) catch |err| { + problem.* = if (err == error.AssistantForAnotherTarget) + try std.fmt.allocPrint(a, "the Gemma 4 MTP assistant at {s} was trained on another target: its backbone width and key heads do not fit Gemma 4 {s} (google/gemma-4-12B-it-assistant drafts for the 12B only)", .{ dir, model.config.geo.name }) + else + try std.fmt.allocPrint(a, "the native Gemma 4 engine cannot load the MTP assistant at {s} ({s}): it takes google/gemma-4-12B-it-assistant as released (bf16, no ordered embeddings)", .{ dir, @errorName(err) }); + return error.Refused; + }; + } else { + h.draft = g.dflash.Drafter.init(gpa, io, model, dir, .{ .capacity = @intCast(window) }) catch |err| { + problem.* = if (err == error.DraftForAnotherTarget) + try std.fmt.allocPrint(a, "the native Gemma 4 drafter cannot load {s}: its width or taps do not fit Gemma 4 {s} (DFlash drafts take z-lab/gemma-4-31B-it-DFlash on the 31B only)", .{ dir, model.config.geo.name }) + else + try std.fmt.allocPrint(a, "the native Gemma 4 drafter cannot load {s} ({s}): it takes z-lab/gemma-4-31B-it-DFlash", .{ dir, @errorName(err) }); + return error.Refused; + }; + }; + errdefer if (h.draft) |d| d.deinit(); + errdefer if (h.assist) |d| d.deinit(); + const drafter: ?lanes.drafter.Drafter = if (h.draft) |d| d.drafter() else if (h.assist) |d| d.drafter() else null; + if (drafter) |d| h.wrap = try lanes.drafted.Drafted.init(gpa, h.metal.backend(), d); + errdefer if (h.wrap) |*w| w.deinit(); + const facts = if (h.wrap) |*w| w.facts(h.metal.facts()) else h.metal.facts(); + h.config = try lanes.Config.init(gpa, facts, g.state.max_rows, g.state.max_rows - 1); + errdefer h.config.deinit(gpa); + h.clock = .{ .io = io }; + h.core = lanes.Engine.init(gpa, &h.config, if (h.wrap) |*w| w.backend() else h.metal.backend(), h.clock.clock()); + errdefer h.core.deinit(); + const vis: ?api.Vision = if (o.vision) |v| .{ .ctx = h, .prepare = prepareImage, .image_token = h.markers[1], .markers = &h.markers, .wrap = .{ h.markers[0], h.markers[2] }, .part_type = "image", .image_tokens = v.image_tokens, .max_images = v.max_images, .explain = g.image.decoderProblem } else null; + h.host = api.LaneHost.init(gpa, io, &h.core, .{ .lanes = o.lanes, .context_window = @intCast(window), .prefill_step = g.state.chunk_rows, .structures = true, .vision = vis }); + h.cache = null; + try enableCache(h, a, o, @intCast(window), problem); + errdefer if (h.cache) |*store| store.deinit(); + h.warm = .{ .queue = model.queue }; + holdResident(h) catch |err| std.log.warn("gemma4: no residency set ({s}); the first request after an idle second re-wires the weights", .{@errorName(err)}); + errdefer if (h.resident) |*r| r.deinit(); + h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; + try h.host.start(); // the host's lists are empty until its thread runs, so a failed start leaves nothing to free + std.log.info("gemma4: Gemma 4 {s}, context {d} tokens, {d} streams{s}{s}", .{ model.config.geo.name, window, @max(o.lanes, 2), if (h.draft != null) ", DFlash drafts" else if (h.assist != null) ", MTP assistant drafts" else "", if (vis != null) ", images" else "" }); + return h; +} + +/// The buffers that outlive a request, in one set the idle keepalive holds wired (a stream's caches come and go). +fn holdResident(h: *Host) !void { + var list: std.ArrayList(mtl.Buffer) = .empty; + defer list.deinit(h.gpa); + const m = h.model; + for (m.checkpoint.shards.items) |sh| try list.append(h.gpa, sh.buffer); // the weights, the tower's among them + try list.appendSlice(h.gpa, m.weights.owned.items); + try list.appendSlice(h.gpa, &.{ m.inv_sliding, m.inv_global }); + try buffersOf(h.gpa, &list, &h.metal.scratch); + try buffersOf(h.gpa, &list, &h.metal.prompt_scratch); + if (h.draft) |d| { + try list.appendSlice(h.gpa, d.w.owned.items); + try buffersOf(h.gpa, &list, d); + } + if (h.assist) |d| { + try list.appendSlice(h.gpa, d.w.owned.items); + try buffersOf(h.gpa, &list, d); + } + if (m.vision) |v| try v.appendBuffers(h.gpa, &list); + h.resident = try Resident.hold(h.gpa, h.io, m.device, m.queue, list.items, "gemma4") orelse return; + h.warm = (&h.resident.?).target(); +} + +/// `x`'s buffer fields, plain, optional or in arrays: a scratch's or the drafter's working buffers. +fn buffersOf(gpa: std.mem.Allocator, list: *std.ArrayList(mtl.Buffer), x: anytype) !void { + const info = @typeInfo(@TypeOf(x.*)).@"struct"; + inline for (info.field_names, info.field_types) |name, T| { + const v = @field(x.*, name); + switch (@typeInfo(T)) { + .array => |a| if (a.child == mtl.Buffer) try list.appendSlice(gpa, &v), + .optional => |o| if (o.child == mtl.Buffer) if (v) |b| try list.append(gpa, b), + else => if (T == mtl.Buffer) try list.append(gpa, v), + } + } +} + +/// The prompt cache's copies of a stream: the target's caches at a mark (backend.zig), and the drafter's context there. +const Snaps = struct { + const M = g.backend.Metal; + fn host(ptr: *anyopaque) *Host { + return @ptrCast(@alignCast(ptr)); + } + fn bytes(ptr: *anyopaque, at: u32) u64 { + const h = host(ptr); + return M.snapBytes(h.metal, at) + if (h.draft) |d| d.keptBytes(at) else 0; + } + fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const h = host(ptr); + const saved = try M.snapSave(h.metal, owner, at); + if (h.draft) |d| d.keepLane(@ptrCast(@alignCast(owner.?)), @ptrCast(@alignCast(saved))) catch |err| { + M.snapDrop(h.metal, saved); + return err; + }; + return saved; + } + fn restore(_: *anyopaque, _: ?*anyopaque, _: *anyopaque) anyerror!void { + return error.BackendRestores; // the prompt pass restores, before its first chunk + } + fn drop(ptr: *anyopaque, saved: *anyopaque) void { + M.snapDrop(host(ptr).metal, saved); + } +}; + +/// Prompt reuse in what 70% of RAM leaves beside the streams' caches (16 GiB by default), at planned chunk starts. +fn enableCache(h: *Host, a: std.mem.Allocator, o: api.Open, window: usize, problem: *[]const u8) !void { + const later = @as(u64, @max(o.lanes, 2)) * g.state.cacheBytes(h.model.config.geo, window); + const room = h.model.device.maxWorkingSet() -| h.model.device.allocated() -| (8 << 30) -| later; + var why: []const u8 = ""; + const plan = cache_fit.budgetBeside(o.prompt_cache_gib, o.prompt_cache_over_cap, room, later, a, &why, "") catch |err| switch (err) { + error.CacheOverCap => { + problem.* = why; + return error.Refused; + }, + else => return err, + }; + if (plan.budget_bytes == 0) return; + // a prompt chunk's rows are its own arithmetic: keep and resume only where the request's chunks start + const rules: api.prompt_cache.Rules = .{ .planned = true, .grid = g.state.chunk_rows, .min_prompt = 0 }; + h.cache = api.prompt_cache.Store.init(h.gpa, .{ .ptr = h, .vtable = &.{ .bytes = Snaps.bytes, .save = Snaps.save, .restore = Snaps.restore, .drop = Snaps.drop } }, rules, plan.budget_bytes); + h.host.cache = &h.cache.?; + h.host.info_.prompt_cache = true; + h.host.info_.prompt_cache_plan = plan; +} + +/// Info.vision's prepare (on the HTTP thread); the template writes <|image|> alone: the server adds boi and eoi. +fn prepareImage(ctx: *anyopaque, a: std.mem.Allocator, bytes: []const u8, max_tokens: u32) anyerror!api.PreparedImage { + const h: *Host = @ptrCast(@alignCast(ctx)); + const p = try g.image.prepare(h.io, a, bytes, max_tokens, h.model.vision.?.layout()); + return .{ .pixels = p.pixels, .gh = p.gh, .gw = p.gw, .tokens = p.tokens, .hash = p.hash }; +} diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index 0fdb9cc5f8..7c2fe17ba5 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -27,6 +27,9 @@ const affine_widths = blk: { break :blk &done; }; +/// Gemma 4's: MLX affine 2-6 or 8 bits in groups of 64, uniform or per module (oQ packs). +const gemma4_widths: []const []const u8 = &.{ "mlx-q2g64", "mlx-q3g64", "mlx-q4g64", "mlx-q5g64", "mlx-q6g64", "mlx-q8g64" }; + pub const families: []const api.Family = &.{ .{ .model_type = "nemotron_h", .formats = affine_widths }, .{ .model_type = "qwen4_exp", .formats = affine_widths }, @@ -34,8 +37,15 @@ pub const families: []const api.Family = &.{ .{ .model_type = "qwen3_5", .formats = affine_widths }, .{ .model_type = "lfm2_moe", .formats = &.{ "mlx-q4g64", "mlx-q6g64", "mlx-q8g64" } }, .{ .model_type = "glm_moe_dsa", .formats = &.{"exlr8"} }, + .{ .model_type = "gemma4", .formats = gemma4_widths }, + .{ .model_type = "gemma4_unified", .formats = gemma4_widths }, }; +/// Gemma 4 31B (gemma4) and 12B (gemma4_unified, its text model): one native family, gemma4.zig's host. +fn isGemma4(model_type: []const u8) bool { + return std.mem.eql(u8, model_type, "gemma4") or std.mem.eql(u8, model_type, "gemma4_unified"); +} + /// The chip class gate entries name ("apple-m5" for an Apple M5 Max); null without an Apple GPU. pub fn chip(a: Allocator) ?[]const u8 { const device = mtl.Device.init() catch return null; @@ -133,15 +143,16 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c problem.* = why; return null; } - if (o.drafter != null and !dense) { - problem.* = "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head)"; + if (o.drafter != null and !dense and !isGemma4(o.model_type)) { + problem.* = "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head) and Gemma 4 dense"; return null; } if (dense) return openQwen27(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "qwen4_exp")) return openFlashNext(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "glm5_next")) return openGlm(a, gpa, io, o, problem); + if (isGemma4(o.model_type)) return @import("gemma4.zig").open(a, gpa, io, o, problem); if (o.vision != null) { - problem.* = try std.fmt.allocPrint(a, "--vision: the native engine reads images for GLM-5.3-Flash and Qwen3.8-27B only, not {s} checkpoints", .{o.model_type}); + problem.* = try std.fmt.allocPrint(a, "--vision: the native engine reads images for GLM-5.3-Flash, Qwen3.8-27B and Gemma 4 only, not {s} checkpoints", .{o.model_type}); return null; } if (std.mem.eql(u8, o.model_type, "glm_moe_dsa")) return openGlmDsa(a, gpa, io, o, problem); @@ -272,13 +283,13 @@ const recall_refusal = "--recall drafts from a record in the lane engine's round fn refusal(a: Allocator, o: api.Open, dense: bool) !?[]const u8 { const t = o.model_type; if (o.slide and !std.mem.eql(u8, t, "nemotron_h")) return try std.fmt.allocPrint(a, "--slide: {s} checkpoints do not learn on Macs yet", .{t}); - const one_mac: ?[]const u8 = if (std.mem.eql(u8, t, "nemotron_h")) "Nemotron" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (!std.mem.eql(u8, t, "qwen3_5")) null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; + const one_mac: ?[]const u8 = if (std.mem.eql(u8, t, "nemotron_h")) "Nemotron" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (!std.mem.eql(u8, t, "qwen3_5")) null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; if (o.speed_up != null) if (one_mac) |name| return try std.fmt.allocPrint(a, "--speed-up joins two Macs for Flash Next, GLM-5.3-Flash and GLM-5.3 only: serve {s} without it", .{name}); if (o.recall and std.mem.eql(u8, t, "qwen4_exp")) return recall_refusal; if (o.learn == null) return null; // the 27B writes conversation states if (o.learn_conversations and std.mem.eql(u8, t, "qwen3_5") and dense) return null; - const name: []const u8 = if (std.mem.eql(u8, t, "glm_moe_dsa")) "GLM-5.3 (EXLR8)" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (!std.mem.eql(u8, t, "qwen3_5")) return null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; + const name: []const u8 = if (std.mem.eql(u8, t, "glm_moe_dsa")) "GLM-5.3 (EXLR8)" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (!std.mem.eql(u8, t, "qwen3_5")) return null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B"; if (o.learn_conversations) return try std.fmt.allocPrint(a, "--learn-conversations keeps conversation states on disk for Qwen3.8-27B, Nemotron, Flash Next and GLM-5.3-Flash: serve {s} without it", .{name}); const or_talks: []const u8 = if (dense) " (Qwen3.8-27B keeps conversation states on disk with --learn-conversations)" else ""; return try std.fmt.allocPrint(a, "--learn learns prompt prefixes for Nemotron, Flash Next and GLM-5.3-Flash only{s}: serve {s} without --learn, --learn-dir or --learn-gib", .{ or_talks, name }); diff --git a/zig/src/server/image_key.zig b/zig/src/server/image_key.zig index 1df9476ff5..6444a26381 100644 --- a/zig/src/server/image_key.zig +++ b/zig/src/server/image_key.zig @@ -4,6 +4,7 @@ const api = @import("engine_api"); const json = @import("json"); const errors = @import("errors.zig"); const images = @import("images.zig"); +const messages = @import("messages.zig"); /// The marker before an image's rows: its identity's low 31 bits, exactly (core/picture/identity.zig keeps 62). pub fn open(hash: u64) u32 { @@ -85,3 +86,28 @@ test "one-token images whose row words collide share no kept state; the same ima try std.testing.expect(store.find(other, &.{}, &.{}) == null); try std.testing.expectEqual(@as(?u32, 4), if (store.find(try keyOf(a, ids[0]), &.{}, &.{})) |e| e.at else null); } + +test "a template that reads \"image\" parts gets them so typed, twice normalized, and the server's wrap around the rows" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var dummy: u8 = 0; + const v: api.Vision = .{ .ctx = &dummy, .prepare = named, .image_token = 9, .markers = &.{ 5, 9, 6 }, .wrap = .{ 5, 6 }, .part_type = "image", .image_tokens = 64, .max_images = 4 }; + const text = + \\[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:,%FF%D8%FF%01%02%03%04%05%06%07%08"}}, {"type": "text", "text": "x"}]}] + ; + const msgs = switch (try json.parseText(a, text)) { + .ok => |m| m, + .err => return error.TestUnexpectedResult, + }; + var cx: errors.Cx = .{ .a = a }; + const once = try messages.normalizeWith(&cx, msgs, "system", false, .of(v)); + const twice = try messages.normalizeWith(&cx, once, "system", false, .of(v)); + const part = twice.array[0].get("content").?.array[0]; + try std.testing.expectEqualStrings("image", part.get("type").?.string); + try std.testing.expect(part.get("image_url") != null); + // the template's placeholder alone (1, 9, 2): the server writes 5 and 6 around the image's row + const x = (try images.expand(v, null, images.MarkText.text(), &cx, twice, &[_]u32{ 1, 9, 2 }, 3)).?; + try std.testing.expectEqualSlices(u32, &.{ 1, 5, 9, 6, 2 }, x.ids); + try std.testing.expectEqual(@as(u32, 2), x.images[0].at); +} diff --git a/zig/src/server/images.zig b/zig/src/server/images.zig index 141424fdc7..ba35c0b699 100644 --- a/zig/src/server/images.zig +++ b/zig/src/server/images.zig @@ -26,7 +26,7 @@ pub const Expanded = struct { const Part = struct { video: bool = false, url: []const u8, low: bool = false }; /// The image and video parts of normalized `messages`, in the order the template meets them. -fn parts(cx: *Cx, messages: Value) errors.Refused![]const Part { +fn parts(cx: *Cx, messages: Value, image_type: []const u8) errors.Refused![]const Part { var out: std.ArrayList(Part) = .empty; if (messages != .array) return out.items; for (messages.array) |m| { @@ -35,7 +35,7 @@ fn parts(cx: *Cx, messages: Value) errors.Refused![]const Part { for (content.array) |p| { const kind = p.get("type") orelse continue; if (kind != .string) continue; - if (std.mem.eql(u8, kind.string, "image_url")) { + if (std.mem.eql(u8, kind.string, "image_url") or std.mem.eql(u8, kind.string, image_type)) { const iu = p.get("image_url").?; const detail = iu.get("detail"); try out.append(cx.a, .{ .url = iu.get("url").?.string, .low = detail != null and std.mem.eql(u8, detail.?.string, "low") }); @@ -222,7 +222,7 @@ pub fn encodeShielded(text: model_text.Text, a: std.mem.Allocator, rendered: []c /// The prompt with image placeholders expanded to tokens and video ones to frame groups, its images and cache ids. pub fn expand(vision: ?api.Vision, fetch: ?Fetch, text: model_text.Text, cx: *Cx, messages: Value, ids: []const u32, history_len: usize) errors.Refused!?Expanded { const v = vision orelse return null; - const ps = try parts(cx, messages); + const ps = try parts(cx, messages, v.part_type); var n_images: usize = 0; for (ps) |p| n_images += @intFromBool(!p.video); const n_videos = ps.len - n_images; diff --git a/zig/src/server/messages.zig b/zig/src/server/messages.zig index dd03af2a31..1c39134f8c 100644 --- a/zig/src/server/messages.zig +++ b/zig/src/server/messages.zig @@ -31,9 +31,10 @@ pub fn normalize(cx: *Cx, messages: ?Value, late_system: []const u8, needs_user_ pub const Media = struct { images: bool = false, videos: bool = false, + image_type: []const u8 = "image_url", // the part type the template reads (Vision.part_type) pub fn of(v: ?api.Vision) Media { - return .{ .images = v != null, .videos = v != null and v.?.video != null }; + return .{ .images = v != null, .videos = v != null and v.?.video != null, .image_type = if (v) |x| x.part_type else "image_url" }; } }; @@ -51,8 +52,8 @@ fn videoPart(cx: *Cx, part: Value) errors.Refused!Value { return .{ .object = out }; } -/// One image part as the template and the image path read it: {"type": "image_url", "image_url": {"url", "detail"?}}. -fn imagePart(cx: *Cx, part: Value) errors.Refused!Value { +/// One image part as the template and the image path read it: {"type": `kind`, "image_url": {"url", "detail"?}}. +fn imagePart(cx: *Cx, part: Value, kind: []const u8) errors.Refused!Value { const iu = part.get("image_url") orelse return cx.refuse("an image_url content part must contain image_url"); var url: []const u8 = ""; var detail: ?Value = null; @@ -70,7 +71,7 @@ fn imagePart(cx: *Cx, part: Value) errors.Refused!Value { try inner.put(cx.a, "url", .{ .string = url }); if (detail) |d| try inner.put(cx.a, "detail", d); const out = try json.newObject(cx.a); - try out.put(cx.a, "type", .{ .string = "image_url" }); + try out.put(cx.a, "type", .{ .string = kind }); try out.put(cx.a, "image_url", .{ .object = inner }); return .{ .object = out }; } @@ -95,10 +96,11 @@ pub fn normalizeWith(cx: *Cx, messages: ?Value, late_system: []const u8, needs_u var has_image = false; for (content.?.array) |part| { const kind = if (part == .object and part.get("type") != null and part.get("type").? == .string) part.get("type").?.string else ""; - if (images and std.mem.eql(u8, kind, "image_url")) { + // an image part, or one this normalization already wrote with the template's type (it runs twice) + if (images and (std.mem.eql(u8, kind, "image_url") or std.mem.eql(u8, kind, media.image_type))) { // tool results carry screenshots too; the template renders them in its response if (!std.mem.eql(u8, role.?.string, "user") and !std.mem.eql(u8, role.?.string, "tool")) return cx.refuse("image parts are accepted in user and tool messages only"); - try parts.append(cx.a, try imagePart(cx, part)); + try parts.append(cx.a, try imagePart(cx, part, media.image_type)); has_image = true; continue; } diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 61db37bcec..fb9073974f 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -36,6 +36,7 @@ pub const flashnext_learned = @import("families/flashnext/learned.zig"); pub const flashnext_think_check = @import("families/flashnext/think_check.zig"); pub const glm = @import("families/glm/glm.zig"); pub const qwen35 = @import("families/qwen3_5/qwen3_5.zig"); +pub const gemma4 = @import("families/gemma4/gemma4.zig"); pub const lfm2_moe = @import("families/lfm2_moe/lfm2_moe.zig"); pub const flashnext_pack = @import("families/flashnext/pack.zig"); @@ -64,6 +65,7 @@ test { _ = @import("families/flashnext/sample.zig"); // a sampled window's draws on the GPU (Metal) _ = @import("families/glm/glm.zig"); _ = @import("families/qwen3_5/config.zig"); + _ = @import("families/gemma4/gemma4.zig"); _ = @import("families/qwen3_5/backend.zig"); _ = @import("families/lfm2_moe/config.zig"); _ = @import("families/lfm2_moe/backend.zig"); diff --git a/zig/tests/affine_ref.zig b/zig/tests/affine_ref.zig index 4b0f7900dc..ca71d67544 100644 --- a/zig/tests/affine_ref.zig +++ b/zig/tests/affine_ref.zig @@ -20,8 +20,9 @@ fn fromBf(v: u16) f64 { return @as(f32, @bitCast(@as(u32, v) << 16)); } -/// Synthetic MLX affine weights [n, k] (codes little-endian in bytes) and `wide` rows of x; `deq` is s q + b exactly. +/// Synthetic MLX affine weights [n, k] (codes little-endian in bytes) and `wide` rows of x; `deq` s q + b exactly. const Synth = struct { + deqr: []f64, // bf16(fma(s, q, b)): the weights as affine_mm's tile dequantizes them n: usize, k: usize, bits: u8, @@ -36,7 +37,7 @@ const Synth = struct { var prng = std.Random.DefaultPrng.init(seed); const r = prng.random(); const sn = n * (k / group); - const self = Synth{ .n = n, .k = k, .bits = bits, .group = group, .x = try device.buffer(wide * k * 2, opts), .w = try device.buffer(n * k * bits / 8, opts), .s = try device.buffer(sn * 2, opts), .b = try device.buffer(sn * 2, opts), .deq = try a.alloc(f64, n * k) }; + const self = Synth{ .n = n, .k = k, .bits = bits, .group = group, .x = try device.buffer(wide * k * 2, opts), .w = try device.buffer(n * k * bits / 8, opts), .s = try device.buffer(sn * 2, opts), .b = try device.buffer(sn * 2, opts), .deq = try a.alloc(f64, n * k), .deqr = try a.alloc(f64, n * k) }; for (self.x.slice(u16, wide * k)) |*v| v.* = toBf(r.float(f32) * 2 - 1); r.bytes(self.w.slice(u8, n * k * bits / 8)); for (self.s.slice(u16, sn), self.b.slice(u16, sn)) |*s, *b| { @@ -45,7 +46,10 @@ const Synth = struct { } for (0..n) |j| for (0..k) |i| { const g = j * (k / group) + i / group; - self.deq[j * k + i] = fromBf(self.s.slice(u16, sn)[g]) * @as(f64, @floatFromInt(self.code(j, i))) + fromBf(self.b.slice(u16, sn)[g]); + const s32: f32 = @floatCast(fromBf(self.s.slice(u16, sn)[g])); + const b32: f32 = @floatCast(fromBf(self.b.slice(u16, sn)[g])); + self.deq[j * k + i] = @as(f64, s32) * @as(f64, @floatFromInt(self.code(j, i))) + @as(f64, b32); + self.deqr[j * k + i] = fromBf(toBf(@mulAdd(f32, s32, @floatFromInt(self.code(j, i)), b32))); }; return self; } @@ -53,6 +57,7 @@ const Synth = struct { fn deinit(self: Synth) void { for ([_]mtl.Buffer{ self.x, self.w, self.s, self.b }) |x| x.deinit(); a.free(self.deq); + a.free(self.deqr); } fn code(self: Synth, j: usize, i: usize) u32 { @@ -65,25 +70,20 @@ const Synth = struct { return (v >> @intCast(at % 8)) & ((@as(u32, 1) << @intCast(self.bits)) - 1); } - /// y = x W^T in float64, and the sizes the fp32 sums carry when codes enter as q + off (affine_mm's 4-bit: 128). - fn reference(self: Synth, off: f64) !Ref { + /// y = x W^T in float64 on deq or (`rounded`) deqr, the sizes its fp32 sums carry, and the scales' share. + fn reference(self: Synth, rounded: bool) !Ref { const n = self.n; const k = self.k; const groups = k / self.group; const ref = Ref{ .want = try a.alloc(f64, wide * n), .mag = try a.alloc(f64, wide * n), .scaled = try a.alloc(f64, wide * n) }; const x = self.x.slice(u16, wide * k); const sc = self.s.slice(u16, n * groups); - const bi = self.b.slice(u16, n * groups); const xd = try a.alloc(f64, k); defer a.free(xd); - const xg = try a.alloc(f64, groups); // each group's sum of |x| - defer a.free(xg); + const bi = self.b.slice(u16, n * groups); + const w = if (rounded) self.deqr else self.deq; for (0..wide) |r| { - @memset(xg, 0); - for (xd, 0..) |*v, i| { - v.* = fromBf(x[r * k + i]); - xg[i / self.group] += @abs(v.*); - } + for (xd, 0..) |*v, i| v.* = fromBf(x[r * k + i]); for (0..n) |j| { var sum: f64 = 0; var mag: f64 = 0; @@ -91,14 +91,17 @@ const Synth = struct { for (0..groups) |g| { var qx: f64 = 0; var qxs: f64 = 0; + var xg: f64 = 0; for (g * self.group..(g + 1) * self.group) |i| { const q: f64 = @floatFromInt(self.code(j, i)); - sum += xd[i] * self.deq[j * k + i]; - qx += (q + off) * @abs(xd[i]); + sum += xd[i] * w[j * k + i]; + if (rounded) mag += @abs(xd[i] * w[j * k + i]); + qx += q * @abs(xd[i]); qxs += q * xd[i]; + xg += @abs(xd[i]); } const s = fromBf(sc[j * groups + g]); - mag += s * qx + @abs(fromBf(bi[j * groups + g]) - off * s) * xg[g]; + if (!rounded) mag += s * qx + @abs(fromBf(bi[j * groups + g])) * xg; scaled += s * qxs; } ref.want[r * n + j] = sum; @@ -135,14 +138,14 @@ const Out = struct { } }; -/// The bound: fp32 sums of a group's products, its row sum and a chain of groups, plus one bf16 ulp of the value. -fn bound(want: f64, mag: f64, s: *const Synth, f32_out: bool, extra: usize) f64 { - const terms: f64 = @floatFromInt(s.group + s.k / s.group + 8 + extra); +/// The bound: fp32 sums over 16-deep tensor-unit slices and their chain, or a group and a chain of groups; a bf16 ulp. +fn bound(want: f64, mag: f64, s: *const Synth, f32_out: bool, extra: usize, tensor: bool) f64 { + const terms: f64 = @floatFromInt((if (tensor) 16 + s.k / 16 else s.group + s.k / s.group) + 8 + extra); return terms * unit * mag + (if (f32_out) 0 else @abs(want) * 0x1p-7) + 1e-30; } /// `rows` rows of `out` against the float64 reference: the worst column's max abs and relative error, and the bound. -fn compare(t: *Tally, name: []const u8, out: Out, rows: usize, ref: Ref, s: *const Synth, extra: usize) void { +fn compare(t: *Tally, name: []const u8, out: Out, rows: usize, ref: Ref, s: *const Synth, extra: usize, tensor: bool) void { const n = s.n; var col: usize = 0; var col_abs: f64 = -1; @@ -156,7 +159,7 @@ fn compare(t: *Tally, name: []const u8, out: Out, rows: usize, ref: Ref, s: *con for (0..rows) |r| { const i = r * n + j; const d = @abs(out.get(i) - ref.want[i]); - const tol = bound(ref.want[i], ref.mag[i], s, out.f32_out, extra); + const tol = bound(ref.want[i], ref.mag[i], s, out.f32_out, extra, tensor); const q = d / tol; ctl = @max(ctl, @abs(out.get(i) - ref.want[i] - 0.01 * ref.scaled[i]) / tol); if (!(d <= abs)) abs = d; @@ -216,17 +219,15 @@ fn commit(cb: anytype, e: mtl.ComputeEncoder) !void { const Kernel = enum { dense, gather }; -/// affine_mm over rows [start, start + m): its row sums, then the dense product or a gather with one expert. -fn affineWindow(e: mtl.ComputeEncoder, p: mm.Pipes, s: *const Synth, sums: mtl.Buffer, offs: mtl.Buffer, y: Out, kernel: Kernel, start: usize, m: usize) void { +/// affine_mm over rows [start, start + m): the dense product or a gather with one expert. +fn affineWindow(e: mtl.ComputeEncoder, p: mm.Pipes, s: *const Synth, offs: mtl.Buffer, y: Out, kernel: Kernel, start: usize, m: usize) void { const x = .{ .buf = s.x, .off = start * s.k * 2 }; - const xs = .{ .buf = sums, .off = start * (s.k / s.group) * 4 }; const yo = .{ .buf = y.buf, .off = start * s.n * y.esize() }; const z: usize = 0; const q = .{ .w = .{ .buf = s.w, .off = z }, .s = .{ .buf = s.s, .off = z }, .b = .{ .buf = s.b, .off = z }, .n = @as(u32, @intCast(s.n)), .k = @as(u32, @intCast(s.k)) }; - mm.rowSums(e, p, @intCast(s.group), x, xs, @intCast(m), @intCast(s.k)); switch (kernel) { - .dense => mm.dense(e, p, x, xs, q, yo, @intCast(m)), - .gather => mm.gather(e, p, x, xs, q, .{ .buf = offs, .off = z }, yo, @intCast(m), 1), + .dense => mm.dense(e, p, x, q, yo, @intCast(m)), + .gather => mm.gather(e, p, x, q, .{ .buf = offs, .off = z }, yo, @intCast(m), 1), } } @@ -239,8 +240,6 @@ fn affine(t: *Tally, device: mtl.Device, queue: mtl.Queue, s: *const Synth, ref: var p: mm.Pipes = undefined; for (mm.names, 0..) |name, i| p[i] = try mtl.Pipeline.init(device, lib, name, false); defer for (p) |x| x.deinit(); - const sums = try device.buffer(wide * (s.k / s.group) * 4, opts); - defer sums.deinit(); const offs = try device.buffer(4, opts); defer offs.deinit(); offs.slice(i32, 1)[0] = 0; @@ -251,9 +250,9 @@ fn affine(t: *Tally, device: mtl.Device, queue: mtl.Queue, s: *const Synth, ref: defer all.buf.deinit(); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - affineWindow(e, p, s, sums, offs, all, kernel, 0, wide); + affineWindow(e, p, s, offs, all, kernel, 0, wide); try commit(cb, e); - compare(t, label, all, wide, ref, s, 0); + compare(t, label, all, wide, ref, s, 0, true); var same = true; for (windows) |m| { const got = try newOut(device, s, f32_out); @@ -261,7 +260,7 @@ fn affine(t: *Tally, device: mtl.Device, queue: mtl.Queue, s: *const Synth, ref: const wb = queue.commandBuffer(); const we = wb.compute(.serial); var start: usize = 0; - while (start + m <= wide) : (start += m) affineWindow(we, p, s, sums, offs, got, kernel, start, m); + while (start + m <= wide) : (start += m) affineWindow(we, p, s, offs, got, kernel, start, m); try commit(wb, we); same = sameRows(t, label, m, all, got, wide / m * m, s.n) and same; } @@ -286,7 +285,7 @@ fn simd(t: *Tally, device: mtl.Device, queue: mtl.Queue, p: *const sq.Pipelines, try simdWindow(e, p, s, all, false, 0, wide); try commit(cb, e); const extra = std.math.log2_int(u32, sq.splits(s.n)) + 1; // the chunks' fixed tree - compare(t, label, all, wide, ref, s, extra); + compare(t, label, all, wide, ref, s, extra, false); var same = true; for (windows) |m| { const got = try newOut(device, s, false); @@ -310,7 +309,7 @@ fn simd(t: *Tally, device: mtl.Device, queue: mtl.Queue, p: *const sq.Pipelines, var start: usize = 0; while (start + m <= wide) : (start += m) try simdWindow(we, p, s, got, true, start, m); try commit(wb, we); - compare(t, twin, got, wide / m * m, ref, s, 0); + compare(t, twin, got, wide / m * m, ref, s, 0, false); const rows = wide / m * m; const eq = std.mem.eql(u8, all.buf.slice(u8, rows * s.n * 2), got.buf.slice(u8, rows * s.n * 2)); std.debug.print("RESULT {s}: {s} the tiles' bits; twinExact says {s}{s}\n", .{ twin, if (eq) "equals" else "differs from", if (exact) "equal" else "differ", if (exact and !eq) " FAIL" else "" }); @@ -332,14 +331,14 @@ pub fn main(_: std.process.Init) !void { seed += 1; const s = try Synth.init(device, shape[0], shape[1], bits, group, seed); defer s.deinit(); - const ref = try s.reference(if (bits == 4) 128 else 0); + const ref = try s.reference(true); defer ref.deinit(); for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); }; for ([_][2]usize{ .{ 192, 1024 }, .{ 256, 4096 } }) |shape| for ([_]usize{ 32, 64 }) |group| { // 2 bits: oQ2's experts const s = try Synth.init(device, shape[0], shape[1], 2, group, 0x2716 + shape[0] + group); defer s.deinit(); - const ref = try s.reference(0); + const ref = try s.reference(true); defer ref.deinit(); for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); }; @@ -356,7 +355,7 @@ pub fn main(_: std.process.Init) !void { 0x3716 + 100 * @as(usize, bits) + shape[0] + group, ); defer s.deinit(); - const ref = try s.reference(0); + const ref = try s.reference(true); defer ref.deinit(); for ([_]bool{ false, true }) |f32_out| try affine(&t, device, queue, &s, ref, f32_out); }; @@ -366,7 +365,7 @@ pub fn main(_: std.process.Init) !void { seed += 1; const s = try Synth.init(device, shape[0], shape[1], 4, 64, seed); defer s.deinit(); - const ref = try s.reference(0); + const ref = try s.reference(false); defer ref.deinit(); try simd(&t, device, queue, &p, &s, ref); } diff --git a/zig/tests/contract/metal_affine.zig b/zig/tests/contract/metal_affine.zig index 844a20f440..b98994179b 100644 --- a/zig/tests/contract/metal_affine.zig +++ b/zig/tests/contract/metal_affine.zig @@ -24,9 +24,9 @@ fn rowDepth(s: c.Shape) usize { return 16 + 3 * ((s.k + 511) / 512) + 8; } -/// A group's products in the tile op, its scale and bias, the groups in turn (tf-affine-ref's bound). -fn groupDepth64(s: c.Shape) usize { - return 64 + s.k / 32 + 8; +/// Sixteen products a tensor step, the steps over all of K in order (tf-affine-ref's bound for affine_mm). +fn tensorDepth(s: c.Shape) usize { + return 16 + s.k / 16 + 8; } /// A group's 64 products in eight MMA steps, its sum, the groups a chunk takes, then the chunks' fixed tree. @@ -39,15 +39,12 @@ fn laneDepth(s: c.Shape) usize { return 128 + 2 * (s.k / 32) + 32; } -/// affine_mm's 4-bit codes enter as q + 128 and its bias as b - 128 s, so an infinite input meets opposite signs. -const offset_inf = c.Known{ .gate = .nonfinite, .finding = "an infinite input gives NaN where MLX's s*sum(q x) + b*sum(x) keeps the infinity (the +128 code offset moves the bias to b - 128 s)" }; - pub fn spec(v: Variant, name: []const u8) c.Spec { const out: c.Elem = if (v.f32_out) .f32 else .bf16; const step: usize = @max(16, v.group); return switch (v.kind) { .row => .{ .name = name, .accum = .f32, .depth = rowDepth, .out = out, .tiles = .{ .{ .size = 2 }, .{ .size = 8, .step = 8, .min = 8 }, .{ .size = 512, .step = step, .min = step } }, .served = &.{ .{ .rows = 16, .n = 2688, .k = 4096 }, .{ .rows = 3, .n = 4608, .k = 2688 } } }, - .affine_mm => .{ .name = name, .accum = .f32, .depth = groupDepth64, .out = out, .known = if (v.bits == 4) offset_inf else null, .tiles = .{ .{ .size = 64 }, .{ .size = 64, .step = 64, .min = 64 }, .{ .size = 128, .step = 64, .min = 64 } }, .served = &.{ .{ .rows = 128, .n = 2048, .k = 4096 }, .{ .rows = 65, .n = 1024, .k = 1024 } } }, + .affine_mm => .{ .name = name, .accum = .f32, .depth = tensorDepth, .out = out, .tiles = .{ .{ .size = 64 }, .{ .size = 64, .step = 64, .min = 64 }, .{ .size = 128, .step = 64, .min = 64 } }, .served = &.{ .{ .rows = 128, .n = 2048, .k = 4096 }, .{ .rows = 65, .n = 1024, .k = 1024 } } }, .simd, .simd_bits => .{ .name = name, .accum = .f32, .depth = simdDepth, .out = out, .tiles = .{ .{ .size = 8 }, .{ .size = 32, .step = 8, .min = 8 }, .{ .size = 512, .step = 64, .min = 64 } }, .served = &.{ .{ .rows = 16, .n = 2688, .k = 4096 }, .{ .rows = 9, .n = 6400, .k = 256 } } }, .lane, .lane_reg => .{ .name = name, .accum = .f32, .depth = laneDepth, .out = out, .replaces = if (v.kind == .lane_reg) "the cooperative lane kernel" else null, .tiles = .{ .{ .size = 16 }, .{ .size = 32, .step = 16, .min = 32 }, .{ .size = 4 * v.group, .step = v.group, .min = v.group } }, .served = &.{ .{ .rows = 16, .n = 2048, .k = 2048 }, .{ .rows = 5, .n = 512, .k = 896 } } }, }; @@ -70,7 +67,6 @@ pub const Proj = struct { w: rig_.Ref = undefined, s: rig_.Ref = undefined, b: rig_.Ref = undefined, - sums: rig_.Ref = undefined, tiled: rig_.Ref = undefined, base: rig_.Ref = undefined, out: rig_.Ref = undefined, @@ -142,7 +138,7 @@ pub const Proj = struct { p.drop(); if (!try p.takes(shape)) return false; var f = try Affine.init(p.a, shape, p.v.bits, p.v.group); - if (p.v.kind == .affine_mm and p.v.bits == 4) f.off = 128; + if (p.v.kind == .affine_mm) f.form = .rounded; p.f = f; try f.fill(fill, seed); const words = try p.a.alloc(u32, shape.n * shape.k * p.v.bits / 32); @@ -155,7 +151,6 @@ pub const Proj = struct { p.w = try p.rig.input(std.mem.sliceAsBytes(words)); p.s = try p.rig.input(std.mem.sliceAsBytes(f.scales)); p.b = try p.rig.input(std.mem.sliceAsBytes(f.biases)); - if (p.v.kind == .affine_mm) p.sums = try p.rig.scratch(shape.rows * (shape.k / p.v.group) * 4); }, .lane, .lane_reg => { const l = p.layout(shape, false); @@ -207,10 +202,8 @@ pub const Proj = struct { }, .affine_mm => { const x = .{ .buf = p.x.buf, .off = p.x.off }; - const xs = .{ .buf = p.sums.buf, .off = p.sums.off }; const q = .{ .w = .{ .buf = p.w.buf, .off = p.w.off }, .s = .{ .buf = p.s.buf, .off = p.s.off }, .b = .{ .buf = p.b.buf, .off = p.b.off }, .n = @as(u32, @intCast(shape.n)), .k = @as(u32, @intCast(shape.k)) }; - mm.rowSums(e, p.mm_p.?, @intCast(p.v.group), x, xs, @intCast(shape.rows), @intCast(shape.k)); - mm.dense(e, p.mm_p.?, x, xs, q, .{ .buf = p.out.buf, .off = p.out.off }, @intCast(shape.rows)); + mm.dense(e, p.mm_p.?, x, q, .{ .buf = p.out.buf, .off = p.out.off }, @intCast(shape.rows)); }, .lane, .lane_reg => try p.lane_p.?.encode(e, ref(p.x), ref(p.w), ref(p.s), ref(p.out), @intCast(shape.rows)), } diff --git a/zig/tests/contract/metal_repro.zig b/zig/tests/contract/metal_repro.zig index d432d9cec0..9c1b8173cd 100644 --- a/zig/tests/contract/metal_repro.zig +++ b/zig/tests/contract/metal_repro.zig @@ -23,7 +23,6 @@ pub fn affineInf(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !vo const wr = try r.input(std.mem.sliceAsBytes(&words)); const sr = try r.input(std.mem.sliceAsBytes(&s)); const br = try r.input(std.mem.sliceAsBytes(&s)); - const sums = try r.scratch(4); const y_mm = try r.output(2 * n); const y_row = try r.output(2 * n); const src = try mm.source(device, a, .{ .bits = 4, .group = 64 }); @@ -37,8 +36,7 @@ pub fn affineInf(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !vo defer rp.deinit(); const pass = r.begin(); const q = .{ .w = .{ .buf = wr.buf, .off = wr.off }, .s = .{ .buf = sr.buf, .off = sr.off }, .b = .{ .buf = br.buf, .off = br.off }, .n = @as(u32, n), .k = @as(u32, k) }; - mm.rowSums(pass.e, pipes, 64, .{ .buf = xr.buf, .off = xr.off }, .{ .buf = sums.buf, .off = sums.off }, 1, k); - mm.dense(pass.e, pipes, .{ .buf = xr.buf, .off = xr.off }, .{ .buf = sums.buf, .off = sums.off }, q, .{ .buf = y_mm.buf, .off = y_mm.off }, 1); + mm.dense(pass.e, pipes, .{ .buf = xr.buf, .off = xr.off }, q, .{ .buf = y_mm.buf, .off = y_mm.off }, 1); const cl = try row.call(&rp, .{ .w = wr.buf, .w_off = wr.off, .scales = sr.buf, .s_off = sr.off, .biases = br.buf, .b_off = br.off, .n = n, .k = k }, 1, false); pass.e.setPipeline(cl.pipeline); for ([_]rig_.Ref{ xr, wr, sr, br }, 0..) |ref, i| pass.e.setBuffer(ref.buf, ref.off, i); @@ -50,7 +48,7 @@ pub fn affineInf(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue) !vo const inf = std.math.inf(f64); std.debug.print("REPRO affine_mm q4 g64, x = [+Inf, 0 x 63], codes 1, scale = bias = {d} (bf16 0x{x:0>4})\n", .{ sv, s[0] }); std.debug.print(" float64, MLX's s*sum(q x) + b*sum(x): {d}; dequantized first, x (s q + b): {d}\n", .{ sv * 1 * inf + sv * inf, inf * (sv + sv) }); - std.debug.print(" float64, the kernel's s*sum((q + 128) x) + (b - 128 s)*sum(x): {d}\n", .{sv * 129 * inf + (sv - 128 * sv) * inf}); + std.debug.print(" float64, the kernel's sum(bf16(s q + b) x): {d}\n", .{inf * c.fromBf16(c.toBf16(@floatCast(sv + sv)))}); for ([_]struct { []const u8, rig_.Ref }{ .{ "affine_mm (tf_affine_mm)", y_mm }, .{ "row projection (tf_row_projection_q4_f32)", y_row } }) |o| { const got = r.bytes(o[1], 2 * n); const h0 = std.mem.readInt(u16, got[0..2], .little); diff --git a/zig/tests/gemma4/bench.zig b/zig/tests/gemma4/bench.zig new file mode 100644 index 0000000000..7f23748472 --- /dev/null +++ b/zig/tests/gemma4/bench.zig @@ -0,0 +1,377 @@ +//! Gemma 4 dense step costs: every layer's projections and head alone by rows, against whole forwards. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); + +const g4 = tf.gemma4; +const cfg = g4.config; +const lane = tf.lane_projection; + +fn gpuMs(cb: mtl.CommandBuffer) f64 { + return cb.gpuSeconds() * 1e3; +} + +/// Every projection of the layers and the head at `rows` rows, chained in one command buffer: GPU ms. +var given_sums: ?mtl.Buffer = null; + +fn projections(m: *g4.Model, rows: usize, x: mtl.Buffer, y: mtl.Buffer, reg: bool) !f64 { + const cb = m.queue.commandBuffer(); + const e = cb.compute(.serial); + for (m.weights.layers[0..m.config.geo.layers]) |l| for ([_]g4.weights.Linear{ l.qkv.runs[0], l.o, l.gate_up.runs[0], l.down }) |lin| { + const p = try m.projections.get(m.gpa, m.device, lin); + if (reg) try p.encodeSums(e, .{ .buf = x }, .{ .buf = lin.w }, .{ .buf = lin.sb }, .{ .ref = .{ .buf = given_sums.? }, .stride = @intCast(rows) }, .{ .buf = y }, @intCast(rows)) else try p.encode(e, .{ .buf = x }, .{ .buf = lin.w }, .{ .buf = lin.sb }, .{ .buf = y }, @intCast(rows)); + }; + const head = try m.projections.get(m.gpa, m.device, m.weights.head); + try head.encode(e, .{ .buf = x }, .{ .buf = m.weights.head.w }, .{ .buf = m.weights.head.sb }, .{ .buf = y }, @intCast(rows)); + e.end(); + cb.commit(); + cb.wait(); + return gpuMs(cb); +} + +const Q = struct { w: lane.Ref, s: lane.Ref, b: lane.Ref, n: u32, k: u32 }; + +fn member(m: *g4.Model, name: []const u8) !Q { + var buf: [3][160]u8 = undefined; + const w = try m.checkpoint.get(try std.fmt.bufPrint(&buf[0], "{s}.weight", .{name})); + const sc = try m.checkpoint.get(try std.fmt.bufPrint(&buf[1], "{s}.scales", .{name})); + const bi = try m.checkpoint.get(try std.fmt.bufPrint(&buf[2], "{s}.biases", .{name})); + return .{ .w = .{ .buf = w.buffer, .off = w.offset }, .s = .{ .buf = sc.buffer, .off = sc.offset }, .b = .{ .buf = bi.buffer, .off = bi.offset }, .n = @intCast(w.shape[0]), .k = @intCast(w.shape[1] * 8) }; +} + +/// A 512-row prompt chunk's layers at a cache depth, by stage: GPU ms for all, the GEMMs, the attention, the glue. +fn promptStages(m: *g4.Model) !void { + const cap: usize = 131072 + 1024; + var ps = try g4.prompt.Scratch.init(m.device, m.config.geo, cap); + defer ps.deinit(); + var c = try g4.state.Cache.init(m.device, m.config.geo, cap); + defer c.deinit(); + const names = [_][]const u8{ "all", "GEMMs", "attention", "glue" }; + const masks = [_]g4.prompt.Stages{ .{}, .{ .attn = false, .glue = false }, .{ .mm = false, .glue = false }, .{ .mm = false, .attn = false } }; + const one: ?usize = if (std.c.getenv("GEMMA_STAGES")) |v| std.fmt.parseInt(usize, std.mem.span(v), 10) catch null else null; + for ([_]usize{ 0, 8192, 32768, 131072 }) |depth| { + if (one) |d| if (d != depth) continue; + for (ps.pos.slice(i32, 512), 0..) |*p, i| p.* = @intCast(depth + i); + for (names, masks) |name, st| { + var best: f64 = 1e9; + for (0..3) |_| { + c.len = depth; + const cb = m.queue.commandBuffer(); + const e = cb.compute(.serial); + for (0..m.config.geo.layers) |i| try g4.prompt.layer(m, e, &ps, &c, i, i % 2, 512, st); + e.end(); + cb.commit(); + cb.wait(); + best = @min(best, gpuMs(cb)); + } + std.debug.print("prompt chunk of 512 rows at depth {d}, {s}: {d:.1} ms GPU\n", .{ depth, name, best }); + } + } +} + +/// A decode window's layers at a cache depth, by stage: GPU ms for all, the projections, the attention, the glue. +fn decodeStages(m: *g4.Model) !void { + const cap: usize = 32768 + 1024; + var s = try g4.state.Scratch.init(m.device, m.config.geo, cap); + defer s.deinit(); + var c = try g4.state.Cache.init(m.device, m.config.geo, cap); + defer c.deinit(); + const names = [_][]const u8{ "all", "projections", "attention", "glue", "no attention" }; + const masks = [_]g4.forward.Stages{ .{}, .{ .attn = false, .glue = false }, .{ .mm = false, .glue = false }, .{ .mm = false, .attn = false }, .{ .attn = false } }; + const fuse = std.c.getenv("GEMMA_NO_FUSE") == null; + const rows: usize = if (std.c.getenv("GEMMA_ROWS")) |v| std.fmt.parseInt(usize, std.mem.span(v), 10) catch 1 else 1; + var ids: [128]u32 = undefined; + for (&ids, 0..) |*t, i| t.* = @intCast(1000 + i); + for ([_]usize{ 0, 8192, 32768 }) |depth| { + for (names, masks) |name, st| { + var best: f64 = 1e9; + for (0..3) |_| { + c.len = depth; + const seg = [_]g4.forward.Segment{.{ .cache = &c, .rows = rows }}; + var chunks: g4.forward.Chunks = undefined; + g4.forward.prepare(&s, &seg, ids[0..rows], &chunks); + const cb = m.queue.commandBuffer(); + const e = cb.compute(.serial); + for (0..m.config.geo.layers) |i| try g4.forward.layer(m, e, &s, &seg, &chunks, i, i % 2, rows, st, fuse and rows == 1 and st.mm and st.glue); + e.end(); + cb.commit(); + cb.wait(); + best = @min(best, gpuMs(cb)); + } + std.debug.print("decode window of {d} rows at depth {d}, {s}{s}: {d:.2} ms GPU\n", .{ rows, depth, name, if (fuse) " (fused)" else "", best }); + } + } +} + +/// Each decode projection kind at one row over every layer's weights, plain and fused: GPU microseconds a call. +fn fusedCosts(m: *g4.Model) !void { + const geo = m.config.geo; + var s = try g4.state.Scratch.init(m.device, geo, 4096); + defer s.deinit(); + const l = &m.weights.layers[1]; + const names = [_][]const u8{ "qkv", "gate_up", "down", "o" }; + const kinds = [_]u8{ 1, 2, 3, 0 }; + for (names, kinds) |name, kind| { + for ([_]bool{ false, true }) |fused| { + if (kind == 0 and fused) continue; + var best: f64 = 1e9; + for (0..6) |_| { + const cb = m.queue.commandBuffer(); + const e = cb.compute(.serial); + for (m.weights.layers[0..geo.layers]) |*ll| { + const lin = if (kind == 1) ll.qkv.runs[0] else if (kind == 2) ll.gate_up.runs[0] else if (kind == 3) ll.down else ll.o; + if (fused) { + const p = try m.projections.getFused(m.gpa, m.device, lin, kind); + const x = if (kind == 3) s.act else s.y; + try p.encodeFused(e, .{ .buf = x }, .{ .buf = lin.w }, .{ .buf = lin.sb }, .{ .buf = if (kind == 2) s.act else s.qkv }, .{ .a = .{ .{ .buf = s.hn }, .{ .buf = s.y }, l.post_feedforward, l.scalar, l.input_norm, l.input_norm }, .o = .{ .{ .buf = s.h[0] }, .{ .buf = s.h[1] } }, .f = .{ cfg.eps, 0, 0, 0 } }, 1); + } else { + const p = try m.projections.get(m.gpa, m.device, lin); + try p.encodeSums(e, .{ .buf = s.act }, .{ .buf = lin.w }, .{ .buf = lin.sb }, .{ .ref = .{ .buf = s.sums }, .stride = 1 }, .{ .buf = s.qkv }, 1); + } + } + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| std.debug.print("GPU failure: {s}\n", .{msg}); + best = @min(best, cb.gpuSeconds() * 1e6 / @as(f64, @floatFromInt(geo.layers))); + } + std.debug.print("{s} {s}: {d:.1} us a call\n", .{ name, if (fused) "fused" else "plain", best }); + } + } + var best: f64 = 1e9; + for (0..6) |_| { + const cb = m.queue.commandBuffer(); + const e = cb.compute(.serial); + const p = try m.projections.get(m.gpa, m.device, m.weights.head); + for (0..8) |_| try p.encodeSums(e, .{ .buf = s.act }, .{ .buf = m.weights.head.w }, .{ .buf = m.weights.head.sb }, .{ .ref = .{ .buf = s.sums }, .stride = 1 }, .{ .buf = s.logits }, 1); + e.end(); + cb.commit(); + cb.wait(); + best = @min(best, cb.gpuSeconds() * 1e6 / 8); + } + std.debug.print("head plain: {d:.1} us a call\n", .{best}); +} + +/// Every projection of the layers on core/affine_mm (prompt GEMMs from the checkpoint's own tensors): GPU ms. +fn affineAll(m: *g4.Model, pipes: tf.affine_mm.Pipes, rows: u32, x: mtl.Buffer, y: mtl.Buffer) !f64 { + const cb = m.queue.commandBuffer(); + const e = cb.compute(.serial); + var name: [128]u8 = undefined; + for (0..m.config.geo.layers) |i| { + const members = if (cfg.global(i)) &[_][]const u8{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.o_proj", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj" } else &[_][]const u8{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj" }; + for (members) |mm| { + const q = try member(m, try std.fmt.bufPrint(&name, "language_model.model.layers.{d}.{s}", .{ i, mm })); + tf.affine_mm.dense(e, pipes, lane.Ref{ .buf = x }, q, lane.Ref{ .buf = y }, rows); + } + } + const head = try member(m, "language_model.model.embed_tokens"); + tf.affine_mm.dense(e, pipes, lane.Ref{ .buf = x }, head, lane.Ref{ .buf = y }, 1); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| std.debug.print("GPU failure: {s}\n", .{msg}); + return gpuMs(cb); +} + +/// Each glue kernel at one row, a layer count back to back in one command buffer: GPU microseconds a call. +fn glueCosts(m: *g4.Model) !void { + const geo = m.config.geo; + var s = try g4.state.Scratch.init(m.device, geo, 4096); + defer s.deinit(); + var c = try g4.state.Cache.init(m.device, geo, 4096); + defer c.deinit(); + const l = &m.weights.layers[0]; + const kinds = [_][]const u8{ "attn_tail", "dense_tail", "geglu", "qkv_prep", "rms", "empty" }; + for (kinds) |name| { + const cb = m.queue.commandBuffer(); + const e = cb.compute(.serial); + const kd = &m.kernels.sliding; + for (0..geo.layers) |_| { + if (std.mem.eql(u8, name, "attn_tail")) { + e.setPipeline(kd.attn_tail); + for ([_]mtl.Buffer{ s.h[0], s.o }, 0..) |b, i| e.setBuffer(b, 0, i); + e.setBuffer(l.post_attention.buf, l.post_attention.off, 2); + e.setBuffer(l.pre_feedforward.buf, l.pre_feedforward.off, 3); + e.setValue(cfg.eps, 4); + for ([_]mtl.Buffer{ s.hn, s.n_mlp, s.sums }, 5..) |b, i| e.setBuffer(b, 0, i); + e.setValue(@as(i32, 1), 8); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 256 }); + } else if (std.mem.eql(u8, name, "dense_tail")) { + e.setPipeline(kd.dense_tail); + for ([_]mtl.Buffer{ s.hn, s.y }, 0..) |b, i| e.setBuffer(b, 0, i); + e.setBuffer(l.post_feedforward.buf, l.post_feedforward.off, 2); + e.setBuffer(l.scalar.buf, l.scalar.off, 3); + e.setBuffer(l.input_norm.buf, l.input_norm.off, 4); + e.setValue(cfg.eps, 5); + for ([_]mtl.Buffer{ s.h[1], s.normed[1], s.sums }, 6..) |b, i| e.setBuffer(b, 0, i); + e.setValue(@as(i32, 1), 9); + e.setBuffer(s.sums, 0, 10); + e.setValue(@as(i32, 0), 11); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 256 }); + } else if (std.mem.eql(u8, name, "geglu")) { + e.setPipeline(kd.geglu); + for ([_]mtl.Buffer{ s.gate_up, s.act }, 0..) |b, i| e.setBuffer(b, 0, i); + e.setValue(@as(u32, @intCast(geo.intermediate)), 2); + e.setBuffer(s.sums, 0, 3); + e.setValue(@as(i32, 1), 4); + e.setValue(@as(i32, 1), 5); + e.dispatchGroups(.{ .width = geo.intermediate / 64, .height = 1 }, .{ .width = 64 }); + } else if (std.mem.eql(u8, name, "qkv_prep")) { + e.setPipeline(kd.qkv); + e.setBuffer(s.qkv, 0, 0); + e.setBuffer(l.q_norm.buf, l.q_norm.off, 1); + e.setBuffer(l.k_norm.buf, l.k_norm.off, 2); + for ([_]mtl.Buffer{ m.inv_sliding, s.pos }, 3..) |b, i| e.setBuffer(b, 0, i); + e.setValue(cfg.eps, 5); + for ([_]mtl.Buffer{ s.q, s.k, s.v, c.k[0], c.v[0] }, 6..) |b, i| e.setBuffer(b, 0, i); + e.setValue([3]i32{ 0, 0, 0 }, 11); + e.dispatchGroups(.{ .width = geo.query_heads + 2 * geo.sliding_kv_heads, .height = 1 }, .{ .width = 32 }); + } else if (std.mem.eql(u8, name, "rms")) { + e.setPipeline(kd.rms); + for ([_]mtl.Buffer{ s.h[0], s.normed[0] }, 0..) |b, i| e.setBuffer(b, 0, if (i == 0) 0 else 2); + e.setBuffer(l.input_norm.buf, l.input_norm.off, 1); + e.setValue(cfg.eps, 3); + e.setBuffer(s.sums, 0, 4); + e.setValue(@as(i32, 1), 5); + e.dispatchGroups(.{ .width = 1 }, .{ .width = 1024 }); + } else { + e.setPipeline(kd.softcap); + e.setBuffer(s.logits, 0, 0); + e.setValue(cfg.softcap, 1); + e.dispatchThreads(.{ .width = 32 }, .{ .width = 32 }); + } + } + e.end(); + cb.commit(); + cb.wait(); + std.debug.print("{s}: {d:.1} us a call\n", .{ name, cb.gpuSeconds() * 1e6 / @as(f64, @floatFromInt(geo.layers)) }); + } +} + +pub fn main(init: std.process.Init) !void { + const gpa = init.gpa; + const io = init.io; + var args = try std.process.Args.Iterator.initAllocator(init.minimal.args, gpa); + defer args.deinit(); + _ = args.next(); + const model_dir = args.next() orelse return error.Usage; + const m = try g4.Model.load(gpa, io, model_dir); + defer m.deinit(); + if (!m.config.widths.allFast(m.config.geo.layers)) { + std.debug.print("tf-gemma4-bench measures the 4-bit groups-of-64 kernels; this checkpoint has other widths\n", .{}); + return error.Usage; + } + m.fuse = std.c.getenv("GEMMA_NO_FUSE") == null; + const geo = m.config.geo; + var bytes: usize = 0; + for (m.weights.layers[0..geo.layers]) |l| for ([_]g4.weights.Linear{ l.qkv.runs[0], l.o, l.gate_up.runs[0], l.down }) |lin| { + bytes += lin.n * lin.k / 2 + lin.n * lin.k / 64 * 4; + }; + bytes += m.weights.head.n * m.weights.head.k / 2 + m.weights.head.n * m.weights.head.k / 64 * 4; + const x = try m.device.buffer(128 * geo.intermediate * 2, mtl.ResourceOptions.private); + defer x.deinit(); + const y = try m.device.buffer(128 * cfg.vocab * 2, mtl.ResourceOptions.private); + defer y.deinit(); + std.debug.print("weights read a step: {d:.2} GB\n", .{@as(f64, @floatFromInt(bytes)) / 1e9}); + if (std.c.getenv("GEMMA_GLUE") != null) return glueCosts(m); + if (std.c.getenv("GEMMA_STAGES") != null) return promptStages(m); + if (std.c.getenv("GEMMA_DSTAGES") != null) return decodeStages(m); + if (std.c.getenv("GEMMA_FUSED") != null) return fusedCosts(m); + for ([_]usize{ 1, 16, 64, 128 }) |rows| { + var best: f64 = 1e9; + for (0..5) |_| best = @min(best, try projections(m, rows, x, y, false)); + std.debug.print("projections, {d} rows: {d:.2} ms GPU ({d:.0} GB/s)\n", .{ rows, best, @as(f64, @floatFromInt(bytes)) / best / 1e6 }); + if (rows <= 16) { + const sb = try m.device.buffer(16 * (geo.intermediate / 64) * 4, mtl.ResourceOptions.private); + defer sb.deinit(); + given_sums = sb; + var best2: f64 = 1e9; + for (0..5) |_| best2 = @min(best2, try projections(m, rows, x, y, true)); + std.debug.print(" with sums given (no sums dispatches): {d:.2} ms GPU\n", .{best2}); + } + } + { + const text = try tf.affine_mm.source(m.device, gpa, .{}); + defer gpa.free(text); + const lib = try mtl.Library.fromSource(m.device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var pipes: tf.affine_mm.Pipes = undefined; + for (&pipes, tf.affine_mm.names) |*p, nm| p.* = try mtl.Pipeline.init(m.device, lib, nm, false); + const big_x = try m.device.buffer(512 * geo.intermediate * 2, mtl.ResourceOptions.private); + defer big_x.deinit(); + const big_y = try m.device.buffer(512 * 2 * geo.intermediate * 2, mtl.ResourceOptions.private); + defer big_y.deinit(); + for ([_]u32{ 128, 256, 512 }) |rows| { + var best: f64 = 1e9; + for (0..3) |_| best = @min(best, try affineAll(m, pipes, rows, big_x, big_y)); + std.debug.print("affine_mm projections, {d} rows: {d:.2} ms GPU ({d:.0} tok/s)\n", .{ rows, best, @as(f64, @floatFromInt(rows)) * 1e3 / best }); + } + } + { + var ss = try g4.state.Scratch.init(m.device, geo, 4096); + defer ss.deinit(); + var ps = try g4.prompt.Scratch.init(m.device, geo, 4096); + defer ps.deinit(); + var tok: [40]u32 = undefined; + for (&tok, 0..) |*t, i| t.* = @intCast(2000 + i * 37); + { + var ca = try g4.state.Cache.init(m.device, geo, 4096); + defer ca.deinit(); + var cbk = try g4.state.Cache.init(m.device, geo, 4096); + defer cbk.deinit(); + const step: usize = if (std.c.getenv("GEMMA_STEP")) |v| std.fmt.parseInt(usize, std.mem.span(v), 10) catch 40 else 40; + var at: usize = 0; + while (at < tok.len) : (at += step) try g4.prompt.run(m, &ps, &ca, tok[at..@min(at + step, tok.len)], &.{}); + try g4.forward.run(m, &ss, &.{.{ .cache = &cbk, .rows = 16 }}, tok[0..16]); + try g4.forward.run(m, &ss, &.{.{ .cache = &cbk, .rows = 16 }}, tok[16..32]); + try g4.forward.run(m, &ss, &.{.{ .cache = &cbk, .rows = 8 }}, tok[32..40]); + const a = ps.logits.slice(u16, cfg.vocab); + const b = ss.logits.slice(u16, 16 * cfg.vocab)[7 * cfg.vocab ..][0..cfg.vocab]; + var worst: f32 = 0; + for (a, b) |x0, y0| worst = @max(worst, @abs(@as(f32, @bitCast(@as(u32, x0) << 16)) - @as(f32, @bitCast(@as(u32, y0) << 16)))); + std.debug.print("prompt path vs decode path, 40 tokens, last-row logits: worst |diff| {d}\n", .{worst}); + } + { + // one-row windows on the fused pipelines against the plain kernels and glue: every logit the same bits + var fused = try g4.state.Cache.init(m.device, geo, 4096); + defer fused.deinit(); + var plain = try g4.state.Cache.init(m.device, geo, 4096); + defer plain.deinit(); + var differ: usize = 0; + for (tok[0..12], 0..) |t, i| { + m.fuse = true; + try g4.forward.run(m, &ss, &.{.{ .cache = &fused, .rows = 1 }}, &.{t}); + const a = try gpa.dupe(u16, ss.logits.slice(u16, cfg.vocab)); + defer gpa.free(a); + m.fuse = false; + try g4.forward.run(m, &ss, &.{.{ .cache = &plain, .rows = 1 }}, &.{t}); + if (!std.mem.eql(u16, a, ss.logits.slice(u16, cfg.vocab))) { + differ += 1; + std.debug.print(" token {d}: fused logits differ from plain\n", .{i}); + } + } + m.fuse = true; + std.debug.print("fused vs plain one-row decode, 12 tokens: {s}\n", .{if (differ == 0) "same bits" else "DIFFER"}); + } + if (std.c.getenv("GEMMA_DEBUG_ONLY") != null) return; + } + var s = try g4.state.Scratch.init(m.device, geo, 33024); + defer s.deinit(); + var c = try g4.state.Cache.init(m.device, geo, 33024); + defer c.deinit(); + var ids: [128]u32 = undefined; + for (&ids, 0..) |*t, i| t.* = @intCast(1000 + i); + for ([_]usize{ 0, 8000, 32000 }) |depth| { + c.len = depth; + for ([_]usize{ 1, 128 }) |rows| { + var best: f64 = 1e9; + for (0..5) |_| { + c.len = depth; + const t0 = std.c.mach_absolute_time(); + try g4.forward.run(m, &s, &.{.{ .cache = &c, .rows = rows }}, ids[0..rows]); + best = @min(best, @as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) / 24e3); + } + std.debug.print("whole forward, {d} rows at position {d}: {d:.2} ms wall ({d:.0} tok/s)\n", .{ rows, depth, best, @as(f64, @floatFromInt(rows)) * 1e3 / best }); + } + } +} diff --git a/zig/tests/gemma4/glue.zig b/zig/tests/gemma4/glue.zig new file mode 100644 index 0000000000..7ade4ed887 --- /dev/null +++ b/zig/tests/gemma4/glue.zig @@ -0,0 +1,356 @@ +//! Gemma 4 dense glue kernels against gemma4_dump.py's fixtures: embed, norm, q/k/v, attention, tails, softcap. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); + +const g4 = tf.gemma4; +const cfg = g4.config; +const npy = tf.npy; +/// The checkpoint's sizes, read from its config in main. +var geo: cfg.Geometry = undefined; +const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + +const Ctx = struct { + gpa: std.mem.Allocator, + io: std.Io, + device: mtl.Device, + queue: mtl.Queue, + ckpt: *tf.checkpoint_metal.Checkpoint, + k: g4.kernels.Kernels, + dump: []const u8, + failed: usize = 0, + owned: std.ArrayList(mtl.Buffer) = .empty, + + fn buf(c: *Ctx, bytes: usize) !mtl.Buffer { + const b = try c.device.buffer(@max(bytes, 16), shared); + try c.owned.append(c.gpa, b); + return b; + } + + /// A dumped array as a GPU buffer: bf16 bits (the dump writes bf16 as exact float32), or its raw f32/i32 data. + fn load(c: *Ctx, name: []const u8, as_bf16: bool) !struct { b: mtl.Buffer, n: usize } { + const path = try std.fmt.allocPrint(c.gpa, "{s}/{s}.npy", .{ c.dump, name }); + defer c.gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(c.io, path, c.gpa, .unlimited); + defer c.gpa.free(bytes); + const a = try npy.parse(bytes); + const n = a.count(); + if (as_bf16) { + if (!std.mem.eql(u8, a.descr, "> 16); + return .{ .b = b, .n = n }; + } + const b = try c.buf(n * 4); + @memcpy(b.slice(u8, n * 4), a.data[0 .. n * 4]); + return .{ .b = b, .n = n }; + } + + fn tensor(c: *Ctx, name: []const u8) !mtl.Buffer { + const t = try c.ckpt.get(name); + const b = try c.buf(t.bytes); + @memcpy(b.slice(u8, t.bytes), t.host(u8)[0..t.bytes]); + return b; + } + + fn run(c: *Ctx, body: anytype) !void { + const cb = c.queue.commandBuffer(); + const e = cb.compute(.serial); + try body.encode(e); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("GPU failure: {s}\n", .{msg}); + return error.GpuFailed; + } + } + + /// Counts bf16 outputs in `got` that differ from the dumped `want`. + fn compare(c: *Ctx, label: []const u8, got: mtl.Buffer, want_name: []const u8) !void { + const want = try c.load(want_name, true); + const a = got.slice(u16, want.n); + const b = want.b.slice(u16, want.n); + var diff: usize = 0; + for (a, b) |x, y| diff += @intFromBool(x != y); + c.failed += diff; + std.debug.print("{s} ({s}): {d} of {d} differ\n", .{ label, want_name, diff, want.n }); + } +}; + +var name_buf: [8][128]u8 = undefined; +fn nm(slot: usize, comptime f: []const u8, args: anytype) []const u8 { + return std.fmt.bufPrint(&name_buf[slot], f, args) catch unreachable; +} + +fn bf16Of(x: f32) f32 { + const b: u32 = @bitCast(x); + return @bitCast((b + 0x7FFF + ((b >> 16) & 1)) & 0xFFFF0000); +} + +fn rowsCheck(c: *Ctx, rows: usize, p0: usize) !void { + const kd = &c.k.sliding; + const tag = nm(0, "r{d}", .{rows}); + const ids = try c.load(nm(1, "{s}.tokens", .{tag}), false); + const emb = try c.buf(rows * geo.hidden * 2); + const ew = try c.tensor("language_model.model.embed_tokens.weight"); + const es = try c.tensor("language_model.model.embed_tokens.scales"); + const eb = try c.tensor("language_model.model.embed_tokens.biases"); + const scale = bf16Of(@sqrt(@as(f32, @floatFromInt(geo.hidden)))); + try c.run(struct { + kd: *const g4.kernels.Kind, + ids: mtl.Buffer, + w: mtl.Buffer, + s: mtl.Buffer, + b: mtl.Buffer, + out: mtl.Buffer, + scale: f32, + rows: usize, + fn encode(x: @This(), e: mtl.ComputeEncoder) !void { + e.setPipeline(x.kd.embed); + for ([_]mtl.Buffer{ x.ids, x.w, x.s, x.b, x.out }, 0..) |bb, i| e.setBuffer(bb, 0, i); + e.setValue(x.scale, 5); + e.dispatchThreads(.{ .width = geo.hidden / 2, .height = x.rows }, .{ .width = 256 }); + } + }{ .kd = kd, .ids = ids.b, .w = ew, .s = es, .b = eb, .out = emb, .scale = scale, .rows = rows }); + try c.compare("embed", emb, nm(2, "{s}.embed", .{tag})); + const normed = try c.buf(rows * geo.hidden * 2); + const w0 = try c.tensor("language_model.model.layers.0.input_layernorm.weight"); + try c.run(struct { + kd: *const g4.kernels.Kind, + x: mtl.Buffer, + w: mtl.Buffer, + out: mtl.Buffer, + rows: usize, + fn encode(s: @This(), e: mtl.ComputeEncoder) !void { + e.setPipeline(s.kd.rms); + for ([_]mtl.Buffer{ s.x, s.w, s.out }, 0..) |bb, i| e.setBuffer(bb, 0, i); + e.setValue(cfg.eps, 3); + e.setBuffer(s.out, 0, 4); + e.setValue(@as(i32, 0), 5); + e.dispatchGroups(.{ .width = s.rows }, .{ .width = 1024 }); + } + }{ .kd = kd, .x = emb, .w = w0, .out = normed, .rows = rows }); + try c.compare("rms looped", normed, nm(2, "{s}.normed0", .{tag})); + for ([_]usize{ 0, 5 }) |layer| try layerCheck(c, rows, p0, layer); + const raw = try c.load(nm(3, "{s}.head_raw", .{tag}), true); + try c.run(struct { + kd: *const g4.kernels.Kind, + x: mtl.Buffer, + n: usize, + fn encode(s: @This(), e: mtl.ComputeEncoder) !void { + e.setPipeline(s.kd.softcap); + e.setBuffer(s.x, 0, 0); + e.setValue(cfg.softcap, 1); + e.dispatchThreads(.{ .width = s.n }, .{ .width = 256 }); + } + }{ .kd = kd, .x = raw.b, .n = raw.n }); + try c.compare("softcap", raw.b, nm(4, "{s}.logits", .{tag})); +} + +fn layerCheck(c: *Ctx, rows: usize, p0: usize, layer: usize) !void { + const global = cfg.global(layer); + const kd = c.k.kind(global); + const dh = cfg.headDim(layer); + const nk = geo.kvHeads(layer); + const pre = nm(0, "r{d}.L{d}", .{ rows, layer }); + var pbuf: [64]u8 = undefined; + const prefix = try std.fmt.bufPrint(&pbuf, "{s}", .{pre}); + const lw = nm(5, "language_model.model.layers.{d}.", .{layer}); + var lbuf: [64]u8 = undefined; + const lp = try std.fmt.bufPrint(&lbuf, "{s}", .{lw}); + // q/k/v: head norms and RoPE at each row's position + const qkv = try c.load(nm(1, "{s}.qkv", .{prefix}), true); + const qw = try c.tensor(nm(2, "{s}self_attn.q_norm.weight", .{lp})); + const kw = try c.tensor(nm(2, "{s}self_attn.k_norm.weight", .{lp})); + const invf = try c.load(nm(2, "L{d}.inv_freq", .{layer}), false); + const pos = try c.buf(rows * 4); + for (pos.slice(i32, rows), 0..) |*p, i| p.* = @intCast(p0 + i); + const q = try c.buf(rows * geo.query_heads * dh * 2); + const k = try c.buf(nk * rows * dh * 2); + const v = try c.buf(nk * rows * dh * 2); + try c.run(struct { + kd: *const g4.kernels.Kind, + bs: [8]mtl.Buffer, + slots: usize, + rows: usize, + fn encode(s: @This(), e: mtl.ComputeEncoder) !void { + e.setPipeline(s.kd.qkv); + for (s.bs[0..5], 0..) |bb, i| e.setBuffer(bb, 0, i); + e.setValue(cfg.eps, 5); + for (s.bs[5..8], 6..) |bb, i| e.setBuffer(bb, 0, i); + for (s.bs[6..8], 9..) |bb, i| e.setBuffer(bb, 0, i); + e.setValue([3]i32{ 0, 0, 0 }, 11); + e.dispatchGroups(.{ .width = s.slots, .height = s.rows }, .{ .width = 32 }); + } + }{ .kd = kd, .bs = .{ qkv.b, qw, kw, invf.b, pos, q, k, v }, .slots = geo.query_heads + 2 * nk, .rows = rows }); + try c.compare("q", q, nm(3, "{s}.q", .{prefix})); + try c.compare("k", k, nm(3, "{s}.k", .{prefix})); + try c.compare("v", v, nm(3, "{s}.v", .{prefix})); + try attentionCheck(c, kd, prefix, rows, p0, global, dh, nk); + try tailsCheck(c, kd, prefix, lp, layer, rows); +} + +fn attentionCheck(c: *Ctx, kd: *const g4.kernels.Kind, prefix: []const u8, rows: usize, p0: usize, global: bool, dh: usize, nk: usize) !void { + const ring: usize = if (global) 0 else cfg.window + 128; + const cap: usize = if (global) p0 + 64 else ring; + const kb = try c.buf(nk * cap * dh * 2); + const vb = try c.buf(nk * cap * dh * 2); + const before_k = try c.load(nm(1, "{s}.keys_before", .{prefix}), true); + const before_v = try c.load(nm(2, "{s}.values_before", .{prefix}), true); + const first: usize = if (global) 0 else (if (p0 + 1 > cfg.window) p0 + 1 - cfg.window else 0); + const n_before = p0 - first; + for (0..nk) |h| for (0..n_before) |j| { + const p = first + j; + const slot = if (global) p else p % ring; + const src = (h * n_before + j) * dh; + const dst = (h * cap + slot) * dh; + @memcpy(kb.slice(u16, nk * cap * dh)[dst..][0..dh], before_k.b.slice(u16, before_k.n)[src..][0..dh]); + @memcpy(vb.slice(u16, nk * cap * dh)[dst..][0..dh], before_v.b.slice(u16, before_v.n)[src..][0..dh]); + }; + const q = try c.load(nm(3, "{s}.q", .{prefix}), true); + const kn = try c.load(nm(4, "{s}.k", .{prefix}), true); + const vn = try c.load(nm(6, "{s}.v", .{prefix}), true); + const pos = try c.buf(rows * 4); + const lo = try c.buf(rows * 4); + var low: usize = std.math.maxInt(usize); + for (0..rows) |i| { + const p = p0 + i; + pos.slice(i32, rows)[i] = @intCast(p); + const l: usize = if (global) 0 else (if (p + 1 > cfg.window) p + 1 - cfg.window else 0); + lo.slice(i32, rows)[i] = @intCast(l); + low = @min(low, l); + } + const first_chunk = low / kd.chunk; + const chunks = (p0 + rows - 1) / kd.chunk - first_chunk + 1; + const meta = try c.buf(32); + meta.slice(i32, 8)[0..5].* = .{ @intCast(first_chunk), @intCast(chunks), @intCast(ring), @intCast(rows), @intCast(p0) }; + const slots = geo.query_heads * rows * chunks; + const pm = try c.buf(slots * 4); + const pl = try c.buf(slots * 4); + const po = try c.buf(slots * dh * 4); + const out = try c.buf(rows * geo.query_heads * dh * 2); + try c.run(struct { + kd: *const g4.kernels.Kind, + bs: [13]mtl.Buffer, + cap: i32, + chunks: usize, + nk: usize, + rows: usize, + fn encode(s: @This(), e: mtl.ComputeEncoder) !void { + e.setPipeline(s.kd.partial); + for (s.bs[0..8], 0..) |bb, i| e.setBuffer(bb, 0, i); + e.setValue(s.cap, 8); + for (s.bs[8..11], 9..) |bb, i| e.setBuffer(bb, 0, i); + for (12..23) |i| e.setBuffer(s.bs[0], 0, i); // the prep's and the direct output's buffers, unread + e.setValue([2]i32{ 0, 0 }, 17); + e.dispatchGroups(.{ .width = s.chunks, .height = s.nk, .depth = s.rows }, .{ .width = 32 * s.kd.heads_per_kv * s.kd.split }); + e.barrier(); + e.setPipeline(s.kd.merge); + for ([_]mtl.Buffer{ s.bs[8], s.bs[9], s.bs[10], s.bs[7], s.bs[11], s.bs[8] }, 0..) |bb, i| e.setBuffer(bb, 0, i); + e.setValue(@as(i32, 0), 6); + e.dispatchGroups(.{ .width = 1, .height = geo.query_heads, .depth = s.rows }, .{ .width = 32 }); + } + }{ .kd = kd, .bs = .{ q.b, kb, vb, kn.b, vn.b, pos, lo, meta, pm, pl, po, out, out }, .cap = @intCast(cap), .chunks = chunks, .nk = nk, .rows = rows }); + try c.compare("attention", out, nm(5, "{s}.o_in", .{prefix})); +} + +fn tailsCheck(c: *Ctx, kd: *const g4.kernels.Kind, prefix: []const u8, lp: []const u8, layer: usize, rows: usize) !void { + const h = try c.load(nm(1, "{s}.h", .{prefix}), true); + const o = try c.load(nm(2, "{s}.o", .{prefix}), true); + const wa = try c.tensor(nm(3, "{s}post_attention_layernorm.weight", .{lp})); + const w1 = try c.tensor(nm(3, "{s}pre_feedforward_layernorm.weight", .{lp})); + const hn = try c.buf(rows * geo.hidden * 2); + const n1 = try c.buf(rows * geo.hidden * 2); + const Tail = struct { + pipe: mtl.Pipeline, + bs: []const mtl.Buffer, + eps_at: usize, + rows: usize, + fn encode(s: @This(), e: mtl.ComputeEncoder) !void { + e.setPipeline(s.pipe); + for (s.bs, 0..) |bb, i| e.setBuffer(bb, 0, if (i < s.eps_at) i else i + 1); + e.setValue(cfg.eps, s.eps_at); + e.setBuffer(s.bs[0], 0, s.bs.len + 1); + e.setValue(@as(i32, 0), s.bs.len + 2); + e.setBuffer(s.bs[0], 0, s.bs.len + 3); // dense_tail: no tap + e.setValue(@as(i32, 0), s.bs.len + 4); + e.dispatchGroups(.{ .width = s.rows }, .{ .width = 256 }); + } + }; + try c.run(Tail{ .pipe = kd.attn_tail, .bs = &.{ h.b, o.b, wa, w1, hn, n1 }, .eps_at = 4, .rows = rows }); + try c.compare("attn_tail hn", hn, nm(4, "{s}.hn", .{prefix})); + try c.compare("attn_tail norm", n1, nm(4, "{s}.n_mlp", .{prefix})); + const gu = try c.load(nm(1, "{s}.gate_up", .{prefix}), true); + const act = try c.buf(rows * geo.intermediate * 2); + try c.run(struct { + kd: *const g4.kernels.Kind, + gu: mtl.Buffer, + act: mtl.Buffer, + rows: usize, + fn encode(s: @This(), e: mtl.ComputeEncoder) !void { + e.setPipeline(s.kd.geglu); + e.setBuffer(s.gu, 0, 0); + e.setBuffer(s.act, 0, 1); + e.setValue(@as(u32, @intCast(geo.intermediate)), 2); + e.setBuffer(s.act, 0, 3); + e.setValue(@as(i32, 0), 4); + e.setValue(@as(i32, 0), 5); + e.dispatchGroups(.{ .width = geo.intermediate / 64, .height = s.rows }, .{ .width = 64 }); + } + }{ .kd = kd, .gu = gu.b, .act = act, .rows = rows }); + try c.compare("geglu", act, nm(4, "{s}.act", .{prefix})); + const hn_in = try c.load(nm(1, "{s}.hn", .{prefix}), true); + const y = try c.load(nm(2, "{s}.y", .{prefix}), true); + const wp = try c.tensor(nm(3, "{s}post_feedforward_layernorm.weight", .{lp})); + const sc = try c.tensor(nm(3, "{s}layer_scalar", .{lp})); + const wn = try c.tensor(nm(3, "language_model.model.layers.{d}.input_layernorm.weight", .{layer + 1})); + const h2 = try c.buf(rows * geo.hidden * 2); + const nx = try c.buf(rows * geo.hidden * 2); + try c.run(Tail{ .pipe = kd.dense_tail, .bs = &.{ hn_in.b, y.b, wp, sc, wn, h2, nx }, .eps_at = 5, .rows = rows }); + try c.compare("dense_tail h", h2, nm(4, "{s}.h2", .{prefix})); + try c.compare("dense_tail next", nx, nm(4, "{s}.normed2", .{prefix})); +} + +pub fn main(init: std.process.Init) !void { + const gpa = init.gpa; + var args = try std.process.Args.Iterator.initAllocator(init.minimal.args, gpa); + defer args.deinit(); + _ = args.next(); + const model = args.next() orelse return error.Usage; + const dump = args.next() orelse return error.Usage; + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var ckpt = tf.checkpoint_metal.Checkpoint.init(gpa); + defer ckpt.deinit(); + const files = try tf.checkpoint_host.shardFiles(gpa, init.io, model); + defer tf.checkpoint_host.freeShardFiles(gpa, files); + for (files) |path| try ckpt.addFileSelected(device, path, "", "language_model."); + const conf = try cfg.Config.read(gpa, init.io, model); + geo = conf.geo; + var c = Ctx{ .gpa = gpa, .io = init.io, .device = device, .queue = queue, .ckpt = &ckpt, .k = try g4.kernels.Kernels.init(gpa, device, geo, conf.widths.embed.bits), .dump = dump }; + defer c.k.deinit(); + defer { + for (c.owned.items) |b| b.deinit(); + c.owned.deinit(gpa); + } + for ([_]bool{ false, true }) |global| { + var mine: [cfg.global_dim / 2]f32 = undefined; + const n: usize = @as(usize, if (global) cfg.global_dim else cfg.sliding_dim) / 2; + try g4.kernels.inverseFrequencies(&c.k, queue, device, global, mine[0..n]); + const ref = try c.load(if (global) "L5.inv_freq" else "L0.inv_freq", false); + var diff: usize = 0; + for (mine[0..n], ref.b.slice(u32, n)) |m, r| diff += @intFromBool(@as(u32, @bitCast(m)) != r); + c.failed += diff; + std.debug.print("rope table {s}: {d} of {d} differ\n", .{ if (global) "global" else "sliding", diff, n }); + } + for ([_]usize{ 1, 2, 16 }) |rows| { + std.debug.print("== rows {d}\n", .{rows}); + try rowsCheck(&c, rows, 1100); + } + std.debug.print("{s}: {d} differing outputs\n", .{ if (c.failed == 0) "all bits equal" else "bits differ", c.failed }); + if (c.failed > 0) std.process.exit(1); +} diff --git a/zig/tests/gemma4/kernels.zig b/zig/tests/gemma4/kernels.zig new file mode 100644 index 0000000000..77404631e3 --- /dev/null +++ b/zig/tests/gemma4/kernels.zig @@ -0,0 +1,240 @@ +//! Gemma 4 dense decode kernels against tools/zig/gemma4_dump.py's fixtures: each output's bits, by window width. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); + +const lane = tf.lane_projection; +const g4 = tf.gemma4; +const npy = tf.npy; +const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + +const Ctx = struct { + gpa: std.mem.Allocator, + io: std.Io, + device: mtl.Device, + queue: mtl.Queue, + ckpt: *tf.checkpoint_metal.Checkpoint, + geo: g4.config.Geometry, + dump: []const u8, + failed: usize = 0, +}; + +/// A dumped array's values as bf16 bits (the dump writes bf16 tensors as exact float32). +fn loadBf16(c: *Ctx, name: []const u8) ![]u16 { + const path = try std.fmt.allocPrint(c.gpa, "{s}/{s}.npy", .{ c.dump, name }); + defer c.gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(c.io, path, c.gpa, .unlimited); + defer c.gpa.free(bytes); + const a = try npy.parse(bytes); + if (!std.mem.eql(u8, a.descr, "> 16); + return out; +} + +/// The members' MLX affine tensors stacked along N, packed into one lane projection's layout. +const Stack = struct { + proj: lane.Projection, + w: mtl.Buffer, + sb: mtl.Buffer, + n: usize, + k: usize, + + /// `decode`: the split the model decodes at (weights.splitK) rather than the reference's (referenceSplit). + fn init(c: *Ctx, names: []const []const u8, precompute: bool, decode: bool) !Stack { + const reg = reg_mode; + var n: usize = 0; + var k: usize = 0; + for (names) |m| { + const t = try c.ckpt.get(try key(c, m, "weight")); + n += t.shape[0]; + k = t.shape[1] * 8; + } + const words = try c.gpa.alloc(u32, n * k / 8); + defer c.gpa.free(words); + const scales = try c.gpa.alloc(u16, n * k / 64); + defer c.gpa.free(scales); + const biases = try c.gpa.alloc(u16, n * k / 64); + defer c.gpa.free(biases); + var at: usize = 0; + for (names) |m| { + const w = (try c.ckpt.get(try key(c, m, "weight"))).host(u32); + const s = (try c.ckpt.get(try key(c, m, "scales"))).host(u16); + const b = (try c.ckpt.get(try key(c, m, "biases"))).host(u16); + @memcpy(words[at * k / 8 ..][0..w.len], w); + @memcpy(scales[at * k / 64 ..][0..s.len], s); + @memcpy(biases[at * k / 64 ..][0..b.len], b); + at += w.len / (k / 8); + } + const sk = if (decode) g4.weights.splitK(c.geo, n, k) else g4.weights.referenceSplit(n, k); + const l = lane.Layout{ .n = n, .k = k, .format = .{ .bits = 4, .group = 64 }, .sk = sk, .precompute_sums = precompute or reg, .reg = reg }; + const wb = try c.device.buffer(l.weightWords() * 4, shared); + const sbb = try c.device.buffer(l.metadataElements() * 2, shared); + try lane.pack(l, words, scales, biases, wb.slice(u32, l.weightWords()), sbb.slice(u16, l.metadataElements())); + if (reg) try lane.regOrder(c.device, &.{.{ .w = .{ .buf = wb }, .n = n, .k = k }}); + return .{ .proj = try lane.Projection.init(c.gpa, c.device, l), .w = wb, .sb = sbb, .n = n, .k = k }; + } + + fn deinit(s: Stack) void { + s.proj.deinit(); + s.w.deinit(); + s.sb.deinit(); + } +}; + +var key_buf: [256]u8 = undefined; +var reg_mode = false; +fn key(_: *Ctx, member: []const u8, part: []const u8) ![]const u8 { + return std.fmt.bufPrint(&key_buf, "language_model.model.{s}.{s}", .{ member, part }); +} + +/// Runs `s` on the dumped input and counts outputs whose bits differ; a decode-split stack's are reported, not counted. +fn checkProjection(c: *Ctx, label: []const u8, s: Stack, rows: usize, in_name: []const u8, out_name: []const u8, decode: bool) !void { + const x = try loadBf16(c, in_name); + defer c.gpa.free(x); + const want = try loadBf16(c, out_name); + defer c.gpa.free(want); + if (x.len != rows * s.k or want.len != rows * s.n) return error.FixtureShape; + const xb = try c.device.buffer(x.len * 2, shared); + defer xb.deinit(); + @memcpy(xb.slice(u16, x.len), x); + const yb = try c.device.buffer(want.len * 2, shared); + defer yb.deinit(); + const cb = c.queue.commandBuffer(); + const e = cb.compute(.serial); + try s.proj.encode(e, .{ .buf = xb }, .{ .buf = s.w }, .{ .buf = s.sb }, .{ .buf = yb }, @intCast(rows)); + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("{s}: GPU failure {s}\n", .{ label, msg }); + return error.GpuFailed; + } + const got = yb.slice(u16, want.len); + var diff: usize = 0; + var worst: f32 = 0; + for (got, want) |g, w| if (g != w) { + diff += 1; + worst = @max(worst, @abs(@as(f32, @bitCast(@as(u32, g) << 16)) - @as(f32, @bitCast(@as(u32, w) << 16)))); + }; + if (decode) { + std.debug.print("{s} rows {d} at the decode split: {d} of {d} outputs differ, worst |diff| {d}\n", .{ label, rows, diff, want.len, worst }); + return; + } + c.failed += diff; + std.debug.print("{s} rows {d}: {d} of {d} outputs differ{s}\n", .{ label, rows, diff, want.len, if (diff == 0) "" else " (first at a column of row 0 or later)" }); +} + +/// One affine_mm product (layer 0's q_proj, from the checkpoint) on the dumped input against the dumped output. +fn checkAffine(c: *Ctx) !void { + const text = try tf.affine_mm.source(c.device, c.gpa, .{}); + defer c.gpa.free(text); + const lib = try mtl.Library.fromSource(c.device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + var pipes: tf.affine_mm.Pipes = undefined; + for (&pipes, tf.affine_mm.names) |*p, nm| p.* = try mtl.Pipeline.init(c.device, lib, nm, false); + const x = try loadBf16(c, "r16.L0.x"); + defer c.gpa.free(x); + const want = try loadBf16(c, "r16.L0.qkv"); + defer c.gpa.free(want); + const xb = try c.device.buffer(x.len * 2, shared); + defer xb.deinit(); + @memcpy(xb.slice(u16, x.len), x); + const qn = c.geo.query_heads * g4.config.sliding_dim; // layer 0's q, then its k and v: the dumped qkv's row + const row = qn + 2 * c.geo.sliding_kv_heads * g4.config.sliding_dim; + const yb = try c.device.buffer(16 * row * 2, shared); + defer yb.deinit(); + const w = try c.ckpt.get("language_model.model.layers.0.self_attn.q_proj.weight"); + const sc = try c.ckpt.get("language_model.model.layers.0.self_attn.q_proj.scales"); + const bi = try c.ckpt.get("language_model.model.layers.0.self_attn.q_proj.biases"); + const Q = struct { w: lane.Ref, s: lane.Ref, b: lane.Ref, n: u32, k: u32 }; + const q = Q{ .w = .{ .buf = w.buffer, .off = w.offset }, .s = .{ .buf = sc.buffer, .off = sc.offset }, .b = .{ .buf = bi.buffer, .off = bi.offset }, .n = @intCast(qn), .k = @intCast(c.geo.hidden) }; + const cb = c.queue.commandBuffer(); + const e = cb.compute(.serial); + tf.affine_mm.denseBatch(e, pipes, lane.Ref{ .buf = xb }, q, lane.Ref{ .buf = yb }, 16, 1, .{ .x_row = @intCast(c.geo.hidden), .y_row = @intCast(row) }); + e.end(); + cb.commit(); + cb.wait(); + const got = yb.slice(u16, 16 * row); + var worst: f32 = 0; + var exact: usize = 0; + for (0..16) |r| for (0..qn) |j| { + const a: f32 = @bitCast(@as(u32, got[r * row + j]) << 16); + const b: f32 = @bitCast(@as(u32, want[r * row + j]) << 16); + worst = @max(worst, @abs(a - b)); + exact += @intFromBool(got[r * row + j] == want[r * row + j]); + }; + std.debug.print("affine_mm q_proj vs lane: {d} of {d} bit-equal, worst |diff| {d}; sample got {x} want {x}\n", .{ exact, 16 * qn, worst, got[0], want[0] }); +} + +pub fn main(init: std.process.Init) !void { + const gpa = init.gpa; + var args = try std.process.Args.Iterator.initAllocator(init.minimal.args, gpa); + defer args.deinit(); + _ = args.next(); + const model = args.next() orelse return error.Usage; + const dump = args.next() orelse return error.Usage; + const flag = args.next(); + const precompute = if (flag) |a| std.mem.eql(u8, a, "--sums") else false; + reg_mode = if (flag) |a| std.mem.eql(u8, a, "--reg") else false; + const device = try mtl.Device.init(); + defer device.deinit(); + const queue = try device.queue(); + defer queue.deinit(); + var ckpt = tf.checkpoint_metal.Checkpoint.init(gpa); + defer ckpt.deinit(); + const files = try tf.checkpoint_host.shardFiles(gpa, init.io, model); + defer tf.checkpoint_host.freeShardFiles(gpa, files); + for (files) |path| try ckpt.addFileSelected(device, path, "", "language_model."); + const geo = (try g4.config.Config.read(gpa, init.io, model)).geo; + var c = Ctx{ .gpa = gpa, .io = init.io, .device = device, .queue = queue, .ckpt = &ckpt, .geo = geo, .dump = dump }; + std.debug.print("device {s}, precomputed sums {}, reg {}\n", .{ device.name(), precompute, reg_mode }); + if (flag) |f| if (std.mem.eql(u8, f, "--affine")) return checkAffine(&c); + for ([_]usize{ 0, 5 }) |layer| { + var names: [4][64]u8 = undefined; + const q = try std.fmt.bufPrint(&names[0], "layers.{d}.self_attn.q_proj", .{layer}); + const k = try std.fmt.bufPrint(&names[1], "layers.{d}.self_attn.k_proj", .{layer}); + const v = try std.fmt.bufPrint(&names[2], "layers.{d}.self_attn.v_proj", .{layer}); + var mlp: [3][64]u8 = undefined; + const gate = try std.fmt.bufPrint(&mlp[0], "layers.{d}.mlp.gate_proj", .{layer}); + const up = try std.fmt.bufPrint(&mlp[1], "layers.{d}.mlp.up_proj", .{layer}); + const o_name = try std.fmt.bufPrint(&names[3], "layers.{d}.self_attn.o_proj", .{layer}); + const down_name = try std.fmt.bufPrint(&mlp[2], "layers.{d}.mlp.down_proj", .{layer}); + for ([_]bool{ false, true }) |decode| { + const qkv = try Stack.init(&c, if (layer % 6 == 5) &.{ q, k } else &.{ q, k, v }, precompute, decode); + defer qkv.deinit(); + const o = try Stack.init(&c, &.{o_name}, precompute, decode); + defer o.deinit(); + const gu = try Stack.init(&c, &.{ gate, up }, precompute, decode); + defer gu.deinit(); + const down = try Stack.init(&c, &.{down_name}, precompute, decode); + defer down.deinit(); + for ([_]usize{ 1, 2, 16 }) |rows| { + if (decode and rows != 1) continue; + var p: [6][48]u8 = undefined; + const pre = try std.fmt.bufPrint(&p[0], "r{d}.L{d}", .{ rows, layer }); + try checkProjection(&c, "qkv", qkv, rows, try cat(&p[1], pre, ".x"), try cat(&p[2], pre, ".qkv"), decode); + try checkProjection(&c, "o", o, rows, try cat(&p[1], pre, ".o_in"), try cat(&p[2], pre, ".o"), decode); + try checkProjection(&c, "gate_up", gu, rows, try cat(&p[1], pre, ".n_mlp"), try cat(&p[2], pre, ".gate_up"), decode); + try checkProjection(&c, "down", down, rows, try cat(&p[1], pre, ".act"), try cat(&p[2], pre, ".y"), decode); + } + } + std.debug.print("layer {d} done\n", .{layer}); + } + for ([_]bool{ false, true }) |decode| { + const head = try Stack.init(&c, &.{"embed_tokens"}, precompute, decode); + defer head.deinit(); + for ([_]usize{ 1, 2, 16 }) |rows| { + if (decode and rows != 1) continue; + var p: [3][48]u8 = undefined; + const pre = try std.fmt.bufPrint(&p[0], "r{d}", .{rows}); + try checkProjection(&c, "head", head, rows, try cat(&p[1], pre, ".head_in"), try cat(&p[2], pre, ".head_raw"), decode); + } + } + std.debug.print("{s}: {d} differing outputs\n", .{ if (c.failed == 0) "all bits equal" else "bits differ", c.failed }); + if (c.failed > 0) std.process.exit(1); +} + +fn cat(buf: []u8, a: []const u8, b: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}{s}", .{ a, b }); +} diff --git a/zig/tests/gemma4/lanes.zig b/zig/tests/gemma4/lanes.zig new file mode 100644 index 0000000000..047b6b1cb3 --- /dev/null +++ b/zig/tests/gemma4/lanes.zig @@ -0,0 +1,197 @@ +//! Gemma 4 dense on the lane core: copy drafts, DFlash or MTP assistant drafts (GEMMA_DRAFTER) against plain replies. +const std = @import("std"); +const tf = @import("tensorfold"); + +const lanes = tf.lanes; +const g4 = tf.gemma4; + +const Prompt = struct { name: []const u8, ids: []const u32, want: []const u32 }; + +fn ms(io: std.Io, since: std.Io.Timestamp) f64 { + return @as(f64, @floatFromInt(since.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds())) / 1e6; +} + +fn readJson(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8, suffix: []const u8) ![]u32 { + var p: [512]u8 = undefined; + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, try std.fmt.bufPrint(&p, "{s}/{s}{s}", .{ dir, name, suffix }), gpa, .limited(1 << 22)); + defer gpa.free(bytes); + const parsed = try std.json.parseFromSlice([]u32, gpa, bytes, .{}); + defer parsed.deinit(); + return gpa.dupe(u32, parsed.value); +} + +/// The engine's own plain reply to `ids` (one stream, no drafts): the reference drafted runs must equal. +fn plain(gpa: std.mem.Allocator, engine: *lanes.Engine, ids: []const u32, count: u32) ![]u32 { + return plainWith(gpa, engine, ids, count, null); +} + +/// `sampling`: draws keyed by position from its seed, the same in a plain and a drafted run. +fn plainWith(gpa: std.mem.Allocator, engine: *lanes.Engine, ids: []const u32, count: u32, sampling: ?lanes.Sampling) ![]u32 { + var s = try lanes.Stream.init(gpa, .{ .id = "plain", .prompt = ids, .max_new = count, .drafts = false, .sampling = sampling }); + defer s.deinit(gpa); + try engine.addStream(&s); + while (engine.activeCount() > 0) try engine.step(); + return gpa.dupe(u32, s.emitted()); +} + +/// Runs `streams` copies of the prompts (round-robin) together; returns how many streams' tokens differ. +fn case(gpa: std.mem.Allocator, io: std.Io, engine: *lanes.Engine, prompts: []const Prompt, streams: usize, copies: bool, label: []const u8) !usize { + return caseWith(gpa, io, engine, prompts, streams, copies, true, null, label); +} + +/// `drafts` false runs the streams plain on a drafted engine; `copies` adds the copy proposer. +fn caseWith(gpa: std.mem.Allocator, io: std.Io, engine: *lanes.Engine, prompts: []const Prompt, streams: usize, copies: bool, drafts: bool, sampling: ?lanes.Sampling, label: []const u8) !usize { + const ss = try gpa.alloc(lanes.Stream, streams); + defer gpa.free(ss); + const looks = try gpa.alloc(lanes.SuffixLookup, streams); + defer gpa.free(looks); + for (ss, looks, 0..) |*s, *l, i| { + const p = prompts[i % prompts.len]; + l.* = try lanes.SuffixLookup.init(gpa, .{ .min_match = 4 }); + s.* = try lanes.Stream.init(gpa, .{ .id = p.name, .prompt = p.ids, .max_new = @intCast(p.want.len), .drafts = drafts, .sampling = sampling, .proposer = if (copies) l.proposer() else null }); + } + defer for (ss, looks) |*s, *l| { + s.deinit(gpa); + l.deinit(); + }; + const t0 = std.Io.Clock.awake.now(io); + for (ss) |*s| try engine.addStream(s); + var steps: usize = 0; + while (engine.activeCount() > 0) : (steps += 1) try engine.step(); + const elapsed = ms(io, t0); + if (std.c.getenv("GEMMA_DRAFT_TRACE") != null) std.debug.print("{s}: {d} engine steps, {d:.2} ms a step\n", .{ label, steps, elapsed / @as(f64, @floatFromInt(@max(steps, 1))) }); + var bad: usize = 0; + var tokens: usize = 0; + var rounds: usize = 0; + for (ss, 0..) |*s, i| { + const p = prompts[i % prompts.len]; + const ok = std.mem.eql(u32, s.emitted(), p.want); + bad += @intFromBool(!ok); + tokens += s.emitted().len; + rounds += s.rounds; + if (!ok) { + const at = for (s.emitted(), p.want, 0..) |a, b, j| { + if (a != b) break j; + } else @min(s.emitted().len, p.want.len); + std.debug.print(" {s} stream {d} ({s}) differs at token {d}\n", .{ label, i, p.name, at }); + } + } + std.debug.print("{s}: {d} of {d} streams equal their reference, {d} tokens in {d} stream rounds ({d:.2} a round), {d:.0} tok/s\n", .{ label, streams - bad, streams, tokens, rounds, @as(f64, @floatFromInt(tokens)) / @as(f64, @floatFromInt(@max(rounds, 1))), @as(f64, @floatFromInt(tokens)) * 1e3 / elapsed }); + return bad; +} + +pub fn main(init: std.process.Init) !void { + const gpa = init.gpa; + const io = init.io; + var args = try std.process.Args.Iterator.initAllocator(init.minimal.args, gpa); + defer args.deinit(); + _ = args.next(); + const model_dir = args.next() orelse return error.Usage; + const oracle = args.next() orelse return error.Usage; + const extra = args.next(); // a directory with repeat.ids.json and edit.ids.json: copy-heavy prompts + const m = try g4.Model.load(gpa, io, model_dir); + defer m.deinit(); + const metal = try g4.backend.Metal.init(gpa, io, m, .{ .capacity = 4096, .streams = 8 }); + defer metal.deinit(); + try metal.measure(); + std.debug.print("window forward: {d:.1} ms at 1 row, {d:.1} at 16, {d:.1} at 32, {d:.1} at 64, {d:.1} at 128\n", .{ metal.costs[0].ms, metal.costs[15].ms, try metal.timeRows(32), try metal.timeRows(64), try metal.timeRows(128) }); + var cfg = try lanes.Config.init(gpa, metal.facts(), g4.state.max_rows, g4.state.max_rows - 1); + defer cfg.deinit(gpa); + var clock = lanes.backend.WallClock{ .io = io }; + var prompts: [4]Prompt = undefined; + for (&prompts, [_][]const u8{ "copy", "arithmetic", "story", "code" }) |*p, name| { + p.* = .{ .name = name, .ids = try readJson(gpa, io, oracle, name, ".ids.json"), .want = &.{} }; + } + defer for (prompts) |p| { + gpa.free(p.ids); + gpa.free(p.want); + }; + var bad: usize = 0; + if (std.c.getenv("GEMMA_DRAFTER")) |dir| return drafted(gpa, io, m, metal, &prompts, std.mem.span(dir), extra); + var engine = lanes.Engine.init(gpa, &cfg, metal.backend(), clock.clock()); + defer engine.deinit(); + // the reference: the engine's own plain one-stream reply (prompts take the tensor-unit path, not the oracle's) + for (&prompts) |*p| p.want = try plain(gpa, &engine, p.ids, 128); + for (prompts) |p| { // and how it compares with the oracle's receipt (informative: the prompt paths differ) + const want = readJson(gpa, io, oracle, p.name, ".reference.json") catch continue; + defer gpa.free(want); + const n = @min(want.len, p.want.len); + const at = for (want[0..n], p.want[0..n], 0..) |a, b, j| { + if (a != b) break j; + } else n; + std.debug.print("plain {s} vs oracle: {s} (first difference at {d} of {d})\n", .{ p.name, if (at == n) "equal" else "DIFFERS", at, n }); + } + for (prompts) |p| bad += try case(gpa, io, &engine, &.{p}, 1, true, "copy drafts, one stream"); + bad += try case(gpa, io, &engine, &prompts, 4, true, "copy drafts, four streams"); + bad += try case(gpa, io, &engine, &prompts, 8, true, "copy drafts, eight streams"); + if (extra) |dir| { + var copy_heavy: [2]Prompt = undefined; + for (©_heavy, [_][]const u8{ "repeat", "edit" }) |*p, name| { + const ids = try readJson(gpa, io, dir, name, ".ids.json"); + p.* = .{ .name = name, .ids = ids, .want = try plain(gpa, &engine, ids, 256) }; + } + defer for (copy_heavy) |p| { + gpa.free(p.ids); + gpa.free(p.want); + }; + for (copy_heavy) |p| bad += try case(gpa, io, &engine, &.{p}, 1, true, "copy drafts on a copying reply, one stream"); + const mixed = [_]Prompt{ copy_heavy[0], prompts[0], copy_heavy[1], prompts[1], copy_heavy[0], prompts[2], copy_heavy[1], prompts[3] }; + bad += try case(gpa, io, &engine, &mixed, 8, true, "copy drafts, eight mixed streams"); + } + std.debug.print("{s}: {d} streams differ\n", .{ if (bad == 0) "all equal" else "differ", bad }); + if (bad > 0) std.process.exit(1); +} + +/// The drafter at `dir` in front of the backend (DFlash, or the MTP assistant by its config): drafted == plain. +fn drafted(gpa: std.mem.Allocator, io: std.Io, m: *g4.Model, metal: *g4.backend.Metal, prompts: []Prompt, dir: []const u8, extra: ?[]const u8) !void { + const t0 = std.Io.Clock.awake.now(io); + var flash: ?*g4.dflash.Drafter = null; + defer if (flash) |d| d.deinit(); + var assist: ?*g4.assistant.Drafter = null; + defer if (assist) |d| d.deinit(); + const head: lanes.drafter.Drafter = if (try g4.assistant_weights.isAssistant(gpa, io, dir)) blk: { + const d = try g4.assistant.Drafter.init(gpa, io, metal, dir); + assist = d; + std.debug.print("MTP assistant loaded in {d:.0} ms; a chain of 1 draft {d:.2} ms, of {d} drafts {d:.2} ms, step {d:.3} ms a draft\n", .{ ms(io, t0), d.chain_ms[0], g4.assistant.depth, d.chain_ms[1], d.step_ms }); + break :blk d.drafter(); + } else blk: { + const d = try g4.dflash.Drafter.init(gpa, io, m, dir, .{ .capacity = 4096 }); + flash = d; + std.debug.print("drafter loaded in {d:.0} ms; a block of 1 draft {d:.2} ms, of 15 drafts {d:.2} ms, step {d:.3} ms a draft\n", .{ ms(io, t0), d.block_ms[0], d.block_ms[1], d.step_ms }); + break :blk d.drafter(); + }; + const name = if (assist != null) "MTP assistant" else "DFlash"; + var wrap = try lanes.drafted.Drafted.init(gpa, metal.backend(), head); + defer wrap.deinit(); + var cfg = try lanes.Config.init(gpa, wrap.facts(metal.facts()), g4.state.max_rows, g4.state.max_rows - 1); + defer cfg.deinit(gpa); + var clock = lanes.backend.WallClock{ .io = io }; + var engine = lanes.Engine.init(gpa, &cfg, wrap.backend(), clock.clock()); + defer engine.deinit(); + var bad: usize = 0; + var label: [128]u8 = undefined; + for (prompts) |*p| p.want = try plain(gpa, &engine, p.ids, 128); + for (prompts) |p| bad += try caseWith(gpa, io, &engine, &.{p}, 1, false, true, null, try std.fmt.bufPrint(&label, "{s}, one stream", .{name})); + bad += try caseWith(gpa, io, &engine, prompts, 4, false, true, null, try std.fmt.bufPrint(&label, "{s}, four streams", .{name})); + bad += try caseWith(gpa, io, &engine, prompts, 8, false, true, null, try std.fmt.bufPrint(&label, "{s}, eight streams", .{name})); + const seeded: lanes.Sampling = .{ .seed = 7, .temperature = 1.0 }; + var sampled: [4]Prompt = undefined; + for (&sampled, prompts) |*s, p| s.* = .{ .name = p.name, .ids = p.ids, .want = try plainWith(gpa, &engine, p.ids, 128, seeded) }; + defer for (sampled) |s| gpa.free(s.want); + for (sampled) |p| bad += try caseWith(gpa, io, &engine, &.{p}, 1, false, true, seeded, try std.fmt.bufPrint(&label, "{s}, seeded sampling, one stream", .{name})); + bad += try caseWith(gpa, io, &engine, &sampled, 4, false, true, seeded, try std.fmt.bufPrint(&label, "{s}, seeded sampling, four streams", .{name})); + if (extra) |x| { + var copy_heavy: [2]Prompt = undefined; + for (©_heavy, [_][]const u8{ "repeat", "edit" }) |*p, n| { + const ids = try readJson(gpa, io, x, n, ".ids.json"); + p.* = .{ .name = n, .ids = ids, .want = try plain(gpa, &engine, ids, 256) }; + } + defer for (copy_heavy) |p| { + gpa.free(p.ids); + gpa.free(p.want); + }; + for (copy_heavy) |p| bad += try caseWith(gpa, io, &engine, &.{p}, 1, false, true, null, try std.fmt.bufPrint(&label, "{s} on a copying reply, one stream", .{name})); + } + std.debug.print("{s}: {d} streams differ\n", .{ if (bad == 0) "all equal" else "differ", bad }); + if (bad > 0) std.process.exit(1); +} diff --git a/zig/tests/gemma4/run.zig b/zig/tests/gemma4/run.zig new file mode 100644 index 0000000000..17784bc47a --- /dev/null +++ b/zig/tests/gemma4/run.zig @@ -0,0 +1,105 @@ +//! Gemma 4 dense greedy decode on the native engine: tokens and top-two margins for gemma4_gate.py check. +const std = @import("std"); +const tf = @import("tensorfold"); + +const g4 = tf.gemma4; + +const Top2 = struct { ids: [2]u32, logits: [2]f32, position: usize }; + +fn top2(row: []const u16, position: usize) Top2 { + var best: [2]u32 = .{ 0, 0 }; + var val: [2]f32 = .{ -std.math.inf(f32), -std.math.inf(f32) }; + for (row, 0..) |bits, i| { + const v: f32 = @bitCast(@as(u32, bits) << 16); + if (v > val[0]) { + best[1] = best[0]; + val[1] = val[0]; + best[0] = @intCast(i); + val[0] = v; + } else if (v > val[1]) { + best[1] = @intCast(i); + val[1] = v; + } + } + return .{ .ids = best, .logits = val, .position = position }; +} + +pub fn main(init: std.process.Init) !void { + const gpa = init.gpa; + const io = init.io; + var args = try std.process.Args.Iterator.initAllocator(init.minimal.args, gpa); + defer args.deinit(); + _ = args.next(); + const model_dir = args.next() orelse return error.Usage; + const oracle = args.next() orelse return error.Usage; + const out = args.next() orelse return error.Usage; + const count = if (args.next()) |a| try std.fmt.parseInt(usize, a, 10) else 128; + const width = if (args.next()) |a| try std.fmt.parseInt(usize, a, 10) else 1; // prompt rows a window + const cwd = std.Io.Dir.cwd(); + try cwd.createDirPath(io, out); + const t0 = std.c.mach_absolute_time(); + const m = try g4.Model.load(gpa, io, model_dir); + defer m.deinit(); + std.debug.print("loaded in {d:.1} s\n", .{@as(f64, @floatFromInt(std.c.mach_absolute_time() - t0)) / 24e6}); + var dir = try cwd.openDir(io, oracle, .{ .iterate = true }); + defer dir.close(io); + var it = dir.iterate(); + while (try it.next(io)) |entry| { + if (!std.mem.endsWith(u8, entry.name, ".ids.json")) continue; + const name = entry.name[0 .. entry.name.len - ".ids.json".len]; + const bytes = try dir.readFileAlloc(io, entry.name, gpa, .limited(1 << 22)); + defer gpa.free(bytes); + const ids = try std.json.parseFromSlice([]u32, gpa, bytes, .{}); + defer ids.deinit(); + const capacity = (ids.value.len + count + 64 + 63) / 64 * 64; + var cache = try g4.state.Cache.init(m.device, m.config.geo, capacity); + defer cache.deinit(); + var s = try g4.state.Scratch.init(m.device, m.config.geo, capacity); + defer s.deinit(); + const t1 = std.c.mach_absolute_time(); + var ps: ?g4.prompt.Scratch = if (width == g4.state.chunk_rows) try g4.prompt.Scratch.init(m.device, m.config.geo, capacity) else null; + defer if (ps) |*p| p.deinit(); + var at: usize = 0; + while (at < ids.value.len) : (at += width) { + const end = @min(at + width, ids.value.len); + if (ps) |*p| try g4.prompt.run(m, p, &cache, ids.value[at..end], &.{}) else try g4.forward.window(m, &s, &cache, ids.value[at..end], at); + } + if (ps) |*p| @memcpy(s.logits.slice(u16, g4.config.vocab), p.logits.slice(u16, g4.config.vocab)); + const last = if (ps != null) 0 else (ids.value.len - 1 - (ids.value.len - 1) / width * width) * g4.config.vocab; + const prefill_s = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t1)) / 24e6; + var lp: [256]u8 = undefined; + const row = s.logits.slice(u16, g4.state.prompt_rows * g4.config.vocab)[last..][0..g4.config.vocab]; + const f32s = try gpa.alloc(f32, g4.config.vocab); + defer gpa.free(f32s); + for (f32s, row) |*f, b| f.* = @bitCast(@as(u32, b) << 16); + try cwd.writeFile(io, .{ .sub_path = try std.fmt.bufPrint(&lp, "{s}/{s}.last.f32", .{ out, name }), .data = std.mem.sliceAsBytes(f32s) }); + @memcpy(s.logits.slice(u16, g4.config.vocab), row); + var tokens: std.ArrayList(u32) = .empty; + defer tokens.deinit(gpa); + var tops: std.ArrayList(Top2) = .empty; + defer tops.deinit(gpa); + var pos = ids.value.len; + while (tokens.items.len < count) : (pos += 1) { + const t = top2(s.logits.slice(u16, g4.config.vocab), pos); + try tops.append(gpa, t); + try tokens.append(gpa, t.ids[0]); + if (tokens.items.len == count) break; + try g4.forward.window(m, &s, &cache, &.{t.ids[0]}, pos); + } + const secs = @as(f64, @floatFromInt(std.c.mach_absolute_time() - t1)) / 24e6; + var path: [256]u8 = undefined; + var w: std.Io.Writer.Allocating = .init(gpa); + defer w.deinit(); + try std.json.Stringify.value(tokens.items, .{}, &w.writer); + try cwd.writeFile(io, .{ .sub_path = try std.fmt.bufPrint(&path, "{s}/{s}.tokens.json", .{ out, name }), .data = w.written() }); + w.clearRetainingCapacity(); + try w.writer.writeByte('['); + for (tops.items, 0..) |t, i| { + if (i > 0) try w.writer.writeAll(", "); + try w.writer.print("{{\"ids\": [{d}, {d}], \"logits\": [{d}, {d}], \"position\": {d}, \"dtype\": \"mlx.core.bfloat16\"}}", .{ t.ids[0], t.ids[1], @as(f64, t.logits[0]), @as(f64, t.logits[1]), t.position }); + } + try w.writer.writeByte(']'); + try cwd.writeFile(io, .{ .sub_path = try std.fmt.bufPrint(&path, "{s}/{s}.top2.json", .{ out, name }), .data = w.written() }); + std.debug.print("{s}: {d} prompt tokens in {d:.1} s ({d:.0} tok/s, windows of {d}), {d} decoded, {d:.1} s in all\n", .{ name, ids.value.len, prefill_s, @as(f64, @floatFromInt(ids.value.len)) / prefill_s, width, count, secs }); + } +} diff --git a/zig/tests/gemma4/vision.zig b/zig/tests/gemma4/vision.zig new file mode 100644 index 0000000000..3d36a89923 --- /dev/null +++ b/zig/tests/gemma4/vision.zig @@ -0,0 +1,202 @@ +//! Gemma 4's image input against tools/zig/gemma4_vision_fixtures.py's dumps: patches bit for bit, then each stage. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const g = tf.gemma4; +const vision = g.vision; +const unified = g.vision_unified; +const ck = tf.checkpoint; +const shards = tf.checkpoint_host; +const affine_mm = tf.affine_mm; + +const usage = + \\tf-gemma4-vision MODEL_DIR REF_DIR... + \\Each REF_DIR (one image's dump) holds meta.json, patches.f32 and the stages' references in bf16 and fp32 (NAME.f32, + \\NAME_fp32.f32): embed (the patch rows), block0, blocks, pooled and features (the tower's output rows). Prepares the + \\image natively and compares its patches with patches.f32 (the PIL processor's), then runs each stage on patches.f32 + \\and prints its relative error and worst row cosine against the fp32 run beside torch's own bf16 run's. A stage fails + \\past 1.1 times torch's error on one image (bf16 drift over 27 layers moves either run by that much), and the + \\features fail when their error averaged over the images is past torch's average. +; + +/// Each stage's rows; a null width: the language model's (the projection's output). `from`: where the rows are left. +const Stage = struct { name: []const u8, stop: ?u32, width: ?usize, from: enum { x, pooled, c, d, out } }; +const stages = [_]Stage{ + .{ .name = "embed", .stop = 0, .width = vision.width, .from = .x }, + .{ .name = "block0", .stop = 1, .width = vision.width, .from = .x }, + .{ .name = "blocks", .stop = vision.depth, .width = vision.width, .from = .x }, + .{ .name = "pooled", .stop = null, .width = vision.width, .from = .pooled }, + .{ .name = "features", .stop = null, .width = null, .from = .out }, +}; +/// The 12B's embedder: patch_ln2's rows, pos_norm's, then the projection's. +const unified_stages = [_]Stage{ + .{ .name = "patch", .stop = 0, .width = unified.width, .from = .c }, + .{ .name = "posed", .stop = 1, .width = unified.width, .from = .d }, + .{ .name = "features", .stop = null, .width = null, .from = .out }, +}; + +const Input = union(enum) { + tower: *vision.Vision, + embedder: *unified.Embedder, + + fn encode(in: Input, e: mtl.ComputeEncoder, pixels: []const f32, gh: u32, gw: u32, out: vision.Ref, stop: ?u32) !void { + switch (in) { + inline else => |v| try v.encode(e, pixels, gh, gw, out, stop), + } + } + + fn rows(in: Input, s: Stage, out: mtl.Buffer) vision.Ref { + return switch (s.from) { + .x => in.tower.x, + .pooled => in.tower.pooled, + .c => in.embedder.c, + .d => in.embedder.d, + .out => .{ .buf = out }, + }; + } +}; + +fn readFile(a: std.mem.Allocator, dir: []const u8, name: []const u8) ![]u8 { + const path = try std.fmt.allocPrintSentinel(a, "{s}/{s}", .{ dir, name }, 0); + const f = try mtl.MappedFile.open(path); + defer f.deinit(); + return a.dupe(u8, f.bytes[0..f.size]); +} + +fn floats(a: std.mem.Allocator, dir: []const u8, name: []const u8) ![]f32 { + const raw = try readFile(a, dir, name); + const out = try a.alloc(f32, raw.len / 4); + @memcpy(std.mem.sliceAsBytes(out), raw[0 .. out.len * 4]); + return out; +} + +/// ||got - want|| / ||want|| and the worst row's cosine, rows `width` wide. +fn compare(got: []const f32, want: []const f32, width: usize) [2]f64 { + var d2: f64 = 0; + var w2: f64 = 0; + var worst: f64 = 1; + for (0..want.len / width) |r| { + var dot: f64 = 0; + var ng: f64 = 0; + var nw: f64 = 0; + for (got[r * width ..][0..width], want[r * width ..][0..width]) |x, y| { + d2 += (x - y) * (x - y); + w2 += y * y; + dot += x * y; + ng += x * x; + nw += y * y; + } + worst = @min(worst, dot / (@sqrt(ng * nw) + 1e-30)); + } + return .{ @sqrt(d2 / w2), worst }; +} + +pub fn main(init: std.process.Init) !void { + const a = init.arena.allocator(); + const gpa = std.heap.page_allocator; + const args = try init.minimal.args.toSlice(a); + if (args.len < 3) { + std.debug.print("{s}\n", .{usage}); + return error.Usage; + } + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const device = try mtl.Device.init(); + const queue = try device.queue(); + var c = ck.Checkpoint.init(gpa); + defer c.deinit(); + const files = try shards.shardFiles(gpa, init.io, args[1]); + defer shards.freeShardFiles(gpa, files); + const conf = try g.config.Config.read(gpa, init.io, args[1]); + const geo = conf.geo; + if (geo.unified()) try unified.addTensors(&c, device, files) else try vision.addTensors(&c, device, files); + var mm: affine_mm.Pipes = undefined; + { + const text = try affine_mm.source(device, gpa, .{ .bits = conf.widths.vision.bits, .group = conf.widths.vision.group }); + defer gpa.free(text); + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + for (&mm, affine_mm.names) |*p, name| p.* = try mtl.Pipeline.init(device, lib, name, false); + } + defer for (mm) |p| p.deinit(); + const in: Input = if (geo.unified()) .{ .embedder = try unified.Embedder.load(gpa, device, &c, geo.hidden, mm, conf.widths.vision) } else .{ .tower = try vision.Vision.load(gpa, device, &c, mm, conf.widths.vision, geo.hidden) }; + defer switch (in) { + inline else => |v| v.deinit(), + }; + const list: []const Stage = if (geo.unified()) &unified_stages else &stages; + const layout: g.image.Layout = if (geo.unified()) .merged else .patches; + const out = try device.buffer(vision.max_patches / 9 * geo.hidden * 2, mtl.ResourceOptions.shared); + defer out.deinit(); + const back = try device.buffer(vision.max_patches * vision.width * 2, mtl.ResourceOptions.shared); + defer back.deinit(); + var failed = false; + var sums: [2]f64 = .{ 0, 0 }; // the features' error, native and torch bf16, summed over the images + for (args[2..]) |ref| { + const meta = try std.json.parseFromSliceLeaky(std.json.Value, a, try readFile(a, ref, "meta.json"), .{}); + const image_path = try std.fmt.allocPrintSentinel(a, "{s}", .{meta.object.get("image").?.string}, 0); + const f = try mtl.MappedFile.open(image_path); + defer f.deinit(); + const max_tokens: u32 = @intCast(meta.object.get("max_tokens").?.integer); + const prep = try g.image.prepare(init.io, a, f.bytes[0..f.size], max_tokens, layout); + const want_patches = try floats(a, ref, "patches.f32"); + var differ: usize = 0; + var largest: f32 = 0; + if (want_patches.len != prep.pixels.len) { + std.debug.print("{s}: native grid {d}x{d} has {d} values, the reference {d}\n", .{ ref, prep.gw, prep.gh, prep.pixels.len, want_patches.len }); + failed = true; + continue; + } + for (prep.pixels, want_patches) |x, y| { + if (x != y) differ += 1; + largest = @max(largest, @abs(x - y)); + } + std.debug.print("{s}: grid {d}x{d}, {d} tokens; patches: {d} of {d} values differ (largest {d:.6})\n", .{ std.fs.path.basename(ref), prep.gw, prep.gh, prep.tokens, differ, want_patches.len, largest }); + if (differ != 0) failed = true; + const P = prep.gh * prep.gw; + for (list) |s| { + const width = s.width orelse geo.hidden; + const cb = queue.commandBuffer(); + const e = cb.compute(.serial); + const t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); + try in.encode(e, want_patches, prep.gh, prep.gw, .{ .buf = out }, s.stop); // the reference's input: the tower's error alone + e.end(); + const rows: usize = if (s.from == .x) P else prep.tokens; + const from = in.rows(s, out); + if (from.buf.id != out.id) { + const b = cb.blit(); + b.copy(from.buf, from.off, back, 0, rows * width * 2); + b.end(); + } + cb.commit(); + cb.wait(); + if (cb.failure()) |msg| { + std.debug.print("command buffer failed: {s}\n", .{msg}); + return error.GpuFailed; + } + const ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(init.io).toNanoseconds() - t0)) / 1e6; + const src: [*]const u16 = @ptrCast(@alignCast(if (from.buf.id != out.id) back.contents() else out.contents())); + const got = try a.alloc(f32, rows * width); + for (got, 0..) |*x, i| x.* = @bitCast(@as(u32, src[i]) << 16); + const want = try floats(a, ref, try std.fmt.allocPrint(a, "{s}_fp32.f32", .{s.name})); + const torch = try floats(a, ref, try std.fmt.allocPrint(a, "{s}.f32", .{s.name})); + if (want.len != got.len) { + std.debug.print(" {s}: native {d} values, reference {d}\n", .{ s.name, got.len, want.len }); + failed = true; + continue; + } + const n = compare(got, want, width); + const t = compare(torch, want, width); + const ok = n[0] <= t[0] * 1.1; + if (s.from == .out) { + sums[0] += n[0]; + sums[1] += t[0]; + } + if (!ok) failed = true; + std.debug.print(" {s:<8} vs fp32: native rel {d:.4} worst cos {d:.4} | torch bf16 rel {d:.4} worst cos {d:.4} {s} ({d:.1} ms)\n", .{ s.name, n[0], n[1], t[0], t[1], if (ok) "ok" else "WORSE", ms }); + } + } + const images: f64 = @floatFromInt(args.len - 2); + const mean_ok = sums[0] <= sums[1]; + std.debug.print("features averaged over {d} images: native rel {d:.4} | torch bf16 rel {d:.4} {s}\n", .{ args.len - 2, sums[0] / images, sums[1] / images, if (mean_ok) "ok" else "WORSE" }); + if (failed or !mean_ok) return error.GemmaVisionMismatch; +} diff --git a/zig/tests/glm_rows_batch.zig b/zig/tests/glm_rows_batch.zig index c45406d8d7..dfe9e0cff1 100644 --- a/zig/tests/glm_rows_batch.zig +++ b/zig/tests/glm_rows_batch.zig @@ -27,9 +27,8 @@ fn headsLaunch( h: u32, io: wide.Io, kx: u32, - sums: Ref, ) void { - wide.headsBatch(p, e, q, h, io, kx, sums); + wide.headsBatch(p, e, q, h, io, kx); } /// A chunk's embed_q / unembed_out (wide.headsBatch): the heads in one affine_mm batch == each head's own product. @@ -43,7 +42,6 @@ pub fn headsCheck(g: *const Gpu, mm: affine_mm.Pipes, bits: u8, n: usize, k: usi const r = prng.random(); const bs = [_]mtl.Buffer{ try buf(g.device, m * ldx * 2), - try buf(g.device, m * nh * (k / 64) * 4), try buf(g.device, m * nh * n * 2), try buf(g.device, m * k * 2), try buf(g.device, m * n * 2), @@ -54,24 +52,23 @@ pub fn headsCheck(g: *const Gpu, mm: affine_mm.Pipes, bits: u8, n: usize, k: usi &g.p, w.q, @as(u32, nh), - wide.Io{ .in = ref(bs[0]), .ldx = @intCast(ldx), .out = ref(bs[2]), .ldo = @intCast(nh * n), .rows = m }, + wide.Io{ .in = ref(bs[0]), .ldx = @intCast(ldx), .out = ref(bs[1]), .ldo = @intCast(nh * n), .rows = m }, @as(u32, @intCast(k)), - ref(bs[1]), }); - const all = bs[2].slice(u16, m * nh * n); + const all = bs[1].slice(u16, m * nh * n); for (0..nh) |h| { for (0..m) |i| @memcpy( - bs[3].slice(u16, m * k)[i * k ..][0..k], + bs[2].slice(u16, m * k)[i * k ..][0..k], bs[0].slice(u16, m * ldx)[i * ldx + h * k ..][0..k], ); - try g.run(mmDense, .{ mm, w.part(h, nh), ref(bs[3]), ref(bs[1]), ref(bs[4]), @as(u32, m) }); - for (0..m) |i| try eql(all[(i * nh + h) * n ..][0..n], bs[4].slice(u16, m * n)[i * n ..][0..n]); + try g.run(mmDense, .{ mm, w.part(h, nh), ref(bs[2]), ref(bs[3]), @as(u32, m) }); + for (0..m) |i| try eql(all[(i * nh + h) * n ..][0..n], bs[3].slice(u16, m * n)[i * n ..][0..n]); t.pairs += m * n; } } -fn denseMm(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, io: wide.Io, sums: Ref) void { - wide.denseMm(p, e, q, io, sums); +fn denseMm(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, io: wide.Io) void { + wide.denseMm(p, e, q, io); } /// A chunk's split stack (wide.denseMm) at the 64-padded pitch == each member alone (affine_mm or the row kernels). @@ -101,7 +98,6 @@ pub fn splitMmCheck(g: *const Gpu, k: usize, ns: []const usize, bits: []const u8 const ld = std.mem.alignForward(usize, at, 64); const bs = [_]mtl.Buffer{ try buf(g.device, m * k * 2), - try buf(g.device, m * (k / 64) * 4), try buf(g.device, m * ld * 2), try buf(g.device, m * ld * 2), }; @@ -110,22 +106,21 @@ pub fn splitMmCheck(g: *const Gpu, k: usize, ns: []const usize, bits: []const u8 try g.run(denseMm, .{ &g.p, q, - wide.Io{ .in = ref(bs[0]), .ldx = @intCast(k), .out = ref(bs[2]), .ldo = @intCast(ld), .rows = m }, - ref(bs[1]), + wide.Io{ .in = ref(bs[0]), .ldx = @intCast(k), .out = ref(bs[1]), .ldo = @intCast(ld), .rows = m }, }); var on_mm: usize = 0; for (ms[0..ns.len], 0..) |x, i| { const run = q.run(i); if (wide.mm(&g.p, x.q) != null and run.n % 64 == 0) { - try g.run(mmDense, .{ wide.mm(&g.p, x.q).?, x.q, ref(bs[0]), ref(bs[1]), ref(bs[3]), @as(u32, m) }); + try g.run(mmDense, .{ wide.mm(&g.p, x.q).?, x.q, ref(bs[0]), ref(bs[2]), @as(u32, m) }); on_mm += 1; } else try g.run( dense, - .{ &g.p, x.q, wide.Io{ .in = ref(bs[0]), .ldx = @intCast(k), .out = ref(bs[3]), .ldo = run.n, .rows = m } }, + .{ &g.p, x.q, wide.Io{ .in = ref(bs[0]), .ldx = @intCast(k), .out = ref(bs[2]), .ldo = run.n, .rows = m } }, ); for (0..m) |y| try eql( - bs[2].slice(u16, m * ld)[y * ld + run.row ..][0..run.n], - bs[3].slice(u16, m * ld)[y * run.n ..][0..run.n], + bs[1].slice(u16, m * ld)[y * ld + run.row ..][0..run.n], + bs[2].slice(u16, m * ld)[y * run.n ..][0..run.n], ); t.pairs += m * run.n; } diff --git a/zig/tests/glm_rows_check.zig b/zig/tests/glm_rows_check.zig index 50748a1d84..bcb54c849e 100644 --- a/zig/tests/glm_rows_check.zig +++ b/zig/tests/glm_rows_check.zig @@ -312,23 +312,19 @@ pub fn mmDense( mm: affine_mm.Pipes, q: affine.Affine, x: Ref, - sums: Ref, y: Ref, n_rows: u32, ) void { - affine_mm.rowSums(e, mm, 64, x, sums, n_rows, q.k); - affine_mm.dense(e, mm, x, sums, q, y, n_rows); + affine_mm.dense(e, mm, x, q, y, n_rows); } /// dims: x rows, sorted pairs, experts, picks a row. -fn mmGlu(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, gate: affine.Affine, up: affine.Affine, x: Ref, sums: Ref, sorted: [2]Ref, act: Ref, dims: [4]u32, lim: f32) void { - affine_mm.rowSums(e, mm, 64, x, sums, dims[0], gate.k); - affine_mm.gatherGlu(e, mm, x, sums, gate, up, sorted[0], sorted[1], act, dims[1], dims[2], dims[3], lim); +fn mmGlu(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, gate: affine.Affine, up: affine.Affine, x: Ref, sorted: [2]Ref, act: Ref, dims: [4]u32, lim: f32) void { + affine_mm.gatherGlu(e, mm, x, gate, up, sorted[0], sorted[1], act, dims[1], dims[2], dims[3], lim); } -fn mmDown(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, q: affine.Affine, x: Ref, sums: Ref, sorted: [2]Ref, y: Ref, slots: u32, experts: u32) void { - affine_mm.rowSums(e, mm, 64, x, sums, slots, q.k); - affine_mm.gatherTo(e, mm, x, sums, q, sorted[0], y, sorted[1], slots, experts); +fn mmDown(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, q: affine.Affine, x: Ref, sorted: [2]Ref, y: Ref, slots: u32, experts: u32) void { + affine_mm.gatherTo(e, mm, x, q, sorted[0], y, sorted[1], slots, experts); } /// oQ2's prompt experts, affine_mm at 2/64: gathered GLU == glm_swiglu(dense gate | up), gathered down == dense. @@ -358,14 +354,13 @@ fn gatherCheck(g: *const Gpu, m: Mat, up: Mat, t: *Tally) !void { const sorted = [2]Ref{ ref(ob), ref(ob).at(parts * 4) }; const bs = [_]mtl.Buffer{ try buf(g.device, rows * (k / 64) * 4), try buf(g.device, 2 * parts * rows * r * 2), try buf(g.device, slots * 2 * r * 2), try buf(g.device, 2 * slots * r * 2), try buf(g.device, slots * k * 2), try buf(g.device, (parts + 1) * slots * r * 2) }; defer for (bs) |b| b.deinit(); - const sums = ref(bs[0]); var gate = m.q; gate.n = @intCast(r); var upq = up.q; upq.n = @intCast(r); // every expert's gate and up, dense on every x row: [2][parts][rows, r] for (0..parts) |e| for ([2]Mat{ m, up }, 0..) |w, j| { - try g.run(mmDense, .{ mm, w.part(e, parts), ref(m.x()), sums, ref(bs[1]).at(((j * parts + e) * rows * r) * 2), @as(u32, rows) }); + try g.run(mmDense, .{ mm, w.part(e, parts), ref(m.x()), ref(bs[1]).at(((j * parts + e) * rows * r) * 2), @as(u32, rows) }); }; const dg = bs[1].slice(u16, 2 * parts * rows * r); const gu = bs[2].slice(u16, slots * 2 * r); @@ -375,7 +370,7 @@ fn gatherCheck(g: *const Gpu, m: Mat, up: Mat, t: *Tally) !void { for (0..2) |j| @memcpy(gu[(2 * i + j) * r ..][0..r], dg[((j * parts + e) * rows + xr) * r ..][0..r]); } const act = ref(bs[3]); - try g.run(mmGlu, .{ mm, gate, upq, ref(m.x()), sums, sorted, act, [4]u32{ rows, slots, parts, repeat }, lim }); + try g.run(mmGlu, .{ mm, gate, upq, ref(m.x()), sorted, act, [4]u32{ rows, slots, parts, repeat }, lim }); try g.run(swiglu, .{ g.swiglu, ref(bs[2]), act.at(slots * r * 2), @as(u32, slots), @as(u32, @intCast(r)), lim }); const both = bs[3].slice(u16, 2 * slots * r); try eql(both[0 .. slots * r], both[slots * r ..]); @@ -383,8 +378,8 @@ fn gatherCheck(g: *const Gpu, m: Mat, up: Mat, t: *Tally) !void { const xs = bs[4].slice(u16, slots * k); for (order, 0..) |o, i| @memcpy(xs[i * k ..][0..k], m.x().slice(u16, rows * k)[@as(usize, @intCast(o)) * k ..][0..k]); const y = bs[5].slice(u16, (parts + 1) * slots * r); - try g.run(mmDown, .{ mm, gate, ref(bs[4]), sums, sorted, ref(bs[5]), @as(u32, slots), @as(u32, parts) }); - for (0..parts) |e| try g.run(mmDense, .{ mm, m.part(e, parts), ref(bs[4]), sums, ref(bs[5]).at((1 + e) * slots * r * 2), @as(u32, slots) }); + try g.run(mmDown, .{ mm, gate, ref(bs[4]), sorted, ref(bs[5]), @as(u32, slots), @as(u32, parts) }); + for (0..parts) |e| try g.run(mmDense, .{ mm, m.part(e, parts), ref(bs[4]), ref(bs[5]).at((1 + e) * slots * r * 2), @as(u32, slots) }); for (order, 0..) |o, i| { const e: usize = @intCast(picks[@intCast(o)]); try eql(y[@as(usize, @intCast(o)) * r ..][0..r], y[((1 + e) * slots + i) * r ..][0..r]); diff --git a/zig/tests/glm_rows_classes.zig b/zig/tests/glm_rows_classes.zig index 4b9f99d4f6..2e2ffa1a73 100644 --- a/zig/tests/glm_rows_classes.zig +++ b/zig/tests/glm_rows_classes.zig @@ -62,8 +62,8 @@ pub fn gamma(k: usize) f64 { return @as(f64, @floatFromInt(k / 8 + 64)) * 0x1p-24; } -/// out row i vs fp64 on every column within the fp32 sums' bound; `scale` != 1: a control, columns past it counted. -fn against(q: affine.Affine, w: []const u8, s: []const u16, b: []const u16, x: []const u16, out: []const u16, pitch: usize, i: usize, t: *Tally, scale: f64) !void { +/// out row i vs fp64 within the fp32 sums' bound, weights bf16(fma(s, q, b)) when `rounded`; `scale` != 1: a control. +fn against(q: affine.Affine, w: []const u8, s: []const u16, b: []const u16, x: []const u16, out: []const u16, pitch: usize, i: usize, t: *Tally, scale: f64, rounded: bool) !void { const k: usize = q.k; const g_n = k / q.group; const rb = k * q.bits / 8; @@ -80,7 +80,7 @@ fn against(q: affine.Affine, w: []const u8, s: []const u16, b: []const u16, x: [ const kk = g * q.group + j; const xv = fromBf(x[i * k + kk]); const cv: f64 = @floatFromInt(code(w[c * rb ..][0..rb], kk, q.bits)); - y += xv * (sc * cv + bi); + y += xv * (if (rounded) fromBf(toBf(@mulAdd(f32, @floatCast(sc), @floatCast(cv), @floatCast(bi)))) else sc * cv + bi); dot += @abs(xv * cv); sum += @abs(xv); } @@ -123,11 +123,9 @@ fn rowsLaunch(e: mtl.ComputeEncoder, p: *const wide.Pipes, q: affine.Affine, x: wide.dense(p, e, q, .{ .in = (affine.Ref{ .buf = x }).at(at), .ldx = q.k, .out = .{ .buf = out }, .ldo = q.n, .rows = m }); } -fn mmLaunch(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, q: affine.Affine, x: mtl.Buffer, sums: mtl.Buffer, out: mtl.Buffer, m: u32) void { +fn mmLaunch(e: mtl.ComputeEncoder, mm: affine_mm.Pipes, q: affine.Affine, x: mtl.Buffer, out: mtl.Buffer, m: u32) void { const xr: affine.Ref = .{ .buf = x }; - const sr: affine.Ref = .{ .buf = sums }; - affine_mm.rowSums(e, mm, 64, xr, sr, m, q.k); - affine_mm.dense(e, mm, xr, sr, q, affine.Ref{ .buf = out }, m); + affine_mm.dense(e, mm, xr, q, affine.Ref{ .buf = out }, m); } /// Each class at its widths (g64): decode rows (alone == windowed) and the prompt's product vs fp64, every column. @@ -141,7 +139,7 @@ pub fn check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue) !Tally const np = std.mem.alignForward(usize, n, 64); // the loader's padded rows (affine_mm reads whole 64-row tiles) const w_bytes = np * k * bits / 8; const tables = np * (k / 64); - const bufs = [_]mtl.Buffer{ try device.buffer(w_bytes, opts), try device.buffer(tables * 2, opts), try device.buffer(tables * 2, opts), try device.buffer(prompt_rows * k * 2, opts), try device.buffer(prompt_rows * np * 2, opts), try device.buffer(prompt_rows * np * 2, opts), try device.buffer(prompt_rows * (k / 64) * 4, opts) }; + const bufs = [_]mtl.Buffer{ try device.buffer(w_bytes, opts), try device.buffer(tables * 2, opts), try device.buffer(tables * 2, opts), try device.buffer(prompt_rows * k * 2, opts), try device.buffer(prompt_rows * np * 2, opts), try device.buffer(prompt_rows * np * 2, opts) }; defer for (bufs) |x| x.deinit(); r.bytes(bufs[0].slice(u8, w_bytes)); for (bufs[1].slice(u16, tables)) |*v| v.* = toBf(r.float(f32) * 0.02 + 0.001); @@ -160,17 +158,17 @@ pub fn check(p: *const wide.Pipes, device: mtl.Device, queue: mtl.Queue) !Tally if (!std.mem.eql(u16, bufs[5].slice(u16, n), bufs[4].slice(u16, rows * n)[i * n ..][0..n])) return error.WindowDiffersFromRowAlone; t.pairs += n; } - for ([_]usize{ 0, rows - 1 }) |i| try against(q, w, s, b, x, bufs[4].slice(u16, rows * n), n, i, &t, 1); + for ([_]usize{ 0, rows - 1 }) |i| try against(q, w, s, b, x, bufs[4].slice(u16, rows * n), n, i, &t, 1, false); const dec = t.worst; - try against(q, w, s, b, x, bufs[4].slice(u16, rows * n), n, 0, &t, 1.01); // scales 1% off must be seen + try against(q, w, s, b, x, bufs[4].slice(u16, rows * n), n, 0, &t, 1.01, false); // scales 1% off must be seen if (t.caught == 0) return error.ControlNotCaught; // prompt: the chunk's product at rows 0 and 32; on the row kernels every row's bits equal the decode's const mm: ?affine_mm.Pipes = wide.mm(p, q); // prompt.zig's (2: the routed experts' gathers) const on_mm = mm != null; - if (mm) |pipes| try submit(queue, mmLaunch, .{ pipes, q, bufs[3], bufs[6], bufs[5], prompt_rows }) else try submit(queue, rowsLaunch, .{ p, q, bufs[3], 0, bufs[5], prompt_rows }); + if (mm) |pipes| try submit(queue, mmLaunch, .{ pipes, q, bufs[3], bufs[5], prompt_rows }) else try submit(queue, rowsLaunch, .{ p, q, bufs[3], 0, bufs[5], prompt_rows }); t.worst = 0; const pitch = if (on_mm) np else n; - for ([_]usize{ 0, prompt_rows - 1 }) |i| try against(q, w, s, b, x, bufs[5].slice(u16, prompt_rows * pitch), pitch, i, &t, 1); + for ([_]usize{ 0, prompt_rows - 1 }) |i| try against(q, w, s, b, x, bufs[5].slice(u16, prompt_rows * pitch), pitch, i, &t, 1, on_mm); if (!on_mm and !std.mem.eql(u16, bufs[5].slice(u16, rows * n), bufs[4].slice(u16, rows * n))) return error.PromptDiffersFromDecode; std.debug.print("{s} [{d}, {d}] q{d}: decode worst err/bound {d:.4}, prompt ({s}) {d:.4} at col {d}; max |err| {e:.3} = {e:.3} of max |y|; scales 1% off caught in {d} of {d} columns\n", .{ cl.name, n, k, bits, dec, if (on_mm) "affine_mm" else "row kernels", t.worst, t.worst_col, t.err, t.err / t.ymax, t.caught, n }); all.values += t.values; diff --git a/zig/tests/glm_rows_speed_prompt.zig b/zig/tests/glm_rows_speed_prompt.zig index 0aab170c6e..b57d2f1db7 100644 --- a/zig/tests/glm_rows_speed_prompt.zig +++ b/zig/tests/glm_rows_speed_prompt.zig @@ -28,14 +28,11 @@ fn headsRows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, pool: Pool, fn headsMm(e: mtl.ComputeEncoder, i: usize, mm: tf.affine_mm.Pipes, pool: Pool, io: Io, ldx: u32, kx: u32) void { var head = pool.at(i); head.n /= H; - tf.affine_mm.rowSumsAt(e, mm, 64, rf(io.x), rf(io.z), PM, H * kx, ldx); - tf.affine_mm.denseBatch(e, mm, rf(io.x), rf(io.z), head, rf(io.y), PM, H, .{ + tf.affine_mm.denseBatch(e, mm, rf(io.x), head, rf(io.y), PM, H, .{ .x_row = @intCast(ldx), .y_row = @intCast(H * head.n), - .sums_row = H, .x_batch = @intCast(kx), .y_batch = @intCast(head.n), - .sums_batch = 1, .w_batch = @intCast(head.n), }); } @@ -51,8 +48,7 @@ fn chunkRows(e: mtl.ComputeEncoder, i: usize, p: *const wide.Pipes, pool: Pool, } fn chunkMm(e: mtl.ComputeEncoder, i: usize, mm: tf.affine_mm.Pipes, pool: Pool, io: Io) void { const q = pool.at(i); - tf.affine_mm.rowSums(e, mm, 64, rf(io.x), rf(io.z), PM, q.k); - tf.affine_mm.dense(e, mm, rf(io.x), rf(io.z), q, rf(io.y), PM); + tf.affine_mm.dense(e, mm, rf(io.x), q, rf(io.y), PM); } /// A prompt chunk's products where oQ's paths differ from uniform's, ms a chunk of 2048 rows. diff --git a/zig/tests/lfm2_fp64.zig b/zig/tests/lfm2_fp64.zig index cf4a3f90fb..423cb6077b 100644 --- a/zig/tests/lfm2_fp64.zig +++ b/zig/tests/lfm2_fp64.zig @@ -129,8 +129,7 @@ const Nax = struct { w: row.Weights, fn encode(t: Nax, e: mtl.ComputeEncoder, p: Probe) !void { const pipes = p.m.kernels.naxPipes(t.w.bits) orelse return error.NoTensorPipelines; // the matrix's own width - am.rowSums(e, pipes, 64, Ref{ .buf = p.x }, Ref{ .buf = p.sums }, rows, @intCast(t.w.k)); - am.dense(e, pipes, Ref{ .buf = p.x }, Ref{ .buf = p.sums }, affine(t.w), Ref{ .buf = p.y[1] }, rows); + am.dense(e, pipes, Ref{ .buf = p.x }, affine(t.w), Ref{ .buf = p.y[1] }, rows); } }; @@ -212,15 +211,13 @@ fn glu(p: Probe, x: q.weights.Moe, worse: *usize) !void { tile.encodeSums(e, t, b.gate.k, pr.x, 0, pr.sums, rows); tile.encodeWork(e, t, b.starts, 64, b.work); try tile.encodeGroupedGlu(e, t, g, u, 64, pr.x, pr.sums, b.ident, pr.y[0], rows, b.starts, b.work); - const sums = Ref{ .buf = pr.sums }; const pipes = pr.m.kernels.naxPipes(b.gate.bits) orelse return error.NoTensorPipelines; - am.rowSums(e, pipes, 64, Ref{ .buf = pr.x }, sums, rows, @intCast(b.gate.k)); const offsets = Ref{ .buf = b.starts }; const order = Ref{ .buf = b.ident }; const y = Ref{ .buf = pr.y[1] }; const ga = affine(b.gate); const ua = affine(b.up); - am.gatherGlu(e, pipes, Ref{ .buf = pr.x }, sums, ga, ua, offsets, order, y, rows, 64, 1, 0); + am.gatherGlu(e, pipes, Ref{ .buf = pr.x }, ga, ua, offsets, order, y, rows, 64, 1, 0); try tile.encodeTwinGlu(e, t, g, u, pr.x, pr.y[2], b.zeros, null, rows, 1); } }; diff --git a/zig/tests/mm_bench.zig b/zig/tests/mm_bench.zig index 1ec779f916..8b52040afd 100644 --- a/zig/tests/mm_bench.zig +++ b/zig/tests/mm_bench.zig @@ -5,16 +5,16 @@ const ks = @import("kernel_sources"); const usage = \\tf-mm-bench M K N FILE[:BM:BN:THREADS[:ENTRY]] ... - \\Each FILE takes tf_affine_mm's buffers (x, w, scales, biases, row sums, args, y, strides) and is compiled with - \\TF_BITS 4, TF_GROUP 64 and nax.h inlined; each file's tf_affine_row_sums fills its row sums (timed too). + \\Each FILE takes tf_affine_mm's buffers (x, w, scales, biases, args, y, strides) and is compiled with + \\TF_BITS 4, TF_GROUP 64 and nax.h inlined. \\MM_REPS (dispatches a timed buffer, 16), MM_TRIALS (5), MM_COPIES (weight copies cycled so reads come from DRAM, 4), \\MM_CHECK_ROWS (rows checked against fp64, 16), MM_F32=1 (fp32 outputs). ; const Args = extern struct { rows: i32, n: i32, k: i32, experts: i32 }; -const Strides = extern struct { x_row: i32 = 0, y_row: i32 = 0, sums_row: i32 = 1, x_batch: i32 = 0, y_batch: i32 = 0, sums_batch: i32 = 0, w_batch: i32 = 0, pad: i32 = 0 }; +const Strides = extern struct { x_row: i32 = 0, y_row: i32 = 0, x_batch: i32 = 0, y_batch: i32 = 0, w_batch: i32 = 0, pad: i32 = 0 }; -const Kernel = struct { path: []const u8, bm: u32 = 64, bn: u32 = 64, threads: u32 = 128, entry: []const u8 = "tf_affine_mm", lib: mtl.Library = undefined, pipe: mtl.Pipeline = undefined, sums_pipe: mtl.Pipeline = undefined, sums: mtl.Buffer = undefined, y: mtl.Buffer = undefined, ms: [64]f64 = undefined, sums_ms: f64 = 0 }; +const Kernel = struct { path: []const u8, bm: u32 = 64, bn: u32 = 64, threads: u32 = 128, entry: []const u8 = "tf_affine_mm", lib: mtl.Library = undefined, pipe: mtl.Pipeline = undefined, y: mtl.Buffer = undefined, ms: [64]f64 = undefined }; fn envInt(name: [:0]const u8, default: usize) usize { const v = std.c.getenv(name) orelse return default; @@ -93,33 +93,12 @@ pub fn main(init: std.process.Init) !void { k.lib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); k.pipe = try mtl.Pipeline.init(device, k.lib, k.entry, false); k.y = try device.buffer(@as(usize, M) * N * ob, opts); - k.sums_pipe = try mtl.Pipeline.init(device, k.lib, "tf_affine_row_sums", false); - k.sums = try device.buffer(@as(usize, M) * G * 4, opts); } defer for (kernels) |*k| { - k.sums.deinit(); - k.sums_pipe.deinit(); k.y.deinit(); k.pipe.deinit(); k.lib.deinit(); }; - for (kernels) |*k| { // each kernel's own row sums, timed over `reps` dispatches - const cb = queue.commandBuffer(); - const e = cb.compute(.serial); - for (0..reps) |_| { - e.setPipeline(k.sums_pipe); - e.setBuffer(x, 0, 0); - e.setValue(Args{ .rows = @intCast(M), .n = 0, .k = @intCast(K), .experts = 0 }, 5); - e.setBuffer(k.sums, 0, 6); - e.dispatchThreads(mtl.Size.of(G, M, 1), mtl.Size.of(@min(G, 64), 1, 1)); - } - e.end(); - cb.commit(); - cb.wait(); - if (cb.failure()) |msg| std.debug.print("row sums failed: {s}\n", .{msg}); - k.sums_ms = (cb.gpuEnd() - cb.gpuStart()) * 1e3 / @as(f64, @floatFromInt(reps)); - } - const encode = struct { fn one(e: mtl.ComputeEncoder, k: *const Kernel, xb: mtl.Buffer, wb: mtl.Buffer, sb: mtl.Buffer, bb: mtl.Buffer, copy: usize, m: u32, kk: u32, n: u32) void { const g = kk / 64; @@ -128,7 +107,6 @@ pub fn main(init: std.process.Init) !void { e.setBuffer(wb, copy * @as(usize, n) * kk / 2, 1); e.setBuffer(sb, copy * @as(usize, n) * g * 2, 2); e.setBuffer(bb, copy * @as(usize, n) * g * 2, 3); - e.setBuffer(k.sums, 0, 4); e.setValue(Args{ .rows = @intCast(m), .n = @intCast(n), .k = @intCast(kk), .experts = 0 }, 5); e.setBuffer(k.y, 0, 6); e.setValue(Strides{ .x_row = @intCast(kk), .y_row = @intCast(n) }, 7); @@ -215,6 +193,6 @@ pub fn main(init: std.process.Init) !void { worst = @max(worst, d); } const rms = @sqrt(err2 / @as(f64, @floatFromInt(check_rows * @as(usize, N)))); - std.debug.print("{s:<40} {d:8.3} ms (min {d:.3} max {d:.3}) {d:6.1} TFLOPS, row sums {d:.3} ms; differ from first {d}/{d}; err rms {e:.3} max {e:.3} (of rms {e:.3})\n", .{ k.path, med, ms[0], ms[trials - 1], flops / (med * 1e-3) / 1e12, k.sums_ms, differ, total, rms / ref_rms, worst / ref_rms, ref_rms }); + std.debug.print("{s:<40} {d:8.3} ms (min {d:.3} max {d:.3}) {d:6.1} TFLOPS; differ from first {d}/{d}; err rms {e:.3} max {e:.3} (of rms {e:.3})\n", .{ k.path, med, ms[0], ms[trials - 1], flops / (med * 1e-3) / 1e12, differ, total, rms / ref_rms, worst / ref_rms, ref_rms }); } } diff --git a/zig/tests/tile_qmm_tensor.zig b/zig/tests/tile_qmm_tensor.zig index 3a1c144423..a74cb1b329 100644 --- a/zig/tests/tile_qmm_tensor.zig +++ b/zig/tests/tile_qmm_tensor.zig @@ -70,18 +70,15 @@ pub fn tensor(g: Gpu, bits: u8, seed: u64, a: std.mem.Allocator) !void { s: Synth, e: usize, x: mtl.Buffer, - sums: mtl.Buffer, y: mtl.Buffer, rows: usize, pub fn encode(d: @This(), enc: mtl.ComputeEncoder, _: *const tile.Pipelines) !void { - am.rowSums(enc, d.p, 64, Ref{ .buf = d.x }, Ref{ .buf = d.sums }, @intCast(d.rows), @intCast(d.s.k)); - const rs = Ref{ .buf = d.sums }; - am.dense(enc, d.p, Ref{ .buf = d.x }, rs, d.s.affine(d.e), Ref{ .buf = d.y }, @intCast(d.rows)); + am.dense(enc, d.p, Ref{ .buf = d.x }, d.s.affine(d.e), Ref{ .buf = d.y }, @intCast(d.rows)); } }; - try g.run(Dense{ .p = p, .s = s, .e = 1, .x = s.x, .sums = sums, .y = wide, .rows = max_rows }); + try g.run(Dense{ .p = p, .s = s, .e = 1, .x = s.x, .y = wide, .rows = max_rows }); for ([_]usize{ 1, 17, 64 }) |rows| { - try g.run(Dense{ .p = p, .s = s, .e = 1, .x = s.x, .sums = sums, .y = part, .rows = rows }); + try g.run(Dense{ .p = p, .s = s, .e = 1, .x = s.x, .y = part, .rows = rows }); const alone = part.slice(u16, rows * cols); if (!std.mem.eql(u16, wide.slice(u16, rows * cols), alone)) return error.TensorRowsDiffer; } @@ -136,29 +133,25 @@ pub fn tensor(g: Gpu, bits: u8, seed: u64, a: std.mem.Allocator) !void { p: am.Pipes, s: Synth, x: mtl.Buffer, - sums: mtl.Buffer, offsets: mtl.Buffer, order: mtl.Buffer, y: mtl.Buffer, glu: bool = false, pub fn encode(d: @This(), enc: mtl.ComputeEncoder, _: *const tile.Pipelines) !void { const rows: u32 = tokens * topk; - const kk: u32 = @intCast(d.s.k); - am.rowSums(enc, d.p, 64, Ref{ .buf = d.x }, Ref{ .buf = d.sums }, if (d.glu) tokens else rows, kk); const x = Ref{ .buf = d.x }; - const rs = Ref{ .buf = d.sums }; const ro = Ref{ .buf = d.offsets }; const y = Ref{ .buf = d.y }; const m = d.s.affine(0); const ex: u32 = @intCast(d.s.experts); - if (d.glu) return am.gatherGlu(enc, d.p, x, rs, m, m, ro, Ref{ .buf = d.order }, y, rows, ex, topk, 0); - am.gather(enc, d.p, x, rs, m, ro, y, rows, ex); + if (d.glu) return am.gatherGlu(enc, d.p, x, m, m, ro, Ref{ .buf = d.order }, y, rows, ex, topk, 0); + am.gather(enc, d.p, x, m, ro, y, rows, ex); } }; - var gather = Gather{ .p = p, .s = s, .x = sorted, .sums = sums, .offsets = offsets, .order = slots_buf, .y = out }; + var gather = Gather{ .p = p, .s = s, .x = sorted, .offsets = offsets, .order = slots_buf, .y = out }; try g.run(gather); for (0..experts) |e| { - try g.run(Dense{ .p = p, .s = s, .e = e, .x = s.x, .sums = sums, .y = wide, .rows = tokens }); + try g.run(Dense{ .p = p, .s = s, .e = e, .x = s.x, .y = wide, .rows = tokens }); const last: usize = if (e + 1 < experts) @intCast(starts[e + 1]) else tokens * topk; for (@as(usize, @intCast(starts[e]))..last) |j| { const want = wide.slice(u16, tokens * cols)[(order[j] / topk) * cols ..][0..cols]; @@ -204,5 +197,6 @@ pub fn tensor(g: Gpu, bits: u8, seed: u64, a: std.mem.Allocator) !void { 32 * cols, }); if (glu_ulp > 2 or glu_exact * 100 < tokens * topk * cols * 99) return error.TensorGluDiffers; - if (worst[1] > 1.0 / 128.0 or worst[1] > worst[0]) return error.TensorFarFromFloat64; + // tensor rows take each weight as bf16(fma(s, q, b)), as MLX's quantized matmul does: held to the bound + if (worst[1] > 1.0 / 128.0) return error.TensorFarFromFloat64; }