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

Add Scaled Dot Product Attention for FP8 #22670

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Add Scaled Dot Product Attention for FP8 #22670

wants to merge 4 commits into from

Conversation

wenscarl
Copy link
Contributor

@wenscarl wenscarl commented Jul 25, 2024 •
edited
Loading

Scaled Dot Product Attention for FP8. @kaixih @zhangqiaorjc

@wenscarl wenscarl marked this pull request as draft July 25, 2024 19:53
@wenscarl wenscarl marked this pull request as ready for review August 22, 2024 15:51
@kaixih
Copy link
Contributor

kaixih commented Aug 22, 2024

@wenscarl is this ready to review? I can help with that. If possible, can you resolve the conflicts and run the pre-commit for the format consistency?

@wenscarl
Copy link
Contributor Author

It's ready for review and conflict fixed. Thanks!

superbobry
superbobry requested changes Sep 17, 2024
View reviewed changes
Copy link
Collaborator

@superbobry superbobry left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like this PR forked jax/jax/_src/cudnn
/fused_attention_stablehlo.py and then modified it for FP8. Does that sound right? If yes, could we instead add support for FP8 to jax/jax/_src/cudnn
/fused_attention_stablehlo.py directly?

@wenscarl
Copy link
Contributor Author

wenscarl commented Sep 17, 2024 •
edited
Loading

Compact all changes into fused_attention_stablehlo.py

@wenscarl wenscarl mentioned this pull request Sep 30, 2024
Closed
@wenscarl wenscarl force-pushed the fp8_sdpa branch 4 times, most recently from db6e2b9 to 3ca52db Compare October 18, 2024 18:28
@wenscarl
Copy link
Contributor Author

wenscarl commented Oct 18, 2024 •
edited
Loading

It looks like this PR forked jax/jax/_src/cudnn /fused_attention_stablehlo.py and then modified it for FP8. Does that sound right? If yes, could we instead add support for FP8 to jax/jax/_src/cudnn /fused_attention_stablehlo.py directly?

Merged into fused_attention_stablehlo.py.
Gentle ping @superbobry :)

mattjj
mattjj approved these changes Oct 24, 2024
View reviewed changes
Copy link
Collaborator

@mattjj mattjj left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you!

@wenscarl wenscarl force-pushed the fp8_sdpa branch 2 times, most recently from e2dde8c to 2f5d884 Compare October 24, 2024 20:59
@wenscarl
Copy link
Contributor Author

wenscarl commented Oct 30, 2024 •
edited
Loading

gently ping @mattjj and @superbobry

@superbobry
Copy link
Collaborator

Hey @wenscarl, sorry for the silence. I'm a bit swamped this week, but I will try to look through the PR over the weekend or first thing next week.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@superbobry superbobry Awaiting requested review from superbobry

@yashk2810 yashk2810 Awaiting requested review from yashk2810

@mattjj mattjj Awaiting requested review from mattjj

Assignees

@superbobry superbobry

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
5 participants
@wenscarl @kaixih @superbobry @mattjj @yashk2810