Layer 3 · Code

Attention

Implement scaled dot-product and multi-head attention from scratch in PyTorch, and see how FlashAttention changes only the implementation, not the math.

15 min read110 XP

raw torch ops
scores, mask, softmax, matmul — visible
wrap in a module
learnable Q/K/V/O projections
split into heads
reshape, don't duplicate compute
swap the kernel
F.scaled_dot_product_attention (FlashAttention)

Same math, four levels of implementation, decreasing memory footprint.

First, the formula with nothing hidden — raw tensor ops, so every step from the math block in Layer 2 has a visible line of code.

Go deeper: L4 Foundations