Description

Level 2 — How It Works (The Mechanics)

1. Linear Projections into Q, K, V: We project using three learned weight matrices :

2. Scaled Dot-Product Attention:
• computes the pairwise similarity score between every token and every token .
• Why divide by ? If components of and are independent random variables with mean and variance , their dot product has mean and variance (standard deviation ). For large (e.g., or ), unscaled dot products grow large in magnitude, pushing softmax into regions where one entry is and all others are , causing vanishing gradients. Dividing by restores unit variance ().

3. Causal Masking ():
• In autoregressive (decoder-only) models like GPT and LLaMA, we train the model to predict the next token for all positions simultaneously in one forward pass.
• To prevent position from cheating by looking at position , we add an upper-triangular mask matrix where if and if . After softmax, , guaranteeing zero information leakage from the future.

4. Multi-Head Attention (MHA):
• Instead of one single attention pass of dimension , we split into heads, each of dimension :

Level 3 — Advanced Depth (MHA vs MQA vs GQA & FlashAttention)

• Why Multi-Head instead of Single-Head? Softmax is a weighted average: in a single head, one dominant similarity score drowns out secondary relationships. Multiple heads project into distinct -dimensional subspaces so the model can attend to syntactic structure, coreference, and positional neighbors simultaneously.

• Grouped-Query Attention (GQA) & Multi-Query Attention (MQA): In standard MHA, every head has its own and tensors, making the KV cache huge during inference. MQA shares a single head across all Query heads (fast, slight quality drop). GQA (used in LLaMA-3, Mistral) groups Query heads into groups (e.g., Query heads sharing KV heads), cutting KV-cache memory by with virtually zero accuracy loss.

• FlashAttention (IO-Aware Exact Attention): Standard attention materializes the matrix in slow GPU High-Bandwidth Memory (HBM), making it memory-bandwidth bound and in memory. FlashAttention tiles into fast on-chip SRAM and uses the online softmax trick (rescaling running max and denominator on the fly) to compute exact attention in HBM memory without ever writing the matrix to HBM.

Learning notes

Multi-Head Causal Self-Attention Architecture

Visualizing the Causal Mask Matrix ()

Raw Scaled Scores (QK^T / sqrt(d_k)):       After Causal Mask (+ M):                After Softmax (Probabilities):
       [The]  [cat]  [sat]  [down]                 [The]  [cat]  [sat]  [down]             [The]  [cat]  [sat]  [down]
[The]  [ 2.1,   0.8,  -0.4,   1.2 ]         [The]  [ 2.1,  -inf,  -inf,  -inf ]     [The]  [1.00,  0.00,  0.00,  0.00]
[cat]  [ 0.9,   2.5,   1.1,   0.3 ]   ==>   [cat]  [ 0.9,   2.5,  -inf,  -inf ] ==> [cat]  [0.17,  0.83,  0.00,  0.00]
[sat]  [ 0.2,   1.8,   2.0,  -0.5 ]         [sat]  [ 0.2,   1.8,   2.0,  -inf ]     [sat]  [0.08,  0.41,  0.51,  0.00]
[down] [ 0.1,   0.9,   1.7,   2.2 ]         [down] [ 0.1,   0.9,   1.7,   2.2 ]     [down] [0.06,  0.14,  0.31,  0.49]

MHA vs. GQA vs. MQA Comparison

Example or implementation

Clean PyTorch Implementation of Causal Multi-Head Self-Attention

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalMultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, max_seq_len: int = 2048):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        # Fused QKV projection for GPU efficiency
        self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)

        # Register lower-triangular causal mask as a non-trainable buffer
        mask = torch.tril(torch.ones(max_seq_len, max_seq_len))
        self.register_buffer("causal_mask", mask.view(1, 1, max_seq_len, max_seq_len))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, C = x.shape  # Batch, SeqLen, d_model

        # 1. Project and split into Q, K, V -> each (B, n_heads, N, d_k)
        qkv = self.qkv_proj(x).chunk(3, dim=-1)
        q, k, v = [t.view(B, N, self.n_heads, self.d_k).transpose(1, 2) for t in qkv]

        # 2. Scaled Dot-Product Scores: (B, n_heads, N, N)
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)

        # 3. Apply Causal Mask: fill future positions (mask == 0) with -inf
        scores = scores.masked_fill(self.causal_mask[:, :, :N, :N] == 0, float("-inf"))

        # 4. Softmax probabilities & weighted sum of Values
        attn_weights = F.softmax(scores, dim=-1)
        context = attn_weights @ v  # (B, n_heads, N, d_k)

        # 5. Re-concatenate heads and apply output projection W_O
        context = context.transpose(1, 2).contiguous().view(B, N, C)
        return self.out_proj(context)

# Quick verification
mha = CausalMultiHeadAttention(d_model=512, n_heads=8)
x = torch.randn(2, 16, 512)
out = mha(x)
print("Input shape:", x.shape, "-> Output shape:", out.shape)

Interview-ready explanation

Questions to revisit

Q1: Why do we set masked positions to BEFORE softmax instead of setting probabilities to AFTER softmax?

Q2: Does Multi-Head Attention with heads cost more compute than Single-Head Attention?

Q3: What is the computational complexity of Self-Attention with respect to sequence length and model dimension ?