tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2026-05-31 20:18:09 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge 637ecf6f12 into dc54a1a307

Browse Source
This commit is contained in:
Mithun Kannaa 2026-05-05 06:39:08 +03:00 committed by GitHub
commit cee0b39642
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
nanochat/gpt.py
View File

@ -119,7 +119,10 @@ class CausalSelfAttention(nn.Module):
# Advance position after last layer processes
if self.layer_idx == kv_cache.n_layers - 1:
kv_cache.advance(T)
# XSA (Exclusive Self Attention)
Vn = F.normalize(v, dim=-1)
Vn = Vn.repeat_interleave(self.n_head // self.n_kv_head, dim=2)
y = y - (y * Vn).sum(dim=-1, keepdim=True) * Vn
# Re-assemble the heads and project back to residual stream
y = y.contiguous().view(B, T, -1)
y = self.c_proj(y)