Optimize MQA Kernel (#452)

2023-07-14 20:06:40 -04:00
parent dbed69058c
commit 96853af5a8
5 changed files with 84 additions and 72 deletions
--- a/csrc/attention.cpp
+++ b/csrc/attention.cpp
@@ -6,6 +6,7 @@ void single_query_cached_kv_attention(
  torch::Tensor& query,
  torch::Tensor& key_cache,
  torch::Tensor& value_cache,
+  torch::Tensor& head_mapping,
  float scale,
  torch::Tensor& block_tables,
  torch::Tensor& context_lens,