Skip to content

test(cuda): opt-in --jobs and --shard for kernel-contract - #3093

Draft
BlivionIaG wants to merge 1 commit into
ashhart:mainfrom
BlivionIaG:pr/kernel-contract-jobs-shard
Draft

BlivionIaG wants to merge 1 commit into
ashhart:mainfrom
BlivionIaG:pr/kernel-contract-jobs-shard

Conversation

@BlivionIaG

@BlivionIaG BlivionIaG commented Nov 13, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

tf-cuda-test kernel-contract runs serially, and a full run takes hours. This adds two opt-in ways to run it faster: --jobs N runs the gates in N child processes and prints one merged SUMMARY, and --shard K/N runs one slice of them. The default run and the existing name filter are unchanged; the filter and --jobs are separate switches and compose (filter first, then shard).

Checks on this head

  • Head commit: 210d696b7bbd428dbc9768dd918fdd5641f2a981.
  • python3 tools/pr_check.py --base 6c5cef8941b27d86b06a986834e13deb9b515e1e: exit 0, 0 errors; lean 345 vs 345 on main, 0 new. zig fmt --check clean. Cross-compile of the CUDA test runner for Linux x86_64 and aarch64: exit 0.
  • Tests run: zig build test-contract-shard-host, 16 of 16 shard tests and 2 of 2 process tests pass, no GPU (the Linux-only parent-death test ran in a Linux VM) (partition covers each index once; filter then shard keeps each named kernel once; merged tallies equal the serial tally; tally lines round-trip and malformed ones are refused; merging refuses a missing, repeated or double-run shard; a crashed or silent shard fails the merge; memory, integrated-GPU, mixed-GPU and argument refusals, each naming the value).
  • CI: runs on this head.
  • Not run, and why: the gates on a chip other than the one below.

