[ET-VK][ez] Ensure that attn_weight buffers do not exceed GPU buffer numel limit

[SS-JIA]

Stack from ghstack (oldest at bottom):

• [ET-VK] buffer implementation of rotary positional embeddings #15620
• -> [ET-VK][ez] Ensure that attn_weight buffers do not exceed GPU buffer numel limit #15651
• [ET-VK] Implementation of to_dim_order_copy #15619

Title says it all!

To give a concrete example, Llama3.2-1B-Instruct will have attn weights with size {1, 32, max_seq_len, max_context_len}. Usually max_seq_len == max_context_len, and if max_context_len = 2048 Then the attention weight tensors will have sizes {1, 32, 2048, 2048} which will contain 134217728 elements. The maxStorageBufferRange for Adreno 750 is also 134217728 (2^27), so using context length of 2048 will produce incorrect results on Adreno 750.

In practice, it is unlikely that the prompt sequence length will be equal to the context length, so the solution is to adjust down the max_seq_len dim of the attention weight tensors to ensure that the GPU buffer numel limit is not hit.

Differential Revision: D86443407

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/15651

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (5 Unrelated Failures)

As of commit 5250269 with merge base d361573 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / unittest / windows / windows-job (gh) (detected as infra flaky with no log or failing log classifier)

BROKEN TRUNK - The following jobs failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-binary-size-linux-gcc / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-openvino-linux / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-setup-linux-gcc / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  backends/xnnpack/test/ops/test_conv1d.py::TestConv1d::test_qs8_conv1d_batchnorm_seq

This comment was automatically generated by Dr. CI and updates every 15 minutes.