Dynamic flash attention
Open
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 333
- Forks
- 150
- Avg merge
- 4d 19h
- Merged PRs (30d)
- 54
Description
Here is what we need to do to get the dynamic flash attention working.
- Lower dynamic flash attention to ref operators. This will help ensure that migraphx can handle such dynamic shapes in the model correctly.
- Chunk the flash attention into 32 chunks. We will need to insert padding to get this working. With this we can resue the current implementations(ie CK, MLIR). It might be faster as well due to better tuning, but we will need to test this though.
- Make a flash attention callable library that uses CK underneath so we can call into it from migraphx.
Step 3 can be done in parallel with step 1 and 2.
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
No files or tests are named. Start by tracing dynamic flash attention through MIGraphX and reviewing the existing CK and MLIR implementations; the work is complete when lowering, 32-chunk padding, and a callable CK-based library are implemented and tested for dynamic shapes.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100