Firehawk52-backup/llama
28
0
Fork 0
mirror of https://github.com/meta-llama/llama.git synced 2026-01-15 08:22:55 -03:00
Code Issues Actions Packages Projects Releases Wiki Activity

Correct KV comment seqlen -> seqlen + cache_len

Browse Source 
Update and add comments about the shape of the key and value
matrices in the attention component. E.g., the second dimension is
of length seqlen + cache_len not seqlen as previously stated.
This commit is contained in:
flu0r1ne
2023-11-13 14:05:24 -06:00
parent 6b3154bfbb
commit cd0719ddb4
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
llama/model.py
View File

@@ -289,12 +289,12 @@ class Attention(nn.Module):
values = self.cache_v[:bsz, : start_pos + seqlen]
# repeat k/v heads if n_kv_heads < n_heads
keys = repeat_kv(keys, self.n_rep) # (bs, seqlen, n_local_heads, head_dim)
values = repeat_kv(values, self.n_rep) # (bs, seqlen, n_local_heads, head_dim)
keys = repeat_kv(keys, self.n_rep) # (bs, cache_len + seqlen, n_local_heads, head_dim)
values = repeat_kv(values, self.n_rep) # (bs, cache_len + seqlen, n_local_heads, head_dim)
xq = xq.transpose(1, 2) # (bs, n_local_heads, seqlen, head_dim)
keys = keys.transpose(1, 2)
values = values.transpose(1, 2)
keys = keys.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim)
values = values.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim)
scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
if mask is not None:
scores = scores + mask # (bs, n_local_heads, seqlen, cache_len + seqlen)