chore: import upstream snapshot with attribution

tensorflow/lite/experimental/genai/BUILD

@@ -0,0 +1,101 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+
+# Custom Ops useful for GenAI models.
+load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
+load("//tensorflow/lite:build_def.bzl", "tflite_copts")
+load("//tensorflow/lite:special_rules.bzl", "tflite_portable_test_suite")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
+
+cc_library(
+    name = "genai_ops",
+    srcs = [
+        "external_kvcache.cc",
+        "genai_ops.cc",
+        "kvcache.cc",
+        "sdpa.cc",
+    ],
+    hdrs = [
+        "genai_ops.h",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core:subgraph",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/experimental/resource",
+        "//tensorflow/lite/experimental/resource:cache_buffer",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/kernels/internal:common",
+        "//tensorflow/lite/kernels/internal:reference_base",
+        "//tensorflow/lite/kernels/internal:tensor",
+        "//tensorflow/lite/kernels/internal:types",
+        "@flatbuffers",
+    ],
+)
+
+cc_test(
+    name = "kvcache_test",
+    srcs = ["kvcache_test.cc"],
+    copts = tflite_copts(),
+    deps = [
+        ":genai_ops",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/kernels:test_util",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "external_kvcache_test",
+    srcs = ["external_kvcache_test.cc"],
+    copts = tflite_copts(),
+    deps = [
+        ":genai_ops",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/kernels:test_util",
+        "//tensorflow/lite/kernels/internal:tensor",
+        "//tensorflow/lite/kernels/internal:types",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+pybind_extension(
+    name = "pywrap_genai_ops",
+    srcs = [
+        "genai_ops_wrapper.cc",
+    ],
+    hdrs = ["genai_ops.h"],
+    additional_exported_symbols = ["GenAIOpsRegisterer"],
+    common_lib_packages = [
+        "litert/python",
+        "tensorflow/lite/python",
+    ],
+    enable_stub_generation = True,
+    link_in_framework = True,
+    module_name = "pywrap_genai_ops",
+    pytype_srcs = [
+        "pywrap_genai_ops.pyi",
+    ],
+    visibility = ["//visibility:public"],
+    wrap_py_init = True,
+    deps = [
+        ":genai_ops",
+        "//tensorflow/lite:framework_stable",
+        "//tensorflow/lite:mutable_op_resolver",
+        "//tensorflow/lite/c:common",
+        "@pybind11",
+        "@xla//third_party/python_runtime:headers",
+    ],
+)
+
+tflite_portable_test_suite()

tensorflow/lite/experimental/genai/external_kvcache.cc

@@ -0,0 +1,165 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <cstring>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+namespace llm {
+
+static const int kKeyTensor = 0;
+static const int kValueTensor = 1;
+static const int kPositionTensor = 2;
+static const int kKeySliceTensor = 3;
+static const int kValueSliceTensor = 4;
+
+static const int kRequiredNumDimensions = 4;
+
+TfLiteStatus ExternalKVCachePrepare(TfLiteContext* context, TfLiteNode* node) {
+  // k_cache, v_cache, position, k_slice, v_slice
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 5);
+  // updated: k_cache, v_cache
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 2);
+
+  const TfLiteTensor* k_cache;
+  const TfLiteTensor* v_cache;
+  const TfLiteTensor* position;
+  const TfLiteTensor* k_slice;
+  const TfLiteTensor* v_slice;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kKeyTensor, &k_cache));
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kValueTensor, &v_cache));
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kPositionTensor, &position));
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kKeySliceTensor, &k_slice));
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kValueSliceTensor, &v_slice));
+
+  TfLiteTensor* updated_k_cache;
+  TfLiteTensor* updated_v_cache;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kKeyTensor, &updated_k_cache));
+  TF_LITE_ENSURE_OK(
+      context, GetOutputSafe(context, node, kValueTensor, &updated_v_cache));
+
+  TF_LITE_ENSURE_EQ(context, k_cache->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_EQ(context, v_cache->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_EQ(context, position->type, kTfLiteInt32);
+  TF_LITE_ENSURE_EQ(context, k_slice->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_EQ(context, v_slice->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_EQ(context, updated_k_cache->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_EQ(context, updated_v_cache->type, kTfLiteFloat32);
+
+  TF_LITE_ENSURE(context, HaveSameShapes(k_cache, v_cache));
+  TF_LITE_ENSURE(context, HaveSameShapes(k_slice, v_slice));
+  TF_LITE_ENSURE(context, HaveSameShapes(updated_k_cache, updated_v_cache));
+  TF_LITE_ENSURE(context, HaveSameShapes(k_cache, updated_k_cache));
+
+  // Support only (B, S, N, H) for now.
+  TF_LITE_ENSURE(context, NumDimensions(k_slice) == kRequiredNumDimensions);
+  TF_LITE_ENSURE(context, NumDimensions(k_cache) == kRequiredNumDimensions);
+  // Ensure Positions correspond to KV sequence length.
+  TF_LITE_ENSURE(context, NumDimensions(position) == 1);
+  TF_LITE_ENSURE(context, GetTensorShape(position).Dims(0) ==
+                              GetTensorShape(k_slice).Dims(1));
+  // Enforce Batch == 1 for now.
+  TF_LITE_ENSURE(context, GetTensorShape(k_slice).Dims(0) == 1);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus ExternalKVCacheEval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* k_cache;
+  const TfLiteTensor* v_cache;
+  const TfLiteTensor* position;
+  const TfLiteTensor* k_slice;
+  const TfLiteTensor* v_slice;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kKeyTensor, &k_cache));
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kValueTensor, &v_cache));
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kPositionTensor, &position));
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kKeySliceTensor, &k_slice));
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kValueSliceTensor, &v_slice));
+
+  TfLiteTensor* updated_k_cache;
+  TfLiteTensor* updated_v_cache;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kKeyTensor, &updated_k_cache));
+  TF_LITE_ENSURE_OK(
+      context, GetOutputSafe(context, node, kValueTensor, &updated_v_cache));
+
+  // Note: For the best performance, the following memcpys should be avoided.
+  // The way to avoid that is to take advantage of CustomAllocation and use
+  // the same buffer for both input and output.
+  if (k_cache->data.raw != updated_k_cache->data.raw) {
+    memcpy(updated_k_cache->data.data, k_cache->data.data, k_cache->bytes);
+  }
+  if (v_cache->data.raw != updated_v_cache->data.raw) {
+    memcpy(updated_v_cache->data.data, v_cache->data.data, v_cache->bytes);
+  }
+
+  // Copy the new slice to the updated cache.
+  const int32_t elements_in_one_entry =
+      GetTensorShape(k_cache).Dims(2) * GetTensorShape(k_cache).Dims(3);
+  const int32_t cache_size = GetTensorShape(k_cache).Dims(1);
+  int32_t last_update_position = -1;
+  for (int i = 0; i < position->bytes / sizeof(int32_t); ++i) {
+    const int32_t update_position = position->data.i32[i];
+    // We are making the assumption that the positions are in increasing order
+    // and a decrease or equal value shows exhaustion of update slices.
+    // This assumption can be relaxed once we switch to dynamic shapes.
+    if (update_position < last_update_position) {
+      break;
+    }
+    last_update_position = update_position;
+    TF_LITE_ENSURE(context, update_position < cache_size);
+    const int32_t cache_offset = update_position * elements_in_one_entry;
+    const int32_t update_offset = i * elements_in_one_entry;
+    TF_LITE_ENSURE(context,
+                   (cache_offset + elements_in_one_entry) * sizeof(float) <=
+                       k_cache->bytes);
+    memcpy(updated_k_cache->data.f + cache_offset,
+           k_slice->data.f + update_offset,
+           elements_in_one_entry * sizeof(float));
+    memcpy(updated_v_cache->data.f + cache_offset,
+           v_slice->data.f + update_offset,
+           elements_in_one_entry * sizeof(float));
+  }
+
+  return kTfLiteOk;
+}
+
+}  // namespace llm
+
+TfLiteRegistration* Register_EXTERNAL_KV_CACHE() {
+  static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr,
+                                 llm::ExternalKVCachePrepare,
+                                 llm::ExternalKVCacheEval};
+  return &r;
+}
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/experimental/genai/external_kvcache_test.cc

