21 lines
458 B
C++
21 lines
458 B
C++
|
#include <torch/extension.h>
|
||
|
|
|
||
|
|
void copy_blocks(
|
||
|
|
torch::Tensor& src,
|
||
|
|
torch::Tensor& dst,
|
||
|
|
const std::map<int64_t, int64_t>& block_mapping);
|
||
|
|
|
||
|
|
void copy_cache_blocks(
|
||
|
|
torch::Tensor& src,
|
||
|
|
torch::Tensor& dst,
|
||
|
|
const std::map<int64_t, int64_t>& block_mapping) {
|
||
|
|
copy_blocks(src, dst, block_mapping);
|
||
|
|
}
|
||
|
|
|
||
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||
|
|
m.def(
|
||
|
|
"copy_cache_blocks",
|
||
|
|
©_cache_blocks,
|
||
|
|
"Copy the cache blocks from src to dst");
|
||
|
|
}
|