From 7cf132ece7441404842a3af0d9ce3567a17cc828 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Sat, 31 Oct 2026 00:31:02 -0700 Subject: [PATCH] [ET-VK][sdpa] Expose shader override Pull Request resolved: https://github.com/pytorch/executorch/pull/24343 ## Ulterior Motive Benchmark SDPA shader routes without duplicating the production operator. ## Rationale Expose `sdpa_impl_with_shader_override()` to assemble the full SDPA operator while selecting its attention-value shader. The registered operator delegates to this function with no override, preserving automatic routing and the existing QK/softmax/AV dispatch chain. ```cpp sdpa_impl_with_shader_override( graph, args, graph.add_scalar(kShaderOverrideForceTile2)); ``` The override is a graph scalar using the existing SDPA route constants. Validation is provided by SDPA JSON cases and route sweeps in the later `benchmark_v2` commit. Authored with Codex (AI). ghstack-source-id: 461636572 @exported-using-ghexport Differential Revision: [D128509576](https://our.internmc.facebook.com/intern/diff/D128509576/) --- backends/vulkan/runtime/graph/ops/impl/SDPA.cpp | 12 ++++++++++-- backends/vulkan/runtime/graph/ops/impl/SDPA.h | 5 +++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index c83a8a318c0..8a2c447f677 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -663,7 +663,10 @@ void update_cache_impl(ComputeGraph& graph, const std::vector& args) { add_sdpa_kv_cache_update_node(graph, input_pos_symint, value, cache); } -void sdpa_impl(ComputeGraph& graph, const std::vector& args) { +void sdpa_impl_with_shader_override( + ComputeGraph& graph, + const std::vector& args, + const ValueRef shader_override) { int arg_idx = 0; const ValueRef q_projected = args[arg_idx++]; const ValueRef k_cache = args[arg_idx++]; @@ -792,7 +795,12 @@ void sdpa_impl(ComputeGraph& graph, const std::vector& args) { /*k=*/kDummyValueRef, input_pos_symint, out, - SDPAMode::LLM); + SDPAMode::LLM, + shader_override); +} + +void sdpa_impl(ComputeGraph& graph, const std::vector& args) { + sdpa_impl_with_shader_override(graph, args, kDummyValueRef); } void sdpa_with_kv_cache_impl( diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.h b/backends/vulkan/runtime/graph/ops/impl/SDPA.h index fa51ebef111..f5324472d8a 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.h +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.h @@ -25,6 +25,11 @@ namespace vkcompute { // enum class SDPAMode { LLM, FUSED }; +void sdpa_impl_with_shader_override( + ComputeGraph& graph, + const std::vector& args, + ValueRef shader_override); + void add_sdpa_compute_attn_weights_node( ComputeGraph& graph, const ValueRef q,