@@ -0,0 +1,279 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <numeric>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/types/span.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/experimental/genai/genai_ops.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+using ::testing::TestWithParam;
+
+enum class TestType {
+  kSharedKV = 0,
+  kPingPongKV = 1,
+};
+
+class ExternalKVSingleOpModel : public SingleOpModel {
+ public:
+  ExternalKVSingleOpModel(const TensorData& k_cache, const TensorData& v_cache,
+                          const TensorData& pos_tensor,
+                          const TensorData& k_slice, const TensorData& v_slice,
+                          TestType test_type)
+      : test_type_(test_type) {
+    k_cache_in_ = AddInput(k_cache);
+    v_cache_in_ = AddInput(v_cache);
+    position_ = AddInput(pos_tensor);
+    k_slice_ = AddInput(k_slice);
+    v_slice_ = AddInput(v_slice);
+    k_cache_out_ = AddOutput(k_cache);
+    v_cache_out_ = AddOutput(v_cache);
+
+    auto get_padded_cache_size = [&](const TensorData& cache) {
+      size_t size = static_cast<size_t>(std::accumulate(
+          cache.shape.begin(), cache.shape.end(), 1.0, std::multiplies<>()));
+      // This added padding is to ensure that we can get an aligned buffer
+      // for the cache.
+      return size + kDefaultTensorAlignment;
+    };
+
+    SetCustomOp("EKV_Cache", {}, ops::custom::Register_EXTERNAL_KV_CACHE);
+    BuildInterpreter({GetShape(k_cache_in_), GetShape(v_cache_in_),
+                      GetShape(position_), GetShape(k_slice_),
+                      GetShape(v_slice_)});
+    k_cache_1_.resize(get_padded_cache_size(k_cache), 0.0);
+    v_cache_1_.resize(get_padded_cache_size(v_cache), 0.0);
+    k_cache_2_.resize(get_padded_cache_size(k_cache), 0.0);
+    v_cache_2_.resize(get_padded_cache_size(v_cache), 0.0);
+  }
+
+  TfLiteStatus Run(absl::Span<const int32_t> position,
+                   absl::Span<const float> k_slice,
+                   absl::Span<const float> v_slice) {
+    if (test_type_ == TestType::kSharedKV) {
+      TF_LITE_ENSURE_STATUS(SharedBufferPrepare());
+    } else {
+      TF_LITE_ENSURE_STATUS(PingPongBufferPrepare());
+    }
+    PopulateTensor(position_, position);
+    PopulateTensor(k_slice_, k_slice);
+    PopulateTensor(v_slice_, v_slice);
+
+    return SingleOpModel::Invoke();
+  };
+
+  std::vector<float> GetKCache() { return ExtractVector<float>(k_cache_out_); }
+
+  std::vector<float> GetVCache() { return ExtractVector<float>(v_cache_out_); }
+
+ protected:
+  TfLiteStatus SharedBufferPrepare() {
+    TF_LITE_ENSURE_STATUS(
+        SetCustomAllocationFromCache(k_cache_1_, k_cache_in_));
+    TF_LITE_ENSURE_STATUS(
+        SetCustomAllocationFromCache(v_cache_1_, v_cache_in_));
+    TF_LITE_ENSURE_STATUS(
+        SetCustomAllocationFromCache(k_cache_1_, k_cache_out_));
+    TF_LITE_ENSURE_STATUS(
+        SetCustomAllocationFromCache(v_cache_1_, v_cache_out_));
+    return interpreter_->AllocateTensors();
+  }
+
+  TfLiteStatus PingPongBufferPrepare() {
+    // The ping-pong buffer is useful for cases where source and destination
+    // buffers cannot be the same (e.g., GPU).
+    std::vector<float>* input_k_caches;
+    std::vector<float>* input_v_caches;
+    std::vector<float>* output_k_caches;
+    std::vector<float>* output_v_caches;
+    if (kv_flop_) {
+      input_k_caches = &k_cache_1_;
+      input_v_caches = &v_cache_1_;
+      output_k_caches = &k_cache_2_;
+      output_v_caches = &v_cache_2_;
+    } else {
+      input_k_caches = &k_cache_2_;
+      input_v_caches = &v_cache_2_;
+      output_k_caches = &k_cache_1_;
+      output_v_caches = &v_cache_1_;
+    }
+    kv_flop_ = !kv_flop_;
+
+    TF_LITE_ENSURE_STATUS(
+        SetCustomAllocationFromCache(*input_k_caches, k_cache_in_));
+    TF_LITE_ENSURE_STATUS(
+        SetCustomAllocationFromCache(*input_v_caches, v_cache_in_));
+    TF_LITE_ENSURE_STATUS(
+        SetCustomAllocationFromCache(*output_k_caches, k_cache_out_));
+    TF_LITE_ENSURE_STATUS(
+        SetCustomAllocationFromCache(*output_v_caches, v_cache_out_));
+    return interpreter_->AllocateTensors();
+  }
+
+  TfLiteStatus SetCustomAllocationFromCache(std::vector<float>& cache,
+                                            int tensor_index) {
+    size_t total_bytes = cache.size() * sizeof(float);
+    size_t required_number_of_bytes = total_bytes - kDefaultTensorAlignment;
+    void* original_buffer = static_cast<void*>(cache.data());
+    void* aligned_buffer =
+        std::align(kDefaultTensorAlignment, required_number_of_bytes,
+                   original_buffer, total_bytes);
+    if (aligned_buffer == nullptr ||
+        reinterpret_cast<intptr_t>(aligned_buffer) % kDefaultTensorAlignment !=
+            0) {
+      return kTfLiteError;
+    }
+    TfLiteCustomAllocation allocation = {.data = aligned_buffer,
+                                         .bytes = required_number_of_bytes};
+    return interpreter_->SetCustomAllocationForTensor(tensor_index, allocation);
+  };
+
+  int position_;
+  int k_cache_in_;
+  int v_cache_in_;
+  int k_slice_;
+  int v_slice_;
+  int k_cache_out_;
+  int v_cache_out_;
+  TestType test_type_;
+  std::vector<float> k_cache_1_;
+  std::vector<float> v_cache_1_;
+  std::vector<float> k_cache_2_;
+  std::vector<float> v_cache_2_;
+  bool kv_flop_ = true;
+};
+
+class EKVCacheTest : public TestWithParam<TestType> {};
+
+TEST_P(EKVCacheTest, SingleSliceUpdateTest) {
+  ExternalKVSingleOpModel m(
+      {TensorType_FLOAT32, {1, 3, 2, 2}}, {TensorType_FLOAT32, {1, 3, 2, 2}},
+      {TensorType_INT32, {1}}, {TensorType_FLOAT32, {1, 1, 2, 2}},
+      {TensorType_FLOAT32, {1, 1, 2, 2}}, GetParam());
+  {
+    ASSERT_EQ(m.Run(/*position=*/{0}, /*k_slice=*/{10, 11, 12, 13},
+                    /*v_slice=*/{20, 21, 22, 23}),
+              kTfLiteOk);
+
+    std::vector<float> k = m.GetKCache();
+    ASSERT_THAT(k, ElementsAreArray({10, 11, 12, 13, 0, 0, 0, 0, 0, 0, 0, 0}));
+
+    std::vector<float> v = m.GetVCache();
+    ASSERT_THAT(v, ElementsAreArray({20, 21, 22, 23, 0, 0, 0, 0, 0, 0, 0, 0}));
+  }
+  {
+    ASSERT_EQ(m.Run(/*position=*/{2}, /*k_slice=*/{50, 51, 52, 53},
+                    /*v_slice=*/{60, 61, 62, 63}),
+              kTfLiteOk);
+
+    std::vector<float> k = m.GetKCache();
+    ASSERT_THAT(k,
+                ElementsAreArray({10, 11, 12, 13, 0, 0, 0, 0, 50, 51, 52, 53}));
+
+    std::vector<float> v = m.GetVCache();
+    ASSERT_THAT(v,
+                ElementsAreArray({20, 21, 22, 23, 0, 0, 0, 0, 60, 61, 62, 63}));
+  }
+  {
+    ASSERT_EQ(m.Run(/*position=*/{1}, /*k_slice=*/{70, 71, 72, 73},
+                    /*v_slice=*/{80, 81, 82, 83}),
+              kTfLiteOk);
+
+    std::vector<float> k = m.GetKCache();
+    ASSERT_THAT(
+        k, ElementsAreArray({10, 11, 12, 13, 70, 71, 72, 73, 50, 51, 52, 53}));
+
+    std::vector<float> v = m.GetVCache();
+    ASSERT_THAT(
+        v, ElementsAreArray({20, 21, 22, 23, 80, 81, 82, 83, 60, 61, 62, 63}));
+  }
+  {
+    ASSERT_EQ(m.Run(/*position=*/{1}, /*k_slice=*/{1, 2, 3, 4},
+                    /*v_slice=*/{1, 2, 3, 4}),
+              kTfLiteOk);
+
+    std::vector<float> k = m.GetKCache();
+    ASSERT_THAT(k,
+                ElementsAreArray({10, 11, 12, 13, 1, 2, 3, 4, 50, 51, 52, 53}));
+
+    std::vector<float> v = m.GetVCache();
+    ASSERT_THAT(v,
+                ElementsAreArray({20, 21, 22, 23, 1, 2, 3, 4, 60, 61, 62, 63}));
+  }
+}
+
+TEST_P(EKVCacheTest, MultipleSliceUpdateTest) {
+  ExternalKVSingleOpModel m(
+      {TensorType_FLOAT32, {1, 3, 2, 2}}, {TensorType_FLOAT32, {1, 3, 2, 2}},
+      {TensorType_INT32, {2}}, {TensorType_FLOAT32, {1, 2, 2, 2}},
+      {TensorType_FLOAT32, {1, 2, 2, 2}}, GetParam());
+  {
+    ASSERT_EQ(m.Run(/*position=*/{0, 1}, /*k_slice=*/{1, 1, 1, 1, 2, 2, 2, 2},
+                    /*v_slice=*/{5, 5, 5, 5, 6, 6, 6, 6}),
+              kTfLiteOk);
+
+    std::vector<float> k = m.GetKCache();
+    ASSERT_THAT(k, ElementsAreArray({1, 1, 1, 1, 2, 2, 2, 2, 0, 0, 0, 0}));
+
+    std::vector<float> v = m.GetVCache();
+    ASSERT_THAT(v, ElementsAreArray({5, 5, 5, 5, 6, 6, 6, 6, 0, 0, 0, 0}));
+  }
+  {
+    ASSERT_EQ(
+        m.Run(/*position=*/{1, 2}, /*k_slice=*/{10, 10, 10, 10, 11, 11, 11, 11},
+              /*v_slice=*/{20, 20, 20, 20, 21, 21, 21, 21}),
+        kTfLiteOk);
+
+    std::vector<float> k = m.GetKCache();
+    ASSERT_THAT(k,
+                ElementsAreArray({1, 1, 1, 1, 10, 10, 10, 10, 11, 11, 11, 11}));
+
+    std::vector<float> v = m.GetVCache();
+    ASSERT_THAT(v,
+                ElementsAreArray({5, 5, 5, 5, 20, 20, 20, 20, 21, 21, 21, 21}));
+  }
+}
+
+TEST_P(EKVCacheTest, FailsOnOutOfBoundPosition) {
+  ExternalKVSingleOpModel m(
+      {TensorType_FLOAT32, {1, 3, 2, 2}}, {TensorType_FLOAT32, {1, 3, 2, 2}},
+      {TensorType_INT32, {1}}, {TensorType_FLOAT32, {1, 1, 2, 2}},
+      {TensorType_FLOAT32, {1, 1, 2, 2}}, GetParam());
+  ASSERT_EQ(m.Run(/*position=*/{3}, /*k_slice=*/{1, 2, 3, 4},
+                  /*v_slice=*/{1, 2, 3, 4}),
+            kTfLiteError);
+}
+
+INSTANTIATE_TEST_SUITE_P(EKVCacheTest, EKVCacheTest,
+                         testing::Values(TestType::kSharedKV,
+                                         TestType::kPingPongKV));
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/experimental/genai/genai_ops.cc

