cache.cpp 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647
  1. #include <cstdint>
  2. #include <torch/extension.h>
  3. #include <map>
  4. #include <vector>
  5. void swap_blocks(
  6. torch::Tensor& src,
  7. torch::Tensor& dst,
  8. const std::map<int64_t, int64_t>& block_mapping);
  9. void copy_blocks(
  10. std::vector<torch::Tensor>& key_caches,
  11. std::vector<torch::Tensor>& value_caches,
  12. const std::map<int64_t, std::vector<int64_t>>& block_mapping);
  13. void reshape_and_cache(
  14. torch::Tensor& key,
  15. torch::Tensor& value,
  16. torch::Tensor& key_cache,
  17. torch::Tensor& value_cache,
  18. torch::Tensor& slot_mapping);
  19. void gather_cached_kv(
  20. torch::Tensor& key,
  21. torch::Tensor& value,
  22. torch::Tensor& key_cache,
  23. torch::Tensor& value_cache,
  24. torch::Tensor& slot_mapping);
  25. PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  26. m.def(
  27. "swap_blocks",
  28. &swap_blocks,
  29. "Swap in (out) the cache blocks from src to dst");
  30. m.def(
  31. "copy_blocks",
  32. &copy_blocks,
  33. "Copy the cache blocks from src to dst");
  34. m.def(
  35. "reshape_and_cache",
  36. &reshape_and_cache,
  37. "Reshape the key and value tensors and cache them");
  38. m.def(
  39. "gather_cached_kv",
  40. &gather_cached_kv,
  41. "Gather key and value from the cache into contiguous KV tensors");
  42. }