Lightning-AI / Lightning-AI/lightning-thunder
A mask conversion for cuDNN execution of scaled dot product attention should be fused
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
🚀 Feature
Currently, cuDNN doesn't support boolean masks for torch.nn.functional.scaled_dot_product_attention and this conversion is done eagerly at runtime:
https://github.com/Lightning-AI/lightning-thunder/blob/70c06666ee8847746f836d178aee541fe749c62f/thunder/executors/cudnnex.py#L286-L290
Instead execution_transform could be modified to do this conversion in Thunder so that it could be fused with previous operations in the trace
https://github.com/Lightning-AI/lightning-thunder/blob/70c06666ee8847746f836d178aee541fe749c62f/thunder/executors/cudnnex.py#L807
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 in thunder/executors/cudnnex.py at the referenced mask-conversion code and the execution_transform entry point. Trace how boolean masks for scaled dot product attention are converted for cuDNN, then verify that the conversion happens in Thunder and can fuse with preceding operations rather than occurring eagerly at runtime.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, machine-learning, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100