trifast 0.1.2

Fast kernel for triangle self attetion.

Fused Triangle Self Attention kernel, written in triton. Basically flash attention, but for triangle self attention. Implementation heavily inspired by FlagAttention and the triton fused attention tutorial.

• n^2 memory complexity (vs n^3 for pure pytorch).
• Faster (~2x) backward pass than next fastest implementation I could find (DS4S evoformer kernel).
• Faster (~4x) forward pass than next fastest implementation I could find (DS4S evoformer kernel).
• As far as I can tell, faster than naieve implementation.

Plots

All done on a 3090 in bfloat16.

Forward

Backward

Todos:

• [] Try to train a model with it.
• [] Can we perform and of dq/db/dkv transposed?
• [] Rewrite autotuner