[FIX] Keep traces from mutating Autotuner/Heuristics or failing their warmup by mark14wu · Pull Request #494 · Deep-Learning-Profiling-Tools/tilelens

mark14wu · 2026-10-04T23:47:42Z

Summary

On main, a trace:

put the interpreted function and a dummy benchmarker into the user's Autotuner/Heuristics. Now it runs on copies of that chain, which also fixes autotune over heuristics and cache_results.
failed every launch of a @triton.heuristics kernel in the warmup. Now each Heuristics copy warms up through fn.warmup.

Changes

`tilelens/core/trace.py`

_rebuild_runner replaces _interpreter_runner and _warmup_runner. It builds the trace's Autotuner/Heuristics chain from shallow copies.
- Before, tracing changed fn and the benchmarker on the user's objects, so later untraced launches ran in the interpreter.
- Before, @triton.autotune over @triton.heuristics failed with "missing 1 required positional argument", because the Heuristics layer was replaced by the interpreted function.
Each Autotuner copy:
- gets its own cache. A shallow copy would otherwise share the user's.
- sets cache_results = False. With it on, check_disk_cache fails on the interpreted chain with "'function' object has no attribute 'fn'".
- sets do_bench to the dummy benchmarker. do_bench is a cached property, so once the user's Autotuner had benchmarked, setting _do_bench did nothing. A two-config kernel ran 111 times in the interpreter instead of 3.
Each Heuristics copy gets a warmup that fills in its values and calls fn.warmup, where the clients vote on whether to compile. Before, KernelInterface.warmup passed warmup= a second time, so every launch failed with "Heuristics.run() got multiple values for keyword argument 'warmup'".

CI and tests

tests/unit/test_multithreading.py no longer sets TRITON_INTERPRET=1 at import, which made every later @triton.jit in the session an InterpretedFunction. .github/workflows/python-app.yml sets it in a separate step for that file and for test_autotune_interpreter_mode, which covers a trace without a JITFunction.
tests/unit/test_trace_lifecycle.py (new):
- autotune over heuristics produces the right output and leaves the user's Autotuner, including its cache, unchanged;
- a profiler launch of a heuristics kernel compiles once, with the heuristic's value filled in.
tests/end_to_end/test_tracer.py reads the copy's do_bench instead of _do_bench.

github-actions · 2026-10-04T23:51:10Z

Performance Benchmark

Benchmark	main (min)	PR (min)	Change	Samples
gemm	0.074s	0.074s	-1.1%	20 / 20
gemm_oob	0.083s	0.082s	-0.6%	20 / 20
indirect_load	0.015s	0.015s	+0.4%	20 / 20
nested_loop	0.153s	0.154s	+0.2%	20 / 20
block_pointer_loop_advance	0.087s	0.087s	+0.2%	20 / 20
liger_jsd	0.103s	0.102s	-0.7%	20 / 20
flaggems_layernorm	0.268s	0.268s	-0.0%	20 / 20
swiglu	0.125s	0.125s	-0.3%	20 / 20
cross_entropy	0.699s	0.702s	+0.5%	20 / 20
fused_linear_jsd	0.155s	0.157s	+0.9%	20 / 20
Total	1.762s	1.765s	+0.2%	N/A

Iterations: 1 warmup + 20 measured
Samples are shown as main / PR; long pytest benchmarks may use fewer samples.

…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.

Keep the fixes and drop the code they did not need: - patch_warmup keeps main's shape and takes one real_compile hook from the trace; _WarmupGate (out-of-order scopes), RealArgs, the compile_context/real_args parameters, _without_warmup and _voted_warmup are gone; - one _rebuild_runner: Autotuner copies get their own cache, cache_results=False (check_disk_cache crashed on the interpreted chain) and do_bench set directly; Triton's restore hooks are rebound only when they close over the original; no nested-trace look-through; - a traced function called outside a traced launch does what the untraced function does (Triton's interpreter, or JITFunction's refusal) instead of raising a custom TypeError; - parameter-default unwrapping is dropped; _is_triton_internal is gone (it hid traced helpers in triton.* modules), and so are the global walker's defensive guards; - finalize and the first-True-vote short circuit are back to main's; collecting finalize errors moves up the stack, where a client can first raise in finalize; - patch_run creates its Launch per launch and sets up inside the try, without the lang_patched flag. test_trace_lifecycle.py keeps 7 tests, each failing on main.

A trace without a JITFunction (decorated while the interpret knob was on) that reaches a real compile fails either way. The skips only changed the error, and for an argument they let the code generator call the trace, which patched triton.language. The walk no longer looks through untraced JITFunctions. The CLI traces every triton.jit, so such a chain needs triton.runtime.jit.jit or an explicit tilelens.trace on a device function; main does not support it either.

