Merge pull request #266 from ufownl:bugfix/kvcache

PiperOrigin-RevId: 645329504
This commit is contained in:
Copybara-Service 2024-06-21 03:06:52 -07:00
commit 82f16087ba
1 changed files with 1 additions and 1 deletions

View File

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