youngkingdom/vllm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
youkaichao
2024-05-08 12:07:05 -07:00
committed by GitHub
parent ad932a221d
commit 20cfcdec99
21 changed files with 137 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
csrc/cpu/cache.cpp
View File

@ -83,7 +83,7 @@ void reshape_and_cache_cpu_impl(
void copy_blocks(std::vector<torch::Tensor> &key_caches,
std::vector<torch::Tensor> &value_caches,
torch::Tensor& block_mapping) {
const torch::Tensor& block_mapping) {
int num_layers = key_caches.size();
TORCH_CHECK(num_layers == value_caches.size());
if (num_layers == 0) {
@ -128,6 +128,6 @@ void reshape_and_cache(torch::Tensor &key, torch::Tensor &value,
}
void swap_blocks(torch::Tensor &src, torch::Tensor &dst,
const std::map<int64_t, int64_t> &block_mapping) {
const torch::Tensor&block_mapping) {
TORCH_CHECK(false, "swap_blocks is unsupported on CPU.")
}