DEV Community: Lewis Won

PagedAttention by hand

Lewis Won — Sat, 22 Aug 2026 06:08:35 +0000

Table of Contents

Motivation
What is PagedAttention
Setup
Forward pass (PagedAttention)
Why divide memory into physical blocks?
Code

Motivation

PagedAttention (the core algorithm behind the vLLM inference engine) is widely used to serve Large Language Models (LLMs) with high throughput. While libraries like vLLM provide out-of-the-box support for it, understanding how PagedAttention is implemented by hand may help with:

Building intuition: understand how it maps logical token sequences to non-contiguous physical memory blocks, mimicking operating system virtual memory.
Understanding the memory-capacity bottleneck: understand why the number of sequences you can batch is set by how much KV cache fits in GPU memory, not by how fast the GPU computes, and why fragmentation limits it.
Optimizing deployments: the trade-off between internal fragmentation and bandwidth in choosing the optimal block_size.

This article can serve as a complementary study guide for understanding the vLLM paper: Efficient Memory Management for Large Language Model Serving with PagedAttention.

This article is written with the assistance of AI.

What is PagedAttention

When an LLM generates a response, it produces tokens one by one. To avoid recomputing the attention states for all past tokens at every step, the model caches the Key and Value vectors for all previously generated tokens. This is called the KV cache. I have a worked example of a KV cache with a single-layer model of embedding dimension d = 2 here as reference.

However, the KV cache has unique characteristics that make it a nightmare for contiguous memory allocators:

It grows dynamically over time.
Its final length is unknown ahead of time.

For systems which pre-allocate contiguous chunks of memory for the maximum possible sequence length, this will lead to three distinct wastes:

Waste type	Description	Does PageAttention eliminate it
Reservation	Slots held for tokens not yet generated (eventually used, but unavailable to other sequences in the meantime)	Eliminated. Nothing is reserved for the max length; blocks are allocated as tokens arrive.
External fragmentation	From the allocator itself	Eliminated. All blocks are the same size, so every free block fits every need.
Internal fragmentation	From over-provisioning to the maximum length when the sequence finishes early	Reduced, not removed. Bounded to ≤ B−1 slots in the last partial block — ~15 slots per sequence at B=16, regardless of length.

Internal vs external fragmentation

Internal fragmentation is waste inside an allocation owned by a sequence, e.g. the sequence asked for a 2048-slot chunk and only used 60. The other 1988 slots are sitting idle and no other sequences can use them.

External fragmentation is waste between allocations. The memory is free, but it is in pieces of the wrong size or in the wrong places, e.g. you can have 10 GB free in total and still fail a 2 GB sequence because no single contiguous run is 2 GB long.

Using the parking lot as an analogy: internal fragmentation is being assigned a bus bay and parking a motorbike in it. External fragmentation is a lot where cars parked at random leave plenty of empty tarmac but no single gap wide enough for a bus.

How PagedAttention solves reservation and fragmentations

PagedAttention solves this by borrowing the concept of virtual memory with paging from operating systems:

Partitions Memory: It divides the KV cache into fixed-size "blocks" (pages), where each block contains the keys and values for a fixed number of tokens.
Non-contiguous Storage: These blocks do not need to be stored next to each other in physical GPU memory.
Block Tables: It uses a "Block Table" (like a page table) to map the continuous logical sequence of tokens to their scattered physical blocks.

Uniform blocks solve external fragmentation because with only one block size, there are no holes of the wrong size. The free list is a stack of interchangeable IDs, and any free block satisfies any sequence.

Why not set B = 1 to solve internal fragmentation?

The paper tests this (§7.2, Fig. 18b of paper) and finds it among the worst-performing settings. The reason is that a block is not only an allocation unit, but also the unit of a coalesced memory read. vLLM assigns one GPU warp per block (§5.1 of paper), so at B=16 those 32 threads fetch 16 tokens of adjacent key vectors in a single wide transaction. At B=1 you get one narrow scattered read per token, and since decoding is memory-bandwidth-bound, that costs far more than the ≤15 slots saved. Block size is a trade between fragmentation and bandwidth.

Estimating size of a block byte

Symbol	What it is	Value here
B	block size — how many tokens fit in one block	2 in this article
d	head dimension — length of one token's key vector	2 in this article
warp	fixed GPU thread group — hardware constant	32
—	bytes per number	2 for FP16

block bytes = B tokens × d numbers × 2 bytes
            = 16 × 128 × 2
            = 4,096 bytes

At B=1 that is 128 × 2 = 256 bytes. 128 is the length of a single token's key vector, and it stays 128 regardless of what B is. Changing B changes how many of those 256-byte vectors sit side by side.

Setup

To keep the scope of the article manageable, this article discusses PagedAttention in the context of autoregressive generation phase of LLM inference (serving).

We will walk through the attention computation for a single attention head generating a new token. Let's assume our LLM has a tiny hidden dimension of d = 2.

Our sequence currently has 5 tokens, and we are generating the 6th token.
Instead of a contiguous tensor, our KV cache is managed by PagedAttention with a Block Size (B) of 2 tokens per block.

Since we have 5 tokens, we require $⌈ 5/2 ⌉ = 3$ logical blocks.

Step 1: Initial Setup - Defining the Tensors and Memory

Our sequence currently has 5 tokens. We are running the decode step that will predict the 6th. The input to this step is token 5, so the query vector q is the query of token 5. Note that token 5's own key and value are already in the cache (vLLM writes them before calling the attention kernel):

q = (12)

We have a pre-allocated pool of 9 physical blocks in GPU memory, shared by every sequence on this worker. The system assigned our sequence to Physical Blocks 2, 0, and 1 (in that exact order) — the first three that happened to be free. Each of the other 8 blocks are either free, or already belong to another sequence.

This mapping is stored in our Block Table:
Block Table = [2, 0, 1]

Here is what is currently stored in those physical KV cache blocks (Keys and Values for past tokens):

Logical Block 0 (Tokens 1 & 2) -> Mapped to Physical Block 2:
$K_{p h y s_{2}} = (1001)$ , $V_{p h y s_{2}} = (1212)$

Logical Block 1 (Tokens 3 & 4) -> Mapped to Physical Block 0:
$K_{p h y s_{0}} = (1210)$ , $V_{p h y s_{0}} = (3434)$

Logical Block 2 (Token 5 & Empty) -> Mapped to Physical Block 1:
Notice that the second slot in this block is currently empty because we have an odd number of tokens (5).
$K_{p h y s_{1}} = (0020)$ , $V_{p h y s_{1}} = (5050)$

Note that each token contributes two vectors — one key and one value. A block with B = 2 therefore holds two rows in its K matrix and two rows in its V matrix, with row i of each belonging to the same token.

Forward pass (Paged Attention)

Standard attention calculates $S core = so f t ma x (q K^{T}) V$ .
Because our K and V matrices are scattered across physical memory, PagedAttention processes them block by block. (Note: For simplicity in mental math, we will omit the standard $\frac{1}{d}$ scaling factor).

Step 2: Compute Unnormalized Attention Scores ( $q K^{T}$ )

We iterate through the Block Table [2, 0, 1].

1. First Logical Block (Physical Block 2)

S cor e_{0} = q K_{p h y s_{2} T} = (12) (1001)^{T}

Score_0[0] = (1*1) + (2*0) = 1
Score_0[1] = (1*0) + (2*1) = 2

Result: [1, 2]

2. Second Logical Block (Physical Block 0)

S cor e_{1} = q K_{p h y s_{0} T} = (12) (1210)^{T}

Score_1[0] = (1*1) + (2*1) = 3
Score_1[1] = (1*2) + (2*0) = 2

Result: [3, 2]

3. Third Logical Block (Physical Block 1)

S cor e_{2} = q K_{p h y s_{1} T} = (12) (0020)^{T}

Score_2[0] = (1*0) + (2*2) = 4
Score_2[1] = (1*0) + (2*0) = 0

Result: [4, 0]

We only have 5 valid tokens. The 6th slot in Physical Block 1 is uninitialised memory, so we must stop the model attending to it. We replace the invalid score with negative infinity (-inf), which sends its softmax weight to exactly zero.

We do not leave the score at 0, because doing so means the slot absorbs about 1% of the probability mass, diluting every other token's weight and shifting the output by roughly the same amount. In a real server that slot holds whatever a previously evicted sequence left behind, so the contribution is arbitrary rather than merely small.

In production, instead of masking, we can bound the so that the CUDA kernel never reads those addresses at all and invalid slots cost zero memory traffic. Since decoding is memory-bandwidth-bound, not reading can save bandwidth on unnecessary reads.

Step 3: Apply Softmax

We concatenate our block scores into a single logical sequence of scores:
Scores = [1, 2, 3, 2, 4, -inf]

To calculate the softmax, PyTorch subtracts the maximum value (which is 4) for numerical stability, then takes the exponent, and divides by the sum of exponents.

Scores - 4 = [-3, -2, -1, -2, 0, -inf]
Exp = [0.0498, 0.1353, 0.3679, 0.1353, 1.0000, 0]
Sum of Exp = 1.6883

Dividing by the sum gives us our attention probabilities:
Probs ≈ [0.0295, 0.0802, 0.2179, 0.0802, 0.5923, 0]

Step 4: Multiply by Values ( $P ro b s \times V$ )

We slice the probabilities back into their respective blocks and multiply them by the physical Value blocks.

1. First Logical Block (Physical Block 2)

Probs_0 = [0.0295, 0.0802]

O u t_{0} = 0.0295 \times (11) + 0.0802 \times (22) = (0.0295 0.0295) + (0.1604 0.1604) = (0.1899 0.1899)

2. Second Logical Block (Physical Block 0)

Probs_1 = [0.2179, 0.0802]

O u t_{1} = 0.2179 \times (33) + 0.0802 \times (44) = (0.6537 0.6537) + (0.3208 0.3208) = (0.9745 0.9745)

3. Third Logical Block (Physical Block 1)

Probs_2 = [0.5923, 0]

O u t_{2} = 0.5923 \times (55) + 0 \times (00) = (2.9615 2.9615)

4. Final Aggregation

We sum the outputs from all blocks to get the final attention output for the new token:

FinalOutput = O u t_{0} + O u t_{1} + O u t_{2} = (0.1899 0.1899) + (0.9745 0.9745) + (2.9615 2.9615) = (4.1259 4.1259)

This tensor is then passed into the rest of the Transformer block (Linear layers, etc.) to predict the 6th token. At the next iteration, token 6 becomes the model's input; its key and value are computed and written into the empty slot in Physical Block 1 before attention runs, and the cycle repeats.

Why divide memory into physical blocks?

Since the numerical output is exactly the same as standard Attention, why go through the headache of block tables and scattered memory?

1. Zero External Fragmentation & Minimal Internal Fragmentation

With contiguous memory allocation, if an LLM supports a max sequence of 2048 tokens, the system pre-allocates contiguous memory for 2048 tokens the moment a sequence arrives. If the user prompt is only 50 tokens, 1998 tokens are wasted to reserved fragmentation. If the model responds with 10 tokens, now 1988 slots of memory are wasted to internal fragmentation.

With PagedAttention, memory is allocated on demand. In our example, we only allocated exactly 3 blocks (6 slots) for a 5-token sequence. The only "wasted" space is the single empty slot in the last block (internal fragmentation).

§7.2 of the paper default block size is 16, which is large enough to keep the GPU's read parallelism busy and small enough to avoid meaningful internal fragmentation. That default guarantees that wasted space is bounded by at most B − 1 slots per sequence, or 15 slots in this case, no matter whether the sequence is 20 tokens or 20,000.

2. Memory Sharing (Copy-on-Write)

Imagine we ask for two samples (n=2), from the prompt prompt "The capital of France is":

Candidate A attempts to generate: "Paris, which is..."
Candidate B attempts to generate: "a city that..."

With contiguous memory, the KV cache for the prompt would have to be duplicated for each candidate.

With PagedAttention, both candidates share the same logical-to-physical block mapping for the prompt. The block manager tracks a reference count for each physical block. When a candidate needs to write a token into a block whose reference count is greater than 1, it performs an OS-style copy-on-write: allocate a fresh physical block, copy that one block's contents, decrement the original's reference count, and write into the copy.

Note: A reference count is a tally on the number of sequences currently using a physical block.

Example of Copy-on-Write

Prompt: 5 tokens. Logical blocks: [t1,t2] [t3,t4] [t5,_]. Both candidates map to physical blocks 2, 0, 1 with refcount 2 on each.

Now both generate token 6, and they disagree.

	Candidate A ("Paris")	Candidate B ("a")
Wants to write	slot 1 of logical block 2	slot 1 of logical block 2
Physical target	block 1, refcount 2	block 1, refcount 2
Action	refcount > 1 → allocate physical block 4, copy [k5,v5] into it, drop block 1's refcount to 1, write k6/v6 into slot 1 of block 4	refcount now 1 → write k6/v6 directly into block 1

Afterwards: blocks 2 and 0 (holding t1–t4) are still shared, reference count is 2. Only one block was duplicated.

Compare the alternatives. Never copy: physically impossible, the numbers differ. Always copy: that's the old system, i.e. duplicate all 3 blocks per candidate, which is the waste vLLM exists to remove. copy-on-write gives you 2 blocks shared, 1 copied.

Why contigous KV cache has difficulty implementing prefix matching

If a sequence's KV cache is a contiguous tensor, the attention kernel takes a base pointer and a length, then strides through memory linearly. Under that constraint, candidate A's cache mus look like:

