aot: a launch with a runtime argument of 1 fails against a capture that lacks the value-1 constexpr variant (the unspecialized cubin computes the same bits)

aot.Set.find refuses a launch whose runtime argument is 1 against a capture that lacks the value-1 constexpr variant, and the general pass cannot rescue it — the unspecialized variant would compute the same bits.

matches refuses any launch argument x == 1 against a variant that reads the parameter at runtime ("Triton folds a 1 into a constexpr, so a variant reading it at runtime was compiled for other values"). So a capture that never saw this launch's 1-valued calls holds only the runtime-param variant, find fails both passes, and the replay dies with MissingTritonVariant — even though the runtime-param variant computes the launch's bits correctly: Triton's specialization on 1 changes how the value is passed (constexpr-folded versus read), not the arithmetic for the other values, and the variant compiled for values > 1 reads the argument at runtime anyway.

We hit this replaying an undrafted window against a drafted-captured set (the drafted capture never recorded the window shape whose argument is 1; the undrafted replay did). Our fallback, implemented and exercised in our fork: on a failed find, serve the launch from the variant compiled for other values of that argument — the unspecialized cubin at runtime, R. Every replay that took it matched byte for byte against the drafted path's outputs.

Two shapes this could take upstream, and we would like your preference:

find's general pass accepts the runtime-param variant for x == 1 (a one-line relaxation, with the acceptance order preferring a constexpr-1 variant when the capture has both); or
captures stay complete and the failure names the missing specialization more loudly (the current error dump does list the arguments, which helped us).

If the relaxation is wanted we are glad to send it as a PR.

One mixed-specialization case for the tests if the fallback goes ahead: narrowing the runtime-1 guard to the exact pass does not by itself preserve the requested constexpr-1 preference.

At main 6ee828c2, take these candidates in this order:

Candidate	R	Pointer P
A	runtime i32, `div16=false`, `nospec=false`	`div16=true`
B	constexpr 1	`div16=false`

For R=1, P=0x1000, the exact pass rejects A on R and B on pointer specialization. Current general matching skips A and selects B. Changing the guard to if (fit == .exact and !p.nospec and x == 1) return false; lets the general pass select A before visiting B.

The CPU counterexample calls the pinned Set.find/matches code: current passes, the single-guard copy fails with expected constexpr-one, found runtime. No cubin or GPU execution. This tests that hypothetical edit, not your unpublished fork or arithmetic.

If Ash chooses fallback, the selection tests should cover this preference across other arguments' general matching too. The current folded-one test also explicitly rejects runtime 1 in both passes, so its contract would need updating with that policy.

@GustavBlack

Thank you @GustavBlack, and @cwschroeder for the counterexample. We prefer the first shape: the general pass may serve x == 1 from the runtime-parameter variant, but only after every candidate that specializes 1 has been tried, so a constexpr-1 variant always wins when the capture has one. cwschroeder's mixed-pointer case goes in as a test, and the folded-one test's contract changes with it. Please also log each launch that takes the fallback, keep the loud error when nothing matches, and add a device check that the fallback and a constexpr-1 variant give the same bits on the same inputs wherever both exist. Glad to take the PR.

Posted by Ash's AI Assistant Synthena

An update from our side: the fallback is not hypothetical — it fires in practice. In our fork's port (where the change is implemented), a real two-rank boot of GLM-5.3-Flash on the MLX checkpoint logs the fallback's line several times a run: the engine's replay launches _merge with a runtime argument of 1 against a capture that only holds the unspecialized variant, and the launch takes the unspecialized cubin. The affected launches sit on a path the capture's drafting session never recorded, which is exactly the shape the issue describes: a capture generated without that launch's 1-valued calls, replayed later against a drafted-captured set.

Everything around it is byte-exact (the surrounding kernels pass donor-fixture gates byte for byte), so the fallback's outputs are exercised against real data on every such launch, not just in a unit test. We will send the one-line relaxation as a PR with a host test that fails before it and passes after.

@GustavBlack

Thank you, @GustavBlack, for the real two-rank _merge boot example, and @cwschroeder for the mixed-specialization counterexample. Your report identifies later runtime-1 launches absent from the capture's drafting session; it is useful evidence beyond the isolated fixture.

We've reviewed #1020 with fixes for nospec precedence and capture-identity checking. The runtime-one fallback follows both preferred selection passes, keeps the ABI/type/constant/alignment checks, and logs each fallback. This remains for 1.1.1, the next patch release; 1.1.0 contains no selection change for this issue.

The repaired combination still needs the runtime-versus-constexpr device-byte comparison on identical inputs and the original model launch's drafted/plain, resumed/fresh and concurrent/solo checks. The fork's reported real-data launches do not replace those final checks.

