cache.h 1.1 KB

12345678910111213141516171819202122232425262728
  1. #pragma once
  2. #include <torch/extension.h>
  3. #include <map>
  4. #include <vector>
  5. void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
  6. const torch::Tensor& block_mapping);
  7. void copy_blocks(std::vector<torch::Tensor>& key_caches,
  8. std::vector<torch::Tensor>& value_caches,
  9. const torch::Tensor& block_mapping);
  10. void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
  11. torch::Tensor& key_cache, torch::Tensor& value_cache,
  12. torch::Tensor& slot_mapping,
  13. const std::string& kv_cache_dtype, const float kv_scale);
  14. void reshape_and_cache_flash(torch::Tensor& key, torch::Tensor& value,
  15. torch::Tensor& key_cache,
  16. torch::Tensor& value_cache,
  17. torch::Tensor& slot_mapping,
  18. const std::string& kv_cache_dtype);
  19. // Just for unittest
  20. void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
  21. const float scale, const std::string& kv_cache_dtype);