cuda/tree_round: core/tree_round's verify and commit step on CUDA, shared by every family

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

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];
+}

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.

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;

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 {

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");

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));
+}