From 9c0c93aefa75dd01ab2135e3855a6c3e57ded0a9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BCrgen=20Schmied?=
 <600316+jschmied@users.noreply.github.com>
Date: Mon, 12 Oct 2026 16:56:24 +0200
Subject: [PATCH 1/2] cuda/tree_round: core/tree_round's verified rounds on
 CUDA, for every family

The Metal engine's tree round (picks, each tree's verified path, the commit
steps that read it) as shared CUDA kernels: tf_round_bf16_argmax,
tf_tree_round_match over a batch of trees (tree-local parents, per-tree
budget and EOS ids), and keep rows / gather rows / keep head for the commit.
Results use core/tree_round's Result layout, so decode() reads them as-is.

tf-cuda-test tree-round checks argmax against rowArgmax (ties, NaN, +-inf,
padding past vocab), 4096 batched random trees against oracle() and decode()
(every stop and status, full 16-row paths), trees alone against the batch,
and the commit kernels against the host; four kernel mutations are caught.
---
 zig/build/cuda.zig                  |   1 +
 zig/kernels/cuda/tree_round.cu      | 195 +++++++++++++++
 zig/src/core/root.zig               |   2 +
 zig/src/cuda/kernels.zig            |   1 +
 zig/src/cuda/root.zig               |   2 +
 zig/src/cuda/tree_round.zig         | 103 ++++++++
 zig/tests/cuda/main.zig             |   3 +
 zig/tests/cuda/tree_round_tests.zig | 356 ++++++++++++++++++++++++++++
 8 files changed, 663 insertions(+)
 create mode 100644 zig/kernels/cuda/tree_round.cu
 create mode 100644 zig/src/cuda/tree_round.zig
 create mode 100644 zig/tests/cuda/tree_round_tests.zig

diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig
index 684d6c3cc9..2689776ce1 100644
--- a/zig/build/cuda.zig
+++ b/zig/build/cuda.zig
@@ -45,6 +45,7 @@ const kernels = [_]Kernel{
     .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: 4-bit affine and bf16 token rows (cuda/embed_rows.zig)
     .{ .name = "lfm2", .flags = glue }, // ours: LFM2's short convolution, QK norm and rotary, gates and label rows
     .{ .name = "vit", .flags = glue }, // ours: vision towers' LayerNorm, Linear epilogues, activations, unshuffle
+    .{ .name = "tree_round", .flags = &.{"-O3"} }, // ours: verified tree rounds for every family (cuda/tree_round.zig)
     .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3
     .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1
     .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts
diff --git a/zig/kernels/cuda/tree_round.cu b/zig/kernels/cuda/tree_round.cu
new file mode 100644
index 0000000000..c0ad1dd2b4
--- /dev/null
+++ b/zig/kernels/cuda/tree_round.cu
@@ -0,0 +1,195 @@
+// Verified tree rounds for every CUDA family (host reference: core/tree_round.zig, the Metal kernels' contract).
+
+// Deterministic BF16 picks, each tree's verified path and emitted tokens, and the commit steps that read the result.
+
+#include <stdint.h>
+
+namespace tf_tree_round {
+
+constexpr int MAX_ROWS = 16;
+constexpr uint32_t INVALID = 0xffffffffu;
+
+struct Pick {
+    uint32_t token, nonfinite;
+};
+
+struct Result {
+    uint32_t status, stop, nonfinite, consumed_count, emitted_count, matched_count;
+    uint32_t bonus_emitted, pending_valid, pending_token, reserved;
+    uint32_t path[MAX_ROWS], tokens[MAX_ROWS];
+};
+
+// One tree of a round: its rows first .. first + rows of tokens, parents (tree-local, -1 at the root) and picks.
+struct Tree {
+    uint32_t first, rows, budget, eos_first, eos_count;
+};
+
+__device__ __forceinline__ bool valid(const Result& r, uint32_t rows) {
+    if (r.status || r.nonfinite || r.reserved || !r.consumed_count || r.consumed_count > rows ||
+        r.consumed_count != r.emitted_count || r.path[0] != 0)
+        return false;
+    for (uint32_t i = 0; i < r.consumed_count; ++i)
+        if (r.path[i] >= rows || (i && r.path[i] <= r.path[i - 1])) return false;
+    return true;
+}
+
+}  // namespace tf_tree_round
+
+using namespace tf_tree_round;
+
+// picks[row]: a row's first largest finite bf16 logit, or INVALID and its nonfinite kinds (1 NaN, 2 +inf, 4 -inf).
+extern "C" __global__ void __launch_bounds__(256) tf_round_bf16_argmax(const uint16_t* __restrict__ logits,
+                                                                       Pick* __restrict__ picks, uint32_t vocab,
+                                                                       uint32_t stride) {
+    __shared__ float values[256];
+    __shared__ uint32_t ids[256], masks[256];
+    const uint32_t row = blockIdx.x, lane = threadIdx.x;
+    float best = -INFINITY;
+    uint32_t id = INVALID, mask = 0;
+    for (uint64_t token = lane; token < vocab; token += 256) {
+        const float value = __uint_as_float(static_cast<uint32_t>(logits[static_cast<uint64_t>(row) * stride + token]) << 16);
+        if (!isfinite(value)) mask |= isnan(value) ? 1u : value > 0 ? 2u : 4u;
+        else if (id == INVALID || value > best || (value == best && token < id)) {
+            best = value;
+            id = static_cast<uint32_t>(token);
+        }
+    }
+    values[lane] = best;
+    ids[lane] = id;
+    masks[lane] = mask;
+    __syncthreads();
+    for (uint32_t step = 128; step > 0; step >>= 1) {
+        if (lane < step) {
+            const float other = values[lane + step];
+            const uint32_t other_id = ids[lane + step];
+            if (other > values[lane] || (other == values[lane] && other_id < ids[lane])) {
+                values[lane] = other;
+                ids[lane] = other_id;
+            }
+            masks[lane] |= masks[lane + step];
+        }
+        __syncthreads();
+    }
+    if (lane == 0) picks[row] = {masks[0] != 0 ? INVALID : ids[0], masks[0]};
+}
+
+// out[t]: tree t's verified round (core/tree_round.oracle), one thread a tree.
+extern "C" __global__ void tf_tree_round_match(const uint32_t* __restrict__ tokens, const int32_t* __restrict__ parents,
+                                               const Pick* __restrict__ picks, const uint32_t* __restrict__ eos,
+                                               const Tree* __restrict__ trees, Result* __restrict__ results,
+                                               uint32_t vocab, uint32_t count) {
+    const uint32_t t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= count) return;
+    const Tree tr = trees[t];
+    Result& out = results[t];
+    const uint32_t* tok = tokens + tr.first;
+    const int32_t* par = parents + tr.first;
+    const Pick* pk = picks + tr.first;
+    const uint32_t* stop_ids = eos + tr.eos_first;
+    out.status = 0;
+    out.stop = 0;
+    out.nonfinite = 0;
+    out.consumed_count = 0;
+    out.emitted_count = 0;
+    out.matched_count = 0;
+    out.bonus_emitted = 0;
+    out.pending_valid = 0;
+    out.pending_token = 0;
+    out.reserved = 0;
+    for (int i = 0; i < MAX_ROWS; ++i) out.path[i] = 0, out.tokens[i] = 0;
+    if (tr.rows == 0 || tr.rows > MAX_ROWS || tr.eos_count > MAX_ROWS) {
+        out.status = 2;
+        out.stop = 3;
+        return;
+    }
+    for (uint32_t row = 0; row < tr.rows; ++row)
+        if (tok[row] >= vocab || (row == 0 ? par[row] != -1 : par[row] < 0 || par[row] >= static_cast<int32_t>(row))) {
+            out.status = 2;
+            out.stop = 3;
+            return;
+        }
+    for (uint32_t i = 0; i < tr.eos_count; ++i)
+        if (stop_ids[i] >= vocab) {
+            out.status = 2;
+            out.stop = 3;
+            return;
+        }
+    bool bad_pick = false;
+    for (uint32_t row = 0; row < tr.rows; ++row) {
+        out.nonfinite |= pk[row].nonfinite & 7u;
+        bad_pick = bad_pick || (pk[row].nonfinite & ~7u) != 0 || (pk[row].nonfinite == 0 && pk[row].token >= vocab);
+    }
+    if (bad_pick || out.nonfinite != 0) {
+        out.status = bad_pick ? 3u : 1u;
+        out.stop = 3;
+        return;
+    }
+    if (tr.budget == 0) {
+        out.stop = 1;
+        return;
+    }
+    out.path[out.consumed_count++] = 0;
+    uint32_t node = 0;
+    for (;;) {
+        const uint32_t token = pk[node].token;
+        uint32_t child = INVALID;
+        for (uint32_t row = node + 1; row < tr.rows; ++row)
+            if (par[row] == static_cast<int32_t>(node) && tok[row] == token) {
+                child = row;
+                break;
+            }
+        out.tokens[out.emitted_count++] = token;
+        if (child == INVALID) out.bonus_emitted = 1;
+        else ++out.matched_count;
+        bool hit_eos = false;
+        for (uint32_t i = 0; i < tr.eos_count; ++i) hit_eos = hit_eos || stop_ids[i] == token;
+        if (hit_eos) {
+            out.stop = 2;
+            return;
+        }
+        if (out.emitted_count == tr.budget) {
+            out.stop = 1;
+            return;
+        }
+        if (child == INVALID) {
+            out.pending_valid = 1;
+            out.pending_token = token;
+            return;
+        }
+        out.path[out.consumed_count++] = child;
+        node = child;
+    }
+}
+
+// keep[t * 16 + i]: tree t's i-th kept row (first + path[i]) or -1 past its consumed count; counts[t]: kept.
+extern "C" __global__ void tf_round_keep_rows(const Result* __restrict__ results, const Tree* __restrict__ trees,
+                                              int32_t* __restrict__ keep, uint32_t* __restrict__ counts, uint32_t count) {
+    const uint32_t t = blockIdx.x, i = threadIdx.x;
+    if (t >= count || i >= MAX_ROWS) return;
+    const Result& r = results[t];
+    const uint32_t n = valid(r, trees[t].rows) ? r.consumed_count : 0;
+    keep[t * MAX_ROWS + i] = i < n ? static_cast<int32_t>(trees[t].first + r.path[i]) : -1;
+    if (i == 0) counts[t] = n;
+}
+
+// out[t, i, :] = rows[first_t + path_t[i], :] (width bf16 a row) for each tree's kept rows: taps, hidden rows.
+extern "C" __global__ void tf_round_gather_rows(const Result* __restrict__ results, const Tree* __restrict__ trees,
+                                                const uint16_t* __restrict__ rows, uint16_t* __restrict__ out,
+                                                uint32_t width, uint32_t count) {
+    const uint32_t t = blockIdx.y, i = blockIdx.z, c = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= count || c >= width) return;
+    const Result& r = results[t];
+    if (!valid(r, trees[t].rows) || i >= r.consumed_count) return;
+    out[(static_cast<uint64_t>(t) * MAX_ROWS + i) * width + c] = rows[static_cast<uint64_t>(trees[t].first + r.path[i]) * width + c];
+}
+
+// head[t, :] = logits[first_t + the tree's last kept row, :]: the row the next round's first pick reads.
+extern "C" __global__ void tf_round_keep_head(const Result* __restrict__ results, const Tree* __restrict__ trees,
+                                              const uint16_t* __restrict__ logits, uint16_t* __restrict__ head,
+                                              uint32_t width, uint32_t stride, uint32_t count) {
+    const uint32_t t = blockIdx.y, c = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= count || c >= width) return;
+    const Result& r = results[t];
+    if (!valid(r, trees[t].rows)) return;
+    head[static_cast<uint64_t>(t) * width + c] = logits[static_cast<uint64_t>(trees[t].first + r.path[r.consumed_count - 1]) * stride + c];
+}
diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig
index 122ffe9e4c..3ab41484e7 100644
--- a/zig/src/core/root.zig
+++ b/zig/src/core/root.zig
@@ -28,6 +28,7 @@ pub const modelopt = @import("modelopt.zig");
 pub const picture = @import("picture/picture.zig");
 /// Qwen-VL image input every Qwen3.5-VL checkpoint shares: settings, geometry, M-RoPE positions, the tower's tables.
 pub const qwen_vl = @import("qwen_vl/qwen_vl.zig");
