fam/nanochat-omni
1
0
Fork 0
Code Issues Pull Requests Actions 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
nanochat/engine.py
+1
View File
@@ -139,6 +139,7 @@ class KVCache:
additional_shape[4] = t_needed - self.kv_cache.size(4) additional_shape[4] = t_needed - self.kv_cache.size(4)
additional_cache = torch.empty(additional_shape, dtype=k.dtype, device=k.device) 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_cache = torch.cat([self.kv_cache, additional_cache], dim=4).contiguous()
self.kv_shape = self.kv_cache.shape
# Insert k, v into the cache # Insert k, v into the cache
self.kv_cache[layer_idx, 0, :, :, t0:t1] = k self.kv_cache[layer_idx, 0, :, :, t0:t1] = k
self.kv_cache[layer_idx, 1, :, :, t0:t1] = v self.kv_cache[layer_idx, 1, :, :, t0:t1] = v