happyz
/
gemma.cpp
mirror of https://github.com/google/gemma.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #172 from szabadka:gemma2 

PiperOrigin-RevId: 629438917
This commit is contained in:
Copybara-Service 2024-04-30 09:33:38 -07:00
parent 374fd7478a f8ccb8e37c
commit 8f04a8346d
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
gemma/gemma.cc
View File

@ -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