[Core][Optimization] change python dict to pytorch tensor for blocks to swap (#4659)

2024-05-08 12:07:05 -07:00
parent ad932a221d
commit 20cfcdec99
21 changed files with 137 additions and 109 deletions
--- a/csrc/cache.h
+++ b/csrc/cache.h
@@ -8,12 +8,12 @@
 void swap_blocks(
  torch::Tensor& src,
  torch::Tensor& dst,
-  const std::map<int64_t, int64_t>& block_mapping);
+  const torch::Tensor& block_mapping);

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

 void reshape_and_cache(
  torch::Tensor& key,