diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7d3f34c0cc3..ae965c88f54 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,11 +4,13 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
## Unreleased
-- Cortheon (`--cortheon`): DuckDuckGo results are kept, and its ads still skipped, when only the spelling differs (either quote style, a `class` attribute in any letter case, an unquoted `href`, spaces around its `=`, or a line break before it); a duplicate whose scheme or host differs only in letter case is dropped while paths and queries keep their case, and the PyPI, npm and crates.io allow-list takes the host in any case with the path exact (#2750). Thanks to @ppardus.
+- CUDA EXL3 grouped experts launch and size their storage at the caller's tile and pair count, so decode no longer launches three quarters of its grid empty, and Qwen3.6-35B-A3B EXL3 sizes its grouped storage for 36,864 pair rows instead of 163,904 (about 3 GB less); outputs are bit-equal (#2843, for #2823).
-- GLM-5.3-Flash on CUDA: with `--vision` and no video library, rank 0 says at startup how much each stream sets aside for `--vision-video-tokens` rows that no request can fill (128 MiB a stream at the default 16,384), and that `--vision-video-tokens 64` on every rank cuts it to 512 KiB (#2813). Thanks to @Bizuayeu.
+- CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney.
+- Cortheon (`--cortheon`): DuckDuckGo results are kept, and its ads still skipped, when only the spelling differs (either quote style, a `class` attribute in any letter case, an unquoted `href`, spaces around its `=`, or a line break before it); a duplicate whose scheme or host differs only in letter case is dropped while paths and queries keep their case, and the PyPI, npm and crates.io allow-list takes the host in any case with the path exact (#2750). Thanks to @ppardus.
+- GLM-5.3-Flash on CUDA: with `--vision` and no video library, rank 0 says at startup how much each stream sets aside for `--vision-video-tokens` rows that no request can fill (128 MiB a stream at the default 16,384), and that `--vision-video-tokens 64` on every rank cuts it to 512 KiB (#2813). Thanks to @Bizuayeu.
- CUDA: `Qwen/Qwen3-Embedding-0.6B`, `-4B` and `-8B` load on the native engine from the published BF16 weights and serve `/v1/embeddings` (last-token pooling, unit-normed, `dimensions` crops). On one GB10 the worst cosine distance to float32 is 1.55e-4, 1.34e-4 and 1.08e-4 (transformers' own BF16: 1.88e-4, 1.79e-4 and 1.77e-4), a request's vectors equal each input sent alone, and a reranker folder or a flag the family does not serve is refused by name (#2285). Thanks to @Arminova.
- Qwen3.8-27B on Metal: the startup line says how many streams draft ("for both streams", "for all 4 streams", "for 8 of 16 streams"), the smaller of the stream count and `TF_Q27_DRAFT_SLOTS`, instead of "for one of them"; log text only (#2806).
diff --git a/README.md b/README.md
index 4469c700e83..33840f0b8dc 100644
--- a/README.md
+++ b/README.md
@@ -177,7 +177,7 @@ Conversions are explicitly marked: their upstream weights need the linked conver
| Model / format | Download | Machine and memory scope | Native recipe / limits |
| --- | --- | --- | --- |
-| dots.mocr, BF16 OCR | Use the local `dots-studio/dots.mocr` checkpoint. | Native HIP; AMD R9700 (`gfx1201`) checked, one GPU. | [Native OCR recipe](docs/recipes/dots-mocr-hip.md); full model-card demo matches HF tokens. No video, MTP head, CUDA, or Metal qualification. |
+| dots.mocr, BF16 OCR | Use the local `dots-studio/dots.mocr` checkpoint. | Native HIP; AMD R9700 (`gfx1201`) checked, one GPU. | [Native OCR recipe](docs/recipes/dots-mocr-hip.md); full model-card demo matches HF tokens. Native CUDA serves text prompts only ([CUDA recipe](docs/recipes/dots-mocr-cuda.md)). No video, MTP head, or Metal qualification. |
| Qwen3.5-2B, MLX 4-bit | `tensorfold pull mlx-community/Qwen3.5-2B-MLX-4bit@93760be4f1f69842a46bc13dbdc0f19e291392a3` | One Mac (M5 Max checked), or one GB10; no RAM minimum published. | [2B](docs/recipes/qwen3.5-2b-native.md); [CUDA release checks](CHANGELOG.md#110). Text only. |
| Qwen3.8-27B, MLX 4-bit | `tensorfold pull TensorFold/Qwen3.8-27B-MLX-4bit`
`tensorfold pull z-lab/Qwen3.8-27B-DFlash2` | One Mac (M5 Max/M3 Ultra), or one GB10; a 32 GB Mac needs `iogpu.wired_limit_mb` raised (the startup refusal names the value). | [27B](docs/recipes/qwen3.8-27b.md). Add `--drafter z-lab/Qwen3.8-27B-DFlash2`; without it decoding is plain. |
| Nemotron 3.5 Lightning 30B-A3B, MLX 4-bit | `tensorfold pull TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit` | One Mac (M1-M5), or one GB10; also checked on RTX 3090 / RTX PRO 6000. | [Lightning](docs/recipes/nemotron-3.5.md). Includes MTP. |
diff --git a/docs/recipes/README.md b/docs/recipes/README.md
index 6abb258dffa..a84de785b19 100644
--- a/docs/recipes/README.md
+++ b/docs/recipes/README.md
@@ -8,6 +8,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL.
| Family | Recipe |
| --- | --- |
| dots.mocr | [Native HIP OCR, BF16, one GPU](dots-mocr-hip.md) |
+| dots.mocr | [Native CUDA text decoder, BF16, one GPU](dots-mocr-cuda.md) |
| Kimi K3 / K3.1 | [Native Metal bring-up, synthetic only, resident budget refusal](kimi-k3-metal.md) |
| Nemotron 3.5 Lightning | [MLX](nemotron-3.5.md), [Intel Arc, experimental](xpu-arc.md) |
| Nemotron 3.5 Parse | [Native HIP OCR, validation and qualification limits](nemotron-parse.md) |
diff --git a/docs/recipes/dots-mocr-cuda.md b/docs/recipes/dots-mocr-cuda.md
new file mode 100644
index 00000000000..3a218382fec
--- /dev/null
+++ b/docs/recipes/dots-mocr-cuda.md
@@ -0,0 +1,37 @@
+# dots.mocr on native CUDA
+
+The native CUDA engine serves dots.mocr's Qwen2 decoder for text prompts
+through the shared dense engine (`zig/src/cuda_dense/`), from the BF16
+`dots-studio/dots.mocr` checkpoint at revision
+`e539fbb52280393adc081b289ec597430a0f9031`. It is qualified on a GB10
+(sm_121) only. Images are not served on CUDA yet: `--vision` is refused
+by name. The HIP engine, with the image tower, is in
+[dots.mocr on native HIP](dots-mocr-hip.md).
+
+```sh
+zig build native -Dnvcc=/usr/local/cuda/bin/nvcc
+zig-out/native/bin/tensorfold-native serve MODEL_DIRECTORY \
+ --context 8192 --parallel 8 --name dots --host 127.0.0.1 --port 8090
+```
+
+Replies carry no reasoning block: all text is answer content. The engine
+keeps no prompt states, so `--prompt-cache-gib` is refused; so are
+`--drafter`, `--slide`, `--kv-cache fp8` and `--prompt-precision int8`.
+The engine's own copy drafts and `--no-drafts` give the same tokens.
+
+| Responsibility | Files |
+|---|---|
+| Lanes, pages, attention, draws, the host's declarations | `zig/src/cuda_dense/` |
+| Geometry, weights, one layer's arithmetic | `zig/src/families/dots_mocr/cuda_recipe.zig`, `cuda_weights.zig` |
+| q, k and v biases with the rotary in Torch's rounding | `zig/kernels/cuda/dots_rope.cuh` (`tf_dots_qkv_rope`) |
+
+The text reference runs the pinned oracle's `text` mode on any HIP or CUDA
+GPU, and `score` classifies every native id by HF's teacher-forced logits
+(exact, one BF16 step, or a failure, which exits nonzero):
+
+```sh
+python3 tools/zig/dots_mocr_reference.py text MODEL_DIRECTORY REFERENCE \
+ --prompt-file PROMPT.txt --pci PCI_ADDRESS
+python3 tools/zig/dots_mocr_reference.py score MODEL_DIRECTORY REFERENCE \
+ NATIVE_IDS.json DIFFERENCES.json
+```
diff --git a/docs/recipes/dots-mocr-hip.md b/docs/recipes/dots-mocr-hip.md
index 221ff743ea7..68d5701c9e9 100644
--- a/docs/recipes/dots-mocr-hip.md
+++ b/docs/recipes/dots-mocr-hip.md
@@ -3,7 +3,8 @@
This implementation serves the BF16 `dots-studio/dots.mocr`
checkpoint at revision `e539fbb52280393adc081b289ec597430a0f9031`.
It runs the image tower and Qwen2 decoder in native Zig/HIP. Python is
-used only for the validation reference. CUDA and Metal are not qualified.
+used only for the validation reference. CUDA serves the text decoder only
+([dots.mocr on native CUDA](dots-mocr-cuda.md)); Metal is not qualified.
Build for the local GPU architecture. The checked host uses `gfx1201`:
diff --git a/tools/zig/dots_mocr_reference.py b/tools/zig/dots_mocr_reference.py
index 6af268c9734..4c0038cfdb9 100644
--- a/tools/zig/dots_mocr_reference.py
+++ b/tools/zig/dots_mocr_reference.py
@@ -14,6 +14,7 @@
REVISION = "e539fbb52280393adc081b289ec597430a0f9031"
ORACLE_VERSIONS = {"transformers": "4.57.6"}
+ROPE_POSITIONS = [0, 1, 4096, 131071] # the rotary values the native kernel's contract compares, short and long
PROCESSOR_VERSIONS = {"transformers": "4.51.0", "pillow": "11.2.1", "numpy": "2.2.5"}
# The revision's non-weight sources: code, configuration, tokenizer, template, processor and the model card.
SOURCES = {
@@ -67,8 +68,11 @@ def example_prompt(checkpoint):
def selected_pci(torch):
- if torch.version.hip is None or not torch.cuda.is_available() or torch.cuda.device_count() != 1:
- raise RuntimeError("the oracle requires exactly one visible HIP GPU")
+ if not torch.cuda.is_available() or torch.cuda.device_count() != 1:
+ raise RuntimeError("the oracle requires exactly one visible GPU")
+ if torch.version.hip is None:
+ p = torch.cuda.get_device_properties(0)
+ return f"{p.pci_domain_id:04x}:{p.pci_bus_id:02x}:{p.pci_device_id:02x}.0"
lib = ctypes.CDLL("libamdhip64.so")
function = lib.hipDeviceGetPCIBusId
function.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int]
@@ -79,6 +83,40 @@ def selected_pci(torch):
return text.value.decode().lower()
+class Scores:
+ """Retain the scores used for greedy selection without a growing device tensor history (a LogitsProcessor)."""
+
+ def __init__(self, path, vocab):
+ self.file = path.open("w+b")
+ self.vocab, self.count = vocab, 0
+ self.header()
+
+ def header(self):
+ text = repr({"descr": " std.math.maxInt(i32) or k > std.math.maxInt(i32) or n > std.math.maxInt(i32))
+ pub fn bytes(k: usize, n: usize, pairs: usize, tile: usize, sk: usize, mats: usize) !usize {
+ _ = try tileBlocks(tile);
+ if (pairs > std.math.maxInt(i32) or k > std.math.maxInt(i32) or n > std.math.maxInt(i32))
return error.Exl3Shape;
const x = try std.math.mul(usize, k, 4);
- const z = try std.math.mul(usize, try std.math.mul(usize, n, sk * mats), 4);
- return std.math.mul(usize, p, try std.math.add(usize, x, z));
+ const z = try std.math.mul(usize, try std.math.mul(usize, n, try std.math.mul(usize, sk, mats)), 4);
+ return std.math.mul(usize, pairs, try std.math.add(usize, x, z));
}
- pub fn run(g: *Grouped, s: Stream, epi: experts.Epi, in: experts.Rows, l: experts.Layer, p: grouped.Plan, out: u64, items: usize, skip: i32, _: bool) !void {
+ pub fn run(g: *Grouped, s: Stream, epi: experts.Epi, in: experts.Rows, l: experts.Layer, p: grouped.Plan, out: u64, items: usize, pairs: usize, tile: usize, skip: i32) !void {
+ const mt = try tileBlocks(tile);
if (epi == .up and g.up_contract != .bf16) return error.UnsupportedEpilogue;
const w = l.exl3 orelse return error.BadExperts;
if (w.dims != l.dims or w.width != l.width or w.experts != l.experts or l.relu2 or
@@ -66,17 +67,17 @@ pub const Grouped = struct {
const up = epi == .up;
const k = if (up) w.dims else w.width;
const n = if (up) w.width else w.dims;
- if (k == 0 or n == 0 or k % 128 != 0 or n % 128 != 0 or in.stride < k) return error.Exl3Shape;
+ if (k == 0 or n == 0 or k % 128 != 0 or n % 128 != 0 or in.stride < k or items > std.math.maxInt(i32))
+ return error.Exl3Shape;
const sk = try ex.plan(k, n, up);
const mats: usize = if (up) 2 else 1;
- const size = try bytes(k, n, items, sk, mats);
- if (items == 0) return;
+ const size = try bytes(k, n, pairs, tile, sk, mats);
+ if (items == 0 or pairs == 0) return;
if (g.storage == null or g.storage.?.len < size) {
const replacement = try memory.DeviceBuffer.alloc(g.d, size);
if (g.storage) |*old| old.free();
g.storage = replacement;
}
- const pairs = items * 64;
const x0 = g.storage.?.ptr;
const x1 = x0 + pairs * k * 2;
const z = x1 + pairs * k * 2;
@@ -90,16 +91,25 @@ pub const Grouped = struct {
try launch.launch(g.kernels.plan_rotate, .{ .grid = .{ .x = @intCast(items), .y = @intCast(k / 128), .z = @intCast(mats) }, .block = .{ .x = 32 } }, s, &rot);
var args: launch.Args = .{};
for ([_]u64{ x0, x1, if (up) w.gate else w.down, if (up) w.up else w.down, if (up) w.gate_k2 else w.down_k2, if (up) w.up_k2 else w.down_k2, p.items, p.counts, p.members, z }) |v| args.add(v);
- for ([_]usize{ k, n, pairs, sk, 64, in.slots, items, mats, if (up) w.gu_stride else 0, 1 }) |v|
+ for ([_]usize{ k, n, pairs, sk, tile, in.slots, items, mats, if (up) w.gu_stride else 0, 1 }) |v|
args.add(@as(i32, @intCast(v)));
args.add(skip);
- try launch.launch(g.kernels.grouped[@backingInt(w.codebook)], .{ .grid = .{ .x = @intCast(items), .y = @intCast(n / 128), .z = @intCast(mats * sk * 4) }, .block = .{ .x = 128 } }, s, &args);
+ try launch.launch(g.kernels.grouped[@backingInt(w.codebook)], .{ .grid = .{ .x = @intCast(items), .y = @intCast(n / 128), .z = @intCast(mats * sk * mt) }, .block = .{ .x = 128 } }, s, &args);
var finish: launch.Args = .{};
finish.add(a);
try launch.launch(if (l.after) g.kernels.plan_finish_after else g.kernels.plan_finish, .{ .grid = .{ .x = @intCast(items), .y = @intCast(n / 128) }, .block = .{ .x = 32 } }, s, &finish);
}
};
+/// Member-tile blocks a grouped launch spans: 16-member decode items or 64-member prompt items.
+fn tileBlocks(tile: usize) !usize {
+ return switch (tile) {
+ 16 => 1,
+ 64 => 4,
+ else => error.TileNotBuilt,
+ };
+}
+
/// The activations plan_finish runs (0 bf16, 1 fp32, 2 SwiGLU-OAI); SwiGLU-OAI needs its clamp and alpha.
fn checkAct(w: ex.Layer, after: bool) !void {
if (w.act_mode < 0 or w.act_mode > 2) return error.BadExperts;
@@ -135,16 +145,96 @@ test "Grouped takes the clamp-after order at fp32 activations only: not at bf16,
try std.testing.expectError(error.OptionNotBuilt, checkAct(w, true));
}
-test "Grouped scratch includes every pair and rejects overflowing geometry" {
- try std.testing.expectEqual(@as(usize, 64 * (128 * 4 + 256 * 4 * 2 * 4)), try Grouped.bytes(128, 256, 1, 4, 2));
- try std.testing.expectError(error.Exl3Shape, Grouped.bytes(std.math.maxInt(usize), 256, 1, 4, 2));
- try std.testing.expectError(error.Exl3Shape, Grouped.bytes(128, 256, std.math.maxInt(i32), 4, 2));
+test "Grouped scratch scales linearly with pairs, refuses unsupported tiles, and rejects overflowing geometry" {
+ const one = try Grouped.bytes(2048, 512, 1, 16, 4, 2);
+ try std.testing.expectEqual(@as(usize, 2048 * 4 + 512 * 4 * 2 * 4), one);
+ try std.testing.expectEqual(144 * one, try Grouped.bytes(2048, 512, 144, 16, 4, 2));
+ try std.testing.expectEqual(144 * one, try Grouped.bytes(2048, 512, 144, 64, 4, 2));
+ for ([_]usize{ 0, 8, 15, 17, 32, 48, 128 }) |bad| {
+ try std.testing.expectError(error.TileNotBuilt, Grouped.bytes(2048, 512, 144, bad, 4, 2));
+ }
+ try std.testing.expectError(error.Exl3Shape, Grouped.bytes(std.math.maxInt(usize), 256, 1, 16, 4, 2));
+ try std.testing.expectError(error.Exl3Shape, Grouped.bytes(128, 256, @as(usize, std.math.maxInt(i32)) + 1, 16, 4, 2));
}
test "Grouped's bf16 up refuses a call that asks for the combined fp32 activation pipeline" {
var g: Grouped = .{ .d = undefined, .kernels = undefined, .up_contract = .combined_fp32 };
const layer: experts.Layer = .{ .format = .exl3, .up = 0, .down = 0, .dims = 128, .width = 128, .experts = 1 };
- try std.testing.expectError(error.UnsupportedEpilogue, g.run(undefined, .up, undefined, layer, undefined, 0, 0, -1, false));
+ for ([_]usize{ 0, 8, 32, 128 }) |bad| {
+ try std.testing.expectError(error.TileNotBuilt, g.run(undefined, .up, undefined, layer, undefined, 0, 1, 1, bad, -1));
+ }
+ try std.testing.expectError(error.UnsupportedEpilogue, g.run(undefined, .up, undefined, layer, undefined, 0, 0, 0, 16, -1));
g.up_contract = .bf16;
- try std.testing.expectError(error.BadExperts, g.run(undefined, .up, undefined, layer, undefined, 0, 0, -1, false));
+ try std.testing.expectError(error.BadExperts, g.run(undefined, .up, undefined, layer, undefined, 0, 0, 0, 16, -1));
+}
+
+test "Grouped.run launches the plan tile's grid and compact pair stride, dropping no member at tile 64" {
+ const abi = @import("abi.zig");
+ const Trace = struct {
+ var n: usize = 0;
+ var grids: [3][3]u32 = undefined;
+ var rot: Stage = undefined;
+ var fin: Stage = undefined;
+ var ptrs: [10]u64 = undefined;
+ var ints: [11]i32 = undefined;
+ fn launchFn(_: abi.Function, gx: c_uint, gy: c_uint, gz: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: abi.Stream, args: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) abi.Result {
+ grids[n] = .{ gx, gy, gz };
+ if (n == 0) rot = @as(*const Stage, @ptrCast(@alignCast(args.?[0].?))).*;
+ if (n == 1) {
+ for (&ptrs, 0..) |*p, i| p.* = @as(*const u64, @ptrCast(@alignCast(args.?[i].?))).*;
+ for (&ints, 0..) |*v, i| v.* = @as(*const i32, @ptrCast(@alignCast(args.?[10 + i].?))).*;
+ }
+ if (n == 2) fin = @as(*const Stage, @ptrCast(@alignCast(args.?[0].?))).*;
+ n += 1;
+ return 0;
+ }
+ };
+ var d: Driver = undefined;
+ d.api.cuLaunchKernel = Trace.launchFn;
+ var ks: ex.Experts = undefined;
+ ks.plan_rotate = .{ .d = &d, .handle = @ptrFromInt(1) };
+ ks.plan_finish = .{ .d = &d, .handle = @ptrFromInt(2) };
+ ks.plan_finish_after = .{ .d = &d, .handle = @ptrFromInt(3) };
+ for (&ks.grouped, 0..) |*f, i| f.* = .{ .d = &d, .handle = @ptrFromInt(10 + i) };
+ const base: u64 = 0x4000_0000;
+ var g: Grouped = .{ .d = &d, .kernels = &ks, .storage = .{ .d = &d, .ptr = base, .len = 1 << 28 } };
+ const w: ex.Layer = .{ .gate = 0x10, .up = 0x20, .down = 0x30, .gate_k2 = 0x40, .up_k2 = 0x50, .down_k2 = 0x60, .suh_g = 1, .suh_u = 2, .suh_d = 3, .svh_g = 4, .svh_u = 5, .svh_d = 6, .dims = 2048, .width = 512, .experts = 257, .codebook = .mul1 };
+ const l: experts.Layer = .{ .format = .exl3, .exl3 = w, .up = 0, .down = 0, .dims = 2048, .width = 512, .experts = 257 };
+ const pl: grouped.Plan = .{ .members = 0x100, .items = 0x200, .counts = 0x300, .rank = 0, .hist = 0 };
+ const s: Stream = .{ .d = &d, .handle = null };
+ const grp: experts.Grouped = .{ .exl3 = &g };
+ for ([_]struct { tile: usize, pairs: usize, items: usize, epi: experts.Epi, prompt: bool }{
+ .{ .tile = 16, .pairs = 9, .items = 9, .epi = .up, .prompt = false },
+ .{ .tile = 16, .pairs = 144, .items = 153, .epi = .down_f32, .prompt = false },
+ .{ .tile = 64, .pairs = 512, .items = 136, .epi = .up, .prompt = false },
+ .{ .tile = 64, .pairs = 576, .items = 266, .epi = .down_bf16, .prompt = true },
+ }) |c| {
+ Trace.n = 0;
+ const up = c.epi == .up;
+ const k: usize = if (up) 2048 else 512;
+ const n: usize = if (up) 512 else 2048;
+ const sk: usize = try ex.plan(k, n, up);
+ const mats: usize = if (up) 2 else 1;
+ const in: experts.Rows = .{ .x = 0x9000, .stride = k, .slots = if (up) 9 else 0 };
+ if (c.prompt) try grp.prompt(s, c.epi, in, l, pl, 0xa000, c.items, c.pairs, c.tile, -1) else try grp.decode(s, c.epi, in, l, pl, 0xa000, c.items, c.pairs, c.tile, -1);
+ try std.testing.expectEqual(@as(usize, 3), Trace.n);
+ try std.testing.expectEqual(@as(u32, @intCast(mats * sk * (c.tile / 16))), Trace.grids[1][2]);
+ try std.testing.expectEqual(@as(i32, @intCast(c.pairs)), Trace.rot.p);
+ try std.testing.expectEqual(@as(i32, @intCast(c.pairs)), Trace.fin.p);
+ try std.testing.expectEqual(@as(i32, @intCast(c.pairs)), Trace.ints[2]);
+ try std.testing.expectEqual(@as(i32, @intCast(c.tile)), Trace.ints[4]);
+ try std.testing.expectEqual(base + c.pairs * k * 2, Trace.rot.x1);
+ try std.testing.expectEqual(base + c.pairs * k * 4, Trace.rot.z);
+ try std.testing.expectEqual(Trace.rot.x1, Trace.ptrs[1]);
+ try std.testing.expectEqual(Trace.rot.z, Trace.ptrs[9]);
+ if (c.tile == 64 and up) {
+ var seen: [64]usize = @splat(0);
+ const mt: u32 = @intCast((@as(usize, @intCast(Trace.ints[4])) + 15) / 16);
+ for (0..Trace.grids[1][2]) |zz| {
+ const mtile = zz % mt;
+ for (0..16) |tx| seen[mtile * 16 + tx] += 1;
+ }
+ for (seen) |hits| try std.testing.expectEqual(mats * sk, hits);
+ }
+ }
}
diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig
index 3c5828bea36..3d70527c1d8 100644
--- a/zig/src/cuda/experts.zig
+++ b/zig/src/cuda/experts.zig
@@ -53,11 +53,11 @@ pub const Grouped = struct {
nvfp4_csf0: ?*const @import("experts32.zig").Nvfp4b32Experts = null,
iq2xxs: ?*const Iq2xxsExperts = null,
- /// Decode rows: `items` bounds the plan's items (grouped.maxItems at tile 16); `skip` an expert left out, or -1.
- pub fn decode(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void {
+ /// Decode rows: `items` bounds the plan, `pairs` its picks; `tile` 16 or 64 (EXL3 only); `skip` an expert or -1.
+ pub fn decode(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, pairs: usize, tile: usize, skip: i32) !void {
if (l.after and l.format != .exl3) return error.OptionNotBuilt; // the clamp-after order is built for EXL3 only
switch (l.format) {
- .exl3 => return (g.exl3 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip, false),
+ .exl3 => return (g.exl3 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, pairs, tile, skip),
.affine4 => {
const a = g.affine4 orelse return error.FormatNotLoaded;
if (epi == .down_bf16) return error.EpilogueNotBuilt;
@@ -80,11 +80,11 @@ pub const Grouped = struct {
}
}
- /// Prompt rows: `items` bounds the plan's items at the format's prompt tile (prompt_tile).
- pub fn prompt(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void {
+ /// Prompt rows: `items` bounds the plan, `pairs` its picks; `tile` 16 or 64 (EXL3 only); `skip` an expert or -1.
+ pub fn prompt(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, pairs: usize, tile: usize, skip: i32) !void {
if (l.after and l.format != .exl3) return error.OptionNotBuilt; // the clamp-after order is built for EXL3 only
switch (l.format) {
- .exl3 => return (g.exl3 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip, true),
+ .exl3 => return (g.exl3 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, pairs, tile, skip),
.affine4 => {
const a = g.affine4 orelse return error.FormatNotLoaded;
if (epi == .down_f32) return error.EpilogueNotBuilt;
@@ -556,8 +556,8 @@ test "nvfp4 rows of either kind are refused when the format is not loaded" {
const l: Layer = .{ .format = .nvfp4, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2, .relu2 = true };
const s: Stream = undefined; // never reached: the refusal comes first
const p: Plan = .{ .members = 0, .items = 0, .counts = 0, .rank = 0, .hist = 0 };
- try std.testing.expectError(error.FormatNotLoaded, g.prompt(s, .up, .{ .x = 0, .stride = 64 }, l, p, 0, 1, -1));
- try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .down_bf16, .{ .x = 0, .stride = 64 }, l, p, 0, 1, -1));
+ try std.testing.expectError(error.FormatNotLoaded, g.prompt(s, .up, .{ .x = 0, .stride = 64 }, l, p, 0, 1, 1, 64, -1));
+ try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .down_bf16, .{ .x = 0, .stride = 64 }, l, p, 0, 1, 1, 16, -1));
try std.testing.expectEqual(@as(usize, 2 * 2 * 3 * 144 * 4), Nvfp4Experts.bytes(2, 64, 96, 1));
try std.testing.expectEqual(@as(u32, 2 * 17408), Nvfp4Experts.promptSmem(1, 8)); // 64 rows x 64 inputs, 8 x 2 blocks
try std.testing.expectEqual(@as(u32, 2 * 17408), Nvfp4Experts.promptSmem(2, 4));
diff --git a/zig/src/cuda/experts_host_test.zig b/zig/src/cuda/experts_host_test.zig
index 1b645b64f84..635eea49d2f 100644
--- a/zig/src/cuda/experts_host_test.zig
+++ b/zig/src/cuda/experts_host_test.zig
@@ -12,15 +12,15 @@ test "a format that is not loaded, or an option it was not built with, is refuse
const l: Layer = .{ .format = .fp8g, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2 };
const s: Stream = undefined; // never reached: the refusal comes first
const p: Plan = .{ .members = 0, .items = 0, .counts = 0, .rank = 0, .hist = 0 };
- try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .up, .{ .x = 0, .stride = 64 }, l, p, 0, 1, -1));
+ try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .up, .{ .x = 0, .stride = 64 }, l, p, 0, 1, 1, 16, -1));
const a4: experts.Affine4Experts = undefined;
const h: Grouped = .{ .affine4 = &a4 };
var la = l;
la.format = .affine4;
- try std.testing.expectError(error.EpilogueNotBuilt, h.decode(s, .down_bf16, .{ .x = 0, .stride = 64 }, la, p, 0, 1, -1));
- try std.testing.expectError(error.EpilogueNotBuilt, h.prompt(s, .down_f32, .{ .x = 0, .stride = 64 }, la, p, 0, 1, -1));
- try std.testing.expectError(error.OptionNotBuilt, h.decode(s, .up, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 3));
- try std.testing.expectError(error.OptionNotBuilt, h.prompt(s, .up, .{ .x = 0, .stride = 64 }, .{ .format = .affine4, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2, .after = true }, p, 0, 1, -1));
+ try std.testing.expectError(error.EpilogueNotBuilt, h.decode(s, .down_bf16, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 1, 16, -1));
+ try std.testing.expectError(error.EpilogueNotBuilt, h.prompt(s, .down_f32, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 1, 64, -1));
+ try std.testing.expectError(error.OptionNotBuilt, h.decode(s, .up, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 1, 16, 3));
+ try std.testing.expectError(error.OptionNotBuilt, h.prompt(s, .up, .{ .x = 0, .stride = 64 }, .{ .format = .affine4, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2, .after = true }, p, 0, 1, 1, 64, -1));
}
test "the clamp-after order: exl3 takes it on both paths, every other format refuses it on both" {
@@ -30,11 +30,11 @@ test "the clamp-after order: exl3 takes it on both paths, every other format ref
const in: experts.Rows = .{ .x = 0, .stride = 64 };
var l: Layer = .{ .format = .exl3, .up = 0, .down = 0, .width = 64, .dims = 64, .experts = 2, .limit = 10, .after = true };
// exl3 accepts the order, then finds its kernels not loaded
- try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .up, in, l, p, 0, 1, -1));
- try std.testing.expectError(error.FormatNotLoaded, g.prompt(s, .up, in, l, p, 0, 1, -1));
+ try std.testing.expectError(error.FormatNotLoaded, g.decode(s, .up, in, l, p, 0, 1, 1, 16, -1));
+ try std.testing.expectError(error.FormatNotLoaded, g.prompt(s, .up, in, l, p, 0, 1, 1, 64, -1));
for (std.enums.values(experts.Format)) |f| if (f != .exl3) {
l.format = f; // nvfp4 among them; not loaded either, so the refusal comes before the format's
- try std.testing.expectError(error.OptionNotBuilt, g.decode(s, .up, in, l, p, 0, 1, -1));
- try std.testing.expectError(error.OptionNotBuilt, g.prompt(s, .up, in, l, p, 0, 1, -1));
+ try std.testing.expectError(error.OptionNotBuilt, g.decode(s, .up, in, l, p, 0, 1, 1, 16, -1));
+ try std.testing.expectError(error.OptionNotBuilt, g.prompt(s, .up, in, l, p, 0, 1, 1, 64, -1));
};
}
diff --git a/zig/src/cuda_dense/backend.zig b/zig/src/cuda_dense/backend.zig
new file mode 100644
index 00000000000..7a1e5ba3be0
--- /dev/null
+++ b/zig/src/cuda_dense/backend.zig
@@ -0,0 +1,350 @@
+//! A dense decoder behind lanes.Backend: a slot and pool pages a stream, a chunk a prompt step, one forward a round.
+const std = @import("std");
+const cuda = @import("cuda");
+const lanes = @import("lanes");
+const kv = @import("core").kv_pool;
+const pass = @import("pass.zig");
+const mdl = @import("model.zig");
+const kern = @import("kernels.zig");
+const pages = @import("pages.zig");
+const draws_ = @import("draws.zig");
+
+const Draws = draws_.Draws;
+const ring = draws_.ring;
+const be = lanes.backend;
+const Stream = lanes.Stream;
+
+/// Prompt rows a step takes: the forward's chunk grid, from position 0.
+pub const prompt_rows = mdl.prompt_rows;
+/// Window rows a round takes over every stream.
+pub const round_rows = mdl.round_rows;
+/// The widest window a stream verifies: one-row steps and every wider window give a row the same bits.
+pub const max_window = pages.window;
+
+const Lane = struct {
+ slot: usize,
+ len: u64 = 0, // the cache length: positions [0, len) hold this stream's keys
+ base: u64 = 0, // the cache length the last verify started from
+ rows: std.ArrayList(u32) = .empty, // the last verify's row tokens
+ last: u64, // the prompt's last row's bf16 logits [V] (the first draw reads them)
+ rule: u64, // sample.cu's Rule, when the stream samples
+ sampled: bool = false,
+ fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts
+ lease: kv.Lease = .{}, // the pages it holds and is promised
+};
+
+pub fn Dense(comptime R: type) type {
+ return struct {
+ const D = @This();
+ gpa: std.mem.Allocator,
+ d: *const cuda.Driver,
+ e: pass.Engine(R),
+ m: *mdl.Model,
+ lanes: std.AutoHashMapUnmanaged(*const Stream, Lane) = .empty,
+ free: []bool, // slot i is free
+ draws: Draws,
+ slot_bufs: []cuda.DeviceBuffer, // a slot's last-row logits, then its rule: two a slot
+ pool: ?kv.Pool = null, // the pages every slot leases (attach)
+
+ /// The device bytes `init` allocates: the draws and a round's slots' last rows and rules.
+ pub fn openBytes(vocab: usize) usize {
+ return Draws.openBytes() + round_rows * (vocab * 2 + @sizeOf(cuda.sample.Rule));
+ }
+
+ pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *mdl.Model, k: *const kern.Kernels, own: *const R.Kernels, w: *const R.Weights, s: cuda.Stream) !D {
+ const free = try gpa.alloc(bool, round_rows);
+ errdefer gpa.free(free);
+ @memset(free, true);
+ var draws = try Draws.init(gpa, d, k, s, m.g.vocab);
+ errdefer draws.deinit(gpa);
+ const bufs = try gpa.alloc(cuda.DeviceBuffer, 2 * round_rows);
+ errdefer gpa.free(bufs);
+ var made: usize = 0;
+ errdefer for (bufs[0..made]) |*b| b.free();
+ for (bufs, 0..) |*b, i| {
+ b.* = try cuda.DeviceBuffer.alloc(d, if (i % 2 == 0) m.g.vocab * 2 else @sizeOf(cuda.sample.Rule));
+ made += 1;
+ }
+ return .{ .gpa = gpa, .d = d, .e = .{ .m = m, .k = k, .own = own, .w = w, .s = s }, .m = m, .free = free, .draws = draws, .slot_bufs = bufs };
+ }
+
+ pub fn deinit(x: *D) void {
+ var it = x.lanes.valueIterator();
+ while (it.next()) |l| {
+ if (x.pool) |*p| p.finish(&l.lease);
+ l.rows.deinit(x.gpa);
+ }
+ if (x.pool) |*p| p.deinit();
+ x.lanes.deinit(x.gpa);
+ for (x.slot_bufs) |*b| b.free();
+ x.gpa.free(x.slot_bufs);
+ x.gpa.free(x.free);
+ x.draws.deinit(x.gpa);
+ }
+
+ /// The model's arena of `n` pages for `extra + 1` slots, and the pool that leases them.
+ pub fn attach(x: *D, extra: usize, n: usize) !void {
+ if (x.pool != null) return error.BadAttach;
+ try x.m.attach(x.e.s, extra, n);
+ x.pool = try kv.Pool.init(x.gpa, n, x.m.pageBytes());
+ }
+
+ /// The pool's free pages in a seeded random order, leased in that order (a check's hook).
+ pub fn shufflePages(x: *D, seed: u64) void {
+ var prng = std.Random.DefaultPrng.init(seed);
+ if (x.pool) |*p| prng.random().shuffle(u32, p.free[0..p.nfree]);
+ }
+
+ pub fn backend(x: *D) be.Backend {
+ return .{ .ptr = x, .vtable = &vtable };
+ }
+
+ const vtable: be.Backend.VTable = .{ .prefill = prefill, .prefill_step = prefillStep, .cancel_prefill = cancelPrefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .reports_fill = true };
+
+ /// The round loop's facts: no draft head; host copies in windows up to 16 rows, streams sharing a forward.
+ pub fn facts(_: *const D) lanes.Model {
+ return .{ .exact_width = max_window, .gpu_tokens = false, .mtp = false, .speculate = false, .hidden_rows = true, .batch_rows = round_rows, .max_streams = round_rows };
+ }
+
+ fn self(ptr: *anyopaque) *D {
+ return @ptrCast(@alignCast(ptr));
+ }
+
+ fn forward(x: *D, chains: []const pass.Chain, prompt: bool, head0: usize, heads: usize) !void {
+ try pass.forward(R, x.e, chains, prompt, head0, heads);
+ }
+
+ fn freeSlot(x: *const D) ?usize {
+ return std.mem.indexOfScalar(bool, x.free[0..x.m.slots], true);
+ }
+
+ /// Whether a request could start now: a free slot and its pages; one that never fits passes (refused later).
+ pub fn admits(x: *const D, prompt: usize, max_new: usize) bool {
+ if (prompt + max_new > x.m.context) return true;
+ const p = x.pool orelse return false;
+ return x.freeSlot() != null and p.fits(pages.pagesFor(prompt, max_new));
+ }
+
+ const Map = struct {
+ x: *D,
+ slot: usize,
+ fn write(ctx: *anyopaque, ids: []const u32, from: usize) anyerror!void {
+ const m: *const Map = @ptrCast(@alignCast(ctx));
+ try m.x.m.mapPages(m.x.e.s, m.slot, ids[from..], from);
+ }
+ };
+
+ /// The lane's pages cover positions [0, end): pages taken from its promise as its context grows.
+ fn reach(x: *D, l: *Lane, end: usize) !void {
+ const p = if (x.pool) |*p| p else return error.NoPages;
+ const n = pages.pagesOf(end);
+ if (n <= l.lease.ids.items.len) return;
+ var map: Map = .{ .x = x, .slot = l.slot };
+ try p.cover(&l.lease, n, .{ .ctx = &map, .write = Map.write });
+ }
+
+ /// A free slot for the request, its pages promised now (taken as its rows need them).
+ fn takeSlot(x: *D, lease: *kv.Lease, prompt: usize, max_new: usize) !usize {
+ const p = if (x.pool) |*p| p else return error.NoPages;
+ const slot = x.freeSlot() orelse return error.NoFreeSlot;
+ try p.promise(lease, pages.pagesFor(prompt, max_new));
+ return slot;
+ }
+
+ /// A prompt pass begins: a slot, its pages, the stream's rule; no prompt cache, so its first chunk starts at 0.
+ fn opening(x: *D, s: *Stream) !*Lane {
+ const p = s.prompt();
+ if (p.len == 0) return error.EmptyPrompt;
+ if (s.images.len > 0) return error.DenseTextOnly;
+ if (p.len + @as(usize, s.max_new) > x.m.context) return error.PastContext;
+ const got = try x.lanes.getOrPut(x.gpa, s);
+ if (!got.found_existing) {
+ var lease: kv.Lease = .{};
+ const slot = x.takeSlot(&lease, p.len, s.max_new) catch |e| {
+ x.lanes.removeByPtr(got.key_ptr);
+ return e;
+ };
+ x.free[slot] = false;
+ got.value_ptr.* = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr, .lease = lease };
+ }
+ const l = got.value_ptr;
+ try x.reach(l, p.len);
+ l.sampled = false;
+ if (s.sampling) |sp| if (sp.temperature > 0) {
+ const r = cuda.sample.Rule.from(lanes.gpu_rule.rule(sp));
+ try mdl.put(x.d, x.e.s, l.rule, std.mem.asBytes(&r));
+ l.sampled = true;
+ };
+ s.cached = 0;
+ if (s.reuse.saved != null) s.reuse_failed = true; // no prompt cache: every pass starts at 0
+ l.fill = 0;
+ s.fill_at.store(0, .monotonic);
+ return l;
+ }
+
+ /// One chunk from `l.fill` to the next grid line or mark; true once the prompt's last row's logits are written.
+ fn chunk(x: *D, s: *Stream, l: *Lane) !bool {
+ const p = s.prompt();
+ const at = l.fill.?;
+ if (s.isCancelled()) {
+ s.keepInterruptedPrefill(@intCast(at));
+ return error.Cancelled;
+ }
+ var end = @min(p.len, (at / prompt_rows + 1) * prompt_rows);
+ for (s.reuse.marks) |mk| if (mk > at and mk < end) {
+ end = mk;
+ };
+ const done = end == p.len;
+ try x.forward(&.{.{ .slot = l.slot, .p0 = at, .tokens = p[at..end] }}, true, end - at - 1, @intFromBool(done));
+ if (done) try x.d.check(x.d.api.cuMemcpyDtoDAsync_v2(l.last, x.m.s.logits, x.m.g.vocab * 2, x.e.s.handle), "cuMemcpyDtoDAsync");
+ s.fill_at.store(@intCast(end), .monotonic);
+ l.fill = if (done) null else end;
+ if (done) l.len = p.len;
+ if (s.reuse.hook) |hk| if (std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(end)) != null) hk.at(hk.ptr, s, @intCast(end));
+ return done;
+ }
+
+ fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void {
+ const x = self(ptr);
+ const l = try x.opening(s);
+ errdefer x.dropLane(s);
+ while (!try x.chunk(s, l)) {}
+ }
+
+ fn prefillStep(ptr: *anyopaque, s: *Stream, _: u32) anyerror!bool {
+ const x = self(ptr);
+ const going = if (x.lanes.getPtr(s)) |l| l.fill != null else false;
+ const l = if (going) x.lanes.getPtr(s).? else try x.opening(s);
+ errdefer x.dropLane(s);
+ return x.chunk(s, l);
+ }
+
+ fn cancelPrefill(ptr: *anyopaque, s: *Stream) !void {
+ const l = self(ptr).lanes.getPtr(s) orelse return;
+ if (l.fill) |at| s.keepInterruptedPrefill(@intCast(at));
+ }
+
+ fn first(ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64 {
+ const x = self(ptr);
+ const l = x.lanes.getPtr(s) orelse return error.UnknownStream;
+ if (position != l.len or l.fill != null) return error.PositionMismatch;
+ const h = x.draws.handle();
+ try x.draws.rows(l.sampled, l.rule, l.last, 0, 1, position, h.at, 0);
+ return h.h;
+ }
+
+ fn value(x: *D, feed: be.Feed) !u32 {
+ return switch (feed) {
+ .value => |v| v,
+ .handle => |h| x.draws.read(h),
+ };
+ }
+
+ fn queue(ptr: *anyopaque, s: *Stream, feed: be.Feed, position: u64) anyerror!u64 {
+ const x = self(ptr);
+ const l = x.lanes.getPtr(s) orelse return error.UnknownStream;
+ const t = try x.value(feed);
+ if (position != l.len + 1) return error.PositionMismatch;
+ try x.reach(l, l.len + 1);
+ try x.forward(&.{.{ .slot = l.slot, .p0 = l.len, .tokens = &.{t} }}, false, 0, 1);
+ l.len += 1;
+ const h = x.draws.handle();
+ try x.draws.rows(l.sampled, l.rule, x.m.s.logits, 0, 1, position, h.at, 0);
+ return h.h;
+ }
+
+ fn read(ptr: *anyopaque, h: u64) anyerror!u32 {
+ return self(ptr).draws.read(h);
+ }
+
+ fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void {
+ const x = self(ptr);
+ if (windows.len > round_rows) return error.RoundTooWide;
+ const chains = try x.gpa.alloc(pass.Chain, windows.len);
+ defer x.gpa.free(chains);
+ var total: usize = 0;
+ for (windows, chains) |w, *c| {
+ if (w.held != 0) return error.NoHeldDrafts; // no draft head: every draft comes from the host
+ if (w.parents != null) return error.ChainsOnly;
+ if (w.masks.len != 0) return error.StructuresUnsupported;
+ const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream;
+ if (w.rows() > max_window) return error.WindowTooWide;
+ l.base = l.len;
+ l.rows.clearRetainingCapacity();
+ try l.rows.append(x.gpa, w.pending);
+ try l.rows.appendSlice(x.gpa, w.tokens);
+ for (w.positions, 0..) |pos, r| if (pos != l.base + r + 1) return error.PositionMismatch;
+ try x.reach(l, l.base + l.rows.items.len);
+ c.* = .{ .slot = l.slot, .p0 = l.base, .tokens = l.rows.items };
+ total += l.rows.items.len;
+ }
+ if (total > round_rows) return error.RoundTooWide;
+ try x.forward(chains, false, 0, total);
+ var at: usize = 0;
+ var sampled = false;
+ for (windows) |w| {
+ const l = x.lanes.getPtr(w.stream).?;
+ try x.draws.pick(l.sampled, x.m.s.logits, at, l.rows.items.len, ring + at);
+ sampled = sampled or l.sampled;
+ at += l.rows.items.len;
+ }
+ if (sampled) try x.draws.finite(ring, total); // one wait a round, before any draw reads a row
+ at = 0;
+ for (windows, 0..) |w, i| {
+ const l = x.lanes.getPtr(w.stream).?;
+ if (l.sampled) try x.draws.draw(l.rule, x.m.s.logits, at, l.rows.items.len, w.positions[0], ring + at, i);
+ at += l.rows.items.len;
+ }
+ var got: [round_rows]u32 = undefined;
+ try x.draws.tokens(ring, got[0..total]);
+ at = 0;
+ for (windows, out) |w, o| {
+ const l = x.lanes.getPtr(w.stream).?;
+ const n = l.rows.items.len;
+ @memcpy(o.sampled[0..n], got[at..][0..n]);
+ @memcpy(o.drafts[0 .. n - 1], l.rows.items[1..]);
+ l.len = l.base + n; // a chain keeps every row unless rolled back
+ at += n;
+ }
+ }
+
+ fn keep(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void {
+ const x = self(ptr);
+ for (windows, paths) |w, path| {
+ const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream;
+ if (path.len == 0 or path.len > l.rows.items.len) return error.EmptyPath;
+ for (path, 0..) |r, j| if (r != j) return error.ChainsOnly; // a chain keeps a prefix of its rows
+ l.len = l.base + path.len; // caches are positional: later rows overwrite what the rejected ones wrote
+ }
+ }
+
+ /// No draft head: absorbing is a no-op, a request for drafts is refused by name (copies come from the host).
+ fn draft(_: *anyopaque, requests: []const be.DraftRequest) anyerror!void {
+ for (requests) |r| if (r.depth != 0 or r.lanes != null) return error.NoDraftHead;
+ }
+
+ fn dropLane(x: *D, s: *const Stream) void {
+ if (x.lanes.fetchRemove(s)) |entry| {
+ var l = entry.value;
+ if (x.pool) |*p| p.finish(&l.lease); // one CUDA stream: a later lease's writes queue behind these reads
+ x.m.unmapPages(l.slot);
+ x.free[l.slot] = true;
+ l.rows.deinit(x.gpa);
+ if (x.lanes.count() == 0) x.report();
+ }
+ }
+
+ /// The idle server's line: nothing promised, every page free.
+ fn report(x: *const D) void {
+ if (@import("builtin").is_test) return;
+ const p = x.pool orelse return;
+ const c = p.census();
+ const ok = p.idle() and c.free == c.total;
+ std.log.info("kv pool idle{s}: promised {d}, free {d} == total {d}", .{ if (ok) "" else " (UNEVEN)", c.promised, c.free, c.total });
+ }
+
+ fn release(ptr: *anyopaque, s: *Stream) void {
+ self(ptr).dropLane(s);
+ }
+ };
+}
diff --git a/zig/src/cuda_dense/device_fail_test.zig b/zig/src/cuda_dense/device_fail_test.zig
new file mode 100644
index 00000000000..7fbc24e48b8
--- /dev/null
+++ b/zig/src/cuda_dense/device_fail_test.zig
@@ -0,0 +1,202 @@
+//! Failure paths on a recording fake device: no buffer freed before completion, no draw over non-finite logits.
+const std = @import("std");
+const cuda = @import("cuda");
+const mdl = @import("model.zig");
+const kern = @import("kernels.zig");
+const draws = @import("draws.zig");
+const Geometry = @import("recipe.zig").Geometry;
+const t = std.testing;
+
+const Fake = struct {
+ const queue: cuda.abi.Stream = @ptrFromInt(0x1800);
+ const base: u64 = 0x100000;
+ var mem: [4096]u8 = undefined; // device bytes from `base` on, for downloads
+ var next: u64 = 0x200000;
+ var live: usize = 0;
+ var pending = false;
+ var freed_pending: usize = 0;
+ var proofs: usize = 0;
+ var fail_wait = false;
+ var argmax: usize = 0;
+ var drawn: usize = 0;
+ fn reset() void {
+ live = 0;
+ pending = false;
+ freed_pending = 0;
+ proofs = 0;
+ fail_wait = false;
+ argmax = 0;
+ drawn = 0;
+ }
+ fn alloc(out: *u64, len: usize) callconv(.c) c_int {
+ out.* = next;
+ next += len + 256;
+ live += 1;
+ return 0;
+ }
+ fn free(_: u64) callconv(.c) c_int {
+ freed_pending += @intFromBool(pending);
+ live -= 1;
+ return 0;
+ }
+ fn fill(_: u64, _: u8, _: usize, _: cuda.abi.Stream) callconv(.c) c_int {
+ pending = true;
+ return 0;
+ }
+ fn host(out: *?*anyopaque, len: usize, _: c_uint) callconv(.c) c_int {
+ out.* = std.c.malloc(len) orelse return cuda.abi.error_not_ready;
+ return 0;
+ }
+ fn hostFree(p: ?*anyopaque) callconv(.c) c_int {
+ std.c.free(p);
+ return 0;
+ }
+ fn copy(_: u64, _: ?*const anyopaque, _: usize, _: cuda.abi.Stream) callconv(.c) c_int {
+ pending = true;
+ return 0;
+ }
+ fn wait(_: cuda.abi.Stream) callconv(.c) c_int {
+ if (fail_wait) return cuda.abi.error_not_ready;
+ pending = false;
+ return 0;
+ }
+ fn context() callconv(.c) c_int {
+ proofs += 1;
+ pending = false;
+ return 0;
+ }
+ fn event(out: *cuda.abi.Event, _: c_uint) callconv(.c) c_int {
+ out.* = @ptrFromInt(1);
+ return 0;
+ }
+ fn record(_: cuda.abi.Event, _: cuda.abi.Stream) callconv(.c) c_int {
+ return 0;
+ }
+ fn query(_: cuda.abi.Event) callconv(.c) c_int {
+ return if (pending) cuda.abi.error_not_ready else 0;
+ }
+ fn destroy(_: cuda.abi.Event) callconv(.c) c_int {
+ return 0;
+ }
+ fn launch(f: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) c_int {
+ argmax += @intFromBool(f == @as(cuda.abi.Function, @ptrFromInt(1)));
+ drawn += @intFromBool(f == @as(cuda.abi.Function, @ptrFromInt(2)));
+ return 0;
+ }
+ fn down(dst: ?*anyopaque, src: u64, len: usize) callconv(.c) c_int {
+ @memcpy(@as([*]u8, @ptrCast(dst.?))[0..len], mem[src - base ..][0..len]);
+ return 0;
+ }
+ fn device(out: *cuda.abi.Device) callconv(.c) c_int {
+ out.* = 0;
+ return 0;
+ }
+ fn attribute(out: *c_int, _: cuda.abi.DeviceAttribute, _: cuda.abi.Device) callconv(.c) c_int {
+ out.* = 0;
+ return 0;
+ }
+ fn driver() cuda.Driver {
+ reset();
+ var d: cuda.Driver = .{ .lib = undefined, .api = undefined };
+ d.api.cuMemAlloc_v2 = alloc;
+ d.api.cuMemFree_v2 = free;
+ d.api.cuMemsetD8Async = fill;
+ d.api.cuMemHostAlloc = host;
+ d.api.cuMemFreeHost = hostFree;
+ d.api.cuMemcpyHtoDAsync_v2 = copy;
+ d.api.cuStreamSynchronize = wait;
+ d.api.cuCtxSynchronize = context;
+ d.api.cuEventCreate = event;
+ d.api.cuEventRecord = record;
+ d.api.cuEventQuery = query;
+ d.api.cuEventDestroy_v2 = destroy;
+ d.api.cuLaunchKernel = launch;
+ d.api.cuMemcpyDtoH_v2 = down;
+ d.api.cuCtxGetDevice = device;
+ d.api.cuDeviceGetAttribute = attribute;
+ return d;
+ }
+};
+
+const tiny: Geometry = .{ .hidden = 256, .layers = 2, .heads = 4, .kv_heads = 2, .head_dim = 64, .vocab = 512, .eps = 1e-6, .context = 4096, .acc = 512, .wide = 1024, .mid = 512, .theta = 1e4 };
+
+fn lane() cuda.qmmf.Lane {
+ var l: cuda.qmmf.Lane = undefined;
+ l.clusters = false; // partBytes reads the tile plan alone
+ return l;
+}
+
+test "a failed scratch upload wait proves completion by the context before init frees its buffers" {
+ var d = Fake.driver();
+ const s: cuda.Stream = .{ .d = &d, .handle = Fake.queue };
+ const l = lane();
+ Fake.fail_wait = true;
+ try t.expectError(error.NotReady, mdl.Model.init(t.allocator, &d, s, tiny, 4096, &l));
+ try t.expectEqual(@as(usize, 0), Fake.freed_pending);
+ try t.expectEqual(@as(usize, 0), Fake.live);
+ try t.expect(Fake.proofs >= 1);
+ Fake.fail_wait = false;
+ try s.synchronize();
+ cuda.copy_source.dropSpares(&d, cuda.copy_source.current());
+}
+
+test "a failed attach wait proves completion before its arena and tables go, and leaves the model attachable" {
+ var d = Fake.driver();
+ const s: cuda.Stream = .{ .d = &d, .handle = Fake.queue };
+ const l = lane();
+ var m = try mdl.Model.init(t.allocator, &d, s, tiny, 4096, &l);
+ const held = Fake.live;
+ Fake.fail_wait = true;
+ try t.expectError(error.NotReady, m.attach(s, 1, 2));
+ try t.expectEqual(@as(usize, 0), Fake.freed_pending);
+ try t.expectEqual(held, Fake.live);
+ try t.expect(m.pages == 0 and m.slots == 0 and m.ids.len == 0 and m.arena == 0);
+ Fake.fail_wait = false;
+ try m.attach(s, 1, 2);
+ try t.expect(m.pages == 2 and m.slots == 2);
+ m.deinit();
+ try t.expectEqual(@as(usize, 0), Fake.live);
+ cuda.copy_source.dropSpares(&d, cuda.copy_source.current());
+}
+
+fn buffer(d: *const cuda.Driver, offset: u64, len: usize) cuda.DeviceBuffer {
+ return .{ .d = d, .ptr = Fake.base + offset, .len = len, .borrowed = true };
+}
+
+fn function(d: *const cuda.Driver, id: usize) cuda.Function {
+ return .{ .d = d, .handle = @ptrFromInt(id) };
+}
+
+test "a sampled row over NaN or an infinity is refused before any draw reads it, as a greedy one is" {
+ var d = Fake.driver();
+ const s: cuda.Stream = .{ .d = &d, .handle = Fake.queue };
+ var k: kern.Kernels = undefined;
+ k.round = .{ .argmax_fn = function(&d, 1), .match_fn = function(&d, 9), .keep_fn = function(&d, 9), .gather_fn = function(&d, 9), .head_fn = function(&d, 9), .ids_fn = function(&d, 9) };
+ k.draws = .{ .draw_fn = function(&d, 2), .draw_ids_fn = function(&d, 9), .logprob_rows_fn = function(&d, 9), .topk_slices_fn = function(&d, 9), .topk_merge_fn = function(&d, 9) };
+ var greedy: [draws.ring + mdl.round_rows]bool = @splat(true);
+ var x: draws.Draws = .{ .d = &d, .k = &k, .s = s, .vocab = 512, .tok = buffer(&d, 0, 64), .picks = buffer(&d, 64, 128), .greedy = &greedy, .meta = buffer(&d, 512, 64) };
+ const Pick = cuda.tree_round.Pick;
+ var picks = [2]Pick{ .{ .token = 3, .nonfinite = 0 }, .{ .token = 4, .nonfinite = 0 } };
+ @memcpy(Fake.mem[64..][0..16], std.mem.sliceAsBytes(&picks)); // what the argmax writes for finite rows
+ @memcpy(Fake.mem[0..8], std.mem.sliceAsBytes(&[2]u32{ 7, 8 }));
+ try x.rows(true, 0x9000, 0xa000, 0, 2, 5, 0, 0);
+ try t.expectEqual(@as(usize, 1), Fake.argmax); // the sampled rows' non-finite flags
+ try t.expectEqual(@as(usize, 1), Fake.drawn);
+ var out: [2]u32 = undefined;
+ try x.tokens(0, &out);
+ try t.expectEqualSlices(u32, &.{ 7, 8 }, &out); // drawn, not picked
+ for ([_]u32{ 1, 2, 4 }) |mask| { // NaN, +inf, -inf (tree_round.cu's mask)
+ picks[1] = .{ .token = 0xFFFFFFFF, .nonfinite = mask };
+ @memcpy(Fake.mem[64..][0..16], std.mem.sliceAsBytes(&picks));
+ try t.expectError(error.NonFiniteLogits, x.rows(true, 0x9000, 0xa000, 0, 2, 5, 0, 0));
+ try t.expectEqual(@as(usize, 1), Fake.drawn); // tf_draw never saw the row
+ }
+ try x.rows(false, 0, 0xa000, 0, 2, 5, 0, 0);
+ try t.expectError(error.NonFiniteLogits, x.tokens(0, &out));
+ try t.expectEqual(@as(usize, 1), Fake.drawn);
+ picks[1] = .{ .token = 4, .nonfinite = 0 };
+ @memcpy(Fake.mem[64..][0..16], std.mem.sliceAsBytes(&picks));
+ try x.tokens(0, &out);
+ try t.expectEqualSlices(u32, &.{ 3, 4 }, &out);
+ cuda.copy_source.dropSpares(&d, cuda.copy_source.current());
+}
diff --git a/zig/src/cuda_dense/draws.zig b/zig/src/cuda_dense/draws.zig
new file mode 100644
index 00000000000..770d2cdd5be
--- /dev/null
+++ b/zig/src/cuda_dense/draws.zig
@@ -0,0 +1,106 @@
+//! The dense engine's draws: a ring of drawn tokens behind handles, greedy picks or keyed samples of bf16 logits.
+const std = @import("std");
+const cuda = @import("cuda");
+const kern = @import("kernels.zig");
+const mdl = @import("model.zig");
+
+const Pick = cuda.tree_round.Pick;
+const round_rows = mdl.round_rows;
+
+/// Drawn tokens a handle names (a ring; read before the round loop could wrap it).
+pub const ring = 1024;
+
+pub const Draws = struct {
+ d: *const cuda.Driver,
+ k: *const kern.Kernels,
+ s: cuda.Stream,
+ vocab: usize,
+ tok: cuda.DeviceBuffer, // drawn tokens (u32): the handle ring, then a round's rows
+ picks: cuda.DeviceBuffer, // argmax picks (Pick), the same entries
+ greedy: []bool, // entry e of the ring (and the round's rows) is a pick, not a draw
+ meta: cuda.DeviceBuffer, // a draw's first position minus one, a window each
+ next: u64 = 0,
+
+ /// The device bytes `init` allocates: the token and pick entries and the metas.
+ pub fn openBytes() usize {
+ return (ring + round_rows) * (4 + @sizeOf(Pick)) + round_rows * 16;
+ }
+
+ pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, k: *const kern.Kernels, s: cuda.Stream, vocab: usize) !Draws {
+ const greedy = try gpa.alloc(bool, ring + round_rows);
+ errdefer gpa.free(greedy);
+ @memset(greedy, true);
+ var tok = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * 4);
+ errdefer tok.free();
+ var picks = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * @sizeOf(Pick));
+ errdefer picks.free();
+ const meta = try cuda.DeviceBuffer.alloc(d, round_rows * 16);
+ return .{ .d = d, .k = k, .s = s, .vocab = vocab, .tok = tok, .picks = picks, .greedy = greedy, .meta = meta };
+ }
+
+ pub fn deinit(x: *Draws, gpa: std.mem.Allocator) void {
+ gpa.free(x.greedy);
+ x.tok.free();
+ x.picks.free();
+ x.meta.free();
+ }
+
+ /// The next handle and its ring entry.
+ pub fn handle(x: *Draws) struct { h: u64, at: usize } {
+ const h = x.next;
+ x.next += 1;
+ return .{ .h = h, .at = @intCast(h % ring) };
+ }
+
+ /// Rows [row0, row0 + rows) of bf16 `logits` into entries [at, at + rows), row r at `position0 + r`.
+ pub fn rows(x: *Draws, sampled: bool, rule: u64, logits: u64, row0: usize, n: usize, position0: u64, at: usize, meta_at: usize) !void {
+ try x.pick(sampled, logits, row0, n, at);
+ if (!sampled) return;
+ try x.finite(at, n);
+ try x.draw(rule, logits, row0, n, position0, at, meta_at);
+ }
+
+ /// Every row's argmax pick into entries [at, at + n): a greedy row's token, and any row's NaN or infinity flag.
+ pub fn pick(x: *Draws, sampled: bool, logits: u64, row0: usize, n: usize, at: usize) !void {
+ const v = x.vocab;
+ @memset(x.greedy[at..][0..n], !sampled);
+ try x.k.round.argmax(x.s, logits + row0 * v * 2, x.picks.ptr + at * @sizeOf(Pick), n, v, v);
+ }
+
+ /// Entries [at, at + n)'s picks read back: a row of NaN or infinite logits is refused before any draw reads it.
+ pub fn finite(x: *Draws, at: usize, n: usize) !void {
+ try x.s.synchronize();
+ var picks: [round_rows]Pick = undefined;
+ try x.picks.download(at * @sizeOf(Pick), std.mem.sliceAsBytes(picks[0..n]));
+ for (picks[0..n]) |p| if (p.nonfinite != 0) return error.NonFiniteLogits;
+ }
+
+ /// Keyed samples of rows whose picks `finite` passed (tf_draw faults on a row with no finite logit).
+ pub fn draw(x: *Draws, rule: u64, logits: u64, row0: usize, n: usize, position0: u64, at: usize, meta_at: usize) !void {
+ const v = x.vocab;
+ const meta = [4]i32{ @intCast(position0 - 1), 0, 0, 0 };
+ const mp = x.meta.ptr + meta_at * 16;
+ try mdl.put(x.d, x.s, mp, std.mem.asBytes(&meta));
+ try x.k.draws.draw(x.s, logits + row0 * v * 2, v, rule, mp, 0, x.tok.ptr + at * 4, n, null, null);
+ }
+
+ /// Entries [at, at + out.len) as tokens after the stream's work; non-finite logits are an error, drawn or picked.
+ pub fn tokens(x: *Draws, at: usize, out: []u32) !void {
+ try x.s.synchronize();
+ var picks: [round_rows]Pick = undefined;
+ try x.tok.download(at * 4, std.mem.sliceAsBytes(out));
+ try x.picks.download(at * @sizeOf(Pick), std.mem.sliceAsBytes(picks[0..out.len]));
+ for (out, picks[0..out.len], x.greedy[at..][0..out.len]) |*t, p, g| {
+ if (p.nonfinite != 0) return error.NonFiniteLogits;
+ if (g) t.* = p.token;
+ }
+ }
+
+ /// The token handle `h` names, once drawn; a handle the ring has passed is refused.
+ pub fn read(x: *Draws, h: u64) !u32 {
+ if (h >= x.next or x.next - h > ring) return error.NoSuchToken;
+ var out: [1]u32 = undefined;
+ try x.tokens(@intCast(h % ring), &out);
+ return out[0];
+ }
+};
diff --git a/zig/src/cuda_dense/kernels.zig b/zig/src/cuda_dense/kernels.zig
new file mode 100644
index 00000000000..32365e5cc62
--- /dev/null
+++ b/zig/src/cuda_dense/kernels.zig
@@ -0,0 +1,145 @@
+//! The dense engine's shared launches: the lane matmul, norms, embedding rows, bias, draws, paged attention, SwiGLU.
+const std = @import("std");
+const cuda = @import("cuda");
+const recipe = @import("recipe.zig");
+const pages = @import("pages.zig");
+
+const Args = cuda.Args;
+const Function = cuda.Function;
+const Stream = cuda.Stream;
+const Geometry = recipe.Geometry;
+
+/// Streams one shared-window attention launch takes (paged_attn.cuh's MAXSEG).
+pub const max_segs = 8;
+
+/// paged_attn.cuh's AttnSegs: each stream's caches, position, page table and first row.
+pub const AttnSegs = extern struct { kc: [max_segs]u64, vc: [max_segs]u64, meta: [max_segs]u64, pt: [max_segs]u64, ps: i64, row0: [max_segs]c_int, n: c_int };
+
+const images = .{ "sample", "tree_round", "prefill_attention", "nemotron_attention", "lfm2" };
+
+fn int(x: usize) c_int {
+ return @intCast(x);
+}
+
+pub const Kernels = struct {
+ mods: [images.len]cuda.Module,
+ g: Geometry,
+ shape: recipe.Shape,
+ lane: cuda.qmmf.Lane,
+ norms: cuda.norms.Norms,
+ embed: cuda.embed_rows.EmbedRows,
+ vit: cuda.vit.Kernels,
+ draws: cuda.sample.Draws,
+ round: cuda.tree_round.Kernels,
+ pattn: Function,
+ kv_write: Function,
+ chunk: Function,
+ merge: Function,
+ silu_mul: Function,
+
+ /// The shared modules, and the attention launches `g`'s shape names; a geometry the kernels lack refuses by name.
+ pub fn load(ctx: *const cuda.Context, g: Geometry) !Kernels {
+ const d = ctx.d;
+ const shape = try recipe.attentionShape(g);
+ if (!cuda.kernels.available) return error.BuiltWithoutKernels;
+ var k: Kernels = undefined;
+ k.g = g;
+ k.shape = shape;
+ var loaded: usize = 0;
+ errdefer for (k.mods[0..loaded]) |*m| m.unload();
+ inline for (images, 0..) |name, i| {
+ k.mods[i] = try cuda.Module.load(d, @field(cuda.kernels, name));
+ loaded += 1;
+ }
+ const m = &k.mods;
+ k.draws = try cuda.sample.Draws.resolve(m[0]);
+ k.round = try cuda.tree_round.Kernels.resolve(m[1]);
+ k.pattn = try m[2].function(shape.pattn);
+ try k.pattn.allowDynamicShared(shape.shared);
+ k.kv_write = try m[3].function("tf_nemo_kv_write_segs");
+ k.chunk = try m[3].function(shape.chunk);
+ k.merge = try m[3].function(shape.merge);
+ k.silu_mul = try m[4].function("tf_lfm2_silu_mul");
+ k.lane = try cuda.qmmf.Lane.load(d, (try ctx.capability()) / 10);
+ errdefer k.lane.unload();
+ k.norms = try cuda.norms.Norms.load(d);
+ errdefer k.norms.unload();
+ k.embed = try cuda.embed_rows.EmbedRows.load(d);
+ errdefer k.embed.unload();
+ k.vit = try cuda.vit.Kernels.load(d);
+ return k;
+ }
+
+ pub fn deinit(k: *Kernels) void {
+ k.vit.unload();
+ k.embed.unload();
+ k.norms.unload();
+ k.lane.unload();
+ for (&k.mods) |*m| m.unload();
+ }
+
+ fn go(f: Function, s: Stream, grid: [3]usize, block: usize, smem: u32, a: *Args) !void {
+ const g: cuda.Dim3 = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) };
+ try cuda.launch.launch(f, .{ .grid = g, .block = .{ .x = @intCast(block) }, .shared = smem }, s, a);
+ }
+
+ /// out [rows, f] = bf16(bf16(silu(gate)) * up) from rows gate | up [rows, 2 f] (lfm2.cu).
+ pub fn siluMul(k: *const Kernels, s: Stream, gu: u64, out: u64, rows: usize, f: usize) !void {
+ var a: Args = .{};
+ a.add(gu);
+ a.add(out);
+ a.add(int(f));
+ try go(k.silu_mul, s, .{ rows, (f + 255) / 256, 1 }, 256, 0, &a);
+ }
+
+ /// Rows k, v (kv [rows, 2 kv width]) into each stream's pages at its position on.
+ pub fn kvWrite(k: *const Kernels, s: Stream, sg: AttnSegs, kv: u64, rows: usize) !void {
+ const kvw = k.g.kvWidth();
+ var a: Args = .{};
+ a.add(sg);
+ a.add(kv);
+ a.add(int(2 * kvw));
+ a.add(@as(c_int, 0));
+ a.add(int(kvw));
+ try go(k.kv_write, s, .{ rows, 1, 1 }, 256, 0, &a);
+ }
+
+ fn scale(k: *const Kernels) f32 {
+ return 1.0 / @sqrt(@as(f32, @floatFromInt(k.g.head_dim)));
+ }
+
+ /// Decode rows: q [rows, heads, dim] over each stream's pages in 512-key chunks, merged in order into out and xs.
+ pub fn attend(k: *const Kernels, s: Stream, sg: AttnSegs, q: u64, po: u64, pm: u64, pl: u64, out: u64, xs: u64, rows: usize, nch: usize, z: usize) !void {
+ const g = k.g;
+ var a: Args = .{};
+ a.add(sg);
+ a.add(q);
+ for ([_]u64{ po, pm, pl }) |v| a.add(v);
+ for ([_]usize{ g.qWidth(), g.heads, g.kv_heads, g.group(), pages.page_rows, nch }) |v| a.add(int(v));
+ a.add(k.scale());
+ try go(k.chunk, s, .{ rows, g.kv_heads, z }, 256, 0, &a);
+ var b: Args = .{};
+ b.add(sg);
+ for ([_]u64{ po, pm, pl, out, xs }) |v| b.add(v);
+ for ([_]usize{ g.heads, g.kv_heads, g.group(), pages.page_rows, nch }) |v| b.add(int(v));
+ try go(k.merge, s, .{ rows, g.kv_heads, 1 }, 256, 0, &b);
+ }
+
+ /// A prompt chunk's rows at positions p0.. against its stream's pages (written through them): out [rows, H, dim].
+ pub fn prompt(k: *const Kernels, s: Stream, q: u64, kc: u64, vc: u64, out: u64, p0: usize, rows: usize, pt: u64, ps: usize) !void {
+ const g = k.g;
+ var a: Args = .{};
+ for ([_]u64{ q, kc, vc, out }) |v| a.add(v);
+ for ([_]usize{ p0, rows, g.heads, g.kv_heads, g.group() }) |v| a.add(int(v));
+ a.add(k.scale());
+ a.add(pt);
+ a.add(@as(i64, @intCast(ps)));
+ try go(k.pattn, s, .{ k.shape.blocks(rows), k.shape.grid_y, 1 }, 256, k.shape.shared, &a);
+ }
+};
+
+test "AttnSegs is paged_attn.cuh's struct, padded to 8 bytes: 304" {
+ try std.testing.expectEqual(@as(usize, 304), @sizeOf(AttnSegs));
+ try std.testing.expectEqual(@as(usize, 256), @offsetOf(AttnSegs, "ps"));
+ try std.testing.expectEqual(@as(usize, 296), @offsetOf(AttnSegs, "n"));
+}
diff --git a/zig/src/cuda_dense/model.zig b/zig/src/cuda_dense/model.zig
new file mode 100644
index 00000000000..06a293223cf
--- /dev/null
+++ b/zig/src/cuda_dense/model.zig
@@ -0,0 +1,251 @@
+//! The dense engine's device state: the forward's scratch, the page arena of every layer's KV, each slot's table.
+const std = @import("std");
+const cuda = @import("cuda");
+const kv = @import("core").kv_pool;
+const recipe = @import("recipe.zig");
+const pages = @import("pages.zig");
+const rope = @import("rope.zig");
+
+const DeviceBuffer = cuda.DeviceBuffer;
+const Geometry = recipe.Geometry;
+
+/// Rows a prompt chunk takes.
+pub const prompt_rows = 1024;
+/// Rows a round of windows takes over every stream.
+pub const round_rows = 128;
+/// Key chunks one decode attention launch walks in parallel at most.
+const chunk_blocks = 64;
+
+/// The forward's buffers, `R` rows of each (a prompt chunk's); a recipe reads and writes them through its Pass.
+pub const Scratch = struct {
+ ids: u64, // int32 [R]: tokens
+ pos: u64, // int32 [R]: each row's position
+ meta: u64, // int32 [round_rows]: each chain's first position (AttnSegs' META)
+ h: u64, // bf16 [R, hidden]: the residual
+ x: u64, // bf16 [R, hidden]: a norm's output
+ r: u64, // bf16 [R, hidden]: a block's output
+ acc: u64, // fp32 [R, acc width]: a projection's sums before its bias
+ wide: u64, // bf16 [R, wide]: gate|up
+ y: u64, // bf16 [R, max(q width, mid)]: attention out, the MLP's activation
+ q: u64, // bf16 [R, q width]: rotated queries
+ kv: u64, // bf16 [R, 2 kv width]: keys then values a row
+ xs: u64, // fp32 [round_rows, q width / 64]: decode attention's 64-value sums
+ po: u64, // decode attention partials: o, m, l
+ pm: u64,
+ pl: u64,
+ logits: u64, // bf16 [round_rows, vocab]
+ part: u64, // the lane matmul's slice partials
+};
+
+/// bytes to dst on `s`: a private pinned copy kept until the stream completes, so the source may change at once.
+pub fn put(d: *const cuda.Driver, s: cuda.Stream, dst: u64, bytes: []const u8) !void {
+ try cuda.copy_source.enqueue(d, dst, bytes, s);
+}
+
+pub const Model = struct {
+ gpa: std.mem.Allocator,
+ d: *const cuda.Driver,
+ g: Geometry,
+ context: usize, // the served window: a request's prompt and reply fit in it
+ slots: usize = 0, // streams with a table row: none until attach
+ pages: usize = 0, // pages in the arena
+ arena: u64 = 0, // [page][layer][K, V][512][kv width] bf16
+ table: u64 = 0, // int32 [slots][perSlot()]: slot s's page ids in position order
+ ids: []u32 = &.{}, // the tables on the host
+ caps: []usize = &.{}, // positions slot s's pages cover
+ inv: u64 = 0, // fp32 [head_dim / 2]: RoPE's inverse frequencies (a geometry with theta)
+ acc_width: usize,
+ part_bytes: usize,
+ s: Scratch = undefined,
+ buffers: std.ArrayList(DeviceBuffer) = .empty,
+ planned: ?usize = null, // a plan: the bytes counted, nothing allocated
+
+ pub fn perSlot(m: *const Model) usize {
+ return pages.perSlot(m.context);
+ }
+
+ /// Positions a slot's table reaches: the window and its verify rows, in whole pages.
+ pub fn capacity(m: *const Model) usize {
+ return m.perSlot() * pages.page_rows;
+ }
+
+ pub fn pageElems(m: *const Model) usize {
+ return pages.pageElems(m.g);
+ }
+
+ pub fn pageBytes(m: *const Model) u64 {
+ return pages.pageBytes(m.g);
+ }
+
+ pub fn layout(m: *const Model) kv.Layout {
+ return pages.layout(m.g, m.context);
+ }
+
+ pub fn streamBytes(m: *const Model) u64 {
+ return pages.streamBytes(m.g, m.context);
+ }
+
+ /// Decode attention's key chunks a row walks at most, and the blocks that walk them.
+ pub fn chunks(m: *const Model) usize {
+ return m.perSlot();
+ }
+
+ pub fn chunkBlocks(m: *const Model) usize {
+ return @min(m.chunks(), chunk_blocks);
+ }
+
+ fn alloc(m: *Model, bytes: usize) !u64 {
+ if (m.planned) |*p| {
+ p.* += (@max(bytes, 256) + 255) / 256 * 256;
+ return 0;
+ }
+ const b = try DeviceBuffer.alloc(m.d, @max(bytes, 256));
+ errdefer {
+ var x = b;
+ x.free();
+ }
+ try m.buffers.append(m.gpa, b);
+ return b.ptr;
+ }
+
+ /// The f32 sums width every biased projection fits, and the lane's largest slice partials over the shapes.
+ fn sizes(g: Geometry, lane: *const cuda.qmmf.Lane) struct { acc: usize, part: usize } {
+ const shapes = [_][2]usize{ .{ g.acc, g.hidden }, .{ g.hidden, g.qWidth() }, .{ g.wide, g.hidden }, .{ g.hidden, g.mid }, .{ g.vocab, g.hidden } };
+ var part: usize = 0;
+ for (shapes) |nk| for (1..prompt_rows + 1) |rows| {
+ part = @max(part, lane.partBytes(rows, cuda.modelopt.bf16Weight(0, nk[0], nk[1])));
+ };
+ return .{ .acc = @max(g.acc, g.wide, g.hidden), .part = part };
+ }
+
+ fn blank(gpa: std.mem.Allocator, d: *const cuda.Driver, g: Geometry, context: usize, lane: *const cuda.qmmf.Lane) Model {
+ const z = sizes(g, lane);
+ return .{ .gpa = gpa, .d = d, .g = g, .context = context, .acc_width = z.acc, .part_bytes = z.part };
+ }
+
+ /// The scratch and the RoPE table; the arena and tables come with attach, once the host has planned the pages.
+ pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, g: Geometry, context: usize, lane: *const cuda.qmmf.Lane) !Model {
+ var m = blank(gpa, d, g, context, lane);
+ errdefer m.deinit();
+ errdefer cuda.completion.drain(d, s, "the dense engine's scratch uploads did not drain before their buffers were freed");
+ try m.build(s);
+ return m;
+ }
+
+ /// The device bytes `init` allocates, counted by the same path without allocating.
+ pub fn openBytes(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, g: Geometry, context: usize, lane: *const cuda.qmmf.Lane) !usize {
+ var m = blank(gpa, d, g, context, lane);
+ m.planned = 0;
+ defer m.deinit();
+ try m.build(s);
+ return m.planned.?;
+ }
+
+ fn build(m: *Model, s: cuda.Stream) !void {
+ const g = m.g;
+ const r = prompt_rows;
+ const parts = round_rows * m.chunks() * g.heads;
+ m.s = .{
+ .ids = try m.alloc(r * 4),
+ .pos = try m.alloc(r * 4),
+ .meta = try m.alloc(round_rows * 4),
+ .h = try m.alloc(r * g.hidden * 2),
+ .x = try m.alloc(r * g.hidden * 2),
+ .r = try m.alloc(r * g.hidden * 2),
+ .acc = try m.alloc(r * m.acc_width * 4),
+ .wide = try m.alloc(r * g.wide * 2),
+ .y = try m.alloc(r * @max(g.qWidth(), g.mid) * 2),
+ .q = try m.alloc(r * g.qWidth() * 2),
+ .kv = try m.alloc(r * 2 * g.kvWidth() * 2),
+ .xs = try m.alloc(round_rows * (g.qWidth() / 64) * 4),
+ .po = try m.alloc(parts * g.head_dim * 4),
+ .pm = try m.alloc(parts * 4),
+ .pl = try m.alloc(parts * 4),
+ .logits = try m.alloc(round_rows * g.vocab * 2),
+ .part = try m.alloc(m.part_bytes),
+ };
+ const theta = g.theta orelse return;
+ m.inv = try m.alloc(rope.count(g.head_dim) * 4);
+ if (m.planned != null) return;
+ const inv = try m.gpa.alloc(f32, rope.count(g.head_dim));
+ defer m.gpa.free(inv);
+ if (g.inv_freq) |table| {
+ if (table.len != inv.len) return error.BadRopeTable;
+ @memcpy(inv, table);
+ } else rope.inverse(theta, g.head_dim, inv);
+ try m.upload(s, m.inv, std.mem.sliceAsBytes(inv));
+ try s.synchronize();
+ }
+
+ fn upload(m: *const Model, s: cuda.Stream, dst: u64, bytes: []const u8) !void {
+ try put(m.d, s, dst, bytes);
+ }
+
+ pub fn deinit(m: *Model) void {
+ for (m.buffers.items) |*b| b.free();
+ m.buffers.deinit(m.gpa);
+ m.gpa.free(m.ids);
+ m.gpa.free(m.caps);
+ }
+
+ /// The arena of `pool` pages (zeroed) and the tables of `extra + 1` slots, no page mapped.
+ pub fn attach(m: *Model, s: cuda.Stream, extra: usize, pool: usize) !void {
+ if (m.pages != 0 or pool == 0 or extra >= round_rows or m.planned != null) return error.BadAttach;
+ const mark = m.buffers.items.len; // the errdefer reads it at failure, when the list has grown
+ errdefer m.undo(s, mark);
+ const bytes = pool * m.pageBytes();
+ m.arena = try m.alloc(bytes);
+ try m.d.check(m.d.api.cuMemsetD8Async(m.arena, 0, bytes, s.handle), "cuMemsetD8Async");
+ const slots = extra + 1;
+ m.table = try m.alloc(slots * m.perSlot() * 4);
+ m.ids = try m.gpa.alloc(u32, slots * m.perSlot());
+ @memset(m.ids, 0);
+ m.caps = try m.gpa.alloc(usize, slots);
+ @memset(m.caps, 0);
+ try m.upload(s, m.table, std.mem.sliceAsBytes(m.ids));
+ try s.synchronize();
+ m.slots = slots;
+ m.pages = pool;
+ }
+
+ /// A failed attach: its fill and table upload proved complete, then its arena, tables and host rows gone.
+ fn undo(m: *Model, s: cuda.Stream, mark: usize) void {
+ cuda.completion.drain(m.d, s, "the dense engine's page arena fill did not drain before it was freed");
+ while (m.buffers.items.len > mark) {
+ var b = m.buffers.pop().?;
+ b.free();
+ }
+ m.gpa.free(m.ids);
+ m.gpa.free(m.caps);
+ m.ids = &.{};
+ m.caps = &.{};
+ m.arena = 0;
+ m.table = 0;
+ }
+
+ /// Slot `slot`'s pages from table entry `from` on are `ids` (arena page numbers in position order).
+ pub fn mapPages(m: *Model, s: cuda.Stream, slot: usize, ids: []const u32, from: usize) !void {
+ const per = m.perSlot();
+ if (slot >= m.slots or from + ids.len > per) return error.BadChain;
+ for (ids) |id| if (id >= m.pages) return error.BadChain;
+ const at = slot * per + from;
+ @memcpy(m.ids[at..][0..ids.len], ids);
+ m.caps[slot] = (from + ids.len) * pages.page_rows;
+ try m.upload(s, m.table + at * 4, std.mem.sliceAsBytes(m.ids[at..][0..ids.len]));
+ }
+
+ /// Slot `slot` holds no pages: its table entries stay, no forward may reach them.
+ pub fn unmapPages(m: *Model, slot: usize) void {
+ if (slot < m.slots) m.caps[slot] = 0;
+ }
+
+ /// Layer `layer`'s key (or value) cache as the kernels take it: its part of page 0.
+ pub fn cache(m: *const Model, layer: usize, value: bool) u64 {
+ return m.arena + ((2 * layer + @intFromBool(value)) * pages.page_rows * m.g.kvWidth()) * 2;
+ }
+
+ /// Slot s's page table row.
+ pub fn tableOf(m: *const Model, slot: usize) u64 {
+ return m.table + slot * m.perSlot() * 4;
+ }
+};
diff --git a/zig/src/cuda_dense/native.zig b/zig/src/cuda_dense/native.zig
new file mode 100644
index 00000000000..818ec0f2bf9
--- /dev/null
+++ b/zig/src/cuda_dense/native.zig
@@ -0,0 +1,152 @@
+//! Native(R): a dense recipe as the native CUDA host's family, its decls generated once for every recipe.
+const std = @import("std");
+const cuda = @import("cuda");
+const lanes = @import("lanes");
+const kv = @import("core").kv_pool;
+const recipe = @import("recipe.zig");
+const kern = @import("kernels.zig");
+const mdl = @import("model.zig");
+const pages = @import("pages.zig");
+const backend = @import("backend.zig");
+
+/// The refusals the engine words itself, for every recipe; a recipe's own come from its `explain`.
+pub const engine_errors = [_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoDraftHead, error.DenseTextOnly, error.StructuresUnsupported, error.NonFiniteLogits, error.CapturedKernels, error.ParallelPastRound, error.NoStreamFits, error.UnsupportedDenseHeadDim, error.UnsupportedDenseGroup };
+
+pub fn Native(comptime R: type) type {
+ comptime recipe.check(R);
+ return struct {
+ const Dense = backend.Dense(R);
+
+ pub const model_type: []const u8 = R.model_type;
+ pub const formats: []const []const u8 = R.formats;
+ pub const default_context: i64 = R.default_context;
+ pub const max_segments: u32 = 1;
+ pub const prompt_rows: u32 = backend.prompt_rows;
+ /// A request's verify reserve: its pages hold that many rows past its prompt and reply.
+ pub const window_rows: u32 = backend.max_window;
+ /// `explain` words this family's open refusals too.
+ pub const explains_open = true;
+ /// No drafter: --drafter is refused before the load.
+ pub const takes_drafter = false;
+ pub const drafter_refusal = std.fmt.comptimePrint("--drafter: {s} on CUDA loads no drafter; serve it without --drafter", .{R.name});
+ /// No prompt states: the host plans no prompt cache for this family and refuses --prompt-cache-gib.
+ pub const keeps_no_prompts = true;
+ /// Whether replies open with a reasoning block (the recipe's word, else the host's default).
+ pub const reasoning: bool = if (@hasDecl(R, "reasoning")) R.reasoning else true;
+
+ /// The server's reading of the room left after the weights; null if unreadable.
+ pub const Room = struct { ctx: *anyopaque, read: *const fn (*anyopaque) ?u64 };
+ /// No vision, precision, kv8, prompt_int8, loops or rank field: the host refuses those flags by name.
+ pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8, lanes_fixed: bool = false, room: ?Room = null };
+
+ pub const Loaded = struct {
+ backend: lanes.backend.Backend,
+ facts: lanes.Model,
+ rows: u32,
+ stream_bytes: usize,
+ free_streams: u32 = 1,
+ ctx: *anyopaque,
+ deinit: *const fn (*anyopaque) void,
+ target: *Dense,
+ /// Every stream's keys and values are pages of the core pool, the first stream's too.
+ layout: ?kv.Layout,
+ engine_paged: bool = true,
+ };
+
+ const Owned = struct { gpa: std.mem.Allocator, s: cuda.Stream, w: R.Weights, k: kern.Kernels, own: R.Kernels, m: mdl.Model, b: Dense };
+
+ /// The checkpoint, the shared and the recipe's kernels, the scratch; streams' pages come with attach.
+ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded {
+ if (kernels != null) return error.CapturedKernels;
+ if (o.lanes_fixed and o.lanes > backend.round_rows) return error.ParallelPastRound;
+ const g = try R.geometry(gpa, io, dir);
+ _ = try recipe.attentionShape(g); // refused by name before the weights load
+ const own = try gpa.create(Owned);
+ errdefer gpa.destroy(own);
+ own.gpa = gpa;
+ own.s = try cuda.Stream.init(ctx.d, false);
+ errdefer own.s.deinit();
+ own.w = try R.load(gpa, io, ctx.d, own.s, dir);
+ errdefer own.w.deinit();
+ own.k = try kern.Kernels.load(ctx, g);
+ errdefer own.k.deinit();
+ own.own = try R.Kernels.load(ctx);
+ errdefer own.own.deinit();
+ const room: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64);
+ const bytes = try mdl.Model.openBytes(gpa, ctx.d, own.s, g, o.context, &own.k.lane) + Dense.openBytes(g.vocab);
+ if (bytes > room) return error.NoStreamFits;
+ own.m = try mdl.Model.init(gpa, ctx.d, own.s, g, o.context, &own.k.lane);
+ errdefer own.m.deinit();
+ own.b = try Dense.init(gpa, ctx.d, &own.m, &own.k, &own.own, &own.w, own.s);
+ return .{
+ .backend = own.b.backend(),
+ .facts = own.b.facts(),
+ .rows = backend.max_window,
+ .stream_bytes = own.m.streamBytes(),
+ .ctx = own,
+ .deinit = release,
+ .target = &own.b,
+ .layout = own.m.layout(),
+ };
+ }
+
+ /// The checkpoint's weight bytes, counted from its tensors before the load.
+ pub fn textWeightBytes(a: std.mem.Allocator, io: std.Io, dir: []const u8) !?u64 {
+ return try R.weightBytes(a, io, dir);
+ }
+
+ fn owned(p: *anyopaque) *Owned {
+ return @ptrCast(@alignCast(p));
+ }
+
+ /// The pool: `streams` more slots beside the first and an arena of `n` pages.
+ pub fn attach(p: *anyopaque, streams: u32, n: usize) !void {
+ try owned(p).b.attach(streams, n);
+ }
+
+ /// Whether a request's pages could be promised now (the lane host's admission, on its engine thread).
+ pub fn fits(p: *anyopaque, prompt: u32, max_tokens: u32) bool {
+ return owned(p).b.admits(prompt, max_tokens);
+ }
+
+ /// The pool the lane host's ledger admits this family's promises by.
+ pub fn ledgerPool(p: *anyopaque) ?*kv.Pool {
+ const b = &owned(p).b;
+ return if (b.pool) |*pool| pool else null;
+ }
+
+ /// A request or load this engine refuses, in words; then the recipe's own; null: none.
+ pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 {
+ return words(err) orelse if (@hasDecl(R, "explain")) R.explain(err) else null;
+ }
+
+ fn words(err: anyerror) ?[]const u8 {
+ const name = R.name;
+ return switch (err) {
+ error.PastContext => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens",
+ error.NoFreeSlot => "every stream's caches are in use: retry once another request ends",
+ error.NoDraftHead => name ++ " on CUDA has no draft head: it drafts from the prompt's own text",
+ error.DenseTextOnly => name ++ " on CUDA reads text prompts only so far: this request has images",
+ error.StructuresUnsupported => "structured output is not built for " ++ name ++ " on CUDA yet",
+ error.NonFiniteLogits => "the model's logits went non-finite at this row: the request is stopped rather than answered from them",
+ error.CapturedKernels => name ++ " runs on the engine's own kernels: unset TENSORFOLD_CUDA_KERNELS",
+ error.ParallelPastRound => std.fmt.comptimePrint("--parallel: {s} on CUDA holds at most {d} streams (a round's rows): serve --parallel {d} or fewer, or auto", .{ name, backend.round_rows, backend.round_rows }),
+ error.NoStreamFits => "the CUDA memory budget fits no " ++ name ++ " stream after the weights: lower --context, free device memory, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB",
+ error.UnsupportedDenseHeadDim => "this " ++ name ++ " checkpoint's attention head dim is not 64 or 128, the two the CUDA dense engine's attention kernels take",
+ error.UnsupportedDenseGroup => "this " ++ name ++ " checkpoint's query heads a KV head are not an even count of at most 16, the groups the CUDA dense engine's attention kernels take",
+ else => null,
+ };
+ }
+
+ fn release(p: *anyopaque) void {
+ const own = owned(p);
+ own.b.deinit();
+ own.m.deinit();
+ own.own.deinit();
+ own.k.deinit();
+ own.w.deinit();
+ own.s.deinit();
+ own.gpa.destroy(own);
+ }
+ };
+}
diff --git a/zig/src/cuda_dense/native_test.zig b/zig/src/cuda_dense/native_test.zig
new file mode 100644
index 00000000000..6788f5675a0
--- /dev/null
+++ b/zig/src/cuda_dense/native_test.zig
@@ -0,0 +1,105 @@
+//! Native(R)'s host contract with a tiny recipe: the decls it emits and omits, its refusals' words, a compiled path.
+const std = @import("std");
+const cuda = @import("cuda");
+const native = @import("native.zig");
+const Pass = @import("pass.zig").Pass;
+const Geometry = @import("recipe.zig").Geometry;
+const Allocator = std.mem.Allocator;
+
+const tiny: Geometry = .{ .hidden = 256, .layers = 2, .heads = 4, .kv_heads = 2, .head_dim = 64, .vocab = 512, .eps = 1e-6, .context = 4096, .acc = 512, .wide = 1024, .mid = 512, .theta = 1e4 };
+
+/// A recipe of the dots shape on tiny sizes: biased q|k|v sums, attention, o, SwiGLU, an untied head.
+const Tiny = struct {
+ pub const name = "Tiny";
+ pub const model_type = "tiny_dense";
+ pub const formats: []const []const u8 = &.{"bf16"};
+ pub const default_context: i64 = 4096;
+ pub const reasoning = false;
+ pub const post_embed: ?fn (*const Pass(Tiny)) anyerror!void = null;
+ pub const Weights = struct {
+ embed: u64 = 0,
+ norm: u64 = 0,
+ bias: u64 = 0,
+ proj: cuda.qmmf.Weight = cuda.modelopt.bf16Weight(0, 512, 256),
+ pub fn deinit(_: *Weights) void {}
+ };
+ pub const Kernels = struct {
+ pub fn load(_: *const cuda.Context) !Kernels {
+ return .{};
+ }
+ pub fn deinit(_: *Kernels) void {}
+ };
+ pub fn geometry(_: Allocator, _: std.Io, _: []const u8) !Geometry {
+ return tiny;
+ }
+ pub fn load(_: Allocator, _: std.Io, _: *const cuda.Driver, _: cuda.Stream, _: []const u8) !Weights {
+ return .{};
+ }
+ pub fn weightBytes(_: Allocator, _: std.Io, _: []const u8) !u64 {
+ return 1 << 20;
+ }
+ pub fn embed(p: *const Pass(Tiny)) !void {
+ try p.k.embed.bf16(p.s, p.m.s.ids, p.w.embed, p.m.s.h, p.rows, p.g.hidden);
+ }
+ pub fn layer(p: *const Pass(Tiny), i: usize) !void {
+ const x = p.m.s;
+ try p.norm(if (i == 0) null else x.r, p.w.norm);
+ try p.projectBias(p.w.proj, p.w.bias, x.x, x.wide);
+ try p.projectF32(p.w.proj, x.x, x.acc);
+ try p.attend(i, x.q, x.kv, x.y);
+ try p.project(p.w.proj, x.y, x.r);
+ try p.siluMul(x.wide, x.y, p.g.mid);
+ }
+ pub fn final(p: *const Pass(Tiny)) !void {
+ try p.norm(p.m.s.r, p.w.norm);
+ }
+ pub fn head(w: *const Weights) cuda.qmmf.Weight {
+ return w.proj;
+ }
+ pub fn explain(err: anyerror) ?[]const u8 {
+ return if (err == error.TinyConfig) "this Tiny checkpoint is not the one the recipe serves" else null;
+ }
+};
+
+const N = native.Native(Tiny);
+
+test "Native emits the host decls a dense family takes and no others" {
+ inline for (.{ "model_type", "formats", "default_context", "max_segments", "prompt_rows", "window_rows", "explains_open", "takes_drafter", "drafter_refusal", "Room", "Options", "Loaded", "open", "textWeightBytes", "attach", "fits", "ledgerPool", "explain", "keeps_no_prompts", "reasoning" }) |d| {
+ if (!@hasDecl(N, d)) return error.MissingDecl;
+ }
+ // each of these switches a host path on (prompt cache, KV sharing, learning, own planning, drafters, ranks)
+ inline for (.{ "snap_save", "snap_restore", "snap_drop", "snap_bytes", "snap_slots", "snap_grid", "snap_lookahead", "share_attach", "share_pool", "share_states", "kvPool", "slide", "learnModel", "linked", "contextCeiling", "pagedLayout", "rankRoom", "request_starts", "request_starts_switch", "openDrafter", "warm_quiet" }) |d| {
+ if (@hasDecl(N, d)) return error.UnexpectedDecl;
+ }
+ inline for (.{ "context", "drafts", "segments", "slide", "lanes", "lanes_fixed", "room" }) |f| try std.testing.expect(@hasField(N.Options, f));
+ inline for (.{ "vision", "precision", "kv8", "prompt_int8", "loops", "rank", "drafter" }) |f| try std.testing.expect(!@hasField(N.Options, f));
+ inline for (.{ "layout", "engine_paged", "free_streams", "target" }) |f| try std.testing.expect(@hasField(N.Loaded, f));
+ try std.testing.expectEqualStrings("tiny_dense", N.model_type);
+ try std.testing.expect(N.keeps_no_prompts and !N.reasoning and !N.takes_drafter);
+ try std.testing.expectEqual(@as(u32, 16), N.window_rows);
+ try std.testing.expectEqual(@as(u32, 1024), N.prompt_rows);
+ try std.testing.expectEqual(@as(u32, 1), N.max_segments);
+ try std.testing.expectStringStartsWith(N.drafter_refusal, "--drafter: Tiny on CUDA");
+}
+
+test "every engine refusal is worded, the recipe's own after it, the rest left to the server" {
+ for (native.engine_errors) |e| try std.testing.expect(N.explain(null, e) != null);
+ try std.testing.expectStringStartsWith(N.explain(null, error.TinyConfig).?, "this Tiny checkpoint");
+ try std.testing.expectStringStartsWith(N.explain(null, error.DenseTextOnly).?, "Tiny on CUDA reads text prompts only");
+ try std.testing.expect(N.explain(null, error.OutOfMemory) == null);
+}
+
+test "open refuses captured kernels and a fixed --parallel past a round before it reads the checkpoint" {
+ const ctx: *const cuda.Context = undefined;
+ const base: N.Options = .{ .context = 4096, .drafts = true };
+ try std.testing.expectError(error.CapturedKernels, N.open(std.testing.allocator, undefined, ctx, "", "captured", base));
+ var wide = base;
+ wide.lanes = 129;
+ wide.lanes_fixed = true;
+ try std.testing.expectError(error.ParallelPastRound, N.open(std.testing.allocator, undefined, ctx, "", null, wide));
+}
+
+test "a recipe's whole path compiles against the engine: open, the lane backend, the forward" {
+ const hooks = .{ &N.open, &N.attach, &N.fits, &N.ledgerPool, &N.textWeightBytes, &N.explain };
+ inline for (hooks) |h| try std.testing.expect(@intFromPtr(h) != 0);
+}
diff --git a/zig/src/cuda_dense/pages.zig b/zig/src/cuda_dense/pages.zig
new file mode 100644
index 00000000000..83698cae7ea
--- /dev/null
+++ b/zig/src/cuda_dense/pages.zig
@@ -0,0 +1,66 @@
+//! KV pages of the dense engine: 512 rows of every layer's keys and values, the host pool's layout, a request's pages.
+const std = @import("std");
+const kv = @import("core").kv_pool;
+const Geometry = @import("recipe.zig").Geometry;
+
+/// Rows a page holds: one decode attention chunk (nemotron_attention.cu's PAGE).
+pub const page_rows = 512;
+/// The widest window a stream verifies: a request's pages hold that many rows past its prompt and reply.
+pub const window = 16;
+
+/// Pages that hold positions [0, positions).
+pub fn pagesOf(positions: usize) usize {
+ return (positions + page_rows - 1) / page_rows;
+}
+
+/// Pages a request holds at most: its prompt, its reply and the verify window.
+pub fn pagesFor(prompt: usize, max_new: usize) usize {
+ return pagesOf(prompt + max_new + window);
+}
+
+/// Table entries a stream takes at a `context`-token window, its verify window included.
+pub fn perSlot(context: usize) usize {
+ return pagesOf(context + window);
+}
+
+/// bf16 elements a page holds: [layer][K, V][512][kv width].
+pub fn pageElems(g: Geometry) usize {
+ return g.layers * 2 * page_rows * g.kvWidth();
+}
+
+pub fn pageBytes(g: Geometry) u64 {
+ return pageElems(g) * 2;
+}
+
+/// Bytes of one stream at a whole `context`-token window: its pages and its table row.
+pub fn streamBytes(g: Geometry, context: usize) u64 {
+ return perSlot(context) * (pageBytes(g) + 4);
+}
+
+/// What the host's pool plans: a token's keys and values over every layer, a stream's table row fixed.
+pub fn layout(g: Geometry, context: usize) kv.Layout {
+ const row: u64 = 2 * g.kvWidth() * 2;
+ return kv.Layout.ofRows(.{ .kind = .full, .row = row, .layers = @intCast(g.layers) }, page_rows, perSlot(context) * 4, 0);
+}
+
+const dots: Geometry = .{ .hidden = 1536, .layers = 28, .heads = 12, .kv_heads = 2, .head_dim = 128, .vocab = 151936, .eps = 1e-6, .context = 131072, .acc = 2048, .wide = 17920, .mid = 8960, .theta = 1e6 };
+
+test "dots' page: 512 rows of 28 layers' keys and values at 2 x 128 dims, 14 MiB" {
+ try std.testing.expectEqual(@as(u64, 28 * 2 * 512 * 256 * 2), pageBytes(dots));
+ try std.testing.expectEqual(@as(u64, 14680064), pageBytes(dots));
+ const l = layout(dots, 32768);
+ try std.testing.expectEqual(pageBytes(dots), l.pageBytes());
+ try std.testing.expectEqual(@as(u64, 65 * 4), l.fixed_bytes);
+ try std.testing.expectEqual(@as(usize, 65), l.pages(32768 + window));
+ try std.testing.expectEqual(@as(u64, 65 * (14680064 + 4)), streamBytes(dots, 32768));
+}
+
+test "a request holds its prompt, its reply and the verify window in whole pages" {
+ try std.testing.expectEqual(@as(usize, 1), pagesFor(1, 1));
+ try std.testing.expectEqual(@as(usize, 1), pagesFor(400, 96));
+ try std.testing.expectEqual(@as(usize, 2), pagesFor(400, 97));
+ try std.testing.expectEqual(@as(usize, 65), pagesFor(32768, 0));
+ try std.testing.expectEqual(@as(usize, 65), perSlot(32768));
+ try std.testing.expectEqual(@as(usize, 1), perSlot(496));
+ try std.testing.expectEqual(@as(usize, 2), perSlot(497));
+}
diff --git a/zig/src/cuda_dense/pass.zig b/zig/src/cuda_dense/pass.zig
new file mode 100644
index 00000000000..63d9a277094
--- /dev/null
+++ b/zig/src/cuda_dense/pass.zig
@@ -0,0 +1,150 @@
+//! One dense forward: the engine's operators a recipe's layers call, and the pass from token ids to bf16 logits.
+const std = @import("std");
+const cuda = @import("cuda");
+const recipe = @import("recipe.zig");
+const kern = @import("kernels.zig");
+const mdl = @import("model.zig");
+
+const Model = mdl.Model;
+const Kernels = kern.Kernels;
+const Stream = cuda.Stream;
+const Weight = cuda.qmmf.Weight;
+const tail: cuda.norms.Tail = .{ .round = true }; // the norms' torch points: bf16(w * bf16(h * inv))
+
+/// A stream's rows in one forward: its slot, first position, tokens.
+pub const Chain = struct { slot: usize, p0: usize, tokens: []const u32 };
+
+/// What a forward runs on: the model's state, the shared and the recipe's kernels, the weights, the CUDA stream.
+pub fn Engine(comptime R: type) type {
+ return struct { m: *const Model, k: *const Kernels, own: *const R.Kernels, w: *const R.Weights, s: Stream };
+}
+
+/// A forward's rows (a prompt chunk, or a round of windows) and the operators a recipe's layers call on them.
+pub fn Pass(comptime R: type) type {
+ return struct {
+ const P = @This();
+ m: *const Model,
+ k: *const Kernels,
+ own: *const R.Kernels,
+ w: *const R.Weights,
+ s: Stream,
+ g: recipe.Geometry,
+ chains: []const Chain,
+ rows: usize,
+ prompt: bool,
+
+ fn part(p: *const P, w: Weight) !u64 {
+ if (p.k.lane.partBytes(p.rows, w) > p.m.part_bytes) return error.LanePartTooSmall;
+ return p.m.s.part;
+ }
+
+ /// out [rows, n] bf16 = x [rows, k] times `w`, fp32 sums rounded once.
+ pub fn project(p: *const P, w: Weight, x: u64, out: u64) !void {
+ try p.k.lane.matmul(p.s, x, w.k, p.rows, w, out, try p.part(w));
+ }
+
+ /// out [rows, n] fp32 = x [rows, k] times `w`: the sums a recipe's own kernel rounds.
+ pub fn projectF32(p: *const P, w: Weight, x: u64, out: u64) !void {
+ try p.k.lane.matmulF32(p.s, x, w.k, p.rows, w, out, try p.part(w));
+ }
+
+ /// out [rows, n] = bf16(x times `w` + b): the fp32 sums and the bias rounded once (vit.cu's bias).
+ pub fn projectBias(p: *const P, w: Weight, b: u64, x: u64, out: u64) !void {
+ if (w.n > p.m.acc_width) return error.AccTooNarrow;
+ try p.projectF32(w, x, p.m.s.acc);
+ try p.k.vit.bias(p.s, p.m.s.acc, b, out, p.rows, w.n);
+ }
+
+ /// h = bf16(h + r) (r null: as is), x = bf16(w * bf16(h * inv_rms)).
+ pub fn norm(p: *const P, r: ?u64, w: u64) !void {
+ const x = p.m.s;
+ try p.k.norms.addRmsTail(p.s, x.h, r, w, x.h, x.x, null, p.rows, p.g.hidden, p.g.eps, tail);
+ }
+
+ /// out [rows, f] = bf16(bf16(silu(gate)) * up) from gate|up [rows, 2 f].
+ pub fn siluMul(p: *const P, gu: u64, out: u64, f: usize) !void {
+ try p.k.siluMul(p.s, gu, out, p.rows, f);
+ }
+
+ /// Layer `layer`'s keys and values (kv [rows, 2 kv width]) into the pages, then q's attention into out.
+ pub fn attend(p: *const P, layer: usize, q: u64, kv: u64, out: u64) !void {
+ if (layer >= p.g.layers) return error.BadLayer;
+ const m = p.m;
+ const x = m.s;
+ const kc = m.cache(layer, false);
+ const vc = m.cache(layer, true);
+ const qw = p.g.qWidth();
+ var row0: usize = 0;
+ var at: usize = 0;
+ while (at < p.chains.len) : (at += kern.max_segs) {
+ const group = p.chains[at..@min(p.chains.len, at + kern.max_segs)];
+ var sg = std.mem.zeroes(kern.AttnSegs);
+ var rows: usize = 0;
+ for (group, 0..) |ch, j| {
+ sg.kc[j] = kc;
+ sg.vc[j] = vc;
+ sg.meta[j] = x.meta + (at + j) * 4;
+ sg.pt[j] = m.tableOf(ch.slot);
+ sg.row0[j] = @intCast(rows);
+ rows += ch.tokens.len;
+ }
+ sg.ps = @intCast(m.pageElems());
+ sg.n = @intCast(group.len);
+ try p.k.kvWrite(p.s, sg, kv + row0 * 2 * p.g.kvWidth() * 2, rows);
+ const qa = q + row0 * qw * 2;
+ const oa = out + row0 * qw * 2;
+ if (p.prompt) {
+ try p.k.prompt(p.s, qa, kc, vc, oa, group[0].p0, rows, m.tableOf(group[0].slot), m.pageElems());
+ } else {
+ try p.k.attend(p.s, sg, qa, x.po, x.pm, x.pl, oa, x.xs + row0 * (qw / 64) * 4, rows, m.chunks(), m.chunkBlocks());
+ }
+ row0 += rows;
+ }
+ }
+ };
+}
+
+/// Every chain's rows (a prompt chunk, or windows): the recipe's layers, then logits rows [head0, head0 + heads).
+pub fn forward(comptime R: type, e: Engine(R), chains: []const Chain, prompt: bool, head0: usize, heads: usize) !void {
+ const m = e.m;
+ var rows: usize = 0;
+ for (chains) |ch| {
+ if (ch.tokens.len == 0 or ch.slot >= m.slots or ch.p0 + ch.tokens.len > m.caps[ch.slot]) return error.BadChain;
+ rows += ch.tokens.len;
+ }
+ if (rows > @as(usize, if (prompt) mdl.prompt_rows else mdl.round_rows) or (prompt and chains.len != 1)) return error.BadWindow;
+ if (chains.len > mdl.round_rows or heads > mdl.round_rows or head0 + heads > rows) return error.BadWindow;
+ try upload(m, e.s, chains, rows);
+ const p: Pass(R) = .{ .m = m, .k = e.k, .own = e.own, .w = e.w, .s = e.s, .g = m.g, .chains = chains, .rows = rows, .prompt = prompt };
+ try R.embed(&p);
+ for (0..m.g.layers) |i| try R.layer(&p, i);
+ try R.final(&p);
+ if (heads == 0) return;
+ const head: Weight = R.head(e.w);
+ const x = m.s.x + head0 * m.g.hidden * 2;
+ if (e.k.lane.partBytes(heads, head) > m.part_bytes) return error.LanePartTooSmall;
+ try e.k.lane.matmul(e.s, x, m.g.hidden, heads, head, m.s.logits, m.s.part);
+}
+
+fn put(m: *const Model, s: Stream, dst: u64, bytes: []const u8) !void {
+ try mdl.put(m.d, s, dst, bytes);
+}
+
+/// Tokens and positions a row, each chain's first position, staged for the launches.
+fn upload(m: *const Model, s: Stream, chains: []const Chain, rows: usize) !void {
+ const ids = try m.gpa.alloc(i32, 2 * rows + chains.len);
+ defer m.gpa.free(ids);
+ var at: usize = 0;
+ for (chains, 0..) |ch, j| {
+ for (ch.tokens, 0..) |t, i| {
+ ids[at + i] = @intCast(t);
+ ids[rows + at + i] = @intCast(ch.p0 + i);
+ }
+ ids[2 * rows + j] = @intCast(ch.p0);
+ at += ch.tokens.len;
+ }
+ const x = m.s;
+ try put(m, s, x.ids, std.mem.sliceAsBytes(ids[0..rows]));
+ try put(m, s, x.pos, std.mem.sliceAsBytes(ids[rows .. 2 * rows]));
+ try put(m, s, x.meta, std.mem.sliceAsBytes(ids[2 * rows ..]));
+}
diff --git a/zig/src/cuda_dense/put_test.zig b/zig/src/cuda_dense/put_test.zig
new file mode 100644
index 00000000000..236f2303c1d
--- /dev/null
+++ b/zig/src/cuda_dense/put_test.zig
@@ -0,0 +1,72 @@
+//! The engine's uploads own their source: a fake driver reads it at the stream's sync, after the caller changed it.
+const std = @import("std");
+const cuda = @import("cuda");
+const put = @import("model.zig").put;
+const t = std.testing;
+
+const Fake = struct {
+ var pinned: ?[]align(16) u8 = null;
+ var pending: ?[]const u8 = null;
+ var result: [16]u8 = undefined;
+ fn host(out: *?*anyopaque, len: usize, _: c_uint) callconv(.c) cuda.abi.Result {
+ pinned = t.allocator.alignedAlloc(u8, .@"16", len) catch return cuda.abi.error_not_ready;
+ out.* = pinned.?.ptr;
+ return 0;
+ }
+ fn free(_: ?*anyopaque) callconv(.c) cuda.abi.Result {
+ std.debug.assert(pending == null);
+ t.allocator.free(pinned.?);
+ pinned = null;
+ return 0;
+ }
+ fn copy(_: u64, ptr: ?*const anyopaque, len: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result {
+ pending = @as([*]const u8, @ptrCast(ptr.?))[0..len];
+ return 0;
+ }
+ fn sync(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result {
+ if (pending) |bytes| @memcpy(result[0..bytes.len], bytes);
+ pending = null;
+ return 0;
+ }
+ fn event(out: *cuda.abi.Event, _: c_uint) callconv(.c) cuda.abi.Result {
+ out.* = @ptrFromInt(1);
+ return 0;
+ }
+ fn record(_: cuda.abi.Event, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result {
+ return 0;
+ }
+ fn query(_: cuda.abi.Event) callconv(.c) cuda.abi.Result {
+ return if (pending != null) cuda.abi.error_not_ready else 0;
+ }
+ fn destroy(_: cuda.abi.Event) callconv(.c) cuda.abi.Result {
+ return 0;
+ }
+};
+
+test "an upload's source may change as soon as put returns: the device gets the bytes it was given" {
+ var d: cuda.Driver = undefined;
+ d.api.cuMemHostAlloc = Fake.host;
+ d.api.cuMemFreeHost = Fake.free;
+ d.api.cuMemcpyHtoDAsync_v2 = Fake.copy;
+ d.api.cuStreamSynchronize = Fake.sync;
+ d.api.cuEventCreate = Fake.event;
+ d.api.cuEventRecord = Fake.record;
+ d.api.cuEventQuery = Fake.query;
+ d.api.cuEventDestroy_v2 = Fake.destroy;
+ const s: cuda.Stream = .{ .d = &d, .handle = @ptrFromInt(0x1800) };
+ Fake.pending = null;
+ Fake.result = @splat(0);
+ const host_before = cuda.memory.usage(false).host;
+ defer {
+ s.synchronize() catch unreachable;
+ cuda.copy_source.dropSpares(&d, cuda.copy_source.current());
+ }
+ var meta = [4]i32{ 41, 0, 0, 0 };
+ try put(&d, s, 0x1000, std.mem.asBytes(&meta));
+ meta[0] = -7; // a draw's next meta, written into the same stack slot before the stream ran
+ try s.synchronize();
+ cuda.copy_source.dropSpares(&d, cuda.copy_source.current());
+ try t.expectEqual(@as(i32, 41), std.mem.bytesToValue(i32, Fake.result[0..4]));
+ try t.expect(Fake.pinned == null);
+ try t.expectEqual(host_before, cuda.memory.usage(false).host);
+}
diff --git a/zig/src/cuda_dense/recipe.zig b/zig/src/cuda_dense/recipe.zig
new file mode 100644
index 00000000000..b0f35964a52
--- /dev/null
+++ b/zig/src/cuda_dense/recipe.zig
@@ -0,0 +1,136 @@
+//! A dense decoder's recipe contract: its geometry, the decls the engine reads, and the attention launch shape.
+const std = @import("std");
+
+/// A dense GQA decoder's sizes, read from its checkpoint's config before the weights load.
+pub const Geometry = struct {
+ hidden: usize,
+ layers: usize,
+ heads: usize,
+ kv_heads: usize,
+ head_dim: usize, // 64 or 128
+ vocab: usize,
+ eps: f32,
+ context: usize, // the checkpoint's own window
+ acc: usize, // q|k|v width: a fused projection's f32 sums
+ wide: usize, // gate|up width
+ mid: usize, // the MLP's inner width
+ theta: ?f64 = null, // 1-D RoPE base (rope.zig's inverse frequencies on the device); null: no rotary positions
+ inv_freq: ?[]const f32 = null, // the reference's own inverse frequencies, when its pow is not correctly rounded
+
+ pub fn qWidth(g: Geometry) usize {
+ return g.heads * g.head_dim;
+ }
+
+ pub fn kvWidth(g: Geometry) usize {
+ return g.kv_heads * g.head_dim;
+ }
+
+ pub fn group(g: Geometry) usize {
+ return g.heads / g.kv_heads;
+ }
+};
+
+/// Query heads a KV head the decode kernels take at most (paged_attn.cuh's GMAX).
+pub const max_group = 16;
+
+/// The attention launches of one geometry: prompt pattn and its block, decode chunk and merge.
+pub const Shape = struct {
+ pattn: [:0]const u8,
+ chunk: [:0]const u8,
+ merge: [:0]const u8,
+ hpc: usize, // query heads a pattn block takes from one KV head
+ rows: usize, // prompt rows a pattn block takes
+ grid_y: usize, // pattn blocks over the heads
+ shared: u32, // pattn's dynamic shared bytes
+
+ pub fn blocks(s: Shape, rows: usize) usize {
+ return (rows + s.rows - 1) / s.rows;
+ }
+};
+
+fn pattnSymbol(comptime d: usize, comptime hpc: usize) [:0]const u8 {
+ return std.fmt.comptimePrint("_ZN20tf_prefill_attention12pattn_kernelILi{d}ELi8ELi{d}ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix", .{ d, hpc });
+}
+
+/// The pattn instances prefill_attention.cu builds at eight warps and eight slots: head dim, query heads a block.
+const Instance = struct { d: usize, hpc: usize, symbol: [:0]const u8 };
+const instances = [_]Instance{
+ .{ .d = 128, .hpc = 8, .symbol = pattnSymbol(128, 8) },
+ .{ .d = 128, .hpc = 2, .symbol = pattnSymbol(128, 2) },
+ .{ .d = 64, .hpc = 4, .symbol = pattnSymbol(64, 4) },
+ .{ .d = 64, .hpc = 2, .symbol = pattnSymbol(64, 2) },
+};
+
+/// pattn's shared bytes: eight staging slots of 32 keys, a key row in 16-byte chunks (prefill_attention_rows.cuh).
+fn pattnShared(d: usize) u32 {
+ const chunks = (2 * ((d + 15) / 16) + 7) / 8 * 8;
+ return @intCast(8 * 32 * chunks * 16);
+}
+
+/// The launches for `g`: the largest built pattn instance whose heads a block divide the group; refusals by name.
+pub fn attentionShape(g: Geometry) error{ UnsupportedDenseHeadDim, UnsupportedDenseGroup }!Shape {
+ if (g.head_dim != 64 and g.head_dim != 128) return error.UnsupportedDenseHeadDim;
+ if (g.kv_heads == 0 or g.heads % g.kv_heads != 0 or g.group() > max_group) return error.UnsupportedDenseGroup;
+ const group_n = g.group();
+ for (instances) |in| if (in.d == g.head_dim and group_n % in.hpc == 0) return .{
+ .pattn = in.symbol,
+ .chunk = if (g.head_dim == 128) "tf_nemo_attn_chunk_segs" else "tf_paged_attn_chunk_segs_64",
+ .merge = if (g.head_dim == 128) "tf_nemo_attn_merge_segs" else "tf_paged_attn_merge_segs_64",
+ .hpc = in.hpc,
+ .rows = 16 * (8 / in.hpc),
+ .grid_y = g.kv_heads * (group_n / in.hpc),
+ .shared = pattnShared(g.head_dim),
+ };
+ return error.UnsupportedDenseGroup;
+}
+
+fn need(comptime R: type, comptime name: []const u8) void {
+ if (!@hasDecl(R, name)) @compileError("a cuda_dense recipe declares " ++ name);
+}
+
+fn absent(comptime R: type, comptime name: []const u8) void {
+ if (@hasDecl(R, name) and @field(R, name) != null) @compileError(name ++ " is not built in the dense engine yet: set it null");
+}
+
+/// The decls a recipe gives the engine, at compile time; the capabilities not built yet must stay null.
+pub fn check(comptime R: type) void {
+ const required = .{ "name", "model_type", "formats", "default_context", "Weights", "Kernels", "geometry", "load", "weightBytes", "embed", "layer", "final", "head" };
+ inline for (required) |name| need(R, name);
+ inline for (.{ "post_embed", "Tower", "scores", "propose" }) |name| absent(R, name);
+}
+
+fn shapeOf(heads: usize, kv_heads: usize, head_dim: usize) !Shape {
+ return attentionShape(.{ .hidden = heads * head_dim, .layers = 1, .heads = heads, .kv_heads = kv_heads, .head_dim = head_dim, .vocab = 1, .eps = 1e-6, .context = 1, .acc = 1, .wide = 1, .mid = 1 });
+}
+
+test "attention shapes: dots 12/2 x 128, 16/8 x 128, 16/8 x 64 and 32/8 x 64" {
+ const dots = try shapeOf(12, 2, 128);
+ try std.testing.expectEqual(@as(usize, 2), dots.hpc);
+ try std.testing.expectEqual(@as(usize, 64), dots.rows);
+ try std.testing.expectEqual(@as(usize, 6), dots.grid_y);
+ try std.testing.expectEqual(@as(u32, 65536), dots.shared);
+ try std.testing.expectEqualStrings("_ZN20tf_prefill_attention12pattn_kernelILi128ELi8ELi2ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix", dots.pattn);
+ try std.testing.expectEqualStrings("tf_nemo_attn_chunk_segs", dots.chunk);
+ try std.testing.expectEqual(@as(usize, 2), dots.blocks(65));
+ const glm = try shapeOf(16, 8, 128);
+ try std.testing.expect(glm.hpc == 2 and glm.rows == 64 and glm.grid_y == 8 and glm.shared == 65536);
+ const parse = try shapeOf(16, 8, 64);
+ try std.testing.expect(parse.hpc == 2 and parse.rows == 64 and parse.grid_y == 8 and parse.shared == 32768);
+ try std.testing.expectEqualStrings("_ZN20tf_prefill_attention12pattn_kernelILi64ELi8ELi2ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix", parse.pattn);
+ try std.testing.expectEqualStrings("tf_paged_attn_merge_segs_64", parse.merge);
+ const lfm = try shapeOf(32, 8, 64);
+ try std.testing.expect(lfm.hpc == 4 and lfm.rows == 32 and lfm.grid_y == 8 and lfm.shared == 32768);
+ try std.testing.expectEqualStrings("_ZN20tf_prefill_attention12pattn_kernelILi64ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix", lfm.pattn);
+}
+
+test "attention shapes refuse groups past 16, odd groups, uneven heads and head dims other than 64 and 128" {
+ try std.testing.expectError(error.UnsupportedDenseGroup, shapeOf(34, 2, 128));
+ try std.testing.expectError(error.UnsupportedDenseGroup, shapeOf(12, 4, 128));
+ try std.testing.expectError(error.UnsupportedDenseGroup, shapeOf(12, 5, 64));
+ try std.testing.expectError(error.UnsupportedDenseGroup, shapeOf(8, 8, 64));
+ try std.testing.expectError(error.UnsupportedDenseHeadDim, shapeOf(16, 8, 80));
+ try std.testing.expectError(error.UnsupportedDenseHeadDim, shapeOf(16, 8, 256));
+ try std.testing.expectEqual(@as(usize, 2), (try shapeOf(32, 8, 128)).hpc); // group 4 at head 128: no 4-head block
+ try std.testing.expectEqual(@as(usize, 8), (try shapeOf(32, 2, 128)).hpc);
+ try std.testing.expectEqual(@as(usize, 8), (try shapeOf(64, 8, 128)).hpc);
+}
diff --git a/zig/src/cuda_dense/root.zig b/zig/src/cuda_dense/root.zig
new file mode 100644
index 00000000000..c20f45e0e5e
--- /dev/null
+++ b/zig/src/cuda_dense/root.zig
@@ -0,0 +1,30 @@
+//! The shared CUDA dense-decoder engine: a recipe's layers over paged GQA attention, lanes, draws and the host decls.
+pub const recipe = @import("recipe.zig");
+pub const pages = @import("pages.zig");
+pub const rope = @import("rope.zig");
+pub const kernels = @import("kernels.zig");
+pub const model = @import("model.zig");
+pub const pass = @import("pass.zig");
+pub const draws = @import("draws.zig");
+pub const backend = @import("backend.zig");
+pub const native = @import("native.zig");
+
+pub const Native = native.Native;
+pub const Pass = pass.Pass;
+pub const Geometry = recipe.Geometry;
+pub const check = recipe.check;
+
+test {
+ _ = recipe;
+ _ = pages;
+ _ = rope;
+ _ = kernels;
+ _ = model;
+ _ = pass;
+ _ = draws;
+ _ = backend;
+ _ = native;
+ _ = @import("native_test.zig");
+ _ = @import("put_test.zig");
+ _ = @import("device_fail_test.zig");
+}
diff --git a/zig/src/cuda_dense/rope.zig b/zig/src/cuda_dense/rope.zig
new file mode 100644
index 00000000000..e49633f0ef4
--- /dev/null
+++ b/zig/src/cuda_dense/rope.zig
@@ -0,0 +1,31 @@
+//! 1-D rotary inverse frequencies as HF's default rope computes them in float32; angles and cos/sin are device work.
+const std = @import("std");
+
+/// Frequencies a head of `dim` dims takes: one a rotated pair.
+pub fn count(dim: usize) usize {
+ return dim / 2;
+}
+
+/// out[i] = 1 / theta^(2i / dim): the exponent an f32 quotient, the power of the f32 operands rounded once to f32.
+pub fn inverse(theta: f64, dim: usize, out: []f32) void {
+ std.debug.assert(out.len == count(dim));
+ const base: f32 = @floatCast(theta);
+ for (out, 0..) |*v, i| {
+ const e: f32 = @as(f32, @floatFromInt(2 * i)) / @as(f32, @floatFromInt(dim));
+ const pw: f32 = @floatCast(std.math.pow(f64, base, e));
+ v.* = @as(f32, 1.0) / pw;
+ }
+}
+
+test "dots' inverse frequencies (theta 1e6, head dim 128) in float32" {
+ var got: [64]f32 = undefined;
+ inverse(1e6, 128, &got);
+ for (got, 0..) |v, i| {
+ const e = @as(f32, @floatFromInt(i * 2)) / 128.0;
+ const want = @as(f32, 1.0) / @as(f32, @floatCast(std.math.pow(f64, @as(f32, 1e6), e)));
+ try std.testing.expectEqual(@as(u32, @bitCast(want)), @as(u32, @bitCast(v)));
+ }
+ try std.testing.expectEqual(@as(f32, 1), got[0]);
+ try std.testing.expectEqual(@as(f32, 1e-3), got[32]); // theta^(1/2): 1000 exactly, its f32 reciprocal
+ try std.testing.expect(got[63] < got[62] and got[63] > 1e-6);
+}
diff --git a/zig/src/families/dots_mocr/cuda.zig b/zig/src/families/dots_mocr/cuda.zig
new file mode 100644
index 00000000000..837f383434b
--- /dev/null
+++ b/zig/src/families/dots_mocr/cuda.zig
@@ -0,0 +1,61 @@
+//! dots.mocr's text decoder on CUDA: the dots recipe behind the shared dense engine (native/cuda.zig's family decls).
+const std = @import("std");
+const dense = @import("cuda_dense");
+const cfg = @import("config.zig");
+const fixtures = @import("config_test.zig");
+
+pub const recipe = @import("cuda_recipe.zig");
+pub const weights = @import("cuda_weights.zig");
+pub const kernels = @import("cuda_kernels.zig");
+pub const native = dense.Native(recipe);
+
+test "dots geometry from the released config: 12 query heads over 2 KV heads of 128, two query heads a pattn block" {
+ const c = try cfg.Config.parse(std.testing.allocator, fixtures.sample, fixtures.generation);
+ const g = recipe.geometryOf(c);
+ var plain = g;
+ plain.inv_freq = null;
+ try std.testing.expectEqual(dense.Geometry{ .hidden = 1536, .layers = 28, .heads = 12, .kv_heads = 2, .head_dim = 128, .vocab = 151936, .eps = 1e-6, .context = 131072, .acc = 2048, .wide = 17920, .mid = 8960, .theta = 1e6 }, plain);
+ try std.testing.expectEqual(@as(usize, 64), g.inv_freq.?.len); // Transformers' own table (cuda_recipe.zig)
+ const shape = try dense.recipe.attentionShape(g);
+ try std.testing.expect(shape.hpc == 2 and shape.grid_y == 6 and shape.rows == 64);
+}
+
+test "dots text bytes: embedding, untied head, final norm and 28 layers, the tower left out" {
+ const c = try cfg.Config.parse(std.testing.allocator, fixtures.sample, fixtures.generation);
+ const h: u64 = 1536;
+ const kv: u64 = 256;
+ const m: u64 = 8960;
+ const layer = 2 * h + (h + 2 * kv) * h + (h + 2 * kv) + h * h + 3 * m * h;
+ const want = (2 * 151936 * h + h + 28 * layer) * 2;
+ try std.testing.expectEqual(@as(u64, 3554176000), want);
+ try std.testing.expectEqual(want, try weights.textBytes(std.testing.allocator, c));
+}
+
+test "dots on CUDA: text only, no reasoning block, its config and tensor refusals worded" {
+ try std.testing.expect(!native.reasoning and native.keeps_no_prompts);
+ // compiles the load and the layer
+ const hooks = .{ &native.open, &native.attach, &native.fits, &native.ledgerPool, &native.textWeightBytes };
+ inline for (hooks) |h| try std.testing.expect(@intFromPtr(h) != 0);
+ try std.testing.expectEqualStrings("dots_ocr", native.model_type);
+ try std.testing.expectStringStartsWith(native.explain(null, error.DenseTextOnly).?, "dots.mocr on CUDA reads text prompts only");
+ for ([_]anyerror{ error.UnsupportedDotsConfig, error.BadDotsConfig, error.BadEos, error.DotsShape, error.DotsMissingTensor, error.DotsTensorCount, error.DotsIndexMismatch }) |e|
+ try std.testing.expect(native.explain(null, e) != null);
+}
+
+test "dots on CUDA reads its geometry from the checkpoint's files and refuses another model's config by name" {
+ const a = std.testing.allocator;
+ const io = std.testing.io;
+ var tmp = std.testing.tmpDir(.{});
+ defer tmp.cleanup();
+ try tmp.dir.writeFile(io, .{ .sub_path = "generation_config.json", .data = fixtures.generation });
+ try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = fixtures.sample });
+ const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path});
+ defer a.free(dir);
+ try std.testing.expectEqual(@as(usize, 2048), (try recipe.geometry(a, io, dir)).acc);
+ try std.testing.expectError(error.FileNotFound, recipe.weightBytes(a, io, dir)); // no shards: nothing to count
+ const other = try std.mem.replaceOwned(u8, a, fixtures.sample, "\"dots_ocr\"", "\"qwen3_5\"");
+ defer a.free(other);
+ try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = other });
+ try std.testing.expectError(error.UnsupportedDotsConfig, recipe.geometry(a, io, dir));
+ try std.testing.expect(native.explain(null, error.UnsupportedDotsConfig) != null);
+}
diff --git a/zig/src/families/dots_mocr/cuda_kernels.zig b/zig/src/families/dots_mocr/cuda_kernels.zig
new file mode 100644
index 00000000000..39a101a2e0c
--- /dev/null
+++ b/zig/src/families/dots_mocr/cuda_kernels.zig
@@ -0,0 +1,31 @@
+//! dots.mocr's own CUDA launch beside the dense engine's: the q|k|v bias and Torch-order RoPE (dots_mocr.cu).
+const cuda = @import("cuda");
+
+/// Threads a (row, head) block takes: dims j and j + 64 of a 128-dim head.
+const threads = 64;
+
+pub const Kernels = struct {
+ mod: cuda.Module,
+ rope: cuda.Function,
+
+ pub fn load(ctx: *const cuda.Context) !Kernels {
+ if (!cuda.kernels.available) return error.BuiltWithoutKernels;
+ var mod = try cuda.Module.load(ctx.d, cuda.kernels.dots_mocr);
+ errdefer mod.unload();
+ return .{ .mod = mod, .rope = try mod.function("tf_dots_qkv_rope") };
+ }
+
+ pub fn deinit(k: *Kernels) void {
+ k.mod.unload();
+ }
+
+ /// q [rows, h, 128] and kv [rows, k hk | v hk] from fp32 acc [rows, (h + 2 hk) 128] plus bias, rotated at pos.
+ pub fn qkvRope(k: *const Kernels, s: cuda.Stream, acc: u64, bias: u64, inv: u64, pos: u64, q: u64, kv: u64, rows: usize, h: usize, hk: usize) !void {
+ if (rows == 0) return;
+ var a: cuda.Args = .{};
+ for ([_]u64{ acc, bias, inv, pos, q, kv }) |v| a.add(v);
+ a.add(@as(c_int, @intCast(h)));
+ a.add(@as(c_int, @intCast(hk)));
+ try cuda.launch.launch(k.rope, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(h + 2 * hk), .z = 1 }, .block = .{ .x = threads } }, s, &a);
+ }
+};
diff --git a/zig/src/families/dots_mocr/cuda_recipe.zig b/zig/src/families/dots_mocr/cuda_recipe.zig
new file mode 100644
index 00000000000..3052c4d69a1
--- /dev/null
+++ b/zig/src/families/dots_mocr/cuda_recipe.zig
@@ -0,0 +1,97 @@
+//! dots.mocr's Qwen2 text decoder as a cuda_dense recipe: HIP's rounding points, its tower refused until it is built.
+const std = @import("std");
+const cuda = @import("cuda");
+const dense = @import("cuda_dense");
+const cfg = @import("config.zig");
+const wts = @import("cuda_weights.zig");
+const kern = @import("cuda_kernels.zig");
+
+const Recipe = @This();
+const P = *const dense.Pass(Recipe);
+
+pub const name = "dots.mocr";
+pub const model_type = "dots_ocr";
+pub const formats: []const []const u8 = &.{"bf16"};
+pub const default_context: i64 = 32768;
+/// Replies carry no reasoning block: the text is content.
+pub const reasoning = false;
+
+pub const Weights = wts.Weights;
+pub const Kernels = kern.Kernels;
+
+/// The engine's geometry of an admitted dots.mocr config.
+pub fn geometryOf(c: cfg.Config) dense.Geometry {
+ const h: usize = c.hidden;
+ return .{ .hidden = h, .layers = c.layers, .heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.headDim(), .vocab = c.vocab, .eps = c.eps, .context = c.context, .acc = h + 2 * @as(usize, c.kvWidth()), .wide = 2 * @as(usize, c.intermediate), .mid = c.intermediate, .theta = c.theta, .inv_freq = if (c.theta == 1e6 and c.headDim() == 128) &hf_inv_freq else null };
+}
+
+/// config.json and generation_config.json admitted as dots.mocr's, before any tensor is read.
+pub fn geometry(a: std.mem.Allocator, io: std.Io, dir: []const u8) !dense.Geometry {
+ return geometryOf(try cfg.Config.read(a, io, dir));
+}
+
+pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, dir: []const u8) !Weights {
+ return wts.load(gpa, io, d, s, dir);
+}
+
+/// The text decoder's bytes on the device, once every shard's header is checked.
+pub fn weightBytes(a: std.mem.Allocator, io: std.Io, dir: []const u8) !u64 {
+ const c = try cfg.Config.read(a, io, dir);
+ _ = try @import("weights.zig").inspect(a, io, dir, c);
+ return wts.textBytes(a, c);
+}
+
+pub fn embed(p: P) !void {
+ try p.k.embed.bf16(p.s, p.m.s.ids, p.w.embed, p.m.s.h, p.rows, p.g.hidden);
+}
+
+/// One Qwen2 layer: RMS, q|k|v sums, bias and RoPE in Torch's order, attention, o, RMS, SwiGLU.
+pub fn layer(p: P, i: usize) !void {
+ const x = p.m.s;
+ const l = p.w.layers[i];
+ try p.norm(if (i == 0) null else x.r, l.ln1);
+ try p.projectF32(l.qkv, x.x, x.acc);
+ try p.own.qkvRope(p.s, x.acc, l.bias, p.m.inv, x.pos, x.q, x.kv, p.rows, p.g.heads, p.g.kv_heads);
+ try p.attend(i, x.q, x.kv, x.y);
+ try p.project(l.o, x.y, x.r);
+ try p.norm(x.r, l.ln2);
+ try p.project(l.gate_up, x.x, x.wide);
+ try p.siluMul(x.wide, x.y, p.g.mid);
+ try p.project(l.down, x.y, x.r);
+}
+
+pub fn final(p: P) !void {
+ try p.norm(p.m.s.r, p.w.norm);
+}
+
+pub fn head(w: *const Weights) cuda.qmmf.Weight {
+ return w.head;
+}
+
+pub fn explain(err: anyerror) ?[]const u8 {
+ return switch (err) {
+ error.UnsupportedDotsConfig, error.BadDotsConfig, error.BadEos => "this checkpoint's config is not dots.mocr's (DotsOCRForCausalLM, Qwen2 1.5B text decoder in BF16), the one this engine serves",
+ error.DotsDType, error.DotsShape, error.DotsExtent, error.DotsMissingTensor, error.DotsTensorCount, error.DotsIndexMismatch => "this dots.mocr checkpoint's tensors are not the released BF16 ones: download it again",
+ else => null,
+ };
+}
+
+/// Transformers' own float32 inv_freq for theta 1e6 at head dim 128 (fixtures/README.md): Torch's CPU pow, bit for bit.
+const hf_inv_freq: [64]f32 = blk: {
+ @setEvalBranchQuota(20_000);
+ const text = @embedFile("fixtures/rope_inv_freq.txt");
+ var out: [64]f32 = undefined;
+ var it = std.mem.tokenizeAny(u8, text, ", \n");
+ for (&out) |*v| v.* = @bitCast(std.fmt.parseInt(u32, it.next().?, 0) catch unreachable);
+ if (it.next() != null) @compileError("fixtures/rope_inv_freq.txt holds more than 64 values");
+ break :blk out;
+};
+
+test "dots' rotary frequencies are Transformers' own: the correctly rounded pow differs at entry 37 only" {
+ var mine: [64]f32 = undefined;
+ @import("cuda_dense").rope.inverse(1e6, 128, &mine);
+ for (hf_inv_freq, mine, 0..) |hf, own, i| {
+ const one_ulp = @as(u32, @bitCast(hf)) -% @as(u32, @bitCast(own)) <= 1 or @as(u32, @bitCast(own)) -% @as(u32, @bitCast(hf)) <= 1;
+ try std.testing.expect(if (i == 37) hf != own and one_ulp else hf == own);
+ }
+}
diff --git a/zig/src/families/dots_mocr/cuda_weights.zig b/zig/src/families/dots_mocr/cuda_weights.zig
new file mode 100644
index 00000000000..edda9806954
--- /dev/null
+++ b/zig/src/families/dots_mocr/cuda_weights.zig
@@ -0,0 +1,231 @@
+//! dots.mocr's text decoder on CUDA: BF16 tensors on the device, q|k|v and gate|up fused into packed lane weights.
+const std = @import("std");
+const cuda = @import("cuda");
+const core = @import("core");
+const cfg = @import("config.zig");
+const inventory = @import("weights.zig");
+
+const DeviceBuffer = cuda.DeviceBuffer;
+const Weight = cuda.qmmf.Weight;
+
+/// The text decoder's names: the tower's (vision_tower.) stay on the host until the tower is built on CUDA.
+const text = [_][]const u8{ "model.", "lm_head." };
+
+pub const Layer = struct {
+ ln1: u64, // bf16 [hidden]: input_layernorm
+ ln2: u64, // bf16 [hidden]: post_attention_layernorm
+ qkv: Weight, // q|k|v [hidden + 2 kv width, hidden]
+ bias: u64, // bf16 [hidden + 2 kv width]: q|k|v biases
+ o: Weight,
+ gate_up: Weight, // gate|up [2 intermediate, hidden]
+ down: Weight,
+};
+
+pub const Weights = struct {
+ gpa: std.mem.Allocator,
+ config: cfg.Config,
+ buffers: std.ArrayList(DeviceBuffer) = .empty,
+ layers: []Layer = &.{},
+ embed: u64 = 0, // bf16 [vocab, hidden]
+ norm: u64 = 0, // bf16 [hidden]: model.norm
+ head: Weight = undefined, // the untied lm_head
+
+ pub fn deinit(w: *Weights) void {
+ for (w.buffers.items) |*b| b.free();
+ w.buffers.deinit(w.gpa);
+ w.gpa.free(w.layers);
+ }
+
+ fn alloc(w: *Weights, d: *const cuda.Driver, bytes: usize) !DeviceBuffer {
+ try w.buffers.ensureUnusedCapacity(w.gpa, 1);
+ const b = try DeviceBuffer.alloc(d, bytes);
+ w.buffers.appendAssumeCapacity(b);
+ return b;
+ }
+};
+
+/// The text tensors' bytes as stored, which their device copies take too (every lane weight's rows a multiple of 128).
+pub fn textBytes(gpa: std.mem.Allocator, c: cfg.Config) !u64 {
+ var arena: std.heap.ArenaAllocator = .init(gpa);
+ defer arena.deinit();
+ var total: u64 = 0;
+ for (try inventory.inventory(arena.allocator(), c)) |s| {
+ for (text) |p| if (std.mem.startsWith(u8, s.name, p)) {
+ total += s.bytes();
+ };
+ }
+ return total;
+}
+
+/// Reads the checkpoint's text tensors into the device on `s`, fusing and packing the projections.
+const Loader = struct {
+ w: *Weights,
+ ck: *core.Checkpoint,
+ d: *const cuda.Driver,
+ s: cuda.Stream,
+ pack: *const cuda.modelopt.Packer,
+ stage: DeviceBuffer,
+
+ /// A bf16 tensor as stored.
+ fn raw(l: *Loader, name: []const u8, shape: []const usize) !u64 {
+ return l.joined(&.{name}, &.{shape[0]}, if (shape.len == 2) shape[1] else 1);
+ }
+
+ /// bf16 tensors [rows[i], k] one after another in a buffer of their own.
+ fn joined(l: *Loader, names: []const []const u8, rows: []const usize, k: usize) !u64 {
+ var n: usize = 0;
+ for (rows) |r| n += r;
+ const b = try l.w.alloc(l.d, n * k * 2);
+ var at: usize = 0;
+ for (names, rows) |name, r| {
+ const t = try l.ck.expect(name, .bf16, if (k == 1) &.{r} else &.{ r, k });
+ try b.upload(at, t.bytes);
+ at += t.bytes.len;
+ }
+ return b.ptr;
+ }
+
+ /// bf16 projections [rows[i], k] stacked into one [sum rows, k] lane weight, packed on the device.
+ fn lane(l: *Loader, names: []const []const u8, rows: []const usize, k: usize) !Weight {
+ var n: usize = 0;
+ var at: usize = 0;
+ for (names, rows) |name, r| {
+ const t = try l.ck.expect(name, .bf16, &.{ r, k });
+ try l.stage.upload(at, t.bytes);
+ at += t.bytes.len;
+ n += r;
+ }
+ const b = try l.w.alloc(l.d, cuda.modelopt.padded(n) * k * 2);
+ try l.pack.bf16Lane(l.s, l.stage.ptr, b.ptr, n, k);
+ try l.s.synchronize(); // the next projection's upload reuses the stage
+ return cuda.modelopt.bf16Weight(b.ptr, n, k);
+ }
+
+ /// A failed load's work proved complete (the packing stream, then the uploads' null stream) before buffers go.
+ fn settle(l: *const Loader) void {
+ cuda.completion.drain(l.d, l.s, "a dots.mocr weight load did not drain before its buffers were freed");
+ cuda.completion.drain(l.d, .{ .d = l.d, .handle = null }, "dots.mocr weight uploads did not drain before their buffers were freed");
+ }
+};
+
+/// The text decoder of the dots.mocr checkpoint at `dir`, every header checked first; a tensor left over refuses.
+pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, dir: []const u8) !Weights {
+ const c = try cfg.Config.read(gpa, io, dir);
+ _ = try inventory.inspect(gpa, io, dir, c);
+ var ck = try core.Checkpoint.openModelPrefixes(gpa, io, dir, &text);
+ defer ck.close();
+ var pack = try cuda.modelopt.Packer.load(d);
+ defer pack.unload();
+ var w: Weights = .{ .gpa = gpa, .config = c };
+ errdefer w.deinit();
+ w.layers = try gpa.alloc(Layer, c.layers);
+ const h: usize = c.hidden;
+ const kvw: usize = c.kvWidth();
+ const m: usize = c.intermediate;
+ var stage = try DeviceBuffer.alloc(d, @as(usize, c.vocab) * h * 2);
+ defer stage.free();
+ var l: Loader = .{ .w = &w, .ck = &ck, .d = d, .s = s, .pack = &pack, .stage = stage };
+ errdefer l.settle(); // runs first on failure: the stage, the weights and the packer go after it
+ w.embed = try l.raw("model.embed_tokens.weight", &.{ c.vocab, h });
+ w.norm = try l.raw("model.norm.weight", &.{h});
+ w.head = try l.lane(&.{"lm_head.weight"}, &.{c.vocab}, h);
+ var names: [7][96]u8 = undefined;
+ for (w.layers, 0..) |*y, i| {
+ var n: [7][]const u8 = undefined;
+ for (&n, &names, [_][]const u8{ "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj" }) |*out, *buf, part|
+ out.* = try std.fmt.bufPrint(buf, "model.layers.{d}.{s}.weight", .{ i, part });
+ var bias: [3][96]u8 = undefined;
+ var b: [3][]const u8 = undefined;
+ for (&b, &bias, n[0..3]) |*out, *buf, weight| out.* = try std.fmt.bufPrint(buf, "{s}.bias", .{weight[0 .. weight.len - ".weight".len]});
+ var norms: [2][96]u8 = undefined;
+ y.* = .{
+ .ln1 = try l.raw(try std.fmt.bufPrint(&norms[0], "model.layers.{d}.input_layernorm.weight", .{i}), &.{h}),
+ .ln2 = try l.raw(try std.fmt.bufPrint(&norms[1], "model.layers.{d}.post_attention_layernorm.weight", .{i}), &.{h}),
+ .qkv = try l.lane(n[0..3], &.{ h, kvw, kvw }, h),
+ .bias = try l.joined(&b, &.{ h, kvw, kvw }, 1),
+ .o = try l.lane(n[3..4], &.{h}, h),
+ .gate_up = try l.lane(n[4..6], &.{ m, m }, h),
+ .down = try l.lane(n[6..7], &.{h}, m),
+ };
+ }
+ if (ck.unused() != 0) return error.DotsTensorCount;
+ try s.synchronize();
+ return w;
+}
+
+/// A device whose packing stream's wait fails: it counts buffers freed while a launch may still write them.
+const Fake = struct {
+ const queue: cuda.abi.Stream = @ptrFromInt(0x1800);
+ var next: u64 = 0x10000;
+ var live: usize = 0;
+ var pending = false;
+ var freed_pending: usize = 0;
+ var proofs: usize = 0;
+ fn alloc(out: *u64, len: usize) callconv(.c) cuda.abi.Result {
+ out.* = next;
+ next += len + 256;
+ live += 1;
+ return 0;
+ }
+ fn free(_: u64) callconv(.c) cuda.abi.Result {
+ freed_pending += @intFromBool(pending);
+ live -= 1;
+ return 0;
+ }
+ fn copy(_: u64, _: ?*const anyopaque, _: usize) callconv(.c) cuda.abi.Result {
+ return 0;
+ }
+ fn launch(_: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) cuda.abi.Result {
+ pending = true;
+ return 0;
+ }
+ fn wait(s: cuda.abi.Stream) callconv(.c) cuda.abi.Result {
+ return if (s == queue) cuda.abi.error_not_ready else 0;
+ }
+ fn context() callconv(.c) cuda.abi.Result {
+ proofs += 1;
+ pending = false;
+ return 0;
+ }
+ fn driver() cuda.Driver {
+ var d: cuda.Driver = .{ .lib = undefined, .api = undefined };
+ d.api.cuMemAlloc_v2 = alloc;
+ d.api.cuMemFree_v2 = free;
+ d.api.cuMemcpyHtoD_v2 = copy;
+ d.api.cuLaunchKernel = launch;
+ d.api.cuStreamSynchronize = wait;
+ d.api.cuCtxSynchronize = context;
+ return d;
+ }
+};
+
+test "a failed pack wait proves completion by the context before the stage or any weight is freed" {
+ const a = std.testing.allocator;
+ const io = std.testing.io;
+ var tmp = std.testing.tmpDir(.{});
+ defer tmp.cleanup();
+ const header = "{\"x.weight\":{\"dtype\":\"BF16\",\"shape\":[2,64],\"data_offsets\":[0,256]}}";
+ var file: [8 + header.len + 256]u8 = @splat(0);
+ std.mem.writeInt(u64, file[0..8], header.len, .little);
+ @memcpy(file[8..][0..header.len], header);
+ try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = &file });
+ const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path});
+ defer a.free(dir);
+ var ck = try core.Checkpoint.openModelPrefixes(a, io, dir, &.{});
+ defer ck.close();
+ var d = Fake.driver();
+ var pack: cuda.modelopt.Packer = undefined;
+ pack.bf16_lane = .{ .d = &d, .handle = @ptrFromInt(1) };
+ var w: Weights = .{ .gpa = a, .config = undefined };
+ var stage = try DeviceBuffer.alloc(&d, 256);
+ const l: Loader = .{ .w = &w, .ck = &ck, .d = &d, .s = .{ .d = &d, .handle = Fake.queue }, .pack = &pack, .stage = stage };
+ var run = l;
+ try std.testing.expectError(error.NotReady, run.lane(&.{"x.weight"}, &.{2}, 64));
+ try std.testing.expect(Fake.pending and Fake.live == 2);
+ l.settle();
+ stage.free();
+ w.deinit();
+ try std.testing.expectEqual(@as(usize, 0), Fake.freed_pending);
+ try std.testing.expectEqual(@as(usize, 0), Fake.live);
+ try std.testing.expectEqual(@as(usize, 1), Fake.proofs);
+}
diff --git a/zig/src/families/dots_mocr/fixtures/README.md b/zig/src/families/dots_mocr/fixtures/README.md
index f75f365ad29..75fbffe9e66 100644
--- a/zig/src/families/dots_mocr/fixtures/README.md
+++ b/zig/src/families/dots_mocr/fixtures/README.md
@@ -50,3 +50,9 @@ FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE \
OUTPUT_DIR --pci PCI_ID --max-new-tokens 64
python tools/zig/dots_mocr_reference.py compare OUTPUT_DIR NATIVE_PREFIX
```
+
+`rope_inv_freq.txt` holds the 64 float32 bit patterns of the checkpoint's
+rotary `inv_freq` as Transformers 4.57.6 computes them on the CPU (theta 1e6,
+head dim 128), dumped by the oracle's `text` mode (`rope-inv_freq.npy`). The
+CUDA recipe uploads them as they are, because Torch's CPU `powf` is not
+correctly rounded: entry 37 is one ulp from the correctly rounded value.
diff --git a/zig/src/families/dots_mocr/fixtures/rope_inv_freq.txt b/zig/src/families/dots_mocr/fixtures/rope_inv_freq.txt
new file mode 100644
index 00000000000..8d31a9bc1db
--- /dev/null
+++ b/zig/src/families/dots_mocr/fixtures/rope_inv_freq.txt
@@ -0,0 +1 @@
+0x3f800000,0x3f4e4bad,0x3f263de0,0x3f05f6ef,0x3ed7e89b,0x3eadfcff,0x3e8c3504,0x3e61f835,0x3e361887,0x3e12bd91,0x3dec7fd6,0x3dbe94c6,0x3d99940d,0x3d778513,0x3d47763f,0x3d20bc1d,0x3d0186e3,0x3cd0c1a8,0x3ca8398b,0x3c879008,0x3c5a7bf2,0x3c301052,0x3c0de12d,0x3be4aa46,0x3bb8449c,0x3b947dae,0x3b6f520e,0x3b40dac5,0x3b1b690d,0x3afa78f0,0x3ac9d75c,0x3aa2a6f7,0x3a83126f,0x3a533f28,0x3a2a3b44,0x3a092e02,0x39dd1725,0x39b229fb,0x398f9272,0x39676492,0x393a7753,0x39164324,0x38f22ce2,0x38c327b5,0x389d43a4,0x387d75d4,0x384c3fbe,0x382497ab,0x3804a2b2,0x37d5c441,0x37ac431d,0x378ad0ed,0x375fba4f,0x37344a0e,0x371148e3,0x36ea2732,0x36bcb0c1,0x36980e02,0x36751070,0x36457bab,0x361f23e4,0x36003dec,0x35ceaf79,0x35a68e4c
diff --git a/zig/src/families/glm/cuda_experts.zig b/zig/src/families/glm/cuda_experts.zig
index 8eca3815e3b..48f4c2b5c71 100644
--- a/zig/src/families/glm/cuda_experts.zig
+++ b/zig/src/families/glm/cuda_experts.zig
@@ -60,13 +60,13 @@ pub fn run(ks: *const Kernels, stream: cuda.Stream, which: Which, l: cuda.expert
};
const in: cuda.experts.Rows = .{ .x = x, .stride = x_stride, .slots = plan.slots };
const pairs = rows * plan.slots;
- if (plan.prompt) return g.prompt(stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, staged_tile), skip);
+ if (plan.prompt) return g.prompt(stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, staged_tile), pairs, staged_tile, skip);
if (plan.fact) {
// gate-up reads two matrices a K tile and gains from the factored tiles; down is short and keeps the waves
if (epi == .up and !l.relu2) return cuda.experts.Nvfp4Experts.runFact(&ks.nvfp4, stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, staged_tile), skip);
return cuda.experts_waves.run(&ks.nvfp4, stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, tile), skip);
}
- return g.decode(stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, tile), skip);
+ return g.decode(stream, epi, in, l, plan.route, out, maxItems(pairs, plan.experts, tile), pairs, tile, skip);
}
test "block 32's kernels are wanted for a csf0_v1 module alone, not for block 16's recipes or a single-format NVFP4" {
diff --git a/zig/src/families/kolibri1/forward.zig b/zig/src/families/kolibri1/forward.zig
index dcfa7d3abb8..28b504135cc 100644
--- a/zig/src/families/kolibri1/forward.zig
+++ b/zig/src/families/kolibri1/forward.zig
@@ -464,11 +464,11 @@ pub const Model = struct {
const up_in: experts.Rows = .{ .x = m.s.x, .stride = c.hidden, .slots = slots_e };
const down_in: experts.Rows = .{ .x = m.s.act, .stride = c.moe_width };
if (prompt) {
- try gx.prompt(s, .up, up_in, L.experts, m.s.plan, m.s.act, items, -1);
- try gx.prompt(s, .down_bf16, down_in, L.experts, m.s.plan, m.s.y, items, -1);
+ try gx.prompt(s, .up, up_in, L.experts, m.s.plan, m.s.act, items, pairs, tile, -1);
+ try gx.prompt(s, .down_bf16, down_in, L.experts, m.s.plan, m.s.y, items, pairs, tile, -1);
} else {
- try gx.decode(s, .up, up_in, L.experts, m.s.plan, m.s.act, items, -1);
- try gx.decode(s, .down_f32, down_in, L.experts, m.s.plan, m.s.y, items, -1);
+ try gx.decode(s, .up, up_in, L.experts, m.s.plan, m.s.act, items, pairs, tile, -1);
+ try gx.decode(s, .down_f32, down_in, L.experts, m.s.plan, m.s.y, items, pairs, tile, -1);
}
try o.combine(m.s.y, prompt, m.s.wts, m.s.mo, n, slots_e, c.hidden);
const after = if (i + 1 < w.layers.len) w.layers[i + 1].input_norm else w.norm;
diff --git a/zig/src/families/lfm2_moe/cuda_forward.zig b/zig/src/families/lfm2_moe/cuda_forward.zig
index 6c2b7ef3577..b0ac16a93af 100644
--- a/zig/src/families/lfm2_moe/cuda_forward.zig
+++ b/zig/src/families/lfm2_moe/cuda_forward.zig
@@ -133,11 +133,11 @@ const Pass = struct {
const rows_in: cuda.experts.Rows = .{ .x = x.x, .stride = c.hidden, .slots = c.topk };
const act_in: cuda.experts.Rows = .{ .x = x.eact, .stride = c.moe_intermediate };
if (p.prompt) {
- try grouped.prompt(p.s, .up, rows_in, e.experts, plan, x.eact, items, -1);
- try grouped.prompt(p.s, .down_bf16, act_in, e.experts, plan, x.ey, items, -1);
+ try grouped.prompt(p.s, .up, rows_in, e.experts, plan, x.eact, items, pairs, tile, -1);
+ try grouped.prompt(p.s, .down_bf16, act_in, e.experts, plan, x.ey, items, pairs, tile, -1);
} else {
- try grouped.decode(p.s, .up, rows_in, e.experts, plan, x.eact, items, -1);
- try grouped.decode(p.s, .down_f32, act_in, e.experts, plan, x.ey, items, -1);
+ try grouped.decode(p.s, .up, rows_in, e.experts, plan, x.eact, items, pairs, tile, -1);
+ try grouped.decode(p.s, .down_f32, act_in, e.experts, plan, x.ey, items, pairs, tile, -1);
}
try k.moe.combine(p.s, x.ey, !p.prompt, x.wts, x.r, p.rows, c.topk, c.hidden, 0);
}
diff --git a/zig/src/families/minimax_m3/cuda_forward.zig b/zig/src/families/minimax_m3/cuda_forward.zig
index 1f1298cd30c..42fbbb2a084 100644
--- a/zig/src/families/minimax_m3/cuda_forward.zig
+++ b/zig/src/families/minimax_m3/cuda_forward.zig
@@ -61,10 +61,11 @@ pub const Modules = struct {
/// The grouped experts' scratch for a window of `rows` (allocated before serving, never mid-forward).
pub fn reserve(m: *const Modules, c: *const cfg.Config, share: anytype, rows: usize, prefill: bool) !void {
- const items = cuda.grouped.maxItems(rows * c.topk, c.experts, if (prefill) 64 else 16);
+ const pairs = rows * c.topk;
+ const tile: usize = if (prefill) 64 else 16;
const ni = share.moe.len();
- const up = try cuda.exl3_grouped.Grouped.bytes(c.hidden, ni, items, try cuda.exl3_experts.plan(c.hidden, ni, true), 2);
- const down = try cuda.exl3_grouped.Grouped.bytes(ni, c.hidden, items, try cuda.exl3_experts.plan(ni, c.hidden, false), 1);
+ const up = try cuda.exl3_grouped.Grouped.bytes(c.hidden, ni, pairs, tile, try cuda.exl3_experts.plan(c.hidden, ni, true), 2);
+ const down = try cuda.exl3_grouped.Grouped.bytes(ni, c.hidden, pairs, tile, try cuda.exl3_experts.plan(ni, c.hidden, false), 1);
const need = @max(up, down);
const g = &m.x3.grouped;
if (g.storage != null and g.storage.?.len >= need) return;
@@ -243,8 +244,8 @@ pub const Forward = struct {
const in_up: cuda.experts.Rows = .{ .x = b.x, .stride = h, .slots = topk };
const in_down: cuda.experts.Rows = .{ .x = b.eact, .stride = m.routed.width };
// one path for a prompt chunk and a window: every pick's down rows in fp32, so a row's bits never depend on it
- try g.decode(f.stream, .up, in_up, m.routed, b.plan(), b.eact, b.max_items, -1);
- try g.decode(f.stream, .down_f32, in_down, m.routed, b.plan(), b.ey, b.max_items, -1);
+ try g.decode(f.stream, .up, in_up, m.routed, b.plan(), b.eact, b.max_items, pairs, tile, -1);
+ try g.decode(f.stream, .down_f32, in_down, m.routed, b.plan(), b.ey, b.max_items, pairs, tile, -1);
try f.mlp(m.shared, f.w.share.shared.len(), R, b.sy, .f32);
try kern.moePartial(&f.m.k, f.stream, b.ey, b.wts, b.sy, f.slot(R, 2), R, topk, h);
}
diff --git a/zig/src/families/nemotron/cuda_forward.zig b/zig/src/families/nemotron/cuda_forward.zig
index 02b01c2c5b5..0bf7a6cc629 100644
--- a/zig/src/families/nemotron/cuda_forward.zig
+++ b/zig/src/families/nemotron/cuda_forward.zig
@@ -274,13 +274,13 @@ pub const Forward = struct {
const g = o.grouped();
const in = f.rowsIn(ex);
if (prompt) {
- try g.prompt(o.s, .up, in[0], ex, b.plan, b.act, items, -1);
- try g.prompt(o.s, .down_bf16, in[1], ex, b.plan, b.ymoe, items, -1);
+ try g.prompt(o.s, .up, in[0], ex, b.plan, b.act, items, pairs, tile, -1);
+ try g.prompt(o.s, .down_bf16, in[1], ex, b.plan, b.ymoe, items, pairs, tile, -1);
if (f.dump) |d| for ([_][]const u8{ "pick", "act", "ymoe" }, [_]u64{ b.pick, b.act, b.ymoe }, [_]usize{ 4, ex.width * 2, ex.dims * 2 }) |l, p, w| try d.input(o, l, p, pairs * w);
} else {
- try g.decode(o.s, .up, in[0], ex, b.plan, b.act, items, -1);
+ try g.decode(o.s, .up, in[0], ex, b.plan, b.act, items, pairs, tile, -1);
try f.mark(.experts_up);
- try g.decode(o.s, .down_f32, in[1], ex, b.plan, b.ymoe, items, -1);
+ try g.decode(o.s, .down_f32, in[1], ex, b.plan, b.ymoe, items, pairs, tile, -1);
try f.mark(.experts_down);
}
try f.sharedRest(o, m, rows, !prompt);
@@ -305,12 +305,12 @@ pub const Forward = struct {
const up_in: cuda.experts.Rows = .{ .x = b.lat, .stride = c.latent, .slots = K };
const down_in: cuda.experts.Rows = .{ .x = b.act, .stride = ex.width };
if (prompt) {
- try g.prompt(o.s, .up, up_in, ex, b.plan, b.act, items, -1);
- try g.prompt(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, -1);
+ try g.prompt(o.s, .up, up_in, ex, b.plan, b.act, items, pairs, tile, -1);
+ try g.prompt(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, pairs, tile, -1);
} else {
- try g.decode(o.s, .up, up_in, ex, b.plan, b.act, items, -1);
+ try g.decode(o.s, .up, up_in, ex, b.plan, b.act, items, pairs, tile, -1);
try f.mark(.experts_up);
- try g.decode(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, -1);
+ try g.decode(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, pairs, tile, -1);
try f.mark(.experts_down);
}
try o.latentSum(b.ymoe, b.wts, b.comb, rows, K, c.latent);
@@ -322,11 +322,11 @@ pub const Forward = struct {
const sh_up: cuda.experts.Rows = .{ .x = b.y, .stride = c.hidden, .slots = 1 };
const sh_down: cuda.experts.Rows = .{ .x = b.act, .stride = sh.width };
if (prompt) {
- try g.prompt(o.s, .up, sh_up, sh, b.plan, b.act, items, -1);
- try g.prompt(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, -1);
+ try g.prompt(o.s, .up, sh_up, sh, b.plan, b.act, items, rows, tile, -1);
+ try g.prompt(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, rows, tile, -1);
} else {
- try g.decode(o.s, .up, sh_up, sh, b.plan, b.act, items, -1);
- try g.decode(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, -1);
+ try g.decode(o.s, .up, sh_up, sh, b.plan, b.act, items, rows, tile, -1);
+ try g.decode(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, rows, tile, -1);
}
try o.addBf16(b.delta, b.delta, b.ymoe, true, rows * c.hidden);
try f.mark(.shared_wait);
@@ -373,8 +373,9 @@ pub const Forward = struct {
try sd.fork.record(o.s);
try sd.s.wait(sd.fork);
const sp = b.sharedPlan(rows);
- try g.decode(sd.s, .up, in[0], ex, sp, b.act, 2, -1);
- try g.decode(sd.s, .down_f32, in[1], ex, sp, b.ymoe, 2, -1);
+ const pairs = rows * c.slots();
+ try g.decode(sd.s, .up, in[0], ex, sp, b.act, 2, pairs, 16, -1);
+ try g.decode(sd.s, .down_f32, in[1], ex, sp, b.ymoe, 2, pairs, 16, -1);
try f.sharedRest(so, m, rows, true);
try sd.join.record(sd.s);
try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.slots(), c.routed_scaling, c.norm_topk);
@@ -382,9 +383,9 @@ pub const Forward = struct {
try o.planRouted(b.pick, rows, c.slots(), c.top_k, c.experts, 16, b.plan);
try f.mark(.plan);
const items = kern.maxItems(rows * c.top_k, c.experts, 16);
- try g.decode(o.s, .up, in[0], ex, b.plan, b.act, items, -1);
+ try g.decode(o.s, .up, in[0], ex, b.plan, b.act, items, pairs, 16, -1);
try f.mark(.experts_up);
- try g.decode(o.s, .down_f32, in[1], ex, b.plan, b.ymoe, items, -1);
+ try g.decode(o.s, .down_f32, in[1], ex, b.plan, b.ymoe, items, pairs, 16, -1);
try f.mark(.experts_down);
try o.s.wait(sd.join);
try f.mark(.shared_wait);
diff --git a/zig/src/families/nemotron/cuda_latent.zig b/zig/src/families/nemotron/cuda_latent.zig
index ffdb4ded492..a119d6cbf5b 100644
--- a/zig/src/families/nemotron/cuda_latent.zig
+++ b/zig/src/families/nemotron/cuda_latent.zig
@@ -339,8 +339,8 @@ fn routed(o: kern.Ops, t: Tables, s: Scratch, picks: u64, lat: u64, rows: usize,
const g = o.grouped();
try o.plan(picks, pairs, t.routed.experts, 16, s.plan);
const items = kern.maxItems(pairs, t.routed.experts, 16);
- try g.decode(o.s, .up, .{ .x = lat, .stride = c.latent, .slots = c.top_k * own }, t.routed, s.plan, s.act, items, -1);
- try g.decode(o.s, .down_f32, .{ .x = s.act, .stride = t.routed.width }, t.routed, s.plan, out, items, -1);
+ try g.decode(o.s, .up, .{ .x = lat, .stride = c.latent, .slots = c.top_k * own }, t.routed, s.plan, s.act, items, pairs, 16, -1);
+ try g.decode(o.s, .down_f32, .{ .x = s.act, .stride = t.routed.width }, t.routed, s.plan, out, items, pairs, 16, -1);
}
/// The shared expert's (row, share) pairs into `out` (partials' second region), on `s`'s fixed plan where it has one.
@@ -350,8 +350,8 @@ fn sharedRows(o: kern.Ops, t: Tables, s: Scratch, hidden: u64, rows: usize, own:
const sp = s.shared_plan orelse s.plan;
if (s.shared_plan == null) try o.plan(s.sids, rows * own, own, 16, sp);
const sh_items = kern.maxItems(rows * own, own, 16);
- try g.decode(o.s, .up, .{ .x = hidden, .stride = c.hidden, .slots = own }, t.shared, sp, s.sact, sh_items, -1);
- try g.decode(o.s, .down_f32, .{ .x = s.sact, .stride = t.shared.width }, t.shared, sp, out, sh_items, -1);
+ try g.decode(o.s, .up, .{ .x = hidden, .stride = c.hidden, .slots = own }, t.shared, sp, s.sact, sh_items, rows * own, 16, -1);
+ try g.decode(o.s, .down_f32, .{ .x = s.sact, .stride = t.shared.width }, t.shared, sp, out, sh_items, rows * own, 16, -1);
}
/// The MoE block's delta into b.delta: the other ranks' shares asked for, ours run, all summed in share order.
diff --git a/zig/src/families/nemotron/cuda_train_fwd.zig b/zig/src/families/nemotron/cuda_train_fwd.zig
index cb4ad780e00..7c0c4d30449 100644
--- a/zig/src/families/nemotron/cuda_train_fwd.zig
+++ b/zig/src/families/nemotron/cuda_train_fwd.zig
@@ -92,8 +92,8 @@ pub fn experts(f: *const fwd.Forward, m: weights.MoE, rows: usize, scale: Scale,
const items = kern.maxItems(pairs, ex.experts, tile);
const g = o.grouped();
const in = f.rowsIn(ex);
- try g.prompt(o.s, .up, in[0], ex, b.plan, b.act, items, -1);
- try g.prompt(o.s, .down_bf16, in[1], ex, b.plan, b.ymoe, items, -1);
+ try g.prompt(o.s, .up, in[0], ex, b.plan, b.act, items, pairs, tile, -1);
+ try g.prompt(o.s, .down_bf16, in[1], ex, b.plan, b.ymoe, items, pairs, tile, -1);
for (m.rest, 0..) |r, h| if (r != 0) try o.restRows(b.act, ex.width, r, b.ymoe, ex.dims, false, rows, c.slots(), c.top_k + h, ex.dims, ex.width);
const t: train_ops.Ops = .{ .k = &o.k.train, .s = o.s };
if (m.pocket) |p| try t.pocketRows(b.y, ex.dims, p.up, p.down, p.n, b.ymoe, false, rows, c.slots(), c.top_k, scale.rows, scale.from[layer]);
diff --git a/zig/src/families/qwen3_5/cuda_exl3_weights.zig b/zig/src/families/qwen3_5/cuda_exl3_weights.zig
index bf1f1c6b1d5..8e001a8ddde 100644
--- a/zig/src/families/qwen3_5/cuda_exl3_weights.zig
+++ b/zig/src/families/qwen3_5/cuda_exl3_weights.zig
@@ -33,14 +33,17 @@ pub const Kit = struct {
k.grouped = .{ .d = d, .kernels = &k.experts };
if (!g.moe()) return k; // a dense checkpoint's MLP is three Dense projections
// the grouped experts' scratch grows on first use: allocate it before any graph capture, for the widest call
- const pairs = engine.max_rows * g.slots();
- const items = cuda.grouped.maxItems(pairs, g.experts + 1, 16);
- const up = try cuda.exl3_grouped.Grouped.bytes(g.hidden, g.intermediate, items, try ex.plan(g.hidden, g.intermediate, true), 2);
- const down = try cuda.exl3_grouped.Grouped.bytes(g.intermediate, g.hidden, items, try ex.plan(g.intermediate, g.hidden, false), 1);
- k.grouped.storage = try cuda.DeviceBuffer.alloc(d, @max(up, down));
+ k.grouped.storage = try cuda.DeviceBuffer.alloc(d, try storageBytes(g));
return k;
}
+ pub fn storageBytes(g: geometry.Geometry) !usize {
+ const pairs = engine.max_rows * g.slots();
+ const up = try cuda.exl3_grouped.Grouped.bytes(g.hidden, g.intermediate, pairs, 64, try ex.plan(g.hidden, g.intermediate, true), 2);
+ const down = try cuda.exl3_grouped.Grouped.bytes(g.intermediate, g.hidden, pairs, 64, try ex.plan(g.intermediate, g.hidden, false), 1);
+ return @max(up, down);
+ }
+
pub fn deinit(k: *Kit, gpa: std.mem.Allocator) void {
k.grouped.deinit();
k.b16.unload();
diff --git a/zig/src/families/qwen3_5/cuda_routed.zig b/zig/src/families/qwen3_5/cuda_routed.zig
index 79c31dcaa16..636d1bfc673 100644
--- a/zig/src/families/qwen3_5/cuda_routed.zig
+++ b/zig/src/families/qwen3_5/cuda_routed.zig
@@ -57,11 +57,11 @@ pub fn rows(e: *E.Engine, router: u64, layer: cuda.experts.Layer, x: u64, out: u
try launch(k, e.stream, .up, rows_in, layer, plan, s.act, items, -1);
try launch(k, e.stream, .down_f32, act_in, layer, plan, s.y, items, -1);
} else if (wide) {
- try grouped.prompt(e.stream, .up, rows_in, layer, plan, s.act, items, -1);
- try grouped.prompt(e.stream, .down_bf16, act_in, layer, plan, s.y, items, -1);
+ try grouped.prompt(e.stream, .up, rows_in, layer, plan, s.act, items, pairs, tile, -1);
+ try grouped.prompt(e.stream, .down_bf16, act_in, layer, plan, s.y, items, pairs, tile, -1);
} else {
- try grouped.decode(e.stream, .up, rows_in, layer, plan, s.act, items, -1);
- try grouped.decode(e.stream, .down_f32, act_in, layer, plan, s.y, items, -1);
+ try grouped.decode(e.stream, .up, rows_in, layer, plan, s.act, items, pairs, tile, -1);
+ try grouped.decode(e.stream, .down_f32, act_in, layer, plan, s.y, items, pairs, tile, -1);
}
try e.moe.combine(e.stream, s.y, !wide or nvfp4, s.wts, out, total, slots, g.hidden, params);
}
diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig
index 10d955392e3..57c34a6bac1 100644
--- a/zig/src/native/cuda.zig
+++ b/zig/src/native/cuda.zig
@@ -490,12 +490,12 @@ fn reasoning(comptime F: type) bool {
return if (@hasDecl(F, "reasoning")) F.reasoning else true;
}
-test "a CUDA family declares replies without reasoning with `pub const reasoning = false`; every served one keeps it" {
+test "a CUDA family declares replies without reasoning with `pub const reasoning = false`; only dots.mocr does" {
try std.testing.expect(reasoning(struct {}));
try std.testing.expect(!reasoning(struct {
pub const reasoning = false;
}));
- inline for (registry) |F| try std.testing.expect(reasoning(F));
+ inline for (registry) |F| try std.testing.expectEqual(!std.mem.eql(u8, F.model_type, "dots_ocr"), reasoning(F));
}
test "the Info openWith serves carries each family's reasoning declaration over the host's" {
diff --git a/zig/src/native/cuda_families.zig b/zig/src/native/cuda_families.zig
index e4bc3c2a493..de591b1b621 100644
--- a/zig/src/native/cuda_families.zig
+++ b/zig/src/native/cuda_families.zig
@@ -14,7 +14,7 @@ const dory = @import("dory");
const qwen3 = @import("qwen3_cuda");
/// CUDA families provide metadata, open their lane backend and explain their refusals.
-pub const registry = .{ nemotron.native, nemotron.omni, glm.native, minimax.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native, @import("lfm2_moe_cuda").native };
+pub const registry = .{ nemotron.native, nemotron.omni, glm.native, minimax.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native, @import("dots_mocr_cuda").native, @import("lfm2_moe_cuda").native };
pub const families: []const api.Family = blk: {
var out: [registry.len + 5]api.Family = undefined;
for (registry, 0..) |F, i| out[i] = .{ .model_type = F.model_type, .formats = F.formats, .decision_formats = if (@hasDecl(F, "decision_formats")) F.decision_formats else &.{} };
@@ -37,6 +37,14 @@ test "every registered family and d1's decision engine are listed for capabiliti
try std.testing.expectEqualStrings("qwen3", families[registry.len + 4].model_type);
}
+test "dots.mocr serves its publisher's BF16 weights on CUDA through the dense engine" {
+ for (families) |f| if (std.mem.eql(u8, f.model_type, "dots_ocr")) {
+ try std.testing.expectEqual(@as(usize, 1), f.formats.len);
+ return std.testing.expectEqualStrings("bf16", f.formats[0]);
+ };
+ return error.NoDotsFamily;
+}
+
test "GLM capabilities: CUDA offers the EXL3 format its loader accepts" {
const config_json = "{\"quantization_config\":{\"quant_method\":\"exl3\",\"bits\":4,\"codebook\":\"mcg\"}}";
var why: @import("core").modelopt.Why = .{};
diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig
index fa601d6bbd2..02f9ef56dc2 100644
--- a/zig/tests/cuda/contract_limits.zig
+++ b/zig/tests/cuda/contract_limits.zig
@@ -30,7 +30,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 {
if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes;
inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple", "gemma4_ple" }) |n|
if (std.mem.eql(u8, name, n)) return @field(kk, n);
- inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "nvfp4_experts32", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image", "lfm2_moe", "dots_mocr" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n);
+ inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "nvfp4_experts32", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image", "lfm2_moe", "lfm2", "dots_mocr" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n);
if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{};
return error.UnknownContractImage;
}
diff --git a/zig/tests/cuda/exl3_grouped_oai_tests.zig b/zig/tests/cuda/exl3_grouped_oai_tests.zig
index f08573d4885..063e27c416d 100644
--- a/zig/tests/cuda/exl3_grouped_oai_tests.zig
+++ b/zig/tests/cuda/exl3_grouped_oai_tests.zig
@@ -76,7 +76,7 @@ fn up(r: *Rig, k: *Kit, l: cuda.experts.Layer, x: u64, pick: u64, m: usize, tile
};
try k.router.route(r.s, pick, pairs, E, tile, plan);
const out = try r.zeros(pairs * I * 2);
- try k.g.run(r.s, .up, .{ .x = x, .stride = D, .slots = SLOTS }, l, plan, out, items, -1, false);
+ try k.g.run(r.s, .up, .{ .x = x, .stride = D, .slots = SLOTS }, l, plan, out, items, pairs, tile, -1);
return r.back(out, pairs * I * 2);
}
diff --git a/zig/tests/cuda/exl3_grouped_tile_tests.zig b/zig/tests/cuda/exl3_grouped_tile_tests.zig
new file mode 100644
index 00000000000..39bb79cc44a
--- /dev/null
+++ b/zig/tests/cuda/exl3_grouped_tile_tests.zig
@@ -0,0 +1,100 @@
+//! Grouped EXL3: maxm 16 and the old maxm 64 are bit-equal; maxm 16 on a tile-64 plan drops members.
+const std = @import("std");
+const cuda = @import("cuda");
+const check = @import("check.zig");
+const quant_ref_tests = @import("quant_ref_tests.zig");
+const f64t = @import("exl3_experts_f64_tests.zig");
+
+const Rig = quant_ref_tests.RigT;
+const E = f64t.E;
+const D = f64t.D;
+const I = f64t.I;
+const SLOTS = f64t.SLOTS;
+
+const Plan = struct { p: cuda.grouped.Plan, items: usize, pairs: usize };
+const Out = struct { dev: u64, bytes: []u8 };
+
+pub fn run(gpu: check.Gpu) !void {
+ var arena = std.heap.ArenaAllocator.init(gpu.gpa);
+ defer arena.deinit();
+ var s = try cuda.Stream.init(gpu.d, false);
+ defer s.deinit();
+ var prng = std.Random.DefaultPrng.init(61);
+ var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() };
+ defer {
+ for (r.bufs.items) |*b| b.free();
+ r.bufs.deinit(gpu.gpa);
+ }
+ var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts);
+ defer mod.unload();
+ const router = try cuda.grouped.Router.resolve(mod);
+ var exp = try cuda.exl3_experts.Experts.load(gpu.d);
+ defer exp.unload();
+ const b = try f64t.buildLayer(&r, &.{ 3, 4, 5, 6, 3, 4, 5, 6, 3 });
+ const l: cuda.experts.Layer = .{ .format = .exl3, .exl3 = b.layer, .up = 0, .down = 0, .width = I, .dims = D, .experts = E };
+ var g: cuda.exl3_grouped.Grouped = .{ .d = gpu.d, .kernels = &exp };
+ defer g.deinit();
+ const grp: cuda.experts.Grouped = .{ .exl3 = &g };
+
+ for ([_]usize{ 1, 4, 16 }) |rows| {
+ const pairs = rows * SLOTS;
+ const p16 = try plan(&r, router, b.dpick, pairs, 16);
+ const p64 = try plan(&r, router, b.dpick, pairs, 64);
+ const up16 = try upBytes(&r, grp, l, b.dx, p16, 16);
+ const upOld = try upBytes(&r, grp, l, b.dx, p16, 64);
+ try check.sameBytes("up: maxm 16 vs the old maxm 64 on a tile-16 plan", upOld.bytes, up16.bytes);
+ const upT64 = try upBytes(&r, grp, l, b.dx, p64, 64);
+ try check.sameBytes("up: a tile-64 plan through decode()", upT64.bytes, up16.bytes);
+ const down16 = try downBytes(&r, grp, l, up16.dev, p16, 16);
+ const downOld = try downBytes(&r, grp, l, up16.dev, p16, 64);
+ try check.sameBytes("down: maxm 16 vs the old maxm 64 on a tile-16 plan", downOld, down16);
+ const downT64 = try downBytes(&r, grp, l, upT64.dev, p64, 64);
+ try check.sameBytes("down: a tile-64 plan through decode()", downT64, down16);
+ check.pass("exl3-grouped-tile {d} rows: maxm 16 and the old 64 bit-equal, tile-64 plan bit-equal", .{rows});
+ }
+
+ // 16 rows' 128 pairs all pick expert 0: two 64-member items, so the maxm-16 mutant drops members 16..63.
+ const skew_host = try r.a.alloc(i32, 16 * SLOTS);
+ @memset(skew_host, 0);
+ const skew = try r.dev(std.mem.sliceAsBytes(skew_host));
+ const mp = try plan(&r, router, skew, 16 * SLOTS, 64);
+ const poison_x = try r.bfs(16 * SLOTS * D, 0.5);
+ const poison_x_dev = try r.dev(std.mem.sliceAsBytes(poison_x));
+ var gm: cuda.exl3_grouped.Grouped = .{ .d = gpu.d, .kernels = &exp };
+ defer gm.deinit();
+ const gmgrp: cuda.experts.Grouped = .{ .exl3 = &gm };
+ const poisoned = try upBytes(&r, gmgrp, l, poison_x_dev, mp, 64);
+ const mutant = try upBytes(&r, gmgrp, l, b.dx, mp, 16);
+ var gc: cuda.exl3_grouped.Grouped = .{ .d = gpu.d, .kernels = &exp };
+ defer gc.deinit();
+ const real = try upBytes(&r, .{ .exl3 = &gc }, l, b.dx, mp, 64);
+ const row: usize = 16 * I * 2;
+ try check.sameBytes("planted mutant keeps the poisoned members 16..31", mutant.bytes[row..][0..row], poisoned.bytes[row..][0..row]);
+ try check.expect(!std.mem.eql(u8, mutant.bytes, real.bytes), "planted maxm 16 on a tile-64 plan dropped members and must differ from maxm 64", .{});
+ check.pass("exl3-grouped-tile planted mutant (maxm 16 on a tile-64 plan) drops members 16..63", .{});
+}
+
+fn plan(r: *Rig, router: cuda.grouped.Router, pick: u64, pairs: usize, tile: usize) !Plan {
+ const items = cuda.grouped.maxItems(pairs, E, tile);
+ const p: cuda.grouped.Plan = .{
+ .members = try r.zeros(pairs * 4 + 16),
+ .items = try r.zeros(items * 3 * 4 + 16),
+ .counts = try r.zeros(16),
+ .rank = try r.zeros(pairs * 4 + 16),
+ .hist = try r.zeros((E + 1) * 4 * ((pairs + 1023) / 1024 + 1)),
+ };
+ try router.route(r.s, pick, pairs, E, tile, p);
+ return .{ .p = p, .items = items, .pairs = pairs };
+}
+
+fn upBytes(r: *Rig, grp: cuda.experts.Grouped, l: cuda.experts.Layer, x: u64, p: Plan, tile: usize) !Out {
+ const out = try r.zeros(p.pairs * I * 2);
+ try grp.decode(r.s, .up, .{ .x = x, .stride = D, .slots = SLOTS }, l, p.p, out, p.items, p.pairs, tile, -1);
+ return .{ .dev = out, .bytes = try r.back(out, p.pairs * I * 2) };
+}
+
+fn downBytes(r: *Rig, grp: cuda.experts.Grouped, l: cuda.experts.Layer, act: u64, p: Plan, tile: usize) ![]u8 {
+ const out = try r.zeros(p.pairs * D * 4);
+ try grp.decode(r.s, .down_f32, .{ .x = act, .stride = I }, l, p.p, out, p.items, p.pairs, tile, -1);
+ return r.back(out, p.pairs * D * 4);
+}
diff --git a/zig/tests/cuda/experts3_tests.zig b/zig/tests/cuda/experts3_tests.zig
index 3c61bdeef20..7e3497976a4 100644
--- a/zig/tests/cuda/experts3_tests.zig
+++ b/zig/tests/cuda/experts3_tests.zig
@@ -195,11 +195,11 @@ fn sameBits(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random, sh: Shape) !voi
const rows_in: cuda.experts.Rows = .{ .x = dx.ptr, .stride = sh.dims, .slots = sh.slots };
const act_in: cuda.experts.Rows = .{ .x = act.ptr, .stride = sh.ni };
if (prompt) {
- try g.prompt(s, .up, rows_in, l, pb.plan(), act.ptr, pb.items_n, -1);
- try g.prompt(s, .down_bf16, act_in, l, pb.plan(), y.ptr, pb.items_n, -1);
+ try g.prompt(s, .up, rows_in, l, pb.plan(), act.ptr, pb.items_n, pairs, 64, -1);
+ try g.prompt(s, .down_bf16, act_in, l, pb.plan(), y.ptr, pb.items_n, pairs, 64, -1);
} else {
- try g.decode(s, .up, rows_in, l, pb.plan(), act.ptr, pb.items_n, -1);
- try g.decode(s, .down_f32, act_in, l, pb.plan(), y.ptr, pb.items_n, -1);
+ try g.decode(s, .up, rows_in, l, pb.plan(), act.ptr, pb.items_n, pairs, 16, -1);
+ try g.decode(s, .down_f32, act_in, l, pb.plan(), y.ptr, pb.items_n, pairs, 16, -1);
}
try s.synchronize();
outs[li] = .{ try check.download(gpu, act), try check.download(gpu, y) };
@@ -256,7 +256,7 @@ fn projections(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random) !void {
var y = try cuda.DeviceBuffer.alloc(gpu.d, pairs * n * 4);
defer y.free();
const in: cuda.experts.Rows = .{ .x = dx.ptr, .stride = kk };
- if (prompt) try g.prompt(s, .down_bf16, in, l, pb.plan(), y.ptr, pb.items_n, -1) else try g.decode(s, .down_f32, in, l, pb.plan(), y.ptr, pb.items_n, -1);
+ if (prompt) try g.prompt(s, .down_bf16, in, l, pb.plan(), y.ptr, pb.items_n, pairs, 64, -1) else try g.decode(s, .down_f32, in, l, pb.plan(), y.ptr, pb.items_n, pairs, 16, -1);
try s.synchronize();
outs[li] = try check.download(gpu, y);
for (0..pairs * n) |i| {
@@ -334,8 +334,8 @@ fn timeDecode(gpu: Gpu, k: Kernels, s: cuda.Stream, r: std.Random, reps: usize)
for (0..reps) |i| {
const p = plans[i % plans.len];
if (form < 2) {
- try g.decode(s, .up, rows_in, l3, p.plan(), act.ptr, p.items_n, -1);
- } else try g.decode(s, .down_f32, act_in, l3, p.plan(), y.ptr, p.items_n, -1);
+ try g.decode(s, .up, rows_in, l3, p.plan(), act.ptr, p.items_n, rows * sh.slots, 16, -1);
+ } else try g.decode(s, .down_f32, act_in, l3, p.plan(), y.ptr, p.items_n, rows * sh.slots, 16, -1);
}
var graph = try cuda.graph.endCapture(s);
defer graph.deinit();
diff --git a/zig/tests/cuda/fp8_experts_tests.zig b/zig/tests/cuda/fp8_experts_tests.zig
index aa74443dee0..3e50068162d 100644
--- a/zig/tests/cuda/fp8_experts_tests.zig
+++ b/zig/tests/cuda/fp8_experts_tests.zig
@@ -140,7 +140,7 @@ fn oneCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, s: cud
const bound = grouped.maxItems(pairs, l.experts, t);
try act.fill8(0xff, s.handle); // a call that writes nothing fails rather than pass on old bytes
const up_in: ex_.Rows = .{ .x = dx.ptr, .stride = l.dims, .slots = slots };
- if (prompt) try g.prompt(s, .up, up_in, l, plan, act.ptr, bound, -1) else try g.decode(s, .up, up_in, l, plan, act.ptr, bound, -1);
+ if (prompt) try g.prompt(s, .up, up_in, l, plan, act.ptr, bound, pairs, t, -1) else try g.decode(s, .up, up_in, l, plan, act.ptr, bound, pairs, t, -1);
try s.synchronize();
const got_act = try check.download(gpu, act);
defer a.free(got_act);
@@ -153,7 +153,7 @@ fn oneCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, s: cud
defer y.free();
try y.fill8(0xff, s.handle);
const down_in: ex_.Rows = .{ .x = act.ptr, .stride = l.width };
- if (prompt) try g.prompt(s, .down_bf16, down_in, l, plan, y.ptr, bound, -1) else try g.decode(s, .down_f32, down_in, l, plan, y.ptr, bound, -1);
+ if (prompt) try g.prompt(s, .down_bf16, down_in, l, plan, y.ptr, bound, pairs, t, -1) else try g.decode(s, .down_f32, down_in, l, plan, y.ptr, bound, pairs, t, -1);
try s.synchronize();
const got_y = try check.download(gpu, y);
defer a.free(got_y);
diff --git a/zig/tests/cuda/ggml_tests.zig b/zig/tests/cuda/ggml_tests.zig
index 849349b7074..a5cd2b28d42 100644
--- a/zig/tests/cuda/ggml_tests.zig
+++ b/zig/tests/cuda/ggml_tests.zig
@@ -150,7 +150,7 @@ fn expertCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, k:
// up: relu^2 through experts.Grouped, then fp32 sums by the same kernels with no activation
try y.fill8(0xff, s.handle); // a call that writes nothing fails rather than pass on old bytes
const call = if (prompt) &ex_.Grouped.prompt else &ex_.Grouped.decode;
- try call(g, s, .up, .{ .x = dx.ptr, .stride = l.dims, .slots = slots }, l, plan, y.ptr, bound, -1);
+ try call(g, s, .up, .{ .x = dx.ptr, .stride = l.dims, .slots = slots }, l, plan, y.ptr, bound, pairs, tile, -1);
try s.synchronize();
try compare(gpu, try std.fmt.bufPrint(&name, "iq2 up relu^2 {s} {d} rows x {d}", .{ form, rows, slots }), y, ref_up, .relu2, tolerance(.relu2));
try y.fill8(0xff, s.handle);
@@ -160,7 +160,7 @@ fn expertCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, k:
// down from pair rows, fp32 and bf16
inline for (.{ .{ ex_.Epi.down_f32, ggml.Out.f32 }, .{ ex_.Epi.down_bf16, ggml.Out.bf16 } }) |c| {
try y.fill8(0xff, s.handle);
- try call(g, s, c[0], .{ .x = dh.ptr, .stride = l.width }, l, plan, y.ptr, bound, -1);
+ try call(g, s, c[0], .{ .x = dh.ptr, .stride = l.width }, l, plan, y.ptr, bound, pairs, tile, -1);
try s.synchronize();
try compare(gpu, try std.fmt.bufPrint(&name, "iq2 down {t} {s} {d} rows x {d}", .{ c[1], form, rows, slots }), y, ref_down, c[1], tolerance(c[1]));
}
@@ -256,8 +256,8 @@ pub fn bench(gpu: Gpu) !void {
for (0..2) |warm| {
try t0.record(s);
for (0..reps) |_| for (ps) |p| {
- if (pi != 1) try g.decode(s, .up, .{ .x = x.ptr, .stride = dims, .slots = 9 }, l, p, h.ptr, bound, -1);
- if (pi != 0) try g.decode(s, .down_f32, .{ .x = h.ptr, .stride = width }, l, p, y.ptr, bound, -1);
+ if (pi != 1) try g.decode(s, .up, .{ .x = x.ptr, .stride = dims, .slots = 9 }, l, p, h.ptr, bound, 9, 16, -1);
+ if (pi != 0) try g.decode(s, .down_f32, .{ .x = h.ptr, .stride = width }, l, p, y.ptr, bound, 9, 16, -1);
};
try t1.record(s);
try t1.synchronize();
@@ -299,9 +299,9 @@ pub fn bench(gpu: Gpu) !void {
defer t2.deinit();
for (0..2) |warm| {
try t0.record(s);
- try g.prompt(s, .up, .{ .x = xp.ptr, .stride = dims, .slots = 9 }, l, pp, hp.ptr, pb, -1);
+ try g.prompt(s, .up, .{ .x = xp.ptr, .stride = dims, .slots = 9 }, l, pp, hp.ptr, pb, pairs, 64, -1);
try t1.record(s);
- try g.prompt(s, .down_bf16, .{ .x = hp.ptr, .stride = width }, l, pp, yp.ptr, pb, -1);
+ try g.prompt(s, .down_bf16, .{ .x = hp.ptr, .stride = width }, l, pp, yp.ptr, pb, pairs, 64, -1);
try t2.record(s);
try t2.synchronize();
if (warm == 0) continue;
diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig
index 941b0c5a80b..0f4d6c266d6 100644
--- a/zig/tests/cuda/main.zig
+++ b/zig/tests/cuda/main.zig
@@ -457,6 +457,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void {
if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null);
if (std.mem.eql(u8, cmd, "exlr8-bench")) return exlr8_tests.bench(gpu, try arg(rest, 0), if (rest.len > 1) rest[1] else "home", if (rest.len > 2) try std.fmt.parseInt(u16, rest[2], 10) else 40, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4);
if (std.mem.eql(u8, cmd, "exl3")) return exl3_tests.run(gpu, if (rest.len > 0) rest[0] else "none");
+ if (std.mem.eql(u8, cmd, "exl3-grouped-tile")) return @import("exl3_grouped_tile_tests.zig").run(gpu);
if (std.mem.eql(u8, cmd, "exl3-tiles")) return exl3_candidate_tests.run(gpu, if (rest.len > 0) rest[0] else null);
if (std.mem.eql(u8, cmd, "exl3-shapes")) return exl3_plan_bench.profileShapes(gpu, if (rest.len > 0) rest[0] else null);
if (std.mem.eql(u8, cmd, "exl3-midm")) return exl3_plan_bench.midm(gpu);
diff --git a/zig/tests/cuda/modelopt_tests.zig b/zig/tests/cuda/modelopt_tests.zig
index 42ec3e8748d..b848dcbd255 100644
--- a/zig/tests/cuda/modelopt_tests.zig
+++ b/zig/tests/cuda/modelopt_tests.zig
@@ -354,11 +354,11 @@ fn experts(gpu: Gpu, rng: std.Random, p: *const modelopt.Packer, s: cuda.Stream)
const up_in: cuda.experts.Rows = .{ .x = dx.ptr, .stride = dims, .slots = slots };
const down_in: cuda.experts.Rows = .{ .x = act.ptr, .stride = width };
if (arm == 0) {
- try g.decode(s, .up, up_in, l, plan, act.ptr, items, -1);
- try g.decode(s, .down_f32, down_in, l, plan, y.ptr, items, -1);
+ try g.decode(s, .up, up_in, l, plan, act.ptr, items, np, tile, -1);
+ try g.decode(s, .down_f32, down_in, l, plan, y.ptr, items, np, tile, -1);
} else {
- try g.prompt(s, .up, up_in, l, plan, act.ptr, items, -1);
- try g.prompt(s, .down_bf16, down_in, l, plan, y.ptr, items, -1);
+ try g.prompt(s, .up, up_in, l, plan, act.ptr, items, np, tile, -1);
+ try g.prompt(s, .down_bf16, down_in, l, plan, y.ptr, items, np, tile, -1);
}
try s.synchronize();
outs[arm] = .{ try check.download(gpu, act), try check.download(gpu, y) };
diff --git a/zig/tests/cuda/moe_live_tests.zig b/zig/tests/cuda/moe_live_tests.zig
index ab15cd4c9a5..29d0e4847d8 100644
--- a/zig/tests/cuda/moe_live_tests.zig
+++ b/zig/tests/cuda/moe_live_tests.zig
@@ -75,8 +75,8 @@ const Chain = struct {
try c.f.route(c.s, x, dims, rows, c.w, c.logits, dims, count, params);
try c.f.pick(c.s, c.logits, c.picks, c.wts, rows, routed, top_k, params);
try c.router.routeRows(c.s, c.picks, pairs, count, 16, c.plan, if (params == 0) 0 else params + 8, slots);
- try c.grouped.decode(c.s, .up, .{ .x = x, .stride = dims, .slots = slots }, c.layer, c.plan, c.act, items, -1);
- try c.grouped.decode(c.s, .down_f32, .{ .x = c.act, .stride = ni }, c.layer, c.plan, c.y, items, -1);
+ try c.grouped.decode(c.s, .up, .{ .x = x, .stride = dims, .slots = slots }, c.layer, c.plan, c.act, items, pairs, 16, -1);
+ try c.grouped.decode(c.s, .down_f32, .{ .x = c.act, .stride = ni }, c.layer, c.plan, c.y, items, pairs, 16, -1);
try c.f.combine(c.s, c.y, true, c.wts, out, rows, slots, dims, params);
}
diff --git a/zig/tests/cuda/moe_tests.zig b/zig/tests/cuda/moe_tests.zig
index b3301a13b4f..809cc2f97a5 100644
--- a/zig/tests/cuda/moe_tests.zig
+++ b/zig/tests/cuda/moe_tests.zig
@@ -339,8 +339,8 @@ fn pairExperts(gpu: Gpu, s: cuda.Stream, r: std.Random) !void {
const layer: cuda.experts.Layer = .{ .format = .affine4, .up = dpair.ptr, .down = ddown.ptr, .width = ni, .dims = dims, .experts = count, .swiglu = true };
const grouped: cuda.experts.Grouped = .{ .affine4 = &ex };
const items_n = cuda.grouped.maxItems(pairs, count, 16);
- try grouped.decode(s, .up, .{ .x = dx.ptr, .stride = dims, .slots = slots }, layer, plan, dact.ptr, items_n, -1);
- try grouped.decode(s, .down_f32, .{ .x = dact.ptr, .stride = ni }, layer, plan, dy.ptr, items_n, -1);
+ try grouped.decode(s, .up, .{ .x = dx.ptr, .stride = dims, .slots = slots }, layer, plan, dact.ptr, items_n, pairs, 16, -1);
+ try grouped.decode(s, .down_f32, .{ .x = dact.ptr, .stride = ni }, layer, plan, dy.ptr, items_n, pairs, 16, -1);
try s.synchronize();
const got = try check.download(gpu, dy);
defer a.free(got);
diff --git a/zig/tests/cuda/nvfp4_experts_tests.zig b/zig/tests/cuda/nvfp4_experts_tests.zig
index 20310178366..10f2e0a2108 100644
--- a/zig/tests/cuda/nvfp4_experts_tests.zig
+++ b/zig/tests/cuda/nvfp4_experts_tests.zig
@@ -138,7 +138,7 @@ fn oneCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, s: cud
var act = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.width * 2);
defer act.free();
try act.fill8(0xff, s.handle); // a call that writes nothing fails rather than pass on old bytes
- try g.decode(s, .up, .{ .x = dx.ptr, .stride = l.dims, .slots = slots }, l, plan, act.ptr, bound, -1);
+ try g.decode(s, .up, .{ .x = dx.ptr, .stride = l.dims, .slots = slots }, l, plan, act.ptr, bound, pairs, 16, -1);
try s.synchronize();
const got_act = try check.download(gpu, act);
defer a.free(got_act);
@@ -150,7 +150,7 @@ fn oneCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, s: cud
var y = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.dims * width);
defer y.free();
try y.fill8(0xff, s.handle);
- try g.decode(s, if (wide) .down_bf16 else .down_f32, .{ .x = act.ptr, .stride = l.width }, l, plan, y.ptr, bound, -1);
+ try g.decode(s, if (wide) .down_bf16 else .down_f32, .{ .x = act.ptr, .stride = l.width }, l, plan, y.ptr, bound, pairs, 16, -1);
try s.synchronize();
const got_y = try check.download(gpu, y);
defer a.free(got_y);
diff --git a/zig/tests/cuda/quant_ref_experts.zig b/zig/tests/cuda/quant_ref_experts.zig
index f841aaf3bbf..98dae8897cb 100644
--- a/zig/tests/cuda/quant_ref_experts.zig
+++ b/zig/tests/cuda/quant_ref_experts.zig
@@ -67,11 +67,11 @@ fn batch(r: *Rig, comptime label: []const u8, g: experts.Grouped, router: groupe
const up_in: experts.Rows = .{ .x = try r.dev(std.mem.sliceAsBytes(x[0 .. rows * D])), .stride = D, .slots = slots };
const down_in: experts.Rows = .{ .x = act, .stride = NI };
if (prompt) {
- try g.prompt(r.s, .up, up_in, l, pl.p, act, pl.items, -1);
- try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y, pl.items, -1);
+ try g.prompt(r.s, .up, up_in, l, pl.p, act, pl.items, pairs, tile, -1);
+ try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y, pl.items, pairs, tile, -1);
} else {
- try g.decode(r.s, .up, up_in, l, pl.p, act, pl.items, -1);
- try g.decode(r.s, .down_f32, down_in, l, pl.p, y, pl.items, -1);
+ try g.decode(r.s, .up, up_in, l, pl.p, act, pl.items, pairs, tile, -1);
+ try g.decode(r.s, .down_f32, down_in, l, pl.p, y, pl.items, pairs, tile, -1);
}
// gate-up: each pair's row of x through its expert, SwiGLU with the kernels' roundings
const want = try r.a.alloc(f64, pairs * NI);
@@ -140,15 +140,15 @@ fn skipped(
const up_in: experts.Rows = .{ .x = xd, .stride = D, .slots = slots };
const a0 = try r.zeros(pairs * NI * 2);
const a1 = try r.dev(fill[0 .. pairs * NI * 2]);
- try g.prompt(r.s, .up, up_in, l, pl.p, a0, pl.items, -1);
- try g.prompt(r.s, .up, up_in, l, pl.p, a1, pl.items, skip);
+ try g.prompt(r.s, .up, up_in, l, pl.p, a0, pl.items, pairs, experts.promptTile(l.format), -1);
+ try g.prompt(r.s, .up, up_in, l, pl.p, a1, pl.items, pairs, experts.promptTile(l.format), skip);
const up_plain = try r.back(a0, pairs * NI * 2);
try leftAlone(label ++ " prompt gate-up", pk, skip, up_plain, try r.back(a1, pairs * NI * 2), NI * 2, fill);
const down_in: experts.Rows = .{ .x = a0, .stride = NI };
const y0 = try r.zeros(pairs * D * 2);
const y1 = try r.dev(fill[0 .. pairs * D * 2]);
- try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y0, pl.items, -1);
- try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y1, pl.items, skip);
+ try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y0, pl.items, pairs, experts.promptTile(l.format), -1);
+ try g.prompt(r.s, .down_bf16, down_in, l, pl.p, y1, pl.items, pairs, experts.promptTile(l.format), skip);
const down_plain = try r.back(y0, pairs * D * 2);
try leftAlone(label ++ " prompt down", pk, skip, down_plain, try r.back(y1, pairs * D * 2), D * 2, fill);
}
diff --git a/zig/tests/exl3_compile.zig b/zig/tests/exl3_compile.zig
index a3abe2e3790..e5e8468f983 100644
--- a/zig/tests/exl3_compile.zig
+++ b/zig/tests/exl3_compile.zig
@@ -21,9 +21,9 @@ export fn exl3Compile(d: *const cuda.Driver, dense: *cuda.exl3.Dense, group: *cu
linear.decodeF32(s.*, weight, .{ .x = 0 }, 0, 1) catch {};
linear.promptF32(s.*, weight, .{ .x = 0 }, 0, 1) catch {};
const grouped: cuda.experts.Grouped = .{ .exl3 = group };
- grouped.decode(s.*, .up, .{ .x = 0, .stride = l.dims }, l.*, p.*, 0, 1, -1) catch {};
- grouped.prompt(s.*, .down_f32, .{ .x = 0, .stride = l.width }, l.*, p.*, 0, 1, -1) catch {};
- grouped.prompt(s.*, .down_bf16, .{ .x = 0, .stride = l.width }, l.*, p.*, 0, 1, -1) catch {};
+ grouped.decode(s.*, .up, .{ .x = 0, .stride = l.dims }, l.*, p.*, 0, 1, 1, 16, -1) catch {};
+ grouped.prompt(s.*, .down_f32, .{ .x = 0, .stride = l.width }, l.*, p.*, 0, 1, 1, 64, -1) catch {};
+ grouped.prompt(s.*, .down_bf16, .{ .x = 0, .stride = l.width }, l.*, p.*, 0, 1, 1, 64, -1) catch {};
}
export fn candidateCompile(d: *const cuda.Driver, m: *const cuda.Module, s: *const cuda.Stream) void {