books/LLMs-from-scratch
1
0
Fork 0
mirror of https://github.com/rasbt/LLMs-from-scratch.git synced 2026-04-10 12:33:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add GPT-2 KV cache to pkg (#687)

Browse Source
This commit is contained in:
Sebastian Raschka
2025-06-21 12:29:04 -05:00
committed by GitHub
parent 9d62ca0598
commit bb57756444
4 changed files with 315 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
ch04/03_kv-cache/gpt_with_kv_cache_optimized.py
View File

@@ -80,8 +80,6 @@ class MultiHeadAttention(nn.Module):
keys, values = keys_new, values_new
self.ptr_cur = 0 # keep pointer sane if you interleave modes
####################################################
# Compute scaled dot-product attention (aka self-attention) with a causal mask
attn_scores = queries @ keys.transpose(2, 3) # Dot product for each head