Refactor trimat

[gordicaleksa]

Made sure we're consistent in our notation:

• Use (B,T,NH,HS) only
• Constant are upper case

Added additional comments to clarify what each of the kernels is doing, including what the indexing operation is doing, e.g.:
annotate flat indexing: const float* q = inp + b * T * C3 + nh * HS;
with Q[b][nh][:][:] = inp[b][:][0][nh][:]

[gordicaleksa]

Additionally I've moved us to the more familiar Q @ K^T form (from previous K^T @ Q).

This is how the attention is always written down in papers and i had a mental overhead parsing the code due to this permutation. :D I made sure there are no differences in terms of perf (they're completely equivalent).