github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

hotfix: that repeat_kv belongs outside the if

Browse Source
This commit is contained in:
George Hotz
2025-05-11 18:41:57 -07:00
parent 98c84a711d
commit 8864ff894b
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
extra/models/llama.py
View File

@@ -82,11 +82,11 @@ class Attention:
keys = self.cache_kv[0, :, 0:start_pos+seqlen, :, :]
values = self.cache_kv[1, :, 0:start_pos+seqlen, :, :]
keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
else:
assert start_pos == 0
keys, values = xk, xv
keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
attn = attn.reshape(bsz, seqlen, -1)