+pub const tree_round = @import("tree_round.zig");
 
 test {
     _ = safetensors;
@@ -41,6 +42,7 @@ test {
     _ = norms_ref;
     _ = act_ref;
     _ = contract;
+    _ = tree_round;
     _ = exl3_format;
     _ = exlr8;
     _ = eaq_container;
diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig
index 18542b8cb4..4406862c61 100644
--- a/zig/src/cuda/kernels.zig
+++ b/zig/src/cuda/kernels.zig
@@ -56,6 +56,7 @@ pub const ck_lane: []const u8 = if (available) &Blob("fatbin_ck_lane").bytes els
 pub const ck_experts: []const u8 = if (available) &Blob("fatbin_ck_experts").bytes else &.{};
 pub const norms: []const u8 = if (available) &Blob("fatbin_norms").bytes else &.{};
 pub const embed_rows: []const u8 = if (available) &Blob("fatbin_embed_rows").bytes else &.{};
+pub const tree_round: []const u8 = if (available) &Blob("fatbin_tree_round").bytes else &.{};
 
 /// Symbols in the gdn image as cuobjdump lists them for the built fatbin (named namespace tf_gdn).
 pub const gdn_symbols = struct {
diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig
index 6c4e0991b6..4baa0aced3 100644
--- a/zig/src/cuda/root.zig
+++ b/zig/src/cuda/root.zig
@@ -34,6 +34,7 @@ pub const qlinear = @import("qlinear.zig");
 pub const norms = @import("norms.zig");
 pub const embed_rows = @import("embed_rows.zig");
 pub const vit = @import("vit.zig");
+pub const tree_round = @import("tree_round.zig");
 pub const experts = @import("experts.zig");
 pub const sample = @import("sample.zig");
 pub const rank_scatter = @import("rank_scatter.zig");
@@ -49,6 +50,7 @@ test {
     _ = @import("norms.zig");
     _ = @import("embed_rows.zig");
     _ = @import("vit.zig");
+    _ = @import("tree_round.zig");
     _ = @import("experts.zig");
     _ = @import("sample.zig");
     _ = @import("rank_scatter.zig");
diff --git a/zig/src/cuda/tree_round.zig b/zig/src/cuda/tree_round.zig
new file mode 100644
index 0000000000..bf8150d31b
--- /dev/null
+++ b/zig/src/cuda/tree_round.zig
@@ -0,0 +1,103 @@
+//! tree_round.cu's verified tree rounds for every CUDA family: picks, each tree's verified path, and its commit steps.
+const std = @import("std");
+const module = @import("module.zig");
+const launch_ = @import("launch.zig");
+const stream_ = @import("stream.zig");
+
+const Module = module.Module;
+const Function = module.Function;
+const Stream = stream_.Stream;
+
+/// core/tree_round.zig's contract, as the kernels lay it out.
+pub const max_rows = 16;
+pub const Pick = extern struct { token: u32, nonfinite: u32 };
+pub const Result = extern struct {
+    status: u32,
+    stop: u32,
+    nonfinite: u32,
+    consumed_count: u32,
+    emitted_count: u32,
+    matched_count: u32,
+    bonus_emitted: u32,
+    pending_valid: u32,
+    pending_token: u32,
+    reserved: u32,
+    path: [max_rows]u32,
+    tokens: [max_rows]u32,
+};
+
+/// One tree of a round: its rows of tokens, parents (tree-local, -1 at the root) and picks, its budget and EOS ids.
+pub const Tree = extern struct { first: u32, rows: u32, budget: u32, eos_first: u32 = 0, eos_count: u32 = 0 };
+
+/// tree_round.cu's kernels, resolved from the loaded `tree_round` module (kernels.tree_round).
+pub const Kernels = struct {
+    argmax_fn: Function,
+    match_fn: Function,
+    keep_fn: Function,
+    gather_fn: Function,
+    head_fn: Function,
+
+    pub fn resolve(mod: Module) !Kernels {
+        return .{ .argmax_fn = try mod.function("tf_round_bf16_argmax"), .match_fn = try mod.function("tf_tree_round_match"), .keep_fn = try mod.function("tf_round_keep_rows"), .gather_fn = try mod.function("tf_round_gather_rows"), .head_fn = try mod.function("tf_round_keep_head") };
+    }
+
+    /// picks[r] for every row of a round's trees: the first largest finite bf16 logit, or none.
+    pub fn argmax(g: *const Kernels, s: Stream, logits: u64, picks: u64, rows: usize, vocab: usize, stride: usize) !void {
+        if (rows == 0) return;
+        if (vocab == 0 or stride < vocab) return error.BadRoundShape;
+        var a: launch_.Args = .{};
+        a.add(logits);
+        a.add(picks);
+        a.add(@as(u32, @intCast(vocab)));
+        a.add(@as(u32, @intCast(stride)));
+        try launch_.launch(g.argmax_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = 256 } }, s, &a);
+    }
+
+    /// results[t] for `count` trees (Tree [count] at `trees`): each tree's verified path and emitted tokens.
+    pub fn match(g: *const Kernels, s: Stream, tokens: u64, parents: u64, picks: u64, eos: u64, trees: u64, results: u64, vocab: usize, count: usize) !void {
+        if (count == 0) return;
+        var a: launch_.Args = .{};
+        for ([_]u64{ tokens, parents, picks, eos, trees, results }) |v| a.add(v);
+        a.add(@as(u32, @intCast(vocab)));
+        a.add(@as(u32, @intCast(count)));
+        try launch_.launch(g.match_fn, .{ .grid = .{ .x = @intCast((count + 63) / 64), .y = 1, .z = 1 }, .block = .{ .x = 64 } }, s, &a);
+    }
+
+    /// keep i32 [count, max_rows]: each tree's kept rows (round row indices), -1 past them; counts u32 [count].
+    pub fn keepRows(g: *const Kernels, s: Stream, results: u64, trees: u64, keep: u64, counts: u64, count: usize) !void {
+        if (count == 0) return;
+        var a: launch_.Args = .{};
+        for ([_]u64{ results, trees, keep, counts }) |v| a.add(v);
+        a.add(@as(u32, @intCast(count)));
+        try launch_.launch(g.keep_fn, .{ .grid = .{ .x = @intCast(count), .y = 1, .z = 1 }, .block = .{ .x = max_rows } }, s, &a);
+    }
+
+    /// out bf16 [count, max_rows, width]: each tree's kept rows of `rows` (bf16 [round rows, width]), in path order.
+    pub fn gatherRows(g: *const Kernels, s: Stream, results: u64, trees: u64, rows: u64, out: u64, width: usize, count: usize) !void {
+        if (count == 0 or width == 0) return;
+        var a: launch_.Args = .{};
+        for ([_]u64{ results, trees, rows, out }) |v| a.add(v);
+        a.add(@as(u32, @intCast(width)));
+        a.add(@as(u32, @intCast(count)));
+        try launch_.launch(g.gather_fn, .{ .grid = .{ .x = @intCast((width + 255) / 256), .y = @intCast(count), .z = max_rows }, .block = .{ .x = 256 } }, s, &a);
+    }
+
+    /// head bf16 [count, width]: each tree's last kept row of `logits` (`stride` apart), the next round's first row.
+    pub fn keepHead(g: *const Kernels, s: Stream, results: u64, trees: u64, logits: u64, head: u64, width: usize, stride: usize, count: usize) !void {
+        if (count == 0 or width == 0) return;
+        var a: launch_.Args = .{};
+        for ([_]u64{ results, trees, logits, head }) |v| a.add(v);
+        a.add(@as(u32, @intCast(width)));
+        a.add(@as(u32, @intCast(stride)));
+        a.add(@as(u32, @intCast(count)));
+        try launch_.launch(g.head_fn, .{ .grid = .{ .x = @intCast((width + 255) / 256), .y = @intCast(count), .z = 1 }, .block = .{ .x = 256 } }, s, &a);
+    }
+};
+
+test "the kernels' layouts" {
+    try std.testing.expectEqual(@as(usize, 8), @sizeOf(Pick));
+    try std.testing.expectEqual(@as(usize, 168), @sizeOf(Result));
+    try std.testing.expectEqual(@as(usize, 40), @offsetOf(Result, "path"));
+    try std.testing.expectEqual(@as(usize, 104), @offsetOf(Result, "tokens"));
+    try std.testing.expectEqual(@as(usize, 20), @sizeOf(Tree));
+}
diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig
index fe98682c86..81e13d8ad3 100644
--- a/zig/tests/cuda/main.zig
+++ b/zig/tests/cuda/main.zig
@@ -38,6 +38,7 @@ const lfm2_vision_tests = @import("lfm2_vision_tests.zig");
 const radio_tests = @import("radio_tests.zig");
 const glm_lanes_tests = @import("glm_lanes_tests.zig");
 const glm_synth = @import("glm_synth.zig");
+const tree_round_tests = @import("tree_round_tests.zig");
 
 const usage =
     \\usage: tf-cuda-test <command>
@@ -59,6 +60,7 @@ const usage =
     \\  triton <dir>              a Triton cubin's bits against the Python oracle's fixture
     \\  norms                     RMS norm and the residual step against float64 references (synthetic)
     \\  norm-rows                 bf16-residual and gated norms, token rows against float64 (synthetic)
+    \\  tree-round                tree_round.cu's picks, batched verified trees and commits against core/tree_round
     \\  sample                    sample.cu's keyed draws against the Metal rule's host references (synthetic rows)
     \\  glue                      Nemotron's glue kernels against their host references at its shapes (synthetic rows)
     \\  lfm2-kernels              LFM2's glue kernels against their host references at d1's shapes (synthetic rows)
@@ -164,6 +166,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void {
     if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu);
     if (std.mem.eql(u8, cmd, "norms")) return norms_tests.run(gpu);
     if (std.mem.eql(u8, cmd, "norm-rows")) return norm_rows_tests.run(gpu);
+    if (std.mem.eql(u8, cmd, "tree-round")) return tree_round_tests.run(gpu);
     if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu);
     if (std.mem.eql(u8, cmd, "lfm2-kernels")) return lfm2_kernel_tests.run(gpu);
     if (std.mem.eql(u8, cmd, "pocket"))
diff --git a/zig/tests/cuda/tree_round_tests.zig b/zig/tests/cuda/tree_round_tests.zig
new file mode 100644
index 0000000000..aaf0d6bc26
--- /dev/null
+++ b/zig/tests/cuda/tree_round_tests.zig
@@ -0,0 +1,356 @@
+//! cuda/tree_round.zig against core/tree_round.zig: picks (rowArgmax), batched trees (oracle, decode), the commits.
+
+const std = @import("std");
+const cuda = @import("cuda");
+const core = @import("core");
+const check = @import("check.zig");
+const Gpu = check.Gpu;
+
+const tr = cuda.tree_round;
+const ref = core.tree_round;
+
+comptime {
+    std.debug.assert(@sizeOf(tr.Result) == @sizeOf(ref.Result) and @sizeOf(tr.Pick) == @sizeOf(ref.Pick));
+    std.debug.assert(tr.max_rows == ref.max_rows);
+}
+
+fn bf16(x: f32) u16 {
+    const b: u32 = @bitCast(x);
+    return @intCast((b +% 0x7FFF +% ((b >> 16) & 1)) >> 16);
+}
+
+const nan: u16 = 0x7FC0;
+const pos_inf: u16 = 0x7F80;
+const neg_inf: u16 = 0xFF80;
+/// A large finite value written past each row's vocab: a kernel reading the padding picks it.
+const decoy: u16 = 0x7F00;
+
+/// Row r by r's kind: spread, coarse ties, a NaN, +inf, -inf, all -inf, the max last, ties at both ends.
+fn fillRow(rng: std.Random, r: usize, row: []u16) void {
+    const v = row.len;
+    for (row) |*h| h.* = switch (r % 8) {
+        1 => bf16(-@floor(rng.float(f32) * 4.0)),
+        5 => neg_inf,
+        else => bf16(rng.floatNorm(f32) * 3.0),
+    };
+    switch (r % 8) {
+        2 => row[rng.uintLessThan(usize, v)] = nan,
+        3 => row[rng.uintLessThan(usize, v)] = pos_inf,
+        4 => row[rng.uintLessThan(usize, v)] = neg_inf,
+        6 => row[v - 1] = bf16(100.0),
+        7 => {
+            row[0] = bf16(100.0);
+            row[v - 1] = bf16(100.0);
+        },
+        else => {},
+    }
+}
+
+fn picks(gpu: Gpu, k: *const tr.Kernels, s: cuda.Stream) !void {
+    const gpa = gpu.gpa;
+    var prng = std.Random.DefaultPrng.init(0x7AEE_0001);
+    const rng = prng.random();
+    const rows = 24; // more than one tree's 16: a round's argmax covers every tree
+    var tested: usize = 0;
+    for ([_]usize{ 1, 7, 255, 256, 257, 50_000, 151_936 }) |vocab| {
+        const stride = vocab + 64;
+        const host = try gpa.alloc(u16, rows * stride);
+        defer gpa.free(host);
+        @memset(host, decoy);
+        for (0..rows) |r| fillRow(rng, r, host[r * stride ..][0..vocab]);
+        var logits = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host));
+        defer logits.free();
+        var out = try cuda.DeviceBuffer.alloc(gpu.d, rows * @sizeOf(tr.Pick));
+        defer out.free();
+        try out.fill8(0xA5, s.handle);
+        try k.argmax(s, logits.ptr, out.ptr, rows, vocab, stride);
+        try s.synchronize();
+        var got: [rows]tr.Pick = undefined;
+        try out.download(0, std.mem.sliceAsBytes(&got));
+        var want: [rows]ref.Pick = undefined;
+        for (&want, 0..) |*w, r| w.* = ref.rowArgmax(host[r * stride ..][0..vocab]);
+        try check.sameBytes("argmax picks", std.mem.sliceAsBytes(&got), std.mem.sliceAsBytes(&want));
+        tested += rows;
+    }
+    check.pass("tree-round argmax: {d} rows over 7 vocabularies equal rowArgmax (ties, NaN, +-inf, padding)", .{tested});
+}
+
+/// A round of `count` random trees in the kernels' layout, with the host reference's per-tree slices.
+const Round = struct {
+    tokens: std.ArrayList(u32) = .empty,
+    parents: std.ArrayList(i32) = .empty,
+    picks: std.ArrayList(tr.Pick) = .empty,
+    eos: std.ArrayList(u32) = .empty,
+    trees: std.ArrayList(tr.Tree) = .empty,
+
+    fn deinit(self: *Round, gpa: std.mem.Allocator) void {
+        self.tokens.deinit(gpa);
+        self.parents.deinit(gpa);
+        self.picks.deinit(gpa);
+        self.eos.deinit(gpa);
+        self.trees.deinit(gpa);
+    }
+
+    /// Tree t's expected result: the oracle's, or a bad tree where the oracle refuses the shape.
+    fn want(self: *const Round, t: usize, vocab: u32) !ref.Result {
+        const x = self.trees.items[t];
+        const p = ref.Match{ .rows = x.rows, .vocab = vocab, .budget = x.budget, .eos_count = x.eos_count };
+        const rows = @min(x.rows, self.tokens.items.len - x.first);
+        return ref.oracle(p, self.tokens.items[x.first..][0..rows], self.parents.items[x.first..][0..rows], @ptrCast(self.picks.items[x.first..][0..rows]), self.eos.items[x.eos_first..][0..x.eos_count]) catch |e| switch (e) {
+            error.BadRoundShape => .{ .status = @backingInt(ref.Status.bad_tree), .stop = @backingInt(ref.Stop.invalid) },
+        };
+    }
+};
+
+const vocab_small: u32 = 6; // small enough that drafts match the picks often
+
+fn buildRound(gpa: std.mem.Allocator, rng: std.Random, count: usize) !Round {
+    var round: Round = .{};
+    errdefer round.deinit(gpa);
+    for (0..count) |t| {
+        const chain = t % 97 == 50; // a full-depth path: every draft matches, 16 rows kept
+        const rows: u32 = if (t % 97 == 13) 0 else if (t % 97 == 14) 17 else if (chain) 16 else 1 + rng.uintLessThan(u32, 16);
+        const first: u32 = @intCast(round.tokens.items.len);
+        for (0..rows) |r| {
+            try round.parents.append(gpa, if (r == 0) -1 else if (chain) @intCast(r - 1) else @intCast(rng.uintLessThan(usize, r)));
+            try round.tokens.append(gpa, rng.uintLessThan(u32, vocab_small));
+        }
+        // picks follow a child most of the time, so paths run deep
+        for (0..rows) |r| {
+            var token = rng.uintLessThan(u32, vocab_small);
+            if (chain or rng.uintLessThan(u32, 10) < 7) {
+                for (r + 1..rows) |c| if (round.parents.items[first + c] == @as(i32, @intCast(r))) {
+                    token = round.tokens.items[first + c];
+                    if (rng.boolean()) break;
+                };
+            }
+            try round.picks.append(gpa, .{ .token = token, .nonfinite = 0 });
+        }
+        const eos_first: u32 = @intCast(round.eos.items.len);
+        const eos_count: u32 = if (chain or rng.boolean()) 0 else 1 + rng.uintLessThan(u32, 2);
+        for (0..eos_count) |_| try round.eos.append(gpa, rng.uintLessThan(u32, vocab_small));
+        if (rows > 0 and !chain) switch (t % 41) {
+            3 => round.picks.items[first + rng.uintLessThan(u32, rows)] = .{ .token = ref.invalid_token, .nonfinite = @as(u32, 1) << @intCast(rng.uintLessThan(u32, 3)) },
+            5 => round.picks.items[first + rng.uintLessThan(u32, rows)].nonfinite = 8,
+            7 => round.picks.items[first + rng.uintLessThan(u32, rows)].token = vocab_small,
+            11 => if (rows > 1) {
+                round.parents.items[first + rows - 1] = @intCast(rows - 1);
+            },
+            17 => round.tokens.items[first + rng.uintLessThan(u32, rows)] = vocab_small + 3,
+            19 => if (eos_count > 0) {
+                round.eos.items[eos_first] = vocab_small;
+            },
+            else => {},
+        };
+        try round.trees.append(gpa, .{ .first = first, .rows = rows, .budget = if (chain) 16 + rng.uintLessThan(u32, 2) else rng.uintLessThan(u32, 18), .eos_first = eos_first, .eos_count = eos_count });
+    }
+    return round;
+}
+
+/// The round on the GPU: its inputs, and results / keep / counts / gathered rows / heads for every tree.
+const Device = struct {
+    tokens: cuda.DeviceBuffer,
+    parents: cuda.DeviceBuffer,
+    picks: cuda.DeviceBuffer,
+    eos: cuda.DeviceBuffer,
+    trees: cuda.DeviceBuffer,
+    results: cuda.DeviceBuffer,
+
+    fn init(gpu: Gpu, round: *const Round) !Device {
+        // a zero-length round still needs addresses: one spare element each
+        const eos_bytes = std.mem.sliceAsBytes(round.eos.items);
+        return .{
+            .tokens = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(round.tokens.items)),
+            .parents = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(round.parents.items)),
+            .picks = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(round.picks.items)),
+            .eos = if (eos_bytes.len > 0) try cuda.DeviceBuffer.fromHost(gpu.d, eos_bytes) else try cuda.DeviceBuffer.alloc(gpu.d, 4),
+            .trees = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(round.trees.items)),
+            .results = try cuda.DeviceBuffer.alloc(gpu.d, round.trees.items.len * @sizeOf(tr.Result)),
+        };
+    }
+
+    fn deinit(self: *Device) void {
+        inline for (.{ "tokens", "parents", "picks", "eos", "trees", "results" }) |f| @field(self, f).free();
+    }
+
+    fn match(self: *const Device, k: *const tr.Kernels, s: cuda.Stream, count: usize) !void {
+        try self.results.fill8(0xA5, s.handle); // every field is written: a missed one shows (on s: a legacy-stream fill races it)
+        try k.match(s, self.tokens.ptr, self.parents.ptr, self.picks.ptr, self.eos.ptr, self.trees.ptr, self.results.ptr, vocab_small, count);
+        try s.synchronize();
+    }
+};
+
+fn results(gpu: Gpu, dev: *const Device, count: usize) ![]tr.Result {
+    const out = try gpu.gpa.alloc(tr.Result, count);
+    try dev.results.download(0, std.mem.sliceAsBytes(out));
+    return out;
+}
+
+/// Tree t's kept round rows: its verified path, for an ok result that consumed any.
+fn kept(r: *const tr.Result) usize {
+    return if (r.status == 0 and r.consumed_count > 0) r.consumed_count else 0;
+}
+
+fn matches(gpu: Gpu, k: *const tr.Kernels, s: cuda.Stream) !void {
+    const gpa = gpu.gpa;
+    var prng = std.Random.DefaultPrng.init(0x7AEE_0002);
+    const count = 4096;
+    var round = try buildRound(gpa, prng.random(), count);
+    defer round.deinit(gpa);
+    var dev = try Device.init(gpu, &round);
+    defer dev.deinit();
+    try dev.match(k, s, count);
+    const got = try results(gpu, &dev, count);
+    defer gpa.free(got);
+    var by_stop: [4]usize = @splat(0);
+    var by_status: [4]usize = @splat(0);
+    var deepest: u32 = 0;
+    for (got, 0..) |*g, t| {
+        const w = try round.want(t, vocab_small);
+        try check.sameBytes("tree result", std.mem.asBytes(g), std.mem.asBytes(&w));
+        const x = round.trees.items[t];
+        if (x.rows >= 1 and x.rows <= tr.max_rows) {
+            const d = ref.decode(@ptrCast(g), .{ .rows = x.rows, .vocab = vocab_small, .budget = x.budget, .eos_count = x.eos_count }) catch |e| {
+                std.debug.print("FAIL tree {d}: decode {s}\n", .{ t, @errorName(e) });
+                return error.TestFailed;
+            };
+            by_stop[@backingInt(d.stop)] += 1;
+            by_status[@backingInt(d.status)] += 1;
+            deepest = @max(deepest, g.matched_count);
+        }
+    }
+    for (by_stop, 0..) |n, i| try check.expect(n > 0, "the round never stopped by {s}", .{@tagName(std.enums.fromInt(ref.Stop, @as(u32, @intCast(i))).?)});
+    for (by_status, 0..) |n, i| try check.expect(n > 0, "the round never had status {s}", .{@tagName(std.enums.fromInt(ref.Status, @as(u32, @intCast(i))).?)});
+    try check.expect(deepest == tr.max_rows - 1, "no path matched more than {d} drafts", .{deepest});
+    check.pass("tree-round match: {d} batched trees equal the oracle and decode (stops bonus/length/eos/invalid {any}, deepest {d})", .{ count, by_stop, deepest });
+
+    // a tree alone gives its batched result: the batch shares nothing between trees
+    for (0..64) |t| {
+        const one = round.trees.items[t];
+        var single = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&one));
+        defer single.free();
+        try dev.results.fill8(0x5A, s.handle);
+        try k.match(s, dev.tokens.ptr, dev.parents.ptr, dev.picks.ptr, dev.eos.ptr, single.ptr, dev.results.ptr, vocab_small, 1);
+        try s.synchronize();
+        var alone: tr.Result = undefined;
+        try dev.results.download(0, std.mem.asBytes(&alone));
+        try check.sameBytes("tree alone", std.mem.asBytes(&alone), std.mem.asBytes(&got[t]));
+    }
+    check.pass("tree-round match: 64 trees alone equal their batched results", .{});
+
+    // mutation: one followed pick changed must change that tree's result, and the comparison must see it
+    const victim = for (got, 0..) |*g, t| {
+        if (g.status == 0 and g.matched_count >= 2) break t;
+    } else unreachable;
+    const at = round.trees.items[victim].first + got[victim].path[1];
+    const saved = round.picks.items[at];
+    var bent = saved;
+    bent.token = (saved.token + 1) % vocab_small;
+    try dev.picks.upload(at * @sizeOf(tr.Pick), std.mem.asBytes(&bent));
+    try dev.match(k, s, count);
+    const mutated = try results(gpu, &dev, count);
+    defer gpa.free(mutated);
+    try check.expect(!std.mem.eql(u8, std.mem.asBytes(&mutated[victim]), std.mem.asBytes(&got[victim])), "a changed pick on tree {d}'s path left its result alone", .{victim});
+    try dev.picks.upload(at * @sizeOf(tr.Pick), std.mem.asBytes(&saved));
+    check.pass("tree-round match: a changed pick on a verified path is caught", .{});
+
+    try commits(gpu, k, s, &round, &dev, got);
+}
+
+fn word(row: usize, c: usize) u16 {
+    return @truncate((row *% 0x9E37 +% c *% 0x85EB +% 1) | 1);
+}
+
+fn commits(gpu: Gpu, k: *const tr.Kernels, s: cuda.Stream, round: *const Round, dev: *const Device, got: []const tr.Result) !void {
+    const gpa = gpu.gpa;
+    const count = got.len;
+    try dev.match(k, s, count);
+    const total = round.tokens.items.len;
+
+    var keep = try cuda.DeviceBuffer.alloc(gpu.d, count * tr.max_rows * 4);
+    defer keep.free();
+    var counts = try cuda.DeviceBuffer.alloc(gpu.d, count * 4);
+    defer counts.free();
+    try keep.fill8(0xA5, s.handle);
+    try counts.fill8(0xA5, s.handle);
+    try k.keepRows(s, dev.results.ptr, dev.trees.ptr, keep.ptr, counts.ptr, count);
+    try s.synchronize();
+    const keep_got = try gpa.alloc(i32, count * tr.max_rows);
+    defer gpa.free(keep_got);
+    const counts_got = try gpa.alloc(u32, count);
+    defer gpa.free(counts_got);
+    try keep.download(0, std.mem.sliceAsBytes(keep_got));
+    try counts.download(0, std.mem.sliceAsBytes(counts_got));
+    const keep_want = try gpa.alloc(i32, count * tr.max_rows);
+    defer gpa.free(keep_want);
+    const counts_want = try gpa.alloc(u32, count);
+    defer gpa.free(counts_want);
+    for (got, 0..) |*r, t| {
+        const n = kept(r);
+        counts_want[t] = @intCast(n);
+        for (0..tr.max_rows) |i| keep_want[t * tr.max_rows + i] = if (i < n) @intCast(round.trees.items[t].first + r.path[i]) else -1;
+    }
+    try check.sameBytes("keep rows", std.mem.sliceAsBytes(keep_got), std.mem.sliceAsBytes(keep_want));
+    try check.sameBytes("keep counts", std.mem.sliceAsBytes(counts_got), std.mem.sliceAsBytes(counts_want));
+
+    const width = 300; // not a multiple of the 256-column blocks
+    const stride = width + 8;
+    const rows_host = try gpa.alloc(u16, total * stride);
+    defer gpa.free(rows_host);
+    for (0..total) |row| for (0..stride) |c| {
+        rows_host[row * stride + c] = word(row, c);
+    };
+    var src = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(rows_host));
+    defer src.free();
+    const sentinel: u16 = 0xA5A5;
+
+    // gather reads packed rows (width apart): the first total * width words of a packed copy
+    const packed_host = try gpa.alloc(u16, total * width);
+    defer gpa.free(packed_host);
+    for (0..total) |row| @memcpy(packed_host[row * width ..][0..width], rows_host[row * stride ..][0..width]);
+    var packed_rows = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(packed_host));
+    defer packed_rows.free();
+    var gathered = try cuda.DeviceBuffer.alloc(gpu.d, count * tr.max_rows * width * 2);
+    defer gathered.free();
+    try gathered.fill8(0xA5, s.handle);
+    try k.gatherRows(s, dev.results.ptr, dev.trees.ptr, packed_rows.ptr, gathered.ptr, width, count);
+    var head = try cuda.DeviceBuffer.alloc(gpu.d, count * width * 2);
+    defer head.free();
+    try head.fill8(0xA5, s.handle);
+    try k.keepHead(s, dev.results.ptr, dev.trees.ptr, src.ptr, head.ptr, width, stride, count);
+    try s.synchronize();
+
+    const gather_got = try gpa.alloc(u16, count * tr.max_rows * width);
+    defer gpa.free(gather_got);
+    try gathered.download(0, std.mem.sliceAsBytes(gather_got));
+    const head_got = try gpa.alloc(u16, count * width);
+    defer gpa.free(head_got);
+    try head.download(0, std.mem.sliceAsBytes(head_got));
+    const gather_want = try gpa.alloc(u16, count * tr.max_rows * width);
+    defer gpa.free(gather_want);
+    const head_want = try gpa.alloc(u16, count * width);
+    defer gpa.free(head_want);
+    @memset(gather_want, sentinel);
+    @memset(head_want, sentinel);
+    var moved: usize = 0;
+    for (got, 0..) |*r, t| {
+        const n = kept(r);
+        const first = round.trees.items[t].first;
+        for (0..n) |i| @memcpy(gather_want[(t * tr.max_rows + i) * width ..][0..width], packed_host[(first + r.path[i]) * width ..][0..width]);
+        if (n > 0) @memcpy(head_want[t * width ..][0..width], rows_host[(first + r.path[n - 1]) * stride ..][0..width]);
+        moved += n;
+    }
+    try check.sameBytes("gathered rows", std.mem.sliceAsBytes(gather_got), std.mem.sliceAsBytes(gather_want));
+    try check.sameBytes("kept heads", std.mem.sliceAsBytes(head_got), std.mem.sliceAsBytes(head_want));
+    check.pass("tree-round commit: keep rows, gather ({d} rows x {d}) and heads equal the host; nothing else written", .{ moved, width });
+}
+
+pub fn run(gpu: Gpu) !void {
+    var m = try cuda.Module.load(gpu.d, cuda.kernels.tree_round);
+    defer m.unload();
+    const k = try tr.Kernels.resolve(m);
+    var s = try cuda.Stream.init(gpu.d, true);
+    defer s.deinit();
+    try picks(gpu, &k, s);
+    try matches(gpu, &k, s);
+}

