Add a Resource for KV Cache buffer storage

PiperOrigin-RevId: 620097541

Author: talumbau

tensorflow/lite/experimental/resource/BUILD

@@ -6,6 +6,32 @@ package(
     licenses = ["notice"],
 )
 
+cc_library(
+    name = "cache_buffer",
+    srcs = ["cache_buffer.cc"],
+    hdrs = [
+        "cache_buffer.h",
+        "//tensorflow/lite/core/c:common.h",
+    ],
+    deps = [
+        ":resource",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "cache_buffer_test",
+    srcs = ["cache_buffer_test.cc"],
+    deps = [
+        ":cache_buffer",
+        "//tensorflow/lite/c:common",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "resource",
     srcs = [

tensorflow/lite/experimental/resource/cache_buffer.cc

@@ -0,0 +1,56 @@
+/* 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/resource/cache_buffer.h"
+
+#include <cstdlib>
+#include <cstring>
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace resource {
+
+constexpr char kCacheBufferTensorName[] = "CacheBuffer";
+
+TfLiteStatus CacheBuffer::Initialize(const TfLiteIntArray &shape,
+                                     const TfLiteType &type) {
+  // Set basic parameters.
+  tensor_.name = kCacheBufferTensorName;
+  tensor_.allocation_type = kTfLiteDynamic;
+  tensor_.type = type;
+
+  // Set the shape and allocate the memory.
+  tensor_.dims = TfLiteIntArrayCopy(&shape);
+  const size_t num_bytes = TfLiteTypeGetSize(type) * NumElements(&tensor_);
+  TfLiteTensorRealloc(num_bytes, &tensor_);
+
+  memset(tensor_.data.raw, 0, tensor_.bytes);
+  is_initialized_ = true;
+  return kTfLiteOk;
+}
+
+size_t CacheBuffer::GetNumEntries() const { return num_entries_; }
+
+void CacheBuffer::SetNumEntries(size_t count) {
+  TFLITE_DCHECK(count <= tensor_.dims->data[2]);
+  num_entries_ = count;
+}
+
+}  // namespace resource
+}  // namespace tflite

tensorflow/lite/experimental/resource/cache_buffer.h

@@ -0,0 +1,51 @@
+/* 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_RESOURCE_CACHE_BUFFER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_CACHE_BUFFER_H_
+
+#include <memory>
+#include <unordered_map>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/resource/resource_variable.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace resource {
+
+/// WARNING: Experimental interface, subject to change.
+// A Cache Buffer class. Useful for keeping the keys and values of a
+// transformer block attention mechanism in autoregressive decode.
+// Ops can access this buffer and add tensors to it. It also keeps track of the
+// number of used entries in the cache.
+class CacheBuffer : public ResourceVariable {
+ public:
+  CacheBuffer() = default;
+  CacheBuffer(const CacheBuffer &) = delete;
+  CacheBuffer &operator=(const CacheBuffer &) = delete;
+  // Initialize tensor of a certain shape using the provided type.
+  TfLiteStatus Initialize(const TfLiteIntArray &shape, const TfLiteType &type);
+  size_t GetNumEntries() const;
+  void SetNumEntries(size_t count);
+
+ private:
+  // The number of entries currently used in the buffer;
+  size_t num_entries_ = 0;
+};
+
+}  // namespace resource
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_CACHE_BUFFER_H_

tensorflow/lite/experimental/resource/cache_buffer_test.cc

@@ -0,0 +1,47 @@
+/* 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/resource/cache_buffer.h"
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/c/common.h"
+
+namespace tflite {
+namespace resource {
+
+TEST(CacheBufferTest, Initialize) {
+  TfLiteIntArray* shape = TfLiteIntArrayCreate(4);
+  shape->data[0] = 1;
+  shape->data[1] = 3;
+  shape->data[2] = 5;
+  shape->data[3] = 7;
+
+  TfLiteType type = kTfLiteFloat32;
+  CacheBuffer cache_buffer;
+  cache_buffer.Initialize(*shape, type);
+
+  EXPECT_EQ(cache_buffer.GetTensor()->type, type);
+  EXPECT_EQ(cache_buffer.GetTensor()->dims->size, 4);
+  EXPECT_EQ(cache_buffer.GetTensor()->dims->data[0], 1);
+  EXPECT_EQ(cache_buffer.GetTensor()->dims->data[1], 3);
+  EXPECT_EQ(cache_buffer.GetTensor()->bytes, 420);
+  ASSERT_NE(cache_buffer.GetTensor()->data.raw, nullptr);
+  EXPECT_EQ(cache_buffer.GetNumEntries(), 0);
+  cache_buffer.SetNumEntries(3);
+  EXPECT_EQ(cache_buffer.GetNumEntries(), 3);
+  TfLiteIntArrayFree(shape);
+}
+
+}  // namespace resource
+}  // namespace tflite

tensorflow/lite/experimental/resource/resource_variable.h

@@ -50,7 +50,7 @@ class ResourceVariable : public ResourceBase {
     return is_initialized_ ? tensor_.bytes : 0;
   }
 
- private:
+ protected:
   // The tensor (and its buffer stored in `tensor_.data` is fully owned by
   // the `ResourceVariable` object.
   TfLiteTensor tensor_;