diff --git a/README.md b/README.md
index 46023ee01a..56aa478ad3 100644
--- a/README.md
+++ b/README.md
@@ -173,7 +173,7 @@ Conversions are explicitly marked: their upstream weights need the linked conver
| GLM-5.3-Flash, FP4-Local | `tensorfold pull nvidia/GLM-5.3-Flash-FP4-Local@76348ddb2be2e35e39ecfc00e75dc07703f4f3d5` | Two GB10s at TP=2, or three at TP=3 over a ring of direct links: @Bizuayeu loaded it as stored and served it at both sizes. 165.3 GiB of weights, 13% less than NVFP4's 190.4 GiB. | [Native CUDA](docs/recipes/glm-5.3-flash-cuda.md), in 1.1.5 (#1050). NVIDIA's mixed-precision checkpoint: NVFP4 routed experts with compressed block scales, MXFP8 and 128×128-block FP8 for most of the rest. As close to BF16 as NVFP4 at its smaller size: at TP=2 its NLL is 1.11% above BF16's against NVFP4's 1.37% ([contributed receipt](docs/receipts/glm-fp4local-80517fc3/README.md) on #1050's head, not rerun on main), a difference the per-document intervals in [XSFP4's card](https://huggingface.co/Bizuayeu/GLM-5.3-Flash-XSFP4#measured) do not resolve. |
| GLM-5.3-Flash, XSFP4 | `tensorfold pull Bizuayeu/GLM-5.3-Flash-XSFP4@88736f916f8939607fde8a4b881dcae749f7551d` | Two GB10s at TP=2, or three at TP=3 over a ring of direct links: on main `8086c205c` @Bizuayeu loaded it in 370.7 s at TP=2 (rank 0's model 96.97 GiB at a 524,288-token window) and in 313.8 s at TP=3 (63.60 GiB at 1,048,576), and both served. 165.7 GiB of weights, 0.4 GiB more than FP4-Local's. | [Native CUDA](docs/recipes/glm-5.3-flash-cuda.md), in 1.1.5 through FP4-Local's path. FP4-Local with its routed experts' NVFP4 blocks fitted again from the BF16 weights, a smoothing folded into them, so it runs on the same kernels. Closer to BF16 than FP4-Local and NVFP4 at TP=2 by all three measures in [its card](https://huggingface.co/Bizuayeu/GLM-5.3-Flash-XSFP4#measured), on #1050's code: top-1 agreement 93.26%, NLL 0.85% above and a KL bound of 0.0431, against FP4-Local's 92.82%, 1.11% and 0.0510. |
| GLM-5.3, EXLR8 `home` | `tensorfold pull kindlingai/glm-5.3-exlr8-k2-k3-k4@f459f5d0272a529403dd9b3fd32f771f4250a0d4 --exlr8 home --node 0/2 --dir ~/models/glm-5.3-exlr8-home` | Two M5 Ultra Macs, 256 GB each; about 160 GB resident per Mac at 8k context / four streams. | [Two-Mac recipe](docs/glm-5.3-exlr8.md). Use `--node 1/2` on the other Mac; Thunderbolt 5 and MCDMA required. No CUDA backend. |
-| LFM2.5-24B-A2B, MLX 4 / 6 / 8-bit | `tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-4bit@667408d494bde191c2cb9dd4f4549827816b9275`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-6bit@1deec53eb49571e21cdcbf82fd5839248f07ed88`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-8bit@67ab417dd02d0c9545802a726df668620eae049d` | One Mac; M4 Pro checked. | [Native operator recipe](docs/recipes/README.md), [release history](CHANGELOG.md). No separate native model recipe yet. |
+| LFM2.5-24B-A2B, MLX 4 / 6 / 8-bit | `tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-4bit@667408d494bde191c2cb9dd4f4549827816b9275`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-6bit@1deec53eb49571e21cdcbf82fd5839248f07ed88`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-8bit@67ab417dd02d0c9545802a726df668620eae049d` | One Mac; M4 Pro checked. | [Native recipe](docs/recipes/lfm2.5-24b-a2b.md), [release history](CHANGELOG.md). |
| Qwen3.6-35B-A3B, MLX with MTP | `tensorfold pull TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP@81169a9bc511a27c1b4eedb77a2cd98ced431847` | One GB10; 20.9 GB checkpoint. | [Qwen3.6](docs/recipes/qwen3.6-moe.md). Native CUDA; the recipe's Metal path belongs to Python. |
| Gemma 4 31B, MLX 4-bit | `tensorfold pull mlx-community/gemma-4-31B-it-4bit@696d436c404745a59f30e4939a658162b0a9e57f`
`tensorfold pull z-lab/gemma-4-31B-it-DFlash@eabd648301ce28583cc14757912e5e0f84e152e1` | One M5 Mac; M5 Ultra checked. M1-M4 refused. | [Native 31B](docs/recipes/gemma-4-native.md). Name the drafter with `--drafter`; images with `--vision`. |
| Gemma 4 12B, converted MLX 4-bit | `hf download google/gemma-4-12B-it --revision 707f0a3b8a3c7ad586ed01e27eafbad8a27dd0f7 --local-dir ~/models/gemma-4-12b-hf`
`hf download google/gemma-4-12B-it-assistant --revision 46d4c6f13f0ac0ad827b915669b8df9b81c64c51 --local-dir ~/models/gemma-4-12b-assistant` | One M5 Mac; converted text weights 6.3 GB, image weights 100 MB, plus assistant/caches. M1-M4 refused. | [Convert and serve](docs/recipes/gemma-4-12b-native.md). Upstream download needs `hf`; conversion needs Python, mlx and mlx-lm. Serving is native. |
diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md
index 0d8241d5c2..52d9014565 100644
--- a/docs/kernel-contract.md
+++ b/docs/kernel-contract.md
@@ -131,6 +131,11 @@ Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), a
formats on the M5's tensor units, where its reg kernel must equal the cooperative one bit for bit), and tile_qmm's
prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles). Gate 5 also holds the rest of
tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10.
+LFM2.5-24B-A2B's shared rounds: the greedy pick (`lfm2_argmax`) and the conv and attention kernels that run every
+stream's rows in one dispatch, each bit-equal to the per-stream kernels it replaces. Four controls edit the source: a tie
+to the highest id, a nonfinite logit not flagged, and every stream reading or writing the first stream's state or cache.
+Their launches are in `lfm2_records.zig` with the M1 and M2 limits from `--offline`; `lfm2_attn_segs` reserves its 128
+threads in its source, as Metal cannot build it for M1 or M2 at another count.
CUDA: the RMS norm rows (16-byte form and the fp32-out tree), the bf16 residual's norm tails (MLX's rounded rows and
the once-rounded ones, each checked on its residual and its output, with four controls each dropping a rounding point), Nemotron's router logits (and their 128-row tiles,
bit-equal to them) and top-k, the latent top-k's decode form (bit-equal to `tf_super_topk`), the latent combine and
diff --git a/docs/recipes/README.md b/docs/recipes/README.md
index 9cbf9293e8..9aabe20db2 100644
--- a/docs/recipes/README.md
+++ b/docs/recipes/README.md
@@ -22,6 +22,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL.
| Qwen3.6-35B-A3B | [MLX with MTP drafts, one-GPU CUDA](qwen3.6-moe.md) |
| Step-5-Preview | [Native Metal on one 256 GB Mac: 2-bit experts, MTP drafts, images](step-5-preview.md) |
| Kolibri-1 | [Native CUDA, block FP8, research checkpoint](kolibri-1.md) |
+| LFM2.5-24B-A2B | [Native Metal, MLX 4-bit, shared rounds](lfm2.5-24b-a2b.md) |
| Spark-X2.6-4B | [Native CUDA](spark-x2.6-4b-cuda.md), [native Metal, the bf16 release, contexts to 1M](spark-x2.6-4b-metal.md) |
## Capability floor
diff --git a/docs/recipes/lfm2.5-24b-a2b.md b/docs/recipes/lfm2.5-24b-a2b.md
new file mode 100644
index 0000000000..7b8b2cdf4b
--- /dev/null
+++ b/docs/recipes/lfm2.5-24b-a2b.md
@@ -0,0 +1,43 @@
+# LFM2.5-24B-A2B, native Metal
+
+This recipe covers [`LiquidAI/LFM2.5-24B-A2B-MLX-4bit`](https://huggingface.co/LiquidAI/LFM2.5-24B-A2B-MLX-4bit)
+at revision `667408d494bde191c2cb9dd4f4549827816b9275`. The numbers below are measured on this 4-bit export only.
+
+```bash
+hf download LiquidAI/LFM2.5-24B-A2B-MLX-4bit \
+ --revision 667408d494bde191c2cb9dd4f4549827816b9275 --local-dir "$HOME/models/lfm2.5-24b-a2b-4bit"
+zig build native -Dcpu=apple_m1
+zig-out/native/bin/tensorfold-native serve "$HOME/models/lfm2.5-24b-a2b-4bit" --port 8080 --parallel 16
+```
+
+A Mac with less memory may not hold 16 streams at the checkpoint's 131,072-token window; the server then names how
+many fit. `--context 40960` holds 16 on a 48 GB M4 Pro.
+
+## Shared rounds
+
+Every stream's tokens go through the lane rounds. A round of several streams is one forward over all their rows.
+In it each layer's short conv (30 layers) and each attention layer's qk, attention and merge (10 layers) run once
+over every row (`lfm2_segments.metal`), from a table of each stream's cache addresses, row offsets and key counts
+(`segments.zig`). Each row's arithmetic is the per-stream kernel's, so a reply equals its solo run. Prompt calls and
+one-stream windows keep the per-stream launches. A greedy row takes the GPU's pick of its logits (`lfm2_argmax`);
+sampled and masked rows draw on the host. The dense twin takes its MMA window only for matrices of 4,096 outputs or
+more (the short conv's input and the head); narrower ones keep the scalar passes.
+
+`zig build test-lfm2` runs the family's tests. `tf-kernel-contract --only lfm2` checks the shared-round kernels and
+the pick, and `tf-lfm2-exact` compares windows and shared rounds with one-row decode on the real weights.
+
+## Measured
+
+Greedy, 256 tokens, 16 public prompts, a different one a stream, 3 reps in each of 3 alternated rounds on an M5 Ultra;
+aggregate tok/s, the range of the rounds' medians. Before is main at 0d24ecb3 (each stream's own launches), after adds
+the shared-round kernels on it. The M4 Pro is not yet measured on a quiet machine.
+
+| Streams | before | after |
+| --- | --- | --- |
+| 1 | 258.5-259.2 | 266.5-267.7 |
+| 4 | 483.5-486.5 | 565.2-566.1 |
+| 8 | 551.5-551.9 | 732.2-735.6 |
+| 16 | 660.7-661.9 | 922.5-923.5 |
+
+Every reply was token for token equal before and after, and equal to its solo run and to `"draft": false`.
+Cold prefill at 2k, 8k and 32k tokens and peak memory did not change.
diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig
index d5ab369dd7..6db7154689 100644
--- a/zig/build/metal_targets.zig
+++ b/zig/build/metal_targets.zig
@@ -455,6 +455,10 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s
});
test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = helpers })).step);
test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = engine })).step);
+ const lfm2_tests = b.addRunArtifact(b.addTest(.{ .root_module = engine, .filters = &.{"lfm2_moe"} }));
+ b.step("test-lfm2", "The LFM2-MoE family's tests: config, draws, state and its shared rounds on this Mac's GPU").dependOn(&lfm2_tests.step);
+ const lfm2_dspark = b.addRunArtifact(b.addTest(.{ .root_module = engine, .filters = &.{"lfm2 dspark:"} }));
+ b.step("test-lfm2-dspark", "LFM2's taps for a drafter: the tap kernel on this Mac's GPU").dependOn(&lfm2_dspark.step);
test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = lanes })).step);
// the fake fabric maps its windows with mmap: the Macs' RDMA, not Windows
if (target.result.os.tag != .windows) test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = fabric })).step);
diff --git a/zig/kernels/metal/lfm2/lfm2.metal b/zig/kernels/metal/lfm2/lfm2.metal
index 0d17e093ae..b2772cf64e 100644
--- a/zig/kernels/metal/lfm2/lfm2.metal
+++ b/zig/kernels/metal/lfm2/lfm2.metal
@@ -105,11 +105,8 @@ inline void lfm2_rms(const thread float* v, const device bfloat* w, float eps, d
}
// A segment's short conv: Bx = B * x, 3 taps over (state [2, D], Bx), y = C * conv; SNAP: rows' states from FIRST on.
-[[kernel]] void lfm2_conv(const device bfloat* bcx [[buffer(0)]], const device bfloat* state [[buffer(1)]],
- const device bfloat* w [[buffer(2)]], device bfloat* y [[buffer(3)]], device bfloat* snap [[buffer(4)]],
- constant uint& first [[buffer(5)]], uint2 at [[thread_position_in_grid]]) {
- const uint ch = at.x, r = at.y;
- if (ch >= D) return;
+inline void lfm2_conv_row(const device bfloat* bcx, const device bfloat* state, const device bfloat* w, device bfloat* y,
+ device bfloat* snap, uint first, uint ch, uint r) {
bfloat bx[3];
for (int tap = 0; tap < 3; ++tap) {
const int src = int(r) + tap - 2;
@@ -127,6 +124,11 @@ inline void lfm2_rms(const thread float* v, const device bfloat* w, float eps, d
snap[(size_t(r - first) * 2) * D + ch] = bx[1];
snap[(size_t(r - first) * 2 + 1) * D + ch] = bx[2];
}
+[[kernel]] void lfm2_conv(const device bfloat* bcx [[buffer(0)]], const device bfloat* state [[buffer(1)]],
+ const device bfloat* w [[buffer(2)]], device bfloat* y [[buffer(3)]], device bfloat* snap [[buffer(4)]],
+ constant uint& first [[buffer(5)]], uint2 at [[thread_position_in_grid]]) {
+ if (at.x < D) lfm2_conv_row(bcx, state, w, y, snap, first, at.x, at.y);
+}
// One head of one row: RMS norm over its 64 values (lane l: l and l + 32), then RoPE on the pair at position `pos`.
inline void lfm2_head(const device bfloat* src, const device bfloat* gain, float eps, float log2_theta, uint pos, uint lane,
@@ -143,11 +145,10 @@ inline void lfm2_head(const device bfloat* src, const device bfloat* gain, float
// A simdgroup a (head, row): queries normed and rotated, keys and values into the [KVH, cap, HD] cache at len + r.
struct lfm2_qk_args { uint len, cap; float eps, log2_theta; };
-[[kernel]] void lfm2_qk(const device bfloat* q [[buffer(0)]], const device bfloat* k [[buffer(1)]], const device bfloat* v [[buffer(2)]],
- const device bfloat* qn [[buffer(3)]], const device bfloat* kn [[buffer(4)]], constant lfm2_qk_args& p [[buffer(5)]],
- device bfloat* queries [[buffer(6)]], device bfloat* keys [[buffer(7)]], device bfloat* values [[buffer(8)]],
- uint2 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) {
- const uint head = tg.x, r = tg.y, pos = p.len + r;
+inline void lfm2_qk_row(const device bfloat* q, const device bfloat* k, const device bfloat* v, const device bfloat* qn,
+ const device bfloat* kn, lfm2_qk_args p, device bfloat* queries, device bfloat* keys,
+ device bfloat* values, uint head, uint r, uint lane) {
+ const uint pos = p.len + r;
bfloat a, b;
if (head < QH) {
lfm2_head(q + (size_t(r) * QH + head) * HD, qn, p.eps, p.log2_theta, pos, lane, a, b);
@@ -164,6 +165,12 @@ struct lfm2_qk_args { uint len, cap; float eps, log2_theta; };
values[at + lane] = v[(size_t(r) * KVH + kh) * HD + lane];
values[at + lane + 32] = v[(size_t(r) * KVH + kh) * HD + lane + 32];
}
+[[kernel]] void lfm2_qk(const device bfloat* q [[buffer(0)]], const device bfloat* k [[buffer(1)]], const device bfloat* v [[buffer(2)]],
+ const device bfloat* qn [[buffer(3)]], const device bfloat* kn [[buffer(4)]], constant lfm2_qk_args& p [[buffer(5)]],
+ device bfloat* queries [[buffer(6)]], device bfloat* keys [[buffer(7)]], device bfloat* values [[buffer(8)]],
+ uint2 tg [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) {
+ lfm2_qk_row(q, k, v, qn, kn, p, queries, keys, values, tg.x, tg.y, lane);
+}
// Keys a chunk: a row's keys split at fixed positions, so prompt tiles and decode rows cut them the same way.
constant constexpr uint PER = 4096;
@@ -226,14 +233,10 @@ inline void lfm2_stage(const thread uint4* pre, threadgroup uint* ks, threadgrou
// Every row's attention partials (one kernel for prompt and decode): a threadgroup a (chunk, kv head, 8 * RB rows).
template // a simdgroup's RB row blocks share each staged K and V fragment, each block's arithmetic its own
-[[kernel]] void lfm2_attn_mma(const device bfloat* queries [[buffer(0)]], const device bfloat* keys [[buffer(1)]],
- const device bfloat* values [[buffer(2)]], constant lfm2_attn_args& p [[buffer(3)]],
- device float* pm [[buffer(4)]], device float* pl [[buffer(5)]], device float* po [[buffer(6)]],
- constant uint& rows [[buffer(7)]], uint3 tg [[threadgroup_position_in_grid]],
- uint tid [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]],
- uint lane [[thread_index_in_simdgroup]]) {
- threadgroup uint ks[2][HD * (KT / 2)];
- threadgroup bfloat vs[2][KB * KP];
+inline void lfm2_attn_rows(const device bfloat* queries, const device bfloat* keys, const device bfloat* values,
+ lfm2_attn_args p, device float* pm, device float* pl, device float* po, uint rows, uint3 tg,
+ uint tid, uint sg, uint lane, threadgroup uint (*ks)[HD * (KT / 2)],
+ threadgroup bfloat (*vs)[KB * KP]) {
// Row tiles first in the grid: the threadgroups in flight walk the same keys together (one chunk, one kv head).
const uint chunk = tg.z / p.subs, sub0 = p.subs == 1 ? 0 : tg.z % SUBS, sub1 = p.subs == 1 ? SUBS : sub0 + 1;
const uint kh = tg.y, r0 = tg.x * 8 * RB, head = kh * (QH / KVH) + sg;
@@ -337,15 +340,23 @@ template // a simdgroup's RB row blocks share each staged K and V frag
for (int d = 0; d < 8; ++d) *(device float2*)(po + at * HD + 8 * d + fn) = O[b][d];
}
}
+template
+[[kernel]] void lfm2_attn_mma(const device bfloat* queries [[buffer(0)]], const device bfloat* keys [[buffer(1)]],
+ const device bfloat* values [[buffer(2)]], constant lfm2_attn_args& p [[buffer(3)]],
+ device float* pm [[buffer(4)]], device float* pl [[buffer(5)]], device float* po [[buffer(6)]],
+ constant uint& rows [[buffer(7)]], uint3 tg [[threadgroup_position_in_grid]],
+ uint tid [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]],
+ uint lane [[thread_index_in_simdgroup]]) {
+ threadgroup uint ks[2][HD * (KT / 2)];
+ threadgroup bfloat vs[2][KB * KP];
+ lfm2_attn_rows(queries, keys, values, p, pm, pl, po, rows, tg, tid, sg, lane, ks, vs);
+}
template [[host_name("lfm2_attn_mma")]] [[kernel]] decltype(lfm2_attn_mma<1>) lfm2_attn_mma<1>;
template [[host_name("lfm2_attn_mma_r2")]] [[kernel]] decltype(lfm2_attn_mma<2>) lfm2_attn_mma<2>;
// The chunks merged in chunk order for (head, row): Y [rows, QH*HD] bf16.
-[[kernel]] void lfm2_attn_merge(const device float* pm [[buffer(0)]], const device float* pl [[buffer(1)]],
- const device float* po [[buffer(2)]], constant lfm2_attn_args& p [[buffer(3)]],
- device bfloat* y [[buffer(4)]], uint2 tg [[threadgroup_position_in_grid]],
- uint lane [[thread_index_in_simdgroup]]) {
- const uint head = tg.x, r = tg.y;
+inline void lfm2_merge_row(const device float* pm, const device float* pl, const device float* po, lfm2_attn_args p,
+ device bfloat* y, uint head, uint r, uint lane) {
const uint count = p.len + r + 1;
const uint per = PER;
const uint chunks = (count + per - 1) / per;
@@ -380,6 +391,12 @@ template [[host_name("lfm2_attn_mma_r2")]] [[kernel]] decltype(lfm2_attn_mma<2>)
out[lane] = static_cast(o0 / l);
out[lane + 32] = static_cast(o1 / l);
}
+[[kernel]] void lfm2_attn_merge(const device float* pm [[buffer(0)]], const device float* pl [[buffer(1)]],
+ const device float* po [[buffer(2)]], constant lfm2_attn_args& p [[buffer(3)]],
+ device bfloat* y [[buffer(4)]], uint2 tg [[threadgroup_position_in_grid]],
+ uint lane [[thread_index_in_simdgroup]]) {
+ lfm2_merge_row(pm, pl, po, p, y, tg.x, tg.y, lane);
+}
// Router logits [rows, E] in fp32: a simdgroup an expert, lane l over inputs 4l..4l+3 of each 128, then simd_sum.
[[kernel]] void lfm2_route_logits(const device bfloat* x [[buffer(0)]], const device float* router [[buffer(1)]],
@@ -494,3 +511,68 @@ constant constexpr uint SORT_CHUNKS = 16; // thread (e, c) counts expert e's pi
}
r[size_t(at.y) * D + at.x] = static_cast(acc);
}
+
+// Row r's greedy pick over `vocab` bf16 logits: the lowest id of the largest, then 1 if any logit is not finite.
+[[kernel]] void lfm2_argmax(const device ushort* logits [[buffer(0)]], constant uint& vocab [[buffer(1)]],
+ device uint* best [[buffer(2)]], uint row [[threadgroup_position_in_grid]],
+ uint t [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]],
+ uint lane [[thread_index_in_simdgroup]]) {
+ threadgroup float tv[32];
+ threadgroup uint ti[32], tn[32];
+ const device ushort* x = logits + size_t(row) * vocab;
+ float bv = -INFINITY;
+ uint bi = 0xffffffffu, bad = 0;
+ // Eight logits a load, ids rising along a thread: a strictly larger value moves the pick (vocab % 8 == 0).
+ for (uint v = 8 * t; v < vocab; v += 8 * 1024) {
+ const uint4 w = *(const device uint4*)(x + v);
+ const uint words[4] = {w.x, w.y, w.z, w.w};
+ for (uint j = 0; j < 8; ++j) {
+ const float f = as_type((j % 2) ? (words[j / 2] & 0xFFFF0000u) : (words[j / 2] << 16));
+ if (!isfinite(f)) bad = 1;
+ else if (f > bv) {
+ bv = f;
+ bi = v + j;
+ }
+ }
+ }
+ for (uint off = 16; off > 0; off >>= 1) {
+ const float ov = simd_shuffle_xor(bv, off);
+ const uint oi = simd_shuffle_xor(bi, off);
+ if (ov > bv || (ov == bv && oi < bi)) {
+ bv = ov;
+ bi = oi;
+ }
+ }
+ bad = simd_max(bad);
+ if (lane == 0) {
+ tv[sg] = bv;
+ ti[sg] = bi;
+ tn[sg] = bad;
+ }
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ if (sg != 0) return;
+ bv = tv[lane];
+ bi = ti[lane];
+ bad = tn[lane];
+ for (uint off = 16; off > 0; off >>= 1) {
+ const float ov = simd_shuffle_xor(bv, off);
+ const uint oi = simd_shuffle_xor(bi, off);
+ if (ov > bv || (ov == bv && oi < bi)) {
+ bv = ov;
+ bi = oi;
+ }
+ }
+ bad = simd_max(bad);
+ if (lane == 0) {
+ best[2 * row] = bi;
+ best[2 * row + 1] = bad;
+ }
+}
+
+// A drafter's tap: H's rows [rows, D] copied into TAPS at column `at.y` of rows `at.x` values apart, 8 a thread.
+[[kernel]] void lfm2_tap(const device bfloat* h [[buffer(0)]], device bfloat* taps [[buffer(1)]], constant uint2& at [[buffer(2)]],
+ uint2 i [[thread_position_in_grid]]) {
+ if (i.x >= D / 8) return;
+ const device vec* from = (const device vec*)(h + size_t(i.y) * D);
+ ((device vec*)(taps + size_t(i.y) * at.x + at.y))[i.x] = from[i.x];
+}
diff --git a/zig/kernels/metal/lfm2/lfm2_segments.metal b/zig/kernels/metal/lfm2/lfm2_segments.metal
new file mode 100644
index 0000000000..09997c53e8
--- /dev/null
+++ b/zig/kernels/metal/lfm2/lfm2_segments.metal
@@ -0,0 +1,66 @@
+
+// A shared round's streams in one dispatch a kernel (after lfm2.metal): each row runs its own stream's arithmetic.
+
+// A stream of a recorded call at one layer: its cache rows, its rows from `base`, its key count and sub-spans.
+struct lfm2_seg {
+ device bfloat* keys;
+ device bfloat* values;
+ device bfloat* conv;
+ uint base, rows, len, cap, subs, pad;
+};
+// Partials a row owns: every stream's at its first row times the most sub-spans, so no two streams' overlap.
+constant constexpr size_t SEG_PART = size_t(QH) * CHUNKS * SUBS;
+
+// lfm2_conv for every row of a recorded call: row y of stream rowseg[y], its state and snapshot rows.
+[[kernel]] void lfm2_conv_segs(const device bfloat* bcx [[buffer(0)]], const device lfm2_seg* segs [[buffer(1)]],
+ const device bfloat* w [[buffer(2)]], device bfloat* y [[buffer(3)]],
+ device bfloat* snap [[buffer(4)]], const device uint* rowseg [[buffer(5)]],
+ uint2 at [[thread_position_in_grid]]) {
+ if (at.x >= D) return;
+ const lfm2_seg s = segs[rowseg[at.y]];
+ lfm2_conv_row(bcx + size_t(s.base) * 3 * D, s.conv, w, y + size_t(s.base) * D, snap + size_t(s.base) * 2 * D, 0, at.x,
+ at.y - s.base);
+}
+
+// lfm2_qk for every row of a recorded call, each stream's keys and values into its own cache.
+[[kernel]] void lfm2_qk_segs(const device bfloat* q [[buffer(0)]], const device bfloat* k [[buffer(1)]],
+ const device bfloat* v [[buffer(2)]], const device bfloat* qn [[buffer(3)]],
+ const device bfloat* kn [[buffer(4)]], constant lfm2_qk_args& a [[buffer(5)]],
+ device bfloat* queries [[buffer(6)]], const device lfm2_seg* segs [[buffer(7)]],
+ const device uint* rowseg [[buffer(8)]], uint2 tg [[threadgroup_position_in_grid]],
+ uint lane [[thread_index_in_simdgroup]]) {
+ const lfm2_seg s = segs[rowseg[tg.y]];
+ const lfm2_qk_args p = {s.len, s.cap, a.eps, a.log2_theta};
+ const size_t b = s.base;
+ lfm2_qk_row(q + b * QH * HD, k + b * KVH * HD, v + b * KVH * HD, qn, kn, p, queries + b * QH * HD, s.keys, s.values,
+ tg.x, tg.y - s.base, lane);
+}
+
+// lfm2_attn_mma over the streams' 8-row tiles: tile tg.x is tiles[tg.x] = (stream << 16) | the stream's tile.
+[[max_total_threads_per_threadgroup(128)]] [[kernel]] void lfm2_attn_segs(const device bfloat* queries [[buffer(0)]], const device lfm2_seg* segs [[buffer(1)]],
+ const device uint* tiles [[buffer(2)]], constant lfm2_attn_args& a [[buffer(3)]],
+ device float* pm [[buffer(4)]], device float* pl [[buffer(5)]],
+ device float* po [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]],
+ uint tid [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]],
+ uint lane [[thread_index_in_simdgroup]]) {
+ threadgroup uint ks[2][HD * (KT / 2)];
+ threadgroup bfloat vs[2][KB * KP];
+ const uint item = tiles[tg.x];
+ const lfm2_seg s = segs[item >> 16];
+ const lfm2_attn_args p = {s.len, s.cap, a.scale, s.subs};
+ const size_t part = size_t(s.base) * SEG_PART;
+ lfm2_attn_rows<1>(queries + size_t(s.base) * QH * HD, s.keys, s.values, p, pm + part, pl + part, po + part * HD,
+ s.rows, uint3(item & 0xFFFFu, tg.y, tg.z), tid, sg, lane, ks, vs);
+}
+
+// lfm2_attn_merge for every row of a recorded call, from its stream's partials.
+[[kernel]] void lfm2_attn_merge_segs(const device float* pm [[buffer(0)]], const device float* pl [[buffer(1)]],
+ const device float* po [[buffer(2)]], constant lfm2_attn_args& a [[buffer(3)]],
+ device bfloat* y [[buffer(4)]], const device lfm2_seg* segs [[buffer(5)]],
+ const device uint* rowseg [[buffer(6)]], uint2 tg [[threadgroup_position_in_grid]],
+ uint lane [[thread_index_in_simdgroup]]) {
+ const lfm2_seg s = segs[rowseg[tg.y]];
+ const lfm2_attn_args p = {s.len, s.cap, a.scale, s.subs};
+ const size_t part = size_t(s.base) * SEG_PART;
+ lfm2_merge_row(pm + part, pl + part, po + part * HD, p, y + size_t(s.base) * QH * HD, tg.x, tg.y - s.base, lane);
+}
diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig
index 51b14b5a40..7355fdea30 100644
--- a/zig/kernels/metal/sources.zig
+++ b/zig/kernels/metal/sources.zig
@@ -19,7 +19,7 @@ pub const gemma4_prompt_attention = @embedFile("gemma4/prompt_attention.metal");
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");
+pub const lfm2 = @embedFile("lfm2/lfm2.metal") ++ @embedFile("lfm2/lfm2_segments.metal");
/// LFM2's prompt attention on the tensor units (nax.h inlined at load, after lfm2.metal's text).
pub const lfm2_attn_nax = @embedFile("lfm2/lfm2_attn_nax.metal");
diff --git a/zig/src/core/contract/lfm2_records.zig b/zig/src/core/contract/lfm2_records.zig
new file mode 100644
index 0000000000..e3f7179595
--- /dev/null
+++ b/zig/src/core/contract/lfm2_records.zig
@@ -0,0 +1,15 @@
+//! LFM2-MoE's shared-round and greedy-pick launches (lfm2.metal, lfm2_segments.metal), as the encoders launch them.
+const std = @import("std");
+const M = @import("limits.zig").MetalLaunch;
+pub const records = [_]M{
+ .{ .name = "metal lfm2 argmax lfm2_argmax", .threads = 1024, .tg_memory = 384, .m1 = 1024, .m2 = 1024 },
+ .{ .name = "metal lfm2 stream conv lfm2_conv_segs", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 },
+ .{ .name = "metal lfm2 stream attention qk lfm2_qk_segs", .threads = 32, .tg_memory = 0, .m1 = 832, .m2 = 832 },
+ .{ .name = "metal lfm2 stream attention lfm2_attn_segs", .threads = 128, .tg_memory = 18688, .reserved = 128 },
+ .{ .name = "metal lfm2 stream attention merge lfm2_attn_merge_segs", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 },
+ .{ .name = "metal lfm2 tap lfm2_tap", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 },
+};
+
+test "the LFM2 launches fit every Apple family" {
+ try std.testing.expectEqual(@as(usize, 0), try @import("registry.zig").checkMetal(std.testing.allocator, &records, true));
+}
diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig
index d8cf0ee6dd..aea7619d0a 100644
--- a/zig/src/core/contract/registry.zig
+++ b/zig/src/core/contract/registry.zig
@@ -566,7 +566,7 @@ test {
_ = paging;
}
/// The Metal collector keeps measured slots and the restored vision declarations.
-pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal;
+pub const all_metal = metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("lfm2_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal;
const learning = @import("registry_learning.zig");
pub const gradient_pending = learning.gradient_pending;
pub const qwen_learning_pending = learning.qwen_learning_pending;
diff --git a/zig/src/core/tile_qmm.zig b/zig/src/core/tile_qmm.zig
index 64bb1a4971..030bd097e1 100644
--- a/zig/src/core/tile_qmm.zig
+++ b/zig/src/core/tile_qmm.zig
@@ -125,6 +125,8 @@ pub const twin_rows = 32;
pub const window_slots = 4 * twin_rows;
/// Rows from which a dense twin call takes the MMA window kernel over the scalar twin.
pub const window_rows = 8;
+/// Outputs from which a matrix takes that window kernel: narrower ones keep the scalar passes' threadgroups.
+pub const window_n = 4096;
/// Window rows from which the by-expert twin (each picked expert read once) beats a threadgroup row a slot.
pub const group_rows = 16;
/// Columns a scalar twin threadgroup (2 simdgroups of 4) and a window threadgroup (2 simdgroups of 8, one MMA block).
@@ -168,7 +170,7 @@ pub fn encodeTwin(
) !void {
const i = try width(m);
if (rows == 0 or rows > twin_rows) return error.NoRows;
- const window = rows >= window_rows;
+ const window = rows >= window_rows and m.n >= window_n;
e.setPipeline(if (window) p.rows[i] else if (rows > 1) p.twin_few[i] else p.twin[i]);
bind(e, m, .{ .k = @intCast(m.k), .n = @intCast(m.n), .r = @intCast(rows) }, x, x_off, y, y_off);
if (window) return e.dispatchGroups(.{ .width = m.n / window_columns }, .{ .width = 64 });
diff --git a/zig/src/families/lfm2_moe/backend.zig b/zig/src/families/lfm2_moe/backend.zig
index 670b9dacee..44427d435a 100644
--- a/zig/src/families/lfm2_moe/backend.zig
+++ b/zig/src/families/lfm2_moe/backend.zig
@@ -23,6 +23,8 @@ pub const Metal = struct {
/// Forward ms by window width, timed at load: the lane core weighs drafts against them.
costs: [st.window_rows]lanes.config.Cost = undefined,
timed: usize = 0,
+ /// Forwards run so far: a stream's tapped rows in the scratch are current while their gen matches.
+ gen: u64 = 0,
pub fn init(gpa: std.mem.Allocator, model: *Model, options: Options) !*Metal {
const out = try gpa.create(Metal);
@@ -88,7 +90,7 @@ pub const Metal = struct {
}
pub fn backend(self: *Metal) be.Backend {
- return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, .first = firstFn, .first_masked = firstMaskedFn, .queue = queueFn, .read = readFn, .verify = verifyFn, .keep = keepFn, .draft = draftFn, .release = releaseFn } };
+ 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 {
@@ -107,8 +109,65 @@ pub const Metal = struct {
return cache;
}
+ /// One forward on the scratch; a stream's tapped rows from earlier forwards are stale after it.
+ fn run(self: *Metal, segments: []const fwd.Segment, tokens: []const u32, record: bool, head: fwd.Head) !void {
+ self.gen += 1;
+ try fwd.run(self.model, &self.scratch, segments, tokens, record, head);
+ }
+
fn prefillFn(ptr: *anyopaque, stream: *lanes.Stream) !void {
+ return cast(ptr).prefillWith(stream, null);
+ }
+
+ /// The prompt pass, `sink` taking each chunk's tapped rows before the chunk's end is kept for reuse.
+ fn prefillFeaturesFn(ptr: *anyopaque, stream: *lanes.Stream, layers: []const u32, sink: be.FeatureSink) !void {
const self = cast(ptr);
+ const t = self.scratch.taps orelse return error.UnpreparedLfmTaps;
+ if (layers.len != t.count) return error.UnpreparedLfmTaps;
+ return self.prefillWith(stream, sink);
+ }
+
+ /// A tapped row's bytes: every tap's hidden row, bf16.
+ fn tapBytes(self: *const Metal) usize {
+ return if (self.scratch.taps) |t| t.count * c.hidden * 2 else 0;
+ }
+
+ /// The drafter's layers: every forward from now on keeps their outputs for each row.
+ fn prepareFn(ptr: *anyopaque, layers: []const u32) !void {
+ const self = cast(ptr);
+ if (layers.len == 0 or layers.len > st.max_taps) return error.BadLfmTaps;
+ for (layers) |l| if (l >= c.layers) return error.BadLfmTaps;
+ if (self.scratch.taps) |t| return if (std.mem.eql(u32, t.layers[0..t.count], layers)) {} else error.BadLfmTaps;
+ var t: st.Taps = .{ .count = @intCast(layers.len), .buf = undefined };
+ @memcpy(t.layers[0..layers.len], layers);
+ t.buf = try self.scratch.memory.alloc(self.scratch.rows * layers.len * c.hidden * 2);
+ self.scratch.taps = t;
+ }
+
+ /// Cache rows [start, start + count) of the stream's latest tapped rows, on the device.
+ fn featuresFn(ptr: *anyopaque, stream: *lanes.Stream, layers: []const u32, start: u64, count: u32) !be.Features {
+ const self = cast(ptr);
+ const t = self.scratch.taps orelse return error.UnpreparedLfmTaps;
+ if (layers.len != t.count) return error.UnpreparedLfmTaps;
+ const cache = try self.cacheOf(stream);
+ const held = cache.tap orelse return error.LfmTapsNotHeld;
+ if (count == 0 or start < held.start or start + count > held.start + held.rows) return error.LfmTapsNotHeld;
+ const width = self.tapBytes();
+ const buf = if (held.gen == 0) cache.last_tap.? else if (held.gen == self.gen) t.buf else return error.StaleLfmTaps;
+ const row = held.row + @as(usize, @intCast(start - held.start));
+ return .{ .buffer = @intFromPtr(buf.id), .offset = row * width, .rows = count, .row_bytes = @intCast(width), .space = .device, .dtype = .bf16 };
+ }
+
+ /// The prompt's last row's taps, copied beside the stream's caches for its first draft (gen 0: kept there).
+ fn keepLastTap(self: *Metal, cache: *st.Cache, row: usize, position: usize) !void {
+ const width = self.tapBytes();
+ if (width == 0) return;
+ if (cache.last_tap == null) cache.last_tap = try cache.memory.alloc(width);
+ @memcpy(cache.last_tap.?.contents()[0..width], (self.scratch.taps.?.buf.contents() + row * width)[0..width]);
+ cache.tap = .{ .start = position, .row = 0, .rows = 1, .gen = 0 };
+ }
+
+ 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.InvalidLfmPrompt;
const capacity = @min(self.options.capacity, ids.len + stream.max_new + st.window_rows);
@@ -121,7 +180,9 @@ pub const Metal = struct {
stream.cached = 0;
if (stream.reuse.saved) |saved| { // a kept state of this prompt's prefix: the pass starts there
const snap: *snapshot.Snap = @ptrCast(@alignCast(saved));
- if (snap.at < ids.len) if (snapshot.restore(cache, snap)) |_| {
+ // a drafting stream resumes only a state kept with its drafter's context
+ const drafter_less = sink != null and snap.draft == null;
+ if (snap.at < ids.len and !drafter_less) if (snapshot.restore(cache, snap)) |_| {
at = snap.at;
} else |_| {};
stream.cached = @intCast(at);
@@ -136,7 +197,13 @@ pub const Metal = struct {
};
while (mark < stream.reuse.marks.len and stream.reuse.marks[mark] <= at) mark += 1;
if (mark < stream.reuse.marks.len) end = @min(end, stream.reuse.marks[mark]);
- try fwd.run(self.model, &self.scratch, &.{.{ .cache = cache, .rows = end - at }}, ids[at..end], false, .last);
+ try self.run(&.{.{ .cache = cache, .rows = end - at }}, ids[at..end], false, .last);
+ if (sink) |k| {
+ const width = self.tapBytes();
+ const f: be.Features = .{ .buffer = @intFromPtr(self.scratch.taps.?.buf.id), .rows = @intCast(end - at), .row_bytes = @intCast(width), .space = .device, .dtype = .bf16 };
+ try k.take(k.ptr, stream, f, at);
+ if (end == ids.len) try self.keepLastTap(cache, end - 1 - at, end - 1);
+ }
at = end;
if (stream.reuse.hook) |hk| if (std.mem.indexOfScalar(u32, stream.reuse.marks, @intCast(at)) != null)
hk.at(hk.ptr, stream, @intCast(at));
@@ -165,7 +232,8 @@ pub const Metal = struct {
.value => |v| v,
.handle => |h| @intCast(h),
};
- try fwd.run(self.model, &self.scratch, &.{.{ .cache = cache, .rows = 1 }}, &.{token}, false, .all);
+ try self.run(&.{.{ .cache = cache, .rows = 1 }}, &.{token}, false, .all);
+ if (self.scratch.taps != null) cache.tap = .{ .start = cache.len - 1, .row = 0, .rows = 1, .gen = self.gen };
return try draw(self.gpa, cache.logits.slice(u16, c.vocab), stream.sampling, position, &.{});
}
@@ -191,13 +259,24 @@ pub const Metal = struct {
@memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens);
total += rows;
}
- try fwd.run(self.model, &self.scratch, segments[0..windows.len], ids[0..total], true, .all);
+ try self.run(segments[0..windows.len], ids[0..total], true, .all);
+ if (self.scratch.taps != null) for (segments[0..windows.len]) |seg| {
+ const last = seg.cache.last.?;
+ seg.cache.tap = .{ .start = last.start, .row = last.base, .rows = last.rows, .gen = self.gen };
+ };
const logits = self.scratch.logits.slice(u16, total * c.vocab);
+ const best = self.scratch.best.slice(u32, 2 * 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.InvalidLfmOutputs;
for (o.sampled, w.positions, 0..) |*t, p, r| {
const mask = if (w.masks.len > 0) w.masks[r * mask_words ..][0..mask_words] else &.{};
+ // A greedy row takes the GPU's pick, the host's draw on the same logits; a nonfinite row errors there.
+ const greedy = if (w.stream.sampling) |s| s.temperature <= 0 else true;
+ if (greedy and mask.len == 0 and best[2 * (base + r) + 1] == 0) {
+ t.* = best[2 * (base + r)];
+ continue;
+ }
t.* = try draw(self.gpa, logits[(base + r) * c.vocab ..][0..c.vocab], w.stream.sampling, p, mask);
}
@memcpy(o.drafts, w.tokens);
@@ -306,6 +385,10 @@ fn factsFor(width: u32, streams: u32, costs: []const lanes.config.Cost) lanes.Mo
};
}
+test {
+ _ = @import("segments.zig"); // a shared round's table, with the family's tests
+}
+
test "after a failed wider-lanes check every forward is one stream's one row, so no window or by-expert twin runs" {
const gpa = std.testing.allocator;
var one = try lanes.Config.init(gpa, factsFor(1, 8, &.{}), st.window_rows, st.window_rows - 1);
diff --git a/zig/src/families/lfm2_moe/dspark_test.zig b/zig/src/families/lfm2_moe/dspark_test.zig
new file mode 100644
index 0000000000..d220b9ea66
--- /dev/null
+++ b/zig/src/families/lfm2_moe/dspark_test.zig
@@ -0,0 +1,77 @@
+//! The DSpark drafter's kernels on the GPU against host references: the target's taps.
+const std = @import("std");
+const mtl = @import("metal");
+const c = @import("config.zig");
+const kernels = @import("kernels.zig");
+
+const shared = mtl.ResourceOptions.shared;
+
+const Rig = struct {
+ device: mtl.Device,
+ queue: mtl.Queue,
+ k: kernels.Kernels,
+
+ fn init() !Rig {
+ const device = mtl.Device.init() catch return error.SkipZigTest;
+ errdefer device.deinit();
+ const queue = try device.queue();
+ errdefer queue.deinit();
+ return .{ .device = device, .queue = queue, .k = try kernels.load(device) };
+ }
+
+ fn deinit(r: *Rig) void {
+ r.k.deinit();
+ r.queue.deinit();
+ r.device.deinit();
+ }
+
+ fn buffer(r: *Rig, comptime T: type, n: usize) !mtl.Buffer {
+ const b = try r.device.buffer(n * @sizeOf(T), shared);
+ @memset(b.slice(u8, n * @sizeOf(T)), 0);
+ return b;
+ }
+
+ fn run(r: *Rig, comptime key: []const u8, bufs: []const mtl.Buffer, value: anytype, at: usize, grid: [2]usize) !void {
+ const cb = r.queue.commandBuffer();
+ const e = cb.compute(.serial);
+ e.setPipeline(r.k.get(key));
+ for (bufs, 0..) |b, i| e.setBuffer(b, 0, if (i >= at) i + 1 else i);
+ e.setValue(value, at);
+ e.dispatchThreads(mtl.Size.of(grid[0], grid[1], 1), mtl.Size.of(@min(grid[0], 256), 1, 1));
+ e.end();
+ cb.commit();
+ cb.wait();
+ if (cb.failure()) |_| return error.GpuFailed;
+ }
+};
+
+/// Round to nearest even, as MLX casts.
+fn toBf16(f: f32) u16 {
+ const b: u32 = @bitCast(f);
+ return @truncate((b + 0x7fff + ((b >> 16) & 1)) >> 16);
+}
+
+fn noise(seed: u64, out: []u16, scale: f32) void {
+ var prng = std.Random.DefaultPrng.init(seed);
+ for (out) |*v| v.* = toBf16((prng.random().float(f32) * 2 - 1) * scale);
+}
+
+test "lfm2 dspark: the tap kernel copies each row's hidden values into its tap's columns, the rest untouched" {
+ var r = try Rig.init();
+ defer r.deinit();
+ const rows = 5;
+ const taps = 3;
+ const h = try r.buffer(u16, rows * c.hidden);
+ defer h.deinit();
+ const out = try r.buffer(u16, rows * taps * c.hidden);
+ defer out.deinit();
+ noise(7, h.slice(u16, rows * c.hidden), 4);
+ try r.run("lfm2_tap", &.{ h, out }, [2]u32{ taps * c.hidden, c.hidden }, 2, .{ c.hidden / 8, rows });
+ const want = h.slice(u16, rows * c.hidden);
+ const got = out.slice(u16, rows * taps * c.hidden);
+ for (0..rows) |row| {
+ try std.testing.expectEqualSlices(u16, want[row * c.hidden ..][0..c.hidden], got[(row * taps + 1) * c.hidden ..][0..c.hidden]);
+ for (got[row * taps * c.hidden ..][0..c.hidden]) |v| try std.testing.expectEqual(@as(u16, 0), v);
+ for (got[(row * taps + 2) * c.hidden ..][0..c.hidden]) |v| try std.testing.expectEqual(@as(u16, 0), v);
+ }
+}
diff --git a/zig/src/families/lfm2_moe/forward.zig b/zig/src/families/lfm2_moe/forward.zig
index 0bbfc3ef02..7b63833192 100644
--- a/zig/src/families/lfm2_moe/forward.zig
+++ b/zig/src/families/lfm2_moe/forward.zig
@@ -8,10 +8,12 @@ const c = @import("config.zig");
const row = @import("../../core/row_projection.zig");
const tile = @import("../../core/tile_qmm.zig");
const am = @import("../../core/affine_mm.zig");
+const segs = @import("segments.zig");
const Buffer = mtl.Buffer;
pub const Segment = struct { cache: *st.Cache, rows: usize };
/// GPU seconds of every call so far (the bench reports a chunk's GPU time beside its wall time).
pub var gpu_seconds: f64 = 0;
+/// `all`: every row's logits and greedy pick (Scratch.best); `last`: the last row's logits alone.
pub const Head = enum { last, all };
/// Rows from which a call takes the tiles (the twin's bits, each weight read once for every row).
const tile_rows = 33;
@@ -29,6 +31,8 @@ const Encoder = struct {
prompt: bool,
/// A recorded window: every row's conv state is kept for a later keep.
record: bool,
+ /// A recorded window of several streams: their conv and attention kernels run once a layer over every row.
+ plan: ?segs.Plan = null,
fn pipe(self: Encoder, comptime key: []const u8) void {
self.e.setPipeline(self.m.kernels.get(key));
@@ -128,8 +132,9 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const
defer pool.pop();
const cb = m.queue.commandBuffer();
const prompt = (!record and head == .last) or total >= tile_rows;
- const enc = Encoder{ .e = cb.compute(.serial), .m = m, .s = s, .prompt = prompt, .record = record };
- const e = enc.e;
+ const e = cb.compute(.serial);
+ const plan = if (segs.batched(record, segments.len)) segs.prepare(e, s, segments) else null;
+ const enc = Encoder{ .e = e, .m = m, .s = s, .prompt = prompt, .record = record, .plan = plan };
const embedding = m.weights.embedding;
e.setPipeline(m.kernels.embed(embedding.bits));
e.setBuffer(s.ids, 0, 0);
@@ -141,6 +146,7 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const
for (m.weights.blocks, 0..) |block, i| {
enc.norm(s, block.operator_norm, i == 0, total, combine);
combine = false;
+ if (i > 0) tap(enc, s, i - 1, total);
switch (block.mixer) {
.conv => |cv| conv(enc, s, cv, i, segments, total),
.attention => |a| attention(enc, s, a, &m.config, segments, i, total),
@@ -157,9 +163,15 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const
}
}
enc.norm(s, m.weights.norm, false, total, combine);
- if (head == .all) enc.projection(embedding, s.x, 0, s.logits, 0, total) else {
- enc.projection(embedding, s.x, (total - 1) * c.hidden * 2, s.logits, 0, 1);
- }
+ tap(enc, s, c.layers - 1, total);
+ if (head == .all) {
+ enc.projection(embedding, s.x, 0, s.logits, 0, total);
+ enc.pipe("lfm2_argmax");
+ e.setBuffer(s.logits, 0, 0);
+ e.setValue(@as(u32, c.vocab), 1);
+ e.setBuffer(s.best, 0, 2);
+ enc.groups(.{ total, 1, 1 }, .{ 1024, 1, 1 });
+ } else enc.projection(embedding, s.x, (total - 1) * c.hidden * 2, s.logits, 0, 1);
e.end();
cb.commit();
cb.wait();
@@ -177,9 +189,32 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const
}
}
+/// Layer `layer`'s output rows (h, just summed by the next norm) into each tap of it, when a drafter reads them.
+fn tap(enc: Encoder, s: *st.Scratch, layer: usize, total: usize) void {
+ const taps = s.taps orelse return;
+ for (taps.layers[0..taps.count], 0..) |l, t| if (l == layer) {
+ enc.pipe("lfm2_tap");
+ enc.e.setBuffer(s.h, 0, 0);
+ enc.e.setBuffer(taps.buf, 0, 1);
+ enc.e.setValue([2]u32{ taps.count * c.hidden, @intCast(t * c.hidden) }, 2);
+ enc.threads(.{ c.hidden / 8, total, 1 }, .{ 256, 1, 1 });
+ };
+}
+
fn conv(enc: Encoder, s: *st.Scratch, cv: wts.Conv, layer: usize, segments: []const Segment, total: usize) void {
const e = enc.e;
enc.projection(cv.in, s.x, 0, s.bcx, 0, total);
+ if (enc.plan != null) {
+ enc.pipe("lfm2_conv_segs");
+ e.setBuffer(s.bcx, 0, 0);
+ e.setBuffer(s.segs, layer * segs.table_bytes, 1);
+ enc.tensor(cv.weight, 2);
+ e.setBuffer(s.y, 0, 3);
+ e.setBuffer(s.snapshots[layer].?, 0, 4);
+ e.setBuffer(s.rowseg, 0, 5);
+ enc.threads(.{ c.hidden, total, 1 }, .{ 256, 1, 1 });
+ return enc.projection(cv.out, s.y, 0, s.r, 0, total);
+ }
var base: usize = 0;
for (segments) |seg| {
enc.pipe("lfm2_conv");
@@ -200,6 +235,10 @@ fn attention(enc: Encoder, s: *st.Scratch, a: wts.Attention, config: *const c.Co
enc.projection(a.q, s.x, 0, s.q, 0, total);
enc.projection(a.k, s.x, 0, s.k, 0, total);
enc.projection(a.v, s.x, 0, s.v, 0, total);
+ if (enc.plan) |plan| {
+ attentionSegs(enc, s, a, config, layer, total, plan);
+ return enc.projection(a.out, s.y, 0, s.r, 0, total);
+ }
var base: usize = 0;
for (segments) |seg| {
const cache = seg.cache.blocks[layer].attention;
@@ -253,6 +292,42 @@ fn attention(enc: Encoder, s: *st.Scratch, a: wts.Attention, config: *const c.Co
enc.projection(a.out, s.y, 0, s.r, 0, total);
}
+/// A recorded call's attention over every stream in one dispatch a kernel: qk, the 8-row tiles, the merge.
+fn attentionSegs(enc: Encoder, s: *st.Scratch, a: wts.Attention, config: *const c.Config, layer: usize, total: usize, plan: segs.Plan) void {
+ const e = enc.e;
+ const table = layer * segs.table_bytes;
+ enc.pipe("lfm2_qk_segs");
+ e.setBuffer(s.q, 0, 0);
+ e.setBuffer(s.k, 0, 1);
+ e.setBuffer(s.v, 0, 2);
+ enc.tensor(a.q_norm, 3);
+ enc.tensor(a.k_norm, 4);
+ e.setBytes(std.mem.asBytes(&QkArgs{ .len = 0, .cap = 0, .eps = c.eps, .log2_theta = std.math.log2(config.theta) }), 5);
+ e.setBuffer(s.queries, 0, 6);
+ e.setBuffer(s.segs, table, 7);
+ e.setBuffer(s.rowseg, 0, 8);
+ enc.groups(.{ c.query_heads + c.kv_heads, total, 1 }, .{ 32, 1, 1 });
+ const args = AttnArgs{ .len = 0, .cap = 0, .scale = 1.0 / @sqrt(@as(f32, c.head_dim)), .subs = 0 };
+ enc.pipe("lfm2_attn_segs");
+ e.setBuffer(s.queries, 0, 0);
+ e.setBuffer(s.segs, table, 1);
+ e.setBuffer(s.tiles, 0, 2);
+ e.setBytes(std.mem.asBytes(&args), 3);
+ e.setBuffer(s.pm, 0, 4);
+ e.setBuffer(s.pl, 0, 5);
+ e.setBuffer(s.po, 0, 6);
+ enc.groups(.{ plan.tiles, c.kv_heads, plan.used }, .{ 128, 1, 1 });
+ enc.pipe("lfm2_attn_merge_segs");
+ e.setBuffer(s.pm, 0, 0);
+ e.setBuffer(s.pl, 0, 1);
+ e.setBuffer(s.po, 0, 2);
+ e.setBytes(std.mem.asBytes(&args), 3);
+ e.setBuffer(s.y, 0, 4);
+ e.setBuffer(s.segs, table, 5);
+ e.setBuffer(s.rowseg, 0, 6);
+ enc.groups(.{ c.query_heads, total, 1 }, .{ 32, 1, 1 });
+}
+
/// True when the picks' combine is left to the next norm (decode rows).
fn moe(enc: Encoder, s: *st.Scratch, x: wts.Moe, scale: f32, total: usize) bool {
const e = enc.e;
diff --git a/zig/src/families/lfm2_moe/kernels.zig b/zig/src/families/lfm2_moe/kernels.zig
index 2cbf18b018..61af45f4c8 100644
--- a/zig/src/families/lfm2_moe/kernels.zig
+++ b/zig/src/families/lfm2_moe/kernels.zig
@@ -22,6 +22,12 @@ const names = [_][:0]const u8{
"lfm2_combine_norm",
"lfm2_sort",
"lfm2_combine_sorted",
+ "lfm2_argmax",
+ "lfm2_conv_segs",
+ "lfm2_qk_segs",
+ "lfm2_attn_segs",
+ "lfm2_attn_merge_segs",
+ "lfm2_tap",
};
pub const Kernels = struct {
diff --git a/zig/src/families/lfm2_moe/segments.zig b/zig/src/families/lfm2_moe/segments.zig
new file mode 100644
index 0000000000..aa2a7b810f
--- /dev/null
+++ b/zig/src/families/lfm2_moe/segments.zig
@@ -0,0 +1,103 @@
+//! A recorded call's streams in one dispatch a kernel and layer: each layer's stream table, the rows' streams, tiles.
+const std = @import("std");
+const mtl = @import("metal");
+const st = @import("state.zig");
+const c = @import("config.zig");
+
+/// lfm2_segments.metal's lfm2_seg: a stream's cache at one layer, its rows from `base`, key count and sub-spans.
+pub const Seg = extern struct { keys: u64 = 0, values: u64 = 0, conv: u64 = 0, base: u32, rows: u32, len: u32, cap: u32, subs: u32, pad: u32 = 0 };
+
+/// Bytes of one layer's table in Scratch.segs.
+pub const table_bytes = st.batch_rows * @sizeOf(Seg);
+
+/// The attention's 8-row tiles over every stream and the grid depth (chunks times sub-spans) the deepest needs.
+pub const Plan = struct { tiles: u32 = 0, used: u32 = 0 };
+
+/// Whether a call takes these kernels: a recorded window of more than one stream.
+pub fn batched(record: bool, streams: usize) bool {
+ return record and streams > 1;
+}
+
+/// The streams' tables, row map and tiles into the scratch, and every cache buffer the kernels reach declared on `e`.
+pub fn prepare(e: anytype, s: *st.Scratch, streams: anytype) Plan {
+ var plan: Plan = .{};
+ const rowseg = s.rowseg.slice(u32, st.batch_rows);
+ const tiles = s.tiles.slice(u32, st.batch_rows);
+ var base: usize = 0;
+ for (streams, 0..) |seg, j| {
+ // A stream of more than split_rows rows folds a chunk's sub-spans in its tiles, as its own call would.
+ const subs: u32 = if (seg.rows > st.split_rows) 1 else st.subs;
+ for (0..(seg.rows + 7) / 8) |t| {
+ tiles[plan.tiles] = @intCast((j << 16) | t);
+ plan.tiles += 1;
+ }
+ const used: u32 = @intCast((seg.cache.len + seg.rows + st.chunk_keys - 1) / st.chunk_keys * subs);
+ plan.used = @max(plan.used, used);
+ @memset(rowseg[base..][0..seg.rows], @intCast(j));
+ for (seg.cache.blocks, 0..) |block, layer| {
+ var entry: Seg = .{ .base = @intCast(base), .rows = @intCast(seg.rows), .len = @intCast(seg.cache.len), .cap = @intCast(seg.cache.capacity), .subs = subs };
+ switch (block) {
+ .conv => |b| entry.conv = b.gpuAddress(),
+ .attention => |a| {
+ entry.keys = a.keys.gpuAddress();
+ entry.values = a.values.gpuAddress();
+ },
+ }
+ s.segs.slice(Seg, c.layers * st.batch_rows)[layer * st.batch_rows + j] = entry;
+ }
+ // The kernels reach the caches through the table: each buffer is declared, in a residency set or not.
+ for (seg.cache.memory.buffers.items) |b| e.useResource(b, mtl.ResourceUsage.read | mtl.ResourceUsage.write);
+ base += seg.rows;
+ }
+ return plan;
+}
+
+/// Buffers a call declares to its encoder (useResource), in order.
+const Declared = struct {
+ ids: [8 * c.layers]?*anyopaque = undefined,
+ count: usize = 0,
+ pub fn useResource(self: *Declared, b: mtl.Buffer, _: usize) void {
+ self.ids[self.count] = @ptrCast(b.id);
+ self.count += 1;
+ }
+ fn has(self: *const Declared, b: mtl.Buffer) bool {
+ return std.mem.indexOfScalar(?*anyopaque, self.ids[0..self.count], @ptrCast(b.id)) != null;
+ }
+};
+
+test "lfm2_moe: a shared round's table reaches only cache buffers it declares, and maps its rows and tiles" {
+ const gpa = std.testing.allocator;
+ const device = mtl.Device.init() catch return error.SkipZigTest;
+ defer device.deinit();
+ var attention: [c.layers]bool = @splat(false);
+ for (&attention, 0..) |*a, i| a.* = i % 4 == 2;
+ var s = try st.Scratch.init(gpa, device, &attention, st.batch_rows);
+ defer s.deinit();
+ var caches: [2]st.Cache = undefined;
+ for (&caches) |*cache| cache.* = try st.Cache.init(gpa, device, &attention, 64);
+ defer for (&caches) |*cache| cache.deinit();
+ caches[0].len = 5;
+ caches[1].len = 40;
+ const streams = [_]struct { cache: *st.Cache, rows: usize }{ .{ .cache = &caches[0], .rows = 9 }, .{ .cache = &caches[1], .rows = 1 } };
+ var declared: Declared = .{};
+ const plan = prepare(&declared, &s, &streams);
+ // Each cache's every buffer: a buffer the kernels reach through the table only is not resident unless declared.
+ var want: usize = 0;
+ for (&caches) |*cache| for (cache.memory.buffers.items) |b| {
+ try std.testing.expect(declared.has(b));
+ want += 1;
+ };
+ try std.testing.expectEqual(want, declared.count);
+ try std.testing.expectEqual(@as(u32, 3), plan.tiles);
+ try std.testing.expectEqual(@as(u32, st.subs), plan.used);
+ try std.testing.expectEqualSlices(u32, &.{ 0, 1, 1 << 16 }, s.tiles.slice(u32, 3));
+ try std.testing.expectEqualSlices(u32, &.{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 1 }, s.rowseg.slice(u32, 10));
+ const table = s.segs.slice(Seg, c.layers * st.batch_rows);
+ for (attention, 0..) |a, layer| for (&caches, 0..) |*cache, j| {
+ const e = table[layer * st.batch_rows + j];
+ try std.testing.expectEqual(@as(u32, if (j == 0) 0 else 9), e.base);
+ try std.testing.expectEqual(@as(u32, if (j == 0) 1 else st.subs), e.subs);
+ try std.testing.expectEqual(@as(u32, @intCast(cache.len)), e.len);
+ if (a) try std.testing.expectEqual(cache.blocks[layer].attention.keys.gpuAddress(), e.keys) else try std.testing.expectEqual(cache.blocks[layer].conv.gpuAddress(), e.conv);
+ };
+}
diff --git a/zig/src/families/lfm2_moe/snapshot.zig b/zig/src/families/lfm2_moe/snapshot.zig
index d1c1a8b3ae..d7289331b7 100644
--- a/zig/src/families/lfm2_moe/snapshot.zig
+++ b/zig/src/families/lfm2_moe/snapshot.zig
@@ -4,7 +4,7 @@ const c = @import("config.zig");
const st = @import("state.zig");
/// One state: conv states by conv layer, then keys and values [kv head][at][head dim] by attention layer.
-pub const Snap = struct { at: u32, bytes: []u8 };
+pub const Snap = struct { at: u32, bytes: []u8, draft: ?[]u8 = null };
/// A row's keys or values for every kv head of one attention layer.
const row_bytes = c.kv_heads * c.head_dim * 2;
@@ -56,6 +56,7 @@ pub fn restore(cache: *st.Cache, snap: *const Snap) !void {
}
pub fn drop(gpa: std.mem.Allocator, snap: *Snap) void {
+ if (snap.draft) |d| gpa.free(d); // the drafter's context at `at`, when one was kept with the state
gpa.free(snap.bytes);
gpa.destroy(snap);
}
diff --git a/zig/src/families/lfm2_moe/state.zig b/zig/src/families/lfm2_moe/state.zig
index 955110fb1e..1bc749ae95 100644
--- a/zig/src/families/lfm2_moe/state.zig
+++ b/zig/src/families/lfm2_moe/state.zig
@@ -51,6 +51,12 @@ pub const Storage = struct {
};
pub const Attention = struct { keys: mtl.Buffer, values: mtl.Buffer };
+/// Layers a drafter reads at most (each layer's output, hidden wide).
+pub const max_taps = 8;
+/// The drafter's layers in its order and every forward row's outputs of them, [rows][count][hidden] bf16.
+pub const Taps = struct { layers: [max_taps]u32 = @splat(0), count: u32 = 0, buf: mtl.Buffer };
+/// Where a stream's latest tapped rows sit in the scratch: cache rows [start, start + rows) from scratch row `row`.
+pub const TapRows = struct { start: usize, row: usize, rows: usize, gen: u64 };
pub const Mixer = union(enum) { conv: mtl.Buffer, attention: Attention };
pub const Cache = struct {
@@ -60,6 +66,9 @@ pub const Cache = struct {
last: ?struct { start: usize, base: usize, rows: usize } = null,
blocks: [c.layers]Mixer,
logits: mtl.Buffer,
+ /// The stream's latest tapped rows in the scratch, and its prompt's last row's taps (kept for the first draft).
+ tap: ?TapRows = null,
+ last_tap: ?mtl.Buffer = null,
pub fn init(gpa: std.mem.Allocator, device: mtl.Device, attention: *const [c.layers]bool, capacity: usize) !Cache {
var out: Cache = .{ .memory = .{ .gpa = gpa, .device = device }, .capacity = capacity, .blocks = undefined, .logits = undefined };
@@ -141,7 +150,15 @@ pub const Scratch = struct {
/// A projection's input rows' group sums for the tiles and their twin (f32 [rows, k / 64]).
sums: mtl.Buffer,
logits: mtl.Buffer,
+ /// Each head row's greedy pick and nonfinite flag (u32 pairs), from a forward whose head takes every row.
+ best: mtl.Buffer,
snapshots: [c.layers]?mtl.Buffer = @splat(null),
+ /// A recorded call's streams at every layer, each row's stream and the attention's row tiles (segments.zig).
+ segs: mtl.Buffer,
+ rowseg: mtl.Buffer,
+ tiles: mtl.Buffer,
+ /// A drafter's taps of every row, once the backend is told its layers (null: none kept).
+ taps: ?Taps = null,
pub fn init(gpa: std.mem.Allocator, device: mtl.Device, attention: *const [c.layers]bool, rows: usize) !Scratch {
var out: Scratch = undefined;
@@ -149,6 +166,7 @@ pub const Scratch = struct {
errdefer out.memory.deinit();
out.rows = rows;
out.snapshots = @splat(null);
+ out.taps = null;
const mem = &out.memory;
const slots = rows * c.topk;
out.ids = try mem.alloc(rows * 4);
@@ -157,8 +175,7 @@ pub const Scratch = struct {
out.bcx = try mem.alloc(rows * 3 * c.hidden * 2);
out.k = try mem.alloc(rows * c.kv_heads * c.head_dim * 2);
out.v = try mem.alloc(rows * c.kv_heads * c.head_dim * 2);
- const slice_rows: usize = @min(rows, attn_slice);
- const partials = @max(slice_rows, split_rows * subs) * c.query_heads * chunks;
+ const partials = partialCount(rows);
out.pm = try mem.alloc(partials * 4);
out.pl = try mem.alloc(partials * 4);
out.po = try mem.alloc(partials * c.head_dim * 4);
@@ -179,6 +196,10 @@ pub const Scratch = struct {
out.sums = try mem.alloc(@max(rows * (c.intermediate / 64), slots * (c.hidden / 64)) * 4);
const logit_rows: usize = @min(rows, batch_rows);
out.logits = try mem.alloc(logit_rows * c.vocab * 2);
+ out.best = try mem.alloc(logit_rows * 2 * 4);
+ out.segs = try mem.alloc(c.layers * batch_rows * @sizeOf(@import("segments.zig").Seg));
+ out.rowseg = try mem.alloc(batch_rows * 4);
+ out.tiles = try mem.alloc(batch_rows * 4);
for (&out.snapshots, attention) |*s, a| if (!a) { // a recorded window's rows, or a prompt call's last row
s.* = try mem.alloc(logit_rows * conv_bytes);
};
@@ -190,6 +211,14 @@ pub const Scratch = struct {
}
};
+/// Attention partials a scratch of `rows` holds: a prompt slice's, and a recorded call's streams kept apart.
+pub fn partialCount(rows: usize) usize {
+ const slice_rows: usize = @min(rows, attn_slice);
+ // usize: @min with a constant narrows its type, and `most * subs` would overflow it.
+ const most: usize = @min(rows, batch_rows);
+ return @max(slice_rows, split_rows * subs, most * subs) * c.query_heads * chunks;
+}
+
/// The served window: --context, else the checkpoint's, never past `max_context`; null when --context is outside it.
pub fn window(checkpoint: usize, requested: ?i64) ?usize {
const most = @min(checkpoint, max_context);
@@ -197,6 +226,13 @@ pub fn window(checkpoint: usize, requested: ?i64) ?usize {
return if (w >= 1 and w <= most) @as(usize, @intCast(w)) else null;
}
+test "a shared round's 32 rows keep their attention partials apart" {
+ // 32 rows of 8 sub-spans each, every row's 32 heads over 32 chunks; a prompt chunk's rows hold more.
+ try std.testing.expectEqual(@as(usize, 32 * subs * c.query_heads * chunks), partialCount(32));
+ try std.testing.expectEqual(@as(usize, split_rows * subs * c.query_heads * chunks), partialCount(1));
+ try std.testing.expectEqual(@as(usize, 2048 * c.query_heads * chunks), partialCount(2048));
+}
+
test "the served window stays inside the attention's 131,072 positions" {
try std.testing.expectEqual(@as(usize, 131072), max_context);
try std.testing.expectEqual(@as(?usize, 131072), window(131072, null));
diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig
index 28e78092f2..5ca1b23d39 100644
--- a/zig/src/tensorfold.zig
+++ b/zig/src/tensorfold.zig
@@ -84,6 +84,7 @@ test {
_ = @import("families/lfm2_moe/config.zig");
_ = @import("families/lfm2_moe/backend.zig");
_ = @import("families/lfm2_moe/state.zig");
+ _ = @import("families/lfm2_moe/dspark_test.zig");
}
pub const lane_projection = @import("core/lane_projection.zig");
diff --git a/zig/tests/contract/metal_lfm2_argmax.zig b/zig/tests/contract/metal_lfm2_argmax.zig
new file mode 100644
index 0000000000..423919447f
--- /dev/null
+++ b/zig/tests/contract/metal_lfm2_argmax.zig
@@ -0,0 +1,125 @@
+//! LFM2's greedy pick of every verify row (lfm2_argmax): the lowest id of a row's largest bf16 logit, and a flag.
+const std = @import("std");
+const mtl = @import("metal");
+const contract = @import("contract");
+const sources = @import("kernel_sources");
+const c = contract.spec;
+const g = contract.gates;
+const rig_ = @import("rig.zig");
+const controls = @import("metal_lfm2_controls.zig");
+
+const vocab = 128000;
+pub const name = "metal lfm2 argmax";
+
+/// rows: a verify's head rows, one to a shared round's 32.
+const shapes = [_]c.Shape{ .{ .rows = 1, .n = vocab, .k = 1 }, .{ .rows = 9, .n = vocab, .k = 1 }, .{ .rows = 32, .n = vocab, .k = 1 } };
+
+fn none(_: c.Shape) usize {
+ return 0;
+}
+
+pub fn spec() c.Spec {
+ return .{ .name = name, .accum = .none, .depth = none, .out = .raw, .served = &shapes, .nan = .{ .stated = "a row with a NaN or an infinity is flagged nonfinite" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "picks compare values: no sums" } };
+}
+
+/// A bf16 logit for element `i`: few distinct values (ties in every row); the nonfinite fill plants NaN and infinities.
+fn value(i: usize, seed: u64, fill: g.Fill) u16 {
+ const h: u64 = (@as(u64, i) *% 0x9e3779b97f4a7c15 +% seed *% 0xbf58476d1ce4e5b9) >> 40;
+ if (fill == .nonfinite and h % 50021 < 3) return ([3]u16{ 0x7fc0, 0x7f80, 0xff80 })[h % 50021];
+ const v = @as(f32, @floatFromInt(h % 61)) / 8 - 4; // -4 to 3.5 in eighths, signed zeros among them
+ return if (v == 0 and (h >> 20) & 1 == 1) 0x8000 else c.toBf16(v);
+}
+
+pub const Probe = struct {
+ a: std.mem.Allocator,
+ rig: rig_.Rig,
+ pipe: mtl.Pipeline,
+ mutation: controls.Mutation,
+ rows: usize = 0,
+ x: []u16 = &.{},
+ args: [2]rig_.Ref = undefined,
+ out: rig_.Ref = undefined,
+
+ pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, mutation: controls.Mutation) !Probe {
+ const text = try controls.source(a, mutation);
+ defer if (mutation != .none) a.free(text);
+ return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, text, "lfm2_argmax"), .mutation = mutation };
+ }
+
+ fn drop(p: *Probe) void {
+ p.a.free(p.x);
+ p.x = &.{};
+ p.rig.reset();
+ }
+
+ pub fn deinit(p: *Probe) void {
+ p.drop();
+ p.rig.deinit();
+ p.pipe.deinit();
+ }
+
+ const vt = g.hooks(Probe);
+ pub fn kernel(p: *Probe) g.Kernel {
+ var s = spec();
+ if (p.mutation != .none) s.name = p.mutation.name();
+ return .{ .ctx = p, .vt = &vt, .spec = s, .control = p.mutation.gate() };
+ }
+
+ pub fn prepare(p: *Probe, s: c.Shape, f: g.Fill, seed: u64) !bool {
+ p.drop();
+ if (s.n != vocab or s.k != 1 or s.rows == 0 or s.rows > 32) return false;
+ p.rows = s.rows;
+ p.x = try p.a.alloc(u16, s.rows * vocab);
+ for (p.x, 0..) |*v, i| v.* = value(i, seed, f);
+ // Row 0's best value sits at two ids: the pick is the lower one.
+ p.x[0 * vocab + 70001] = c.toBf16(5);
+ p.x[0 * vocab + 1234] = c.toBf16(5);
+ p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.x));
+ const n: u32 = vocab;
+ p.args[1] = try p.rig.input(std.mem.asBytes(&n));
+ p.out = try p.rig.output(p.count());
+ return true;
+ }
+
+ pub fn count(p: *Probe) usize {
+ return p.rows * 8;
+ }
+
+ pub fn run(p: *Probe) !void {
+ const pass = p.rig.begin();
+ const e = pass.e;
+ e.setPipeline(p.pipe);
+ e.setBuffer(p.args[0].buf, p.args[0].off, 0);
+ e.setBuffer(p.args[1].buf, p.args[1].off, 1);
+ e.setBuffer(p.out.buf, p.out.off, 2);
+ e.dispatchGroups(mtl.Size.of(p.rows, 1, 1), mtl.Size.of(1024, 1, 1));
+ try p.rig.finish(pass);
+ }
+
+ pub fn reference(_: *Probe, want: []f64, mag: []f64) !void {
+ @memset(want, 0); // the picks check themselves (check)
+ @memset(mag, 0);
+ }
+
+ /// Each row's pick against the host's greedy draw: the first id of the largest value, or the row flagged.
+ pub fn check(p: *Probe, _: g.Gate) !?[]const u8 {
+ const out = p.rig.bytes(p.out, p.count());
+ for (0..p.rows) |r| {
+ const token = std.mem.readInt(u32, out[8 * r ..][0..4], .little);
+ const flag = std.mem.readInt(u32, out[8 * r + 4 ..][0..4], .little);
+ var best: ?usize = null;
+ var bad = false;
+ for (p.x[r * vocab ..][0..vocab], 0..) |w, id| {
+ const v = c.fromBf16(w);
+ if (!std.math.isFinite(v)) bad = true else if (best == null or v > c.fromBf16(p.x[r * vocab + best.?])) best = id;
+ }
+ if (bad) {
+ if (flag != 1) return "a row with a nonfinite logit is not flagged";
+ continue;
+ }
+ if (flag != 0) return "a finite row is flagged nonfinite";
+ if (token != best.?) return "a pick differs from the first id of the row's largest logit";
+ }
+ return null;
+ }
+};
diff --git a/zig/tests/contract/metal_lfm2_controls.zig b/zig/tests/contract/metal_lfm2_controls.zig
new file mode 100644
index 0000000000..4d875b0fe8
--- /dev/null
+++ b/zig/tests/contract/metal_lfm2_controls.zig
@@ -0,0 +1,57 @@
+//! Planted controls for LFM2's shared-round kernels: lfm2.metal's text with one fault, which one gate must report.
+const std = @import("std");
+const contract = @import("contract");
+const sources = @import("kernel_sources");
+
+pub const Mutation = enum {
+ none,
+ argmax_last_tie,
+ argmax_no_flag,
+ conv_first_state,
+ attention_first_cache,
+ tap_packed_rows,
+ tap_clobbers_column,
+
+ pub fn name(m: Mutation) []const u8 {
+ return switch (m) {
+ .none => "",
+ .argmax_last_tie => "control: metal lfm2 argmax, a tie goes to the highest id",
+ .argmax_no_flag => "control: metal lfm2 argmax, a nonfinite logit is not flagged",
+ .conv_first_state => "control: metal lfm2 stream conv, every stream reads the first stream's state",
+ .attention_first_cache => "control: metal lfm2 stream attention, every stream writes the first stream's cache",
+ .tap_packed_rows => "control: metal lfm2 tap copy, rows packed at the hidden size, not the taps' row",
+ .tap_clobbers_column => "control: metal lfm2 tap copy, a column of another tap is overwritten",
+ };
+ }
+
+ /// The gate that must go red.
+ pub fn gate(m: Mutation) ?contract.gates.Gate {
+ return switch (m) {
+ .none => null,
+ .argmax_no_flag => .nonfinite,
+ else => .shapes,
+ };
+ }
+
+ const Edit = struct { from: []const u8, to: []const u8 };
+
+ fn edit(m: Mutation) Edit {
+ return switch (m) {
+ .none => unreachable,
+ .argmax_last_tie => .{ .from = "(ov == bv && oi < bi)", .to = "(ov == bv && oi > bi)" },
+ .argmax_no_flag => .{ .from = "if (!isfinite(f)) bad = 1;", .to = "if (false) bad = 1;" },
+ .conv_first_state => .{ .from = "lfm2_conv_row(bcx + size_t(s.base) * 3 * D, s.conv,", .to = "lfm2_conv_row(bcx + size_t(s.base) * 3 * D, segs[0].conv," },
+ .attention_first_cache => .{ .from = "queries + b * QH * HD, s.keys, s.values,", .to = "queries + b * QH * HD, segs[0].keys, segs[0].values," },
+ .tap_packed_rows => .{ .from = "(taps + size_t(i.y) * at.x + at.y))[i.x]", .to = "(taps + size_t(i.y) * D + at.y))[i.x]" },
+ .tap_clobbers_column => .{ .from = "[i.x] = from[i.x];\n}\n", .to = "[i.x] = from[i.x];\n *(device bfloat*)(taps + size_t(i.y) * at.x) = 0;\n}\n" },
+ };
+ }
+};
+
+/// lfm2.metal's text as built for the kernel under test; a control's one edit must change it, or it proves nothing.
+pub fn source(a: std.mem.Allocator, m: Mutation) ![]const u8 {
+ if (m == .none) return sources.lfm2;
+ const e = m.edit();
+ if (std.mem.indexOf(u8, sources.lfm2, e.from) == null) return error.ControlEditMissing;
+ return std.mem.replaceOwned(u8, a, sources.lfm2, e.from, e.to);
+}
diff --git a/zig/tests/contract/metal_lfm2_streams.zig b/zig/tests/contract/metal_lfm2_streams.zig
new file mode 100644
index 0000000000..e9a59ef4f8
--- /dev/null
+++ b/zig/tests/contract/metal_lfm2_streams.zig
@@ -0,0 +1,478 @@
+//! LFM2's shared-round kernels (lfm2_segments.metal): every stream's conv or attention rows in one dispatch a kernel.
+const std = @import("std");
+const mtl = @import("metal");
+const contract = @import("contract");
+const sources = @import("kernel_sources");
+const c = contract.spec;
+const g = contract.gates;
+const rig_ = @import("rig.zig");
+const controls = @import("metal_lfm2_controls.zig");
+
+const d = 2048;
+const qh = 32;
+const kvh = 8;
+const hd = 64;
+const subs = 8; // sub-spans a 4,096-key chunk for a stream of at most 8 rows (state.subs)
+const chunks = 32;
+const chunk_keys = 4096;
+const most = 32; // a shared round's rows (state.batch_rows)
+const part = qh * chunks * subs; // partials a row owns (the kernels' SEG_PART)
+pub const Which = enum { attention, conv };
+const names = [_][]const u8{ "metal lfm2 stream attention", "metal lfm2 stream conv" };
+
+/// rows: a round's rows over its streams (sizes from `sizes` in turn); k: the first stream's cached keys.
+const attn_shapes = [_]c.Shape{ .{ .rows = 2, .n = qh * hd, .k = 5 }, .{ .rows = 13, .n = qh * hd, .k = 300 }, .{ .rows = 32, .n = qh * hd, .k = 4100 } };
+const conv_shapes = [_]c.Shape{ .{ .rows = 2, .n = d, .k = 3 }, .{ .rows = 13, .n = d, .k = 3 }, .{ .rows = 32, .n = d, .k = 3 } };
+const sizes = [_]usize{ 1, 9, 3, 16, 2, 8, 5 };
+
+fn none(_: c.Shape) usize {
+ return 0;
+}
+
+pub fn spec(w: Which) c.Spec {
+ return switch (w) {
+ .attention => .{ .name = names[0], .accum = .none, .depth = none, .out = .bf16, .served = &attn_shapes, .replaces = "lfm2_qk, lfm2_attn_mma and lfm2_attn_merge a stream", .nan = .{ .stated = "a NaN key of one stream's kv head makes that stream's heads of it NaN, every other output stays finite" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "the per-stream kernels' bits; the check bounds outputs by float64 attention" } },
+ .conv => .{ .name = names[1], .accum = .none, .depth = none, .out = .bf16, .served = &conv_shapes, .replaces = "lfm2_conv a stream", .nan = .{ .stated = "outputs equal the host's replay of the kernel's bf16 steps, NaN for NaN" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "three fp32 products a channel, replayed exactly on the host" } },
+ };
+}
+
+/// lfm2_segments.metal's lfm2_seg.
+const Seg = extern struct { keys: u64 = 0, values: u64 = 0, conv: u64 = 0, base: u32, rows: u32, len: u32, cap: u32, subs: u32, pad: u32 = 0 };
+
+fn value(i: usize, seed: u64, scale: f32) u16 {
+ const h: u64 = (@as(u64, i) *% 0x9e3779b97f4a7c15 +% seed *% 0xbf58476d1ce4e5b9) >> 40;
+ return c.toBf16((@as(f32, @floatFromInt(h % 2048)) / 1024 - 1) * scale);
+}
+
+pub const Probe = struct {
+ a: std.mem.Allocator,
+ rig: rig_.Rig,
+ which: Which,
+ mutation: controls.Mutation,
+ pipes: [8]mtl.Pipeline,
+ made: usize = 0,
+ rows: usize = 0,
+ nan: bool = false,
+ first_len: usize = 0,
+ streams: [most]usize = undefined, // rows a stream
+ count_streams: usize = 0,
+ host: []u16 = &.{}, // attention: the cached keys and values as prepared; conv: bcx, then the states
+ weights: []u16 = &.{},
+ args: [6]rig_.Ref = undefined,
+ caches: [2][most]rig_.Ref = undefined, // attention: each stream's [keys | values], batched and per stream
+ y: rig_.Ref = undefined,
+ queries: rig_.Ref = undefined,
+ work: [9]rig_.Ref = undefined, // segs, rowseg, tiles, pm, pl, po; the per-stream queries, y and partials' base
+ golden: []u8 = &.{},
+ snap: rig_.Ref = undefined,
+
+ pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which, mutation: controls.Mutation) !Probe {
+ const list: []const [:0]const u8 = switch (w) {
+ .attention => &.{ "lfm2_qk_segs", "lfm2_attn_segs", "lfm2_attn_merge_segs", "lfm2_qk", "lfm2_attn_mma", "lfm2_attn_mma_r2", "lfm2_attn_merge" },
+ .conv => &.{ "lfm2_conv_segs", "lfm2_conv" },
+ };
+ var p: Probe = .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = w, .mutation = mutation, .pipes = undefined };
+ const text = try controls.source(a, mutation);
+ defer if (mutation != .none) a.free(text);
+ errdefer for (p.pipes[0..p.made]) |x| x.deinit();
+ for (list) |n| {
+ p.pipes[p.made] = try rig_.pipeline(device, text, n);
+ p.made += 1;
+ }
+ return p;
+ }
+
+ fn drop(p: *Probe) void {
+ p.a.free(p.golden);
+ p.golden = &.{};
+ p.a.free(p.host);
+ p.a.free(p.weights);
+ p.host = &.{};
+ p.weights = &.{};
+ p.rig.reset();
+ }
+
+ pub fn deinit(p: *Probe) void {
+ p.drop();
+ p.rig.deinit();
+ for (p.pipes[0..p.made]) |x| x.deinit();
+ }
+
+ const vt = g.hooks(Probe);
+ pub fn kernel(p: *Probe) g.Kernel {
+ var s = spec(p.which);
+ if (p.mutation != .none) s.name = p.mutation.name();
+ return .{ .ctx = p, .vt = &vt, .spec = s, .control = p.mutation.gate() };
+ }
+
+ fn len(p: *const Probe, j: usize) usize {
+ return p.first_len + 37 * j;
+ }
+
+ fn cap(p: *const Probe, j: usize) usize {
+ return (p.len(j) + p.streams[j] + 63) / 64 * 64 + 32;
+ }
+
+ fn base(p: *const Probe, j: usize) usize {
+ var b: usize = 0;
+ for (p.streams[0..j]) |r| b += r;
+ return b;
+ }
+
+ pub fn prepare(p: *Probe, s: c.Shape, f: g.Fill, seed: u64) !bool {
+ if (s.rows < 2 or s.rows > most or s.n != @as(usize, if (p.which == .attention) qh * hd else d)) return false;
+ if (p.which == .attention and (s.k == 0 or s.k + 37 * most + most > chunks * chunk_keys)) return false;
+ p.drop();
+ p.rows = s.rows;
+ p.first_len = s.k;
+ p.nan = f == .nonfinite;
+ p.count_streams = 0;
+ var left = s.rows;
+ while (left > 0) : (p.count_streams += 1) {
+ const r = @min(left, sizes[p.count_streams % sizes.len]);
+ p.streams[p.count_streams] = r;
+ left -= r;
+ }
+ return if (p.which == .attention) p.prepareAttention(f, seed) else p.prepareConv(f, seed);
+ }
+
+ fn prepareAttention(p: *Probe, f: g.Fill, seed: u64) !bool {
+ const n = p.rows;
+ var q = try p.a.alloc(u16, n * (qh + 2 * kvh) * hd);
+ defer p.a.free(q);
+ for (q, 0..) |*x, i| x.* = value(i, seed, if (f == .stress) 4 else 1);
+ p.weights = try p.a.alloc(u16, 2 * hd);
+ for (p.weights, 0..) |*x, i| x.* = value(i, seed +% 9, 1.5);
+ p.args[0] = try p.rig.input(std.mem.sliceAsBytes(q[0 .. n * qh * hd]));
+ p.args[1] = try p.rig.input(std.mem.sliceAsBytes(q[n * qh * hd ..][0 .. n * kvh * hd]));
+ p.args[2] = try p.rig.input(std.mem.sliceAsBytes(q[n * (qh + kvh) * hd ..][0 .. n * kvh * hd]));
+ p.args[3] = try p.rig.input(std.mem.sliceAsBytes(p.weights));
+ var total: usize = 0;
+ for (0..p.count_streams) |j| total += 2 * kvh * p.cap(j) * hd;
+ p.host = try p.a.alloc(u16, total);
+ var at: usize = 0;
+ for (0..p.count_streams) |j| {
+ const kv = p.host[at..][0 .. 2 * kvh * p.cap(j) * hd];
+ for (kv, 0..) |*x, i| x.* = value(i, seed +% 3 +% j, 1);
+ if (p.nan and j == 0) kv[(0 * p.cap(j) + p.len(j) / 2) * hd + 7] = 0x7fc0; // stream 0, kv head 0
+ for (0..2) |copy| p.caches[copy][j] = try p.rig.scratch(kv.len * 2);
+ at += kv.len;
+ }
+ p.y = try p.rig.output(n * d * 2);
+ p.queries = try p.rig.scratch(n * d * 2);
+ p.work[0] = try p.rig.scratch(most * @sizeOf(Seg));
+ p.work[1] = try p.rig.scratch(most * 4);
+ p.work[2] = try p.rig.scratch(most * 4);
+ p.work[3] = try p.rig.scratch(n * part * 4);
+ p.work[4] = try p.rig.scratch(n * part * 4);
+ p.work[5] = try p.rig.scratch(n * part * hd * 4);
+ p.work[6] = try p.rig.scratch(n * d * 2);
+ p.work[7] = try p.rig.scratch(n * d * 2);
+ return true;
+ }
+
+ fn prepareConv(p: *Probe, f: g.Fill, seed: u64) !bool {
+ const n = p.rows;
+ p.host = try p.a.alloc(u16, n * 3 * d + p.count_streams * 2 * d);
+ for (p.host, 0..) |*x, i| x.* = value(i, seed, if (f == .stress) 8 else 1);
+ if (p.nan) {
+ p.host[1 * 3 * d + 11] = 0x7fc0; // row 1's B at channel 11
+ p.host[n * 3 * d + 2 * d + 5] = 0x7f80; // stream 1's state at channel 5
+ }
+ p.weights = try p.a.alloc(u16, d * 3);
+ for (p.weights, 0..) |*x, i| x.* = value(i, seed +% 5, 1);
+ p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.host[0 .. n * 3 * d]));
+ p.args[1] = try p.rig.input(std.mem.sliceAsBytes(p.weights));
+ for (0..p.count_streams) |j| p.caches[0][j] = try p.rig.input(std.mem.sliceAsBytes(p.host[n * 3 * d + j * 2 * d ..][0 .. 2 * d]));
+ p.y = try p.rig.output(n * d * 2);
+ p.snap = try p.rig.output(n * 2 * d * 2);
+ return true;
+ }
+
+ pub fn count(p: *Probe) usize {
+ return p.rows * d * @as(usize, if (p.which == .conv) 3 else 1);
+ }
+
+ /// Each stream's cache as prepared (its cached keys before the call's).
+ fn fill(p: *Probe, copy: usize) void {
+ var at: usize = 0;
+ for (0..p.count_streams) |j| {
+ const n = 2 * kvh * p.cap(j) * hd;
+ @memcpy(std.mem.bytesAsSlice(u16, p.rig.bytes(p.caches[copy][j], n * 2)), p.host[at..][0..n]);
+ at += n;
+ }
+ }
+
+ fn table(p: *Probe, copy: usize) ![]Seg {
+ const t = try p.a.alloc(Seg, most);
+ for (0..p.count_streams) |j| {
+ const r = p.caches[copy][j];
+ const addr = r.buf.gpuAddress() + r.off;
+ const wide = p.streams[j] > 8;
+ t[j] = .{ .base = @intCast(p.base(j)), .rows = @intCast(p.streams[j]), .len = @intCast(p.len(j)), .cap = @intCast(p.cap(j)), .subs = if (wide) 1 else subs };
+ if (p.which == .attention) {
+ t[j].keys = addr;
+ t[j].values = addr + kvh * p.cap(j) * hd * 2;
+ } else t[j].conv = addr;
+ }
+ return t;
+ }
+
+ pub fn run(p: *Probe) !void {
+ if (p.which == .conv) return p.runConv(p.y, p.snap, false);
+ p.fill(0);
+ const n = p.rows;
+ const t = try p.table(0);
+ defer p.a.free(t);
+ var rowseg: [most]u32 = undefined;
+ var tiles: [most]u32 = undefined;
+ var count_tiles: u32 = 0;
+ var used: u32 = 0;
+ for (0..p.count_streams) |j| {
+ @memset(rowseg[p.base(j)..][0..p.streams[j]], @intCast(j));
+ for (0..(p.streams[j] + 7) / 8) |i| {
+ tiles[count_tiles] = @intCast((j << 16) | i);
+ count_tiles += 1;
+ }
+ const s: usize = if (p.streams[j] > 8) 1 else subs;
+ used = @max(used, @as(u32, @intCast((p.len(j) + p.streams[j] + chunk_keys - 1) / chunk_keys * s)));
+ }
+ const segs = p.work[0];
+ @memcpy(p.rig.bytes(segs, most * @sizeOf(Seg)), std.mem.sliceAsBytes(t));
+ const rs = p.work[1];
+ @memcpy(p.rig.bytes(rs, most * 4), std.mem.asBytes(&rowseg));
+ const ts = p.work[2];
+ @memcpy(p.rig.bytes(ts, most * 4), std.mem.asBytes(&tiles));
+ const queries = p.queries;
+ const pm = p.work[3];
+ const pl = p.work[4];
+ const po = p.work[5];
+ const pass = p.rig.begin();
+ const e = pass.e;
+ for (p.caches[0][0..p.count_streams]) |r| e.useResource(r.buf, mtl.ResourceUsage.read | mtl.ResourceUsage.write);
+ e.setPipeline(p.pipes[0]);
+ for (p.args[0..3], 0..) |r, i| e.setBuffer(r.buf, r.off, i);
+ e.setBuffer(p.args[3].buf, p.args[3].off, 3);
+ e.setBuffer(p.args[3].buf, p.args[3].off + hd * 2, 4);
+ const qk = [4]u32{ 0, 0, @bitCast(@as(f32, 1e-5)), @bitCast(std.math.log2(@as(f32, 1e6))) };
+ e.setBytes(std.mem.asBytes(&qk), 5);
+ e.setBuffer(queries.buf, queries.off, 6);
+ e.setBuffer(segs.buf, segs.off, 7);
+ e.setBuffer(rs.buf, rs.off, 8);
+ e.dispatchGroups(mtl.Size.of(qh + kvh, n, 1), mtl.Size.of(32, 1, 1));
+ const args = [4]u32{ 0, 0, @bitCast(1.0 / @sqrt(@as(f32, hd))), 0 };
+ e.setPipeline(p.pipes[1]);
+ e.setBuffer(queries.buf, queries.off, 0);
+ e.setBuffer(segs.buf, segs.off, 1);
+ e.setBuffer(ts.buf, ts.off, 2);
+ e.setBytes(std.mem.asBytes(&args), 3);
+ for ([_]rig_.Ref{ pm, pl, po }, 4..) |r, i| e.setBuffer(r.buf, r.off, i);
+ e.dispatchGroups(mtl.Size.of(count_tiles, kvh, used), mtl.Size.of(128, 1, 1));
+ e.setPipeline(p.pipes[2]);
+ for ([_]rig_.Ref{ pm, pl, po }, 0..) |r, i| e.setBuffer(r.buf, r.off, i);
+ e.setBytes(std.mem.asBytes(&args), 3);
+ e.setBuffer(p.y.buf, p.y.off, 4);
+ e.setBuffer(segs.buf, segs.off, 5);
+ e.setBuffer(rs.buf, rs.off, 6);
+ e.dispatchGroups(mtl.Size.of(qh, n, 1), mtl.Size.of(32, 1, 1));
+ try p.rig.finish(pass);
+ }
+
+ /// conv: the batched kernel, or (`alone`) lfm2_conv a stream, into y and snap.
+ fn runConv(p: *Probe, y: rig_.Ref, snap: rig_.Ref, alone: bool) !void {
+ const pass = p.rig.begin();
+ const e = pass.e;
+ if (alone) {
+ e.setPipeline(p.pipes[1]);
+ for (0..p.count_streams) |j| {
+ const b = p.base(j);
+ e.setBuffer(p.args[0].buf, p.args[0].off + b * 3 * d * 2, 0);
+ e.setBuffer(p.caches[0][j].buf, p.caches[0][j].off, 1);
+ e.setBuffer(p.args[1].buf, p.args[1].off, 2);
+ e.setBuffer(y.buf, y.off + b * d * 2, 3);
+ e.setBuffer(snap.buf, snap.off + b * 2 * d * 2, 4);
+ e.setValue(@as(u32, 0), 5);
+ e.dispatchThreads(mtl.Size.of(d, p.streams[j], 1), mtl.Size.of(256, 1, 1));
+ }
+ return p.rig.finish(pass);
+ }
+ const t = try p.table(0);
+ defer p.a.free(t);
+ var rowseg: [most]u32 = undefined;
+ for (0..p.count_streams) |j| @memset(rowseg[p.base(j)..][0..p.streams[j]], @intCast(j));
+ const segs = try p.rig.scratch(most * @sizeOf(Seg));
+ @memcpy(p.rig.bytes(segs, most * @sizeOf(Seg)), std.mem.sliceAsBytes(t));
+ const rs = try p.rig.scratch(most * 4);
+ @memcpy(p.rig.bytes(rs, most * 4), std.mem.asBytes(&rowseg));
+ for (p.caches[0][0..p.count_streams]) |r| e.useResource(r.buf, mtl.ResourceUsage.read);
+ e.setPipeline(p.pipes[0]);
+ e.setBuffer(p.args[0].buf, p.args[0].off, 0);
+ e.setBuffer(segs.buf, segs.off, 1);
+ e.setBuffer(p.args[1].buf, p.args[1].off, 2);
+ e.setBuffer(y.buf, y.off, 3);
+ e.setBuffer(snap.buf, snap.off, 4);
+ e.setBuffer(rs.buf, rs.off, 5);
+ e.dispatchThreads(mtl.Size.of(d, p.rows, 1), mtl.Size.of(256, 1, 1));
+ try p.rig.finish(pass);
+ }
+
+ /// The per-stream kernels' bytes on the same inputs.
+ pub fn replaced(p: *Probe) ![]const u8 {
+ const n = p.rows;
+ if (p.which == .conv) {
+ const y = try p.rig.scratch(n * d * 2);
+ const snap = try p.rig.scratch(n * 2 * d * 2);
+ try p.runConv(y, snap, true);
+ const out = try p.a.alloc(u8, n * 3 * d * 2);
+ @memcpy(out[0 .. n * d * 2], p.rig.bytes(y, n * d * 2));
+ @memcpy(out[n * d * 2 ..], p.rig.bytes(snap, n * 2 * d * 2));
+ return p.keep(out);
+ }
+ p.fill(1);
+ const queries = p.work[6];
+ const y = p.work[7];
+ const pm = p.work[3];
+ const pl = p.work[4];
+ const po = p.work[5];
+ const pass = p.rig.begin();
+ const e = pass.e;
+ for (0..p.count_streams) |j| {
+ const b = p.base(j);
+ const rows = p.streams[j];
+ const cache = p.caches[1][j];
+ const kv_bytes = kvh * p.cap(j) * hd * 2;
+ e.setPipeline(p.pipes[3]);
+ e.setBuffer(p.args[0].buf, p.args[0].off + b * qh * hd * 2, 0);
+ e.setBuffer(p.args[1].buf, p.args[1].off + b * kvh * hd * 2, 1);
+ e.setBuffer(p.args[2].buf, p.args[2].off + b * kvh * hd * 2, 2);
+ e.setBuffer(p.args[3].buf, p.args[3].off, 3);
+ e.setBuffer(p.args[3].buf, p.args[3].off + hd * 2, 4);
+ const qk = [4]u32{ @intCast(p.len(j)), @intCast(p.cap(j)), @bitCast(@as(f32, 1e-5)), @bitCast(std.math.log2(@as(f32, 1e6))) };
+ e.setBytes(std.mem.asBytes(&qk), 5);
+ e.setBuffer(queries.buf, queries.off + b * d * 2, 6);
+ e.setBuffer(cache.buf, cache.off, 7);
+ e.setBuffer(cache.buf, cache.off + kv_bytes, 8);
+ e.dispatchGroups(mtl.Size.of(qh + kvh, rows, 1), mtl.Size.of(32, 1, 1));
+ const wide = rows > 8;
+ const s: u32 = if (wide) 1 else subs;
+ const args = [4]u32{ @intCast(p.len(j)), @intCast(p.cap(j)), @bitCast(1.0 / @sqrt(@as(f32, hd))), s };
+ e.setPipeline(p.pipes[if (wide) 5 else 4]);
+ e.setBuffer(queries.buf, queries.off + b * d * 2, 0);
+ e.setBuffer(cache.buf, cache.off, 1);
+ e.setBuffer(cache.buf, cache.off + kv_bytes, 2);
+ e.setBytes(std.mem.asBytes(&args), 3);
+ for ([_]rig_.Ref{ pm, pl, po }, 4..) |r, i| e.setBuffer(r.buf, r.off, i);
+ e.setValue(@as(u32, @intCast(rows)), 7);
+ const per: usize = if (wide) 16 else 8;
+ const used = (p.len(j) + rows + chunk_keys - 1) / chunk_keys * s;
+ e.dispatchGroups(mtl.Size.of((rows + per - 1) / per, kvh, used), mtl.Size.of(128, 1, 1));
+ e.setPipeline(p.pipes[6]);
+ for ([_]rig_.Ref{ pm, pl, po }, 0..) |r, i| e.setBuffer(r.buf, r.off, i);
+ e.setBytes(std.mem.asBytes(&args), 3);
+ e.setBuffer(y.buf, y.off + b * d * 2, 4);
+ e.dispatchGroups(mtl.Size.of(qh, rows, 1), mtl.Size.of(32, 1, 1));
+ }
+ try p.rig.finish(pass);
+ return p.keep(try p.a.dupe(u8, p.rig.bytes(y, n * d * 2)));
+ }
+
+ fn keep(p: *Probe, bytes: []u8) []const u8 {
+ p.a.free(p.golden);
+ p.golden = bytes;
+ return bytes;
+ }
+
+ pub fn reference(_: *Probe, want: []f64, mag: []f64) !void {
+ @memset(want, 0); // the check computes its own references
+ @memset(mag, 0);
+ }
+
+ pub fn check(p: *Probe, _: g.Gate) !?[]const u8 {
+ return if (p.which == .conv) p.checkConv() else p.checkAttention();
+ }
+
+ /// Each output against float64 attention over its stream's keys, from the queries and cache the call left.
+ fn checkAttention(p: *Probe) !?[]const u8 {
+ const out = std.mem.bytesAsSlice(u16, p.rig.bytes(p.y, p.rows * d * 2));
+ const queries = std.mem.bytesAsSlice(u16, p.rig.bytes(p.queries, p.rows * d * 2));
+ const w = try p.a.alloc(f64, chunks * chunk_keys);
+ defer p.a.free(w);
+ for (0..p.count_streams) |j| {
+ const cp = p.cap(j);
+ const kv = std.mem.bytesAsSlice(u16, p.rig.bytes(p.caches[0][j], 2 * kvh * cp * hd * 2));
+ for (0..p.streams[j]) |r| for (0..qh) |h| {
+ const row = p.base(j) + r;
+ const kh = h / (qh / kvh);
+ const keys = p.len(j) + r + 1;
+ var top: f64 = -std.math.inf(f64);
+ for (0..keys) |k| {
+ var dot: f64 = 0;
+ for (0..hd) |i| dot += c.fromBf16(queries[(row * qh + h) * hd + i]) * c.fromBf16(kv[(kh * cp + k) * hd + i]);
+ w[k] = dot / @sqrt(@as(f64, hd));
+ top = @max(top, w[k]);
+ }
+ var sum: f64 = 0;
+ for (w[0..keys]) |*x| {
+ x.* = @exp(x.* - top);
+ sum += x.*;
+ }
+ for (0..hd) |i| {
+ var o: f64 = 0;
+ var big: f64 = 0;
+ for (0..keys) |k| {
+ const v: f64 = c.fromBf16(kv[((kvh + kh) * cp + k) * hd + i]);
+ o += w[k] / sum * v;
+ big = @max(big, @abs(v));
+ }
+ const got: f64 = c.fromBf16(out[(row * qh + h) * hd + i]);
+ if (p.nan and j == 0 and kh == 0) {
+ if (!std.math.isNan(got)) return "a head behind its stream's NaN key gave a number";
+ continue;
+ }
+ if (!std.math.isFinite(got)) return "a finite head gave a non-finite output";
+ if (@abs(got - o) > big * 0x1p-7 + @abs(o) * 0x1p-8) return "an output is past its float64 bound";
+ }
+ };
+ }
+ return null;
+ }
+
+ /// Each output against the host's replay of lfm2_conv's steps (bf16 products, fp32 taps, the bf16 gate).
+ fn checkConv(p: *Probe) !?[]const u8 {
+ const n = p.rows;
+ const y = std.mem.bytesAsSlice(u16, p.rig.bytes(p.y, n * d * 2));
+ const snap = std.mem.bytesAsSlice(u16, p.rig.bytes(p.snap, n * 2 * d * 2));
+ for (0..p.count_streams) |j| {
+ const b = p.base(j);
+ const state = p.host[n * 3 * d + j * 2 * d ..][0 .. 2 * d];
+ for (0..p.streams[j]) |r| for (0..d) |ch| {
+ var bx: [3]u16 = undefined;
+ for (0..3) |tap| {
+ const src = @as(isize, @intCast(r + tap)) - 2;
+ if (src >= 0) {
+ const row = p.host[(b + @as(usize, @intCast(src))) * 3 * d ..];
+ bx[tap] = c.toBf16(c.fromBf16(row[ch]) * c.fromBf16(row[2 * d + ch]));
+ } else bx[tap] = state[@as(usize, @intCast(src + 2)) * d + ch];
+ }
+ // The taps' sum as products then adds, or as fused steps: the compiler may contract either way.
+ var acc: f32 = 0;
+ var fused: f32 = 0;
+ for (0..3) |tap| {
+ acc += c.fromBf16(bx[tap]) * c.fromBf16(p.weights[ch * 3 + tap]);
+ fused = @mulAdd(f32, c.fromBf16(bx[tap]), c.fromBf16(p.weights[ch * 3 + tap]), fused);
+ }
+ const gate = c.fromBf16(p.host[(b + r) * 3 * d + d + ch]);
+ const got = y[(b + r) * d + ch];
+ if (!same(got, c.toBf16(gate * c.fromBf16(c.toBf16(acc)))) and !same(got, c.toBf16(gate * c.fromBf16(c.toBf16(fused)))))
+ return "a conv output differs from the host's replay";
+ if (!same(snap[((b + r) * 2) * d + ch], bx[1]) or !same(snap[((b + r) * 2 + 1) * d + ch], bx[2])) return "a kept state differs from the host's replay";
+ };
+ }
+ return null;
+ }
+};
+
+fn same(x: u16, y: u16) bool {
+ const nx = std.math.isNan(c.fromBf16(x));
+ return if (nx) std.math.isNan(c.fromBf16(y)) else x == y;
+}
diff --git a/zig/tests/contract/metal_lfm2_tap.zig b/zig/tests/contract/metal_lfm2_tap.zig
new file mode 100644
index 0000000000..021baffc9e
--- /dev/null
+++ b/zig/tests/contract/metal_lfm2_tap.zig
@@ -0,0 +1,137 @@
+//! LFM2's tap copy for an external drafter (lfm2_tap): a layer's rows into their tap's columns, the others kept.
+const std = @import("std");
+const mtl = @import("metal");
+const contract = @import("contract");
+const sources = @import("kernel_sources");
+const c = contract.spec;
+const g = contract.gates;
+const rig_ = @import("rig.zig");
+const controls = @import("metal_lfm2_controls.zig");
+
+const hidden = 2048;
+const taps = 5;
+pub const name = "metal lfm2 tap copy";
+
+/// rows: one decode row, a verify window, a shared round's 32 and a 2,048-row prompt chunk.
+const shapes = [_]c.Shape{ .{ .rows = 1, .n = taps * hidden, .k = 1 }, .{ .rows = 10, .n = taps * hidden, .k = 1 }, .{ .rows = 32, .n = taps * hidden, .k = 1 }, .{ .rows = 2048, .n = taps * hidden, .k = 1 } };
+
+fn none(_: c.Shape) usize {
+ return 0;
+}
+
+pub fn spec() c.Spec {
+ return .{ .name = name, .accum = .none, .depth = none, .out = .bf16, .served = &shapes, .sums = .{ .f16 = 0, .bf16 = 0, .why = "a copy moves values: no sums" } };
+}
+
+/// A bf16 value for element `i` over signs and binades; the nonfinite fill plants NaN and infinities.
+fn value(i: usize, seed: u64, fill: g.Fill) u16 {
+ if (fill == .nonfinite) switch ((@as(u64, i) *% 2654435761 +% seed) % 9973) {
+ 0 => return 0x7fc0,
+ 1 => return 0x7f80,
+ 2 => return 0xff80,
+ else => {},
+ };
+ const h: u64 = (@as(u64, i) *% 0x9e3779b97f4a7c15 +% seed *% 0xbf58476d1ce4e5b9) >> 40;
+ const mant: f32 = @as(f32, @floatFromInt(h % 1024)) / 1024 + 1;
+ const exp: i32 = @as(i32, @intCast((h >> 10) % 8)) - 4;
+ const v = std.math.ldexp(mant, exp) * @as(f32, if (((h >> 15) & 1) == 0) 1 else -1);
+ return c.toBf16(v);
+}
+
+pub const Probe = struct {
+ a: std.mem.Allocator,
+ rig: rig_.Rig,
+ pipe: mtl.Pipeline,
+ mutation: controls.Mutation,
+ rows: usize = 0,
+ src: []u16 = &.{}, // H rows
+ prior: []u16 = &.{}, // the taps rows as the call finds them (the other taps' columns must stay)
+ args: [2]rig_.Ref = undefined,
+ y: rig_.Ref = undefined,
+
+ pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, mutation: controls.Mutation) !Probe {
+ const text = try controls.source(a, mutation);
+ defer if (mutation != .none) a.free(text);
+ return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, text, "lfm2_tap"), .mutation = mutation };
+ }
+
+ fn drop(p: *Probe) void {
+ p.a.free(p.src);
+ p.a.free(p.prior);
+ p.src = &.{};
+ p.prior = &.{};
+ p.rig.reset();
+ }
+
+ pub fn deinit(p: *Probe) void {
+ p.drop();
+ p.rig.deinit();
+ p.pipe.deinit();
+ }
+
+ const vt = g.hooks(Probe);
+ pub fn kernel(p: *Probe) g.Kernel {
+ var s = spec();
+ if (p.mutation != .none) s.name = p.mutation.name();
+ return .{ .ctx = p, .vt = &vt, .spec = s, .control = p.mutation.gate() };
+ }
+
+ fn fill(out: []u16, seed: u64, f: g.Fill) void {
+ for (out, 0..) |*v, i| v.* = value(i, seed, f);
+ }
+
+ pub fn prepare(p: *Probe, s: c.Shape, f: g.Fill, seed: u64) !bool {
+ p.drop();
+ if (s.n != taps * hidden or s.rows == 0) return false;
+ p.rows = s.rows;
+ p.src = try p.a.alloc(u16, s.rows * hidden);
+ fill(p.src, seed, f);
+ p.prior = try p.a.alloc(u16, s.rows * taps * hidden);
+ fill(p.prior, seed +% 3, .random);
+ p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.src));
+ const at = [2]u32{ taps * hidden, 2 * hidden }; // the third tap's columns
+ p.args[1] = try p.rig.input(std.mem.asBytes(&at));
+ p.y = try p.rig.output(p.count() * 2);
+ return true;
+ }
+
+ pub fn count(p: *Probe) usize {
+ return p.rows * taps * hidden;
+ }
+
+ pub fn run(p: *Probe) !void {
+ @memcpy(p.rig.bytes(p.y, p.count() * 2), std.mem.sliceAsBytes(p.prior));
+ const pass = p.rig.begin();
+ const e = pass.e;
+ e.setPipeline(p.pipe);
+ e.setBuffer(p.args[0].buf, p.args[0].off, 0);
+ e.setBuffer(p.y.buf, p.y.off, 1);
+ e.setBuffer(p.args[1].buf, p.args[1].off, 2);
+ e.dispatchThreads(mtl.Size.of(hidden / 8, p.rows, 1), mtl.Size.of(256, 1, 1));
+ return p.rig.finish(pass);
+ }
+
+ pub fn reference(p: *Probe, want: []f64, mag: []f64) !void {
+ for (want, mag, p.prior) |*w, *m, v| { // the other taps' columns keep what the call found
+ w.* = c.fromBf16(v);
+ m.* = @abs(w.*);
+ }
+ for (0..p.rows) |r| for (0..hidden) |i| {
+ const v = c.fromBf16(p.src[r * hidden + i]);
+ want[r * taps * hidden + 2 * hidden + i] = v;
+ mag[r * taps * hidden + 2 * hidden + i] = @abs(v);
+ };
+ }
+
+ /// The copy is exact: the third tap's columns hold H's bits, the others what the call found.
+ pub fn check(p: *Probe, _: g.Gate) !?[]const u8 {
+ const out = std.mem.bytesAsSlice(u16, p.rig.bytes(p.y, p.count() * 2));
+ for (out, p.prior, 0..) |got, before, i| {
+ const r = i / (taps * hidden);
+ const col = i % (taps * hidden);
+ const want = if (col >= 2 * hidden and col < 3 * hidden) p.src[r * hidden + col - 2 * hidden] else before;
+ if (got != want) return "an output differs from the copy's bits";
+ }
+ return null;
+ }
+};
diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig
index 26166bba9b..3e56849984 100644
--- a/zig/tests/contract/metal_limits.zig
+++ b/zig/tests/contract/metal_limits.zig
@@ -52,6 +52,7 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La
text = tf.flashnext_vision_contract.source(w);
return .{ .decl = d, .function = fun, .here = text, .pre_m5 = text };
}
+ if (std.mem.startsWith(u8, d.name, "metal lfm2 ")) return .{ .decl = d, .function = fun, .here = sources.lfm2, .pre_m5 = sources.lfm2 };
if (std.mem.startsWith(u8, d.name, "drama3 ")) {
const text_d3 = try @import("metal_drama3.zig").source(a, d.name);
return .{ .decl = d, .function = fun, .here = text_d3, .pre_m5 = text_d3 };
diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig
index a4535a92b7..ee31a9c0bd 100644
--- a/zig/tests/contract/metal_main.zig
+++ b/zig/tests/contract/metal_main.zig
@@ -23,6 +23,10 @@ const flash = @import("metal_flashnext.zig");
const gemma = @import("metal_gemma4.zig");
const gemma_embed = @import("metal_gemma4_embed.zig");
const gemma_mixed = @import("metal_gemma4_mixed.zig");
+const lfm2_argmax = @import("metal_lfm2_argmax.zig");
+const lfm2_streams = @import("metal_lfm2_streams.zig");
+const lfm2_tap = @import("metal_lfm2_tap.zig");
+const lfm2_controls = @import("metal_lfm2_controls.zig");
const g = contract.gates;
const usage = "usage: tf-kernel-contract [--only TEXT] [--no-gates] [--no-limits] [--offline DIR] [--seed N] [--repro] [--rotary-long] [--rotary-accuracy FILE]";
@@ -151,6 +155,58 @@ pub fn main(init: std.process.Init) !u8 {
if (!tensor) return error.TensorUnitsRequired;
try kernels.append(a, mixed.kernel());
}
+ var lfm2_list: std.ArrayList(*lfm2_argmax.Probe) = .empty;
+ var lfm2_streams_list: std.ArrayList(*lfm2_streams.Probe) = .empty;
+ defer {
+ for (lfm2_list.items) |p| {
+ p.deinit();
+ a.destroy(p);
+ }
+ for (lfm2_streams_list.items) |p| {
+ p.deinit();
+ a.destroy(p);
+ }
+ lfm2_list.deinit(a);
+ lfm2_streams_list.deinit(a);
+ }
+ for (std.enums.values(lfm2_controls.Mutation)) |m| {
+ const argmax_side = m == .none or m == .argmax_last_tie or m == .argmax_no_flag;
+ if (argmax_side) {
+ if (!wanted(o, if (m == .none) lfm2_argmax.name else m.name())) continue;
+ const p = try a.create(lfm2_argmax.Probe);
+ errdefer a.destroy(p);
+ p.* = try lfm2_argmax.Probe.init(a, device, queue, m);
+ try lfm2_list.append(a, p);
+ try kernels.append(a, p.kernel());
+ if (m != .none) continue;
+ }
+ for (std.enums.values(lfm2_streams.Which)) |w| {
+ const control_for = (m == .conv_first_state and w == .conv) or (m == .attention_first_cache and w == .attention);
+ if (m != .none and !control_for) continue;
+ if (!wanted(o, if (m == .none) lfm2_streams.spec(w).name else m.name())) continue;
+ const p = try a.create(lfm2_streams.Probe);
+ errdefer a.destroy(p);
+ p.* = try lfm2_streams.Probe.init(a, device, queue, w, m);
+ try lfm2_streams_list.append(a, p);
+ try kernels.append(a, p.kernel());
+ }
+ }
+ var tap_list: std.ArrayList(*lfm2_tap.Probe) = .empty;
+ defer {
+ for (tap_list.items) |p| {
+ p.deinit();
+ a.destroy(p);
+ }
+ tap_list.deinit(a);
+ }
+ for ([_]lfm2_controls.Mutation{ .none, .tap_packed_rows, .tap_clobbers_column }) |m| {
+ if (!wanted(o, if (m == .none) lfm2_tap.name else m.name())) continue;
+ const p = try a.create(lfm2_tap.Probe);
+ errdefer a.destroy(p);
+ p.* = try lfm2_tap.Probe.init(a, device, queue, m);
+ try tap_list.append(a, p);
+ try kernels.append(a, p.kernel());
+ }
var rows_list: std.ArrayList(*rows_k.Row) = .empty;
var proj_list: std.ArrayList(*affine_k.Proj) = .empty;
defer for (rows_list.items) |r| r.deinit();
diff --git a/zig/tests/contract/metal_tile.zig b/zig/tests/contract/metal_tile.zig
index 98fb49ae1d..5acdb678c6 100644
--- a/zig/tests/contract/metal_tile.zig
+++ b/zig/tests/contract/metal_tile.zig
@@ -19,7 +19,8 @@ fn tileDepth(s: c.Shape) usize {
/// LFM2.5-24B-A2B's expert projections, which #822 will run on the tiles: gate or up, 1,536 from 2,048, and down.
const prompt_rows = [_]c.Shape{ .{ .rows = 64, .n = 1536, .k = 2048 }, .{ .rows = 33, .n = 2048, .k = 1536 } };
-const decode_rows = [_]c.Shape{ .{ .rows = 32, .n = 1536, .k = 2048 }, .{ .rows = 5, .n = 2048, .k = 1536 } };
+/// The twin's rows: those projections, and the short conv's 6,144 from 2,048 at 16 rows (the twin's MMA window).
+const decode_rows = [_]c.Shape{ .{ .rows = 32, .n = 1536, .k = 2048 }, .{ .rows = 5, .n = 2048, .k = 1536 }, .{ .rows = 16, .n = 6144, .k = 2048 } };
pub fn spec(kind: Kind, name: []const u8) c.Spec {
const k_tile = c.Tile{ .size = 128, .step = 64, .min = 64 };
diff --git a/zig/tests/lfm2_exact.zig b/zig/tests/lfm2_exact.zig
index f9a8d6b66c..6fe9a3980b 100644
--- a/zig/tests/lfm2_exact.zig
+++ b/zig/tests/lfm2_exact.zig
@@ -33,6 +33,16 @@ fn prompt(m: *q.Model, scratch: *st.Scratch, cache: *st.Cache, ids: []const u32)
}
}
+/// Each of a forward's `rows` head rows: the GPU's greedy pick equals the host's draw on its logits.
+fn picks(m: *q.Model, s: *st.Scratch, rows: usize) !void {
+ const best = s.best.slice(u32, 2 * rows);
+ for (0..rows) |r| {
+ const want = try q.backend.draw(m.gpa, s.logits.slice(u16, rows * c.vocab)[r * c.vocab ..][0..c.vocab], null, 0, &.{});
+ try std.testing.expectEqual(want, best[2 * r]);
+ try std.testing.expectEqual(@as(u32, 0), best[2 * r + 1]);
+ }
+}
+
/// A prompt of `prefix` rows, then plain decode to `len`: what a window's kept rows equal.
fn decoded(m: *q.Model, solo: *st.Scratch, cache: *st.Cache, ids: []const u32, prefix: usize, len: usize) !void {
try prompt(m, solo, cache, ids[0..prefix]);
@@ -48,6 +58,7 @@ fn window(m: *q.Model, wide: *st.Scratch, solo: *st.Scratch, ids: []const u32, p
try prompt(m, solo, &b, ids[0..prefix]);
try equal(&a, &b);
try fwd.run(m, wide, &.{.{ .cache = &a, .rows = rows }}, ids[prefix..][0..rows], true, .all);
+ try picks(m, wide, rows);
for (0..rows) |r| {
try fwd.run(m, solo, &.{.{ .cache = &b, .rows = 1 }}, ids[prefix + r ..][0..1], false, .all);
const want = solo.logits.contents()[0 .. c.vocab * 2];
@@ -124,6 +135,7 @@ fn round(
total += r;
}
try fwd.run(m, wide, segs[0..n], tokens[0..total], true, .all);
+ try picks(m, wide, total);
var base: usize = 0;
for (caches, prompts, lens, rows) |*cache, start, len, r| {
if (r == 0) continue;