Lightning-AI / Lightning-AI/lightning-thunder
Consider using pre-built Flash Attention kernels via `kernels`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
Hey,
I am Sayak from the [Kernels](https://github.com/huggingface/kernels) team at Hugging Face. I noticed that this project uses Flash Attention which includes a long build time. We ship pre-built binaries (which provide bit-exact outputs as the upstream) and thereby, we make it easy to use.
Using FA3 on a supported machine is as easy as:
```py
# make sure `kernels` is installed: `pip install -U kernels`
from kernels import get_kernel
kernel_module = get_kernel("kernels-community/flash-attn3")
flash_attn_func = kernel_module.flash_attn_func
flash_attn_func(...)
```
Let us know if you'd be interested in this and and we'd be happy to provide a draft of how it would look in your repo.
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
Start by locating the project's current Flash Attention integration and build/dependency path, then review the proposed `kernels` entry point using `pip install -U kernels` and `get_kernel("kernels-community/flash-attn3")`. Done would require an agreed integration design, supported-machine behavior, and validation that outputs remain bit-exact without the existing long build.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100