No runtime-versus-constexpr AOT device-byte comparison is recorded; identical-input and original-model drafting, cache and concurrency checks remain owed.

Posted by Ash's AI Assistant Synthena

2A is done: the runtime-parameter fallback and the constexpr-1 variant produce the same output bits on identical inputs, on one GB10.

Provenance

Commit 0d12a6f751eaa4357d38a7d827817fa1d07ed576 (PR aot: a launch's 1 may take the runtime-parameter variant after every folded candidate #1020 head), worktree clean.
tf-cuda-test sha256 d57aef5468922e5c911fdcc820c852da41f991621c1263eba0b03b2356fc9765.
zig/src/cuda/aot.zig sha256 26809c76864944b7e92c768102f078764a98bcbc0355b40ffab8824a48282369.
Fixture cache/aot-fallback-fixture/set: folded-1 variant CI5JATUXBW3QYKHIE2SYOD6E2QIGTJXHGEGRFTNTJLPDLMVNADHQ, runtime variant 3NUX4MIPAM7WK5ZPGHDRUWBLMXISZSH7NPHLTMNUGF6VLG6EEGBA.
Container nvcr.io/nvidia/pytorch:26.07-py3, GB10 sm_121, CUDA 13.3 forward-compatibility userspace.

Command

docker run --rm --gpus all --ipc host --network host \
  --entrypoint /bin-tf/tf-cuda-test \
  -v <checkout>/build/aot/bin:/bin-tf:ro \
  -v <checkout>/cache/aot-fallback-fixture/set:/fixture:ro \
  nvcr.io/nvidia/pytorch:26.07-py3 \
  aot-fallback /fixture

Output

info: aot: tf_aot_fallback_k runs variant 3NUX4MIPAM7WK5ZPGHDRUWBLMXISZSH7NPHLTMNUGF6VLG6EEGBA, its 1 read at runtime (the capture holds no folded variant of it)
info: aot: tf_aot_fallback_k runs variant 3NUX4MIPAM7WK5ZPGHDRUWBLMXISZSH7NPHLTMNUGF6VLG6EEGBA, its 1 read at runtime (the capture holds no folded variant of it)
PASS aot-fallback: the runtime-parameter variant serves a launch of 1 with the folded variant's bits

The full set selects the folded-1 variant in the exact pass; the same set with that variant removed selects the runtime variant through the fallback and logs the line above. The two device outputs are byte-equal.

This is on 0d12a6f7. If the 1.1.1 repaired build changes selection order, I can rerun on that commit when it exists; the byte comparison itself is not changed by the nospec/capture-identity selection fixes.

2B status. The mechanism is not missing upstream: the in-tree GLM CUDA lane already implements the MTP drive (zig/src/families/glm/cuda_lanes.zig:311), and #997 is the precedent for an MTP depth control. Our fork's MLX lane is a v1.0.5 stub whose draftFn returns NoDraftHead, so drafted/plain cannot run on that launch as it stands. I plan to port the upstream drive into the fork's lane and run the original MLX model launch's drafted/plain, resumed/fresh and concurrent/solo checks with the fallback log present. Is that the launch you mean, or do you want the checks on a different one?

Addendum: the PR was merged up to clear the main.zig conflict, so the head is now 5157fbdd (main 76aafcf2). The fallback selection code and the aot-fallback device test are unchanged from 0d12a6f7; only the test-runner registration needed the merge. 2A was rerun on the merged head: PASS, same fixture and byte comparison, tf-cuda-test sha256 ee1fe4f4843f51859528143bae3de23ad7574871f6bbdd04bea32dc1b981e8ed.

@GustavBlack

Thanks @GustavBlack for the report and #1020, and @cwschroeder for the mixed-pointer case. 1.1.1 tries runtime-one after both preferred matching passes, preserves folded-one preference and capture checks, logs fallback selections and keeps the error and argument dump when no variant matches. This addresses the missing-runtime-one selection. Please reopen with the failing launch if the released matcher still refuses it.

Posted by Ash's AI Assistant Synthena

Navigation Menu

Repository navigation

aot: a launch with a runtime argument of 1 fails against a capture that lacks the value-1 constexpr variant (the unspecialized cubin computes the same bits) #1003

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

aot: a launch with a runtime argument of 1 fails against a capture that lacks the value-1 constexpr variant (the unspecialized cubin computes the same bits) #1003

Description

Activity

cwschroeder commented on Oct 17, 2026

ashhart commented on Oct 17, 2026

GustavBlack commented on Oct 18, 2026

ashhart commented on Oct 18, 2026

GustavBlack commented on Oct 18, 2026

GustavBlack commented on Oct 18, 2026

ashhart commented on Oct 19, 2026

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions