Repository navigation
Conversation
…nt hooks
TritonTrace built its interpreted runner by assigning to the user's
Autotuner/Heuristics (runner.fn, _do_bench) and its warmup runner by
deep-copying the result. On Triton 3.8 this crashes three ways on main:
- a @triton.heuristics kernel: "Heuristics.run() got multiple values for
keyword argument 'warmup'" (the launch's warmup=False reached
KernelInterface.warmup, which adds warmup=True);
- @triton.autotune over @triton.heuristics: "k() missing 1 required
positional argument" (the Heuristics layer was replaced by the
interpreted function, so its constexpr never reached the kernel);
- a client that votes for a real compile (the profiler) on a kernel that
calls a traced helper, as under the CLI wrappers: the compile got the
TritonTrace. Called by name, it was rejected by the dependency walk
that keys the compile ("Unsupported function referenced:
<TritonTrace ...>"); passed as a constexpr argument, it was called by
the code generator, which ran the interpreter and left
triton.language patched.
Every Autotuner/Heuristics layer is now shallow-copied over the
interpreted function or the JITFunction. An Autotuner copy gets its own
config cache, never uses Triton's on-disk autotune cache, has Triton's
reset_to_zero/restore_value hooks rebound to it, and (interpreted) uses
the dummy benchmarker even if the user's autotuner cached a real one. A
Heuristics copy warms up through fn.warmup, so the clients' warmup vote
applies to it. During a voted real compile, the traced helpers the
kernel's code reaches by name are bound back to their JITFunctions, and
traced arguments and parameter defaults are passed as JITFunctions. A
TritonTrace called outside any interpreter (neither a traced launch's
nor Triton's own under TRITON_INTERPRET=1) raises a TypeError instead of
running the interpreter; an untraced kernel run by Triton's interpreter
still calls it as before.
Client hooks:
- every traced launch gets its own Launch; two launches used to append
the same object to `launches`;
- patch_warmup asks every client before compiling (main stopped at the
first True vote, yet every client got post_warmup_callback). On exit a
scope takes its own gate out wherever it sits, so overlapping scopes
(e.g. on two host threads) may close in any order, and jit_fn ends
with exactly the warmup it had (main left a bound method in its
instance dict);
- patch_run undoes partial patching when registering client hooks fails:
the sanitizer plus the race detector on one kernel refuse with two
loop overriders and used to leave the interpreter's ops patched;
- ClientManager.finalize finalizes every client even if one raises, then
re-raises the first failure.
mark14wu
added this pull request to stack #496
October 4, 2026 23:48
mark14wu
changed the title
This was referenced Oct 4, 2026
Performance Benchmark
Iterations: 1 warmup + 20 measured |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes in the existing trace path; no IR-mode code. This is the bottom of the stack #494 → #479 → #495 → #480 → #482 and can be reviewed and merged on its own.
Bugs fixed (reproduced on main, Triton 3.8)
@triton.heuristicskernel: "Heuristics.run() got multiple values for keyword argument 'warmup'".@triton.autotuneover@triton.heuristics: "missing 1 required positional argument" (the constexpr the Heuristics layer adds never reached the kernel).triton.languagepatched.Launchobject tolaunches.patch_warmupstopped asking clients at the first True vote, and left a bound method in the JITFunction's instance dict.finalizekept the others from finalizing.How
TritonTraceno longer assigns to the user'sAutotuner/Heuristics. Each layer is shallow-copied over the interpreted function or the JITFunction: an Autotuner copy gets its own config cache, never uses Triton's on-disk autotune cache, hasreset_to_zero/restore_valuerebound to it and (interpreted) uses the dummy benchmarker. During a voted real compile, the traced helpers the kernel reaches are bound back to their JITFunctions. A warmup scope takes its own gate out on exit, so overlapping scopes can close in any order.Behaviour change
A
TritonTracecalled outside any interpreter (neither a traced launch's nor Triton's own underTRITON_INTERPRET=1) raisesTypeErrorinstead of silently running the interpreter. UnderTRITON_INTERPRET=1it runs as before.Testing
tests/unit/test_trace_lifecycle.py: 19 tests; 18 of them fail on main.tile-*scripts); 198 passed at this branch's tip.