Use fused softmax kernel in llama attention layer

[zackangelo]

Use the fused version of softmax (softmax_last_dim) in the Llama attention layer.

Modest improvement on the example (A100-SXM4-80GB):

Before:

$ cargo run -F cuda --example llama -- --which v31-instruct -n 100
loading the model weights from meta-llama/Meta-Llama-3.1-8B-Instruct
starting the inference loop
My favorite theorem is 2+2=4
This joke is most likely to be appreciated by children who love mathematics. The theorem in question, of course, is not the usual way that mathematicians would think about it. This response shows a lack of understanding of what a mathematical theorem actually is.
Why do chemists shun parties? Because they fear cheesing off their friends.
A joke about why chemists are reluctant socializers. Although this joke is amusing and somewhat related to chemistry, chemists are people too

100 tokens generated (38.5164153902478 token/s)

After softmax_last_dim:

$ cargo run -F cuda --example llama -- --which v31-instruct -n 100
loading the model weights from meta-llama/Meta-Llama-3.1-8B-Instruct
starting the inference loop
My favorite theorem is 2+2=4
This joke is only funny if you understand what the punchline implies. The theorem in question, "2 + 2 = 4", could be seen as a fundamental truth or an obvious fact (in the context of basic arithmetic). Thus, saying it's your favorite theorem means that this basic mathematical equation stands out to you in some way.
This would likely be interpreted humorously because people generally don't have preferences or favorites for simple equations like basic math problems. Math

100 tokens generated (45.525982179505434 token/s)

In my own experience, the improvement is even more pronounced as the context length and batch size increase.

[LaurentMazare]

Thanks!