ros/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2026-01-30 04:22:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

update the kv_shape

Browse Source
This commit is contained in:
Haowei Zhang
2025-10-27 02:47:13 -07:00
parent b062b422ac
commit 2b9c085559
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
nanochat/engine.py
View File

@@ -139,6 +139,7 @@ class KVCache:
additional_shape[4] = t_needed - self.kv_cache.size(4)
additional_cache = torch.empty(additional_shape, dtype=k.dtype, device=k.device)
self.kv_cache = torch.cat([self.kv_cache, additional_cache], dim=4).contiguous()
self.kv_shape = self.kv_cache.shape
# Insert k, v into the cache
self.kv_cache[layer_idx, 0, :, :, t0:t1] = k
self.kv_cache[layer_idx, 1, :, :, t0:t1] = v