aot: a launch's 1 may take the runtime-parameter variant after every folded candidate

[GustavBlack]

Closes #1003, with @cwschroeder's mixed-specialization case from that thread as a test.

What this changes

find()'s second pass may now serve a launch's 1 from the variant that reads the argument at runtime — the
fallback #1003 proposes, in the shape Ash preferred there: only after every candidate that specialized the
launch's 1s has had its own pass, so a constexpr-1 variant always wins when the capture holds one. The
matcher's guard narrows to the exact pass (fit == .exact), and the general pass splits in two:

• the variants that folded every one of the launch's 1s (Triton's own pick) — cwschroeder's case: their
  constexpr-1 candidate wins even though the runtime variant also matches under the general rules;
• then the rest, where a runtime-parameter variant takes the launch's 1. Each launch served here logs one line
  (aot: <fn> runs variant <hash>, its 1 read at runtime (the capture holds no folded variant of it)); the loud
  MissingTritonVariant error and its argument dump are unchanged.

The rule aot.zig opens with ("each launch picks the variant Triton itself would have picked") holds in the
exact pass and in the new first pass; the second pass is the fallback this issue argued for, and the one our
fork's GLM port exercises on every two-rank boot (the captured set's drafting session never records _merge's
one-row launches).

Tests

Host (aot.zig, no GPU — the set parses without its cubins):

• the folded-1 test's contract updated: the exact pass still refuses a runtime variant for a launch of 1, and
  the folded variant wins when the set holds it (was: refused "in either pass");
• a launch of 1 with no folded variant in the set takes the runtime-parameter variant — fails on main
  (9ccf4727) with MissingTritonVariant, passes here;
• cwschroeder's mixed case (runtime variant first in the set, folded second, an aligned pointer): the folded
  one is picked — passes on main too, kept as the guard against the naive single-guard narrowing;
• the folded variant that cannot serve (an unaligned pointer's div16) leaves the launch to the runtime
  variant.

Device (tf-cuda-test aot-fallback <dir>, new): zig/tests/cuda/aot_fallback_kernels.py compiles one tiny
kernel's two specializations on the running device (R=1 folds, R=2 stays a runtime parameter) into an aot set;
the check runs the launch from the full set and from the set with the folded variant removed, and compares the
outputs byte for byte.

Receipt

• Environment: this branch at 0d12a6f7 on 9ccf4727; the build and the device check on a GB10 (sm_121,
  CUDA 13.3, the nvcc container pytorch 26.07); the host tests in the same container without GPUs.
• zig build test (the three-phase build's second phase): 1379/1386 passed, 7 skipped (the GPU-gated ones, on
  their usual skip), none failed.
• tf-cuda-test aot-fallback on the GB10: PASS, the fallback's log line firing on the fallback launches, the
  folded variant's bits byte for byte.
• python3 tools/zig/lean_check.py: 517 problems, identical to the same run on 9ccf4727 (517); the one
  finding in aot.zig (a /// line past 120 columns) is inherited and shifted.
• zig fmt --check clean.

Checklist

• Every token still goes through the lane rounds — unchanged; this touches variant selection only.
• No precision traded for speed — no arithmetic change; the fallback's cubin is the same one the engine
  already launches for that variant's other launches.
• New tests fail before the change, pass after it, and skip cleanly without their dependency (the device
  check's fixture needs Triton in the image; the host tests parse aot.json text with no GPU and no cubins).
• Comments one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[ashhart]

Thanks @GustavBlack for the runtime-one fallback and your 1,379-passing-test receipt, and @cwschroeder for the mixed-specialization example. We're taking this with fixes for 1.1.1. Both preferred selection passes retain priority before the runtime-1 fallback. Review added protection for nospec arguments and kept the capture-identity checker on its preferred-specialization policy, so the fallback cannot weaken that audit. The loud missing-variant error and fallback function/hash diagnostic remain. Host controls pass for those repairs; the combined build still owes your folded-versus-runtime device comparison and the original model launch's drafted/plain, resumed/fresh and concurrent/solo checks. Your original device receipt does not replace that gate.

Posted by Ash's AI Assistant Synthena

[GustavBlack]

2A evidence is posted on #1003: #1003 (comment)

On 0d12a6f7, the runtime-parameter fallback and the constexpr-1 variant give byte-equal outputs on identical inputs, and the fallback log line fires. 2B is still pending the launch question in that comment; I have not substituted the _merge boot.

[GustavBlack]

PR conflict cleared. I merged main 76aafcf2 into the branch; head is now 5157fbdd.

• The only conflict was zig/tests/cuda/main.zig: main added many new test registrations, and this PR adds the aot-fallback import/help/dispatch lines. The merge keeps both; the PR diff against main is still the same four files (251 insertions, 5 deletions).
• The merged tree builds on a GB10: 3-phase build rc=0, FMT-CLEAN, host suite green.
• 2A was rerun on the merged head: PASS, same fixture and byte comparison, tf-cuda-test sha256 ee1fe4f4843f51859528143bae3de23ad7574871f6bbdd04bea32dc1b981e8ed. The fallback code and the device test are unchanged from 0d12a6f7; only the test-runner registration was merged.
• GitHub now reports the PR mergeable. The unstable state is the CI checks still running, not a conflict.

[GustavBlack]

Checks finished on the merged head: fmt, mobile, desktop and zig pass; zig-windows fails.

That job is new in main (.github/workflows/no-gpu.yml, added after 9ccf4727) and runs the Windows host suite. It did not run on the pre-merge head 0d12a6f7, where the four existing checks passed. The errors are in main's Windows paths (core/digest_key.zig:34, server/speech.zig:279, std.dynamic_library, families/flashnext/cuda_*, families/h3/cuda_runtime_fake.zig, server/benchmark*); none are in aot.zig, aot_fallback_tests.zig, or the main.zig registration. The PR is conflict-free and its own checks are green; the Windows job is a main-side failure.