Implement cache ops

2023-02-16 07:47:03 +00:00
parent a1c67e6db8
commit 6f058c7ba8
4 changed files with 109 additions and 6 deletions
--- a/csrc/cache.cpp
+++ b/csrc/cache.cpp
@@ -0,0 +1,20 @@
+#include <torch/extension.h>
+
+void copy_blocks(
+  torch::Tensor& src,
+  torch::Tensor& dst,
+  const std::map<int64_t, int64_t>& block_mapping);
+
+void copy_cache_blocks(
+  torch::Tensor& src,
+  torch::Tensor& dst,
+  const std::map<int64_t, int64_t>& block_mapping) {
+  copy_blocks(src, dst, block_mapping);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def(
+    "copy_cache_blocks",
+    &copy_cache_blocks,
+    "Copy the cache blocks from src to dst");
+}