Optimize MQA Kernel (#452)

2026-09-13 13:10:13 +08:00 · 2023-07-14 20:06:40 -04:00
parent dbed69058c
commit 96853af5a8
5 changed files with 84 additions and 72 deletions
@@ -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,