From 2b9c0855592dbd128e5a44109a95ba261c1e0201 Mon Sep 17 00:00:00 2001 From: Haowei Zhang Date: Mon, 27 Oct 2025 02:47:13 -0700 Subject: [PATCH] update the kv_shape --- nanochat/engine.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nanochat/engine.py b/nanochat/engine.py index 307590b..44ed16b 100644 --- a/nanochat/engine.py +++ b/nanochat/engine.py @@ -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