Skip to content

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

Closed

Description

@zweck

Note: the same change already exists in #726 as ff7cc22 (@BobClawblaw), inside that stack. This issue and its draft pull request track the matcher fix for main on its own; if you would rather it land as part of #726, the pull request can simply be closed.

The matcher compares a pointer's alignment with the captured variant's div16 by equality. Current line, main, zig/src/cuda/aot.zig:196:

if (!std.mem.eql(u8, p.type, x.ty) or p.div16 != (x.addr % 16 == 0)) return false;

The same equality is applied to i32 arguments (:192) and u64 arguments (:204).

div16 is Triton's tt.divisibility 16 at capture, meaning the cubin may assume it. A variant without it assumes nothing about the address, so it computes the same for any pointer; one with it needs an aligned pointer. Equality gets one of the four cases wrong:

launch pointer variant div16 equality correct
16-aligned false refused serve: the cubin assumes nothing
16-aligned true served serve
not aligned false served serve
not aligned true refused refuse: the cubin may use 16-byte accesses on it

The safe rule is an implication: a variant captured without divisibility serves any pointer; one captured with it needs an aligned pointer. The refusal in the last row must stay.

Two launches on a GB10 (sm_121) against #608's engine show the two cases apart:

  • _embed (zig/src/families/qwen27/cuda/triton.zig:51 there). Every pointer the Zig caller passes is 16-aligned, but the variant captured from 0.6.6 was compiled without div16 on at least one of them. Equality refused it with MissingTritonVariant; under the implication the variant serves the launch.
  • _paths (triton.zig:96 there). PARENTS was not 16-aligned at all — round.zig:179 computes it as attn_plan + (W + 4 + 3 * n_items) * 4, aligned only when W + 3 * n_items is a multiple of 4. It was refused, so the captured variant has div16 on PARENTS, and that refusal is required under either rule. The fix belongs to the caller: start each plan segment on a 16-byte boundary, as ff7cc22 in Qwen3.8-27B on CUDA: image input and --kv-cache fp8 (256K with vision on a 5090) #726 does (pad16, AttnAt).

The same equality refuses an int argument Triton would have folded into a constexpr (the value 1) when the captured variant has it as a runtime parameter. That variant also computes the same result.

Activity

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions