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

Support various block sizes & Change default block size to 16 (#38)

Browse Source
This commit is contained in:
Woosuk Kwon
2023-04-15 09:03:24 -07:00
committed by GitHub
parent 84eee24e20
commit 0f4b32199e
7 changed files with 594 additions and 611 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
csrc/attention.cpp
View File

@@ -11,25 +11,9 @@ void single_query_cached_kv_attention(
int block_size,
int max_context_len);
void multi_query_cached_kv_attention(
torch::Tensor& cu_query_lens,
torch::Tensor& out,
torch::Tensor& query,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
float scale,
torch::Tensor& block_tables,
torch::Tensor& context_lens,
int block_size,
int max_context_len);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def(
"single_query_cached_kv_attention",
&single_query_cached_kv_attention,
"Compute the attention between an input query and the cached key/value tensors");
m.def(
"multi_query_cached_kv_attention",
&multi_query_cached_kv_attention,
"Compute the attention between multiple input queries and the cached key/value tensors");
}