init

csrc_musa/cache.h

@@ -0,0 +1,38 @@
+#pragma once
+
+#include <torch/extension.h>
+
+#include <map>
+#include <vector>
+
+void swap_blocks(
+  torch::Tensor& src,
+  torch::Tensor& dst,
+  const std::map<int64_t, int64_t>& block_mapping);
+
+void copy_blocks(
+  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(
+  torch::Tensor& key,
+  torch::Tensor& value,
+  torch::Tensor& key_cache,
+  torch::Tensor& value_cache,
+  torch::Tensor& slot_mapping,
+  const std::string& kv_cache_dtype,
+  const float kv_scale);
+
+void reshape_and_cache_flash(
+  torch::Tensor& key,
+  torch::Tensor& value,
+  torch::Tensor& key_cache,
+  torch::Tensor& value_cache,
+  torch::Tensor& slot_mapping,
+  const std::string& kv_cache_dtype);
+
+// Just for unittest
+void convert_fp8(
+  torch::Tensor& src_cache,
+  torch::Tensor& dst_cache);