[Llama2] Optimize prefill performance

[cglagovich]

High level issue for prefill performance improvements.

• [x] 8x8 grid for mlp matmuls
• [x] block-sharded eltwise mul in mlp
• [x] 8x8 projection matmuls
• [x] 8x8 grid for rmsnorm, and sweep chunk sizes
• [x] 8x8 grid for attn ops, and sweep chunk sizes

[cglagovich]

single layer latency with 8x8 mlp matmuls: 29469281 ns

[cglagovich]

with block-sahrded eltwise mul: 29176066 ns

[cglagovich]

with 8x8 projection matmuls: 28052277 ns

[cglagovich]

with 8x8 rms @ 256 chunk size: 27283850 ns

[cglagovich]

with 8x8 rms @ 512 chunk size: 27160413 ns

[cglagovich]

1024 chunk size OOM

[cglagovich]

8x8 spda @ chunk size 256: 25084698 ns

[cglagovich]

sdpa chunk size 512 OOM

[cglagovich]

fyi @johanna-rock-tt