[jschmied]

What this changes

CUDA has no shared attention for drafted rows. A family that drafts a tree on CUDA writes its own attention over the
committed keys plus each row's path, and a DFlash-style drafter writes its own block attention. This adds both as
shared modules, operator-only (no model calls them yet), next to cuda/tree_round (#734), which does the verify and
commit step of the same round:

• zig/kernels/cuda/tree_attention.cu, zig/src/cuda/tree_attention.zig: drafted rows of any number of streams in
  one call. Each row sees its stream's committed keys and then its root-to-row path (paths() builds the paths from
  tree-local parents, depth up to 128). Keys run in 512-key groups on tensor cores, merged in group order, at head
  dims 128 and 256 with grouped KV heads. A stream's fully committed groups are computed once for all its rows (16
  (row, head) pairs as one mma tile, the same tiles and merge as a row's own block), so every partial keeps its bits
  whether a row runs alone or beside others. The grid is fixed: at most 16 group slots per (row, key head), each
  walking its groups, so a captured graph does not grow with the context. A group's key tiles go to 4 slices; a
  four-warp block combines its slices in order; where four-warp blocks would leave SMs idle (TreeAttention.sms),
  each slice takes its own one-warp block and the merge combines the 4 in the same order (a small fixed slice area
  in the scratch), so the bits do not depend on the device's size. A one-warp block fetches its tiles 2 ahead (3
  cp.async stages where the one-warp grid fits one wave of the SMs, else 1 ahead in 2, two blocks an SM; stages
  never change bits) while it computes one. Partials hold the query heads only. The merge runs a block per (row, key
  head, head).
• zig/kernels/cuda/block_attention.cu, zig/src/cuda/block_attention.zig: a DFlash-style drafter's block rows
  over their stream's context (each stream's K and V cache through a pointer table, within a sliding window) and then the block itself,
  causal or not. One warp per head on tensor cores.
• zig/src/core/attention_ref.zig: one query head's attention in float64 with a per-element tolerance (2^-7 of
  sum p|v| / sum p, plus a bf16 half-ulp), the reference both tests check against.

NaN and infinities (found by the kernel contract's gate 1, fixed here):

• tree_attention.cu: the flag that lets a tile move the running max was taken from the first of the two mma rows
  a thread holds and applied to both, so a head whose scores were all NaN stopped another head's max (dim 128: head
  8 all NaN because head 0 had a NaN query). Each row now has its own flag.
• block_attention.cu: a key a row cannot see (the window's lower edge, a causal block's later keys) sits in the
  same 32-key tile and the same mma as rows that do see it, with P = 0, so an infinite value there gave 0 x inf =
  NaN. Its V cannot be zeroed for one row only, and the mma cannot select per row. So the kernel checks its output:
  a block whose output holds a NaN or an infinity runs again, with every NaN and infinite value of V set to 0 in
  shared memory before the mma, and each row adding back the ones it sees after the last tile, at the key's float64
  weight: the key's score and the row's max recomputed in float64 in attention_ref's order, and an infinite value
  whose weight exp(score - max) is 0 adds 0 x inf = NaN, as float64 does (only where its weights sum is positive, so
  a row with no weight stays NaN). The block's own tile runs after the context's values are added back, and leaves
  an infinite o as it is (attn_odd::keep): when a block key holds the row's max and the context scores sit more
  than ~103 below it, its rescale is 0 in float32. The shared step is attn_odd.cuh.
• tree_attention.cu: the same replay. A visible infinite value whose weight underflows in float32 (more than ~103
  below the max) but not in float64 (745) went through the mma as 0 x inf, or met a merge factor of 0: NaN, where
  float64 gives inf. A block whose partial holds a NaN or an infinity now runs again, out of line: its V tiles' NaN
  and infinite values set to 0 before the mma, and each (row, head) adding them back after its last tile at the
  key's float64 weight against its row's float64 max over all its keys (committed, then its path). The four-warp
  combine and the merge fold as before; only a NaN result folds again with those infinities unscaled
  (attn_odd::fold, keep). tf_tree_shared takes the window's keys (kn) as a new last parameter, for the path
  keys' weights.

Users we know of: the Qwen3.8-27B family in #608 (target tree attention at dim 256 and its DFlash2 drafter), and any
CUDA family that drafts trees with #734.

Receipt

• Environment: TensorFold 46a08c18 (1.0.12) plus these commits; one DGX Spark (GB10, sm_121, CUDA 13). No
  checkpoint for the tests.
• Exactness:
  • tf-cuda-test tree: dim 128 (48 query / 4 KV heads) and 256 (24 / 4), several streams with different committed
    lengths and trees, every element within the float64 tolerance (worst 0.47 of it). A stream's rows run alone
    equal their bits in the batch; one shared block walking every (row, group), a grid sized for a 262k context,
    and one-warp blocks equal the default grid.
  • tf-cuda-test block: windowed contexts then the block, causal and not, block widths 7 and 16, within
    tolerance (worst 0.38). A stream run alone equals its bits in the batch.
• Kernel contract (docs/kernel-contract.md), CUDA on GB10 (sm_121), zig build kernel-contract -Dnvcc=/usr/local/cuda/bin/nvcc, exit 0:

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

  • Before the two fixes: SUMMARY 21 kernels: 16 pass gates 1-4, 5 fail, all five at gate 1 (the three dim-128
    tree forms, both block forms).
  • Registered (zig/tests/cuda/contract_tree_attn.zig, contract_block_attn.zig, contract_attn.zig):
    tf_tree_paths; tree attention in all six forms (dim 128 with 48/4 heads and dim 256 with 24/4, four-warp
    blocks, one-warp blocks with 2 stages, one-warp blocks with 3 stages; each form's shared, part and merge launches
    as one call); block attention as the DFlash2 drafter serves it (blocks of 16, 32/8 heads, window 2,048) and causal
    blocks of 7.
  • Served shapes: one tree of 16 rows over 2,100 committed keys and four trees (64 rows) over 700, at both dims;
    one stream's block over 2,100 context keys (past the window) and four streams over 700; paths for 16, 128 and 600
    rows. Tile edges: the 16-row tree, the 512-key group, the 32-key block tile, the 128-row path block.
  • NaN rule: as float64 (propagate). Gate 1 plants a NaN query, a NaN key, and +inf and -inf values: in the
    context, in a block's key (which earlier causal rows cannot see), at the window's edge (a key only block row 0
    sees), and +inf and -inf in one column. For block attention, +inf values row 0 sees at a float64 weight of 0
    (scores 1,448 apart: a key in the window's edge tile, and the block's first key) and at weights below float32's
    least but not float64's (scores 362 apart, every key of a key head, edge tiles and the others); on stream 1, key
    head 7, block key 0 holding the max over keys all -8 with +inf, queried all 2 (inf) and all 8 (NaN). For tree
    attention, key head 1, row 0: committed key p0 - 1 all 8, every other key (committed and its path) all -8 with
    +inf, queried all 2 (inf, below float32's least) and all 8 (NaN, float64 weight 0). tf_tree_paths
    states its rule (integers only: depths and the nearest 128 ancestors, other slots untouched).
  • Bounds: the framework's bound takes fp32 terms, and it cannot express P rounded to bf16 before P x V. So the
    attention entries check gates 1 and 3 against core/attention_ref.zig's float64 bound (2^-7 of sum p|v| / sum p
    plus a bf16 half-ulp, the bound tf-cuda-test tree and block use) through their own check hook; gate 3's line
    reads "exact by its own check". Gate 4 builds sums where every key scores alike (p = 1): exact at 2,100 keys,
    where fp16 sums would miss by inf and bf16 by 1.6 to 2.1 million times.
  • Gate 5: all 15 launches (paths, 12 tree, merge, block) 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). The fixes moved register counts, re-recorded: tf_block_attention 166 to 230 (sm_80/86/89) and
    170 to 228 (sm_120/121), its 256-thread block still one an SM on every target; nine tree launches by -10 to +12
    on sm_80/86/89 only. The tree replay re-records the tree launches: dim 128 from 162-168 to 190-228, dim 256 from
    208-243 to 249-255 (255 on sm_80/86/89; no local memory on any target), the merge 148 to 128 (sm_120/121) and 97
    to 114 (sm_80/86/89). On the GB10 the live check reads PASS limits: 41 registered CUDA launches against sm_80/86/89/120/121, 0 problems.
  • Mutations: each fix reverted alone turns gate 1 red (the shared flag: the three dim-128 tree forms, "head 8
    element 0: nan where float64 gives finite"; the block kernel as it was: both block forms). The replay as
    reviewed (ee81000d, adding a value back without its key's weight, and only for the edge tiles): both block forms,
    "head 20 element 17: pos_inf where float64 gives nan"; with only the below-float32 plant, "head 28 element 21: nan
    where float64 gives pos_inf"; with only the block's first key, causal, "head 24 element 19: pos_inf where float64
    gives nan". On ff6e5c0e the plants added with the tree replay are red: the six tree forms, "head 12 element 21:
    nan where float64 gives pos_inf" (dim 256: head 6); both block forms, "head 540 element 23: nan where float64
    gives pos_inf" (causal: head 252). The float64 weight test forced true: the six tree forms, "head 13 element 21:
    pos_inf where float64 gives nan" (dim 256: head 7), and both block forms. Holding the attention
    bound 64x tighter turns gate 3 red for all eight attention entries. All 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.
• The fixes keep every finite output's bytes. The fork heads before them (591065e0) and after, the same inputs:
  every changed kernel over the registration's served shapes and tile edges with 12 seeds in two finite fills
  (normal values, and gate 4's alike-key fill), plus sweeps with 2 seeds: trees of 1 to 64 rows over 1 to 4,096
  committed keys (both sides of each 512-key group), blocks of 1 to 4 streams over 1 to 3,000 context keys (1, 7, 16,
  31 to 33, 63 to 65, and 2,032 to 2,081 around the 2,048 window plus 16), causal and not. 3,548 runs, 851,591,168
  output bytes, 0 differ. A one-ulp change to the tree's score scale differs in 708,973 and 1,032,086 bytes (dims 128
  and 256), so the comparison sees it.
  • Why the tree's bits cannot move: a tile's mask is per key column, the same for both rows a thread holds, so on any
    input where a row's visible scores are numbers both rows' flags agree and every operation is the old one.
  • The block kernel's first pass is the old kernel; a finite output skips the second.
  • The float64-weight replay (ff6e5c0e) against the reviewed head (ee81000d): both block forms over the served
    shapes with 12 seeds in both finite fills, and 1 to 4 streams over 1, 7, 16, 31 to 33, 63 to 65, 700, 2,032 to
    2,081, 2,100 and 3,000 context keys with 2 seeds: 2,080 runs, 489,881,600 output bytes, 0 differ. A one-ulp change
    to the kernel's log2(e) differs in 245,828 bytes. Block attention 86.0 to 87.6 µs (one stream, 2,100 keys), 34.7
    to 35.4 µs (four streams, 700), causal 84.0 to 84.5 and 33.9 to 33.6 µs (best of 5 x 100 launches with sync).
  • The tree replay and the block tile's keep (e0726dc1, 942e2fe6; b502747c changes comments only, the same
    SASS on all five targets) against ff6e5c0e: the six tree forms over
    the served shapes with 6 seeds, and 1, 5, 16, 17, 33 and 64 rows over 1, 15 to 17, 31 to 33, 511 to 513, 1,023 to
    1,025, 1,536, 2,100 and 3,000 committed keys; both block forms over the served shapes with 6 seeds, and 1 to 4
    streams over 1, 7, 16, 31 to 33, 63 to 65, 700, 2,047 to 2,050, 2,079 to 2,081, 2,100 and 3,000 context keys; both
    finite fills: 1,632 runs, 461,963,264 output bytes, 0 differ. A one-ulp change to the tree's log2(e) differs in
    74,217 bytes. Speed (best of 7 x 100 launches with sync, 3 alternating runs, GB10): block attention -0.2 to
    +0.9 %; the twelve tree cells 0.0 to 15.6 % faster (dim 128 one-warp, 64 rows: 129.8 to 110.2 µs), the merge
    now at 128 registers instead of 148.
  • The reference's constructor (contract_attn.zig, Ref.init) frees what it took when a later allocation fails:
    a host test fails each allocation in turn (std.testing.checkAllAllocationFailures); on ee81000d it reports
    a leak, now it passes (in zig build test-cuda-host).
  • tf-cuda-test tree and block give the same numbers as before (worst 0.470 and 0.380).
  • Speed of the fixes, GB10, 200 launches: block attention 88.1 to 83.9 µs (one stream, 2,100 keys, either form),
    34.8 to 32.8 µs (four streams, 700); the six tree forms 0.3 to 2.4 % slower at the served shapes (launch included; runs
    spread about 1 %).
• Mutations: each of these fails its test: paths cut at depth 3, the score scale off by 1/16, the window bound
  off by one, the causal mask off by one, the merge leaving out a group's fourth slice (19x the tolerance, and the chunked merge's fourth slice), a tile fetched into the stage being read.
• In a model (fork only, not in this PR): with Qwen3.8-27B on CUDA: the qwen27 family's engine for NVIDIA GPUs, every stream drafted #608's 27B engine wired to these kernels instead of its Triton
  attention, drafted long prompt: attention 46.1 ms against Triton's 49.0 ms; 16 streams at 32k context, code and
  chat at temperature 0 and 1: every bench_concurrent --alone --serial cell exact, throughput within -4..+2 % of
  Triton. Replies equal 0.6.6's except two at bf16 near-ties (an exact top-2 tie, and one ulp), where drafted still
  equals serial. Teacher-forced on 1,023 rows of real text: NLL 0.8058 vs 0.8029, argmax agreement 99.12 %.
  The slices, one-warp blocks and stages (de31ba9b, 6c3d49d5, 197cacc7) change no bits: 4,092 teacher-forced
  logit rows byte-identical to the build before, drafted and 8-stream shared rounds with the same shas and round
  counts. 512-key groups (403549d1) give the 27B new bits once, on every device: drafted == serial on all four
  prompts in both block forms, and against the Triton set on 4,092 teacher-forced rows dNLL +0.0009, argmax
  agreement 99.32 %, mean KL 0.00066 (2,048-key groups: +0.0011, 99.12 %, 0.00071). Serial decode after a 4,722-token
  prompt, 128 tokens, GB10: attention 275.8 ms with the first version, 209.1 ms now (one-warp blocks; four-warp
  217.5 ms; Triton 244.3 ms). On an RTX 5090 (Qwen3.8-27B on CUDA: the qwen27 family's engine for NVIDIA GPUs, every stream drafted #608): attention 34.6 µs a call against Triton's 59.7, long 2.92 vs
  2.94 s. The merge issuing a chunk's loads first (591065e0) changes no bits (the same 4,092 logit rows, shas and
  round counts) and takes the merge from 10.0 to 4.6 µs a call on GB10, attention 199.2 ms in all.
• Speed in this PR: none claimed, nothing calls it.
• Tests run: zig build test from an empty cache on b502747c (86/86 steps, 1,378 of 1,385 tests, 7 skipped) and zig build test-cuda-host (41/41 steps), all passing, with the fixes and the registration. tf-cuda-test tree and block
  pass with the numbers above. zig fmt --check clean. lean_check.py: 517 findings, as on main.
• Not run, and why: Metal (unchanged); no model on main uses it.

Registration lines in build/cuda.zig, cuda/kernels.zig, cuda/root.zig and tests/cuda/main.zig sit where
#734, #736 and #740 add theirs; whichever lands second takes both.

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.
• 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 591065e0. 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 ee81000d. Gate 1 found two real NaN bugs here, and both are fixed (90aebfcb, fc48c41c). On finite inputs the outputs are byte-identical to 591065e0. The CUDA SUMMARY line, the identity sweep and the mutation checks are in the description.

[ashhart]

Thanks @jschmied for the registrations and both NaN repairs. Two source blockers remain.

The replay adds a visible infinite value back without checking that key's weight. When its softmax weight underflows to zero, the declared FP64 reference gives NaN but the helper returns infinity. A host control reproduces this through the actual helper; it is not a GPU run. The shared reference constructor also leaks earlier arrays when a later allocation fails.

Please repair both with failing controls, keep the fifteen launch forms pending, then rerun finite-output identity against the reviewed head and all numerical/resource gates. The current receipts do not close those cases.

Posted by Ash's AI Assistant Synthena

[jschmied]

Both are fixed, in ff6e5c0e and f4de6c3b.

• Replay weight: attn_odd::add_back now recomputes the key's score and the row's max in float64, in attention_ref's order. It adds an infinite value back only where exp(score - max) is above 0, and at weight 0 it adds 0 x inf = NaN. Fixing this showed the reverse case too: an infinite value in a tile the replay did not clean went through the mma, where a weight that underflows in float32 but not in float64 gave NaN and float64 gives inf. So the replay now takes every NaN and infinite value out of the mma, not only the edge tiles' values.

• Controls (gate 1 plants in contract_block_attn.zig), red on ee81000d:

  • A weight-0 key in the window's edge tile, both forms: "head 20 element 17: pos_inf where float64 gives nan".
  • The block's first key alone, causal: "head 24 element 19: pos_inf where float64 gives nan".
  • Below-float32 weights over every key of a key head, alone, both forms: "head 28 element 21: nan where float64 gives pos_inf".

  All of them pass on ff6e5c0e.

• Constructor: Ref.init now frees the arrays it already took when a later allocation fails (errdefer). A host test uses std.testing.checkAllAllocationFailures: it reports 1 leak on ee81000d and passes now, in zig build test-cuda-host.

• Reruns on ff6e5c0e, GB10:

  • Kernel contract: SUMMARY 21 kernels: 21 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 0 limit problems.
  • Limits: kernel-contract-limits passes for sm_80/86/89/120/121. tf_block_attention is re-recorded at 230/228 registers, still one block per SM.
  • Finite-output identity against ee81000d: 2,080 runs, 489,881,600 bytes, 0 differ. A one-ulp change to log2(e) changes 245,828 bytes, so the comparison does see changes.
  • Tests: tf-cuda-test block (0.380) and tree (0.470) pass, and zig build test from an empty cache passes 86/86 steps.
  • Speed: block attention changes by -1 to +2.6 %.

• The fifteen launch forms stay pending.

Not changed: tree attention has no replay. A visible infinite value whose weight underflows in float32 but not in float64 still gives NaN there, where float64 gives inf. The same holds for #798's decode kernels. Gate 1 does not plant that case for them.

[jschmied]

The open case is closed: tree attention now has the replay too, in 942e2fe6, and the new head is b502747c.

• Tree attention: a block whose partial holds a NaN or an infinity runs again, out of line. V's NaN and infinite values leave the mma, and each (row, head) adds them back after its last tile at the key's float64 weight, against its row's float64 max over all its keys. The merges fold as before, and only a NaN result folds again with those infinities unscaled. tf_tree_shared takes kn as a new last parameter, for the path keys' weights.

• Block attention (e0726dc1): the replay's block tile rescales after the context's values are added back. With a block key holding the max, that factor is 0 in float32 and turned an infinity into NaN, so the tile now leaves an infinite value as it is.

• Controls (gate 1 plants), red on ff6e5c0e:

  • Tree, key head 1, row 0: "head 12 element 21: nan where float64 gives pos_inf" in all six forms (dim 256: head 6).
  • Block, stream 1, key head 7, the max in the block: "head 540 element 23: nan where float64 gives pos_inf" (causal: head 252).

  All of them pass now. The weight-0 heads (query all 8) pass on both heads, and they turn red when the float64 weight test is forced to true ("head 13 element 21: pos_inf where float64 gives nan").

• Reruns on the new head, GB10 (b502747c only joins comments: the same SASS as 942e2fe6 on all five targets, where the identity and speed runs were):

  • Kernel contract: SUMMARY 21 kernels: 21 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 0 limit problems.
  • Limits: kernel-contract-limits passes for all five targets. The tree launches are re-recorded, with no local memory on any target.
  • Finite-output identity against ff6e5c0e: 1,632 runs, 461,963,264 bytes, 0 differ. A one-ulp change to log2(e) changes 74,217 bytes.
  • Tests: tf-cuda-test tree (0.470) and block (0.380) pass, and zig build test from an empty cache passes 86/86 steps.
  • Speed: block attention changes by -0.2 to +0.9 %, and tree attention is 0 to 15.6 % faster (the merge now uses 128 registers instead of 148).

• The fifteen launch forms stay pending.

[ashhart]

Thank you, @jschmied, for CUDA tree attention, DFlash block attention and the float64 checks.

We merged the reviewed head 591065e09 for 1.1.0 in 307d686, retaining both sets of reference and runner imports. The later NaN-repair increment is separate work for 1.1.1, the next patch release.

That increment still adds infinity where a visible key's zero weight makes the reference NaN, and its new reference constructor leaks after an allocation failure. Please repair both with failing controls before we carry the increment; the affected helper and constructor are absent from the 1.1.0 candidate.

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

[jschmied]

Both are already repaired on this branch, with failing controls, for the 1.1.1 increment. The head is b502747c.

• Zero-weight infinity: ff6e5c0e re-adds a visible infinite value only where its float64 weight is above 0, so a weight of 0 gives NaN as the reference does. 942e2fe6 and e0726dc1 extend this to tree attention and block attention's replay tile.
• Constructor leak: f4de6c3b adds an errdefer per array, tested with checkAllAllocationFailures.

Controls, finite-output identity and the gate results are in #746 (comment) and #746 (comment).

[ashhart]

Thanks @jschmied for the follow-up and failing controls. The constructor cleanup, weighted non-finite replay and tree replay address the earlier findings in source. This remains open for 1.1.2 integration: Kolibri now uses these operators, so we need the merged launch-argument/resource review and Kolibri drafted/plain, cached/fresh and concurrent/solo regressions. Please attach the raw final-head numerical, mutation, finite-byte and compiled-resource receipts; the description's GB10 results remain reported evidence. We will carry the shared changes once with #798, retaining the existing registrations and bounds.

Posted by Ash's AI Assistant Synthena

[jschmied]

A correction to my last comment first. I wrote that the tree launches have "no local memory on any target". I read that from cuobjdump's LOCAL field, which is 0. ptxas -v shows small register spills, into the stack frame, in eight tree attention entry points on some targets. Block, window, moe and dense have no spills. The contract's resource gate (gate 5) records registers, shared memory and bounds, not spills, so it passes either way.

