[jschmied]

What this changes

Metal runs a drafted tree round on the GPU through core/tree_round.zig's contract: one pick per row, each tree's verified path and emitted tokens, then the commit steps that read that result. CUDA has tree attention and the GDN tree kernel, but no verify or commit step, so every CUDA family that drafts a tree has to write its own.

This adds that step as a shared module, zig/kernels/cuda/tree_round.cu with zig/src/cuda/tree_round.zig:

• tf_round_bf16_argmax: each row's first largest finite bf16 logit, at any row count and stride. A row with NaN or ±inf gets no token and the mask core/tree_round uses (1, 2, 4).
• tf_tree_round_match: oracle() for a batch of trees, one thread per tree. Each Tree{first, rows, budget, eos_first, eos_count} names its rows, its budget and its EOS ids. Parents are tree-local. Results use core/tree_round.Result's layout, so decode() reads them unchanged.
• Commit: tf_round_keep_rows (each tree's kept round rows and their count), tf_round_gather_rows (those rows of any bf16 buffer, in path order, for hidden rows or drafter taps) and tf_round_keep_head (the last kept row's logits, which the next round's first pick reads). Each one ignores a tree whose result is not ok.

No family calls it yet. Nemotron's chain path is unchanged. The 27B in #608 and any other tree drafter can take it as it is. core/root.zig now exports tree_round so the CUDA test runner can check against it.

Receipt

• Environment: TensorFold 46a08c18 (1.0.12) plus this commit; one DGX Spark (GB10, sm_121, driver 580.173.02, CUDA 13). No checkpoint needed.
• Exactness, tf-cuda-test tree-round:
  • Picks: 168 rows over 7 vocabularies (1 to 151,936) equal rowArgmax bit for bit. That covers ties, NaN, +inf, -inf and all -inf rows. Each row's padding past vocab holds a large finite decoy that must never be picked.
  • Match: 4,096 random trees in one launch equal oracle() byte for byte, and decode() accepts every result. The trees hit every stop (bonus, length, EOS, invalid) and every status, and include full 16-row paths, budget 0, and malformed trees and picks. 64 trees run alone equal their batched results.
  • Commit: kept rows, counts, gathered rows (width 300, not a multiple of the block) and heads equal the host's. Everything outside the kept rows stays untouched.
• Kernel contract (docs/kernel-contract.md), CUDA on GB10 (sm_121), zig build kernel-contract -Dnvcc=/usr/local/cuda/bin/nvcc, exit 0:

  SUMMARY 17 kernels: 17 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 0 limit problems; gates 43597 ms, limits 1 ms
  

  • All five kernels are registered (zig/tests/cuda/contract_tree_round.zig): tf_round_bf16_argmax, tf_tree_round_match, tf_round_keep_rows, tf_round_gather_rows and tf_round_keep_head. None of them sums anything, so each declares accum = none (gate 4 is N/A) and checks its own bytes: picks against rowArgmax, results against oracle(), and the commits against the host's kept rows, with every byte nothing writes still holding the rig's poison.
  • Each states its NaN rule: a row with NaN or ±inf picks no token and is flagged 1, 2 or 4; a flagged pick or malformed tree gives the oracle's invalid result; a tree that is not ok moves nothing, and kept rows move bit for bit, NaN and infinities included. Gate 1 plants those values.
  • Served shapes: picks for a 27B tree (16 × 248,320), eight streams' trees (128 × 248,320) and Nemotron (48 × 131,072); 8 and 200 trees for match and keep; hidden rows of 5,120 and 2,688 for gather; logits rows of 248,320 and 131,072 for the head. Tile edges are the 256-thread vocab and column blocks and the 64-tree match blocks.
  • Gate 5: all five launches are in registry.zig, with compile records from zig build fatbins kernel-contract-fatbins -Dsm=80,86,89,120,121 and tf-cuda-test kernel-contract-limits zig-out/fatbin (PASS kernel-contract-limits: every registered launch fits sm_80, sm_86, sm_89, sm_120 and sm_121, every control red). On the GB10 the live check reads PASS limits: 31 registered CUDA launches against sm_80/86/89/120/121, 0 problems.
  • Mutations: reversing the argmax tie order fails gates 1 and 3 for tf_round_bf16_argmax. Taking the head from the path's first row instead of its last fails gates 1 and 3 for tf_round_keep_head. Both were restored before the run above.
  • Not run: Metal, which has no kernels in this PR. sm_80, sm_86, sm_89 and sm_120 are covered by compile records only, not run on a GPU.
• Mutations: five changes to the kernels each fail the test: tie order, a skipped EOS id, the keep count at depth 16, the head row, and a child matched without checking its parent. So does a changed pick on a verified path.
• Speed: none claimed. Nothing calls it yet.
• Tests run: zig build test (84/84 steps, 1,370 of 1,377 tests, 7 skipped; 84/84 steps again with the registration) and zig build test-cuda-host (39/39 steps), all passing. tf-cuda-test tree-round passes with the numbers above. tf-cuda-test sample still matches 2,880 of 2,880. zig fmt --check clean. lean_check.py: 517 findings, as on main.
• Not run, and why: Metal (unchanged); no model, since no family is wired to it in this PR.

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (No serving path changes.)
• No precision traded for speed. If the bits change, the description says which sums change and why. (Picks are exact; nothing is summed.)
• A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt. (CUDA; no Metal kernels here.)
• Prompt processing is no slower than the last release.
• The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
• New tests fail before the change, pass after it, and skip cleanly without their dependency.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[jschmied]

Rebased onto main 46a08c18 (1.0.12), now at 9c0c93ae. The conflicts were only in the build and test registration lists, and the diff against main is otherwise unchanged. The GPU checks give the same numbers as before, and the description has the new test and lint counts.

[jschmied]

The kernels are now registered with the kernel contract in 4bc13eea (all five, launch records for sm_80 to sm_121). The CUDA SUMMARY line and the mutation checks are in the description.

[ashhart]

Thanks @jschmied for registering all five kernels and supplying the tie/head mutation receipts. The registration increment can land with our prepared integration fixes.

The CUDA result still holds 16 rows while the common oracle now holds 64. We kept the CUDA layout in the adapter and its existing oracle conversion; the host control fails at 168 versus 552 bytes before that repair. All five launch records stay pending until measured on the integrated images. Production kernels are unchanged by this increment. The host checks pass, but we have not rerun the CUDA numerical, finite-bit or resource gates.

Posted by Ash's AI Assistant Synthena

[ashhart]

Thank you, @jschmied, for verified CUDA tree rounds and the 1.0.12 rebase.

We merged the reviewed head 9c0c93aef for 1.1.0 in 01bfea6. It adds the shared keep, gather and head rounds; no family calls them yet.

Your newer registration increment 4bc13eeaa is for 1.1.1, the next patch release. It needs the paired fixture repair that keeps CUDA's 16-row result layout beside the common 64-row oracle, and all five launches stay pending until the integrated images have their receipts.

NVIDIA GB10/RTX PRO 6000 release contracts passed 55/55 with five controls red; these operator checks do not establish family adoption.

Posted by Ash's AI Assistant Synthena

[ashhart]

Thanks @jschmied for #734's registration follow-up. 1.1.1 includes the five tree-round launch registrations and their numerical fixtures, with the fixture adapted to the common 64-row oracle while retaining the CUDA result layout. This increment adds contract coverage; it does not change production arithmetic.

Posted by Ash's AI Assistant Synthena