--- title: "Building a Transformer from Scratch: Attention, Architecture, and Memory" description: "A ground-up implementation of the Transformer: scaled dot-product attention in NumPy and PyTorch, the full decoder block with RoPE, RMSNorm, and causal masking, then the memory optimizations that make large models practical — FlashAttention, KV caching, GQA, and LoRA." date: "May 10, 2024" url: "https://himanshuat.com/blogs/understanding-transformers-implementation" --- # Building a Transformer from Scratch: Attention, Architecture, and Memory An RNN folds the entire input sequence into a single fixed-size vector before making a prediction. Attention removes that constraint — the model computes similarity weights between tokens and retrieves information dynamically across the full sequence at each step. Mathematically, scaled dot-product attention maps query vectors against key-value pairs: $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$ This post builds the mechanism from raw tensor operations, assembles it into a full decoder block, then works through the optimizations that make it run at scale. --- ## Part 1: Scaled dot-product attention ### The query–key–value abstraction Every token simultaneously acts as a query, a key, and a value. When the model processes token $i$, its query vector $q_i$ is compared against every key vector $k_j$ via an inner product. Softmax normalizes those scores into a probability distribution, which weights a sum over all value vectors $v_j$. ```mermaid flowchart LR E[Embeddings] --> Q[Query Q] E --> K[Key K] E --> V[Value V] Q --> S[Dot Product Scores] K --> S S --> SC[Scale by 1 / sqrt dk] SC --> SM[Softmax] SM --> O[Weighted Sum] V --> O ``` Every arrow is a batched matrix multiplication — the whole operation parallelizes across GPU tensor cores with no recurrent loop. ```python import numpy as np import torch import torch.nn as nn import math seq_len = 5 d_model = 8 d_k = d_model input_embeddings = np.random.rand(1, seq_len, d_model) ``` ### Step 1: Project into Q, K, V spaces ```python # Learned projection weights W_Q = np.random.randn(d_model, d_k) * 0.01 W_K = np.random.randn(d_model, d_k) * 0.01 W_V = np.random.randn(d_model, d_k) * 0.01 X = input_embeddings[0] # (seq_len, d_model) Q = X @ W_Q # (seq_len, d_k) K = X @ W_K V = X @ W_V print(f"Q shape: {Q.shape}") # (5, 8) ``` ### Step 2: Compute attention scores and apply causal mask $$ \text{scores} = \frac{QK^T}{\sqrt{d_k}} $$ Without scaling, dot products grow with $d_k$, pushing softmax into regions where gradients vanish. ```python scores = Q @ K.T / np.sqrt(d_k) # (seq_len, seq_len) # Causal mask: token i cannot attend to j > i mask = np.triu(np.ones((seq_len, seq_len)), k=1) * -1e9 scores += mask def softmax(x): x -= x.max(axis=-1, keepdims=True) e = np.exp(x) return e / e.sum(axis=-1, keepdims=True) attention_weights = softmax(scores) # (seq_len, seq_len) output = attention_weights @ V # (seq_len, d_k) print(f"Attention output shape: {output.shape}") # (5, 8) ``` ### Step 3: Multi-head attention in PyTorch Multi-head attention runs $h$ attention heads in parallel, each projecting into a lower-dimensional subspace ($d_k = d_{\text{model}} / h$), then concatenates and projects the results: $$ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O $$ ```python class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, n_heads: int): super().__init__() assert d_model % n_heads == 0 self.d_model = d_model self.n_heads = n_heads self.head_dim = d_model // n_heads self.q_proj = nn.Linear(d_model, d_model, bias=False) self.k_proj = nn.Linear(d_model, d_model, bias=False) self.v_proj = nn.Linear(d_model, d_model, bias=False) self.out_proj = nn.Linear(d_model, d_model, bias=False) def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: b, s, _ = x.shape Q = self.q_proj(x).view(b, s, self.n_heads, self.head_dim).transpose(1, 2) K = self.k_proj(x).view(b, s, self.n_heads, self.head_dim).transpose(1, 2) V = self.v_proj(x).view(b, s, self.n_heads, self.head_dim).transpose(1, 2) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) weights = torch.softmax(scores, dim=-1) out = torch.matmul(weights, V) out = out.transpose(1, 2).contiguous().view(b, s, self.d_model) return self.out_proj(out) ``` --- ## Part 2: The full decoder block ### RMSNorm and Pre-LN residuals Modern LLMs use Pre-LN placement — normalize before each sub-layer, not after. This keeps gradient flow stable at large depths. RMSNorm replaces LayerNorm's mean subtraction with root-mean-square normalization only: $$ \text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \cdot \gamma, \quad \text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^{d} x_i^2} $$ ```python class RMSNorm(nn.Module): def __init__(self, d_model: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(d_model)) def forward(self, x: torch.Tensor) -> torch.Tensor: rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) return (x / rms) * self.weight ``` ### Rotary Position Embeddings (RoPE) RoPE encodes position by rotating Q and K vectors in 2D subspaces. Unlike absolute positional embeddings, rotation preserves the relative distance between tokens and extends naturally to longer sequences: $$ q_m \cdot k_n = \text{Re}\left[\sum_{j} q_{m,j} k_{n,j}^* e^{i(m-n)\theta_j}\right] $$ ```python def precompute_rope_freqs(head_dim: int, max_seq_len: int, base: float = 10000.0) -> torch.Tensor: theta = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) t = torch.arange(max_seq_len) freqs = torch.outer(t, theta) return torch.polar(torch.ones_like(freqs), freqs) def apply_rope(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) freqs_cis = freqs_cis[:x.shape[-2], :].unsqueeze(0).unsqueeze(0) x_rotated = torch.view_as_real(x_complex * freqs_cis).flatten(-2) return x_rotated.type_as(x) ``` ### Feed-forward with SwiGLU The SwiGLU activation replaces ReLU in modern decoder FFNs. It uses a gating mechanism that allows the network to suppress activations selectively: $$ \text{SwiGLU}(x) = \text{SiLU}(W_1 x) \otimes W_3 x $$ ```python class FeedForward(nn.Module): def __init__(self, d_model: int, d_ff: int): super().__init__() self.gate_proj = nn.Linear(d_model, d_ff, bias=False) self.up_proj = nn.Linear(d_model, d_ff, bias=False) self.down_proj = nn.Linear(d_ff, d_model, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down_proj( torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x) ) ``` ### Full decoder block ```python class TransformerDecoderBlock(nn.Module): def __init__(self, d_model: int, n_heads: int, d_ff: int): super().__init__() self.attn = MultiHeadAttention(d_model, n_heads) self.ff = FeedForward(d_model, d_ff) self.norm1 = RMSNorm(d_model) self.norm2 = RMSNorm(d_model) def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: # Pre-LN: normalize before each sub-layer, add residual after x = x + self.attn(self.norm1(x), mask) x = x + self.ff(self.norm2(x)) return x ``` --- ## Part 3: Memory and compute optimizations ### KV caching During autoregressive generation, each new token recomputes attention over the full sequence history. KV caching stores past key and value tensors and appends to them, reducing per-step compute from $O(S^2)$ to $O(S)$: ```python class KVCache: def __init__(self): self.k_cache: list[torch.Tensor] = [] self.v_cache: list[torch.Tensor] = [] def update(self, k: torch.Tensor, v: torch.Tensor): self.k_cache.append(k) self.v_cache.append(v) def get(self) -> tuple[torch.Tensor, torch.Tensor]: return torch.cat(self.k_cache, dim=2), torch.cat(self.v_cache, dim=2) def clear(self): self.k_cache.clear() self.v_cache.clear() ``` ### Grouped Query Attention (GQA) Multi-head attention with $h$ heads stores $h$ key and value tensors per layer. GQA reduces this by sharing one key-value head across $G$ query heads. Llama 2 70B uses 8 KV heads for 64 query heads — an 8x reduction in KV cache memory at inference: ```python class GroupedQueryAttention(nn.Module): def __init__(self, d_model: int, n_query_heads: int, n_kv_heads: int): super().__init__() assert n_query_heads % n_kv_heads == 0 self.n_query_heads = n_query_heads self.n_kv_heads = n_kv_heads self.n_rep = n_query_heads // n_kv_heads self.head_dim = d_model // n_query_heads self.q_proj = nn.Linear(d_model, d_model, bias=False) self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False) self.out_proj = nn.Linear(d_model, d_model, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: b, s, _ = x.shape Q = self.q_proj(x).view(b, s, self.n_query_heads, self.head_dim).transpose(1, 2) K = self.k_proj(x).view(b, s, self.n_kv_heads, self.head_dim).transpose(1, 2) V = self.v_proj(x).view(b, s, self.n_kv_heads, self.head_dim).transpose(1, 2) # Repeat KV heads to match query head count K = K.repeat_interleave(self.n_rep, dim=1) V = V.repeat_interleave(self.n_rep, dim=1) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) weights = torch.softmax(scores, dim=-1) out = torch.matmul(weights, V) out = out.transpose(1, 2).contiguous().view(b, s, -1) return self.out_proj(out) ``` ### FlashAttention: fused tiled SRAM kernels Standard attention materializes the full $S \times S$ attention matrix in HBM (high-bandwidth memory). For a 4,096-token sequence with 32 heads, that's ~2GB of intermediate tensors per layer, with multiple round trips between HBM and SRAM. FlashAttention fuses the three attention operations (QKT, softmax, AV) into a single kernel that tiles Q, K, V into SRAM blocks, computes attention locally, and accumulates the running softmax normalization without writing the full attention matrix back to HBM. The kernel implements a numerically stable online softmax: $$ m_i^{(j)} = \max(m_i^{(j-1)}, \max_l S_{il}^{(j)}) $$ $$ \ell_i^{(j)} = e^{m_i^{(j-1)} - m_i^{(j)}} \ell_i^{(j-1)} + \sum_l e^{S_{il}^{(j)} - m_i^{(j)}} $$ In practice: `torch.nn.functional.scaled_dot_product_attention` uses FlashAttention kernels automatically when inputs are in the right format. No manual kernel writing needed for most applications. Peak memory drops from $O(S^2)$ to $O(S)$. On an A100 with a 2,048-token sequence, FlashAttention 2 achieves approximately 2.2x speedup and 5-10x memory reduction compared to standard attention. ### Low-Rank Adaptation (LoRA) LoRA freezes the pre-trained weight matrix $W_0 \in \mathbb{R}^{d \times k}$ and injects a low-rank decomposition $\Delta W = BA$ where $B \in \mathbb{R}^{d \times r}$, $A \in \mathbb{R}^{r \times k}$, with $r \ll \min(d, k)$: $$ h = W_0 x + \Delta W x = W_0 x + B A x $$ Only $A$ and $B$ are trained. A rank-16 adapter on a 4,096-dimensional projection reduces trainable parameters from 16.7M to 131K — a 128x reduction. ```python class LoRALinear(nn.Module): def __init__(self, in_features: int, out_features: int, rank: int = 16, alpha: float = 16.0): super().__init__() self.weight = nn.Parameter(torch.empty(out_features, in_features), requires_grad=False) nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) self.lora_A = nn.Parameter(torch.randn(rank, in_features) * 0.01) self.lora_B = nn.Parameter(torch.zeros(out_features, rank)) self.scale = alpha / rank def forward(self, x: torch.Tensor) -> torch.Tensor: base_out = nn.functional.linear(x, self.weight) lora_out = nn.functional.linear(nn.functional.linear(x, self.lora_A), self.lora_B) return base_out + self.scale * lora_out ``` Scale with `alpha / rank` instead of a raw learning rate because it decouples adapter sensitivity from rank choice — you can change rank without re-tuning the learning rate. --- ### Memory and compute summary | Optimization | Memory impact | Compute impact | |---|---|---| | KV caching | KV grows linearly with sequence length | Per-step compute drops from $O(S^2)$ to $O(S)$ | | GQA | KV cache reduced by $n\_query\_heads / n\_kv\_heads$ | Minimal compute reduction | | FlashAttention | $O(S^2) \to O(S)$ peak SRAM | 2–3x throughput on A100 for long sequences | | LoRA | ~1% of full fine-tune parameters | Same forward pass cost; backward over $A, B$ only | | INT8 quantization | ~50% model size reduction | Minor throughput gain; small accuracy drop on some tasks | Part 2 of this series covers encoder vs. decoder architecture differences — BERT's bidirectional masking vs. GPT's causal masking, their pre-training objectives, and the structural consequences for downstream tasks. --- Source: https://himanshuat.com/blogs/understanding-transformers-implementation