Lightning-AI / Lightning-AI/lightning-thunder

can I use thunder to autotune and select kernel implementations based on input shapes.

Open
#2,151 1 comment 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.