diff --git a/README.md b/README.md
index 46023ee01a..41328a20e1 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). LiquidAI's DSpark drafter with `--drafter`. |
| 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. |
@@ -287,7 +287,7 @@ The binary's `capabilities --json` response lists its supported flags and platfo
| `--loop-guard` | Close reasoning that repeats itself verbatim (a short cycle, or a long passage) where the engine supports it. |
| `--no-drafts` | Produce the plain reference through the same engine. |
| `--drafter DRAFTER` | An external drafter, as a directory or a repo id in the Hugging Face cache, for models that take one: Qwen3.8-27B uses DFlash2 and Gemma 4 31B the drafter named in its recipe. Without one they decode plain. |
-| `--drafter-bits 0\|4` | Drafter weight precision where the model offers a choice: 4 bits (the default), or 0 to keep bf16. Qwen3.8-27B on CUDA takes 4 only. |
+| `--drafter-bits 0\|4\|8` | Drafter weight precision where the model offers a choice: 4 bits (the default), or 0 to keep bf16. Qwen3.8-27B on CUDA takes 4 only. LFM2.5-24B-A2B's DSpark drafter runs bf16 as released (0, the default when not given) or 8-bit codes quantized at load (8); it has no 4-bit mode. |
| `--record DIR` | Write each stream's committed tokens, positions and tapped rows under DIR, for drafter training. CUDA Nemotron only; a single stream runs through the lane rounds, and recording stops below 25 GiB free disk. |
| `--keep-warm SECONDS` | Keep the Metal GPU active while idle for this long after a request, default 900; zero disables it. |
| `--cfg-gamma GAMMA` | Metal Drama speech: semantic guidance strength, finite 1–5; 0/default disables it. Only reference-conditioned requests are guided. Enabled replies report `X-Cfg-Gamma`; codec heads stay conditional. |
diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md
index 0d8241d5c2..1367e0c5a8 100644
--- a/docs/kernel-contract.md
+++ b/docs/kernel-contract.md
@@ -131,6 +131,13 @@ 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.
+The DSpark drafter's tap copy, Markov scan and pick, draft ids, block attention and confidence logit are there too, with seven more controls.
+Its bf16 linears (few rows, 64x64 tiles) and the Markov scan over a bf16 W2 have four and two more.
+Their launches are in `lfm2_records.zig` with the M1 and M2 limits from `--offline`; `lfm2_attn_segs` and
+`lfm2_attn_block` reserve their 128 threads in their source, as Metal cannot build them 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..966aecc11b 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, DSpark drafts](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..0417184862
--- /dev/null
+++ b/docs/recipes/lfm2.5-24b-a2b.md
@@ -0,0 +1,110 @@
+# 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.
+
+## DSpark drafts
+
+LiquidAI's DSpark drafter for this model is
+[`LiquidAI/LFM2.5-24B-A2B-DSpark`](https://huggingface.co/LiquidAI/LFM2.5-24B-A2B-DSpark/tree/e052e5b1dd103c38cdd4f8ddc068d7dea4586f55)
+at revision `e052e5b1dd103c38cdd4f8ddc068d7dea4586f55` (327.7M bf16 parameters, 655 MB). Its weights are under the
+[LFM Open License v1.0](https://huggingface.co/LiquidAI/LFM2.5-24B-A2B-DSpark/blob/e052e5b1dd103c38cdd4f8ddc068d7dea4586f55/LICENSE);
+TensorFold does not ship them, so read the license before use.
+
+```bash
+hf download LiquidAI/LFM2.5-24B-A2B-DSpark --revision e052e5b1dd103c38cdd4f8ddc068d7dea4586f55 \
+ config.json model.safetensors LICENSE README.md --local-dir "$HOME/models/LFM2.5-24B-A2B-DSpark"
+zig-out/native/bin/tensorfold-native serve "$MODEL_DIR" --port 8080 --parallel 16 \
+ --drafter "$HOME/models/LFM2.5-24B-A2B-DSpark"
+```
+
+Without `--drafter` the engine serves as before. `"draft": false` on a request, or `--no-drafts`, turns the drafter
+off; the lane core verifies every draft, so a drafted reply equals its `"draft": false` reply token for token.
+
+The drafter has no embedding and no head of its own:
+
+- Context: the target's rows after layers 2, 10, 18, 26 and 34 (the backend's taps, `lfm2_tap`), through `fc` and
+ `hidden_norm`, give each of its five layers' context keys and values (k_norm, then RoPE at the row's position).
+- Block: the pending token and mask tokens run five Qwen3 layers with non-causal attention over the context and the
+ block, then the target's head. A block is as long as the deepest chain its pass holds, at most the nine rows the
+ drafter was trained on.
+- Markov head: row i's logits plus `W2 W1[previous token]` (rank 256); draft i is their argmax.
+
+The drafter runs the released bf16 matrices as they are (`--drafter-bits 0`, the default). `--drafter-bits 8` instead
+quantizes them at load to 8-bit MLX affine codes in groups of 64, about 5% faster on one stream and 361 MiB of weights
+against 625 MiB; the server logs which mode runs, and 4 bits is refused. Both run on the family's kernels
+(bf16 on `lfm2_bf16_rows`, `lfm2_bf16_mma` and `lfm2_markov_scan_bf16`, 8-bit on `core/tile_qmm`); the q and k rows are reordered so that the target's half-split RoPE kernel gives the
+checkpoint's interleaved RoPE. Its own kernels are `lfm2_tap`, `lfm2_attn_block` and `lfm2_attn_block_merge`
+(non-causal block attention, the target's attention kernels with a template flag), `lfm2_markov_scan` and
+`lfm2_markov_pick` (a Markov step) and `lfm2_draft_ids` (a hold's drafts into the next verify's rows on the GPU, so
+the host does not wait for a hold that fits one pass). Each kept prompt state keeps the drafter's context too, so a
+drafting stream resumes prompts. Shared rounds take the lane core's shared-round guard: they draft only while
+drafted rounds measure faster than plain ones at that stream count.
+
+`tf-kernel-contract --only lfm2` checks these kernels. `tf-lfm2-dspark MODEL DRAFTER TOKENS [ROWS]` runs the drafter
+against an fp32 host reference of the released checkpoint (the same taps and the target's embedding and head),
+prints each draft against the reference's pick with the margin between the reference's two best tokens, and times
+the drafter's passes; a near tie can pick the other token.
+
+### Verify schedule
+
+The drafter also carries a confidence head: a linear map of a draft's final block row and `W1[previous token]`
+(2,304 inputs) to one logit, through a sigmoid (`lfm2_dspark_confidence`). It does not change the drafts.
+The verify rows follow it: a chain's chance at draft j is the product of its first j confidences. Each round the
+lane core's window allocator takes the likeliest drafts over all streams while the expected tokens a ms rise, so
+each stream gets its own verify length, and the next hold drafts as deep as the stream's grant. A hold still
+running on the GPU is priced by its lane's last confidences, so the host does not wait for it. A checkpoint
+without the head verifies by the core's depth rule, from each stream's measured acceptance by depth.
+
+A stream with a copy ahead verifies the copy and asks no hold, and a drafting stream prices its copy
+windows by its own copies. The lane core also weighs the hold pass itself: the drafter times holds of 1 to 16
+blocks at load and corrects those prices toward each settled hold's GPU time. A lone stream holds drafts only while
+its drafted rounds measure faster than plain ones (`zig/src/core/lanes/lone_guard.zig`); shared rounds keep the
+shared-round guard.
+
+`tf-lfm2-dspark` also prints the confidence
+head against its fp32 reference and times shared holds of 2 to 16 streams.
+
+The loader reads the checkpoint one matrix at a time and keeps only the 8-bit codes; a stream's drafter lane is one
+Metal buffer (its layers' keys and values, then its staged rows), and its bytes join the streams the load limit
+admits. At load the drafter logs its weights, scratch and a stream's bytes. `tf-lfm2-dspark` also counts host
+allocations over twenty served drafter rounds, which must be none.
diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig
index d5ab369dd7..727daf7e2e 100644
--- a/zig/build/metal_targets.zig
+++ b/zig/build/metal_targets.zig
@@ -292,6 +292,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s
.{ .name = "tf-lfm2-forward", .path = "zig/tests/lfm2_forward.zig", .about = "LFM2-MoE native teacher-forced logits" },
.{ .name = "tf-lfm2-bench", .path = "zig/tests/lfm2_bench.zig", .about = "LFM2-MoE prompt-chunk and one-row decode timing" },
.{ .name = "tf-lfm2-exact", .path = "zig/tests/lfm2_exact.zig", .about = "LFM2-MoE prompt chunks, windows and committed-state exactness" },
+ .{ .name = "tf-lfm2-dspark", .path = "zig/tests/lfm2_dspark.zig", .about = "LFM2's DSpark drafter on the GPU against an fp32 host reference of the released drafter" },
.{ .name = "tf-lfm2-fp64", .path = "zig/tests/lfm2_fp64.zig", .about = "LFM2 projection classes and attention on the tensor units, the tiles and the twin, each against float64" },
.{ .name = "tensorfold", .path = "zig/src/main.zig", .about = "The native engine's command line" },
.{ .name = "tf-nemotron-fixtures", .path = "zig/tests/nemotron_fixtures.zig", .about = "Nemotron kernels against the Python engine's captured ops" },
@@ -455,6 +456,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 DSpark drafter and its taps: config, quantizer and kernels 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..2da8578970 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;
@@ -225,22 +232,18 @@ 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];
+template // RB row blocks share each staged K and V fragment; FULL: all see [0, len + 1)
+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;
const uint qid = lane / 4, fm = (qid & 4) + ((lane / 2) % 4), fn = (qid & 2) * 2 + (lane % 2) * 2;
uint count[RB]; // this lane's rows' visible keys
- for (int b = 0; b < RB; ++b) count[b] = p.len + min(r0 + 8 * b + fm, rows - 1) + 1;
- const uint top = p.len + min(r0 + 8 * RB - 1, rows - 1) + 1; // the tile's most
+ for (int b = 0; b < RB; ++b) count[b] = FULL ? p.len + 1 : p.len + min(r0 + 8 * b + fm, rows - 1) + 1;
+ const uint top = FULL ? p.len + 1 : p.len + min(r0 + 8 * RB - 1, rows - 1) + 1; // the tile's most
if (chunk * PER + sub0 * SUB >= top) return;
vec q[RB][8];
for (int b = 0; b < RB; ++b)
@@ -337,16 +340,25 @@ 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;
- const uint count = p.len + r + 1;
+template
+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 = FULL ? p.len + 1 : p.len + r + 1;
const uint per = PER;
const uint chunks = (count + per - 1) / per;
const size_t at = (size_t(r) * QH + head) * (CHUNKS * p.subs);
@@ -380,6 +392,15 @@ 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);
}
+template
+[[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);
+}
+template [[host_name("lfm2_attn_merge")]] [[kernel]] decltype(lfm2_attn_merge) lfm2_attn_merge;
+template [[host_name("lfm2_attn_block_merge")]] [[kernel]] decltype(lfm2_attn_merge) lfm2_attn_merge;
// 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 +515,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_dspark.metal b/zig/kernels/metal/lfm2/lfm2_dspark.metal
new file mode 100644
index 0000000000..7440805007
--- /dev/null
+++ b/zig/kernels/metal/lfm2/lfm2_dspark.metal
@@ -0,0 +1,258 @@
+// DSpark's Markov step, compiled after lfm2.metal's text (its includes and constants).
+
+// DSpark's Markov scan, W2 as 8-bit codes (BF16 false) or bf16 as released (true): a threadgroup's best of 512 tokens.
+struct lfm2_markov_args { uint stride, rank, vocab, prev_stride; };
+constant constexpr uint MARKOV_SPAN = 512;
+
+// Eight bf16 in a uint4 as floats (a bf16 is the high half of its float).
+inline void lfm2_unpack8(uint4 v, thread float* f) {
+ for (int j = 0; j < 4; ++j) {
+ f[2 * j] = as_type(v[j] << 16);
+ f[2 * j + 1] = as_type(v[j] & 0xffff0000u);
+ }
+}
+template
+inline void lfm2_markov_scan_t(const device bfloat* logits, const device bfloat* w1, const device uint8_t* w2,
+ const device uint* prev, constant lfm2_markov_args& p, device float* best,
+ device uint* at, const device bfloat* w2s, const device bfloat* w2b, uint2 tg,
+ uint2 grid, uint sg, uint lane, threadgroup float* e, threadgroup float* sv,
+ threadgroup uint* si) {
+ const uint s = tg.y, t = sg * 32 + lane, id = prev[size_t(s) * p.prev_stride];
+ for (uint i = t; i < p.rank; i += 256) e[i] = id < p.vocab ? float(w1[size_t(id) * p.rank + i]) : 0.0f;
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ const device bfloat* row = logits + size_t(s) * p.stride;
+ float mv = -INFINITY;
+ uint mi = 0xffffffffu;
+ bool nan = false;
+ const uint first = tg.x * MARKOV_SPAN, last = min(first + MARKOV_SPAN, p.vocab);
+ for (uint v = first + sg; v < last; v += 8) {
+ float d = 0.0f;
+ for (uint i = 8 * lane; i < p.rank; i += 256) {
+ if (BF16) {
+ float f[8];
+ lfm2_unpack8(*(const device uint4*)((const device bfloat*)w2 + size_t(v) * p.rank + i), f);
+ for (int j = 0; j < 8; ++j) d = fma(f[j], e[i + j], d);
+ continue;
+ }
+ const size_t g = (size_t(v) * p.rank + i) / 64;
+ const float sc = w2s[g], bi = w2b[g];
+ const uchar4 lo = *(const device uchar4*)(w2 + size_t(v) * p.rank + i);
+ const uchar4 hi = *(const device uchar4*)(w2 + size_t(v) * p.rank + i + 4);
+ for (int j = 0; j < 4; ++j) d = fma(fma(sc, float(lo[j]), bi), e[i + j], d);
+ for (int j = 0; j < 4; ++j) d = fma(fma(sc, float(hi[j]), bi), e[i + 4 + j], d);
+ }
+ d = simd_sum(d);
+ const float sc = float(static_cast(float(row[v]) + float(static_cast(d))));
+ nan = nan || isnan(sc);
+ if (sc > mv || (mi == 0xffffffffu && sc == mv)) {
+ mv = sc;
+ mi = v;
+ }
+ }
+ if (lane == 0) {
+ sv[sg] = nan ? NAN : mv;
+ si[sg] = mi;
+ }
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ if (t != 0) return;
+ float bv = sv[0];
+ uint bi = si[0];
+ for (uint g = 1; g < 8; ++g) {
+ if (isnan(bv) || isnan(sv[g])) bv = NAN;
+ else if (sv[g] > bv || (sv[g] == bv && si[g] < bi)) {
+ bv = sv[g];
+ bi = si[g];
+ }
+ }
+ best[size_t(s) * grid.x + tg.x] = bv;
+ at[size_t(s) * grid.x + tg.x] = bi;
+}
+
+[[kernel]] void lfm2_markov_scan(const device bfloat* logits [[buffer(0)]], const device bfloat* w1 [[buffer(1)]],
+ const device uint8_t* w2 [[buffer(2)]], const device uint* prev [[buffer(3)]],
+ constant lfm2_markov_args& p [[buffer(4)]], device float* best [[buffer(5)]],
+ device uint* at [[buffer(6)]], const device bfloat* w2s [[buffer(7)]],
+ const device bfloat* w2b [[buffer(8)]], uint2 tg [[threadgroup_position_in_grid]],
+ uint2 grid [[threadgroups_per_grid]], uint sg [[simdgroup_index_in_threadgroup]],
+ uint lane [[thread_index_in_simdgroup]]) {
+ threadgroup float e[256];
+ threadgroup float sv[8];
+ threadgroup uint si[8];
+ lfm2_markov_scan_t(logits, w1, w2, prev, p, best, at, w2s, w2b, tg, grid, sg, lane, e, sv, si);
+}
+
+// The same scan over the released bf16 W2 [vocab, rank] (buffers 7 and 8 unused).
+[[kernel]] void lfm2_markov_scan_bf16(const device bfloat* logits [[buffer(0)]], const device bfloat* w1 [[buffer(1)]],
+ const device uint8_t* w2 [[buffer(2)]], const device uint* prev [[buffer(3)]],
+ constant lfm2_markov_args& p [[buffer(4)]], device float* best [[buffer(5)]],
+ device uint* at [[buffer(6)]], uint2 tg [[threadgroup_position_in_grid]],
+ uint2 grid [[threadgroups_per_grid]], uint sg [[simdgroup_index_in_threadgroup]],
+ uint lane [[thread_index_in_simdgroup]]) {
+ threadgroup float e[256];
+ threadgroup float sv[8];
+ threadgroup uint si[8];
+ lfm2_markov_scan_t(logits, w1, w2, prev, p, best, at, w1, w1, tg, grid, sg, lane, e, sv, si);
+}
+
+// DSpark's Markov step pick: the scan's best over a stream's threadgroups (ties to the lowest token), as a Pick.
+[[kernel]] void lfm2_markov_pick(const device float* best [[buffer(0)]], const device uint* at [[buffer(1)]],
+ constant uint& groups [[buffer(2)]], device uint* picks [[buffer(3)]],
+ uint s [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]]) {
+ threadgroup float tv[256];
+ threadgroup uint ti[256];
+ threadgroup uint tn[256];
+ float bv = -INFINITY;
+ uint bi = 0xffffffffu, nan = 0;
+ for (uint g = t; g < groups; g += 256) {
+ const float v = best[size_t(s) * groups + g];
+ const uint i = at[size_t(s) * groups + g];
+ if (isnan(v)) nan = 1;
+ else if (v > bv || (v == bv && i < bi)) {
+ bv = v;
+ bi = i;
+ }
+ }
+ tv[t] = bv;
+ ti[t] = bi;
+ tn[t] = nan;
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ for (uint w = 128; w > 0; w >>= 1) {
+ if (t < w) {
+ const float v = tv[t + w];
+ const uint i = ti[t + w];
+ if (v > tv[t] || (v == tv[t] && i < ti[t])) {
+ tv[t] = v;
+ ti[t] = i;
+ }
+ tn[t] |= tn[t + w];
+ }
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ }
+ if (t == 0) {
+ picks[2 * s] = ti[0];
+ picks[2 * s + 1] = tn[0];
+ }
+}
+
+// Held drafts from picks[i * stride] into ids[i]; a nonfinite or past-vocabulary pick is token 0 and sets bad.
+[[kernel]] void lfm2_draft_ids(const device uint* picks [[buffer(0)]], constant uint4& p [[buffer(1)]],
+ device uint* ids [[buffer(2)]], device atomic_uint* bad [[buffer(3)]],
+ uint i [[thread_position_in_grid]]) {
+ if (i >= p.y) return;
+ const uint token = picks[size_t(i) * p.x];
+ const bool ok = token < p.z && picks[size_t(i) * p.x + 1] == 0u;
+ ids[i] = ok ? token : 0u;
+ if (!ok) atomic_store_explicit(bad, 1u, memory_order_relaxed);
+}
+
+// DSpark's confidence logit of draft i of stream s: proj . [x row s * n + i, W1[previous token]] in fp32, no bias.
+struct lfm2_confidence_args { uint n, pass, vocab, rank, hidden; };
+[[kernel]] void lfm2_dspark_confidence(const device bfloat* x [[buffer(0)]], const device bfloat* w1 [[buffer(1)]],
+ const device uint* ids [[buffer(2)]], const device uint* picks [[buffer(3)]],
+ const device bfloat* proj [[buffer(4)]],
+ constant lfm2_confidence_args& p [[buffer(5)]], device float* out [[buffer(6)]],
+ uint2 tg [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]],
+ uint sg [[simdgroup_index_in_threadgroup]],
+ uint lane [[thread_index_in_simdgroup]]) {
+ threadgroup float part[8];
+ const uint i = tg.x, s = tg.y;
+ // the previous token: the block's pending id for the first draft, else the Markov step's pick before it
+ const uint id = i == 0 ? ids[size_t(s) * p.n] : picks[2 * (size_t(i - 1) * p.pass + s)];
+ const device bfloat* row = x + (size_t(s) * p.n + i) * p.hidden;
+ float d = 0.0f;
+ for (uint k = t; k < p.hidden; k += 256) d = fma(float(row[k]), float(proj[k]), d);
+ if (id < p.vocab)
+ for (uint k = t; k < p.rank; k += 256) d = fma(float(w1[size_t(id) * p.rank + k]), float(proj[p.hidden + k]), d);
+ d = simd_sum(d);
+ if (lane == 0) part[sg] = d;
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ if (t != 0) return;
+ float sum = 0.0f;
+ for (uint g = 0; g < 8; ++g) sum += part[g];
+ out[size_t(i) * p.pass + s] = sum;
+}
+
+// DSpark's bf16 linears (the released matrices as stored): y [rows, n] = x [rows, k] W [n, k]^T, fp32 sums, bf16 out.
+struct lfm2_bf16_args { uint rows, n, k; };
+constant constexpr uint BF16_ROWS = 16;
+
+// Few rows: a simdgroup a column of W over 16 rows, so W is read once a 16 rows. k a multiple of 8.
+[[kernel]] void lfm2_bf16_rows(const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]],
+ constant lfm2_bf16_args& p [[buffer(2)]], device bfloat* y [[buffer(3)]],
+ uint2 tg [[threadgroup_position_in_grid]], uint sg [[simdgroup_index_in_threadgroup]],
+ uint lane [[thread_index_in_simdgroup]]) {
+ const uint col = tg.x * 8 + sg, r0 = tg.y * BF16_ROWS;
+ if (col >= p.n) return;
+ const uint rows = min(BF16_ROWS, p.rows - r0);
+ float acc[BF16_ROWS];
+ for (uint r = 0; r < BF16_ROWS; ++r) acc[r] = 0.0f;
+ const device bfloat* wr = w + size_t(col) * p.k;
+ const device bfloat* xr = x + size_t(r0) * p.k;
+ for (uint i = 8 * lane; i < p.k; i += 256) {
+ float wv[8];
+ lfm2_unpack8(*(const device uint4*)(wr + i), wv);
+ for (uint r = 0; r < BF16_ROWS; ++r) {
+ if (r >= rows) break;
+ float xv[8];
+ lfm2_unpack8(*(const device uint4*)(xr + size_t(r) * p.k + i), xv);
+ for (int j = 0; j < 8; ++j) acc[r] = fma(wv[j], xv[j], acc[r]);
+ }
+ }
+ for (uint r = 0; r < BF16_ROWS; ++r) {
+ const float sum = simd_sum(acc[r]);
+ if (lane == 0 && r < rows) y[size_t(r0 + r) * p.n + col] = static_cast(sum);
+ }
+}
+
+// Many rows: 64x64 tiles of y, four simdgroups of 32x32, k in steps of 32. k a multiple of 32.
+[[kernel]] void lfm2_bf16_mma(const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]],
+ constant lfm2_bf16_args& p [[buffer(2)]], device bfloat* y [[buffer(3)]],
+ uint2 tg [[threadgroup_position_in_grid]], uint sg [[simdgroup_index_in_threadgroup]],
+ uint t [[thread_index_in_threadgroup]]) {
+ threadgroup bfloat xs[64 * 32];
+ threadgroup bfloat ws[64 * 32];
+ threadgroup float out[4][32 * 32];
+ const uint c0 = tg.x * 64, r0 = tg.y * 64;
+ simdgroup_matrix acc[4][4];
+ for (uint i = 0; i < 4; ++i)
+ for (uint j = 0; j < 4; ++j) acc[i][j] = simdgroup_matrix(0.0f);
+ const uint sr = (sg / 2) * 32, sc = (sg % 2) * 32;
+ for (uint k0 = 0; k0 < p.k; k0 += 32) {
+ for (uint i = t; i < 64 * 4; i += 128) { // 64 rows of 32 bf16 as 4 chunks of 8, 16 bytes a load
+ const uint r = i / 4, q = (i % 4) * 8;
+ const uint4 xv = r0 + r < p.rows ? *(const device uint4*)(x + size_t(r0 + r) * p.k + k0 + q) : uint4(0);
+ const uint4 wv = c0 + r < p.n ? *(const device uint4*)(w + size_t(c0 + r) * p.k + k0 + q) : uint4(0);
+ *(threadgroup uint4*)(xs + r * 32 + q) = xv;
+ *(threadgroup uint4*)(ws + r * 32 + q) = wv;
+ }
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ for (uint kk = 0; kk < 32; kk += 8) {
+ simdgroup_matrix a[4], b[4];
+ for (uint i = 0; i < 4; ++i) simdgroup_load(a[i], xs + (sr + 8 * i) * 32 + kk, 32);
+ for (uint j = 0; j < 4; ++j) simdgroup_load(b[j], ws + (sc + 8 * j) * 32 + kk, 32, ulong2(0, 0), true);
+ for (uint i = 0; i < 4; ++i)
+ for (uint j = 0; j < 4; ++j) simdgroup_multiply_accumulate(acc[i][j], a[i], b[j], acc[i][j]);
+ }
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ }
+ for (uint i = 0; i < 4; ++i)
+ for (uint j = 0; j < 4; ++j) simdgroup_store(acc[i][j], out[sg] + 8 * i * 32 + 8 * j, 32);
+ simdgroup_barrier(mem_flags::mem_threadgroup);
+ for (uint i = t % 32; i < 32 * 32; i += 32) {
+ const uint r = i / 32, c = i % 32;
+ const uint row = r0 + sr + r, col = c0 + sc + c;
+ if (row < p.rows && col < p.n) y[size_t(row) * p.n + col] = static_cast(out[sg][i]);
+ }
+}
+
+// A drafter block's attention: lfm2_attn_mma where every row sees all keys, its 128 threads reserved for M1 and M2.
+[[max_total_threads_per_threadgroup(128)]] [[kernel]] void lfm2_attn_block(
+ 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<1, true>(queries, keys, values, p, pm, pl, po, rows, tg, tid, sg, lane, ks, vs);
+}
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..bec767cf18 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_dspark.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..2e7587fa3d
--- /dev/null
+++ b/zig/src/core/contract/lfm2_records.zig
@@ -0,0 +1,25 @@
+//! 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;
+// The drafter kernels carry M1 and M2 limits too, recorded from the offline compile; no run covers those chips.
+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 },
+ .{ .name = "metal lfm2 markov scan lfm2_markov_scan", .threads = 256, .tg_memory = 1088, .m1 = 1024, .m2 = 1024 },
+ .{ .name = "metal lfm2 markov scan bf16 lfm2_markov_scan_bf16", .threads = 256, .tg_memory = 1088, .m1 = 1024, .m2 = 1024 },
+ .{ .name = "metal lfm2 bf16 rows lfm2_bf16_rows", .threads = 256, .tg_memory = 0, .m1 = 1024, .m2 = 1024 },
+ .{ .name = "metal lfm2 bf16 tile lfm2_bf16_mma", .threads = 128, .tg_memory = 24576, .m1 = 832, .m2 = 832 },
+ .{ .name = "metal lfm2 markov pick lfm2_markov_pick", .threads = 256, .tg_memory = 3072, .m1 = 1024, .m2 = 1024 },
+ .{ .name = "metal lfm2 draft ids lfm2_draft_ids", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 },
+ .{ .name = "metal lfm2 block attention lfm2_attn_block", .threads = 128, .tg_memory = 18688, .reserved = 128 },
+ .{ .name = "metal lfm2 block attention merge lfm2_attn_block_merge", .threads = 32, .tg_memory = 0, .m1 = 1024, .m2 = 1024 },
+ .{ .name = "metal lfm2 dspark confidence lfm2_dspark_confidence", .threads = 256, .tg_memory = 32, .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/engine_api.zig b/zig/src/core/engine_api.zig
index c44a338894..7191dabbff 100644
--- a/zig/src/core/engine_api.zig
+++ b/zig/src/core/engine_api.zig
@@ -285,9 +285,11 @@ pub const Open = struct {
/// --parallel named a number: an engine whose memory fits fewer streams refuses instead of serving fewer.
lanes_fixed: bool = false,
drafts: bool = true,
- /// --drafter: a draft model's directory (Qwen3.8-27B's DFlash2; null: none), and --drafter-bits: 4, or 0 for bf16.
+ /// --drafter: a draft model's directory (null: none); --drafter-bits: 4, 0 for bf16, or 8 (LFM2 only).
drafter: ?[]const u8 = null,
drafter_bits: u8 = 4,
+ /// Whether --drafter-bits was given: LFM2's drafter runs bf16 as released unless it is asked for 8-bit codes.
+ drafter_bits_given: bool = false,
/// --record: each stream's committed rows and tapped features written under this folder, for drafter training.
record: ?[]const u8 = null,
speed_up: ?[]const u8 = null,
diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig
index cb2159b03a..3db209405b 100644
--- a/zig/src/core/lanes/config.zig
+++ b/zig/src/core/lanes/config.zig
@@ -14,6 +14,7 @@ pub const default_prior = [_]f64{ 0.85, 0.75, 0.7, 0.65, 0.6, 0.55, 0.5, 0.5, 0.
pub const Model = struct {
root_copies: bool = false, // verified context copies precede head drafts and skip their proposal pass
copy_drafts: bool = true, // global serving switch, independent of shared verification capacity
+ own_copy_rates: bool = false, // a copy window is priced by its stream's own copies, not copy_rate's prior
exact_width: u32 = 1, // the widest window that reproduces one-token steps bit for bit
first_copy_rows: u32 = 0, // a lone stream's first copy window (0: the exact width)
@@ -35,7 +36,12 @@ pub const Model = struct {
max_streams: u32 = 0, // a shared forward's streams at most (0: 32)
shared_costs: []const Cost = &.{}, // shared forwards' ms by total rows (timed at load)
draft_probabilities: bool = false, // the head gives each drafted node's chance of landing
+ granted_depth: bool = false, // with draft_probabilities: every stream drafts its last grant plus a slack
+ draft_slack: u32 = 0, // drafts past a stream's last grant, short or not (0: the Constants' own)
draft_streams: bool = false, // the head drafts every stream of a shared round in one batch
+ shared_guard: bool = false, // a shared round drafts only while drafted rounds measure faster (shared_guard.zig)
+ copy_first: bool = false, // a copy ahead takes the round: the head holds no drafts for it (a drafter's own pass)
+ hold_costs: []const @import("hold_price.zig").HoldCost = &.{}, // a drafter's own hold pass, timed at load
head_trees: bool = false, // the head drafts trees of lanes (DraftRequest.lanes) and holds them
join_tail: bool = false, // a prompt in pieces leaves its last token for the next shared round's window (Backend.join)
piece_runs: bool = false, // a round's pieces go to the backend in one call (Backend.pieces: one forward over them)
@@ -52,6 +58,7 @@ pub const Config = struct {
root_copies: bool = false,
family_width: u32,
max_copy: u32,
+ own_copy_rates: bool = false, // Model.own_copy_rates
first_copy: u32,
base_width: u32,
family_mtp: bool,
@@ -64,10 +71,12 @@ pub const Config = struct {
batch_rows: u32,
batch_streams: u32,
node_probabilities: bool,
+ granted_depth: bool = false, // Model.granted_depth
fill_lanes: u32 = 0, // rows a window may grow to with suffix-match branches (0: off; lanes/fill.zig)
fill_match: u32 = 8, // context tokens a suffix-match branch must match (shorter ones are coincidental phrases)
fill_floor: f64 = 0.35, // the share of grafted rows a stream must keep to go on reading its chain back every round
draft_streams: bool,
+ shared_guard: bool = false,
head_trees: bool,
/// a prompt in pieces leaves its last token as the pending row of its first shared round's window
join_tail: bool = false,
@@ -134,6 +143,7 @@ pub const Config = struct {
.lone_cost_tile = m.lone_cost_tile,
.family_width = width,
.max_copy = max_copy,
+ .own_copy_rates = m.own_copy_rates,
.first_copy = first_copy,
.base_width = @min(width, first_copy + 1),
.family_mtp = m.mtp and width >= 2 and m.speculate,
@@ -146,7 +156,9 @@ pub const Config = struct {
.batch_rows = batch_rows,
.batch_streams = if (streams) (if (m.max_streams != 0) m.max_streams else 32) else 1,
.node_probabilities = m.draft_probabilities or m.draft_confidence != null,
+ .granted_depth = m.granted_depth and m.draft_probabilities,
.draft_streams = m.draft_streams,
+ .shared_guard = m.shared_guard,
.head_trees = m.head_trees,
.join_tail = m.join_tail,
.piece_runs = m.piece_runs,
@@ -157,7 +169,14 @@ pub const Config = struct {
.shared_costs = try alloc.extendCosts(gpa, timed.items, batch_rows),
.lane_costs = try alloc.extendCosts(gpa, m.lane_costs, 64),
.mtp_step_ms = m.mtp_step_ms,
- .k = .{ .depth_rate = if (m.depth_rate > 0) m.depth_rate else (Constants{}).depth_rate },
+ .k = .{
+ .depth_rate = if (m.depth_rate > 0) m.depth_rate else (Constants{}).depth_rate,
+ .draft_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).draft_slack,
+ .grant_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).grant_slack,
+ .short_slack = if (m.draft_slack > 0) m.draft_slack else (Constants{}).short_slack,
+ .copy_first = m.copy_first and !m.context_copies,
+ .hold_costs = m.hold_costs,
+ },
.draft_step_ms = m.draft_step_ms,
};
}
diff --git a/zig/src/core/lanes/copied_rate_test.zig b/zig/src/core/lanes/copied_rate_test.zig
new file mode 100644
index 0000000000..d037e3bf7c
--- /dev/null
+++ b/zig/src/core/lanes/copied_rate_test.zig
@@ -0,0 +1,56 @@
+//! Model.own_copy_rates: a shared round's copy windows priced by each stream's own copies, else by copy_rate.
+const std = @import("std");
+const Config = @import("config.zig").Config;
+const Cost = @import("config.zig").Cost;
+const Engine = @import("engine.zig").Engine;
+const Stream = @import("stream.zig").Stream;
+const win = @import("windows.zig");
+const fake = @import("fake.zig");
+
+const gpa = std.testing.allocator;
+const copied = [_]u32{ 5, 6, 7, 8, 9, 10, 11 };
+
+/// The copy rows allocation grants each of four streams' seven-row copy windows in one shared round.
+fn granted(e: *Engine, streams: []Stream, out: []usize) !void {
+ _ = e.arena.reset(.retain_capacity);
+ var plans: [4]win.Plan = undefined;
+ for (plans[0..streams.len], streams) |*p, *s| p.* = .{ .stream = s, .position = s.context.items.len, .kind = .copy, .tokens = &copied };
+ try win.allocate(e, plans[0..streams.len]);
+ for (plans[0..streams.len], out) |p, *o| o.* = p.count();
+}
+
+test "own_copy_rates: copies that keep breaking are cut, by source; without it every copy is priced at copy_rate" {
+ // plain rows 10 ms, 0.6 ms a row past four: seven copy rows a stream pay at the prior, not landing 1 a copy
+ var costs: [32]Cost = undefined;
+ for (&costs, 1..) |*c, r| c.* = .{ .width = @intCast(r), .ms = 10 + 0.6 * @as(f64, @floatFromInt(@max(r, 4) - 4)) };
+ for ([_]bool{ false, true }) |own| {
+ var cfg = try Config.init(gpa, .{ .exact_width = 16, .hidden_rows = true, .window_costs = &costs, .batch_rows = 32, .own_copy_rates = own }, 16, 15);
+ defer cfg.deinit(gpa);
+ var target: fake.Fake = .{ .gpa = gpa };
+ defer target.deinit();
+ var clock: fake.FixedClock = .{};
+ var e = Engine.init(gpa, &cfg, target.backend(), clock.clock());
+ defer e.deinit();
+ var streams: [4]Stream = undefined;
+ var made: usize = 0;
+ defer for (streams[0..made]) |*s| s.deinit(gpa);
+ for (&streams) |*s| {
+ s.* = try Stream.init(gpa, .{ .id = "copied", .prompt = &.{ 1, 2 }, .max_new = 64 });
+ made += 1;
+ }
+ var rows: [4]usize = undefined;
+ try granted(&e, &streams, &rows);
+ try std.testing.expectEqualSlices(usize, &.{ 7, 7, 7, 7 }, &rows);
+ for (0..8) |_| win.observeCopied(&e, &streams[0], false, 1, 7); // stream 0's context copies land one row of seven
+ try granted(&e, &streams, &rows);
+ if (own) {
+ try std.testing.expect(rows[0] <= 1);
+ try std.testing.expectEqualSlices(usize, &.{ 7, 7, 7 }, rows[1..]);
+ try std.testing.expect(win.copiedRate(&e, &streams[0], false) < 0.6);
+ try std.testing.expectApproxEqAbs(cfg.copy_rate, win.copiedRate(&e, &streams[0], true), 1e-12); // Recall's keep theirs
+ } else {
+ try std.testing.expectEqualSlices(usize, &.{ 7, 7, 7, 7 }, &rows);
+ try std.testing.expectEqual(@as(f64, -1), streams[0].copied_landed[0]); // no history kept
+ }
+ }
+}
diff --git a/zig/src/core/lanes/depth.zig b/zig/src/core/lanes/depth.zig
index 8595a67aab..d5eca5704a 100644
--- a/zig/src/core/lanes/depth.zig
+++ b/zig/src/core/lanes/depth.zig
@@ -12,8 +12,16 @@ pub const Constants = struct {
plain_margin: f64 = 1.05, // a drafted depth must beat plain rounds' tokens a ms by this much
plain_wait_most: u64 = 128, // plain rounds between probes at most
draft_slack: i64 = 2, // drafts a stream offers past what its last shared round granted it
+ grant_slack: i64 = 1, // granted depths: drafts past a grant that verified every draft the stream held
+ short_slack: i64 = 0, // granted depths: drafts past a grant that left some of them unverified
+ copy_first: bool = false, // a copy ahead takes the next round whole: the head drafts none for it (else one)
+ hold_costs: []const HoldCost = &.{}, // a drafter's own hold pass by blocks and depth: granted holds must pay for it
};
+pub const HoldCost = @import("hold_price.zig").HoldCost;
+const holdMs = @import("hold_price.zig").holdMs;
+const lone_guard = @import("lone_guard.zig");
+
/// A stream's acceptance by depth and probe counters (Python `_depth_state[stream]`).
pub const State = struct {
p: []f64,
@@ -21,6 +29,7 @@ pub const State = struct {
plain: ?u64 = null,
wait: ?u64 = null,
ms: ?Table = null, // plain_guard: the stream's own round times
+ lone: lone_guard.Lone = .{}, // lone_guard: whether its lone holds draft
pub fn deinit(st: *State, gpa: Allocator) void {
gpa.free(st.p);
@@ -35,8 +44,11 @@ pub const Who = struct {
finished: bool,
forced: bool,
granted: ?u32 = null,
+ grant_short: bool = false, // the last grant left some held drafts unverified (the width did not bind)
+ gain: ?f64 = null, // the drafts its last grant expected to keep (null: none priced)
stream: ?*anyopaque = null,
sampled: bool = false, // keyed draws: host work a row that a greedy round does not have
+ shared: bool = false, // a shared round's budget: the lone guard stays out
};
/// Whether a copied continuation is ahead of a stream (the round loop proposes; Python `_copy_proposal`).
@@ -59,6 +71,10 @@ pub const Rule = struct {
draft_step_ms: f64 = 0, // a shared round's draft past a stream's first (budgets price it)
plain_guard: bool,
node_probabilities: bool,
+ granted_depth: bool = false, // node_probabilities: a stream drafts its last grant plus draft_slack, lone ones too
+ lone_guard: bool = false, // granted_depth: a lone stream holds drafts only while its drafted rounds measure faster
+ lone_k: lone_guard.Constants = .{},
+ lone_plain: lone_guard.Plain = .{}, // plain rounds' tokens a ms, every lone stream's
confidence_floor: ?f64 = null,
batch_rows: u32,
batched_heads: bool = false, // one batch drafts every stream of a shared round (its steps: the deepest stream's)
@@ -192,17 +208,123 @@ pub const Rule = struct {
return true;
}
- /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 with a copy ahead).
+ /// Drafts past a grant: none when it stopped short of the held drafts (the width did not bind), else grant_slack.
+ fn slack(r: *const Rule, who: Who) i64 {
+ return if (who.grant_short) r.k.short_slack else r.k.grant_slack;
+ }
+
+ /// The last grant plus its slack (the deepest chain before any); after a grant of none, a probe of draft_slack.
+ fn grantDepth(r: *Rule, who: Who) !i64 {
+ const most = @min(@as(i64, r.most_drafts), @max(1, who.draft_room - 1));
+ if (r.lone_guard and !who.shared) {
+ const l = &(try r.state(who)).lone;
+ if (!l.drafts()) return 0;
+ if (l.fresh) {
+ l.fresh = false;
+ return @min(most, 1 + r.k.grant_slack);
+ }
+ }
+ const granted = who.granted orelse return most;
+ const st = try r.state(who);
+ if (granted == 0) return if (!r.plain_guard or r.probePlain(st)) @min(most, r.k.draft_slack) else 0;
+ st.plain = 0;
+ st.wait = r.k.depth_probe_every;
+ return @min(most, @as(i64, granted) + r.slack(who));
+ }
+
+ fn state(r: *Rule, who: Who) !*State {
+ _ = try r.rates(who);
+ return &who.state.*.?;
+ }
+
+ /// A lone round's committed tokens and wall ms at the stream's context band, for the lone guard.
+ pub fn observeLone(r: *Rule, who: Who, band: usize, tokens: usize, ms: f64) !void {
+ if (!r.lone_guard) return;
+ const memo = r.lone_plain.memo(who.sampled, band);
+ (try r.state(who)).lone.observeBand(r.lone_k, r.lone_plain.at(who.sampled, band), memo, tokens, ms);
+ }
+
+ /// Python `_head_depth`: the next round's head drafts for a stream (0 when forced, 1 or none with a copy ahead).
pub fn headDepth(r: *Rule, who: *const Who, budget: ?i64, probe: Probe) !i64 {
- const d: i64 = if (who.finished) 0 else if (budget) |b| b else if (r.node_probabilities)
+ const d: i64 = if (who.finished) 0 else if (budget) |b| b else if (r.granted_depth)
+ try r.grantDepth(who.*)
+ else if (r.node_probabilities)
@min(@as(i64, r.most_drafts), @max(1, who.draft_room - 1))
else
try r.depth(who.*);
if (who.forced) return 0;
- if (d != 0 and probe.copy(probe.ptr, who)) return 1;
+ if (d != 0 and probe.copy(probe.ptr, who)) return @intFromBool(!r.k.copy_first);
return d;
}
+ /// A shared hold pads every block to its deepest: depths cut so the drafting streams' blocks fit `rows` (one pass).
+ fn capPass(out: []i64, rows: u32) void {
+ var n: i64 = 0;
+ for (out) |d| n += @intFromBool(d > 0);
+ if (n == 0) return;
+ const most = @max(1, @divFloor(@as(i64, rows), n));
+ for (out) |*d| d.* = @min(d.*, most);
+ }
+
+ /// Granted holds kept, most expected drafts a row first, while they raise the next round's expected tokens a ms.
+ fn payHolds(r: *Rule, whos: []const Who, sampled: bool, out: []i64) !void {
+ const order = try r.gpa.alloc(usize, whos.len);
+ defer r.gpa.free(order);
+ var m: usize = 0;
+ var probes: usize = 0; // a probe after a grant of none keeps its hold: it measures the stream again
+ var deepest: i64 = 0;
+ for (whos, out, 0..) |w, d, i| {
+ if (d <= 0) continue;
+ if (w.granted == @as(?u32, 0)) {
+ probes += 1;
+ deepest = @max(deepest, d);
+ } else {
+ order[m] = i;
+ m += 1;
+ }
+ }
+ if (m == 0) return;
+ const Ctx = struct {
+ r: *Rule,
+ whos: []const Who,
+ fn less(c: @This(), x: usize, y: usize) bool {
+ const gx = c.r.gainOf(c.whos[x]) catch 0;
+ const gy = c.r.gainOf(c.whos[y]) catch 0;
+ return gx * rowsOf(c.whos[y]) > gy * rowsOf(c.whos[x]) or (gx * rowsOf(c.whos[y]) == gy * rowsOf(c.whos[x]) and x < y);
+ }
+ };
+ std.mem.sort(usize, order[0..m], Ctx{ .r = r, .whos = whos }, Ctx.less);
+ const fixed = r.fixedMs(whos.len, sampled);
+ var rows: f64 = @floatFromInt(whos.len);
+ var tokens = rows;
+ const plain = r.shared_costs.get(@intCast(whos.len)) orelse return;
+ const probe_ms = if (probes > 0) holdMs(r.k.hold_costs, probes, deepest) else 0;
+ var best = tokens / (plain + fixed + probe_ms);
+ var keep: usize = 0;
+ for (order[0..m], 1..) |i, k| {
+ rows += rowsOf(whos[i]);
+ tokens += try r.gainOf(whos[i]);
+ deepest = @max(deepest, out[i]);
+ const verify = r.shared_costs.get(@intFromFloat(@ceil(rows))) orelse break;
+ const now = tokens / (verify + fixed + holdMs(r.k.hold_costs, k + probes, deepest));
+ if (now > best) {
+ best = now;
+ keep = k;
+ }
+ }
+ for (order[keep..m]) |i| out[i] = 0;
+ }
+
+ /// The drafts a granted stream's next round is expected to keep: its last grant's, or its first draft's rate.
+ fn gainOf(r: *Rule, w: Who) !f64 {
+ if (w.gain) |g| if (w.granted != null) return g;
+ return (try r.rates(w))[0];
+ }
+
+ fn rowsOf(w: Who) f64 {
+ return @floatFromInt(@max(1, w.granted orelse 1));
+ }
+
/// Head drafts for the streams of a shared round, sized by allocation over their chain chances.
pub fn budgets(r: *Rule, whos: []Who, probe: Probe, sampled: bool, out: []i64) !void {
if (r.confidence_floor != null) {
@@ -211,10 +333,13 @@ pub const Rule = struct {
}
if (r.node_probabilities) {
for (whos, out) |*w, *o| {
+ w.shared = true;
const inner = try r.headDepth(w, null, probe);
const granted: i64 = if (w.granted) |g| g else r.most_drafts;
- o.* = try r.headDepth(w, @min(inner, granted + r.k.draft_slack), probe);
+ o.* = try r.headDepth(w, if (r.granted_depth) inner else @min(inner, granted + r.k.draft_slack), probe);
}
+ if (r.granted_depth) capPass(out, r.batch_rows);
+ if (r.granted_depth and r.k.hold_costs.len > 0) try r.payHolds(whos, sampled, out);
return;
}
const most = try r.gpa.alloc(i64, whos.len);
@@ -381,3 +506,81 @@ test "a shared round's drafts come out the same whichever depths its overhead wa
}
try std.testing.expectEqualSlices(i64, &got[0], &got[1]);
}
+
+test "granted depths in a shared round fit one hold pass: every block is as deep as the deepest" {
+ const empty: Table = .{};
+ var rule: Rule = .{ .gpa = std.testing.allocator, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &empty, .mtp_step_ms = 0.5, .plain_guard = false, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 };
+ defer rule.deinit();
+ var states: [8]?State = @splat(null);
+ defer for (&states) |*st| if (st.*) |*x| x.deinit(std.testing.allocator);
+ var whos: [8]Who = undefined;
+ for (&whos, &states, 0..) |*w, *st, i| w.* = .{ .state = st, .draft_room = 100, .finished = false, .forced = false, .granted = @intCast(i) };
+ const No = struct {
+ fn copy(_: *anyopaque, _: *const Who) bool {
+ return false;
+ }
+ };
+ var unused: u8 = 0;
+ var out: [8]i64 = undefined;
+ try rule.budgets(&whos, .{ .ptr = &unused, .copy = No.copy }, false, &out);
+ // a grant of none probes two; grants 1..7 that took every held draft, plus one; cut to 32 rows over eight blocks
+ try std.testing.expectEqualSlices(i64, &.{ 2, 2, 3, 4, 4, 4, 4, 4 }, &out);
+ try rule.budgets(whos[0..2], .{ .ptr = &unused, .copy = No.copy }, false, out[0..2]);
+ try std.testing.expectEqualSlices(i64, &.{ 2, 2 }, out[0..2]); // two blocks fit whole
+ for (whos[0..2]) |*w| w.grant_short = true;
+ try rule.budgets(whos[0..2], .{ .ptr = &unused, .copy = No.copy }, false, out[0..2]);
+ try std.testing.expectEqualSlices(i64, &.{ 2, 1 }, out[0..2]); // a grant short of the held drafts: no slack
+}
+
+test "a copy ahead: one head draft beside it, or none when the head's drafts cost a pass of their own" {
+ const empty: Table = .{};
+ const Yes = struct {
+ fn copy(_: *anyopaque, _: *const Who) bool {
+ return true;
+ }
+ };
+ var unused: u8 = 0;
+ const probe: Probe = .{ .ptr = &unused, .copy = Yes.copy };
+ var st: ?State = null;
+ defer if (st) |*x| x.deinit(std.testing.allocator);
+ const who: Who = .{ .state = &st, .draft_room = 100, .finished = false, .forced = false, .granted = 3 };
+ for ([_]bool{ false, true }) |first| {
+ var rule: Rule = .{ .gpa = std.testing.allocator, .k = .{ .copy_first = first }, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &empty, .mtp_step_ms = 0.5, .plain_guard = false, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 };
+ defer rule.deinit();
+ try std.testing.expectEqual(@as(i64, @intFromBool(!first)), try rule.headDepth(&who, null, probe));
+ var out: [1]i64 = undefined;
+ var whos = [_]Who{who};
+ try rule.budgets(&whos, probe, false, &out);
+ try std.testing.expectEqual(@as(i64, @intFromBool(!first)), out[0]);
+ }
+}
+
+test "a granted hold stays only while its stream's expected drafts pay for the verify rows and the hold pass" {
+ const gpa = std.testing.allocator;
+ var shared: Table = .{};
+ defer shared.deinit(gpa);
+ for (1..33) |w| try shared.put(gpa, @intCast(w), 10.0 + @as(f64, @floatFromInt(w)));
+ const empty: Table = .{};
+ const cheap = [_]HoldCost{ .{ .streams = 1, .depth = 1, .ms = 1.0 }, .{ .streams = 8, .depth = 1, .ms = 4.5 } };
+ var states: [4]?State = @splat(null);
+ defer for (&states) |*st| if (st.*) |*x| x.deinit(gpa);
+ const granted = [_]?u32{ 2, 2, 0, 1 };
+ const gains = [_]?f64{ 1.8, 0.2, null, 0.9 };
+ var whos: [4]Who = undefined;
+ for (&whos, &states, granted, gains) |*w, *st, g, x| w.* = .{ .state = st, .draft_room = 100, .finished = false, .forced = false, .granted = g, .gain = x };
+ var rule: Rule = .{ .gpa = gpa, .k = .{ .hold_costs = &cheap }, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &shared, .mtp_step_ms = 0.5, .plain_guard = true, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 };
+ defer rule.deinit();
+ // 8 ms overhead: plain 4 / 22; stream 0 then 3 raise the rate (5.8 / 25, 6.7 / 26.5), stream 1 would lower it
+ var out = [_]i64{ 4, 4, 0, 3 };
+ try rule.payHolds(&whos, false, &out);
+ try std.testing.expectEqualSlices(i64, &.{ 4, 0, 0, 3 }, &out);
+ const dear = [_]HoldCost{.{ .streams = 1, .depth = 1, .ms = 20.0 }};
+ rule.k.hold_costs = &dear;
+ out = .{ 4, 4, 0, 3 };
+ try rule.payHolds(&whos, false, &out);
+ try std.testing.expectEqualSlices(i64, &.{ 0, 0, 0, 0 }, &out); // no gain pays a 20 ms pass
+ whos[2].granted = 0;
+ out = .{ 4, 4, 2, 3 };
+ try rule.payHolds(&whos, false, &out);
+ try std.testing.expectEqualSlices(i64, &.{ 0, 0, 2, 0 }, &out); // a probe after a grant of none keeps its hold
+}
diff --git a/zig/src/core/lanes/drafted.zig b/zig/src/core/lanes/drafted.zig
index 4c509b103a..ca4a4f7986 100644
--- a/zig/src/core/lanes/drafted.zig
+++ b/zig/src/core/lanes/drafted.zig
@@ -29,6 +29,8 @@ pub const Drafted = struct {
absorbs: std.ArrayList(dr.Absorb) = .empty,
holds: std.ArrayList(dr.Hold) = .empty,
firsts: std.ArrayList(u32) = .empty,
+ /// The target reads held drafts from the drafter itself (Window.held), so windows without fillers pass as they are.
+ held_on_target: bool = false,
/// The wrapper, with the target told the drafter's taps before any forward.
pub fn init(gpa: Allocator, target: be.Backend, drafter: dr.Drafter) !Drafted {
@@ -103,6 +105,16 @@ pub const Drafted = struct {
m.draft_streams = d.batched;
m.draft_probabilities = false;
m.head_trees = d.trees;
+ m.shared_guard = d.shared_guard;
+ m.copy_first = d.copy_first;
+ m.hold_costs = d.hold_costs;
+ if (d.chances and x.drafter.vtable.chances != null) {
+ // a chain's drafts priced by the drafter's own chances; the next hold as deep as the grant and its slack
+ m.draft_probabilities = d.confidence_floor == null;
+ m.draft_confidence = d.confidence_floor;
+ m.granted_depth = d.confidence_floor == null;
+ m.draft_slack = d.slack;
+ }
if (d.tree_nodes > 0) {
// the drafter grows its own trees, best first, and prices their nodes
m.drafts = @min(d.tree_nodes, if (target.exact_width > 1) target.exact_width - 1 else 0);
@@ -217,11 +229,13 @@ pub const Drafted = struct {
lane.tree_first = null;
};
var total: usize = 0;
+ var through = x.held_on_target;
for (windows) |w| {
if (w.held > 0 and w.tokens.len != 0) return error.MixedDrafts;
total += w.held;
+ if (through and w.held > 0 and (x.lanes.getPtr(w.stream) orelse return error.NotDrafting).filler > 0) through = false;
}
- if (total == 0) return x.target.verify(x.windows.items, out);
+ if (total == 0 or through) return x.target.verify(x.windows.items, out);
try x.tokens.resize(x.gpa, total);
x.streams.clearRetainingCapacity();
x.outs.clearRetainingCapacity();
diff --git a/zig/src/core/lanes/drafted_chances_test.zig b/zig/src/core/lanes/drafted_chances_test.zig
new file mode 100644
index 0000000000..d3f248e8da
--- /dev/null
+++ b/zig/src/core/lanes/drafted_chances_test.zig
@@ -0,0 +1,266 @@
+//! A drafter that prices its chains (Facts.chances): verify rows go by its chances and replies stay one-token decode.
+const std = @import("std");
+const Config = @import("config.zig").Config;
+const Model = @import("config.zig").Model;
+const Cost = @import("config.zig").Cost;
+const Engine = @import("engine.zig").Engine;
+const be = @import("backend.zig");
+const sm = @import("stream.zig");
+const fake = @import("fake.zig");
+const dr = @import("drafter.zig");
+const depth = @import("depth.zig");
+const Table = @import("table.zig").Table;
+const Drafted = @import("drafted.zig").Drafted;
+const dt = @import("drafted_test.zig");
+
+const gpa = std.testing.allocator;
+const most_streams = 8;
+
+/// drafted_test's drafter with a chance for each held draft: `sure` where its guess is right, `unsure` where wrong.
+pub const Priced = struct {
+ head: dt.FakeDrafter,
+ sure: f64 = 0.7,
+ unsure: f64 = 0.7,
+ vtable: dr.Drafter.VTable = undefined,
+ inner_hold: *const fn (*anyopaque, []const dr.Hold) anyerror!void = undefined,
+ at: std.AutoHashMapUnmanaged(*sm.Stream, u64) = .empty, // each stream's last hold's position
+ asked: u32 = 0,
+ held: u64 = 0, // drafts its holds asked for
+
+ fn of(ptr: *anyopaque) *Priced {
+ return @fieldParentPtr("head", @as(*dt.FakeDrafter, @ptrCast(@alignCast(ptr))));
+ }
+
+ fn hold(ptr: *anyopaque, items: []const dr.Hold) anyerror!void {
+ const p = of(ptr);
+ for (items) |it| {
+ try p.at.put(gpa, it.stream, it.position);
+ p.held += it.depth;
+ }
+ return p.inner_hold(ptr, items);
+ }
+
+ /// The running product along the chain, as the DSpark confidence head's survivals are.
+ fn chances(ptr: *anyopaque, s: *sm.Stream, out: []f64) anyerror!bool {
+ const p = of(ptr);
+ const lane = p.head.lanes.getPtr(s) orelse return false;
+ if (out.len > lane.held.items.len) return error.TooFewDrafts;
+ const position = p.at.get(s) orelse return error.NoHold;
+ p.asked += 1;
+ var reach: f64 = 1;
+ for (out, 0..) |*o, j| {
+ const at = position + j;
+ reach *= if ((at *% 2654435761 + j) % 5 == 0) p.unsure else p.sure; // drafted_test's wrong guesses
+ o.* = reach;
+ }
+ return true;
+ }
+
+ pub fn drafter(p: *Priced) dr.Drafter {
+ const base_ = p.head.drafter();
+ p.vtable = base_.vtable.*;
+ p.inner_hold = base_.vtable.hold;
+ p.vtable.hold = hold;
+ p.vtable.chances = chances;
+ return .{ .ptr = &p.head, .vtable = &p.vtable };
+ }
+
+ pub fn deinit(p: *Priced) void {
+ p.at.deinit(gpa);
+ p.head.deinit();
+ }
+};
+
+/// Every shared verify's window rows: whether one round gave its streams unequal rows (a ragged verify).
+const Seen = struct {
+ var inner: *const fn (*anyopaque, []const be.Window, []be.Verified) anyerror!void = undefined;
+ var ragged: u32 = 0;
+ var rows: u64 = 0;
+
+ fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void {
+ for (windows) |w| {
+ rows += w.rows();
+ if (w.rows() != windows[0].rows()) {
+ ragged += 1;
+ break;
+ }
+ }
+ return inner(ptr, windows, out);
+ }
+};
+
+var costs: [16]Cost = undefined;
+
+pub fn base() Model {
+ for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) };
+ return .{ .exact_width = 16, .gpu_tokens = true, .window_costs = &costs, .hidden_rows = true, .batch_rows = 32, .max_streams = most_streams };
+}
+
+const Run = struct { tokens: [most_streams][64]u32 = undefined, lens: [most_streams]usize = @splat(0), drafted: u64 = 0, holds: u32 = 0, asked: u32 = 0, ragged: u32 = 0, rows: u64 = 0, held: u64 = 0 };
+
+/// The drafter's Facts.slack for the next runs (0: the depth rule's own; 9: every hold the whole block).
+var slack: u32 = 0;
+
+/// The drafter's chances: equal everywhere, or knowing its wrong guesses (a confidence head that ranks well).
+const Head = enum { flat, knows };
+
+/// Each prompt decoded together (one stream: alone); `floor` null: the chances price rows, else a fixed window.
+fn run(prompts: []const []const u32, drafts: bool, floor: ?f64, head: Head) !Run {
+ var target: fake.Fake = .{ .gpa = gpa };
+ defer target.deinit();
+ var p: Priced = .{ .head = .{ .facts_ = .{ .depth = 9, .step_ms = 0.5, .batched = true, .chances = true, .confidence_floor = floor, .slack = slack } } };
+ if (head == .knows) {
+ p.sure = 0.95;
+ p.unsure = 0.05;
+ }
+ defer p.deinit();
+ var strict: dt.Strict = .{ .inner = target.backend() };
+ defer strict.deinit();
+ var seen = strict.backend();
+ var table = seen.vtable.*;
+ Seen.inner = table.verify;
+ Seen.ragged = 0;
+ Seen.rows = 0;
+ table.verify = Seen.verify;
+ seen.vtable = &table;
+ var wrapped = try Drafted.init(gpa, seen, p.drafter());
+ defer wrapped.deinit();
+ var cfg = try Config.init(gpa, if (drafts) wrapped.facts(base()) else base(), 16, 15);
+ defer cfg.deinit(gpa);
+ var clock: fake.FixedClock = .{};
+ var engine = Engine.init(gpa, &cfg, if (drafts) wrapped.backend() else target.backend(), clock.clock());
+ defer engine.deinit();
+ var streams: [most_streams]sm.Stream = undefined;
+ const n = prompts.len;
+ for (prompts, streams[0..n]) |prompt, *s| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = prompt, .max_new = 48, .eos = &.{96}, .drafts = drafts });
+ defer for (streams[0..n]) |*s| s.deinit(gpa);
+ for (streams[0..n]) |*s| try engine.addStream(s);
+ while (engine.activeCount() > 0) try engine.step();
+ var out: Run = .{ .holds = p.head.holds, .asked = p.asked, .ragged = Seen.ragged, .rows = Seen.rows, .held = p.held };
+ for (streams[0..n], 0..) |*s, i| {
+ out.lens[i] = s.emitted().len;
+ @memcpy(out.tokens[i][0..out.lens[i]], s.emitted());
+ out.drafted += s.drafted;
+ }
+ return out;
+}
+
+fn expectSame(a: Run, b: Run, n: usize) !void {
+ for (0..n) |i| try std.testing.expectEqualSlices(u32, a.tokens[i][0..a.lens[i]], b.tokens[i][0..b.lens[i]]);
+}
+
+/// `n` streams of unequal prompts, each a prefix of drafted_test's long prompt.
+fn unequal(buf: *[most_streams][]const u32, n: usize) []const []const u32 {
+ for (buf[0..n], 0..) |*p, i| p.* = dt.long[0 .. 11 + i];
+ return buf[0..n];
+}
+
+test "a drafter's chances price its chain's rows, and drafted replies equal one-token decode alone and shared" {
+ const sets = [_][]const []const u32{ &.{&dt.p1}, &.{ &dt.p1, &dt.long, &dt.p1 } };
+ for (sets) |ps| {
+ const plain = try run(ps, false, null, .flat);
+ const priced = try run(ps, true, null, .flat);
+ try expectSame(plain, priced, ps.len);
+ try std.testing.expect(priced.asked > 0 and priced.drafted > 0); // the window allocator read the chances
+ const fixed = try run(ps, true, 0, .flat);
+ try expectSame(plain, fixed, ps.len);
+ try std.testing.expect(fixed.drafted > priced.drafted); // a fixed window verifies every draft the rows fit
+ }
+}
+
+test "at 1, 4 and 8 streams the scheduled and the fixed verify equal one-token decode; shared rounds are ragged" {
+ var buf: [most_streams][]const u32 = undefined;
+ for ([_]usize{ 1, 4, 8 }) |n| {
+ const ps = unequal(&buf, n);
+ const plain = try run(ps, false, null, .flat);
+ const scheduled = try run(ps, true, null, .knows);
+ try expectSame(plain, scheduled, n);
+ const fixed = try run(ps, true, 0, .knows);
+ try expectSame(plain, fixed, n);
+ try std.testing.expect(scheduled.asked > 0 and scheduled.drafted > 0);
+ // a confident head's survivals cut each stream's chain where its guess goes wrong: fewer rows verified
+ try std.testing.expect(scheduled.rows < fixed.rows);
+ if (n > 1) try std.testing.expect(scheduled.ragged > 0); // streams' verify lengths differ within a round
+ }
+}
+
+test "the drafter's chances reach the lane facts: rows by chances and holds by grants, or a fixed window" {
+ var p: Priced = .{ .head = .{ .facts_ = .{ .depth = 9, .step_ms = 0.5, .chances = true } } };
+ defer p.deinit();
+ var target: fake.Fake = .{ .gpa = gpa };
+ defer target.deinit();
+ var wrapped: Drafted = .{ .gpa = gpa, .target = target.backend(), .drafter = p.drafter() };
+ defer wrapped.deinit();
+ const m = wrapped.facts(base());
+ try std.testing.expect(m.draft_probabilities and m.granted_depth and m.draft_confidence == null);
+ var cfg = try Config.init(gpa, m, 16, 15);
+ defer cfg.deinit(gpa);
+ try std.testing.expectEqual(@as(i64, 2), cfg.k.draft_slack); // the depth rule's own
+ try std.testing.expect(cfg.k.grant_slack == 1 and cfg.k.short_slack == 0); // a granted width's own
+ p.head.facts_.slack = 1;
+ var one = try Config.init(gpa, wrapped.facts(base()), 16, 15);
+ defer one.deinit(gpa);
+ try std.testing.expectEqual(@as(i64, 1), one.k.draft_slack); // a drafter's own slack past its grant
+ try std.testing.expect(one.k.grant_slack == 1 and one.k.short_slack == 1); // short or not
+ p.head.facts_.confidence_floor = 0;
+ const f = wrapped.facts(base());
+ try std.testing.expect(!f.draft_probabilities and !f.granted_depth and f.draft_confidence.? == 0);
+ p.head.facts_.chances = false;
+ const none = wrapped.facts(base());
+ try std.testing.expect(!none.draft_probabilities and none.draft_confidence == null);
+}
+
+test "a granted depth: the deepest chain first, then the grant plus its slack, and after a grant of none a probe" {
+ const empty: Table = .{};
+ var rule: depth.Rule = .{ .gpa = gpa, .prior = &.{0.8}, .most_drafts = 9, .family_costs = &empty, .shared_costs = &empty, .mtp_step_ms = 0.5, .plain_guard = true, .node_probabilities = true, .granted_depth = true, .batch_rows = 32 };
+ defer rule.deinit();
+ var st: ?depth.State = null;
+ defer if (st) |*s| s.deinit(gpa);
+ const No = struct {
+ fn copy(_: *anyopaque, _: *const depth.Who) bool {
+ return false;
+ }
+ };
+ var unused: u8 = 0;
+ const probe: depth.Probe = .{ .ptr = &unused, .copy = No.copy };
+ var who: depth.Who = .{ .state = &st, .draft_room = 100, .finished = false, .forced = false };
+ try std.testing.expectEqual(@as(i64, 9), try rule.headDepth(&who, null, probe));
+ who.granted = 3;
+ try std.testing.expectEqual(@as(i64, 4), try rule.headDepth(&who, null, probe)); // every held draft granted: one more
+ who.grant_short = true;
+ try std.testing.expectEqual(@as(i64, 3), try rule.headDepth(&who, null, probe)); // drafts left unverified: the grant
+ who.grant_short = false;
+ who.granted = 8;
+ try std.testing.expectEqual(@as(i64, 9), try rule.headDepth(&who, null, probe));
+ rule.k.grant_slack = 2; // a drafter's slack sets both
+ rule.k.short_slack = 2;
+ who.grant_short = true;
+ who.granted = 3;
+ try std.testing.expectEqual(@as(i64, 5), try rule.headDepth(&who, null, probe));
+ who.granted = 0;
+ var drafted: usize = 0;
+ for (0..8) |_| drafted += @intFromBool(try rule.headDepth(&who, null, probe) > 0);
+ try std.testing.expectEqual(@as(usize, 1), drafted); // one probe of draft_slack drafts in the first eight rounds
+}
+
+test "holds as wide as the grants: replies equal one-token decode at 1, 4 and 8 streams, fewer drafts than whole blocks" {
+ var buf: [most_streams][]const u32 = undefined;
+ defer slack = 0;
+ for ([_]usize{ 1, 4, 8 }) |n| {
+ const ps = unequal(&buf, n);
+ const plain = try run(ps, false, null, .flat);
+ for ([_]Head{ .flat, .knows }) |head| {
+ slack = 0;
+ const narrow = try run(ps, true, null, head);
+ slack = 2;
+ const two = try run(ps, true, null, head);
+ slack = 9;
+ const whole = try run(ps, true, null, head);
+ try expectSame(plain, narrow, n);
+ try expectSame(plain, two, n);
+ try expectSame(plain, whole, n);
+ try std.testing.expect(narrow.drafted > 0);
+ try std.testing.expect(narrow.held < two.held and two.held < whole.held);
+ }
+ }
+}
diff --git a/zig/src/core/lanes/drafted_copy_test.zig b/zig/src/core/lanes/drafted_copy_test.zig
new file mode 100644
index 0000000000..329731f69c
--- /dev/null
+++ b/zig/src/core/lanes/drafted_copy_test.zig
@@ -0,0 +1,135 @@
+//! A drafter whose hold costs a pass of its own: no hold for a stream whose next round is a copy.
+const std = @import("std");
+const Config = @import("config.zig").Config;
+const Engine = @import("engine.zig").Engine;
+const be = @import("backend.zig");
+const sm = @import("stream.zig");
+const fake = @import("fake.zig");
+const dr = @import("drafter.zig");
+const Proposer = @import("proposer.zig").Proposer;
+const Drafted = @import("drafted.zig").Drafted;
+const dt = @import("drafted_test.zig");
+
+const gpa = std.testing.allocator;
+
+/// Copies the fake target's own continuation (it lands whole) whenever the context's length is a multiple of three.
+const Copier = struct {
+ out: [3]u32 = undefined,
+ match: i64 = 0,
+
+ fn propose(ptr: *anyopaque, context: []const u32, most: i64) anyerror![]const u32 {
+ const c: *Copier = @ptrCast(@alignCast(ptr));
+ c.match = 0;
+ if (context.len % 3 != 0 or most < 2) return &.{};
+ var history: [128]u32 = undefined;
+ @memcpy(history[0..context.len], context);
+ const n: usize = @intCast(@min(3, most));
+ for (c.out[0..n], 0..) |*t, k| {
+ t.* = fake.next(history[0 .. context.len + k], null, 0);
+ history[context.len + k] = t.*;
+ }
+ c.match = 100;
+ return c.out[0..n];
+ }
+
+ fn lastMatch(ptr: *anyopaque) i64 {
+ return @as(*Copier, @ptrCast(@alignCast(ptr))).match;
+ }
+
+ const vtable: Proposer.VTable = .{ .propose = propose, .last_match = lastMatch };
+};
+
+/// Each stream's last hold position, and the copy windows verified at a position a hold was asked for.
+const Watch = struct {
+ var inner_hold: *const fn (*anyopaque, []const dr.Hold) anyerror!void = undefined;
+ var inner_verify: *const fn (*anyopaque, []const be.Window, []be.Verified) anyerror!void = undefined;
+ var at: [4]?u64 = @splat(null);
+ var streams: [4]?*sm.Stream = @splat(null);
+ var wasted: u32 = 0;
+ var copies: u32 = 0;
+ var blocks: u32 = 0; // the streams' blocks held, over every hold
+
+ fn slot(s: *sm.Stream) usize {
+ for (&streams, 0..) |*x, i| if (x.* == null or x.* == s) {
+ x.* = s;
+ return i;
+ };
+ unreachable;
+ }
+
+ fn hold(ptr: *anyopaque, items: []const dr.Hold) anyerror!void {
+ blocks += @intCast(items.len);
+ // a joining stream's first hold comes before its first token is read: no copy can be foreseen for it
+ for (items) |it| at[slot(it.stream)] = if (it.stream.emitted().len > 0) it.position else null;
+ return inner_hold(ptr, items);
+ }
+
+ fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void {
+ for (windows) |w| if (w.held == 0 and w.tokens.len > 0) {
+ copies += 1;
+ if (at[slot(w.stream)] == w.positions[0]) wasted += 1;
+ };
+ return inner_verify(ptr, windows, out);
+ }
+};
+
+const Run = struct { tokens: [3][64]u32 = undefined, lens: [3]usize = @splat(0), holds: u32 = 0, wasted: u32 = 0, copies: u32 = 0 };
+
+fn run(n: usize, drafts: bool, copy_first: bool) !Run {
+ Watch.at = @splat(null);
+ Watch.streams = @splat(null);
+ Watch.wasted = 0;
+ Watch.copies = 0;
+ Watch.blocks = 0;
+ var target: fake.Fake = .{ .gpa = gpa };
+ defer target.deinit();
+ var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .copy_first = copy_first } };
+ defer head.deinit();
+ var hooked = head.drafter();
+ var table = hooked.vtable.*;
+ Watch.inner_hold = table.hold;
+ table.hold = Watch.hold;
+ hooked.vtable = &table;
+ var wrapped = try Drafted.init(gpa, target.backend(), hooked);
+ defer wrapped.deinit();
+ // the round loop's windows, before the wrapper turns held drafts into host tokens
+ var seen = if (drafts) wrapped.backend() else target.backend();
+ var verified = seen.vtable.*;
+ Watch.inner_verify = verified.verify;
+ verified.verify = Watch.verify;
+ seen.vtable = &verified;
+ const m: @import("config.zig").Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 };
+ var cfg = try Config.init(gpa, if (drafts) wrapped.facts(m) else m, 16, 15);
+ defer cfg.deinit(gpa);
+ var clock: fake.FixedClock = .{};
+ var engine = Engine.init(gpa, &cfg, seen, clock.clock());
+ defer engine.deinit();
+ var copiers: [3]Copier = @splat(.{});
+ var streams: [3]sm.Stream = undefined;
+ for (streams[0..n], copiers[0..n], 0..) |*s, *c, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 12 + i], .max_new = 40, .eos = &.{96}, .drafts = drafts, .proposer = .{ .ptr = c, .vtable = &Copier.vtable } });
+ defer for (streams[0..n]) |*s| s.deinit(gpa);
+ for (streams[0..n]) |*s| try engine.addStream(s);
+ while (engine.activeCount() > 0) try engine.step();
+ var out: Run = .{ .holds = Watch.blocks, .wasted = Watch.wasted, .copies = Watch.copies };
+ for (streams[0..n], 0..) |*s, i| {
+ out.lens[i] = s.emitted().len;
+ @memcpy(out.tokens[i][0..out.lens[i]], s.emitted());
+ }
+ return out;
+}
+
+test "a copy ahead takes the round: with copy_first no hold is asked for it, and replies equal one-token decode" {
+ for ([_]usize{ 1, 3 }) |n| {
+ const plain = try run(n, false, false);
+ const before = try run(n, true, false);
+ const after = try run(n, true, true);
+ for (0..n) |i| {
+ try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], before.tokens[i][0..before.lens[i]]);
+ try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], after.tokens[i][0..after.lens[i]]);
+ }
+ try std.testing.expect(before.copies > 0 and after.copies > 0);
+ try std.testing.expect(before.wasted > 0); // a hold of one draft beside each copy, dropped by the copy round
+ try std.testing.expectEqual(@as(u32, 0), after.wasted);
+ try std.testing.expect(after.holds < before.holds);
+ }
+}
diff --git a/zig/src/core/lanes/drafted_guard_test.zig b/zig/src/core/lanes/drafted_guard_test.zig
new file mode 100644
index 0000000000..3715bec357
--- /dev/null
+++ b/zig/src/core/lanes/drafted_guard_test.zig
@@ -0,0 +1,155 @@
+//! A drafter that asks for the shared-round guard (Facts.shared_guard): which shared rounds draft, replies unchanged.
+const std = @import("std");
+const Config = @import("config.zig").Config;
+const Model = @import("config.zig").Model;
+const Engine = @import("engine.zig").Engine;
+const sm = @import("stream.zig");
+const fake = @import("fake.zig");
+const Drafted = @import("drafted.zig").Drafted;
+const dt = @import("drafted_test.zig");
+
+const gpa = std.testing.allocator;
+
+test "a stream joining shared rounds that the guard runs plain holds no first drafts; a lone one does" {
+ var target: fake.Fake = .{ .gpa = gpa };
+ defer target.deinit();
+ var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = true } };
+ defer head.deinit();
+ var wrapped = try Drafted.init(gpa, target.backend(), head.drafter());
+ defer wrapped.deinit();
+ const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 };
+ var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15);
+ defer cfg.deinit(gpa);
+ var clock: fake.FixedClock = .{};
+ var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock());
+ defer engine.deinit();
+ var streams: [2]sm.Stream = undefined;
+ for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 12 + i], .max_new = 24, .eos = &.{96} });
+ defer for (&streams) |*s| s.deinit(gpa);
+ try engine.addStream(&streams[0]);
+ try std.testing.expectEqual(@as(u32, 1), head.holds); // alone: its first drafts are held
+ for (0..3) |_| try engine.step();
+ const before = head.holds;
+ try std.testing.expect(!try engine.guard.drafts(2)); // the guard measures plain shared rounds first
+ try engine.addStream(&streams[1]);
+ try std.testing.expectEqual(before, head.holds); // joining: none held
+ while (engine.activeCount() > 0) try engine.step();
+}
+
+test "without the guard's flag a stream joining shared rounds holds its first drafts as before" {
+ var target: fake.Fake = .{ .gpa = gpa };
+ defer target.deinit();
+ var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true } };
+ defer head.deinit();
+ var wrapped = try Drafted.init(gpa, target.backend(), head.drafter());
+ defer wrapped.deinit();
+ const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 };
+ var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15);
+ defer cfg.deinit(gpa);
+ try std.testing.expect(!cfg.shared_guard);
+ var clock: fake.FixedClock = .{};
+ var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock());
+ defer engine.deinit();
+ var streams: [2]sm.Stream = undefined;
+ for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 12 + i], .max_new = 24, .eos = &.{96} });
+ defer for (&streams) |*s| s.deinit(gpa);
+ try engine.addStream(&streams[0]);
+ for (0..3) |_| try engine.step();
+ const before = head.holds;
+ try engine.addStream(&streams[1]);
+ try std.testing.expectEqual(before + 1, head.holds); // joining: its first drafts held
+ while (engine.activeCount() > 0) try engine.step();
+}
+
+const Run = struct { tokens: [3][64]u32 = undefined, lens: [3]usize = @splat(0), holds: u32 = 0 };
+
+/// Three streams decoded together: drafts off, or the fake drafter with or without the guard.
+fn run(drafts: bool, guard: bool) !Run {
+ var target: fake.Fake = .{ .gpa = gpa };
+ defer target.deinit();
+ var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = guard } };
+ defer head.deinit();
+ var wrapped = try Drafted.init(gpa, target.backend(), head.drafter());
+ defer wrapped.deinit();
+ const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 };
+ var cfg = try Config.init(gpa, if (drafts) wrapped.facts(m) else m, 16, 15);
+ defer cfg.deinit(gpa);
+ var clock: fake.FixedClock = .{};
+ var engine = Engine.init(gpa, &cfg, if (drafts) wrapped.backend() else target.backend(), clock.clock());
+ defer engine.deinit();
+ var streams: [3]sm.Stream = undefined;
+ for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 11 + i], .max_new = 48, .eos = &.{96}, .drafts = drafts });
+ defer for (&streams) |*s| s.deinit(gpa);
+ for (&streams) |*s| try engine.addStream(s);
+ while (engine.activeCount() > 0) try engine.step();
+ var out: Run = .{ .holds = head.holds };
+ for (&streams, 0..) |*s, i| {
+ out.lens[i] = s.emitted().len;
+ @memcpy(out.tokens[i][0..out.lens[i]], s.emitted());
+ }
+ return out;
+}
+
+test "the guard changes only which shared rounds draft: replies equal one-token decode with it on and off" {
+ const plain = try run(false, false);
+ const off = try run(true, false);
+ const on = try run(true, true);
+ for (0..3) |i| {
+ try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], off.tokens[i][0..off.lens[i]]);
+ try std.testing.expectEqualSlices(u32, plain.tokens[i][0..plain.lens[i]], on.tokens[i][0..on.lens[i]]);
+ }
+ try std.testing.expect(on.holds < off.holds); // its plain stretch holds no drafts
+}
+
+/// A clock that reads the same ms for every span: two of them with different ms stand for a fast and a slow machine.
+const FlatClock = struct {
+ ms: f64,
+
+ fn clock(c: *FlatClock) @import("backend.zig").Clock {
+ return .{ .ptr = c, .vtable = &.{ .start = start, .elapsed_ms = elapsed } };
+ }
+
+ fn start(_: *anyopaque) void {}
+
+ fn elapsed(ptr: *anyopaque, _: @import("backend.zig").Mark) f64 {
+ const c: *FlatClock = @ptrCast(@alignCast(ptr));
+ return c.ms;
+ }
+};
+
+/// Three streams, the guard asked for or not, timed or not, on a clock reading `ms` for every span.
+fn runClock(timed: bool, guard: bool, ms: f64) !Run {
+ var target: fake.Fake = .{ .gpa = gpa };
+ defer target.deinit();
+ var head: dt.FakeDrafter = .{ .facts_ = .{ .depth = 4, .step_ms = 0.5, .batched = true, .shared_guard = guard } };
+ defer head.deinit();
+ var wrapped = try Drafted.init(gpa, target.backend(), head.drafter());
+ defer wrapped.deinit();
+ const m: Model = .{ .exact_width = 16, .gpu_tokens = true, .hidden_rows = true, .batch_rows = 32, .max_streams = 4 };
+ var cfg = try Config.init(gpa, wrapped.facts(m), 16, 15);
+ defer cfg.deinit(gpa);
+ cfg.timed = timed;
+ var clock: FlatClock = .{ .ms = ms };
+ var engine = Engine.init(gpa, &cfg, wrapped.backend(), clock.clock());
+ defer engine.deinit();
+ var streams: [3]sm.Stream = undefined;
+ for (&streams, 0..) |*s, i| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = dt.long[0 .. 11 + i], .max_new = 48, .eos = &.{96}, .drafts = true });
+ defer for (&streams) |*s| s.deinit(gpa);
+ for (&streams) |*s| try engine.addStream(s);
+ while (engine.activeCount() > 0) try engine.step();
+ var out: Run = .{ .holds = head.holds };
+ for (&streams, 0..) |*s, i| {
+ out.lens[i] = s.emitted().len;
+ @memcpy(out.tokens[i][0..out.lens[i]], s.emitted());
+ }
+ return out;
+}
+
+test "with timed off the shared guard keeps no plain stretch: the same schedule on a fast and a slow clock" {
+ const fast = try runClock(false, true, 0.5);
+ const slow = try runClock(false, true, 400.0);
+ try std.testing.expectEqual(fast.holds, slow.holds);
+ for (0..3) |i| try std.testing.expectEqualSlices(u32, fast.tokens[i][0..fast.lens[i]], slow.tokens[i][0..slow.lens[i]]);
+ const unasked = try runClock(false, false, 0.5);
+ try std.testing.expectEqual(unasked.holds, fast.holds); // as if the guard were not asked for
+}
diff --git a/zig/src/core/lanes/drafted_held_test.zig b/zig/src/core/lanes/drafted_held_test.zig
new file mode 100644
index 0000000000..c8a6010744
--- /dev/null
+++ b/zig/src/core/lanes/drafted_held_test.zig
@@ -0,0 +1,78 @@
+//! Drafted.verify with held_on_target: windows of drafting streams reach the target with Window.held as asked.
+const std = @import("std");
+const be = @import("backend.zig");
+const dr = @import("drafter.zig");
+const Stream = @import("stream.zig").Stream;
+const Drafted = @import("drafted.zig").Drafted;
+
+/// A target that records each verify's held rows and host tokens, and a drafter that only counts host readbacks.
+const Fake = struct {
+ seen: u32 = 0,
+ tokens: [4]u32 = @splat(0),
+ readbacks: u32 = 0,
+
+ fn cast(ptr: *anyopaque) *Fake {
+ return @ptrCast(@alignCast(ptr));
+ }
+ fn verify(ptr: *anyopaque, windows: []const be.Window, _: []be.Verified) anyerror!void {
+ const f = cast(ptr);
+ f.seen = windows[0].held;
+ @memcpy(f.tokens[0..windows[0].tokens.len], windows[0].tokens);
+ }
+ fn held(ptr: *anyopaque, _: []const *Stream, out: []const []u32) anyerror!void {
+ cast(ptr).readbacks += 1;
+ for (out) |o| @memset(o, 5);
+ }
+ fn none(_: *anyopaque) []const u32 {
+ return &.{};
+ }
+ fn facts(_: *anyopaque) dr.Facts {
+ return .{ .depth = 3, .step_ms = 0, .prior = &.{}, .stream_bytes = 0 };
+ }
+ fn unused(_: *anyopaque, _: *Stream) anyerror!void {
+ return error.Unused;
+ }
+ fn absorb(_: *anyopaque, _: []const dr.Absorb) anyerror!void {
+ return error.Unused;
+ }
+ fn hold(_: *anyopaque, _: []const dr.Hold) anyerror!void {
+ return error.Unused;
+ }
+ fn release(_: *anyopaque, _: *Stream) void {}
+ fn first(_: *anyopaque, _: *Stream, _: u64) anyerror!u64 {
+ return error.Unused;
+ }
+ fn queue(_: *anyopaque, _: *Stream, _: be.Feed, _: u64) anyerror!u64 {
+ return error.Unused;
+ }
+ fn read(_: *anyopaque, _: u64) anyerror!u32 {
+ return error.Unused;
+ }
+ fn keep(_: *anyopaque, _: []const be.Window, _: []const []const u32) anyerror!void {}
+ fn draft(_: *anyopaque, _: []const be.DraftRequest) anyerror!void {}
+};
+
+test "held_on_target passes held windows through, and a stream with fillers still gets host tokens" {
+ var f: Fake = .{};
+ const target: be.Backend = .{ .ptr = &f, .vtable = &.{ .prefill = Fake.unused, .first = Fake.first, .queue = Fake.queue, .read = Fake.read, .verify = Fake.verify, .keep = Fake.keep, .draft = Fake.draft, .release = Fake.release } };
+ const drafter: dr.Drafter = .{ .ptr = &f, .vtable = &.{ .taps = Fake.none, .facts = Fake.facts, .open = Fake.unused, .absorb = Fake.absorb, .hold = Fake.hold, .held = Fake.held, .release = Fake.release } };
+ var x: Drafted = .{ .gpa = std.testing.allocator, .target = target, .drafter = drafter, .held_on_target = true };
+ defer x.deinit();
+ var s: Stream = undefined;
+ try x.lanes.put(std.testing.allocator, &s, .{ .rows = true });
+ var sampled: [4]u32 = undefined;
+ var drafts: [3]u32 = undefined;
+ const w = [_]be.Window{.{ .stream = &s, .pending = 9, .held = 3, .tokens = &.{}, .parents = null, .positions = &.{ 1, 2, 3, 4 } }};
+ var o = [_]be.Verified{.{ .sampled = &sampled, .drafts = &drafts }};
+ try x.backend().verify(&w, &o);
+ try std.testing.expectEqual(@as(u32, 3), f.seen);
+ try std.testing.expectEqual(@as(u32, 0), f.readbacks);
+ x.lanes.getPtr(&s).?.filler = 3; // fillers: the pending token as host tokens, no readback
+ try x.backend().verify(&w, &o);
+ try std.testing.expectEqual(@as(u32, 0), f.seen);
+ try std.testing.expectEqualSlices(u32, &.{ 9, 9, 9 }, f.tokens[0..3]);
+ x.held_on_target = false; // off: the drafter's drafts as host tokens
+ try x.backend().verify(&w, &o);
+ try std.testing.expectEqual(@as(u32, 1), f.readbacks);
+ try std.testing.expectEqualSlices(u32, &.{ 5, 5, 5 }, f.tokens[0..3]);
+}
diff --git a/zig/src/core/lanes/drafted_test.zig b/zig/src/core/lanes/drafted_test.zig
index a421dcc4d3..1e425d3f29 100644
--- a/zig/src/core/lanes/drafted_test.zig
+++ b/zig/src/core/lanes/drafted_test.zig
@@ -596,4 +596,5 @@ test "the drafter wrapper sees the committed callback before its first hold" {
test {
_ = @import("drafted_role_test.zig");
_ = @import("drafted_step_test.zig");
+ _ = @import("drafted_held_test.zig");
}
diff --git a/zig/src/core/lanes/drafter.zig b/zig/src/core/lanes/drafter.zig
index 80577e0edf..3ba9feb671 100644
--- a/zig/src/core/lanes/drafter.zig
+++ b/zig/src/core/lanes/drafter.zig
@@ -18,6 +18,12 @@ pub const Facts = struct {
window_costs: []const Cost = &.{}, // the target's window ms by rows (timed at load), for the depth rule
lane_costs: []const Cost = &.{}, // the same for the lane planner (Model.lane_costs)
tree_nodes: u32 = 0, // hold() without lanes grows its own tree of up to `depth` nodes, at most this many (tree, chances)
+ shared_guard: bool = false, // shared rounds draft only while they measure faster than plain ones (Model.shared_guard)
+ chances: bool = false, // chances() gives each held chain draft's chance: verify rows go by them (Model.granted_depth)
+ confidence_floor: ?f64 = null, // with chances: every round drafts the deepest chain and verifies drafts this likely
+ slack: u32 = 0, // with chances: drafts a hold offers past the stream's last grant (0: the depth rule's own)
+ copy_first: bool = false, // a hold costs a pass of its own: none is asked for a stream with a copy ahead
+ hold_costs: []const @import("hold_price.zig").HoldCost = &.{}, // that pass's ms by blocks and depth (timed at load)
};
/// Target rows [start, start + features.rows) and the token after each; `rows`, when given, picks which to absorb.
diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig
index 7894a19757..31fef8331b 100644
--- a/zig/src/core/lanes/engine.zig
+++ b/zig/src/core/lanes/engine.zig
@@ -3,6 +3,7 @@ const std = @import("std");
const Allocator = std.mem.Allocator;
const Config = @import("config.zig").Config;
const depth = @import("depth.zig");
+const shared_guard = @import("shared_guard.zig");
const accept = @import("accept.zig");
const be = @import("backend.zig");
const sm = @import("stream.zig");
@@ -16,6 +17,7 @@ const lone_costs = @import("lone_costs.zig");
const open = @import("open.zig");
const handover = @import("handover.zig");
const trim = @import("trim.zig");
+const pipelined = @import("pipelined.zig");
const Stream = sm.Stream;
const Feed = be.Feed;
const LogRow = @import("logprob.zig").Row;
@@ -24,7 +26,7 @@ const f = ev.f;
const str = trail.str;
const int = trail.int;
const Outcome = struct { got: []const u32, path: []const u32, cut: ?usize, follow: []const u32 };
-const Result = struct { rows: u32, keep: u32, got: []const u32 };
+pub const Result = struct { rows: u32, keep: u32, got: []const u32 };
const Unread = struct { event: usize, handle: u64 };
pub const Engine = struct {
const pieces_ = @import("pieces.zig");
@@ -42,11 +44,13 @@ pub const Engine = struct {
log: ?*ev.Log = null,
trim: ?trim.Trim = null, // a family's row policy in place of the built-in allocation (lanes/trim.zig)
rule: depth.Rule,
+ guard: shared_guard.Guard,
live: std.ArrayList(*Stream) = .empty, // admitted streams in admission order (Python `_live`)
drafted: u64 = 0,
accepted: u64 = 0,
shared_rounds: u64 = 0,
alone: bool = true, // this step has one live stream: its copies may take their ramp
+ plain_round: bool = false, // a shared round the guard keeps plain: no stream copies or drafts, as with drafts off
lone_kept: [2][lone_costs.bands]?lone_costs.Kept = @splat(@splat(null)), // the newest lone calibrations, greedy then sampled, a band each
steps: u64 = 0,
arena: std.heap.ArenaAllocator, // a round's temporaries
@@ -70,10 +74,13 @@ pub const Engine = struct {
.draft_step_ms = cfg.draft_step_ms,
.plain_guard = cfg.plain_guard,
.node_probabilities = cfg.node_probabilities,
+ .granted_depth = cfg.granted_depth,
+ .lone_guard = cfg.granted_depth and cfg.shared_guard and cfg.timed,
.confidence_floor = cfg.draft_confidence,
.batch_rows = cfg.batch_rows,
.batched_heads = cfg.draft_streams,
},
+ .guard = .{ .gpa = gpa },
};
}
@@ -81,6 +88,7 @@ pub const Engine = struct {
e.live.deinit(e.gpa);
e.unread.deinit(e.gpa);
e.rule.deinit();
+ e.guard.deinit();
for (&e.lone_kept) |*mode| for (mode) |*k| if (k.*) |*x| x.deinit(e.gpa);
e.arena.deinit();
}
@@ -124,7 +132,7 @@ pub const Engine = struct {
if (live.items.len > 1 and e.cfg.family_streams) {
// a shared round is synchronous: a stream that ran a step ahead lands its queued token first
for (live.items) |s| {
- if (s.inflight != null) try e.landInflight(s);
+ if (s.inflight != null) try pipelined.land(e, s);
}
var kept: usize = 0;
for (live.items) |s| {
@@ -147,7 +155,7 @@ pub const Engine = struct {
for (live.items) |s| {
// logprobs and grammars take verify windows: rows' logits are read, or bits made, before a forward
const windowed = (e.cfg.family_mtp and s.drafts) or !e.cfg.pipelined or s.logprobs != null or s.grammar != null;
- const r = if (windowed) try e.familyRound(s, null) else try e.pipelinedRound(s);
+ const r = if (windowed) try e.familyRound(s, null) else try pipelined.round(e, s);
s.min_rows = if (s.min_rows == 0) r.rows else @min(s.min_rows, r.rows);
}
}
@@ -195,7 +203,7 @@ pub const Engine = struct {
}
/// One stream's round: verify pending and drafted tokens together, keep drafts to the first mismatch.
- fn familyRound(e: *Engine, s: *Stream, copied: ?win.Copy) !Result {
+ pub fn familyRound(e: *Engine, s: *Stream, copied: ?win.Copy) !Result {
e.clock.start();
const priced = s.priced;
s.priced = null;
@@ -231,6 +239,7 @@ pub const Engine = struct {
lone_costs.observe(s, plan.cost_probe_rows, rows, ms);
}
}
+ if (e.rule.lone_guard and e.alone and s.drafts and (plan.kind == .head or plan.kind == .none)) try e.rule.observeLone(win.who(s), lone_costs.band(s.context.items.len), o.got.len, e.clock.elapsedMs(.cost));
if (e.cfg.timed and s.graft_room != null and s.rounds > 1) {
// the stream's round time beyond its window and head levels, for the lane planner
if (e.cfg.lane_costs.get(@intCast(rows))) |w| {
@@ -255,6 +264,8 @@ pub const Engine = struct {
e.clock.start();
const a = e.arena.allocator();
const plans = try a.alloc(Plan, entries.len);
+ e.plain_round = e.cfg.shared_guard and e.cfg.timed and !try e.guard.drafts(entries.len);
+ defer e.plain_round = false;
for (entries, plans) |s, *p| p.* = try win.plan(e, s, null);
try win.allocate(e, plans);
for (plans) |*p| try win.constrain(e, p);
@@ -284,7 +295,9 @@ pub const Engine = struct {
for (whos, heads.items) |*w, i| w.* = win.who(plans[i].stream);
const budgets = try a.alloc(i64, heads.items.len);
const sampled = win.anySampled(plans);
- try e.rule.budgets(whos, win.probe(e), sampled, budgets);
+ // the guard's plain stretches: no stream drafts, so the round costs what it would without a head
+ const guarded = e.cfg.shared_guard and e.cfg.timed and !try e.guard.draftsFresh(plans.len, whos);
+ if (guarded) @memset(budgets, 0) else try e.rule.budgets(whos, win.probe(e), sampled, budgets);
const ids = try a.alloc(ev.Value, heads.items.len);
for (ids, heads.items) |*id, i| id.* = str(plans[i].stream.id);
try trail.event(e, &.{ f("ev", str("budgets")), f("streams", .{ .list = ids }), f("depths", .{ .i64s = budgets }) });
@@ -313,6 +326,7 @@ pub const Engine = struct {
try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(plans.len)), f("rows", int(total)), f("ms", ev.bits(ms)) });
try e.rule.observeOverhead(@intCast(plans.len), total, ms, sampled, e.rule.headMs(budgets));
}
+ if (e.cfg.shared_guard and e.cfg.timed) try e.guard.observe(plans.len, shared_guard.committed(outcomes), e.clock.elapsedMs(.overhead));
if (e.backend.round_observer != null) e.backend.observed(paths);
}
@@ -344,9 +358,7 @@ pub const Engine = struct {
for (path[1..]) |r| own_kept += @intFromBool(r < own_end);
const own: usize = own_end - 1;
if (p.kind == .copy) {
- if (p.recalled) {
- if (s.recall) |r| r.observe(@intCast(own), @intCast(own_kept), s.rounds);
- } else if (s.proposer) |pr| pr.observe(@intCast(own), @intCast(own_kept));
+ win.observeCopied(e, s, p.recalled, own_kept, own);
if (e.alone) {
const width = s.copy_width orelse e.cfg.first_copy;
s.copy_width = if (own_kept == own) @min(e.cfg.max_copy, 2 * width + 1) else e.cfg.first_copy;
@@ -492,64 +504,6 @@ pub const Engine = struct {
}
}
- /// One token, the next forward queued first; a copied continuation ahead switches to verify windows.
- fn pipelinedRound(e: *Engine, s: *Stream) !Result {
- const before = s.mode;
- // an adopted stream holds neither a mode nor a queued step: it starts where a drain step leaves one
- const mode: sm.Mode = s.mode orelse if (s.inflight == null) .verify else .pipe;
- var result: Result = .{ .rows = 1, .keep = 1, .got = &.{} };
- if (mode == .verify or mode == .exit) {
- const verifying = mode == .verify and e.cfg.family_width >= 2;
- const proposal: win.Copy = if (verifying) try win.copyProposal(e, s) else .{};
- if (proposal.tokens.len > 0) {
- const copied = try e.arena.allocator().dupe(u32, proposal.tokens);
- result = try e.familyRound(s, .{ .tokens = copied, .recalled = proposal.recalled });
- if (result.keep == 1 or s.force.items.len > 0) s.mode = .exit;
- } else {
- // nothing to copy: back to steps queued ahead (this one lands next round)
- s.mode = .pipe;
- try e.queueNext(s, .{ .value = s.pending.? });
- }
- try trail.pipe(e, s, before, result.got);
- return result;
- }
- var current: Feed = .{ .handle = s.inflight.? };
- s.inflight = null;
- if (try e.forcedNext(s)) |t| current = .{ .value = t };
- var token: u32 = undefined;
- if (mode == .drain) {
- token = try e.readFeed(current); // the last queued step: no new one
- s.mode = .verify;
- } else {
- try e.queueNext(s, current); // the GPU starts the next step first
- token = try e.readFeed(current);
- }
- s.rounds += 1;
- const landed = try s.commit(e.gpa, &.{token}, &.{});
- s.pending = token;
- if (s.recall) |r| r.noteHead(1); // a step lands one token: the bar a copy must beat
- if (e.cfg.root_copies) s.root_copy_miss = false;
- const piping = e.cfg.family_width >= 2 and (s.mode orelse .pipe) == .pipe and !s.finished and s.drafts;
- if (piping and (try win.copyProposal(e, s)).tokens.len > 0)
- s.mode = .drain; // a copy window is ahead: land the queued step
- result.got = try e.arena.allocator().dupe(u32, s.emitted()[s.emitted().len - landed ..]);
- try trail.pipe(e, s, before, result.got);
- return result;
- }
-
- /// Land the queued token without queuing another (a shared round is synchronous).
- fn landInflight(e: *Engine, s: *Stream) !void {
- var current: Feed = .{ .handle = s.inflight.? };
- s.inflight = null;
- if (try e.forcedNext(s)) |t| current = .{ .value = t };
- const token = try e.readFeed(current);
- s.rounds += 1;
- const landed = try s.commit(e.gpa, &.{token}, &.{});
- s.pending = token;
- s.mode = .verify;
- try trail.event(e, &.{ f("ev", str("land")), f("stream", str(s.id)), f("got", .{ .u32s = s.emitted()[s.emitted().len - landed ..] }) });
- }
-
/// The first token's draw: among the grammar's allowed tokens when the stream has one that constrains it.
pub fn drawFirst(e: *Engine, s: *Stream, position: u64) !u64 {
const g = s.grammar orelse return e.backend.first(s, position);
diff --git a/zig/src/core/lanes/engine_test.zig b/zig/src/core/lanes/engine_test.zig
index 36167e30d6..6995de7dbe 100644
--- a/zig/src/core/lanes/engine_test.zig
+++ b/zig/src/core/lanes/engine_test.zig
@@ -463,6 +463,9 @@ test "a family that keeps copies to shared rounds proposes none for a lone strea
try std.testing.expect((try win.copyProposal(&engine, &s)).tokens.len >= 2);
engine.alone = true;
try std.testing.expectEqual(lone_copies, (try win.copyProposal(&engine, &s)).tokens.len >= 2);
+ engine.alone = false;
+ engine.plain_round = true; // a shared round the guard keeps plain: no copies either, as with drafts off
+ try std.testing.expectEqual(@as(usize, 0), (try win.copyProposal(&engine, &s)).tokens.len);
}
}
diff --git a/zig/src/core/lanes/hold_price.zig b/zig/src/core/lanes/hold_price.zig
new file mode 100644
index 0000000000..e4c854410e
--- /dev/null
+++ b/zig/src/core/lanes/hold_price.zig
@@ -0,0 +1,58 @@
+//! A drafter's hold pass priced from its timed passes: by the blocks it holds and their depth.
+const std = @import("std");
+
+/// A drafter's hold pass of `streams` blocks, each `depth` drafts deep, in ms (timed at load).
+pub const HoldCost = struct { streams: u32, depth: u32, ms: f64 };
+
+/// A hold pass's ms: linear in depth at each timed stream count, then in streams (past the widest, by its slope).
+pub fn holdMs(costs: []const HoldCost, streams: usize, depth: i64) f64 {
+ var below: ?u32 = null;
+ var above: ?u32 = null;
+ for (costs) |c| {
+ if (c.streams <= streams and (below == null or c.streams > below.?)) below = c.streams;
+ if (c.streams >= streams and (above == null or c.streams < above.?)) above = c.streams;
+ }
+ const lo = below orelse above.?;
+ const hi = above orelse lo;
+ const at_lo = atDepth(costs, lo, depth);
+ if (hi == lo) {
+ if (above != null) return at_lo;
+ // past the widest timed count: the slope from the count below it
+ var under: ?u32 = null;
+ for (costs) |c| if (c.streams < lo and (under == null or c.streams > under.?)) {
+ under = c.streams;
+ };
+ const u = under orelse return at_lo * @as(f64, @floatFromInt(streams)) / @as(f64, @floatFromInt(lo));
+ const slope = (at_lo - atDepth(costs, u, depth)) / @as(f64, @floatFromInt(lo - u));
+ return at_lo + slope * @as(f64, @floatFromInt(streams - lo));
+ }
+ const t = @as(f64, @floatFromInt(streams - lo)) / @as(f64, @floatFromInt(hi - lo));
+ return at_lo + t * (atDepth(costs, hi, depth) - at_lo);
+}
+
+/// One stream count's hold ms at `depth`: linear between its two nearest timed depths (one alone: as timed).
+fn atDepth(costs: []const HoldCost, streams: u32, depth: i64) f64 {
+ var a: ?HoldCost = null;
+ var b: ?HoldCost = null;
+ for (costs) |c| {
+ if (c.streams != streams) continue;
+ if (a == null or c.depth < a.?.depth) {
+ b = a;
+ a = c;
+ } else if (b == null or c.depth < b.?.depth) b = c;
+ }
+ const x = a.?;
+ const y = b orelse return x.ms;
+ const t = (@as(f64, @floatFromInt(depth)) - @as(f64, @floatFromInt(x.depth))) / @as(f64, @floatFromInt(y.depth - x.depth));
+ return x.ms + t * (y.ms - x.ms);
+}
+
+test "a hold's ms between and past the timed passes" {
+ const t = [_]HoldCost{ .{ .streams = 1, .depth = 1, .ms = 1.0 }, .{ .streams = 1, .depth = 4, .ms = 2.5 }, .{ .streams = 4, .depth = 1, .ms = 3.0 }, .{ .streams = 4, .depth = 4, .ms = 4.5 }, .{ .streams = 8, .depth = 1, .ms = 4.0 } };
+ try std.testing.expectApproxEqAbs(@as(f64, 1.0), holdMs(&t, 1, 1), 1e-12);
+ try std.testing.expectApproxEqAbs(@as(f64, 2.0), holdMs(&t, 1, 3), 1e-12); // linear in depth
+ try std.testing.expectApproxEqAbs(@as(f64, 3.0), holdMs(&t, 1, 5), 1e-12); // and past the deepest
+ try std.testing.expectApproxEqAbs(@as(f64, 1.0 + 2.0 / 3.0), holdMs(&t, 2, 1), 1e-12); // linear in streams
+ try std.testing.expectApproxEqAbs(@as(f64, 4.0), holdMs(&t, 8, 3), 1e-12); // one depth timed: as timed
+ try std.testing.expectApproxEqAbs(@as(f64, 5.0), holdMs(&t, 12, 1), 1e-12); // past the widest: the slope below it
+}
diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig
index 43b9a78756..b1f5121613 100644
--- a/zig/src/core/lanes/lanes.zig
+++ b/zig/src/core/lanes/lanes.zig
@@ -18,6 +18,7 @@ pub const cost_rule = @import("cost_rule.zig");
pub const gpu_rule = @import("gpu_rule.zig");
pub const gpu_full = @import("gpu_full.zig");
pub const depth = @import("depth.zig");
+pub const hold_price = @import("hold_price.zig");
pub const mtp_policy = @import("mtp_policy.zig");
pub const trim = @import("trim.zig");
pub const config = @import("config.zig");
@@ -56,17 +57,24 @@ test {
_ = @import("recorder_test.zig");
_ = @import("gpu_full_test.zig");
_ = @import("drafted_test.zig");
+ _ = @import("drafted_chances_test.zig");
+ _ = @import("drafted_copy_test.zig");
_ = @import("root_copy_test.zig");
_ = @import("drafted_grammar_test.zig");
_ = @import("windows_grammar_test.zig");
_ = @import("pricing_test.zig");
_ = @import("depth_probe_test.zig");
+ _ = @import("shared_guard.zig");
+ _ = @import("lone_guard.zig");
+ _ = @import("lone_guard_test.zig");
+ _ = @import("drafted_guard_test.zig");
_ = @import("proposal_external_test.zig");
_ = @import("lone_costs_test.zig");
_ = @import("lone_backoff_test.zig");
_ = @import("fill_test.zig");
_ = @import("fill_trim_test.zig");
_ = @import("fair_turns_test.zig");
+ _ = @import("copied_rate_test.zig");
}
pub const learned_disk = @import("learned_disk.zig");
diff --git a/zig/src/core/lanes/lone_guard.zig b/zig/src/core/lanes/lone_guard.zig
new file mode 100644
index 0000000000..4657ffe829
--- /dev/null
+++ b/zig/src/core/lanes/lone_guard.zig
@@ -0,0 +1,242 @@
+//! Whether a lone stream holds drafts: its drafted rounds' tokens a ms against plain rounds', the loser re-measured.
+const std = @import("std");
+
+pub const Constants = struct {
+ drafted: u32 = 16, // measured drafted rounds a verdict takes, again and again while drafting wins
+ plain: u32 = 16, // measured plain rounds that take the plain rate when there is none yet
+ wait: u32 = 16, // plain rounds before drafting is measured again; doubles each time plain wins again
+ wait_most: u32 = 512,
+ margin: f64 = 1.05, // drafted rounds must beat plain ones' tokens a ms by this much
+ probe_cost: f64 = 0.03, // after two losses a probe costs at most this share of plain time
+ confirm: u32 = 2, // losing verdicts in a row (a remembered one counts) before the wait follows probe_cost
+ blend: f64 = 0.5, // the newest stretch's weight in a rate
+};
+
+pub const Mode = enum { drafted, plain };
+
+/// Plain rounds' tokens a ms by sampling mode and context band, kept from stream to stream (no draft changes them).
+pub const Plain = struct {
+ rates: [2][bands]?f64 = @splat(@splat(null)),
+ verdicts: [2][bands]?Verdict = @splat(@splat(null)), // each band's last losing verdict, until drafting wins there
+
+ pub const Verdict = struct { drafted: f64, wait: u32 };
+
+ pub fn memo(p: *Plain, sampled: bool, band: usize) *?Verdict {
+ return &p.verdicts[@intFromBool(sampled)][@min(band, bands - 1)];
+ }
+
+ pub const bands = 6;
+
+ pub fn at(p: *Plain, sampled: bool, band: usize) *?f64 {
+ return &p.rates[@intFromBool(sampled)][@min(band, bands - 1)];
+ }
+};
+
+/// One lone stream's guard: the mode its holds follow, the stretch measured and its drafted rate.
+pub const Lone = struct {
+ mode: Mode = .drafted,
+ left: u32 = 0,
+ tokens: f64 = 0,
+ ms: f64 = 0,
+ skip: bool = true, // a mode's first round still verifies the other mode's hold (or none): not measured
+ started: bool = false,
+ pending: bool = false, // a drafted verdict waits for the plain rate this stretch takes
+ fresh: bool = false, // a drafted stretch's first hold: one draft and its slack, whatever the last grant
+ drafted: ?f64 = null,
+ drafted_ms: f64 = 0, // the last drafted stretch's wall ms
+ drafted_now: f64 = 0, // the last drafted stretch's own tokens a ms (unblended)
+ losses: u32 = 0, // losing verdicts in a row
+ wait: u32 = 0,
+
+ /// Whether the next hold may draft.
+ pub fn drafts(l: *const Lone) bool {
+ return l.mode == .drafted;
+ }
+
+ /// A lone round of held drafts or none (copies left out): its tokens and wall ms; `plain` the band's rate.
+ pub fn observe(l: *Lone, k: Constants, plain: *?f64, tokens: usize, ms: f64) void {
+ l.observeBand(k, plain, null, tokens, ms);
+ }
+
+ /// observe, with the band's last verdict: a stream starting where drafting lost starts plain for its wait.
+ pub fn observeBand(l: *Lone, k: Constants, plain: *?f64, memo: ?*?Plain.Verdict, tokens: usize, ms: f64) void {
+ var switched = false;
+ if (!l.started) {
+ l.started = true;
+ l.left = k.drafted;
+ l.wait = k.wait;
+ if (memo) |m| if (m.*) |v| if (plain.*) |p| if (v.drafted < p * k.margin) {
+ l.drafted = v.drafted;
+ l.losses = k.confirm - 1; // one more losing probe confirms the band's loss again
+ l.switchTo(.plain, v.wait);
+ l.wait = @min(v.wait *| 2, k.wait_most);
+ switched = true;
+ };
+ }
+ if (l.skip) {
+ l.skip = switched; // the next round still verifies this round's drafted hold
+ return;
+ }
+ l.tokens += @floatFromInt(tokens);
+ l.ms += ms;
+ l.left -|= 1;
+ if (l.left > 0) return;
+ const now = if (l.ms > 0) l.tokens / l.ms else 0;
+ const stretch_ms = l.ms;
+ l.tokens = 0;
+ l.ms = 0;
+ switch (l.mode) {
+ .drafted => {
+ l.drafted = blend(k, l.drafted, now);
+ l.drafted_ms = stretch_ms;
+ l.drafted_now = now;
+ if (plain.* == null) {
+ l.pending = true;
+ l.switchTo(.plain, k.plain);
+ } else l.verdict(k, plain.*.?, memo);
+ },
+ .plain => {
+ plain.* = blend(k, plain.*, now);
+ if (l.pending) {
+ l.pending = false;
+ l.verdict(k, plain.*.?, memo);
+ } else l.switchTo(.drafted, k.drafted);
+ },
+ }
+ }
+
+ /// Drafting goes on while it beats plain by the margin, else plain rounds for a wait that doubles.
+ fn verdict(l: *Lone, k: Constants, plain: f64, memo: ?*?Plain.Verdict) void {
+ if (l.drafted.? >= plain * k.margin) {
+ l.wait = k.wait;
+ l.losses = 0;
+ l.switchTo(.drafted, k.drafted);
+ if (memo) |m| m.* = null;
+ return;
+ }
+ var wait = l.wait;
+ if (l.drafted_now >= plain * k.margin) {
+ // the blend lost but this stretch won: no loss counted, the band forgets its loss
+ l.losses = 0;
+ if (memo) |m| m.* = null;
+ } else {
+ l.losses += 1;
+ if (l.losses >= k.confirm) {
+ const lost_ms = l.drafted_ms * @max(0, 1 - l.drafted_now / plain);
+ const need: f64 = @ceil(lost_ms * plain / k.probe_cost); // plain rounds of about one token each
+ wait = @min(k.wait_most, @max(wait, @as(u32, @intFromFloat(@min(need, 1e6)))));
+ }
+ if (memo) |m| m.* = .{ .drafted = l.drafted.?, .wait = wait };
+ }
+ l.switchTo(.plain, wait);
+ l.wait = @min(wait *| 2, k.wait_most);
+ }
+
+ fn switchTo(l: *Lone, mode: Mode, rounds: u32) void {
+ if (mode != l.mode) {
+ l.skip = true;
+ l.fresh = mode == .drafted;
+ }
+ l.mode = mode;
+ l.left = @max(1, rounds);
+ }
+};
+
+fn blend(k: Constants, old: ?f64, now: f64) f64 {
+ return if (old) |r| r + k.blend * (now - r) else now;
+}
+
+/// Rounds of one stream whose drafted rounds give `drafted` and plain ones `plain` tokens a ms: how many drafted.
+fn run(l: *Lone, k: Constants, rate: *?f64, rounds: usize, plain: f64, drafted: f64) usize {
+ var n: usize = 0;
+ for (0..rounds) |_| {
+ const d = l.drafts();
+ n += @intFromBool(d);
+ l.observe(k, rate, if (d) 2 else 1, if (d) 2 / drafted else 1 / plain);
+ }
+ return n;
+}
+
+test "a lone stream whose drafted rounds lose runs plain but for probes that grow rarer" {
+ var l: Lone = .{};
+ var rate: ?f64 = null;
+ const k: Constants = .{};
+ const drafted = run(&l, k, &rate, 1000, 0.128, 0.122);
+ // the first drafted stretch, the plain rate's own, then probes after 16, 32, 64, 128, 256 and 512 plain rounds
+ try std.testing.expect(drafted <= 7 * (k.drafted + 1));
+ try std.testing.expect(drafted > k.drafted);
+ try std.testing.expect(!l.drafts() or l.left <= k.drafted);
+ try std.testing.expectApproxEqAbs(@as(f64, 0.128), rate.?, 1e-9);
+}
+
+test "a lone stream whose drafted rounds win keeps drafting; a later stream reuses the plain rate" {
+ var rate: ?f64 = null;
+ const k: Constants = .{};
+ var first: Lone = .{};
+ const n = run(&first, k, &rate, 1000, 0.128, 0.15);
+ try std.testing.expect(n > 1000 - k.plain - 2); // one plain stretch takes the plain rate
+ var second: Lone = .{};
+ try std.testing.expectEqual(@as(usize, 300), run(&second, k, &rate, 300, 0.128, 0.15)); // no plain stretch again
+ var third: Lone = .{};
+ const lost = run(&third, k, &rate, 300, 0.128, 0.11);
+ try std.testing.expect(lost < 300 / 3); // plain from its first verdict, without measuring plain again
+}
+
+test "a drafted stretch's first round is not measured and its first hold is fresh" {
+ var l: Lone = .{};
+ var rate: ?f64 = 0.2;
+ const k: Constants = .{ .drafted = 2, .wait = 2 };
+ l.observe(k, &rate, 9, 1); // skipped: the stream's first round
+ l.observe(k, &rate, 1, 10);
+ l.observe(k, &rate, 1, 10);
+ try std.testing.expect(!l.drafts()); // 0.1 against 0.2
+ try std.testing.expect(!l.fresh);
+ l.observe(k, &rate, 9, 1); // skipped: the last drafted round's hold
+ l.observe(k, &rate, 1, 5);
+ l.observe(k, &rate, 1, 5);
+ try std.testing.expect(l.drafts() and l.fresh); // the probe after the wait
+ try std.testing.expectApproxEqAbs(@as(f64, 0.2), rate.?, 1e-12);
+}
+
+test "a probe that loses a third waits long enough that probes cost at most probe_cost; a band's verdict carries over" {
+ const k: Constants = .{};
+ var rate: ?f64 = null;
+ var memo: ?Plain.Verdict = null;
+ var l: Lone = .{};
+ // 130k on the M5: plain 8.67 ms a token, drafted 15.7 ms for 1.2 tokens
+ var drafted: usize = 0;
+ for (0..200) |_| {
+ const d = l.drafts();
+ drafted += @intFromBool(d);
+ l.observeBand(k, &rate, &memo, if (d) 12 else 10, if (d) 157 else 86.7);
+ }
+ try std.testing.expect(drafted <= 2 * (1 + k.drafted)); // two losing stretches, then the probe_cost wait
+ try std.testing.expect(memo != null and memo.?.wait > 200);
+ var next: Lone = .{};
+ var later: usize = 0;
+ for (0..200) |_| {
+ const d = next.drafts();
+ later += @intFromBool(d);
+ next.observeBand(k, &rate, &memo, if (d) 12 else 10, if (d) 157 else 86.7);
+ }
+ try std.testing.expect(later <= 1);
+}
+
+test "a winning stretch the blend still calls a loss counts no loss: no probe_cost wait" {
+ const k: Constants = .{ .drafted = 2, .wait = 2 };
+ var rate: ?f64 = 0.2;
+ var memo: ?Plain.Verdict = null;
+ var l: Lone = .{};
+ l.observeBand(k, &rate, &memo, 9, 1); // skipped
+ l.observeBand(k, &rate, &memo, 1, 10);
+ l.observeBand(k, &rate, &memo, 1, 10); // 0.1: a loss, wait 2
+ try std.testing.expect(!l.drafts() and l.losses == 1 and memo != null);
+ l.observeBand(k, &rate, &memo, 9, 1); // skipped
+ l.observeBand(k, &rate, &memo, 1, 5);
+ l.observeBand(k, &rate, &memo, 1, 5);
+ try std.testing.expect(l.drafts());
+ l.observeBand(k, &rate, &memo, 9, 1); // skipped
+ l.observeBand(k, &rate, &memo, 2, 8);
+ l.observeBand(k, &rate, &memo, 2, 8); // 0.25 alone wins; blended 0.175 loses
+ try std.testing.expect(!l.drafts() and l.losses == 0 and memo == null and l.left == 4);
+}
diff --git a/zig/src/core/lanes/lone_guard_test.zig b/zig/src/core/lanes/lone_guard_test.zig
new file mode 100644
index 0000000000..2becd90158
--- /dev/null
+++ b/zig/src/core/lanes/lone_guard_test.zig
@@ -0,0 +1,115 @@
+//! The lone guard through the round loop: a lone stream drafts while drafting measures faster, else runs plain.
+const std = @import("std");
+const Config = @import("config.zig").Config;
+const Engine = @import("engine.zig").Engine;
+const be = @import("backend.zig");
+const sm = @import("stream.zig");
+const fake = @import("fake.zig");
+const dr = @import("drafter.zig");
+const Drafted = @import("drafted.zig").Drafted;
+const dt = @import("drafted_test.zig");
+const ct = @import("drafted_chances_test.zig");
+
+const gpa = std.testing.allocator;
+
+/// A round's ms from its verify's rows and whether a hold ran before it, as a GPU that runs both in turn.
+const Costs = struct { row: f64, extra: f64, hold: f64 };
+
+const Cost = struct {
+ var costs: Costs = undefined;
+ var inner_verify: *const fn (*anyopaque, []const be.Window, []be.Verified) anyerror!void = undefined;
+ var inner_hold: *const fn (*anyopaque, []const dr.Hold) anyerror!void = undefined;
+ var rows: usize = 1;
+ var held: bool = false; // a hold since the last verify
+ var before: bool = false; // the last verify ran after a hold
+ var holds: u32 = 0;
+ var verifies: u32 = 0;
+
+ fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void {
+ rows = windows[0].rows();
+ before = held;
+ held = false;
+ verifies += 1;
+ return inner_verify(ptr, windows, out);
+ }
+
+ fn hold(ptr: *anyopaque, items: []const dr.Hold) anyerror!void {
+ for (items) |it| held = held or it.depth > 0;
+ holds += @intFromBool(held);
+ return inner_hold(ptr, items);
+ }
+
+ fn clock() be.Clock {
+ return .{ .ptr = @ptrCast(&costs), .vtable = &.{ .start = start, .elapsed_ms = elapsed } };
+ }
+
+ fn start(_: *anyopaque) void {}
+
+ fn elapsed(_: *anyopaque, _: be.Mark) f64 {
+ const extra: f64 = @floatFromInt(rows - 1);
+ return costs.row + costs.extra * extra + (if (before) costs.hold else 0);
+ }
+};
+
+const Run = struct { tokens: [400]u32 = undefined, len: usize = 0, holds: u32 = 0, rounds: u32 = 0, drafted: u64 = 0 };
+
+/// One lone stream of 300 tokens; `guard` the drafter's guard switch, `costs` null: drafts off.
+fn run(costs: ?Costs, guard: bool) !Run {
+ var target: fake.Fake = .{ .gpa = gpa };
+ defer target.deinit();
+ var p: ct.Priced = .{ .head = .{ .facts_ = .{ .depth = 9, .step_ms = 0.5, .batched = true, .chances = true, .shared_guard = guard } }, .sure = 0.95, .unsure = 0.05 };
+ defer p.deinit();
+ var strict: dt.Strict = .{ .inner = target.backend() };
+ defer strict.deinit();
+ var seen = strict.backend();
+ var table = seen.vtable.*;
+ Cost.inner_verify = table.verify;
+ table.verify = Cost.verify;
+ seen.vtable = &table;
+ var drafter = p.drafter();
+ var hold_table = drafter.vtable.*;
+ Cost.inner_hold = hold_table.hold;
+ hold_table.hold = Cost.hold;
+ drafter.vtable = &hold_table;
+ Cost.costs = costs orelse .{ .row = 8, .extra = 0, .hold = 0 };
+ Cost.holds = 0;
+ Cost.verifies = 0;
+ Cost.held = false;
+ var wrapped = try Drafted.init(gpa, seen, drafter);
+ defer wrapped.deinit();
+ const drafts = costs != null;
+ var cfg = try Config.init(gpa, if (drafts) wrapped.facts(ct.base()) else ct.base(), 16, 15);
+ defer cfg.deinit(gpa);
+ var engine = Engine.init(gpa, &cfg, if (drafts) wrapped.backend() else target.backend(), Cost.clock());
+ defer engine.deinit();
+ var s = try sm.Stream.init(gpa, .{ .id = "s", .prompt = &dt.p1, .max_new = 300, .eos = &.{}, .drafts = drafts });
+ defer s.deinit(gpa);
+ try engine.addStream(&s);
+ while (engine.activeCount() > 0) try engine.step();
+ var out: Run = .{ .len = s.emitted().len, .holds = Cost.holds, .rounds = Cost.verifies, .drafted = s.drafted };
+ @memcpy(out.tokens[0..out.len], s.emitted());
+ return out;
+}
+
+const slow: Costs = .{ .row = 8, .extra = 4, .hold = 9 }; // a hold and a verify row cost more than drafts save
+const fast: Costs = .{ .row = 8, .extra = 1, .hold = 1 };
+
+test "a lone stream whose drafted rounds lose runs plain rounds and probes now and then, replies unchanged" {
+ const plain = try run(null, true);
+ const unguarded = try run(slow, false);
+ const guarded = try run(slow, true);
+ try std.testing.expectEqualSlices(u32, plain.tokens[0..plain.len], unguarded.tokens[0..unguarded.len]);
+ try std.testing.expectEqualSlices(u32, plain.tokens[0..plain.len], guarded.tokens[0..guarded.len]);
+ try std.testing.expect(unguarded.holds * 10 > unguarded.rounds * 9); // without the guard nearly every round holds
+ try std.testing.expect(guarded.holds * 3 < guarded.rounds); // with it most rounds are plain
+ try std.testing.expect(guarded.holds > 16); // the first drafted stretch and probes
+}
+
+test "a lone stream whose drafted rounds win keeps drafting with the guard on" {
+ const plain = try run(null, true);
+ const unguarded = try run(fast, false);
+ const guarded = try run(fast, true);
+ try std.testing.expectEqualSlices(u32, plain.tokens[0..plain.len], guarded.tokens[0..guarded.len]);
+ try std.testing.expect(guarded.holds * 10 > guarded.rounds * 8); // drafting but for the plain rate's one stretch
+ try std.testing.expect(guarded.rounds * 10 < unguarded.rounds * 12);
+}
diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig
index 4c4cb585e5..153da72e53 100644
--- a/zig/src/core/lanes/open.zig
+++ b/zig/src/core/lanes/open.zig
@@ -135,6 +135,9 @@ fn openedCommitted(e: *Engine, s: *Stream, reach: Reach, hook: ?be.Committed) !v
/// The first drafts' depth for `reach`.
pub fn firstDepth(e: *Engine, s: *Stream, reach: Reach) !u32 {
var who = win.who(s);
+ // a stream joining shared rounds the guard runs plain holds no drafts, as the guard measured it faster
+ const others = e.activeCount();
+ if (others > 0 and e.cfg.family_streams and e.cfg.shared_guard and e.cfg.timed and !try e.guard.drafts(others + 1)) return 0;
if (reach == .head and e.cfg.node_probabilities) return @intCast(try e.rule.headDepth(&who, null, win.probe(e)));
return @intCast(try e.rule.depth(who));
}
diff --git a/zig/src/core/lanes/pipelined.zig b/zig/src/core/lanes/pipelined.zig
new file mode 100644
index 0000000000..99de33304b
--- /dev/null
+++ b/zig/src/core/lanes/pipelined.zig
@@ -0,0 +1,70 @@
+//! A lone stream's pipelined steps: the next forward queued before the last one's token is read.
+const be = @import("backend.zig");
+const sm = @import("stream.zig");
+const ev = @import("events.zig");
+const win = @import("windows.zig");
+const trail = @import("trail.zig");
+const Engine = @import("engine.zig").Engine;
+const Result = @import("engine.zig").Result;
+const Stream = sm.Stream;
+const Feed = be.Feed;
+const f = ev.f;
+const str = trail.str;
+
+/// One token, the next forward queued first; a copied continuation ahead switches to verify windows.
+pub fn round(e: *Engine, s: *Stream) !Result {
+ const before = s.mode;
+ // an adopted stream holds neither a mode nor a queued step: it starts where a drain step leaves one
+ const mode: sm.Mode = s.mode orelse if (s.inflight == null) .verify else .pipe;
+ var result: Result = .{ .rows = 1, .keep = 1, .got = &.{} };
+ if (mode == .verify or mode == .exit) {
+ const verifying = mode == .verify and e.cfg.family_width >= 2;
+ const proposal: win.Copy = if (verifying) try win.copyProposal(e, s) else .{};
+ if (proposal.tokens.len > 0) {
+ const copied = try e.arena.allocator().dupe(u32, proposal.tokens);
+ result = try e.familyRound(s, .{ .tokens = copied, .recalled = proposal.recalled });
+ if (result.keep == 1 or s.force.items.len > 0) s.mode = .exit;
+ } else {
+ // nothing to copy: back to steps queued ahead (this one lands next round)
+ s.mode = .pipe;
+ try e.queueNext(s, .{ .value = s.pending.? });
+ }
+ try trail.pipe(e, s, before, result.got);
+ return result;
+ }
+ var current: Feed = .{ .handle = s.inflight.? };
+ s.inflight = null;
+ if (try e.forcedNext(s)) |t| current = .{ .value = t };
+ var token: u32 = undefined;
+ if (mode == .drain) {
+ token = try e.readFeed(current); // the last queued step: no new one
+ s.mode = .verify;
+ } else {
+ try e.queueNext(s, current); // the GPU starts the next step first
+ token = try e.readFeed(current);
+ }
+ s.rounds += 1;
+ const landed = try s.commit(e.gpa, &.{token}, &.{});
+ s.pending = token;
+ if (s.recall) |r| r.noteHead(1); // a step lands one token: the bar a copy must beat
+ if (e.cfg.root_copies) s.root_copy_miss = false;
+ const piping = e.cfg.family_width >= 2 and (s.mode orelse .pipe) == .pipe and !s.finished and s.drafts;
+ if (piping and (try win.copyProposal(e, s)).tokens.len > 0)
+ s.mode = .drain; // a copy window is ahead: land the queued step
+ result.got = try e.arena.allocator().dupe(u32, s.emitted()[s.emitted().len - landed ..]);
+ try trail.pipe(e, s, before, result.got);
+ return result;
+}
+
+/// Land the queued token without queuing another (a shared round is synchronous).
+pub fn land(e: *Engine, s: *Stream) !void {
+ var current: Feed = .{ .handle = s.inflight.? };
+ s.inflight = null;
+ if (try e.forcedNext(s)) |t| current = .{ .value = t };
+ const token = try e.readFeed(current);
+ s.rounds += 1;
+ const landed = try s.commit(e.gpa, &.{token}, &.{});
+ s.pending = token;
+ s.mode = .verify;
+ try trail.event(e, &.{ f("ev", str("land")), f("stream", str(s.id)), f("got", .{ .u32s = s.emitted()[s.emitted().len - landed ..] }) });
+}
diff --git a/zig/src/core/lanes/shared_guard.zig b/zig/src/core/lanes/shared_guard.zig
new file mode 100644
index 0000000000..52dc848b95
--- /dev/null
+++ b/zig/src/core/lanes/shared_guard.zig
@@ -0,0 +1,176 @@
+//! Whether a shared round drafts at all: drafted and plain rounds' measured tokens a ms, by stream count.
+const std = @import("std");
+const Allocator = std.mem.Allocator;
+const depth = @import("depth.zig");
+
+pub const Constants = struct {
+ explore: [2]u32 = .{ 16, 4 }, // measured rounds plain and drafted rounds run when their rate is (re)taken
+ exploit: u32 = 512, // measured rounds the faster mode runs before the slower one is measured again
+ longest: u5 = 6, // each time the same mode wins again its run doubles, up to exploit << longest
+ margin: f64 = 1.05, // drafted rounds must beat plain ones' tokens a ms by this much
+ blend: f64 = 0.5, // the newest stretch's weight in a mode's rate
+};
+
+/// One stream count's state: each mode's rate (0 plain, 1 drafted), the mode running and its stretch.
+const Entry = struct {
+ streams: u32,
+ rate: [2]?f64 = .{ null, null },
+ mode: u1 = 0, // plain first: drafted rounds that lose cost more than plain ones, so they are measured short
+ exploring: bool = true,
+ left: u32,
+ tokens: f64 = 0,
+ ms: f64 = 0,
+ skip: bool = true, // a mode's first round still carries the other mode's drafts or holds: not measured
+ best: ?u1 = null, // the mode the last measurement chose
+ streak: u5 = 0, // how many times in a row it chose that mode
+};
+
+pub const Guard = struct {
+ gpa: Allocator,
+ k: Constants = .{},
+ list: std.ArrayList(Entry) = .empty,
+
+ pub fn deinit(g: *Guard) void {
+ g.list.deinit(g.gpa);
+ }
+
+ fn entry(g: *Guard, streams: usize) !*Entry {
+ const n: u32 = @intCast(streams);
+ for (g.list.items) |*e| if (e.streams == n) return e;
+ try g.list.append(g.gpa, .{ .streams = n, .left = g.k.explore[0] });
+ return &g.list.items[g.list.items.len - 1];
+ }
+
+ /// Whether a shared round of `streams` streams may draft (false: every stream runs its row alone).
+ pub fn drafts(g: *Guard, streams: usize) !bool {
+ return (try g.entry(streams)).mode == 1;
+ }
+
+ /// drafts(); a drafted stretch's first round grants every stream one draft, so it measures drafted rounds.
+ pub fn draftsFresh(g: *Guard, streams: usize, whos: []depth.Who) !bool {
+ const e = try g.entry(streams);
+ if (e.mode == 1 and e.skip) for (whos) |*w| {
+ w.granted = 1;
+ w.grant_short = false;
+ w.gain = null; // priced by its first draft's rate, not a grant of none
+ };
+ return e.mode == 1;
+ }
+
+ /// A finished shared round: its committed tokens and wall ms, in the mode drafts() gave it.
+ pub fn observe(g: *Guard, streams: usize, tokens: u64, ms: f64) !void {
+ const e = try g.entry(streams);
+ if (e.skip) {
+ e.skip = false;
+ return;
+ }
+ e.tokens += @floatFromInt(tokens);
+ e.ms += ms;
+ e.left -|= 1;
+ if (e.left > 0) return;
+ if (e.ms > 0) {
+ const now = e.tokens / e.ms;
+ e.rate[e.mode] = if (e.rate[e.mode]) |r| r + g.k.blend * (now - r) else now;
+ }
+ e.tokens = 0;
+ e.ms = 0;
+ const before = e.mode;
+ if (e.exploring and e.rate[e.mode ^ 1] == null) {
+ e.mode ^= 1; // the other mode has no rate yet: measure it next
+ } else if (e.exploring) {
+ e.mode = if (e.rate[1].? >= e.rate[0].? * g.k.margin) 1 else 0;
+ e.streak = if (e.best == e.mode) @min(e.streak + 1, g.k.longest) else 0;
+ e.best = e.mode;
+ e.exploring = false;
+ } else {
+ e.mode ^= 1; // the slower mode is measured again, in case the streams' text changed
+ e.exploring = true;
+ }
+ e.left = if (e.exploring) g.k.explore[e.mode] else g.k.exploit << e.streak;
+ e.skip = e.mode != before;
+ }
+};
+
+/// The tokens a round's outcomes committed (each outcome's `got`).
+pub fn committed(outcomes: anytype) u64 {
+ var n: u64 = 0;
+ for (outcomes) |o| n += o.got.len;
+ return n;
+}
+
+/// Rounds of a guard at `streams` with rates `plain` and `drafted` tokens a ms: how many it drafted.
+fn run(g: *Guard, streams: usize, rounds: usize, plain: f64, drafted: f64) !usize {
+ var n: usize = 0;
+ for (0..rounds) |_| {
+ const d = try g.drafts(streams);
+ n += @intFromBool(d);
+ try g.observe(streams, if (d) 3 else 1, if (d) 3 / drafted else 1 / plain);
+ }
+ return n;
+}
+
+test "a shared round drafts only while drafted rounds beat plain ones, measuring the slower mode again now and then" {
+ var g: Guard = .{ .gpa = std.testing.allocator };
+ defer g.deinit();
+ // drafted rounds slower at four streams: after one stretch of each, plain but for a re-measure a period
+ const slow = try run(&g, 4, 1000, 0.25, 0.20);
+ try std.testing.expect(slow <= 4 * (g.k.explore[1] + 1)); // the first measure and three re-measures, short ones
+ try std.testing.expect(slow > 0);
+ // the same verdict again and again: the slower mode is measured ever more rarely, under 1% of rounds
+ const later = try run(&g, 4, 20000, 0.25, 0.20);
+ try std.testing.expect(later * 500 < 20000);
+ // faster at one stream: drafted but for the plain re-measures
+ const fast = try run(&g, 1, 1000, 0.13, 0.15);
+ try std.testing.expect(fast > 800);
+ // each stream count keeps its own rates
+ try std.testing.expect(!try g.drafts(4) or (try g.entry(4)).exploring);
+ try std.testing.expect(try g.drafts(1) or (try g.entry(1)).exploring);
+}
+
+test "a mode's first round is not measured: it carries the other mode's drafts" {
+ var g: Guard = .{ .gpa = std.testing.allocator, .k = .{ .explore = .{ 2, 2 }, .exploit = 4 } };
+ defer g.deinit();
+ try std.testing.expect(!try g.drafts(2)); // plain first
+ try g.observe(2, 0, 1); // skipped: the stream count's first round
+ try g.observe(2, 1, 4);
+ try g.observe(2, 1, 4);
+ try std.testing.expect(try g.drafts(2)); // plain measured at 0.25 a ms; drafted next
+ try g.observe(2, 100, 1); // skipped: plain's last held rows, a fantasy rate that would keep drafting
+ try g.observe(2, 1, 10);
+ try g.observe(2, 1, 10);
+ try std.testing.expect(!try g.drafts(2)); // plain at 0.25 beats drafted at 0.1
+ try std.testing.expectEqual(@as(?f64, 0.1), (try g.entry(2)).rate[1]);
+}
+
+test "a drafted stretch starts from a grant of one: a plain stretch's grants of none would hold nothing to measure" {
+ var g: Guard = .{ .gpa = std.testing.allocator, .k = .{ .explore = .{ 2, 2 }, .exploit = 4 } };
+ defer g.deinit();
+ var st: ?depth.State = null;
+ var whos = [_]depth.Who{.{ .state = &st, .draft_room = 100, .finished = false, .forced = false, .granted = 0 }};
+ try std.testing.expect(!try g.draftsFresh(2, &whos)); // plain first: grants as they are
+ try std.testing.expectEqual(@as(?u32, 0), whos[0].granted);
+ for (0..3) |_| try g.observe(2, 1, 4);
+ try std.testing.expect(try g.draftsFresh(2, &whos)); // the drafted stretch's first round
+ try std.testing.expectEqual(@as(?u32, 1), whos[0].granted);
+ try g.observe(2, 1, 4); // its first round, skipped
+ whos[0].granted = 0;
+ try std.testing.expect(try g.draftsFresh(2, &whos)); // later rounds keep the allocation's grants
+ try std.testing.expectEqual(@as(?u32, 0), whos[0].granted);
+}
+
+test "a guard allocates once a stream count: its rounds after that take no allocation" {
+ var counted = std.testing.FailingAllocator.init(std.testing.allocator, .{});
+ var g: Guard = .{ .gpa = counted.allocator(), .k = .{ .explore = .{ 2, 2 }, .exploit = 4 } };
+ defer g.deinit();
+ for ([_]usize{ 1, 4, 8, 16 }) |streams| _ = try run(&g, streams, 1, 0.25, 0.2);
+ const first = .{ counted.allocations, counted.resize_index };
+ var st: [16]?depth.State = @splat(null);
+ var whos: [16]depth.Who = undefined;
+ for (&whos, &st) |*w, *s| w.* = .{ .state = s, .draft_room = 100, .finished = false, .forced = false, .granted = 1 };
+ for (0..200) |i| for ([_]usize{ 1, 4, 8, 16 }) |streams| {
+ _ = try g.draftsFresh(streams, whos[0..streams]);
+ try g.observe(streams, i % 3 + 1, 4);
+ };
+ try std.testing.expectEqual(first[0], counted.allocations);
+ try std.testing.expectEqual(first[1], counted.resize_index);
+}
diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig
index 1d76f9e7bf..b98302df55 100644
--- a/zig/src/core/lanes/stream.zig
+++ b/zig/src/core/lanes/stream.zig
@@ -203,6 +203,8 @@ pub const Stream = struct {
prices_seeded: bool = false, // its own prices started from the engine's kept calibration
copy_landed: [2]f64 = .{ -1, -1 }, // a lone context copy's rows that landed, decayed, by short and long match (-1: unprimed)
copy_broke: [2]f64 = .{ 0, 0 }, // and its chains that broke before their end: their ratio is a copied row's landing rate
+ copied_landed: [2]f64 = .{ -1, -1 }, // a copy window's rows that landed, decayed, by source: context, Recall (-1: unprimed)
+ copied_broke: [2]f64 = .{ 0, 0 }, // and its copies that broke before their end
drafted: u64 = 0,
accepted: u64 = 0,
branch_rows: u64 = 0, // suffix-match branch rows verified (lanes/fill.zig)
@@ -228,6 +230,8 @@ pub const Stream = struct {
served: ?u64 = null,
proposal_depth: ?u32 = null, // null: legacy; otherwise the last routed request, zero until it succeeds
granted: ?u32 = null,
+ grant_short: bool = false, // the last grant verified fewer drafts than the stream held
+ grant_gain: ?f64 = null, // the drafts its last grant expected to keep (the chances it was priced by)
copy_width: ?u32 = null,
pub fn init(gpa: Allocator, spec: Spec) !Stream {
@@ -291,6 +295,8 @@ pub const Stream = struct {
s.inflight = null;
s.served = null;
s.granted = null;
+ s.grant_short = false;
+ s.grant_gain = null;
s.copy_width = null;
s.proposal_depth = null;
}
diff --git a/zig/src/core/lanes/windows.zig b/zig/src/core/lanes/windows.zig
index e8e48be2ad..e4957cddc6 100644
--- a/zig/src/core/lanes/windows.zig
+++ b/zig/src/core/lanes/windows.zig
@@ -197,7 +197,7 @@ fn headTrunk(e: *Engine, s: *Stream, p: Plan, levels: usize) !?fill.Trunk {
/// A copy from the context backed by `enter_match` ids (Python `_copy_proposal`), else Recall's when it should pay.
pub fn copyProposal(e: *Engine, s: *Stream) !Copy {
- if (e.cfg.max_copy == 0 or s.force.items.len > 0 or (e.cfg.root_copies and s.root_copy_miss)) return .{};
+ if (e.cfg.max_copy == 0 or e.plain_round or s.force.items.len > 0 or (e.cfg.root_copies and s.root_copy_miss)) return .{};
const width: u32 = if (e.alone) (s.copy_width orelse e.cfg.first_copy) else e.cfg.first_copy;
const most = @min(@as(i64, width), s.draftRoom() - 1);
if (s.proposer) |p| if (!e.alone or e.cfg.lone_copies) {
@@ -226,7 +226,7 @@ pub fn probe(e: *Engine) depth.Probe {
}
pub fn who(s: *Stream) depth.Who {
- return .{ .state = &s.depth, .draft_room = s.draftRoom(), .finished = s.finished, .forced = s.force.items.len > 0, .granted = s.granted, .stream = s, .sampled = s.sampling != null };
+ return .{ .state = &s.depth, .draft_room = s.draftRoom(), .finished = s.finished, .forced = s.force.items.len > 0, .granted = s.granted, .grant_short = s.grant_short, .gain = s.grant_gain, .stream = s, .sampled = s.sampling != null };
}
/// Whether any of a round's streams draws keyed samples (its rows carry their host work).
@@ -248,7 +248,8 @@ pub fn allocate(e: *Engine, plans: []Plan) !void {
pr.* = &.{};
} else if (p.kind == .copy) {
fx.* = 1;
- pr.* = try alloc.chainProbabilities(a, &.{e.cfg.copy_rate}, n);
+ const rate = if (e.cfg.own_copy_rates) copiedRate(e, p.stream, p.recalled) else e.cfg.copy_rate;
+ pr.* = try alloc.chainProbabilities(a, &.{rate}, n);
} else {
fx.* = 1;
const chances = try planChances(e, p, n);
@@ -281,8 +282,16 @@ pub fn allocate(e: *Engine, plans: []Plan) !void {
const ms = own.get(1 + @as(i64, counts[0]));
try trail.event(e, &.{ f("ev", trail.str("lone-price")), f("stream", trail.str(plans[0].stream.id)), f("rows", trail.int(counts[0] + 1)), f("expected", ev.bits(expected)), f("ms", if (ms) |value| ev.bits(value) else .null), f("probe", .{ .bool = trial != null }) });
}
- for (plans, counts) |*p, c| {
- if (p.kind == .head) p.stream.granted = c;
+ for (plans, counts, probs) |*p, c, pr| {
+ if (p.kind == .head) {
+ p.stream.granted = c;
+ if (e.cfg.granted_depth) {
+ p.stream.grant_short = c < p.count();
+ var gain: f64 = 0;
+ for (pr[0..@min(c, pr.len)]) |chance| gain += chance;
+ p.stream.grant_gain = gain;
+ }
+ }
if (p.kind == .forced or p.count() == 0 or c == p.count()) continue;
try cut.prefix(a, p, c);
}
@@ -362,6 +371,31 @@ pub fn observeCopy(e: *const Engine, s: *Stream, trusted: bool, kept: u64, offer
}
}
+/// A copy window's history weighs this many rows of cfg.copy_rate's prior before the stream's own copies teach it.
+const copied_prior = 4;
+
+/// A copy window's row landing rate for this stream and source: landed rows over those plus breaks, decayed a copy.
+pub fn copiedRate(e: *const Engine, s: *Stream, recalled: bool) f64 {
+ const i = @intFromBool(recalled);
+ if (s.copied_landed[i] < 0) {
+ s.copied_landed[i] = copied_prior * e.cfg.copy_rate;
+ s.copied_broke[i] = copied_prior * (1 - e.cfg.copy_rate);
+ }
+ return s.copied_landed[i] / (s.copied_landed[i] + s.copied_broke[i]);
+}
+
+/// A verified copy window of `offered` rows, `kept` landing: its source's counters, then its own rate's history.
+pub fn observeCopied(e: *const Engine, s: *Stream, recalled: bool, kept: usize, offered: usize) void {
+ if (recalled) {
+ if (s.recall) |r| r.observe(@intCast(offered), @intCast(kept), s.rounds);
+ } else if (s.proposer) |pr| pr.observe(@intCast(offered), @intCast(kept));
+ if (!e.cfg.own_copy_rates) return;
+ _ = copiedRate(e, s, recalled);
+ const i = @intFromBool(recalled);
+ s.copied_landed[i] = copy_decay * s.copied_landed[i] + @as(f64, @floatFromInt(kept));
+ s.copied_broke[i] = copy_decay * s.copied_broke[i] + @as(f64, @floatFromInt(@intFromBool(kept < offered)));
+}
+
/// A context copy's appended rows: each its parent's chance times the stream's copy rate; the drafts keep their own.
fn planChances(e: *Engine, p: Plan, n: usize) ![]const f64 {
if (!p.context_copy or p.branch_from == 0 or p.parents == null or n < p.branch_from) return nodeChances(e, p.stream, n);
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..8af90dc43b 100644
--- a/zig/src/families/lfm2_moe/backend.zig
+++ b/zig/src/families/lfm2_moe/backend.zig
@@ -10,6 +10,14 @@ const snapshot = @import("snapshot.zig");
const Model = @import("model.zig").Model;
const be = lanes.backend;
+/// Where a drafter holds a stream's drafts on the device, for verify windows that take them there (Window.held).
+pub const HeldSource = struct {
+ ptr: *anyopaque,
+ get: *const fn (ptr: *anyopaque, s: *lanes.Stream) ?Device,
+ host: *const fn (ptr: *anyopaque, s: *lanes.Stream, out: []u32) anyerror!void, // where `get` gives none
+};
+pub const Device = struct { src: mtl.Buffer, off: usize, stride: u32, wait: fwd.Wait };
+
pub const Options = struct { capacity: usize, chunk: usize = 128, streams: u32 = 8 };
pub const Metal = struct {
@@ -23,6 +31,10 @@ 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,
+ /// A drafter whose drafts verify windows read on the device (null: windows carry host tokens only).
+ held_source: ?HeldSource = null,
pub fn init(gpa: std.mem.Allocator, model: *Model, options: Options) !*Metal {
const out = try gpa.create(Metal);
@@ -88,7 +100,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 +119,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 +190,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 +207,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 +242,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, &.{});
}
@@ -178,29 +256,54 @@ pub const Metal = struct {
const self = cast(ptr);
var segments: [st.batch_rows]fwd.Segment = undefined;
var ids: [st.batch_rows]u32 = undefined;
+ var held: [st.batch_rows]fwd.Held = undefined;
+ var count: usize = 0;
+ var wait: ?fwd.Wait = null;
var total: usize = 0;
if (windows.len > segments.len or windows.len != out.len) return error.InvalidLfmWindows;
for (windows, 0..) |w, i| {
const rows = w.rows();
- if (w.parents != null or w.held != 0 or rows > st.window_rows or total + rows > st.batch_rows or w.positions.len != rows) return error.UnsupportedLfmWindow;
+ if (w.held != 0) { // the drafter's picks, copied into the window's rows on the GPU
+ const source = self.held_source orelse return error.UnsupportedLfmWindow;
+ if (total + rows > st.batch_rows) return error.UnsupportedLfmWindow;
+ if (source.get(source.ptr, w.stream)) |dev| {
+ held[count] = .{ .row = total + 1, .src = dev.src, .off = dev.off, .stride = dev.stride, .count = w.held };
+ count += 1;
+ wait = if (wait) |x| (if (x.value >= dev.wait.value) x else dev.wait) else dev.wait;
+ @memset(ids[total + 1 ..][0..w.held], 0);
+ } else try source.host(source.ptr, w.stream, ids[total + 1 ..][0..w.held]);
+ }
+ if (w.parents != null or rows > st.window_rows or total + rows > st.batch_rows or w.positions.len != rows) return error.UnsupportedLfmWindow;
if (w.masks.len != 0 and w.masks.len != rows * mask_words) return error.InvalidLfmMasks;
const cache = try self.cacheOf(w.stream);
for (w.positions, 0..) |p, r| if (p != cache.len + r + 1) return error.InvalidLfmPosition;
segments[i] = .{ .cache = cache, .rows = rows };
ids[total] = w.pending;
- @memcpy(ids[total + 1 ..][0..w.tokens.len], w.tokens);
+ @memcpy(ids[total + 1 + w.held ..][0..w.tokens.len], w.tokens);
total += rows;
}
- try fwd.run(self.model, &self.scratch, segments[0..windows.len], ids[0..total], true, .all);
+ self.gen += 1;
+ try fwd.runHeld(self.model, &self.scratch, segments[0..windows.len], ids[0..total], true, .all, held[0..count], wait);
+ 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;
+ if (o.sampled.len != w.rows() or o.drafts.len != w.rows() - 1) 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);
+ @memcpy(o.drafts, self.scratch.ids.slice(u32, base + w.rows())[base + 1 ..][0 .. w.rows() - 1]);
base += w.rows();
}
}
@@ -306,6 +409,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.zig b/zig/src/families/lfm2_moe/dspark.zig
new file mode 100644
index 0000000000..24b7ae5aba
--- /dev/null
+++ b/zig/src/families/lfm2_moe/dspark.zig
@@ -0,0 +1,531 @@
+//! LiquidAI's DSpark drafter on LFM2's taps: a stream's context keys, a block of rows, a Markov chain of drafts.
+const std = @import("std");
+const mtl = @import("metal");
+const lanes = @import("lanes");
+const c = @import("config.zig");
+const dc = @import("dspark_config.zig");
+const dw = @import("dspark_weights.zig");
+const block = @import("dspark_block.zig");
+const st = @import("state.zig");
+const snapshot = @import("snapshot.zig");
+const Model = @import("model.zig").Model;
+const backend = @import("backend.zig");
+
+const dr = lanes.drafter;
+const Stream = lanes.Stream;
+const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked;
+
+/// Block rows one hold pass takes (the decode twin's rows), and context rows one absorb pass takes.
+pub const pass_rows = 32;
+pub const absorb_rows = 2048;
+/// Verified rows a lane keeps on the host until its next hold absorbs them (a round's rows for no GPU work).
+pub const stage_rows = 128;
+/// The blocks a hold is timed with at load, each at one draft and at four (two at 16 blocks: one 32-row pass).
+const hold_grid = [_]u32{ 1, 2, 4, 8, 16 };
+const kw = c.kv_heads * c.head_dim;
+
+/// `bits`: 0 runs the bf16 matrices as stored, 8 quantizes them at load.
+pub const Options = struct { capacity: usize, bits: u8 = 0 };
+
+/// One stream's context keys and values a layer ([kv head][cap][head dim], as the target's attention reads them).
+pub const Lane = struct {
+ mem: mtl.Buffer, // each layer's keys then values, then the staged rows (shared storage)
+ cap: usize, // positions each holds: the context and one block past it
+ stage_at: usize, // the staged rows' byte offset in `mem`
+ end: u64 = 0, // context rows absorbed: positions [0, end)
+ staged: u32 = 0, // tapped rows [end, end + staged), not absorbed yet
+ held: [16]u32 = @splat(0),
+ held_len: u32 = 0,
+ conf: [16]f32 = @splat(0), // each held draft's chance given the drafts before it landed (the confidence head)
+ conf_len: u32 = 0, // the drafts `conf` covers: the last settled hold's depth
+ slot: ?u32 = null, // its block's place in the last signalled hold, whose picks still hold its drafts
+
+ /// Layer `layer`'s keys, or its values when `values`, in `mem`.
+ pub fn kv(l: *const Lane, layer: usize, values: bool) block.Ref {
+ return .{ .buf = l.mem, .off = (2 * layer + @intFromBool(values)) * l.cap * kw * 2 };
+ }
+};
+
+/// The passes' working buffers, one pass at a time on the target's queue.
+pub const Scratch = struct {
+ ids: mtl.Buffer,
+ h: mtl.Buffer,
+ x: mtl.Buffer,
+ r: mtl.Buffer,
+ y: mtl.Buffer,
+ q: mtl.Buffer,
+ k: mtl.Buffer,
+ v: mtl.Buffer,
+ queries: mtl.Buffer,
+ pm: mtl.Buffer,
+ pl: mtl.Buffer,
+ po: mtl.Buffer,
+ gate: mtl.Buffer,
+ up: mtl.Buffer,
+ act: mtl.Buffer,
+ logits: mtl.Buffer,
+ sums: mtl.Buffer,
+ fcx: mtl.Buffer, // an absorb pass's context rows: fc's output, normed in place
+ dump: mtl.Buffer, // the query heads lfm2_qk rotates from context rows, which no one reads
+ best: mtl.Buffer, // the Markov scan's best score and token of each threadgroup, by stream
+ at: mtl.Buffer,
+ picks: mtl.Buffer,
+ conf: mtl.Buffer, // the confidence logits, step by step, a pass's blocks each
+};
+
+pub const Drafter = struct {
+ gpa: std.mem.Allocator,
+ io: std.Io,
+ target: *Model,
+ config: dc.Config,
+ w: dw.Weights,
+ options: Options,
+ s: Scratch,
+ owned: std.ArrayList(mtl.Buffer) = .empty,
+ lanes: std.AutoHashMapUnmanaged(*Stream, *Lane) = .empty,
+ /// One stream's hold wall ms at full depth, and that over the block (the depth rule's cost a draft).
+ hold_ms: f64 = 0,
+ step_ms: f64 = 0,
+ holds: [hold_grid.len * 2]lanes.hold_price.HoldCost = undefined, // hold passes timed at load, then by settled holds
+ holds_len: usize = 0,
+ event: mtl.Event = undefined,
+ signalled: u64 = 0,
+ pending: ?mtl.CommandBuffer = null,
+ pass: [pass_rows]*Lane = undefined, // the signalled hold's lanes and depths, in its blocks' order
+ depths: [pass_rows]u32 = undefined,
+ pass_len: usize = 0,
+
+ pub fn init(gpa: std.mem.Allocator, io: std.Io, target: *Model, dir: []const u8, options: Options) !*Drafter {
+ const d = try gpa.create(Drafter);
+ errdefer gpa.destroy(d);
+ d.* = .{ .gpa = gpa, .io = io, .target = target, .config = try dc.Config.read(gpa, io, dir), .w = undefined, .options = options, .s = undefined };
+ d.event = try target.device.event();
+ errdefer d.event.deinit();
+ d.w = try dw.load(gpa, target.device, io, dir, &d.config, options.bits);
+ errdefer d.w.deinit(gpa);
+ errdefer for (d.owned.items) |b| b.deinit();
+ errdefer d.owned.deinit(gpa);
+ try d.scratch();
+ try d.warm();
+ return d;
+ }
+
+ pub fn deinit(d: *Drafter) void {
+ d.settle() catch {};
+ d.event.deinit();
+ var it = d.lanes.valueIterator();
+ while (it.next()) |l| d.dropLane(l.*);
+ d.lanes.deinit(d.gpa);
+ for (d.owned.items) |b| b.deinit();
+ d.owned.deinit(d.gpa);
+ d.w.deinit(d.gpa);
+ d.gpa.destroy(d);
+ }
+
+ fn buffer(d: *Drafter, bytes: usize) !mtl.Buffer {
+ const b = try d.target.device.buffer(@max(bytes, 16), shared);
+ errdefer b.deinit();
+ try d.owned.append(d.gpa, b);
+ return b;
+ }
+
+ fn scratch(d: *Drafter) !void {
+ const r = pass_rows;
+ const wide = @max(pass_rows, absorb_rows);
+ const i = d.config.intermediate;
+ const partials = @max(st.split_rows, d.config.block) * st.subs * c.query_heads * st.chunks;
+ const sums = @max(absorb_rows * (d.config.tap_count * c.hidden / 64), r * (i / 64));
+ d.s = .{
+ .ids = try d.buffer(r * 4),
+ .h = try d.buffer(r * c.hidden * 2),
+ .x = try d.buffer(r * c.hidden * 2),
+ .r = try d.buffer(r * c.hidden * 2),
+ .y = try d.buffer(r * c.hidden * 2),
+ .q = try d.buffer(r * c.hidden * 2),
+ .k = try d.buffer(wide * kw * 2),
+ .v = try d.buffer(wide * kw * 2),
+ .queries = try d.buffer((r + st.attn_rows) * c.hidden * 2),
+ .pm = try d.buffer(partials * 4),
+ .pl = try d.buffer(partials * 4),
+ .po = try d.buffer(partials * c.head_dim * 4),
+ .gate = try d.buffer(r * i * 2),
+ .up = try d.buffer(r * i * 2),
+ .act = try d.buffer(r * i * 2),
+ .logits = try d.buffer(r * c.vocab * 2),
+ .sums = try d.buffer(sums * 4),
+ .fcx = try d.buffer(absorb_rows * c.hidden * 2),
+ .dump = try d.buffer(absorb_rows * c.hidden * 2),
+ .best = try d.buffer(r * block.scan_groups * 4),
+ .at = try d.buffer(r * block.scan_groups * 4),
+ .picks = try d.buffer(d.config.block * r * 8),
+ .conf = try d.buffer(d.config.block * r * 4),
+ };
+ }
+
+ /// The drafter's buffers that every round reads (weights and scratch), for the server's residency set.
+ pub fn buffers(d: *const Drafter, gpa: std.mem.Allocator) ![]mtl.Buffer {
+ return std.mem.concat(gpa, mtl.Buffer, &.{ d.w.buffers(), d.owned.items });
+ }
+
+ pub fn drafter(d: *Drafter) dr.Drafter {
+ return .{ .ptr = d, .vtable = &.{ .taps = tapsFn, .facts = factsFn, .open = openFn, .absorb = absorbFn, .hold = holdFn, .held = heldFn, .release = releaseFn, .chances = chancesFn } };
+ }
+
+ /// Whether holds compute the confidence head: the verify rows follow its chances, else the core's depth rule.
+ pub fn confident(d: *const Drafter) bool {
+ return d.w.confidence != null;
+ }
+
+ fn cast(ptr: *anyopaque) *Drafter {
+ return @ptrCast(@alignCast(ptr));
+ }
+
+ fn tapsFn(ptr: *anyopaque) []const u32 {
+ const d = cast(ptr);
+ return d.config.taps[0..d.config.tap_count];
+ }
+
+ /// Acceptance a depth until a stream has its own: the card's 5.16 tokens a step at block 9 is 0.84 a depth.
+ const prior: [16]f64 = @splat(0.84);
+
+ fn factsFn(ptr: *anyopaque) dr.Facts {
+ const d = cast(ptr);
+ var f: dr.Facts = .{ .depth = d.config.block, .step_ms = d.step_ms, .prior = prior[0..d.config.block], .batched = true, .shared_guard = true, .stream_bytes = d.laneBytes(d.options.capacity), .copy_first = true };
+ f.chances = d.confident();
+ f.hold_costs = d.holds[0..d.holds_len];
+ return f;
+ }
+
+ /// A held chain's chances: the confidence head's running product, or for a fixed window, ones that fall by depth.
+ fn chancesFn(ptr: *anyopaque, s: *Stream, out: []f64) anyerror!bool {
+ const d = cast(ptr);
+ const l = d.lanes.get(s) orelse return false;
+ if (!d.confident()) return false;
+ if (d.pending != null and l.slot != null and l.conf_len > 0) {
+ // the hold still runs: the lane's last settled confidences price it, so the host does not wait on the GPU
+ if (out.len > d.depths[l.slot.?]) return error.DraftHeldFewer;
+ var reach: f64 = 1;
+ for (out, 0..) |*o, j| {
+ reach *= l.conf[@min(j, l.conf_len - 1)];
+ o.* = reach;
+ }
+ return true;
+ }
+ try d.settle();
+ if (out.len > l.held_len) return error.DraftHeldFewer;
+ var reach: f64 = 1;
+ for (out, l.conf[0..out.len]) |*o, p| {
+ reach *= p;
+ o.* = reach;
+ }
+ return true;
+ }
+
+ /// A stream's lane at `capacity` positions: its context keys and values a layer, and its staged rows.
+ pub fn laneBytes(d: *const Drafter, capacity: usize) usize {
+ return @intCast(d.keptBytes(capacity + d.config.block) + stage_rows * d.tapBytes());
+ }
+
+ /// A tapped row's bytes: every tap's hidden row, bf16.
+ pub fn tapBytes(d: *const Drafter) usize {
+ return d.config.tap_count * c.hidden * 2;
+ }
+
+ /// The lane's staged rows as an absorb piece (null: none staged).
+ pub fn stagedPiece(l: *Lane) ?block.Piece {
+ if (l.staged == 0) return null;
+ return .{ .lane = l, .src = .{ .buf = l.mem, .off = l.stage_at }, .rows = l.staged, .p0 = l.end };
+ }
+
+ /// The lane's staged rows absorbed now, in a command buffer of their own.
+ fn flush(d: *Drafter, l: *Lane) !void {
+ const pc = stagedPiece(l) orelse return;
+ const pool = mtl.objc.Pool.push();
+ defer pool.pop();
+ const run = block.begin(d);
+ try block.absorbPass(d, run.e, &.{pc});
+ try block.end(run.cb, run.e);
+ l.end += l.staged;
+ l.staged = 0;
+ }
+
+ fn openFn(ptr: *anyopaque, s: *Stream) anyerror!void {
+ const d = cast(ptr);
+ if (d.lanes.contains(s)) return error.DraftStreamOpen;
+ const want = s.prompt_len + s.max_new + 2 * d.config.block;
+ const l = try d.newLane(@min(d.options.capacity, (want + 63) / 64 * 64));
+ errdefer d.dropLane(l);
+ try d.lanes.put(d.gpa, s, l);
+ }
+
+ /// A lane in one buffer of laneBytes(capacity).
+ pub fn newLane(d: *Drafter, capacity: usize) !*Lane {
+ const l = try d.gpa.create(Lane);
+ errdefer d.gpa.destroy(l);
+ const at: usize = @intCast(d.keptBytes(capacity + d.config.block));
+ l.* = .{ .mem = try d.target.device.buffer(d.laneBytes(capacity), shared), .cap = capacity + d.config.block, .stage_at = at };
+ return l;
+ }
+
+ pub fn dropLane(d: *Drafter, l: *Lane) void {
+ l.mem.deinit();
+ d.gpa.destroy(l);
+ }
+
+ fn releaseFn(ptr: *anyopaque, s: *Stream) void {
+ const d = cast(ptr);
+ d.unpass() catch {};
+ if (d.lanes.fetchRemove(s)) |e| d.dropLane(e.value);
+ }
+
+ fn heldFn(ptr: *anyopaque, streams: []const *Stream, out: []const []u32) anyerror!void {
+ const d = cast(ptr);
+ try d.settle();
+ for (streams, out) |s, o| {
+ const l = d.lanes.get(s) orelse return error.DraftStreamUnknown;
+ if (o.len > l.held_len) return error.DraftHeldFewer;
+ @memcpy(o, l.held[0..o.len]);
+ }
+ }
+
+ /// Bytes a lane's context at `at` takes beside a kept prompt state.
+ pub fn keptBytes(d: *const Drafter, at: usize) u64 {
+ return 2 * @as(u64, d.config.layers) * at * kw * 2;
+ }
+
+ /// Stream `s`'s context copied into `snap` when its lane stands at the snapshot's position (a prompt chunk's end).
+ pub fn keepLane(d: *Drafter, s: *Stream, snap: *snapshot.Snap) !void {
+ const l = d.lanes.get(s) orelse return;
+ try d.settle();
+ try d.flush(l);
+ if (l.end != snap.at) return;
+ const bytes = try d.gpa.alloc(u8, @intCast(d.keptBytes(snap.at)));
+ copyLane(d, l, snap.at, bytes, true);
+ snap.draft = bytes;
+ }
+
+ /// The lane takes the context kept with the state its stream's prompt pass resumed at `start`.
+ fn resumeLane(d: *Drafter, s: *Stream, l: *Lane, start: u64) !void {
+ const saved = s.reuse.saved orelse return error.DraftContextGap;
+ const snap: *const snapshot.Snap = @ptrCast(@alignCast(saved));
+ const bytes = snap.draft orelse return error.DraftContextGap;
+ if (s.cached != start or snap.at != start or l.end != 0 or start + d.config.block > l.cap) return error.DraftContextGap;
+ copyLane(d, l, snap.at, bytes, false);
+ l.end = start;
+ }
+
+ /// Rows [0, at) of every layer's keys and values and `bytes` (`out`: the lane's into bytes), kv head by kv head.
+ fn copyLane(d: *const Drafter, l: *const Lane, at: usize, bytes: []u8, out: bool) void {
+ const n = at * c.head_dim * 2;
+ var off: usize = 0;
+ for (0..d.config.layers) |i| for ([_]block.Ref{ l.kv(i, false), l.kv(i, true) }) |r| for (0..c.kv_heads) |h| {
+ const live = (r.buf.contents() + r.off + h * l.cap * c.head_dim * 2)[0..n];
+ if (out) @memcpy(bytes[off..][0..n], live) else @memcpy(live, bytes[off..][0..n]);
+ off += n;
+ };
+ }
+
+ /// Rows from `start` replace the context there: rows that fit the stage wait on the host for the next hold.
+ fn absorbFn(ptr: *anyopaque, items: []const dr.Absorb) anyerror!void {
+ const d = cast(ptr);
+ try d.settle();
+ const width = d.tapBytes();
+ const pool = mtl.objc.Pool.push();
+ defer pool.pop();
+ var run: ?@TypeOf(block.begin(d)) = null;
+ var pieces: [32]block.Piece = undefined;
+ var count: usize = 0;
+ var rows: usize = 0;
+ for (items) |it| {
+ const l = d.lanes.get(it.stream) orelse return error.DraftStreamUnknown;
+ const f = it.features;
+ if (f.space != .device or f.dtype != .bf16 or f.row_bytes != width or it.rows != null) return error.DraftFeaturesUnsupported;
+ if (it.start > l.end + l.staged) try d.resumeLane(it.stream, l, it.start); // a prompt pass that resumed
+ if (it.start + f.rows + d.config.block > l.cap) return error.DraftContextGap;
+ l.held_len = 0;
+ l.slot = null;
+ const src = mtl.Buffer{ .id = @ptrFromInt(f.buffer) };
+ // rows from `start` replace what was staged past it; rows before the absorbed end go to the GPU now
+ const keep: u32 = if (it.start >= l.end) @intCast(it.start - l.end) else 0;
+ l.staged = @min(l.staged, keep);
+ if (it.start >= l.end and l.staged + f.rows <= stage_rows) {
+ const from = (src.contents() + f.offset)[0 .. f.rows * width];
+ @memcpy((l.mem.contents() + l.stage_at + l.staged * width)[0 .. f.rows * width], from);
+ l.staged += f.rows;
+ continue;
+ }
+ if (run == null) run = block.begin(d);
+ const e = run.?.e;
+ if (stagedPiece(l)) |pc| {
+ if (count == pieces.len or rows + pc.rows > absorb_rows) {
+ try block.absorbPass(d, e, pieces[0..count]);
+ count = 0;
+ rows = 0;
+ }
+ pieces[count] = pc;
+ count += 1;
+ rows += pc.rows;
+ }
+ var at: usize = 0;
+ while (at < f.rows) {
+ if (rows == absorb_rows or count == pieces.len) { // the pass is full: encode it
+ try block.absorbPass(d, e, pieces[0..count]);
+ count = 0;
+ rows = 0;
+ }
+ const n = @min(absorb_rows - rows, f.rows - at);
+ pieces[count] = .{ .lane = l, .src = .{ .buf = src, .off = f.offset + at * width }, .rows = n, .p0 = it.start + at };
+ count += 1;
+ rows += n;
+ at += n;
+ }
+ l.end = it.start + f.rows;
+ l.staged = 0;
+ }
+ const r = run orelse return;
+ if (count > 0) try block.absorbPass(d, r.e, pieces[0..count]);
+ try block.end(r.cb, r.e);
+ }
+
+ /// Each stream's chain of `depth` drafts after `pending` at `position`, its block's rows in a pass of pass_rows.
+ fn holdFn(ptr: *anyopaque, items: []const dr.Hold) anyerror!void {
+ const d = cast(ptr);
+ try d.unpass();
+ var at: usize = 0;
+ while (at < items.len) {
+ // as many streams as their blocks' rows fit a pass: each block is as long as the pass's deepest chain
+ var deepest: usize = 0;
+ var end = at;
+ while (end < items.len) : (end += 1) {
+ const rows = @max(deepest, items[end].depth);
+ if ((end - at + 1) * rows > pass_rows) break;
+ deepest = rows;
+ }
+ if (end == at) return error.DraftDepth;
+ var blocks: [pass_rows]block.Block = undefined;
+ const group = items[at..end];
+ at = end;
+ for (group, blocks[0..group.len]) |it, *b| {
+ const l = d.lanes.get(it.stream) orelse return error.DraftStreamUnknown;
+ if (it.depth == 0 or it.depth > d.config.block or it.position == 0) return error.DraftDepth;
+ const anchor = it.position - 1;
+ if (l.end + l.staged != anchor or anchor + d.config.block > l.cap) return error.DraftContextGap;
+ b.* = .{ .lane = l, .pending = it.pending, .anchor = anchor, .depth = it.depth };
+ }
+ if (group.len == items.len) { // one pass: signalled, not waited on
+ d.signalled += 1;
+ d.pending = try block.hold(d, blocks[0..group.len], .{ .event = d.event, .value = d.signalled });
+ for (blocks[0..group.len], 0..) |b, bi| {
+ d.pass[bi] = b.lane;
+ d.depths[bi] = b.depth;
+ b.lane.slot = @intCast(bi);
+ b.lane.held_len = 0;
+ }
+ d.pass_len = group.len;
+ continue;
+ }
+ _ = try block.hold(d, blocks[0..group.len], null);
+ }
+ }
+
+ /// The signalled hold waited on, its picks read into its lanes' held chains (nothing pending: nothing done).
+ fn settle(d: *Drafter) !void {
+ const cb = d.pending orelse return;
+ d.pending = null;
+ defer cb.release();
+ cb.wait();
+ if (cb.failure()) |msg| {
+ std.log.err("lfm2 dspark: a hold failed on the GPU: {s}", .{msg});
+ return error.DraftGpuFailed;
+ }
+ d.priceHold((cb.gpuEnd() - cb.gpuStart()) * 1000);
+ try block.readPicks(d, d.pass[0..d.pass_len], d.depths[0..d.pass_len]);
+ }
+
+ /// A settled hold's GPU ms against its price: the timed passes at its block counts move a tenth of the way to it.
+ pub fn priceHold(d: *Drafter, ms: f64) void {
+ const timed = d.holds[0..d.holds_len];
+ if (timed.len == 0 or d.pass_len == 0 or !(ms > 0)) return;
+ const deepest = std.mem.max(u32, d.depths[0..d.pass_len]);
+ const priced = lanes.hold_price.holdMs(timed, d.pass_len, deepest);
+ if (!(priced > 0)) return;
+ const step = 1 + 0.1 * (ms / priced - 1);
+ var lo: u32 = 0;
+ var hi: u32 = std.math.maxInt(u32);
+ for (timed) |h| {
+ if (h.streams <= d.pass_len) lo = @max(lo, h.streams);
+ if (h.streams >= d.pass_len) hi = @min(hi, h.streams);
+ }
+ for (timed) |*h| if (h.streams >= lo and h.streams <= hi) {
+ h.ms *= step;
+ };
+ }
+
+ /// The signalled hold settled and its lanes' picks forgotten (a lane may be dropped, or a hold may overwrite).
+ fn unpass(d: *Drafter) !void {
+ defer {
+ for (d.pass[0..d.pass_len]) |l| l.slot = null;
+ d.pass_len = 0;
+ }
+ try d.settle();
+ }
+
+ /// Where stream `s`'s drafts from the last signalled hold lie in the picks, for a verify to copy on the GPU.
+ pub fn device(ptr: *anyopaque, s: *Stream) ?backend.Device {
+ const d = cast(ptr);
+ const l = d.lanes.get(s) orelse return null;
+ const slot = l.slot orelse return null;
+ return .{ .src = d.s.picks, .off = slot * 8, .stride = pass_rows * 2, .wait = .{ .event = d.event, .value = d.signalled } };
+ }
+
+ /// The target's verify reading the held drafts on the device, or as host tokens where it cannot.
+ pub fn heldSource(d: *Drafter) backend.HeldSource {
+ return .{ .ptr = d, .get = device, .host = hostHeld };
+ }
+
+ fn hostHeld(ptr: *anyopaque, s: *Stream, out: []u32) anyerror!void {
+ return heldFn(ptr, &.{s}, &.{out});
+ }
+
+ fn mibOf(bufs: []const mtl.Buffer) f64 {
+ var n: usize = 0;
+ for (bufs) |b| n += b.length();
+ return @as(f64, @floatFromInt(n)) / (1024.0 * 1024.0);
+ }
+
+ /// Holds on empty contexts: the pipelines compiled, then holds timed by blocks and depth (best of three).
+ fn warm(d: *Drafter) !void {
+ var made: [hold_grid[hold_grid.len - 1]]*Lane = undefined;
+ var n: usize = 0;
+ defer for (made[0..n]) |l| d.dropLane(l);
+ while (n < made.len) : (n += 1) made[n] = try d.newLane(64);
+ const l = made[0];
+ var best: [2]f64 = @splat(std.math.floatMax(f64));
+ for (0..4) |i| for ([_]u32{ 1, d.config.block }, &best) |depth, *b| {
+ const t0 = std.Io.Clock.awake.now(d.io);
+ _ = try block.hold(d, &.{.{ .lane = l, .pending = 1, .anchor = 0, .depth = depth }}, null);
+ const ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(d.io)).toNanoseconds())) / 1e6;
+ if (i > 0) b.* = @min(b.*, ms);
+ };
+ d.hold_ms = best[1];
+ d.step_ms = best[1] / @as(f64, @floatFromInt(d.config.block));
+ std.log.info("lfm2 dspark: a 1-draft hold takes {d:.2} ms, a {d}-draft hold {d:.2} ms", .{ best[0], d.config.block, best[1] });
+ std.log.info("lfm2 dspark: weights {d:.1} MiB, scratch {d:.1} MiB, a stream {d} KiB a position and {d} KiB staged", .{ mibOf(d.w.buffers()), mibOf(d.owned.items), d.keptBytes(1) / 1024, stage_rows * d.tapBytes() / 1024 });
+ var blocks: [hold_grid[hold_grid.len - 1]]block.Block = undefined;
+ for (made, &blocks) |lane, *b| b.* = .{ .lane = lane, .pending = 1, .anchor = 0, .depth = 1 };
+ d.holds_len = 0;
+ for (hold_grid) |streams| for ([_]u32{ 1, @min(4, pass_rows / streams) }, 0..) |depth, j| {
+ if (j == 1 and depth == 1) continue;
+ for (blocks[0..streams]) |*b| b.depth = depth;
+ var ms_best = std.math.floatMax(f64);
+ for (0..4) |i| {
+ const t0 = std.Io.Clock.awake.now(d.io);
+ _ = try block.hold(d, blocks[0..streams], null);
+ const ms = @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(d.io)).toNanoseconds())) / 1e6;
+ if (i > 0) ms_best = @min(ms_best, ms);
+ }
+ d.holds[d.holds_len] = .{ .streams = streams, .depth = depth, .ms = ms_best };
+ d.holds_len += 1;
+ };
+ }
+};
diff --git a/zig/src/families/lfm2_moe/dspark_bf16_test.zig b/zig/src/families/lfm2_moe/dspark_bf16_test.zig
new file mode 100644
index 0000000000..293549bde1
--- /dev/null
+++ b/zig/src/families/lfm2_moe/dspark_bf16_test.zig
@@ -0,0 +1,121 @@
+//! The DSpark drafter's bf16 kernels on the GPU against host references: the linears and the Markov scan over bf16 W2.
+const std = @import("std");
+const mtl = @import("metal");
+const quant = @import("dspark_quant.zig");
+const base = @import("dspark_test.zig");
+
+const Args = extern struct { rows: u32, n: u32, k: u32 };
+
+/// y = x W^T through `key` on random bf16 data, each output within a bf16 step of the fp64 product.
+fn linear(comptime key: []const u8, rows: usize, n: usize, k: usize, groups: [2]usize, threads: usize) !void {
+ var r = try base.Rig.init();
+ defer r.deinit();
+ const x = try r.buffer(u16, rows * k);
+ defer x.deinit();
+ const w = try r.buffer(u16, n * k);
+ defer w.deinit();
+ const y = try r.buffer(u16, rows * n);
+ defer y.deinit();
+ base.noise(31, x.slice(u16, rows * k), 2);
+ base.noise(32, w.slice(u16, n * k), 0.5);
+ const cb = r.queue.commandBuffer();
+ const e = cb.compute(.serial);
+ e.setPipeline(r.k.get(key));
+ e.setBuffer(x, 0, 0);
+ e.setBuffer(w, 0, 1);
+ e.setValue(Args{ .rows = @intCast(rows), .n = @intCast(n), .k = @intCast(k) }, 2);
+ e.setBuffer(y, 0, 3);
+ e.dispatchGroups(mtl.Size.of(groups[0], groups[1], 1), mtl.Size.of(threads, 1, 1));
+ e.end();
+ cb.commit();
+ cb.wait();
+ if (cb.failure()) |_| return error.GpuFailed;
+ const xs = x.slice(u16, rows * k);
+ const ws = w.slice(u16, n * k);
+ const ys = y.slice(u16, rows * n);
+ for (0..rows) |i| for (0..n) |j| {
+ var dot: f64 = 0;
+ for (0..k) |t| dot += @as(f64, quant.toF32(xs[i * k + t])) * @as(f64, quant.toF32(ws[j * k + t]));
+ const got: f64 = quant.toF32(ys[i * n + j]);
+ try std.testing.expect(@abs(got - dot) <= @abs(dot) * 0x1p-7 + 1e-3);
+ };
+}
+
+test "lfm2 dspark: the bf16 rows kernel is x W^T for one row, a full 16 and a ragged 19 (two row groups)" {
+ for ([_]usize{ 1, 16, 19 }) |rows| try linear("lfm2_bf16_rows", rows, 100, 288, .{ 13, (rows + 15) / 16 }, 256);
+}
+
+test "lfm2 dspark: the bf16 tile kernel is x W^T for a ragged 70 rows and 100 columns, k a multiple of 32" {
+ try linear("lfm2_bf16_mma", 70, 100, 288, .{ 2, 2 }, 128);
+}
+
+test "lfm2 dspark: the bf16 tile kernel at the drafter's widths" {
+ try linear("lfm2_bf16_mma", 33, 512, 2048, .{ 8, 1 }, 128);
+}
+
+test "lfm2 dspark: a Markov step over bf16 W2 picks the argmax of bf16(base + bf16(W2 . W1[prev])), a NaN row flagged" {
+ var r = try base.Rig.init();
+ defer r.deinit();
+ const rank = 256;
+ const vocab = 1300;
+ const streams = 4;
+ const groups = (vocab + 511) / 512;
+ const w1 = try r.buffer(u16, vocab * rank);
+ defer w1.deinit();
+ base.noise(11, w1.slice(u16, vocab * rank), 1);
+ const w2 = try r.buffer(u16, vocab * rank);
+ defer w2.deinit();
+ base.noise(12, w2.slice(u16, vocab * rank), 0.5);
+ const ids = try r.buffer(u32, 2 * streams);
+ defer ids.deinit();
+ @memcpy(ids.slice(u32, 2 * streams), &[_]u32{ 17, 9, vocab - 1, 9, vocab + 5, 0, 3, 3 });
+ const logits = try r.buffer(u16, streams * 2 * vocab);
+ defer logits.deinit();
+ base.noise(14, logits.slice(u16, streams * 2 * vocab), 4);
+ logits.slice(u16, streams * 2 * vocab)[3 * 2 * vocab + 700] = 0x7fc0;
+ const best = try r.buffer(f32, streams * groups);
+ defer best.deinit();
+ const at = try r.buffer(u32, streams * groups);
+ defer at.deinit();
+ const picks = try r.buffer(u32, 2 * streams);
+ defer picks.deinit();
+ const cb = r.queue.commandBuffer();
+ const e = cb.compute(.serial);
+ e.setPipeline(r.k.get("lfm2_markov_scan_bf16"));
+ for ([_]mtl.Buffer{ logits, w1, w2, ids }, 0..) |b, i| e.setBuffer(b, 0, i);
+ e.setValue([4]u32{ 2 * vocab, rank, vocab, 2 }, 4);
+ for ([_]mtl.Buffer{ best, at }, 5..) |b, i| e.setBuffer(b, 0, i);
+ e.dispatchGroups(mtl.Size.of(groups, streams, 1), mtl.Size.of(256, 1, 1));
+ e.setPipeline(r.k.get("lfm2_markov_pick"));
+ e.setBuffer(best, 0, 0);
+ e.setBuffer(at, 0, 1);
+ e.setValue(@as(u32, groups), 2);
+ e.setBuffer(picks, 0, 3);
+ e.dispatchGroups(mtl.Size.of(streams, 1, 1), mtl.Size.of(256, 1, 1));
+ e.end();
+ cb.commit();
+ cb.wait();
+ const got = picks.slice(u32, 2 * streams);
+ const logit = logits.slice(u16, streams * 2 * vocab);
+ for (0..streams) |s| {
+ if (s == 3) {
+ try std.testing.expectEqual(@as(u32, 1), got[2 * s + 1]);
+ continue;
+ }
+ try std.testing.expectEqual(@as(u32, 0), got[2 * s + 1]);
+ const id = ids.slice(u32, 2 * streams)[2 * s];
+ var top: f64 = -std.math.inf(f64);
+ var mine: f64 = 0;
+ for (0..vocab) |v| {
+ var dot: f64 = 0;
+ for (0..rank) |i| {
+ const e1: f64 = if (id < vocab) quant.toF32(w1.slice(u16, vocab * rank)[id * rank + i]) else 0;
+ dot += @as(f64, quant.toF32(w2.slice(u16, vocab * rank)[v * rank + i])) * e1;
+ }
+ const score: f64 = quant.toF32(quant.toBf16(quant.toF32(logit[s * 2 * vocab + v]) + quant.toF32(quant.toBf16(@floatCast(dot)))));
+ top = @max(top, score);
+ if (v == got[2 * s]) mine = score;
+ }
+ try std.testing.expect(got[2 * s] < vocab and mine >= top - @abs(top) * 0x1p-7);
+ }
+}
diff --git a/zig/src/families/lfm2_moe/dspark_block.zig b/zig/src/families/lfm2_moe/dspark_block.zig
new file mode 100644
index 0000000000..c668c67109
--- /dev/null
+++ b/zig/src/families/lfm2_moe/dspark_block.zig
@@ -0,0 +1,356 @@
+//! DSpark's passes on the target's queue: context rows into a lane's keys, and a block's rows into Markov drafts.
+const std = @import("std");
+const mtl = @import("metal");
+const c = @import("config.zig");
+const st = @import("state.zig");
+const fwd = @import("forward.zig");
+const tile = @import("../../core/tile_qmm.zig");
+const am = @import("../../core/affine_mm.zig");
+const round = @import("tree_round_gpu");
+const ds = @import("dspark.zig");
+const dw = @import("dspark_weights.zig");
+const Drafter = ds.Drafter;
+const Lane = ds.Lane;
+
+const kw = c.kv_heads * c.head_dim;
+const QkArgs = extern struct { len: u32, cap: u32, eps: f32, log2_theta: f32 };
+const AttnArgs = extern struct { len: u32, cap: u32, scale: f32, subs: u32 };
+const Pick = round.contract.Pick;
+const Bf16Args = extern struct { rows: u32, n: u32, k: u32 };
+const MarkovArgs = extern struct { stride: u32, rank: u32, vocab: u32, prev_stride: u32 };
+const ConfidenceArgs = extern struct { n: u32, pass: u32, vocab: u32, rank: u32, hidden: u32 };
+/// Tokens a Markov scan threadgroup takes (the kernel's MARKOV_SPAN), and threadgroups a stream's row takes.
+pub const scan_span = 512;
+pub const scan_groups = (c.vocab + scan_span - 1) / scan_span;
+
+/// A buffer and byte offset.
+pub const Ref = struct { buf: mtl.Buffer, off: usize = 0 };
+
+/// One stream's block in a hold pass: rows [pending, mask x (block - 1)] at positions anchor .. anchor + block - 1.
+pub const Block = struct { lane: *Lane, pending: u32, anchor: u64, depth: u32 };
+
+const Enc = struct {
+ d: *Drafter,
+ e: mtl.ComputeEncoder,
+
+ fn pipe(x: Enc, comptime key: []const u8) void {
+ x.e.setPipeline(x.d.target.kernels.get(key));
+ }
+
+ /// y[rows] = x[rows] W^T: the twin up to its rows, past them the M5's tensor-unit kernels, else the tiles.
+ fn project(x: Enc, lin: dw.Linear, in: Ref, out: Ref, rows: usize) !void {
+ const w = switch (lin) {
+ .codes => |m| m,
+ .bf16 => |m| return x.bf16(m, in, out, rows),
+ };
+ const t = &x.d.target.kernels.tiles;
+ if (rows > tile.twin_rows) if (x.d.target.kernels.naxPipes(w.bits)) |p| {
+ const sums = Ref{ .buf = x.d.s.sums };
+ am.rowSums(x.e, p, c.group, in, sums, @intCast(rows), @intCast(w.k));
+ const q = .{ .w = Ref{ .buf = w.w, .off = w.w_off }, .s = Ref{ .buf = w.scales, .off = w.s_off }, .b = Ref{ .buf = w.biases, .off = w.b_off }, .n = @as(u32, @intCast(w.n)), .k = @as(u32, @intCast(w.k)) };
+ return am.dense(x.e, p, in, sums, q, out, @intCast(rows));
+ };
+ if (rows <= tile.twin_rows and x.d.target.kernels.twin) return tile.encodeTwin(x.e, t, w, in.buf, in.off, out.buf, out.off, rows);
+ tile.encodeSums(x.e, t, w.k, in.buf, in.off, x.d.s.sums, rows);
+ try tile.encode(x.e, t, w, in.buf, in.off, x.d.s.sums, 0, out.buf, out.off, rows);
+ }
+
+ /// y[rows] = x[rows] W^T over the released bf16 matrix: 16 rows a pass of W up to 32 rows, else 64x64 tiles.
+ fn bf16(x: Enc, m: dw.Bf16, in: Ref, out: Ref, rows: usize) void {
+ const args = Bf16Args{ .rows = @intCast(rows), .n = @intCast(m.n), .k = @intCast(m.k) };
+ const few = rows <= 32;
+ if (few) x.pipe("lfm2_bf16_rows") else x.pipe("lfm2_bf16_mma");
+ x.e.setBuffer(in.buf, in.off, 0);
+ x.e.setBuffer(m.w, 0, 1);
+ x.e.setBytes(std.mem.asBytes(&args), 2);
+ x.e.setBuffer(out.buf, out.off, 3);
+ if (few) {
+ x.e.dispatchGroups(mtl.Size.of((m.n + 7) / 8, (rows + 15) / 16, 1), mtl.Size.of(256, 1, 1));
+ } else x.e.dispatchGroups(mtl.Size.of((m.n + 63) / 64, (rows + 63) / 64, 1), mtl.Size.of(128, 1, 1));
+ }
+
+ /// h += r when `add`, then out = norm(h) * gain: LFM2's own norm (the drafter's width and eps are the target's).
+ fn norm(x: Enc, h: mtl.Buffer, r: mtl.Buffer, gain: mtl.Buffer, add: bool, out: mtl.Buffer, rows: usize) void {
+ x.pipe("lfm2_norm");
+ x.e.setBuffer(h, 0, 0);
+ x.e.setBuffer(r, 0, 1);
+ x.e.setBuffer(gain, 0, 2);
+ x.e.setValue(c.eps, 3);
+ x.e.setValue(@as(u32, @intFromBool(add)), 4);
+ x.e.setBuffer(out, 0, 5);
+ x.e.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(c.hidden / 4, 1, 1));
+ }
+
+ /// Rows' queries normed and rotated into `queries`, keys and values into the lane's layer at positions from `len`.
+ fn qk(x: Enc, q: Ref, k: Ref, v: Ref, layer: usize, lane: *const Lane, len: u64, queries: Ref, rows: usize) void {
+ const ly = &x.d.w.layers[layer];
+ x.pipe("lfm2_qk");
+ x.e.setBuffer(q.buf, q.off, 0);
+ x.e.setBuffer(k.buf, k.off, 1);
+ x.e.setBuffer(v.buf, v.off, 2);
+ x.e.setBuffer(ly.q_norm, 0, 3);
+ x.e.setBuffer(ly.k_norm, 0, 4);
+ const args = QkArgs{ .len = @intCast(len), .cap = @intCast(lane.cap), .eps = c.eps, .log2_theta = std.math.log2(x.d.config.theta) };
+ x.e.setBytes(std.mem.asBytes(&args), 5);
+ x.e.setBuffer(queries.buf, queries.off, 6);
+ x.e.setBuffer(lane.mem, lane.kv(layer, false).off, 7);
+ x.e.setBuffer(lane.mem, lane.kv(layer, true).off, 8);
+ x.e.dispatchGroups(mtl.Size.of(c.query_heads + c.kv_heads, rows, 1), mtl.Size.of(32, 1, 1));
+ }
+
+ /// Query rows [row, row + rows) against the lane's keys [0, keys), all of them visible to every row.
+ fn attend(x: Enc, layer: usize, lane: *const Lane, row: usize, rows: usize, keys: u64) void {
+ const s = &x.d.s;
+ const args = AttnArgs{ .len = @intCast(keys - 1), .cap = @intCast(lane.cap), .scale = 1.0 / @sqrt(@as(f32, c.head_dim)), .subs = st.subs };
+ const used = (keys + st.chunk_keys - 1) / st.chunk_keys * st.subs;
+ x.pipe("lfm2_attn_block");
+ x.e.setBuffer(s.queries, row * c.hidden * 2, 0);
+ x.e.setBuffer(lane.mem, lane.kv(layer, false).off, 1);
+ x.e.setBuffer(lane.mem, lane.kv(layer, true).off, 2);
+ x.e.setBytes(std.mem.asBytes(&args), 3);
+ x.e.setBuffer(s.pm, 0, 4);
+ x.e.setBuffer(s.pl, 0, 5);
+ x.e.setBuffer(s.po, 0, 6);
+ x.e.setValue(@as(u32, @intCast(rows)), 7);
+ x.e.dispatchGroups(mtl.Size.of((rows + 7) / 8, c.kv_heads, @intCast(used)), mtl.Size.of(128, 1, 1));
+ x.pipe("lfm2_attn_block_merge");
+ x.e.setBuffer(s.pm, 0, 0);
+ x.e.setBuffer(s.pl, 0, 1);
+ x.e.setBuffer(s.po, 0, 2);
+ x.e.setBytes(std.mem.asBytes(&args), 3);
+ x.e.setBuffer(s.y, row * c.hidden * 2, 4);
+ x.e.dispatchGroups(mtl.Size.of(c.query_heads, rows, 1), mtl.Size.of(32, 1, 1));
+ }
+};
+
+/// One stream's tapped rows in an absorb pass: `rows` rows at `src`, the context from position `p0`.
+pub const Piece = struct { lane: *Lane, src: Ref, rows: usize, p0: u64 };
+
+/// Each piece's rows into its lane's keys and values (fc, hidden_norm, k, v, k_norm, RoPE) in passes on `e`.
+pub fn absorbPass(d: *Drafter, e: mtl.ComputeEncoder, pieces: []const Piece) !void {
+ const s = &d.s;
+ const x = Enc{ .d = d, .e = e };
+ var rows: usize = 0;
+ for (pieces) |pc| {
+ try x.project(d.w.fc, pc.src, .{ .buf = s.fcx, .off = rows * c.hidden * 2 }, pc.rows);
+ rows += pc.rows;
+ }
+ if (rows == 0 or rows > ds.absorb_rows) return error.DraftContextGap;
+ x.norm(s.fcx, s.fcx, d.w.hidden_norm, false, s.fcx, rows); // in place: a thread reads its values first
+ for (0..d.config.layers) |i| {
+ const ly = &d.w.layers[i];
+ try x.project(ly.k, .{ .buf = s.fcx }, .{ .buf = s.k }, rows);
+ try x.project(ly.v, .{ .buf = s.fcx }, .{ .buf = s.v }, rows);
+ var at: usize = 0;
+ for (pieces) |pc| {
+ const q = Ref{ .buf = s.fcx, .off = at * c.hidden * 2 };
+ x.qk(q, .{ .buf = s.k, .off = at * kw * 2 }, .{ .buf = s.v, .off = at * kw * 2 }, i, pc.lane, pc.p0, .{ .buf = s.dump }, pc.rows);
+ at += pc.rows;
+ }
+ }
+}
+
+/// A command buffer for absorb passes, and its wait.
+pub fn begin(d: *Drafter) struct { cb: mtl.CommandBuffer, e: mtl.ComputeEncoder } {
+ const cb = d.target.queue.commandBuffer();
+ return .{ .cb = cb, .e = cb.compute(.serial) };
+}
+
+pub fn end(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder) !void {
+ e.end();
+ cb.commit();
+ cb.wait();
+ if (cb.failure()) |msg| {
+ std.log.err("lfm2 dspark: command buffer failed: {s}", .{msg});
+ return error.GpuFailed;
+ }
+}
+
+/// A GPU event a hold signals instead of being waited on, and the value it sets.
+pub const Signal = struct { event: mtl.Event, value: u64 };
+
+/// The blocks' drafts (backbone, then Markov chain) into the lanes' held chains; with `signal`, left for readPicks.
+pub fn hold(d: *Drafter, blocks: []const Block, signal: ?Signal) !?mtl.CommandBuffer {
+ const pool = mtl.objc.Pool.push();
+ defer pool.pop();
+ const s = &d.s;
+ var depth: u32 = 0;
+ for (blocks) |b| depth = @max(depth, b.depth);
+ // a block of the deepest chain's rows, at most the 9 the drafter was trained on
+ const n: u32 = depth;
+ const rows = blocks.len * n;
+ if (rows == 0 or rows > ds.pass_rows) return error.DraftTooManyStreams;
+ const ids = s.ids.slice(u32, ds.pass_rows);
+ for (0..blocks.len) |i| {
+ ids[i * n] = blocks[i].pending;
+ @memset(ids[i * n + 1 ..][0 .. n - 1], d.config.mask);
+ }
+ const steps = depth;
+ // past the last layer's attention a lone block carries only the rows its chain reads; several keep them all
+ const tail = if (blocks.len == 1) steps else rows;
+ const cb = d.target.queue.commandBuffer();
+ const x = Enc{ .d = d, .e = cb.compute(.serial) };
+ const e = x.e;
+ // the blocks' staged context rows first, in the same command buffer
+ var staged: [ds.pass_rows]Piece = undefined;
+ var count: usize = 0;
+ var held: usize = 0;
+ for (blocks) |b| if (ds.Drafter.stagedPiece(b.lane)) |pc| {
+ if (held + pc.rows > ds.absorb_rows) {
+ try absorbPass(d, e, staged[0..count]);
+ count = 0;
+ held = 0;
+ }
+ staged[count] = pc;
+ count += 1;
+ held += pc.rows;
+ };
+ if (count > 0) try absorbPass(d, e, staged[0..count]);
+ for (blocks) |b| {
+ b.lane.end += b.lane.staged;
+ b.lane.staged = 0;
+ }
+ const t = d.target;
+ const emb = t.weights.embedding;
+ e.setPipeline(t.kernels.embed(emb.bits));
+ e.setBuffer(s.ids, 0, 0);
+ e.setBuffer(emb.w, emb.w_off, 1);
+ e.setBuffer(emb.scales, emb.s_off, 2);
+ e.setBuffer(emb.biases, emb.b_off, 3);
+ e.setBuffer(s.h, 0, 4);
+ const units = c.hidden / @as(usize, if (emb.bits == 6) 4 else 8 / emb.bits);
+ e.dispatchThreads(mtl.Size.of(units, rows, 1), mtl.Size.of(256, 1, 1));
+ for (0..d.config.layers) |i| {
+ const ly = &d.w.layers[i];
+ const out = if (i + 1 == d.config.layers) tail else rows;
+ x.norm(s.h, s.r, ly.input_norm, i > 0, s.x, rows);
+ try x.project(ly.q, .{ .buf = s.x }, .{ .buf = s.q }, rows);
+ try x.project(ly.k, .{ .buf = s.x }, .{ .buf = s.k }, rows);
+ try x.project(ly.v, .{ .buf = s.x }, .{ .buf = s.v }, rows);
+ for (blocks, 0..) |b, bi| {
+ const at = bi * n;
+ x.qk(.{ .buf = s.q, .off = at * c.hidden * 2 }, .{ .buf = s.k, .off = at * kw * 2 }, .{ .buf = s.v, .off = at * kw * 2 }, i, b.lane, b.anchor, .{ .buf = s.queries, .off = at * c.hidden * 2 }, n);
+ x.attend(i, b.lane, at, if (blocks.len == 1) out else n, b.anchor + n);
+ }
+ try x.project(ly.o, .{ .buf = s.y }, .{ .buf = s.r }, out);
+ x.norm(s.h, s.r, ly.post_norm, true, s.x, out);
+ try x.project(ly.gate, .{ .buf = s.x }, .{ .buf = s.gate }, out);
+ try x.project(ly.up, .{ .buf = s.x }, .{ .buf = s.up }, out);
+ x.pipe("lfm2_swiglu");
+ e.setBuffer(s.gate, 0, 0);
+ e.setBuffer(s.up, 0, 1);
+ e.setBuffer(s.act, 0, 2);
+ e.setValue(@as(u32, @intCast(out * d.config.intermediate)), 3);
+ e.dispatchThreads(mtl.Size.of(out * d.config.intermediate / 8, 1, 1), mtl.Size.of(256, 1, 1));
+ try x.project(ly.down, .{ .buf = s.act }, .{ .buf = s.r }, out);
+ }
+ x.norm(s.h, s.r, d.w.norm, true, s.x, tail);
+ try x.project(.{ .codes = fwd.matrix(emb) }, .{ .buf = s.x }, .{ .buf = s.logits }, tail);
+ try markov(x, blocks.len, n, steps);
+ if (d.confident()) confidence(x, blocks.len, n, steps);
+ if (signal) |sg| { // committed unwaited, the command buffer returned retained
+ e.end();
+ cb.signalEvent(sg.event, sg.value);
+ const kept = cb.retain();
+ cb.commit();
+ return kept;
+ }
+ try end(cb, e);
+ const picks = s.picks.slice(Pick, d.config.block * ds.pass_rows);
+ for (blocks, 0..) |b, bi| {
+ for (b.lane.held[0..steps], 0..) |*h, i| {
+ const p = picks[i * ds.pass_rows + bi];
+ if (p.nonfinite != 0 or p.token >= c.vocab) return error.DraftNonfinite;
+ h.* = p.token;
+ }
+ b.lane.held_len = b.depth;
+ }
+ for (blocks, 0..) |b, bi| readConfidence(d, b.lane, b.depth, bi);
+ return null;
+}
+
+/// Lane `l`'s drafts' chances from slot `bi` of the last hold's logits: sigmoid(logit + bias); none: left as they are.
+fn readConfidence(d: *Drafter, l: *Lane, depth: u32, bi: usize) void {
+ if (!d.confident()) return;
+ const logits = d.s.conf.slice(f32, d.config.block * ds.pass_rows);
+ for (l.conf[0..depth], 0..) |*p, i| p.* = 1 / (1 + @exp(-(logits[i * ds.pass_rows + bi] + d.w.confidence_bias)));
+ l.conf_len = depth;
+}
+
+/// A finished hold's picks into the lanes' held chains; a nonfinite or invalid pick is refused, as when synchronous.
+pub fn readPicks(d: *Drafter, lanes: []const *Lane, depths: []const u32) !void {
+ const picks = d.s.picks.slice(Pick, d.config.block * ds.pass_rows);
+ var chains: [ds.pass_rows][]u32 = undefined;
+ for (lanes, chains[0..lanes.len]) |l, *ch| ch.* = &l.held;
+ try copyPicks(picks, chains[0..lanes.len], depths);
+ for (lanes, depths, 0..) |l, depth, bi| {
+ l.held_len = depth;
+ readConfidence(d, l, depth, bi);
+ }
+}
+
+/// Step i of block bi is picks[i * ds.pass_rows + bi]: each block's first `depth` picks into its chain.
+pub fn copyPicks(picks: []const Pick, chains: []const []u32, depths: []const u32) !void {
+ for (chains, depths, 0..) |chain, depth, bi| {
+ for (chain[0..depth], 0..) |*h, i| {
+ const p = picks[i * ds.pass_rows + bi];
+ if (p.nonfinite != 0 or p.token >= c.vocab) return error.DraftNonfinite;
+ h.* = p.token;
+ }
+ }
+}
+
+/// Each stream's draft i confidence logit (the final norm's row i, W1 of the token before draft i), for readConfidence.
+fn confidence(x: Enc, streams: usize, n: u32, depth: u32) void {
+ const d = x.d;
+ const e = x.e;
+ x.pipe("lfm2_dspark_confidence");
+ e.setBuffer(d.s.x, 0, 0);
+ e.setBuffer(d.w.w1, 0, 1);
+ e.setBuffer(d.s.ids, 0, 2);
+ e.setBuffer(d.s.picks, 0, 3);
+ e.setBuffer(d.w.confidence.?, 0, 4);
+ const args = ConfidenceArgs{ .n = n, .pass = ds.pass_rows, .vocab = c.vocab, .rank = d.config.rank, .hidden = c.hidden };
+ e.setBytes(std.mem.asBytes(&args), 5);
+ e.setBuffer(d.s.conf, 0, 6);
+ e.dispatchGroups(mtl.Size.of(depth, streams, 1), mtl.Size.of(256, 1, 1));
+}
+
+/// Step i of each stream: its row i's logits plus W2 W1[previous draft] (the pending token first), then the argmax.
+fn markov(x: Enc, streams: usize, n: u32, depth: u32) !void {
+ const d = x.d;
+ const s = &d.s;
+ const e = x.e;
+ for (0..depth) |i| {
+ // step i's previous token: the pending one (each block's first id), then step i - 1's pick
+ switch (d.w.w2) {
+ .codes => x.pipe("lfm2_markov_scan"),
+ .bf16 => x.pipe("lfm2_markov_scan_bf16"),
+ }
+ e.setBuffer(s.logits, i * c.vocab * 2, 0);
+ e.setBuffer(d.w.w1, 0, 1);
+ e.setBuffer(switch (d.w.w2) {
+ .codes => |m| m.w,
+ .bf16 => |m| m.w,
+ }, switch (d.w.w2) {
+ .codes => |m| m.w_off,
+ .bf16 => 0,
+ }, 2);
+ if (i == 0) e.setBuffer(s.ids, 0, 3) else e.setBuffer(s.picks, (i - 1) * ds.pass_rows * @sizeOf(Pick), 3);
+ const args = MarkovArgs{ .stride = n * c.vocab, .rank = d.config.rank, .vocab = c.vocab, .prev_stride = if (i == 0) n else @sizeOf(Pick) / 4 };
+ e.setBytes(std.mem.asBytes(&args), 4);
+ e.setBuffer(s.best, 0, 5);
+ e.setBuffer(s.at, 0, 6);
+ if (d.w.w2 == .codes) {
+ e.setBuffer(d.w.w2.codes.scales, d.w.w2.codes.s_off, 7);
+ e.setBuffer(d.w.w2.codes.biases, d.w.w2.codes.b_off, 8);
+ }
+ e.dispatchGroups(mtl.Size.of(scan_groups, streams, 1), mtl.Size.of(256, 1, 1));
+ x.pipe("lfm2_markov_pick");
+ e.setBuffer(s.best, 0, 0);
+ e.setBuffer(s.at, 0, 1);
+ e.setValue(@as(u32, scan_groups), 2);
+ e.setBuffer(s.picks, i * ds.pass_rows * @sizeOf(Pick), 3);
+ e.dispatchGroups(mtl.Size.of(streams, 1, 1), mtl.Size.of(256, 1, 1));
+ }
+}
diff --git a/zig/src/families/lfm2_moe/dspark_config.zig b/zig/src/families/lfm2_moe/dspark_config.zig
new file mode 100644
index 0000000000..79e518aa48
--- /dev/null
+++ b/zig/src/families/lfm2_moe/dspark_config.zig
@@ -0,0 +1,140 @@
+//! LiquidAI's DSpark drafter for LFM2.5-24B-A2B: its config admitted, any other shape refused with its reason.
+const std = @import("std");
+const c = @import("config.zig");
+
+/// The geometry the drafter's kernels take: the target's attention shape, five taps at most eight layers.
+pub const Config = struct {
+ taps: [8]u32 = @splat(0),
+ tap_count: u32,
+ layers: u32,
+ intermediate: u32,
+ /// Rows a block takes: the pending token and block - 1 mask tokens, one draft a row.
+ block: u32,
+ mask: u32,
+ rank: u32,
+ theta: f32,
+
+ pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Config {
+ const path = try std.fs.path.join(gpa, &.{ dir, "config.json" });
+ defer gpa.free(path);
+ const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20));
+ defer gpa.free(bytes);
+ return parse(gpa, bytes);
+ }
+
+ pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Config {
+ var parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{});
+ defer parsed.deinit();
+ const o = if (parsed.value == .object) parsed.value.object else return error.BadDraftConfig;
+ const archs = o.get("architectures") orelse return error.BadDraftConfig;
+ if (archs != .array or archs.array.items.len != 1 or archs.array.items[0] != .string) return error.BadDraftConfig;
+ if (!std.mem.eql(u8, archs.array.items[0].string, "Lfm2DSparkDraftModel")) return error.NotLfmDSpark;
+ const d = if (o.get("dflash_config")) |v| (if (v == .object) v.object else return error.BadDraftConfig) else return error.BadDraftConfig;
+ var out: Config = .{
+ .tap_count = 0,
+ .layers = try small(o, "num_hidden_layers"),
+ .intermediate = try small(o, "intermediate_size"),
+ .block = try small(o, "block_size"),
+ .mask = try small(d, "mask_token_id"),
+ .rank = try small(o, "markov_rank"),
+ .theta = @floatCast(try number(o, "rope_theta")),
+ };
+ // the target's widths: the drafter shares its embedding, head and attention kernels
+ if (try small(o, "hidden_size") != c.hidden or try small(o, "vocab_size") != c.vocab) return error.DraftForAnotherTarget;
+ if (try small(d, "num_target_layers") != c.layers) return error.DraftForAnotherTarget;
+ const heads = try small(o, "num_attention_heads");
+ const kv = try small(o, "num_key_value_heads");
+ if (heads != c.query_heads or kv != c.kv_heads or try small(o, "head_dim") != c.head_dim) return error.UnsupportedDraftShape;
+ if (@as(f32, @floatCast(try number(o, "rms_norm_eps"))) != c.eps) return error.UnsupportedDraftShape;
+ const taps = d.get("target_layer_ids") orelse return error.BadDraftConfig;
+ if (taps != .array or taps.array.items.len == 0 or taps.array.items.len > out.taps.len) return error.BadDraftConfig;
+ for (taps.array.items, 0..) |t, i| {
+ if (t != .integer or t.integer < 0) return error.BadDraftConfig;
+ if (t.integer >= c.layers) return error.DraftForAnotherTarget;
+ out.taps[i] = @intCast(t.integer);
+ }
+ out.tap_count = @intCast(taps.array.items.len);
+ try kinds(o, out.layers);
+ // the block's MLP is the fixed SiLU gate
+ const act = o.get("hidden_act") orelse return error.BadDraftConfig;
+ if (act != .string or !std.mem.eql(u8, act.string, "silu")) return error.UnsupportedDraftActivation;
+ // the vendored defaults decide the block: non-causal attention, interleaved RoPE, a plain Markov head
+ if (o.get("is_causal") != null) return error.UnsupportedDraftMask;
+ const neox = o.get("rope_is_neox_style") orelse return error.BadDraftConfig;
+ if (neox != .bool or neox.bool) return error.UnsupportedDraftRope;
+ const head = o.get("markov_head_type") orelse return error.BadDraftConfig;
+ if (head != .string or !std.mem.eql(u8, head.string, "vanilla")) return error.UnsupportedMarkovHead;
+ if (o.get("sample_from_anchor")) |v| if (v != .bool or !v.bool) return error.UnsupportedDraftAnchor;
+ if (out.layers == 0 or out.layers > 8 or out.block < 2 or out.block > 16) return error.UnsupportedDraftShape;
+ if (out.intermediate == 0 or out.intermediate % 64 != 0 or out.rank == 0 or out.rank % 64 != 0) return error.UnsupportedDraftShape;
+ if (out.mask >= c.vocab or !(out.theta > 0)) return error.BadDraftConfig;
+ return out;
+ }
+};
+
+/// Every layer full attention (the block's mask has no sliding window).
+fn kinds(o: std.json.ObjectMap, layers: u32) !void {
+ const v = o.get("layer_types") orelse return;
+ if (v != .array or v.array.items.len != layers) return error.BadDraftConfig;
+ for (v.array.items) |k| if (k != .string or !std.mem.eql(u8, k.string, "full_attention")) return error.UnsupportedDraftMask;
+}
+
+fn small(o: std.json.ObjectMap, key: []const u8) !u32 {
+ const v = o.get(key) orelse return error.BadDraftConfig;
+ if (v != .integer or v.integer < 0 or v.integer > std.math.maxInt(u32)) return error.BadDraftConfig;
+ return @intCast(v.integer);
+}
+
+fn number(o: std.json.ObjectMap, key: []const u8) !f64 {
+ const v = o.get(key) orelse return error.BadDraftConfig;
+ return switch (v) {
+ .integer => |i| @floatFromInt(i),
+ .float => |f| f,
+ else => error.BadDraftConfig,
+ };
+}
+
+/// LiquidAI/LFM2.5-24B-A2B-DSpark's config.json at e052e5b1dd10, as released.
+pub const stock =
+ \\{"architectures":["Lfm2DSparkDraftModel"],"model_type":"qwen3","dtype":"bfloat16","hidden_size":2048,
+ \\"num_hidden_layers":5,"num_attention_heads":32,"num_key_value_heads":8,"head_dim":64,"intermediate_size":6144,
+ \\"hidden_act":"silu","rms_norm_eps":1e-05,"vocab_size":128000,"rope_theta":10000000.0,
+ \\"max_position_embeddings":131072,"layer_types":["full_attention","full_attention","full_attention",
+ \\"full_attention","full_attention"],"block_size":9,"dflash_config":{"mask_token_id":125017,
+ \\"target_layer_ids":[2,10,18,26,34],"num_target_layers":40},"markov_rank":256,"rope_is_neox_style":false,
+ \\"enable_confidence_head":true,"markov_head_type":"vanilla"}
+;
+
+fn edited(from: []const u8, to: []const u8) !Config {
+ const a = std.testing.allocator;
+ const text = try std.mem.replaceOwned(u8, a, stock, from, to);
+ defer a.free(text);
+ return Config.parse(a, text);
+}
+
+test "lfm2 dspark: the released config is admitted with its taps, block, mask token and Markov rank" {
+ const cfg = try Config.parse(std.testing.allocator, stock);
+ try std.testing.expectEqualSlices(u32, &.{ 2, 10, 18, 26, 34 }, cfg.taps[0..cfg.tap_count]);
+ try std.testing.expectEqual(@as(u32, 5), cfg.layers);
+ try std.testing.expectEqual(@as(u32, 9), cfg.block);
+ try std.testing.expectEqual(@as(u32, 125017), cfg.mask);
+ try std.testing.expectEqual(@as(u32, 256), cfg.rank);
+ try std.testing.expectEqual(@as(u32, 6144), cfg.intermediate);
+ try std.testing.expectEqual(@as(f32, 1e7), cfg.theta);
+}
+
+test "lfm2 dspark: other drafters and shapes are refused with their reason" {
+ try std.testing.expectError(error.NotLfmDSpark, edited("Lfm2DSparkDraftModel", "DFlashDraftModel"));
+ try std.testing.expectError(error.DraftForAnotherTarget, edited("\"hidden_size\":2048", "\"hidden_size\":4096"));
+ try std.testing.expectError(error.DraftForAnotherTarget, edited("\"num_target_layers\":40", "\"num_target_layers\":24"));
+ try std.testing.expectError(error.DraftForAnotherTarget, edited("26,34]", "26,44]"));
+ try std.testing.expectError(error.UnsupportedDraftShape, edited("\"head_dim\":64", "\"head_dim\":128"));
+ try std.testing.expectError(error.UnsupportedDraftActivation, edited("\"silu\"", "\"gelu\""));
+ try std.testing.expectError(error.UnsupportedDraftRope, edited("\"rope_is_neox_style\":false", "\"rope_is_neox_style\":true"));
+ try std.testing.expectError(error.UnsupportedMarkovHead, edited("\"vanilla\"", "\"gated\""));
+ try std.testing.expectError(error.UnsupportedDraftMask, edited("\"full_attention\"]", "\"sliding_attention\"]"));
+ try std.testing.expectError(error.UnsupportedDraftMask, edited("\"block_size\"", "\"is_causal\":true,\"block_size\""));
+ try std.testing.expectError(error.UnsupportedDraftAnchor, edited("\"block_size\"", "\"sample_from_anchor\":false,\"block_size\""));
+ try std.testing.expectError(error.UnsupportedDraftShape, edited("\"block_size\":9", "\"block_size\":17"));
+ try std.testing.expectError(error.BadDraftConfig, edited("\"mask_token_id\":125017", "\"mask_token_id\":128000"));
+}
diff --git a/zig/src/families/lfm2_moe/dspark_quant.zig b/zig/src/families/lfm2_moe/dspark_quant.zig
new file mode 100644
index 0000000000..a6a5baccd6
--- /dev/null
+++ b/zig/src/families/lfm2_moe/dspark_quant.zig
@@ -0,0 +1,118 @@
+//! BF16 drafter rows quantized at load to MLX's affine codes in groups of 64, the layout core/tile_qmm reads.
+const std = @import("std");
+
+/// MLX's affine codes of 64 bf16 values at `bits`: bias the larger edge, scale |edge| / round(|edge| / step).
+pub fn group(x: []const u16, bits: u8, codes: []u8, scale: *u16, bias: *u16) void {
+ var lo: f32 = std.math.inf(f32);
+ var hi: f32 = -std.math.inf(f32);
+ for (x) |v| {
+ lo = @min(lo, toF32(v));
+ hi = @max(hi, toF32(v));
+ }
+ const bins: f32 = @floatFromInt((@as(u32, 1) << @intCast(bits)) - 1);
+ const step = @max(hi - lo, 1e-7) / bins;
+ const from_min = @abs(lo) > @abs(hi);
+ const edge = if (from_min) lo else hi;
+ const q0 = @round(@abs(edge) / step);
+ const s_f: f32 = if (q0 == 0) step else (if (from_min) @abs(edge) else -@abs(edge)) / q0;
+ scale.* = toBf16(s_f);
+ bias.* = toBf16(if (q0 == 0) lo else edge);
+ const s = toF32(scale.*);
+ const b = toF32(bias.*);
+ @memset(codes, 0);
+ for (x, 0..) |v, i| {
+ const q: u8 = @intFromFloat(std.math.clamp(@round((toF32(v) - b) / s), 0, bins));
+ if (bits == 8) codes[i] = q else codes[i / 2] |= q << @intCast((i % 2) * 4);
+ }
+}
+
+pub fn toF32(v: u16) f32 {
+ return @bitCast(@as(u32, v) << 16);
+}
+
+/// Round to nearest even, as MLX casts.
+pub fn toBf16(f: f32) u16 {
+ const b: u32 = @bitCast(f);
+ if ((b & 0x7f80_0000) == 0x7f80_0000) return @truncate(b >> 16);
+ return @truncate((b + 0x7fff + ((b >> 16) & 1)) >> 16);
+}
+
+/// A matrix's codes [n][k * bits / 8], scales and biases [n][k / 64], each written by one thread's rows.
+pub const Job = struct {
+ src: []const u16,
+ k: usize,
+ bits: u8,
+ codes: []u8,
+ scales: []u16,
+ biases: []u16,
+ fn run(j: Job, from: usize, to: usize) void {
+ const groups = j.k / 64;
+ const per = 64 * @as(usize, j.bits) / 8;
+ for (from..to) |r| for (0..groups) |g| {
+ const at = r * groups + g;
+ group(j.src[r * j.k + g * 64 ..][0..64], j.bits, j.codes[at * per ..][0..per], &j.scales[at], &j.biases[at]);
+ };
+ }
+
+ /// The job over its rows on every CPU (at most 16 threads).
+ pub fn all(j: Job, gpa: std.mem.Allocator, rows: usize) !void {
+ const threads = @min(std.Thread.getCpuCount() catch 1, 16);
+ const per = (rows + threads - 1) / threads;
+ const handles = try gpa.alloc(std.Thread, threads);
+ defer gpa.free(handles);
+ var spawned: usize = 0;
+ defer for (handles[0..spawned]) |h| h.join();
+ var at: usize = 0;
+ while (at < rows) : (at += per) {
+ handles[spawned] = try std.Thread.spawn(.{}, run, .{ j, at, @min(rows, at + per) });
+ spawned += 1;
+ }
+ }
+};
+
+/// Head `h`'s interleaved RoPE pairs (2i, 2i + 1) as half-split pairs (i, i + dim / 2): row h * dim + j's source.
+pub fn halfSplit(row: usize, dim: usize) usize {
+ const h = row / dim;
+ const j = row % dim;
+ const half = dim / 2;
+ return h * dim + if (j < half) 2 * j else 2 * (j - half) + 1;
+}
+
+test "lfm2 dspark: 4- and 8-bit codes keep the larger edge exact, as MLX quantizes" {
+ var x: [64]u16 = undefined;
+ for (&x, 0..) |*v, i| v.* = toBf16(@as(f32, @floatFromInt(i)) * 0.125 - 2.0);
+ for ([_]u8{ 4, 8 }) |bits| {
+ var codes: [64]u8 = undefined;
+ var s: u16 = undefined;
+ var b: u16 = undefined;
+ group(&x, bits, codes[0 .. 64 * bits / 8], &s, &b);
+ try std.testing.expect(toF32(s) < 0 and toF32(b) == 5.875); // the max is the edge: codes descend from it
+ const last: u8 = if (bits == 8) codes[63] else codes[31] >> 4;
+ try std.testing.expectEqual(@as(u8, 0), last);
+ for (x, 0..) |v, i| {
+ const q: u8 = if (bits == 8) codes[i] else (codes[i / 2] >> @intCast((i % 2) * 4)) & 0xf;
+ const back = toF32(s) * @as(f32, @floatFromInt(q)) + toF32(b);
+ try std.testing.expect(@abs(back - toF32(v)) <= @abs(toF32(s)) * 0.51);
+ }
+ }
+}
+
+test "lfm2 dspark: half-split rows rotated by the half-split rule equal interleaved pairs rotated, row for row" {
+ const dim = 64;
+ var x: [2 * dim]f32 = undefined;
+ for (&x, 0..) |*v, i| v.* = @sin(@as(f32, @floatFromInt(i)) * 0.37) * 3;
+ var moved: [2 * dim]f32 = undefined;
+ for (&moved, 0..) |*v, r| v.* = x[halfSplit(r, dim)];
+ const pos: f32 = 1234;
+ for (0..2) |h| for (0..dim / 2) |i| {
+ const angle = pos * std.math.pow(f32, 1e7, -@as(f32, @floatFromInt(i)) / (dim / 2));
+ const cs = @cos(angle);
+ const sn = @sin(angle);
+ const a = x[h * dim + 2 * i];
+ const b = x[h * dim + 2 * i + 1];
+ const l = moved[h * dim + i];
+ const r = moved[h * dim + i + dim / 2];
+ try std.testing.expectEqual(a * cs - b * sn, l * cs - r * sn);
+ try std.testing.expectEqual(a * sn + b * cs, l * sn + r * cs);
+ };
+}
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..043a590353
--- /dev/null
+++ b/zig/src/families/lfm2_moe/dspark_test.zig
@@ -0,0 +1,272 @@
+//! The DSpark drafter's kernels on the GPU against host references: taps, the Markov step, half-split RoPE.
+const std = @import("std");
+const mtl = @import("metal");
+const c = @import("config.zig");
+const kernels = @import("kernels.zig");
+const quant = @import("dspark_quant.zig");
+const block = @import("dspark_block.zig");
+const Pick = @import("tree_round_gpu").contract.Pick;
+const ds = @import("dspark.zig");
+const lanes = @import("lanes");
+
+const shared = mtl.ResourceOptions.shared;
+
+pub const Rig = struct {
+ device: mtl.Device,
+ queue: mtl.Queue,
+ k: kernels.Kernels,
+
+ pub 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) };
+ }
+
+ pub fn deinit(r: *Rig) void {
+ r.k.deinit();
+ r.queue.deinit();
+ r.device.deinit();
+ }
+
+ pub 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;
+ }
+};
+
+pub fn noise(seed: u64, out: []u16, scale: f32) void {
+ var prng = std.Random.DefaultPrng.init(seed);
+ for (out) |*v| v.* = quant.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);
+ }
+}
+
+test "lfm2 dspark: a Markov step picks the argmax of bf16(base + bf16(W2 . W1[prev])) over every threadgroup" {
+ var r = try Rig.init();
+ defer r.deinit();
+ const rank = 256;
+ const vocab = 1300; // three scan threadgroups, the last one short
+ const streams = 4;
+ const groups = (vocab + 511) / 512;
+ const w1 = try r.buffer(u16, vocab * rank);
+ defer w1.deinit();
+ noise(11, w1.slice(u16, vocab * rank), 1);
+ const codes = try r.buffer(u8, vocab * rank);
+ defer codes.deinit();
+ for (codes.slice(u8, vocab * rank), 0..) |*q, i| q.* = @truncate(i *% 2654435761 >> 7);
+ const sc = try r.buffer(u16, vocab * rank / 64);
+ defer sc.deinit();
+ const bi = try r.buffer(u16, vocab * rank / 64);
+ defer bi.deinit();
+ noise(12, sc.slice(u16, vocab * rank / 64), 0.01);
+ noise(13, bi.slice(u16, vocab * rank / 64), 0.5);
+ // a token every second word (the picks' layout): inside W1, its last row, past it (a zero row), and a NaN row
+ const ids = try r.buffer(u32, 2 * streams);
+ defer ids.deinit();
+ @memcpy(ids.slice(u32, 2 * streams), &[_]u32{ 17, 9, vocab - 1, 9, vocab + 5, 0, 3, 3 });
+ const logits = try r.buffer(u16, streams * 2 * vocab);
+ defer logits.deinit();
+ noise(14, logits.slice(u16, streams * 2 * vocab), 4);
+ logits.slice(u16, streams * 2 * vocab)[3 * 2 * vocab + 700] = 0x7fc0;
+ const best = try r.buffer(f32, streams * groups);
+ defer best.deinit();
+ const at = try r.buffer(u32, streams * groups);
+ defer at.deinit();
+ const picks = try r.buffer(u32, 2 * streams);
+ defer picks.deinit();
+ const cb = r.queue.commandBuffer();
+ const e = cb.compute(.serial);
+ e.setPipeline(r.k.get("lfm2_markov_scan"));
+ for ([_]mtl.Buffer{ logits, w1, codes, ids }, 0..) |b, i| e.setBuffer(b, 0, i);
+ e.setValue([4]u32{ 2 * vocab, rank, vocab, 2 }, 4);
+ for ([_]mtl.Buffer{ best, at, sc, bi }, 5..) |b, i| e.setBuffer(b, 0, i);
+ e.dispatchGroups(mtl.Size.of(groups, streams, 1), mtl.Size.of(256, 1, 1));
+ e.setPipeline(r.k.get("lfm2_markov_pick"));
+ e.setBuffer(best, 0, 0);
+ e.setBuffer(at, 0, 1);
+ e.setValue(@as(u32, groups), 2);
+ e.setBuffer(picks, 0, 3);
+ e.dispatchGroups(mtl.Size.of(streams, 1, 1), mtl.Size.of(256, 1, 1));
+ e.end();
+ cb.commit();
+ cb.wait();
+ const got = picks.slice(u32, 2 * streams);
+ const base = logits.slice(u16, streams * 2 * vocab);
+ const id_list = ids.slice(u32, 2 * streams);
+ for (0..streams) |s| {
+ if (s == 3) { // the planted NaN: flagged
+ try std.testing.expectEqual(@as(u32, 1), got[2 * s + 1]);
+ continue;
+ }
+ try std.testing.expectEqual(@as(u32, 0), got[2 * s + 1]);
+ const id = id_list[2 * s];
+ var top: f64 = -std.math.inf(f64);
+ var mine: f64 = 0;
+ for (0..vocab) |v| {
+ var dot: f64 = 0;
+ for (0..rank) |i| {
+ const g = (v * rank + i) / 64;
+ const w = @as(f64, quant.toF32(sc.slice(u16, vocab * rank / 64)[g])) * @as(f64, @floatFromInt(codes.slice(u8, vocab * rank)[v * rank + i])) + quant.toF32(bi.slice(u16, vocab * rank / 64)[g]);
+ dot += w * (if (id < vocab) quant.toF32(w1.slice(u16, vocab * rank)[id * rank + i]) else 0);
+ }
+ const score: f64 = quant.toF32(quant.toBf16(quant.toF32(base[s * 2 * vocab + v]) + quant.toF32(quant.toBf16(@floatCast(dot)))));
+ top = @max(top, score);
+ if (v == got[2 * s]) mine = score;
+ }
+ // the pick is the best score, or within two bf16 steps of it where fp32 sums round apart
+ try std.testing.expect(got[2 * s] < vocab and mine >= top - @abs(top) * 0x1p-7);
+ }
+}
+
+test "lfm2 dspark: lfm2_qk on half-split rows and gains gives interleaved RoPE's keys, reordered" {
+ var r = try Rig.init();
+ defer r.deinit();
+ const hd = c.head_dim;
+ const pos: u32 = 777;
+ const theta: f32 = 1e7;
+ var k_il: [c.kv_heads * hd]u16 = undefined;
+ var gain_il: [hd]u16 = undefined;
+ noise(21, &k_il, 2);
+ noise(22, &gain_il, 1.5);
+ const q = try r.buffer(u16, c.query_heads * hd);
+ defer q.deinit();
+ const k = try r.buffer(u16, c.kv_heads * hd);
+ defer k.deinit();
+ const v = try r.buffer(u16, c.kv_heads * hd);
+ defer v.deinit();
+ const qn = try r.buffer(u16, hd);
+ defer qn.deinit();
+ const kn = try r.buffer(u16, hd);
+ defer kn.deinit();
+ for (k.slice(u16, c.kv_heads * hd), 0..) |*x, i| x.* = k_il[quant.halfSplit(i, hd)];
+ for (kn.slice(u16, hd), 0..) |*x, i| x.* = gain_il[quant.halfSplit(i, hd)];
+ @memset(qn.slice(u16, hd), quant.toBf16(1));
+ const cap = 1024;
+ const queries = try r.buffer(u16, c.query_heads * hd);
+ defer queries.deinit();
+ const keys = try r.buffer(u16, c.kv_heads * cap * hd);
+ defer keys.deinit();
+ const values = try r.buffer(u16, c.kv_heads * cap * hd);
+ defer values.deinit();
+ const Args = extern struct { len: u32, cap: u32, eps: f32, log2_theta: f32 };
+ const cb = r.queue.commandBuffer();
+ const e = cb.compute(.serial);
+ e.setPipeline(r.k.get("lfm2_qk"));
+ for ([_]mtl.Buffer{ q, k, v, qn, kn }, 0..) |b, i| e.setBuffer(b, 0, i);
+ const args = Args{ .len = pos, .cap = cap, .eps = c.eps, .log2_theta = std.math.log2(theta) };
+ e.setBytes(std.mem.asBytes(&args), 5);
+ for ([_]mtl.Buffer{ queries, keys, values }, 6..) |b, i| e.setBuffer(b, 0, i);
+ e.dispatchGroups(mtl.Size.of(c.query_heads + c.kv_heads, 1, 1), mtl.Size.of(32, 1, 1));
+ e.end();
+ cb.commit();
+ cb.wait();
+ const got = keys.slice(u16, c.kv_heads * cap * hd);
+ for (0..c.kv_heads) |h| {
+ var ss: f32 = 0;
+ for (0..hd) |i| ss += quant.toF32(k_il[h * hd + i]) * quant.toF32(k_il[h * hd + i]);
+ const inv = 1 / @sqrt(ss / hd + c.eps);
+ for (0..hd / 2) |i| {
+ const angle = @as(f32, pos) * std.math.pow(f32, theta, -@as(f32, @floatFromInt(i)) / (hd / 2));
+ const a = quant.toF32(k_il[h * hd + 2 * i]) * inv * quant.toF32(gain_il[2 * i]);
+ const b = quant.toF32(k_il[h * hd + 2 * i + 1]) * inv * quant.toF32(gain_il[2 * i + 1]);
+ const want = [2]f32{ a * @cos(angle) - b * @sin(angle), a * @sin(angle) + b * @cos(angle) };
+ const at = (h * cap + pos) * hd;
+ const have = [2]f32{ quant.toF32(got[at + i]), quant.toF32(got[at + i + hd / 2]) };
+ for (want, have) |w, x| try std.testing.expect(@abs(w - x) <= 0.02 * @max(1, @abs(w)));
+ }
+ }
+}
+test "lfm2 dspark: a settled hold moves the timed passes at its block counts a tenth of the way to its GPU ms" {
+ var d: ds.Drafter = undefined;
+ const grid = [_]lanes.hold_price.HoldCost{ .{ .streams = 1, .depth = 1, .ms = 1 }, .{ .streams = 4, .depth = 1, .ms = 2 }, .{ .streams = 4, .depth = 4, .ms = 4 }, .{ .streams = 8, .depth = 1, .ms = 3 } };
+ @memcpy(d.holds[0..grid.len], &grid);
+ d.holds_len = grid.len;
+ d.pass_len = 4;
+ d.depths[0..4].* = .{ 1, 1, 1, 1 };
+ d.priceHold(4); // priced 2 ms: the 4-block passes rise by a tenth of the doubling
+ try std.testing.expectApproxEqAbs(@as(f64, 2.2), d.holds[1].ms, 1e-12);
+ try std.testing.expectApproxEqAbs(@as(f64, 4.4), d.holds[2].ms, 1e-12);
+ try std.testing.expectEqual(@as(f64, 1), d.holds[0].ms);
+ try std.testing.expectEqual(@as(f64, 3), d.holds[3].ms);
+ d.pass_len = 6; // between 4 and 8 blocks: both move
+ d.depths[0..6].* = .{ 1, 1, 1, 1, 1, 1 };
+ d.priceHold(0);
+ try std.testing.expectEqual(@as(f64, 2.2), d.holds[1].ms); // no GPU time: nothing learned
+ d.priceHold(2.6 * 2);
+ try std.testing.expect(d.holds[1].ms > 2.2 and d.holds[3].ms > 3 and d.holds[0].ms == 1);
+}
+
+/// lfm2_draft_ids on one pick (token, nonfinite flag) at a stride of 4 words: the id it wrote and the bad flag.
+fn draftIds(r: *Rig, pick: Pick) ![2]u32 {
+ const picks = try r.buffer(u32, 4);
+ defer picks.deinit();
+ const ids = try r.buffer(u32, 1);
+ defer ids.deinit();
+ const bad = try r.buffer(u32, 1);
+ defer bad.deinit();
+ picks.slice(u32, 4)[0] = pick.token;
+ picks.slice(u32, 4)[1] = pick.nonfinite;
+ try r.run("lfm2_draft_ids", &.{ picks, ids, bad }, [4]u32{ 4, 1, c.vocab, 0 }, 1, .{ 1, 1 });
+ return .{ ids.slice(u32, 1)[0], bad.slice(u32, 1)[0] };
+}
+
+test "lfm2 dspark: held ids carry a flagged or invalid pick as token 0 and the bad flag, a good pick untouched" {
+ var r = try Rig.init();
+ defer r.deinit();
+ try std.testing.expectEqual([2]u32{ 5, 0 }, try draftIds(&r, .{ .token = 5 }));
+ try std.testing.expectEqual([2]u32{ c.vocab - 1, 0 }, try draftIds(&r, .{ .token = c.vocab - 1 }));
+ for ([_]u32{ 1, 2, 4, 7 }) |kind| try std.testing.expectEqual([2]u32{ 0, 1 }, try draftIds(&r, .{ .token = 9, .nonfinite = kind }));
+ try std.testing.expectEqual([2]u32{ 0, 1 }, try draftIds(&r, .{ .token = std.math.maxInt(u32) }));
+ try std.testing.expectEqual([2]u32{ 0, 1 }, try draftIds(&r, .{ .token = c.vocab }));
+}
+
+test "lfm2 dspark: a settled hold refuses a NaN, infinity or invalid pick as the synchronous path does" {
+ var picks: [4 * ds.pass_rows]Pick = @splat(.{ .token = 3 });
+ var a: [4]u32 = undefined;
+ var b: [4]u32 = undefined;
+ const chains = [_][]u32{ &a, &b };
+ try block.copyPicks(&picks, &chains, &.{ 2, 2 });
+ try std.testing.expectEqualSlices(u32, &.{ 3, 3 }, a[0..2]);
+ const bad = [_]Pick{ .{ .token = 3, .nonfinite = 1 }, .{ .token = 3, .nonfinite = 2 }, .{ .token = 3, .nonfinite = 4 }, .{ .token = std.math.maxInt(u32) }, .{ .token = c.vocab } };
+ for (bad) |p| {
+ picks[ds.pass_rows + 1] = p; // the second block's second step
+ try std.testing.expectError(error.DraftNonfinite, block.copyPicks(&picks, &chains, &.{ 2, 2 }));
+ picks[ds.pass_rows + 1] = .{ .token = 3 };
+ }
+ picks[3 * ds.pass_rows] = .{ .nonfinite = 1 }; // a step past the block's depth is not read
+ try block.copyPicks(&picks, &chains, &.{ 2, 2 });
+}
diff --git a/zig/src/families/lfm2_moe/dspark_weights.zig b/zig/src/families/lfm2_moe/dspark_weights.zig
new file mode 100644
index 0000000000..9d87dd6eed
--- /dev/null
+++ b/zig/src/families/lfm2_moe/dspark_weights.zig
@@ -0,0 +1,178 @@
+//! The DSpark drafter's weights: the released bf16 linears as stored, or on request quantized to 8-bit affine codes.
+const std = @import("std");
+const mtl = @import("metal");
+const c = @import("config.zig");
+const dc = @import("dspark_config.zig");
+const quant = @import("dspark_quant.zig");
+const ckpt = @import("../../core/checkpoint_metal.zig");
+const shards = @import("../../core/checkpoint.zig");
+const tile = @import("../../core/tile_qmm.zig");
+
+const shared = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked;
+
+/// An [n, k] bf16 matrix as released (rows only reordered for half-split RoPE where the checkpoint is read).
+pub const Bf16 = struct { w: mtl.Buffer, n: usize, k: usize };
+
+/// A linear: bf16 as stored (bits 0), or MLX affine codes at the loader's width.
+pub const Linear = union(enum) { bf16: Bf16, codes: tile.Matrix };
+
+/// A bf16 vector in its own buffer (norm gains; RoPE-paired ones in the half-split order the kernels rotate).
+pub const Vector = mtl.Buffer;
+
+pub const Layer = struct {
+ q: Linear,
+ k: Linear,
+ v: Linear,
+ o: Linear,
+ gate: Linear,
+ up: Linear,
+ down: Linear,
+ input_norm: Vector,
+ post_norm: Vector,
+ q_norm: Vector,
+ k_norm: Vector,
+};
+
+pub const Weights = struct {
+ fc: Linear,
+ hidden_norm: Vector,
+ norm: Vector,
+ layers: [8]Layer,
+ /// The Markov head: W1 [vocab, rank] bf16 as stored (rows gathered by token), W2 [vocab, rank] as the linears.
+ w1: mtl.Buffer,
+ w2: Linear,
+ /// The confidence head: its [hidden + rank] bf16 weights and fp32 bias (null: the checkpoint has none).
+ confidence: ?mtl.Buffer = null,
+ confidence_bias: f32 = 0,
+ owned: std.ArrayList(mtl.Buffer) = .empty,
+
+ pub fn deinit(w: *Weights, gpa: std.mem.Allocator) void {
+ for (w.owned.items) |b| b.deinit();
+ w.owned.deinit(gpa);
+ }
+
+ /// Every buffer the weights hold, for a residency set.
+ pub fn buffers(w: *const Weights) []const mtl.Buffer {
+ return w.owned.items;
+ }
+};
+
+const Loader = struct {
+ gpa: std.mem.Allocator,
+ device: mtl.Device,
+ check: *const ckpt.Checkpoint,
+ stage: *ckpt.Stage, // a matrix read from the file, one at a time, into one reused buffer
+ owned: *std.ArrayList(mtl.Buffer),
+ bits: u8,
+
+ fn buffer(l: Loader, bytes: usize) !mtl.Buffer {
+ const b = try l.device.buffer(bytes, shared);
+ errdefer b.deinit();
+ try l.owned.append(l.gpa, b);
+ return b;
+ }
+
+ /// A vector as loaded, or a matrix read into the stage (valid until the next matrix), of `shape` in bf16.
+ fn tensor(l: Loader, name: []const u8, shape: []const usize) !ckpt.Tensor {
+ l.stage.reset();
+ if (shape.len == 2) try l.stage.reserve(l.check, &.{name});
+ const t = if (shape.len == 2) try l.stage.read(l.check, name) else try l.check.get(name);
+ if (t.dtype != .bf16 or t.rank != shape.len or !std.mem.eql(usize, t.shape[0..t.rank], shape)) {
+ std.log.err("{s}: unexpected DSpark tensor dtype or shape", .{name});
+ return error.BadDraftWeight;
+ }
+ return t;
+ }
+
+ /// A bf16 vector copied out, its pairs reordered for half-split RoPE when `paired`.
+ fn vector(l: Loader, name: []const u8, len: usize, paired: bool) !Vector {
+ const src = (try l.tensor(name, &.{len})).host(u16);
+ const b = try l.buffer(len * 2);
+ const dst = b.slice(u16, len);
+ for (dst, 0..) |*v, i| v.* = src[if (paired) quant.halfSplit(i, len) else i];
+ return b;
+ }
+
+ /// An [n, k] bf16 matrix as it is or as affine codes; rows reordered for half-split RoPE when `paired`.
+ fn matrix(l: Loader, name: []const u8, n: usize, k: usize, paired: bool) !Linear {
+ const t = try l.tensor(name, &.{ n, k });
+ var rows = t.host(u16);
+ var moved: ?[]u16 = null;
+ defer if (moved) |m| l.gpa.free(m);
+ if (paired) {
+ const m = try l.gpa.alloc(u16, n * k);
+ moved = m;
+ for (0..n) |r| @memcpy(m[r * k ..][0..k], rows[quant.halfSplit(r, c.head_dim) * k ..][0..k]);
+ rows = m;
+ }
+ if (l.bits == 0) {
+ const b = try l.buffer(n * k * 2);
+ @memcpy(b.slice(u16, n * k), rows);
+ return .{ .bf16 = .{ .w = b, .n = n, .k = k } };
+ }
+ const codes = try l.buffer(n * k * l.bits / 8);
+ const scales = try l.buffer(n * k / 64 * 2);
+ const biases = try l.buffer(n * k / 64 * 2);
+ const job: quant.Job = .{ .src = rows, .k = k, .bits = l.bits, .codes = codes.slice(u8, n * k * l.bits / 8), .scales = scales.slice(u16, n * k / 64), .biases = biases.slice(u16, n * k / 64) };
+ try job.all(l.gpa, n);
+ return .{ .codes = .{ .w = codes, .scales = scales, .biases = biases, .n = n, .k = k, .bits = l.bits } };
+ }
+};
+
+fn part(buf: []u8, layer: usize, name: []const u8) ![]const u8 {
+ return std.fmt.bufPrint(buf, "layers.{d}.{s}.weight", .{ layer, name });
+}
+
+/// Matrices stay in the file until the loader reads them; vectors load with the header.
+fn matrices(_: []const u8, rank: usize) bool {
+ return rank == 2;
+}
+
+/// The checkpoint read as released (`bits` 0) or quantized to 8-bit codes, a matrix at a time; buffers freed after.
+pub fn load(gpa: std.mem.Allocator, device: mtl.Device, io: std.Io, dir: []const u8, cfg: *const dc.Config, bits: u8) !Weights {
+ if (bits != 0 and bits != 8) return error.UnsupportedDraftBits;
+ var check = ckpt.Checkpoint.init(gpa);
+ defer check.deinit();
+ var stage: ckpt.Stage = .{ .device = device };
+ defer stage.deinit();
+ const files = try shards.shardFiles(gpa, io, dir);
+ defer shards.freeShardFiles(gpa, files);
+ for (files) |path| try check.addFileSplit(device, path, "", null, matrices);
+ var w: Weights = undefined;
+ w.owned = .empty;
+ errdefer w.deinit(gpa);
+ const l: Loader = .{ .gpa = gpa, .device = device, .check = &check, .stage = &stage, .owned = &w.owned, .bits = bits };
+ const qw = c.query_heads * c.head_dim;
+ const kw = c.kv_heads * c.head_dim;
+ w.fc = try l.matrix("fc.weight", c.hidden, cfg.tap_count * c.hidden, false);
+ w.hidden_norm = try l.vector("hidden_norm.weight", c.hidden, false);
+ w.norm = try l.vector("norm.weight", c.hidden, false);
+ var name: [128]u8 = undefined;
+ for (w.layers[0..cfg.layers], 0..) |*ly, i| {
+ ly.q = try l.matrix(try part(&name, i, "self_attn.q_proj"), qw, c.hidden, true);
+ ly.k = try l.matrix(try part(&name, i, "self_attn.k_proj"), kw, c.hidden, true);
+ ly.v = try l.matrix(try part(&name, i, "self_attn.v_proj"), kw, c.hidden, false);
+ ly.o = try l.matrix(try part(&name, i, "self_attn.o_proj"), c.hidden, qw, false);
+ ly.gate = try l.matrix(try part(&name, i, "mlp.gate_proj"), cfg.intermediate, c.hidden, false);
+ ly.up = try l.matrix(try part(&name, i, "mlp.up_proj"), cfg.intermediate, c.hidden, false);
+ ly.down = try l.matrix(try part(&name, i, "mlp.down_proj"), c.hidden, cfg.intermediate, false);
+ ly.input_norm = try l.vector(try part(&name, i, "input_layernorm"), c.hidden, false);
+ ly.post_norm = try l.vector(try part(&name, i, "post_attention_layernorm"), c.hidden, false);
+ ly.q_norm = try l.vector(try part(&name, i, "self_attn.q_norm"), c.head_dim, true);
+ ly.k_norm = try l.vector(try part(&name, i, "self_attn.k_norm"), c.head_dim, true);
+ }
+ const w1 = (try l.tensor("markov_head.markov_w1.weight", &.{ c.vocab, cfg.rank })).host(u16);
+ w.w1 = try l.buffer(c.vocab * cfg.rank * 2);
+ @memcpy(w.w1.slice(u16, c.vocab * cfg.rank), w1);
+ w.w2 = try l.matrix("markov_head.markov_w2.weight", c.vocab, cfg.rank, false);
+ w.confidence = null;
+ w.confidence_bias = 0;
+ if (check.has("confidence_head.proj.weight")) {
+ const width = c.hidden + cfg.rank; // the final norm's row, then W1 of the token before the draft
+ const proj = (try l.tensor("confidence_head.proj.weight", &.{ 1, width })).host(u16);
+ w.confidence = try l.buffer(width * 2);
+ @memcpy(w.confidence.?.slice(u16, width), proj);
+ w.confidence_bias = quant.toF32((try l.tensor("confidence_head.proj.bias", &.{1})).host(u16)[0]);
+ }
+ return w;
+}
diff --git a/zig/src/families/lfm2_moe/forward.zig b/zig/src/families/lfm2_moe/forward.zig
index 0bbfc3ef02..51408e2fd4 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));
@@ -114,7 +118,17 @@ pub fn matrix(w: row.Weights) tile.Matrix {
return .{ .w = w.w, .w_off = w.w_off, .scales = w.scales, .s_off = w.s_off, .biases = w.biases, .b_off = w.b_off, .n = w.n, .k = w.k, .bits = w.bits };
}
+/// A window's held drafts on the device: `count` tokens at `src` + `off`, `stride` words apart, into ids from `row`.
+pub const Held = struct { row: usize, src: Buffer, off: usize, stride: u32, count: u32 };
+/// A GPU event the forward waits on before it reads held drafts, and the value it waits for.
+pub const Wait = struct { event: mtl.Event, value: u64 };
+
pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const u32, record: bool, head: Head) !void {
+ return runHeld(m, s, segments, tokens, record, head, &.{}, null);
+}
+
+/// `run` with held drafts copied into the rows' ids on the GPU first (after `wait`): their host ids are ignored.
+pub fn runHeld(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const u32, record: bool, head: Head, held: []const Held, wait: ?Wait) !void {
if (tokens.len == 0 or tokens.len > s.rows or (record and tokens.len > st.batch_rows)) return error.LfmRowsExceeded;
var total: usize = 0;
for (segments) |seg| {
@@ -127,9 +141,20 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const
const pool = mtl.objc.Pool.push();
defer pool.pop();
const cb = m.queue.commandBuffer();
+ if (wait) |w| cb.waitForEvent(w.event, w.value);
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 };
+ if (held.len > 0) @memset(s.bad.contents()[0..4], 0);
+ for (held) |h| {
+ enc.pipe("lfm2_draft_ids");
+ e.setBuffer(h.src, h.off, 0);
+ e.setValue([4]u32{ h.stride, h.count, c.vocab, 0 }, 1);
+ e.setBuffer(s.ids, h.row * 4, 2);
+ e.setBuffer(s.bad, 0, 3);
+ enc.threads(.{ h.count, 1, 1 }, .{ 32, 1, 1 });
+ }
const embedding = m.weights.embedding;
e.setPipeline(m.kernels.embed(embedding.bits));
e.setBuffer(s.ids, 0, 0);
@@ -141,6 +166,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 +183,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();
@@ -168,6 +200,7 @@ pub fn run(m: *Model, s: *st.Scratch, segments: []const Segment, tokens: []const
std.debug.print("LFM2 forward: {s}\n", .{msg});
return error.GpuFailed;
}
+ if (held.len > 0 and std.mem.readInt(u32, s.bad.contents()[0..4], .little) != 0) return error.DraftNonfinite;
var base: usize = 0;
for (segments) |seg| {
const logit_row = if (head == .all) base + seg.rows - 1 else 0;
@@ -177,9 +210,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 +256,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 +313,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..c83c9a0b4a 100644
--- a/zig/src/families/lfm2_moe/kernels.zig
+++ b/zig/src/families/lfm2_moe/kernels.zig
@@ -16,12 +16,27 @@ const names = [_][:0]const u8{
"lfm2_attn_mma",
"lfm2_attn_mma_r2",
"lfm2_attn_merge",
+ "lfm2_attn_block",
+ "lfm2_attn_block_merge",
"lfm2_route_logits",
"lfm2_route_select",
"lfm2_swiglu",
"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",
+ "lfm2_markov_scan",
+ "lfm2_markov_scan_bf16",
+ "lfm2_bf16_rows",
+ "lfm2_bf16_mma",
+ "lfm2_markov_pick",
+ "lfm2_draft_ids",
+ "lfm2_dspark_confidence",
};
pub const Kernels = struct {
diff --git a/zig/src/families/lfm2_moe/lfm2_moe.zig b/zig/src/families/lfm2_moe/lfm2_moe.zig
index 97f9cbd5d8..8f861e31b4 100644
--- a/zig/src/families/lfm2_moe/lfm2_moe.zig
+++ b/zig/src/families/lfm2_moe/lfm2_moe.zig
@@ -7,3 +7,7 @@ pub const state = @import("state.zig");
pub const forward = @import("forward.zig");
pub const backend = @import("backend.zig");
pub const snapshot = @import("snapshot.zig");
+pub const dspark = @import("dspark.zig");
+pub const dspark_config = @import("dspark_config.zig");
+pub const dspark_block = @import("dspark_block.zig");
+pub const dspark_quant = @import("dspark_quant.zig");
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..883f670fbe 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 };
@@ -110,6 +119,7 @@ pub const Scratch = struct {
memory: Storage,
rows: usize,
ids: mtl.Buffer,
+ bad: mtl.Buffer, // set by lfm2_draft_ids when a held pick is flagged nonfinite or past the vocabulary
h: mtl.Buffer,
x: mtl.Buffer,
r: mtl.Buffer,
@@ -141,7 +151,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,16 +167,17 @@ 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);
+ out.bad = try mem.alloc(4);
inline for (.{ "h", "x", "r", "y", "q" }) |name| @field(out, name) = try mem.alloc(rows * c.hidden * 2);
out.queries = try mem.alloc((rows + attn_rows) * c.hidden * 2); // a prompt attention tile reads attn_rows from its first
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 +198,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 +213,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 +228,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/native/cuda_options.zig b/zig/src/native/cuda_options.zig
index d6a6eefe28..0c07152bdb 100644
--- a/zig/src/native/cuda_options.zig
+++ b/zig/src/native/cuda_options.zig
@@ -36,6 +36,10 @@ pub fn set(comptime F: type, o: api.Open, opts: *F.Options, problem: *[]const u8
return false;
};
}
+ if (o.drafter_bits_given and o.drafter_bits == 8) {
+ problem.* = "--drafter-bits 8: only LFM2.5-24B-A2B's Metal DSpark drafter takes it; this family takes 0 (bf16) or 4";
+ return false;
+ }
if (comptime @hasField(F.Options, "drafter")) opts.drafter = o.drafter;
if (comptime @hasField(F.Options, "drafter_bits")) opts.drafter_bits = o.drafter_bits;
if (@hasField(F.Options, "kv8")) opts.kv8 = o.kv8;
diff --git a/zig/src/native/lfm2_moe.zig b/zig/src/native/lfm2_moe.zig
index 5340d807de..d5114d2a0c 100644
--- a/zig/src/native/lfm2_moe.zig
+++ b/zig/src/native/lfm2_moe.zig
@@ -1,4 +1,4 @@
-//! LFM2-MoE served by the existing lane host and scheduler.
+//! LFM2-MoE served by the existing lane host and scheduler, with LiquidAI's DSpark drafter from --drafter.
const std = @import("std");
const mtl = @import("metal");
const api = @import("engine_api");
@@ -8,11 +8,16 @@ const lanes = tf.lanes;
const cache_fit = @import("cache_fit.zig");
const metal_admit = @import("metal_admit.zig");
+/// --drafter loads LiquidAI's DSpark drafter (LiquidAI/LFM2.5-24B-A2B-DSpark) in front of the engine.
+pub const takes_drafter = true;
+
const Host = struct {
resident: ?mtl.resident.Resident = null,
gpa: std.mem.Allocator,
model: *q.Model,
metal: *q.backend.Metal,
+ draft: ?*q.dspark.Drafter = null,
+ wrap: ?lanes.drafted.Drafted = null,
warm: mtl.keepalive.Target = undefined, // the model's queue, for the lane host's idle ticker
config: lanes.Config,
clock: lanes.backend.WallClock,
@@ -25,7 +30,9 @@ const Host = struct {
h.host.stop();
if (h.cache) |*store| store.deinit();
h.core.deinit();
+ if (h.wrap) |*w| w.deinit();
mtl.resident.streams.drop(h);
+ if (h.draft) |d| d.deinit();
h.config.deinit(h.gpa);
h.metal.deinit();
h.model.deinit();
@@ -35,6 +42,13 @@ const Host = struct {
const cannot = "the native LFM2 engine cannot load this checkpoint ({s})";
+/// A refused open's backend, model and host, freed before it reports why.
+fn release(h: *Host) void {
+ h.metal.deinit();
+ h.model.deinit();
+ h.gpa.destroy(h);
+}
+
pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened {
const pool = mtl.objc.Pool.push();
defer pool.pop();
@@ -65,18 +79,44 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope
const chunk = 2048;
h.metal = try q.backend.Metal.init(gpa, model, .{ .capacity = window, .chunk = chunk, .streams = @max(o.lanes, 2) });
errdefer h.metal.deinit();
- const per_stream = q.state.cacheBytes(&model.config.attention, window);
+ h.draft = null;
+ h.wrap = null;
+ if (if (o.drafts) o.drafter else null) |dir| {
+ const bits: u8 = if (o.drafter_bits_given) o.drafter_bits else 0;
+ if (bits != 0 and bits != 8) {
+ problem.* = "--drafter-bits: the LFM2 DSpark drafter takes 0 (bf16 as released, the default) or 8 (8-bit codes quantized at load)";
+ release(h);
+ return null;
+ }
+ h.draft = q.dspark.Drafter.init(gpa, io, model, dir, .{ .capacity = window, .bits = bits }) catch |err| {
+ problem.* = try std.fmt.allocPrint(a, "the native LFM2 drafter cannot load {s} ({s}): it takes LiquidAI/LFM2.5-24B-A2B-DSpark", .{ dir, @errorName(err) });
+ release(h);
+ return null;
+ };
+ }
+ errdefer if (h.draft) |d| d.deinit();
+ // a stream's caches, and its drafter lane when one drafts
+ const per_stream = q.state.cacheBytes(&model.config.attention, window) + if (h.draft) |d| d.laneBytes(window) else 0;
const served = try metal_admit.streams(a, "lfm2", per_stream, metal_admit.room(model.device), o, window, problem) orelse {
- h.metal.deinit();
- model.deinit();
- gpa.destroy(h);
+ if (h.draft) |d| d.deinit();
+ release(h);
return null;
};
h.metal.options.streams = @max(served, 2);
- h.config = try lanes.Config.init(gpa, h.metal.facts(), q.state.window_rows, q.state.window_rows - 1);
+ if (h.draft) |d| {
+ h.wrap = try lanes.drafted.Drafted.init(gpa, h.metal.backend(), d.drafter());
+ // the verify copies the drafts on the GPU after the unwaited hold
+ h.metal.held_source = d.heldSource();
+ h.wrap.?.held_on_target = true;
+ }
+ errdefer if (h.wrap) |*w| w.deinit();
+ var facts = if (h.wrap) |*w| w.facts(h.metal.facts()) else h.metal.facts();
+ // a drafting stream prices its copy windows by its own copies: on distinct prompts they keep breaking
+ if (h.wrap != null) facts.own_copy_rates = true;
+ h.config = try lanes.Config.init(gpa, facts, q.state.window_rows, q.state.window_rows - 1);
errdefer h.config.deinit(gpa);
h.clock = .{ .io = io };
- h.core = lanes.Engine.init(gpa, &h.config, h.metal.backend(), h.clock.clock());
+ h.core = lanes.Engine.init(gpa, &h.config, if (h.wrap) |*w| w.backend() else h.metal.backend(), h.clock.clock());
errdefer h.core.deinit();
h.host = api.LaneHost.init(gpa, io, &h.core, .{ .lanes = served, .context_window = @intCast(window), .prefill_step = chunk, .structures = true });
h.cache = null;
@@ -97,8 +137,8 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope
if (plan.budget_bytes > 0) {
const S = q.backend.Snaps;
h.cache = api.prompt_cache.Store.init(gpa, .{
- .ptr = h.metal,
- .vtable = &.{ .bytes = S.bytes, .save = S.save, .restore = S.restore, .drop = S.drop },
+ .ptr = h,
+ .vtable = &.{ .bytes = Snaps.bytes, .save = Snaps.save, .restore = S.restore, .drop = Snaps.drop },
}, .{ .chunk = chunk }, plan.budget_bytes);
h.host.cache = &h.cache.?;
h.host.info_.prompt_cache = true;
@@ -106,7 +146,40 @@ pub fn open(a: std.mem.Allocator, gpa: std.mem.Allocator, io: std.Io, o: api.Ope
errdefer if (h.cache) |*store| store.deinit();
try mtl.resident.streams.hold(h, io, "lfm2 moe");
errdefer mtl.resident.streams.drop(h);
+ if (h.draft) |d| if (h.resident) |*r| {
+ const held = try d.buffers(gpa);
+ defer gpa.free(held);
+ r.add(held);
+ };
h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick };
try h.host.start();
+ if (h.draft) |d| {
+ const mode = if (d.options.bits == 8) "8-bit codes quantized at load (--drafter-bits 8), not the released bf16" else "bf16 matrices as released";
+ std.log.info("lfm2: DSpark drafts from --drafter, {d} a round at most, {s}", .{ d.config.block, mode });
+ }
return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h };
}
+
+/// The prompt cache's copies of a stream: the target's state at a mark (backend.zig), and the drafter's context there.
+const Snaps = struct {
+ const S = q.backend.Snaps;
+ fn host(ptr: *anyopaque) *Host {
+ return @ptrCast(@alignCast(ptr));
+ }
+ fn bytes(ptr: *anyopaque, at: u32) u64 {
+ const h = host(ptr);
+ return S.bytes(h.metal, at) + if (h.draft) |d| d.keptBytes(at) else 0;
+ }
+ fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque {
+ const h = host(ptr);
+ const saved = try S.save(h.metal, owner, at);
+ if (h.draft) |d| d.keepLane(@ptrCast(@alignCast(owner.?)), @ptrCast(@alignCast(saved))) catch |err| {
+ S.drop(h.metal, saved);
+ return err;
+ };
+ return saved;
+ }
+ fn drop(ptr: *anyopaque, saved: *anyopaque) void {
+ S.drop(host(ptr).metal, saved);
+ }
+};
diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig
index 4466b42513..41dc28577d 100644
--- a/zig/src/native/metal.zig
+++ b/zig/src/native/metal.zig
@@ -173,8 +173,8 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c
problem.* = why;
return null;
}
- if (o.drafter != null and !dense and !isGemma4(o.model_type)) {
- problem.* = "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head) and Gemma 4 dense";
+ if (o.drafter != null and !dense and !isGemma4(o.model_type) and !std.mem.eql(u8, o.model_type, "lfm2_moe")) {
+ problem.* = "--drafter is supported for Qwen3.8-27B-class checkpoints (qwen3_5 with its own output head), Gemma 4 dense and LFM2.5-24B-A2B";
return null;
}
if (dense) return openQwen27(a, gpa, io, o, problem);
@@ -336,7 +336,6 @@ fn refusal(a: Allocator, o: api.Open, dense: bool) !?[]const u8 {
if (o.recall and std.mem.eql(u8, t, "qwen4_exp")) return recall_refusal;
if (spill.refusal(o)) |why| return why;
if (o.learn == null) return null;
- // the 27B writes conversation states
if (o.learn_conversations and std.mem.eql(u8, t, "qwen3_5") and dense) return null;
const name: []const u8 = if (std.mem.eql(u8, t, "glm_moe_dsa")) "GLM-5.3 (EXLR8)" else if (std.mem.eql(u8, t, "lfm2_moe")) "LFM2.5-24B-A2B" else if (isGemma4(t)) "Gemma 4" else if (std.mem.eql(u8, t, "step5_vl")) "Step-5" else if (std.mem.eql(u8, t, "spark2_5")) "Spark-X2.6" else if (!std.mem.eql(u8, t, "qwen3_5")) return null else if (dense) "Qwen3.8-27B" else "Qwen3.5-2B";
if (o.learn_conversations) return try std.fmt.allocPrint(a, "{s} keeps conversation states on disk for Qwen3.8-27B, Nemotron, Flash Next and GLM-5.3-Flash: serve {s} without it", .{ spill.flag(o), name });
diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig
index d36a933aa7..5a38d9fd3a 100644
--- a/zig/src/server/cli.zig
+++ b/zig/src/server/cli.zig
@@ -100,7 +100,7 @@ pub const flags = [_]Flag{
.{ .name = "--slide", .kind = .store_true, .native = true },
.{ .name = "--slide-graph", .native = true },
.{ .name = "--drafter", .native = true },
- .{ .name = "--drafter-bits", .choices = &.{ "0", "4" }, .native = true },
+ .{ .name = "--drafter-bits", .choices = &.{ "0", "4", "8" }, .native = true },
.{ .name = "--record", .native = true },
.{ .name = "--kv-cache", .choices = &.{ "bf16", "fp8" }, .native = true },
.{ .name = "--prompt-precision", .choices = &.{ "bf16", "int8" }, .native = true },
@@ -210,7 +210,7 @@ pub const Args = struct {
compact_keep: ?u32 = null,
compact_memory: ?[]const u8 = null,
drafter: ?[]const u8 = null,
- drafter_bits: u8 = 4,
+ drafter_bits: ?u8 = null, // null: not given (a family without its own default takes 4)
record: ?[]const u8 = null, // drafter training rows written under this folder
kv_cache: []const u8 = "bf16", // fp8: e4m3 attention caches with a scale a row, about twice the context
prompt_precision: []const u8 = "bf16", // int8: prompt matmuls quantize their inputs to int8 a 64-input group
@@ -540,7 +540,8 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args
.lanes_fixed = parallelFixed(args.parallel),
.drafts = !args.no_drafts,
.drafter = args.drafter,
- .drafter_bits = args.drafter_bits,
+ .drafter_bits = args.drafter_bits orelse 4,
+ .drafter_bits_given = args.drafter_bits != null,
.record = args.record,
.kv8 = std.mem.eql(u8, args.kv_cache, "fp8"),
.prompt_int8 = std.mem.eql(u8, args.prompt_precision, "int8"),
diff --git a/zig/src/server/cli_test.zig b/zig/src/server/cli_test.zig
index f85890a632..f35a44e631 100644
--- a/zig/src/server/cli_test.zig
+++ b/zig/src/server/cli_test.zig
@@ -163,8 +163,10 @@ test "parse and capabilities share the table" {
try std.testing.expectEqual(@as(usize, 2), args.api_key.len);
try std.testing.expect(!args.thinking);
try std.testing.expectEqualStrings("x", (try parse(a, &.{ "m", "--drafter", "x" }, &u)).drafter.?);
- try std.testing.expectEqual(@as(u8, 0), (try parse(a, &.{ "m", "--drafter-bits", "0" }, &u)).drafter_bits);
- try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--drafter-bits", "8" }, &u));
+ try std.testing.expectEqual(@as(?u8, 0), (try parse(a, &.{ "m", "--drafter-bits", "0" }, &u)).drafter_bits);
+ try std.testing.expectEqual(@as(?u8, 8), (try parse(a, &.{ "m", "--drafter-bits", "8" }, &u)).drafter_bits);
+ try std.testing.expectEqual(@as(?u8, null), (try parse(a, &.{ "m", "--drafter", "x" }, &u)).drafter_bits);
+ try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--drafter-bits", "3" }, &u));
try std.testing.expectEqualStrings("/r", (try parse(a, &.{ "m", "--record", "/r" }, &u)).record.?);
try std.testing.expectEqualStrings("fp8", (try parse(a, &.{ "m", "--kv-cache", "fp8" }, &u)).kv_cache);
try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--kv-cache", "int4" }, &u));
diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig
index 5dab0c4568..12fbfd93f2 100644
--- a/zig/src/server/main.zig
+++ b/zig/src/server/main.zig
@@ -25,7 +25,7 @@ const usage_line = "usage: tensorfold serve [-h] [--host HOST] [--port PORT] [--
" [--thinking | --no-thinking] [--reasoning-effort {low,medium,high,xhigh}]" ++
" [--thinking-budget THINKING_BUDGET] [--thinking-reserve THINKING_RESERVE] [--loop-guard] [--no-drafts]" ++
" [--drafter DRAFTER]" ++
- " [--drafter-bits {0,4}] [--record RECORD] [--keep-warm SECONDS] [--idle-unload SECONDS]" ++
+ " [--drafter-bits {0,4,8}] [--record RECORD] [--keep-warm SECONDS] [--idle-unload SECONDS]" ++
" [--silence-floor-dbfs DBFS] [--silence-seconds SECONDS]" ++
" [--silence-onset-s SECONDS] [--silence-onset-voiced-s SECONDS]" ++
" [--silence-mid-s SECONDS] [--silence-mid-min-voiced-s SECONDS] [--silence-redraws N]" ++
diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig
index 28e78092f2..fc49c99405 100644
--- a/zig/src/tensorfold.zig
+++ b/zig/src/tensorfold.zig
@@ -84,6 +84,10 @@ test {
_ = @import("families/lfm2_moe/config.zig");
_ = @import("families/lfm2_moe/backend.zig");
_ = @import("families/lfm2_moe/state.zig");
+ _ = @import("families/lfm2_moe/dspark_config.zig");
+ _ = @import("families/lfm2_moe/dspark_quant.zig");
+ _ = @import("families/lfm2_moe/dspark_test.zig");
+ _ = @import("families/lfm2_moe/dspark_bf16_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_attn.zig b/zig/tests/contract/metal_lfm2_attn.zig
new file mode 100644
index 0000000000..60fa6939ae
--- /dev/null
+++ b/zig/tests/contract/metal_lfm2_attn.zig
@@ -0,0 +1,172 @@
+//! LFM2's block attention (lfm2_attn_block, then lfm2_attn_block_merge): every row of a drafter's block sees all keys.
+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 qh = 32;
+const kvh = 8;
+const hd = 64;
+const subs = 8; // the drafter's sub-spans a 4,096-key chunk (state.subs)
+const chunks = 32;
+const chunk_keys = 4096;
+pub const name = "metal lfm2 block attention";
+
+/// rows: a block's rows; k: the keys every row sees (a context and the block past it).
+const shapes = [_]c.Shape{ .{ .rows = 1, .n = qh * hd, .k = 9 }, .{ .rows = 3, .n = qh * hd, .k = 700 }, .{ .rows = 9, .n = qh * hd, .k = 4500 } };
+
+fn depth(s: c.Shape) usize {
+ return s.k;
+}
+
+pub fn spec() c.Spec {
+ return .{ .name = name, .accum = .none, .depth = depth, .out = .bf16, .served = &shapes, .nan = .{ .stated = "a NaN key of a kv head makes its query heads' outputs NaN, the others stay finite" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "the check bounds each output by its bf16 probabilities' rounding, not a sum's bits" } };
+}
+
+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,
+ mutation: controls.Mutation,
+ pipes: [2]mtl.Pipeline,
+ shape: c.Shape = undefined,
+ cap: usize = 0,
+ nan: bool = false,
+ q: []u16 = &.{},
+ k: []u16 = &.{},
+ v: []u16 = &.{},
+ args: [5]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);
+ const first = try rig_.pipeline(device, text, "lfm2_attn_block");
+ errdefer first.deinit();
+ return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipes = .{ first, try rig_.pipeline(device, text, "lfm2_attn_block_merge") }, .mutation = mutation };
+ }
+
+ fn drop(p: *Probe) void {
+ inline for (.{ "q", "k", "v" }) |f| {
+ p.a.free(@field(p, f));
+ @field(p, f) = &.{};
+ }
+ p.rig.reset();
+ }
+
+ pub fn deinit(p: *Probe) void {
+ p.drop();
+ p.rig.deinit();
+ for (p.pipes) |x| x.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 {
+ if (s.n != qh * hd or s.rows == 0 or s.rows > 16 or s.k < s.rows or s.k > chunks * chunk_keys) return false;
+ p.drop();
+ p.shape = s;
+ p.nan = f == .nonfinite;
+ p.cap = (s.k + 63) / 64 * 64 + 32; // keys past `k` hold finite values, as a lane's cache does
+ const padded = (s.rows + 7) / 8 * 8 + 8; // the kernel reads a tile's 8 rows whatever its row count
+ p.q = try p.a.alloc(u16, padded * qh * hd);
+ p.k = try p.a.alloc(u16, kvh * p.cap * hd);
+ p.v = try p.a.alloc(u16, kvh * p.cap * hd);
+ for (p.q, 0..) |*x, i| x.* = value(i, seed, 1);
+ for (p.k, 0..) |*x, i| x.* = value(i, seed +% 1, if (f == .stress) 4 else 1);
+ for (p.v, 0..) |*x, i| x.* = value(i, seed +% 2, 2);
+ if (p.nan) p.k[(0 * p.cap + s.k / 2) * hd + 5] = 0x7fc0; // kv head 0, a visible key
+ p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.q));
+ p.args[1] = try p.rig.input(std.mem.sliceAsBytes(p.k));
+ p.args[2] = try p.rig.input(std.mem.sliceAsBytes(p.v));
+ const Args = extern struct { len: u32, cap: u32, scale: f32, subs: u32 };
+ const args = Args{ .len = @intCast(s.k - 1), .cap = @intCast(p.cap), .scale = 1.0 / @sqrt(@as(f32, hd)), .subs = subs };
+ p.args[3] = try p.rig.input(std.mem.asBytes(&args));
+ const n: u32 = @intCast(s.rows);
+ p.args[4] = try p.rig.input(std.mem.asBytes(&n));
+ p.y = try p.rig.output(p.count() * 2);
+ return true;
+ }
+
+ pub fn count(p: *Probe) usize {
+ return p.shape.rows * qh * hd;
+ }
+
+ pub fn run(p: *Probe) !void {
+ const rows = p.shape.rows;
+ const partials = @max(rows, 8) * subs * qh * chunks;
+ const pm = try p.rig.scratch(partials * 4);
+ const pl = try p.rig.scratch(partials * 4);
+ const po = try p.rig.scratch(partials * hd * 4);
+ const pass = p.rig.begin();
+ const e = pass.e;
+ e.setPipeline(p.pipes[0]);
+ for (p.args[0..4], 0..) |r, i| e.setBuffer(r.buf, r.off, i);
+ for ([_]rig_.Ref{ pm, pl, po, p.args[4] }, 4..) |r, i| e.setBuffer(r.buf, r.off, i);
+ const used = (p.shape.k + chunk_keys - 1) / chunk_keys * subs;
+ e.dispatchGroups(mtl.Size.of((rows + 7) / 8, kvh, used), mtl.Size.of(128, 1, 1));
+ e.setPipeline(p.pipes[1]);
+ for ([_]rig_.Ref{ pm, pl, po, p.args[3], p.y }, 0..) |r, i| e.setBuffer(r.buf, r.off, i);
+ e.dispatchGroups(mtl.Size.of(qh, rows, 1), mtl.Size.of(32, 1, 1));
+ try p.rig.finish(pass);
+ }
+
+ pub fn reference(_: *Probe, want: []f64, mag: []f64) !void {
+ @memset(want, 0); // the check computes its own references
+ @memset(mag, 0);
+ }
+
+ /// Every output against float64 attention: within 2^-7 of its values' largest magnitude and 2^-8 of itself.
+ pub fn check(p: *Probe, _: g.Gate) !?[]const u8 {
+ const s = p.shape;
+ const out = std.mem.bytesAsSlice(u16, p.rig.bytes(p.y, p.count() * 2));
+ const w = try p.a.alloc(f64, s.k);
+ defer p.a.free(w);
+ for (0..s.rows) |r| for (0..qh) |h| {
+ const kv = h / (qh / kvh);
+ const nan_head = p.nan and kv == 0;
+ var top: f64 = -std.math.inf(f64);
+ for (0..s.k) |j| {
+ var dot: f64 = 0;
+ for (0..hd) |d| dot += c.fromBf16(p.q[(r * qh + h) * hd + d]) * c.fromBf16(p.k[(kv * p.cap + j) * hd + d]);
+ w[j] = dot / @sqrt(@as(f64, hd));
+ top = @max(top, w[j]);
+ }
+ var sum: f64 = 0;
+ for (w) |*x| {
+ x.* = @exp(x.* - top);
+ sum += x.*;
+ }
+ for (0..hd) |d| {
+ var o: f64 = 0;
+ var big: f64 = 0;
+ for (0..s.k) |j| {
+ const vv: f64 = c.fromBf16(p.v[(kv * p.cap + j) * hd + d]);
+ o += w[j] / sum * vv;
+ big = @max(big, @abs(vv));
+ }
+ const got: f64 = c.fromBf16(out[(r * qh + h) * hd + d]);
+ if (nan_head) {
+ if (!std.math.isNan(got)) return "a head behind a 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;
+ }
+};
diff --git a/zig/tests/contract/metal_lfm2_bf16.zig b/zig/tests/contract/metal_lfm2_bf16.zig
new file mode 100644
index 0000000000..0a191a48fa
--- /dev/null
+++ b/zig/tests/contract/metal_lfm2_bf16.zig
@@ -0,0 +1,133 @@
+//! LFM2's DSpark bf16 linears under gates 1-4: x W^T with bf16 inputs and fp32 sums, few rows or 64x64 tiles.
+const std = @import("std");
+const mtl = @import("metal");
+const contract = @import("contract");
+const rig_ = @import("rig.zig");
+const controls = @import("metal_lfm2_controls.zig");
+
+const c = contract.spec;
+const g = contract.gates;
+const Dense = contract.dense.Dense;
+
+pub const Which = enum { rows, tile };
+const names = [_][]const u8{ "metal lfm2 bf16 rows", "metal lfm2 bf16 tile" };
+
+/// The drafter's matrices: q and o (2,048 from 2,048), k and v (512), gate and up (6,144), down (2,048 from 6,144).
+const rows_served = [_]c.Shape{
+ .{ .rows = 1, .n = 2048, .k = 2048 }, .{ .rows = 9, .n = 512, .k = 2048 }, .{ .rows = 16, .n = 6144, .k = 2048 },
+ .{ .rows = 32, .n = 2048, .k = 6144 }, .{ .rows = 4, .n = 2048, .k = 10240 },
+};
+/// A cold prompt's absorb: the taps' projection (2,048 from 10,240) and k and v, up to 2,048 rows a pass.
+const tile_served = [_]c.Shape{
+ .{ .rows = 33, .n = 2048, .k = 2048 }, .{ .rows = 128, .n = 512, .k = 2048 },
+ .{ .rows = 2048, .n = 2048, .k = 10240 }, .{ .rows = 70, .n = 6144, .k = 2048 },
+};
+
+/// A lane's k / 32 fp32 fmas in turn, then the simdgroup's five-level sum.
+fn rowsDepth(s: c.Shape) usize {
+ return s.k / 32 + 8;
+}
+
+/// One 8-wide MMA step a k step of 8, accumulating in fp32.
+fn tileDepth(s: c.Shape) usize {
+ return s.k / 8 + 8;
+}
+
+pub fn spec(w: Which) c.Spec {
+ return switch (w) {
+ // rows in groups of 16, a column a simdgroup (any n), k in 256-wide lane steps of 8
+ .rows => .{ .name = names[0], .accum = .f32, .depth = rowsDepth, .out = .bf16, .tiles = .{ .{ .size = 16 }, .{ .size = 8 }, .{ .size = 256, .step = 8, .min = 8 } }, .served = &rows_served },
+ .tile => .{ .name = names[1], .accum = .f32, .depth = tileDepth, .out = .bf16, .tiles = .{ .{ .size = 64 }, .{ .size = 64 }, .{ .size = 32, .step = 32, .min = 32 } }, .served = &tile_served },
+ };
+}
+
+const Args = extern struct { rows: u32, n: u32, k: u32 };
+
+pub const Probe = struct {
+ a: std.mem.Allocator,
+ rig: rig_.Rig,
+ which: Which,
+ mutation: controls.Mutation,
+ pipe: mtl.Pipeline,
+ d: ?Dense = null,
+ x: rig_.Ref = undefined,
+ w: rig_.Ref = undefined,
+ out: rig_.Ref = undefined,
+
+ pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which, mutation: controls.Mutation) !Probe {
+ const text = try controls.source(a, mutation);
+ defer if (mutation != .none) a.free(text);
+ const key: []const u8 = if (w == .rows) "lfm2_bf16_rows" else "lfm2_bf16_mma";
+ return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = w, .mutation = mutation, .pipe = try rig_.pipeline(device, text, key) };
+ }
+
+ pub fn deinit(p: *Probe) void {
+ if (p.d) |d| d.deinit();
+ p.rig.deinit();
+ p.pipe.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 .{ .spec = s, .ctx = p, .vt = &vt, .control = p.mutation.gate() };
+ }
+
+ pub fn prepare(p: *Probe, shape: c.Shape, fill: g.Fill, seed: u64) !bool {
+ if (shape.rows == 0 or shape.n == 0 or shape.k == 0) return false;
+ if (p.which == .rows and (shape.k % 8 != 0 or shape.rows > 32)) return false;
+ if (p.which == .tile and shape.k % 32 != 0) return false;
+ if (p.d) |d| d.deinit();
+ p.d = null;
+ p.rig.reset();
+ const d = try Dense.init(p.a, shape, 1);
+ p.d = d;
+ d.fillX(fill, seed, 2);
+ // W as bf16 bits in [-1, 1), or gate 4's mirrored integers whose halves cancel (partials under 2^24)
+ const bits = try p.a.alloc(u16, shape.n * shape.k);
+ defer p.a.free(bits);
+ var prng = std.Random.DefaultPrng.init(seed ^ 0x6d);
+ const rng = prng.random();
+ const k = shape.k;
+ for (bits) |*b| b.* = c.toBf16(if (fill == .stress) 0 else rng.float(f32) * 2 - 1);
+ if (fill == .stress) for (0..shape.n) |j| for (0..k / 2) |i| {
+ const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 7));
+ bits[j * k + i] = c.toBf16(v);
+ bits[j * k + k / 2 + i] = c.toBf16(-v);
+ };
+ for (d.w, bits) |*wv, b| wv.* = c.fromBf16(b);
+ p.x = try p.rig.input(std.mem.sliceAsBytes(d.x));
+ p.w = try p.rig.input(std.mem.sliceAsBytes(bits));
+ p.out = try p.rig.output(d.outputs() * c.size(spec(p.which).out));
+ return true;
+ }
+
+ pub fn run(p: *Probe) !void {
+ const d = p.d.?;
+ const pass = p.rig.begin();
+ const e = pass.e;
+ e.setPipeline(p.pipe);
+ e.setBuffer(p.x.buf, p.x.off, 0);
+ e.setBuffer(p.w.buf, p.w.off, 1);
+ const args = Args{ .rows = @intCast(d.shape.rows), .n = @intCast(d.shape.n), .k = @intCast(d.shape.k) };
+ e.setBytes(std.mem.asBytes(&args), 2);
+ e.setBuffer(p.out.buf, p.out.off, 3);
+ if (p.which == .rows) {
+ e.dispatchGroups(mtl.Size.of((d.shape.n + 7) / 8, (d.shape.rows + 15) / 16, 1), mtl.Size.of(256, 1, 1));
+ } else e.dispatchGroups(mtl.Size.of((d.shape.n + 63) / 64, (d.shape.rows + 63) / 64, 1), mtl.Size.of(128, 1, 1));
+ try p.rig.finish(pass);
+ }
+
+ pub fn count(p: *Probe) usize {
+ return p.d.?.outputs();
+ }
+
+ pub fn reference(p: *Probe, want: []f64, mag: []f64) !void {
+ return p.d.?.reference(want, mag);
+ }
+
+ pub fn simulate(p: *Probe, acc: c.Accum, order: g.Order, out: []f64) !void {
+ return p.d.?.simulate(acc, order, out);
+ }
+};
diff --git a/zig/tests/contract/metal_lfm2_confidence.zig b/zig/tests/contract/metal_lfm2_confidence.zig
new file mode 100644
index 0000000000..ebf74ea5f2
--- /dev/null
+++ b/zig/tests/contract/metal_lfm2_confidence.zig
@@ -0,0 +1,110 @@
+//! LFM2's DSpark confidence logit (lfm2_dspark_confidence): a dot of a row and a W1 row with the head's 2,304 weights.
+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 rows = contract.rows;
+const rig_ = @import("rig.zig");
+const controls = @import("metal_lfm2_controls.zig");
+
+const hidden = 2048;
+const rank = 256;
+pub const name = "metal lfm2 dspark confidence";
+
+/// rows: one stream's drafts (row i reads the block's row i and the W1 row of the token before draft i).
+const shapes = [_]c.Shape{ .{ .rows = 1, .n = 1, .k = hidden + rank }, .{ .rows = 4, .n = 1, .k = hidden + rank }, .{ .rows = 9, .n = 1, .k = hidden + rank } };
+
+/// A thread's chain over every 256th input, a 32-way simd_sum tree, then the eight simdgroups' partials in order.
+fn depth(s: c.Shape) usize {
+ return (s.k + 255) / 256 + 5 + 8;
+}
+
+pub fn spec() c.Spec {
+ return .{ .name = name, .accum = .f32, .depth = depth, .out = .f32, .served = &shapes };
+}
+
+pub const Probe = struct {
+ a: std.mem.Allocator,
+ rig: rig_.Rig,
+ pipe: mtl.Pipeline,
+ mutation: controls.Mutation,
+ r: ?rows.Rows = null,
+ args: [6]rig_.Ref = undefined, // x, W1 rows, ids, picks, the head's weights, the shape
+ 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_dspark_confidence"), .mutation = mutation };
+ }
+
+ pub fn deinit(p: *Probe) void {
+ if (p.r) |r| r.deinit();
+ 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() };
+ }
+
+ /// Row i's first 2,048 inputs as the block's row i, its last 256 as W1 row i: the token before draft i is i.
+ pub fn prepare(p: *Probe, s: c.Shape, f: g.Fill, seed: u64) !bool {
+ if (s.k != hidden + rank or s.rows == 0 or s.rows > 16) return false;
+ if (p.r) |r| r.deinit();
+ p.r = null;
+ p.rig.reset();
+ const r = try rows.Rows.init(p.a, .dot, s);
+ p.r = r;
+ r.fill(f, seed);
+ const x = try p.a.alloc(u16, s.rows * hidden);
+ defer p.a.free(x);
+ const w1 = try p.a.alloc(u16, s.rows * rank);
+ defer p.a.free(w1);
+ for (0..s.rows) |i| {
+ @memcpy(x[i * hidden ..][0..hidden], r.x[i * s.k ..][0..hidden]);
+ @memcpy(w1[i * rank ..][0..rank], r.x[i * s.k + hidden ..][0..rank]);
+ }
+ // the block's pending id is token 0; the pick before draft i (i > 0) is token i, a flag word after each
+ var picks: [32]u32 = @splat(0xdead);
+ for (1..s.rows) |i| picks[2 * (i - 1)] = @intCast(i);
+ const ids = [1]u32{0};
+ const args = [5]u32{ @intCast(s.rows), 1, @intCast(s.rows), rank, hidden };
+ p.args[0] = try p.rig.input(std.mem.sliceAsBytes(x));
+ p.args[1] = try p.rig.input(std.mem.sliceAsBytes(w1));
+ p.args[2] = try p.rig.input(std.mem.asBytes(&ids));
+ p.args[3] = try p.rig.input(std.mem.asBytes(&picks));
+ p.args[4] = try p.rig.input(std.mem.sliceAsBytes(r.w));
+ p.args[5] = try p.rig.input(std.mem.asBytes(&args));
+ p.out = try p.rig.output(s.rows * 4);
+ return true;
+ }
+
+ pub fn run(p: *Probe) !void {
+ const n = p.r.?.shape.rows;
+ const pass = p.rig.begin();
+ const e = pass.e;
+ e.setPipeline(p.pipe);
+ for (p.args, 0..) |ref, i| e.setBuffer(ref.buf, ref.off, i);
+ e.setBuffer(p.out.buf, p.out.off, 6);
+ e.dispatchGroups(mtl.Size.of(n, 1, 1), mtl.Size.of(256, 1, 1));
+ try p.rig.finish(pass);
+ }
+
+ pub fn count(p: *Probe) usize {
+ return p.r.?.outputs();
+ }
+
+ pub fn reference(p: *Probe, want: []f64, mag: []f64) !void {
+ return p.r.?.reference(want, mag);
+ }
+
+ pub fn simulate(p: *Probe, acc: c.Accum, order: g.Order, out: []f64) !void {
+ return p.r.?.simulate(acc, order, out);
+ }
+};
diff --git a/zig/tests/contract/metal_lfm2_controls.zig b/zig/tests/contract/metal_lfm2_controls.zig
new file mode 100644
index 0000000000..bd8a890529
--- /dev/null
+++ b/zig/tests/contract/metal_lfm2_controls.zig
@@ -0,0 +1,90 @@
+//! 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,
+ markov_nan_dropped,
+ markov_first_prev,
+ draft_ids_unzeroed,
+ block_attention_causal,
+ confidence_no_w1,
+ markov_bf16_nan_dropped,
+ markov_bf16_first_prev,
+ bf16_rows_short_k,
+ bf16_rows_last_row,
+ bf16_tile_untransposed,
+ bf16_tile_last_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",
+ .markov_nan_dropped => "control: metal lfm2 markov step, a NaN score is not flagged",
+ .markov_first_prev => "control: metal lfm2 markov step, every stream reads the first stream's previous token",
+ .draft_ids_unzeroed => "control: metal lfm2 draft ids, a pick past the vocabulary keeps its token",
+ .confidence_no_w1 => "control: metal lfm2 dspark confidence, the previous token's W1 row is not added",
+ .block_attention_causal => "control: metal lfm2 block attention, a row misses the last key",
+ .markov_bf16_nan_dropped => "control: metal lfm2 markov step bf16, a NaN score is not flagged",
+ .markov_bf16_first_prev => "control: metal lfm2 markov step bf16, every stream reads the first stream's previous token",
+ .bf16_rows_short_k => "control: metal lfm2 bf16 rows, a lane skips every second step of k",
+ .bf16_rows_last_row => "control: metal lfm2 bf16 rows, a group's last row is not computed",
+ .bf16_tile_untransposed => "control: metal lfm2 bf16 tile, W's tile is read untransposed",
+ .bf16_tile_last_column => "control: metal lfm2 bf16 tile, the last column is not written",
+ };
+ }
+
+ /// The gate that must go red.
+ pub fn gate(m: Mutation) ?contract.gates.Gate {
+ return switch (m) {
+ .none => null,
+ .argmax_no_flag, .markov_nan_dropped, .markov_bf16_nan_dropped => .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]" },
+ .markov_nan_dropped => .{ .from = "nan = nan || isnan(sc);", .to = "nan = false;" },
+ .markov_first_prev => .{ .from = "id = prev[size_t(s) * p.prev_stride];", .to = "id = prev[0];" },
+ .draft_ids_unzeroed => .{ .from = "ids[i] = ok ? token : 0u;", .to = "ids[i] = token;" },
+ .confidence_no_w1 => .{ .from = "if (id < p.vocab)\n for (uint k = t; k < p.rank", .to = "if (false)\n for (uint k = t; k < p.rank" },
+ .block_attention_causal => .{ .from = "count[b] = FULL ? p.len + 1 :", .to = "count[b] = FULL ? p.len :" },
+ .markov_bf16_nan_dropped => .{ .from = "nan = nan || isnan(sc);", .to = "nan = false;" },
+ .markov_bf16_first_prev => .{ .from = "id = prev[size_t(s) * p.prev_stride];", .to = "id = prev[0];" },
+ .bf16_rows_short_k => .{ .from = "for (uint i = 8 * lane; i < p.k; i += 256) {", .to = "for (uint i = 8 * lane; i < p.k; i += 512) {" },
+ .bf16_rows_last_row => .{ .from = "if (r >= rows) break;", .to = "if (r + 1 >= rows) break;" },
+ .bf16_tile_untransposed => .{ .from = "32, ulong2(0, 0), true);", .to = "32, ulong2(0, 0), false);" },
+ .bf16_tile_last_column => .{ .from = "if (row < p.rows && col < p.n) y[", .to = "if (row < p.rows && col + 1 < p.n) y[" },
+ .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_dspark.zig b/zig/tests/contract/metal_lfm2_dspark.zig
new file mode 100644
index 0000000000..04b7fc489a
--- /dev/null
+++ b/zig/tests/contract/metal_lfm2_dspark.zig
@@ -0,0 +1,264 @@
+//! LFM2's DSpark kernels: the Markov step (scan then pick), and held drafts into a verify's ids.
+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");
+
+pub const Which = enum { markov_step, draft_ids, markov_step_bf16 };
+const rank = 256;
+const vocab = 128000;
+const span = 512; // tokens a scan threadgroup takes (the kernel's MARKOV_SPAN)
+const groups = (vocab + span - 1) / span;
+const table = 1000; // W1's rows in the check: ids are inside it or past the vocabulary (a zero row)
+const names = [_][]const u8{ "metal lfm2 markov step", "metal lfm2 draft ids", "metal lfm2 markov step bf16" };
+const stride = 64; // words between a stream's picks: a hold pass's 32 blocks, a token and a flag each
+
+const ids_shapes = [_]c.Shape{ .{ .rows = 1, .n = 1, .k = 1 }, .{ .rows = 4, .n = 1, .k = 1 }, .{ .rows = 15, .n = 1, .k = 1 } };
+const step_shapes = [_]c.Shape{ .{ .rows = 1, .n = vocab, .k = rank }, .{ .rows = 4, .n = vocab, .k = rank }, .{ .rows = 9, .n = vocab, .k = rank } };
+
+fn none(_: c.Shape) usize {
+ return 0;
+}
+
+fn depth(_: c.Shape) usize {
+ return rank;
+}
+
+pub fn spec(w: Which) c.Spec {
+ return switch (w) {
+ .draft_ids => .{ .name = names[1], .accum = .none, .depth = none, .out = .raw, .served = &ids_shapes, .sums = .{ .f16 = 0, .bf16 = 0, .why = "a copy moves tokens: no sums" } },
+ .markov_step, .markov_step_bf16 => .{ .name = if (w == .markov_step) names[0] else names[2], .accum = .none, .depth = depth, .out = .raw, .served = &step_shapes, .nan = .{ .stated = "a row with a NaN score is flagged nonfinite" }, .sums = .{ .f16 = 0, .bf16 = 0, .why = "picks are tokens: the check bounds each pick's score by gamma(256), not sum bits" } },
+ };
+}
+
+/// 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,
+ which: Which,
+ mutation: controls.Mutation,
+ pipes: [2]mtl.Pipeline,
+ rows: usize = 0,
+ src: []u16 = &.{}, // the base logits rows
+ w1: []u16 = &.{},
+ codes: []u8 = &.{},
+ scales: []u16 = &.{},
+ biases: []u16 = &.{},
+ w2: []u16 = &.{}, // the bf16 step's W2
+ ids: []u32 = &.{},
+ args: [8]rig_.Ref = undefined,
+ y: rig_.Ref = undefined,
+ bad: rig_.Ref = undefined, // draft_ids: the flag word it sets for a flagged or invalid pick, zeroed each run
+
+ pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, w: Which, mutation: controls.Mutation) !Probe {
+ const text = try controls.source(a, mutation);
+ defer if (mutation != .none) a.free(text);
+ const kernels: [2][:0]const u8 = switch (w) {
+ .markov_step => .{ "lfm2_markov_scan", "lfm2_markov_pick" },
+ .markov_step_bf16 => .{ "lfm2_markov_scan_bf16", "lfm2_markov_pick" },
+ .draft_ids => .{ "lfm2_draft_ids", "lfm2_draft_ids" },
+ };
+ const first = try rig_.pipeline(device, text, kernels[0]);
+ errdefer first.deinit();
+ const second = try rig_.pipeline(device, text, kernels[1]);
+ return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .which = w, .mutation = mutation, .pipes = .{ first, second } };
+ }
+
+ fn drop(p: *Probe) void {
+ inline for (.{ "src", "w1", "codes", "scales", "biases", "w2", "ids" }) |f| {
+ p.a.free(@field(p, f));
+ @field(p, f) = &.{};
+ }
+ p.rig.reset();
+ }
+
+ pub fn deinit(p: *Probe) void {
+ p.drop();
+ p.rig.deinit();
+ for (p.pipes) |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 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();
+ p.rows = s.rows;
+ if (p.which == .draft_ids) { // picks a stride apart: tokens inside the vocabulary, its last, and past it
+ if (s.n != 1 or s.rows == 0 or s.rows > 16) return false;
+ p.ids = try p.a.alloc(u32, s.rows * stride);
+ for (p.ids, 0..) |*id, i| id.* = switch ((i / stride) % 3) {
+ 0 => @intCast((@as(u64, i) * 7919 + seed) % vocab),
+ 1 => vocab - 1,
+ else => vocab + @as(u32, @intCast(i % 5)),
+ };
+ // the flag word after each pick: zero, but planted nonzero on the last of every four streams
+ for (0..s.rows) |r| p.ids[r * stride + 1] = if (r % 4 == 3) 1 else 0;
+ p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.ids));
+ const at = [4]u32{ stride, @intCast(s.rows), vocab, 0 };
+ p.args[1] = try p.rig.input(std.mem.asBytes(&at));
+ p.y = try p.rig.output(p.count());
+ return true;
+ }
+ if (s.n != vocab or s.k != rank or s.rows == 0) return false;
+ // each stream's step row sits a block (two rows here) apart, as a hold's rows of one step do
+ p.src = try p.a.alloc(u16, 2 * s.rows * vocab);
+ fill(p.src, seed, f);
+ p.w1 = try p.a.alloc(u16, table * rank);
+ fill(p.w1, seed +% 5, .random);
+ p.codes = try p.a.alloc(u8, vocab * rank);
+ for (p.codes, 0..) |*q, i| q.* = @truncate((@as(u64, i) *% 2654435761 +% seed) >> 11);
+ p.scales = try p.a.alloc(u16, vocab * rank / 64);
+ p.biases = try p.a.alloc(u16, vocab * rank / 64);
+ for (p.scales, p.biases, 0..) |*sc, *bi, i| {
+ sc.* = c.toBf16(@as(f32, @floatFromInt(i % 7 + 1)) / 4096);
+ bi.* = c.toBf16(@as(f32, @floatFromInt(i % 5)) / 64 - 0.03);
+ }
+ if (p.which == .markov_step_bf16) {
+ p.w2 = try p.a.alloc(u16, vocab * rank);
+ fill(p.w2, seed +% 9, .random);
+ }
+ // the previous token every second word (the picks' layout): inside W1, its last row, or past the vocabulary
+ p.ids = try p.a.alloc(u32, 2 * s.rows);
+ for (p.ids, 0..) |*id, i| id.* = if (i % 2 == 1) 0xdead else switch ((i / 2) % 3) {
+ 0 => @intCast((@as(u64, i) * 37 + seed) % table),
+ 1 => table - 1,
+ else => vocab + @as(u32, @intCast(i)),
+ };
+ p.args[0] = try p.rig.input(std.mem.sliceAsBytes(p.src));
+ p.args[1] = try p.rig.input(std.mem.sliceAsBytes(p.w1));
+ p.args[2] = if (p.which == .markov_step_bf16) try p.rig.input(std.mem.sliceAsBytes(p.w2)) else try p.rig.input(p.codes);
+ p.args[3] = try p.rig.input(std.mem.sliceAsBytes(p.ids));
+ const args = [4]u32{ 2 * vocab, rank, vocab, 2 };
+ p.args[4] = try p.rig.input(std.mem.asBytes(&args));
+ p.args[5] = try p.rig.input(std.mem.sliceAsBytes(p.scales));
+ p.args[6] = try p.rig.input(std.mem.sliceAsBytes(p.biases));
+ const n: u32 = groups;
+ p.args[7] = try p.rig.input(std.mem.asBytes(&n));
+ p.y = try p.rig.output(p.count());
+ return true;
+ }
+
+ pub fn count(p: *Probe) usize {
+ return switch (p.which) {
+ .markov_step, .markov_step_bf16 => p.rows * 8,
+ .draft_ids => p.rows * 4,
+ };
+ }
+
+ pub fn run(p: *Probe) !void {
+ const pass = p.rig.begin();
+ const e = pass.e;
+ if (p.which == .draft_ids) {
+ e.setPipeline(p.pipes[0]);
+ 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.y.buf, p.y.off, 2);
+ p.bad = try p.rig.scratch(4);
+ @memset(p.rig.bytes(p.bad, 4), 0);
+ e.setBuffer(p.bad.buf, p.bad.off, 3);
+ e.dispatchThreads(mtl.Size.of(p.rows, 1, 1), mtl.Size.of(32, 1, 1));
+ return p.rig.finish(pass);
+ }
+ const best = try p.rig.scratch(p.rows * groups * 4);
+ const at = try p.rig.scratch(p.rows * groups * 4);
+ e.setPipeline(p.pipes[0]);
+ for (p.args[0..5], 0..) |r, i| e.setBuffer(r.buf, r.off, i);
+ e.setBuffer(best.buf, best.off, 5);
+ e.setBuffer(at.buf, at.off, 6);
+ if (p.which == .markov_step) {
+ e.setBuffer(p.args[5].buf, p.args[5].off, 7);
+ e.setBuffer(p.args[6].buf, p.args[6].off, 8);
+ }
+ e.dispatchGroups(mtl.Size.of(groups, p.rows, 1), mtl.Size.of(256, 1, 1));
+ e.setPipeline(p.pipes[1]);
+ e.setBuffer(best.buf, best.off, 0);
+ e.setBuffer(at.buf, at.off, 1);
+ e.setBuffer(p.args[7].buf, p.args[7].off, 2);
+ e.setBuffer(p.y.buf, p.y.off, 3);
+ e.dispatchGroups(mtl.Size.of(p.rows, 1, 1), mtl.Size.of(256, 1, 1));
+ try p.rig.finish(pass);
+ }
+
+ pub fn reference(_: *Probe, want: []f64, mag: []f64) !void {
+ @memset(mag, 0);
+ @memset(want, 0); // the step and the ids check themselves (check)
+ }
+
+ /// Each stream's pick against float64 scores with the kernel's bf16 roundings: the best, but for near ties.
+ pub fn check(p: *Probe, _: g.Gate) !?[]const u8 {
+ const out = p.rig.bytes(p.y, p.count());
+ if (p.which == .draft_ids) { // exact: each stream's first pick, or 0 for a token past the vocabulary or a flag
+ var any = false;
+ for (0..p.rows) |r| {
+ const token = p.ids[r * stride];
+ const ok = token < vocab and p.ids[r * stride + 1] == 0;
+ any = any or !ok;
+ if (std.mem.readInt(u32, out[4 * r ..][0..4], .little) != (if (ok) token else 0)) return "an id differs from its pick";
+ }
+ const raised = std.mem.readInt(u32, p.rig.bytes(p.bad, 4)[0..4], .little);
+ if ((raised != 0) != any) return "the bad flag does not follow the flagged and invalid picks";
+ return null;
+ }
+ for (0..p.rows) |s| {
+ const token = std.mem.readInt(u32, out[8 * s ..][0..4], .little);
+ const flag = std.mem.readInt(u32, out[8 * s + 4 ..][0..4], .little);
+ const id = p.ids[2 * s];
+ var top: f64 = -std.math.inf(f64);
+ var mine: ?f64 = null;
+ var nan = false;
+ var bound: f64 = 0;
+ for (0..vocab) |v| {
+ var dot: f64 = 0;
+ var size: f64 = 0;
+ if (id < table) for (0..rank) |i| {
+ const k = (v * rank + i) / 64;
+ const w: f64 = if (p.which == .markov_step_bf16) c.fromBf16(p.w2[v * rank + i]) else @as(f64, c.fromBf16(p.scales[k])) * @as(f64, @floatFromInt(p.codes[v * rank + i])) + c.fromBf16(p.biases[k]);
+ dot += w * c.fromBf16(p.w1[id * rank + i]);
+ size += @abs(w * c.fromBf16(p.w1[id * rank + i]));
+ };
+ const base: f32 = c.fromBf16(p.src[2 * s * vocab + v]);
+ const score: f64 = c.fromBf16(c.toBf16(base + c.fromBf16(c.toBf16(@floatCast(dot)))));
+ if (std.math.isNan(score)) nan = true;
+ if (score > top) top = score;
+ if (v == token) mine = score;
+ // two bf16 steps of the score and the fp32 sum's own error
+ bound = @max(bound, @abs(score) * 0x1p-7 + c.gamma(rank, 0x1p-24) * size);
+ }
+ if (nan) {
+ if (flag != 1) return "a row with a NaN score is not flagged";
+ continue;
+ }
+ if (flag != 0) return "a finite row is flagged nonfinite";
+ const got = mine orelse return "the pick is past the vocabulary";
+ if (got < top - bound) return "the pick's score is below the row's best by more than its bound";
+ }
+ return null;
+ }
+};
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..4007a7566a 100644
--- a/zig/tests/contract/metal_main.zig
+++ b/zig/tests/contract/metal_main.zig
@@ -23,6 +23,14 @@ 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_dspark = @import("metal_lfm2_dspark.zig");
+const lfm2_bf16 = @import("metal_lfm2_bf16.zig");
+const lfm2_attn = @import("metal_lfm2_attn.zig");
+const lfm2_conf = @import("metal_lfm2_confidence.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 +159,137 @@ 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 drafter_list: std.ArrayList(*lfm2_dspark.Probe) = .empty;
+ var block_list: std.ArrayList(*lfm2_attn.Probe) = .empty;
+ var bf16_list: std.ArrayList(*lfm2_bf16.Probe) = .empty;
+ defer {
+ for (drafter_list.items) |p| {
+ p.deinit();
+ a.destroy(p);
+ }
+ for (block_list.items) |p| {
+ p.deinit();
+ a.destroy(p);
+ }
+ for (bf16_list.items) |p| {
+ p.deinit();
+ a.destroy(p);
+ }
+ drafter_list.deinit(a);
+ block_list.deinit(a);
+ bf16_list.deinit(a);
+ }
+ for (std.enums.values(lfm2_controls.Mutation)) |m| {
+ const drafter_which: ?lfm2_dspark.Which = switch (m) {
+ .markov_nan_dropped, .markov_first_prev => .markov_step,
+ .markov_bf16_nan_dropped, .markov_bf16_first_prev => .markov_step_bf16,
+ .draft_ids_unzeroed => .draft_ids,
+ else => null,
+ };
+ if (m == .none or drafter_which != null) {
+ for (std.enums.values(lfm2_dspark.Which)) |w| {
+ if (m != .none and w != drafter_which.?) continue;
+ if (!wanted(o, if (m == .none) lfm2_dspark.spec(w).name else m.name())) continue;
+ const p = try a.create(lfm2_dspark.Probe);
+ errdefer a.destroy(p);
+ p.* = try lfm2_dspark.Probe.init(a, device, queue, w, m);
+ try drafter_list.append(a, p);
+ try kernels.append(a, p.kernel());
+ }
+ }
+ const bf16_which: ?lfm2_bf16.Which = switch (m) {
+ .bf16_rows_short_k, .bf16_rows_last_row => .rows,
+ .bf16_tile_untransposed, .bf16_tile_last_column => .tile,
+ else => null,
+ };
+ if (m == .none or bf16_which != null) {
+ for (std.enums.values(lfm2_bf16.Which)) |w| {
+ if (m != .none and w != bf16_which.?) continue;
+ if (!wanted(o, if (m == .none) lfm2_bf16.spec(w).name else m.name())) continue;
+ const p = try a.create(lfm2_bf16.Probe);
+ errdefer a.destroy(p);
+ p.* = try lfm2_bf16.Probe.init(a, device, queue, w, m);
+ try bf16_list.append(a, p);
+ try kernels.append(a, p.kernel());
+ }
+ }
+ if (m == .none or m == .block_attention_causal) {
+ if (!wanted(o, if (m == .none) lfm2_attn.name else m.name())) continue;
+ const p = try a.create(lfm2_attn.Probe);
+ errdefer a.destroy(p);
+ p.* = try lfm2_attn.Probe.init(a, device, queue, m);
+ try block_list.append(a, p);
+ try kernels.append(a, p.kernel());
+ }
+ }
+ var conf_list: std.ArrayList(*lfm2_conf.Probe) = .empty;
+ defer {
+ for (conf_list.items) |p| {
+ p.deinit();
+ a.destroy(p);
+ }
+ conf_list.deinit(a);
+ }
+ for ([_]lfm2_controls.Mutation{ .none, .confidence_no_w1 }) |m| {
+ if (!wanted(o, if (m == .none) lfm2_conf.name else m.name())) continue;
+ const p = try a.create(lfm2_conf.Probe);
+ errdefer a.destroy(p);
+ p.* = try lfm2_conf.Probe.init(a, device, queue, m);
+ try conf_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_dspark.zig b/zig/tests/lfm2_dspark.zig
new file mode 100644
index 0000000000..4fb630fb13
--- /dev/null
+++ b/zig/tests/lfm2_dspark.zig
@@ -0,0 +1,398 @@
+//! LFM2's DSpark drafter on the GPU against an fp32 host reference of the released model, on the real weights.
+const std = @import("std");
+const mtl = @import("metal");
+const tf = @import("tensorfold");
+const q = tf.lfm2_moe;
+const c = q.config;
+const st = q.state;
+const ck = tf.checkpoint;
+const quant = q.dspark_quant;
+
+const hd = c.head_dim;
+const kw = c.kv_heads * hd;
+const qw = c.query_heads * hd;
+
+/// The drafter's tensors as fp32 on the host, in the checkpoint's own layout (interleaved RoPE pairs).
+const Ref = struct {
+ gpa: std.mem.Allocator,
+ check: ck.Checkpoint,
+ cfg: q.dspark_config.Config,
+
+ fn get(r: *const Ref, name: []const u8) ![]f32 {
+ const t = try r.check.get(name);
+ const src = t.host(u16);
+ const out = try r.gpa.alloc(f32, src.len);
+ for (out, src) |*o, v| o.* = quant.toF32(v);
+ return out;
+ }
+
+ fn layer(r: *const Ref, buf: []u8, i: usize, part: []const u8) ![]f32 {
+ return r.get(try std.fmt.bufPrint(buf, "layers.{d}.{s}.weight", .{ i, part }));
+ }
+};
+
+/// y[n] = W[n, k] x[k] for each of `rows` rows, on every core.
+fn matmul(gpa: std.mem.Allocator, w: []const f32, x: []const f32, y: []f32, n: usize, k: usize, rows: usize) !void {
+ const Job = struct {
+ fn run(w_: []const f32, x_: []const f32, y_: []f32, n_: usize, k_: usize, rows_: usize, from: usize, to: usize) void {
+ for (0..rows_) |r| for (from..to) |col| {
+ var s: f32 = 0;
+ const wr = w_[col * k_ ..][0..k_];
+ const xr = x_[r * k_ ..][0..k_];
+ for (wr, xr) |a, b| s += a * b;
+ y_[r * n_ + col] = s;
+ };
+ }
+ };
+ const threads = @min(std.Thread.getCpuCount() catch 1, 16);
+ const per = (n + threads - 1) / threads;
+ const hs = try gpa.alloc(std.Thread, threads);
+ defer gpa.free(hs);
+ var made: usize = 0;
+ defer for (hs[0..made]) |h| h.join(); // a failed spawn still joins the workers already running
+ var at: usize = 0;
+ while (at < n) : (at += per) {
+ hs[made] = try std.Thread.spawn(.{}, Job.run, .{ w, x, y, n, k, rows, at, @min(n, at + per) });
+ made += 1;
+ }
+}
+
+fn rms(x: []f32, gain: []const f32) void {
+ var ss: f32 = 0;
+ for (x) |v| ss += v * v;
+ const inv = 1 / @sqrt(ss / @as(f32, @floatFromInt(x.len)) + c.eps);
+ for (x, gain) |*v, g| v.* = v.* * inv * g;
+}
+
+/// Interleaved RoPE at `pos` over each head of `x` (pairs 2i, 2i + 1).
+fn rope(x: []f32, pos: u64, theta: f32) void {
+ var h: usize = 0;
+ while (h < x.len) : (h += hd) for (0..hd / 2) |i| {
+ const angle = @as(f32, @floatFromInt(pos)) * std.math.pow(f32, theta, -@as(f32, @floatFromInt(i)) / (hd / 2));
+ const a = x[h + 2 * i];
+ const b = x[h + 2 * i + 1];
+ x[h + 2 * i] = a * @cos(angle) - b * @sin(angle);
+ x[h + 2 * i + 1] = a * @sin(angle) + b * @cos(angle);
+ };
+}
+
+/// The target's embedding row `id`, dequantized from its codes at the checkpoint's own width (scale * q + bias).
+fn embedRow(m: *q.Model, id: u32, out: []f32) void {
+ const e = m.weights.embedding;
+ const bits: usize = e.bits;
+ const row_bytes = c.hidden * bits / 8;
+ const codes = (e.w.contents() + e.w_off)[@as(usize, id) * row_bytes ..][0..row_bytes];
+ const sc = @as([*]const u16, @ptrCast(@alignCast(e.scales.contents() + e.s_off)))[@as(usize, id) * c.hidden / 64 ..][0 .. c.hidden / 64];
+ const bi = @as([*]const u16, @ptrCast(@alignCast(e.biases.contents() + e.b_off)))[@as(usize, id) * c.hidden / 64 ..][0 .. c.hidden / 64];
+ for (0..c.hidden) |j| {
+ // codes are a little-endian bit stream: code j takes bits [j * bits, (j + 1) * bits)
+ const at = j * bits;
+ var word: u32 = codes[at / 8];
+ if (at / 8 + 1 < row_bytes) word |= @as(u32, codes[at / 8 + 1]) << 8;
+ const qv = (word >> @intCast(at % 8)) & ((@as(u32, 1) << @intCast(bits)) - 1);
+ out[j] = quant.toF32(sc[j / 64]) * @as(f32, @floatFromInt(qv)) + quant.toF32(bi[j / 64]);
+ }
+}
+
+/// GPU bf16 row against the reference's row + bias: cosine and worst difference.
+fn compare(label: []const u8, gpu: []const u16, row: []const f32, bias: []const f32) void {
+ var dot: f64 = 0;
+ var gg: f64 = 0;
+ var rr: f64 = 0;
+ var worst: f32 = 0;
+ for (gpu, row, bias) |g, b, m2| {
+ const x = quant.toF32(g);
+ const y = b + m2;
+ dot += x * y;
+ gg += x * x;
+ rr += y * y;
+ worst = @max(worst, @abs(x - y));
+ }
+ std.debug.print("{s}: cosine {d:.5}, worst difference {d:.3}\n", .{ label, dot / @sqrt(gg * rr), worst });
+}
+
+/// A hold's GPU ms: the blocks held in a signalled command buffer, waited on and read for its GPU times.
+fn gpuHold(d: *q.dspark.Drafter, blocks: []const q.dspark_block.Block) !f64 {
+ d.signalled += 1;
+ const cb = (try q.dspark_block.hold(d, blocks, .{ .event = d.event, .value = d.signalled })).?;
+ defer cb.release();
+ cb.wait();
+ if (cb.failure() != null) return error.GpuFailed;
+ return (cb.gpuEnd() - cb.gpuStart()) * 1000;
+}
+
+fn absorb(d: *q.dspark.Drafter, piece: q.dspark_block.Piece) !void {
+ const run = q.dspark_block.begin(d);
+ try q.dspark_block.absorbPass(d, run.e, &.{piece});
+ try q.dspark_block.end(run.cb, run.e);
+}
+
+/// Served rounds through the drafter's interface: after four warm rounds, no host allocation and one hold pass a round.
+fn steady(gpa: std.mem.Allocator, d: *q.dspark.Drafter, taps: mtl.Buffer, ids: []const u32, ctx: usize) !void {
+ const lanes = tf.lanes;
+ var s = try lanes.Stream.init(gpa, .{ .id = "steady", .prompt = ids[0 .. ctx + 1], .max_new = 160 });
+ defer s.deinit(gpa);
+ const dr = d.drafter();
+ try dr.open(&s);
+ defer dr.release(&s);
+ const width: u32 = @intCast(d.tapBytes());
+ const feats = struct {
+ fn at(buf: mtl.Buffer, row: usize, rows: usize, w: u32) lanes.backend.Features {
+ return .{ .buffer = @intFromPtr(buf.id), .offset = row * w, .rows = @intCast(rows), .row_bytes = w, .space = .device, .dtype = .bf16 };
+ }
+ };
+ try dr.absorb(&.{.{ .stream = &s, .features = feats.at(taps, 0, ctx, width), .start = 0, .follow = &.{} }});
+ var counted = std.testing.FailingAllocator.init(gpa, .{});
+ var anchor: u64 = ctx;
+ const rounds = 24;
+ const warm = 4;
+ var holds: u64 = 0;
+ for (0..rounds) |r| {
+ if (r == warm) {
+ d.gpa = counted.allocator();
+ holds = d.signalled;
+ }
+ try dr.hold(&.{.{ .stream = &s, .pending = ids[@intCast(anchor % ctx)], .position = anchor + 1, .depth = 3 }});
+ var chances: [3]f64 = undefined;
+ _ = try dr.chances(&s, &chances);
+ var drafts: [3]u32 = undefined;
+ try dr.held(&.{&s}, &.{&drafts});
+ const rows = r % 3 + 1;
+ const at: usize = @intCast(anchor % (ctx - 3));
+ try dr.absorb(&.{.{ .stream = &s, .features = feats.at(taps, at, rows, width), .start = anchor, .follow = &.{}, .prompt = false }});
+ anchor += rows;
+ }
+ d.gpa = gpa;
+ const n = rounds - warm;
+ std.debug.print("steady rounds: {d} rounds of hold, chances, held and a verified absorb: {d} host allocations, {d} resizes, {d} unwaited holds\n", .{ n, counted.allocations, counted.resize_index, d.signalled - holds });
+ if (counted.allocations != 0 or counted.resize_index != 0 or d.signalled - holds != n) return error.SteadyRoundsAllocate;
+}
+
+pub fn main(init: std.process.Init) !void {
+ const gpa = init.gpa;
+ const args = try init.minimal.args.toSlice(init.arena.allocator());
+ if (args.len < 4 or args.len > 6) return error.ExpectedModelDrafterPromptLengthRowsBits;
+ const bits: u8 = if (args.len == 6) try std.fmt.parseInt(u8, args[5], 10) else 0;
+ const length = try std.fmt.parseInt(usize, args[3], 10);
+ const pool = mtl.objc.Pool.push();
+ defer pool.pop();
+ const m = try q.Model.load(gpa, init.io, args[1]);
+ defer m.deinit();
+ const d = try q.dspark.Drafter.init(gpa, init.io, m, args[2], .{ .capacity = length + 64, .bits = bits });
+ defer d.deinit();
+ const cfg = d.config;
+ // the compared block's rows (a hold of depth d runs a d-row block)
+ const block: u32 = if (args.len >= 5) try std.fmt.parseInt(u32, args[4], 10) else cfg.block;
+ if (block == 0 or block > cfg.block) return error.BadBlockRows;
+ var scratch = try st.Scratch.init(gpa, m.device, &m.config.attention, length);
+ defer scratch.deinit();
+ const width = cfg.tap_count * c.hidden;
+ var taps: st.Taps = .{ .count = cfg.tap_count, .buf = try scratch.memory.alloc(length * width * 2) };
+ @memcpy(taps.layers[0..cfg.tap_count], cfg.taps[0..cfg.tap_count]);
+ scratch.taps = taps;
+ var cache = try st.Cache.init(gpa, m.device, &m.config.attention, length + 16);
+ defer cache.deinit();
+ const ids = try gpa.alloc(u32, length);
+ defer gpa.free(ids);
+ for (ids, 0..) |*id, i| id.* = @intCast((i * 7919 + 13) % 120000);
+ try q.forward.run(m, &scratch, &.{.{ .cache = &cache, .rows = length }}, ids, false, .last);
+ // the context: rows [0, length - 1); the pending token is the last id, at position length - 1
+ const ctx = length - 1;
+ const lane = try d.newLane(length + 32);
+ defer d.dropLane(lane);
+ var at: usize = 0;
+ while (at < ctx) : (at += q.dspark.absorb_rows) {
+ const rows = @min(q.dspark.absorb_rows, ctx - at);
+ try absorb(d, .{ .lane = lane, .src = .{ .buf = taps.buf, .off = at * width * 2 }, .rows = rows, .p0 = at });
+ }
+ lane.end = ctx;
+ _ = try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = block }}, null);
+ var held: [16]u32 = undefined;
+ @memcpy(held[0..block], lane.held[0..block]);
+ var conf: [16]f32 = undefined;
+ @memcpy(conf[0..block], lane.conf[0..block]);
+ const gpu = held[0..block];
+ for ([_]usize{ 1, 4, 16, 128, q.dspark.absorb_rows }) |rows| {
+ if (rows > ctx) continue;
+ var best: f64 = std.math.floatMax(f64);
+ for (0..5) |_| {
+ const t0 = std.Io.Clock.awake.now(init.io);
+ try absorb(d, .{ .lane = lane, .src = .{ .buf = taps.buf, .off = (ctx - rows) * width * 2 }, .rows = rows, .p0 = ctx - rows });
+ best = @min(best, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(init.io)).toNanoseconds())) / 1e6);
+ }
+ std.debug.print("absorb of {d} rows: {d:.2} ms (best of 5)\n", .{ rows, best });
+ }
+ const depths = [_]u32{ 1, 3, cfg.block };
+ for (depths) |depth| {
+ var best: f64 = std.math.floatMax(f64);
+ for (0..5) |_| {
+ const t0 = std.Io.Clock.awake.now(init.io);
+ _ = try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = depth }}, null);
+ best = @min(best, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(init.io)).toNanoseconds())) / 1e6);
+ }
+ std.debug.print("hold at depth {d}: {d:.2} ms (best of 5)\n", .{ depth, best });
+ }
+ // a shared round's hold: S streams' blocks of d rows in one pass (S x d at most pass_rows), each on this context
+ var lanes_: [16]*q.dspark.Lane = undefined;
+ lanes_[0] = lane;
+ var made: usize = 1;
+ defer for (lanes_[1..made]) |l| d.dropLane(l);
+ while (made < lanes_.len) : (made += 1) {
+ const l = try d.newLane(length + 32);
+ lanes_[made] = l;
+ var from: usize = 0;
+ while (from < ctx) : (from += q.dspark.absorb_rows) {
+ const rows = @min(q.dspark.absorb_rows, ctx - from);
+ try absorb(d, .{ .lane = l, .src = .{ .buf = taps.buf, .off = from * width * 2 }, .rows = rows, .p0 = from });
+ }
+ l.end = ctx;
+ }
+ for ([_]usize{ 2, 4, 8, 16 }) |streams| for ([_]u32{ 1, 2, 3, 4, 8 }) |depth| {
+ if (streams * depth > q.dspark.pass_rows) continue;
+ var blocks: [16]q.dspark_block.Block = undefined;
+ for (blocks[0..streams], lanes_[0..streams]) |*b, l| b.* = .{ .lane = l, .pending = ids[ctx], .anchor = ctx, .depth = depth };
+ var best: f64 = std.math.floatMax(f64);
+ var gpu_best: f64 = std.math.floatMax(f64);
+ for (0..5) |_| {
+ const t0 = std.Io.Clock.awake.now(init.io);
+ _ = try q.dspark_block.hold(d, blocks[0..streams], null);
+ best = @min(best, @as(f64, @floatFromInt(t0.durationTo(std.Io.Clock.awake.now(init.io)).toNanoseconds())) / 1e6);
+ gpu_best = @min(gpu_best, try gpuHold(d, blocks[0..streams]));
+ }
+ std.debug.print("hold of {d} streams at depth {d}: {d:.2} ms, GPU {d:.2} ms (best of 5)\n", .{ streams, depth, best, gpu_best });
+ };
+ try steady(gpa, d, taps.buf, ids, ctx);
+ _ = try q.dspark_block.hold(d, &.{.{ .lane = lane, .pending = ids[ctx], .anchor = ctx, .depth = block }}, null);
+ // the host reference, fp32 throughout
+ var r: Ref = .{ .gpa = gpa, .check = ck.Checkpoint.init(gpa), .cfg = cfg };
+ defer r.check.deinit();
+ const path = try std.fs.path.joinZ(gpa, &.{ args[2], "model.safetensors" });
+ defer gpa.free(path);
+ try r.check.addFile(m.device, path, "");
+ var arena = std.heap.ArenaAllocator.init(gpa);
+ defer arena.deinit();
+ r.gpa = arena.allocator();
+ const a = r.gpa;
+ const n = block;
+ const tap_rows = taps.buf.slice(u16, length * width);
+ const feats = try a.alloc(f32, ctx * width);
+ for (feats, tap_rows[0 .. ctx * width]) |*f, v| f.* = quant.toF32(v);
+ const hctx = try a.alloc(f32, ctx * c.hidden);
+ try matmul(gpa, try r.get("fc.weight"), feats, hctx, c.hidden, width, ctx);
+ const hn = try r.get("hidden_norm.weight");
+ for (0..ctx) |j| rms(hctx[j * c.hidden ..][0..c.hidden], hn);
+ const h = try a.alloc(f32, n * c.hidden);
+ for (0..n) |i| embedRow(m, if (i == 0) ids[ctx] else cfg.mask, h[i * c.hidden ..][0..c.hidden]);
+ const keys = ctx + n;
+ var name: [96]u8 = undefined;
+ const x = try a.alloc(f32, n * c.hidden);
+ for (0..cfg.layers) |li| {
+ const kmat = try r.layer(&name, li, "self_attn.k_proj");
+ const vmat = try r.layer(&name, li, "self_attn.v_proj");
+ const qn = try r.layer(&name, li, "self_attn.q_norm");
+ const kn = try r.layer(&name, li, "self_attn.k_norm");
+ const kk = try a.alloc(f32, keys * kw);
+ const vv = try a.alloc(f32, keys * kw);
+ try matmul(gpa, kmat, hctx, kk, kw, c.hidden, ctx);
+ try matmul(gpa, vmat, hctx, vv, kw, c.hidden, ctx);
+ @memcpy(x, h);
+ for (0..n) |i| rms(x[i * c.hidden ..][0..c.hidden], try r.layer(&name, li, "input_layernorm"));
+ const qq = try a.alloc(f32, n * qw);
+ try matmul(gpa, try r.layer(&name, li, "self_attn.q_proj"), x, qq, qw, c.hidden, n);
+ try matmul(gpa, kmat, x, kk[ctx * kw ..], kw, c.hidden, n);
+ try matmul(gpa, vmat, x, vv[ctx * kw ..], kw, c.hidden, n);
+ for (0..keys) |j| for (0..c.kv_heads) |kh| {
+ const row = kk[j * kw + kh * hd ..][0..hd];
+ rms(row, kn);
+ rope(row, j, cfg.theta);
+ };
+ for (0..n) |i| for (0..c.query_heads) |qh| {
+ const row = qq[i * qw + qh * hd ..][0..hd];
+ rms(row, qn);
+ rope(row, ctx + i, cfg.theta);
+ };
+ const o = try a.alloc(f32, n * qw);
+ const sc = try a.alloc(f32, keys);
+ for (0..n) |i| for (0..c.query_heads) |qh| {
+ const kh = qh / (c.query_heads / c.kv_heads);
+ var top: f32 = -std.math.inf(f32);
+ for (0..keys) |j| {
+ var s: f32 = 0;
+ for (0..hd) |t| s += qq[i * qw + qh * hd + t] * kk[j * kw + kh * hd + t];
+ sc[j] = s / @sqrt(@as(f32, hd));
+ top = @max(top, sc[j]);
+ }
+ var sum: f32 = 0;
+ for (sc) |*s| {
+ s.* = @exp(s.* - top);
+ sum += s.*;
+ }
+ for (0..hd) |t| {
+ var acc: f32 = 0;
+ for (0..keys) |j| acc += sc[j] * vv[j * kw + kh * hd + t];
+ o[i * qw + qh * hd + t] = acc / sum;
+ }
+ };
+ const y = try a.alloc(f32, n * c.hidden);
+ try matmul(gpa, try r.layer(&name, li, "self_attn.o_proj"), o, y, c.hidden, qw, n);
+ for (h, y) |*v, w| v.* += w;
+ @memcpy(x, h);
+ for (0..n) |i| rms(x[i * c.hidden ..][0..c.hidden], try r.layer(&name, li, "post_attention_layernorm"));
+ const gate = try a.alloc(f32, n * cfg.intermediate);
+ const up = try a.alloc(f32, n * cfg.intermediate);
+ try matmul(gpa, try r.layer(&name, li, "mlp.gate_proj"), x, gate, cfg.intermediate, c.hidden, n);
+ try matmul(gpa, try r.layer(&name, li, "mlp.up_proj"), x, up, cfg.intermediate, c.hidden, n);
+ for (gate, up) |*g, u| g.* = g.* / (1 + @exp(-g.*)) * u;
+ try matmul(gpa, try r.layer(&name, li, "mlp.down_proj"), gate, y, c.hidden, cfg.intermediate, n);
+ for (h, y) |*v, w| v.* += w;
+ }
+ const fin = try r.get("norm.weight");
+ for (0..n) |i| rms(h[i * c.hidden ..][0..c.hidden], fin);
+ const head = try a.alloc(f32, c.vocab * c.hidden);
+ for (0..c.vocab) |v| embedRow(m, @intCast(v), head[v * c.hidden ..][0..c.hidden]);
+ const base = try a.alloc(f32, n * c.vocab);
+ try matmul(gpa, head, h, base, c.vocab, c.hidden, n);
+ const w1 = try r.get("markov_head.markov_w1.weight");
+ const w2 = try r.get("markov_head.markov_w2.weight");
+ const bias = try a.alloc(f32, c.vocab);
+ const proj = try r.get("confidence_head.proj.weight");
+ const proj_bias = (try r.get("confidence_head.proj.bias"))[0];
+ var conf_worst: f32 = 0;
+ var prev = ids[ctx];
+ var same: usize = 0;
+ var worst: f32 = 0;
+ for (0..n) |i| {
+ try matmul(gpa, w2, w1[@as(usize, prev) * cfg.rank ..][0..cfg.rank], bias, c.vocab, cfg.rank, 1);
+ const row = base[i * c.vocab ..][0..c.vocab];
+ var best: usize = 0;
+ var second: usize = 1;
+ for (row, bias, 0..) |b, m2, v| {
+ const s = b + m2;
+ if (s > row[best] + bias[best]) {
+ second = best;
+ best = v;
+ } else if (v != best and s > row[second] + bias[second]) second = v;
+ }
+ const margin = (row[best] + bias[best]) - (row[second] + bias[second]);
+ if (i == 0) { // the GPU keeps the base logits (the Markov step adds its bias on the way to the pick)
+ const zero = try a.alloc(f32, c.vocab);
+ @memset(zero, 0);
+ compare("row 0 base logits", d.s.logits.slice(u16, c.vocab), row, zero);
+ }
+ const ok = gpu[i] == best;
+ same += @intFromBool(ok);
+ std.debug.print("step {d}: reference {d} (margin {d:.3}, second {d}), GPU {d} {s}\n", .{ i, best, margin, second, gpu[i], if (ok) "equal" else "DIFFERENT" });
+ if (!ok) worst = @max(worst, margin);
+ // the confidence head: [final norm row i, W1 of the token before draft i] . proj + bias, through a sigmoid
+ var logit: f32 = proj_bias;
+ for (h[i * c.hidden ..][0..c.hidden], proj[0..c.hidden]) |x0, w0| logit += x0 * w0;
+ for (w1[@as(usize, prev) * cfg.rank ..][0..cfg.rank], proj[c.hidden..]) |x0, w0| logit += x0 * w0;
+ const want = 1 / (1 + @exp(-logit));
+ conf_worst = @max(conf_worst, @abs(want - conf[i]));
+ std.debug.print("step {d}: confidence reference {d:.4}, GPU {d:.4}\n", .{ i, want, conf[i] });
+ prev = gpu[i]; // teacher-forced: step i + 1 follows the GPU's pick, so each step is compared alone
+ }
+ std.debug.print("the reference's largest margin over a different GPU pick: {d:.3}\n", .{worst});
+ std.debug.print("the confidences' worst difference from the fp32 reference: {d:.5}\n", .{conf_worst});
+ std.debug.print("{d} of {d} drafts equal the fp32 reference's picks on the same chain ({d}-token context)\n", .{ same, n, ctx });
+}
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;