yujunjun/LLMs-from-scratch
1
0
Fork 0
mirror of https://github.com/rasbt/LLMs-from-scratch.git synced 2025-10-28 16:30:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix bug in masking when kv cache is used.

Browse Source
This commit is contained in:
martinzwm 2025-06-22 14:06:00 -07:00
parent 01be5a42e4
commit ffc5e4e5d6
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
ch04/03_kv-cache/gpt_with_kv_cache.py
View File

@ -72,7 +72,9 @@ class MultiHeadAttention(nn.Module):
attn_scores = queries @ keys.transpose(2, 3) # Dot product for each head
# Original mask truncated to the number of tokens and converted to boolean
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
num_tokens_Q = queries.shape[-2]
num_tokens_K = keys.shape[-2]
mask_bool = self.mask.bool()[:num_tokens_Q, :num_tokens_K]
# Use the mask to fill attention scores
attn_scores.masked_fill_(mask_bool, -torch.inf)