From f8ccb8e37c661861e4ab33ba09a8fbb102fb3242 Mon Sep 17 00:00:00 2001 From: Zoltan Szabadka Date: Tue, 30 Apr 2024 16:19:14 +0000 Subject: [PATCH] Fix kv offset computation for MHA config. --- gemma/gemma.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/gemma/gemma.cc b/gemma/gemma.cc index a494b3b..05d1aaf 100644 --- a/gemma/gemma.cc +++ b/gemma/gemma.cc @@ -787,11 +787,12 @@ HWY_NOINLINE void Attention(size_t batch_start, size_t batch_idx, size_t layer, ProjQ(head, q_offset); const size_t kv_offset = - cache_pos * kCachePosSize + layer * kCacheLayerSize + head * kQKVDim; + cache_pos * kCachePosSize + layer * kCacheLayerSize + + head * kQKVDim * 2; ProjKV(k_offset, v_offset, kv_offset); - Attn(head, head * kQKVDim); + Attn(head, head * kQKVDim * 2); }); } else { // Multi-Query Attention