Replaces the standard attention $O = \text{softmax}(QK^T)V$ (which materializes a full $N \times N$ attention matrix in HBM) with an IO-aware tiled algorithm that streams Q, K, V through on-chip SRAM and never writes the $N \times N$ matrix out.

Motivation: standard attention is bottlenecked by memory, not FLOPs. For a sequence of length $N$ with head dim $d$, the $N \times N$ softmax matrix costs $O(N^2)$ memory and forces $O(N^2)$ HBM reads/writes on GPU. Flash Attention 2 cuts this to $O(N)$ HBM accesses (linear in sequence length) by recomputing attention in tiles and using the online softmax trick to merge partial results.

Standard Attention

Components

• $Q, K, V \in R^{N \times d}$: query, key, value matrices for a single attention head, where $N$ is sequence length and $d$ is head dimension.
• $S \in R^{N \times N}$: attention scores, one scalar per (query, key) pair.
• $P \in R^{N \times N}$: attention weights after row-wise softmax. Each row sums to 1.
• $O \in R^{N \times d}$: output, one $d$-dim vector per query.

Why this is inefficient

• Quadratic Memory Scaling $O(N^2)$: $S$ and $P$ are $N \times N$. For $N = 8192$ and fp16 that’s $8192^2 \cdot 2\text{B} \approx 128\text{ MB}$ per head, per layer. With 32 layers and 32 heads, this is already enormous amounts of memory. Since N is our sequence length, our memory scales quadratically in longer conversations, which is the main bottleneck for long-context LLMs.

• Bandwidth bound $O(N^2)$: the GPU makes 3 trips to HBM.

  1. Read $Q, K$ to compute $S$. Write $S$ back to HBM.
  2. Read $S$ from HBM to compute the row-max for softmax, then write $P$ back.
  3. Read $P$ again to multiply with $V$. Three full passes over an $N \times N$ matrix that mostly lives in slow HBM, not fast SRAM.

  HBM might higher bandwidth than VRAM, but it’s still nothing compared to on-chip SRAM. We want to minimize HBM trips.

Flash Attention 2:

Break down the Q, K, V matrices into tiles that fit in SRAM, and compute attention block by block.

for now, we will assume a tile has size $m_r \times m_c$, where $m_r$ is the number of rows of Q, V and $m_c$ is the number of rows of $K^T$. This size is dynamically computed based on SRAM in the hardware, during compilation.

The full forward pass for a single Q block $Q_i$ loops over all K, V blocks $(K_j, V_j)$ where $K_j$ and $V_j$ are matrices of size $B_c \times d$, $B_r \times d$ respectively and accumulates an output $O_i$ using an online softmax. Concretely:

Inside the inner loop (for a fixed Query block $i$ and streaming Key/Value block $j$):

1. Compute raw scores for the tile: \(S_{ij} = Q_i K_j^T \in \mathbb{R}^{B_r \times B_c}\)

2. Find the local row maximums and compute stable exponentials: \(m_{ij} = \text{rowmax}(S_{ij}) \in \mathbb{R}^{B_r}\) \(\~{P_{ij}} = e^{S_{ij} - m_{ij}} \in \mathbb{R}^{B_r \times B_c}\)

3. Compute the local row sums of the exponentials: \(\ell_{ij} = \text{rowsum}(\~{P_{ij}}) \in \mathbb{R}^{B_r}\)

4. Update the global running statistics (Max and Denominator) using the correction factor: \(m_i^{\text{new}} = \max(m_i^{\text{old}}, m_{ij})\) \(\ell_i^{\text{new}} = \left( e^{m_i^{\text{old}} - m_i^{\text{new}}} \odot \ell_i^{\text{old}} \right) + \left( e^{m_{ij} - m_i^{\text{new}}} \odot \ell_{ij} \right)\)

5. Update the running, unnormalized Output accumulator ($\tilde{O}_i$): \(\tilde{O}_i^{\text{new}} = \left( e^{m_i^{\text{old}} - m_i^{\text{new}}} \odot \tilde{O}_i^{\text{old}} \right) + \left( e^{m_{ij} - m_i^{\text{new}}} \odot (\~{P_{ij}} V_j) \right) \in \mathbb{R}^{B_r \times d}\)

After the inner loop finishes (all Key/Value blocks $j$ have been processed):

1. Perform the final Safe Normalization to get the true Attention Output ($O_i$): \(O_i = \tilde{O}_i^{\text{final}} \oslash \ell_i^{\text{final}}\)

here, $O_i$ is rescaled to compensate for the running max update. This is an element-wise division.

Components

• $~{P_{ij}}$: unnormalized attention weights, not yet divided by the row sum for softmax. We normalize after the multiplication with $V_j$ to save memory and compute.
• $Q_i \in R^{B_r \times d}$: tile $i$ of the query matrix, of block size $B_r$ rows. $i = 1 \dots T_r$ where $T_r = \lceil N / B_r \rceil$.
• $K_j, V_j \in R^{B_c \times d}$: tile $j$ of the key and value matrices, of block size $B_c$ rows. $j = 1 \dots T_c$ where $T_c = \lceil N / B_c \rceil$.
• $S_{ij} \in R^{B_r \times B_c}$: block of attention scores between $Q_i$ and $K_j$. Lives only in SRAM, never written to HBM.
• $m_i \in R^{B_r}$: running row-max for $Q_i$ across all $K$ blocks seen so far. Initialized to $-\infty$ (or just the first block’s rowmax, same thing).
• $\ell_i \in R^{B_r}$: running row-sum of unnormalized exp scores for $Q_i$. Initialized to $0$.
• $O_i \in R^{B_r \times d}$: running output for $Q_i$. Initialized to $0$. After all $T_c$ inner steps it is divided by $\ell_i$ to give the true softmax-weighted output.
• $\text{rowmax}(\cdot), \text{rowsum}(\cdot)$: per-row max and sum. Cheap, on chip.
• $B_r, B_c$: block sizes. Typical: $B_r = 64$ or $128$, $B_c = 64$ or $128$. Picked so that $Q_i, K_j, V_j, S_{ij}, P_{ij}$ all fit in SRAM together.

Why

• Tiling $B_r, B_c$: instead of computing the full $N \times N$ softmax matrix, we only ever materialize a $B_r \times B_c$ block in fast on-chip SRAM. Once the inner step is done, $S_{ij}$ and $P_{ij}$ are discarded. HBM only ever sees the $N \times d$ inputs and outputs.
• Online softmax (Milakov & Gimelshein 2018): the standard softmax needs the full row to know the max. The online trick fuses a new $K$ block into the running $(m_i, \ell_i, O_i)$ by rescaling the previous accumulator with $e^{m_i - m_i^{new}}$. This is the same identity you’d get if you had computed the full row at once, just stitched together block by block. Read more here: online softmax.
• No HBM round trip: standard attention reads/writes $S$ and $P$ three times. Flash Attention 2 only loads $Q_i, K_j, V_j$ once each, does all the math in SRAM, and writes the final $O_i$ back. Total HBM traffic is $O(N d)$ instead of $O(N^2 d)$.
• Recomputation in backward pass: normally the backward pass would need the saved $S$ and $P$ to compute gradients w.r.t. $Q, K, V$. Flash Attention 2 throws them away and re-tiles the forward on the fly during backward. Recomputing is cheap (it’s the same flops we were already doing) and saves the $O(N^2)$ memory that would have been needed to store $S, P$.
• FLOPs are unchanged: Flash Attention 2 is exactly the same math as standard attention, just rearranged. The speedup is purely from memory access patterns, not from fewer operations.

Worked Example

since this is a hard concept, I will walk through a tiny example with $N = 6$, $d = 2$, $B_r = 2$, $B_c = 2$. That gives $T_r = T_c = 3$ tiles of $2$ rows each. I will fix a single query tile $Q_1$ (rows $1$–$2$ of $Q$) and walk all three inner steps $j = 1, 2, 3$, then do the final step-$6$ normalization. The arithmetic below is also reproduced and verified by verify_fa_example.py in this directory.

We start with a row-$Q$, a transposed $K^T$ (so the $\text{row} \times \text{col}$ shapes line up with the algorithm on line $31$), and a row-$V$: \(Q = \begin{bmatrix} 2 & -1 \\ -2 & 5 \\ 4 & 1 \\ 0 & 2 \\ 3 & 0 \\ 1 & -3 \end{bmatrix} \qquad K^T = \begin{bmatrix} 1 & 0 & 4 & -1 & 5 & 2 \\ -2 & 3 & 1 & 2 & 0 & -4 \end{bmatrix} \qquad V = \begin{bmatrix} 3 & 0 \\ 1 & 2 \\ -2 & 1 \\ 5 & -3 \\ 4 & 0 \\ -1 & 2 \end{bmatrix}\) (so $K = (K^T)^\top$ is $6 \times 2$, e.g. the first two rows of $K$ are $[1,-2]$ and $[0, 3]$, matching $K^T$’s first two columns).

Initialization

The running statistics are vectors of length $B_r = 2$ (one entry per row of $Q_1$), and the running output $O_1$ is a $B_r \times d$ matrix. All start at the “empty product” / “empty sum” identities for softmax: \(m_0 = \begin{bmatrix} -\infty \\ -\infty \end{bmatrix}, \qquad \ell_0 = \begin{bmatrix} 0 \\ 0 \end{bmatrix}, \qquad O_0 = \begin{bmatrix} 0 & 0 \\ 0 & 0 \end{bmatrix}\)

Tiles for $Q_1$

Row-tile of $Q$ and the three column-tiles of $K^T$ (i.e. the three row-tiles of $K$): \(Q_1 = \begin{bmatrix} 2 & -1 \\ -2 & 5 \end{bmatrix}, \qquad K_1^T = \begin{bmatrix} 1 & 0 \\ -2 & 3 \end{bmatrix}, \quad K_2^T = \begin{bmatrix} 4 & -1 \\ 1 & 2 \end{bmatrix}, \quad K_3^T = \begin{bmatrix} 5 & 2 \\ 0 & -4 \end{bmatrix}\) \(V_1 = \begin{bmatrix} 3 & 0 \\ 1 & 2 \end{bmatrix}, \quad V_2 = \begin{bmatrix} -2 & 1 \\ 5 & -3 \end{bmatrix}, \quad V_3 = \begin{bmatrix} 4 & 0 \\ -1 & 2 \end{bmatrix}\)

Inner step $(i = 1,\ j = 1)$: tile $(Q_1, K_1, V_1)$

Step 1. $S_{11} = Q_1 K_1^T \in \mathbb{R}^{B_r \times B_c}$: \(S_{11} = \begin{bmatrix} 2 & -1 \\ -2 & 5 \end{bmatrix} \begin{bmatrix} 1 & 0 \\ -2 & 3 \end{bmatrix} = \begin{bmatrix} 2\cdot 1 + (-1)\cdot(-2) & 2\cdot 0 + (-1)\cdot 3 \\ -2\cdot 1 + 5\cdot(-2) & -2\cdot 0 + 5\cdot 3 \end{bmatrix} = \begin{bmatrix} 4 & -3 \\ -12 & 15 \end{bmatrix}\)

Step 2. $m_{11} = \text{rowmax}(S_{11})$, $P_{11} = e^{S_{11} - m_{11}}$ (subtracting the row max for numerical stability, the standard softmax trick): \(m_{11} = \begin{bmatrix} 4 \\ 15 \end{bmatrix}, \qquad P_{11} = \begin{bmatrix} e^{4-4} & e^{-3-4} \\ e^{-12-15} & e^{15-15} \end{bmatrix} = \begin{bmatrix} 1 & e^{-7} \\ e^{-27} & 1 \end{bmatrix}\)

Step 3. $\ell_{11} = \text{rowsum}(P_{11})$: \(\ell_{11} = \begin{bmatrix} 1 + e^{-7} \\ e^{-27} + 1 \end{bmatrix}\)

Step 4. Update running $m_1, \ell_1$. Since $m_0 = -\infty$ the old terms vanish, and $m_1^{\text{new}}$ just inherits the new tile’s max: \(m_1^{\text{new}} = \max(m_0, m_{11}) = \begin{bmatrix} 4 \\ 15 \end{bmatrix}, \qquad e^{m_0 - m_1^{\text{new}}} = \begin{bmatrix} 0 \\ 0 \end{bmatrix}, \qquad e^{m_{11} - m_1^{\text{new}}} = \begin{bmatrix} 1 \\ 1 \end{bmatrix}\) \(\ell_1^{\text{new}} = \begin{bmatrix} 0 \\ 0 \end{bmatrix} \odot \ell_0 + \begin{bmatrix} 1 \\ 1 \end{bmatrix} \odot \ell_{11} = \ell_{11} = \begin{bmatrix} 1 + e^{-7} \\ 1 + e^{-27} \end{bmatrix}\)

Step 5. Update running $O_1$. Same story: $O_0 = 0$, so only the $P_{11} V_1$ term survives: \(P_{11} V_1 = \begin{bmatrix} 1 & e^{-7} \\ e^{-27} & 1 \end{bmatrix} \begin{bmatrix} 3 & 0 \\ 1 & 2 \end{bmatrix} = \begin{bmatrix} 3 + e^{-7} & 2e^{-7} \\ 3e^{-27} + 1 & 2 \end{bmatrix} = O_1^{\text{new}}\)

Inner step $(i = 1,\ j = 2)$: tile $(Q_1, K_2, V_2)$

Step 1. \(S_{12} = Q_1 K_2^T = \begin{bmatrix} 2 & -1 \\ -2 & 5 \end{bmatrix} \begin{bmatrix} 4 & -1 \\ 1 & 2 \end{bmatrix} = \begin{bmatrix} 7 & -4 \\ -3 & 12 \end{bmatrix}\)

Step 2. \(m_{12} = \begin{bmatrix} 7 \\ 12 \end{bmatrix}, \qquad P_{12} = \begin{bmatrix} 1 & e^{-11} \\ e^{-15} & 1 \end{bmatrix}\)

Step 3. \(\ell_{12} = \begin{bmatrix} 1 + e^{-11} \\ 1 + e^{-15} \end{bmatrix}\)

Step 4. This is the first time the online-softmax correction actually fires — row $0$’s running max moves from $4$ to $7$, so the old $O_1$ has to be rescaled by $e^{4 - 7} = e^{-3}$ before the new $P_{12} V_2$ term is added: \(m_1^{\text{new}} = \max\!\left( \begin{bmatrix} 4 \\ 15 \end{bmatrix},\ \begin{bmatrix} 7 \\ 12 \end{bmatrix} \right) = \begin{bmatrix} 7 \\ 15 \end{bmatrix}\) \(e^{m_1^{\text{old}} - m_1^{\text{new}}} = \begin{bmatrix} e^{4-7} \\ e^{15-15} \end{bmatrix} = \begin{bmatrix} e^{-3} \\ 1 \end{bmatrix}, \qquad e^{m_{12} - m_1^{\text{new}}} = \begin{bmatrix} e^{7-7} \\ e^{12-15} \end{bmatrix} = \begin{bmatrix} 1 \\ e^{-3} \end{bmatrix}\) \(\ell_1^{\text{new}} = \begin{bmatrix} e^{-3} \\ 1 \end{bmatrix} \odot \ell_1^{\text{old}} + \begin{bmatrix} 1 \\ e^{-3} \end{bmatrix} \odot \ell_{12} = \begin{bmatrix} e^{-3}(1 + e^{-7}) + (1 + e^{-11}) \\ 1\cdot(1 + e^{-27}) + e^{-3}(1 + e^{-15}) \end{bmatrix} = \begin{bmatrix} 1 + e^{-3} + e^{-10} + e^{-11} \\ 1 + e^{-3} + e^{-18} + e^{-27} \end{bmatrix}\)

