[Vulkan] Add the off-graph flat KV cache and update_and_attend op

[kiymetakdemir]

Adds the off-graph KV cache to the Vulkan backend, behind the neutral kvcache::update_and_attend op. The cache is runtime state the host installs before the graph is built; it owns one buffer pool per layer for K and one for V, and the op function wraps them using the external-storage path. Pools are buffers.

Added the missing shader variants — [texture3d, buffer] for the SDPA shaders and [buffer, texture3d] for the cache update so a buffer pool works inside the default texture-storage graph.

sdpa_impl asserted scale was None and then hardcoded 1/sqrt(head_dim). The op contract makes scale a required float, so it now uses the supplied value when there is one.

Files

• backends/vulkan/runtime/graph/VulkanCache.h — graph-facing interface: pool shape, dtype, layer count, pool buffers.
• backends/vulkan/runtime/graph/VulkanSequenceCache.h — the byte layer; allocates pools at construction from the global context. Flat layers only.
• backends/vulkan/runtime/graph/ComputeGraph.h — the installed cache, set before the graph is built.
• backends/vulkan/runtime/graph/ops/impl/SDPA.cpp — update_and_attend_impl, and scale in sdpa_impl.
• backends/vulkan/runtime/graph/ops/glsl/sdpa_*.{glsl,yaml} — the missing storage-type variants.
• backends/vulkan/test/op_tests/sdpa_test.cpp — OFFGRAPH mode.

Testing

vulkan_sdpa_test — OFFGRAPH runs through the existing driver, so it covers all four configurations and both storage types: prefill, single-token decode, mid-sequence chunks, four GQA ratios and four head dims, each step compared against the ATen reference.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 6 Unrelated Failures

As of commit 275e70d with merge base 2c88f57 ():

NEW FAILURE - The following job has failed:

• Test Vulkan Backend (specialized runners) / test-vulkan-nvidia / linux-job (gh)
  undefined reference to testing::internal::PrintBytesInObjectTo(unsigned char const*, unsigned long, std::ostream*)'`

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

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

• pull / unittest / linux / linux-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest / macos / macos-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest / windows / windows-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest-editable / linux / linux-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest-editable / macos / macos-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions

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.