From 4bc13eeaa5b32677a3b57884e9cbdb68e0a8fb2a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BCrgen=20Schmied?=
 <600316+jschmied@users.noreply.github.com>
Date: Sat, 17 Oct 2026 06:49:32 +0200
Subject: [PATCH 2/2] cuda/tree_round: registered with the kernel contract

---
 zig/src/core/contract/registry.zig     |   5 +
 zig/tests/cuda/contract_tests.zig      |  10 +-
 zig/tests/cuda/contract_tree_round.zig | 211 +++++++++++++++++++++++++
 zig/tests/cuda/tree_round_tests.zig    |  12 +-
 4 files changed, 231 insertions(+), 7 deletions(-)
 create mode 100644 zig/tests/cuda/contract_tree_round.zig

diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig
index b79337b5c5..f3b46831db 100644
--- a/zig/src/core/contract/registry.zig
+++ b/zig/src/core/contract/registry.zig
@@ -45,6 +45,11 @@ pub const cuda = [_]CudaRecord{
     .{ .launch = .{ .name = "vit tf_vit_pointwise", .symbol = "tf_vit_pointwise", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 }, .{ .regs = 14, .static = 0, .bound = 256 } } },
     .{ .launch = .{ .name = "vit tf_vit_bias", .symbol = "tf_vit_bias", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 }, .{ .regs = 12, .static = 0, .bound = 256 } } },
     .{ .launch = .{ .name = "vit tf_vit_unshuffle", .symbol = "tf_vit_unshuffle", .image = "vit", .threads = 256 }, .usage = .{ .{ .regs = 29, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 }, .{ .regs = 34, .static = 0, .bound = 256 } } },
