Skip to content

[FIX] Keep the JITFunction reachable from traced autotuners - #489

Open
mark14wu wants to merge 1 commit into
mainfrom
claude/autotune-listener-attributeerror-ded222
Open

mark14wu wants to merge 1 commit into
mainfrom
claude/autotune-listener-attributeerror-ded222

Conversation

@mark14wu

@mark14wu mark14wu commented Oct 3, 2026

Copy link
Copy Markdown
Collaborator

Summary

On Triton 3.8, launching a traced autotuned kernel on the interpreted path raises:

AttributeError: 'function' object has no attribute 'fn'
  at triton/runtime/autotuner.py:256 (Autotuner.run)

This happens whenever knobs.autotuning.listener is set. Before calling the listener, Autotuner.run walks self.fn down to the JITFunction it tunes. The interpreted path replaces autotuner.fn with an InterpretedFunction, whose .fn is the plain Python function, so the walk fails. check_disk_cache does the same walk, so cache_results=True and TRITON_CACHE_AUTOTUNING=1 hit the same error even without a listener. That part already failed before the 3.8 requirement.

Changes in tilelens/core/trace.py:

  • _InterpretedLeaf. Traced Autotuner/Heuristics layers now sit on an _InterpretedLeaf instead of a bare interpreted function. Its run interprets the kernel, other attributes forward to the interpreted function, and fn is the kernel's JITFunction. The listener therefore receives the user's own JITFunction, the same object an untraced launch reports. The leaf mirrors _IRLeaf from the IR-mode work.
  • Disk cache off. Traced autotuners set cache_results = False. Triton only turns its disk cache off for autotuners decorated in interpreter mode, but TileLens enables interpreter mode at launch time, after that decision. Without this, the now-successful walk would either fail on driver.active (CPU) or write the dummy (1.0, 1.0, 1.0) timings into Triton's cache (GPU), where a later real run would read them.
  • Shallow warmup copy. _warmup_runner uses copy instead of deepcopy, because the chain now reaches a JITFunction, which holds an RLock. Autotuner.warmup only rebinds nargs, so the shared state is safe.
  • Re-tracing. unpack_kernel in both TritonTrace and GluonTrace recognizes a leaf left by an earlier trace of the same autotuner. Re-tracing therefore keeps jit_fn; before, the second trace lost it.
  • Gluon. GluonTrace hit the same crash; it now builds the same leaf and records jit_fn.

Not changed: stock Triton 3.8 with TRITON_INTERPRET=1 and a listener set crashes the same way without TileLens, because no JITFunction exists anywhere in that chain. That is an upstream interpreter issue, so sources that are already InterpretedFunctions keep the old behaviour.

Test Plan

New tests in tests/end_to_end/test_core.py:

  • test_autotune_listener_sees_jit_function and test_gluon_autotune_listener_sees_jit_function: the listener fires once with the traced jit_fn, and the kernel output is correct.
  • test_autotune_retrace_keeps_jit_function and test_gluon_autotune_retrace_keeps_jit_function: tracing the same autotuner twice keeps jit_fn and still launches.
  • test_autotune_interpreter_skips_disk_cache: a cache_results=True autotuner launches and writes no *.autotune.json.
  • The two Gluon wrapper structure tests now assert on the leaf.

The Triton test kernels are decorated inside knobs.runtime.scope() with interpret = False. Otherwise a worker that also collected tests/unit/test_multithreading.py (which sets TRITON_INTERPRET=1 at import) would give an InterpretedFunction and silently test the upstream case instead.

Runs (CPU, Triton 3.8):

  • Without the fix, all 7 new or changed tests fail. Both listener tests and the disk-cache test fail with the AttributeError above.
  • With the fix, pytest tests/ -n auto has no new failures. The only failures are 6 CLI tests that also fail on main in this environment because the tile-* entry points are not installed.
  • The new tests also pass when run together with tests/unit/test_multithreading.py.
  • Before rebasing onto [COMPAT] Require Triton 3.8 #485, the same comparison on Triton 3.6 had no new failures either, and the disk-cache and re-trace tests went from failing to passing there.
  • On an RTX 4090, the CUDA-only autotune tests in test_core.py / test_tracer.py pass.

Related Issues

None filed. Stock TRITON_INTERPRET=1 + knobs.autotuning.listener fails in Triton itself and should be reported upstream separately.

Breaking Changes

None intended. After tracing, the user's Autotuner/Heuristics object now holds an _InterpretedLeaf at .fn rather than an InterpretedFunction. Attribute access and autotuner.warmup(...) on that object behave as before, and .fn.fn is now the JITFunction. Traced autotuners no longer use Triton's autotune disk cache.

Checklist

  • I added tests to all new functionality I added/bugs I fixed.
  • I verified that a human has reviewed all code in this PR.
  • I ran npm run build:frontend if the PR modified any TypeScript code. (no TypeScript changes)
  • I made sure that my code is well documented (comments explaining strange code, docstrings for functions, website modified if new functionality added).

Triton 3.8's Autotuner walks its .fn chain down to a JITFunction before
calling knobs.autotuning.listener, and check_disk_cache does the same
walk. The interpreted path replaced autotuner.fn with an
InterpretedFunction, whose .fn is the plain Python function, so a traced
autotuned launch raised AttributeError: 'function' object has no
attribute 'fn' when a listener was set (and with cache_results=True or
TRITON_CACHE_AUTOTUNING=1 even without one).

Put an _InterpretedLeaf under traced Autotuner/Heuristics layers instead:
it runs the interpreted kernel, forwards other attributes to it, and
exposes the JITFunction as .fn. Also turn the traced autotuner's disk
cache off so dummy interpreter timings are never persisted, shallow-copy
the warmup runner (a JITFunction holds an RLock), and unpack a leaf left
by an earlier trace so re-tracing keeps jit_fn. GluonTrace gets the same
leaf and now records jit_fn.
@github-actions

github-actions Bot commented Oct 3, 2026

Copy link
Copy Markdown

Performance Benchmark

Benchmark main (min) PR (min) Change Samples
gemm 0.054s 0.057s +6.2% ⚠️ 20 / 20
gemm_oob 0.061s 0.063s +4.3% 20 / 20
indirect_load 0.011s 0.011s +6.2% ⚠️ 20 / 20
nested_loop 0.114s 0.115s +0.9% 20 / 20
block_pointer_loop_advance 0.062s 0.064s +2.4% 20 / 20
liger_jsd 0.077s 0.079s +3.1% 20 / 20
flaggems_layernorm 0.197s 0.201s +2.0% 20 / 20
swiglu 0.093s 0.090s -3.5% 20 / 20
cross_entropy 0.526s 0.516s -1.8% 20 / 20
fused_linear_jsd 0.115s 0.119s +3.9% 20 / 20
Total 1.308s 1.315s +0.5% N/A

Threshold: >5% regression flagged with ⚠️
Iterations: 1 warmup + 20 measured
Samples are shown as main / PR; long pytest benchmarks may use fewer samples.

This branch has not been deployed

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant