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

Improve MHA einsum (#781)

Browse Source 
Efficiency update for einsum as mentioned in #772
This commit is contained in:
Jestine Paul
2025-08-23 04:12:26 +08:00
committed by GitHub
parent 670f7a4dd0
commit a3a62c509a
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
ch03/02_bonus_efficient-multihead-attention/mha-implementations.ipynb
View File

@@ -519,7 +519,7 @@
" scores = torch.einsum(\"bhnd,bhmd->bhnm\", Q, K) / (self.head_dim ** 0.5)\n",
"\n",
" # Apply mask\n",
" mask = self.mask[:n, :n].unsqueeze(0).unsqueeze(1).expand(b, self.num_heads, n, n)\n",
" mask = self.mask[:n, :n]\n",
" scores = scores.masked_fill(mask.bool(), -torch.inf)\n",
"\n",
" # Softmax and dropout\n",