From 275e70d9346c3207929f5d92cd417dc7e24e1786 Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Wed, 19 Aug 2026 15:41:50 -0700 Subject: [PATCH] [Vulkan] Add the off-graph KV cache and update_and_attend op --- backends/vulkan/runtime/graph/ComputeGraph.h | 12 ++ backends/vulkan/runtime/graph/VulkanCache.h | 40 ++++++ .../runtime/graph/VulkanSequenceCache.h | 122 ++++++++++++++++++ .../glsl/sdpa_compute_attn_weights_coop.glsl | 1 + .../glsl/sdpa_compute_attn_weights_coop.yaml | 1 + .../glsl/sdpa_compute_attn_weights_tiled.glsl | 1 + .../glsl/sdpa_compute_attn_weights_tiled.yaml | 1 + .../graph/ops/glsl/sdpa_compute_out_coop.glsl | 1 + .../graph/ops/glsl/sdpa_compute_out_coop.yaml | 1 + .../ops/glsl/sdpa_compute_out_tiled.glsl | 1 + .../ops/glsl/sdpa_compute_out_tiled.yaml | 1 + .../graph/ops/glsl/sdpa_kv_cache_update.glsl | 4 +- .../graph/ops/glsl/sdpa_kv_cache_update.yaml | 1 + .../vulkan/runtime/graph/ops/impl/SDPA.cpp | 60 ++++++++- backends/vulkan/test/op_tests/sdpa_test.cpp | 42 +++++- 15 files changed, 284 insertions(+), 5 deletions(-) create mode 100644 backends/vulkan/runtime/graph/VulkanCache.h create mode 100644 backends/vulkan/runtime/graph/VulkanSequenceCache.h diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h index de85f13a89a..e7b6af2759c 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.h +++ b/backends/vulkan/runtime/graph/ComputeGraph.h @@ -17,6 +17,7 @@ #include #include +#include #include #include @@ -224,6 +225,9 @@ class ComputeGraph final { // Flag to indicate if re-encoding is required bool requires_reencode_ = false; + // Off-graph KV cache, installed by the backend before the graph is built. + VulkanCache* kv_cache_ = nullptr; + protected: size_t values_in_use_ = 0; size_t execute_count_ = 0; @@ -1180,6 +1184,14 @@ class ComputeGraph final { requires_reencode_ = true; } + inline void set_kv_cache(VulkanCache* cache) noexcept { + kv_cache_ = cache; + } + + inline VulkanCache* kv_cache() const noexcept { + return kv_cache_; + } + // // Miscellaneous Utilities // diff --git a/backends/vulkan/runtime/graph/VulkanCache.h b/backends/vulkan/runtime/graph/VulkanCache.h new file mode 100644 index 00000000000..223ac971185 --- /dev/null +++ b/backends/vulkan/runtime/graph/VulkanCache.h @@ -0,0 +1,40 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +#include +#include + +namespace vkcompute { + +// Graph-facing view of the off-graph KV cache. The host installs one before the +// graph is built; the op function reads each layer's pool shape and wraps the +// pool buffer. Pools are [B, S, H, D], the layout the SDPA shaders index. +class VulkanCache { + public: + virtual ~VulkanCache() = default; + + // Pool shape for `layer`, in this backend's layout. + virtual std::vector pool_sizes(int layer) const = 0; + + // Element type the pools store K/V in. + virtual vkapi::ScalarType pool_dtype() const = 0; + + // Layers the cache was built for; the op function bounds `layer` by this. + virtual int num_layers() const = 0; + + // The pools themselves, for the graph to wrap. + virtual const vkapi::VulkanBuffer& k_buffer(int layer) const = 0; + virtual const vkapi::VulkanBuffer& v_buffer(int layer) const = 0; +}; + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/VulkanSequenceCache.h b/backends/vulkan/runtime/graph/VulkanSequenceCache.h new file mode 100644 index 00000000000..1435ff30c64 --- /dev/null +++ b/backends/vulkan/runtime/graph/VulkanSequenceCache.h @@ -0,0 +1,122 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// The Vulkan byte layer behind the neutral SequenceCache. Holds one buffer +// pool per layer for K and one for V, allocated at construction. +// +// Flat layers only: every dispatch here assumes slot == position, which a ring +// layer breaks. + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace vkcompute { + +namespace cache = ::executorch::extension::llm::cache; +namespace runtime = ::executorch::runtime; + +class VulkanSequenceCache : public cache::SequenceCache, public VulkanCache { + public: + static runtime::Result> create( + const cache::CacheConfig& cfg) { + ET_CHECK_OR_RETURN_ERROR( + cache::valid(cfg), + InvalidArgument, + "VulkanSequenceCache: invalid config"); + for (const cache::LayerConfig& lc : cfg.layers) { + ET_CHECK_OR_RETURN_ERROR( + lc.policy.kind == cache::LayerPolicy::Kind::Flat, + NotSupported, + "VulkanSequenceCache: only flat layers are supported"); + ET_CHECK_OR_RETURN_ERROR( + lc.n_kv_heads > 0 && lc.head_dim > 0, + InvalidArgument, + "VulkanSequenceCache: n_kv_heads and head_dim must be positive"); + } + // The SDPA shaders are only generated for fp32 and fp16. + using EtScalarType = ::executorch::runtime::etensor::ScalarType; + vkapi::ScalarType pool_dtype; + switch (static_cast(cfg.kv_dtype)) { + case EtScalarType::Float: + pool_dtype = vkapi::kFloat; + break; + case EtScalarType::Half: + pool_dtype = vkapi::kHalf; + break; + default: + ET_LOG(Error, "VulkanSequenceCache: unsupported kv_dtype"); + return runtime::Error::NotSupported; + } + return std::unique_ptr( + new VulkanSequenceCache(cfg, pool_dtype)); + } + + // [B, S, H, D], the layout the SDPA shaders index. S is the allocated depth. + std::vector pool_sizes(int layer) const override { + const cache::LayerConfig& lc = layers_[static_cast(layer)]; + return {1, capacity(), lc.n_kv_heads, lc.head_dim}; + } + + vkapi::ScalarType pool_dtype() const override { + return pool_dtype_; + } + + int num_layers() const override { + return static_cast(layers_.size()); + } + + const vkapi::VulkanBuffer& k_buffer(int layer) const override { + return kpool_[static_cast(layer)].buffer(); + } + const vkapi::VulkanBuffer& v_buffer(int layer) const override { + return vpool_[static_cast(layer)].buffer(); + } + + private: + VulkanSequenceCache( + const cache::CacheConfig& cfg, + const vkapi::ScalarType pool_dtype) + : cache::SequenceCache(cfg), pool_dtype_(pool_dtype) { + layers_.reserve(static_cast(cfg.n_layers)); + for (int l = 0; l < cfg.n_layers; ++l) { + // layers size 1 = one config broadcast to every layer, else per-layer. + layers_.push_back( + cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l]); + } + // The global context outlives every graph. + kpool_.reserve(static_cast(cfg.n_layers)); + vpool_.reserve(static_cast(cfg.n_layers)); + for (int l = 0; l < cfg.n_layers; ++l) { + for (auto* pool : {&kpool_, &vpool_}) { + pool->emplace_back( + api::context(), + pool_sizes(l), + pool_dtype_, + utils::kBuffer, + utils::kWidthPacked); + } + } + } + + std::vector layers_; + vkapi::ScalarType pool_dtype_; + std::vector kpool_; + std::vector vpool_; +}; + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl index b7f14f435fa..967450d0c8b 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl @@ -32,6 +32,7 @@ $if K_CACHE_STORAGE == "buffer": #define NUM_WORKERS_PER_OUT 64 ${define_required_extensions(IO_STORAGE, DTYPE)} +${define_required_extensions(K_CACHE_STORAGE, DTYPE)} layout(std430) buffer; diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml index d5cadc36060..a61643c2187 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml @@ -18,6 +18,7 @@ sdpa_compute_attn_weights_coop: - parameter_values: [texture3d, texture3d] - parameter_values: [buffer, texture3d] - parameter_values: [buffer, buffer] + - parameter_values: [texture3d, buffer] DTYPE: - VALUE: float - VALUE: half diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl index 662d5edde68..6e6432d0909 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl @@ -46,6 +46,7 @@ $if HAS_BIAS: #define TILE_N ${TILE_N4 * 4} ${define_required_extensions(IO_STORAGE, [IN_DTYPE, OUT_DTYPE])} +${define_required_extensions(K_CACHE_STORAGE, [IN_DTYPE, OUT_DTYPE])} layout(std430) buffer; diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.yaml index 24494b408fa..7ffed75a311 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.yaml @@ -22,6 +22,7 @@ sdpa_compute_attn_weights_tiled: - parameter_values: [texture3d, texture3d] - parameter_values: [buffer, texture3d] - parameter_values: [buffer, buffer] + - parameter_values: [texture3d, buffer] combination1: parameter_names: [IN_DTYPE, OUT_DTYPE] combos: diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl index f6f00c9bfe5..a75e5ba1592 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl @@ -43,6 +43,7 @@ $if GQA: #define MAX_GROUP_SIZE 8 ${define_required_extensions(IO_STORAGE, DTYPE)} +${define_required_extensions(V_CACHE_STORAGE, DTYPE)} layout(std430) buffer; diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml index 546eb9da7c0..dc4745f5386 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml @@ -19,6 +19,7 @@ sdpa_compute_out_coop: - parameter_values: [texture3d, texture3d] - parameter_values: [buffer, texture3d] - parameter_values: [buffer, buffer] + - parameter_values: [texture3d, buffer] DTYPE: - VALUE: float - VALUE: half diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.glsl index 9f8f2dbc231..3a9caf92d98 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.glsl @@ -38,6 +38,7 @@ $else: #define TILE_N ${TILE_N4 * 4} ${define_required_extensions(IO_STORAGE, DTYPE)} +${define_required_extensions(V_CACHE_STORAGE, DTYPE)} layout(std430) buffer; diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.yaml index ba91114ae92..abdbf226d55 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.yaml @@ -20,6 +20,7 @@ sdpa_compute_out_tiled: - parameter_values: [texture3d, texture3d] - parameter_values: [buffer, texture3d] - parameter_values: [buffer, buffer] + - parameter_values: [texture3d, buffer] DTYPE: - VALUE: float - VALUE: half diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.glsl index a7d94d5309c..004e290b5a6 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.glsl @@ -3,6 +3,7 @@ #define PRECISION ${PRECISION} #define IN_VEC4_T ${texel_load_type(DTYPE, INPUT_STORAGE)} +#define OUT_VEC4_T ${texel_load_type(DTYPE, OUTPUT_STORAGE)} #define T ${buffer_scalar_type(DTYPE)} $if OUTPUT_STORAGE == "buffer": @@ -11,6 +12,7 @@ $if INPUT_STORAGE == "buffer": #define INPUT_BUFFER ${define_required_extensions(INPUT_STORAGE, DTYPE)} +${define_required_extensions(OUTPUT_STORAGE, DTYPE)} layout(std430) buffer; @@ -65,7 +67,7 @@ void write_cache_d4( const int C, const int H) { #ifdef OUTPUT_BUFFER - t_cache[(c * H * D4) + (h * D4) + d4] = texel; + t_cache[(c * H * D4) + (h * D4) + d4] = OUT_VEC4_T(texel); #else imageStore(t_cache, ivec3(d4, h, c), texel); #endif diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.yaml index 5ec2f3e190c..351c3628b31 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.yaml @@ -16,6 +16,7 @@ sdpa_kv_cache_update: - parameter_values: [texture3d, texture3d] - parameter_values: [texture3d, buffer] - parameter_values: [buffer, buffer] + - parameter_values: [buffer, texture3d] DTYPE: - VALUE: half - VALUE: float diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index c4259d4c530..f75650ccda6 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -705,7 +705,6 @@ void sdpa_impl(ComputeGraph& graph, const std::vector& args) { VK_CHECK_COND( graph.val_is_none(dropout_p) || graph.extract_scalar(dropout_p) == 0); - VK_CHECK_COND(graph.val_is_none(scale)); // is_causal is assumed to be true in the current implementation. VK_CHECK_COND( graph.val_is_none(is_causal) || graph.extract_scalar(is_causal)); @@ -772,7 +771,9 @@ void sdpa_impl(ComputeGraph& graph, const std::vector& args) { utils::kWidthPacked); const int32_t head_dim_size = graph.size_at(-1, q_projected); - const float scale_val = 1.0f / std::sqrt(static_cast(head_dim_size)); + const float scale_val = graph.val_is_none(scale) + ? 1.0f / std::sqrt(static_cast(head_dim_size)) + : static_cast(graph.extract_scalar(scale)); add_sdpa_compute_attn_weights_node( graph, @@ -847,6 +848,60 @@ void sdpa_with_kv_cache_impl( out}); } +// Writes this step's K/V into the installed cache's pools and attends over +// them. Pool shape and dtype come from the cache, so the graph carries neither. +void update_and_attend_impl( + ComputeGraph& graph, + const std::vector& args) { + int arg_idx = 0; + const ValueRef q_projected = args[arg_idx++]; + const ValueRef k_projected = args[arg_idx++]; + const ValueRef v_projected = args[arg_idx++]; + const ValueRef input_pos_symint = args[arg_idx++]; + const ValueRef layer_id = args[arg_idx++]; + const ValueRef scale = args[arg_idx++]; + // `out` already carries this dtype; unpacked to keep the arg positions. + const ValueRef out_dtype = args[arg_idx++]; + const ValueRef out = args[arg_idx++]; + (void)out_dtype; + + VulkanCache* cache = graph.kv_cache(); + VK_CHECK_COND(cache != nullptr, "update_and_attend: no KV cache installed"); + + const int layer = graph.extract_scalar(layer_id); + VK_CHECK_COND( + layer >= 0 && layer < cache->num_layers(), + "update_and_attend: layer out of range"); + + const std::vector cache_sizes = cache->pool_sizes(layer); + + const ValueRef k_cache = graph.add_tensor( + cache_sizes, + cache->pool_dtype(), + utils::kWidthPacked, + cache->k_buffer(layer)); + const ValueRef v_cache = graph.add_tensor( + cache_sizes, + cache->pool_dtype(), + utils::kWidthPacked, + cache->v_buffer(layer)); + + update_cache_impl(graph, {k_projected, k_cache, input_pos_symint, -1}); + update_cache_impl(graph, {v_projected, v_cache, input_pos_symint, -1}); + + sdpa_impl( + graph, + {q_projected, + k_cache, + v_cache, + input_pos_symint, + kDummyValueRef, // attn_mask: LLM mode derives causality from input_pos + kDummyValueRef, // dropout_p + kDummyValueRef, // is_causal + scale, + out}); +} + void compute_attn_weight_with_kv_cache_impl( ComputeGraph& graph, const std::vector& args) { @@ -1006,6 +1061,7 @@ REGISTER_OPERATORS { testing.compute_attn_weight_with_kv_cache.default, compute_attn_weight_with_kv_cache_impl); VK_REGISTER_OP(et_vk.sdpa.default, fused_sdpa_impl); + VK_REGISTER_OP(kvcache.update_and_attend.default, update_and_attend_impl); } } // namespace vkcompute diff --git a/backends/vulkan/test/op_tests/sdpa_test.cpp b/backends/vulkan/test/op_tests/sdpa_test.cpp index c1c252165eb..e9a0cda9b40 100644 --- a/backends/vulkan/test/op_tests/sdpa_test.cpp +++ b/backends/vulkan/test/op_tests/sdpa_test.cpp @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -27,7 +28,7 @@ // SDPA Mode Enum // -enum class SDPAMode { DECOMPOSED, FUSED, ATTN_WEIGHT_ONLY }; +enum class SDPAMode { DECOMPOSED, FUSED, ATTN_WEIGHT_ONLY, OFFGRAPH }; std::ostream& operator<<(std::ostream& os, const SDPAMode& mode) { switch (mode) { @@ -37,6 +38,8 @@ std::ostream& operator<<(std::ostream& os, const SDPAMode& mode) { return os << "FUSED"; case SDPAMode::ATTN_WEIGHT_ONLY: return os << "ATTN_WEIGHT_ONLY"; + case SDPAMode::OFFGRAPH: + return os << "OFFGRAPH"; } return os; } @@ -327,9 +330,27 @@ void test_vulkan_sdpa( // Build Vulkan SDPA graph using namespace vkcompute; + std::unique_ptr offgraph_cache; + if (mode == SDPAMode::OFFGRAPH) { + cache::CacheConfig cfg; + cfg.capacity = max_seq_len + 128; + cfg.n_layers = 1; + cfg.layers = {cache::LayerConfig{{}, num_kv_heads, head_dim}}; + cfg.kv_dtype = static_cast( + dtype == at::kFloat ? ::executorch::runtime::etensor::ScalarType::Float + : ::executorch::runtime::etensor::ScalarType::Half); + auto created = VulkanSequenceCache::create(cfg); + ASSERT_EQ(created.error(), executorch::runtime::Error::Ok); + offgraph_cache = std::move(created.get()); + } + GraphConfig config; ComputeGraph graph(config); + if (offgraph_cache) { + graph.set_kv_cache(offgraph_cache.get()); + } + // "Data" variant for vulkan initialization at::Tensor k_cache_data = at::zeros_like(k_cache); @@ -358,6 +379,20 @@ void test_vulkan_sdpa( out.sizes().vec(), from_at_scalartype(out.scalar_type()), storage_type); switch (mode) { + case SDPAMode::OFFGRAPH: + VK_GET_OP_FN("kvcache.update_and_attend.default") + (graph, + { + r_q.value, + r_k.value, + r_v.value, + r_input_pos_symint, + graph.add_scalar(0), // layer_id + graph.add_scalar(1.0 / std::sqrt((double)head_dim)), // scale + kDummyValueRef, // out_dtype + r_out, + }); + break; case SDPAMode::DECOMPOSED: { const ValueRef r_k_cache = graph.add_tensor( k_cache_data.sizes().vec(), @@ -589,7 +624,10 @@ void test_vulkan_sdpa( const int batch_size, at::ScalarType dtype = at::kFloat) { for (SDPAMode mode : - {SDPAMode::ATTN_WEIGHT_ONLY, SDPAMode::DECOMPOSED, SDPAMode::FUSED}) { + {SDPAMode::ATTN_WEIGHT_ONLY, + SDPAMode::DECOMPOSED, + SDPAMode::FUSED, + SDPAMode::OFFGRAPH}) { // Test texture test_vulkan_sdpa( start_input_pos,