jax Add Scaled Dot Product Attention for FP8

Scaled Dot Product Attention for FP8. @kaixih @zhangqiaorjc

Jul 25 '24 19:07 wenscarl

@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?

Aug 22 '24 16:08 kaixih

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

Aug 25 '24 13:08 wenscarl

Compact all changes into fused_attention_stablehlo.py

Sep 17 '24 18:09 wenscarl

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

Oct 18 '24 18:10 wenscarl

gently ping @mattjj and @superbobry

Oct 30 '24 19:10 wenscarl

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.

Nov 06 '24 21:11 superbobry

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. Gently ping @superbobry . Thanks!

Dec 03 '24 04:12 wenscarl

Thanks, can you squash the commits please?

Dec 13 '24 14:12 superbobry

Thanks, can you squash the commits please?

Done.

Dec 13 '24 18:12 wenscarl