Spills from ptxas -v (nvcc 13.0, -O3, the build's flags), as registers and bytes stored/loaded. Main's register count is in parentheses, and main spills nowhere here except tf_tree_shared_w1/w1s3 (4/4 B on sm_120/121). The SASS is identical at both PR heads and in the merge.

kernel	sm_80	sm_86	sm_89	sm_120	sm_121
tf_tree_merge	114 regs (97), none	114 regs (97), none	114 regs (97), none	128 regs (148), 24/32 B	128 regs (148), 24/32 B
tf_tree_part_256_w1s3	255 regs (218), 16/32 B	255 regs (218), 16/32 B	255 regs (218), 16/32 B	251 regs (231), none	251 regs (231), none
tf_tree_part_256_w1	255 regs (208), 16/24 B	255 regs (208), 16/24 B	255 regs (208), 16/24 B	253 regs (233), none	253 regs (233), none
tf_tree_part_w1	197 regs (168), none	197 regs (168), none	197 regs (168), none	168 regs (168), 56/52 B	168 regs (168), 56/52 B
tf_tree_part_256	255 regs (229), 16/12 B	255 regs (229), 16/12 B	255 regs (229), 16/12 B	255 regs (231), 44/40 B	255 regs (231), 44/40 B
tf_tree_shared_256_w1s3	255 regs (220), 12/32 B	255 regs (220), 4/0 B	255 regs (220), 4/0 B	250 regs (236), none	250 regs (236), none
tf_tree_shared_256_w1	255 regs (210), 12/24 B	255 regs (210), 4/0 B	255 regs (210), 4/0 B	249 regs (237), none	249 regs (237), none
tf_tree_shared_256	255 regs (230), 36/76 B	255 regs (230), 36/76 B	255 regs (230), 36/76 B	255 regs (243), 48/48 B	255 regs (243), 48/48 B

On the GB10 (sm_121), the tf_tree_merge spill is one 8-byte store at the kernel's start and one load after the main loop, per thread: outside the loops. Earlier timing showed tree attention 0 to 15.6 % faster than the head before the tree replay, and the merge times the same as the heads (below), so I see no slowdown. If you want the dim-256 and merge kernels back under the register cap before 1.1.2, we can do that.

A test fix only the merge needs (PR branches unchanged): main's test "the six served MoE and MTP launches are in the live registry" looks records up by name alone. tf_moe_combine now exists in both moe (main's Qwen engine) and moe_glue (#798), so it finds two and fails. With name and image matched (045985b8 in my scratch merge), zig build test passes 154/154 steps and test-cuda-host 82/82:

--- a/zig/src/core/contract/qwen_served_records.zig
+++ b/zig/src/core/contract/qwen_served_records.zig
@@ -20,7 +20,7 @@
     for (records) |record| {
         var found: usize = 0;
-        for (registry.all_cuda) |active| if (std.mem.eql(u8, active.launch.symbol, record.launch.symbol)) {
+        for (registry.all_cuda) |active| if (std.mem.eql(u8, active.launch.symbol, record.launch.symbol) and std.mem.eql(u8, active.launch.image, record.launch.image)) {
             try std.testing.expectEqualDeep(record, active);
             found += 1;
         };

Merged build: main d46d17a (1.1.0) + #746 b502747c + #798 74b749fa, merged locally (not pushed; branches untouched). For Kolibri, #991 d2f62303 on top.

Conflicts and how I resolved them:

• cuda: tree attention and DFlash block attention for every family, checked against float64 #746 into main
  • zig/build/cuda.zig, hostTests: main's longer module list kept; cuda: tree attention and DFlash block attention for every family, checked against float64 #746's contract_attn.zig host test added after it.
  • zig/src/core/contract/registry.zig: cuda: tree attention and DFlash block attention for every family, checked against float64 #746's 15 tree/block records go into main's cuda array, before the ++ registry_media/nemotron_prompt/glm imports. Main's 13.2 and 13.3 sets are unchanged.
  • zig/tests/cuda/contract_tests.zig: imports, the image list and the run body are a union (main's kernels first, then tree/block).
• cuda: sliding-window attention, sigmoid top-k routing and combine, dense rows for every family, checked against float64 #798 on top
  • zig/kernels/cuda/attn_odd.cuh (add/add): cuda: sliding-window attention, sigmoid top-k routing and combine, dense rows for every family, checked against float64 #798's file is cuda: tree attention and DFlash block attention for every family, checked against float64 #746's plus the run-time d overload of score64, so cuda: sliding-window attention, sigmoid top-k routing and combine, dense rows for every family, checked against float64 #798's is taken.
  • cuda.zig, registry.zig, contract_tests.zig: as above, plus cuda: sliding-window attention, sigmoid top-k routing and combine, dense rows for every family, checked against float64 #798's 12 window/moe/dense records and their kernel setup.
  • Every kernel, src/cuda and contract test file is byte-identical to 74b749fa. The registry has every cuda: sliding-window attention, sigmoid top-k routing and combine, dense rows for every family, checked against float64 #798 record line. The one exception is vit full_kernel<72>, where main's newer record wins.
• A semantic conflict that git did not flag: main's test "the six served MoE and MTP launches are in the live registry" (qwen_served_records.zig) looks up records by symbol alone. tf_moe_combine is now in two images, moe (main's Qwen engine) and moe_glue (cuda: sliding-window attention, sigmoid top-k routing and combine, dense rows for every family, checked against float64 #798), so the test found 2 and failed. In the merge I made the lookup match symbol and image. With that, zig build test 154/154 steps and test-cuda-host 82/82 pass. This fix has to go into the merge commit.
• kolibri1: Kolibri-1 on CUDA through the lane host, drafts off (#481) #991 on top (Kolibri only)
  • native/cuda.zig: the family list is main's plus kolibri1.native. Main's kv_pages paging stays. kolibri1: Kolibri-1 on CUDA through the lane host, drafts off (#481) #991's growRoom feeds the unpaged admit and the TooMany count, with served.stream_bytes as on main.
  • cuda_refusal.zig: main's vision, tp and tp-draft refusals stay, and the drafter refusal takes kolibri1: Kolibri-1 on CUDA through the lane host, drafts off (#481) #991's drafter_refusal.
  • root.zig: main's exports. tests/cuda/main.zig: plus the two kolibri suites. build/cuda.zig: plus the kolibri1 import and its host test. registry.zig: plus cuda: SwiGLU, output gates, Gated DeltaNet inputs, QK norm and RoPE for every family, checked against float64 #797's 15 act/gdn/qk records. contract_tests.zig: plus their setup. docs/recipes/README.md: both rows.

Results, GB10 (all raw outputs below):

• Kernel contract, ReleaseFast: merge 80/80 pass, 5/5 controls red; merge with kolibri1: Kolibri-1 on CUDA through the lane host, drafts off (#481) #991 103/103. At the heads, cuda: tree attention and DFlash block attention for every family, checked against float64 #746 21/21 and cuda: sliding-window attention, sigmoid top-k routing and combine, dense rows for every family, checked against float64 #798 37/37. Our kernels' verdicts are the same as at the heads; only the random-input ratios move, because each kernel's seed depends on its place in the list. Debug and ReleaseSafe still panic at contract_ck_lane.zig:108 on main (a read past the end that is already there). I left that alone.
• Limits: kernel-contract-limits passes on sm_80, 86, 89, 120 and 121, every record current. Only nvcc 13.0 is on this machine, so these launches have no 13.2 records yet.
• SASS of tree, block, window, moe and dense is identical to the heads on all five targets.
• Finite bytes: the merge equals cuda: sliding-window attention, sigmoid top-k routing and combine, dense rows for every family, checked against float64 #798's head (3,224 runs, 795,961,344 bytes, 0 differ) and cuda: tree attention and DFlash block attention for every family, checked against float64 #746's for tree and block (1,632 runs, 461,963,264 bytes, 0 differ). NaN heads: 2,080 runs, 0 bad.
• Mutations, each red: float64 weight test forced true (11 kernels and 8 limit records), a moe combine tail slot dropped (3), the dense rows lane sum halved (4), log2(e) one ulp up (170,756 bytes differ).
• Timing: merge vs heads within -0.8 to +2.8 %. One 11 us decode launch reads -7.1 %, but the merge's own three rounds spread 10.4 to 11.4 us.

Kolibri-1 on the GB10, --context 32768 --parallel 8, #991's head against the merge with #991:

Kolibri-1, GB10, --context 32768 --parallel 8	#991 head d2f62303	merged build (main + #746 + #798 + #991)
engine: kolibri-reference, kolibri-layers, kolibri-full	pass	pass, the same result lines
drafted ("draft": true) vs plain, 7 prompts × 160 tokens	7/7 equal	7/7 equal, the same 7 SHAs as d2f62303
cached/resumed vs fresh (reuse_served.py, cache 8 GiB vs 0)	35/35 equal	35/35 equal
kept states at 2048/3001/4096 vs fresh (kolibri-full)	equal	equal
concurrent (4) vs alone, t=0 and t=1.0	12/12 equal per cell	12/12 equal per cell
2,048-token prompt, first request TTFT	1.007 s	1.005 s
8,192-token prompt, first request TTFT	3.755 s	3.760 s
decode, 1 stream (code / chat, t=0)	52.3 / 50.6 tok/s	52.0 / 50.3 tok/s
decode, 4 streams aggregate (code / chat, t=0)	124.8 / 126.6 tok/s	124.6 / 126.3 tok/s
load line	model 85.23 GiB, 8 streams × 1.29 GiB	the same

The cache benchmark (bench_cache.py) reports 0 cached tokens for repeated raw-token prompts on both builds, so its --require-reuse rows are not qualified (outputs equal). Resumed vs fresh is covered by the reuse_served.py rows in the table.

kernel-contract SUMMARY lines

merge (main + #746 + #798):
SUMMARY 80 kernels: 80 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 0 limit problems; gates 318538 ms, limits 7089 ms
merge + #991:
SUMMARY 103 kernels: 103 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 0 limit problems; gates 317467 ms, limits 7378 ms
#746 head b502747c:
SUMMARY 21 kernels: 21 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 0 limit problems; gates 86132 ms, limits 2 ms
#798 head 74b749fa:
SUMMARY 37 kernels: 37 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 0 limit problems; gates 96612 ms, limits 2 ms

kernel-contract, merge, this PR's kernels

PASS nonfinite cuda tree attention tf_tree_paths: the stated rule holds: integers only: each row's depth (at most 128) and its nearest 128 ancestors root first; other slots untouched
PASS repeat cuda tree attention tf_tree_paths: five runs byte-identical (8256 bytes), inputs and guards untouched, no overlap
PASS shapes cuda tree attention tf_tree_paths: 6 shapes exact by its own check (worst 0.000 of the bound), 0 edges refused by the kernel
N/A  sums cuda tree attention tf_tree_paths: declares none sums
PASS nonfinite cuda tree attention dim 128, four-warp blocks (tf_tree_shared, tf_tree_part, tf_tree_merge): the stated rule holds: outputs equal the float64 reference's bytes
PASS repeat cuda tree attention dim 128, four-warp blocks (tf_tree_shared, tf_tree_part, tf_tree_merge): five runs byte-identical (196608 bytes), inputs and guards untouched, no overlap
PASS shapes cuda tree attention dim 128, four-warp blocks (tf_tree_shared, tf_tree_part, tf_tree_merge): 10 shapes exact by its own check (worst 0.000 of the bound), 0 edges refused by the kernel
PASS sums cuda tree attention dim 128, four-warp blocks (tf_tree_shared, tf_tree_part, tf_tree_merge): exact sums at k 2100, where fp16 sums would miss by infx and bf16 by 1839101x
PASS nonfinite cuda tree attention dim 128, one-warp blocks (tf_tree_shared_w1, tf_tree_part_w1, sliced merge): the stated rule holds: outputs equal the float64 reference's bytes
PASS repeat cuda tree attention dim 128, one-warp blocks (tf_tree_shared_w1, tf_tree_part_w1, sliced merge): five runs byte-identical (196608 bytes), inputs and guards untouched, no overlap
PASS shapes cuda tree attention dim 128, one-warp blocks (tf_tree_shared_w1, tf_tree_part_w1, sliced merge): 9 shapes exact by its own check (worst 0.000 of the bound), 1 edges refused by the kernel
PASS sums cuda tree attention dim 128, one-warp blocks (tf_tree_shared_w1, tf_tree_part_w1, sliced merge): exact sums at k 2100, where fp16 sums would miss by infx and bf16 by 1642027x
PASS nonfinite cuda tree attention dim 128, one-warp blocks, 3 stages (tf_tree_shared_w1s3, tf_tree_part_w1s3): the stated rule holds: outputs equal the float64 reference's bytes
PASS repeat cuda tree attention dim 128, one-warp blocks, 3 stages (tf_tree_shared_w1s3, tf_tree_part_w1s3): five runs byte-identical (196608 bytes), inputs and guards untouched, no overlap
PASS shapes cuda tree attention dim 128, one-warp blocks, 3 stages (tf_tree_shared_w1s3, tf_tree_part_w1s3): 9 shapes exact by its own check (worst 0.000 of the bound), 1 edges refused by the kernel
PASS sums cuda tree attention dim 128, one-warp blocks, 3 stages (tf_tree_shared_w1s3, tf_tree_part_w1s3): exact sums at k 2100, where fp16 sums would miss by infx and bf16 by 1544002x
PASS nonfinite cuda tree attention dim 256, four-warp blocks (tf_tree_shared_256, tf_tree_part_256, tf_tree_merge): the stated rule holds: outputs equal the float64 reference's bytes
PASS repeat cuda tree attention dim 256, four-warp blocks (tf_tree_shared_256, tf_tree_part_256, tf_tree_merge): five runs byte-identical (196608 bytes), inputs and guards untouched, no overlap
PASS shapes cuda tree attention dim 256, four-warp blocks (tf_tree_shared_256, tf_tree_part_256, tf_tree_merge): 10 shapes exact by its own check (worst 0.000 of the bound), 0 edges refused by the kernel
PASS sums cuda tree attention dim 256, four-warp blocks (tf_tree_shared_256, tf_tree_part_256, tf_tree_merge): exact sums at k 2100, where fp16 sums would miss by infx and bf16 by 1708231x
PASS nonfinite cuda tree attention dim 256, one-warp blocks (tf_tree_shared_256_w1, tf_tree_part_256_w1): the stated rule holds: outputs equal the float64 reference's bytes
PASS repeat cuda tree attention dim 256, one-warp blocks (tf_tree_shared_256_w1, tf_tree_part_256_w1): five runs byte-identical (196608 bytes), inputs and guards untouched, no overlap
PASS shapes cuda tree attention dim 256, one-warp blocks (tf_tree_shared_256_w1, tf_tree_part_256_w1): 9 shapes exact by its own check (worst 0.000 of the bound), 1 edges refused by the kernel
PASS sums cuda tree attention dim 256, one-warp blocks (tf_tree_shared_256_w1, tf_tree_part_256_w1): exact sums at k 2100, where fp16 sums would miss by infx and bf16 by 1806768x
PASS nonfinite cuda tree attention dim 256, one-warp blocks, 3 stages (tf_tree_shared_256_w1s3): the stated rule holds: outputs equal the float64 reference's bytes
PASS repeat cuda tree attention dim 256, one-warp blocks, 3 stages (tf_tree_shared_256_w1s3): five runs byte-identical (196608 bytes), inputs and guards untouched, no overlap
PASS shapes cuda tree attention dim 256, one-warp blocks, 3 stages (tf_tree_shared_256_w1s3): 9 shapes exact by its own check (worst 0.000 of the bound), 1 edges refused by the kernel
PASS sums cuda tree attention dim 256, one-warp blocks, 3 stages (tf_tree_shared_256_w1s3): exact sums at k 2100, where fp16 sums would miss by infx and bf16 by 1741076x
PASS nonfinite cuda block attention tf_block_attention (DFlash2 blocks of 16): the stated rule holds: outputs equal the float64 reference's bytes
PASS repeat cuda block attention tf_block_attention (DFlash2 blocks of 16): five runs byte-identical (131072 bytes), inputs and guards untouched, no overlap
PASS shapes cuda block attention tf_block_attention (DFlash2 blocks of 16): 6 shapes exact by its own check (worst 0.000 of the bound), 0 edges refused by the kernel
PASS sums cuda block attention tf_block_attention (DFlash2 blocks of 16): exact sums at k 2100, where fp16 sums would miss by infx and bf16 by 1644519x
PASS nonfinite cuda block attention tf_block_attention, causal blocks of 7: the stated rule holds: outputs equal the float64 reference's bytes
PASS repeat cuda block attention tf_block_attention, causal blocks of 7: five runs byte-identical (57344 bytes), inputs and guards untouched, no overlap
PASS shapes cuda block attention tf_block_attention, causal blocks of 7: 6 shapes exact by its own check (worst 0.000 of the bound), 0 edges refused by the kernel
PASS sums cuda block attention tf_block_attention, causal blocks of 7: exact sums at k 2100, where fp16 sums would miss by infx and bf16 by 1868403x
PASS limits live tree attention tf_tree_paths: 128 threads launched, the driver allows 1024; 30 registers, 0 + 0 shared bytes (recorded 30, 0)
PASS limits live tree attention tf_tree_shared: 128 threads launched, the driver allows 128; 212 registers, 0 + 65536 shared bytes (recorded 212, 0)
PASS limits live tree attention tf_tree_shared_w1: 32 threads launched, the driver allows 32; 195 registers, 0 + 32768 shared bytes (recorded 195, 0)
PASS limits live tree attention tf_tree_shared_w1s3: 32 threads launched, the driver allows 32; 198 registers, 0 + 49152 shared bytes (recorded 198, 0)
PASS limits live tree attention tf_tree_part: 128 threads launched, the driver allows 128; 228 registers, 0 + 65536 shared bytes (recorded 228, 0)
PASS limits live tree attention tf_tree_part_w1: 32 threads launched, the driver allows 32; 168 registers, 0 + 32768 shared bytes (recorded 168, 0)
PASS limits live tree attention tf_tree_part_w1s3: 32 threads launched, the driver allows 32; 202 registers, 0 + 49152 shared bytes (recorded 202, 0)
PASS limits live tree attention tf_tree_shared_256: 128 threads launched, the driver allows 128; 255 registers, 0 + 73728 shared bytes (recorded 255, 0)
PASS limits live tree attention tf_tree_shared_256_w1: 32 threads launched, the driver allows 32; 249 registers, 0 + 40960 shared bytes (recorded 249, 0)
PASS limits live tree attention tf_tree_shared_256_w1s3: 32 threads launched, the driver allows 32; 250 registers, 0 + 57344 shared bytes (recorded 250, 0)
PASS limits live tree attention tf_tree_part_256: 128 threads launched, the driver allows 128; 255 registers, 0 + 73728 shared bytes (recorded 255, 0)
PASS limits live tree attention tf_tree_part_256_w1: 32 threads launched, the driver allows 32; 253 registers, 0 + 40960 shared bytes (recorded 253, 0)
PASS limits live tree attention tf_tree_part_256_w1s3: 32 threads launched, the driver allows 32; 251 registers, 0 + 57344 shared bytes (recorded 251, 0)
PASS limits live tree attention tf_tree_merge: 256 threads launched, the driver allows 256; 128 registers, 0 + 0 shared bytes (recorded 128, 0)
PASS limits live block attention tf_block_attention (up to 8 query heads a key head): 256 threads launched, the driver allows 256; 228 registers, 0 + 16384 shared bytes (recorded 228, 0)

kernel-contract-limits, merge, this PR's records (sm_80, 86, 89, 120, 121)

RECORD tree attention tf_tree_paths: .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 30, .static = 0, .bound = 0 }, .{ .regs = 30, .static = 0, .bound = 0 }, }
RECORD tree attention tf_tree_shared: .usage = .{ .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, }
RECORD tree attention tf_tree_shared_w1: .usage = .{ .{ .regs = 193, .static = 0, .bound = 32 }, .{ .regs = 193, .static = 0, .bound = 32 }, .{ .regs = 193, .static = 0, .bound = 32 }, .{ .regs = 195, .static = 0, .bound = 32 }, .{ .regs = 195, .static = 0, .bound = 32 }, }
RECORD tree attention tf_tree_shared_w1s3: .usage = .{ .{ .regs = 190, .static = 0, .bound = 32 }, .{ .regs = 190, .static = 0, .bound = 32 }, .{ .regs = 190, .static = 0, .bound = 32 }, .{ .regs = 198, .static = 0, .bound = 32 }, .{ .regs = 198, .static = 0, .bound = 32 }, }
RECORD tree attention tf_tree_part: .usage = .{ .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 202, .static = 0, .bound = 128 }, .{ .regs = 228, .static = 0, .bound = 128 }, .{ .regs = 228, .static = 0, .bound = 128 }, }
RECORD tree attention tf_tree_part_w1: .usage = .{ .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, }
RECORD tree attention tf_tree_part_w1s3: .usage = .{ .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 197, .static = 0, .bound = 32 }, .{ .regs = 202, .static = 0, .bound = 32 }, .{ .regs = 202, .static = 0, .bound = 32 }, }
RECORD tree attention tf_tree_shared_256: .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, }
RECORD tree attention tf_tree_shared_256_w1: .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 249, .static = 0, .bound = 32 }, .{ .regs = 249, .static = 0, .bound = 32 }, }
RECORD tree attention tf_tree_shared_256_w1s3: .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 250, .static = 0, .bound = 32 }, .{ .regs = 250, .static = 0, .bound = 32 }, }
RECORD tree attention tf_tree_part_256: .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, }
RECORD tree attention tf_tree_part_256_w1: .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 253, .static = 0, .bound = 32 }, .{ .regs = 253, .static = 0, .bound = 32 }, }
RECORD tree attention tf_tree_part_256_w1s3: .usage = .{ .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 255, .static = 0, .bound = 32 }, .{ .regs = 251, .static = 0, .bound = 32 }, .{ .regs = 251, .static = 0, .bound = 32 }, }
RECORD tree attention tf_tree_merge: .usage = .{ .{ .regs = 114, .static = 0, .bound = 256 }, .{ .regs = 114, .static = 0, .bound = 256 }, .{ .regs = 114, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, }
RECORD block attention tf_block_attention (up to 8 query heads a key head): .usage = .{ .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 230, .static = 0, .bound = 256 }, .{ .regs = 228, .static = 0, .bound = 256 }, .{ .regs = 228, .static = 0, .bound = 256 }, }
INFO limits fatbins: images built by nvcc 13.0, checked against nvcc 13.0's records
PASS kernel-contract-limits: every registered launch fits sm_80, sm_86, sm_89, sm_120 and sm_121, every control red

