From d7787c8f6cfe58a4f1171c664f4ba1ae02cb663d Mon Sep 17 00:00:00 2001
From: RangerUFO <ufownl@gmail.com>
Date: Fri, 21 Jun 2024 13:06:26 +0800
Subject: [PATCH] Fix KV cache size calculation error

---
 gemma/gemma.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gemma/gemma.cc b/gemma/gemma.cc
index f32cde5..1c07c9e 100644
--- a/gemma/gemma.cc
+++ b/gemma/gemma.cc
@@ -114,7 +114,7 @@ struct CreateKVCache {
   KVCache operator()() const {
     KVCache kv_cache = {};
 
-    const size_t size_cache_pos = CacheLayerSize<TConfig>()();
+    const size_t size_cache_pos = CachePosSize<TConfig>()();
     if (size_cache_pos != 0) {
       const size_t seq_len = TConfig::kSeqLen + kPrefillBatchSize;
       kv_cache.kv_cache = hwy::AllocateAligned<float>(seq_len * size_cache_pos);