The autotune test runs both with Triton's restore_value post_hook and with a user post_hook that closes over one list, so the hook rebinding and its identity check are both pinned. The voted-compile test passes FN by position.

Python 3.12 inlines comprehensions, so a kernel's co_names already holds the names its tuple comprehensions call; _code_names' recursion into nested code objects only served 3.10/3.11. README and AGENTS.md follow, and the cp310/cp311 NKI wheel sources go.

The version files go back to main's. #507 requires Python 3.12, and this PR depends on it.

Rebasing onto #507 replayed "Move the Python 3.12 requirement to #507", which put the 3.10 version files back.

When #507 merged, GitHub rebased #494 onto main and replayed such a commit as a revert of #507.

When #507 merged, GitHub rebased #494 onto main instead of only retargeting it, so the merge -s ours step moves up to the PR above the rebased one.

Jokeren · 2026-10-14T22:07:53Z

-            self._populate_loop_hooks(all_loop_callbacks)
-            patch_for_loop(frontend_name)
-            patch_lang(fn, frontend_name, client_manager=self)
+            # Inside the try: a failing setup (e.g. two loop overriders) must


How could we fail?

Removed. client.py is now unchanged from main.

Jokeren · 2026-10-14T22:13:23Z

+        restore_value=["out_ptr"],
+        # Triton's own restore_value post_hook, or a user's that closes over
+        # just `posts` and must stay bound to it.
+        post_hook=(lambda nargs, exception: posts.append(exception))


This is not a typical use

Agreed. Removed the hook rebinding along with this test.

mark14wu · 2026-10-14T22:23:42Z

Update: #494 no longer fixes this (it fails as on main); #513 fixes it.

Minimal reproducer: an ordinary kernel that calls a @triton.jit helper, run under tile-profiler (Triton 3.8).

import torch, triton, triton.language as tl


@triton.jit
def add_one(x):
    return x + 1


@triton.jit
def kernel(x_ptr, BLOCK: tl.constexpr):
    offs = tl.arange(0, BLOCK)
    tl.store(x_ptr + offs, add_one(tl.load(x_ptr + offs)))


x = torch.zeros(8, device="cuda")
kernel[(1,)](x, BLOCK=8)
print("ok")

tile-profiler repro.py:

main: RuntimeError: Unsupported function referenced: <tilelens.core.trace.TritonTrace object at 0x...>
this PR: the profiler report, then ok

The CLI traces every @triton.jit function, so add_one is a TritonTrace. The profiler asks for a real compile, and Triton accepts only JITFunction callees. This PR binds traced helpers back to their JITFunctions for that compile (_real_compile).

Its setup only fails when two clients with loop overriders trace one kernel, which is not a typical use.

A launch passes warmup=False, so it no longer compiles first. That is the profiler's need, not every trace's, and with it go the fixes that made the per-launch compile work: traced helpers, constexpr traces and the Heuristics warmup.

Without it, Triton's hooks leave restore_copies on the user's Autotuner, as an untraced run does; the kernel's output is the same.

Only a host-side call of a traced device function reaches it outside a launch, which is not a typical use.

The autotune-over-heuristics test fails the same way when a launch warms up.

The custom client is down to no-op callbacks, which the tracer covers.

Jokeren · 2026-10-15T18:12:00Z

+
+@pytest.fixture(autouse=True)
+def _real_jit():
+    # tests/unit/test_multithreading.py sets TRITON_INTERPRET=1 at import


I think the right way is to avoid using TRITON_INTERPRET=1 for certain tests in the CI configuration file

Done. CI now sets TRITON_INTERPRET=1 only for test_multithreading.py and test_autotune_interpreter_mode, and this fixture is gone.

Jokeren · 2026-10-15T18:14:34Z

        frontend = get_frontend(frontend_name)
        namespaces = frontend.namespaces
+        # A fresh Launch per launch, so `launches` gets distinct entries.
+        self.launch = Launch()


I'm not sure if this is correct. I thought the user is responsible for cleaning launches, otherwise we will lose records

You're right. The Tracer's records accumulate across launches, so a fresh Launch lost the earlier launches' tensors and the visualizer hit a KeyError. Reverted.

Jokeren · 2026-10-15T18:15:21Z

-        with self.client_manager.patch_warmup(self.jit_fn):
-            if self.warmup_runner:
-                self.warmup_runner.warmup(*args, **kwargs)
+        if kwargs.get("warmup"):


We no longer call patch_warmup? How about pre and post warm up callbacks?

Restored main's run, so every launch goes through patch_warmup and the warmup callbacks again. The remaining fix is the Heuristics warmup, which passed warmup= twice.