ptxas -v per entry point, merge: regs/stack/spill store/spill load (bytes)

function: sm_80 | sm_86 | sm_89 | sm_120 | sm_121, each regs/stack/spill-store/spill-load (bytes)
tf_tree_merge: 114/0/0/0 | 114/0/0/0 | 114/0/0/0 | 128/24/24/32 | 128/24/24/32
tf_tree_part_256_w1s3: 255/64/16/32 | 255/48/16/32 | 255/48/16/32 | 251/0/0/0 | 251/0/0/0
tf_tree_part_w1s3: 197/0/0/0 | 197/0/0/0 | 197/0/0/0 | 202/0/0/0 | 202/0/0/0
tf_tree_part_256_w1: 255/72/16/24 | 255/48/16/24 | 255/48/16/24 | 253/0/0/0 | 253/0/0/0
tf_tree_part_w1: 197/0/0/0 | 197/0/0/0 | 197/0/0/0 | 168/64/56/52 | 168/64/56/52
tf_tree_part_256: 255/80/16/12 | 255/80/16/12 | 255/80/16/12 | 255/48/44/40 | 255/48/44/40
tf_tree_part: 202/0/0/0 | 202/0/0/0 | 202/0/0/0 | 228/0/0/0 | 228/0/0/0
tf_tree_shared_256_w1s3: 255/88/12/32 | 255/32/4/0 | 255/32/4/0 | 250/0/0/0 | 250/0/0/0
tf_tree_shared_w1s3: 190/0/0/0 | 190/0/0/0 | 190/0/0/0 | 198/0/0/0 | 198/0/0/0
tf_tree_shared_256_w1: 255/88/12/24 | 255/32/4/0 | 255/32/4/0 | 249/0/0/0 | 249/0/0/0
tf_tree_shared_w1: 193/0/0/0 | 193/0/0/0 | 193/0/0/0 | 195/0/0/0 | 195/0/0/0
tf_tree_shared_256: 255/128/36/76 | 255/128/36/76 | 255/128/36/76 | 255/48/48/48 | 255/48/48/48
tf_tree_shared: 202/0/0/0 | 202/0/0/0 | 202/0/0/0 | 212/0/0/0 | 212/0/0/0
tf_tree_paths: 32/0/0/0 | 40/0/0/0 | 40/0/0/0 | 30/0/0/0 | 30/0/0/0
tf_block_attention: 230/0/0/0 | 230/0/0/0 | 230/0/0/0 | 228/0/0/0 | 228/0/0/0