@@ -0,0 +1,35 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/experimental/genai/genai_ops.h"
+
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+extern "C" void GenAIOpsRegisterer(::tflite::MutableOpResolver* resolver) {
+  resolver->AddCustom("odml.update_kv_cache",
+                      tflite::ops::custom::Register_KV_CACHE());
+  resolver->AddCustom("odml.scaled_dot_product_attention",
+                      tflite::ops::custom::Register_SDPA());
+  resolver->AddCustom("odml.update_external_kv_cache",
+                      tflite::ops::custom::Register_EXTERNAL_KV_CACHE());
+}
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/experimental/genai/genai_ops.h

@@ -0,0 +1,36 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_GENAI_GENAI_OPS_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_GENAI_GENAI_OPS_H_
+
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+TfLiteRegistration* Register_KV_CACHE();
+TfLiteRegistration* Register_EXTERNAL_KV_CACHE();
+TfLiteRegistration* Register_SDPA();
+
+extern "C" void GenAIOpsRegisterer(::tflite::MutableOpResolver* resolver);
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_GENAI_GENAI_OPS_H_

tensorflow/lite/experimental/genai/genai_ops_wrapper.cc

@@ -0,0 +1,38 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "tensorflow/lite/experimental/genai/genai_ops.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+PYBIND11_MODULE(pywrap_genai_ops, m) {
+  m.doc() = R"pbdoc(
+    pywrap_genai_ops
+    -----
+  )pbdoc";
+  m.def(
+      "GenAIOpsRegisterer",
+      [](uintptr_t resolver) {
+        tflite::ops::custom::GenAIOpsRegisterer(
+            reinterpret_cast<tflite::MutableOpResolver*>(resolver));
+      },
+      R"pbdoc(
+        GenAI op registerer function with the correct signature.
+        Registers GenAI custom ops.
+      )pbdoc");
+}

tensorflow/lite/experimental/genai/kvcache.cc

