Repository navigation
tools/zig: one manifest → kernel set step for every family, reading a Recorder dump as well as a manifest #792
Description
What exists
tools/zig/triton_aot_manifest.pyrecords 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.pyturns one or more manifests into the set the engine loads:aot.jsonin the
KernelJsonschema (zig/src/cuda/aot.zig:32-46) pluscubins/<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) andRUNBOOK.md(TENSORFOLD_CUDA_KERNELS) describe the set's layout,
sm<capability>/aot.jsonandcubins/<captured-hash>.cubin, but not how a set is made.
What is missing
- 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. - It reads only the manifest shape:
kernelsas a list,cubinrelative to the cache,abi,cubin_sha256. A
Recorder dump has a different shape:kernelsis a mapping keyed by hash,paramslists every parameter
including the constexprs, and the cubin is an absolute path infiles['<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 runningbuildagainst the same cache. --jitis required, and only the in-tree capture scripts writejit.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.--cachemust 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 thesignatureentries whose type is not
constexpr, inparamsorder when the input hasparams. - Find each cubin by
cubinunder the cache, or byfiles['<name>.cubin']. When that absolute path was the run's own
(a container mount), look for it under--cache. - Make
--jitoptional. Without it no parameter is markednospec. Bymatches(), 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
--cacheto 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'sparams(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
- Triton's hidden launch pointers are not
params. Triton's launcher appendsglobal_scratchand
profile_scratch; the caller never names them.matches()ends withreturn runtime == k.params.len
(aot.zig:209), so a variant that lists either one can never match.aot_pack.pyalready 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 hasabi, and it buildsparamsfrom the
signature for both shapes. A set whose kernels need scratch is refused at load (aot.zig:100). div16records what Triton proved, not what the caller will pass. The step writes it as captured: a
tt.divisibilityof 16 inattrs. 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 thantt.divisibility16.aot_pack.py:53counts 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