Lightning-AI / Lightning-AI/lightning-thunder
can I use thunder to autotune and select kernel implementations based on input shapes.
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Enable `thunder` to autotune and select kernel implementations based on input shapes.
### Motivation
Different kernel implementations for an op can perform better on different input shapes. Manually switching kernels is tedious. `thunder` could automate this for better performance, especially with dynamic shapes.
For instance, if I have `my_op_impl_A` (for small inputs) and `my_op_impl_B` (for large inputs), I'd like `thunder` to pick the right one automatically at runtime based on the actual input tensor shapes.
### Pitch
1. **Register Implementations:** Allow users to provide multiple kernel versions for an operation.
2. **Autotune:** `thunder` profiles these kernels with varying input shapes (user-provided or inferred).
3. **Auto-Select:** Based on tuning results, `thunder` dynamically dispatches to the best kernel for the current input shape during execution.
4. **Cache Results:** Store tuning outcomes to avoid re-profiling for similar shapes.
This makes `thunder` an intelligent dispatcher, boosting performance without manual kernel management.
### Alternatives
* **Manual `if/else` in Python:** Clunky and doesn't integrate well with JIT.
* **Single "good enough" kernel:** Misses optimization opportunities.
An integrated autotuning/selection feature in `thunder` would be superior.
### Additional context
This capability is present in libraries like TVM and TensorRT. Adding it to `thunder` would significantly enhance its utility for PyTorch acceleration, particularly for users with custom ops or those seeking peak performance.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
The issue names no files, tests, or entry points to begin with. Start by locating Thunder's kernel registration and dispatch mechanisms, then define how shape-based profiling, runtime selection, and cached tuning results should work; done would require an agreed design and implementation covering those requirements.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100