[llama] Generate causal mask better

[Groverkss]

Previously the causal mask was being generated as::

causal_mask = triu(context_len, context_len)
causal_mask = causal_mask[:batch_seqlen, :batch_seqlen]

This is not a good thing to do. This is simply making it harder for the compiler to fuse this "fill like" computation to a dispatch. The slicing with a dynamic dimension is harder to move around.

Instead, it's a better idea to generate the attention mask as:

causal_mask = triu(batch_seqlen, batch_seqlen)

This PR also removes the ability to put the mask in a buffer. This is again a really bad idea. The causal mask computation should never actually materialize outside a dispatch. If the compiler fails to do this, we should see memory usage spikes and fix the compiler.

[Groverkss]

Turning into draft since I want to see if we can make the compiler propagate the slice.