Skip to content

tools/zig: one manifest → kernel set step for every family, reading a Recorder dump as well as a manifest #792

Closed

Description

@zweck

What exists

  • tools/zig/triton_aot_manifest.py records a run's launches (Recorder.dump → launches.json) and maps each
    launched specialisation to its cached cubin (build → manifest.json).
  • zig/tests/cuda/nemotron/aot_pack.py turns one or more manifests into the set the engine loads: aot.json in the
    KernelJson schema (zig/src/cuda/aot.zig:32-46) plus cubins/<hash>.cubin. It is documented only in step 3 of
    zig/tests/cuda/nemotron/box/README.md. Qwen3.8-27B on CUDA: the qwen27 family's engine for NVIDIA GPUs, every stream drafted #608 packs the 27B's set with it as well.
  • packaging/README.md (-Ddist-cuda-aot) and RUNBOOK.md (TENSORFOLD_CUDA_KERNELS) describe the set's layout,
    sm<capability>/aot.json and cubins/<captured-hash>.cubin, but not how a set is made.

What is missing

  1. Every CUDA family's set goes through this step, but it is filed under one family's tests, and the docs that
    describe a set do not point at it.
  2. It reads only the manifest shape: kernels as a list, cubin relative to the cache, abi, cubin_sha256. A
    Recorder dump has a different shape: kernels is a mapping keyed by hash, params lists every parameter
    including the constexprs, and the cubin is an absolute path in files['<name>.cubin']. The 0.6.6 capture made on a
    GB10 for the 27B came in that shape (from a newer copy of the Recorder, keyed by cubin sha256), and it cannot be
    packed without first running build against the same cache.
  3. --jit is required, and only the in-tree capture scripts write jit.json. A capture made outside the tree (the
    10 Oct rule, see Qwen3.8-27B on CUDA: image input and --kv-cache fp8 (256K with vision on a 5090) #726) has none.
  4. --cache must be repeated beside every --manifest, although each manifest records its own cache
    (triton_aot_manifest.py:265).

Proposal

Move aot_pack.py to tools/zig/triton_aot_pack.py, beside the manifest step, and change it as follows:

  • Accept either shape for each --manifest. The runtime arguments are the signature entries whose type is not
    constexpr, in params order when the input has params.
  • Find each cubin by cubin under the cache, or by files['<name>.cubin']. When that absolute path was the run's own
    (a container mount), look for it under --cache.
  • Make --jit optional. Without it no parameter is marked nospec. By matches(), that can only refuse a launch
    (an int argument of 1, or a multiple of 16, that Triton was told not to specialise: aot.zig:191-192). It can
    never serve the launch a variant compiled for other values.
  • Default --cache to the manifest's own.
  • Write a constexpr only when its name is not a runtime parameter. The matcher reads a variant's constant for an
    argument only when that argument is not one of the variant's params (aot.zig:183-187).

These stay as they are: the fields written, the PTX check wherever the input has abi, the sha256 check wherever it
has cubin_sha256, and the sort order. The Nemotron box README and test_aot_manifest.py follow the file to its new
path.

Two rules any manifest → set step must keep

  1. Triton's hidden launch pointers are not params. Triton's launcher appends global_scratch and
    profile_scratch; the caller never names them. matches() ends with return runtime == k.params.len
    (aot.zig:209), so a variant that lists either one can never match. aot_pack.py already leaves them out, and it
    refuses a cubin whose PTX parameters are not the runtime arguments followed by exactly those two
    (aot_pack.py:54-57). The new file keeps that check wherever the input has abi, and it builds params from the
    signature for both shapes. A set whose kernels need scratch is refused at load (aot.zig:100).
  2. div16 records what Triton proved, not what the caller will pass. The step writes it as captured: a
    tt.divisibility of 16 in attrs. Whether a variant without it may serve an aligned pointer is up to the
    matcher. Main refuses it (aot.zig:196, strict equality); Qwen3.8-27B on CUDA: image input and --kv-cache fp8 (256K with vision on a 5090) #726 (ff7cc22) accepts it in a second pass. See
    aot.zig: a captured variant compiled without 16-divisibility refuses an aligned launch #791.

Effect on existing sets

None. The engine reads kernels and ignores every other field (aot.zig:91). Sets already packed are not touched,
and tools/release/package.sh copies sets without calling this step. For input that aot_pack.py accepts today, the
output differs only in three places:

  • generator.
  • div16, where an argument carries an attribute other than tt.divisibility 16. aot_pack.py:53 counts any
    attribute as divisibility.
  • Constants that share a runtime parameter's name, which are no longer written.

Gate: python3 -B zig/tests/cuda/nemotron/test_aot_manifest.py, run against the new path and given a second case
that packs a Recorder dump directly. It is host-only and needs no GPU or model.

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