QKV calculation improvements in attention mechanism

[benraha]

Motivation and Context

This PR has two changes:

1. Replace costly einsum with simple matrix multiplication - this version is faster in both PyTorch and ONNX. IMO, we do miss some of the clarity, but since this happens in every forward pass, it is worth it.
2. Add a fast-path for the broadcast_kv_across_heads function - since the use case where we don't have to change anything happens 50% of the time, this simple fix has great performance benefits.

---

Public API Changes

• No Public API changes
• Yes, Public API changes (Details below)

---

How Has This Been Tested?

Locally.

---

Checklist

• The changes have been tested locally.
• Documentation has been updated (if the public API or usage changes).
• A entry has been added to CHANGELOG.md (if relevant for users).
• The code follows the project's style guidelines.
• I have considered the impact of these changes on the public API.

---

[gemini-code-assist]

Code Review

This pull request introduces two performance optimizations to the attention mechanism. The first replaces a costly einsum operation with a more efficient matrix multiplication using torch.mm, and the second adds a fast-path to broadcast_kv_across_heads for a common use case. Both changes are logical and aim to improve performance. My review includes a suggestion to further refine the matrix multiplication implementation for better readability and to use a more idiomatic PyTorch function.

[benraha]

@priorphil can you please review? :)

[priorphil]

Thanks! Could you share some of the benchmarks (including hardware, datasets shapes, dtype, timings) you ran for each of these changes so I can get a feeling for the magnitude of the speedup? :)

[benraha]

Sure! I wrap the model in ONNX, and I measure a 15% reduction in the number of nodes in the graph and ~10% faster inference.

I work on CPU inference, with a fit of 1000 rows, running inference on a few lines at a time.

[priorphil]

LGTM, thanks for the changes.