ptxas -v per entry point, main d46d17a

function: sm_80 | sm_86 | sm_89 | sm_120 | sm_121, each regs/stack/spill-store/spill-load (bytes)
tf_tree_merge: 97/0/0/0 | 97/0/0/0 | 97/0/0/0 | 148/0/0/0 | 148/0/0/0
tf_tree_part_256_w1s3: 218/0/0/0 | 218/0/0/0 | 218/0/0/0 | 231/0/0/0 | 231/0/0/0
tf_tree_part_w1s3: 168/0/0/0 | 168/0/0/0 | 168/0/0/0 | 167/0/0/0 | 167/0/0/0
tf_tree_part_256_w1: 208/0/0/0 | 208/0/0/0 | 208/0/0/0 | 233/0/0/0 | 233/0/0/0
tf_tree_part_w1: 168/0/0/0 | 168/0/0/0 | 168/0/0/0 | 168/0/0/0 | 168/0/0/0
tf_tree_part_256: 229/0/0/0 | 229/0/0/0 | 229/0/0/0 | 231/0/0/0 | 231/0/0/0
tf_tree_part: 168/0/0/0 | 168/0/0/0 | 168/0/0/0 | 162/0/0/0 | 162/0/0/0
tf_tree_shared_256_w1s3: 220/0/0/0 | 220/0/0/0 | 220/0/0/0 | 236/0/0/0 | 236/0/0/0
tf_tree_shared_w1s3: 168/0/0/0 | 168/0/0/0 | 168/0/0/0 | 168/8/4/4 | 168/8/4/4
tf_tree_shared_256_w1: 210/0/0/0 | 210/0/0/0 | 210/0/0/0 | 237/0/0/0 | 237/0/0/0
tf_tree_shared_w1: 168/0/0/0 | 168/0/0/0 | 168/0/0/0 | 168/8/4/4 | 168/8/4/4
tf_tree_shared_256: 230/0/0/0 | 230/0/0/0 | 230/0/0/0 | 243/0/0/0 | 243/0/0/0
tf_tree_shared: 168/0/0/0 | 168/0/0/0 | 168/0/0/0 | 167/0/0/0 | 167/0/0/0
tf_tree_paths: 32/0/0/0 | 40/0/0/0 | 40/0/0/0 | 30/0/0/0 | 30/0/0/0
tf_block_attention: 166/0/0/0 | 166/0/0/0 | 166/0/0/0 | 170/0/0/0 | 170/0/0/0