@@ -0,0 +1,356 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+
+#include "flatbuffers/flexbuffers.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/experimental/resource/cache_buffer.h"
+#include "tensorflow/lite/experimental/resource/resource_base.h"
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+namespace llm {
+
+static const int kPositionTensor = 0;
+static const int kKeyTensor = 1;
+static const int kValueTensor = 2;
+static const int kFullKeyTensor = 0;
+static const int kFullValueTensor = 1;
+static const int kRequiredNumDimensions = 4;
+static const int kDefaultMaxNumCacheEntries = 2048;
+static const int kDefaultNumTransformerLayers = 32;
+static const int kDefaultTransformerLayerId = 0;
+
+static const int KVCACHE_KEY_RESOURCE = 42;
+static const int KVCACHE_VALUE_RESOURCE = 43;
+
+struct OpData {
+  int num_layers;
+  int layer_index;
+  int max_num_entries;
+  int first_slot_index;
+  // Pointers to the key and value cache buffers that this Op doesn't own
+  // (and therefore does not free on destruction of this Op).
+  resource::CacheBuffer* key_cache_buffer;
+  resource::CacheBuffer* value_cache_buffer;
+  bool is_initialized;
+  uint8_t* key_cache_ptr;
+  uint8_t* value_cache_ptr;
+};
+
+void* KVCacheInit(TfLiteContext* context, const char* buffer, size_t length) {
+  OpData* op_data = new OpData();
+  // TODO(b/333891673) Reset this value via ClearCaches in
+  // InternalBackendContext.
+  op_data->max_num_entries = -1;
+  op_data->num_layers = -1;
+  op_data->layer_index = -1;
+  op_data->first_slot_index = -1;
+  op_data->key_cache_buffer = nullptr;
+  op_data->value_cache_buffer = nullptr;
+  op_data->is_initialized = false;
+  op_data->key_cache_ptr = nullptr;
+  op_data->value_cache_ptr = nullptr;
+  return op_data;
+}
+
+TfLiteStatus KVCachePrepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 2);
+
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+
+  if (!op_data->is_initialized) {
+    const uint8_t* buffer =
+        reinterpret_cast<const uint8_t*>(node->custom_initial_data);
+    const size_t length = node->custom_initial_data_size;
+    auto flexbuffer_map = flexbuffers::GetRoot(buffer, length).AsMap();
+    int32_t max_num_entries = flexbuffer_map["kv_cache_max"].AsInt32();
+    int32_t num_layers = flexbuffer_map["num_layers"].AsInt32();
+    int32_t layer_index = flexbuffer_map["layer_index"].AsInt32();
+    op_data->max_num_entries =
+        max_num_entries > 0 ? max_num_entries : kDefaultMaxNumCacheEntries;
+    op_data->num_layers =
+        num_layers > 0 ? num_layers : kDefaultNumTransformerLayers;
+    op_data->layer_index =
+        layer_index > 0 ? layer_index : kDefaultTransformerLayerId;
+    op_data->first_slot_index = 0;
+    op_data->is_initialized = true;
+  }
+
+  // Prepare the inputs.
+  const TfLiteTensor* position;
+  const TfLiteTensor* key;
+  const TfLiteTensor* value;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kPositionTensor, &position));
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kKeyTensor, &key));
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kValueTensor, &value));
+
+  TF_LITE_ENSURE_EQ(context, position->type, kTfLiteInt64);
+  TF_LITE_ENSURE_EQ(context, key->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_EQ(context, value->type, kTfLiteFloat32);
+  // Ensure Positions correspond to KV sequence length.
+  TF_LITE_ENSURE(context, NumDimensions(position) == 1);
+  TF_LITE_ENSURE(
+      context, GetTensorShape(position).Dims(0) == GetTensorShape(key).Dims(1));
+  // Support only (B, S, N, H) for now.
+  TF_LITE_ENSURE(context, NumDimensions(key) == kRequiredNumDimensions);
+  // Enforce Batch == 1 for now.
+  TF_LITE_ENSURE(context, GetTensorShape(key).Dims(0) == 1);
+  TF_LITE_ENSURE(context, HaveSameShapes(key, value));
+
+  // Create the key and value caches. Currently statically sized.
+  TfLiteTensor* kfull;
+  TfLiteTensor* vfull;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kFullKeyTensor, &kfull));
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kFullValueTensor, &vfull));
+  // Custom data pointer to the resource cache buffer.
+  kfull->allocation_type = kTfLiteCustom;
+  vfull->allocation_type = kTfLiteCustom;
+
+  kfull->type = kTfLiteFloat32;
+  vfull->type = kTfLiteFloat32;
+
+  TfLiteIntArray* input_dims = key->dims;
+  TfLiteIntArray* kcache_dims = TfLiteIntArrayCopy(input_dims);
+  TfLiteIntArray* vcache_dims = TfLiteIntArrayCopy(input_dims);
+  kcache_dims->data[1] = op_data->max_num_entries;
+  vcache_dims->data[1] = op_data->max_num_entries;
+
+  TfLiteIntArray* kcache_buffer_dims = TfLiteIntArrayCreate(5);
+  // Batch
+  kcache_buffer_dims->data[0] = input_dims->data[0];
+  // Number of layers
+  kcache_buffer_dims->data[1] = op_data->num_layers;
+  // Sequence Length
+  kcache_buffer_dims->data[2] = op_data->max_num_entries;
+  // Num heads
+  kcache_buffer_dims->data[3] = input_dims->data[2];
+  // Head dim
+  kcache_buffer_dims->data[4] = input_dims->data[3];
+
+  TfLiteIntArray* vcache_buffer_dims = TfLiteIntArrayCopy(kcache_buffer_dims);
+
+  // Get the pointer to the tensor for our buffer storage.
+  Subgraph* subgraph = reinterpret_cast<Subgraph*>(context->impl_);
+  auto& resources = subgraph->resources();
+
+  if (resources.count(KVCACHE_KEY_RESOURCE) == 0) {
+    auto* cbuffer = new resource::CacheBuffer();
+    cbuffer->Initialize(*kcache_buffer_dims);
+    resources.emplace(KVCACHE_KEY_RESOURCE, cbuffer);
+    op_data->key_cache_buffer = cbuffer;
+  } else {
+    resource::ResourceBase* resourcePtr =
+        resources.at(KVCACHE_KEY_RESOURCE).get();
+    resource::CacheBuffer* cbuffer = (resource::CacheBuffer*)(resourcePtr);
+    op_data->key_cache_buffer = cbuffer;
+  }
+  if (resources.count(KVCACHE_VALUE_RESOURCE) == 0) {
+    auto* cbuffer = new resource::CacheBuffer();
+    cbuffer->Initialize(*vcache_buffer_dims);
+    resources.emplace(KVCACHE_VALUE_RESOURCE, cbuffer);
+    op_data->value_cache_buffer = cbuffer;
+  } else {
+    resource::ResourceBase* resourcePtr =
+        resources.at(KVCACHE_VALUE_RESOURCE).get();
+    resource::CacheBuffer* cbuffer = (resource::CacheBuffer*)(resourcePtr);
+    op_data->value_cache_buffer = cbuffer;
+  }
+
+  // Get the pointers to the individual caches for a layer.
+  RuntimeShape shape(GetTensorShape(key));
+  const int elements_in_one_entry = shape.Dims(2) * shape.Dims(3);
+  const int elements_in_one_block =
+      op_data->max_num_entries * elements_in_one_entry;
+  uint8_t* k_ptr =
+      reinterpret_cast<uint8_t*>(op_data->key_cache_buffer->GetBuffer());
+  uint8_t* v_ptr =
+      reinterpret_cast<uint8_t*>(op_data->value_cache_buffer->GetBuffer());
+  k_ptr = k_ptr + sizeof(float) * op_data->layer_index * elements_in_one_block;
+  v_ptr = v_ptr + sizeof(float) * op_data->layer_index * elements_in_one_block;
+
+  size_t kcache_dims_flatsize = kcache_dims->data[0] * kcache_dims->data[1] *
+                                kcache_dims->data[2] * kcache_dims->data[3];
+  size_t vcache_dims_flatsize = vcache_dims->data[0] * vcache_dims->data[1] *
+                                vcache_dims->data[2] * vcache_dims->data[3];
+  RuntimeShape kfull_shape(GetTensorShape(kfull));
+  RuntimeShape vfull_shape(GetTensorShape(vfull));
+  // Some testing utils don't fully set the output tensor shape
+  if (kfull_shape.FlatSize() > 1 && vfull_shape.FlatSize() > 1) {
+    TF_LITE_ENSURE_EQ(context, kfull_shape.FlatSize(), kcache_dims_flatsize);
+    TF_LITE_ENSURE_EQ(context, vfull_shape.FlatSize(), vcache_dims_flatsize);
+  }
+  TF_LITE_ENSURE_EQ(context, elements_in_one_block, kcache_dims_flatsize);
+  TF_LITE_ENSURE_EQ(context, elements_in_one_block, vcache_dims_flatsize);
+
+  kfull->data.data = k_ptr;
+  vfull->data.data = v_ptr;
+  op_data->key_cache_ptr = k_ptr;
+  op_data->value_cache_ptr = v_ptr;
+
+  TF_LITE_ENSURE_OK(context,
+                    context->ResizeTensor(context, kfull, kcache_dims));
+  TF_LITE_ENSURE_OK(context,
+                    context->ResizeTensor(context, vfull, vcache_dims));
+
+  TfLiteIntArrayFree(kcache_buffer_dims);
+  TfLiteIntArrayFree(vcache_buffer_dims);
+  return kTfLiteOk;
+}
+
+void KVCacheFree(TfLiteContext* context, void* buffer) {
+  delete static_cast<OpData*>(buffer);
+}
+
+TfLiteStatus KVCacheEval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* position;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kPositionTensor, &position));
+  const TfLiteTensor* key;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kKeyTensor, &key));
+  const TfLiteTensor* value;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kValueTensor, &value));
+
+  // Prepare the outputs.
+  TfLiteTensor* kfull;
+  TfLiteTensor* vfull;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kFullKeyTensor, &kfull));
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kFullValueTensor, &vfull));
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+
+  float* key_cache_ptr = op_data->key_cache_buffer->GetBuffer();
+  float* value_cache_ptr = op_data->value_cache_buffer->GetBuffer();
+  const int layer_index = op_data->layer_index;
+  const int64_t max_num_entries = op_data->max_num_entries;
+  int current_num_entries =
+      op_data->key_cache_buffer->GetNumEntries(layer_index);
+
+  // Compute some constants for various pieces of the cache.
+  RuntimeShape shape(GetTensorShape(key));
+  const int64_t num_slots_needed = shape.Dims(1);
+  const int elements_in_one_entry = shape.Dims(2) * shape.Dims(3);
+  const int elements_in_one_block =
+      op_data->max_num_entries * elements_in_one_entry;
+  const int64_t num_bytes_per_tensor = sizeof(float) * elements_in_one_entry;
+
+  // Get the pointers to the individual caches for a layer.
+  uint8_t* k_ptr = reinterpret_cast<uint8_t*>(key_cache_ptr);
+  uint8_t* v_ptr = reinterpret_cast<uint8_t*>(value_cache_ptr);
+  k_ptr = k_ptr + sizeof(float) * op_data->layer_index * elements_in_one_block;
+  v_ptr = v_ptr + sizeof(float) * op_data->layer_index * elements_in_one_block;
+
+  // 0. Ensure output ptr is pointing to the cache data
+  TF_LITE_ENSURE(context, k_ptr == op_data->key_cache_ptr);
+  TF_LITE_ENSURE(context, v_ptr == op_data->value_cache_ptr);
+  TF_LITE_ENSURE(context, k_ptr == kfull->data.data);
+  TF_LITE_ENSURE(context, v_ptr == vfull->data.data);
+
+  // 1. Determine which slots the inputs take up, and which slots are in the
+  //    existing span of the cache.
+
+  // Compute the span of the inputs.
+  const int64_t input_first_idx = position->data.i64[0];
+  const int64_t input_last_idx = input_first_idx + num_slots_needed - 1;
+
+  // Compute the span of the cache.
+  const int64_t cache_first_slot_idx = op_data->first_slot_index;
+  const int64_t cache_last_slot_idx =
+      cache_first_slot_idx + op_data->max_num_entries - 1;
+
+  // Compute if a shift is needed.
+  const int slots_to_shift = std::min(
+      std::max(static_cast<int64_t>(0), input_last_idx - cache_last_slot_idx),
+      max_num_entries);
+
+  // These values determine how we will write to the output tensor:
+  // first_slot := the first cache entry that we will write to in the output
+  int64_t first_slot = input_first_idx - op_data->first_slot_index;
+  if (first_slot < 0) {
+    TF_LITE_KERNEL_LOG(
+        context,
+        "Can not specify a position before this cache's first slot index of %d",
+        op_data->first_slot_index);
+    return kTfLiteError;
+  }
+
+  // byte_offset_for_output := the byte offset for the first slot.
+  int64_t byte_offset_for_output = first_slot * num_bytes_per_tensor;
+  // num_slots_for_output := the number of slots we write in the output
+  int64_t num_slots_for_output = num_slots_needed;
+
+  // 3. If we need more slots, make room in the cache by writing over oldest
+  //    entries.
+  if (slots_to_shift > 0 && slots_to_shift < max_num_entries) {
+    // If we are shifting the cache, we need to start writing from the
+    // beginning.
+    byte_offset_for_output = 0;
+    // And we need to write the entire cache.
+    num_slots_for_output = max_num_entries;
+    const int bytes_offset =
+        sizeof(float) * elements_in_one_entry * slots_to_shift;
+    const int size_bytes_to_shift = sizeof(float) * elements_in_one_entry *
+                                    (max_num_entries - slots_to_shift);
+    // TODO(b/333893996): This is O(cache_size) data motion. Consider optimizing
+    // with a circular buffer or similar.
+    memmove(k_ptr, k_ptr + bytes_offset, size_bytes_to_shift);
+    memmove(v_ptr, v_ptr + bytes_offset, size_bytes_to_shift);
+  }
+
+  // Update the first slot this cache now covers.
+  op_data->first_slot_index = op_data->first_slot_index + slots_to_shift;
+
+  // Recompute the first slot in case any shifting occurred.
+  first_slot = input_first_idx - op_data->first_slot_index;
+  const int64_t bytes_offset_for_cache = first_slot * num_bytes_per_tensor;
+
+  // 4. Put the key and value in their respective caches.
+  memcpy(k_ptr + bytes_offset_for_cache, key->data.data, key->bytes);
+  memcpy(v_ptr + bytes_offset_for_cache, value->data.data, value->bytes);
+
+  // Update counts.
+  current_num_entries =
+      std::min(first_slot + num_slots_needed, max_num_entries);
+  op_data->key_cache_buffer->SetNumEntries(layer_index, current_num_entries);
+  op_data->value_cache_buffer->SetNumEntries(layer_index, current_num_entries);
+
+  return kTfLiteOk;
+}
+
+}  // namespace llm
+
+TfLiteRegistration* Register_KV_CACHE() {
+  static TfLiteRegistration r = {llm::KVCacheInit, llm::KVCacheFree,
+                                 llm::KVCachePrepare, llm::KVCacheEval};
+  return &r;
+}
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/experimental/genai/kvcache_test.cc

@@ -0,0 +1,241 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/experimental/genai/genai_ops.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+static const int kDefaultMaxNumCacheEntries = 2048;
+
+class SimpleCacheOpModel : public SingleOpModel {
+ public:
+  SimpleCacheOpModel(const TensorData& pos_tensor, const TensorData& k_tensor,
+                     const TensorData& v_tensor) {
+    pos_ = AddInput(pos_tensor);
+    k_ = AddInput(k_tensor);
+    v_ = AddInput(v_tensor);
+    kfull_ = AddOutput(k_tensor.type);
+    vfull_ = AddOutput(v_tensor.type);
+    SetCustomOp("KV_Cache", {}, ops::custom::Register_KV_CACHE);
+
+    BuildInterpreter({GetShape(pos_), GetShape(k_), GetShape(v_)});
+  }
+
+  void SetPosition(const std::vector<int64_t>& data) {
+    PopulateTensor(pos_, data);
+  }
+  void SetKey(const std::vector<float>& data) { PopulateTensor(k_, data); }
+  void SetValue(const std::vector<float>& data) { PopulateTensor(v_, data); }
+
+  void ResizePosition(const std::vector<int>& dims) {
+    interpreter_->ResizeInputTensor(pos_, dims);
+  }
+  void ResizeKey(const std::vector<int>& dims) {
+    interpreter_->ResizeInputTensor(k_, dims);
+  }
+  void ResizeValue(const std::vector<int>& dims) {
+    interpreter_->ResizeInputTensor(v_, dims);
+  }
+
+  std::vector<float> GetFullK() {
+    const auto output = ExtractVector<float>(kfull_);
+    return output;
+  }
+
+  std::vector<float> GetFullV() {
+    const auto output = ExtractVector<float>(vfull_);
+    return output;
+  }
+
+  TfLiteStatus ReAllocate() { return interpreter_->AllocateTensors(); }
+
+ protected:
+  int pos_;
+  int k_;
+  int v_;
+  int kfull_;
+  int vfull_;
+};
+
+TEST(SimpleCacheOp1Test, BasicTest) {
+  SimpleCacheOpModel m({TensorType_INT64, {2}},
+                       {TensorType_FLOAT32, {1, 2, 2, 3}},
+                       {TensorType_FLOAT32, {1, 2, 2, 3}});
+
+  m.SetPosition({0, 1});
+  m.SetKey({{1, 0, -6, 2, 4, 3, 1, 0, -6, 2, 4, 3}});
+  m.SetValue({{4, 2, -4, 2, 4, 2, 4, 2, -4, 2, 4, 2}});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  std::vector<float> fullk = m.GetFullK();
+  std::vector<float> fullv = m.GetFullV();
+  ASSERT_EQ(fullk.size(), 2 * 3 * kDefaultMaxNumCacheEntries);
+  ASSERT_EQ(fullv.size(), 2 * 3 * kDefaultMaxNumCacheEntries);
+}
+
+TEST(SimpleCacheOp2Test, AddToCache) {
+  SimpleCacheOpModel m({TensorType_INT64, {2}},
+                       {TensorType_FLOAT32, {1, 2, 2, 3}},
+                       {TensorType_FLOAT32, {1, 2, 2, 3}});
+
+  m.SetPosition({0, 1});
+  std::vector<float> key = {1, 5, -6, 2, 4, 3, 8, 9, -8, 7, 2, 11};
+  m.SetKey(key);
+  std::vector<float> value = {2, 3, -4, 5, 6, 7, 1, 8, -12, 11, 14, 21};
+  m.SetValue(value);
+  const int key_size = 2 * 3;
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  std::vector<float> fullk = m.GetFullK();
+  std::vector<float> fullv = m.GetFullV();
+  for (int i = 0; i < key.size(); ++i) {
+    ASSERT_EQ(fullk[i], key[i]);
+    ASSERT_EQ(fullv[i], value[i]);
+  }
+  for (int i = key.size(); i < fullk.size(); ++i) {
+    ASSERT_EQ(fullk[i], 0.);
+    ASSERT_EQ(fullv[i], 0.);
+  }
+
+  ASSERT_EQ(fullk.size(), 2 * 3 * kDefaultMaxNumCacheEntries);
+  ASSERT_EQ(fullv.size(), 2 * 3 * kDefaultMaxNumCacheEntries);
+
+  for (int i = 0; i < 510; i++) {
+    int offset = 2 * i + 2;
+    m.SetPosition({offset, offset + 1});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  }
+
+  fullk = m.GetFullK();
+  fullv = m.GetFullV();
+
+  for (int i = 0; i < 1022 * key_size; ++i) {
+    ASSERT_NE(fullv[i], 0);
+  }
+
+  for (int i = 1022 * key_size; i < fullk.size(); ++i) {
+    ASSERT_EQ(fullv[i], 0);
+  }
+}
+
+TEST(SimpleCacheOp2Test, ShiftSlotsInCache) {
+  SimpleCacheOpModel m({TensorType_INT64, {2}},
+                       {TensorType_FLOAT32, {1, 2, 2, 3}},
+                       {TensorType_FLOAT32, {1, 2, 2, 3}});
+
+  m.SetPosition({0, 1});
+  std::vector<float> key = {1, 5, -6, 2, 4, 3, 2, 6, -7, 3, 5, 4};
+  m.SetKey(key);
+  std::vector<float> value = {4, 2, -4, 2, 4, 2, 9, 8, -9, 8, 9, 1};
+  m.SetValue(value);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  std::vector<float> fullk = m.GetFullK();
+  std::vector<float> fullv = m.GetFullV();
+  for (int i = 0; i < key.size(); ++i) {
+    ASSERT_EQ(fullk[i], key[i]);
+    ASSERT_EQ(fullv[i], value[i]);
+  }
+  for (int i = key.size(); i < fullk.size(); ++i) {
+    ASSERT_EQ(fullk[i], 0.);
+    ASSERT_EQ(fullv[i], 0.);
+  }
+  ASSERT_EQ(fullk.size(), 2 * 3 * kDefaultMaxNumCacheEntries);
+  ASSERT_EQ(fullv.size(), 2 * 3 * kDefaultMaxNumCacheEntries);
+
+  // Now fill up the cache
+  for (int i = 0; i < 1023; i++) {
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    int offset = 2 * i + 2;
+    m.SetPosition({offset, offset + 1});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  }
+
+  fullk = m.GetFullK();
+  fullv = m.GetFullV();
+
+  for (int i = 0; i < fullk.size(); ++i) {
+    ASSERT_NE(fullk[i], 0);
+    ASSERT_NE(fullv[i], 0);
+  }
+
+  for (int j = 0; j < 6; ++j) {
+    int idxfull = fullk.size() - 6 + j;
+    int idx = 6 + j;
+    ASSERT_EQ(fullk[idxfull], key[idx]);
+    ASSERT_EQ(fullv[idxfull], value[idx]);
+  }
+
+  std::vector<float> key2 = {1, 1, 1, 1, 1, 1, 7, 7, 7, 7, 7, 7};
+  m.SetKey(key2);
+  std::vector<float> value2 = {8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9};
+  m.SetValue(value2);
+  m.SetPosition({2048, 2049});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  fullk = m.GetFullK();
+  fullv = m.GetFullV();
+
+  for (int j = 0; j < 12; ++j) {
+    int idxfull = fullk.size() - 12 + j;
+    ASSERT_EQ(fullk[idxfull], key2[j]);
+    ASSERT_EQ(fullv[idxfull], value2[j]);
+  }
+
+  // Resize to a single entry. Add to a full cache and verify
+  // the cached contents.
+  m.ResizeKey({1, 1, 2, 3});
+  m.ResizeValue({1, 1, 2, 3});
+  m.ResizePosition({1});
+  m.ReAllocate();
+
+  std::vector<float> key3 = {4, 4, 4, 4, 4, 4};
+  m.SetKey(key3);
+  std::vector<float> value3 = {2, 2, 2, 2, 2, 2};
+  m.SetValue(value3);
+  m.SetPosition({2050});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  fullk = m.GetFullK();
+  fullv = m.GetFullV();
+
+  for (int j = 0; j < 6; ++j) {
+    int idxfull = fullk.size() - 6 + j;
+    ASSERT_EQ(fullk[idxfull], key3[j]);
+    ASSERT_EQ(fullv[idxfull], value3[j]);
+  }
+
+  // Verify that other cache entries got shifted up 1.
+  for (int j = 0; j < 6; ++j) {
+    int idxfull = fullk.size() - 12 + j;
+    ASSERT_EQ(fullk[idxfull], key2[6 + j]);
+    ASSERT_EQ(fullv[idxfull], value2[6 + j]);
+  }
+
+  std::vector<float> key4 = {5, 5, 5, 5, 5, 5};
+  m.SetKey(key3);
+  std::vector<float> value4 = {3, 3, 3, 3, 3, 3};
+  m.SetValue(value3);
+  m.SetPosition({0});
+  ASSERT_EQ(m.Invoke(), kTfLiteError);
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/experimental/genai/pywrap_genai_ops.pyi

@@ -0,0 +1,16 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+def GenAIOpsRegisterer(arg0: int) -> None: ...

tensorflow/lite/experimental/genai/sdpa.cc

@@ -0,0 +1,674 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <math.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+
+#include "flatbuffers/flexbuffers.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/reference/add.h"
+#include "tensorflow/lite/kernels/internal/reference/batch_matmul.h"
+#include "tensorflow/lite/kernels/internal/reference/fully_connected.h"
+#include "tensorflow/lite/kernels/internal/reference/softmax.h"
+#include "tensorflow/lite/kernels/internal/reference/transpose.h"
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+namespace llm {
+
+static const int kQueryTensor = 0;
+static const int kKeyTensor = 1;
+static const int kValueTensor = 2;
+static const int kAttentionMaskTensor = 3;
+static const int kOutputTensor = 0;
+
+static const int kNumTempTensors = 10;
+static const int kTransposeQueryTempTensorIndex = 0;
+static const int kTransposeKeyTempTensorIndex = 1;
+static const int kMatMul1TempTensorIndex = 2;
+static const int kAddTempTensorIndex = 3;
+static const int kTransposeValueTempTensorIndex = 4;
+static const int kMatMul2TempTensorIndex = 5;
+static const int kReshape1TempTensorIndex = 6;
+static const int kReshape2TempTensorIndex = 7;
+static const int kBroadcastKTempTensorIndex = 8;
+static const int kBroadcastVTempTensorIndex = 9;
+
+struct OpData {
+  float scale;
+  int scratch_tensor_index;
+};
+
+void* SDPAInit(TfLiteContext* context, const char* buffer, size_t length) {
+  OpData* op_data = new OpData();
+  op_data->scale = 0.0f;
+  context->AddTensors(context, kNumTempTensors, &op_data->scratch_tensor_index);
+  return op_data;
+}
+
+TfLiteStatus SDPAPrepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 4);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+
+  const TfLiteTensor* q_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kQueryTensor, &q_tensor));
+  const TfLiteTensor* k_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kKeyTensor, &k_tensor));
+  const TfLiteTensor* v_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kValueTensor, &v_tensor));
+  const TfLiteTensor* mask_tensor;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node, kAttentionMaskTensor, &mask_tensor));
+  TF_LITE_ENSURE_EQ(context, NumDimensions(q_tensor), NumDimensions(k_tensor));
+  TF_LITE_ENSURE_EQ(context, NumDimensions(k_tensor), NumDimensions(v_tensor));
+  TF_LITE_ENSURE_EQ(context, NumDimensions(v_tensor),
+                    NumDimensions(mask_tensor));
+  TF_LITE_ENSURE_EQ(context, NumDimensions(mask_tensor), 4);
+
+  // Get custom op params
+  const uint8_t* buffer =
+      reinterpret_cast<const uint8_t*>(node->custom_initial_data);
+  const size_t length = node->custom_initial_data_size;
+  auto flexbuffer_map = flexbuffers::GetRoot(buffer, length).AsMap();
+  float scale = flexbuffer_map["scale"].AsFloat();
+  op_data->scale = scale > 0.0f ? scale : 0.0f;
+
+  // If scale is not set, use sqrt(q_tensor->dims->data[3])
+  if (op_data->scale == 0.0f)
+    op_data->scale = 1 / sqrt(q_tensor->dims->data[3]);
+
+  TfLiteIntArrayFree(node->temporaries);
+  node->temporaries = TfLiteIntArrayCreate(kNumTempTensors);
+  bool mqa = k_tensor->dims->data[2] == 1;
+
+  // Temp tensor for Transposed Q;
+  {
+    node->temporaries->data[kTransposeQueryTempTensorIndex] =
+        op_data->scratch_tensor_index + kTransposeQueryTempTensorIndex;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node,
+                                       /*index=*/kTransposeQueryTempTensorIndex,
+                                       &scratch_buffer));
+    TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(4);
+    for (int i = 0; i < 4; ++i) {
+      scratch_buffer_size->data[i] = q_tensor->dims->data[i];
+    }
+    // Swap middle two dimensions.
+    scratch_buffer_size->data[1] = q_tensor->dims->data[2];
+    scratch_buffer_size->data[2] = q_tensor->dims->data[1];
+
+    scratch_buffer->type = kTfLiteFloat32;
+    scratch_buffer->allocation_type = kTfLiteArenaRw;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+  }
+
+  // Temp tensor for Transposed K;
+  {
+    node->temporaries->data[kTransposeKeyTempTensorIndex] =
+        op_data->scratch_tensor_index + kTransposeKeyTempTensorIndex;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node,
+                                       /*index=*/kTransposeKeyTempTensorIndex,
+                                       &scratch_buffer));
+    TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(4);
+    for (int i = 0; i < 4; ++i) {
+      scratch_buffer_size->data[i] = k_tensor->dims->data[i];
+    }
+    // Swap to middle two dimensions.
+    scratch_buffer_size->data[1] = k_tensor->dims->data[2];
+    scratch_buffer_size->data[2] = k_tensor->dims->data[1];
+
+    scratch_buffer->type = kTfLiteFloat32;
+    scratch_buffer->allocation_type = kTfLiteArenaRw;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+  }
+
+  TfLiteIntArray* add_broadcast_shape = nullptr;
+  // Temp tensor for Matmul1 output;
+  {
+    node->temporaries->data[kMatMul1TempTensorIndex] =
+        op_data->scratch_tensor_index + kMatMul1TempTensorIndex;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(
+        context,
+        GetTemporarySafe(context, node,
+                         /*index=*/kMatMul1TempTensorIndex, &scratch_buffer));
+    TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(4);
+    // mha/gqa: [permute_q[0], permute_q[1], permute_q[2], permute_k[2]]
+    int matmul_out_shape[4] = {q_tensor->dims->data[0], q_tensor->dims->data[2],
+                               q_tensor->dims->data[1],
+                               k_tensor->dims->data[1]};
+    for (int i = 0; i < 4; ++i) {
+      scratch_buffer_size->data[i] = matmul_out_shape[i];
+    }
+
+    scratch_buffer->type = kTfLiteFloat32;
+    scratch_buffer->allocation_type = kTfLiteArenaRw;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+    // get dims from attention_mask, matmul1_out for add broadcast
+    CalculateShapeForBroadcast(context, mask_tensor, scratch_buffer,
+                               &add_broadcast_shape);
+  }
+
+  // Temp tensor for add output;
+  int add_out_shape[4];
+  {
+    node->temporaries->data[kAddTempTensorIndex] =
+        op_data->scratch_tensor_index + kAddTempTensorIndex;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node,
+                                                /*index=*/kAddTempTensorIndex,
+                                                &scratch_buffer));
+    TfLiteIntArray* scratch_buffer_size = add_broadcast_shape;
+    for (int i = 0; i < 4; ++i) {
+      add_out_shape[i] = scratch_buffer_size->data[i];
+    }
+
+    scratch_buffer->type = kTfLiteFloat32;
+    scratch_buffer->allocation_type = kTfLiteArenaRw;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+  }
+
+  // Temp tensor for Transposed V;
+  {
+    node->temporaries->data[kTransposeValueTempTensorIndex] =
+        op_data->scratch_tensor_index + kTransposeValueTempTensorIndex;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node,
+                                       /*index=*/kTransposeValueTempTensorIndex,
+                                       &scratch_buffer));
+    TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(4);
+    // Swap to {0, 2, 3, 1} dimensions.
+    scratch_buffer_size->data[0] = v_tensor->dims->data[0];
+    scratch_buffer_size->data[1] = v_tensor->dims->data[2];
+    scratch_buffer_size->data[2] = v_tensor->dims->data[3];
+    scratch_buffer_size->data[3] = v_tensor->dims->data[1];
+
+    scratch_buffer->type = kTfLiteFloat32;
+    scratch_buffer->allocation_type = kTfLiteArenaRw;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+  }
+
+  // Temp tensor for Matmul2 output;
+  {
+    node->temporaries->data[kMatMul2TempTensorIndex] =
+        op_data->scratch_tensor_index + kMatMul2TempTensorIndex;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(
+        context,
+        GetTemporarySafe(context, node,
+                         /*index=*/kMatMul2TempTensorIndex, &scratch_buffer));
+    TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(4);
+    // logits_out_shape = add_out_shape
+    // mha/gqa: [logits_out[0], logits_out[1], logits_out[2], permute_v[2]]
+    scratch_buffer_size->data[0] = add_out_shape[0];
+    scratch_buffer_size->data[1] = add_out_shape[1];
+    scratch_buffer_size->data[2] = add_out_shape[2];
+    scratch_buffer_size->data[3] = v_tensor->dims->data[3];
+
+    scratch_buffer->type = kTfLiteFloat32;
+    scratch_buffer->allocation_type = kTfLiteArenaRw;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+  }
+
+  // Temp tensor for Reshape K / Transpose Q;
+  {
+    node->temporaries->data[kReshape1TempTensorIndex] =
+        op_data->scratch_tensor_index + kReshape1TempTensorIndex;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(
+        context,
+        GetTemporarySafe(context, node,
+                         /*index=*/kReshape1TempTensorIndex, &scratch_buffer));
+    TfLiteIntArray* scratch_buffer_size;
+    if (mqa)
+      scratch_buffer_size = TfLiteIntArrayCreate(2);
+    else
+      scratch_buffer_size = TfLiteIntArrayCreate(4);
+    if (mqa) {
+      scratch_buffer_size->data[0] = k_tensor->dims->data[1];
+      scratch_buffer_size->data[1] = k_tensor->dims->data[3];
+    } else {
+      scratch_buffer_size->data[0] = q_tensor->dims->data[0];
+      scratch_buffer_size->data[1] = q_tensor->dims->data[2];
+      scratch_buffer_size->data[2] = q_tensor->dims->data[3];
+      scratch_buffer_size->data[3] = q_tensor->dims->data[1];
+    }
+
+    scratch_buffer->type = kTfLiteFloat32;
+    scratch_buffer->allocation_type = kTfLiteArenaRw;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+  }
+
+  // Temp tensor for Reshape V / Add_out (softmax_out);
+  {
+    node->temporaries->data[kReshape2TempTensorIndex] =
+        op_data->scratch_tensor_index + kReshape2TempTensorIndex;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(
+        context,
+        GetTemporarySafe(context, node,
+                         /*index=*/kReshape2TempTensorIndex, &scratch_buffer));
+    TfLiteIntArray* scratch_buffer_size;
+    if (mqa)
+      scratch_buffer_size = TfLiteIntArrayCreate(2);
+    else
+      scratch_buffer_size = TfLiteIntArrayCreate(4);
+    if (mqa) {
+      scratch_buffer_size->data[0] = v_tensor->dims->data[3];
+      scratch_buffer_size->data[1] = v_tensor->dims->data[1];
+    } else {
+      scratch_buffer_size->data[0] = add_out_shape[0];
+      scratch_buffer_size->data[1] = add_out_shape[1];
+      scratch_buffer_size->data[2] = add_out_shape[3];
+      scratch_buffer_size->data[3] = add_out_shape[2];
+    }
+
+    scratch_buffer->type = kTfLiteFloat32;
+    scratch_buffer->allocation_type = kTfLiteArenaRw;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+  }
+
+  // Temp tensor for Broadcast K
+  {
+    node->temporaries->data[kBroadcastKTempTensorIndex] =
+        op_data->scratch_tensor_index + kBroadcastKTempTensorIndex;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node,
+                                       /*index=*/kBroadcastKTempTensorIndex,
+                                       &scratch_buffer));
+    TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(4);
+
+    scratch_buffer_size->data[0] = k_tensor->dims->data[0];
+    scratch_buffer_size->data[1] = q_tensor->dims->data[2];  // num_heads
+    scratch_buffer_size->data[2] = k_tensor->dims->data[1];
+    scratch_buffer_size->data[3] = k_tensor->dims->data[3];
+
+    scratch_buffer->type = kTfLiteFloat32;
+    scratch_buffer->allocation_type = kTfLiteArenaRw;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+  }
+
+  // Temp tensor for Broadcast V
+  {
+    node->temporaries->data[kBroadcastVTempTensorIndex] =
+        op_data->scratch_tensor_index + kBroadcastVTempTensorIndex;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node,
+                                       /*index=*/kBroadcastVTempTensorIndex,
+                                       &scratch_buffer));
+    TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(4);
+
+    scratch_buffer_size->data[0] = v_tensor->dims->data[0];
+    scratch_buffer_size->data[1] = q_tensor->dims->data[2];  // num_heads
+    scratch_buffer_size->data[2] = v_tensor->dims->data[3];
+    scratch_buffer_size->data[3] = v_tensor->dims->data[1];
+
+    scratch_buffer->type = kTfLiteFloat32;
+    scratch_buffer->allocation_type = kTfLiteArenaRw;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+  }
+
+  return kTfLiteOk;
+}
+
+void SDPAFree(TfLiteContext* context, void* buffer) {
+  delete static_cast<OpData*>(buffer);
+}
+
+TfLiteStatus SDPAEval(TfLiteContext* context, TfLiteNode* node) {
+  /*
+  Simple implementation of Scaled Dot Product Attention.
+  Takes query_proj, key_proj, value_proj, mask tensors as inputs, and
+  outputs the attention result.
+
+  Notes:
+  Scale is computed using 1/sqrt(head_dim),
+  head_dim = q[-1] = embedding_dim // num_q_heads
+  Only support for FLOAT32 inputs for now.
+  Only support static tensors for now (k/v[1] = max sequence length)
+  */
+
+  const TfLiteTensor* query_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kQueryTensor, &query_tensor));
+  auto query_shape = GetTensorShape(query_tensor);
+  auto query_data = GetTensorData<float>(query_tensor);
+  const TfLiteTensor* key_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kKeyTensor, &key_tensor));
+  auto key_shape = GetTensorShape(key_tensor);
+  auto key_data = GetTensorData<float>(key_tensor);
+  const TfLiteTensor* value_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kValueTensor, &value_tensor));
+  auto value_shape = GetTensorShape(value_tensor);
+  auto value_data = GetTensorData<float>(value_tensor);
+  const TfLiteTensor* attention_mask_tensor;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kAttentionMaskTensor,
+                                          &attention_mask_tensor));
+  auto attention_mask_shape = GetTensorShape(attention_mask_tensor);
+  auto attention_mask_data = GetTensorData<float>(attention_mask_tensor);
+  TfLiteTensor* output_tensor;
+  TF_LITE_ENSURE_OK(
+      context, GetOutputSafe(context, node, kOutputTensor, &output_tensor));
+  auto output_shape = GetTensorShape(output_tensor);
+  auto output_data = GetTensorData<float>(output_tensor);
+
+  // temporaries
+  TfLiteTensor* transpose_q_out_tensor;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetTemporarySafe(context, node, /*index=*/kTransposeQueryTempTensorIndex,
+                       &transpose_q_out_tensor));
+  auto transpose_q_out_shape = GetTensorShape(transpose_q_out_tensor);
+  auto transpose_q_out_data = GetTensorData<float>(transpose_q_out_tensor);
+  TfLiteTensor* transpose_k_out_tensor;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetTemporarySafe(context, node, /*index=*/kTransposeKeyTempTensorIndex,
+                       &transpose_k_out_tensor));
+  auto transpose_k_out_shape = GetTensorShape(transpose_k_out_tensor);
+  auto transpose_k_out_data = GetTensorData<float>(transpose_k_out_tensor);
+  TfLiteTensor* matmul1_out_tensor;
+  TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node,
+                                              /*index=*/kMatMul1TempTensorIndex,
+                                              &matmul1_out_tensor));
+  auto matmul1_out_shape = GetTensorShape(matmul1_out_tensor);
+  auto matmul1_out_data = GetTensorData<float>(matmul1_out_tensor);
+  TfLiteTensor* add_out_tensor;
+  TF_LITE_ENSURE_OK(
+      context, GetTemporarySafe(context, node, /*index=*/kAddTempTensorIndex,
+                                &add_out_tensor));
+  auto add_out_shape = GetTensorShape(add_out_tensor);
+  auto add_out_data = GetTensorData<float>(add_out_tensor);
+  TfLiteTensor* transpose_v_out_tensor;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetTemporarySafe(context, node, /*index=*/kTransposeValueTempTensorIndex,
+                       &transpose_v_out_tensor));
+  auto transpose_v_out_shape = GetTensorShape(transpose_v_out_tensor);
+  auto transpose_v_out_data = GetTensorData<float>(transpose_v_out_tensor);
+  TfLiteTensor* matmul2_out_tensor;
+  TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node,
+                                              /*index=*/kMatMul2TempTensorIndex,
+                                              &matmul2_out_tensor));
+  auto matmul2_out_shape = GetTensorShape(matmul2_out_tensor);
+  auto matmul2_out_data = GetTensorData<float>(matmul2_out_tensor);
+  TfLiteTensor* reshape_k_or_q_out_tensor;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetTemporarySafe(context, node, /*index=*/kReshape1TempTensorIndex,
+                       &reshape_k_or_q_out_tensor));
+  auto reshape_k_or_q_out_shape = GetTensorShape(reshape_k_or_q_out_tensor);
+  auto reshape_k_or_q_out_data =
+      GetTensorData<float>(reshape_k_or_q_out_tensor);
+  TfLiteTensor* reshape_v_or_add_out_tensor;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetTemporarySafe(context, node, /*index=*/kReshape2TempTensorIndex,
+                       &reshape_v_or_add_out_tensor));
+  auto reshape_v_or_add_out_shape = GetTensorShape(reshape_v_or_add_out_tensor);
+  auto reshape_v_or_add_out_data =
+      GetTensorData<float>(reshape_v_or_add_out_tensor);
+  TfLiteTensor* broadcast_k_out_tensor;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetTemporarySafe(context, node, /*index=*/kBroadcastKTempTensorIndex,
+                       &broadcast_k_out_tensor));
+  auto broadcast_k_out_shape = GetTensorShape(broadcast_k_out_tensor);
+  auto broadcast_k_out_data = GetTensorData<float>(broadcast_k_out_tensor);
+  TfLiteTensor* broadcast_v_out_tensor;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetTemporarySafe(context, node, /*index=*/kBroadcastVTempTensorIndex,
+                       &broadcast_v_out_tensor));
+  auto broadcast_v_out_shape = GetTensorShape(broadcast_v_out_tensor);
+  auto broadcast_v_out_data = GetTensorData<float>(broadcast_v_out_tensor);
+
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+
+  bool mqa = key_tensor->dims->data[2] == 1;
+  bool gqa = !mqa && (key_tensor->dims->data[2] != query_tensor->dims->data[2]);
+
+  // scale * q
+  float scale = op_data->scale;
+  int flat_size = query_shape.FlatSize();
+  float output_min = -std::numeric_limits<float>::infinity();
+  float output_max = std::numeric_limits<float>::infinity();
+  for (int i = 0; i < flat_size; ++i) {
+    query_tensor->data.f[i] = ActivationFunctionWithMinMax(
+        query_tensor->data.f[i] * scale, output_min, output_max);
+  }
+
+  // permute q {0, 2, 1, 3}
+  tflite::TransposeParams transpose_q_params;
+  transpose_q_params.perm_count = 4;
+  transpose_q_params.perm[0] = 0;
+  transpose_q_params.perm[1] = 2;
+  transpose_q_params.perm[2] = 1;
+  transpose_q_params.perm[3] = 3;
+  reference_ops::Transpose(transpose_q_params, query_shape, query_data,
+                           transpose_q_out_shape, transpose_q_out_data);
+
+  // permute k {0, 2, 1, 3}
+  tflite::TransposeParams transpose_k_params;
+  transpose_k_params.perm_count = 4;
+  transpose_k_params.perm[0] = 0;
+  transpose_k_params.perm[1] = 2;
+  transpose_k_params.perm[2] = 1;
+  transpose_k_params.perm[3] = 3;
+  reference_ops::Transpose(transpose_k_params, key_shape, key_data,
+                           transpose_k_out_shape, transpose_k_out_data);
+
+  // broadcast k to match num_heads
+  // broadcasting similar to torch.repeat_interleave
+  if (gqa) {
+    float* transpose_k_ptr = transpose_k_out_data;
+    float* broadcast_k_ptr = broadcast_k_out_data;
+    int num_elements =
+        transpose_k_out_shape.Dims(2) * transpose_k_out_shape.Dims(3);
+    int num_repeat =
+        broadcast_k_out_shape.Dims(1) / transpose_k_out_shape.Dims(1);
+    for (int i = 0; i < transpose_k_out_shape.Dims(0); ++i) {
+      for (int j = 0; j < transpose_k_out_shape.Dims(1); ++j) {
+        for (int k = 0; k < num_repeat; ++k) {
+          memcpy(broadcast_k_ptr, transpose_k_ptr,
+                 num_elements * sizeof(float));
+          broadcast_k_ptr += num_elements;
+        }
+        transpose_k_ptr += num_elements;
+      }
+    }
+  }
+
+  // reshape k for MQA, or transpose q for MHA
+  if (mqa) {
+    TF_LITE_ENSURE_EQ(context, transpose_k_out_tensor->bytes,
+                      reshape_k_or_q_out_tensor->bytes);
+    memcpy(reshape_k_or_q_out_tensor->data.data,
+           transpose_k_out_tensor->data.data, transpose_k_out_tensor->bytes);
+  } else {
+    // permute q2 {0, 1, 3, 2}
+    tflite::TransposeParams transpose_q2_params;
+    transpose_q2_params.perm_count = 4;
+    transpose_q2_params.perm[0] = 0;
+    transpose_q2_params.perm[1] = 1;
+    transpose_q2_params.perm[2] = 3;
+    transpose_q2_params.perm[3] = 2;
+    reference_ops::Transpose(transpose_q2_params, transpose_q_out_shape,
+                             transpose_q_out_data, reshape_k_or_q_out_shape,
+                             reshape_k_or_q_out_data);
+  }
+
+  // mqa FC (q, squeezed_k)
+  // mha BMM(q, k) transpose_b = true
+  if (mqa) {
+    tflite::FullyConnectedParams fc_params;
+    fc_params.float_activation_min = output_min;
+    fc_params.float_activation_max = output_max;
+    reference_ops::FullyConnected(
+        fc_params, transpose_q_out_shape, transpose_q_out_data,
+        reshape_k_or_q_out_shape, reshape_k_or_q_out_data, RuntimeShape(),
+        nullptr, matmul1_out_shape, matmul1_out_data);
+  } else if (gqa) {
+    // pass rhs first (this is why we transpose q above)
+    reference_ops::BatchMatMul(
+        broadcast_k_out_shape, broadcast_k_out_data, reshape_k_or_q_out_shape,
+        reshape_k_or_q_out_data, matmul1_out_shape, matmul1_out_data);
+  } else {
+    reference_ops::BatchMatMul(
+        transpose_k_out_shape, transpose_k_out_data, reshape_k_or_q_out_shape,
+        reshape_k_or_q_out_data, matmul1_out_shape, matmul1_out_data);
+  }
+
+  // add matmul_out + mask
+  tflite::ArithmeticParams add_params;
+  SetActivationParams(output_min, output_max, &add_params);
+  reference_ops::BroadcastAdd6DSlow(
+      add_params, attention_mask_shape, attention_mask_data, matmul1_out_shape,
+      matmul1_out_data, add_out_shape, add_out_data);
+
+  // softmax, can do in-place
+  tflite::SoftmaxParams softmax_params;
+  softmax_params.beta = 1.0f;
+  reference_ops::Softmax(softmax_params, add_out_shape, add_out_data,
+                         add_out_shape, add_out_data);
+
+  // permute v {0, 2, 3, 1}
+  tflite::TransposeParams transpose_v_params;
+  transpose_v_params.perm_count = 4;
+  transpose_v_params.perm[0] = 0;
+  transpose_v_params.perm[1] = 2;
+  transpose_v_params.perm[2] = 3;
+  transpose_v_params.perm[3] = 1;
+  reference_ops::Transpose(transpose_v_params, value_shape, value_data,
+                           transpose_v_out_shape, transpose_v_out_data);
+
+  // broadcast v to match num_heads
+  // broadcasting similar to torch.repeat_interleave
+  if (gqa) {
+    float* transpose_v_ptr = transpose_v_out_data;
+    float* broadcast_v_ptr = broadcast_v_out_data;
+    int num_elements =
+        transpose_v_out_shape.Dims(2) * transpose_v_out_shape.Dims(3);
+    int num_repeat =
+        broadcast_v_out_shape.Dims(1) / transpose_v_out_shape.Dims(1);
+    for (int i = 0; i < transpose_v_out_shape.Dims(0); ++i) {
+      for (int j = 0; j < transpose_v_out_shape.Dims(1); ++j) {
+        for (int k = 0; k < num_repeat; ++k) {
+          memcpy(broadcast_v_ptr, transpose_v_ptr,
+                 num_elements * sizeof(float));
+          broadcast_v_ptr += num_elements;
+        }
+        transpose_v_ptr += num_elements;
+      }
+    }
+  }
+
+  // reshape v for MQA, or add_out (softmax_out)
+  if (mqa) {
+    TF_LITE_ENSURE_EQ(context, transpose_v_out_tensor->bytes,
+                      reshape_v_or_add_out_tensor->bytes);
+    memcpy(reshape_v_or_add_out_tensor->data.data,
+           transpose_v_out_tensor->data.data, transpose_v_out_tensor->bytes);
+  } else {
+    // permute softmax_out {0, 1, 3, 2}
+    tflite::TransposeParams transpose_softmax_out_params;
+    transpose_softmax_out_params.perm_count = 4;
+    transpose_softmax_out_params.perm[0] = 0;
+    transpose_softmax_out_params.perm[1] = 1;
+    transpose_softmax_out_params.perm[2] = 3;
+    transpose_softmax_out_params.perm[3] = 2;
+    reference_ops::Transpose(transpose_softmax_out_params, add_out_shape,
+                             add_out_data, reshape_v_or_add_out_shape,
+                             reshape_v_or_add_out_data);
+  }
+
+  // mqa FC (softmax_out, squeezed_v)
+  // mha BMM(softmax_out, v) transpose_b = true
+  if (mqa) {
+    tflite::FullyConnectedParams fc_params;
+    fc_params.float_activation_min = output_min;
+    fc_params.float_activation_max = output_max;
+    reference_ops::FullyConnected(fc_params, add_out_shape, add_out_data,
+                                  reshape_v_or_add_out_shape,
+                                  reshape_v_or_add_out_data, RuntimeShape(),
+                                  nullptr, matmul2_out_shape, matmul2_out_data);
+  } else if (gqa) {
+    // pass rhs first (this is why we transpose add_out above)
+    reference_ops::BatchMatMul(
+        broadcast_v_out_shape, broadcast_v_out_data, reshape_v_or_add_out_shape,
+        reshape_v_or_add_out_data, matmul2_out_shape, matmul2_out_data);
+  } else {
+    reference_ops::BatchMatMul(
+        transpose_v_out_shape, transpose_v_out_data, reshape_v_or_add_out_shape,
+        reshape_v_or_add_out_data, matmul2_out_shape, matmul2_out_data);
+  }
+
+  // permute out {0, 2, 1, 3}
+  tflite::TransposeParams transpose_out_params;
+  transpose_out_params.perm_count = 4;
+  transpose_out_params.perm[0] = 0;
+  transpose_out_params.perm[1] = 2;
+  transpose_out_params.perm[2] = 1;
+  transpose_out_params.perm[3] = 3;
+  reference_ops::Transpose(transpose_out_params, matmul2_out_shape,
+                           matmul2_out_data, output_shape, output_data);
+
+  return kTfLiteOk;
+}
+
+}  // namespace llm
+
+TfLiteRegistration* Register_SDPA() {
+  static TfLiteRegistration r = {llm::SDPAInit, llm::SDPAFree, llm::SDPAPrepare,
+                                 llm::SDPAEval};
+  return &r;
+}
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite