This is an automated email from the ASF dual-hosted git repository.
hongyij pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new aec773d690 [Unity][VM] Add Attention KV cache builtin (#14478)
aec773d690 is described below
commit aec773d6908924288a722b04078faf6de55e9ac0
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Apr 4 11:44:57 2023 -0400
[Unity][VM] Add Attention KV cache builtin (#14478)
This PR provides a simple implementation of inplace attention
kv cache for relax runtime. The main goal here is to help
us enable auto-regressive decoding quickly in relax.
This is likely not the only way to support attention kv-cache.
We keep the implementation private for now and will
continue to evolve the relevant code.
---
src/runtime/relax_vm/attention_kv_cache.cc | 162 +++++++++++++++++++++++++++++
tests/python/relax/test_pipeline.py | 58 ++++++++++-
tests/python/relax/test_runtime_builtin.py | 17 +++
web/emcc/wasm_runtime.cc | 1 +
4 files changed, 237 insertions(+), 1 deletion(-)
diff --git a/src/runtime/relax_vm/attention_kv_cache.cc
b/src/runtime/relax_vm/attention_kv_cache.cc
new file mode 100644
index 0000000000..f94b233975
--- /dev/null
+++ b/src/runtime/relax_vm/attention_kv_cache.cc
@@ -0,0 +1,162 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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.
+ */
+/*!
+ * \file src/runtime/relax_vm/attention_kv_cache.cc
+ * \brief A simple implementation of inplace attention kv cache for runtime.
+ *
+ * This file provides a simple implementation of inplace attention
+ * kv cache for relax runtime. The main goal here is to help us enable
+ * auto-regressive decoding quickly in relax.
+ *
+ * This is not the only way to support attention kv-cache.
+ * Our support of attention kv-cache can subject to future
+ * changes as we build more LM verticals.
+ *
+ * We will keep the impact minimum by puting it as a private
+ * runtime builtin provide as in this file.
+ *
+ * We can evolve this implementation as we build more LM verticals.
+ */
+
+#include <tvm/runtime/container/shape_tuple.h>
+#include <tvm/runtime/device_api.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/runtime/memory.h>
+#include <tvm/runtime/ndarray.h>
+#include <tvm/runtime/relax_vm/vm.h>
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+//-------------------------------------------
+// We keep the implementation private as
+// they may subject to future changes.
+//
+// Users can interact with it through the
+// runtime API function calls
+//-------------------------------------------
+/*!
+ * \brief An object representing an attention kv cache.
+ */
+class AttentionKVCacheObj : public Object {
+ public:
+ /*!
+ * \brief Underlying support data.
+ */
+ NDArray data;
+
+ /*!
+ * \brief number of slots already filled.
+ */
+ int64_t fill_count{0};
+
+ /*!
+ * \brief View all current cached values as one array.
+ * \param shape The cached values.
+ */
+ NDArray View(const ShapeTuple& shape) {
+ CHECK_EQ(shape[0], fill_count) << "Requested shape do not match the filled
count";
+ for (int i = 1; i < this->data->ndim; ++i) {
+ CHECK_EQ(shape[i], data->shape[i]) << "Dimension " << i << " mismatch";
+ }
+ return data.CreateView(shape, data->dtype);
+ }
+
+ /*!
+ * \brief Append value to the cache.
+ * \param value The value to be appended.
+ */
+ void Append(NDArray value) {
+ CHECK(data.DataType() == value.DataType()) << "dtype mismatch";
+ // reallocate cache
+ int64_t reserved_slots = data->shape[0];
+ while (fill_count + value->shape[0] > reserved_slots) {
+ reserved_slots *= 2;
+ }
+ if (reserved_slots != data->shape[0]) {
+ std::vector<int64_t> new_shape(data->shape, data->shape + data->ndim);
+ new_shape[0] = reserved_slots;
+ NDArray new_data = NDArray::Empty(new_shape, data->dtype, data->device);
+ new_data.CreateView(data.Shape(), data->dtype).CopyFrom(data);
+ this->data = new_data;
+ }
+ // copy into the fill count position.
+ ICHECK_LE(fill_count + value->shape[0], data->shape[0]);
+ ICHECK(data.IsContiguous());
+
+ int64_t num_filled_elements = fill_count;
+ for (int i = 1; i < data->ndim; ++i) {
+ CHECK_EQ(value->shape[i], data->shape[i]) << "Dimension " << i << "
mismatch";
+ num_filled_elements *= data->shape[i];
+ }
+ // create a view of copy dest to copy the value into it.
+ DLTensor copy_dst = *(data.operator->());
+ copy_dst.byte_offset = num_filled_elements * ((data->dtype.bits *
data->dtype.lanes + 7) / 8);
+ copy_dst.shape = value->shape;
+ NDArray::CopyFromTo(value.operator->(), ©_dst);
+ this->fill_count += value->shape[0];
+ }
+
+ static constexpr const uint32_t _type_index = TypeIndex::kDynamic;
+ static constexpr const char* _type_key = "relax.vm.AttentionKVCache";
+ TVM_DECLARE_FINAL_OBJECT_INFO(AttentionKVCacheObj, Object);
+};
+
+/*! \brief reference to closure. */
+class AttentionKVCache : public ObjectRef {
+ public:
+ /*!
+ * \brief Create the attention kv cache.
+ * \param init_reserve The initial reserved.
+ */
+ static AttentionKVCache Create(NDArray init_data) {
+ auto n = make_object<AttentionKVCacheObj>();
+ n->data = std::move(init_data);
+ n->fill_count = 0;
+ return AttentionKVCache(n);
+ }
+
+ TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(AttentionKVCache, ObjectRef,
AttentionKVCacheObj);
+};
+
+TVM_REGISTER_OBJECT_TYPE(AttentionKVCacheObj);
+
+//-------------------------------------------------
+// Register runtime functions
+//-------------------------------------------------
+TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_create")
+ .set_body_typed(AttentionKVCache::Create);
+
+AttentionKVCache AttentionKVCacheAppend(AttentionKVCache cache, NDArray value)
{
+ cache->Append(value);
+ return cache;
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_append").set_body_typed(AttentionKVCacheAppend);
+
+NDArray AttentionKVCacheView(AttentionKVCache cache, ShapeTuple shape) {
+ return cache->View(shape);
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_view").set_body_typed(AttentionKVCacheView);
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
diff --git a/tests/python/relax/test_pipeline.py
b/tests/python/relax/test_pipeline.py
index c66066f8f8..3c97d6b701 100644
--- a/tests/python/relax/test_pipeline.py
+++ b/tests/python/relax/test_pipeline.py
@@ -16,8 +16,9 @@
# under the License.
import numpy as np
import tvm
+import tvm.testing
from tvm import relax
-from tvm.script import relax as R
+from tvm.script import relax as R, tir as T
def test_pipeline_compile():
@@ -43,3 +44,58 @@ def test_pipeline_compile():
vm = relax.VirtualMachine(ex, tvm.cpu())
z = vm["main"](x, y)
tvm.testing.assert_allclose(z.numpy(), x_np + y_np, rtol=1e-7, atol=1e-7)
+
+
+def test_pipeline_with_kv_cache():
+ """A dummy pipline that simulates KV update."""
+ pipeline = relax.get_pipeline()
+
+ @tvm.script.ir_module
+ class Mod:
+ @R.function
+ def main(
+ x: R.Tensor((1, 4), "float32"),
+ y: R.Tensor((1, 4), "float32"),
+ shape: R.Shape(["L", 4]),
+ kv_cache: R.Object,
+ ):
+ L = T.int64()
+ # computation of the current value
+ curr_value = R.add(x, y)
+ # update cache
+ kv_cache = R.call_packed(
+ "vm.builtin.attention_kv_cache_append", kv_cache, curr_value,
sinfo_args=[R.Object]
+ )
+ # return the updated cache view
+ kv = R.call_packed(
+ "vm.builtin.attention_kv_cache_view",
+ kv_cache,
+ shape,
+ sinfo_args=[R.Tensor((L, 4), "float32")],
+ )
+ return (kv, kv_cache)
+
+ mod = Mod
+ mod = pipeline(mod)
+
+ target = tvm.target.Target("llvm", host="llvm")
+
+ ex = relax.build(mod, target)
+
+ num_steps = 8
+ cache_np = np.empty((num_steps, 4), dtype="float32")
+
+ fcreate_cache = tvm.get_global_func("vm.builtin.attention_kv_cache_create")
+ kv_cache = fcreate_cache(tvm.nd.empty((2, 4), device=tvm.cpu(),
dtype="float32"))
+
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ for i in range(num_steps):
+ x_np = np.random.rand(1, 4).astype(np.float32)
+ y_np = np.random.rand(1, 4).astype(np.float32)
+ x = tvm.nd.array(x_np)
+ y = tvm.nd.array(y_np)
+ np_shape = (i + 1, 4)
+ kv, kv_cache = vm["main"](x, y, tvm.runtime.ShapeTuple(np_shape),
kv_cache)
+
+ cache_np[i, :] = x_np + y_np
+ tvm.testing.assert_allclose(kv.numpy(), cache_np[: np_shape[0], :],
rtol=1e-7, atol=1e-7)
diff --git a/tests/python/relax/test_runtime_builtin.py
b/tests/python/relax/test_runtime_builtin.py
index b4ba54b455..bb513eb357 100644
--- a/tests/python/relax/test_runtime_builtin.py
+++ b/tests/python/relax/test_runtime_builtin.py
@@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
import tvm
+import tvm.testing
import pytest
import numpy as np
@@ -149,5 +150,21 @@ def test_tuple_getitem():
assert tuple_getitem(t, 1) == y
+def test_attention_kv_cache():
+ fcreate = tvm.get_global_func("vm.builtin.attention_kv_cache_create")
+ fappend = tvm.get_global_func("vm.builtin.attention_kv_cache_append")
+ fview = tvm.get_global_func("vm.builtin.attention_kv_cache_view")
+
+ cache = fcreate(tvm.nd.empty((2, 2), dtype="int32"))
+ num_steps = 0
+ for i in range(num_steps):
+ cache = fappend(cache, tvm.nd.array(i * np.ones((1,
2).astype("int32"))))
+
+ res = fview(cache, tvm.runtime.ShapeTuple((num_steps, 2))).numpy()
+ for i in range(num_steps):
+ assert res[i][0] == i
+ assert res[i][1] == i
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc
index 8f16365eee..9b2b6c180c 100644
--- a/web/emcc/wasm_runtime.cc
+++ b/web/emcc/wasm_runtime.cc
@@ -53,6 +53,7 @@
#include "src/runtime/system_library.cc"
#include "src/runtime/workspace_pool.cc"
// relax setup
+#include "src/runtime/relax_vm/attention_kv_cache.cc"
#include "src/runtime/relax_vm/builtin.cc"
#include "src/runtime/relax_vm/bytecode.cc"
#include "src/runtime/relax_vm/executable.cc"
