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

Basic attention kernel that supports cached KV + (multi-)prompts (#24)

Browse Source
This commit is contained in:
Siyuan (Ryans) Zhuang
2023-04-04 20:34:46 -07:00
committed by GitHub
parent 897cb2ae28
commit 21b3671bbc
3 changed files with 622 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
csrc/attention.cpp
View File

@ -11,9 +11,25 @@ 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");
}