[WIP] exp with SDPA instead of einsum

[ramithuh]

TL;DR: Looks like we can rely on SDPA even with the pair bias track. This will only benefit on longer sequences / more pair-biased attention layers.

Todos:

• see if performance improvement is tied to a specific pytorch version
• see which sequence lengths might benefit from this
• is self.proj_z(z) anyway dominating the runtime so thsi does not matter?

More details:

• When a dense pair bias is given to scaled_dot_product_attention it will fall to Efficient Attention, and it should be competitive with flashattention.

More info about the backends that exist in F.scaled_dot_product_attention can be found in the these slides

---

I ran two quick experiments to check if the implementations are providing the same answer wandb report comparing the two runs. I did not see a difference in training speed at this scale though :|

[ramithuh]

We need to optimize `self.proj_z(z)` because that's the bottleneck. SDPA / einsum debate takes only small fraction in the total time since `self.proj_z(z)` dominates:

some directions:

1. don't have different bias values for each head.
2. use flex attention to maybe fuse this computation

in general: training with +trainer.precision=bf16-mixed helps to speed up transformers