+    .{ .launch = .{ .name = "tree_round tf_round_bf16_argmax", .symbol = "tf_round_bf16_argmax", .image = "tree_round", .threads = 256 }, .usage = .{ .{ .regs = 25, .static = 3072, .bound = 256 }, .{ .regs = 25, .static = 3072, .bound = 256 }, .{ .regs = 25, .static = 3072, .bound = 256 }, .{ .regs = 23, .static = 3072, .bound = 256 }, .{ .regs = 23, .static = 3072, .bound = 256 } } },
+    .{ .launch = .{ .name = "tree_round tf_tree_round_match", .symbol = "tf_tree_round_match", .image = "tree_round", .threads = 64 }, .usage = .{ .{ .regs = 25, .static = 0, .bound = 0 }, .{ .regs = 26, .static = 0, .bound = 0 }, .{ .regs = 26, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 } } },
+    .{ .launch = .{ .name = "tree_round tf_round_keep_rows", .symbol = "tf_round_keep_rows", .image = "tree_round", .threads = 16 }, .usage = .{ .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 16, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 } } },
+    .{ .launch = .{ .name = "tree_round tf_round_gather_rows", .symbol = "tf_round_gather_rows", .image = "tree_round", .threads = 256 }, .usage = .{ .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 } } },
+    .{ .launch = .{ .name = "tree_round tf_round_keep_head", .symbol = "tf_round_keep_head", .image = "tree_round", .threads = 256 }, .usage = .{ .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 } } },
 };
 
 /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has.
diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig
index e87a2e90e5..64644a03f1 100644
--- a/zig/tests/cuda/contract_tests.zig
+++ b/zig/tests/cuda/contract_tests.zig
@@ -8,6 +8,7 @@ const rows_k = @import("contract_rows.zig");
 const router_k = @import("contract_router.zig");
 const mm_k = @import("contract_mm.zig");
 const rank_k = @import("contract_rank.zig");
+const round_k = @import("contract_tree_round.zig");
 
 const contract = core.contract;
 const g = contract.gates;
@@ -21,7 +22,7 @@ const controls = struct {
 
 /// The embedded image a registration names (the controls' only when this binary was built with -Dnvcc).
 fn image(name: []const u8) []const u8 {
-    inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n);
+    inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit", "tree_round" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n);
     return if (options.controls) &controls.bytes else &.{};
 }
 
@@ -89,6 +90,13 @@ pub fn run(gpu: check.Gpu) !void {
     var rank = rank_k.Rank{ .rig = .{ .gpu = gpu, .s = s }, .sum = &rank_sum };
     defer rank.deinit();
     try kernels.append(a, rank.kernel());
+    var round_mod = try cuda.Module.load(gpu.d, kk.tree_round);
+    defer round_mod.unload();
+    const round_fns = try cuda.tree_round.Kernels.resolve(round_mod);
+    var rounds: [std.enums.values(round_k.Which).len]round_k.Round = undefined;
+    for (std.enums.values(round_k.Which), &rounds) |w, *r| r.* = .{ .rig = .{ .gpu = gpu, .s = s }, .which = w, .k = &round_fns };
+    defer for (&rounds) |*r| r.deinit();
+    for (&rounds) |*r| try kernels.append(a, r.kernel());
     var t: g.Tally = .{};
     const t0 = check.now(gpu.io);
     try g.runAll(a, kernels.items, 0x6b63, &t);
diff --git a/zig/tests/cuda/contract_tree_round.zig b/zig/tests/cuda/contract_tree_round.zig
new file mode 100644
index 0000000000..5e0bf30846
--- /dev/null
+++ b/zig/tests/cuda/contract_tree_round.zig
@@ -0,0 +1,211 @@
+//! tree_round.cu (#734) under the contract on CUDA: picks, verified trees and commits, byte for byte against the host.
+const std = @import("std");
+const cuda = @import("cuda");
+const core = @import("core");
+const rig_ = @import("contract_rig.zig");
+const rt = @import("tree_round_tests.zig");
+
+const contract = core.contract;
+const c = contract.spec;
+const g = contract.gates;
+const tr = cuda.tree_round;
+const ref = core.tree_round;
+
+pub const Which = enum { argmax, match, keep, gather, head };
+
+fn noSums(_: c.Shape) usize {
+    return 0;
+}
+
+const pick_rule = "a row holding NaN or an infinity picks no token and flags it (1 NaN, 2 +inf, 4 -inf); ties go to the lower token";
+const match_rule = "a flagged pick or a malformed tree gives core/tree_round.oracle's invalid result, byte for byte";
+const move_rule = "a tree whose result is not ok moves nothing; kept rows move bit for bit, NaN and infinities included";
+
+/// argmax: rows x vocab (k); match: trees (rows) of at most 16 rows (k); commits: trees (rows) x row width (n).
+pub fn spec(which: Which) c.Spec {
+    return switch (which) {
+        // a 27B tree (16 rows), eight streams' trees, Nemotron's 131,072 tokens
+        .argmax => .{ .name = "cuda tree_round tf_round_bf16_argmax", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = pick_rule }, .tiles = .{ .{}, .{}, .{ .size = 256 } }, .served = &.{ .{ .rows = 16, .n = 1, .k = 248320 }, .{ .rows = 128, .n = 1, .k = 248320 }, .{ .rows = 48, .n = 1, .k = 131072 } } },
+        .match => .{ .name = "cuda tree_round tf_tree_round_match", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = match_rule }, .tiles = .{ .{ .size = 64 }, .{}, .{} }, .served = &.{ .{ .rows = 8, .n = 1, .k = 16 }, .{ .rows = 200, .n = 1, .k = 16 } } },
+        .keep => .{ .name = "cuda tree_round tf_round_keep_rows", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = move_rule }, .served = &.{ .{ .rows = 8, .n = 16, .k = 16 }, .{ .rows = 200, .n = 16, .k = 16 } } },
+        // hidden rows of the 27B (5,120) and Nemotron (2,688)
+        .gather => .{ .name = "cuda tree_round tf_round_gather_rows", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = move_rule }, .tiles = .{ .{}, .{ .size = 256 }, .{} }, .served = &.{ .{ .rows = 8, .n = 5120, .k = 16 }, .{ .rows = 64, .n = 2688, .k = 16 } } },
+        // the next round's first logits row
+        .head => .{ .name = "cuda tree_round tf_round_keep_head", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = move_rule }, .tiles = .{ .{}, .{ .size = 256 }, .{} }, .served = &.{ .{ .rows = 8, .n = 248320, .k = 16 }, .{ .rows = 48, .n = 131072, .k = 16 } } },
+    };
+}
+
+/// A large finite value past each row's vocab: a kernel reading the padding picks it.
+const decoy: u16 = 0x7F00;
+const pad = 8;
+
+pub const Round = struct {
+    rig: rig_.Rig,
+    which: Which,
+    k: *const tr.Kernels,
+    shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 },
+    round: ?rt.Round = null,
+    results: []ref.Result = &.{},
+    words: []u16 = &.{},
+    in: [6]u64 = @splat(0),
+    out: [2]u64 = @splat(0),
+    msg: [200]u8 = undefined,
+
+    pub fn deinit(self: *Round) void {
+        self.drop();
+        self.rig.deinit();
+    }
+
+    fn drop(self: *Round) void {
+        const a = self.rig.gpu.gpa;
+        if (self.round) |*r| r.deinit(a);
+        self.round = null;
+        a.free(self.results);
+        self.results = &.{};
+        a.free(self.words);
+        self.words = &.{};
+    }
+
+    pub fn kernel(self: *Round) g.Kernel {
+        return .{ .spec = spec(self.which), .ctx = self, .vt = &vt };
+    }
+
+    const vt = g.hooks(Round);
+
+    fn stride(self: *const Round) usize {
+        return self.shape.k + pad;
+    }
+
+    pub fn prepare(self: *Round, shape: c.Shape, fill: g.Fill, seed: u64) !bool {
+        self.rig.reset();
+        self.drop();
+        if (shape.rows == 0 or shape.n == 0 or shape.k == 0) return false;
+        self.shape = shape;
+        const a = self.rig.gpu.gpa;
+        var prng = std.Random.DefaultPrng.init(seed);
+        const rng = prng.random();
+        if (self.which == .argmax) {
+            self.words = try a.alloc(u16, shape.rows * self.stride());
+            @memset(self.words, decoy);
+            // gate 1: rows 0-3 hold NaN, +inf, -inf, and every -inf; the other gates also see them every eighth row
+            for (0..shape.rows) |r| rt.fillRow(rng, if (fill == .nonfinite) r + 2 else r, self.words[r * self.stride() ..][0..shape.k]);
+            self.in[0] = try self.rig.input(std.mem.sliceAsBytes(self.words));
+            self.out[0] = try self.rig.output(shape.rows * @sizeOf(tr.Pick));
+            return true;
+        }
+        self.round = try rt.buildRound(a, rng, shape.rows);
+        const rd = &self.round.?;
+        if (fill == .nonfinite) for (rd.trees.items[0..@min(4, shape.rows)], 0..) |t, i| if (t.rows > 0 and t.rows <= tr.max_rows) {
+            rd.picks.items[t.first + t.rows - 1] = .{ .token = ref.invalid_token, .nonfinite = ([_]u32{ 1, 2, 4, 6 })[i] };
+        };
+        if (self.which == .match) {
+            const lists = [_][]const u8{ std.mem.sliceAsBytes(rd.tokens.items), std.mem.sliceAsBytes(rd.parents.items), std.mem.sliceAsBytes(rd.picks.items), std.mem.sliceAsBytes(rd.eos.items), std.mem.sliceAsBytes(rd.trees.items) };
+            for (lists, 0..) |l, i| self.in[i] = try self.rig.input(l);
+            self.out[0] = try self.rig.output(shape.rows * @sizeOf(tr.Result));
+            return true;
+        }
+        self.results = try a.alloc(ref.Result, shape.rows);
+        for (self.results, 0..) |*r, t| r.* = try rd.want(t, rt.vocab_small);
+        self.in[0] = try self.rig.input(std.mem.sliceAsBytes(self.results));
+        self.in[1] = try self.rig.input(std.mem.sliceAsBytes(rd.trees.items));
+        const total = rd.tokens.items.len;
+        switch (self.which) {
+            .keep => {
+                self.out[0] = try self.rig.output(shape.rows * tr.max_rows * 4);
+                self.out[1] = try self.rig.output(shape.rows * 4);
+            },
+            .gather, .head => {
+                const width = if (self.which == .head) shape.n + pad else shape.n;
+                self.words = try a.alloc(u16, @max(total, 1) * width);
+                for (self.words) |*w| w.* = rng.int(u16);
+                if (fill == .nonfinite) for (self.words, 0..) |*w, i| switch (i % 7) {
+                    1 => w.* = 0x7FC0,
+                    2 => w.* = 0x7F80,
+                    3 => w.* = 0xFF80,
+                    else => {},
+                };
+                self.in[2] = try self.rig.input(std.mem.sliceAsBytes(self.words));
+                const rows: usize = if (self.which == .head) 1 else tr.max_rows;
+                self.out[0] = try self.rig.output(shape.rows * rows * shape.n * 2);
+            },
+            else => unreachable,
+        }
+        return true;
+    }
+
+    pub fn run(self: *Round) !void {
+        const s = self.rig.s;
+        const n = self.shape.rows;
+        switch (self.which) {
+            .argmax => try self.k.argmax(s, self.in[0], self.out[0], n, self.shape.k, self.stride()),
+            .match => try self.k.match(s, self.in[0], self.in[1], self.in[2], self.in[3], self.in[4], self.out[0], rt.vocab_small, n),
+            .keep => try self.k.keepRows(s, self.in[0], self.in[1], self.out[0], self.out[1], n),
+            .gather => try self.k.gatherRows(s, self.in[0], self.in[1], self.in[2], self.out[0], self.shape.n, n),
+            .head => try self.k.keepHead(s, self.in[0], self.in[1], self.in[2], self.out[0], self.shape.n, self.shape.n + pad, n),
+        }
+        try s.synchronize();
+    }
+
+    pub fn count(self: *Round) usize {
+        const n = self.shape.rows;
+        return switch (self.which) {
+            .argmax => n * @sizeOf(tr.Pick),
+            .match => n * @sizeOf(tr.Result),
+            .keep => n * (tr.max_rows + 1) * 4,
+            .gather => n * tr.max_rows * self.shape.n * 2,
+            .head => n * self.shape.n * 2,
+        };
+    }
+
+    pub fn reference(_: *Round, want: []f64, mag: []f64) !void {
+        @memset(want, 0);
+        @memset(mag, 0);
+    }
+
+    /// The host's bytes (core/tree_round's rowArgmax and oracle, the kept rows); 0xff where nothing writes.
+    pub fn check(self: *Round, _: g.Gate) !?[]const u8 {
+        const a = self.rig.gpu.gpa;
+        const got = try self.rig.outputBytes();
+        const want = try a.alloc(u8, got.len);
+        defer a.free(want);
+        @memset(want, 0xff);
+        const n = self.shape.rows;
+        switch (self.which) {
+            .argmax => for (0..n) |r| {
+                const p = ref.rowArgmax(self.words[r * self.stride() ..][0..self.shape.k]);
+                @memcpy(want[r * 8 ..][0..8], std.mem.asBytes(&p));
+            },
+            .match => for (0..n) |t| {
+                const w = try self.round.?.want(t, rt.vocab_small);
+                @memcpy(want[t * @sizeOf(tr.Result) ..][0..@sizeOf(tr.Result)], std.mem.asBytes(&w));
+            },
+            .keep => {
+                const keep = std.mem.bytesAsSlice(i32, want[0 .. n * tr.max_rows * 4]);
+                const counts = std.mem.bytesAsSlice(u32, want[n * tr.max_rows * 4 ..]);
+                for (self.results, 0..) |*r, t| {
+                    const kept = keptRows(r);
+                    counts[t] = @intCast(kept);
+                    for (0..tr.max_rows) |i| keep[t * tr.max_rows + i] = if (i < kept) @intCast(self.round.?.trees.items[t].first + r.path[i]) else -1;
+                }
+            },
+            .gather, .head => {
+                const w = self.shape.n;
+                const words = std.mem.bytesAsSlice(u16, want);
+                for (self.results, 0..) |*r, t| {
+                    const kept = keptRows(r);
+                    const first = self.round.?.trees.items[t].first;
+                    if (self.which == .gather) {
+                        for (0..kept) |i| @memcpy(words[(t * tr.max_rows + i) * w ..][0..w], self.words[(first + r.path[i]) * w ..][0..w]);
+                    } else if (kept > 0) @memcpy(words[t * w ..][0..w], self.words[(first + r.path[kept - 1]) * (w + pad) ..][0..w]);
+                }
+            },
+        }
+        if (std.mem.indexOfDiff(u8, got, want)) |at| return try std.fmt.bufPrint(&self.msg, "byte {d} of {d} is 0x{x:0>2} where the host gives 0x{x:0>2}", .{ at, got.len, got[at], want[at] });
+        return null;
+    }
+};
+
+/// A tree's kept round rows: its verified path, for an ok result that consumed any.
+fn keptRows(r: *const ref.Result) usize {
+    return if (r.status == 0 and r.consumed_count > 0) r.consumed_count else 0;
+}
diff --git a/zig/tests/cuda/tree_round_tests.zig b/zig/tests/cuda/tree_round_tests.zig
index aaf0d6bc26..793d7f76e4 100644
--- a/zig/tests/cuda/tree_round_tests.zig
+++ b/zig/tests/cuda/tree_round_tests.zig
@@ -26,7 +26,7 @@ const neg_inf: u16 = 0xFF80;
 const decoy: u16 = 0x7F00;
 
 /// Row r by r's kind: spread, coarse ties, a NaN, +inf, -inf, all -inf, the max last, ties at both ends.
-fn fillRow(rng: std.Random, r: usize, row: []u16) void {
+pub fn fillRow(rng: std.Random, r: usize, row: []u16) void {
     const v = row.len;
     for (row) |*h| h.* = switch (r % 8) {
         1 => bf16(-@floor(rng.float(f32) * 4.0)),
@@ -76,14 +76,14 @@ fn picks(gpu: Gpu, k: *const tr.Kernels, s: cuda.Stream) !void {
 }
 
 /// A round of `count` random trees in the kernels' layout, with the host reference's per-tree slices.
-const Round = struct {
+pub const Round = struct {
     tokens: std.ArrayList(u32) = .empty,
     parents: std.ArrayList(i32) = .empty,
     picks: std.ArrayList(tr.Pick) = .empty,
     eos: std.ArrayList(u32) = .empty,
     trees: std.ArrayList(tr.Tree) = .empty,
 
-    fn deinit(self: *Round, gpa: std.mem.Allocator) void {
+    pub fn deinit(self: *Round, gpa: std.mem.Allocator) void {
         self.tokens.deinit(gpa);
         self.parents.deinit(gpa);
         self.picks.deinit(gpa);
@@ -92,7 +92,7 @@ const Round = struct {
     }
 
     /// Tree t's expected result: the oracle's, or a bad tree where the oracle refuses the shape.
-    fn want(self: *const Round, t: usize, vocab: u32) !ref.Result {
+    pub fn want(self: *const Round, t: usize, vocab: u32) !ref.Result {
         const x = self.trees.items[t];
         const p = ref.Match{ .rows = x.rows, .vocab = vocab, .budget = x.budget, .eos_count = x.eos_count };
         const rows = @min(x.rows, self.tokens.items.len - x.first);
@@ -102,9 +102,9 @@ const Round = struct {
     }
 };
 
-const vocab_small: u32 = 6; // small enough that drafts match the picks often
+pub const vocab_small: u32 = 6; // small enough that drafts match the picks often
 
-fn buildRound(gpa: std.mem.Allocator, rng: std.Random, count: usize) !Round {
+pub fn buildRound(gpa: std.mem.Allocator, rng: std.Random, count: usize) !Round {
     var round: Round = .{};
     errdefer round.deinit(gpa);
     for (0..count) |t| {