Repository navigation
aot.zig: a launch with no exact captured variant takes one that assumes less of its arguments - #794
Conversation
…es less of its arguments
…untime (ashhart#794) Triton folds an int argument of 1 into a constexpr, so a captured variant that takes the argument at runtime was compiled for other values. The exact pass refused it; the general pass now refuses it too, so a launch of 1 is served only by the variant compiled for it. A host test holds both variants: a launch of 1 takes the folded one and neither pass takes the runtime one, and a launch of 32 takes the runtime variant only in the general pass. The test set's parse is wrapped at 120 columns. Co-authored-by: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com>
…sumes less, as checked (ashhart#794)
|
Thank you, @zweck, and @BobClawblaw, whose #726 has the same matcher change. A launch that no variant of a captured Triton set matches exactly now takes one that assumes less of its arguments. The exact match still comes first, so every launch that matched before takes the same variant. An integer of 1, which Triton folds into the kernel, still takes only the variant compiled for it. Your commit is in 1.0.10 under your name ( On a GB10 with Nemotron 3.5 Lightning (MLX 4-bit) and the qualified sm_121 set, replies were the same tokens as before, drafted and with |
- aot.zig takes main's version: ashhart#794 already holds this branch's general pass, and also refuses a launch of 1 a runtime variant - lane_host.zig: main's born lengths and host spans, with this branch's cacheKey so an image prompt is still matched by its images' content - glm.zig: main's affine and wide modules beside vision and image
Closes #791.
This is the
aot.zigpart of ff7cc22 in #726 (@BobClawblaw), proposed formainon its own: it is family-neutral(Nemotron's sets go through the same matcher), and #726 is stacked on #608, #598 and #614. If you would rather it
land inside #726, close this and it will wait — no work lost either way.
What this changes
find()makes two passes over the function's variants. The first,.exact, is main's current rule unchanged. Onlywhen no variant matches does the second pass,
.general, accept a variant that assumes less than the launch gives:The two passes keep
aot.zig:1's "each launch picks the variant Triton itself would have picked". A singleimplication pass would not:
aot.jsonis sorted by(fn, hash), so the first variant to match could be a lessspecialised one even when an exact variant is in the set. The added comment is a single line, per CONTRIBUTING.
Effect: every launch that matches on
mainpicks the same variant. Only launches thatmainrefuses can now beserved, and only by a variant that assumed less of an argument than the launch gives; a pointer that is not
16-aligned is still refused by every variant that has
div16on it.Tests added in
aot.zig(onmainthey fail to build — there is no second pass to name; with the change they pass):div16serves a 16-aligned pointer;div16refuses an unaligned pointer in both passes;div16variant.Receipt
0efbd447; Linux x86_64; Zig 0.17.0 (.zig-version);no GPU attached — host tests only, as the change is host-side variant selection.
zig build test-cuda-host --summary all— 25/25 steps succeeded, 72/72 tests passed; the CUDA module's testbinary alone runs 37/37, including the three added.
zig build test --summary all— 44/44 steps succeeded, 257/262 tests passed, 5 skipped.python3 tools/zig/lean_check.py— 539 problems, identical to the same run onmain(539). The onefinding in
aot.zig(a///line past 120 columns,Launches …) is inherited and unchanged.were observed in earlier work, not repeated here.
Serving checklist: not applicable beyond the items below — this changes variant selection only, with no decode path,
no precision change, and no launch that matches on
mainselecting a different variant.