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

Add Scaled Dot Product Attention for FP8 #22670

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

wenscarl
Copy link
Contributor

@wenscarl wenscarl commented Jul 25, 2024

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!

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

Yes. I will refactor changes to be in fused_attention_stablehlo.py

@wenscarl
Copy link
Contributor Author

wenscarl commented Oct 18, 2024

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 :)

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.

3 participants