From c6477f06098b271a1725e37d3ce0a94c9e69d42a Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Fri, 30 Oct 2026 08:02:39 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- 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,