jax
405c79b0 - [Mosaic GPU] Use warpgroup semantics for the FlashAttention 3 pipelined kernel.

Commit

1 year ago

[Mosaic GPU] Use warpgroup semantics for the FlashAttention 3 pipelined kernel. As a result we can remove all transforms and layout casts. PiperOrigin-RevId: 754897809

References

#28382 - [Mosaic GPU] Use warpgroup semantics for the FlashAttention 3 pipelined kernel.

Author

dimitar-asenov

Committer

Google-ML-Automation

Parents

63003bb5

jax 405c79b0 - [Mosaic GPU] Use warpgroup semantics for the FlashAttention 3 pipelined kernel.

jax
405c79b0 - [Mosaic GPU] Use warpgroup semantics for the FlashAttention 3 pipelined kernel.