Implement block copy kernel to optimize beam search (#32)

2023-04-07 17:45:07 -07:00
parent a490aafa36
commit 0f40557af6
6 changed files with 154 additions and 48 deletions
--- a/csrc/cache.cpp
+++ b/csrc/cache.cpp
@@ -9,8 +9,8 @@ void swap_blocks(
  const std::map<int64_t, int64_t>& block_mapping);

 void copy_blocks(
-  torch::Tensor& src,
-  torch::Tensor& dst,
+  std::vector<torch::Tensor>& key_caches,
+  std::vector<torch::Tensor>& value_caches,
  const std::map<int64_t, std::vector<int64_t>>& block_mapping);

 void reshape_and_cache(