Transformers from scratch

LayerNorm with gamma and beta and RMSNorm with a weight, both over the last axis, with eps inside the square root.

From The Full Transformer — Encoder + Decoder

Build the sin and cos position table with interleaved columns and base 10000, and add it to a batch of token embeddings.

From Positional Encoding — Sinusoidal, RoPE, ALiBi

Append keys and values to a preallocated cache, attend one query over the valid prefix, and reject an append past the capacity.

From KV Cache, Flash Attention & Inference Optimization

Rotate adjacent coordinate pairs of queries and keys by position times base^(-2i/d), so attention scores depend only on the position gap.

From Positional Encoding — Sinusoidal, RoPE, ALiBi

Batched softmax(QK^T / sqrt(d_k)) V with an optional boolean mask, plus the causal mask that a decoder uses.

From Self-Attention from Scratch

GQA forward with H query heads and G shared key and value groups, plus mean pooling of a multi-head checkpoint into groups.

From Multi-Head Attention

Route tokens to the top k experts with renormalized softmax gates, and compute the Switch Transformer auxiliary loss E times sum f_i P_i.

From Mixture of Experts (MoE)

Reshape (B, T, D) into H heads and back, then run batched multi-head self-attention with given projection matrices.

From Multi-Head Attention

FlashAttention in numpy: a running max and a running sum over key blocks give exact attention for any block size.

From KV Cache, Flash Attention & Inference Optimization