tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2025-12-06 04:12:13 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

minor fix comment in engine

Browse Source
This commit is contained in:
Andrej Karpathy 2025-11-13 15:28:18 +00:00
parent adb5d4a16c
commit 91f09ccd0d
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
nanochat/engine.py
View File

@ -107,8 +107,9 @@ class KVCache:
assert self.kv_cache is None, "Cannot prefill a non-empty KV cache" assert self.kv_cache is None, "Cannot prefill a non-empty KV cache"
assert other.kv_cache is not None, "Cannot prefill with a None KV cache" assert other.kv_cache is not None, "Cannot prefill with a None KV cache"
for ix, (dim1, dim2) in enumerate(zip(self.kv_shape, other.kv_shape)): for ix, (dim1, dim2) in enumerate(zip(self.kv_shape, other.kv_shape)):
# ix 0: num_layers, 1: k/v, 2: batch_size, 3: num_heads, 4: seq_len, 5: head_dim
if ix in [0, 1, 3, 5]: if ix in [0, 1, 3, 5]:
# num_layers, batch_size, num_heads, head_dim must match # num_layers, k/v, num_heads, head_dim must match
assert dim1 == dim2, f"Dim {ix} mismatch: {dim1} != {dim2}" assert dim1 == dim2, f"Dim {ix} mismatch: {dim1} != {dim2}"
elif ix == 2: elif ix == 2:
# batch_size can be expanded # batch_size can be expanded