Step 5. \(P_{12} V_2 = \begin{bmatrix} 1 & e^{-11} \\ e^{-15} & 1 \end{bmatrix} \begin{bmatrix} -2 & 1 \\ 5 & -3 \end{bmatrix} = \begin{bmatrix} -2 + 5e^{-11} & 1 - 3e^{-11} \\ 5 - 2e^{-15} & -3 + e^{-15} \end{bmatrix}\) \(O_1^{\text{new}} = \begin{bmatrix} e^{-3} \\ 1 \end{bmatrix} \odot O_1^{\text{old}} + \begin{bmatrix} 1 \\ e^{-3} \end{bmatrix} \odot (P_{12} V_2) = \begin{bmatrix} 3e^{-3} + e^{-10} - 2 + 5e^{-11} & 2e^{-10} + 1 - 3e^{-11} \\ 1 + 3e^{-27} + 5e^{-3} - 2e^{-18} & 2 - 3e^{-3} + e^{-18} \end{bmatrix}\)

Inner step $(i = 1,\ j = 3)$: tile $(Q_1, K_3, V_3)$

Step 1. \(S_{13} = Q_1 K_3^T = \begin{bmatrix} 2 & -1 \\ -2 & 5 \end{bmatrix} \begin{bmatrix} 5 & 2 \\ 0 & -4 \end{bmatrix} = \begin{bmatrix} 10 & 8 \\ -10 & -24 \end{bmatrix}\)

Step 2. \(m_{13} = \begin{bmatrix} 10 \\ 8 \end{bmatrix}, \qquad P_{13} = \begin{bmatrix} 1 & e^{-2} \\ e^{-18} & e^{-32} \end{bmatrix}\)

Step 3. \(\ell_{13} = \begin{bmatrix} 1 + e^{-2} \\ e^{-18} + e^{-32} \end{bmatrix}\)

Step 4. Row $0$’s running max moves again, from $7$ to $10$: \(m_1^{\text{new}} = \max\!\left( \begin{bmatrix} 7 \\ 15 \end{bmatrix},\ \begin{bmatrix} 10 \\ 8 \end{bmatrix} \right) = \begin{bmatrix} 10 \\ 15 \end{bmatrix}\) \(e^{m_1^{\text{old}} - m_1^{\text{new}}} = \begin{bmatrix} e^{7-10} \\ e^{15-15} \end{bmatrix} = \begin{bmatrix} e^{-3} \\ 1 \end{bmatrix}, \qquad e^{m_{13} - m_1^{\text{new}}} = \begin{bmatrix} e^{10-10} \\ e^{8-15} \end{bmatrix} = \begin{bmatrix} 1 \\ e^{-7} \end{bmatrix}\) \(\ell_1^{\text{new}} = \begin{bmatrix} e^{-3} \\ 1 \end{bmatrix} \odot \ell_1^{\text{old}} + \begin{bmatrix} 1 \\ e^{-7} \end{bmatrix} \odot \ell_{13} = \begin{bmatrix} 1 + e^{-2} + e^{-3} + e^{-6} + e^{-13} + e^{-14} \\ 1 + e^{-3} + e^{-18} + e^{-25} + e^{-27} + e^{-39} \end{bmatrix} \approx \begin{bmatrix} 1.1876 \\ 1.0498 \end{bmatrix}\)

Step 5. \(P_{13} V_3 = \begin{bmatrix} 1 & e^{-2} \\ e^{-18} & e^{-32} \end{bmatrix} \begin{bmatrix} 4 & 0 \\ -1 & 2 \end{bmatrix} = \begin{bmatrix} 4 - e^{-2} & 2e^{-2} \\ 4e^{-18} - e^{-32} & 2e^{-32} \end{bmatrix}\) \(O_1^{\text{new}} = \begin{bmatrix} e^{-3} \\ 1 \end{bmatrix} \odot O_1^{\text{old}} + \begin{bmatrix} 1 \\ e^{-7} \end{bmatrix} \odot (P_{13} V_3) = \begin{bmatrix} 4 - e^{-2} - 2e^{-3} + 3e^{-6} + e^{-13} + 5e^{-14} & 2e^{-2} + e^{-3} + 2e^{-13} - 3e^{-14} \\ 1 + 5e^{-3} - 2e^{-18} + 3e^{-27} + 4e^{-25} - e^{-39} & 2 - 3e^{-3} + e^{-18} + 2e^{-39} \end{bmatrix} \approx \begin{bmatrix} 3.768 & 0.320 \\ 1.249 & 1.851 \end{bmatrix}\)

After all $T_c = 3$ inner steps we have \(m_1^{\text{final}} = \begin{bmatrix} 10 \\ 15 \end{bmatrix}, \qquad \ell_1^{\text{final}} = \begin{bmatrix} 1.1876 \\ 1.0498 \end{bmatrix}, \qquad \tilde{O}_1^{\text{final}} \approx \begin{bmatrix} 3.768 & 0.320 \\ 1.249 & 1.851 \end{bmatrix}\)

Step 6: final safe normalization

$\tilde{O}_1^{\text{final}}$ is still the unnormalized weighted sum of $V$ rows — each row was scaled by $e^{S - m_1^{\text{final}}}$, with no denominator applied yet. To recover the true softmax-weighted average we divide each row of $\tilde{O}_1$ by the corresponding entry of $\ell_1$. The $\oslash$ on line $62$ is the element-wise version of this: $\ell_1$ is broadcast across the $d$ columns of $\tilde{O}_1$ and each entry in row $k$ gets divided by $\ell_1[k]$.

Plugging in $\tilde{O}_1^{\text{final}} \approx \begin{bmatrix} 3.768 & 0.320 \ 1.249 & 1.851 \end{bmatrix}$ and $\ell_1^{\text{final}} \approx \begin{bmatrix} 1.1876 \ 1.0498 \end{bmatrix}$:

Sanity check. Computing $\text{softmax}(Q_1 K^T) V$ the standard (non-tiled) way on the same matrices gives the same $O_1$ to machine precision (max abs diff $\approx 2.2 \times 10^{-16}$ in the script) — the algorithm is exact, only the memory access pattern changed.

Work partitioning (FA2 vs FA1)

FA1 parallelized by also splitting the inner loop over $K$ blocks across thread blocks, then reducing. This meant multiple thread blocks were redundantly touching the same $Q_i$ tile and writing partial outputs to HBM, which the reduction then had to merge.

FA2 fixes this: each thread block is assigned a slice of $Q$ (a few $Q_i$ tiles) and is responsible for the full inner loop over $K, V$. The result is written directly to HBM with no inter-block reduction.

Why

• No redundant HBM writes: in FA1, two thread blocks working on overlapping $Q_i$ both wrote to HBM. FA2 assigns each query row to exactly one thread block, so each output is written once.
• Better arithmetic intensity per thread block: FA2 thread blocks do more compute per byte loaded from HBM (the same $K_j, V_j$ is reused across many $Q_i$ inside the block). The GPU is finally compute-bound instead of memory-bound.
• Cross-block communication gone: no more atomic adds, no more partial buffers, no more reduction kernel. The output $O$ comes out ready in HBM.
• Same SRAM budget, more useful work: with the saved HBM bandwidth, FA2 can afford to keep $Q_i, K_j, V_j$ resident longer and overlap softmax epilogue with the matmul. This is the second key win.

Causal masking

For decoder attention, we generally mask the $P$ matrix. Fundamentally, Attention works, by comparing each token to the next token, and forming an $\R^{N\times N}$ matrix. But in generative transformers, we don’t want the query to cheat, and look ahead at the answer. We want it to predict based only on past tokens. Thus we mask the $P$ matrix.

where

