Enable fast attention in nanoPET

[frostedoyster]

---

📚 Documentation preview 📚: https://metatrain--454.org.readthedocs.build/en/454/

[spozdn]

I think backward compatible proper implementation is on the way. We can discuss it soon.

upd. not in the backward compatible sense as later clarified, I meant supporting backward pass

[abmazitov]

Does it give the same result as the original version of the code? If so, I would be quite surprised

[abmazitov]

What do you mean by saying that it’s not backward compatible? I don’t see why backward pass should not work

[Luthaf]

What do you mean by saying that it’s not backward compatible? I don’t see why backward pass should not work

Backward compatible in the API-sense, i.e. that the same checkpoint will produce the same results without retraining.

[frostedoyster]

@spozdn I think that using a custom attn_mask here

attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)

will prevent torch from using flash attention

You'll find more details here, where they say that they only implement full (all True) and causal masks https://github.com/Dao-AILab/flash-attention

And here (https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html), where they say that
"All implementations are enabled by default. Scaled dot product attention attempts to automatically select the most optimal implementation based on the inputs."

Anyway, this can be used to check which one is used:
https://pytorch.org/docs/stable/generated/torch.nn.attention.sdpa_kernel.html#torch.nn.attention.sdpa_kernel

EDIT: Flash attention doesn't support a custom mask, but CUDNN attention does, so we'll go with it