The Tracer's records accumulate across launches, so a fresh Launch per launch held earlier launches' records without their tensors, and the visualizer failed with a KeyError.

Every launch calls the warmup again, and the clients' pre_warmup_callback decides whether to compile. Each Heuristics copy fills in its values and calls fn.warmup, instead of KernelInterface.warmup, which passed warmup= a second time.

test_multithreading.py set it at import, which made every later @triton.jit in the session an InterpretedFunction.

7d36ce5 is dcd8610 rebased onto #507, which the previous merge brought in. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

#494 was trimmed before it landed; this layer keeps what it builds on: - client.py: main's patch_warmup and patch_run (#494), with op/loop registration for interpreting clients only. _WarmupGate is gone. begin_launch still gives every launch its own Launch, and finalize still finalizes every client when one raises. - trace.py: IR helpers on #494's _rebuild_runner (ir=True adds the conflicting-meta-parameter check); main's TritonTrace.__call__, comparing interpreting clients. The IR capture keeps mapping traced arguments to their JITFunctions (_untraced_call_args); the global unwrapping (_unwrapped_trace_globals) is dropped with #494's warmup path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

main's squash of #494 has the tree of #494's head, merged above. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

mark14wu added this pull request to stack #496 October 4, 2026 23:48

mark14wu changed the title ~~split/trace fixes~~ [FIX] Keep traces from mutating Autotuner/Heuristics and isolate client hooks Oct 4, 2026

This was referenced Oct 4, 2026

[FEAT] Add the host-compiled IR lifecycle for clients #479

Open

[FEAT] Add the TTIR reader #495

Open

[FEAT] Add a compiled mode to the sanitizer #480

Open

[REFACTOR] Share the Term to Z3 lowering between compiled clients #482

Open

mark14wu force-pushed the split/trace-fixes branch from d675311 to dcd8610 Compare October 12, 2026 18:02

This was referenced Oct 12, 2026

[FEAT] Extend the eager symbolic engine for the race detector #501

Open

[FEAT] Add the race detector's two-copy happens-before solver #503

Open

[FEAT] Add the eager race detector #504

Open

mark14wu changed the title ~~[FIX] Keep traces from mutating Autotuner/Heuristics and isolate client hooks~~ [FIX] Keep traces from mutating Autotuner/Heuristics or leaking patches Oct 13, 2026

This was referenced Oct 14, 2026

[FIX] Close two false proofs in the race detector's two-copy model #505

Open

[COMPAT] Require Python 3.12 #507

Merged

mark14wu removed this pull request from stack #496 October 14, 2026 20:10

mark14wu changed the base branch from main to compat/require-python-3.12 October 14, 2026 20:10

mark14wu added this pull request to stack #508 October 14, 2026 20:10

mark14wu mentioned this pull request Oct 14, 2026

[FEAT] Add the race detector's TTIR reader facts #509

Open

Base automatically changed from compat/require-python-3.12 to main October 14, 2026 20:59

mark14wu added 6 commits October 14, 2026 16:59

[TEST] Pin hook rebinding and positional untracing

c03d7bc

The autotune test runs both with Triton's restore_value post_hook and with a user post_hook that closes over one list, so the hook rebinding and its identity check are both pinned. The voted-compile test passes FN by position.

[COMPAT] Require Python 3.12

45920c6

Python 3.12 inlines comprehensions, so a kernel's co_names already holds the names its tuple comprehensions call; _code_names' recursion into nested code objects only served 3.10/3.11. README and AGENTS.md follow, and the cp310/cp311 NKI wheel sources go.

[COMPAT] Move the Python 3.12 requirement to #507

ad9517f

The version files go back to main's. #507 requires Python 3.12, and this PR depends on it.

Jokeren force-pushed the split/trace-fixes branch from 30f3aa6 to ad9517f Compare October 14, 2026 20:59

[COMPAT] Restore main's Python 3.12 requirement

e7de857

Rebasing onto #507 replayed "Move the Python 3.12 requirement to #507", which put the 3.10 version files back.

mark14wu added a commit that referenced this pull request Oct 14, 2026

[DOCS] Never set files back to the base's version in a stacked PR

e8e9fe3

When #507 merged, GitHub rebased #494 onto main and replayed such a commit as a revert of #507.

mark14wu mentioned this pull request Oct 14, 2026

[DOCS] Keep stacked PRs up to date with merges; run test_multithreading.py without a GPU #500

Open

mark14wu added a commit that referenced this pull request Oct 14, 2026

[DOCS] Describe the rebase GitHub does after the bottom PR merges

6443393

