[rootsbymenda]

What this changes

On an RTX 5090 the kernel contract's nvfp4_expert_waves<2, 2, 4> cutoff control stays green (#1778). The waves kinds' two-wave shape was a fixed 1,024 rows, and the loaded gate-up waves entry holds all of that shape's units in one wave, so the control drops nothing past the first wave. Each waves kind's largest served shape is now sized from the residency the loader measured for the entry it launches, at its 4- or 1-warp CTAs: 8 x (resident warps / column blocks + 1) rows, never fewer than 1,024. That is one item more than a full wave holds.

What the loader measures here, with nvcc 13.0 (the issue's run used 13.3):

RTX 5090, 170 SMs   nvfp4_expert_waves<2, 2, 4>: 680 CTAs = 2,720 resident warps; <1, 0, 4>: 850 CTAs = 3,400 warps
                    (1-warp CTAs: 2,720 and 3,400)
                    1,024 rows route 130 items x 20 column blocks = 2,600 gate-up units: one wave
                    with this change: 1,096 rows, 140 items; the first wave holds 136, so the control drops 4 and goes red
GB10, 48 SMs        gate-up 192 CTAs = 768 warps, down 240 CTAs = 960 warps (1-warp CTAs: 768 and 960)
                    every waves shape stays 1,024 rows, as before

• Only the test fixture changes. No kernel, record, gate, tolerance or control changes, and nothing is skipped or marked known.
• Each control keeps its unmodified twin's shapes: cut_waves runs up_waves's, and cut_waves_one runs down_waves_one's. The half entry keeps its shapes, because it serves only shapes that fit one wave.
• The sizing uses checked arithmetic, so Experts.kernel() now returns an error union and Group.create passes the error on. The sized shapes live in the fixture (Experts.served), which the group owns for the whole run.
• 1,024 rows stays the floor, so devices where it already takes two waves, like the GB10, run the shapes they run today. If you'd rather use the smallest two-wave shape everywhere, removing the floor is a one-line change.

Fixes #1778

Stacked on: nothing (based on main a8c02eef2; merges cleanly with main 1dfc78c55).

Receipt

• Environment, with the TensorFold commit, MLX or PyTorch version, machine and GPU, checkpoint and revision: head 331744663 on main a8c02eef2; Zig 0.17.0; kernels built with nvcc 13.0.88, the release compiler. x86_64: Ubuntu 24.04.5, RTX 5090 (sm_120, 170 SMs), driver 595.99.02. aarch64: DGX Spark, GB10 (sm_121, 48 SMs), driver 580.178.04. No checkpoint is needed, because only the contract fixture changes.

• Exactness, from tools/bench_concurrent.py --alone --serial or token_sha against "draft": false: n/a. No kernel, sum or token path changes.

• Decode speed before and after, from tools/bench_openai.py: n/a, for the same reason.

• Prompt speed before and after: n/a, because the change can't reach prompt processing.

• Host test: contract_nvfp4_waves_test.zig, imported from contract_arch_test.zig, so test-cuda-host runs it. It builds the real fixture's Experts.kernel() over loaded kernels with synthetic residency. For every waves kind at its 4- or 1-warp CTAs, the largest served shape must route more units than the entry's resident warps hold, counting at least rows / 8 items, however the router spreads them. Each control's shapes must also equal its twin's. It runs four residencies:

  • 768 and 960 warps, the GB10's (stays at 1,024 rows);
  • 2,720 and 3,400 warps, the 5090's;
  • 8,160 warps, which is 170 SMs at 48;
  • 2,560 and 10,240 warps, which 1,024 rows' units fill exactly (1,032 rows).

  On main, the same file fails three of them:

  error: 'contract_nvfp4_waves_test.test.nvfp4 waves: 2,720 gate-up and 3,400 down resident warps, 170 SMs at 16 and 20 warps' failed:
         up_waves: 1024 rows route at least 2560 units, and 2720 resident warps hold them in one wave
  error: 'contract_nvfp4_waves_test.test.nvfp4 waves: 8,160 resident warps, 170 SMs at 48 warps' failed:
         up_waves: 1024 rows route at least 2560 units, and 8160 resident warps hold them in one wave
  error: 'contract_nvfp4_waves_test.test.nvfp4 waves: warps that 1,024 rows' units fill exactly take one more item, 1,032 rows' failed:
         up_waves: 1024 rows route at least 2560 units, and 2560 resident warps hold them in one wave
  

• The controls on the RTX 5090. kernel-contract has no per-set filter, and on the 5090 it's CPU-bound here (float64 references), so a full run didn't fit. Instead, a local command that isn't part of this PR runs contract_nvfp4_tile.Group's 18 registrations alone through gates.runAll, with the contract's seed. It also prints the loader's residency and each waves kind's real item count. Results:

  main (a8c02eef2)
  PASS shapes control: nvfp4_expert_waves<2, 2, 4> with its items dropped past the first wave: gate 3 at the two-wave shape refuses it: 11 shapes bit-equal to the kernel it replaces (worst 0.000 of the bound), 0 edges refused by the kernel
  CONTROL MISSED control: nvfp4_expert_waves<2, 2, 4> with its items dropped past the first wave: gate 3 at the two-wave shape refuses it: shapes did not go red
  CONTROL control: nvfp4_expert_waves<1, 0, 1> with its items dropped past the first wave: gate 3 at the two-wave shape refuses it: red at shapes, as planted
  the group: 14 kernels: 14 pass gates 1-4, 0 fail, 0 known; 4 controls: 3 red as planted, 1 missed
  
  head (331744663)
  PASS shapes cuda nvfp4 experts decode gate-up nvfp4_expert_waves<2, 2, 4> (equal waves): 11 shapes bit-equal to the kernel it replaces (worst 0.000 of the bound), 0 edges refused by the kernel
  PASS shapes cuda nvfp4 experts decode down nvfp4_expert_waves<1, 0, 1> (TF_NVFP4_DECODE_WARPS=1): 11 shapes bit-equal to the kernel it replaces (worst 0.000 of the bound), 0 edges refused by the kernel
  FAIL shapes control: nvfp4_expert_waves<2, 2, 4> with its items dropped past the first wave: gate 3 at the two-wave shape refuses it: at 1096x640x2560 its bytes differ from nvfp4_expert_kernel<2, 2, 4>'s from byte 2632960
  CONTROL control: nvfp4_expert_waves<2, 2, 4> with its items dropped past the first wave: gate 3 at the two-wave shape refuses it: red at shapes, as planted
  FAIL shapes control: nvfp4_expert_waves<1, 0, 1> with its items dropped past the first wave: gate 3 at the two-wave shape refuses it: at 1024x2560x640 its bytes differ from nvfp4_expert_kernel<1, 0, 4>'s from byte 20480
  CONTROL control: nvfp4_expert_waves<1, 0, 1> with its items dropped past the first wave: gate 3 at the two-wave shape refuses it: red at shapes, as planted
  the group: 14 kernels: 14 pass gates 1-4, 0 fail, 0 known; 4 controls: 4 red as planted, 0 missed
  

  On the GB10 the same command at head gives 14 of 14 kernels passing and 4 of 4 controls red, at 1,024 rows, the shapes main runs there.

• Mutation: this change with the sizing returning the old 1,024 rows. The host test fails the same three cases, and on the 5090 CONTROL MISSED comes back for nvfp4_expert_waves<2, 2, 4> (130 of 130 items in the first wave).

• The full contract on the GB10 (tf-cuda-test kernel-contract, built with -Dnvcc=... -Dsm=121):

  head  SUMMARY 527 kernels: 527 pass gates 1-4, 0 fail, 6 known findings; 344 controls: 344 red as planted, 0 missed; 181 limit problems; gates 2826044 ms, limits 2405588 ms
        FAIL shapes control: nvfp4_expert_waves<2, 2, 4> with its items dropped past the first wave: gate 3 at the two-wave shape refuses it: at 1024x640x2560 its bytes differ from nvfp4_expert_kernel<2, 2, 4>'s from byte 2560
        FAIL shapes control: nvfp4_expert_waves<1, 0, 1> with its items dropped past the first wave: gate 3 at the two-wave shape refuses it: at 1024x2560x640 its bytes differ from nvfp4_expert_kernel<1, 0, 4>'s from byte 10240
  main  SUMMARY 527 kernels: 527 pass gates 1-4, 0 fail, 6 known findings; 344 controls: 344 red as planted, 0 missed; 181 limit problems; gates 2828731 ms, limits 2133636 ms
        (the same two FAIL lines, from bytes 2560 and 10240)
  

  The run exits 1 on the limit problems alone. Every gate FAIL line pairs with its CONTROL line (344) or KNOWN line (6). The 6 known findings are the ones registered on main.
  The 181 limit problems are main's. On this GB10, main's run and head's give the same verdict lines, line for line, timings aside: every PASS, FAIL, CONTROL, KNOWN and limits line matches. The waves shapes on this device are unchanged, so that's expected. This change touches no kernel, record or build file. They are:

  • 168 launches with no compiled usage recorded: qwen4_hc 48, qwen4_ple 40, qwen4_gdn 20, qwen4_projection 20, qwen35 15, qwen35_kv8 15, nvfp4_experts_tile 10;
  • 13 live launches whose usage differs from the records: qwen35_kv8 6, qwen35 4, nvfp4_experts_tile 2, qwen35_keep 1.

• Tests run, with pass, fail and skip counts (x86_64 Linux, each step from an empty cache):

  zig fmt --check build.zig zig                  exit 0
  git diff --check a8c02eef2..HEAD               exit 0
  python3 tools/zig/lean_check.py                base 424 problems, head 424, new 0 (the same 424 lines)
  zig build test --test-timeout 10m              head: 6009/6036 passed, 27 skipped, 0 failed
                                                 main: 6004/6032 passed, 27 skipped, 1 failed
  zig build test-cuda-host                       head: 4191/4200 passed, 9 skipped, 0 failed (4 of them new)
                                                 main: 4186/4196 passed, 9 skipped, 1 failed
  zig build test-kernel-contract                 exit 0, 144/144 (head and main)
  zig build test-golden                          exit 0, 5/5 steps
  zig build fuzz-cortheon                        exit 0, 19/19
  zig build                                      exit 0, 29/29 steps
  zig build -Dnvcc=... -Dsm=120                  exit 0, 169/169 steps
  zig build -Dnvcc=... -Dsm=121 (aarch64)        exit 0, 169/169 steps
  

  Main's one failure is exl3_load.test.a scan that fails an allocation frees every group prefix it held, which returns error.NondeterministicMemoryUsage from checkAllAllocationFailures. It's intermittent. It failed in both of main's runs above and in both runs of this change with the sizing mutated back. It passed in two more runs of main that had only the new test added, and in head's runs. This change doesn't touch exl3_load.zig.

• atc overlap: contract_nvfp4_fact.zig is shared with Flash Next shared down: remove the scatter with direct slot output #1792, and contract_arch_test.zig with cuda: tune dense 27B MLP tiles on RTX PRO 4500 #1730 and exl3 cuda: the 27B's EXL3 speed series on the native Dense module (#260), stacked on #1465 #1467. Checked with git merge-tree: our files merge cleanly with Flash Next shared down: remove the scatter with direct slot output #1792 and cuda: tune dense 27B MLP tiles on RTX PRO 4500 #1730. exl3 cuda: the 27B's EXL3 speed series on the native Dense module (#260), stacked on #1465 #1467 conflicts with main in contract_arch_test.zig and three other files, and the conflict list is the same with or without this branch.

• Not run, and why:

  • the 5090's complete kernel-contract: it's CPU-bound here, so the full receipt is the GB10's, and the 5090 receipt is the group run above;
  • nvcc 13.3, the issue's compiler: we reproduced the miss with 13.0. The sizing reads whatever residency the loaded build reports, so a compiler that changes registers changes the shape with it;
  • Metal and HIP: this CUDA-only fixture change doesn't touch them.

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. (Test fixture only.)
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (Not measured; no token path changes.)
• No precision traded for speed. If the bits change, the description says which sums change and why. (No bits change.)
• 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. (No kernel changes; the GB10 SUMMARY line is above.)
• Prompt processing is no slower than the last release. (Not reached by the change.)
• 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. (The host test needs no GPU.)
• 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.

[ashhart]

Thank you @rootsbymenda. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.