[ET-VK][ez] Fuse update_cache + custom_sdpa into sdpa_with_kv_cache

[SS-JIA]

Stack from ghstack (oldest at bottom):

• [ET-VK] buffer implementation of rotary positional embeddings #15620
• [ET-VK] Implementation of to_dim_order_copy #15619
• -> [ET-VK][ez] Fuse update_cache + custom_sdpa into sdpa_with_kv_cache #15618
• [ET-VK] Implement select_at_dim_as_symint #15617
• [ET-VK][ez] Add lint rule to check no debug mode in shaders + clean up debug mode usage #15616

SDPA used to be handled by a custom op sdpa_with_kv_cache, but it was eventually split (D62301837) into update_cache and custom_sdpa ops.

However, having a single fused op is useful for Vulkan since it allows more control over how the cache tensors are stored and represented. Essentially, it makes it easier to manage the cache tensors and opens up opportunities for future optimizations. This diff introduces a fusion pass that does 2 things:

1. Combine update_cache and custom_sdpa back into sdpa_with_kv_cache
2. Ensure all references to the cache_pos symint use the same node - this prevents the select_at_dim_as_symint op from being called every time it is used.

Differential Revision: D86340339

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 Cancelled Job, 5 Unrelated Failures

As of commit 9bdd6dc with merge base 2b02316 ():

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / test-openvino-linux / linux-job (gh)

BROKEN TRUNK - The following jobs failed but were 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-setup-linux-gcc / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / unittest / windows / windows-job (gh) (trunk failure)
  backends/xnnpack/test/ops/test_conv1d.py::TestConv1d::test_qs8_conv1d_batchnorm_seq
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  backends/xnnpack/test/ops/test_conv1d.py::TestConv1d::test_qs8_conv1d_batchnorm_seq
• Test Metal Backend / export-voxtral-metal-artifact / macos-job (gh) (trunk failure)
  /Users/ec2-user/runner/_work/executorch/executorch/pytorch/executorch/pip-out/temp.macosx-11.0-arm64-cpython-311/cmake-out/configurations/optimized_native_cpu_ops_lib/NativeFunctions.h:72:228: error: no template named 'optional' in namespace 'torch::executor'; did you mean 'executorch::aten::optional'?

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

[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.