From b062b422ac6c4578b986eae487645aef42eb9152 Mon Sep 17 00:00:00 2001 From: Haowei Zhang Date: Mon, 27 Oct 2025 02:23:08 -0700 Subject: [PATCH] Fix kv cache, given resize will destroys the logical structure --- nanochat/engine.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/nanochat/engine.py b/nanochat/engine.py index fee06a1..307590b 100644 --- a/nanochat/engine.py +++ b/nanochat/engine.py @@ -135,9 +135,10 @@ class KVCache: if t1 > self.kv_cache.size(4): t_needed = t1 + 1024 # as much as we need plus buffer of 1024 t_needed = (t_needed + 1023) & ~1023 # then round up to the nearest multiple of 1024 - current_shape = list(self.kv_cache.shape) - current_shape[4] = t_needed - self.kv_cache.resize_(current_shape) + additional_shape = list(self.kv_cache.shape) + 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() # Insert k, v into the cache self.kv_cache[layer_idx, 0, :, :, t0:t1] = k self.kv_cache[layer_idx, 1, :, :, t0:t1] = v