for example, if $N = 4$, the mask would look like \(M = \begin{bmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 \end{bmatrix}\)

this is because, $P_{ij}-\infty = -\infty$ and $e^{-\infty} = 0$, so the masked-out entries contribute nothing to the softmax output.

There two main options for how to implement this in Flash Attention 2:

Option A: skip masked blocks. in the inner loop, if the entire $K_j$ block lies in the future, skip it. If it straddles the diagonal, compute it but mask out the upper triangle:

then proceed with the same online softmax. The $-\infty$ entries become $0$ after exp, so they contribute nothing to $\ell_i$ or $O_i$.

Option B: no mask needed. in the strict causal case, you can schedule $Q_i$ tiles so that the inner loop only iterates over $K_j$ with $j \le i$ (i.e. the $K, V$ blocks that are already “in the past”). No mask, no wasted compute, just a careful loop bound.

Why

• Option A is general: handles arbitrary masks (sliding window, document boundaries, prefix LM, etc.). The mask is just a $0 / -\infty$ add before softmax, doesn’t change the algorithm.
• Option B is fast for pure causal: the FA2 paper uses this. You can split the sequence into chunks and only feed the relevant $K, V$ blocks. No masked-out FLOPs.
• Both are exact: masking is applied to $S_{ij}$ before the online softmax merge, so the result is bit-identical (modulo FP order) to standard causal attention.

Quick reference

• Standard attention: $O = \text{softmax}(QK^T)V$, costs $O(N^2)$ HBM and $O(N^2 d)$ FLOPs.
• Flash Attention 2: same FLOPs, $O(N d)$ HBM, $O(N^2)$ compute stays the same. The optimization is memory bound.
• Inner loop: for each $Q_i$ tile, iterate over all $K_j, V_j$ tiles, maintain $(m_i, \ell_i, O_i)$, apply online softmax merge.
• Typical shapes: $Q, K, V:\ (N, d)$, $Q_i:\ (B_r, d)$, $K_j, V_j:\ (B_c, d)$, $S_{ij}, P_{ij}:\ (B_r, B_c)$, $O_i:\ (B_r, d)$, $m_i, \ell_i:\ (B_r,)$.
• Backward pass: re-tile and recompute $S, P$ on the fly. No $N \times N$ saved.

Notes

• “FA2” sometimes refers to the original 2022 Flash Attention paper and sometimes to the 2023 revision. The 2023 revision is the one with the work-partitioning fix and is what people mean in practice. Both are exact, both use online softmax, both are IO-aware.
• The output $O$ and the running stats $(m_i, \ell_i)$ must all live in SRAM simultaneously for the merge to work without an HBM round trip. This is what sets the upper bound on $B_r$ and $B_c$ (along with $d$). The Flash Attention 2 paper gives a SRAM budgeting rule of thumb: $B_r \cdot d + B_c \cdot d + B_r \cdot B_c \lesssim \text{SRAM budget}$.
• For very small $N$ (e.g. $N \le 512$), the $N \times N$ attention matrix fits in SRAM anyway and standard attention is competitive. FA2 wins for long sequences where $N^2$ blows past SRAM.
• $d$ here is per-head dim, not the full embedding. Different heads are computed in parallel, so the SRAM budget is per-head. With $d = 128$ and 64 KB of SRAM per CTA, $B_r = B_c = 128$ is comfortable.
• Drop-in replacement for standard attention: same math, same backward, same hyper-parameters. Just call the kernel.

Why this is the modern default

• Memory: attention’s $O(N^2)$ cost is the single biggest bottleneck for long-context LLMs. FA2 turns it into $O(N)$ in HBM, which is what enables $32k$, $128k$, $1M$ context windows.
• Speed: typically 2-4x faster than standard attention wall-clock, even at $N = 2048$ where you’d think it wouldn’t matter. The crossover is around $N = 512$.
• No quality loss: exact attention. Perplexity, downstream accuracy, attention maps, all bit-identical (modulo float reduction order). You don’t pay a quality tax for the speedup.
• Hardware friendly: the algorithm is just nested matmuls with an in-register softmax epilogue. Maps cleanly onto tensor cores. FA3 (2024) takes this further with warp-specialization and async copies, but FA2 is still the baseline everyone implements.

Read the original paper here: FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.

this document is my post-mortem on the atlas module refactor. I moved from a single MaiaDAO class that owned every database operation, to a set of small Repository classes that each own one entity. I’ll explain why I did it, what the patterns mean, and what the end state looks like.

disclaimer: Deepseek v4 Flash (provided by opencode zen/free) was used to assist with the mechanical parts of this refactor, the actual editing of 40+ files.

What is Pleiades?

Initially named: Youtube-ML-Pipeline, it was a semester project with some ridiculous requirements. I was a student, and the project was a solo effort. I had to learn async Python, Prefect, Postgres, SQLAlchemy, and HuggingFace in 3 weeks. A live data ingestion -> Training -> Inference -> Visualization pipeline was the goal. Seeing this, I decided “hey youtube has multi-modal data and LLMs research has been exploring vision language models!”. It seems that a google engineer also had the same idea- because the “ask” feature on youtube is a polished, and genuinely amazing implementation of parsing youtube using an LLM. This is actually why I had abandoned this repository, as it felt redundant, but now it acts as my playground for “what doesn’t work and how do i fix it?”

this repository is quite old now. As time was short at the time, and I was unfamiliar with the technologies listed, not to mention it was a solo project; I “vibe-coded” the project. Then I kept adding to the technical debt, when a new model came out with a free tier, I’d throw it at “Pleades”, hoping it was good enough. And finally, I have now grown enough as a developer, to sit down, and design a maintainable refactor.

this isn’t to say no thought went into the architecture. But, at the time, I wasn’t a good software designer. I knew basic patterns, singleton, dao, and made a model blindly implement what looked good from a deep learning perspective. All I cared about was that the right kind of data, was being stored in a way somewhat familiar to me.

Why Pleiades?

The naming scheme was so that the project would be cool and memorable. I wanted to make 7 modules. One named after each of the 7 stars in the Pleiades star cluster. This would make the codebase easeier to navigate, since I could asociate each module with a concept, a greek goddess and her story.

What is Atlas?

Atlas isn’t even a pleiades. But he’s their dad! And so, in the repository, he is the father of the modules, holding up the database access, just like he holds up the sky in greek mythology. The atlas module is the database access layer for the entire pipeline. It contains adapters, DAOs, and now repositories.

Before: What Atlas Used To Be

the old atlas module centered around MaiaDAO. this was a single class, about 400 lines, that had methods for every database operation in the entire pipeline:

class MaiaDAO:
    async def ingest_video_metadata(self, data): ...
    async def fetch_tracker_targets(self, batch_size): ...
    async def log_video_stats_batch(self, stats): ...
    async def get_channel_by_id(self, channel_id): ...
    async def add_to_search_queue(self, term): ...
    async def archive_cold_stats(self, days): ...
    # ... 30+ more methods

every agent: hunter, tracker, scribe, painter, janitor, archeologist, imported MaiaDAO and used it for everything. Here is what the tracker agent used to look like at the call site:

from atlas.adapters.maia import MaiaDAO

dao = MaiaDAO()
targets = await dao.fetch_tracker_targets(batch_size)

this looks innocent enough. the problems were elsewhere.

problem 1: vault switching was fragile

the DAO had methods like archive_cold_stats that touched both Postgres and the vault (HF or GCP). the switching logic was:

def archive_cold_stats(self, ...):
    stats = self._fetch_all(...)
    v = get_vault()  # global singleton, provider chosen at import time
    v.append_metrics(stats, date=...)

get_vault() returned a HuggingFace vault or a Google Cloud vault based on an environment variable at startup. swapping meant restarting the process. you couldn’t have both. you couldn’t test one path without patching the global.

this was not a “database abstraction.” it was a monolith that happened to talk to two databases.

problem 2: tests were a patch-matching minefield

because MaiaDAO was a single class used everywhere, unit tests had to patch it at the module level:

@patch("maia.hunter.flow.MaiaDAO")
@patch("maia.tracker.flow.MaiaDAO")

if a consumer file renamed its import or changed which method it called, the test broke silently. integration tests were worse — they imported MaiaDAO directly and every fixture created a “real” instance, which meant spinning up Postgres even when testing a single method.

Table 1 — How tests touched the DAO before the refactor

problem 3: no domain boundaries

MaiaDAO.ingest_video_metadata also upserted channels. MaiaDAO.archive_cold_stats also wrote to vault. MaiaDAO.fetch_tracker_targets returned raw tuples that the caller had to know the column order of.

this made it impossible to reason about where data lived. the DAO was a god object — it knew everything about the database schema, and every consumer knew everything about the DAO.

DAO vs Repository: What These Patterns Mean

I want to be precise about these two patterns because they were quite confusing to me at first. A good resource is cosmic python’s chapter

DAO (Data Access Object)

a DAO wraps raw database access. it talks in SQL, cursors, and connection pools. its methods look like this:

class VideoDao:
    async def insert_row(self, values: tuple) -> None: ...
    async def select_by_id(self, id: str) -> Optional[Dict[str, Any]]: ...

the return type is a dict or a tuple. the caller knows it’s talking to a database. this is fine for a thin wrapper, but it leaks storage details everywhere.

Table 2 — DAO characteristics

Repository

a Repository sits on top of the DAO (or adapter) and speaks the language of the problem. it returns domain objects:

class VideoRepository:
    async def fetch_tracker_targets(self, batch_size: int) -> List[Video]: ...
    async def get_by_id(self, video_id: str) -> Optional[Video]: ...

the caller gets a Video model with .title, .status, .model_dump(). it does not know about SQL. it does not know about connection pools. it does not care.

here is the critical distinction in one sentence:

a DAO abstracts the database. a Repository abstracts the data.

Table 3 — Repository characteristics

why both matter

in this refactor, DatabaseAdapter is the DAO. it provides _execute, _fetch_one, _fetch_all. the Repositories inherit from it and add domain methods. the consumer (tracker, hunter, scribe) only sees the Repository.

the layering is:

Consumer Code (maia agents)
    ↓  calls domain methods
Repository (VideoRepository)
    ↓  uses SQL primitives
DatabaseAdapter (DAO)
    ↓  talks to connection pool
Postgres / Neon

Post-Refactor: (What the repo looks like now)

there are five repositories, each in its own file under atlas/repositories/:

Table 4 — Repository layout

each repo owns its table. no cross-table concerns unless they’re genuinely coupled (e.g. ingest_video_metadata still upserts a channel skeleton, because that happens in a single transaction).

before vs after: the tracker

this is the best concrete example. here is what the tracker’s fetch_targets_task looked like before:

# BEFORE — MaiaDAO
from atlas.adapters.maia import MaiaDAO

dao = MaiaDAO()
targets = await dao.fetch_tracker_targets(batch_size)
# targets is List[Dict] — you pray the keys are right

and now:

# AFTER — VideoRepository
from atlas.repositories import VideoRepository

video_repo = VideoRepository()
targets = await video_repo.fetch_tracker_targets(batch_size)
return [t.model_dump() for t in targets]
# targets is List[Video] — typed, validated, has .model_dump()

the difference is the return type. the old code returned raw dicts from SQL. the new code returns Video domain models — Pydantic objects with typed fields, validation, and serialization methods.

before vs after: the stats update

the old update_stats used dao.update_video_stats_batch and dao.log_video_stats_batch, which were two separate methods on MaiaDAO that each took raw dicts:

# BEFORE
dao = MaiaDAO()
await dao.update_video_stats_batch(items)

# AFTER
video_repo = VideoRepository()
stats_list = [
    VideoStats(
        video_id=item["id"],
        views=int(stats.get("viewCount", 0)),
        likes=int(stats.get("likeCount", 0)),
        timestamp=datetime.now(timezone.utc),
    )
    for item in items
]
await video_repo.log_stats_batch(stats_list)

yes, it is more lines. but every field is named and typed. the old version was dao.log_video_stats_batch(items) where items was a blob of YouTube API JSON. you had to read the DAO to know what keys it expected.

before vs after: the unit tests

old tests patched MaiaDAO at the module level. one wrong import path and the patch silently did nothing:

# BEFORE — fragile patch path
@patch("maia.tracker.flow.MaiaDAO")
def test_tracker(mock_dao):
    mock_dao.return_value.fetch_tracker_targets = AsyncMock(...)

new tests patch individual repositories. the patch target matches the actual import:

# AFTER — patches match the real import
@patch("maia.tracker.flow.VideoRepository")
def test_tracker(mock_video_repo):
    mock_video_repo.return_value.fetch_tracker_targets = AsyncMock(...)

and integration tests that used MaiaDAO() now use VideoRepository():

# BEFORE
@pytest_asyncio.fixture
async def dao(fresh_db):
    from atlas.adapters.maia import MaiaDAO
    yield MaiaDAO()  # raises ImportError after removal

# AFTER
@pytest_asyncio.fixture
async def video_repo(fresh_db):
    from atlas.repositories import VideoRepository
    yield VideoRepository()

Why These Patterns Fit This Problem

I chose Repository over DAO here for reasons specific to this codebase. not every project needs this split.

the pipeline has distinct entity lifecycles

the maia pipeline moves a video through stages: discovered → queue → ingested → transcribed → visualized → tracked → archived.

each stage touches different tables. hunter touches search_queue, videos, channels. tracker touches videos, video_stats_log. scribe touches videos, vault. janitor touches video_stats_log, vault, videos.

grouping all of these into one class meant the DAO had to know about every stage of every pipeline. that is not cohesion — it is coupling.

Table 5 — Which entities each agent touches

each agent now imports only the repositories it needs. the dependency graph is explicit.

the async + Prefect architecture rewards stateless objects

repositories are stateless. you instantiate one per task, use it, and drop it. this works naturally with Prefect’s task model:

@task
async def fetch_targets_task(batch_size: int):
    repo = VideoRepository()  # created and thrown away per task
    return await repo.fetch_tracker_targets(batch_size)

you do not need a shared DAO instance. you do not need a connection pool passed through every task. the default constructor uses the global pool, but you can inject a different one per task if you need to.

vault is not a database backend

the old code treated vault (HF/GCP) as a “backend” for the DAO, with methods like archive_cold_stats that wrote to both Postgres and vault in one method. this was the wrong abstraction.

vault is a cold-tier object store. it stores frames, transcripts, and parquet archives. Postgres stores the hot index. they serve different purposes and should be in different classes.

in the new design, VideoRepository.archive_cold_stats still touches vault (that’s a PHASE-4 TODO to extract to a service), but the boundary is clear: the repository handles SQL, and vault is accessed through atlas.vault.get_vault(). they are separate concerns that happen to be called from the same workflow method.

Features Unlocked By This Refactor

swappable backends

DatabaseAdapter currently imports the global db singleton at module level. a follow-up change (already discussed) would make the pool injectable via the constructor:

class DatabaseAdapter:
    def __init__(self, db_pool=None):
        if db_pool is None:
            from atlas.db import db
            db_pool = db
        self._db = db_pool

this would let us:

• use a mock pool in unit tests (no Postgres needed at all)
• point VideoRepository at a read replica for analytics queries
• use SQLite for local development without changing any repo code
• spin up an ephemeral test database per test function

the repos themselves do not need to change. they inherit the adapter. the swap happens in one place — the constructor call.

better MLOps extensibility

because each repository is a concrete class with a well-defined interface, you can wrap them:

class LoggedVideoRepository(VideoRepository):
    async def fetch_tracker_targets(self, batch_size):
        start = time.time()
        result = await super().fetch_tracker_targets(batch_size)
        logger.info(f"fetch_tracker_targets({batch_size}) took {time.time() - start:.2f}s")
        return result

this pattern works for:

• metrics — wrap repos with Prometheus counters per method
• caching — CachedVideoRepository checks Redis before hitting Postgres
• audit logging — AuditedEventRepository logs every event dispatch
• circuit breakers — ResilientVideoRepository fails fast when Postgres is down
• MLflow tracking — log query shapes and response sizes to MLflow for cost analysis

none of these require changing the repos themselves. they are decorators in the classic sense, enabled by the fact that each repo is a single class with a focused interface.

granular testing

before this refactor, testing the tracker’s stats update logic required mocking MaiaDAO — which had 30+ methods. you either mocked all of them (brittle) or used a real DAO (slow).

now you mock VideoRepository — which has methods only about videos. the mock surface is smaller and the intent is clearer:

@patch("maia.tracker.flow.VideoRepository")
def test_update_stats_handles_api_error(mock_repo):
    mock_repo.return_value.fetch_tracker_targets = AsyncMock(return_value=[...])
    mock_repo.return_value.update_stats_batch = AsyncMock()
    mock_repo.return_value.log_stats_batch = AsyncMock()

    result = await update_stats_task(videos, executor)
    assert result == 0  # graceful failure

you can read this test and know exactly what it’s testing. you do not need to understand the entire DAO to understand one test.

independent schema evolution

videos table schema can change without touching channels or search_queue. the VideoRepository owns its migration. the ChannelRepository owns its migration.

this sounds trivial, but with a monolithic DAO, a schema change meant editing the one class that touched every table. you were one bad merge away from breaking the entire pipeline.

now each repo is a separate file with separate tests. you can refactor VideoRepository without worrying about whether ChannelRepository still works.

What I Would Do Differently Next Time

honestly, I would not build MaiaDAO at all. I would start with per-entity repositories from day one.

the monolithic DAO felt like the right choice at the start because the project was small and “one class for all database access” is the simplest possible design. but it did not scale past about 15 methods. the inflection point was when I added vault operations inside the DAO — that is when it became clear the abstraction was wrong.

a better rule of thumb is: if a class has methods that touch two different storage systems, split it. if a class has methods that serve two different agents, split it. if a class has more than 20 methods, split it.

the repository pattern is not fancy. it is just “one class per entity.” the entities in this system are Video, Channel, SearchQueueItem, WatchlistItem, and SystemEvent. five classes. each with a single responsibility. the code is easier to read, easier to test, and easier to change.

This document summarizes my study of the GPU kernels used for the hyper-connection (mHC) preprocessing in FlashInfer’s DeepSeek V4. It covers my exploration of the reference implementation, the GPU optimization techniques involved (like split-K and tiling), and notes on how these kernels work to support the mHC design.

Figure 1 — mHC dataflow diagram

this diagram i found on flashinfer’s discussion forum is a useful visual reference for understanding the data flow and operations in the mHC preprocessing.

disclaimer: Claude Opus 4.8 was used to assist my research, as I am not an expert in GPU programming. However this document is 100% my own work and I have done my best to understand the GPU kernels and optimizations used in FlashInfer’s DeepSeek V4.

mhc_pre_big_fuse.cu

I will be going through my understanding of the mhc_pre_big_fuse kernel. this is an implementation that was merged to flashinfer very recently. I will then explain the optimizations that were made to this kernel and how they improve performance.

the math operation

please reference my math document for a better understanding on why these operations are needed, and what mHC (manifold-constrained hyper connection) is. it is available here

for now, I will be assuming you have already went through that resource. Here is how the variables in the math document map to the variables in the kernel code:

Table 1 — Variable map

refresher on $W_l^{\mathrm{fused}}$:

this is a weight matrix, made by concatonating the three weight matrices $W_l^{\text{pre}}, W_l^{\text{post}}, W_l^{\text{comb}}$ along the first dimension. Instead of computing three separate GEMM operations, we can compute one fused GEMM operation to get the same result. This is a common optimization technique in deep learning.

this is included in the math document, but I thought I’d give a refresher.

template parameters:

template <int NUM_SPLITS, int BLOCK_SIZE, bool COMPUTE_SQRSUM>

passing these as template parameters, allows them to be locked at compile time, and the compiler can unroll loops and hardwire memory offsets.

launch bounds:

__launch_bounds__(BLOCK_SIZE)

this is an upper bound on the number of threads per block that will be launched for this kernel. the compiler can aggressively optimize for this number of threads instead of being conservative.

function declaration:

__global__ void mhc_pre_big_fuse_kernel(...)

standard cuda kernel declaration. the variables were already explained in the table above.

special flags used:

const float* __restrict__ dot_mix, 
const __nv_bfloat16* __restrict__ residual,

restrict promises the compiler that these pointers are not aliased. meaning, they’ll never point at the same memory. Without this, the compiler won’t pull data from these addresses to read only caches for optimization.

hardware data types:

__nv_bfloat16, float

Table 2 — Hardware data types

as documented in Table 1, params like the mix vectors, and rstd are kept in high precision, as they are used in operations that require full numerical accuracy. but where possible ((for example the residual itself), we use the lower-precision __nv_bfloat16 type to save memory bandwidth and storage, while still maintaining a large dynamic range.

static asserts

static_assert(NUM_SPLITS == 1 || !COMPUTE_SQRSUM,
              "internal sqrsum path only supports NUM_SPLITS=1");

static_assert acts as a fail-safe. It stops the code from even compiling if someone tries to launch the kernel with an unsupported combination of optimizations.

high level overview of the kernel and the optimizations applied

one token = one block

each thread block is responsible for processing one token.

this makes things simpler, because my math document (and the deepseek paper) both describe the mHC operations on a per-token basis. so we can directly map the math to the code without worrying about how multiple tokens are handled within a block.

warp 0 vs warps 1+

the kernel is designed such that warp 0 (specifically the first 4 threads in the block) is responsible for computing the metadata (the mix vectors and the routing matrix), while warps 1+ are responsible for the reduction to compute the final output.

const int warp_id = threadIdx.x / kWarpThreads;
const int lane = threadIdx.x & (kWarpThreads - 1);
const bool meta_lane = (warp_id == 0 && lane < kHc4);

kHc4 = 4, meaning this kernel is specialized for exactly 4 hyper-connection lanes.

constants in the kernel

at the top of the file we have:

static constexpr int kHc4 = 4;
static constexpr int kMixHc4 = kHc4 * (2 + kHc4);
static constexpr int kWarpThreads = 32;
static constexpr int kVecWidth = 8;

these constants tell us a lot about how the kernel is designed.

Table 3 — Kernel constants

why kMixHc4 = 24?

this is exactly:

so for every token, dot_mix stores 24 float values generated by the previous fused GEMM.

vectorized bf16 loads and stores

the kernel defines a small helper struct:

struct Bf16x8 {
  uint4 raw;
};

uint4 is 4 unsigned 32-bit values, so total size is:

since bf16 is 16 bits, 16 bytes contains:

this is why it is called Bf16x8.

asm volatile("ld.global.v4.u32 {%0, %1, %2, %3}, [%4];\n" ...)
asm volatile("st.global.v4.u32 [%4], {%0, %1, %2, %3};\n" ...)

Here, we use inline PTX assembly. We force the compiler to load 128-bits of data (in the form of a vector of 8 bf16 values)from the global memory, utilizing the whole bus. this is known as “manual emission”.

square sum and RMSNorm

there are two ways the kernel can get the RMSNorm denominator.

path 1: sqrsum is already provided

in the normal mhc_pre_big_fuse function, sqrsum is passed in from outside:

const float* __restrict__ sqrsum

if COMPUTE_SQRSUM == false, then the kernel does:

sq_total = sqrsum[token];

or if split-K was used:

sq_total += sqrsum[s * total_tokens + token];

so the square sum was computed before this kernel, probably during the previous projection/reduction stage.

path 2: compute sqrsum inside this kernel

there is also a mhc_pre_big_fuse_with_prenorm function. In that path, the kernel launches with:

dispatch_pre_big_fuse<1, true>(...)

so COMPUTE_SQRSUM = true.

then the kernel runs:

const float local_sq = residual_square_sum_vec8<BLOCK_SIZE>(residual_token, kHc4 * H);
sq_total = block_sum<BLOCK_SIZE>(local_sq, sq_partials);

this streams through the residual tensor, squares every element, and reduces it inside the block. “reducing” here means taking the sum across the feature dimension, so we end up with the total sum of squares for that token. This is callde a “reduction” operation in parallel programming, because we are reducing a large set of values down to a single value (the total sum).

mathematically this is computing:

then the reciprocal RMS is:

rstd = rsqrtf(sq_total / static_cast<float>(K) + rms_eps);

which maps to:

this is exactly the denominator of RMSNorm.

block reduction details

there are two reduction helpers:

warp_sum(float v)
block_sum<BLOCK_SIZE>(float v, float* shared)

warp_sum uses warp shuffle instructions:

v += __shfl_down_sync(0xffffffffu, v, offset);

this reduces values inside one warp without shared memory. This is faster than writing every partial sum to shared memory, because warp lanes can exchange register values directly.

block_sum then does a two-stage reduction:

1. every warp reduces its own lanes using warp_sum
2. lane 0 of each warp writes a partial sum to shared memory
3. warp 0 reduces those warp-level partial sums
4. thread 0 publishes the final result back into shared memory

this is a standard fast block reduction pattern.

original path vs my h100 path

there are two versions of the metadata generation in the same kernel. The first one is the original/general implementation, and the second one is the H100-specific path I added after looking through the generated code and profiling behavior.

#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 900)
  // Hopper fast path
#else
  // Portable path
#endif

__CUDA_ARCH__ == 900 means the code is being compiled for sm90, which is H100/Hopper. This lets us keep the original path available for other architectures, while specializing the hot path for the hardware I was targeting.

the original/general path

the original code took a general path, with decent portability and readable structure. here, warp 0 uses a local array:

float y_local[kMixHc4];

that is 24 floats.

then it loads the full metadata row:

for (int c = 0; c < kMixHc4; ++c) {
  y_local[c] = y_row[c];
}

and calls:

write_token_metadata(y_local, rstd, ...)

this code is clean, and it matches the math nicely. But the issue I found is that the helper function uses lane-dependent indexing:

y_local[lane]
y_local[kHc4 + lane]
y_local[2 * kHc4 + lane * kHc4 + k]

from the programmer’s perspective this is obvious: lane 0 handles row 0, lane 1 handles row 1, etc.

but from the compiler’s perspective, lane is dynamic. It is not a compile-time constant. The compiler may not be able to keep the whole y_local array in registers, because indexing into an array with a dynamic index often forces it to put that array into local memory.

Note: CUDA “local memory” is not the same as a CPU local variable. It usually spills to global memory and is cached. So if y_local spills, the metadata warp now has extra global-memory traffic just to access 96 bytes of temporary values.

this was in fact occouring when I inspected the cuobjdump output STACK:96. I could see the register spill/local-memory behavior there, even though I could not dynamically inspect it with Nsight on an H100. That was the first clue that the nice modular helper function was costing us real performance.

I verified after the fix and STACK:0, LOCAL:0 there was no local memory spill, and the kernel was now fully register-resident for the metadata work.

The second clue was the warp scheduling. Warps 1+ were blocked at the barrier while warp 0 finished all the metadata work, including the Sinkhorn loop. So the kernel had two problems:

1. warp 0 was slower than it needed to be because of the local-memory spill.
2. the rest of the block was waiting for warp 0, even though it only needed pre_mix.

I used claude to confirm these observations, and breakdown how to fix the issue.

my h100 path

on H100, I changed the path to avoid the y_local[24] temporary and instead keep only what each lane actually needs.

for each metadata lane:

float y_pre = 0.0f;
float y_post = 0.0f;
float cmv[kHc4];

so lane 0 only holds the row needed for lane 0, lane 1 holds the row needed for lane 1, and so on.

for NUM_SPLITS == 1, the load is:

y_pre = y_row[lane];
y_post = y_row[kHc4 + lane];
for (int k = 0; k < kHc4; ++k) {
  cmv[k] = y_row[2 * kHc4 + lane * kHc4 + k];
}

this is much more register friendly. cmv[k] is accessed with unrolled constant k, and there is no 24-element array indexed by a dynamic lane id. The goal was very direct: keep the metadata slice in registers and stop paying local-memory traffic for a tiny 96-byte temporary.

for split-K, it accumulates the same slice across splits:

for (int s = 0; s < NUM_SPLITS; ++s) {
  sq_total += sqrsum[s * total_tokens + token];
  const float* y_row = dot_mix + (s * total_tokens + token) * kMixHc4;
  y_pre += y_row[lane];
  y_post += y_row[kHc4 + lane];
  for (int k = 0; k < kHc4; ++k) {
    cmv[k] += y_row[2 * kHc4 + lane * kHc4 + k];
  }
}

this is the same idea for split-K. each split produces a partial dot_mix, and this path accumulates only the values this lane owns instead of rebuilding the full 24-float row.

explaining split-K

for one token, the fused projection is:

but $\hat{X_l}$ is huge because it is the flattened 4-lane residual:

and for DeepSeek-style hidden sizes, $d$ can be very large.

Split-K means the previous GEMM can split the long K dimension into chunks. Each chunk computes a partial dot product, and then this kernel adds those partial dot products:

same for the square sum:

why this helps: if the GEMM is too skinny or has too few output columns (only 24 values per token), split-K can create more parallel work for the GPU. It trades an extra reduction for better occupancy/utilization in the projection stage.

computing pre_mix early

this was the second optimization I made. After fixing the register/local-memory issue, the next obvious waste was that most of the block was just sitting at the barrier.

The old logical order was:

1. load all metadata
2. compute pre_mix
3. compute post_mix
4. run Sinkhorn for comb_mix
5. __syncthreads()
6. compute layer_input = A_l X_l

that is easy to read, but bad for overlap. Warps 1+ are idle while warp 0 does all the scalar metadata math.

The thing I noticed is that warps 1+ do not need all metadata. They only need pre_mix, because their job is just:

so I changed the order:

1. warp 0 loads metadata slice
2. warp 0 computes only pre_mix
3. __syncthreads() publishes pre_mix
4. warps 1+ start computing layer_input
5. meanwhile, warp 0 computes post_mix and comb_mix

the code now reflects that split:

if (meta_lane) {
  ...
  rstd = rsqrtf(...);
  const float v = y_pre * rstd * mhc_scale[0] + mhc_base[lane];
  pre_mix[lane] = 1.0f / (1.0f + expf(-v)) + mhc_pre_eps;
}

__syncthreads();

if (meta_lane) {
  // post_mix and Sinkhorn happen here
}

write_layer_input<BLOCK_SIZE>(...);

after this, the kernel no longer waits for the Sinkhorn loop before starting the heavy memory streaming work.

main idea: move the barrier to the first real dependency. not after all metadata is done, only after pre_mix is ready.

execution timeline

here is the optimized execution flow:

flowchart TD
    A[block starts: one token] --> B{COMPUTE_SQRSUM?}
    B -->|yes| C[all threads compute residual square sum]
    B -->|no| D[warp 0 reads sqrsum]
    C --> E[warp 0 lanes 0 to 3 load dot_mix slice]
    D --> E
    E --> F[warp 0 computes rstd and pre_mix]
    F --> G[__syncthreads publishes pre_mix]
    G --> H[warps 1+ compute layer_input = A_l X_l]
    G --> I[warp 0 computes post_mix]
    I --> J[warp 0 runs Sinkhorn]
    J --> K[warp 0 writes comb_mix]
    H --> L[block done]
    K --> L

this diagram is a bit simplified, but the key idea is that after the barrier, two independent pieces of work happen at the same time:

• warp 0 does scalar-heavy metadata work
• warps 1+ do bandwidth-heavy residual compression

computing $A_l$ (pre_mix)

The pre-mix gate is:

in code:

const float v = y_pre * rstd * mhc_scale[0] + mhc_base[lane];
pre_mix[lane] = 1.0f / (1.0f + expf(-v)) + mhc_pre_eps;

mapping:

so:

computing $C_l$ (post_mix)

After the barrier, warp 0 computes post-mix:

const float vp = y_post * rstd * mhc_scale[1] + mhc_base[kHc4 + lane];
post_mix[token * kHc4 + lane] =
    1.0f / (1.0f + expf(-vp)) * mhc_post_mult_value;

mapping:

so:

usually mhc_post_mult_value should be 2, matching the math document.

computing $B_l$ (comb_mix)

The comb-mix matrix is the most expensive scalar part of the kernel.

first the raw values are scaled and shifted:

for (int k = 0; k < kHc4; ++k) {
  cmv[k] = cmv[k] * rstd * mhc_scale[2] + mhc_base[2 * kHc4 + lane * kHc4 + k];
}

this maps to:

then the code applies softmax row-wise:

const float row_max = fmaxf(fmaxf(cmv[0], cmv[1]), fmaxf(cmv[2], cmv[3]));
for (int k = 0; k < kHc4; ++k) {
  cmv[k] = expf(cmv[k] - row_max);
}
float row_sum = cmv[0] + cmv[1] + cmv[2] + cmv[3];
for (int k = 0; k < kHc4; ++k) {
  cmv[k] = cmv[k] / row_sum + mhc_sinkhorn_eps;
}

subtracting row_max is the standard numerically stable softmax trick. It avoids expf overflow.

then it does column normalization using warp shuffles:

float col_sum = cmv[k];
col_sum += __shfl_xor_sync(kLaneMask, col_sum, 1);
col_sum += __shfl_xor_sync(kLaneMask, col_sum, 2);
cmv[k] /= (col_sum + mhc_sinkhorn_eps);

here kLaneMask = (1u << kHc4) - 1, which is:

so only lanes 0..3 participate in the shuffle.

why xor 1 and xor 2?

with four lanes, xor patterns let each lane exchange values with the other lanes:

• lane 0 exchanges with lane 1 using xor 1
• lane 0 exchanges with lane 2 using xor 2
• after both steps, the value has been combined across lanes 0,1,2,3

this is a tiny 4-lane reduction inside warp 0. Since each lane represents one row of the 4x4 matrix, the shuffle lets the rows communicate to compute column sums.

Then the fixed Sinkhorn loop runs:

for (int it = 1; it < sinkhorn_repeat; ++it) {
  row_sum = cmv[0] + cmv[1] + cmv[2] + cmv[3] + mhc_sinkhorn_eps;
  for (int k = 0; k < kHc4; ++k) {
    cmv[k] /= row_sum;
  }
  for (int k = 0; k < kHc4; ++k) {
    float col_sum = cmv[k];
    col_sum += __shfl_xor_sync(kLaneMask, col_sum, 1);
    col_sum += __shfl_xor_sync(kLaneMask, col_sum, 2);
    cmv[k] /= (col_sum + mhc_sinkhorn_eps);
  }
}

this produces an approximate doubly-stochastic matrix:

• rows sum to about 1
• columns sum to about 1
• entries are non-negative

finally it writes:

comb_mix[token * kHc4 * kHc4 + lane * kHc4 + k] = cmv[k];

so each metadata lane writes one row of the 4x4 comb_mix matrix.

computing layer_input $A_l X_l$

this is done by write_layer_input.

if (warp_id == 0) {
  return;
}

warp 0 does not participate. That warp is reserved for metadata. Warps 1+ stream through hidden dimension H.

const int p2_tid = tid - kWarpThreads;
constexpr int p2_threads = BLOCK_SIZE - kWarpThreads;
for (int h = p2_tid * kVecWidth; h < H; h += p2_threads * kVecWidth) {
  ...
}

p2_threads means “phase 2 threads” or the threads after warp 0. If block size is 256, then:

so 224 threads are used to sweep across the hidden dimension.

For each position h, the kernel loads the same hidden slice from all 4 lanes:

for (int j = 0; j < kHc4; ++j) {
  const Bf16x8 raw = load_bf16x8(residual + j * H + h);
  ...
  acc[...] += pre_mix[j] * f.x;
}

mathematically:

or in the notation from the math document:

finally it converts the fp32 accumulator back to bf16:

pairs[v] = __float22bfloat162_rn(make_float2(acc[2 * v], acc[2 * v + 1]));
store_bf16x8(layer_input + h, out);

important detail: accumulation is done in fp32 even though inputs/outputs are bf16. This is the usual mixed precision pattern: bf16 for bandwidth, fp32 for accumulation accuracy.

elaborating on the barrier in pre_mix

there is still a __syncthreads() after pre_mix is written.

this is required because pre_mix is in shared memory:

__shared__ float pre_mix[kHc4];

warp 0 writes it, and warps 1+ read it.

without the barrier, warps 1+ could read uninitialized or partially written pre_mix values. So the barrier cannot be removed entirely.

but the key optimization is that the barrier was moved earlier. It only waits for pre_mix, not for the full metadata generation.

memory layout assumptions

The shape checks at the bottom tell us the expected tensor layouts.

residual:    [total_tokens, 4, H] bf16
post_mix:    [total_tokens, 4] float
comb_mix:    [total_tokens, 4, 4] float
layer_input: [total_tokens, H] bf16
mhc_scale:   [3] float
mhc_base:    [24] float

and H must be divisible by 8:

TVM_FFI_ICHECK_EQ(*H % kVecWidth, 0) << "hidden size must be divisible by 8";

this is because all global memory loads/stores in the layer input path use Bf16x8 vectorization. If H was not divisible by 8, the last vector load could run past the end.

block size choices

there are two default block size selectors:

static int select_pre_big_fuse_block_size(int /* total_tokens */) { return 256; }
static int select_pre_big_fuse_with_prenorm_block_size(int /* total_tokens */) { return 128; }

normal path uses 256 threads. prenorm path uses 128 threads.

my understanding:

• in the normal path, sqrsum is already available, so most of the work after metadata is the layer_input bandwidth sweep. 256 threads gives 224 useful memory threads after reserving warp 0.
• in the prenorm path, the kernel also computes sqrsum internally. This adds a block-wide reduction and may increase register/shared-memory pressure. 128 threads can be a better tradeoff.

also, the dispatch supports 128, 256, and 512 block sizes if the caller overrides it.

if (block_size >= 512) launch 512
else if (block_size >= 256) launch 256
else launch 128

launch and split dispatch

num_splits is dispatched at compile time:

switch (num_splits) {
  case 1:  DISPATCH_SPLITS(1);  break;
  case 2:  DISPATCH_SPLITS(2);  break;
  case 4:  DISPATCH_SPLITS(4);  break;
  case 8:  DISPATCH_SPLITS(8);  break;
  case 16: DISPATCH_SPLITS(16); break;
}

this is important because NUM_SPLITS is a template parameter. If it was a runtime loop bound, the compiler would have a harder time unrolling the accumulation loop and optimizing address arithmetic.

supported split counts are powers of two up to 16.

why this kernel is called pre

from the full mHC equation:

this kernel computes the things needed before the expensive block $F_l$:

1. pre_mix = A_l
2. layer_input = A_l X_l
3. post_mix = C_l
4. comb_mix = B_l

The next heavy block (attention or MoE) consumes layer_input. Later, a post kernel can use post_mix, comb_mix, the old residual, and the block output to form:

so “pre” does not mean it only computes $A_l$. It also prepares B_l and C_l, but the critical output for the immediate next compute block is layer_input.

performance intuition on H100

H100 / sm90 is very good at keeping many operations in flight, but only if the kernel exposes independent work.

The original schedule made the work more serial:

metadata math -> barrier -> residual compression

My H100 path changes this into:

pre_mix -> barrier -> metadata math and residual compression in parallel

breaking down the work happening in parallel:

so the scheduling change gives the SM something useful to do while warp 0 is still finishing the Sinkhorn math.

The other big improvement is the register/local-memory issue. A 96-byte local array sounds tiny, but in the cuobjdump output it showed up as a real spill. That spill was extra traffic in the exact part of the kernel that everyone else was waiting on. So fixing it mattered more than the size of the array suggests.

Results

The two changes — spill elimination and Sinkhorn overlap — were applied to the mhc_pre_big_fuse kernel in FlashInfer PR #3285. All 26 existing tests continued to pass.

The pure path — the primary production path, which consumes a precomputed square sum from the fused GEMM — now reaches 88–91% of the H100’s HBM roofline (3350 GB/s), up from 67–78%. The spill fix dominates at smaller batch sizes and hidden dimensions (where metadata cost is proportionally larger), while the overlap recovers the rest.

summary of concepts

Table 4 — Concepts in mhc_pre_big_fuse.cu

final mental model

The kernel is doing two conceptually different jobs at the same time.

The math job:

The data movement job:

The important realization is that the data movement job only needs $A_l$. It does not care about $B_l$ or $C_l$ yet.

so the optimized kernel computes $A_l$ first, synchronizes once, and then lets the GPU overlap:

• computing $B_l$ and $C_l$ on warp 0
• computing $A_lX_l$ on the rest of the block

this was a good example of GPU programming being less about “doing fewer operations” and more about “placing the operations in an order where the hardware is not waiting around”.

notes / possible followups

things I still want to study further:

• profile the kernel with Nsight Compute and compare local memory transactions between the original path and my H100 path.
• study the matching mhc_post kernel to complete the full equation $X_{l+1} = B_lX_l + C_lF_l(A_lX_l)$.
• Verify my improvement on other architectures (B100, B200, etc.) and see if the same register/local-memory issue exists there.

The standard residual connection

exists, as an innovation in backpropagation. As my previous blog post covered, backpropogation requires smooth gradients, close to 1 (or in the case of matrices, close to the identity matrix I). Now, observe the derivative of the above equation:

you might notice, even if $F’(x_l) \approx 0$, the identity path $I$ preserves gradient flow. This lets us train networks dozens of layers deep while diminishing the vanishing gradient problem.

The Problem with Standard Residuals

Modern deep networks have multiple channels (or lanes) of $X$, each with a full embedding. This allows the model to learn different features in different lanes, and fuse them at the end.

However, this creates challenges in designing skip connections. As discussed above, a skip connection acts as a path for gradient to flow. The model might learn to only use one channel, and ignore the rest.

mHC addresses this by (1) projecting multiple lanes into a smaller compute lane, (2) applying a heavy block (MoE / Attention) there, (3) redistributing the result back to the full lane dimension, and (4) imposing strict constraints on the skip connection to ensure stability across 100+ layers. This creates a structured “hyper-connection” that mixes and routes multiple embedding lanes.

The Hyper-Connection Equation

The HC layer is written compactly as:

Components

• $X_l, X_{l+1}$: input and output of layer $l$ respectively, where $X_l \in \mathbb{R}^{n_{hc} \times d}$.
• $n_{hc}$: the number of hyper-connection lanes.
• $d$: the embedding width.
• $A_l$: projection that squeezes the $n_{hc}$ lanes into a single lane. Shape $A_l \in \mathbb{R}^{1 \times n_{hc}}$, so $A_l X_l \in \mathbb{R}^{1 \times d}$.
• $F_l(\cdot)$: the expensive compute block (e.g. MoE or multi-head attention) that operates on the compressed $\mathbb{R}^{1 \times d}$ lane and returns $\mathbb{R}^{1 \times d}$.
• $C_l$: expansion (or coloring) matrix that maps the $(1, d)$ output of $F_l$ back to $(n_{hc}, d)$. Shape $C_l \in \mathbb{R}^{n_{hc} \times 1}$, so $C_l F_l(A_l X_l) \in \mathbb{R}^{n_{hc} \times d}$.
• $B_l$: mixes the original $n_{hc}$ input lanes. Shape $B_l \in \mathbb{R}^{n_{hc} \times n_{hc}}$; input and output retain shape $X_l, B_l X_l \in \mathbb{R}^{n_{hc} \times d}$ before adding to the expanded path, creating the bypass connection.

Why

• We make a massive dimension $d = 7168$, each value in this vector is a feature of the token.
• We use multi-head attention: each head in the $7168$ embedding tries to learn a single group of features (grammar, semantics, context), and they fuse at the end using some kind of projection.
• We use multiple channels (e.g. 4), each with a full $7168$ embedding. Each channel is a global stream that learns a different feature. Before this, the single vector was overburdened, and prone to losing knowledge as it passed through attention layers.

Notes

• The ordering of linear ops shown above is convenient for exposition; implementations often fuse projections or use learned channel-wise gates.
• Using a single compressed lane for the heavy block reduces compute and memory compared with applying the block to all lanes.
• Replace or tune the shapes ($n_{hc}$ and $d$) to match your model’s lane count and embedding dimension.

Quick Reference

• Equation: $X_{l+1} = B_l X_l + C_l F_l(A_l X_l)$
• Typical shapes: $X_l: (n_{hc}, d)$, $A_l: (1, n_{hc})$, $F_l: (1, d) \to (1, d)$, $C_l: (n_{hc}, 1)$, $B_l: (n_{hc}, n_{hc})$

Constraints on the Parameter Matrices

Constraints from the paper:

Each $B_l$ belongs to a family of matrices $M$ — $n \times n$ matrices where each element $\ge 0$ and the sum of each row and each column is 1.

Hardware activations (implementation notes):

Thus:

where $\tilde A_l$ and $\tilde C_l$ are the raw, unnormalized parameters, $\sigma(\cdot)$ is element-wise sigmoid, and $\epsilon > 0$ is a small floor to avoid exact zeros in hardware.

Why

Non-Explosiveness $|B_l|_2 \le 1$: spectral norm of the matrix. This ensures multiplication by $B$ can never be increasing. Safe for forward pass and gradient backprop, avoids exploding values after 100+ layers of transformation.

Closure Under Multiplication $\forall M_1, M_2 \in M \Rightarrow M_1 M_2 \in M$: this means that as multiplication by various matrices $M$ stack, the resulting matrix is from the set $M$, and has the same restrictions and properties. This ensures $B$ is well behaved over hundreds of layers, where these transformations by $B$ will stack many times.

Non-Negative Bounding $A_{ij}, C_{ij}, B_{ij} \ge 0$: since these matrices also stack additively, this constraint ensures they don’t cause destructive interference over hundreds of layers. Without this, features could overwrite or negate each other.

Bounded Magnitudes $\sigma(\cdot) \le 1$: the sigmoid function enforces strict upper boundaries. $A_l$ is bounded near $1$, and $C_l$ is bounded at $2$. This gives the network the ability to safely scale the lanes without risking infinite amplification.

Identity Initialization $2 \cdot \sigma(\cdot)$: initially during training, weight matrices are initialized to small values close to zero. $\sigma(0) \approx 0.5$. Thus multiplication by $2$ initializes $C_l$ close to identity.

Non-Zero Floor $+\epsilon$: $A_l$ projects $n_{hc}$ lanes into 1 lane. If an element in $A_l$ is ever exactly $0$ — possible through floating-point rounding, especially at low-bit quantizations — the lane is disconnected from the next layers during forward pass, and previous layers during backprop. Thus we add a small offset to ensure this doesn’t occur.

The Sinkhorn-Knopp Algorithm

To enforce the doubly-stochastic constraint on $B_l$, we use the iterative Sinkhorn-Knopp algorithm. Very simple.

1. Start with $M^{(0)} = \text{Softmax}(\tilde B_l)$ — ensures non-negativity, one of our constraints.
2. Iterate for $t = 1, \dots, t_{\max}$:

1. Set $B_l = M^{(t_{\max})}$.

After row normalization, the matrix is normalized across rows but not columns. After column normalization, the matrix is normalized across columns but not rows. After each iteration, the matrix converges toward full constraint satisfaction.

Dynamic Parameterization

The parameters $A_l$, $B_l$, $C_l$ are not simple learned constants. They are dynamically generated from the input $X_l$ itself, using static and dynamic components.

The input is flattened and normalized:

where $\text{vec}(X_l) \in \mathbb{R}^{1 \times n_{hc} \cdot d}$ is a flattening operation on $X_l \in \mathbb{R}^{n_{hc} \times d}$, and RMSNorm is:

Then each raw parameter is a combination of static and dynamic components:

where:

these $W$ matrices are learnable weights for generating dynamic components.

Components

• $\tilde A_l, \tilde B_l, \tilde C_l$: the raw, unnormalized parameters for the $A$, $B$, and $C$ matrices.
• $\alpha_l^{\text{pre}}, \alpha_l^{\text{comb}}, \alpha_l^{\text{post}}$: learnable scalar parameters that control the influence of the dynamic components for $A$, $B$, and $C$ respectively.
• $S_l^{\text{pre}}, S_l^{\text{comb}}, S_l^{\text{post}}$: learnable static parameters for $A$, $B$, and $C$ respectively.
• $W_l^{\text{pre}}, W_l^{\text{comb}}, W_l^{\text{post}}$: learnable weight matrices that project the normalized input $\hat X_l$ into the dynamic components for $A$, $B$, and $C$ respectively.
• $\text{Mat}(\cdot)$: a reshaping operation that converts $\hat X_l W_l^{\text{comb}} \in \mathbb{R}^{1 \times n_{hc}^2}$ into $\text{Mat}(\hat X_l W_l^{\text{comb}}) \in \mathbb{R}^{n_{hc} \times n_{hc}}$.

Why

• Dynamic Components $(\hat X_l W_l)$: these allow the model to make token-specific routings for each hyper-connection. For example, for a token that is a verb, the model might learn to route more information through lane 1, and for a token that is a noun, it might route more through lane 2.
• Static Components $S_l$: these allow the model to learn general patterns of routing that are useful across all tokens.
• Scaling Factors $\alpha_l$: these allow the model to control the overall influence of the dynamic components.
• Training Stability $\alpha_l \approx 0$: the model can use static components as a baseline, and organically rely on dynamic components as training progresses.
• RMSNorm: normalizing the input to the dynamic parameter generation helps stabilize training and ensures the dynamic parameters are generated from a consistent scale of input features — geometrically only the direction is considered, the magnitude is normalized.
• Flattening $\text{vec}(X_l)$: this allows the dynamic parameter generation to consider interactions across all lanes and embedding dimensions, enabling more complex and informed routing decisions.
• Reshape $\text{Mat}(\hat X_l W_l^{\text{comb}})$: this is needed to match the shape of the $B$ matrix ($n_{hc} \times n_{hc}$), allowing it to properly mix the hyper-connection lanes based on the dynamic input.

Hardware Considerations

The three separate matrix multiplications for generating $\tilde A_l$, $\tilde B_l$, and $\tilde C_l$ can be fused into a single large matrix multiplication:

with $W_l^{\text{fused}} = [W_l^{\text{pre}}; W_l^{\text{comb}}; W_l^{\text{post}}]$ concatenating the individual weight matrices. Instead of three separate GEMM operations, a single fused GEMM produces all the raw parameters at once. This is more efficient on GPU hardware, as it reduces the number of separate operations and allows better utilization of the GPU’s parallel processing capabilities.

For $n_{hc} = 4$, the output dimension is $4 + 16 + 4 = 24$ values per token. This small output dimension creates a challenge for the GEMM — it’s very skinny. The GPU can have difficulty saturating its compute units when the output has only 24 columns per token, which is why techniques like Split-K (dividing the input dimension into chunks) are used to create more parallelism.

Summary

mHC replaces the simple additive residual with a structured multi-lane routing mechanism. The key innovations are:

• Compress-Compute-Expand architecture: reduces compute by applying the expensive block on a single compressed lane.
• Mathematical constraints on $A_l, B_l, C_l$: ensure stability across hundreds of layers through non-explosiveness, closure under multiplication, and non-negative bounding.
• Dynamic parameterization: generates token-specific routing weights by combining input-dependent and static components.
• Sinkhorn-Knopp projection: enforces doubly-stochastic constraints on the routing matrix with a fixed-iteration GPU-friendly algorithm.

Small Authors Note:

This is the first of two posts on mHC. The next post will cover the practical implementation details of the mhc_pre_big_fuse kernel (from the repository flashinfer, a frontier gpu inference library), which computes the dynamic parameters and applies the Sinkhorn-Knopp algorithm, on GPU. The kernel is a critical part of the mHC block, and optimizing it was a major focus of my work on this project. I will discuss the challenges of implementing the Sinkhorn-Knopp algorithm efficiently on GPU, and how I overcame some short comings of the existing kernel to achieve high performance. Also, my health has improved, I am now on a consistent routine of studying, practicing and documenting. Along with exercise. I hope the summer will bring enough positive change, that I can get employed soon. Stay tuned! </small>

This post walks through backprop on a two-layer FFN, step by step. The same pattern applies to any differentiable computation graph: convolutions, attention, transformers.

The Chain Rule (Refresher)

For scalar functions, the chain rule says:

For vector functions $f: \mathbb{R}^n \to \mathbb{R}^m$ and $g: \mathbb{R}^m \to \mathbb{R}^o$, the Jacobian form is:

Each entry $(i, j)$ of the product is $\sum_k \frac{\partial f_i}{\partial g_k} \cdot \frac{\partial g_k}{\partial x_j}$: the chain rule summed over the intermediate dimension $k$. This is the mathematical basis of backprop: multiply Jacobians from the output to each parameter, reusing intermediate products.

Forward Pass: The 2-Layer FFN

Use the same two-layer FFN from the feed-forward networks post. Input $x \in \mathbb{R}^d$, hidden width $h$, output dimension $o$:

The forward pass computes each step in order and caches the intermediate values $z_1$, $a_1$, $z_2$ (they will be needed during the backward pass). Without caching, the backward pass would have to recompute every intermediate, doubling the total cost.

Backward Pass

Start at the loss $L$ and walk backward through each operation, computing the gradient of $L$ with respect to each intermediate and each parameter.

Step 1: Output gradient

This is the “error signal” that will propagate backward.

Step 2: Output layer ($z_2 = W_2 a_1 + b_2$)

Given $\frac{\partial L}{\partial z_2}$, compute gradients for $W_2$ and $b_2$, plus the error signal for $a_1$:

Each entry $(i, j)$ of $\partial L / \partial W_2$ is $\frac{\partial L}{\partial z_2}_i \cdot (a_1)_j$: the error for output $i$ multiplied by the activation of hidden unit $j$. If hidden unit $j$ was never active, $(a_1)_j = 0$, and its weights get zero gradient.

Intuition: “The gradient of a weight is the product of two things: the error signal from above and the input that fed into it from below. This is a rank-1 update: outer product of error and activation.”

Step 3: Activation ($a_1 = \sigma(z_1)$)

Pass the error signal through the activation derivative:

For ReLU, $\sigma’(z) = 1$ if $z > 0$, $0$ otherwise. If $z_1$ was negative during the forward pass, $\sigma’(z_1) = 0$ and the error signal is completely blocked: this is the “dead ReLU” problem.

Intuition: “The activation derivative acts as a gate. If $\sigma’(z) \approx 0$, the error signal cannot pass through and the gradient for everything below is zero. Smooth activations like GeLU and SiLU keep this gate open for all inputs.”

Step 4: Hidden layer ($z_1 = W_1 x + b_1$)

Same pattern as step 2, using $x$ as the input:

Step 5: Update with gradient descent

Now we have all four gradients: $\partial L / \partial W_1$, $\partial L / \partial b_1$, $\partial L / \partial W_2$, $\partial L / \partial b_2$. Use the gradient descent update from the GD post:

and the same for $W_2$, $b_2$.

The Full Chain Rule in One Expression

Writing the entire chain from $L$ to $W_1$ as a single product:

Each “$\cdot$” is a Jacobian multiplication. Backprop computes this product from right to left, one factor at a time. The intermediate products ($\partial L / \partial z_2$, $\partial L / \partial a_1$, $\partial L / \partial z_1$) are cached and reused for multiple parameters at the same level, which is why the total cost is $O(P)$ and not $O(P^2)$.

Why Activation Functions Matter

The derivative $\sigma’(z)$ determines how much of the error signal reaches the layers below.

• ReLU has a hard kink at $x=0$ where the derivative jumps from 0 to 1. The derivative is 0 for $x < 0$, meaning any neuron with a negative pre-activation completely blocks gradient flow. Once many neurons are “dead”, the network stops learning.

• Sigmoid has a smooth bell-shaped derivative. For large $

  x

  $, the derivative is close to 0, causing vanishing gradients in deep networks.

• Tanh has the same saturation problem as sigmoid. Its derivative peaks at $x=0$ and decays to 0 at the extremes.

This is why modern networks use ReLU (or its smooth variants GeLU/SiLU): the derivative is 1 for positive inputs, giving a clean gradient signal through most of the network.

Components

• $L$: loss function. A scalar that measures how wrong the prediction is.
• $\hat{y}$: prediction (output of the forward pass).
• $z_\ell$: pre-activation at layer $\ell$ (before the non-linearity).
• $a_\ell$: activation at layer $\ell$ (after the non-linearity).
• $W_\ell$, $b_\ell$: weights and biases at layer $\ell$. The parameters we want gradients for.
• $\sigma$: activation function, applied element-wise.
• $\partial L / \partial \cdot$: gradient of $L$ with respect to some tensor. These flow backward through the graph.

Quick Reference

Notes

• The total cost of backprop is approximately $2\times$ the forward pass cost. Each parameter gets used once in the forward direction and once in the backward direction.
• Memory is the real bottleneck: all cached intermediates ($z_1$, $a_1$, $z_2$) must be stored until the backward pass finishes. This is why large models use gradient checkpointing (recompute some intermediates instead of storing them).
• Modern frameworks (PyTorch, JAX) implement automatic differentiation: you define the forward pass and the framework traces the computation graph, then builds the backward pass automatically. The gradient formulas are exactly the ones derived above.
• The same chain rule pattern applies to any architecture: convolutions, attention, skip connections. Each operation defines a forward function and a backward “gradient function” that multiplies the incoming error by the local Jacobian transpose.
• Skip connections (residual networks) help gradient flow by providing a path where the Jacobian is the identity matrix, so the error signal bypasses the activation derivative entirely.

Small Authors Note: I was quite demotivated during this semester… and my grades suffered, I didn’t score any internships, I had concerning symptoms from chronic intestinal ulcers, and I was unable to keep this blog going. With summer here, I am no longer pursuing corporate work for the time being. I will focus on self development, personal projects, life style improvements, and this blog will take a shift. This post was half-written, I have finished it and it will be followed by one of my greatest accomplishments, a PR on flashinfer (maintainers haven’t noticed it yet, but I’m still proud of myself). I have a large backlog of writing, that I need to format (I will likely be taking help from opencode’s generous deepseek v4 flash access). I will cover some more ML topics (FFN, MoE, etc.). And generally document my journey towards becoming a better developer and researcher. Hopefully I find sustainable employment in the future, due to these efforts.

$d_0$ is the input dimension, $d_1, \dots, d_{L-1}$ are the hidden dimensions (widths), and $d_L$ is the output dimension. Each $W_\ell$ mixes the entries of the previous layer, $b_\ell$ shifts the result, and $\sigma_\ell$ breaks the linearity. The output layer usually skips the activation so the output can be any real value.

Why You Need the Activation

If you drop all $\sigma_\ell$, the whole chain collapses:

No matter how deep, a stack of linear layers is just one linear transformation. Without non-linearity the network can only learn straight-line decision boundaries.

The textbook example is XOR. The four points $(0, 0)$, $(0, 1)$, $(1, 0)$, $(1, 1)$ have labels $0, 1, 1, 0$ (true when exactly one input is 1). No single line separates the 0s from the 1s. But a two-layer FFN can: the hidden layer projects the input into a space where the classes are linearly separable, and the output layer draws the separating line.

Intuition: “Each hidden neuron detects a pattern: it fires (activates) when the input matches a certain weighted combination. The output layer then combines these detected patterns into a decision. More hidden neurons means more patterns can be memorized, and more layers means patterns can build on simpler patterns.”

Two-Layer Case and Universal Approximation

The simplest useful FFN has one hidden layer:

$h$ is the hidden width (number of feature detectors). Each column of $W_1$ is a direction in input space, and $\sigma$ decides whether the input projects onto that direction.

The Universal Approximation Theorem (Cybenko 1989, Hornik 1991) says: with enough hidden neurons, a single-hidden-layer FFN can approximate any continuous function on a bounded domain, to any accuracy you want. This is why FFNs are everywhere: they are general-purpose function approximators.

Intuition: “Each hidden neuron contributes a ‘bump’ to the function. The output layer is a weighted sum of these bumps. With enough bumps you can build any shape. It is the same idea as a Fourier series or a Taylor expansion, except the network learns the basis functions from data instead of using fixed sines or polynomials.”

Activation Functions

Every activation is a trade-off between gradient flow, computational cost, and output range:

ReLU is the default for most networks: one comparison, gradient is either 0 or 1, and it produces sparse activations (many neurons output exactly 0). The downside is the “dead ReLU” problem: once a neuron’s weights push its input below zero for every training example, its gradient is always 0 and it never recovers.

GeLU and SiLU are smooth alternatives. They avoid the hard kink at 0 and gradient flows through them more consistently. They cost more to compute but produce better results at the same model size. All modern LLMs use SiLU inside a gated FFN (SwiGLU).

Sigmoid and tanh were the standard before ReLU. They saturate at large positive and negative values, which kills gradients in deep networks. They are still used when a bounded output is needed (e.g., LSTM gates need values in $(0, 1)$).

Where FFNs Show Up

FFNs are the decision-making backbone of most architectures. The feature extractor (convolution, attention, recurrence) produces a representation; the FFN maps that representation to the final output.

• CNN classifier head: after convolutional layers extract spatial features, a few fully-connected layers (an FFN) map the feature vector to class scores. This is the standard pattern in ResNet, VGG, and EfficientNet.
• Transformer FFN sub-layer: every transformer block has a two-layer (or gated three-layer) FFN after attention. Attention handles inter-token mixing; the FFN processes each token independently through an expansion-activation-contraction pattern. Every major LLM (GPT, Llama, DeepSeek) uses this design.
• LSTM/GRU gates: the forget, input, and output gates are linear projections followed by sigmoid (to get a value in $(0, 1)$). Each gate is a small FFN with a fixed sigmoid activation.
• RNN output projection: the hidden-to-output mapping at each timestep is a linear layer (a depth-1 FFN without activation).

Components

• $x \in \mathbb{R}^{d_0}$: input vector (e.g., flattened image pixels, token embedding).
• $W_\ell \in \mathbb{R}^{d_\ell \times d_{\ell-1}}$: weight matrix for layer $\ell$. Entry $(i, j)$ controls how much unit $j$ from layer $\ell-1$ contributes to unit $i$ in layer $\ell$.
• $b_\ell \in \mathbb{R}^{d_\ell}$: bias vector for layer $\ell$. Shifts the activation threshold.
• $\sigma_\ell$: activation function, applied element-wise. The only source of non-linearity.
• $h_\ell$: hidden state after layer $\ell$ (post-activation).
• $L$: depth (number of weight matrices). Deeper networks can represent hierarchical features: earlier layers detect simple edges or patterns, later layers combine them into concepts.
• $d_1, \dots, d_{L-1}$: hidden widths. Wider layers can detect more features per layer. Most models use a hidden width larger than the input (e.g., $4\times$ in Transformers).

Quick Reference

Notes

• The number of parameters in an FFN is $\sum_{\ell=1}^L (d_\ell d_{\ell-1} + d_\ell)$ (weights + biases). Most of the cost is in the weight matrices.
• Training uses backpropagation: the chain rule computes $\partial f / \partial W_\ell$ and $\partial f / \partial b_\ell$ from the output layer backward. The gradient flows through the linear transform and the activation derivative at each step.
• In practice, FFNs are almost never trained alone. They are stacked with convolutions (CNN), attention (Transformer), or recurrence (RNN). The FFN provides the non-linear mapping; the surrounding layers handle structure (spatial, sequential, pairwise).
• The two-layer FFN with ReLU has a geometric interpretation: the first layer partitions the input space into convex polytopes, and the second layer assigns a value to each region. Deep FFNs create a hierarchy of partitions: earlier layers carve coarse regions, later layers subdivide them.

You have a function $f: \mathbb{R}^d \to \mathbb{R}$ (the loss), parameters $\theta \in \mathbb{R}^d$, and a learning rate $\alpha > 0$. At each step you compute the gradient $\nabla f(\theta_k)$, which points uphill, and move opposite to it.

Intuition: “You are standing on a foggy mountain. You cannot see the whole terrain, but you can feel the slope under your feet. To get to the valley, you step in the direction where the ground drops the most. The learning rate $\alpha$ is how big a step you take: too small and you never get there, too large and you stumble past the valley and have to turn back.”

The Learning Rate

The learning rate $\alpha$ is the single most important hyperparameter. A small $\alpha$ means reliable but slow progress. A large $\alpha$ means fast progress but risk of overshooting.

For the simple function $f(x) = x^2$ starting at $x = 1.8$:

• Too small ($\alpha = 0.01$): each step barely moves. You will reach the minimum eventually, but it takes hundreds of steps.
• Good ($\alpha = 0.1$): the steps shrink naturally as the gradient gets smaller near the minimum. You converge in 5-10 steps.
• Too large ($\alpha = 1.6$): the step overshoots past the minimum, and the next step overshoots back. You oscillate and may never settle.

Intuition: “The learning rate determines how much you trust each gradient measurement. A small rate is cautious but slow. A large rate is aggressive but unstable. The sweet spot depends on the curvature of your loss surface: flatter surfaces can handle larger rates.”

GD on a 2D Surface

Real problems have more than one parameter. The gradient is a vector, and each component tells you the slope along that parameter’s axis.

This is $f(\theta_1, \theta_2) = \theta_1^2 + 2\theta_2^2$, an elliptic paraboloid. The gradient is $\nabla f = [2\theta_1, 4\theta_2]^\mathsf{T}$. Because the surface is steeper along $\theta_2$ (coefficient 2 vs 1), GD takes bigger steps in that direction at first, then zig-zags as it approaches the minimum.

Intuition: “The gradient tells you the slope in every direction at once. You take a single step that combines all those directions. If one direction is steeper, you move more in that direction. This is why GD naturally follows the steepest path down, which is not always a straight line to the minimum.”

The 3D view makes it clear: the path starts on the steep wall, drops quickly into the valley, then follows the valley floor to the minimum.

Connection to Linear Regression

From the linear regression post, the loss is $L(\theta) = |X\theta - y|_2^2$ with gradient $\nabla L = 2X^T(X\theta - y)$. The exact solution via normal equations costs $O(nd^2 + d^3)$. Each gradient descent step costs $O(nd)$.

The tradeoff is:

• Exact solve (normal equations / QR): one shot, $O(nd^2 + d^3)$. Best for small $d$ (under a few thousand).
• Gradient descent: iterative, $O(nd)$ per step. Best for large $d$ (millions of parameters). You never form the $d \times d$ matrix $X^T X$.

This is the fundamental reason gradient descent dominates modern deep learning: models have millions of parameters, and forming a $10^6 \times 10^6$ matrix is not feasible.

Batch vs SGD vs Mini-Batch

The gradient $\nabla L = 2X^T(X\theta - y)$ uses all $n$ data points. This is batch gradient descent: accurate but expensive per step.

SGD (Stochastic Gradient Descent) approximates the gradient using a single randomly-chosen data point $i$:

The gradient estimate is noisy but unbiased. Each step costs $O(d)$ instead of $O(nd)$, and the noise can actually help escape shallow local minima.

Mini-batch GD is the practical middle ground: use 32-1024 random points per step. Noisy enough to generalize, cheap enough to be fast, and easy to parallelize on GPUs. This is what every practical deep learning system uses.

Intuition: “Batch GD is like surveying the entire mountain before each step: accurate but slow. SGD is like glancing at one random pebble under your foot: noisy but fast. Mini-batch is like looking at a small patch of ground around you: good enough to make progress, fast enough to keep moving.”

Components

• $\theta \in \mathbb{R}^d$: parameters being optimized (weights of a model).
• $f: \mathbb{R}^d \to \mathbb{R}$: loss function, a scalar measure of how bad the current parameters are.
• $\nabla f(\theta) \in \mathbb{R}^d$: gradient of $f$ at $\theta$. Vector of partial derivatives, one per parameter.
• $\alpha$: learning rate. Step size. The most important hyperparameter to tune.
• $n$: number of data points (batch size for batch GD, 1 for SGD, $b$ for mini-batch).

Quick Reference

Notes

• Gradient descent only guarantees convergence to a local minimum for convex functions. In deep learning, the loss surface is highly non-convex, but GD with mini-batches consistently finds solutions that generalize well.
• Momentum (Polyak, 1964) improves GD by adding a fraction of the previous update: $\theta_{k+1} = \theta_k - \alpha \nabla f(\theta_k) + \beta(\theta_k - \theta_{k-1})$. This smooths out oscillations and accelerates through flat regions.
• Adam (Kingma & Ba, 2015) combines momentum with per-parameter adaptive learning rates using estimates of first and second moments of the gradient. It is the default optimizer for most deep learning.
• The directional derivative (notation page) gives the mathematical foundation: $D_v f(\theta) = \nabla f(\theta) \cdot v$ is the slope along direction $v$. The negative gradient gives the steepest descent direction because it maximizes $\lvert \nabla f \cdot v \rvert$ among all unit vectors $v$.
• Learning rate schedules (step decay, cosine annealing, warmup) are essential in practice. Start with a small rate, increase it during warmup, then decay it as training progresses.

Problem Setup

We are looking for a line $\hat{y} = mx + b$. For each observed point $(x_i, y_i)$, the error (residual) is $r_i = y_i - (m x_i + b)$. We want to minimize the sum of squared errors:

Write this in matrix form. Stack all $y_i$ into a vector, and all $x_i$ into a matrix with a column of ones for the bias:

The loss becomes $L(\theta) = |X\theta - y|_2^2$, the squared Euclidean norm of the residual vector.

Minimizing the Error

Take the gradient of $L(\theta)$ and set it to zero:

This gives the normal equations:

Intuition: “We are projecting $y$ onto the column space of $X$. The residual $X\theta - y$ is orthogonal to every column of $X$: that is what $X^T (X\theta - y) = 0$ means. The solution is the point in $X$’s column space closest to $y$.”

The cost of forming $X^T X$ and solving is $O(n d^2 + d^3)$ where $d$ is the number of parameters (here $d = 2$). For small problems this is instant. But $X^T X$ squares the condition number: if $X$ is poorly conditioned, $X^T X$ is much worse.

For large problems ($n$ in the millions, $d$ in the thousands), even forming $X^T X$ is too expensive. Instead, use gradient descent: start from a random $\theta$, compute the gradient $\nabla_\theta L = 2 X^T (X\theta - y)$, step in the opposite direction, repeat. Each step costs $O(n d)$ instead of $O(n d^2 + d^3)$. The exact solution via normal equations or QR gives you the answer in one shot, but gradient descent scales to problems where one-shot methods do not fit in memory.

Modern Machine Learning problems are built around gradient descent, because numbers have gotten too large, but it is important to understand the exact solution they are approximating.

QR Is Better

Instead of the normal equations, use the QR decomposition from the previous post. Factor $X = QR$ where $Q \in \mathbb{R}^{n \times d}$ has orthonormal columns and $R \in \mathbb{R}^{d \times d}$ is upper triangular. Then:

No squaring of the condition number, and solving $R\theta = Q^T y$ with back-substitution costs $O(d^2)$. QR is backward stable; the normal equations are not.

Intuition: “Normal equations compute $X^T X$, which squares the condition number (bad for ill-conditioned problems). QR works directly on $X$ and is stable. For a simple 2-parameter line fit the difference is academic, but for wide problems ($n \gg d$) the normal equations can lose digits to rounding.”

Beyond One Line

The same idea extends to multiple inputs and multiple outputs. Let each input be $x \in \mathbb{R}^i$ and each output be $y \in \mathbb{R}^o$. The linear model becomes:

Each output is a linear combination of all inputs plus a bias. If you stack multiple inputs as rows of $X \in \mathbb{R}^{n \times i}$, the matrix form is $Y = X W^T + \mathbf{1} b^T$, and the least squares solution for $W$ and $b$ follows the same pattern.

This is a linear layer: one matrix multiply and one bias add, no activation. It is the simplest building block of a neural network.

the line seems 1d from this angle, but it is actually a plane in 3D space. The concept of a best fit line, can be generalized to a best fit plane between points in 3d.

and then a best fit hyperplane between points in higher dimensions. And it can all be built by stacking linear layers on top of each other, and adding non-linear activations in between to break the linearity and turn lines into curves. That is the feed-forward network, the topic of the next post. This will be covered soon in a Feed Forward Networks post.

Intuition: “For one input and one output, the model is a line. For two inputs and one output, the model is a plane. For $i$ inputs and $o$ outputs, the model is a hyperplane in $(i + o)$-dimensional space: a linear mapping from input to output.”

Components

• $X \in \mathbb{R}^{n \times d}$: design matrix. Rows are data points, columns are features (plus a bias column of ones).
• $y \in \mathbb{R}^n$: observed outputs (targets).
• $\theta \in \mathbb{R}^d$: parameters (slope $m$ and intercept $b$ in the 1D case; all weights and biases in the multi-dimensional case).
• $L(\theta) = |X\theta - y|_2^2$: least squares loss. Minimizing this gives the maximum likelihood estimate under Gaussian noise.
• $\hat{y} = X\theta$: predictions: the orthogonal projection of $y$ onto the column space of $X$.

Quick Reference

Notes

• The normal equations are fine when $d$ is small and $X$ is well-conditioned. For anything else, use QR (or the SVD for rank-deficient problems).
• Adding a bias column of ones is equivalent to centering the data: subtract the mean of $x$ and $y$, fit $m$, then recover $b = \bar{y} - m \bar{x}$.
• The linear layer $y = Wx + b$ is the foundation of every neural network. The next post adds a non-linearity between layers to break the linearity and turn lines into curves: that is the feed-forward network.

LU Decomposition

The Algorithm

$L$ is unit lower triangular ($L_{ii} = 1$, zeros above the diagonal), and $U$ is upper triangular (zeros below the diagonal).

Intuition: “We write $A$ as a product of a lower and an upper triangular matrix. Once you have this, solving $Ax = b$ is two cheap triangular solves: forward then backward.”

Derivation (Gaussian Elimination)

Elimination turns $A$ into $U$ by subtracting multiples of each row from the rows below. Each subtraction is an elementary matrix $E_k$. For a $3 \times 3$:

where each $E_k$ is unit lower triangular. Move them to the other side:

The product $L = (E_3 E_2 E_1)^{-1}$ is unit lower triangular, and it turns out its entries are exactly the multipliers we used during elimination, placed below the diagonal.

Concretely, for $A \in \mathbb{R}^{3 \times 3}$. During step 1 we compute multipliers $m_{21} = a_{21} / a_{11}$, $m_{31} = a_{31} / a_{11}$ and subtract $m_{21} \cdot \text{row}1$ from row 2 and $m{31} \cdot \text{row}1$ from row 3. During step 2 we compute $m{32} = a_{32} / a_{22}$ (using the updated $a_{22}$) and subtract $m_{32} \cdot \text{row}_2$ from row 3. The resulting $L$ and $U$ are:

The multipliers go directly into $L$ at the positions where they zeroed entries in $A$. This is why $LU$ is essentially just Gaussian elimination with a different accounting.

In pseudocode, the algorithm is:

for k = 1..n-1:
    for i = k+1..n:
        m = A[i,k] / A[k,k]          // multiplier
        A[i,k] = m                    // store in L
        for j = k+1..n:
            A[i,j] = A[i,j] - m * A[k,j]  // eliminate

After this, $A$ is overwritten by $U$ in the upper triangle and $L$ (minus the unit diagonal) in the lower triangle.

Components

• $A \in \mathbb{R}^{n \times n}$: the original square matrix.
• $L \in \mathbb{R}^{n \times n}$: unit lower triangular. $L_{ii} = 1$, $L_{ij} = 0$ for $j > i$. The sub-diagonal entries are the elimination multipliers.
• $U \in \mathbb{R}^{n \times n}$: upper triangular. $U_{ij} = 0$ for $i > j$. This is the result of elimination.
• $P \in \mathbb{R}^{n \times n}$: permutation matrix (for partial pivoting). $P$ reorders rows so that the pivot is never zero.

Solving $Ax = b$ with $LU$

Once $A = LU$ (or $PA = LU$):

1. Forward substitution: solve $Ly = Pb$ for $y$. Since $L$ is lower triangular, you start from the first row and work down ($O(n^2)$).
2. Back substitution: solve $Ux = y$ for $x$. Since $U$ is upper triangular, you start from the last row and work up ($O(n^2)$).

The factorization costs $\frac{2}{3}n^3$ flops. Each solve costs $O(n^2)$. If you have many right-hand sides $b$, you pay the cubic cost once and the quadratic cost per solve.

Pivoting ($PA = LU$)

If $a_{11} = 0$ at any step, elimination divides by zero. More subtly, even if $a_{kk}$ is nonzero but tiny, dividing by it produces huge multipliers, amplifying rounding errors.

Partial pivoting fixes this: before eliminating column $k$, swap row $k$ with the row below it that has the largest entry in column $k$. Track the swaps in a permutation matrix $P$:

Why

• Direct, not iterative: LU gives the exact solution (up to floating-point precision) in a fixed number of operations. No convergence criteria, no iteration count to tune.
• Cheap per RHS: after the $O(n^3)$ factorization, each additional solve is $O(n^2)$. This is the main advantage of factorization over computing $A^{-1}$ or using iterative methods for many right-hand sides.
• Well-understood stability: with partial pivoting, the growth factor is bounded in practice. Pathological examples exist (the Wilkinson matrix), but they rarely appear in real applications.

Stability

Without pivoting, error can grow exponentially. The famous example:

The multipliers grow as $2^{k-1}$, and $U$ develops entries as large as $2^{n-1}$. For $n = 60$, this overflows double precision.

With partial pivoting ($PA = LU$), the multipliers are bounded by $1$, and the growth factor is typically $O(n)$. Backward error is proportional to $\text{cond}(A) \cdot \varepsilon$, where $\varepsilon$ is machine epsilon. For well-conditioned $A$, LU with pivoting is perfectly safe.

QR Decomposition

The Algorithm

$Q$ is orthogonal ($Q^\mathsf{T}Q = I$), and $R$ is upper triangular.

Thin (economy) QR: when $m > n$, we only need the first $n$ columns of $Q$:

Intuition: “QR writes $A$ as a rigid rotation/reflection ($Q$) times a triangular factor ($R$). Because $Q$ preserves lengths ($|Qx| = |x|$), multiplying by $Q$ never amplifies errors; this is what makes QR numerically stable.”

Components

• $A \in \mathbb{R}^{m \times n}$: the matrix to factorize. Can be square ($m = n$) or rectangular ($m > n$).
• $Q \in \mathbb{R}^{m \times m}$: orthogonal matrix. $Q^\mathsf{T}Q = QQ^\mathsf{T} = I$. Columns are orthonormal: each has unit length and is perpendicular to every other column.
• $R \in \mathbb{R}^{m \times n}$: upper triangular. $R_{ij} = 0$ for $i > j$. In thin QR, the bottom $m-n$ rows of $R$ are zero.
• Thin $Q \in \mathbb{R}^{m \times n}$: only the first $n$ columns, still orthonormal ($Q^\mathsf{T}Q = I_n$ but $QQ^\mathsf{T} \neq I_m$).

Approach 1: Gram-Schmidt

Classical Gram-Schmidt (CGS)

This is the intuitive approach: orthogonalize one column at a time.

Formal:

Intuition: “At step $k$, take column $a_k$ and subtract off its projection onto every previous $q_i$ in one shot. What’s left is orthogonal to all of them. Then normalize.”

The $R$ entries fall out naturally: $R_{ik} = q_i^\mathsf{T} a_k$ for $i < k$, and $R_{kk} = |\tilde{q}_k|$.

Why CGS is unstable: The subtraction $\tilde{q}_k = a_k - \Sigma(q_i^\mathsf{T} a_k) q_i$ suffers from catastrophic cancellation when $\tilde{q}_k$ is small: the projection terms nearly sum to $a_k$ itself. Since $a_k$ is the original column and the $q_i$ already have O(ε) errors from earlier steps, the cancellation magnifies those errors. The computed $q_k$ drifts from orthogonality; $Q^\mathsf{T}Q$ can have off-diagonals as large as $O(1)$.

Modified Gram-Schmidt (MGS)

Formal:

Intuition: “MGS orthogonalizes against the running residual $v_k$ at each inner step. Instead of computing all dot products from the original $a_k$ (which lets errors accumulate), each $q_i$ only sees what’s left after previous projections. The difference: CGS computes $q_i^\mathsf{T} a_k$ once and subtracts everything. MGS updates $v_k$ after each $q_i$, so later $q_j$ see a cleaner signal.”

MGS is backward stable for least squares: the computed $Q$ satisfies $Q^\mathsf{T}Q \approx I + O(\varepsilon)$, and $(Q, R)$ is the exact QR of $A + \Delta A$ with $|\Delta A| \leq O(\varepsilon)|A|$.

Cost: $2mn^2$ flops for both CGS and MGS. But MGS is BLAS-2 (matrix-vector), so it’s memory-bound on modern hardware.

Approach 2: Householder Reflections

This is the method used in production: LAPACK, MATLAB, SciPy, cuSOLVER all use it.

The Algorithm

A Householder reflector is a matrix that reflects a vector across a hyperplane:

$H$ is orthogonal and symmetric ($H^\mathsf{T} = H$, $H^\mathsf{T}H = I$).

Intuition: “$H$ reflects any vector across the hyperplane orthogonal to $v$. If we want to zero out everything below the diagonal in column $k$, we find $v$ such that $H \cdot x = |x| e_1$: a vector with zeros below the first entry. Geometrically: take the column $x$, find the mirror image across the plane that bisects $x$ and the first basis vector $e_1$, and reflect. The zeros fall out automatically.”

The Householder Vector

Given $x \in \mathbb{R}^{m-k+1}$ (the portion of column $k$ from row $k$ down), we want $v$ such that $Hx = |x| e_1$.

The formula:

where $e_1 = [1, 0, \dots, 0]^\mathsf{T}$.

Intuition: “We take $x$ and add $|x| e_1$ (pointing in the direction of the first axis, scaled by $x$’s length). The sum of these two vectors points along the normal of the reflection plane. The sign of $x_1$ ensures we pick the mirror that keeps $v$ away from zero: if $x_1$ is positive, we reflect using $|x| e_1$ (same direction); if negative, we use $-|x| e_1$ (opposite direction). This avoids $v \approx 0$, which would cause catastrophic cancellation.”

After computing $v$, normalize it so $v^\mathsf{T}v = 2$ (i.e. replace $v$ by $\frac{\sqrt{2} \, v}{|v|}$). This rescaling makes the scalar factor $\frac{2}{v^\mathsf{T}v} = 1$, collapsing the reflector formula to $H = I - v v^\mathsf{T}$ and saving a multiply per application. The reflector can then be applied to any matrix $A$ as:

which is a rank-1 update: cheap and stable.

QR via Householder

For $A \in \mathbb{R}^{m \times n}$:

for k = 1..min(m, n):
    x = A[k:m, k]
    v = x + sign(x₁) * ‖x‖ * e₁
    v = v / ‖v‖
    A[k:m, k:n] = A[k:m, k:n] - v * (vᵀ * A[k:m, k:n])

After the loop, $A$ has been overwritten by $R$ (upper triangular part). The $Q$ matrix is the product of all reflectors: $Q = H_1 H_2 \cdots H_n$, but it’s never formed explicitly; it’s stored as the list of $v$ vectors and applied on demand.

Cost: $2mn^2 - \frac{2}{3}n^3$ flops. About $2\times$ the cost of LU for square matrices.

Why QR is Numerically Stable

This is the central selling point of QR, and it comes down to one inequality:

$Q$ is orthogonal. Multiplying by $Q$ preserves the 2-norm: it’s a rigid rotation or reflection. Errors in $A$ or $b$ are never amplified when multiplied by $Q^\mathsf{T}$.

Compare solving $Ax = b$ using QR vs. the normal equations $A^\mathsf{T}Ax = A^\mathsf{T}b$:

Intuition: “Normal equations form $A^\mathsf{T}A$, which squares the condition number. If $\text{cond}(A) = 10^8$ (not uncommon for ill-conditioned problems), then $A^\mathsf{T}A$ has $\text{cond} = 10^{16}$. IEEE double precision gives exactly 16 digits. You lose all of them. QR never forms $A^\mathsf{T}A$. It operates on the original $A$ and uses orthogonal transformations that preserve norms. You lose only ~8 digits. This is not subtle: it’s the difference between a working solver and complete garbage.”

Backward stability: QR is backward stable. The computed solution $\hat{x}$ is the exact solution of a nearby problem:

where $c$ is a small constant and $\varepsilon$ is machine epsilon. Normal equations do not achieve this: the formation of $A^\mathsf{T}A$ introduces O($\varepsilon$) errors that are then amplified by $\text{cond}(A)^2$.

Approach 3: Givens Rotations (brief)

Givens rotations zero out a single entry at a time using a $2 \times 2$ rotation embedded in the identity. They’re more expensive per entry than Householder ($6$ flops per zero vs. $4$), but they allow selective zeroing (useful for sparse matrices or when you need to add a row to an already-factorized QR). Not covered in depth here.

Blocked Householder Reflections

The Problem with Standard Householder

Each Householder reflector $H_k$ modifies the entire trailing submatrix $A[k:m, k:n]$ as a rank-1 update:

This is a matrix-vector operation (BLAS-2): one dot product $v_k^\mathsf{T}A$ followed by an outer product update. On modern hardware, BLAS-2 is memory bandwidth bound: the CPU spends most of its time waiting for data from main memory. For large matrices, the unblocked Householder QR runs at $5$–$10\times$ below peak flop/s.

The WY Representation

Bischof and Van Loan (1987) showed that $k$ consecutive reflectors can be compacted into a compact form:

Intuition: “Instead of applying $k$ reflectors one at a time ($k$ passes over the matrix), we compact them into a single rank-$k$ update: $A := A + W(Y^\mathsf{T} A)$. The term $Y^\mathsf{T}A$ is $(k \times n)$; it fits in cache. Computing $W(Y^\mathsf{T}A)$ is a general matrix multiply (BLAS-3). Cache efficiency goes through the roof.”

Constructing $W$ and $Y$

Given reflectors $v_1, \dots, v_k$ normalized so $v_i^\mathsf{T}v_i = 2$:

After the loop: $W = W_k$, $Y = Y_k$, and $I + WY^\mathsf{T} = H_1 \cdots H_k$.

Mini example ($k = 2$)

Take two reflectors $H_1 = I - v_1 v_1^\mathsf{T}$ and $H_2 = I - v_2 v_2^\mathsf{T}$ with $v_1^\mathsf{T}v_1 = v_2^\mathsf{T}v_2 = 2$.

Compute:

Verify: $I + W_2 Y_2^\mathsf{T} = I + (-v_1)v_1^\mathsf{T} + (-v_2 + t v_1)v_2^\mathsf{T} = I - v_1 v_1^\mathsf{T} - v_2 v_2^\mathsf{T} + t v_1 v_2^\mathsf{T}$.

Meanwhile, $H_1 H_2 = (I - v_1 v_1^\mathsf{T})(I - v_2 v_2^\mathsf{T}) = I - v_1 v_1^\mathsf{T} - v_2 v_2^\mathsf{T} + (v_1^\mathsf{T} v_2) v_1 v_2^\mathsf{T}$.

Since $t = v_1^\mathsf{T} v_2$, these are identical. The compact representation is exact.

Applying the Block

Once $W$ and $Y$ are computed, the trailing matrix update is:

1. Compute $C = Y^\mathsf{T} A$: $(k \times m) \cdot (m \times n) = k \times n$, two small matrices with high arithmetic intensity.
2. Compute $A := A + W C$: $A$ plus a rank-$k$ update (another gemm).

Both steps are BLAS-3 (matrix-matrix). The factorization still computes reflectors one column at a time (that part isn’t blocked), but the application to the trailing matrix is deferred and batched into a single gemm call.

Why

• BLAS-3 performance: The blocked version runs at near-peak flop/s on modern CPUs and GPUs. Unblocked Householder is memory-bound and typically $5$–$10\times$ slower.
• Cache friendly: $Y^\mathsf{T}A$ operates on an $(m \times n)$ matrix but only produces a $(k \times n)$ result. The trailing matrix is streamed through cache once per block instead of once per reflector.
• Same math: $I + WY^\mathsf{T}$ is bitwise equivalent (modulo floating-point associativity) to applying the reflectors sequentially. $Q$ and $R$ are identical.
• Block size $k$: typically $32$–$64$, chosen so $W$ and $Y$ fit in L2 cache. LAPACK’s DGEQRF uses ILAENV to pick $k$ automatically based on the matrix size and cache geometry.

Worked Example

Let’s work through a concrete $3 \times 2$ matrix to see both LU and QR in action.

LU Decomposition ($PA = LU$)

Since $A$ is $3 \times 2$, we need a square system. Let’s work with the square $2 \times 2$ submatrix formed by the first two rows:

Step 1 (column 1): Pivot is $a_{11} = 3$. No swap needed. Multipliers: $m_{21} = 1/3$.

Step 2: The multipliers go into $L$:

Verify: $LU = \begin{bmatrix} 1 \cdot 3 + 0 \cdot 0 & 1 \cdot 1 + 0 \cdot \frac{5}{3} \ \frac{1}{3} \cdot 3 + 1 \cdot 0 & \frac{1}{3} \cdot 1 + 1 \cdot \frac{5}{3} \end{bmatrix} = \begin{bmatrix} 3 & 1 \ 1 & 2 \end{bmatrix} = A_{2 \times 2}$.

No pivoting needed since the pivot was already the largest in its column.

QR Decomposition via Gram-Schmidt

Step 1: $q_1 = a_1 / |a_1|$:

$R_{11} = |a_1| = \sqrt{14} \approx 3.7417$.

Step 2: Project $a_2$ onto $q_1$ and subtract:

$R_{22} = |\tilde{q}_2| = \frac{1}{14} \sqrt{(-19)^2 + 17^2 + 20^2} = \frac{\sqrt{1170}}{14} \approx 2.4431$.

Result:

Verify $Q^\mathsf{T}Q$:

Orthonormal to machine precision.

Verify $QR = A$:

Matches $A$ exactly (up to floating-point rounding).

QR Decomposition via Householder

Step 1: Take the first column $x = a_1 = [3, 1, 2]^\mathsf{T}$.

$x_1 = 3 > 0$, so $\operatorname{sign}(x_1) = +1$.

Normalize so $v^\mathsf{T}v = 2$: $|v| = \sqrt{6.7417^2 + 1^2 + 2^2} \approx 7.1362$, $v \leftarrow \frac{v}{|v|/\sqrt{2}} = \frac{\sqrt{2} v}{|v|}$.

This gives $v_1 \approx [0.9445, 0.1401, 0.2802]^\mathsf{T}$ with $v^\mathsf{T}v = 2$.

Apply $H_1 = I - v_1 v_1^\mathsf{T}$ to $A$:

$v_1^\mathsf{T} A = [3.7417, 3.1128]$. After the update:

The first column is now $[-3.7417, 0, 0]^\mathsf{T}$: zero below the diagonal, as promised. The rest of $R$ continues to fill in.

Step 2: Take the subcolumn $x = \tilde{a}_2[2:3]$ (rows 2–3 of the updated second column) and repeat. The final $R$ matches the Gram-Schmidt result (up to sign flips in rows of $R$, since Householder can introduce row sign changes that Gram-Schmidt doesn’t).

Quick Reference

Notes

• LU is the default for square $Ax = b$ when the matrix is well-conditioned (roughly $2\times$ cheaper than QR). For ill-conditioned systems ($\text{cond}(A) > 10^6$), QR is safer even for square matrices.
• For least squares $\min |Ax - b|_2$, QR is the standard choice. The normal equations $A^\mathsf{T}Ax = A^\mathsf{T}b$ should be avoided unless $\text{cond}(A)$ is tiny, because squaring the condition number loses digits rapidly.
• LAPACK’s DGEQRF uses blocked Householder with the WY representation. DGELS solves least squares via QR. MATLAB’s \ operator picks QR automatically for rectangular $A$.
• Modified Gram-Schmidt is backward stable but BLAS-2 (pedagogical but not used in production libraries). However, MGS is useful when you need to maintain orthogonality incrementally (e.g., in Krylov subspace methods like GMRES).
• The WY blocking trick generalizes beyond QR. The same idea (compact a panel’s transformations into a rank-$k$ update) is used in blocked LU and blocked Cholesky.
• Block size $k$ is typically $32$–$64$ in LAPACK. Larger $k$ gives better BLAS-3 performance but increases the cost of forming $W$ and $Y$. The optimal $k$ depends on the cache hierarchy and is determined by ILAENV.
• For very large matrices on distributed memory, communication-avoiding QR (CAQR, Demmel et al., 2008) reorganizes the panel factorization itself (not just the trailing matrix update) to reduce network traffic by a factor of $O(\sqrt{n})$.

Iterative/Recursive Formula

Components

• $x_n$: current approximation at step $n$. Initialized to $x_0$, the user’s best guess.

• $f(x_n)$: value of the target function at $x_n$. The method converges when $

  f(x_n)

  $ is small.

• $f’(x_n)$: derivative of $f$ at $x_n$. This is the slope of the tangent line. If $f’(x_n) = 0$, the method fails at this step (division by zero).
• $x_{n+1}$: next approximation. The $x$-intercept of the line through $(x_n, f(x_n))$ with slope $f’(x_n)$.

Derivation

The tangent line to $f$ at $x_n$ is

\(L(x) = f(x_n) + f'(x_n)(x - x_n)\) where $y = f(x_n)$ is the y-intercept and $f’(x_n)$ is the slope. We used the known point $(x_n, f(x_n))$ to write the line in point-slope form.

Set $L(x) = 0$ (find where the tangent crosses the $x$-axis)

this is because, the tangent line is a linear approximation to $f$ near $x_n$, so we hope its root is a better guess than $x_n$ itself. Solving for $x$ gives

now solve for $x$:

That $x$ becomes $x_{n+1}$. The same result falls out of a first-order Taylor expansion: $0 = f(x^) \approx f(x_n) + f’(x_n)(x^ - x_n)$ gives $x^* \approx x_n - f(x_n)/f’(x_n)$.

Convergence

Near a simple root ($f(x^) = 0$, $f’(x^) \neq 0$), Newton’s method converges quadratically:

This means the number of correct digits roughly doubles at each step, which is much faster than bisection or secant methods.

Why

• Quadratic convergence is the main draw. With a good initial guess, you reach machine precision in 4–6 iterations. Bisection, by contrast, gains one bit per iteration (linear convergence, $C = 0.5$). For engineering, you pay for the speed with the need for $f’(x_n)$.
• Self-correcting: a bad step is not catastrophic: the next iteration still uses the true $f$ and $f’$ at the new point, so the method can recover if the new point is closer.
• Generalizes naturally to systems of equations (multivariate Newton’s method with the Jacobian) and to optimization (Newton’s method for finding critical points, using the Hessian). This is why Newton-style updates appear in second-order optimizers.

Improved Variant: Halley’s Method

If the second derivative $f’‘(x_n)$ is available, we can do better than Newton. Instead of a linear tangent approximation, fit a local quadratic model using the first two terms of the Taylor expansion:

Solving this quadratic for $(x^* - x_n)$ gives Halley’s update:

Why

• Cubic convergence: near a simple root, the error shrinks as $

  x_{n+1} - x^*

  \le C

  x_n - x^*

  ^3$. Correct digits roughly triple each iteration instead of doubling. For the $\sqrt{2}$ example above, Halley’s method reaches $10^{-12}$ in 3 iterations instead of 4.

• More robust to curvature: the denominator has a built-in correction for sharp bends in $f$. Halley’s method converges on some problems where Newton oscillates or overshoots.
• Cost tradeoff: you need $f’‘(x_n)$ per step. For analytic functions this is one extra derivative, but for black-box functions the finite-difference cost doubles. In practice, Halley is used when $f$ is cheap and high precision is needed quickly (e.g. computing elementary functions (exp, log, trig) in math libraries).

The same idea extends to even higher derivatives (Householder methods), but Halley is the most practical second-order variant.

Failure Modes

Newton’s method is not robust out of the box. It fails when:

1. Zero derivative $f’(x_n) = 0$: the tangent is horizontal, never crosses the axis. The update divides by zero. This is addressed by halley’s method, which uses the second derivative to avoid this issue. In practice, a small damping factor or line search can also help prevent large jumps when $f’(x_n)$ is small but nonzero.
2. Poor initial guess: if $x_0$ is far from the root, the tangent may point in the wrong direction and diverge.
3. Oscillation: for some functions (e.g. $f(x) = x^3 - 2x + 2$ with $x_0 = 0$), the iterates cycle without converging.
4. Multiple roots: at a repeated root where $f’(x^*) = 0$, convergence drops to linear and the constant $C$ grows large.

Why

• No global guarantee: unlike bisection (which always converges given a sign change), Newton’s method makes no such promise. It is a local method, provably convergent only within a “basin of attraction” around the root.
• Derivative cost: you must compute $f’(x_n)$ at every step. For analytic functions this is cheap; for black-box simulations it may require finite differences, adding cost and truncation error.

Worked Example

Find $\sqrt{2}$ (i.e. solve $f(x) = x^2 - 2 = 0$).

Here $f’(x) = 2x$, and the Newton update is

Start with $x_0 = 2$:

After 4 iterations the error is at $10^{-12}$. Doubling the correct digits each step is the hallmark of quadratic convergence. For comparison, bisection on the interval $[1, 2]$ would need about 40 iterations to reach the same precision (one bit per iteration).

Quick Reference

Notes

• The update $x_{n+1} = x_n - f(x_n)/f’(x_n)$ is the same as applying Newton’s method in optimization to find $f’(x) = 0$, where the update becomes $x_{n+1} = x_n - f’(x_n)/f’‘(x_n)$. The root-finding and optimization versions are the same idea at different orders of differentiation.
• In practice, Newton’s method is often hybridized: bisection or bracketing ensures robustness for the first few steps, then Newton takes over for fast convergence near the root. This is what sqrt implementations in hardware and GLIBC do.
• For high-dimensional systems, the Jacobian replaces $f’(x_n)$ and a linear solve replaces division. The cost scales as $O(d^3)$ for dense systems, which is why quasi-Newton methods (BFGS, L-BFGS) are preferred in ML.
• The denominator $f’(x_n)$ can be zero in exact arithmetic but often just very small in floating point; the iteration can still produce a huge jump, overshooting the root. A damping factor or line search is the standard fix.