[Enhancement] Add eager-mode support for tilelang.autotune#1906
Merged
SiriusNEO merged 3 commits intotile-ai:mainfrom Mar 12, 2026
Merged
[Enhancement] Add eager-mode support for tilelang.autotune#1906SiriusNEO merged 3 commits intotile-ai:mainfrom
SiriusNEO merged 3 commits intotile-ai:mainfrom
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
Contributor
📝 WalkthroughWalkthroughAdds a new end-to-end TileLang autotuning test for a tiled matmul and enhances the autotuner with argument normalization, a public compile() API, lazy/eager execution modes, and persistent out_idx save/load behavior. Changes
Sequence Diagram(s)sequenceDiagram
participant Test as Test Runner
participant AT as AutoTuneImpl
participant JIT as JITCompiler/Cache
participant FS as Disk (out_idx.json)
Test->>AT: call (args, kwargs) or compile(...)
AT->>AT: _normalize_value(args/kwargs)
AT->>AT: compute cache key
AT->>JIT: check cache / compile kernel
JIT-->>AT: kernel + selected_config
AT->>FS: save out_idx.json (on save)
FS-->>AT: load out_idx.json (on load)
AT-->>Test: return kernel (compile mode) or execute and return result
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@testing/python/autotune/test_tilelang_autotune_eager_mode.py`:
- Around line 110-116: The test only exercises the direct-call path for symbolic
dimensions, leaving matmul.compile uncovered; update the test
(test_autotune_matmul_symbolic_m) to run both paths: wrap both
matmul.compile(M=M, N=N, K=K) (when return_kernel True) and matmul(a, b) (when
return_kernel False) inside the same set_autotune_inputs([a, b]) block so the
compile path is exercised with symbolic dims too, ensuring the compile
invocation uses the explicit M/N/K kwargs while the direct call uses inference.
In `@tilelang/autotuner/param.py`:
- Around line 439-443: The code currently opens path / OUT_IDX_PATH and assigns
out_idx_override from the JSON, which will raise FileNotFoundError for older
caches; modify the block around OUT_IDX_PATH to catch FileNotFoundError (or
check existence) and on missing file set out_idx_override =
compile_args.out_idx, while keeping the verbose logger.debug message and only
attempting json.load when the file exists; reference the OUT_IDX_PATH constant,
the out_idx_override variable, and compile_args.out_idx so you update the
correct assignment logic.
In `@tilelang/autotuner/tuner.py`:
- Around line 344-347: _kernel_parameters currently has an incorrect type and
inconsistent structure across lazy and eager paths; update its annotation from
tuple[str, ...] to a 2-tuple representing (args_tuple, kwargs_tuple) (e.g.,
tuple[tuple[Any, ...], tuple[tuple[str, Any], ...]]) and normalize both code
paths so set_kernel_parameters always stores (args,
tuple(sorted(kwargs.items()))) (ensure args are converted to a tuple and kwargs
to a sorted tuple of items) — change the lazy-mode call that currently passes
raw (args, tuple(sorted(kwargs.items()))) if needed to the normalized form and
ensure eager-mode already using (norm_args, norm_kwargs) matches that exact
format, and update any docstring/comments for _kernel_parameters and
set_kernel_parameters to reflect the normalized structure.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 3d9a0326-3395-4f40-a638-ee4166e54a76
📒 Files selected for processing (3)
testing/python/autotune/test_tilelang_autotune_eager_mode.pytilelang/autotuner/param.pytilelang/autotuner/tuner.py
Contributor
There was a problem hiding this comment.
Pull request overview
Extends @tilelang.autotune to support eager-mode @tilelang.jit kernels by enabling eager compilation/execution workflows and improving autotuner artifact persistence.
Changes:
- Update autotune invocation to handle eager vs. lazy JIT modes and add an
AutoTuneImpl.compile()API. - Improve autotuner caching to store both the best kernel and its best config, and adjust cache-key generation for callable config providers.
- Persist and restore eager-mode
out_idx_overridein autotune disk cache; add eager-mode autotune matmul tests.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
tilelang/autotuner/tuner.py |
Adds eager-mode path to autotune wrapper, normalizes call keys, caches best kernel+config, and introduces compile() helper. |
tilelang/autotuner/param.py |
Saves/loads an out_idx_override value alongside autotune artifacts to restore eager-mode output behavior. |
testing/python/autotune/test_tilelang_autotune_eager_mode.py |
New test validating eager-mode autotune execution and the new compile() workflow for a matmul kernel. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
SiriusNEO
requested changes
Mar 10, 2026
Collaborator
SiriusNEO
left a comment
SiriusNEO
left a comment
There was a problem hiding this comment.
LGTM, just some comments. And you can also have a look at the comments from Copilot (You can ignore some Copilot comments if you think it's OK, since it just gives suggestions based on common programming knowledge)
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tilelang/autotuner/tuner.py`:
- Around line 351-353: AutoTuneImpl is overwriting raw invocation data when
normalizing kernel parameters: instead of reusing self._kernel_parameters (which
holds the normalized cache key) to call the configs callable and assigning its
materialized result to self.configs, keep raw invocation args separate from the
normalized key; call the configs callable using a local variable (e.g., configs
= self._kernel_parameters_raw or the local kernel_args/kernel_kwargs passed into
__call__), set key_data["configs"] to that local (raw) configs value, and ensure
run() uses the local/config-callable (not self.configs) so you do not replace
the callable with a frozen materialized list; alternatively introduce a distinct
attribute (like _kernel_key_params) for the normalized key and leave
self.configs as the original callable so subsequent run() calls receive raw
tensors/Vars rather than normalized tuples/strings.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 762bbe26-5e88-4c8a-ad69-56c734c909ab
📒 Files selected for processing (4)
testing/python/autotune/test_tilelang_autotune.pytesting/python/autotune/test_tilelang_autotune_eager_mode.pytesting/python/autotune/test_tilelang_autotune_with_inputs.pytilelang/autotuner/tuner.py
✅ Files skipped from review due to trivial changes (1)
- testing/python/autotune/test_tilelang_autotune.py
SiriusNEO
approved these changes
Mar 12, 2026
1 task
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
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.
Currently
@tilelang.autotuneonly works with lazy mode kernels. This PR extends@tilelang.autotuneto better support eager mode.Summary by CodeRabbit
New Features
Tests