[prompt tokens 1-5][A's tokens 6,7,8...]

all in one unbroken run. Candidate B needs the same thing with its own tokens appended. If both point at one shared prompt region, then B's token 6 has to occupy the address immediately after the prompt, which is where A's token 6 already lives. Two different values, one address, which makes prefix matching. There are alternatives to PagedAttention, such as RadixAttention, that can also enable prefix matching.

Code

For those who prefer following the code, the manual walk-through of PagedAttention above is replicated below using PyTorch. In production, vLLM uses custom CUDA kernels to execute these block reads and matrix multiplications faster.

import torch

# --------------------------------------------------------------------------
# 1. SETUP: Define tensors and physical memory pool
# --------------------------------------------------------------------------

BLOCK_SIZE = 2   # tokens per block. vLLM's default is 16; try changing it.
HEAD_DIM = 2
NUM_PHYSICAL_BLOCKS = 9

# New Query token: shape (1, 2)
q = torch.tensor([[1.0, 2.0]])

# We simulate our GPU physical memory pool. Our sequence will only use 3 of
# these 9 blocks -- the rest are free, or belong to other sequences.
# Shape: (total_physical_blocks, block_size, head_dim) = (9, 2, 2)
#
# Note: vLLM's real cache uses a more intricate layout, reshaping the key
# cache to interleave the head dimension so each block read coalesces across
# a GPU warp (§5.1, "fused reshape and block write"). We flatten it here for
# clarity -- the algorithm is identical, only the memory layout differs.
K_cache = torch.zeros(NUM_PHYSICAL_BLOCKS, BLOCK_SIZE, HEAD_DIM)
V_cache = torch.zeros(NUM_PHYSICAL_BLOCKS, BLOCK_SIZE, HEAD_DIM)

# Fill Physical Block 2 (Logical Block 0: tokens 1, 2)
K_cache[2] = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
V_cache[2] = torch.tensor([[1.0, 1.0], [2.0, 2.0]])

# Fill Physical Block 0 (Logical Block 1: tokens 3, 4)
K_cache[0] = torch.tensor([[1.0, 1.0], [2.0, 0.0]])
V_cache[0] = torch.tensor([[3.0, 3.0], [4.0, 4.0]])

# Fill Physical Block 1 (Logical Block 2: token 5, empty slot)
K_cache[1] = torch.tensor([[0.0, 2.0], [0.0, 0.0]])
V_cache[1] = torch.tensor([[5.0, 5.0], [0.0, 0.0]])

# Blocks 3-8 remain free for other sequences. Notice our sequence's blocks are
# neither contiguous nor in order -- that is the entire point.
block_table = [2, 0, 1]
context_length = 5   # Number of valid tokens we currently have

print("--- PagedAttention Forward Pass ---")

# --------------------------------------------------------------------------
# 2. Compute Unnormalized Scores
# --------------------------------------------------------------------------
scores = []
for i, physical_block_idx in enumerate(block_table):
    # Retrieve the physical keys for this block
    K_block = K_cache[physical_block_idx]      # shape (BLOCK_SIZE, HEAD_DIM)

    # Calculate q * K^T
    score = torch.matmul(q, K_block.T)         # shape (1, BLOCK_SIZE)

    # Mask out unfilled slots in the last block
    if i == len(block_table) - 1:
        valid_in_last_block = context_length % BLOCK_SIZE
        if valid_in_last_block != 0:
            score[0, valid_in_last_block:] = float('-inf')

    scores.append(score)

# Concatenate block scores into a single sequence
scores = torch.cat(scores, dim=-1)
print(f"Unnormalized Scores: {scores}")

# --------------------------------------------------------------------------
# 3. Softmax
# --------------------------------------------------------------------------
probs = torch.softmax(scores, dim=-1)
print(f"Attention Probabilities: {probs}")

# --------------------------------------------------------------------------
# 4. Multiply by Values
# --------------------------------------------------------------------------
final_output = torch.zeros(1, HEAD_DIM)
start_idx = 0

for i, physical_block_idx in enumerate(block_table):
    V_block = V_cache[physical_block_idx]

    # Slice the probabilities corresponding to this block
    p_block = probs[0, start_idx:start_idx + BLOCK_SIZE].unsqueeze(0)

    # Multiply and accumulate
    final_output += torch.matmul(p_block, V_block)

    start_idx += BLOCK_SIZE

print(f"\nFinal Attention Output: {final_output}")

Output

--- PagedAttention Forward Pass ---
Unnormalized Scores: tensor([[1., 2., 3., 2., 4., -inf]])
Attention Probabilities: tensor([[0.0295, 0.0802, 0.2179, 0.0802, 0.5923, 0.0000]])

Final Attention Output: tensor([[4.1256, 4.1256]])

Note: our hand-calculation gives 4.1259 because we rounded probabilities to four decimals at each step. The code, working in full precision, returns 4.1256.

KV Cache by hand

Lewis Won — Fri, 21 Aug 2026 14:47:59 +0000

Table of Contents

Motivation
What is the KV Cache?
Setup
Scenario 1: Generation WITHOUT KV Cache
Scenario 2: Generation WITH KV Cache
The Compute vs. Memory Trade-off
Code
Appendix A: Worked example with 2 layers

Motivation

The Key-Value (KV) cache is an important optimization in Large Language Model (LLM) inference. Working through the math helps with:

Building intuition: Understand why autoregressive generation (predicting the next word) contains massive amounts of redundant math without caching.
Understanding the compute-memory trade-off: The cache turns each step's projections from a matrix-matrix multiply into a matrix-vector multiply. That's far fewer FLOPs — but it also means you're now loading far more memory per unit of compute, which is exactly what makes decoding memory-bandwidth-bound. And the cache itself competes with the weights for VRAM.
Debugging Performance: When estimating serving costs, knowing the size and behavior of the KV cache is essential for calculating maximum batch sizes and throughput.

This article is written with the assistance of AI.

What is the KV Cache?

LLMs generate text autoregressively, meaning they predict the sequence one token at a time. To predict token N, the model needs to look at all previous tokens from 1 to N-1.

In the Transformer attention mechanism, every token is projected into a Query (Q), Key (K), and Value (V) vector.

The Query represents what the current token is "looking for".
The Key represents what a token "contains".
The Value represents the actual "content" the token will contribute to the output.

When generating a new token, the model calculates a new Query, Key, and Value for that specific token. However, its new Query needs to attend to the Keys and Values of all past tokens.

If we don't use a KV cache, we have to recalculate the Keys and Values for every single previous token at every single generation step. The KV cache simply stores the Keys and Values of previous tokens in memory so we only ever have to compute the projection for the newest token.

Quick question: why do we not cache Q? Ans: because of causal masking, a token's query is used exactly once, at the step that token is processed, and is never needed again. Keys and values, by contrast, are read at every subsequent step.

Setup

We will walk through the attention calculation for the decode step that predicts the 4th token.

Assume our embedding dimension is $d = 2$ . We have three weight matrices for our attention head: $W_{q}$ , $W_{k}$ , $W_{v}$ .

Let our weight matrices be:

W_{q} = (1001), W_{k} = (1011), W_{v} = (2002)

Assume our sequence currently consists of 2 tokens, $x_{1}$ and $x_{2}$ , and we just generated the newest token $x_{3}$ :

$x_{1} = [1, 0]$
$x_{2} = [0, 1]$
$x_{3} = [1, 1]$ (The new token)

We now need to perform the attention calculation for $x_{3}$ to predict the next token, $x_{4}$

Scenario 1: Generation WITHOUT KV Cache

Without a KV cache, the model has amnesia. It only knows that the current sequence is three tokens long. To calculate attention, it must process the entire sequence from scratch.

Let our input matrix X be the stack of all three tokens:

X = ⎝ ⎛ 101011 ⎠ ⎞

Step 1: Calculate Q, K, and V for the whole sequence

We must multiply the entire input sequence by our weight matrices.

Calculate Q:

Q = X \times W_{q} = ⎝ ⎛ 101011 ⎠ ⎞ \times (1001) = ⎝ ⎛ 101011 ⎠ ⎞

Calculate K:

K = X \times W_{k} = ⎝ ⎛ 101011 ⎠ ⎞ \times (1011) = ⎝ ⎛ 101112 ⎠ ⎞

Calculate V:

V = X \times W_{v} = ⎝ ⎛ 101011 ⎠ ⎞ \times (2002) = ⎝ ⎛ 202022 ⎠ ⎞

Notice what just happened. The first two rows of K and V are exactly the same as they were during the previous generation step. We wasted valuable GPU FLOPs recalculating them.

Step 2: Compute Attention for the new token

To predict token 4, we only care about the attention output of token 3. We take the 3rd row of Q ( $q_{3}$ ) and multiply it by the transposed K matrix.

(Two caveats about that "only." First, we can drop rows 1 and 2 of Q here because attention is causal — each token attends only to itself and the tokens before it, so position 3's output is unaffected by anything at positions 4 onward, and positions 1 and 2 produce outputs we already have. In the prompt phase, where all positions are computed at once, a causal mask sets the scores for future positions to −∞ to enforce this. Second, we can discard those rows only because our toy example has a single layer. In a real model, layer 2's attention at position 3 needs layer 1's output at every position — so without a cache you must recompute all N positions at all L layers, and only then throw away everything but the last row. The waste is far worse than this example makes it look. This is worked through in Appendix A.)

q_{3} = (11)

Unnormalized Attention Scores:

S cores = q_{3} \times K^{T} = (11) \times (110112)

S cores = ((1 * 1 + 1 * 1) (1 * 0 + 1 * 1) (1 * 1 + 1 * 2)) = (213)

Softmax (approximate for readability):

e^{2} \approx 7.39

e^{1} \approx 2.72

e^{3} \approx 20.09

sum \approx 30.2

P ro b s = (0.24 0.09 0.67)

For readability we omit the standard 1/√d scaling factor from the softmax. Real attention computes softmax(q·Kᵀ/√d), where d is the head dimension; the scaling keeps the dot products from growing large enough to saturate the softmax, but it changes nothing about the caching argument.

Multiply by V:

O u tp u t = P ro b s \times V = (0.24 0.09 0.67) \times ⎝ ⎛ 202022 ⎠ ⎞

O u tp u t = ((0.48 + 0 + 1.34) (0 + 0.18 + 1.34)) = (1.82 1.52)

We arrived at our output, but computing $X \times W_{k}$ and $X \times W_{v}$ scaled linearly with the length of our sequence. If our context was 10,000 tokens long, we would have done matrix multiplication for 10,000 tokens just to generate one new one. Check out Appendix A for a deeper dive into the calculations needed when there are two layers in the model.

Scenario 2: Generation WITH KV Cache

Now, let's assume we are caching our states. During the previous generation step (when we predicted $x_{3}$ ), we saved the Keys and Values for $x_{1}$ and $x_{2}$ in GPU memory.

KV Cache in Memory:

K_{p a s t} = (1011), V_{p a s t} = (2002)

When token $x_{3}$ arrives, we do not feed the whole sequence into the weight matrices. We only pass the brand-new token.

x_{n e w} = (11)

Step 1: Calculate Q, K, and V for ONLY the new token

This is a massive compute saving. We are now doing a simple Vector-Matrix multiplication instead of a Matrix-Matrix multiplication.

q_{n e w} = x_{n e w} \times W_{q} = (11) \times (1001) = (11)

k_{n e w} = x_{n e w} \times W_{k} = (11) \times (1011) = (12)

v_{n e w} = x_{n e w} \times W_{v} = (11) \times (2002) = (22)

Step 2: Update the Cache

We append our newly calculated $k_{n e w}$ and $v_{n e w}$ to our existing cache in memory.

K = concat (K_{p a s t}, k_{n e w}) = ⎝ ⎛ 101112 ⎠ ⎞

V = concat (V_{p a s t}, v_{n e w}) = ⎝ ⎛ 202022 ⎠ ⎞

Step 3: Compute Attention

From here, the math is exactly the same as Step 2 in the previous scenario. We multiply our $q_{n e w}$ by the updated $K$ cache, apply softmax, and multiply by the updated $V$ cache.

S cores = q_{n e w} \times K^{T} = (11) \times (110112) = (213)

The final output is identical: $(1.82 1.52)$ .

The Compute vs. Memory Trade-off

By using the KV cache, we completely eliminated the redundant calculations of previous tokens.

Without Cache: Generating the N-th token takes O(N²). The model recomputes the whole sequence from scratch — and in a multi-layer model that means running attention at every position, not just the last one: N queries against N keys.
With Cache: Generating the N-th token takes O(N). Only the new token is projected, and its single query attends to the N−1 cached keys and values.

The Price of the KV Cache
This massive speedup comes at a steep cost: GPU Memory (VRAM). As a sequence gets longer, the $K$ and $V$ matrices grow linearly.

The prompt length is known when the request arrives, and the maximum length is bounded by the context window — but the output length is unknowable until generation terminates. A system must therefore either over-reserve for the worst case (wasting memory) or allocate incrementally (which is what PagedAttention enables).

This exact problem—managing the massive, dynamic memory footprint of the KV cache—is what paved the way for memory management innovations like PagedAttention, which stores this cache in scattered, fixed-size pages. I will discuss PagedAttention in my next article.

There are also model-side innovations to reduce demand on KV cache, such as grouped-query attention (GQA) implemented in Llama models, and multi-head latent attention (MLA) implemented by Deepseek.

Code

Here is the PyTorch implementation of both scenarios. Notice how the three projection matmuls stay fixed at (1, 2) × (2, 2) no matter how long the sequence gets. The attention matmuls, q_new @ K_cache.T and probs @ V_cache, still grow with the cache length. This is the cost the KV cache cannot remove.

q_new @ K_cache.T is (1, 2) x (2, N-1)
probs @ V_cache is (1, N-1) x (T, 2)

import torch
import torch.nn.functional as F

# --------------------------------------------------------------------------
# 1. SETUP
# --------------------------------------------------------------------------
# Weight matrices (2x2)
W_q = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
W_k = torch.tensor([[1.0, 1.0], [0.0, 1.0]])
W_v = torch.tensor([[2.0, 0.0], [0.0, 2.0]])

# Past tokens and the new token
x1 = torch.tensor([[1.0, 0.0]])
x2 = torch.tensor([[0.0, 1.0]])
x3 = torch.tensor([[1.0, 1.0]]) # The new token

print("=== SCENARIO 1: WITHOUT KV CACHE ===")
# Stack all tokens into a single input matrix X (shape: 3x2)
X = torch.cat([x1, x2, x3], dim=0)
print(f"Input X shape: {X.shape}")

# Matrix multiply the ENTIRE sequence (Wasted compute!)
Q_full = X @ W_q
K_full = X @ W_k
V_full = X @ W_v
print(f"K_full shape: {K_full.shape} (Computed from scratch)")

# We only want the attention for the latest token (q3)
q3 = Q_full[-1:] # Shape 1x2

# Attention calculation
scores_no_cache = q3 @ K_full.T
probs_no_cache = F.softmax(scores_no_cache, dim=-1)
output_no_cache = probs_no_cache @ V_full

print(f"Output without cache:\n{output_no_cache}\n")


print("=== SCENARIO 2: WITH KV CACHE ===")
# Assume we have K and V from the previous step saved in memory
# In a real system this would already be in VRAM from the previous
# step
K_past = torch.cat([x1, x2], dim=0) @ W_k 
V_past = torch.cat([x1, x2], dim=0) @ W_v
print(f"K_past shape in memory: {K_past.shape}")

# We ONLY project the new token (Massive compute savings!)
q_new = x3 @ W_q
k_new = x3 @ W_k
v_new = x3 @ W_v
print(f"k_new shape: {k_new.shape} (Only computed for 1 token)")

# Update the cache
K_cache = torch.cat([K_past, k_new], dim=0)
V_cache = torch.cat([V_past, v_new], dim=0)

# Attention calculation
scores_cache = q_new @ K_cache.T
probs_cache = F.softmax(scores_cache, dim=-1)
output_cache = probs_cache @ V_cache

print(f"Output with cache:\n{output_cache}\n")

# Verify they are mathematically identical
assert torch.allclose(output_no_cache, output_cache)
print("SUCCESS: Both methods are mathematically identical!")

Output:

=== SCENARIO 1: WITHOUT KV CACHE ===
Input X shape: torch.Size([3, 2])
K_full shape: torch.Size([3, 2]) (Computed from scratch)
Output without cache:
tensor([[1.8199, 1.5105]])

=== SCENARIO 2: WITH KV CACHE ===
K_past shape in memory: torch.Size([2, 2])
k_new shape: torch.Size([1, 2]) (Only computed for 1 token)
Output with cache:
tensor([[1.8199, 1.5105]])

SUCCESS: Both methods are mathematically identical!

Appendix A: Worked example with 2 layers

What happens if our model has two layers instead of just one?

The Setup:
We have 2 tokens ( $x_{1}, x_{2}$ ). We just received the new token $x_{3}$ .
Our model has two layers: L1 and L2. Each has its own weights (e.g., $W_{q (L 1)}$ for Layer 1, $W_{q (L 2)}$ for Layer 2).

Step 1: The "Shortcut" Attempt at Layer 1

We want to save compute, so we decide to only calculate the Layer 1 attention output for our newest token, $x_{3}$ .

We project the raw embeddings into $q_{3 (L 1)}$ .
Because we have no cache, we recalculate the keys and values for all tokens from scratch: $k_{1 (L 1)}, v_{1 (L 1)}, k_{2 (L 1)}, v_{2 (L 1)}, k_{3 (L 1)}, v_{3 (L 1)}$ .
We calculate the attention for token 3 and get our result: $O u tp u t_{3 (L 1)}$ .

We take our single vector, $O u tp u t_{3 (L 1)}$ , and pass it up to Layer 2.

Step 2: Hitting the Wall at Layer 2

Layer 2 receives $O u tp u t_{3 (L 1)}$ and begins its work.

It calculates the query for the new token: $q_{3 (L 2)} = O u tp u t_{3 (L 1)} \times W_{q (L 2)}$
Now, Layer 2 needs to perform attention. To do this, it needs the Keys and Values for the past context. Let us just focus on getting the keys for now, i.e. $k_{1 (L 2)}, k_{2 (L 2)}$
It tries to calculate $k_{1 (L 2)}$

The formula for a key at Layer 2 is to multiply the output of Layer 1 by Layer 2's Key weights:

k_{1 (L 2)} = O u tp u t_{1 (L 1)} \times W_{k (L 2)}

CRITICAL ERROR: We don't have $O u tp u t_{1 (L 1)}$ .

In Step 1, we took a shortcut and only calculated the Layer 1 output for token 3. We completely ignored the Layer 1 outputs for tokens 1 and 2. Because Layer 2's Keys and Values are derived from Layer 1's outputs, Layer 2 is now paralyzed.

Step 3: The Cascading Compute Penalty

To fix this, we are forced to go back to Layer 1 and do the full attention calculation for every single past token just to get their $O u tp u t^{(L 1)}$ states.

Calculate for Token 1:

$q_{1 (L 1)} = x_{1} \times W_{q (L 1)}$ $O u tp u t_{1 (L 1)} = softmax (q_{1 (L 1)} \times [k_{1 (L 1)}]^{T}) \times [v_{1 (L 1)}]$
Calculate for Token 2:

$q_{2 (L 1)} = x_{2} \times W_{q (L 1)}$ $O u tp u t_{2 (L 1)} = softmax (q_{2 (L 1)} \times [k_{1 (L 1)}, k_{2 (L 1)}]^{T}) \times [v_{1 (L 1)}, v_{2 (L 1)}]$

Step 4: Finally Executing Layer 2

Now that Layer 1 has re-processed the entire sequence, we pass the full stack of outputs ( $O u tp u t_{1 (L 1)}, O u tp u t_{2 (L 1)}, O u tp u t_{3 (L 1)}$ ) to Layer 2.

Layer 2 can finally calculate the past Keys and Values: $k_{1 (L 2)} = O u tp u t_{1 (L 1)} \times W_{k (L 2)}$ $k_{2 (L 2)} = O u tp u t_{2 (L 1)} \times W_{k (L 2)}$
Layer 2 calculates attention for the 3rd token and outputs the final result.

The Takeaway:

Without a KV cache, you cannot just pass the newest token through the network. Layer 1 must compute attention for ALL tokens, to pass ALL outputs to Layer 2, which computes attention for ALL tokens, to pass ALL outputs to Layer 3...

If your model has 32 layers, this $O (N^{2})$ recomputation penalty happens 32 times per generated word, which is inefficient at long contexts.

FP8, FP4, MXFP4 and NVFP4 by hand

Lewis Won — Sun, 16 Aug 2026 16:09:54 +0000

Table of Contents

Why by hand?
The cast of characters
Everything is a grid and a scale
The E2M1 grid, enumerated by hand
Our running example tensor
Format 1: FP8 (E4M3) by hand
Aside: FP8 is two formats — and FP4 is only one
Format 2: FP4 by hand
Format 3: MXFP4 by hand
Format 4: NVFP4 by hand
The scoreboard
Appendix A: the E4M3 grid by hand
Appendix B: why scale factors round up, not to nearest
Appendix C: the full 32-value table
Appendix D: Why E2M1 stops at 6, E4M3 stops at 448

Why by hand?

I was interested understand how NVFP4 works, why it is different from MXFP4, and why I should care about NVFP4. The diagram in Introducing NVFP4 for Efficient and Accurate Low-Precision Inference is not the easiest to follow either, hence this guide, which will work out NFVP4, MXFP4, and FP8 by hand.

FP4, MXFP4 and NVFP4 look almost identical in that each stores value in 4 bits, i.e. 1 sign bit, 2 exponent bits, and 1 mantissa bit (E2M1), and the only difference is a footnote about scaling factors. This footnote makes all the difference, and is the subject of this article.

By the end you should be able to look at any block-scaled format and immediately answer: what is the grid spacing here, and which values fall through the cracks?

This article was written with the assistance of AI.

The cast of characters

	FP8 (E4M3)	FP4 (E2M1)	MXFP4	NVFP4
Bits per element	8 (1S, 4E, 3M)	4 (1S, 2E, 1M)	4 (1S, 2E, 1M)	4 (1S, 2E, 1M)
Element range	±448	±6 × `s`	±6 × `s`	±6 × `s`
How many scale factors (granularity)	none needed †	1 per tensor	1 per 32 values	1 per 16 values
What format the scale is stored in (precision)	—	FP32 (software)	E8M0 (power of two)	E4M3 (fractional)
Second-level scale	—	—	—	1 per tensor, FP32

† FP8's 4 exponent bits per element already reach ±448, so this article casts to FP8 directly with no scale factor at all. (Appendix A works through the E4M3 grid in full, addressing questions such as where does the figure 448 come from.) Production FP8 inference often adds one FP32 scale per tensor to use that range well — still just one number for the whole tensor.

How the bits work: sign, exponent, mantissa

Four bits give $2^{4} = 16$ distinct bit patterns, and E2M1 reserves none of them for NaN or infinity, so all 16 decode to finite numbers, which are laid out as S EE M — sign, 2 exponent bits, 1 mantissa bit.

The exponent field picks a bracket: A bracket is the stretch of number line between one power of two and the next, $[2^{n}, 2^{n + 1})$ (formally called a binade). Every number inside the same bracket shares the same exponent. A bracket's base is its lower endpoint, i.e. $2^{n}$ , where the mantissa bit is zero.

Why exactly four? Because the exponent field is 2 bits wide, and 2 bits have $2^{2} = 4$ possible settings — 00, 01, 10, 11. Each setting names one bracket, so the count of brackets is fixed by the width of the exponent field:

number of brackets = 2^{(exponent bits)}

Consequence of the exponent width: E2M1 with 2 exponent bits and 1 mantissa bit stops at 6, E4M3 with 4 exponent bits and 3 manitssa bits reaches reaches 448. See (Appendix D for an explanation.

What "width" means, and why it doubles: A bracket's width is its length on the number line: $[1, 2)$ is 1 wide, $[2, 4)$ is 2 wide, $[4, 8)$ is 4 wide. Each is twice the previous one because moving up a bracket doubles both endpoints, and doubling both ends of an interval doubles the distance between them. Concretely:

width = 2^{n + 1} - 2^{n} = 2^{n} (2 - 1) = 2^{n}

The mantissa field splits the bracket: One mantissa bit gives two points per bracket, at 1.0× and 1.5× the base. Four brackets × two points each = 8 magnitudes:

`EE`	bracket	base	width	M = 0 (×1.0)	M = 1 (×1.5)	spacing
`00`	[0, 1)	—	1	0	0.5	0.5
`01`	[1, 2)	1	1	1.0	1.5	0.5
`10`	[2, 4)	2	2	2.0	3.0	1.0
`11`	[4, 8)	4	4	4.0	6.0	2.0

Every bracket holds the same number of points, but each is twice as wide as the one below — so the gap between points doubles too. The grid is non-uniform because exponents step in powers of two.

spacing = \frac{width}{points per bracket} = \frac{2 ^{n}}{2 ^{M}}

Applying the sign bit to those 8 magnitudes gives the 16 codes. Since 0000 and 1000 both mean zero, that is fifteen distinct values:

0, \pm 0.5, \pm 1, \pm 1.5, \pm 2, \pm 3, \pm 4, \pm 6

Putting the three fields together, each of those 16 codes answers three questions in order — which sign, which bracket, which of the two points inside it:

number	`S` — sign	`EE` — which bracket	`M` — which point	reading it off
3.0	`0` positive	`10` → [2, 4), base 2	`1` → ×1.5	+1.5 × 2 = 3.0
2.0	`0` positive	`10` → [2, 4), base 2	`0` → ×1.0	+1.0 × 2 = 2.0
−1.5	`1` negative	`01` → [1, 2), base 1	`1` → ×1.5	−1.5 × 1 = −1.5
−6.0	`1` negative	`11` → [4, 8), base 4	`1` → ×1.5	−1.5 × 4 = −6.0
0.5	`0` positive	`00` → subnormal	`1`	+0.5

In the above table, first two rows have the same sign and bracket, and differ only on the mantissa bits, which is what "the mantissa splits the bracket" means. Rows three and four differ only in the exponent field, and the value quadruples, because two brackets up is two doublings. The sign bit flips the result and touches nothing else, which is why the grid is perfectly symmetric about zero and why 0000 and 1000 collide on zero.

The same three fields, on real numbers

Below are some more examples of how nubmers convert to E2M1 and E4M3:

number	format	bits	significand × base	decodes to	error
3.0	E2M1	`0 10 1`	1.5 × 2	3.0	0
−6.0	E2M1	`1 11 1`	−1.5 × 4	−6.0	0
0.5	E2M1	`0 00 1`	subnormal, 0.5 × 1	0.5	0
5.05	E2M1	`0 11 1`	1.5 × 4	6.0	+0.95
−3.54	E4M3	`1 1000 110`	−1.75 × 2	−3.50	+0.04
20.848	E4M3	`0 1011 011`	1.375 × 16	22	+1.15

Row four is where 4-bit precision bites. 5.05 sits in bracket $[4, 8)$ with base 4, so its significand is $5.05/4 = 1.2625$ — but one mantissa bit offers only 1.0 or 1.5, and nothing in between. It rounds to 1.5 and comes back as 6.0. Rows five and six show what three mantissa bits buy: eight significand settings per bracket instead of two, so −3.54 lands on −3.50 rather than being flung to −4.

What is the scale for?

Notice what every one of those rows has in common: whatever bits you feed in, the result is a member of that same fixed set of fifteen. E2M1 has no other numbers available. It cannot represent 0.0003, and it cannot represent 4700.

However, real tensors contain numbers like 0.0003 and 4700. So the scale applies a separate multiplier s that stretches or shrinks the fixed grid onto wherever the data actually lives.

One mental image is to think of a ruler with fifteen tick marks. The 4 stored bits say which tick; the scale says how far apart the ticks are. The same ruler measures a molecule or a bridge depending only on the units you print on it. I will explain how to apply s with specific examples later.

What does it mean to "share" a scale?

Every element's exponent field is a private per-element scale. For example, FP8 E4M3 is a 3-bit mantissa grid, and a private 4-bit power-of-two scale for every element. When we say "share" a scale in the context of MXFP4 or NVFP4, we are referring to scales shared across elements:

	private exponent (per element)	shared scale (per group)	bits/value
FP8 E4M3	4 bits → 16 brackets	none	8.0
MXFP4	2 bits → 4 brackets	8 bits per 32 values	4.25
NVFP4	2 bits → 4 brackets	8 bits per 16 values	4.5

Note that block scaling does not replace the per-element exponent. Instead, it replaces the exponent's range.** For MXFP4 and NVFP4, the shared scale exists to slide that narrow 4-bracket window onto wherever the data lives. Each element still picks its own bracket inside the window. The bargain made with block scaling is to let a group of values share one scale, and amortise its cost across the group. One 8-bit scale per 32 values costs 0.25 bits per value; per 16 values the 8-bit scale costs 0.5 bits each.

How many scale factors are there, really?

Here is a count of the number of bits for a 32-value tensor:

Format	element bits	scale factors stored	total bits
FP8	32 × 8 = 256	none	256
FP4	32 × 4 = 128	1 FP32 = 32 bits	160
MXFP4	32 × 4 = 128	1 E8M0 = 8 bits	136
NVFP4	32 × 4 = 128	2 E4M3 = 16 bits, plus 1 FP32 = 32 bits	176

NVFP4 stores three numbers of metadata for thirty-two values, and MXFP4 stores exactly one. At 32 values the FP32 per-tensor scale is enormous relative to the data, so NVFP4 works out to 5.5 bits per value here rather than its headline 4.5. Real tensors are where the amortisation happens — a 4096 × 4096 weight matrix holds 16.7 million values, so that same 32-bit scale is spread across all of them:

Format	scale factors in a 4096 × 4096 tensor	bits per value
FP8	0	8.000000
FP4	1	4.000002
MXFP4	524,288	4.250000
NVFP4	1,048,576 block + 1 tensor	4.500002

Because every number in the group is stretched by the same amount, so the ruler has to be long enough for the group's largest member. A single outlier forces a long ruler with widely spaced ticks, and every one of its groupmates is measured on that coarse ruler whether it needed to be or not.

This is why group size matters:

Bigger groups → less metadata, but one outlier contaminates more neighbours.
Smaller groups → more metadata, but damage stays local.

Per-tensor scaling (plain FP4) is the extreme case where the group is everything, so a single outlier anywhere sets the ruler for the whole tensor. MXFP4 shrinks the group to 32, NVFP4 to 16.

Why is scale encoding important

There are two independent precisions in play:

Element precision — how finely you can pick among the ticks. This is the 4 bits, E2M1. It is identical across FP4, MXFP4 and NVFP4. Lack of element precision results in element rounding error, which is per-value and roughly random. Some values round up, some down, and across a dot product the errors partially cancel.
Scale precision — how finely you can set the tick spacing itself. This is E8M0 vs E4M3 vs FP32, and it is where the three formats genuinely differ. Lack of scale precision leads to scale error which is systematic. Get s wrong and all 16 or 32 values in the block are stretched wrongly in the same direction, at once. It is a bias, not noise, and biases do not cancel — they accumulate through the matmul.

So a bit spent on the scale buys more accuracy than a bit spent on elements. That is the central design insight of MXFP4 and NVFP4.

Everything is a grid and a scale

Every format here does the same three things. Strip away the names and each one is:

x^= q \times s, q = Q (\frac{x}{s})

where s is a scaling factor shared by some group of values, and q is the nearest point on a small fixed grid. $Q$ is shorthand for snap this number to the nearest value the format can actually represent. The subscript names which set you are snapping to — for all three 4-bit formats here, the E2M1 grid $0, \pm 0.5, \pm 1, \pm 1.5, \pm 2, \pm 3, \pm 4, \pm 6$ . Written out, it is picking the closest grid point:

Below are some examples:

$x$	$x / s$	$round_{grid}$	ordinary $round$
0.12	1.452	1.5	1
2.71	1.610	1.5	2
10.10	5.050	6	5
0.31	3.750	4	4

The first row lands on 1.5 — a value ordinary rounding cannot even produce — because 1.452 sits 0.048 from 1.5 but 0.452 from 1.0. The third row is the MXFP4 outlier we will meet shortly: 5.05 snaps all the way up to 6 because the grid holds nothing between 4 and 6. (Exact ties go to the even grid point where the mantissa bit is 0.)

The natural choice of scale for a group is the one that stretches the grid to exactly cover the largest magnitude in that group, i.e. absolute maximum (or amax):

s = \frac{amax}{6}, amax = i max ∣ x_{i} ∣

because 6 is the largest magnitude E2M1 can represent. Any smaller s clips the outlier; any larger s wastes grid points on a range the data never visits.

The E2M1 grid, enumerated by hand

Before quantizing anything, let us write out the entire FP4 grid. It has 16 codes and we can list all of them.

E2M1 lays each code out as S EE M — 1 sign bit, a 2-bit exponent field, 1 mantissa bit — and decodes with the usual floating-point formula:

value = (- 1)^{S} \times significand \times 2^{(E - bias)}

Three rules turn the bits into numbers, and a fourth explains why the grid reaches as high as it does:

The mantissa bit gives 1.0 or 1.5. Normal floats carry an implicit leading 1 that is never stored, so with one stored bit the significand is $1 + M /2$ — binary 1.0 or 1.1. This is why every magnitude below is a power of two or 1.5× one.
The bias is 1. For a 2-bit exponent field the bias is $2^{E - 1} - 1 = 2^{2 - 1} - 1 = 1$ , so a stored field $e$ means an actual exponent of $e - 1$ . (Where that formula comes from is worth its own section — see just below.)
Field 00 means subnormal. The implicit leading 1 is dropped (significand becomes $M /2$ ) and the exponent is pinned at $1 - bias = 0$ . This is where 0.5 comes from — it is not a normal value.
Field 11 is not reserved. IEEE formats spend their top exponent field on infinity and NaN. E2M1 has too few codes to spare, so 11 encodes ordinary numbers. That is why 4.0 and 6.0 exist, and why all 16 bit patterns decode to finite values.

Running all four exponent fields against both mantissa values:

exponent field	mantissa field	interpretation	magnitude
00	0	subnormal: $0 \times 2^{0}$	0
00	1	subnormal: $0.5 \times 2^{0}$	0.5
01	0	$1.0 \times 2^{0}$	1.0
01	1	$1.5 \times 2^{0}$	1.5
10	0	$1.0 \times 2^{1}$	2.0
10	1	$1.5 \times 2^{1}$	3.0
11	0	$1.0 \times 2^{2}$	4.0
11	1	$1.5 \times 2^{2}$	6.0

Four exponent fields × two mantissa values = 8 magnitudes. Applying the sign bit gives 16 codes, and since 0000 and 1000 both mean zero, 15 distinct values:

0, \pm 0.5, \pm 1, \pm 1.5, \pm 2, \pm 3, \pm 4, \pm 6

Here is a rewritten, drop-in replacement for your text. It keeps your original formatting (including the math tags and table) so you can copy and paste it directly, but it translates the dense technical jargon into easy-to-understand concepts and analogies.

Where the bias comes from

Earlier, we used a formula for the bias ( $2^{E - 1} - 1$ ) that might have looked like it was pulled out of thin air. Here is what it actually means.

Why do we need a bias at all?
In a floating-point number, the "exponent" part is stored as a simple, positive-only counting number (starting from 0). There is no plus or minus sign attached to it. However, we need negative exponents to represent tiny fractions below 1.
To solve this without wasting space on a minus sign, we "shift the goalposts." By automatically subtracting a fixed number (the bias) from whatever positive number is stored, we instantly slide the lower half of our numbers down into negative territory.

Why use that specific math formula?
That specific formula perfectly centers the window. By subtracting exactly $2^{E - 1} - 1$ , our available exponents are split right down the middle—roughly half fall below zero, and half sit above zero.
More memorably, it puts 1.0 right in the middle. To get an actual exponent of 0 (which is required to make the number 1.0), the computer just stores the bias number itself. This same elegant rule works across all sizes of floats:

format	$E$	bias $= 2^{E - 1} - 1$	field encoding of 1.0	normal exponents
E2M1	2	1	`01`	0 … 2
E3M0	3	3	`011`	−2 … 4
E4M3	4	7	`0111`	−6 … 8
E5M2	5	15	`01111`	−14 … 15
FP32	8	127	`01111111`	−126 … 127
FP64	11	1023	`01111111111`	−1022 … 1023

Why use a bias instead of the computer's normal negative numbers?
The answer is sorting. Using a bias keeps all the binary 1s and 0s in perfect numerical order. If you take positive floating-point numbers and tell the computer to read their raw bits as if they were just basic, simple integers, they will naturally sort correctly from smallest to largest. As Appendix B shows, because everything is perfectly ordered, rounding a number up to the next available slot doesn't require complex math or a lookup table—the computer literally just adds 1 to the raw integer value. This perfect sorting is what makes NVFP4's "scale selection" so incredibly cheap to run.

You can see this smoothly working in the E4M3 format: as the actual numbers get larger, their raw underlying bit patterns just step up like a normal counter—0, 1, 2, 3... climbing perfectly in step with the values themselves. The creators of FP8 deliberately kept this feature because it allows hardware and software to compare and sort these complex decimal numbers using basic, lightning-fast integer operations.

Conversion by hand with a sample tensor

We will use one 32-element tensor, which is exactly one MXFP4 block and exactly two NVFP4 blocks. The values are split into two halves with deliberately different characters, to represent activation tensors in LLMs which routinely mix quiet channels with channels carrying huge outliers.

Block A — a "quiet" group (elements 1–16), amax = 0.47:

 0.31  -0.47   0.12   0.28  -0.19   0.41  -0.35   0.07
 0.22  -0.44   0.38  -0.11   0.16   0.29  -0.25   0.44

Block B — a "loud" group with an outlier (elements 17–32), amax = 10.10:

-3.54  -3.14   1.33  10.10  -0.07   2.71  -1.62   0.88
 3.01  -2.35   0.15   1.97  -0.53   2.24  -1.11   0.66

The tensor-wide amax is 10.10, driven entirely by one element in block B. Block A's largest magnitude is more than 21 times smaller.

Format 1: FP8 (E4M3) by hand

FP8 E4M3 gives us 1 sign, 4 exponent and 3 mantissa bits, with a maximum finite magnitude of 448. Our values all sit comfortably inside that, so no scaling is needed.

To quantize by hand, find the binade (the power-of-two bracket) each value falls in, then round to the nearest of the 8 mantissa steps within it.

Element 20, x = 10.10. It lies in $[8, 16)$ , so the exponent is $2^{3} = 8$ and the mantissa step is $8/8 = 1.0$ . Available points: 8, 9, 10, 11, 12.

10.10 \to 10.00, error 0.10

Element 17, x = -3.54. It lies in $[2, 4)$ , exponent $2^{1} = 2$ , mantissa step $2/8 = 0.25$ . Available points: 3.00, 3.25, 3.50, 3.75.

- 3.54 \to - 3.50, error 0.04

Element 8, x = 0.07. It lies in $[0.0625, 0.125)$ , exponent $2^{- 4}$ , mantissa step $2^{- 4} /8 = 0.0078125$ .

0.07 \to 0.0703, error 0.0003

The tiny value got a tiny absolute error, because FP8 carries its own exponent per element. It does not care that another element in the tensor is 144 times larger. Every value gets roughly 6% relative precision regardless of magnitude.

FP8 result: MSE = 0.00127, max absolute error = 0.110. No value is driven to zero.

Aside: FP8 is two formats

The FP8 standard defines two encodings, and which one you use depends on what the tensor is for. The convention, from NVIDIA's FP8 Formats for Deep Learning and baked into Transformer Engine's Format.HYBRID:

E4M3 for the forward pass (weights and activations). E5M2 for the backward pass (gradients).

Why that split

Both are 8 bits. They differ only in how those bits are divided between exponent and mantissa — which is the same range-versus-resolution trade we saw between E8M0 and E4M3 scales, now applied to the elements themselves:

	E4M3	E5M2
Layout	1S, 4E, 3M	1S, 5E, 2M
Brackets	$2^{4} = 16$	$2^{5} = 32$
Points per bracket	8	4
Max value	448	57,344
Smallest subnormal	$2^{- 9} \approx 0.00195$	$2^{- 16} \approx 0.0000153$
Relative step	12.5%	25%
Dynamic range	$2^{17.8}$	$2^{31.8}$
Distinct values	253	247

Deriving E5M2's maximum the same way we did E4M3's: 5 exponent bits give 32 fields with a bias of 15, and the top field is reserved for infinity and NaN. So the largest usable field is $e = 30$ , exponent $30 - 15 = 15$ , base $2^{15} = 32, 768$ . Two mantissa bits reach a significand of $1 + \frac{3}{4} = 1.75$ :

max = 1.75 \times 2^{15} = 57, 344

That is 128× E4M3's reach, bought by moving exactly one bit from the mantissa to the exponent — one bit doubles the bracket count while halving the points inside each, so range grows exponentially while resolution drops by half.

Why E5M2 for backward pass: Weights and activations are well-behaved in magnitude but need precision; gradients are the opposite. Gradient magnitudes swing across many orders of magnitude during training, and the small ones matter — a gradient that underflows to zero contributes nothing to the update. Gradients also tolerate coarse values well, since they are averaged over a batch and noisy to begin with.

Our tensor demonstrates the forward-pass half directly. It looks like activations — moderate magnitudes, no extreme range — and E4M3 is the clear winner:

	MSE	max abs error
E4M3	0.00127	0.110
E5M2	0.00611	0.240

Flip to gradient-like values and the verdict reverses, because E4M3's floor is $2^{- 9}$ :

value	E4M3	E5M2
$1 0^{- 3}$	0.00195	0.00098
$1 0^{- 5}$	0 (underflow)	0.0000153

A gradient of $1 0^{- 5}$ simply ceases to exist in E4M3. In E5M2 it survives, imprecisely, but still preferable to underflow.

Why this matters for the rest of the article

Because FP8 uses one scaling factor for the whole tensor, E5M2 is needed for backward pass. With a single per-tensor scale, every value must fit inside one format's dynamic range at once, so a tensor spanning many orders of magnitude forces us to spend bits on range — meaning E5M2 and its coarser mantissa.

Block scaling dissolves that constraint. NVIDIA's Transformer Engine documentation notes that MXFP8, by giving each block of 32 values its own scale, reduces the dynamic range any single block has to cover — so MXFP8 uses E4M3 everywhere, forward and backward. The recipe is Format.E4M3 for MXFP8 versus Format.HYBRID for plain FP8.

(The split is a convention, not a law. DeepSeek-V3 trained with E4M3 in both directions. Recent work on FP8 reinforcement learning found the hybrid recipe tracked a BF16 baseline while pure-E4M3 destabilised around 500 steps, gradient norms spiking as the range ran out — so the trade is real, and it depends on how wide your gradients actually spread.)

And for FP4? There is only one

The natural follow-up is whether FP4 has an equivalent pair. It does not. The OCP microscaling specification defines two element encodings for MXFP8 (E4M3 and E5M2) and two for MXFP6 (E2M3 and E3M2) — but for MXFP4 it lists exactly one: E2M1. NVFP4 uses the same scales for both forward and backward pass too.

Two reasons for not splitting the precisions between forward and backward pass:

First, four bits leaves nothing to trade. After the sign bit, three bits must be divided between exponent and mantissa, which gives only four possible layouts — and the ones on either side of E2M1 are degenerate:

layout	magnitudes it can represent	character
E3M0	0, 0.25, 0.5, 1, 2, 4, 8, 16	pure powers of two — no mantissa at all, like E8M0
E2M1	0, 0.5, 1, 1.5, 2, 3, 4, 6	the standard
E1M2	0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5	uniformly spaced — INT4 in disguise
E0M3	—	no exponent; this is integer format

Quantizing our tensor with each, holding the scaling identical (NVFP4-style 16-value blocks, ideal fractional scale) so that only the element layout varies:

element format	MSE (overall)	MSE on block A (quiet)	MSE on block B (outlier)
E3M0	0.07416	0.00262	0.14571
E2M1	0.02535	0.00091	0.04980
E1M2	0.08571	0.00038	0.17103
INT4	0.08571	0.00038	0.17103

E2M1 has the lowest overall MSE. On the quiet block, the uniform layouts are best — E1M2 beats E2M1 by more than 2×, because when values are evenly spread, evenly spaced grid points fit them better. E2M1 wins on the strength of block B, which is a non-uniform grid with points bunched near zero and stretched at the top.

(Note E1M2 and INT4 score identically. That is not a coincidence — E1M2's grid is 0, 0.5, 1.0 … 3.5, which is INT4's 0…7 scaled by one half, and a free per-block scale absorbs the half. E1M2 is INT4. This is also why INT4 remains genuinely competitive for weight-only quantization, where distributions are well behaved, and loses to FP4 on activations, where outliers are the norm.)

Second, block scaling removes the reason a second format would exist. With one scale per tensor, gradients spanning many orders of magnitude need a format that carries its own range. Give every 16 or 32 values their own scale and that pressure vanishes; the scale supplies the range, so element bits are better spent on mantissa.

AMD's engineers put the same point about FP6, observing that E2M3 usually outperforms E3M2 in practice because the scaling factor already captures most of the dynamic range, which makes the extra mantissa bit the more valuable one. The identical logic pushes FP4 toward mantissa — but with only three bits to allocate, pushing one step past E2M1 lands on E1M2, which surrenders the non-uniform grid that outliers require. E2M1 is not one option among several; it is the only balance point that fits.

So the progression across the MX family tells a consistent story: 8 bits leave enough room for two encodings, 6 bits still leave two, and by 4 bits the trade space has collapsed to a single point. The fewer element bits you have, the more the work shifts from the format to the scale.

Format 2: FP4 by hand

Plain FP4 is E2M1 with a single FP32 scaling factor applied in software across the whole tensor.

Step 1: find the scale.

s = \frac{amax}{6} = \frac{10.10}{6} = 1.68333

Step 2: read off the effective grid. Multiply the E2M1 grid by s:

0    0.84   1.68   2.53   3.37       5.05       6.73            10.10
|      |      |      |      |          |          |                |

The dead zone is everything below $0.25 \times 1.68333 = 0.42$ .

Step 3: quantize a few elements.

Element 20, $x = 10.10$ :

\frac{10.10}{1.68333} = 6.000 \to q = 6 \to x^= 6 \times 1.68333 = 10.10, error 0

The outlier is perfect, by construction. Now block A.

Element 1, $x = 0.31$ :

\frac{0.31}{1.68333} = 0.1842 \to q = 0 \to x^= 0, error 0.31

Element 8, $x = 0.07$ :

\frac{0.07}{1.68333} = 0.0416 \to q = 0 \to x^= 0

Element 2, $x = - 0.47$ :

\frac{- 0.47}{1.68333} = - 0.2792 \to q = - 0.5 \to x^= - 0.8417, error 0.372

That last one is worth staring at. The value -0.47 was nearly in the dead zone, and rounding it out to the first grid point nearly doubled its magnitude. When a value sits near the dead-zone boundary, FP4 has only two options and both are badly wrong.

Thirteen of block A's sixteen values collapse to zero. A single scale chosen for a tensor whose amax is 10.10 simply cannot see values around 0.3.

FP4 result: MSE = 0.06735, max absolute error = 0.410. Block A alone has MSE 0.0849; block B has 0.0498.

Format 3: MXFP4 by hand

MXFP4 introduces the first innovation: hardware-supported block scaling. Instead of one scale for the tensor, each 32-value block gets its own scale — and the Tensor Core handles the grouping and rescaling natively.

The catch is how the scale is stored. MXFP4 uses E8M0: 8 exponent bits, no sign bit, and no mantissa bits at all. An E8M0 value is exactly a power of two, $2^{n}$ . It is compact and rescaling becomes an exponent addition, but it means the scale can never be 1.68 — only 1, or 2, or 0.5.

Step 1: compute the shared exponent. The OCP microscaling spec sets

X = ⌊ lo g_{2} (amax)⌋ - emax_{elem}

where $emax elem = 2$ for E2M1, because $2^{2} = 4$ is the largest power of two on the E2M1 grid. Note that $emax elem$ refers to the exponent of the largest normal number. In computing, a normal number is a non-zero floating-point number that falls within the standard, balanced range supported by a hardware format without losing precision. $⌊ \cdot ⌋$ denotes floor operation, which is used because the scaling is E8M0.

Our tensor is exactly one 32-value block, amax = 10.10:

⌊ lo g_{2} (10.10)⌋ = ⌊ 3.336 ⌋ = 3

X = 3 - 2 = 1 \Rightarrow s = 2^{1} = 2

Step 2: read off the effective grid.

MXFP4 (s = 2):    0     1     2     3     4        6        8              12
FP4   (s = 1.68): 0    0.84  1.68  2.53  3.37     5.05     6.73            10.10

Look at where the top of the grid sits: 12.0, when our data only reaches 10.10. The power-of-two constraint forced us to overshoot.

Step 3: quantize the outlier.

\frac{10.10}{2} = 5.05

The neighbouring grid points are 4 and 6. Distances: $∣5.05 - 4∣ = 1.05$ and $∣5.05 - 6∣ = 0.95$ . So it rounds up:

q = 6 \to x^= 6 \times 2 = 12.00, error 1.90

The largest value in the tensor is reconstructed 19% too large. This is the failure mode NVIDIA's Figure 3 illustrates, and it is a direct consequence of the coarse scale: with a fractional scale we could have placed a grid point exactly on 10.10.

Step 4: quantize block A. With a dead zone of 0.50 and block A's largest magnitude being 0.47:

\frac{0.47}{2} = 0.235 \to q = 0

All sixteen values of block A quantize to zero. The entire quiet half of the tensor is deleted.

MXFP4 result: MSE = 0.19549, max absolute error = 1.900. Block A MSE 0.0941; block B MSE 0.2969.

MXFP4 scored worse than plain FP4 on this tensor. With only 32 values, our tensor is a single MXFP4 block, so MXFP4 gets none of the block-scaling benefit it was designed for while paying the full cost of the coarse E8M0 scale. On a large real tensor with many blocks, MXFP4 comfortably beats a single per-tensor scale. What this example isolates cleanly is the cost of the power-of-two constraint, and we will pay that debt off properly in the ablation section.

(Some round the shared exponent up rather than using the floor rule, which trades the overshoot on amax for a wider dead zone.)

Format 4: NVFP4 by hand

NVFP4 changes both knobs at once.

Knob one: the block shrinks from 32 values to 16. Our tensor now has two blocks instead of one, and block A finally gets a scale chosen for block A.

Knob two: the scale is stored in E4M3 instead of E8M0. E4M3 has 3 mantissa bits, so it can express fractional scales — 1.75, 22, 3.5 — not just powers of two. This is what NVIDIA means by "high-precision scale encoding."

But E4M3 tops out at 448 and bottoms out around 0.00195, which is a much narrower range than E8M0's $2^{\pm 127}$ . A tensor whose blocks need wildly different scales could run off the end. This is what the second-level FP32 per-tensor scale is for: it pre-normalizes the whole tensor so that every block's scale lands inside E4M3's range.

Step 1: the per-tensor FP32 scale

The standard recipe picks the global scale so that the block containing the tensor's amax maps to E4M3's maximum, 448:

s_{global} = \frac{448 \times 6}{amax _{tensor}} = \frac{2688}{10.10} = 266.1386

Conceptually we are working with $x^{'} = x \times s_{global}$ , and each block's scale is computed on the stretched values. The 6 accounts for the E2M1 grid maximum; the 448 for the E4M3 maximum. Together they guarantee every block scale fits.

Step 2: block A's scale

amax_{A} = 0.47

s_{A'} = \frac{0.47 \times 266.1386}{6} = 20.848

Note that the microscaling rule with E4M3 is to divide amax by 6, which is what is used in the formula above.

Then round 20.848 into E4M3. It lies in $[16, 32)$ , so the mantissa step is $16/8 = 2$ . The candidates are 20 and 22.

Nearest would be 20 — but scale factors are rounded up, not to nearest, because a scale that is too small clips the block's outlier. (Appendix B works through why.) So:

s_{A'} = 22

Decode side, the effective scale is:

s_{A} = \frac{22}{266.1386} = 0.082664

Block A's grid, obtained by multiplying E2M1 by 0.082664:

0   0.0413  0.0827  0.1240  0.1653  0.2480  0.3307       0.4960
|      |       |       |       |       |       |            |

The dead zone is now $0.25 \times 0.082664 = 0.0207$ . Compare that to MXFP4's 0.50 — it is 24 times narrower. Block A is no longer invisible.

Step 3: quantize block A by hand

Element 1, $x = 0.31$ :

\frac{0.31}{0.082664} = 3.750 \to q = 4 \to x^= 4 \times 0.082664 = 0.3307, error 0.021

Element 2, $x = - 0.47$ :

\frac{- 0.47}{0.082664} = - 5.686 \to q = - 6 \to x^= - 0.4960, error 0.026

Element 8, $x = 0.07$ :

\frac{0.07}{0.082664} = 0.847 \to q = 1 \to x^= 0.0827, error 0.013

Element 3, $x = 0.12$ :

\frac{0.12}{0.082664} = 1.452 \to q = 1.5 \to x^= 0.1240, error 0.004

Every one of these was exactly 0.00 under MXFP4. Zero of block A's sixteen values are lost.

Step 4: block B's scale

amax_{B} = 10.10 \Rightarrow s_{B'} = \frac{10.10 \times 266.1386}{6} = 448.000

Exactly 448 — no rounding at all. This is because we defined $s_{global}$ so the block holding the tensor amax would land precisely on E4M3's ceiling. The effective scale is

s_{B} = \frac{448}{266.1386} = 1.68333 = \frac{10.10}{6}

which is the ideal scale from our formula, recovered exactly.

\frac{10.10}{1.68333} = 6.000 \to q = 6 \to x^= 10.10, error 0

The outlier that MXFP4 reconstructed as 12.00 is now exact.

Other elements in block B still take real damage, because a block whose amax is 10.10 has a coarse grid no matter how precisely you encode the scale:

- 3.54 \to \frac{- 3.54}{1.68333} = - 2.103 \to q = - 2 \to - 3.3667, error 0.173

2.71 \to \frac{2.71}{1.68333} = 1.610 \to q = 1.5 \to 2.5250, error 0.185

- 0.07 \to \frac{- 0.07}{1.68333} = - 0.042 \to q = 0 \to 0, error 0.070

Based on the calculations for block B, we can see that NVFP4 does not rescue a block that genuinely contains a 144:1 dynamic range. Block B's MSE is 0.0498 under both plain FP4 and NVFP4. What NVFP4 buys is that the outlier's damage is contained within its own 16 values instead of spreading across 32.

NVFP4 result: MSE = 0.02555, max absolute error = 0.357, SQNR = 23.0 dB. Block A MSE 0.0013; block B MSE 0.0498.

The scoreboard

Format	MSE	Max abs error	Block A MSE	Block B MSE	Block A values zeroed
FP8 (E4M3)	0.00127	0.110	0.00001	0.00252	0 / 16
FP4 (per-tensor)	0.06735	0.410	0.08490	0.04980	13 / 16
MXFP4	0.19549	1.900	0.09411	0.29688	16 / 16
NVFP4	0.02555	0.357	0.00131	0.04980	0 / 16

NVFP4 delivers 7.6× lower MSE than MXFP4 and 2.6× lower than plain FP4, at 4 bits per element in all three cases.

The gap between NVFP4 and FP8 is still significant — NVFP4 has x20 more MSE than FP8. But most of NVFP4's errors are in block B, and specifically in the values that share a block with a 21× outlier. On a tensor without that pathology, NVFP4 closes a lot of the remaining distance.

Appendix A: the E4M3 grid by hand

Understanding NVFP4's scale rounding requires knowing what E4M3 can actually represent. It has 1 sign bit, 4 exponent bits (bias 7), and 3 mantissa bits.

Within any binade (bracket) $[2^{e}, 2^{e + 1})$ there are exactly 8 representable values, spaced $2^{e} /8$ apart — three mantissa bits split each bracket eight ways instead of E2M1's two. So the absolute spacing doubles with each binade while the relative spacing stays at about 12.5%.

Here is the complete list — all 16 exponent fields, not a sample:

`EEEE`	$e$	exponent $e - 7$	bracket	spacing
`0000`	0	−6 (subnormal)	[0, 0.015625)	0.00195
`0001`	1	−6	[0.015625, 0.03125)	0.00195
`0010`	2	−5	[0.03125, 0.0625)	0.0039
`0011`	3	−4	[0.0625, 0.125)	0.0078
`0100`	4	−3	[0.125, 0.25)	0.0156
`0101`	5	−2	[0.25, 0.5)	0.031
`0110`	6	−1	[0.5, 1)	0.0625
`0111`	7	0	[1, 2)	0.125
`1000`	8	1	[2, 4)	0.25
`1001`	9	2	[4, 8)	0.5
`1010`	10	3	[8, 16)	1
`1011`	11	4	[16, 32)	2
`1100`	12	5	[32, 64)	4
`1101`	13	6	[64, 128)	8
`1110`	14	7	[128, 256)	16
`1111`	15	8	[256, 448]	32

Why the brackets do not start at 1. The list runs down to $[0, 0.015625)$ . Fractional values get brackets exactly like large ones, because the bias of 7 makes stored exponent fields decode to negative exponents: field 1 means $1 - 7 = - 6$ , giving the bracket $[2^{- 6}, 2^{- 5})$ . Bias exists precisely so that a non-negative bit pattern can express a negative exponent; it is $2^{E - 1} - 1 = 2^{3} - 1 = 7$ for E4M3's four exponent bits, and §4 derives that formula. The smallest positive value E4M3 can hold is a subnormal at $2^{- 9} \approx 0.00195$ , not 1.

Where the 0000 row's numbers come from. Both follow from the subnormal rule. When the exponent field is 0, two things change: the implicit leading 1 is dropped, so the significand becomes $m /8$ instead of $1 + m /8$ , and the exponent is pinned at $1 - bias = - 6$ rather than $0 - bias$ . Enumerating the eight mantissa settings:

\frac{m}{8} \times 2^{- 6} for m = 0 \dots 7

giving 0, 0.00195, 0.00391, … 0.01367. The spacing is the step between consecutive $m$ , which is one eighth of $2^{- 6}$ :

spacing = \frac{1}{8} \times 2^{- 6} = 2^{- 3} \times 2^{- 6} = 2^{- 9} \approx 0.00195

The upper limit is then forced: the largest subnormal is $\frac{7}{8} \times 2^{- 6} = 0.01367$ , and one step further lands on $\frac{8}{8} \times 2^{- 6} = 2^{- 6} = 0.015625$ — which is exactly the smallest normal value, the 0001 row's base. So the subnormal region ends precisely where normals begin, which is why the row is written $[0, 0.015625)$ .

Why the exponent is pinned at $1 - bias$ , not $0 - bias$ . This is the part that looks arbitrary and is not. Using $- 6$ makes the subnormal spacing $2^{- 9}$ identical to the spacing inside the first normal bracket $[2^{- 6}, 2^{- 5})$ , which is also $2^{- 6} /8 = 2^{- 9}$ . The two regions therefore join seamlessly: every value from 0 up through $2^{- 5}$ is a multiple of $2^{- 9}$ , a single uniform ladder of 17 rungs with no discontinuity at the boundary.

Had the exponent been $0 - bias = - 7$ instead, the largest subnormal would be $\frac{7}{8} \times 2^{- 7} = 0.00684$ , leaving a gap of 0.0088 before the smallest normal at 0.015625 — a hole nine times wider than the local step size. Dropping the implicit leading 1 is what lets the significand reach 0 at all (so zero is representable, and values shrink gradually toward it instead of falling off a cliff), and pinning the exponent one notch up is what stops that fix from tearing a hole in the grid.

E2M1 uses the identical rule with its own numbers: exponent $1 - 1 = 0$ , significand $m /2$ , giving 0 and 0.5 — and 0.5 is exactly the spacing inside $[1, 2)$ as well. That is why the spacing column in the bracket table stalls at 0.5 for the bottom two rows instead of continuing to halve.

Why the top bracket is $e = 15$ , not $e = 16$ . Four bits count from zero: the sixteen patterns are 0 through 15, and 1111 is 15. There is no bit pattern for 16 — that would need a fifth bit ( $16 = 10000$ ). So the field spans $e = 0 \dots 15$ , which after subtracting the bias gives exponents $- 6$ through $+ 8$ . Sixteen fields, one spent on subnormals, fifteen normal brackets.

Why the range ends at 448 rather than infinity. Because E4M3 has no encoding that means infinity at all. IEEE formats spend their entire top exponent field on Inf and NaN; E4M3 spends one code on NaN and none on Inf, so there is simply no bit pattern left over to represent "unbounded." Values that exceed 448 saturate to 448 instead of becoming Inf. Note also that every float format ends at some finite maximum — in FP32 it is about $3.4 \times 1 0^{38}$ , and infinity is a separate special encoding beyond it, not the endpoint of the numeric range. E4M3 just omits that special encoding. The final row is written $[256, 448]$ with a closed bracket because it is a truncated bracket: it would have run to 512, but the NaN reservation removes its top code.

Where 448 comes from. The top bracket is $e = 15$ , so the exponent is $15 - 7 = 8$ and the base is $2^{8} = 256$ .

Why the significand reaches 1.875. Three mantissa bits hold an integer $m$ from 0 to 7, and the significand is $1 + m /8$ — the denominator being $2^{M}$ , one for each mantissa bit. Read as a binary fraction it is simply 1. followed by the three stored bits:

`MMM`	$m$	$1 + m /8$	as binary
`000`	0	1.000	1.000
`100`	4	1.500	1.100
`110`	6	1.750	1.110
`111`	7	1.875	1.111

So 111 gives $1 + \frac{7}{8} = 1.875$ , the largest significand any bracket can hold. Note it never reaches 2.0 — it stops exactly one step short, which is what keeps a value inside its own bracket instead of spilling into the next.

Why the top bracket stops at 1.75. The single code reserved for NaN is exponent 1111 with mantissa 111 — which is precisely the code that would otherwise have meant $1.875 \times 256 = 480$ . With 111 unavailable at $e = 15$ , the largest usable significand there is 110, i.e. 1.75:

max = 1.75 \times 2^{8} = 448

NVFP4 needs the second-level FP32 scale because E4M3 spans about $2^{18}$ in magnitude, as compared to E8M0 which spans $2^{254}$ . Without pre-normalizing the tensor, a block whose amax is very small relative to the tensor could need a scale below E4M3's floor and would be unrepresentable.

Appendix B: why scale factors round up, not to nearest

When we computed block A's scale, the ideal was 20.848 and the E4M3 candidates were 20 and 22. Nearest-rounding says 20, but we chose 22.

The block's job is to represent amax = 0.47. The grid's largest point is $6 \times s$ , so we need:

6 \times s \geq amax

With $s^{'} = 20$ , the effective scale is $20/266.1386 = 0.075147$ , and the grid maximum is:

6 \times 0.075147 = 0.4509 < 0.47

The block's largest value clips. It reconstructs as 0.4509, an error of 0.019 — and clipping is worse than it looks, because it is a systematic bias rather than symmetric rounding noise. Every value near the top of the block gets pulled inward in the same direction.

With $s^{'} = 22$ , the effective scale is 0.082664 and the grid maximum is:

6 \times 0.082664 = 0.4960 > 0.47

No clipping. The cost is that the grid is now about 5% wider than necessary, so every value carries slightly more rounding error.

The tradeoff is asymmetric. A slightly coarse grid spreads a little extra error evenly across all 16 values; a clipped grid inflicts concentrated, one-directional error on the largest values — which, in a matrix multiply, are the ones contributing most to the output. Rounding scales up is the safe direction.

Note that block B avoided this question entirely: its ideal scale was exactly 448. That is a designed property of the global-scale recipe, not a coincidence, and it means the block containing the tensor's amax always gets a mathematically perfect scale.

Why rounding up is cheap. Choosing the next representable value above a target is arithmetic because of monotonic bit ordering. Read a positive float's encoding as a plain unsigned integer and you get its position in the sorted list of representable values — E4M3's 127 non-negative values encode as exactly 0, 1, 2, … 126, with no gaps:

position	value	bits
0	0	`00000000`
1	0.00195	`00000001`
…	…	…
90	20	`01011010`
91	22	`01011011`
…	…	…
126	448	`01111110`

Think of them as house numbers on a street. The houses sit at wildly uneven distances — 0.002 apart near zero, 32 apart near 448 — but the numbers still run 1, 2, 3, 4, so "the next house up" is always just add one, regardless of the physical gap.

Hence, vonverting 20.848 lands on house 90, the value 20. It is too small, so move to house 91: the value 22. Conversion hardware offers a round-toward- $+ \infty$ mode that does it directly. Rounding scales up is the numerically safe choice and the cheap one, which is a large part of why the convention exists.

Appendix C: the full 32-value table

Block A (quiet, amax = 0.47)

#	original	FP8 E4M3	FP4	NVFP4
1	0.31	0.3125	0.0000	0.3307
2	-0.47	-0.4688	-0.8417	-0.4960
3	0.12	0.1172	0.0000	0.1240
4	0.28	0.2812	0.0000	0.2480
5	-0.19	-0.1875	0.0000	-0.1653
6	0.41	0.4062	0.0000	0.3307
7	-0.35	-0.3438	0.0000	-0.3307
8	0.07	0.0703	0.0000	0.0827
9	0.22	0.2188	0.0000	0.2480
10	-0.44	-0.4375	-0.8417	-0.4960
11	0.38	0.3750	0.0000	0.3307
12	-0.11	-0.1094	0.0000	-0.1240
13	0.16	0.1562	0.0000	0.1653
14	0.29	0.2812	0.0000	0.3307
15	-0.25	-0.2500	0.0000	-0.2480
16	0.44	0.4375	0.8417	0.4960

Block B (loud, amax = 10.10)

#	original	FP8 E4M3	FP4	MXFP4	NVFP4
17	-3.54	-3.5000	-3.3667	-4.0	-3.3667
18	-3.14	-3.2500	-3.3667	-3.0	-3.3667
19	1.33	1.3750	1.6833	1.0	1.6833
20	10.10	10.0000	10.1000	12.0	10.1000
21	-0.07	-0.0703	0.0000	0.0	0.0000
22	2.71	2.7500	2.5250	3.0	2.5250
23	-1.62	-1.6250	-1.6833	-2.0	-1.6833
24	0.88	0.8750	0.8417	1.0	0.8417
25	3.01	3.0000	3.3667	3.0	3.3667
26	-2.35	-2.2500	-2.5250	-2.0	-2.5250
27	0.15	0.1562	0.0000	0.0	0.0000
28	1.97	2.0000	1.6833	2.0	1.6833
29	-0.53	-0.5000	-0.8417	-1.0	-0.8417
30	2.24	2.2500	2.5250	2.0	2.5250
31	-1.11	-1.1250	-0.8417	-1.0	-0.8417
32	0.66	0.6875	0.8417	1.0	0.8417

Scale factors used

Format	Block(s)	Scale(s)	Grid step	Grid max	Dead zone
FP4	all 32	1.68333	0.8417	10.100	0.421
MXFP4	all 32	2.0 ( $2^{1}$ )	1.0000	12.000	0.500
NVFP4	A (16)	0.082664	0.0413	0.496	0.021
NVFP4	B (16)	1.683333	0.8417	10.100	0.421

Appendix D: Why E2M1 stops at 6, E4M3 stops at 448

Think of floating-point math like a machine with two dials:

The Multiplier Dial (The Exponent): This takes huge power-of-two jumps (1, 2, 4, 8, 16, 32...).
The Fine-Tuning Dial (The Mantissa): This is a decimal slider that always sits between 1.0 and just under 2.0.

To get your final number, the simply multiplies the two dials together:
Final Number = Multiplier × Fine-Tuner

To find the maximum possible number in any format, we turn both dials all the way up to their maximum settings.

1. Why E2M1 stops at `6`

E2M1 only has 4 bits.

1 bit is for the plus/minus sign.
2 bits are for the Multiplier dial.
1 bit is for the Fine-Tuning dial.

The Multiplier Dial (2 bits):
With 2 bits, you only have 4 possible settings (00, 01, 10, 11). The computer assigns these to the multipliers: 0.5, 1, 2, and 4.

Maximum Multiplier = **4**

The Fine-Tuning Dial (1 bit):
With 1 bit, the dial only has two clicks: 0 or 1. The computer assigns these to the decimals 1.0 and 1.5.

Maximum Fine-Tuner = **1.5**

The Math:
To get the biggest number, multiply the maximums together.

4 x 1.5 = 6.0

With only 4 bits, the machine simply doesn't have enough notches on its dials to reach anything higher than 6.

2. Why E4M3 stops at `448`

E4M3 (used for FP8 and NVFP4 metadata) has 8 bits. It gives us much bigger dials.

1 bit is for the plus/minus sign.
4 bits are for the Multiplier dial.
3 bits are for the Fine-Tuning dial.

Let's turn these dials to the max:

The Multiplier Dial (4 bits):
With 4 bits, we have 16 settings. The computer scales these powers of two much higher: 1, 2, 4, 8, 16... all the way up to 256.

Maximum Multiplier = **256**

The Fine-Tuning Dial (3 bits):
With 3 bits, we have 8 clicks. The computer slices the space between 1.0 and 2.0 into eighths: 1.0, 1.125, 1.250, etc. The very highest click on this dial is 1.875.

The Math (and the Twist):
If we multiply the maximums together, we get:
256 x 1.875 = 480

Why is the limit 448 and not 480?
Because of a safety rule: computers need a special code to represent an "Error" or "Not a Number" (NaN) — for example, if you try to divide by zero.

Instead of adding an extra bit for errors, the engineers designed E4M3 to sacrifice its absolute highest setting. The computer is not allowed to turn both dials to their absolute max at the same time. That specific combination (256 + 1.875) is locked and reads as "Error".

So, to get the highest real number, we keep the Multiplier at 256, but we have to click the Fine-Tuning dial down by exactly one notch.
The notch just below 1.875 is 1.75.

So the final math is:
256 x 1.75 = 448

Summary

E2M1 maxes out at 6 because its biggest multiplier is 4 and its biggest fine-tuner is 1.5. (4 x 1.5 = 6).
E4M3 maxes out at 448 because its biggest multiplier is 256, and its biggest legal fine-tuner before triggering an error code is 1.75. (256 x 1.75 = 448).

References

Introducing NVFP4 for Efficient and Accurate Low-Precision Inference — NVIDIA Technical Blog
OCP Microscaling Formats (MX) Specification — for the E8M0 shared-exponent rule
FP8 Formats for Deep Learning — Micikevicius et al., the paper defining E4M3 and E5M2
Using FP8 and FP4 with Transformer Engine — the hybrid recipe, and why MXFP8 needs only E4M3
NVIDIA TensorRT Model Optimizer and LLM Compressor — for quantizing your own models to NVFP4

How does low-rank adaptation for large language models work

Lewis Won — Sat, 13 Sep 2025 12:50:55 +0000

Why study LoRA? The challenge of fine-tuning massive models
Conceptual overview: Full Fine-Tuning
Link back to the Math Equation on Full Fine-tuning
Conceptual Overview: Low-Rank Adaptation (LoRA)
Mathematical Walk-through of Forward Pass with LoRA
Mathematical Walk-through of the Backward Pass
Link Back to the Math Equation on LoRA
LoRA in the Self-Attention Module
Appendix A: The "No Additional Inference Latency" Trick
Appendix B: Why Does the Low-Rank Hypothesis Make Sense?

Why study LoRA? The challenge of fine-tuning massive models

Large Language Models (LLMs), because of their large size, present a significant challenge for training and deployment. For example, a model like GPT-3 has 175 billion parameters, would consume 350GB of storage (assuming FP16 quantization), and can consume up to 1.2TB of VRAM during training. (Note: I use GPT-3 as an example throughout this article to stay close to the original paper on LoRA.)

Fine-tuning is a way to adapt these pre-trained models to specific downstream tasks. Full fine-tuning involves updating all of the model's weights, which means that for every new task, you create a new, massive version of the model. Imagine a company wanting to offer 100 different specialized models to its customers. With full fine-tuning, this would require storing 100 separate 350GB models, consuming 35 terabytes of storage. This could make it difficult to deploy and manage customized LLMs at scale.

Low-Rank Adaptation (LoRA), introduced in the paper "LoRA: Low-Rank Adaptation of Large Language Models" by Hu et al., allows for the adaptation of massive models with a tiny fraction of the trainable parameters, significantly reducing storage costs and training overhead while maintaining satisfactory performance.

This article provides an intuitive walkthrough of how LoRA achieves this. This article was written with the assistance of Google Gemini 2.5 Pro.

Conceptual overview: Full Fine-Tuning

To understand the innovation of LoRA, we first need to look at the standard fine-tuning process. A neural network, at its core, is composed of layers, many of which perform matrix multiplication using weight matrices.

Let's imagine a single weight matrix in a pre-trained model, which we'll call $W_{0}$ . This matrix might have thousands of rows and columns.

Pre-trained Weights (W_{0})

When we fine-tune this model on a new task, we update these weights based on the new data. The process learns a "delta" or change matrix, $Δ W$ , which is added to the original weights.

The new, fine-tuned weight matrix, $W_{f t}$ , is:

W_{f t} = W_{0} + Δ W

The crucial point here is that the change matrix $Δ W$ has the exact same dimensions as the original matrix $W_{0}$ . If $W_{0}$ has 100 million parameters, then we are training a $Δ W$ that also has 100 million parameters. To save the fine-tuned model, we must save the entire $W_{f t}$ matrix, which is just as large as the original.

A Simple Example

Imagine our pre-trained weight matrix is a simple 2x3 matrix:

W_{0} = [0.8 0.2 0.1 0.7 0.3 0.5]

After fine-tuning, we might learn the following update matrix:

Δ W = [0.1 - 0.05 - 0.2 0.15 0.05 0.1]

The new, fully fine-tuned weight matrix would be:

W_{f t} = W_{0} + Δ W = [0.9 0.15 - 0.1 0.85 0.35 0.6]

To make this change, we had to train and store 6 new values for $Δ W$ . For a model like GPT-3, this means training and storing 175 billion new values for every single task.

Link back to the Math Equation on Full Fine-tuning

The paper provided the equation below the describe the standard process of full fine-tuning:

Φ max (x, y) \in Z \sum t = 1 \sum ∣ y ∣ lo g (P_{Φ} (y_{t} ∣ x, y_{< t}))

What it means: This formula states that we want to find the best possible set of model parameters, denoted by $Φ$ , that maximizes the probability of generating the correct output sequences (y) given the input sequences (x) and all the preceding correct tokens $y_{< t}$ . We are optimizing over the entire set of parameters in the model.

Let's dissect each component:

$max_{Φ}$ : This means our goal is to maximize the following expression by changing the model's parameters, denoted by the set $Φ$ . In full fine-tuning, $Φ$ represents every single weight and bias in the entire model. For GPT-3, this is a set of ~175 billion parameters.
$\sum_{(x, y) \in Z}$ : This tells us to sum the results over our entire training dataset, which is a set $Z$ of context-target pairs $(x, y)$ . For example, in a summarization task, x would be a long article and y would be its short summary.
$\sum_{t = 1 ∣ y ∣}$ : This is for the autoregressive nature of language models. For each target sequence y, we sum over every single token from the beginning (t=1) to the end (|y|). The model tries to predict each token correctly, one by one.
$lo g (\cdot)$ : We use the logarithm of the probability. This is a standard technique that makes the math more stable and turns a long product of probabilities into a more manageable sum (the "log-likelihood"). Maximizing the log-probability is the same as maximizing the probability itself.
$P_{Φ} (y_{t} ∣ x, y_{< t})$ : It represents the probability assigned by the model $P$ (with its current parameters $Φ$ ) to the correct next token $y_{t}$ , given the input context $x$ and all the preceding correct tokens $y_{< t}$ .

In simple terms, this equation says: Adjust all 175 billion parameters ( $Φ$ ) to make the model as good as possible at predicting the next correct word in the sequence for all the examples in our training data. The model starts with pre-trained weights $Φ_{0}$ and learns an update $ΔΦ$ , resulting in final weights of $Φ_{0} + ΔΦ$ . The problem is that $ΔΦ$ is just as large as $Φ_{0}$ .

Link to this article: This equation is the mathematical representation of the concept described in the section "Conceptual overview: Full Fine-Tuning".
- The initial pre-trained weights, $Φ_{0}$ , correspond to our simple matrix example $W_{0}$ .
- The final, optimized parameters, $Φ$ , are equivalent to our fine-tuned matrix $W_{f t}$ .
- The update, $ΔΦ$ , which is what we learn during training, corresponds to $Δ W$ .
- The paper's key point that |ΔΦ| equals |Φ₀| is exactly what we illustrated: the update matrix $Δ W$ has the same large dimensions as the original matrix $W_{0}$ , making it expensive to train and store.

Conceptual Overview: Low-Rank Adaptation (LoRA)

LoRA is built on a key insight: the update matrix $Δ W$ does not need to have full rank to be effective. The authors hypothesize that the change in weights during adaptation has a low "intrinsic rank". This means that the large $Δ W$ matrix can be approximated with high fidelity by multiplying two much smaller matrices.

Instead of learning $Δ W$ directly, LoRA learns two smaller matrices, which we'll call A and B.

Δ W \approx B \cdot A

This is a low-rank decomposition. The "rank" (r) is a small number we choose (like 1, 2, 8, or 64) that determines the inner dimension of these thin matrices.

If $W_{0}$ is a $d \times k$ matrix, then $Δ W$ is also $d \times k$ .
With LoRA, matrix A will have dimensions $r \times k$ , and matrix B will have dimensions $d \times r$ .

The number of trainable parameters is now the sum of the parameters in A and B ( $(d \times r) + (r \times k)$ ), which is dramatically smaller than the $d \times k$ parameters in $Δ W$ , especially when $r$ is much smaller than $d$ and $k$ .

Crucially, during training with LoRA, the original weights $W_{0}$ are frozen and do not receive gradient updates. We only train the much smaller A and B matrices.

Figure 1 from the paper "LoRA: Low-Rank Adaptation of Large Language Models" by Hu et al. The pretrained weights W are frozen. Only the low-rank matrices A and B are trained.

The Same Example with LoRA

Let's use our 2x3 weight matrix $W_{0}$ again.

W_{0} = [0.8 0.2 0.1 0.7 0.3 0.5]

Here, $d = 2$ and $k = 3$ . Let's choose a tiny rank, $r = 1$ .

Matrix B will be $d \times r \to 2 \times 1$ .
Matrix A will be $r \times k \to 1 \times 3$ .

Suppose after training, we learn the following A and B:

B = [0.4 0.2], A = [0.25 - 0.5 0.1]

Now, let's compute our low-rank approximation of $Δ W$ :

Δ W \approx B \cdot A = [0.4 0.2] \cdot [0.25 - 0.5 0.1] = [0.1 0.05 - 0.2 - 0.1 0.04 0.02]

The key comparison is the number of parameters we had to train:

Full Fine-Tuning: The $Δ W$ matrix had $2 \times 3 = 6$ parameters.
LoRA (r=1): We trained a $2 \times 1$ matrix B and a $1 \times 3$ matrix A. The total is $(2 \times 1) + (1 \times 3) = 5$ parameters.

While the savings seem small here, for a large matrix in a real model, the difference is significant. The paper notes that for GPT-3, LoRA can reduce the number of trainable parameters by 10,000 times (from 175B to about 17M). This means the storage for each new task drops from 350GB to just a few dozen megabytes.

Mathematical Walk-through of Forward Pass with LoRA

The standard forward pass for a layer is $h = W \cdot x$ , where h is the output, W is the weight matrix, and x is the input.

With LoRA, the output of the frozen pre-trained weights is computed as usual, and the output of the LoRA matrices is added to it.

The modified forward pass is:

h = W_{0} \cdot x + B \cdot A \cdot x

This can be written as $h = (W_{0} + B \cdot A) \cdot x$ . Let's use our example matrices to walk through the calculation.

1. Define Inputs

Pre-trained weights $W_{0}$ : $W_{0} = [0.8 0.2 0.1 0.7 0.3 0.5]$
Trained LoRA matrices $A$ and $B$ (with r=1): $B = [0.4 0.2], A = [0.25 - 0.5 0.1]$
An input vector $x$ : $x = ⎣ ⎡ 102030 ⎦ ⎤$

2. Calculate the original path
First, we compute the output from the frozen, pre-trained weights.

h_{0} = W_{0} \cdot x = [0.8 0.2 0.1 0.7 0.3 0.5] \cdot ⎣ ⎡ 102030 ⎦ ⎤ = [(8 + 2 + 9) (2 + 14 + 15)] = [1931]

3. Calculate the LoRA path
Next, we compute the update from our trained LoRA matrices. It's more efficient to multiply $A \cdot x$ first.

A \cdot x = [0.25 - 0.5 0.1] \cdot ⎣ ⎡ 102030 ⎦ ⎤ = [(2.5 - 10 + 3)] = [- 4.5]

Now multiply that result by B:

Δ h = B \cdot (A \cdot x) = [0.4 0.2] \cdot [- 4.5] = [- 1.8 - 0.9]

4. Combine the outputs
Finally, we add the two results together to get the final output h.

h = h_{0} + Δ h = [1931] + [- 1.8 - 0.9] = [17.2 30.1]

This process—keeping $W_{0}$ frozen and only passing gradients through the $B \cdot A$ path—is how LoRA achieves its remarkable parameter efficiency during training.

Mathematical Walk-through of the Backward Pass

The "backward pass," or backpropagation, is the core mechanism by which a neural network learns. Its goal is to calculate how much each trainable parameter in the model contributed to the final error (or "loss"). Once we know this, we can adjust the parameters slightly to reduce that error.

In LoRA, the key efficiency gain comes from the fact that we only need to calculate these adjustments for the tiny A and B matrices. The massive pre-trained weight matrix, W_0, is frozen, so we can completely ignore it during the backward pass, saving immense amounts of computation and memory.

Let's walk through how the gradients for A and B are calculated.

1. The Setup

First, let's recall the forward pass equation:

h = W_{0} \cdot x + B \cdot A \cdot x

The backward pass starts with a gradient signal coming from the next layer of the network. This signal, which we'll call $grad_h$ , tells us how the final loss $(L)$ would change with respect to a small change in our output, $h$ . Mathematically, this is the derivative $\frac{\partial L}{\partial h}$ . Our task is to use this incoming gradient to figure out $\frac{\partial L}{\partial A}$ and $\frac{\partial L}{\partial B}$ .

We will use the same matrices and input vector from our forward pass example:

$A = [0.25 - 0.5 0.1]$
$B = [0.4 0.2]$
$x = ⎣ ⎡ 102030 ⎦ ⎤$

Let's assume the incoming gradient $grad_h$ is:

grad_h = \frac{\partial L}{\partial h} = [0.5 - 0.2]

2. The Chain Rule in Action

The gradient for the LoRA path $(Δ h = B \cdot A \cdot x)$ is the same as the gradient for the total output $h$ , because $W_{0} \cdot x$ is treated as a constant. The gradient flows back only through the parts of the computation that involve our trainable parameters.

3. Calculating the Gradient for `B` (`grad_B`)

To find how B affects the loss, we use the chain rule. The gradient of the loss with respect to B is found by multiplying the gradient of the output (grad_h) by how B affects the output.

The update term is $Δ h = B \cdot (A \cdot x)$ . The derivative of this with respect to B involves the term it was multiplied by, which is $(A \cdot x)$ .

The formula is:

\frac{\partial L}{\partial B} = grad_h \cdot (A \cdot x)^{T}

Let's calculate this:

First, we need the term $(A \cdot x)$ . We already computed this in the forward pass: $A \cdot x = [- 4.5]$
The transpose, $(A \cdot x)^{T}$ , is $[- 4.5]$ .
Now we multiply: $grad_B = \frac{\partial L}{\partial B} = [0.5 - 0.2] \cdot [- 4.5] = [0.5 \times - 4.5 - 0.2 \times - 4.5] = [- 2.25 0.9]$ This grad_B matrix, which has the same shape as B, tells us how to adjust each element in B to reduce the loss.

4. Calculating the Gradient for `A` (`grad_A`)

Similarly, we find how A affects the loss. This time, the gradient has to pass back through B first.

The formula is:

\frac{\partial L}{\partial A} = (B^{T} \cdot grad_h) \cdot x^{T}

Let's calculate this step-by-step:

First, we need the transpose of B: $B^{T} = [0.4 0.2]$
Next, we multiply $B^{T}$ by grad_h: $B^{T} \cdot grad_h = [0.4 0.2] \cdot [0.5 - 0.2] = [(0.4 \times 0.5) + (0.2 \times - 0.2)] = [0.2 - 0.04] = [0.16]$
Finally, we multiply this result by the transpose of the input, $x^{T}$ : $grad_A = \frac{\partial L}{\partial A} = [0.16] \cdot [102030] = [1.6 3.2 4.8]$ This grad_A matrix, which has the same shape as A, tells us how to adjust A.

5. Updating the Weights

After calculating the gradients, the optimizer performs a weight update. Using a simple learning rate (lr), the update rule is:

A_{new} = A_{old} - lr \cdot grad_A

B_{new} = B_{old} - lr \cdot grad_B

And that's it! The crucial part is that W_0 is never updated. It is not part of the gradient calculation and requires no memory to store its gradients or optimizer states (like momentum). This is the source of LoRA's efficiency during the training process.

Link Back to the Math Equation on LoRA

The equation below from the paper describes the optimization objective of LoRA:

Θ max (x, y) \in Z \sum t = 1 \sum ∣ y ∣ lo g (P_{Φ_{0} + ΔΦ (Θ)} (y_{t} ∣ x, y_{< t}))

The key changes of this equation compared to the previous equation for full-model fine-tuning:

$max_{Θ}$ : This is the most important change. Instead of optimizing over the massive set of parameters $Φ$ , we are now optimizing over a much smaller set of parameters, denoted by $Θ$ . As the paper notes, the size of $Θ$ can be as small as 0.01% of the size of the original parameters ( $∣Θ∣ ≪ ∣ Φ_{0} ∣$ ). In our LoRA explainer, $Θ$ represents the collection of all the trainable values in our small A and B matrices.
$P_{Φ_{0} + ΔΦ (Θ)}$ : This shows how the model's weights are constructed.
- $Φ_{0}$ : These are the original, pre-trained weights of the large model. They are treated as a fixed constant and are not trained. This corresponds to our frozen $W_{0}$ matrix.
- $ΔΦ (Θ)$ : This represents the weight update, but it's no longer a huge matrix of trainable parameters. Instead, it's a function that generates the large update matrix from the small set of parameters $Θ$ . For LoRA, this function is the matrix multiplication of our small matrices: $ΔΦ (Θ) = B \cdot A$ . The parameters $Θ$ are the entries of B and A.
What it means:
1. The optimization is no longer over the massive parameter set $Φ$ , but over a much smaller set of parameters denoted by $Θ$ . The paper states $∣Θ∣ ≪ ∣ Φ_{0} ∣$ .
2. The update to the weights, $ΔΦ$ , is now a function of this small parameter set, written as $ΔΦ (Θ)$ . The original weights $Φ_{0}$ remain frozen.
Link to this article: This equation is the mathematical foundation for the section "Conceptual Overview: Low-Rank Adaptation (LoRA)".
- The small, trainable parameter set $Θ$ represents the collection of all the elements in our low-rank matrices, A and B.
- The function that generates the large update, $ΔΦ (Θ)$ , corresponds directly to the matrix multiplication $B \cdot A$ . This operation takes the small number of parameters in A and B and uses them to produce the full-sized update matrix $Δ W$ .
- The term $Φ_{0} + ΔΦ (Θ)$ in the equation is precisely what we illustrated in the forward pass: $h = (W_{0} + B \cdot A) \cdot x$ .

LoRA in the Self-Attention Module

The self-attention mechanism is a cornerstone of the Transformer architecture. For each input token, it computes Query (Q), Key (K), and Value (V) vectors. These are generated by multiplying the input embedding (x) with three distinct weight matrices: $W_{q}, W_{k}, W_{v}$ . After the attention scores are calculated and applied to the Value vectors, the result is passed through a final output projection matrix, $W_{o}$ , to produce the layer's output.

While LoRA can be applied to any weight matrix, the paper's authors focus their study on only the weight matrices in the self-attention module ( $W_{q}, W_{k}, W_{v}, W_{o}$ ). They find that for maximum parameter-efficiency, it is often sufficient to adapt only a subset of these matrices. For example, adapting only the query ( $W_{q}$ ) and value ( $W_{v}$ ) matrices can yield strong performance. For the purpose of a complete illustration, we will describe the process as if LoRA is applied to all four:

For $W_{q}$ , we add $B_{q} \cdot A_{q}$
For $W_{k}$ , we add $B_{k} \cdot A_{k}$
For $W_{v}$ , we add $B_{v} \cdot A_{v}$
For $W_{o}$ , we add $B_{o} \cdot A_{o}$

All eight of these LoRA matrices (A_q, B_q, A_k, B_k, etc.) are the only parameters that are trained. The original W matrices remain frozen. The calculations for each of the four paths are independent and can be performed in parallel.

Mathematical Walk-through of the Forward Pass

To keep the example clear, we will focus on the generation of a single Query vector. The exact same logic applies simultaneously to the Key and Value vectors.

1. Define Inputs

Let's assume our model has an embedding dimension (d_model) of 4, and the dimension of the Query/Key vectors (d_q or d_k) is 3.

Input Embedding (x), for a single token:

$x = ⎣ ⎡ 21 0.5 1.5 ⎦ ⎤ (shape 4 \times 1)$
Pre-trained Query Weight Matrix (W_q):

$W_{q} = ⎣ ⎡ 0.1 0.5 0.9 0.8 0.3 0.2 0.2 0.7 0.3 0.4 0.1 0.6 ⎦ ⎤ (shape 3 \times 4)$
Trainable LoRA Matrices for Query (A_q, B_q), with rank r=1:

$B_{q} = ⎣ ⎡ 0.5 - 0.2 0.1 ⎦ ⎤ (shape 3 \times 1) A_{q} = [0.2 - 0.1 0.4 0.3] (shape 1 \times 4)$

2. Calculate the Original Path (Frozen)

First, we compute the output from the pre-trained W_q matrix.

q_{0} = W_{q} \cdot x = ⎣ ⎡ 0.1 0.5 0.9 0.8 0.3 0.2 0.2 0.7 0.3 0.4 0.1 0.6 ⎦ ⎤ \cdot ⎣ ⎡ 21 0.5 1.5 ⎦ ⎤ = ⎣ ⎡ 0.2 + 0.8 + 0.1 + 0.6 1.0 + 0.3 + 0.35 + 0.15 1.8 + 0.2 + 0.15 + 0.9 ⎦ ⎤ = ⎣ ⎡ 1.7 1.8 3.05 ⎦ ⎤

3. Calculate the LoRA Path (Trainable)

Next, we compute the update from our LoRA matrices.

A_{q} \cdot x = [0.2 - 0.1 0.4 0.3] \cdot ⎣ ⎡ 21 0.5 1.5 ⎦ ⎤ = [0.4 - 0.1 + 0.2 + 0.45] = [0.95]

Δ q = B_{q} \cdot (A_{q} \cdot x) = ⎣ ⎡ 0.5 - 0.2 0.1 ⎦ ⎤ \cdot [0.95] = ⎣ ⎡ 0.475 - 0.19 0.095 ⎦ ⎤

4. Combine for the Final Query Vector

Finally, we add the outputs of the two paths.

q_{f ina l} = q_{0} + Δ q = ⎣ ⎡ 1.7 1.8 3.05 ⎦ ⎤ + ⎣ ⎡ 0.475 - 0.19 0.095 ⎦ ⎤ = ⎣ ⎡ 2.175 1.61 3.145 ⎦ ⎤

This q_final is the Query vector that will be used in the attention score calculation.

5. The Bigger Picture

This exact process happens in parallel for W_k (with A_k, B_k) and W_v (with A_v, B_v) to produce k_final and v_final. Those three vectors then proceed to the standard attention calculation. The final output of the attention mechanism is then passed through the W_o layer, which has its own LoRA update ( $h_{out} = (W_{o} + B_{o} A_{o}) \cdot h_{in}$ ).

Mathematical Walk-through of the Backward Pass

During backpropagation, the gradients flow backward from the loss function. The attention mechanism will provide an incoming gradient for q_final, k_final, and v_final. We'll demonstrate the process using the incoming gradient for our Query vector, grad_q.

1. The Setup

Let's assume the incoming gradient from the rest of the network for our Query vector is:

grad_q = \frac{\partial L}{\partial q _ f ina l} = ⎣ ⎡ 0.8 0.1 - 0.4 ⎦ ⎤

Our goal is to use grad_q to calculate the gradients for A_q and B_q. The gradient does not flow back into W_q.

2. Calculating the Gradient for `B_q`

The gradient for B_q is found by multiplying the incoming gradient grad_q by the term that B_q was multiplied by in the forward pass: (A_q * x).

grad_B_q = \frac{\partial L}{\partial B _{q}} = grad_q \cdot (A_{q} \cdot x)^{T}

We already know from the forward pass that A_q * x = [0.95].

grad_B_q = ⎣ ⎡ 0.8 0.1 - 0.4 ⎦ ⎤ \cdot [0.95] = ⎣ ⎡ 0.76 0.095 - 0.38 ⎦ ⎤

This is the adjustment signal for the B_q matrix.

3. Calculating the Gradient for `A_q`

The gradient for A_q requires the gradient to pass back through B_q first.

grad_A_q = \frac{\partial L}{\partial A _{q}} = B_{q T} \cdot grad_q \cdot x^{T}

First, let's compute B_q^T * grad_q:

B_{q T} \cdot grad_q = [0.5 - 0.2 0.1] \cdot ⎣ ⎡ 0.8 0.1 - 0.4 ⎦ ⎤ = [0.4 - 0.02 - 0.04] = [0.34]

Now, multiply this by x^T:

grad_A_q = [0.34] \cdot [21 0.5 1.5] = [0.68 0.34 0.17 0.51]

This is the adjustment signal for the A_q matrix.

4. Updating All LoRA Weights

The optimizer will use these computed gradients (grad_B_q, grad_A_q) to update the weights of B_q and A_q.

Crucially, this same backward pass logic is applied independently and simultaneously for the other LoRA pairs. The incoming gradient grad_k is used to update A_k and B_k, grad_v is used to update A_v and B_v, and the gradient from the next layer is used to update A_o and B_o. The massive W_q, W_k, W_v, W_o matrices are completely bypassed during the gradient computation, leading to massive savings in time and memory.

Appendix A: The "No Additional Inference Latency" Trick

A common drawback of other adaptation methods (like adding "adapter" layers) is that they introduce extra layers or computations that permanently increase the model's inference latency. LoRA cleverly avoids this.

During training, the forward pass involves two paths, as we saw: $h = W_{0} \cdot x + B \cdot A \cdot x$ . This does add a small amount of extra computation.

However, once training is complete, we can prepare the model for deployment (inference). Since matrices A and B are now fixed, we can perform their matrix multiplication once to get our final update matrix $Δ W$ .

Δ W = B \cdot A

Then, we can add this final update matrix directly to the original pre-trained weights to create a new, fully merged weight matrix, $W_{f t}$ .

W_{f t} = W_{0} + Δ W

Using our example numbers:

Δ W = [0.1 0.05 - 0.2 - 0.1 0.04 0.02]

W_{f t} = W_{0} + Δ W = [0.8 0.2 0.1 0.7 0.3 0.5] + [0.1 0.05 - 0.2 - 0.1 0.04 0.02] = [0.9 0.25 - 0.1 0.6 0.34 0.52]

When the model is deployed, we just use this merged $W_{f t}$ matrix. The forward pass becomes $h = W_{f t} \cdot x$ , which is a single matrix multiplication. This has the exact same computational cost and latency as the original, non-fine-tuned model.

This also makes switching between tasks incredibly fast. As the paper notes, to switch from Task 1 (with matrices $B_{1}, A_{1}$ ) to Task 2 (with $B_{2}, A_{2}$ ), you can compute $W_{0} + B_{2} \cdot A_{2}$ , which is a very fast operation compared to loading an entirely new 350GB model from disk.

Appendix B: Why Does the Low-Rank Hypothesis Make Sense?

The idea that a massive, over-parameterized model can be adapted by changing only a small number of parameters might seem counter-intuitive, but it is supported by research into the dynamics of deep learning.

The LoRA paper builds on work by Aghajanyan et al. (2020) and others, which showed that pre-trained language models have a low "intrinsic dimension." This suggests that even though they exist in a very high-dimensional parameter space (e.g., 175 billion dimensions), they can learn new tasks effectively by moving along a much smaller, lower-dimensional manifold within that space.

LoRA takes this idea a step further by hypothesizing that the change in weights during adaptation ( $Δ W$ ) also has a low "intrinsic rank." This means that the adjustments needed for a new task are not a complex, high-rank transformation but a simpler, low-rank one. The remarkable empirical success of LoRA, even with ranks as low as 1 or 2, provides strong evidence for this hypothesis. By decomposing the update into two small matrices, LoRA is essentially forcing the model to learn an update within this low-rank subspace, which proves to be a very effective and efficient constraint.

FlashAttention by hand

Lewis Won — Sat, 06 Sep 2025 13:42:58 +0000

Table of content

Introduction to FlashAttention
Vanilla self-attention mechanism without FlashAttention
Conceptual Overview: Fused Tiled Attention
FlashAttention by hand
Walkthrough of the FlashAttention diagram
Appendix A - How to derive activation matrices from weight matrices
Appendix B - The Paper's pseudocode vs. this walkthrough

Introduction to FlashAttention

FlashAttention builds upon the principles of online softmax to create an efficient single-pass algorithm for the entire self-attention mechanism. Its key innovation is to compute the final output O (where $O = softmax (Q K^{T}) V$ ) directly, without ever forming the full attention matrix A. This is achieved by fusing the matrix multiplications ( $Q K^{T}$ and $A V$ ) and the online softmax calculation into a single GPU kernel.

By avoiding the need to write and read the large $L \times L$ attention matrix (where L represents the number of input tokens) to and from global memory (DRAM), FlashAttention significantly reduces memory access, which is often the primary bottleneck in attention calculations. This makes it substantially faster and more memory-efficient, especially for long sequences.

This explanation will walk through the FlashAttention algorithm by hand, using the same tiled approach from the online softmax example in my previous article. This walkthrough is based on "From Online Softmax to FlashAttention" by Ye Zihao (2023).

This article was written with the assistance of Google Gemini 2.5 Pro.

Vanilla self-attention mechanism without FlashAttention

In order to better appreciate the efficiency gains with FlashAttention, we first begin with a vanilla implementation of self-attention. Let's illustrate this with our familiar 1 x 6 example used in the previous article on online softmax. Imagine this vector represents the dot products of a single query vector q with six key vectors k.

Dot Products (Logits): $X = q K^{T} = [123621]$

We also need a corresponding V matrix. For simplicity, let's assume V has 6 rows (one for each key) and a dimension of 2.

Value Matrix: $V = ⎣ ⎡ v_{1} v_{2} v_{3} v_{4} v_{5} v_{6} ⎦ ⎤ = ⎣ ⎡ 123456123456 ⎦ ⎤$

The standard process is broken down into three distinct, sequential stages:

Calculate Logits: Compute $X = Q K^{T}$ .
Calculate Attention Scores: Compute $A = softmax (X)$ . This creates a dense attention score matrix.
Calculate Final Output: Compute $O = A \cdot V$ .

Since the logits X are already given, our walkthrough will start from Step 2.

(Note: Q, K and V here are activation matrices which are result of multiplying the input embeddings by the weight matrices W_Q, W_K and W_V. These activation matrices are different for every input sequence and are the actual inputs to the attention calculation kernel that FlashAttention replaces. Refer to Appendix A to understand how these activation matrices are derived from the weight matrices.)

Step 1: Calculate the Full Attention Score Matrix (`A = softmax(X)`)

This step computes the softmax function over the entire logit vector X. This is a multi-pass process, requiring a pass to find the maximum, a pass to compute the denominator, and a pass to compute the final probabilities. The key difference from FlashAttention is that we must compute and store this entire attention vector A before we can even begin to use the V matrix.

1a. Find the Global Maximum (`m`)

First, we scan the entire logit vector X to find its maximum value for numerical stability (the "safe softmax" trick).

m = max ([123621]) = 6

1b. Compute the Exponentials and the Denominator (`d`)

Next, we subtract the maximum from each logit, exponentiate the result, and sum them all up to get the denominator.

Subtract max: $X - m = [1 - 6 2 - 6 3 - 6 6 - 6 2 - 6 1 - 6] = [- 5 - 4 - 3 0 - 4 - 5]$
Exponentiate: $e^{(X - m)} = [e^{- 5} e^{- 4} e^{- 3} e^{0} e^{- 4} e^{- 5}] \approx [0.0067 0.0183 0.0498 1 0.0183 0.0067]$
Sum to get the denominator d: $d = \sum e^{(X - m)} \approx 0.0067 + 0.0183 + 0.0498 + 1 + 0.0183 + 0.0067 = 1.0998$

1c. Normalize to get the Final Attention Vector `A`

Finally, we divide the exponentiated values by the denominator d to get the final probabilities. The result is the complete attention score vector A.

A = \frac{e ^{(X - m)}}{d} \approx \frac{[ 0.0067 0.0183 0.0498 1 0.0183 0.0067 ]}{1.0998} \approx [0.0061 0.0167 0.0453 0.9092 0.0167 0.0061]

At this point, the 1 x 6 vector A is fully computed and stored in memory.

Step 2: Multiply Attention Scores by the Value Matrix ( $O = A \cdot V$ )

Now that we have the attention scores, we can perform the final step: a matrix multiplication between A and V. The attention scores in A act as weights for the corresponding value vectors in V. The output O is the weighted sum of the value vectors.

The operation is $(1 \times 6) \cdot (6 \times 2) \to (1 \times 2)$ .

O = A \cdot V \approx [0.0061 0.0167 0.0453 0.9092 0.0167 0.0061] \cdot ⎣ ⎡ 123456123456 ⎦ ⎤

Let's calculate the weighted sum:

O \approx (0.0061 \cdot [11]) + (0.0167 \cdot [22]) + (0.0453 \cdot [33]) + (0.9092 \cdot [44]) + (0.0167 \cdot [55]) + (0.0061 \cdot [66])

Now we compute each term and then sum them up:

O \approx [0.0061 0.0061] + [0.0334 0.0334] + [0.1359 0.1359] + [3.6368 3.6368] + [0.0835 0.0835] + [0.0366 0.0366]

Summing the vectors component-wise:

O \approx [3.9323 3.9323]

This result, [3.932, 3.932].

Comparison and Key Differences

Intermediate Matrix: The standard method materialized the full 1 x 6 attention vector A. In a real-world scenario with a sequence length L, this would be a large L x L matrix. This is the main bottleneck. I will demonstrate how FlashAttention computes the final output without ever creating or storing this matrix.
Memory Access: The standard method requires at least two major memory operations: writing the entire A matrix to memory (DRAM), and then reading it all back in to multiply with V. I will show how FlashAttention fuses these operations, keeping tiles of Q, K, and V in fast SRAM and avoiding the slow roundtrip to DRAM.
Computation Flow: The process is strictly sequential. You cannot start the $A \cdot V$ multiplication until the entire softmax calculation for A is complete. We will show how FlashAttention integrates these steps, updating a running output vector as it iterates through tiles of the K and V matrices.

Conceptual Overview: Fused Tiled Attention

FlashAttention processes the input matrices Q, K, and V in a tiled manner. For each row of the output matrix O, it iterates through the corresponding rows of K and V in blocks. At each step, it calculates the attention scores for just that block, updates the running statistics (the maximum and the denominator, just like in online softmax), and immediately applies these scores to the corresponding block of V to update a running output vector.

The core idea is to maintain three running statistics for each row of the output:

m_running: The running maximum of the dot products (q · k).
d_running: The running denominator of the softmax.
o_running: The running output vector, which is a weighted sum of the V vectors, scaled by the current (and incomplete) softmax probabilities.

We will use the exact same input data as before with the 1 x 6 example. This vector represents the dot products of a single query vector q with six key vectors k.

Dot Products (Logits): $X = q K^{T} = [123621]$

We also need a corresponding V matrix. For simplicity, let's assume V has 6 rows (one for each key) and a dimension of 2.

Value Matrix: $V = ⎣ ⎡ v_{1} v_{2} v_{3} v_{4} v_{5} v_{6} ⎦ ⎤ = ⎣ ⎡ 123456123456 ⎦ ⎤$

We will process this with a tile size of 3, breaking X and V into two tiles:

$T_{1}$ : Logits $[123]$ and Values $⎣ ⎡ v_{1} v_{2} v_{3} ⎦ ⎤$
$T_{2}$ : Logits $[621]$ and Values $⎣ ⎡ v_{4} v_{5} v_{6} ⎦ ⎤$

The algorithm follows the logic of "Algorithm FlashAttention (Tiling)" on page 6 of "From Online Softmax to FlashAttention" by Ye Zihao (2023).

(Note: For intuitive clarity, this walkthrough maintains an un-normalized running output o_running and performs a single normalization at the end. The paper's pseudocode maintains a normalized running output o' at each step. The final result is mathematically identical. An explanation of why this is true is in Appendix B.)

FlashAttention by hand

Before the main loop begins, the running statistics are initialized.

m_running = -∞
d_running = 0
o_running = [0, 0] (a zero vector of the same dimension as a v vector)

Link back to the FlashAttention algorithm

These initial values correspond to the state before the for loop, where $m_{0} = - \infty$ , $d_{0'} = 0$ , and $o_{0'} = 0$ .

Step 1: Process Tile 1 (`i=1`)

We now begin the first iteration of the for i ← 1, #tiles do loop.

1a. Find the New Maximum

First, we calculate the dot products (logits) for this tile and find the maximum value.

Logits for Tile 1: $x_{1} = [123]$
Local max of Tile 1: $m_{T_{1}} = max ([123]) = 3$
New overall maximum: $m_{n e w} = max (m_{r u nnin g}, m_{T_{1}}) = max (- \infty, 3) = 3$

Link back to the FlashAttention algorithm

The calculation of logits corresponds to the first line inside the loop:

x_{i} \leftarrow Q [k, :] K^{T} [:, (i - 1) b : ib]

Finding the local max corresponds to the second line:

m_{i (local)} = j = 1.. b max (x_{i} [j])

Updating the running max corresponds to the third line:

m_{i} \leftarrow max (m_{i - 1}, m_{i (local)})

(Here, $m_{i - 1}$ is our m_running from before the step).

1b. Calculate Local Denominator and Local Output

Next, we compute the un-normalized attention scores for this tile using m_new, and then use them to calculate a local denominator d_local and a local weighted output o_local.

Un-normalized scores: $P_{1} = ⎣ ⎡ e^{1 - 3} e^{2 - 3} e^{3 - 3} ⎦ ⎤ = ⎣ ⎡ e^{- 2} e^{- 1} e^{0} ⎦ ⎤ \approx ⎣ ⎡ 0.1353 0.3679 1 ⎦ ⎤$
Local denominator: $d_{T_{1}} = \sum P_{1} \approx 0.1353 + 0.3679 + 1 = 1.5032$
Local output (un-normalized weighted sum of V vectors): $o_{T_{1}} = (0.1353 \cdot v_{1}) + (0.3679 \cdot v_{2}) + (1 \cdot v_{3}) \approx (0.1353 \cdot [11]) + (0.3679 \cdot [22]) + (1 \cdot [33]) \approx [0.1353 0.1353] + [0.7358 0.7358] + [33] = [3.8711 3.8711]$

Link back to the FlashAttention algorithm

The calculation of $d_{T_{1}}$ corresponds to the summation part of the denominator update rule:

j = 1 \sum b e^{x_{i} [j] - m_{i}}

The calculation of $o_{T_{1}}$ corresponds to the numerator of the second term in the output update rule (before normalization):

j = 1 \sum b e^{x_{i} [j] - m_{i}} V [j + (i - 1) b, :]

1c. Update Running Statistics

Now, we update our global running statistics.

d_{n e w} o_{n e w} = d_{o l d} \cdot e^{m_{o l d} - m_{n e w}} + d_{l oc a l} = o_{o l d} \cdot e^{m_{o l d} - m_{n e w}} + o_{l oc a l}

m_old = -∞, d_old = 0, o_old = [0, 0]
m_new = 3
d_local = 1.5032, o_local = [3.8711, 3.8711]

d_{r u nnin g} o_{r u nnin g} = (0 \cdot e^{- \infty - 3}) + 1.5032 = 1.5032 = ([00] \cdot e^{- \infty - 3}) + [3.8711 3.8711] = [3.8711 3.8711]

After processing the first tile, our statistics are: m_running = 3, d_running ≈ 1.5032, o_running ≈ [3.8711, 3.8711].

Link back to the FlashAttention algorithm

This entire step corresponds to the final two lines of the loop body for i=1.

Denominator update:

d_{i'} \leftarrow d_{i - 1'} e^{m_{i - 1} - m_{i}} + j = 1 \sum b e^{x_{i} [j] - m_{i}}

Since 
 $d_{0'} = 0$ 
, the first term is zero, and 
 $d_{1'}$ 
 becomes just the local sum, matching our result.

Output update:

o_{i'} \leftarrow o_{i - 1'} \frac{d _{i - 1'} e ^{m_{i - 1} - m_{i}}}{d _{i'}} + \frac{\sum e ^{x_{i} [j] - m_{i}} V [ \dots ]}{d _{i'}}

Similarly, since 
 $o_{0'}$ 
 and 
 $d_{0'}$ 
 are zero, the first term vanishes. Our un-normalized `o_running` is equivalent to 
 $o_{1'} \cdot d_{1'}$ 
, which is simply the numerator of the second term, matching our calculation of 
 $o_{T_{1}}$ 
.

Step 2: Process Tile 2 (`i=2`)

We now proceed to the second and final iteration of the loop.

2a. Find the New Maximum

Logits for Tile 2: $x_{2} = [621]$
Local max of Tile 2: $m_{T_{2}} = max ([621]) = 6$
New overall maximum: $m_{n e w} = max (m_{r u nnin g}, m_{T_{2}}) = max (3, 6) = 6$

Link back to the FlashAttention algorithm

This again maps to the first three lines of the loop body for i=2. This time, $m_{i - 1} = m_{1} = 3$ , so the new maximum is correctly found as $max (3, 6) = 6$ .

2b. Calculate Local Denominator and Local Output

We repeat the process for the second tile's data, using the new max, 6.

Un-normalized scores: $P_{2} = ⎣ ⎡ e^{6 - 6} e^{2 - 6} e^{1 - 6} ⎦ ⎤ = ⎣ ⎡ e^{0} e^{- 4} e^{- 5} ⎦ ⎤ \approx ⎣ ⎡ 1 0.0183 0.0067 ⎦ ⎤$
Local denominator: $d_{T_{2}} = \sum P_{2} \approx 1 + 0.0183 + 0.0067 = 1.025$
Local output: $o_{T_{2}} = (1 \cdot v_{4}) + (0.0183 \cdot v_{5}) + (0.0067 \cdot v_{6}) \approx (1 \cdot [44]) + (0.0183 \cdot [55]) + (0.0067 \cdot [66]) \approx [44] + [0.0915 0.0915] + [0.0402 0.0402] = [4.1317 4.1317]$

Link back to the FlashAttention algorithm

This again corresponds to the summation parts of the update rules for i=2.

2c. Update Running Statistics

Now we perform the final update, including the crucial rescaling step.

m_old = 3, d_old ≈ 1.5032, o_old ≈ [3.8711, 3.8711]
m_new = 6
d_local ≈ 1.025, o_local ≈ [4.1317, 4.1317]

First, update the denominator:

d_{r u nnin g} = d_{o l d} \cdot e^{m_{o l d} - m_{n e w}} + d_{l oc a l} \approx (1.5032 \cdot e^{3 - 6}) + 1.025 \approx (1.5032 \cdot 0.04979) + 1.025 \approx 0.0748 + 1.025 = 1.0998

Next, update the output vector:

o_{r u nnin g} = o_{o l d} \cdot e^{m_{o l d} - m_{n e w}} + o_{l oc a l} \approx ([3.8711 3.8711] \cdot e^{3 - 6}) + [4.1317 4.1317] \approx ([3.8711 3.8711] \cdot 0.04979) + [4.1317 4.1317] \approx [0.1927 0.1927] + [4.1317 4.1317] = [4.3244 4.3244]

Link back to the FlashAttention algorithm

This step corresponds to the full update rules for i=2.

Denominator update:

d_{2'} \leftarrow d_{1'} e^{m_{1} - m_{2}} + j = 1 \sum b e^{x_{2} [j] - m_{2}}

The term 
 $d_{1'} e^{m_{1} - m_{2}}$ 
 is the crucial rescaling factor, which perfectly matches the 
 $(1.5032 \cdot e^{3 - 6})$ 
 part of our calculation.

Output update:

o_{2'} \leftarrow o_{1'} \frac{d _{1'} e ^{m_{1} - m_{2}}}{d _{2'}} + \frac{\sum e ^{x_{2} [j] - m_{2}} V [ \dots ]}{d _{2'}}

Again, our un-normalized `o_running` is equivalent to 
 $o_{2'} \cdot d_{2'}$ 
. If you multiply the paper's update rule by 
 $d_{2'}$ 
, you get 
 $(o_{1'} d_{1'}) e^{m_{1} - m_{2}} + local sum$ 
. This exactly matches our formula:

o_{o l d} \cdot e^{m_{o l d} - m_{n e w}} + o_{l oc a l}

Final Result

After the loop finishes, we have the final un-normalized output and the final denominator.

d_final ≈ 1.0998
o_final ≈ [4.3244, 4.3244]

The last step is to normalize the output vector by dividing by the final denominator.

O_{f ina l} = \frac{o _{f ina l}}{d _{f ina l}} \approx \frac{[ 4.3244 4.3244 ]}{1.0998} \approx [3.932 3.932]

Link back to the FlashAttention algorithm

This final normalization step is implicitly the result of the algorithm. The final output of the loop, $o_{N / b'}$ , is the correctly normalized output row. Our method simply defers this division to the very end for clarity. The final output vector $O [k, :]$ in the algorithm is this final, normalized value.

O [k, :] \leftarrow o_{N / b'}

Walkthrough of the FlashAttention diagram

Having completed the step-by-step walkthrough of the FlashAttention algorithm, we can now walkthrough the FlashAttention schematics from the original paper ["FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"(https://arxiv.org/abs/2205.14135) by Tri Dao et. al. (2022).

High-Level Overview

The diagram shows a tiled computation where the goal is to compute one block of the final output matrix, softmax(QK^T)V, at a time. The loops control which tiles (blocks) of the input matrices (Q, K, V) are loaded from slow global memory (HBM) into fast local memory (SRAM) to perform a piece of the calculation.

Our walkthrough, where we computed a single output row vector, corresponds to one full pass of the "Inner Loop" in this diagram.

Mapping the Diagram Components to Our Walkthrough

Let's look at each element of the diagram and connect it to our example.

1. The Loops

Inner Loop (Blue Arrows): This loop iterates over the queries (rows of Q) and the corresponding output rows. In our example, we only had a single query vector q that produced a single output row o. Therefore, our entire walkthrough represents a single iteration of the Inner Loop.
- The Q matrix block being copied is our single query q.
- The Output to HBM block at the bottom is our o_running vector $[o_{1} o_{2}]$ , which is being progressively built.
Outer Loop (Red Arrows): This loop iterates over the key-value pairs (columns of K^T and rows of V). This loop is the core of the FlashAttention mechanism and maps directly to the steps of our walkthrough.
- Iteration 1 of the Outer Loop corresponds to our "Step 1: Process Tile 1".
- Iteration 2 of the Outer Loop corresponds to our "Step 2: Process Tile 2".

2. The Memory Hierarchy (SRAM vs. HBM)

HBM (High-Bandwidth Memory): This represents the GPU's main global memory (DRAM). This is where the full, large Q, K, V, and final O matrices reside.
SRAM (Fast On-Chip Memory): This is the small but extremely fast "workbench" memory. The diagram shows that we only ever copy small blocks (orange squares) into SRAM to work on them. In our walkthrough, the SRAM would hold:
- Our query vector q.
- The current block of keys and values we are processing (e.g., $k_{1}, k_{2}, k_{3}$ and $v_{1}, v_{2}, v_{3}$ for Tile 1).
- The running statistics: m_running, d_running, and o_running.

3. Putting it all Together: Tracing Our Walkthrough on the Diagram

Let's trace the flow for our single query q.

Initialization:

Before the loops start, the Output to HBM block (our o_running vector) is initialized to [0, 0]. The running statistics m_running = -∞ and d_running = 0 are initialized in SRAM.

Inner Loop Begins (One and only one iteration for our example):

Copy from Q: Our query vector q is loaded from HBM into SRAM. It will stay in SRAM for the entire duration of the Outer Loop.

Outer Loop - Iteration 1 (Our "Process Tile 1"):

Copy from K^T and V: The first block of K^T (corresponding to logits 1, 2, 3) and the first block of V (v_1, v_2, v_3) are loaded from HBM into SRAM.
Compute Block on SRAM: This is the central computation.
- The dot product $q \cdot K_{block 1 T}$ is calculated.
- The local max, local denominator, and local output are computed.
- The running statistics m_running, d_running, and o_running (which live in SRAM) are updated. After this step, o_running is $\approx [3.8711, 3.8711]$ . The + sign with the purple dotted arrow signifies this update step.

Outer Loop - Iteration 2 (Our "Process Tile 2"):

Copy from K^T and V: The previous blocks of K^T and V are discarded. The second block (logits 6, 2, 1 and values v_4, v_5, v_6) is loaded into SRAM.
Compute Block on SRAM: The computation is repeated.
- A new, larger global maximum (m_new = 6) is found.
- This triggers the crucial rescaling of the existing d_running and o_running vectors.
- The local contributions are calculated and added. After this step, the final un-normalized o_running is $\approx [4.3244, 4.3244]$ .

Outer Loop Finishes:

The loop is complete. The final normalization is performed in SRAM ( $o_{running} / d_{running}$ ) to get the final output vector $\approx [3.932, 3.932]$ .

Output to HBM:

The final, correct output vector is written from SRAM back to its designated row in the main output matrix in HBM.

Inner Loop Finishes:

If there are more rows in Q, the Inner Loop will continue to the next row and repeat the entire Outer Loop process.

Appendix A - How to derive activation matrices from weight matrices

For the purpose of focusing on the core concepts of FlashAttention, the walkthrough focused exclusively on the core attention calculation step: $softmax (Q K^{T}) V$ . This is the specific operation that FlashAttention optimizes.

This appendix provides the "prequel" step that was omitted to derive the activation matrices.

The Two Sets of Q, K, V

It's crucial to distinguish between:

The Weight Matrices (W_Q, W_K, W_V): These are the trainable parameters learned during the training of the Transformer model. They are part of a standard Linear layer. Their job is to project the input token embeddings into the query, key, and value spaces. They are the same for every input sequence.
The Activation Matrices (Q, K, V): These are the intermediate representations or activations. They are the result of multiplying the input embeddings by the weight matrices. These matrices are different for every input sequence and are the actual inputs to the attention calculation kernel that FlashAttention replaces. They are not trainable parameters themselves.

Think of it like a recipe:

The weight matrices (W_Q, W_K, W_V) are the instructions in the recipe book (fixed, learned).
The activation matrices (Q, K, V) are the actual ingredients you've prepared for one specific meal (changes every time you cook).

FlashAttention's innovation is in how to efficiently combine the prepared ingredients (Q, K, V), not in the initial preparation step itself.

The Omitted "Prequel" Step: Creating Q, K, and V

Here is the step that happens before our walkthrough begins.

Let's assume we have an input sequence of 6 tokens, and each token has an embedding dimension of 4. This is our input matrix X (not to be confused with the logit vector X from the walkthrough).

Input Embeddings X (size 6x4): $X = ⎣ ⎡ embedding for token 1 embedding for token 2 ⋮ embedding for token 6 ⎦ ⎤$

Now, let's define our trainable weight matrices. Let's say the head dimension d is 2. The weight matrices must project from the embedding dimension (4) to the head dimension (2). So, they will all be size 4x2.

Weight Matrices W_Q, W_K, W_V (size 4x2, trainable): $W_{Q} = [\dots], W_{K} = [\dots], W_{V} = [\dots]$

The Q, K, and V activation matrices are created with standard matrix multiplication:

Q K V = X \cdot W_{Q} (results in a 6 \times 2 matrix) = X \cdot W_{K} (results in a 6 \times 2 matrix) = X \cdot W_{V} (results in a 6 \times 2 matrix)

The V matrix we get from this calculation is precisely the V matrix we used in the walkthrough:

V = ⎣ ⎡ 123456123456 ⎦ ⎤

Connecting to the Logits

In our walkthrough, we focused on a single query vector q attending to all the keys. This corresponds to taking the first row of the Q matrix.

q = Q [0, :] (the first row, size 1 \times 2)

The logit vector X from the walkthrough would then be calculated as:

Logits = q \cdot K^{T}

This multiplication would result in the 1 x 6 vector we started with:

Logits = [123621]

Summary

So, the full, un-omitted process is:

Linear Projections (The Omitted Prequel):
- Start with input embeddings X.
- Compute $Q = X \cdot W_{Q}$ , $K = X \cdot W_{K}$ , $V = X \cdot W_{V}$ using the trainable weight matrices. This is done with standard, highly optimized matrix multiplication libraries (GEMM).
FlashAttention Calculation (The Walkthrough):
- Take the resulting Q, K, and V activation matrices as input.
- Efficiently compute $O = softmax (Q K^{T}) V$ in a single kernel without materializing the full attention matrix.

Appendix B - The Paper's pseudocode vs. this walkthrough

The goal of my walkthrough was to make the calculation as intuitive as possible to follow by hand. To do this, I slightly rearranged the math to keep the intermediate numbers as simple as possible, while still being mathematically equivalent to the paper's algorithm.

1. The Paper's "Algorithm FlashAttention (Tiling)"

The algorithm in the PDF maintains a normalized output vector o' at every step. Let's look closely at the update rule for the output:

o_{i'} \leftarrow o_{i - 1'} \frac{d _{i - 1'} e ^{m_{i - 1} - m_{i}}}{d _{i'}} + \frac{\sum _{j = 1 b} e ^{x_{i} [j] - m_{i}} V [ j + ( i - 1 ) b , : ]}{d _{i'}}

Notice that both terms are divided by $d_{i'}$ , the new total denominator. This means that at the end of each iteration i, the vector o'_i is the correctly normalized attention output for all the data processed up to that tile.

2. My Walkthrough's Method

Doing this normalization (division) at every single step can be cumbersome for a manual example. It's easier to work with un-normalized sums and perform a single division at the very end.

So, this walkthrough maintains an un-normalized running output o_running. Our update rule was:

o_{running} = o_{old} \cdot e^{m_{old} - m_{new}} + o_{local}

This is mathematically equivalent to the numerator of the paper's formula. You can see this if you multiply the paper's entire update rule by $d_{i'}$ :

o_{i'} \cdot d_{i'} o_{i'} \cdot d_{i'} = (o_{i - 1'} \frac{d _{i - 1'} e ^{m_{i - 1} - m_{i}}}{d _{i'}} + \frac{\sum \dots V [ \dots ]}{d _{i'}}) \cdot d_{i'} = (o_{i - 1'} d_{i - 1'}) e^{m_{i - 1} - m_{i}} + \sum \dots V [\dots]

If we define the $o_{running}$ as the paper's $o_{i'} \cdot d_{i'}$ , then this equation is exactly the one used in the walkthrough:

My $o_{new}$ is $o_{i'} \cdot d_{i'}$
My $o_{old}$ is $o_{i - 1'} \cdot d_{i - 1'}$
My $o_{local}$ is $\sum \dots V [\dots]$

The note was an attempt to explain this simplification: we chose to track the un-normalized numerator throughout the process for clarity and then perform the final division $o_{final} / d_{final}$ only once.

Why normalizing only at the last step works

The reason this works is that the update rules for the un-normalized numerator (o_running) and the denominator (d_running) are designed to maintain a consistent relationship. At every step i, the correctly normalized output is simply the ratio of the running numerator to the running denominator at that step. By deferring the division to the end, we arrive at the same final ratio. The proof by induction below demonstrates this formally.

The Two Methods

Let's formally define the two methods we are comparing. We will use the paper's notation where $o_{i'}$ is the normalized output after tile i, and we'll introduce $O_{i}$ (capital O) as our un-normalized running output numerator from the walkthrough.

Method 1: Normalize at Each Step (The Paper's Algorithm)
The state after tile i is defined by:

$d_{i'} = d_{i - 1'} e^{m_{i - 1} - m_{i}} + ⎝ ⎛ j \in Tile_{i} \sum e^{x_{j} - m_{i}} ⎠ ⎞$
$o_{i'} = o_{i - 1'} \frac{d _{i - 1'} e ^{m_{i - 1} - m_{i}}}{d _{i'}} + \frac{\sum _{j \in Tile_{i}} e ^{x_{j} - m_{i}} V _{j}}{d _{i'}}$

The final result is $o_{N'}$ after the last tile N.

Method 2: Normalize Only at the End (The Walkthrough's Method)
The state after tile i is defined by an un-normalized numerator $O_{i}$ and the same denominator $d_{i'}$ :

$d_{i'} = d_{i - 1'} e^{m_{i - 1} - m_{i}} + ⎝ ⎛ j \in Tile_{i} \sum e^{x_{j} - m_{i}} ⎠ ⎞$
$O_{i} = O_{i - 1} e^{m_{i - 1} - m_{i}} + j \in Tile_{i} \sum e^{x_{j} - m_{i}} V_{j}$

The final result is calculated as $O_{N} / d_{N'}$ at the very end.

The Proof of Equivalence

We want to prove that $o_{N'} = O_{N} / d_{N'}$ . We can prove this by induction, showing that the relationship $o_{i'} = O_{i} / d_{i'}$ holds true for every step i.

1. Base Case (i=1)

Let's check the first tile. Both methods start with $d_{0'} = 0$ , $o_{0'} = 0$ , and $O_{0} = 0$ .

Method 1:

$o_{1'} = o_{0'} \cdot (\dots) 0 + \frac{\sum j \in T _{1} e ^{x_{j} - m_{1}} V _{j}}{d _{1'}} = \frac{\sum _{j \in T_{1}} e ^{x_{j} - m_{1}} V _{j}}{d _{1'}}$
Method 2:

$O_{1} = O_{0} \cdot (\dots) 0 + \sum j \in T_{1} e^{x_{j} - m_{1}} V_{j} = j \in T_{1} \sum e^{x_{j} - m_{1}} V_{j}$ The final result would be $O_{1} / d_{1'}$ .

Comparing the two, we see they are identical. The base case holds.

2. Inductive Hypothesis

Assume that the relationship is true for step i-1. That is, assume:

o_{i - 1'} = \frac{O _{i - 1}}{d _{i - 1'}} ⟹ O_{i - 1} = o_{i - 1'} d_{i - 1'}

3. Inductive Step

Now we must prove that the relationship holds for step i. Let's start with the formula for $o_{i'}$ from Method 1 and show it equals $O_{i} / d_{i'}$ .

Start with the definition of $o_{i'}$ :

o_{i'} = o_{i - 1'} \frac{d _{i - 1'} e ^{m_{i - 1} - m_{i}}}{d _{i'}} + \frac{\sum _{j \in T_{i}} e ^{x_{j} - m_{i}} V _{j}}{d _{i'}}

Let's combine the two fractions over the common denominator $d_{i'}$ :

o_{i'} = \frac{( o _{i - 1'} d _{i - 1'} ) e ^{m_{i - 1} - m_{i}} + \sum _{j \in T_{i}} e ^{x_{j} - m_{i}} V _{j}}{d _{i'}}

Now, look at the term in the parentheses: $(o_{i - 1'} d_{i - 1'})$ . According to our Inductive Hypothesis, this is exactly equal to $O_{i - 1}$ . Let's substitute it in:

o_{i'} = \frac{O _{i - 1} e ^{m_{i - 1} - m_{i}} + \sum _{j \in T_{i}} e ^{x_{j} - m_{i}} V _{j}}{d _{i'}}

Now, look at the entire numerator: $O_{i - 1} e^{m_{i - 1} - m_{i}} + \sum_{j \in T_{i}} e^{x_{j} - m_{i}} V_{j}$ . This is precisely the definition of $O_{i}$ from Method 2.

So, we can substitute $O_{i}$ for the numerator:

o_{i'} = \frac{O _{i}}{d _{i'}}

This completes the proof. We have shown that if the relationship holds for step i-1, it must hold for step i. Since it holds for the base case i=1, it holds for all steps.

Intuitive Analogy: Calculating a Weighted Average

Think about calculating the final grade in a class. You have different assignments with different weights (scores).

Method 1 (Normalize at Each Step): After you get your first grade, you calculate your current average. When you get your second grade, you update your average based on the new grade and its weight relative to the first. You keep re-calculating your "running average" after every single assignment.
Method 2 (Normalize at the End): You collect all your points scored for each assignment (score * weight) in one column. You collect the total possible points (sum of all weights) in another column. You do this for the whole semester. At the very end, you do one single division: (total points scored) / (total possible points).

Both methods give you the exact same final grade. The second method simply defers the division. We can apply the same principle in FlashAttention: accumulate the un-normalized numerator (o_running) and the un-normalized denominator (d_running) separately and effectively perform the division at the end.

Online softmax by hand

Lewis Won — Thu, 28 Aug 2025 07:00:53 +0000

Table of Contents

Why study online softmax? The key to flash attention
Conceptual overview: tiled "safe" softmax
Conceptual Overview: Online Softmax
Mathematical conversion from element-wise to tile-wise formula
Appendix A: A simple example of numerical stability
Appendix B: Proof of correctness of online softmax

Why study online softmax? The key to flash attention

Softmax normalization is critical in the self-attention mechanism of transformers because it helps to manage "extreme values and offers more favorable gradient properties during training" (source: Building LLM from Scratch, chapter 3.3.1). For numerical stability, a "safe" softmax is widely used where in order to compute the final probability for any single element in every row of a matrix, we need to calculate two "global" statistics from the entirety of each row:
a. The maximum value of the row, $m = max (S_{i})$ .
b. The normalization denominator, $d = \sum_{j} exp (S_{ij} - m)$ , which is the sum over all elements in the row.

(Note: Please refer to Appendix A for an illustration of the importance of the absolute maximum value (max(x)) for numerical stability.)

With a naive implementation of "safe" softmax, the two global statistics could potentially force a "multi-pass" approach where for each row, the following needs to be computed sequentially:

Read the entire row to find the max.
Calculate the denominator.
Compute the softmax values.

With online softmax, steps 1 (find the max) and 2 (calculate the denominator) can be carried out within the same read of the matrix, instead of two sequential reads. Online softmax achieves this by calculating the necessary global statistics in a single pass over the data. This contrasts with a naive tiled approach that requires two passes to find the statistics, followed by a third pass for the final calculation. By reducing the number of memory passes from three to two, online softmax offers significant efficiency gains, especially since matrix operations in GPUs tend to be memory-bound.

This article provides an intuitive understanding of online softmax by demonstrating by hand how both the max value and denominator of the "safe" softmax can be calculated in a single pass. This walkthrough is based on the article Online normalizer calculation for softmax by Maxim Milakov and Natalia Gimelshein, 2018.

Online softmax also provides the basis for Flash Attention to fuse the entire attention calculation into a single GPU kernel, which I will discuss in my next article.

Conceptual overview: tiled "safe" softmax

The "safe" softmax is inherently a multi-pass algorithm. This is because calculating the final probability for any single element requires knowing two global properties of the entire input vector:

The absolute maximum value (max(x)) for numerical stability.
The sum of all exponentiated values, which serves as the normalization denominator.

Tiling is a technique used to break down large matrices or vectors into smaller blocks, or "tiles". This is crucial for efficiency, especially on hardware like GPUs, as it allows smaller chunks of data to be loaded into fast, local memory (SRAM) for processing, rather than repeatedly accessing slower global memory (DRAM).

When applying tiling to the normal softmax, we process the input tile-by-tile. However, because we need the global maximum and the global sum, we cannot complete the calculation for any tile in a single go. This results in a multi-pass approach over the tiled data. (Note: The algorithm below describes an element-wise process, I will show its relationship to a tile-based approach.)

Let's illustrate this with a simple example. Consider the following 1x6 input vector (which can be thought of as a single row in a larger matrix):

X = [123621]

We will process this vector with a tile size of 3. This breaks our input X into two tiles:

T_{1} = [123]

T_{2} = [621]

The process follows the logic of "Algorithm 2: Safe softmax" from the image, but adapted for tiles.

Step 1: First Pass — Find the Global Maximum

In the first pass, we iterate through each tile to find its local maximum. Then, we find the maximum among all local maximums to determine the global maximum.

1a. Find Local Maximums

For each tile, we compute the maximum value, which we'll call m_local.

For Tile 1 ( $T_{1}$ ):

$m_{1} = max ([123]) = 3$
For Tile 2 ( $T_{2}$ ):

$m_{2} = max ([621]) = 6$

1b. Find Global Maximum

Now, we find the maximum of the local maximums to get the global maximum, m_global.

m_{g l o ba l} = max (m_{1}, m_{2}) = max (3, 6) = 6

This value, m_global, corresponds to m_V in the provided "Algorithm 2".

Link back to the safe softmax algorithm

Line 1: m_0 ← -∞

Explanation: This initializes the variable that will hold the maximum value. It's set to negative infinity to ensure the first element of the vector x_1 becomes the first maximum.
In the Example: Before "Pass 1," we start with a conceptual maximum of $- \infty$ .

Lines 2-4: for k ← 1, V do ... m_k ← max(m_{k-1}, x_k)

Explanation: This is the first pass. The loop iterates through every element of the input vector X to find the single largest value. After the loop finishes, m_V holds the global maximum.
In the Example: This corresponds to our "Step 1: First Pass".
- We first found the local max for Tile 1 ( $m_{1} = 3$ ) and Tile 2 ( $m_{2} = 6$ ).
- We then found the maximum of these, yielding the global max m_global = 6. This is the final value of m_V after this loop completes.

Step 2: Second Pass — Calculate the Global Denominator

In the second pass, we again iterate through each tile. This time, we use the m_global to calculate a local sum of exponentials for each tile. These local sums are then aggregated to get the global denominator.

The formula for each element within a tile is $e^{(x_{j} - m_{g l o ba l})}$ .

2a. Calculate Local Denominators

For Tile 1 ( $T_{1} = [123]$ ):

d_{1} = e^{(1 - 6)} + e^{(2 - 6)} + e^{(3 - 6)} = e^{- 5} + e^{- 4} + e^{- 3} \approx 0.006738 + 0.018316 + 0.049787 = 0.074841

For Tile 2 ( $T_{2} = [621]$ ):

d_{2} = e^{(6 - 6)} + e^{(2 - 6)} + e^{(1 - 6)} = e^{0} + e^{- 4} + e^{- 5} \approx 1 + 0.018316 + 0.006738 = 1.025054

2b. Calculate Global Denominator

The global denominator, d_global, is the sum of all local denominators. This corresponds to d_V in the algorithm.

d_{g l o ba l} = d_{1} + d_{2} \approx 0.074841 + 1.025054 = 1.099895

Link back to the safe softmax algorithm

Line 5: d_0 ← 0

Explanation: This initializes the variable for the denominator (the sum of exponentials).
In the Example: Before "Pass 2," our sum is zero.

Lines 6-8: for j ← 1, V do ... d_j ← d_{j-1} + e^(x_j - m_V)

Explanation: This is the second pass. The loop iterates through every element x_j again. For each element, it subtracts the global maximum m_V (found in the first pass), exponentiates the result, and adds it to the running sum. After the loop, d_V holds the global denominator.
In the Example: This corresponds to our "Step 2: Second Pass".
- For Tile 1, we calculated $d_{1} = e^{(1 - 6)} + e^{(2 - 6)} + e^{(3 - 6)} \approx 0.0748$ .
- For Tile 2, we calculated $d_{2} = e^{(6 - 6)} + e^{(2 - 6)} + e^{(1 - 6)} \approx 1.0251$ .
- The loop's final result, d_V, is the sum $d_{1} + d_{2} \approx 1.0999$ .

Step 3: Third Pass — Calculate Final Softmax Probabilities

In the final pass, we iterate through the tiles one last time. For each element x_i, we compute its exponentiated value (scaled by m_global) and then divide by the d_global to get the final softmax probability y_i.

The formula is:

y_{i} = \frac{e ^{(x_{i} - m_{g l o ba l})}}{d _{g l o ba l}}

3a. Compute Softmax for Each Tile

For Tile 1 ( $T_{1} = [123]$ ):
- $y_{1} = \frac{e ^{(1 - 6)}}{1.099895} = \frac{e ^{- 5}}{1.099895} \approx \frac{0.006738}{1.099895} \approx 0.006126$
- $y_{2} = \frac{e ^{(2 - 6)}}{1.099895} = \frac{e ^{- 4}}{1.099895} \approx \frac{0.018316}{1.099895} \approx 0.016652$
- $y_{3} = \frac{e ^{(3 - 6)}}{1.099895} = \frac{e ^{- 3}}{1.099895} \approx \frac{0.049787}{1.099895} \approx 0.045266$
For Tile 2 ( $T_{2} = [621]$ ):
- $y_{4} = \frac{e ^{(6 - 6)}}{1.099895} = \frac{e ^{0}}{1.099895} \approx \frac{1}{1.099895} \approx 0.909177$
- $y_{5} = \frac{e ^{(2 - 6)}}{1.099895} = \frac{e ^{- 4}}{1.099895} \approx \frac{0.018316}{1.099895} \approx 0.016652$
- $y_{6} = \frac{e ^{(1 - 6)}}{1.099895} = \frac{e ^{- 5}}{1.099895} \approx \frac{0.006738}{1.099895} \approx 0.006126$

Link back to the safe softmax algorithm

Lines 9-11: for i ← 1, V do ... y_i ← e^(x_i - m_V) / d_V

Explanation: This is the third pass. This final loop computes the softmax probability y_i for each element. It re-calculates the exponentiated value for x_i (just as in the second pass) and then divides it by the global denominator d_V.
In the Example: This corresponds to our "Step 3: Third Pass".
- $y_{1} = e^{(1 - 6)} /1.0999 \approx 0.0061$
- $y_{2} = e^{(2 - 6)} /1.0999 \approx 0.0167$
- ...and so on for all six elements.

Final Result

Combining the results from all tiles gives the final softmax output vector Y:

Y \approx [0.0061 0.0167 0.0453 0.9092 0.0167 0.0061]

The sum of these probabilities is approximately 1.0, as expected. This multi-pass process is the standard, "normal" way of computing softmax in a tiled fashion, and stands in contrast to the "online softmax" which is designed to reduce the number of passes over the data.

Conceptual Overview: Online Softmax

The "Online Softmax," as described in the paper by Milakov and Gimelshein, is a clever single-pass algorithm. Its key innovation is the ability to calculate the softmax probabilities tile-by-tile without needing to first compute the global maximum and global sum over the entire input. It achieves this by maintaining a running maximum and a running denominator. When a new tile is processed, these running statistics are updated.

The core of the method lies in how the running denominator is updated when a new, larger maximum value is found. It uses a "telescoping sum" property to rescale the previous sum, effectively correcting it to be consistent with the new maximum. This avoids the need for a second or third pass over the data, which is a significant advantage for memory-bound operations on hardware like GPUs. Proof that online softmax is correct is in Appendix B. Note that the original proof was found in Online normalizer calculation for softmax, pages 3 - 4, but I instead referenced Yi Wang's FlashAttention (Part 2): Online Softmax as I found his writing to be more intuitive. (Note: The algorithm below describes an element-wise process, I will show its relationship to a tile-based approach.)

Let's use the same example vector and tile size:

X = [123621]

With a tile size of 3, we have two tiles:

T_{1} = [123]

T_{2} = [621]

The process will follow the logic of "Algorithm 3: Safe softmax with online normalizer calculation". We will process the vector tile by tile, from left to right, maintaining and updating two running values:

$m_{r u nnin g}$ : The running maximum found so far.
$d_{r u nnin g}$ : The running denominator (the sum of exponentials) scaled by the current $m_{r u nnin g}$ .

Initialization

Before processing the first tile, we initialize our running statistics as per the algorithm:

$m_{0} = m_{r u nnin g} = - \infty$
$d_{0} = d_{r u nnin g} = 0$

These are represented in lines 1 and 2 of the online softmax algorithm.

Step 1: Process Tile 1 ( $T_{1} = [123]$ )

We now process the first tile. For this tile, we will compute a new maximum ( $m_{n e w}$ ) and a new denominator ( $d_{n e w}$ ).

1a. Find the New Maximum for the Tile

First, find the maximum value within the current tile and compare it to our running maximum.

Local max of $T_{1}$ :

m_{T_{1}} = max ([123]) = 3

The new overall maximum is:

m_{n e w} = max (m_{r u nnin g}, m_{T_{1}}) = max (- \infty, 3) = 3

1b. Update the Running Denominator

Next, we update the running denominator. The formula from Algorithm 3 (line 5) combines the previous denominator with the contribution from the current tile.

d_{n e w} \leftarrow d_{o l d} \times e^{(m_{o l d} - m_{n e w})} + x_{j} \in current tile \sum e^{(x_{j} - m_{n e w})}

Let's apply this:

$m_{o l d} = m_{r u nnin g} = - \infty$
$d_{o l d} = d_{r u nnin g} = 0$
$m_{n e w} = 3$

d_{n e w} = (0 \times e^{(- \infty - 3)}) + (e^{(1 - 3)} + e^{(2 - 3)} + e^{(3 - 3)}) = 0 + (e^{- 2} + e^{- 1} + e^{0}) \approx 0.1353 + 0.3679 + 1 = 1.5032

1c. Update Running Statistics

After processing the first tile, our running statistics are:

$m_{r u nnin g} = 3$
$d_{r u nnin g} \approx 1.5032$

At this point, we could calculate intermediate (and incorrect) softmax values for the first tile using these stats, but the key insight of the online method is to wait until all tiles are processed. The final values will need to be rescaled.

Step 1: Process Tile 2 ( $T_{2} = [621]$ )

Now we move to the second tile and repeat the process, using the statistics from the previous step as our "old" values.

1d. Find the New Maximum

Local max of $T_{2}$ :

m_{T_{2}} = max ([621]) = 6

The new overall maximum is: $m_{n e w} = max (m_{r u nnin g}, m_{T_{2}}) = max (3, 6) = 6$

We have found a new, larger maximum. This will trigger the rescaling part of the update formula.

1e. Update the Running Denominator

We use the same formula as before, but now our old values are the results from Tile 1.

$m_{o l d} = 3$
$d_{o l d} \approx 1.5032$
$m_{n e w} = 6$

d_{n e w} = (d_{o l d} \times e^{(m_{o l d} - m_{n e w})}) + (x_{j} \in T_{2} \sum e^{(x_{j} - m_{n e w})}) \approx (1.5032 \times e^{(3 - 6)}) + (e^{(6 - 6)} + e^{(2 - 6)} + e^{(1 - 6)}) \approx (1.5032 \times e^{- 3}) + (e^{0} + e^{- 4} + e^{- 5}) \approx (1.5032 \times 0.04979) + (1 + 0.01832 + 0.00674) \approx 0.07484 + 1.02506 = 1.0999

The first term, $1.5032 \times e^{(3 - 6)}$ , is the crucial rescaling step. It takes the previous sum, which was calculated relative to the old maximum of 3, and correctly scales it to be relative to the new maximum of 6.

1f. Update Final Running Statistics

After processing all tiles, our final global statistics are:

$m_{f ina l} = 6$
$d_{f ina l} \approx 1.0999$

These values match the global maximum and denominator we found in the multi-pass "normal" softmax example.

Link back to the online softmax algorithm

Lines 3-6: for j ← 1, V do ...

Explanation: This is the main single pass of the algorithm. It iterates through the input vector element by element, updating the running max and denominator at each step.
In the Example: We processed this tile-by-tile, which is a block-wise version of this element-wise loop.

Line 4: $m_{j} \leftarrow ma x (m_{j - 1}, x_{j})$

Explanation: At each step j, this line updates the running maximum.
In the Example:
- When processing Tile 1 ( $x_{1}, x_{2}, x_{3}$ ), the running max m_j becomes 1, then 2, then finally 3. After Tile 1, our example had m_running = 3.
- When processing $x_{4} = 6$ , the max max(3, 6) becomes 6. Our example showed this change when moving from Tile 1 to Tile 2, where m_new became 6.

Line 5: $d_{j} \leftarrow d_{j - 1} \times e^{(m_{j - 1} - m_{j})} + e^{(x_{j} - m_{j})}$

Explanation: This is the core of the online update.
- The second term, $e^{(x_{j} - m_{j})}$ , adds the contribution of the current element, scaled by the new running max.
- The first term, $d_{j - 1} * e^{(m_{j - 1} - m_{j})}$ , is the crucial rescaling factor. If a new maximum is found ( $m_{j} > m_{j - 1}$ ), this term correctly scales down the entire previous sum to be relative to the new, larger maximum. If the maximum doesn't change ( $m_{j} = m_{j - 1}$ ), this term simply becomes $d_{j - 1} \times e^{0} = d_{j - 1}$ , so we just keep adding to the sum.
In the Example:
- For Tile 1, the max was always updating, but $d_{j - 1}$ started at 0. The final sum for the tile was $d_{1} = e^{(1 - 3)} + e^{(2 - 3)} + e^{(3 - 3)} \approx 1.5032$ .
- The most important step was processing Tile 2. Our running denominator was d_old ≈ 1.5032 and m_old = 3. The new max became m_new = 6. The update $(1.5032 \times e^{(3 - 6)})$ is the direct application of $d_{j - 1} \times e^{(m_{j - 1} - m_{j})}$ , rescaling the entire sum from the first tile. We then added the contributions from Tile 2's elements.

Mathematical conversion from element-wise to tile-wise formula

You may notice that the online softmax pseudocode does not have a summation Σ which appears in the illustrative calculations. This is because the pseudocode describes the update rule on a per-element basis. The summation Σ appears in the example because we are performing a per-tile update. The summation is simply the result of applying the single-element rule to every element within the tile at the same time. It is a practical optimization that is mathematically equivalent.

Let's see how the tile-wise formula is derived directly from the element-wise formula.

The core element-wise update rule is:

d_{j} = d_{j - 1} \cdot e^{m_{j - 1} - m_{j}} + e^{x_{j} - m_{j}}

Imagine we have just finished processing Tile 1. Our running statistics are $m_{o l d}$ and $d_{o l d}$ . Now we want to process Tile 2, which contains the elements $x_{a}, x_{b}, x_{c}$ .

Instead of processing the whole tile at once, let's apply the element-wise rule three times in a row.

1. Processing the first element of the tile, $x_{a}$ :

m_{n e w_{a}} = max (m_{o l d}, x_{a})

d_{n e w_{a}} = d_{o l d} \cdot e^{m_{o l d} - m_{n e w_{a}}} + e^{x_{a} - m_{n e w_{a}}}

2. Processing the second element, $x_{b}$ (using the results from step 1):

m_{n e w_{b}} = max (m_{n e w_{a}}, x_{b})

d_{n e w_{b}} = d_{n e w_{a}} \cdot e^{m_{n e w_{a}} - m_{n e w_{b}}} + e^{x_{b} - m_{n e w_{b}}}

Let's substitute the expression for $d_{n e w_{a}}$ into this equation:

d_{n e w_{b}} = (d_{o l d} \cdot e^{m_{o l d} - m_{n e w_{a}}} + e^{x_{a} - m_{n e w_{a}}}) \cdot e^{m_{n e w_{a}} - m_{n e w_{b}}} + e^{x_{b} - m_{n e w_{b}}}

d_{n e w_{b}} = d_{o l d} \cdot e^{m_{o l d} - m_{n e w_{b}}} + e^{x_{a} - m_{n e w_{b}}} + e^{x_{b} - m_{n e w_{b}}}

Notice the pattern: the original $d_{o l d}$ is now scaled by the newest max, and the contributions from the elements processed so far ( $x_{a}, x_{b}$ ) are summed up, also scaled by the newest max.

3. Processing the third element, $x_{c}$ (using the results from step 2):

m_{n e w_{c}} = max (m_{n e w_{b}}, x_{c})

d_{n e w_{c}} = d_{n e w_{b}} \cdot e^{m_{n e w_{b}} - m_{n e w_{c}}} + e^{x_{c} - m_{n e w_{c}}}

Substituting the expression for $d_{n e w_{b}}$ :

d_{n e w_{c}} = (d_{o l d} \cdot e^{m_{o l d} - m_{n e w_{b}}} + e^{x_{a} - m_{n e w_{b}}} + e^{x_{b} - m_{n e w_{b}}}) \cdot e^{m_{n e w_{b}} - m_{n e w_{c}}} + e^{x_{c} - m_{n e w_{c}}}

d_{n e w_{c}} = d_{o l d} \cdot e^{m_{o l d} - m_{n e w_{c}}} + e^{x_{a} - m_{n e w_{c}}} + e^{x_{b} - m_{n e w_{c}}} + e^{x_{c} - m_{n e w_{c}}}

The Tile-Wise Optimization

In a practical implementation on parallel hardware (like a GPU), it's far more efficient to do the following:

Load the entire tile $x_{a}, x_{b}, x_{c}$ into fast local memory.
Find the local maximum for the tile: $m_{t i l e} = max (x_{a}, x_{b}, x_{c})$ .
Calculate the new overall maximum in one step: $m_{n e w} = max (m_{o l d}, m_{t i l e})$ . Note that this final m_new is mathematically identical to the m_new_c we found after processing the last element in the step-by-step derivation above.
Apply the update rule in one go.

Looking at our final equation from the element-wise derivation:

d_{n e w_{c}} = d_{o l d} \cdot e^{m_{o l d} - m_{n e w_{c}}} + (e^{x_{a} - m_{n e w_{c}}} + e^{x_{b} - m_{n e w_{c}}} + e^{x_{c} - m_{n e w_{c}}})

We can rewrite the part in the parentheses using a summation:

x_{j} \in Tile \sum e^{x_{j} - m_{n e w_{c}}}

If we substitute $m_{n e w_{c}}$ with our tile-wise $m_{n e w}$ , we get the exact formula used in the example:

d_{n e w} = d_{o l d} \cdot e^{m_{o l d} - m_{n e w}} + x_{j} \in Tile \sum e^{x_{j} - m_{n e w}}

Step 2: Final Pass — Calculate Final Softmax Probabilities

Because the online algorithm only makes a single pass to find the statistics of max and sum, a final pass is still required to compute the output probabilities using the final global statistics.

The formula is the same as in the normal softmax: $y_{i} = \frac{e ^{(x_{i} - m_{f ina l})}}{d _{f ina l}}$

For Tile 1 ( $T_{1} = [123]$ ):
- $y_{1} = \frac{e ^{(1 - 6)}}{1.0999} = \frac{e ^{- 5}}{1.0999} \approx 0.0061$
- $y_{2} = \frac{e ^{(2 - 6)}}{1.0999} = \frac{e ^{- 4}}{1.0999} \approx 0.0167$
- $y_{3} = \frac{e ^{(3 - 6)}}{1.0999} = \frac{e ^{- 3}}{1.0999} \approx 0.0453$
For Tile 2 ( $T_{2} = [621]$ ):
- $y_{4} = \frac{e ^{(6 - 6)}}{1.0999} = \frac{e ^{0}}{1.0999} \approx 0.9092$
- $y_{5} = \frac{e ^{(2 - 6)}}{1.0999} = \frac{e ^{- 4}}{1.0999} \approx 0.0167$
- $y_{6} = \frac{e ^{(1 - 6)}}{1.0999} = \frac{e ^{- 5}}{1.0999} \approx 0.0061$

Link back to the online softmax algorithm

Lines 7-9: for i ← 1, V do ... y_i ← e^(x_i - m_V) / d_V

Explanation: This is a final pass to compute the results. After the main loop (lines 3-6) is complete, m_V and d_V hold the final, correct global statistics. This loop then uses those final values to calculate each y_i.
In the Example: This corresponds to our "Step 3: Final Pass". We took the final m_final = 6 and d_final ≈ 1.0999 and applied them to all original x_i values to get the final probabilities, which is exactly what this loop does.

Final Result

The final softmax vector is identical to the one produced by the normal softmax method:

Y \approx [0.0061 0.0167 0.0453 0.9092 0.0167 0.0061]

The critical advantage is that the online method calculated the necessary global statistics ( $m_{f ina l}$ and $d_{f ina l}$ ) in a single pass over the data, whereas the normal tiled method required two passes just to get the statistics, followed by a third pass for the final calculation. This reduction in memory passes provides for efficiency gains with online softmax.

Appendix A: A simple example of numerical stability

Let's illustrate why this stabilization is crucial with a simple vector, especially when using lower precision floating-point numbers like float16, which is common in modern GPUs for accelerating deep learning workloads. The maximum value for a float16 is 65,504.

Consider the input vector:

X = [2412]

Case 1: The Naive (Unstable) Approach

Using the direct definition of softmax, we compute $exp (x_{i})$ for each element:

$exp (2) \approx 7.389$
$exp (4) \approx 54.598$
$exp (12) \approx 162, 754.79$

Now, let's analyze this from a float16 perspective. With $X = [1212]$ , we proceed to calculate the denominator:
$d = \sum_{j = 1 N} exp (x_{j}) \approx 7.389 + 54.598 + 162, 754.79 = 162, 816.777$

As the maximum finite positive value for a float16 is approximately 65,604, the sum 162,816.777 exceeds the maximum value, and hence causes a catastrophic overflow in float16, resulting in inf (infinity). As a result, the final probabilities would be computed as:

$y_{1} = \frac{7.389}{\infty} = 0$
$y_{2} = \frac{54.598}{\infty} = 0$
$y_{3} = \frac{162 , 754.79}{\infty} = NaN$ (Not a Number, from $\frac{\infty}{\infty}$ )

The result is a vector of [0, 0, NaN], which is completely incorrect and would destroy any subsequent training or inference.

Case 2: The Safe (Stable) Approach

Now, let's use the numerically stable formula.

Step 1: Find the maximum value, m.

m = max ([1212]) = 12

Step 2: Subtract m from each element.

X - m = [1 - 12 2 - 12 12 - 12] = [- 11 - 10 0]

Step 3: Compute the exponentials of the new values.

$exp (- 11) \approx 0.0000167$
$exp (- 10) \approx 0.0000454$
$exp (0) = 1$

Notice that all these values are small, well-behaved numbers between 0 and 1. There is absolutely no risk of overflow.

Step 4: Calculate the denominator, d'.

d^{'} = j = 1 \sum N exp (x_{j} - m) \approx 0.0000167 + 0.0000454 + 1 = 1.0000621

Step 5: Compute the final probabilities.

$y_{1} = \frac{e x p ( - 11 )}{1.0000621} \approx \frac{0.0000167}{1.0000621} \approx 0.0000167$
$y_{2} = \frac{e x p ( - 10 )}{1.0000621} \approx \frac{0.0000454}{1.0000621} \approx 0.0000454$
$y_{3} = \frac{e x p ( 0 )}{1.0000621} = \frac{1}{1.0000621} \approx 0.9999379$

The resulting softmax vector is approximately $[0.000017 0.000045 0.999938]$ . The sum is 1.0, and we have avoided any numerical errors. This demonstrates that subtracting the maximum value is not just a theoretical trick—it is an essential step for making softmax work in practice.

Appendix B: Proof of correctness of online softmax

I have reproduced the proof by Yi Wang in FlashAttention (Part 2): Online Softmax below.

The definition of softmax is as follows:

{\frac{exp ( x _{i} )}{\sum _{j = 1 N} exp ( x _{j} )}}_{i = 1 N}

If any $x_{i}$ is large (e.g. $x_{i} \geq 11$ ), $exp (x_{i})$ exceeds the maximum value of float16. To address this numerical instability, we compute an alternative form which gives equivalent result but numerically stable:

{\frac{exp ( x _{i} - m )}{\sum _{j = 1 N} exp ( x _{j} - m )}}_{i = 1 N}

where $m = max_{j = 1 N} x_{j}$ . This form is safe because $x_{i} - m \leq 0$ , ensuring that $0 < exp (x_{i} - m) \leq 1$ .

As mentioned in this article, online softmax seeks to allow the calculation of the max value $m$ and the denominator $\sum_{j = 1 N} exp (x_{j} - m)$ in parallel. More concretely, for any integer $1 \leq i \leq N$ , we want to be able to:

Calculate an intermediate $δ_{i} = \sum_{j = 1 i} exp (x_{j} - m_{i})$ so that $δ_{N} = \sum_{j = 1 N} exp (x_{j} - m_{N})$
Since $δ_{i}$ is inductive, it should depend on $δ_{i - 1}$ .
To allow parallel execution, $δ_{i}$ must not depend on future values such as $x_{i + 1}, \dots$ or $m_{i + 1}, \dots$ .

We begin by considering:

δ_{i} = j = 1 \sum i exp (x_{j} - m_{i})

To ensure $δ_{i}$ depends on $δ_{i - 1}$ , which is:

δ_{i - 1} = j = 1 \sum i - 1 exp (x_{j} - m_{i - 1})

we need to split $δ_{i}$ into two parts: one involving $δ_{i - 1}$ (which should not depend on $x_{i}$ or $m_{i}$ ), and the remaining terms that depend on $x_{i}$ and $m_{i}$ . The first step is straightforward – we separate the last term in the summation:

δ_{i} = j = 1 \sum i - 1 exp (x_{j} - m_{i}) + exp (x_{i} - m_{i})

Now, $x_{i}$ only appears in the second term. However, $m_{i}$ still appears in the summation. Let's take the next step:

δ_{i} = j = 1 \sum i - 1 exp (x_{j} - m_{i - 1} + m_{i - 1} - m_{i}) + exp (x_{i} - m_{i}) = [j = 1 \sum i - 1 exp (x_{j} - m_{i - 1})] exp (m_{i - 1} - m_{i}) + exp (x_{i} - m_{i})

The expression inside the square brackets is exactly $δ_{i - 1}$ . Therefore, we have:

δ_{i} = δ_{i - 1} exp (m_{i - 1} - m_{i}) + exp (x_{i} - m_{i})

This allows us to compute $δ_{i}$ inductively in parallel with $m_{i}$ .

Tensor parallelism by hand

Lewis Won — Sat, 23 Aug 2025 11:18:13 +0000

Table of Contents

Motivation
What is tensor parallelism
When do we implement tensor parallelism
Setup
Forward pass
Why split Matrix A by columns, Matrix B by rows
Backward pass (backpropagation)
Communication cost of tensor parallelism
Code

Motivation

Tensor parallelism is widely used in training and deploying massive neural networks such as LLMs across multiple GPUs. While libraries such as DeepSpeed provide out-of-the-box solutions for tensor parallelism, understanding how to implement tensor parallelism by hand may help with:

Building intuition: instead of just knowing tensor parallelism splits tensors, you will understand how the choice of splitting by "columns" in one layer and "rows" in the next is essential for the math to work out.
Understanding performance bottlenecks: You will gain an understanding why tensor parallelism is communication intensive and generally recommended to be done over high-speed interconnects such as NVLink.
Debugging: An intuitive mental model of tensor parallelism may be useful for debugging when distributed training jobs fail and the error messages are cryptic.

This article can also serve as a complementary study guide for Stanford's CS336 Lecture 7 on Parallelism. This article was written with assistance of Gemini 2.5 Pro.

What is tensor parallelism

When a model becomes too massive to fit into a single GPU's memory, tensor parallelism offers a way to split up the workload. Tensor parallelism accomplishes two goals simultaneously:

Distributes Memory Load: It enables the combining of the memory of several GPUs to store a single, massive layer that would be too large for any one device.
Parallelizes Computation: Each GPU performs a smaller portion of the matrix multiplication using its local shard of the weight matrix. This allows the GPUs to work on the same forward or backward pass of a single layer at the same time, speeding up the computation.

When do we implement tensor parallelism

Tensor parallelism comes at a significant cost of communication overhead. This article will demonstrate why tensor parallelism requires a high-volume, synchronous communication step (All-Reduce) twice for every single block in the model—once during the forward pass and once during the backward pass. This creates a fundamental trade-off:

Pipeline Parallelism: In pipeline parallelism, you split groups of layers (stages) across GPUs. Communication is less frequent, happening only at the boundaries between stages. However, pipeline parallelism introduces a "bubble" of idle time at the beginning and end of each batch as the pipeline fills and drains, reducing hardware utilization.
Tensor Parallelism: In tensor parallelism, you split the work within each layer. This helps to keep all GPUs perfectly synchronized and busy, completely eliminating the pipeline "bubble." But the price for this high utilization is constant, layer-by-layer communication across all participating GPUs.

The intense communication pattern of tensor parallelism over pipeline parallelism leads to a rule of thumb for designing large-scale training setups:

Use Tensor Parallelism for Intra-Node Scaling: When you have a group of GPUs (e.g., the 8 A100s in a DGX node) connected with a high-speed backplane, tensor parallelism is the ideal way to scale up to handle massive layers.
Use Pipeline and Data Parallelism for Inter-Node Scaling: To scale beyond a single node, you typically combine tensor parallelism with pipeline parallelism, where each stage of the pipeline is a tensor-parallel group of GPUs. This minimizes the slow communication across nodes to only what is necessary between pipeline stages.

Setup

We will walk through the forward pass of a two-layer Multi-Layer Perceptron (MLP) block found in a Transformer, which consists of two main operations:

Y = GeLU(X * A)
Z = Dropout(Y * B)

We'll assume we have two GPUs, GPU 1 and GPU 2.

Source: Stanford CS 336 Language Modeling from Scratch, Lecture 7: Parallelism 1, 58 min 12 sec

Forward pass

We first begin with the forward pass, and the next section will walk through the backpropagation.

The forward pass in a neural network is the process of taking an input and propagating it through the layers of the model to generate a final output. In the context of tensor parallelism, this process is modified because the model's large weight matrices are too big for a single GPU and are instead sharded, or split, across multiple GPUs.

To execute a matrix multiplication, each GPU performs a smaller calculation in parallel using its local shard of the weights. This often follows a specific pattern where one layer's weights are split by columns (requiring no communication after the multiplication) and the next layer's weights are split by rows. This row-parallel approach produces partial results on each GPU that must then be aggregated through a collective communication operation, such as an All-Reduce, to ensure the final activation tensor is mathematically equivalent to the result that would have been computed on a single, monolithic GPU.

Step 1: Initial Setup - Defining and Splitting Matrices

First, let's define our input matrix X and the weight matrices A and B for our two linear layers. We will use small, simple numbers for clarity.

Let our input X be a 2x4 matrix:

X = (15263748)

Let the first weight matrix A be a 4x6 matrix:

A = ⎝ ⎛ 174128523963417452856396 ⎠ ⎞

Let the second weight matrix B be a 6x4 matrix:

B = ⎝ ⎛ 629415518326476237385148 ⎠ ⎞

Now, we split these matrices across our two GPUs.

Matrix A is split vertically (by columns). GPU 1 gets the first 3 columns (A1), and GPU 2 gets the last 3 columns (A2). This is known as column parallelism.

$A_{1} = ⎝ ⎛ 174128523963 ⎠ ⎞$ (on GPU 1)

$A_{2} = ⎝ ⎛ 417452856396 ⎠ ⎞$ (on GPU 2)

Matrix B is split horizontally (by rows). GPU 1 gets the first 3 rows (B1), and GPU 2 gets the last 3 rows (B2). This is known as row parallelism.

$B_{1} = ⎝ ⎛ 629518476385 ⎠ ⎞$ (on GPU 1)

$B_{2} = ⎝ ⎛ 415326237148 ⎠ ⎞$ (on GPU 2)

We split Matrix A by columns, and Matrix B by rows, so as to minimise communication costs incurred by the all-reduce operation. I will go into more detail at the next section on why we split Matrix A by columns, and Matrix B by rows.

Step 2: Forward Pass for `Y = GeLU(X * A)`

This operation is parallelized using the column-sharded matrix A. The input X is broadcast to both GPUs.

On GPU 1: Calculate `Y_intermediate_1 = X * A1`

We perform matrix multiplication. The output will be a 2x3 matrix.

Y_{in t er m e d ia t e_{1}} = (15263748) \times ⎝ ⎛ 174128523963 ⎠ ⎞

Calculating each element:

Y_intermediate_1 = (1*1) + (2*7) + (3*4) + (4*1) = 1 + 14 + 12 + 4 = 31
Y_intermediate_1 = (1*2) + (2*8) + (3*5) + (4*2) = 2 + 16 + 15 + 8 = 41
Y_intermediate_1 = (1*3) + (2*9) + (3*6) + (4*3) = 3 + 18 + 18 + 12 = 51
Y_intermediate_1 = (5*1) + (6*7) + (7*4) + (8*1) = 5 + 42 + 28 + 8 = 83
Y_intermediate_1 = (5*2) + (6*8) + (7*5) + (8*2) = 10 + 48 + 35 + 16 = 109
Y_intermediate_1 = (5*3) + (6*9) + (7*6) + (8*3) = 15 + 54 + 42 + 24 = 135

So, the result of the matrix multiplication on GPU 1 is:

Y_{in t er m e d ia t e_{1}} = (31834110951135)

On GPU 2: Calculate `Y_intermediate_2 = X * A2`

Similarly, we perform the multiplication on GPU 2.

Y_{in t er m e d ia t e_{2}} = (15263748) \times ⎝ ⎛ 417452856396 ⎠ ⎞

Calculating each element:

Y_intermediate_2 = (1*4) + (2*1) + (3*7) + (4*4) = 4 + 2 + 21 + 16 = 43
Y_intermediate_2 = (1*5) + (2*2) + (3*8) + (4*5) = 5 + 4 + 24 + 20 = 53
Y_intermediate_2 = (1*6) + (2*3) + (3*9) + (4*6) = 6 + 6 + 27 + 24 = 63
Y_intermediate_2 = (5*4) + (6*1) + (7*7) + (8*4) = 20 + 6 + 49 + 32 = 107
Y_intermediate_2 = (5*5) + (6*2) + (7*8) + (8*5) = 25 + 12 + 56 + 40 = 133
Y_intermediate_2 = (5*6) + (6*3) + (7*9) + (8*6) = 30 + 18 + 63 + 48 = 159

So, the result of the matrix multiplication on GPU 2 is:

Y_{in t er m e d ia t e_{2}} = (431075313363159)

Apply GeLU Activation Function

Next, the Gaussian Error Linear Unit (GeLU) activation function is applied element-wise to the intermediate results on each GPU. The GeLU function is defined as:

GE LU (x) = x \cdot Φ (x)

where $Φ (x)$ is the Cumulative Distribution Function (CDF) for the standard normal distribution. A common and precise formula is:

GE LU (x) = 0.5 \cdot x \cdot (1 + er f (x / 2))

For simplicity in this manual example, let's assume GeLU(x) ≈ x for large positive x and GeLU(x) ≈ 0 for large negative x. Since all our numbers are large and positive, we'll approximate GeLU(x) = x. In a real scenario, the exact formula would be applied.

(Note: The erf function produces an "S"-shaped sigmoid curve that passes through the origin. As the input x gets very large and positive, erf(x) gets very close to 1. As the input gets very large and negative, erf(x) gets very close to -1.)

So, the outputs on each GPU are:

$Y_{1} = G e LU (Y_{in t er m e d ia t e_{1}}) \approx (31834110951135)$ (on GPU 1)

$Y_{2} = G e LU (Y_{in t er m e d ia t e_{2}}) \approx (431075313363159)$ (on GPU 2)

These two matrices, Y1 and Y2, represent the sharded output of the first layer. They remain on their respective GPUs and become the input for the next step.

Step 3: Forward Pass for `Z = Dropout(Y * B)`

This operation is parallelized using the row-sharded matrix B. The inputs Y1 and Y2 are already on the correct GPUs.

On GPU 1: Calculate `Z_partial_1 = Y1 * B1`

We multiply the local activation Y1 with the local weight shard B1.

Z_{p a r t ia l_{1}} = (31834110951135) \times ⎝ ⎛ 629518476385 ⎠ ⎞

Calculating each element:

Z_partial_1 = (31*6) + (41*2) + (51*9) = 186 + 82 + 459 = 727
Z_partial_1 = (31*5) + (41*1) + (51*8) = 155 + 41 + 408 = 604
Z_partial_1 = (31*4) + (41*7) + (51*6) = 124 + 287 + 306 = 717
Z_partial_1 = (31*3) + (41*8) + (51*5) = 93 + 328 + 255 = 676
Z_partial_1 = (83*6) + (109*2) + (135*9) = 498 + 218 + 1215 = 1931
Z_partial_1 = (83*5) + (109*1) + (135*8) = 415 + 109 + 1080 = 1604
Z_partial_1 = (83*4) + (109*7) + (135*6) = 332 + 763 + 810 = 1905
Z_partial_1 = (83*3) + (109*8) + (135*5) = 249 + 872 + 675 = 1796

Result on GPU 1:

Z_{p a r t ia l_{1}} = (7271931604160471719056761796)

On GPU 2: Calculate `Z_partial_2 = Y2 * B2`

Simultaneously, on GPU 2:

Z_{p a r t ia l_{2}} = (431075313363159) \times ⎝ ⎛ 415326237148 ⎠ ⎞

Calculating each element:

Z_partial_2 = (43*4) + (53*1) + (63*5) = 172 + 53 + 315 = 540
Z_partial_2 = (43*3) + (53*2) + (63*6) = 129 + 106 + 378 = 613
Z_partial_2 = (43*2) + (53*3) + (63*7) = 86 + 159 + 441 = 686
Z_partial_2 = (43*1) + (53*4) + (63*8) = 43 + 212 + 504 = 759
Z_partial_2 = (107*4) + (133*1) + (159*5) = 428 + 133 + 795 = 1356
Z_partial_2 = (107*3) + (133*2) + (159*6) = 321 + 266 + 954 = 1541
Z_partial_2 = (107*2) + (133*3) + (159*7) = 214 + 399 + 1113 = 1726
Z_partial_2 = (107*1) + (133*4) + (159*8) = 107 + 532 + 1272 = 1911

Result on GPU 2:

Z_{p a r t ia l_{2}} = (5401356613154168617267591911)

All-Reduce Communication Step

The logic behind splitting B by rows is that the final output Z is the sum of the partial results from each GPU. This requires a communication step called an All-Reduce. Each GPU sends its partial result to the others and receives their partial results, summing them all together.

Z_{u na c t i v a t e d} = Z_{p a r t ia l_{1}} + Z_{p a r t ia l_{2}}

Z_{u na c t i v a t e d} = (7271931604160471719056761796) + (5401356613154168617267591911)

Adding each element:

Z_{u na c t i v a t e d} = (727 + 540 1931 + 1356 604 + 613 1604 + 1541 717 + 686 1905 + 1726 676 + 759 1796 + 1911)

Z_{u na c t i v a t e d} = (12673287121731451403363114353707)

After the All-Reduce, the complete Z_unactivated matrix is present on both GPUs.

Apply Dropout

Dropout is a regularization technique used during training to prevent overfitting. It randomly sets a fraction of the output elements to zero at each update step.

Let's assume a dropout rate of 50%. We create a "dropout mask" with the same dimensions as Z_unactivated, where each position is randomly 0 or 1.

Example dropout mask:

M a s k = (10010110)

We multiply our Z_unactivated by this mask:

Z_{ma s k e d} = Z_{u na c t i v a t e d} ⊙ M a s k = (1267 \cdot 1 3287 \cdot 0 1217 \cdot 0 3145 \cdot 1 1403 \cdot 0 3631 \cdot 1 1435 \cdot 1 3707 \cdot 0)

Z_{ma s k e d} = (12670031450363114350)

Finally, to ensure the expected sum of outputs remains the same, the remaining non-zero elements are scaled up by a factor of 1 / (1 - dropout_rate). In our case, this is 1 / (1 - 0.5) = 2. This is called "inverted dropout".

Z = Z_{ma s k e d} \times 2 = (25340062900726228700)

This final matrix Z is the output of our tensor-parallel multi-layer perceptron block.

Summary of the forward pass

Start: We have an input X and two GPUs.
Layer 1 (Column Parallelism):
- The weight matrix A is split into columns (A1, A2).
- Each GPU calculates a part of the output: X * A1 on GPU 1 and X * A2 on GPU 2.
- An activation function (GeLU) is applied independently on each GPU to these partial results (Y1, Y2). No communication is needed.
Layer 2 (Row Parallelism):
- The weight matrix B is split into rows (B1, B2).
- Each GPU multiplies its local activation from the previous step with its local weight shard: Y1 * B1 on GPU 1 and Y2 * B2 on GPU 2.
- An All-Reduce operation sums the partial results from all GPUs. This is the critical communication step.
- The final operation (Dropout) is applied to the complete, summed tensor, resulting in the final output Z.

Why split Matrix A by columns, Matrix B by rows

The goal is to sequence the operations to ensure the expensive All-Reduce step (where GPUs must synchronize and sum their results) happens as infrequently as possible. To understand this, let's define the two operations and their communication costs.

1. Column Parallelism (What we do to `A`)

When we split a weight matrix A by columns, the input X is broadcast to all GPUs. Each GPU then calculates a vertical "slice" of the output.

Mathematics: Y = X * [A1, A2] = [X * A1, X * A2] = [Y1, Y2]
Result: The output Y is naturally sharded (split) across the GPUs. Y1 is on GPU 1 and Y2 is on GPU 2.
Communication Cost (Forward Pass): Zero. No communication is needed to produce the sharded output Y. Each GPU works independently. The output is left "as is" on each device, ready for the next step.

2. Row Parallelism (What we do to `B`)

When we split a weight matrix B by rows, the input must already be sharded in a corresponding way (which is exactly what Column Parallelism just produced!). Each GPU multiplies its local input shard with its local weight shard.

Mathematics: Z = Y * B = [Y1, Y2] * [B1; B2] = (Y1 * B1) + (Y2 * B2) (Note: ; denotes stacking rows)
Result: Each GPU calculates a partial result. To get the final, correct output Z, these partial results must be summed together.
Communication Cost (Forward Pass): One All-Reduce operation. This is the expensive step where Z_partial_1 and Z_partial_2 are added across all GPUs.

Analyzing the Standard `Column -> Row` Pattern

Now, let's see how these two pieces fit together in our MLP block:

First Layer: Y = GeLU(X * A) (Column Parallelism)
- We split A by columns.
- The multiplication X * A produces a sharded activation [Y1, Y2] without any need for communication.
- Forward Communication Cost: 0
Second Layer: Z = Dropout(Y * B) (Row Parallelism)
- We split B by rows.
- The sharded activation [Y1, Y2] from the previous step is the perfect input for this layer.
- The multiplication Y * B produces partial results Z_partial_1 and Z_partial_2.
- To get the final output Z, we must perform an All-Reduce to sum the partial results.
- Forward Communication Cost: 1 All-Reduce

By sequencing the operations this way (Column -> Row), we have successfully computed the output for a full two-layer block with only one communication step.

What Happens If We Switch? (Row -> Column)

Let's imagine we tried to do it the other way around.

First Layer: Y = GeLU(X * A) (Row Parallelism)
- We split A by rows (A1, A2).
- For the math X * A to work, we would also have to split the input X by columns (X1, X2).
- The operation becomes (X1 * A1) + (X2 * A2).
- This produces partial results which must be summed.
- Forward Communication Cost: 1 All-Reduce. The output Y is now a complete (un-sharded) tensor, identical on all GPUs.
Second Layer: Z = Dropout(Y * B) (Column Parallelism)
- We split B by columns (B1, B2).
- The input Y is a complete tensor from the previous step.
- The operation Y * [B1, B2] produces a sharded output [Z1, Z2] without communication.
- Forward Communication Cost: 0

Why column then row?

As you can see, both patterns (Column -> Row and Row -> Column) result in the exact same amount of communication: one All-Reduce per two-layer block.

So why is the Column -> Row pattern the standard convention used in frameworks like Megatron-LM?

The reason is modularity and efficiency in the computational graph. The standard Transformer block consists of an Attention block followed by an MLP block. The Column -> Row pattern creates a perfect compositional unit:

The MLP block takes a complete tensor as input (like X).
It performs Column Parallel -> Row Parallel.
It outputs a complete tensor (Z, after the All-Reduce).

This complete tensor Z can then be fed directly into the next Attention block, which will also start with a column-parallel operation. This creates a clean, predictable, and efficient computational flow where the "sharded" nature of the activations is kept internal to the block, and communication is neatly packaged.

In short, the Column -> Row sequence is a deliberate design pattern that cleverly delays the necessary communication to the very last step within a block, creating a modular and efficient structure for building large neural networks.

Backward pass (backpropagation)

The goal of the backward pass is to calculate the gradients of the loss function with respect to all the learnable parameters (the weights A and B) and the input (X). These gradients tell us how to adjust the weights to reduce the loss.

The process starts with the gradient of the loss with respect to the final output, which we'll call dL/dZ or simply dZ. This initial gradient propagates backward through the network. For this example, we'll assume the process has already begun and the initial gradient dZ is a 2x4 matrix of ones.

Initial Gradient (assumed):

d Z = (11111111)

Step 1: Backward Pass through Dropout

The first step is to backpropagate through the Dropout layer. The gradient is only allowed to pass through the neurons that were not "dropped out" during the forward pass. This means we apply the same dropout mask and scaling factor.

Mathematics

The gradient with respect to the pre-dropout output (Z_unactivated) is calculated by element-wise multiplying the incoming gradient dZ with the dropout mask and scaling it by the same factor used in the forward pass.

\frac{\partial L}{\partial Z _{u na c t i v a t e d}} = \frac{\partial L}{\partial Z} ⊙ M a s k \times \frac{1}{1 - dropout_rate}

Calculation

From our forward pass, our mask and scaling factor were:
$M a s k = (10010110)$ , Scale = 2

The gradient dZ is passed then through this inverted dropout layer:

d Z_{ma s k e d} = d Z ⊙ M a s k = (11111111) ⊙ (10010110) = (10010110)

d Z_{u na c t i v a t e d} = d Z_{ma s k e d} \times 2 = (20020220)

This gradient, dZ_unactivated, is now available on both GPUs, as the All-Reduce in the forward pass made the output Z_unactivated identical on both.

Step 2: Backward Pass for `Z = Dropout(Y * B)`

Now we backpropagate through the second linear layer. This involves calculating the gradients with respect to the weights (dB) and the activations from the previous layer (dY).

Mathematics for Gradients `dY` and `dB`

Given the operation Z_unactivated = Y * B, the gradients are:

Gradient w.r.t. Y: $\frac{\partial L}{\partial Y} = \frac{\partial L}{\partial Z _{u na c t i v a t e d}} \cdot B^{T}$
Gradient w.r.t. B: $\frac{\partial L}{\partial B} = Y^{T} \cdot \frac{\partial L}{\partial Z _{u na c t i v a t e d}}$

In our tensor parallel setup (row parallelism), Y = [Y1, Y2] and B is split into B1 and B2. The forward pass involved local multiplications (Y1*B1, Y2*B2) followed by an All-Reduce. The backward pass reverses this.

For dY (Activation Gradients): The gradient dZ_unactivated is multiplied by the transpose of the local weight shard. This is a local computation on each GPU and requires no communication.
- On GPU 1: $d Y_{1} = d Z_{u na c t i v a t e d} \cdot B_{1 T}$
- On GPU 2: $d Y_{2} = d Z_{u na c t i v a t e d} \cdot B_{2 T}$
For dB (Weight Gradients): The transpose of the local activation Y is multiplied by the incoming gradient.
- On GPU 1: $d B_{1} = Y_{1 T} \cdot d Z_{u na c t i v a t e d}$
- On GPU 2: $d B_{2} = Y_{2 T} \cdot d Z_{u na c t i v a t e d}$

Calculation for `dY`

On GPU 1: Calculate dY1 = dZ_unactivated * B1^T

B_{1 T} = ⎝ ⎛ 654321789865 ⎠ ⎞

d Y_{1} = (20020220) \times ⎝ ⎛ 654321789865 ⎠ ⎞ = ((2 * 6 + 2 * 3) (2 * 5 + 2 * 4) (2 * 2 + 2 * 8) (2 * 1 + 2 * 7) (2 * 9 + 2 * 5) (2 * 8 + 2 * 6)) = (181820162828)

On GPU 2: Calculate dY2 = dZ_unactivated * B2^T

B_{2 T} = ⎝ ⎛ 432112345678 ⎠ ⎞

d Y_{2} = (20020220) \times ⎝ ⎛ 432112345678 ⎠ ⎞ = ((2 * 4 + 2 * 1) (2 * 3 + 2 * 2) (2 * 1 + 2 * 4) (2 * 2 + 2 * 3) (2 * 5 + 2 * 8) (2 * 6 + 2 * 7)) = (101010102626)

Calculation for `dB`

On GPU 1: Calculate dB1 = Y1^T * dZ_unactivated

Y_{1 T} = ⎝ ⎛ 31415183109135 ⎠ ⎞

d B_{1} = ⎝ ⎛ 31415183109135 ⎠ ⎞ \times (20020220) = ⎝ ⎛ (31 * 2) (41 * 2) (51 * 2) (83 * 2) (109 * 2) (135 * 2) (83 * 2) (109 * 2) (135 * 2) (31 * 2) (41 * 2) (51 * 2) ⎠ ⎞ = ⎝ ⎛ 62821021662182701662182706282102 ⎠ ⎞

On GPU 2: Calculate dB2 = Y2^T * dZ_unactivated

Y_{2 T} = ⎝ ⎛ 435363107133159 ⎠ ⎞

d B_{2} = ⎝ ⎛ 435363107133159 ⎠ ⎞ \times (20020220) = ⎝ ⎛ (43 * 2) (53 * 2) (63 * 2) (107 * 2) (133 * 2) (159 * 2) (107 * 2) (133 * 2) (159 * 2) (43 * 2) (53 * 2) (63 * 2) ⎠ ⎞ = ⎝ ⎛ 8610612621426631821426631886106126 ⎠ ⎞

Step 3: Backward Pass through GeLU Activation

Next, the gradients dY1 and dY2 are passed backward through the GeLU activation function.

Mathematics

The chain rule dictates that we multiply the incoming gradient by the derivative of the activation function evaluated at its original input.

\frac{\partial L}{\partial Y _{in t er m e d ia t e}} = \frac{\partial L}{\partial Y} ⊙ G e L U^{'} (Y_{in t er m e d ia t e})

The derivative of GeLU is $G e L U^{'} (x) = 0.5 \cdot (1 + er f (x / 2)) + \frac{x}{2 π} e^{- x^{2} /2}$ . As in the forward pass, for the large positive numbers in our example, this derivative is very close to 1. We will use this approximation for simplicity.

G e L U^{'} (x) \approx 1

Calculation

We simply pass the gradients through, as multiplying by 1 does not change them.
On GPU 1:

d Y_{in t er m e d ia t e_{1}} = d Y_{1} ⊙ 1 = (181820162828)

On GPU 2:

d Y_{in t er m e d ia t e_{2}} = d Y_{2} ⊙ 1 = (101010102626)

Step 4: Backward Pass for `Y = GeLU(X * A)`

Finally, we backpropagate through the first linear layer. This computes the gradients dA (for the weights) and dX (for the input).

Mathematics for Gradients `dX` and `dA`

Given Y_intermediate = X * A, the gradients are:

Gradient w.r.t. X: $\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y _{in t er m e d ia t e}} \cdot A^{T}$
Gradient w.r.t. A: $\frac{\partial L}{\partial A} = X^{T} \cdot \frac{\partial L}{\partial Y _{in t er m e d ia t e}}$

In our tensor parallel setup (column parallelism), the input X was broadcast, and the output Y_intermediate was sharded (Y_intermediate_1, Y_intermediate_2).

For dA (Weight Gradients): The calculation is local.
- On GPU 1: $d A_{1} = X^{T} \cdot d Y_{in t er m e d ia t e_{1}}$
- On GPU 2: $d A_{2} = X^{T} \cdot d Y_{in t er m e d ia t e_{2}}$
For dX (Input Gradients): Each GPU calculates a partial gradient for dX. These partial gradients must be summed together using an All-Reduce operation.
- On GPU 1: $d X_{p a r t ia l_{1}} = d Y_{in t er m e d ia t e_{1}} \cdot A_{1 T}$
- On GPU 2: $d X_{p a r t ia l_{2}} = d Y_{in t er m e d ia t e_{2}} \cdot A_{2 T}$
- Final Gradient: $d X = d X_{p a r t ia l_{1}} + d X_{p a r t ia l_{2}}$

Calculation for `dA`

On GPU 1: Calculate dA1 = X^T * dY_intermediate_1

X^{T} = ⎝ ⎛ 12345678 ⎠ ⎞

d A_{1} = ⎝ ⎛ 12345678 ⎠ ⎞ \times (181820162828) = ⎝ ⎛ (18 + 90) (36 + 108) (54 + 126) (72 + 144) (20 + 80) (40 + 96) (60 + 112) (80 + 128) (28 + 140) (56 + 168) (84 + 196) (112 + 224) ⎠ ⎞ = ⎝ ⎛ 108144180216100136172208168224280336 ⎠ ⎞

On GPU 2: Calculate dA2 = X^T * dY_intermediate_2

d A_{2} = ⎝ ⎛ 12345678 ⎠ ⎞ \times (101010102626) = ⎝ ⎛ (10 + 50) (20 + 60) (30 + 70) (40 + 80) (10 + 50) (20 + 60) (30 + 70) (40 + 80) (26 + 130) (52 + 156) (78 + 182) (104 + 208) ⎠ ⎞ = ⎝ ⎛ 60801001206080100120156208260312 ⎠ ⎞

Calculation for `dX` (with All-Reduce)

On GPU 1: Calculate dX_partial_1 = dY_intermediate_1 * A1^T

A_{1 T} = ⎝ ⎛ 123789456123 ⎠ ⎞

d X_{p a r t ia l_{1}} = (181820162828) \times ⎝ ⎛ 123789456123 ⎠ ⎞ = ((18 + 40 + 84) (18 + 32 + 84) (126 + 160 + 252) (126 + 128 + 252) (72 + 100 + 168) (72 + 80 + 168) (18 + 40 + 84) (18 + 32 + 84)) = (142134538506340320142134)

On GPU 2: Calculate dX_partial_2 = dY_intermediate_2 * A2^T

A_{2 T} = ⎝ ⎛ 456123789456 ⎠ ⎞

d X_{p a r t ia l_{2}} = (101010102626) \times ⎝ ⎛ 456123789456 ⎠ ⎞ = ((40 + 50 + 156) (40 + 50 + 156) (10 + 20 + 78) (10 + 20 + 78) (70 + 80 + 234) (70 + 80 + 234) (40 + 50 + 156) (40 + 50 + 156)) = (246246108108384384246246)

All-Reduce Communication Step
Finally, the partial gradients for X are summed across GPUs.

d X = d X_{p a r t ia l_{1}} + d X_{p a r t ia l_{2}} = (142134538506340320142134) + (246246108108384384246246) = (388380646614724704388380)

This final dX is the gradient with respect to the input of the MLP block, ready to be passed further backward to the previous layer in the network. The weight gradients dA1, dA2, dB1, and dB2 are used by the optimizer to update the model's weights on each respective GPU.

The Duality of Communication in Forward vs. Backward Pass

A fundamental principle of backpropagation in tensor parallelism is that the communication patterns are reversed.

If a layer requires an All-Reduce in the forward pass, its corresponding backward pass for the input gradient requires no communication (Identity).
If a layer requires no communication (Identity) in the forward pass, its corresponding backward pass for the input gradient requires an All-Reduce.

Here is a simple table summarizing the communication for the input gradients (dX and dY):

Layer Type (Forward)	Forward Pass Communication (`f`)	Backward Pass Communication (`g`)
Column Parallelism	`Identity` (No communication)	`All-Reduce` (Sum partial results)
Row Parallelism	`All-Reduce` (Sum partial results)	`Identity` (No communication)

(Note: The gradients for the weights, dA and dB, are always local calculations and never require an All-Reduce, so we can focus on the activation gradients.)

Backpropagation Through the Standard `Column -> Row` Pattern

Now, let's trace the backward pass through our standard two-layer block, keeping the table above in mind. The process starts with dZ, the gradient from the layer above. Because the forward pass ended with an All-Reduce, Z was a complete tensor on all GPUs, which means dZ is also a complete tensor.

Backprop through Layer 2: Z = Dropout(Y * B) (Row Parallelism)
- Forward Pass required: All-Reduce.
- Backward Pass requires: Identity (No communication).
- Calculation (dY = dZ * B^T): The complete gradient dZ is multiplied by the local weight shard B^T on each GPU. This directly produces the correct sharded activation gradient dY ([dY1, dY2]) without any need for communication.
Backprop through Layer 1: Y = GeLU(X * A) (Column Parallelism)
- Forward Pass required: Identity.
- Backward Pass requires: All-Reduce.
- Calculation (dX = dY * A^T): The sharded gradient dY from the previous step is multiplied by the local weight shard A^T. This produces a partial gradient for the input, dX_partial. To get the final, complete gradient dX, these partial results must be summed across all GPUs. This is the All-Reduce step.

As you can see, the backward pass for our Column -> Row block has exactly one All-Reduce operation, just like the forward pass.

What if we used `Row -> Column`?

The result would be perfectly symmetric:

Backprop through Layer 2 (Column Parallelism): The incoming gradient dZ would be sharded. This step would require an All-Reduce to produce a complete gradient dY.
Backprop through Layer 1 (Row Parallelism): The complete gradient dY would be the input. This step would require no communication (Identity) to produce the sharded input gradient dX.

The total communication remains one All-Reduce in the forward pass and one in the backward pass.

The Preference is for Modularity, Not Cost

The total cost is the same. So why do we prefer Column -> Row?

The reason is that this pattern creates a perfectly self-contained, modular block.

Input: The block takes a complete tensor (X).
Internal State: It creates a sharded tensor (Y) internally.
Output: It performs an All-Reduce at the end and outputs a complete tensor (Z).

The backward pass mirrors this perfectly:

Input: It takes a complete gradient (dZ).
Internal State: It creates a sharded gradient (dY) internally.
Output: It performs an All-Reduce at the end and outputs a complete gradient (dX).

This means you can stack these Transformer blocks one after another without worrying about whether the tensor you're passing is complete or sharded. Each block handles its own internal sharding and communication, always presenting a clean, complete tensor at its boundaries. This makes the overall architecture of a large model much simpler to design, implement, and debug.

Communication cost of tensor parallelism

In our walkthrough, communication between GPUs did not happen after every mathematical operation. It only occurred at very specific points where partial results needed to be combined to form a complete tensor. Let's pinpoint them:

Forward Pass Communication: This happened at the end of Step 3. To calculate the final Z_unactivated, we needed to sum the partial results from each GPU.
$Z_{u na c t i v a t e d} = Z_{p a r t ia l_{1}} + Z_{p a r t ia l_{2}}$
This summation is an all-reduce operation. GPU 1 sends its Z_partial_1 and receives Z_partial_2; GPU 2 sends Z_partial_2 and receives Z_partial_1. Both then compute the sum.
Backward Pass Communication: This happened at the end of Step 4. To calculate the final input gradient dX, we needed to sum the partial gradients from each GPU.
$d X = d X_{p a r t ia l_{1}} + d X_{p a r t ia l_{2}}$
This summation is also an all-reduce operation.

For a single pass through this two-layer block, we perform two full all-reduce operations.

Analyzing the Communication Volume

The formula for the communication volume in tensor parallelism per layer:

8 b s h (\frac{n _{devices} - 1}{n _{devices}})

Source: Stanford CS 336 Language Modeling from Scratch, Lecture 7: Parallelism 1, 1 hour, 1 min 8 sec

b: This is the batch size. In our example, the input matrix X had 2 rows, so $b = 2$ .
s: This is the sequence length. In a simple MLP, we can consider this to be 1, so $s = 1$ .
h: This is the hidden size of the tensor being communicated. The tensor we communicated in the forward pass was Z_partial, which had a shape of (2, 4). The size of the last dimension is 4, so $h = 4$ .
n_devices: The number of GPUs, which is $n_{devices} = 2$ .

In short, bsh represents the total number of elements in the activation tensor being communicated.

Volume of data in one all-reduce:

In the forward pass all-reduce, the tensor being communicated is Z_partial. The number of elements in this tensor is:

V o l u m e_{e l e m e n t s} = b \times s \times h = 2 \times 1 \times 4 = 8 elements

The all-reduce operation is a collective communication where each GPU sends its 8 elements and receives 8 elements from its peer.

The term $(\frac{n _{devices} - 1}{n _{devices}})$ in the formula characterizes the volume of a standard ring all-reduce algorithm. For our 2 GPUs, this factor is:

(\frac{2 - 1}{2}) = \frac{1}{2}

This signifies that each GPU sends and receives data chunks in a way that scales with the number of devices.

The factor of 8 in the formula often accounts for two things: 4 bytes for a standard 32-bit floating-point number and a factor of 2 because communication happens in both the forward and backward passes.

Why the communication cost of tensor parallelism is high

The communication cost of tensor parallelism is high due to the combination of volume and frequency:

High-Cost Operation: The all-reduce is a "collective" operation that is much more complex and latency-sensitive than a simple point-to-point send/receive. It requires synchronization across all GPUs in the group.
Large Data Volume: The amount of data communicated in each all-reduce is the full activation tensor (bsh), which can be very large in modern LLMs.
High Frequency: This expensive, high-volume operation is not performed once, but twice for every single transformer block in the model during a full training step (forward + backward pass).

Therefore, while tensor parallelism is excellent at keeping all GPUs utilized (eliminating the "bubble" of pipeline parallelism), it comes at the cost of constant, high-volume communication across all participating devices. This is why the general recommendation is to use tensor parallel whenever we have low-latency, high-bandwidth interconnects like NVIDIA's NVLink, which are specifically designed to handle this intense communication pattern efficiently.

Code

For those who prefer following the code, the walk through guide for both the forward and backward pass are replicated below using Python.

Setup

import torch
import torch.nn.functional as F

# --------------------------------------------------------------------------
# 1. SETUP: Define tensors from the manual example
# --------------------------------------------------------------------------
# Set requires_grad=True to track gradients
X = torch.tensor([[1, 2, 3, 4],
                  [5, 6, 7, 8]], dtype=torch.float32, requires_grad=True)

A = torch.tensor([[1, 2, 3, 4, 5, 6],
                  [7, 8, 9, 1, 2, 3],
                  [4, 5, 6, 7, 8, 9],
                  [1, 2, 3, 4, 5, 6]], dtype=torch.float32, requires_grad=True)

B = torch.tensor([[6, 5, 4, 3],
                  [2, 1, 7, 8],
                  [9, 8, 6, 5],
                  [4, 3, 2, 1],
                  [1, 2, 3, 4],
                  [5, 6, 7, 8]], dtype=torch.float32, requires_grad=True)

# NOTE: For verification, we use an identity function to match the manual
# calculation's approximation of GeLU(x) ≈ x. In a real model, you would
# use F.gelu(x).
gelu = lambda x: x

# Define the fixed dropout mask from the manual example
dropout_mask = torch.tensor([[1, 0, 0, 1],
                             [0, 1, 1, 0]], dtype=torch.float32)
dropout_rate = 0.5
dropout_scale = 1 / (1 - dropout_rate)

print("--- Initial Tensors ---")
print(f"X:\n{X}\n")
print(f"A:\n{A}\n")
print(f"B:\n{B}\n")

Output

--- Initial Tensors ---
X:
tensor([[1., 2., 3., 4.],
        [5., 6., 7., 8.]], requires_grad=True)

A:
tensor([[1., 2., 3., 4., 5., 6.],
        [7., 8., 9., 1., 2., 3.],
        [4., 5., 6., 7., 8., 9.],
        [1., 2., 3., 4., 5., 6.]], requires_grad=True)

B:
tensor([[6., 5., 4., 3.],
        [2., 1., 7., 8.],
        [9., 8., 6., 5.],
        [4., 3., 2., 1.],
        [1., 2., 3., 4.],
        [5., 6., 7., 8.]], requires_grad=True)

Forward pass

# --------------------------------------------------------------------------
# 2. FORWARD PASS
# --------------------------------------------------------------------------
print("--- FORWARD PASS ---")

# --- Simulate Tensor Parallelism on 2 GPUs ---

# Split A by columns (Column Parallelism)
# "GPU 1" gets A1, "GPU 2" gets A2
A1 = A[:, :3]
A2 = A[:, 3:]

# Split B by rows (Row Parallelism)
# "GPU 1" gets B1, "GPU 2" gets B2
B1 = B[:3, :]
B2 = B[3:, :]

# --- Layer 1: Y = GeLU(X * A) ---
# Operation is local to each GPU, no communication needed.

# "GPU 1" computes its part
Y_intermediate_1 = X @ A1
Y1 = gelu(Y_intermediate_1)

# "GPU 2" computes its part
Y_intermediate_2 = X @ A2
Y2 = gelu(Y_intermediate_2)

print(f"Y1 (on GPU 1):\n{Y1}\n")
print(f"Y2 (on GPU 2):\n{Y2}\n")

# --- Layer 2: Z = Dropout(Y * B) ---

# Local matrix multiplication on each GPU
# Note: Y is implicitly [Y1, Y2]
Z_partial_1 = Y1 @ B1
Z_partial_2 = Y2 @ B2

print(f"Z_partial_1 (on GPU 1):\n{Z_partial_1}\n")
print(f"Z_partial_2 (on GPU 2):\n{Z_partial_2}\n")

# *** All-Reduce Communication Step ***
# The partial results are summed across GPUs.
Z_unactivated = Z_partial_1 + Z_partial_2
print(f"Z_unactivated (after All-Reduce):\n{Z_unactivated}\n")

# Apply Dropout
Z_masked = Z_unactivated * dropout_mask
Z = Z_masked * dropout_scale

print(f"Final Output Z:\n{Z}\n")

Output

--- FORWARD PASS ---
Y1 (on GPU 1):
tensor([[ 31.,  41.,  51.],
        [ 83., 109., 135.]], grad_fn=<MmBackward0>)

Y2 (on GPU 2):
tensor([[ 43.,  53.,  63.],
        [107., 133., 159.]], grad_fn=<MmBackward0>)

Z_partial_1 (on GPU 1):
tensor([[ 727.,  604.,  717.,  676.],
        [1931., 1604., 1905., 1796.]], grad_fn=<MmBackward0>)

Z_partial_2 (on GPU 2):
tensor([[ 540.,  613.,  686.,  759.],
        [1356., 1541., 1726., 1911.]], grad_fn=<MmBackward0>)

Z_unactivated (after All-Reduce):
tensor([[1267., 1217., 1403., 1435.],
        [3287., 3145., 3631., 3707.]], grad_fn=<AddBackward0>)

Final Output Z:
tensor([[2534.,    0.,    0., 2870.],
        [   0., 6290., 7262.,    0.]], grad_fn=<MulBackward0>)

Backward pass

# --------------------------------------------------------------------------
# 3. BACKWARD PASS
# --------------------------------------------------------------------------
print("--- BACKWARD PASS ---")

# Define a dummy loss. Z.sum() creates an initial gradient dZ of all ones.
loss = Z.sum()
print(f"Loss (for backprop):\n{loss}\n")

# Automatically compute gradients for all tensors with requires_grad=True
loss.backward()

Output

--- BACKWARD PASS ---
Loss (for backprop):
18956.0

Verification

# --------------------------------------------------------------------------
# 4. VERIFICATION: Check gradients against manual calculations
# --------------------------------------------------------------------------
print("--- VERIFYING GRADIENTS ---\n")

# Gradient for X (dX)
# This required an All-Reduce in the backward pass
print(f"Gradient for X (dX):\n{X.grad}\n")

# Gradients for A (dA1 and dA2)
# A.grad contains the gradients for the entire A matrix. We can verify
# by slicing it into the parts corresponding to A1 and A2.
print(f"Gradient for A1 (dA1):\n{A.grad[:, :3]}\n")
print(f"Gradient for A2 (dA2):\n{A.grad[:, 3:]}\n")

# Gradients for B (dB1 and dB2)
# Similarly, B.grad contains the full gradient.
print(f"Gradient for B1 (dB1):\n{B.grad[:3, :]}\n")
print(f"Gradient for B2 (dB2):\n{B.grad[3:, :]}\n")

Output

--- VERIFYING GRADIENTS ---

Gradient for X (dX):
tensor([[388., 646., 724., 388.],
        [380., 614., 704., 380.]])

Gradient for A1 (dA1):
tensor([[108., 100., 168.],
        [144., 136., 224.],
        [180., 172., 280.],
        [216., 208., 336.]])

Gradient for A2 (dA2):
tensor([[ 60.,  60., 156.],
        [ 80.,  80., 208.],
        [100., 100., 260.],
        [120., 120., 312.]])

Gradient for B1 (dB1):
tensor([[ 62., 166., 166.,  62.],
        [ 82., 218., 218.,  82.],
        [102., 270., 270., 102.]])

Gradient for B2 (dB2):
tensor([[ 86., 214., 214.,  86.],
        [106., 266., 266., 106.],
        [126., 318., 318., 126.]])

DEV Community: Lewis Won

PagedAttention by hand

Motivation

What is PagedAttention

Internal vs external fragmentation

How PagedAttention solves reservation and fragmentations

Why not set B = 1 to solve internal fragmentation?

Estimating size of a block byte

Setup

Step 1: Initial Setup - Defining the Tensors and Memory

Forward pass (Paged Attention)

Step 2: Compute Unnormalized Attention Scores ( qKT )

1. First Logical Block (Physical Block 2)

2. Second Logical Block (Physical Block 0)

3. Third Logical Block (Physical Block 1)

Step 3: Apply Softmax

Step 4: Multiply by Values ( Probs×V )

1. First Logical Block (Physical Block 2)

2. Second Logical Block (Physical Block 0)

3. Third Logical Block (Physical Block 1)

4. Final Aggregation

Why divide memory into physical blocks?

1. Zero External Fragmentation & Minimal Internal Fragmentation

2. Memory Sharing (Copy-on-Write)

Example of Copy-on-Write

Why contigous KV cache has difficulty implementing prefix matching

Code

KV Cache by hand

Motivation

What is the KV Cache?

Setup

Scenario 1: Generation WITHOUT KV Cache

Step 1: Calculate Q, K, and V for the whole sequence

Step 2: Compute Attention for the new token

Scenario 2: Generation WITH KV Cache

Step 1: Calculate Q, K, and V for ONLY the new token

Step 2: Update the Cache

Step 3: Compute Attention

The Compute vs. Memory Trade-off

Code

Appendix A: Worked example with 2 layers

FP8, FP4, MXFP4 and NVFP4 by hand

Why by hand?

The cast of characters

How the bits work: sign, exponent, mantissa

The same three fields, on real numbers

What is the scale for?

What does it mean to "share" a scale?

How many scale factors are there, really?

Why is scale encoding important

Everything is a grid and a scale

The E2M1 grid, enumerated by hand

Where the bias comes from

Conversion by hand with a sample tensor

Format 1: FP8 (E4M3) by hand

Aside: FP8 is two formats

Why that split

Why this matters for the rest of the article

And for FP4? There is only one

Format 2: FP4 by hand

Format 3: MXFP4 by hand

Format 4: NVFP4 by hand

Step 1: the per-tensor FP32 scale

Step 2: block A's scale

Step 3: quantize block A by hand

Step 4: block B's scale

The scoreboard

Appendix A: the E4M3 grid by hand

Appendix B: why scale factors round up, not to nearest

Appendix C: the full 32-value table

Appendix D: Why E2M1 stops at 6, E4M3 stops at 448

1. Why E2M1 stops at 6

2. Why E4M3 stops at 448

Summary

How does low-rank adaptation for large language models work

Table of Contents

Why study LoRA? The challenge of fine-tuning massive models

Conceptual overview: Full Fine-Tuning

A Simple Example

Link back to the Math Equation on Full Fine-tuning

Step 2: Compute Unnormalized Attention Scores ( $q K^{T}$ )

Step 4: Multiply by Values ( $P ro b s \times V$ )

1. Why E2M1 stops at `6`

2. Why E4M3 stops at `448`

3. Calculating the Gradient for `B` (`grad_B`)

4. Calculating the Gradient for `A` (`grad_A`)

2. Calculating the Gradient for `B_q`

3. Calculating the Gradient for `A_q`

Step 1: Calculate the Full Attention Score Matrix (`A = softmax(X)`)

1a. Find the Global Maximum (`m`)

1b. Compute the Exponentials and the Denominator (`d`)

1c. Normalize to get the Final Attention Vector `A`

Step 2: Multiply Attention Scores by the Value Matrix ( $O = A \cdot V$ )

Step 1: Process Tile 1 (`i=1`)

Step 2: Process Tile 2 (`i=2`)