Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Feature：Support flashMLA decoding via flashAttn2 #36

Closed
wants to merge 1 commit into from
Closed

Feature：Support flashMLA decoding via flashAttn2 #36

wants to merge 1 commit into from

Conversation

zhangjinnan
Copy link

@zhangjinnan zhangjinnan commented Feb 24, 2025 •
edited
Loading

Changes: (#29)

  1. Implement flashMLA with matrix absorption algorithm via flashAttn2
  2. Add golden test on MXMACA platform

FlashAttention-2 currently supports:

  1. Datatype fp16 and bf16.
  2. Multi-Token Parallelism = 1
  3. Paged kvcache with block size equal to 2^n (n >= 0)

@zhangjinnan
Feature：Support flashMLA decoding via flashAttn2(#29)
e0557de 
Changes:
1. Implement flashMLA with matrix absorption algorithm via flashAttn2
2. Add golden test on MXMACA platform
@zhangjinnan zhangjinnan marked this pull request as draft February 24, 2025 16:17
@zhangjinnan zhangjinnan marked this pull request as ready for review February 24, 2025 16:18
@zhangjinnan zhangjinnan marked this pull request as draft February 24, 2025 16:18
@zhangjinnan zhangjinnan marked this pull request as ready for review February 24, 2025 16:23
@yiakwy-xpu-ml-framework-team
Copy link

Will it be better to move Flashinfer related function from vendors to flashinfer repo ?

@zhangjinnan
Copy link
Author

https://github.com/[MetaX-MACA/FlashMLA](https://github.com/MetaX-MACA/FlashMLA)

@sanggusti
Copy link

I think cutlass that is called via submodules is handling what flashAttn2 functionalities that you mentioned

@zhangjinnan
Copy link
Author

I think cutlass that is called via submodules is handling what flashAttn2 functionalities that you mentioned

you can watch update in https://github.com/MetaX-MACA/FlashMLA repo.

@zhangjinnan zhangjinnan mentioned this pull request Feb 25, 2025
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@zhangjinnan @yiakwy-xpu-ml-framework-team @sanggusti