Skip to content

aot.zig: a launch with no exact captured variant takes one that assumes less of its arguments - #794

Merged
ashhart merged 1 commit into
ashhart:mainfrom
zweck:fix/aot-general-pass
Oct 14, 2026
Merged

ashhart merged 1 commit into
ashhart:mainfrom
zweck:fix/aot-general-pass

Conversation

@zweck

@zweck zweck commented Oct 13, 2026

Copy link
Copy Markdown
Contributor

Closes #791.

This is the aot.zig part of ff7cc22 in #726 (@BobClawblaw), proposed for main on 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. Only
when no variant matches does the second pass, .general, accept a variant that assumes less than the launch gives:

.general => if (p.div16 and !div) return false,

The two passes keep aot.zig:1's "each launch picks the variant Triton itself would have picked". A single
implication pass would not: aot.json is sorted by (fn, hash), so the first variant to match could be a less
specialised 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 main picks the same variant. Only launches that main refuses can now be
served, 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 div16 on it.

Tests added in aot.zig (on main they fail to build — there is no second pass to name; with the change they pass):

  • a variant without div16 serves a 16-aligned pointer;
  • a variant with div16 refuses an unaligned pointer in both passes;
  • when the set holds both, an aligned launch takes the div16 variant.

Receipt

  • Environment, with the TensorFold commit: this branch on 0efbd447; Linux x86_64; Zig 0.17.0 (.zig-version);
    no GPU attached — host tests only, as the change is host-side variant selection.
  • Tests run, with pass, fail and skip counts:
    • zig build test-cuda-host --summary all — 25/25 steps succeeded, 72/72 tests passed; the CUDA module's test
      binary alone runs 37/37, including the three added.
    • zig build test --summary all — 44/44 steps succeeded, 257/262 tests passed, 5 skipped.
  • Lean code: python3 tools/zig/lean_check.py — 539 problems, identical to the same run on main (539). The one
    finding in aot.zig (a /// line past 120 columns, Launches …) is inherited and unchanged.
  • Not run, and why: an on-device/model check — no sm_121 GPU on this host. The GB10 launches quoted in the issue
    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 main selecting a different variant.

  • 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. No AI attribution lines.

@zweck
zweck marked this pull request as ready for review October 13, 2026 16:36
@ashhart
ashhart merged commit 8a876d7 into ashhart:main Oct 14, 2026
ThinkOffApp pushed a commit to ThinkOffApp/TensorFold that referenced this pull request Oct 14, 2026
…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>
ThinkOffApp pushed a commit to ThinkOffApp/TensorFold that referenced this pull request Oct 14, 2026
@ashhart

ashhart commented Oct 14, 2026

Copy link
Copy Markdown
Owner

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. 462f1a15f puts that refusal in the general pass as well as the exact pass, and is co-authored by @BobClawblaw.

Your commit is in 1.0.10 under your name (25690c3e4). The CHANGELOG line is 9253e54f8. Your head 072217288 is merged at 8a876d7d3 with no change to the tree.

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 "draft": false, alone and concurrent. With a copy of the set missing the 11 variants that assume more alignment than a sibling, the new rule served 34,786 launches of _kv_write, _chunk, _merge and _topk with the remaining variants and every reply was still the same tokens, where the build before it refused to load.

@ashhart ashhart mentioned this pull request Oct 14, 2026
BobClawblaw added a commit to BobClawblaw/TensorFold that referenced this pull request Oct 14, 2026
- 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
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.

aot.zig: a captured variant compiled without 16-divisibility refuses an aligned launch

2 participants