dianaml0

@frank-wei yes the breaking test is unrelated, fixing that separately. Could you just make sure that you use the pre-commit hook for this (run `flake8` and `black`) before submitting?

Could you also update the PR with a link to the PR you're reverting and a quick description?

Yes that sounds like a good idea! Should we also have the full LRA run regularly as well not necessarily part of CI?

> hey @dianaml0 this is great but you would need a modern triton, right ? I've had[ a branch up](https://github.com/blefaudeux/xformers/tree/triton_update) on that for a while, some of the API changed...

Updated Numbers Performance Compared to Vanilla FWD ``` [--------- attention (attn_bias=) --------] | optimized | vanilla 1 threads: -------------------------------------------------- f16 B=384, M=197, H=1, K=88 | 1258.8 | 372.9 b16 B=384,...

Thanks a lot for the reviews @fmassa and @danthe3rd ! I've made some changes and added an op for Triton fwd with Flash bwd

Forwards for Triton fwd and Flash bwd: ``` [--------- attention (attn_bias=) --------] | optimized | eager 1 threads: -------------------------------------------------- f16 B=384, M=197, H=1, K=64 | 91.7 | 292.3 b16 B=384,...

@fmassa Thanks a lot for the helpful comments and for taking another pass! I've updated with related changes. Okay to merge for now?