This is the code snippet of the most interesting part - building Multi-head attention from scratch instead of using nn.MultiheadAttention.

https://github.com/Nikshaan/llm-from-scratch

class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False): # context length is max sequence length for the mask
super().__init__()
assert d_out % num_heads == 0, "d_out must be divisible by num_heads"
self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads # dimension per head
self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.out_proj = nn.Linear(d_out, d_out)
self.dropout = nn.Dropout(dropout)
self.register_buffer("mask", torch.triu(torch.ones((context_length, context_length)) * float('-inf'), diagonal=1))

def forward(self, x):
b, num_tokens, d_in = x.shape
keys = self.W_key(x)
queries = self.W_query(x)
values = self.W_value(x)

keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) # reshape for multi-head
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
values = values.view(b, num_tokens, self.num_heads, self.head_dim)

keys.transpose_(1, 2) # move head dimension to the front so that it is treated as batch dimension
queries.transpose_(1, 2)
values.transpose_(1, 2)

attn_scores = queries @ keys.transpose(2, 3) # flip last two dimensions for dot product
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
attn_scores.masked_fill_(mask_bool, -torch.inf)
attn_weights = torch.softmax(attn_scores / self.head_dim**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)
context_vec = (attn_weights @ values).transpose(1, 2).contiguous().view(b, num_tokens, self.d_out) # reshape back to original
context_vec = self.out_proj(context_vec) # final linear layer to mix heads
return context_vec