SASS, merge vs heads

tree_attention: sass 7f110c7348c15ebc vs 7f110c7348c15ebc IDENTICAL (5 images)
block_attention: sass d462152e749a07a2 vs d462152e749a07a2 IDENTICAL (5 images)
tree_attention: sass 7f110c7348c15ebc vs 7f110c7348c15ebc IDENTICAL (5 images)
block_attention: sass d462152e749a07a2 vs d462152e749a07a2 IDENTICAL (5 images)
window_attention: sass 61640a82459ce061 vs 61640a82459ce061 IDENTICAL (5 images)
moe_glue: sass af4323610f70b73f vs af4323610f70b73f IDENTICAL (5 images)
dense_rows: sass 9729df25627cd387 vs 9729df25627cd387 IDENTICAL (5 images)

numerical suites, merge

$ tf-cuda-test tree   (M)
PASS tree attention: dim 128 and 256, committed keys and paths within float64 tolerance (worst 0.470), a stream's bits alone, any group slots, four-warp or one-warp blocks
$ tf-cuda-test block   (M)
PASS block attention: windowed contexts then the block, within float64 tolerance (worst 0.380), a stream's bits alone

finite bytes: merge vs heads, and the log2(e) mutation

merge vs #746 head:
tree0: runs 216 bytes 65273856 differ 0 nonfinite elems 0
tree1: runs 212 bytes 62128128 differ 0 nonfinite elems 0
tree2: runs 212 bytes 62128128 differ 0 nonfinite elems 0
tree3: runs 216 bytes 65273856 differ 0 nonfinite elems 0
tree4: runs 212 bytes 62128128 differ 0 nonfinite elems 0
tree5: runs 212 bytes 62128128 differ 0 nonfinite elems 0
block0: runs 176 bytes 57671680 differ 0 nonfinite elems 0
block1: runs 176 bytes 25231360 differ 0 nonfinite elems 0
TOTAL runs 1632 bytes 461963264 differing 0 nonfinite 0

merge vs #798 head:
block0: runs 176 bytes 57671680 differ 0 nonfinite elems 0
block1: runs 176 bytes 25231360 differ 0 nonfinite elems 0
tree0: runs 216 bytes 65273856 differ 0 nonfinite elems 0
tree1: runs 212 bytes 62128128 differ 0 nonfinite elems 0
tree2: runs 212 bytes 62128128 differ 0 nonfinite elems 0
tree3: runs 216 bytes 65273856 differ 0 nonfinite elems 0
tree4: runs 212 bytes 62128128 differ 0 nonfinite elems 0
tree5: runs 212 bytes 62128128 differ 0 nonfinite elems 0
window0: runs 316 bytes 72376320 differ 0 nonfinite elems 0
window1: runs 200 bytes 14286848 differ 0 nonfinite elems 0
window2: runs 316 bytes 6031360 differ 0 nonfinite elems 0
window3: runs 380 bytes 183730176 differ 0 nonfinite elems 0
window4: runs 380 bytes 57573376 differ 0 nonfinite elems 0
TOTAL runs 3224 bytes 795961344 differing 0 nonfinite 0

merge vs merge with LOG2E one ulp up (tree, block, window):
block0: runs 176 bytes 57671680 differ 18604 nonfinite elems 0
block1: runs 176 bytes 25231360 differ 8096 nonfinite elems 0
tree0: runs 216 bytes 65273856 differ 14107 nonfinite elems 0
tree1: runs 212 bytes 62128128 differ 12618 nonfinite elems 0
tree2: runs 212 bytes 62128128 differ 12618 nonfinite elems 0
tree3: runs 216 bytes 65273856 differ 12618 nonfinite elems 0
tree4: runs 212 bytes 62128128 differ 11128 nonfinite elems 0
tree5: runs 212 bytes 62128128 differ 11128 nonfinite elems 0
window0: runs 316 bytes 72376320 differ 6534 nonfinite elems 0
window1: runs 200 bytes 14286848 differ 763 nonfinite elems 0
window2: runs 316 bytes 6031360 differ 0 nonfinite elems 0
window3: runs 380 bytes 183730176 differ 41728 nonfinite elems 0
window4: runs 380 bytes 57573376 differ 20814 nonfinite elems 0
TOTAL runs 3224 bytes 795961344 differing 170756 nonfinite 0

mutations, merge (controls lines omitted)