When #507 merged, GitHub rebased #494 onto main instead of only retargeting it, so the merge -s ours step moves up to the PR above the rebased one.

Jokeren reviewed Oct 14, 2026

View reviewed changes

mark14wu added 2 commits October 14, 2026 19:17

[REFACTOR] Drop the patch_run setup cleanup

6f95f27

Its setup only fails when two clients with loop overriders trace one kernel, which is not a typical use.

[REFACTOR] Warm up only when run is asked to

7a12e53

A launch passes warmup=False, so it no longer compiles first. That is the profiler's need, not every trace's, and with it go the fixes that made the per-launch compile work: traced helpers, constexpr traces and the Heuristics warmup.

mark14wu mentioned this pull request Oct 15, 2026

[FEAT] Compile each launch the profiler traces #512

Closed

mark14wu added 3 commits October 15, 2026 11:58

[REFACTOR] Drop the restore_value hook rebinding

fc73f42

Without it, Triton's hooks leave restore_copies on the user's Autotuner, as an untraced run does; the kernel's output is the same.

[REFACTOR] Keep main's TritonTrace.__call__

3863ae3

Only a host-side call of a traced device function reaches it outside a launch, which is not a typical use.

[TEST] Drop test_a_launch_does_not_warm_up

56a4290

The autotune-over-heuristics test fails the same way when a launch warms up.

mark14wu changed the title ~~[FIX] Keep traces from mutating Autotuner/Heuristics or leaking patches~~ [FIX] Keep traces from mutating Autotuner/Heuristics or compiling each launch Oct 15, 2026

mark14wu changed the title ~~[FIX] Keep traces from mutating Autotuner/Heuristics or compiling each launch~~ [FIX] Keep traces from mutating Autotuner/Heuristics or warming up every launch Oct 15, 2026

[TEST] Trace the lifecycle tests with the tracer client

4af77ae

The custom client is down to no-op callbacks, which the tracer covers.

Jokeren requested changes Oct 15, 2026

View reviewed changes

mark14wu added 3 commits October 15, 2026 14:49

[REFACTOR] Keep main's shared Launch

260e0dc

The Tracer's records accumulate across launches, so a fresh Launch per launch held earlier launches' records without their tensors, and the visualizer failed with a KeyError.

[TEST] Set TRITON_INTERPRET=1 in CI only for the tests that need it

6c85503

test_multithreading.py set it at import, which made every later @triton.jit in the session an InterpretedFunction.

mark14wu changed the title ~~[FIX] Keep traces from mutating Autotuner/Heuristics or warming up every launch~~ [FIX] Keep traces from mutating Autotuner/Heuristics or failing their warmup Oct 15, 2026

Jokeren approved these changes Oct 15, 2026

View reviewed changes

mark14wu merged commit 5bb5cfb into main Oct 15, 2026
4 checks passed

mark14wu deleted the split/trace-fixes branch October 15, 2026 20:06

mark14wu added a commit that referenced this pull request Oct 15, 2026

Record #494's rebased first commit as merged

ce6aa4a

7d36ce5 is dcd8610 rebased onto #507, which the previous merge brought in. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

mark14wu added a commit that referenced this pull request Oct 15, 2026

Merge main at #494 into ir-mode-core

c336c10

main's squash of #494 has the tree of #494's head, merged above. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Navigation Menu

Repository navigation

[FIX] Keep traces from mutating Autotuner/Heuristics or failing their warmup - #494

[FIX] Keep traces from mutating Autotuner/Heuristics or failing their warmup#494
mark14wu merged 16 commits into
mainfrom
split/trace-fixes

mark14wu commented Oct 4, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 4, 2026 •

edited

Loading

Uh oh!

Jokeren Oct 14, 2026

Uh oh!

mark14wu Oct 15, 2026

Uh oh!

Jokeren Oct 14, 2026

Uh oh!

mark14wu Oct 15, 2026

Uh oh!

mark14wu commented Oct 14, 2026 •

edited

Loading

Uh oh!

Jokeren Oct 15, 2026

Uh oh!

mark14wu Oct 15, 2026

Uh oh!

Jokeren Oct 15, 2026

Uh oh!

mark14wu Oct 15, 2026

Uh oh!

Jokeren Oct 15, 2026

Uh oh!

mark14wu Oct 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

mark14wu commented Oct 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

tilelens/core/trace.py

CI and tests

Uh oh!

github-actions Bot commented Oct 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Performance Benchmark

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mark14wu commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

mark14wu commented Oct 4, 2026 •

edited

Loading

`tilelens/core/trace.py`

github-actions Bot commented Oct 4, 2026 •

edited

Loading

mark14wu commented Oct 14, 2026 •

edited

Loading