pytorch · kiymetakdemir · Aug 19, 2026
@@ -17,6 +17,7 @@
 #include <executorch/backends/vulkan/runtime/api/api.h>
 
 #include <executorch/backends/vulkan/runtime/graph/GraphConfig.h>
+#include <executorch/backends/vulkan/runtime/graph/VulkanCache.h>
 
 #include <executorch/backends/vulkan/runtime/graph/containers/SharedObject.h>
 #include <executorch/backends/vulkan/runtime/graph/containers/Value.h>
@@ -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
   //

@@ -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 <cstdint>
+#include <vector>
+
+#include <executorch/backends/vulkan/runtime/vk_api/Types.h>
+#include <executorch/backends/vulkan/runtime/vk_api/memory/Buffer.h>
+
+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<int64_t> 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
@@ -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 <memory>
+#include <vector>
+
+#include <executorch/backends/vulkan/runtime/api/containers/Tensor.h>
+#include <executorch/backends/vulkan/runtime/graph/VulkanCache.h>
+#include <executorch/extension/llm/cache/cache.h>
+#include <executorch/extension/llm/cache/sequence_cache.h>
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/portable_type/scalar_type.h>
+#include <executorch/runtime/core/result.h>
+
+namespace vkcompute {
+
+namespace cache = ::executorch::extension::llm::cache;
+namespace runtime = ::executorch::runtime;
+
+class VulkanSequenceCache : public cache::SequenceCache, public VulkanCache {
+ public:
+  static runtime::Result<std::unique_ptr<VulkanSequenceCache>> 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<EtScalarType>(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<VulkanSequenceCache>(
+        new VulkanSequenceCache(cfg, pool_dtype));
+  }
+
+  // [B, S, H, D], the layout the SDPA shaders index. S is the allocated depth.
+  std::vector<int64_t> pool_sizes(int layer) const override {
+    const cache::LayerConfig& lc = layers_[static_cast<size_t>(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<int>(layers_.size());
+  }
+
+  const vkapi::VulkanBuffer& k_buffer(int layer) const override {
+    return kpool_[static_cast<size_t>(layer)].buffer();
+  }
+  const vkapi::VulkanBuffer& v_buffer(int layer) const override {
+    return vpool_[static_cast<size_t>(layer)].buffer();
+  }
+
+ private:
+  VulkanSequenceCache(
+      const cache::CacheConfig& cfg,
+      const vkapi::ScalarType pool_dtype)
+      : cache::SequenceCache(cfg), pool_dtype_(pool_dtype) {
+    layers_.reserve(static_cast<size_t>(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<size_t>(cfg.n_layers));
+    vpool_.reserve(static_cast<size_t>(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<cache::LayerConfig> layers_;
+  vkapi::ScalarType pool_dtype_;
+  std::vector<api::vTensor> kpool_;
+  std::vector<api::vTensor> vpool_;
+};
+
+} // namespace vkcompute
@@ -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;
 

@@ -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

@@ -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;
 

@@ -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:

@@ -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;
 

@@ -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

@@ -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;
 

@@ -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

@@ -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

@@ -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

@@ -705,7 +705,6 @@ void sdpa_impl(ComputeGraph& graph, const std::vector<ValueRef>& args) {
   VK_CHECK_COND(
       graph.val_is_none(dropout_p) ||
       graph.extract_scalar<double>(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<bool>(is_causal));
@@ -772,7 +771,9 @@ void sdpa_impl(ComputeGraph& graph, const std::vector<ValueRef>& args) {
       utils::kWidthPacked);
 
   const int32_t head_dim_size = graph.size_at<int32_t>(-1, q_projected);
-  const float scale_val = 1.0f / std::sqrt(static_cast<float>(head_dim_size));
+  const float scale_val = graph.val_is_none(scale)
+      ? 1.0f / std::sqrt(static_cast<float>(head_dim_size))
+      : static_cast<float>(graph.extract_scalar<double>(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<ValueRef>& 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<int>(layer_id);
+  VK_CHECK_COND(
+      layer >= 0 && layer < cache->num_layers(),
+      "update_and_attend: layer out of range");
+
+  const std::vector<int64_t> 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<ValueRef>& 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