weight test forced true (attn_odd.cuh: `>= least() || true`):
FAIL nonfinite cuda tree attention dim 128, four-warp blocks (tf_tree_shared, tf_tree_part, tf_tree_merge): head 13 element 21: pos_inf where float64 gives nan
FAIL nonfinite cuda tree attention dim 128, one-warp blocks (tf_tree_shared_w1, tf_tree_part_w1, sliced merge): head 13 element 21: pos_inf where float64 gives nan
FAIL nonfinite cuda tree attention dim 128, one-warp blocks, 3 stages (tf_tree_shared_w1s3, tf_tree_part_w1s3): head 13 element 21: pos_inf where float64 gives nan
FAIL nonfinite cuda tree attention dim 256, four-warp blocks (tf_tree_shared_256, tf_tree_part_256, tf_tree_merge): head 7 element 21: pos_inf where float64 gives nan
FAIL nonfinite cuda tree attention dim 256, one-warp blocks (tf_tree_shared_256_w1, tf_tree_part_256_w1): head 7 element 21: pos_inf where float64 gives nan
FAIL nonfinite cuda tree attention dim 256, one-warp blocks, 3 stages (tf_tree_shared_256_w1s3): head 7 element 21: pos_inf where float64 gives nan
FAIL nonfinite cuda block attention tf_block_attention (DFlash2 blocks of 16): head 20 element 17: pos_inf where float64 gives nan
FAIL nonfinite cuda block attention tf_block_attention, causal blocks of 7: head 20 element 17: pos_inf where float64 gives nan
FAIL nonfinite cuda window decode tf_window_decode_tc, 48/4 heads, window 513, ring 640: head 361 element 21: pos_inf where float64 gives nan
FAIL nonfinite cuda window decode tf_window_decode_tc, 16/1 heads, no window, ring 300: head 65 element 7: pos_inf where float64 gives nan
FAIL nonfinite cuda window prompt wattn_kernel<128, 8, 4, 8>, 48/4 heads, window 513, ring 2624: head 252 element 17: pos_inf where float64 gives nan
FAIL limits live tree attention tf_tree_shared: 128 threads launched, the driver allows 128; 208 registers, 0 + 65536 shared bytes (recorded 212, 0)
FAIL limits live tree attention tf_tree_part: 128 threads launched, the driver allows 128; 224 registers, 0 + 65536 shared bytes (recorded 228, 0)
FAIL limits live tree attention tf_tree_part_w1: 32 threads launched, the driver allows 32; 191 registers, 0 + 32768 shared bytes (recorded 168, 0)
FAIL limits live tree attention tf_tree_part_w1s3: 32 threads launched, the driver allows 32; 216 registers, 0 + 49152 shared bytes (recorded 202, 0)
FAIL limits live tree attention tf_tree_shared_256_w1s3: 32 threads launched, the driver allows 32; 255 registers, 0 + 57344 shared bytes (recorded 250, 0)
FAIL limits live tree attention tf_tree_part_256_w1: 32 threads launched, the driver allows 32; 255 registers, 0 + 40960 shared bytes (recorded 253, 0)
FAIL limits live tree attention tf_tree_part_256_w1s3: 32 threads launched, the driver allows 32; 255 registers, 0 + 57344 shared bytes (recorded 251, 0)
FAIL limits live block attention tf_block_attention (up to 8 query heads a key head): 256 threads launched, the driver allows 256; 173 registers, 0 + 16384 shared bytes (recorded 228, 0)
FAIL kernel-contract: TestFailed

timing, best of 3 rounds x 7 x 100 calls, #798 head -> merge (us a call, launch + sync)

  +0.2%     92.45 ->    92.60  cuda tree attention dim 128, four-warp blocks (tf_tree_shared, tf_tree_part, tf_tree_merge) rows 16 k 2100
  +0.2%    129.62 ->   129.85  cuda tree attention dim 128, four-warp blocks (tf_tree_shared, tf_tree_part, tf_tree_merge) rows 64 k 700
  -0.8%     85.63 ->    84.97  cuda tree attention dim 128, one-warp blocks (tf_tree_shared_w1, tf_tree_part_w1, sliced merge) rows 16 k 2100
  -0.2%    109.78 ->   109.61  cuda tree attention dim 128, one-warp blocks (tf_tree_shared_w1, tf_tree_part_w1, sliced merge) rows 64 k 700
  +1.4%    101.99 ->   103.44  cuda tree attention dim 128, one-warp blocks, 3 stages (tf_tree_shared_w1s3, tf_tree_part_w1s3) rows 16 k 2100
  -0.3%    130.81 ->   130.48  cuda tree attention dim 128, one-warp blocks, 3 stages (tf_tree_shared_w1s3, tf_tree_part_w1s3) rows 64 k 700
  -0.5%     93.31 ->    92.85  cuda tree attention dim 256, four-warp blocks (tf_tree_shared_256, tf_tree_part_256, tf_tree_merge) rows 16 k 2100
  +0.0%    144.46 ->   144.48  cuda tree attention dim 256, four-warp blocks (tf_tree_shared_256, tf_tree_part_256, tf_tree_merge) rows 64 k 700
  +0.3%    100.95 ->   101.30  cuda tree attention dim 256, one-warp blocks (tf_tree_shared_256_w1, tf_tree_part_256_w1) rows 16 k 2100
  -0.1%    155.32 ->   155.19  cuda tree attention dim 256, one-warp blocks (tf_tree_shared_256_w1, tf_tree_part_256_w1) rows 64 k 700
  +0.7%    170.46 ->   171.73  cuda tree attention dim 256, one-warp blocks, 3 stages (tf_tree_shared_256_w1s3) rows 16 k 2100
  +0.3%    253.08 ->   253.94  cuda tree attention dim 256, one-warp blocks, 3 stages (tf_tree_shared_256_w1s3) rows 64 k 700
  +0.2%     87.63 ->    87.82  cuda block attention tf_block_attention (DFlash2 blocks of 16) rows 1 k 2100
  -0.8%     35.35 ->    35.05  cuda block attention tf_block_attention (DFlash2 blocks of 16) rows 4 k 700
  +0.2%     84.28 ->    84.47  cuda block attention tf_block_attention, causal blocks of 7 rows 1 k 2100
  +0.7%     33.35 ->    33.57  cuda block attention tf_block_attention, causal blocks of 7 rows 4 k 700

Kolibri raw comparison

reference.log: 7 result lines, identical; 1 PASS 0 FAIL
layers.log: 6 result lines, identical; 1 PASS 0 FAIL
full.log: 18 result lines, identical; 1 PASS 0 FAIL
K991: drafted==plain 7/7, shas ['34e0f8a1b2c0', 'c1657fc2f1b6', '3b3596014050', 'cc206f8c50ad', '24865c85fd4c', '947b0722b165', 'e619651b8e23']
MK: drafted==plain 7/7, shas ['34e0f8a1b2c0', 'c1657fc2f1b6', '3b3596014050', 'cc206f8c50ad', '24865c85fd4c', '947b0722b165', 'e619651b8e23']
plain SHAs baseline==merged: True
concurrent vs solo (alone equal/unequal, serial), aggregate tok/s baseline -> merged:
  code t=0.0 c=1: alone 3/3 serial 1/1; 52.3 -> 52.0 (-0.6%), ttft max 0.07 -> 0.07
  code t=0.0 c=4: alone 12/12 serial 1/1; 124.8 -> 124.6 (-0.2%), ttft max 0.36 -> 0.35
  chat t=0.0 c=1: alone 3/3 serial 1/1; 50.6 -> 50.3 (-0.6%), ttft max 0.05 -> 0.05
  chat t=0.0 c=4: alone 12/12 serial 1/1; 126.6 -> 126.3 (-0.2%), ttft max 0.28 -> 0.28
  code t=1.0 c=1: alone 3/3 serial 4/4; 50.1 -> 49.8 (-0.6%), ttft max 0.07 -> 0.07
  code t=1.0 c=4: alone 12/12 serial 4/4; 111.4 -> 111.2 (-0.2%), ttft max 0.36 -> 0.36
  chat t=1.0 c=1: alone 3/3 serial 4/4; 50.5 -> 50.2 (-0.6%), ttft max 0.05 -> 0.05
  chat t=1.0 c=4: alone 12/12 serial 4/4; 115.2 -> 115.0 (-0.2%), ttft max 0.27 -> 0.27
K991 reuse: ['PASS: 35/35 replies equal and as cached as planned', 'timing: 0 resumed replies no sooner than fresh']
MK reuse: ['PASS: 35/35 replies equal and as cached as planned', 'timing: 0 resumed replies no sooner than fresh']
fresh_2k: prompt 2048, first-request ttft 1.007 -> 1.005 s (-0.2%), decode 49.3 -> 49.4 tok/s; repeats+concurrent equal first & baseline tokens: True
fresh_8k: prompt 8192, first-request ttft 3.755 -> 3.760 s (+0.1%), decode 45.9 -> 45.9 tok/s; repeats+concurrent equal first & baseline tokens: True
cached_8k: prompt 8192, first-request ttft 3.747 -> 3.754 s (+0.2%), decode 46.0 -> 45.7 tok/s; repeats+concurrent equal first & baseline tokens: True

[ashhart]

The shared CUDA tree-attention and DFlash block-attention operators and their contract fixtures are included in 1.1.2 (https://github.com/ashhart/TensorFold/releases/tag/v1.1.2). Closing this as carried, with your coauthor credit retained. Thank you for the contribution.

It landed as c7233df.

- Synthema, Ash's AI Assistant.