Receipt

  • Main against branch: n/a for inference. The runner change was measured by running the same command serially and with --jobs 2 on one machine.

  • Boxes it reaches: the CUDA kernel-contract runner only.

  • New switch? --jobs N and --shard K/N [--device D] on the runner command line, off by default; --help prints both with the per-process memory plan. No environment variable.

  • Environment: Zig 0.17.0, nvcc 13.4, two 16 GB sm_120 GPUs, 64 threads, about 125 GiB free host memory. The filtered serial and --jobs 2 rows below come from an earlier revision of this change (same sharding and merge logic; it still printed timing lines, since removed) merged onto a Flash Next EXL3 branch; the unfiltered --jobs 4 run is on this head.

  • Paths and platforms the change reaches: the Linux CUDA runner. Which I ran: the two commands below. Which I could not: other chips.

  • Exactness: the merged SUMMARY counts equal the serial run's. Same commands, same binary, filtered to exl3 flashnext:

    Run Wall Peak RSS SUMMARY
    serial 8:15 4.9 GiB 61 kernels, 61 pass gates 1-4; 22 controls, 22 red, 0 missed; 137 limit problems
    --jobs 2 4:37 3.4 GiB (parent) identical counts

    The limit problems are pre-existing and equal in both runs. The gates time is the slowest shard's (494 s serial, 276 s with two shards).

  • Full unfiltered run with --jobs 4 on this exact head (built from 210d696b7bbd, whose parent is the main in "Stacked on"; the change touches no kernel), two 16 GB sm_120 GPUs, 125 GiB or more of free host memory, one run: wall 1:02:36 (/usr/bin/time -v), SUMMARY 625 kernels: 625 pass gates 1-4, 0 fail, 8 known findings; 410 controls: 410 red as planted, 0 missed; 546 limit problems. The runner exits 1 because of the limit problems, as the serial run does. I did not run main separately; the limit problems come from the kernels and limits main registers, not from this change.

    • Per child (1 s sampler on /proc and nvidia-smi): shard 0/4 17.2 GiB peak RSS, 1956 MiB peak GPU; 1/4 6.6 GiB, 5650 MiB; 2/4 8.5 GiB, 5648 MiB; 3/4 13.9 GiB, 2164 MiB. All children summed at one instant: peak 24.0 GiB; lowest host MemAvailable 126.2 GiB. Each GPU's device memory peaked at about 5.8 GiB of 16.
    • Shard 0 is the straggler: its limit checks took 1559 s against 57 to 367 s on the others, so it ended at about 3755 s (gates 2196 s plus limits 1559 s) while shard 3, the slowest on gates, ended at about 2989 s. Splitting the limit checks across shards would shorten the run; this change does not.
    • The last serial unfiltered run, on an older head and before this change, took 3.2 h with the same three counts then (632 / 421 / 564 on that head; the counts differ from this head's because the trees differ). I have not run serial unfiltered on this head, so the 3.2 h to 1.04 h comparison is across heads.
  • Decode speed, prompt speed: n/a, the change cannot reach inference.

  • Licence of any new weights or drafter: n/a.

Checklist

  • The title and description match the head commit, including the receipt and the stack.
  • 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 in the diff or its commits. No AI attribution lines.

@ashhart

ashhart commented Nov 13, 2026

Copy link
Copy Markdown
Owner

Thanks for this. It's a careful change, and the default path is untouched. We read it closely and found no way for a sharded run to report green when a serial run wouldn't. A crashed, killed or silent shard, an exit its counts don't explain, and a list run more or less than once all fail the merged run. Per-kernel seeds come from kernel names, so shards test the same inputs as a serial run. On our side zig fmt --check, zig build test-contract-shard-host (14 of 14) and the Linux x86_64 and aarch64 cross-compiles of the CUDA runner all pass at 9486f3d.

A few asks before we run it on hardware:

  1. If only the parent process is killed, the shards keep running, holding GPU and host memory, and their output goes nowhere. Could each --shard child set PR_SET_PDEATHSIG to SIGKILL and check its parent is still alive, or could the parent pass SIGTERM and SIGINT on to its children?
  2. The mixed-GPU check compares compute capability only, but registration reads each GPU's multiprocessor count to choose kernel configurations. Two sm_120 cards with different SM counts would test different configurations under one merged SUMMARY. Please also refuse when the SM count (or the device name) differs.
  3. Please post the full unfiltered run when it finishes: wall time, each child's peak RSS and peak GPU memory, and the merged SUMMARY next to the serial one. That's what should back the 28 GiB and 7 GiB plan. For reference, on a 128 GB unified-memory GB10 that plan allows at most three shards for a full run, so --jobs 4 is refused there.
  4. Small things, optional: --jobs 1 prints the gradient-suites PASS line twice, a child output line longer than 64 KiB stops that child, and stopping the remaining shards once one has failed would save time.

On overlaps: #2620 adds a standalone suite to contract_after.zig. Whichever of the two lands second needs that suite as an entry in the new served list, or it quietly stops running. The #2008 conflict is one line in the image name list.

Once these are in and your numbers are posted, we'll run serial against --jobs on a GB10 with the same binary, and compare per-gate lines, merged counts and wall time.

  • Synthena (Ash's Agent)

Shards split kernels, standalone suites and launch records; the parent merges each child's TALLY line into one SUMMARY.
@BlivionIaG
BlivionIaG force-pushed the pr/kernel-contract-jobs-shard branch from 9486f3d to 210d696 Compare November 13, 2026 01:51
@BlivionIaG

Copy link
Copy Markdown
Contributor Author

Thanks. Pushed 210d696b7b (same single commit, amended) with the code asks:

  1. Orphans. Each --jobs child gets an internal --parent PID, sets PR_SET_PDEATHSIG to SIGKILL before anything else and compares getppid() with it, refusing with ParentGone on a mismatch. The parent also forwards SIGINT and SIGTERM to every live shard. I ran the Linux pieces in a Linux VM: the death signal is armed, a wrong parent pid is refused, and after kill -9 of the parent a child is gone. The real --jobs binary needs CUDA, so those paths are not exercised end to end yet.
  2. Mixed GPUs. fit now also compares the multiprocessor count and the device name, and the refusal names the differing field and both values.
  3. Small ones. --jobs 1 prints the gradient PASS line once (host-tested for serial, 1, 2, 4, filtered and unfiltered). A child line over the cap is shown cut and no longer stops the child, and a later TALLY line still parses. When a shard fails to finish, the others are stopped and the output says which ones were stopped.
  4. qwen3_5 cuda: the GDN kernels read the state, its rollback copy and the log through a stream's state table (#2444 P3) #2620 overlap. Confirmed: a suite goes in list or served in contract_after.zig, and one called from anywhere else runs in neither the serial run nor --jobs. The docs now say so. I can't write a host test for a suite missing from served, because there is one list and the suite modules need CUDA.

The full unfiltered --jobs 4 result is in the description (wall 1:06:06, SUMMARY 632 / 421 / 564). It doesn't have each child's peak RSS and peak GPU memory, so I'll run it again on this head with a sampler for both and post the table here, next to the serial SUMMARY. That run waits behind the EXL3 chain's device cells on the same machine.

@BlivionIaG

Copy link
Copy Markdown
Contributor Author

Full unfiltered kernel-contract --jobs 4 on this head (210d696b7bbd, main parent, no other branch), two 16 GB sm_120 GPUs, one run, under /usr/bin/time -v with a 1 s sampler:

  • Wall 1:02:36. SUMMARY 625 kernels: 625 pass gates 1-4, 0 fail, 8 known findings; 410 controls: 410 red as planted, 0 missed; 546 limit problems; exit 1 because of the limit problems.
  • Peak RSS per child: 17.2 / 6.6 / 8.5 / 13.9 GiB (shards 0-3); all children together 24.0 GiB at one instant; lowest host MemAvailable 126.2 GiB. Peak GPU memory per child 1956 / 5650 / 5648 / 2164 MiB; each card's total peaked at about 5.8 GiB of 16.
  • Shard 0 is the straggler: its limit checks took 1559 s against 57 to 367 s elsewhere, ending at about 3755 s; shard 3 (slowest on gates) ended at about 2989 s.
  • Not measured: serial unfiltered on this head. The 3.2 h serial figure is from an older head.

The body is updated; it no longer says this revision is unmeasured.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants