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.

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

Enable fwd and varlen_fwd on AMD #60

Merged
merged 4 commits into from
Jun 19, 2024
Merged

Conversation

micmelesse
Copy link
Collaborator

@micmelesse micmelesse commented Jun 5, 2024

Enable flash_attn_func using amd's triton flash attention kernel.

This enables flash_attn_func. There will be follow up prs to enable the forward pass of the other flash attention apis.

we skip building the cuda kernels and use the triton kernel on AMD. This is fine for our fork but upstream will probably require a way to have one wheel.

@micmelesse micmelesse marked this pull request as ready for review June 7, 2024 16:12
@micmelesse micmelesse changed the title Enable AMD Trition Fwd Kernel Enable flash_attn_func on AMD Jun 7, 2024
@micmelesse micmelesse requested a review from vgokhale June 7, 2024 16:14
Compress

This is a combination of 12 commits.

add scripts

save

add our kernel

import our kernel

round trip

use bshd layout

figure out segfault

fix

show backward failure with prints

save backward work

run forward only

test smallest config on everything

add test

fix

remove pre commit

install triton

skip dropout

pin d

32 factor d

just run power of 2

remove timeout

run serially

clean up

clean up 2
@micmelesse micmelesse force-pushed the micmelesse/enable_fwd branch from 74654c7 to 2e2cc43 Compare June 7, 2024 16:17
This is a combination of 6 commits.

save

some tests passing

enable more

enable everything

move around

alibi works
@micmelesse micmelesse changed the title Enable flash_attn_func on AMD Enable fwd and varlen_fwd on AMD Jun 12, 2024
@micmelesse
Copy link
Collaborator Author

@vgokhale Can I merge this?

@micmelesse micmelesse merged commit 8fdb1bc into main_perf Jun 19, 2024
2 checks passed
@micmelesse micmelesse deleted the micmelesse/enable_fwd branch August 5, 2024 20:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants