DEV Community: Lewis Won

How to profile and optimise kernels

Lewis Won — Sat, 03 Oct 2026 09:50:39 +0000

Table of Contents

Where the last article left off
The setup
How I profile
Round 0: profile the baseline
Round 1: tile sizes, num_warps and num_stages
- The shared-memory budget for an Ada SM
- The sweep
- The heatmaps
- The winner, through the unchanged harness and tests
Round 2: profile the tuned kernel
Round 3: wrap the winner in @triton.autotune, and what it costs
Round 4: inside the inner loop
- Causal tile-skipping
- exp2
Round 5: re-tune, then the final profile
What to look for, as a checklist

In my previous article, I wrote a FlashAttention-2 (FA2) kernel which achieved about 22 TFLOP/s, 32% of PyTorch's scaled_dot_product_attention (SDPA) at N = 4096, and 18% for causal attention.

In this article, I will optimise my kernel with the following loop:

Profile the kernel to identify bottlenecks.
Change one thing to address a bottleneck.
Measure the performance of the amended kernel against the original kernel.
Profile again, to confirm the change, and identify the next bottleneck.

In this article, I first optimise the outer loop of FA2 by changing the tile sizes, num_warps and num_stages, working through the shared-memory budget of an Ada SM, sweeping the configuration space, drawing the heatmap, and wrapping the winner in @triton.autotune and measuring the compile time.

I then optimise the inner loop by skipping the tiles that causal masking throws away, and switching exp to exp2. I will also include screenshots of Nsight Systems and Nsight Compute and analysis of the roofline chart.

What is P?
P is the kernel's P_ij, i.e. the tile of unnormalised attention probabilities, i.e. the softmax numerators before they are divided by the row sum. For query tile i and key tile j:

S = Q_i K_jᵀ · scale is the tile of scores.
P = exp(S − m), where m is the running maximum of each row.

S_ij = tl.dot(Q_i, tl.trans(K_j))                    # scores, fp32
m_new = tl.maximum(m_acc, tl.max(S_ij, axis=1, keep_dims=True) * qk_scale)
P_ij = tl.math.exp2(S_ij * qk_scale - m_new)         # P, in fp32
# ...
l_acc = alpha * l_acc + tl.sum(P_ij, axis=1, keep_dims=True)   # row sums of P go into l
P_ij = P_ij.to(V_j.dtype)                            # "the cast of P to bf16"
O_acc = tl.dot(P_ij, V_j, acc=O_acc)                 # P V goes into O

I did not change the benchmark harness bench_flashattention.py and the test file test_flashattention_triton.py from the previous article, so every number here can be compared with the numbers there.

In summary, below is how the optimised FA2 in this article compare to the FA2 in SDPA and the previous article's kernel (median of ten runs):

	Previous article's kernel	This article's kernel	SDPA
Non-causal	6.39 ms, 21.5 TFLOP/s, 32% of SDPA	1.95 ms, 70.3 TFLOP/s, 103% of SDPA	2.01 ms, 68.4 TFLOP/s
Causal	6.38 ms, 10.8 TFLOP/s, 18% of SDPA	1.08 ms, 63.8 TFLOP/s, 104% of SDPA	1.12 ms, 61.5 TFLOP/s

The improvment in kernel performance came from:

Tuning of tile size and warp count;
Not wasting TFLOPs on computing the tiles covered by causal masking; and
Switching exp to exp2, and re-tuning with autotune.

Using Nsight Compute, I profiled my optimised non-causal kernel at 92% of the tensor cores' peak, and SDPA at 89%.

All numbers are from the same RTX 4070 SUPER under Ubuntu. This article assumes you have read my previous article, am comfortable with FlashAttention-2 forward pass and basic Triton.

This article was written with the assistance of AI. If you spot any mistakes, please let me know.

Where the last article left off

In my previous article, my unoptimised kernel gives each program one tile of 16 query rows and loops over all keys in tiles of 16, keeps the running max m, the running sum l and the unnormalised output O in fp32. My unoptimised kernel used Triton's defaults of 4 warps per program and 3 pipeline stages.

I listed the following problems in my previous article:

No profile.
The tiles are too small.
Causal attention computes every tile and masks half of them away.
tl.exp instead of tl.math.exp2.
The mask is computed on every tile.

The setup

The GPU

The goal of kernel optimisation here is to keep the GPU cores fully utilised. Below are the properties of RTX 4070 SUPER. These properties are derived from torch.cuda.get_device_properties, Triton's driver.active.utils.get_device_properties() and NVIDIA's documentation for compute capability 8.9. The last three rows are measured.

	RTX 4070 SUPER
Architecture	Ada Lovelace (AD104), compute capability 8.9, `sm_89`
Streaming multiprocessors (SMs)	56
Registers per SM	65,536 32-bit registers, at most 255 per thread
Resident threads per SM	up to 1,536 (48 warps) in up to 24 blocks
L1 + shared memory per SM	128 KB, of which up to 100 KB (102,400 bytes) can be shared memory
Shared memory per block (program)	up to 99 KB (101,376 bytes), plus 1 KB the driver reserves for every block
L2 cache	48 MB
DRAM	12 GB GDDR6X, 504 GB/s peak
Tensor-core peak, bf16 inputs with fp32 accumulation (measured)	512 FLOP per clock per SM: 72.2 TFLOP/s at the 2.52 GHz Nsight Compute locks the clock to, 75 to 77 TFLOP/s at the 2.61 to 2.68 GHz at which the final kernel and SDPA ran during a harness run at N = 4096
cuBLAS bf16 matmul, 8192 × 8192 × 8192 (measured)	69.4 TFLOP/s; 70.9 TFLOP/s in a later 12-second run at 2.66 GHz, 93% of the peak at that clock
Device-to-device copy (measured)	414 to 427 GB/s, counting the bytes read and the bytes written

Details to note:

Tensor core peak: On GeForce Ada cards, the tensor cores' peak for fp16 inputs is twice as high with fp16 accumulation as with fp32 accumulation. Nsight Compute lists 1,024 operations per clock per SM for "fp16 in, fp16 out" against 512 for "fp16 in, fp32 out". bf16 inputs accumulating in fp32 also has 512 operations per clock per SM. Because the attention kernel takes bf16 inputs and needs fp32 accumulation, so the ceiling for this article is 512 operations per clock per SM.
TFLOP/s ceiling changes with clock speed: This is why Nsight Compute locks the clock at 2.52 GHz, where the ceiling is 72.2 TFLOP/s. During actual runs outside of Nsight Compute, the optimised kernel and SDPA ran at 2.61 to 2.68 GHz at N = 4096, where the ceiling is 75 to 77 TFLOP/s. SDPA, for example, ran at 2.62 GHz and reached 66.7 TFLOP/s, 88.6% of the ceiling at that clock.
Size of shared memory: Ada has 100KB per SM and 99KB per program, compared to an A100 which has 164 KB of shared memory per SM and an H100 228 KB. The size of shared memory influences the tuning parameters of the optimised kernel.

Software and files

These are the software versions I used:

PyTorch 2.11 (CUDA 13.0)
Triton 3.6.0
Nsight Systems 2026.1.3
Nsight Compute 2026.2.1
Nvidia driver 595.91.

Each version of the kernel lives in its own folder with the same file name, flashattention_autograd_function_triton.py, so the unchanged harness and the unchanged test file can import whichever one is on the path:

PYTHONPATH=kernels/v1_tiles python bench_flashattention.py
PYTHONPATH=kernels/v1_tiles python -m pytest -q test_flashattention_triton.py

Folder	What changes
`kernels/v0_baseline`	the kernel from the previous article, unchanged
`kernels/v1_tiles`	the same kernel body, with the tile sizes, `num_warps` and `num_stages` from the sweep
`kernels/v2_autotune`	v1 wrapped in `@triton.autotune`
`kernels/v3_causal_skip`	the key loop split into unmasked tiles, diagonal tiles, and nothing above the diagonal
`kernels/v4_exp2`	v3 with the softmax scale and log₂e folded into one multiply-add, and `exp2`
`kernels/v5_final`	v4 wrapped in `@triton.autotune`, with configurations from a second sweep

All other artifacts including the sweep, profiling drivers, compile-cost measurements, plotting scripts and raw results are in 2026_10_03_followup/.

The baseline, re-measured

Before changing anything, I re-ran the unchanged harness on the unchanged kernel, to check that today's GPU gives the same numbers as the day the previous article was written. Each result is the median of ten runs. We can see that today's measurement and the previous article's agree to within ~5%.

N	causal	previous article (ms)	today (ms)
512	no	0.107	0.114
1024	no	0.386	0.387
2048	no	1.534	1.540
4096	no	6.195	6.386
4096	yes	6.155	6.383

How I profile

There are two NVIDIA profilers, and they answer different questions.

Nsight Systems (nsys) records a timeline of the whole program. This involves tracking Python threads, CUDA API calls, and every kernel on the GPU. It answers "where does the time go, and is the GPU busy?". It is cheap to run and rarely changes the program's behaviour.
Nsight Compute (ncu) profiles individual kernel launches. It replays each launch dozens of times, reading a different set of hardware counters on each pass, and turns them into utilisation, stall and memory-traffic numbers. It answers "why is this kernel as slow as it is?".

I will run Nsight Systems first to identify kernels with bottlenecks, then Nsight Compute on a specific kernel.

A third tool is the compiler itself. Every compiled Triton kernel keeps its intermediate representations, so I can see what the compiler did with my code without a profiler, and this include looking at register count, spills, shared memory, the layout it chose for each tensor, and the final machine code.

A small program to profile

Profiling the benchmark harness directly is possible, but it launches each kernel hundreds of times between cache-flushing kernels. I profile a short driver instead. It compiles the kernel outside the region of interest, then runs it three times inside NVTX ranges (named regions that show up on the timeline), next to SDPA on the same inputs:

# profile_driver.py (abridged)
# B = 4 sequences in a batch, H = 8 attention heads, and D = 64 numbers per head per token
B, H, D = 4, 8, 64  # the harness shape.
# Three random tensors of shape (B·H, N, D) = (32, N, 64).
# The (... for _ in range(3)) part calls torch.randn three times, so q, k and v hold different values. 
q, k, v = (torch.randn(B * H, args.N, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))
# The same memory, relabelled as (4, 8, N, 64). .view copies nothing, and it only changes the shape and strides. 
q4, k4, v4 = (t.view(B, H, args.N, D) for t in (q, k, v))

# Warm-up: compile (and autotune, if the version autotunes) outside the ranges, then keep
# the GPU busy for about a second so its clocks have ramped up before the capture starts.
t0 = time.perf_counter()
while time.perf_counter() - t0 < 1.0:
    FlashAttentionTriton.apply(q, k, v, causal)
    F.scaled_dot_product_attention(q4, k4, v4, is_causal=causal)
    torch.cuda.synchronize()

# nsys --capture-range=cudaProfilerApi records only what runs between start and stop.
torch.cuda.cudart().cudaProfilerStart()
for _ in range(args.iters):
    torch.cuda.nvtx.range_push(f"ours N={args.N} causal={causal}")
    FlashAttentionTriton.apply(q, k, v, causal)
    torch.cuda.nvtx.range_pop()
    torch.cuda.nvtx.range_push(f"sdpa N={args.N} causal={causal}")
    F.scaled_dot_product_attention(q4, k4, v4, is_causal=causal)
    torch.cuda.nvtx.range_pop()
torch.cuda.synchronize()
torch.cuda.cudart().cudaProfilerStop()

compare_driver.py does the same for several kernel versions in one process, so one timeline or one Nsight Compute report can hold all of them side by side.

The commands

# Timeline: only the region between cudaProfilerStart and cudaProfilerStop.
nsys profile --trace=cuda,nvtx --capture-range=cudaProfilerApi --capture-range-end=stop \
    -o results/nsys/round0 python profile_driver.py kernels/v0_baseline --N 4096 --causal 0

# Counters: every kernel in the same region, with the full set of sections.
ncu --set full --import-source yes --profile-from-start off \
    -o results/ncu/v0_N4096_nc python profile_driver.py kernels/v0_baseline --N 4096 --causal 0 --iters 1 --sdpa 0

Then open the .nsys-rep file in nsys-ui and the .ncu-rep file in ncu-ui. --set full collects nearly every section, including the roofline charts. --import-source yes copies the Python source into the report, so the Source page can show which line of the kernel each instruction came from.

Three points to take note:

Nsight Compute changes the conditions it measures. (--clock-control boost) locks the GPU clock speed, which on my card means 2.52 GHz, below the 2.6 to 2.7 GHz the card runs the harness at. It also flushes the caches before every replay pass, and runs the kernel about 45 times.
Triton's cache can point the profiler at the wrong file. To ensure Triton is pointing at the correct file, point TRITON_CACHE_DIR at an empty directory before profiling. Otherwise, Triton may reuse the cached binary, and the binary's line information may still be pointing at the old file path.
Nsight Compute's suggestions are hypotheses, not diagnoses. The Summary page ranks "optimisation opportunities" with estimated speedups. For the baseline below, the top two were "theoretical occupancy" (20%) and "shared store bank conflicts" (9%), which may not necessarily be the actual bottlenecks.

What the compiler did

inspect_kernel.py compiles one configuration without launching it and prints what the compiler produced. The TTGIR (Triton GPU IR) is the last stage where the program still looks like the Python, but every tensor now carries a layout that says which thread and which warp holds which element, and most trips through shared memory are explicit operations.

compiled = flash_fwd_kernel.warmup(q, k, v, o, L, *strides, N, N, scale,
                                   D=64, Q_TILE_SIZE=16, K_TILE_SIZE=16, is_causal=False,
                                   num_warps=4, num_stages=3, grid=(N // 16, B))
compiled._init_handles()                  # loads the module, which fills in n_regs and n_spills
print(compiled.n_regs, compiled.n_spills, compiled.metadata.shared)
print(compiled.asm["ttgir"])              # Triton GPU IR: layouts and shared-memory buffers
print(compiled.asm["ptx"])                # PTX
cubin = compiled.asm["cubin"]             # SASS via: nvdisasm -c -g kernel.cubin

Round 0: profile the baseline

By profling the baseline unoptimised kernel from the previous article, we can find out why was the kernel so slow. Everything here is at N = 4096, non-causal, bf16 input.

Step 1: the timeline

What I look for in a timeline, in this order:

Gaps in the GPU row. The row labelled CUDA HW shows when the GPU is running a kernel. If it is idle between kernels, the bottleneck is outside of the kernel, e.g. Python, the launcher, synchronisation.
How many kernels each call launches, and which of them takes the time.
The kernel's duration next to a reference that does the same work on the same inputs.

What the timeline above shows:

No gaps. The GPU row is solid from the first kernel to the last. On the CPU thread, the NVTX ranges are tiny ticks at the very start, followed by one long cudaDeviceSynchronize: Python queued all six launches under a millisecond and then waited for the GPU. Hovering a kernel shows its launch latency, the time between the launch call and the kernel starting. For the SDPA kernel under the cursor it is 14.1 ms. The timeline shows this is a GPU-bound program, so the kernel is the target for optimisation.
One kernel per call. flash_fwd_kernel is the only thing the Triton version launches (for both the unoptimised baseline kernel and the SDPA's FA2), so there is nothing else to look at.
3× slower than the reference. Each of my baseline unoptimised kernels takes 5.8 to 6.5 ms, as compared to each of SDPA's FA2 which takes about 2.0 ms on the same inputs.

Hovering the mouse over the SDPA kernel shows a tooltip with its launch configuration: 128 threads per block (4 warps), 49,152 bytes of shared memory, 255 registers per thread and a theoretical occupancy of 16.7%. Its grid, 32 × 4 × 8, is FA2's query blocks × batch × heads, so each block takes 4,096 / 32 = 128 queries. Note that N = 4,096. The tooltip names the kernel flash_fwd_kernel, the same name as my Triton kernel. The full name, from nsys stats --report cuda_gpu_kern_sum, says what SDPA runs on this GPU: pytorch_flash::flash_fwd_kernel>. That is PyTorch's implementation of the FlashAttention-2 kernel for head dimension 64, with 128 × 128 tiles and 4 warps per block.

One warning about durations in a timeline.
The three runs of my kernel above differ by up to 11%, and my first capture of this timeline, without the one-second warm-up in the driver, showed SDPA at 3.4 ms. Nsight Systems can show whether the GPU's clock is to blame: nsys profile --gpu-metrics-devices=0 samples it alongside the kernels. Here are two short captures of the same work, one taken straight after 3 seconds idle, the other after the driver's one-second warm-up:

Straight after idle, the GPU ran at 2.52 GHz for the whole capture, and SDPA took 2.13 to 2.14 ms. After the warm-up it ran at 2.69 to 2.73 GHz and SDPA took 1.97 ms, until the clock dropped to 2.54 GHz for the last two calls and SDPA took 2.11 ms. That looks like a warm-up worth 7 to 8%, but longer captures tell a different story:

The clock waits for sustained work (A, B). From the idle clock, the first kernel takes the GPU straight to 2.52 GHz, the rated boost clock of this factory-overclocked ASUS card (NVIDIA specifies 2.48 GHz). The clock stays there until the GPU has been kept busy, more or less without a break, for 0.13 to 0.3 s. Idle time does not count, and short bursts do not add up. Woken and then left idle for a second, the GPU was still at 2.52 GHz, and it boosted 0.22 s after the work began. Work in 5 ms bursts every 50 ms, 0.15 s of work in all, did not move the clock in 1.5 s.
Then the clock goes as high as the limits allow (A). A low-power load holds 2.73 GHz. v0 and SDPA together reach the card's 220 W power limit, which holds the clock near 2.5 GHz.
Once boosted, the clock rises while the GPU idles (C). Within 40 to 80 ms of idle it is at about 2.7 GHz. It stays there until the GPU has been idle for about 1.8 s, then falls back to 2.52 GHz. Work that starts inside that window runs fast until the power limit pulls the clock down again, 30 to 130 ms later for v0 and SDPA. After 16 to 17 s of idle the clock heads down to the idle clock, which it reaches 19 to 25 s after the work ended.

The warm capture caught the last of these. To start a capture, Nsight Systems holds the program in cudaProfilerStart for about 0.4 s, or 0.6 s with GPU metrics on. So the GPU sat idle after the warm-up, and its clock climbed to 2.73 GHz. The drop to 2.54 GHz 35 ms into the work was the power limit catching up: in sustained runs of this work, the driver reports its power cap as active throughout. The cold capture had the same pause, but no sustained work before it, so its clock stayed at 2.52 GHz. Sustained, this work runs near 2.5 GHz whether the GPU starts cold or warm, and SDPA takes 2.12 to 2.15 ms either way. The Round 0 timeline above was almost certainly taken in the fast window too. Its six kernels ran within the first 24 ms after the pause, and SDPA's 1.97 ms there matches the fast start, about 8% faster than it sustains. The 3× between my kernel and SDPA holds either way, because both ran in the same window.

The clock does not explain the 3.4 ms. Calls like it recur: with no profiler attached, up to 1 SDPA call in 20 took 3.0 to 3.6 ms, though the rate varies widely from run to run. In the longer captures, the SMs sit almost idle for 1.3 to 1.9 ms inside each of these calls, at an unchanged clock, while Nsight Systems shows graphics work (vertex and pixel warps) on the GPU. This card also drives my desktop, and its drawing interrupts the kernels. The SDPA calls that overlapped graphics work took 2.3 to 3.9 ms, and their duration tracks how long the graphics work ran (a correlation of 0.97 to 0.98). No other call took more than 2.37 ms. The 11% spread in the Round 0 timeline above may have the same cause, but that capture has no GPU metrics to check. A single kernel in a timeline is not a benchmark. The timeline is for the shape of the program; the harness is for the time.

Step 2: the Speed of Light section

Opening the Nsight Compute report on the Details page, the first section is GPU Speed Of Light Throughput. "Speed of light" is NVIDIA's name for the theoretical maximum of each hardware unit, and every number in the section is a percentage of it. The two headline numbers are Compute (SM) Throughput and Memory Throughput. Each is the maximum over many sub-units, so the headline alone does not say which unit is busy. The drop-down on the right of the section title switches the chart to GPU Throughput Breakdown, which lists the sub-units:

Both headline numbers read 79.9%, and the rule underneath says "Compute and Memory are well-balanced". That sounds healthy. The breakdown says otherwise:

The top compute entry is SM: Inst Executed Pipe Lsu at 79.9%. The LSU (load/store unit) pipe issues memory instructions, and inside an SM that mostly means shared-memory loads and stores. The tensor pipe, which does the matrix multiplies, is at 22.8%.
The top memory entry is L1: Lsuin Requests at 79.9%: the same LSU traffic, seen from the memory side. DRAM: Cycles Active is at 3.1%.

So the kernel is busy, but busy moving data around inside each SM, not multiplying it.

Step 3: which instructions keep the LSU busy

Further down the Details page, the Compute Workload Analysis section names the busiest pipeline and, under "Most frequently executed instructions for pipeline LSU", the source lines responsible are listed.

The top two lines are standard.py lines 293 and 191. These are Triton's own implementations of tl.sum and tl.max, which the kernel calls for the running row sum and the row maximum. Their opcodes, LDS, SHFL.BFLY and STS, are shared-memory loads, warp shuffles and shared-memory stores, and each of the two lines executed 92 million of them. A row reduction that stays inside one warp needs only shuffles. One that needs shared memory is combining partial results across warps, which means the rows of the score tile are split across warps. The third line is the load of Q, as LDSM (load from shared memory into the matrix-multiply registers), 34 million times.

Step 4: why the warps wait

The Warp State Statistics section shows, for every instruction issued, how many cycles a warp spent waiting and why:

The longest waits, per issued instruction:

Stall reason	Cycles	What it means
Short scoreboard	4.5	waiting for the result of a shared-memory access (or a special-function op like `exp`)
Barrier	4.4	waiting at a `__syncthreads()` for the other warps of the program
MIO throttle	3.7	the queue for shared-memory instructions is full
Wait	1.8	a fixed-latency dependency between two instructions
Math pipe throttle	0.6	the target math pipe (here mostly the tensor pipe) is busy

For comparison, the same section for SDPA's kernel has one dominant reason, math pipe throttle at 6.8 cycles. That is what a healthy compute-bound kernel looks like, i.e. its warps wait because the tensor cores are busy. The baseline unoptimised kernel's warps spend most time waiting on shared memory and on each other.

Step 5: the roofline

The same drop-down on the Speed Of Light section offers several roofline charts. For a kernel built on tl.dot with bf16 inputs, the one that matters is Roofline Tensor Core:

How to read it:

The horizontal axis is arithmetic intensity: operations per byte moved. The vertical axis is achieved operations per second. Both are logarithmic.
The flat line is the tensor-core peak, 72.2 TOP/s at the locked 2.52 GHz. The three sloped lines are bandwidth limits for L1, L2 and DRAM, from left to right: at a given intensity, no kernel can be faster than bandwidth × intensity. The corner where each sloped line meets the flat line is that level's ridge point.
Each dot is the kernel measured against one level of the memory hierarchy: the same operation count, divided by the bytes that moved through L1, L2 or DRAM. Hovering a dot names its level. Here the orange dot is L2, the green one DRAM and the purple one L1.
A dot sitting on a sloped line is bound by that level's bandwidth. A dot on the flat line is compute-bound. A dot below both lines, like these, is limited by something else, usually latency or instruction issue, which is what Steps 2 to 4 found.

Two numbers that are key to diagnosing the problem:

The number of operations. Working from first principle, the algorithm needs 4 × (B·H) × N² × D = 4 × 32 × 4096² × 64 = 137.4 G. However, the table under the chart says the tensor cores did 206G operations. This means the hardware did 1.5 times the work the algorithm needs, which is what happens if one of the two multiplies is done twice. Step 6 below shows that this is what happens in the unoptimised baseline kernel. In fact, SDPA's report shows exactly 137.4 G number of operations. Therefore, a third of this kernel's tensor work is discarded.
The bytes. A dot's position on the horizontal axis is operations divided by bytes. The L2 dot is at 23.9 operations per byte, so the kernel moved 206 G / 23.9 = 8.6 GB from L2 into the SMs. That is 8.6 GB / 6.27 ms = 1.37 TB/s, which is well within the 3.1 TB/s maximum throughput of the L2 cache, hence transfer from L2 cache is not a bottleneck at this point. DRAM is also not a bottleneck because it only needs to transfer about 71 MB of data into the L2 cache (reads), and 23.9 MB out (writes, mostly the output O), hence 95.5 MB in total. Time to move 95.5 MB at DRAM's 504 GB/s means 95.5 MB / 504 GB/s = 0.19 ms, which is well under the 6.27 ms time which the kernel took.

Why is maximum throughput of L2 cache at 3.1 TB/s?
The L2 cache is split into 24 slices. Each slice can send 64 bytes to the SMs per clock, and under Nsight Compute the L2 clock runs at 2.02 GHz. That gives at most 24 × 64 × 2.02 G = 3.1 TB/s.

Why is the DRAM traffic so small?
Without L2, every program's 1.05 MB of K and V would come from DRAM: 8.6 GB / 504 GB/s = 17 ms. That is almost three times slower than the kernel runs today, and DRAM would then be the bottleneck. Because the 256 programs of a head run together, each byte of K and V comes from DRAM once and is then served from L2 another 255 times. That reuse is what turns 8.6 GB of reads into 50 MB from DRAM.

When will DRAM become a bottleneck?
DRAM's minimum traffic is reading Q, K and V and writing O, once each: 4 × 32 heads × N × 64 × 2 bytes. The useful work is 4 × 32 × N² × 64 operations. Dividing one by the other gives N / 2 useful operations per DRAM byte: 2,048 at N = 4,096. DRAM's ridge is 72.2 T / 504 G = 143 operations per byte, so DRAM would only become the limit below about N = 290. The work grows with N², while the DRAM traffic grows only with N.

One correction to the "71 MB": only 50.3 MB of the reads are the kernel's own Q, K and V (the L2 misses). It is the size of Q, K and V read once each: 3 tensors × 32 heads × 4,096 tokens × 64 dims × 2 bytes = 3 × 16.78 MB ≈ 50.3 MB. The other ~21 MB doesn't come from the kernel's loads. L2 also saw 21.5 MB of traffic from units other than the SMs.

Step 6: ask the compiler why

Steps 2 to 5 left two symptoms unexplained:

Too much tensor-core work. The tensor cores did 206.2 G operations where the algorithm needs 137.4 G, 1.5 times too many (Step 5).
Too much shared-memory traffic. The LSU pipe is at 79.9%, tl.max and tl.sum each execute 92 million shuffle and shared-memory instructions, and the warps stall on short scoreboard, memory input/output (MIO) throttle and barriers (Steps 2 to 4).

What is short scoreboard?
The warp needs the result of an earlier shared-memory load, shuffle or special-function op (like exp), and it has not arrived yet. Specifically, that earlier instruction went through the MIO path: shared memory, special-function math or shuffles.

What is MIO throttle?
The warp’s next instruction is a shared-memory access, shuffle or special-function op, and the queue that feeds those units is full, so the instruction cannot issue yet.

Both come from one decision the compiler made: with a 16 × 16 tile and 4 warps, Triton places the warps side by side along the key axis, so no warp owns a whole row of the tile. Three problems follow from that one decision, and each one matches a symptom:

Problem	Where it showed up	Size
1. Two of the four warps compute `Q Kᵀ` a second time	roofline operation count (Step 5)	206.2 G operations instead of 137.4 G
2. Every row maximum and row sum has to combine values from two warps	LSU pipe, the `tl.max` and `tl.sum` lines (Steps 2 and 3)	92 million instructions per line
3. `P` has to move between warps before `P V`	shared-memory and barrier stalls (Step 4)	512 bytes stored and reloaded on every iteration

The rest of this step shows where that layout comes from and then counts each problem, using what the compiler printed.

Five facts to start from

1. The sizes. The baseline uses Q_TILE_SIZE = 16 and K_TILE_SIZE = 16, and the head dimension is D = 64. Each iteration of the loop over the keys works on:

Tensor	What it holds	Shape
`Q_i`	16 queries, 64 numbers each	16 × 64
`K_j`, `V_j`	16 keys and their 16 values, 64 numbers each	16 × 64
`S` = `Q_i K_jᵀ`	the scores: one row per query, one column per key	16 × 16
`P`	`S` after the exponential	16 × 16
`O` = `O` + `P V_j`	the output: one row per query, 64 numbers each	16 × 64

2. A program has 4 warps. A warp is 32 threads that execute the same instruction together. The baseline uses Triton's default num_warps = 4, so one program is 128 threads.

3. The tiles live in registers. S, P and O are not stored in memory. Their elements are spread over the registers of the 128 threads, a few per thread, and a thread can read only its own registers.

4. Sharing inside a warp is cheap; sharing between warps is not. Two threads in the same warp exchange a value register to register with one shuffle (SHFL). Two threads in different warps must go through shared memory: one thread stores (STS), all the warps of the program wait at a barrier (BAR.SYNC, CUDA's __syncthreads()) until the slowest arrives, then the other thread loads (LDS).

5. One tensor-core instruction produces one 16 × 8 block. The instruction is HMMA.16816, run by a whole warp. It multiplies a 16 × 16 piece of the left matrix by a 16 × 8 piece of the right one and adds the result into a 16 × 8 block of the output: 16 × 8 × 16 = 2,048 multiplies and 16 × 8 × 16 = 2,048 adds, hence the roofline counts as 2 x 2,048 = 4,096 operations. Each output element is a sum of products, and one HMMA adds 16 terms of that sum, so a block whose sum has K terms needs K / 16 HMMAs.

Why is each output element a sum of products?
That is simply how matrix multiplication is defined. To get one number in the result, you take one row of the left matrix and one column of the right matrix, multiply them pair by pair, and add up the products.

A tiny example, C = A × B:

A = [ 1  2  3 ]      B = [ 7 ]
                         [ 8 ]
                         [ 9 ]

C[0][0] = 1·7 + 2·8 + 3·9 = 7 + 16 + 27 = 50

The row of A has 3 numbers and the column of B has 3 numbers, so you get 3 products, and C[0][0] is their sum. Every element of a matrix product is built this way. The number of terms in the sum equals the length of the “shared” dimension, the one that disappears in the multiplication. People usually call it K.

In the attention kernel:

S = Q Kᵀ: one score is one query (64 numbers) paired with one key (64 numbers). That gives 64 products, so the sum has 64 terms (K = 64, the head dimension).
O = P V: one output number pairs a row of P (16 values, one for each key in the tile) with a column of V (16 values). That gives 16 terms (K = 16, the key tile size).

Why does one HMMA add exactly 16 terms?
Because the tensor core is a fixed-size circuit, and 16 is the size Nvidia built it for. The name says so where HMMA.16816 means M = 16, N = 8, K = 16.

M = 16: it produces 16 rows of output.
N = 8: it produces 8 columns of output.
K = 16: for each output element, it handles 16 terms of the sum.

Think of it as a cookie cutter of a fixed shape. One instruction does this:

 left piece        right piece         output block
 16 rows ×      ×  16 rows ×       →   16 rows × 8 cols
 16 cols           8 cols              (added onto what's already there)
     ↑                ↑
     └── these 16 are the 16 terms of the sum ──┘

For each of the 16 × 8 = 128 output elements, the hardware takes 16 numbers from a row of the left piece and 16 numbers from a column of the right piece. It multiplies them pair by pair, adds the 16 products, and adds that total onto the running value in the output. That is 128 elements × 16 multiply-adds = 2,048 multiply-adds. Counting each multiply and each add as one operation gives 4,096 operations.

What the compiler chose

The TTGIR gives every tensor a layout: which warp, and which thread in that warp, holds each element (fact 3). This is the layout of S, the result of the first tl.dot:

#mma = #ttg.nvidia_mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [1, 4], instrShape = [16, 8]}>

instrShape = [16, 8] is fact 5: each warp builds its share of the result out of 16 × 8 blocks.
warpsPerCTA = [1, 4] places the 4 warps over the tile: 1 along the rows, 4 along the columns. The warps sit side by side, each taking different columns. (CTA, cooperative thread array, is CUDA's name for a thread block; one Triton program runs as one block.)

Why side by side? The two tl.dots are chained: S feeds P, and P is the left input of the second tl.dot. For chained dots, Triton 3.6 puts all the warps along one axis and uses the same placement for both (warpsPerTileV2 in AccelerateMatmul.cpp). Here the axis follows the shape of O:

at least as many rows as columns: warps stacked, each taking different rows;
fewer rows than columns: warps side by side, each taking different columns.

O is 16 × 64, so the warps go side by side. Stacking would not have helped anyway because the tile is 16 rows and a block is 16 rows, so there is only one row of blocks to share out.

How the 4 warps divide each tile

Counting 16 × 8 blocks:

S is 16 × 16: 1 block tall and 2 wide, so 2 blocks for 4 warps.
O is 16 × 64: 1 block tall and 8 wide, so 8 blocks, 2 per warp.

Warps 0 and 1 take the two blocks of S. There is nothing left for warps 2 and 3, but all four warps run the same instruction stream, so a warp cannot sit out. The layout wraps around and gives warps 2 and 3 the same two blocks again: they compute them a second time. For O, warp w takes columns 8*w* to 8*w* + 7 and, wrapping around, columns 32 + 8*w* to 32 + 8*w* + 7.

There are three problems with the above setup that affects the kernel's efficiency.

Problem 1: the first multiply is done twice

Using the rule from fact 5:

S = Q Kᵀ: each score sums over the 64 numbers of a query and a key, so 64 / 16 = 4 HMMAs per block. Two blocks need 8.
O = P V: each output number sums over the 16 keys of the tile, so 16 / 16 = 1 HMMA per block. Eight blocks need 8.

Multiply	Blocks	`HMMA`s per block	`HMMA`s needed	`HMMA`s each warp runs	`HMMA`s all 4 warps run
`Q Kᵀ`	2	4	8	1 block × 4 = 4	16
`P V`	8	1	8	2 blocks × 1 = 2	8
Both			16	6	24

24 / 16 = 1.5, the ratio from the roofline. The Source page of Nsight Compute agrees: per warp and per iteration, the Q Kᵀ line executes 4 HMMAs and the P V line 2.

Scaling up to the whole kernel

Each program runs 4,096 / 16 = 256 iterations, one per key tile.
There are 4,096 / 16 = 256 programs per head and 32 heads: 8,192 programs.
So 8,192 × 256 = 2,097,152 iterations in total.
Needed: 2,097,152 × 16 HMMAs × 4,096 operations = 137.4 G operations.
Executed: 2,097,152 × 24 HMMAs × 4,096 operations = 206.2 G operations.

The extra 68.7 G operations are warps 2 and 3 computing S a second time.

Problem 2: every row maximum and row sum goes through shared memory

tl.max(S, axis=1) needs the largest of each row's 16 scores, and tl.sum(P, axis=1) the total of each row's 16 values of P. With the warps side by side, every row is cut in two: warp 0 holds keys 0–7 and warp 1 holds keys 8–15 (warps 2 and 3 hold copies). Joining the halves means crossing between warps, so it means shared memory and barriers (fact 4):

Only the first step, 4 shuffles per warp, would be needed if one warp held the whole row. Everything after it exists only because the row is split. tl.sum does the same with an add in place of the max.

The counts match Step 3 exactly. Per warp and per iteration, each of tl.max and tl.sum executes 11 shuffle and shared-memory instructions (4 SHFL, then 3 STS, 3 LDS and 1 SHFL). Over the whole run that is 11 × 4 warps × 2,097,152 iterations = 92.3 million, the "92 million" that Compute Workload Analysis reported for each line.

Problem 3: `P` also goes through shared memory

P is the left input of P V, and each output number sums over all 16 keys of the tile: O[i][c] = P[i][0]·V[0][c] + … + P[i][15]·V[15][c]. So every warp needs P for all 16 keys, but each warp only computed P for 8 of them. The other 8 are in another warp's registers.

In the TTGIR this is a ttg.local_alloc of P in the #mma layout followed by a ttg.local_load into the operand layout of the second tl.dot. On every iteration the warps store all of P (16 × 16 bf16 values, 512 bytes) and every warp loads all of it back, with a barrier before the store and one after. I return to this in the P cast.

What are BAR, STS, and LDSM

These are SASS instructions, the GPU's actual machine code (what Triton's output finally becomes). All three deal with shared memory, a small, fast scratch area inside each SM that all the warps of a program can read and write.

Each thread keeps its working values in its own registers, which no other thread can see. When one warp needs a value that another warp holds, the value has to go through shared memory. These three instructions are how that happens.

STS: Store to Shared memory

A thread copies a value from its registers into shared memory, so other warps can get at it.

thread's registers ──STS──► shared memory

BAR.SYNC: Barrier (synchronise)

A "wait for everyone" point. When a warp reaches it, the warp stops until every warp in the program has also reached it. It is CUDA's __syncthreads().

It is needed because warps run at slightly different speeds. Without a barrier, warp 1 might try to read a value before warp 0 has finished storing it, and get garbage.

warp 0  ───STS──────►│
warp 1  ──────STS───►│ BAR.SYNC: nobody continues until all 4 arrive
warp 2  ─STS────────►│
warp 3  ────────STS─►│──► now it is safe to read

The cost is that the fastest warp sits idle waiting for the slowest one.

LDSM: Load from Shared memory, as a Matrix

A thread copies data from shared memory back into its registers, specifically in the arrangement the tensor cores want for an HMMA input. It is CUDA's ldmatrix. A whole warp runs it together, and each thread ends up holding exactly the pieces of the matrix it needs to feed the next HMMA.

shared memory ──LDSM──► registers, arranged ready for HMMA

Its plain cousin is LDS (Load from Shared memory). LDS loads ordinary values, as in the tl.max and tl.sum reductions, while LDSM loads matrix pieces for the tensor cores.

How they fit together: P's round trip in the baseline

This is exactly what Problem 3 in the article describes. Each warp holds only half of each row of P, but the second tl.dot needs whole rows, so every loop iteration does:

1. BAR.SYNC   wait until everyone is done with the old contents of the buffer
2. STS        each warp writes its half of P into shared memory
3. BAR.SYNC   wait until every half has been written
4. LDSM       each warp reads back all of P, ready for the P V multiply

That is four steps of shuffling data that does no arithmetic. With 64-row tiles, each warp already owns whole rows of P, so all four steps disappear and P goes straight from one HMMA into the next.

Counting the barriers

Each iteration has nine BAR.SYNCs, and at each one all 4 warps stop and wait for the slowest:

Where	`BAR.SYNC`s per iteration	Why
Copying the next K and V tiles into shared memory	2	they guard the shared-memory buffers the asynchronous copies write into
`tl.max`	2	the two in the row-maximum figure
`tl.sum`	3	the same two, plus one before it starts, because it reuses the shared memory `tl.max` has just read
`P`'s trip through shared memory	2	one before the store, one after
Total	9

Only the first 2 are needed in any layout. A program runs 256 iterations, so it stops at 256 × 9 = 2,304 barriers to produce 16 rows of output.

What a bigger tile changes

All three problems come from the tile being one block tall. With Q_TILE_SIZE = 64, the tile is 64 / 16 = 4 blocks tall, one row of blocks per warp. O becomes 64 × 64, with as many rows as columns, so Triton stacks the warps (warpsPerCTA = [4, 1]):

Each warp now owns 16 whole rows of S, P and O, and each problem disappears:

Problem 1: every warp has its own rows, so no warp repeats another's blocks.
Problem 2: a whole row sits in one warp, so the row maximum and row sum need only shuffles.
Problem 3: the warp that computes a row of P is the warp that uses it, so P stays in registers.

By compiling the same kernel with 64 × 32 tiles (64 queries, 32 keys) and 4 warps, one iteration now covers 64 × 32 = 2,048 scores, 8 times the 256 of the baseline, so it should need 8 × 16 = 128 HMMAs:

In one loop iteration	16 × 16 tile, warps side by side	64 × 32 tile, warps stacked
`HMMA`s needed	16	128
`HMMA`s all 4 warps run	24	128
`STS` instructions (stores to shared memory)	8	0
`BAR.SYNC` instructions	9	2

Nothing is computed twice, nothing is stored to shared memory, and the 2 remaining barriers are the ones for the K and V copies. Round 1 finds this configuration with a sweep.

The diagnosis

Everything points to one cause: a 16 × 16 tile is too small for 4 warps. The 4 warps cannot each get their own rows, so they duplicate tensor-core work and synchronise through shared memory on every row reduction. On top of that, 16-row tiles make every program re-read 1 MB of K and V from L2 for 16 rows of output.

Nsight Compute's own top suggestion, theoretical occupancy (the 66.7% limit from shared memory), would not have helped: more resident warps of a kernel that is limited by its own shared-memory instructions would only queue more shared-memory instructions. The tuned kernel in Round 1 runs at half that occupancy and three times the speed.

What are resident warps, why more warps do not help, and why Nsight Compute gave the wrong suggestion

What resident warps are

An SM (one of the 56 processing units on the GPU) can hold several warps at once. A warp is resident when it has been loaded onto the SM: its registers and shared memory are reserved, and it is ready to run as soon as the SM picks it.

On Nvidia RTX 4070 SUPER an SM can hold at most 48 warps. Occupancy is how many it actually holds, as a share of that maximum.

In the baseline, each program is 4 warps, and shared memory limits each SM to 8 programs at a time:

8 programs × 4 warps = 32 resident warps
32 ÷ 48 = 66.7% occupancy

Why more warps usually help

An SM can issue an instruction from only a few warps per clock cycle. When a warp is stuck waiting (for data to arrive, say), the SM simply switches to another resident warp that is ready. More resident warps means more candidates to switch to, so the SM is less likely to sit idle. This is called hiding latency.

warp A: ██ wait........... ██
warp B:    ██ wait......... ██
warp C:       ██ wait....... ██     ← while A waits, B and C keep the SM busy

This works when warps are waiting on something slow but not overloaded, like DRAM. The waiting is just time passing, and other warps can fill it.

Why it would not help here

The baseline's warps were not waiting on something slow. They were waiting on something full. Round 0 found that:

the load/store unit (which handles shared memory) was already 80% busy;
the top stalls were MIO throttle (the queue for shared-memory instructions is full), short scoreboard (waiting for shared-memory results) and barrier (waiting for other warps in the same program).

Every warp in this kernel does the same thing: lots of shared-memory traffic for tl.max, tl.sum and P's round trip. Adding more warps adds more shared-memory requests to a queue that is already full.

Think of a kitchen with one sink. If the cooks are slow because they all queue at the sink, hiring more cooks just makes the queue longer. What helps is needing the sink less.

That is what Round 1 did. Bigger tiles let each warp own whole rows, which removed almost all the shared-memory traffic. The tuned kernel runs at 33.3% occupancy, half the baseline's, and is three times faster, because its warps now wait on the tensor cores, which is useful work.

Why Nsight Compute suggested it anyway

Nsight Compute's suggestions come from general rules, not an understanding of your kernel. The occupancy rule goes roughly like this:

Warps spend a lot of time stalled, so the SM often has nothing ready to issue.
Occupancy is below the maximum.
Therefore more warps might fill the gaps, worth an estimated ~20%.

That reasoning is right for many kernels. But the rule does not check what the warps are waiting for. It cannot tell "waiting on a slow unit that has spare capacity" from "waiting on a unit that is already saturated." So it named a symptom (lots of stalling, low occupancy) instead of the cause (too much shared-memory work from tiles that are too small).

This is why the article treats the suggestions as hypotheses to check, not diagnoses. The breakdown and the stall reasons showed what was actually busy, and that pointed to the real fix.

Round 1: tile sizes, `num_warps` and `num_stages`

Round 0 says the tile is too small for its warps. Before sweeping, it helps to know what each knob changes, and what limits how far each one can go.

Q_TILE_SIZE: query rows per program. Every program streams all of K and V through the SM once, so the L2 traffic is proportional to N² / Q_TILE_SIZE. Bigger is better for traffic, but the fp32 output accumulator O_acc (Q_TILE_SIZE × D) and the score tile (Q_TILE_SIZE × K_TILE_SIZE) live in registers, and the program count drops, which matters at small N.
K_TILE_SIZE: keys per step of the inner loop. Bigger tiles mean fewer loop iterations and bigger matrix multiplies, at the cost of registers for the score tile and shared memory for each K and V buffer.
num_warps: how many warps (groups of 32 threads) share one program's tile. Round 0 showed what happens when the tile cannot give each warp its own rows.
num_stages: the depth of the software pipeline for the K and V loads. Triton's software pipeliner turns the loads in the loop into asynchronous copies into shared memory, issued num_stages - 1 iterations ahead, so the next tiles arrive while the current one is being multiplied. More stages hide more latency and cost more shared memory.

How to think about num_stages

Why more stages isn't always better. Each stage adds one K buffer and one V buffer in shared memory. For the article's 64 × 32 tile that's 8 KB per stage per program, on top of the 8 KB Q tile. An Ada SM has only about 100 KB of shared memory, and every program running on it takes its share. More stages per program means fewer programs fit on the SM at once:

3 stages: 24 KB each, so 4 programs fit side by side.
5 stages: 40 KB each, so only 2 fit.

Fewer resident programs means fewer warps to switch to when one stalls. Once the pipeline is deep enough to keep the tensor cores fed, adding more stages costs memory and occupancy without speeding anything up.

That's why the sweep found 3 stages best (63.5 TFLOP/s, against 57.9 with 1 stage), and 4 and 5 slightly slower again. The article's padding experiment shows that about half of that slowdown is the lost occupancy, and the rest is the deeper pipeline itself.

Why the best number changes. The ideal depth depends on how long a fetch takes compared with a multiply:

Long fetch, short multiply: you need more tiles in flight to cover the wait.
Bigger key tile (more work per iteration): fewer stages suffice. That's why Round 5's 64 × 128 tile was happy with just 2 stages, since each multiply is long enough to hide most of the next fetch.

Two caveats about the toy model:

The 4 : 2 ratio between fetch and multiply time is made up to make the picture clear. Real latencies vary, which is why you sweep rather than calculate.
On the real GPU, the 1-stage kernel isn't as bad as the toy suggests (57.9 vs 63.5 TFLOP/s, not 3× slower). Each SM runs several programs at once, so while one waits for its fetch, another can use the tensor cores. Stages hide latency within a program; occupancy hides it across programs. That's the same trade-off the shared-memory budget forces you to balance.

The shared-memory budget for an Ada SM

Shared memory is what decides which configurations can launch at all, so it is worth knowing exactly what the kernel keeps there. The TTGIR of the configuration that eventually wins (64 × 32 tiles, 4 warps, 3 stages) allocates three buffers:

%Q_i = ttg.local_alloc %Q_i_53 : (tensor<64x64xbf16, #blocked>) -> !ttg.memdesc<64x64xbf16, #shared, #smem>
%K_j = ttg.local_alloc : () -> !ttg.memdesc<2x32x64xbf16, #shared, #smem, mutable>
%V_j = ttg.local_alloc : () -> !ttg.memdesc<2x32x64xbf16, #shared, #smem, mutable>

The Q tile, Q_TILE_SIZE × D bf16 values. The kernel loads Q_i once, but Triton parks it in shared memory and re-reads it in every iteration (ttg.local_load, LDSM in the machine code) as the left-hand operand of Q Kᵀ, rather than holding it in registers for the whole loop.
The K and V buffers, num_stages - 1 of each, K_TILE_SIZE × D bf16 values per buffer. 2x32x64 is two buffers of a 32 × 64 tile: three stages means two tiles in flight while the third is in registers being multiplied. The copies into them are ttg.async_copy_global_to_local, which becomes LDGSTS (load global, store shared, without passing through registers).
Nothing for S, P, O, m or l: those live in registers, provided the warps own whole rows. When they do not, P makes a round trip through shared memory, which is the extra ttg.local_alloc ... #mma from Round 0.

For bf16 inputs (2 bytes) and D = 64, that gives a budget equation:

shared bytes = 2 \cdot D \cdot Q_{tile} Q + (stages - 1) \cdot 2 \cdot 2 \cdot D \cdot K tile K and V = 128 \cdot (Q tile + 2 (stages - 1) K_{tile})

and two limits from the table at the top:

One program must fit: shared bytes ≤ 101,376. Triton checks this when it loads the kernel and raises OutOfResources if not.
Programs per SM: the 102,400 bytes of an SM are shared by every resident program, and each also costs 1,024 bytes of driver reservation, so at most ⌊102,400 / (shared bytes + 1,024)⌋ programs fit at once.

Why the budget equation matters

Shared memory is one of the scarcest resources on this GPU. Every program (thread block) needs some to hold its Q tile and its K and V pipeline buffers. On this Ada card there isn't much to go around: about 100 KB per SM, against 164 KB on an A100 and 228 KB on an H100.

The equation lets you work out before compiling anything how much shared memory a given choice of tile sizes and stages will need. From that one number you can answer two practical questions:

Will this configuration run at all? If it asks for too much, Triton refuses to launch it (OutOfResources). This is why 128 × 128 tiles with 4 stages, the textbook FlashAttention-2 choice, simply doesn't work on this card.
How many programs can share one SM at a time? More programs per SM means that while one program waits for data, another can use the tensor cores. The winning configuration is fast partly because it is small: at 24,576 bytes, four programs fit on each SM at once.

It also explains the trade-offs the sweep found. Each extra pipeline stage hides more memory latency but adds 8 KB. That can push the count from 4 programs per SM down to 3 or 2, which is part of why 4 and 5 stages ran slower than 3.

The two numbers

They are two different limits, one for the whole SM and one for a single program.

102,400 bytes (100 KB): the SM's total pool. Each SM has 128 KB of on-chip memory, split between L1 cache and shared memory. At most 100 KB of that can be used as shared memory. All the programs running on that SM share this pool.
101,376 bytes (99 KB): the most any single program may ask for. On top of what it asks for, the driver reserves 1,024 bytes (1 KB) for every program for its own bookkeeping.

The two numbers fit together exactly: one program at the maximum (101,376) plus its 1 KB reservation is 102,400, the whole pool. That's why a single program can't ask for the full 100 KB. The last 1 KB is always taken by the driver.

An analogy: the SM is a building with 100 rooms. Each tenant (program) can rent at most 99 rooms, and every tenant is also charged 1 room for the lobby.

Using both limits

The first limit decides whether a configuration launches. The second decides how many copies of it fit on one SM.

Configuration	Shared memory	Fits one program? (≤ 101,376)	Programs per SM: ⌊102,400 / (shared + 1,024)⌋
64 × 32, 3 stages	24,576	yes	102,400 / 25,600 = 4
64 × 32, 5 stages	40,960	yes	102,400 / 41,984 = 2
256 × 64, 5 stages	98,304	yes	102,400 / 99,328 = 1
128 × 128, 4 stages	114,688	no	doesn't launch

Shared memory is only one of the two limits on programs per SM. Registers can bind first: 128 × 64 tiles with 8 warps would fit two programs by shared memory, but needs too many registers for more than one. The article covers this right after the budget equation.

Registers give a second, independent limit. Each SM has 65,536 registers, handed out in blocks of 8 per thread, so at most ⌊65,536 / (registers rounded up to a multiple of 8 × 32 × num_warps)⌋ programs fit. The registers a configuration needs are only known after compiling. The largest single block of them holds the two fp32 accumulators, (Q_TILE_SIZE × K_TILE_SIZE + Q_TILE_SIZE × D) values spread over 32 × num_warps threads, 18 to 63% of the total in the table below. When that alone passes 255 per thread, the compiler always spills to local memory (which lives in DRAM, behind the caches), and the kernel slows down sharply; wide key tiles can spill well before that.

How to interpret the three lines of TTGIR-generated code
These three lines are Triton telling you what it keeps in shared memory for the winning configuration (64 × 32 tiles, 4 warps, 3 stages). Shared memory is the small, fast scratchpad inside each SM that all warps of a program can see. Each line creates one buffer there.

How to read the notation

A TTGIR line has the form result = operation input : (input type) -> output type.

%Q_i, %K_j, %V_j are names for the results. Triton reuses your Python variable names. Suffixes like _53 are just the compiler's numbering of intermediate values.
ttg.local_alloc means "allocate a buffer in shared memory". "Local" means local to one program (thread block). It is not about per-thread local memory.
tensor<...> is a value held in registers, spread across the threads.
!ttg.memdesc<...> ("memory descriptor") is not a value. It is a handle to a region of shared memory, something like a typed pointer.
64x64xbf16 is the shape and element type.
#blocked, #shared are layouts: rules for which thread or which address holds which element.
- #blocked is a simple register layout that Triton uses for loading from global memory.
- #shared describes how the data is arranged in shared memory. This arrangement is usually swizzled so that threads reading together don't collide on the same memory bank.
#smem says the memory space is shared memory.
mutable means the buffer will be overwritten repeatedly.

Line 1: the Q tile

%Q_i = ttg.local_alloc %Q_i_53 : (tensor<64x64xbf16, #blocked>) -> !ttg.memdesc<64x64xbf16, #shared, #smem>

In plain English: "Take the Q tile that was just loaded into registers (%Q_i_53, 64 query rows × 64 head-dim values in bf16), copy it into a new shared-memory buffer of the same shape, and call that buffer %Q_i."

Because it has an input, this allocation is created and filled in one step.
It has no mutable, so it is written once and only read afterwards.
Size: 64 × 64 × 2 bytes = 8,192 bytes.

Why it matters: Python loads Q_i once, before the loop. Instead of staying in registers. Triton parks it in shared memory and re-reads it on every loop iteration (ttg.local_load, which becomes LDSM in machine code) to feed the tensor cores for Q Kᵀ. So Q costs shared memory for the whole kernel, and the budget formula has to include it.

Line 2: the K buffers

%K_j = ttg.local_alloc : () -> !ttg.memdesc<2x32x64xbf16, #shared, #smem, mutable>

In plain English: "Reserve an empty shared-memory buffer with room for 2 K tiles, each 32 keys × 64 values in bf16. Its contents will change over time."

() means there is no input, so the buffer starts empty and is filled later. The filling is done by asynchronous copies (ttg.async_copy_global_to_local, which becomes LDGSTS). These copy straight from GPU memory into shared memory without passing through registers.
2 = num_stages − 1 = 3 − 1. This is the software pipeline: while the warps multiply the current tile, the next two tiles are already being copied into these two slots. The buffer is used as a ring, with each slot overwritten in turn, hence mutable.
32 = K_TILE_SIZE and 64 = D.
Size: 2 × 32 × 64 × 2 bytes = 8,192 bytes.

Line 3: the V buffers

%V_j = ttg.local_alloc : () -> !ttg.memdesc<2x32x64xbf16, #shared, #smem, mutable>

This is identical to K, but for the V tiles: 2 slots of 32 × 64 bf16, another 8,192 bytes.

Significance of the three lines of code

1. They justify the budget equation. Adding the three buffers gives:

8,192 (Q) + 8,192 (K) + 8,192 (V) = 24,576 bytes

That is exactly what the formula predicts: 128 × (64 + 2 × (3−1) × 32) = 24,576. It also matches compiled.metadata.shared byte for byte. So the formula isn't a guess. It is read straight off what the compiler allocates, and it can predict whether any configuration will launch before compiling it.

2. Shared memory decides what can run, and how many copies at once.

Fit limit: one program may use at most 101,376 bytes. Configurations over that limit crash with OutOfResources, such as 128 × 128 with 4 stages, or the fp32/D = 128 test failure.
Programs per SM: each SM has 102,400 bytes to share, plus 1,024 bytes reserved per program. That gives ⌊102,400 / (24,576 + 1,024)⌋ = 4 programs per SM. This is part of why the winner wins: it uses only a quarter of the budget, so four programs share an SM and hide each other's waiting.
Why more stages hurt: each extra stage adds one more 32 × 64 slot to both K and V, which is 8 KB. Four programs per SM then drops to 3, then 2, which is part of the slowdown seen at 4 and 5 stages.

3. What's missing is as informative as what's there. There is no buffer for S, P or O. Those stay in registers, which only happens when each warp owns whole rows of the tile. In the 16 × 16 baseline there was a fourth local_alloc, which put P into shared memory on every iteration. That cost 512 extra bytes, plus stores, loads and barriers. Seeing only these three allocations confirms the layout problem from Round 0 is gone. This is why the checklist (item 17) says to scan the TTGIR for extra ttg.local_allocs inside the loop.

4. They explain the remaining barriers. Because the K and V buffers are mutable and reused as a ring, the warps must synchronise so that no slot is overwritten while another warp is still reading it. Those are the two BAR.SYNCs per iteration that Round 2 traced to the tl.load of K. Round 5 then cut their cost by using 128-key tiles, which means fewer iterations and therefore fewer barriers.

Here is the arithmetic for a few configurations, next to what Triton reported in compiled.metadata.shared, n_regs and n_spills, and the speed the sweep below measured at N = 4096:

Q × K tile, warps, stages	Shared memory (formula)	Measured	Programs / SM (shared)	Registers / thread	Programs / SM (registers)	TFLOP/s
16 × 16, 4, 3 (baseline)	10,240	10,752 (+512 for `P`)	8	56	9	21.8
64 × 32, 4, 3	24,576	24,576	4	127	4	63.5
64 × 32, 4, 5	40,960	40,960	2	127	4	61.9
128 × 64, 8, 4	65,536	65,536	1	146	1	61.9
128 × 128, 8, 3	81,920	81,920	1	207	1	60.5
128 × 128, 8, 4	114,688	does not fit	0
256 × 64, 16, 5	98,304	98,304	1	128 (4 spilled)	1	61.3
64 × 256, 4, 2	73,728	73,728	1	255 (152 spilled)	2	33.4

How are the columns for Programs/SM (shared), Registers/thread and Programs/SM (registers) calculated?

The three columns come from two different sources. Programs / SM (shared) and Programs / SM (registers) are arithmetic. Registers / thread is not calculated at all: the compiler reports it.

Programs / SM (shared)

An Ada SM has 102,400 bytes of shared memory. Every resident program uses its own shared memory plus 1,024 bytes that the driver reserves per block. So:

programs per SM = ⌊102,400 / (shared bytes + 1,024)⌋

The "shared bytes" here is the Measured column, which is compiled.metadata.shared. The formula column is only used to predict it.

Config	Measured + 1,024	102,400 ÷ that	Programs
16×16, 4, 3	11,776	8.70	8
64×32, 4, 3	25,600	4.00	4
64×32, 4, 5	41,984	2.44	2
128×64, 8, 4	66,560	1.54	1
128×128, 8, 3	82,944	1.23	1
256×64, 16, 5	99,328	1.03	1
64×256, 4, 2	74,752	1.37	1

The baseline row shows why the measured value matters. The formula's 10,240 bytes would give ⌊102,400 / 11,264⌋ = 9. The extra 512 bytes for the P buffer brings it down to 8.

For 128×128, 8, 4, the 114,688 bytes exceed the 101,376-byte limit for a single block. The program can't launch, so the value is 0.

The formula column itself is 128 · (Q_tile + 2·(stages − 1)·K_tile). For example, 64×32 with 3 stages gives 128 × (64 + 2·2·32) = 128 × 192 = 24,576.

Registers / thread

This is compiled.n_regs, read after compiled._init_handles() (see inspect_kernel.py). The spill counts come from compiled.n_spills. These are calculated after compilation.

To estimate the lower bound, the two fp32 accumulators hold Q×K + Q×D values spread over 32 × num_warps threads:

64×32, 4 warps: (2,048 + 4,096) / 128 = 48 registers, out of 127 actual.
Baseline 16×16, 4 warps: (256 + 1,024) / 128 = 10, out of 56.
64×256, 4 warps: (16,384 + 4,096) / 128 = 160, out of 255. That is the 63%. The rest of the kernel doesn't fit under the 255 cap, hence the 152 spilled registers.

The other registers hold pointers, m, l, loop state and temporaries. That's why the real number has to come from the compiler.

Programs / SM (registers)

An SM has 65,536 registers. They are allocated per thread in multiples of 8, which is 256 per warp. So:

programs per SM = ⌊65,536 / (round_up_8(regs) × 32 × num_warps)⌋

Config	Regs → rounded	× 32 × warps	65,536 ÷ that	Programs
16×16, 4 warps	56 → 56	7,168	9.14	9
64×32, 4 warps	127 → 128	16,384	4.00	4
128×64, 8 warps	146 → 152	38,912	1.68	1
128×128, 8 warps	207 → 208	53,248	1.23	1
256×64, 16 warps	128 → 128	65,536	1.00	1
64×256, 4 warps	255 → 256	32,768	2.00	2

The 64×32, 5-stage row has the same registers and warps as the 3-stage one, so it also gets 4.

Putting them together

The number of programs that actually run on an SM at once is the minimum of the two columns. For example:

Baseline: min(8, 9) = 8, so shared memory binds.
64×32, 4, 5: min(2, 4) = 2, so shared memory binds. This is why extra stages cost occupancy.
64×256, 4, 2: min(1, 2) = 1, so shared memory binds, though it is slow for a different reason (spills).

There is also a third cap of 48 warps and 24 blocks per SM. It never binds in this table: the most warps resident is the baseline's 8 × 4 = 32.

If limit is 255 for registers, why 128 (4 spilled)?

The 255 is only the per-thread ceiling. A second limit applies at the same time: all the threads of one program must fit in the SM's 65,536 registers together, otherwise the program could never launch. The compiler enforces whichever limit is lower.

max registers per thread = min(255, 65,536 / threads per program)

Triton tells the compiler the block size by putting .maxntid (num_warps × 32 threads) in the PTX. ptxas then caps the register count so that one block fits:

`num_warps`	Threads per program	65,536 ÷ threads	Effective cap
1, 2, 4	32, 64, 128	2,048, 1,024, 512	255
8	256	256	255
16	512	128	128
32	1,024	64	64

The 256×64 configuration uses 16 warps, so its cap is 128 registers per thread, not 255. It needed slightly more than 128, so the compiler kept the excess in local memory: the "4 spilled". (Triton's n_spills is the local-memory size divided by 4 bytes, so roughly four 32-bit values per thread.)

The accumulator estimate shows why it lands right at that edge:

S: 256 × 64 = 16,384 values
O_acc: 256 × 64 = 16,384 values
32,768 values ÷ 512 threads = 64 registers per thread for the accumulators alone

That leaves 64 registers for pointers, m, l, operand fragments and temporaries. The kernel needed slightly more, hence a small spill.

This is also why the same row shows exactly 1 program per SM by registers: 128 × 512 = 65,536, the whole register file. This program has 512 threads, and each thread uses 128 slots. So the program takes 512 × 128 = 65,536 slots, which is the entire box. Nothing is left over, so a second program can’t run on that SM until the first one finishes. That’s why the table says only 1 program fits per SM.

Across the whole sweep, the formula matches compiled.metadata.shared byte for byte for every configuration with 2 to 5 stages whose warps own whole 16-row slices of the query tile. The exceptions are the ones where Triton spreads the warps along the key axis, and there the difference is the buffer for P. With 1 stage nothing is in flight, but Triton still passes each V tile through a shared-memory buffer of its own, so the formula comes out one tile short.

Explain the two exceptions in detail

The formula is a packing list for shared memory. It assumes the kernel stores exactly three kinds of things there:

one Q tile, and
(stages − 1) K tiles, and
(stages − 1) V tiles.

For most configurations, that list is complete, so the predicted bytes match what Triton actually allocates. The paragraph names two cases where the list is wrong.

Exception 1: the warps sit side by side, so `P` needs its own buffer

Usually each warp owns a horizontal strip of the tile: whole rows of 16 queries. A warp then computes its own rows of P and uses them in the second multiply (P @ V), so P never leaves the warp's registers.

If the query tile is short compared to the head dimension (for example 16 rows with D = 64), Triton places the warps side by side along the keys instead. Each warp then holds only part of every row of P. The second multiply needs whole rows, so the warps have to swap pieces, and the only place they can swap is shared memory.

That swap space is a buffer the formula doesn't count. The baseline shows it: the formula predicts 10,240 bytes, Triton allocates 10,752. The extra 512 bytes is one 16 × 16 tile of P in bf16 (16 × 16 × 2 bytes).

Exception 2: one stage, but V still uses shared memory

With num_stages = 1 there's no prefetching. Nothing is loaded ahead of time, so the formula's (stages − 1) term becomes zero and predicts that K and V use no shared memory at all.

In practice, Triton still routes each V tile through a shared-memory buffer on its way to the second multiply, even though no tile is being loaded early. The likely reason is that V has to be rearranged into the layout the tensor cores expect. So the real usage is the formula's number plus one V tile. For a 64 × 32 tile with D = 64, that's 8,192 predicted and 8,192 + 4,096 = 12,288 actual.

Three things stand out. First, the classic FlashAttention-2 choice of 128 × 128 tiles fits on Ada only up to 3 stages, and then a single program fills the SM. Second, the configurations that run fastest are not the ones that use the most shared memory: 64 × 32 with 3 stages uses a quarter of the budget, which lets four programs share each SM and hide each other's latency. Third, registers bind as often as shared memory does: 128 × 64 with 8 warps and 3 stages would fit two programs by shared memory (49,152 bytes each), but its 146 registers per thread allow only one, and a 256-key tile with 64 or more query rows spills no matter how much shared memory is left.

The sweep

The configuration space is small enough to measure exhaustively: 5 query tiles × 5 key tiles × 5 warp counts × 5 stage counts = 625 configurations, at each of the harness's 8 shapes (N = 512 to 4096, causal and not). sweep.py times each one exactly the way the harness does: FlashAttentionTriton.apply under triton.testing.do_bench, median, same inputs. Every configuration is also checked against SDPA once per shape, so a fast-but-wrong configuration cannot win.

TILES = (16, 32, 64, 128, 256)
WARPS = (1, 2, 4, 8, 16)
STAGES = (1, 2, 3, 4, 5)
CONFIGS = list(itertools.product(TILES, TILES, WARPS, STAGES))

for causal in (False, True):
    for N in (512, 1024, 2048, 4096):
        q, k, v = (torch.randn(B * H, N, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))
        ref = F.scaled_dot_product_attention(q, k, v, is_causal=causal)
        for bq, bk, nw, ns in CONFIGS:
            try:
                ck = launch(q, k, v, causal, bq, bk, nw, ns)     # returns the compiled kernel
            except triton.runtime.errors.OutOfResources:
                record("out_of_smem"); continue
            FlashAttentionTriton.Q_TILE_SIZE, FlashAttentionTriton.K_TILE_SIZE = bq, bk
            FlashAttentionTriton.NUM_WARPS, FlashAttentionTriton.NUM_STAGES = nw, ns
            o = FlashAttentionTriton.apply(q, k, v, causal)
            if not torch.allclose(o, ref, atol=2e-2, rtol=2e-2):
                record("wrong"); continue
            ms = triton.testing.do_bench(lambda: FlashAttentionTriton.apply(q, k, v, causal),
                                         return_mode="median")
            record("ok", ms, ck.n_regs, ck.n_spills, ck.metadata.shared)

For the sweep, the only change to the kernel file is that forward reads the four values from class attributes and passes num_warps and num_stages to the launch, which the previous version left at Triton's defaults of 4 and 3.

Two practical details to speed up the sweep:

Compilation dominates when every configuration is new, so sweep.py first compiles all of them in 6 worker processes with flash_fwd_kernel.warmup(...), which compiles into Triton's on-disk cache without launching, and only then benchmarks them one at a time in a single process (90 seconds of compiling, then benchmarking). 2. Configurations whose accumulators alone would need more than 512 fp32 registers per thread are skipped without compiling.

Per shape, 465 configurations ran, 100 did not fit in shared memory, and 60 were skipped by the register estimate. None produced a wrong answer.

The heatmaps

Each cell is the best of the 25 num_warps × num_stages combinations for that tile shape, at N = 4096, non-causal. The winner is 64 × 32 with 4 warps and 3 stages at 63.5 TFLOP/s, against 21.8 for the baseline configuration. The cell at the top-left is suprising because it has the same 16 × 16 tile as the baseline while reaching 47.5 TFLOP/s with 1 warp and 1 stage instead of 4 warps and 3 stages (1 warp alone, with the baseline's 3 stages, gives 37.7). More than half of the gap to SDPA was the launch configuration, not the tile size.

Taking the best result over key tiles and stages for each query tile and warp count:

Every query tile is fastest with about one warp for every 16 rows: num_warps = Q_TILE_SIZE / 16, or one step below. That means 1 warp for 16 rows, 4 for 64, 8 for 128 and 16 for 256. The only tile that takes the step below is the 32-row one, which is fastest with 1 warp, not 2. Doubling the warps beyond this halves the speed of the larger tiles. 64-row tiles drop from 63.5 to 33.0 TFLOP/s with 8 warps, and 128-row tiles drop from 61.9 to 30.7 with 16.

The 16 comes from the tensor-core instruction. One HMMA makes a block 16 rows tall (Fact 5), so a query tile splits naturally into slices of 16 rows. The best split gives each warp slices of its own. A 64-row tile has four slices, so with 4 warps each one gets a slice:

64-row tile, 4 warps stacked: one warp per 16-row slice

              ┌──────────────┐
rows  0–15    │    warp 0    │
rows 16–31    │    warp 1    │
rows 32–47    │    warp 2    │
rows 48–63    │    warp 3    │
              └──────────────┘

Whether the warps actually land like this is up to Triton, which follows the rule from Round 0. If the tile has at least as many rows as the head dimension (64 here), Triton stacks the warps one under another, and each gets different query rows. If not, it puts them side by side, and each gets different columns. That rule held across the whole sweep and in every other case I compiled. It leads to three ways of losing speed.

Extra warps on a short tile. A 16- or 32-row tile is shorter than 64, so a second warp goes side by side, and each query's scores are split between two warps:

16-row tile, 2 warps side by side

               keys 0–7     keys 8–15
           ┌─────────────┬─────────────┐
 16 rows   │   warp 0    │   warp 1    │
           └─────────────┴─────────────┘

Each query row is cut in two: warp 0 holds its first 8 scores,
warp 1 the other 8.

The softmax needs the maximum and the sum of each whole row, so the two halves have to be combined through shared memory, behind a barrier (Fact 4). This is the Round 0 problem; with 4 warps on a 16 × 16 tile, the scores are also computed twice. A second warp costs 16- and 32-row tiles 7 to 15%. It is also why the 32-row tile prefers 1 warp, because its 2 warps would go side by side instead of being stacked.

Too many warps on a tall tile. A 64-row tile with 8 warps is tall enough to be stacked, and Triton stacks all 8 (warpsPerCTA = [8, 1]). But 8 stacked warps cover 8 × 16 = 128 rows, and the tile has only 64, so the extra warps wrap around onto rows that are already taken:

64-row tile, 8 warps stacked: two warps per 16-row slice

              ┌──────────────────────────────┐
rows  0–15    │   warp 0 + warp 4 (repeat)   │
rows 16–31    │   warp 1 + warp 5 (repeat)   │
rows 32–47    │   warp 2 + warp 6 (repeat)   │
rows 48–63    │   warp 3 + warp 7 (repeat)   │
              └──────────────────────────────┘

The tensor cores now do every multiply twice, and the speed halves. Round 2 confirms this with the operation counter.

Too few warps. With fewer warps than slices, each warp owns several slices and has to hold all of their results in its registers. The output O alone is (rows per warp) × 64 numbers, shared among the warp's 32 threads:

128-row tile = eight 16-row slices

                   8 warps       2 warps       1 warp
                ┌──────────┐  ┌──────────┐  ┌──────────┐
rows   0–15     │  warp 0  │  │  warp 0  │  │  warp 0  │
rows  16–31     │  warp 1  │  │  warp 0  │  │  warp 0  │
rows  32–47     │  warp 2  │  │  warp 0  │  │  warp 0  │
rows  48–63     │  warp 3  │  │  warp 0  │  │  warp 0  │
rows  64–79     │  warp 4  │  │  warp 1  │  │  warp 0  │
rows  80–95     │  warp 5  │  │  warp 1  │  │  warp 0  │
rows  96–111    │  warp 6  │  │  warp 1  │  │  warp 0  │
rows 112–127    │  warp 7  │  │  warp 1  │  │  warp 0  │
                └──────────┘  └──────────┘  └──────────┘

O registers          32           128           256
  per thread
TFLOP/s             61.9          51.8          17.7

A thread can have at most 255 registers. With 2 warps, O takes 128 of them, and the scores and everything else push the kernel to the 255 limit (51.8 TFLOP/s). With 1 warp, O alone needs 256, which is over the limit. Values spill out to memory, and the speed drops to 17.7.

For the winning tile, the warps and stages:

The 4-warp row is the best, with 2 warps close behind (61.6 TFLOP/s at best); 1, 8 and 16 warps are far behind. Along the 4-warp row, going from 1 stage (no pipelining) to 3 is worth 10% (57.9 to 63.5 TFLOP/s); 4 and 5 stages are slightly slower again.

A deeper pipeline is slower because adding a stage changes two things at once:

The code changes. The kernel keeps one more K tile and one more V tile in flight, and the compiler schedules the loop differently.
Fewer programs fit on each SM. Each extra stage adds 8 KB of shared memory to every program. An SM has 100 KB, so with 3 stages 4 programs fit, with 4 stages 3 fit, and with 5 stages only 2 fit. Fewer programs per SM means fewer warps to cover for each other while one of them waits.

Timing the 5-stage kernel measures both effects added together. To pull them apart, I need a kernel that has the second effect without the first: the 3-stage code, running with only 2 programs per SM.

Triton makes that possible. When it launches a kernel, it hands the GPU two things separately: the compiled code, and how much shared memory to reserve for each program. So I can launch the unchanged 3-stage code but ask the GPU to reserve as much shared memory as the 5-stage kernel needs. The program never touches the extra bytes, but the GPU sets them aside anyway, so only 2 programs fit on each SM, as with 5 stages. I call this the padded kernel. It has the 5-stage kernel's occupancy and the 3-stage kernel's code.

One SM has 100 KB of shared memory. Each box is one program, including the
1 KB the driver reserves for it.

3 stages: 25 KB each, 4 programs fit
┌───────────┬───────────┬───────────┬───────────┐
│  3-stage  │  3-stage  │  3-stage  │  3-stage  │
└───────────┴───────────┴───────────┴───────────┘

5 stages: 41 KB each, 2 programs fit
┌───────────────────┬───────────────────┬────────┐
│      5-stage      │      5-stage      │ unused │
└───────────────────┴───────────────────┴────────┘

3 stages, padded to 41 KB: same code as the first row, but only 2 programs fit
┌───────────┬───────┬───────────┬───────┬────────┐
│  3-stage  │padding│  3-stage  │padding│ unused │
└───────────┴───────┴───────────┴───────┴────────┘

The same trick works for 4 stages: pad the 3-stage kernel to 33 KB, and 3 programs fit, as with 4 stages. Timed at N = 4096 (ten interleaved rounds), against the plain 3-stage kernel:

Kernel	Code	Programs per SM	Slower than 3 stages by
3 stages	3-stage	4	(reference)
4 stages	4-stage	3	1.1%
3 stages, padded to 4-stage size	3-stage	3	0.4%
5 stages	5-stage	2	3.5%
3 stages, padded to 5-stage size	3-stage	2	2.0%

The padded kernel only loses occupancy, so its slowdown is the cost of occupancy alone, and whatever the real deeper kernel loses beyond that is the cost of the deeper pipeline itself:

5 stages: 3.5% in total = 2.0% from occupancy + 1.5% from the deeper pipeline. Occupancy is about half of the loss.
4 stages: 1.1% in total = 0.4% from occupancy + 0.7% from the deeper pipeline. Occupancy is about a third of the loss.

So a deeper pipeline loses twice here: it crowds programs off the SM, and the deeper pipeline is slower on its own terms too. For this tile on this card, 3 stages is the sweet spot.

Finally, does the winner depend on N? Repeating the first heatmap at every sequence length:

The 64 × 32 tile wins at every N. Only the number of stages changes: at N = 512, 4 stages beat 3. The main reason is how the programs are shared out across the GPU, in what are called waves.

What is a wave?
Each SM can hold only a limited number of programs at once. Here the limit is shared memory: each SM has 102,400 bytes. Every program needs its own buffers, plus the 1,024 bytes the driver reserves. Multiply the programs per SM by the 56 SMs and you get the number of programs the whole GPU can run at the same time. Think of these as slots.

If the grid has more programs than there are slots, the GPU fills every slot and the rest wait. Whenever a program finishes, a waiting one takes its place. In the non-causal kernel every program does the same amount of work, so the programs in the slots all finish at about the same time. The kernel therefore runs in rounds: fill every slot, finish, fill again. Each round is a wave.

The number of waves is the number of programs divided by the number of slots. A whole number, such as 1.0 or 2.0, means every wave is full. A number just above a whole number, such as 1.14, means one full wave followed by a tiny last wave. During that last wave most of the GPU sits idle while a few programs finish. This is called the tail effect.

At N = 512 there are only 256 programs (because 512 / 64 = 8 query tiles per head, times 32 heads). How many fit at once depends on the stage count, because each extra stage adds 8 KB of shared memory per program:

3 stages: 24,576 bytes per program, so 4 programs fit on each SM, giving 4 × 56 = 224 slots. 256 / 224 = 1.14 waves.
4 stages: 32,768 bytes per program, so only 3 fit on each SM, giving 3 × 56 = 168 slots. 256 / 168 = 1.52 waves.

3 stages: 4 programs per SM × 56 SMs = 224 slots

Wave 1  ████████████████████████████████  224 programs, every slot busy
Wave 2  █████░░░░░░░░░░░░░░░░░░░░░░░░░░░   32 programs, 14% of slots busy

4 stages: 3 programs per SM × 56 SMs = 168 slots

Wave 1  ████████████████████████████████  168 programs, every slot busy
Wave 2  █████████████████░░░░░░░░░░░░░░░   88 programs, 52% of slots busy

█ = busy slot   ░ = idle slot

Both versions need two waves, but they waste different amounts of the GPU in the second one. With 3 stages the last 32 programs run on an almost empty GPU. With 4 stages the second wave is half full, and each program also shares its SM with only two others instead of three, so it gets a bigger share of the SM. The work is spread more evenly, so the kernel finishes sooner, even though fewer programs fit at once.

The padding experiment from above separates the two effects. At N = 512, 4 stages are 2.3% faster than 3. Padding the 3-stage kernel's shared memory so that only 3 programs fit per SM, with the code unchanged, recovers 1.6 of those 2.3 points. So about two thirds of the gain comes from the wave shape, and only the remaining third from the deeper pipeline.

The causal sweep tells the same story at half the TFLOP/s. This version of the kernel still computes every tile, including the ones the mask throws away, while the harness counts only the useful work. Its winners are 64 × 32 at N ≤ 2048 and 64 × 16 at N = 4096, all with 4 warps. So at these shapes on this GPU, one configuration comes within about 2% of the best everywhere. That will matter when deciding whether autotuning is worth its cost.

The winner, through the unchanged harness and tests

kernels/v1_tiles hard-codes the winner. The kernel body is unchanged; the class attributes and the launch change:

class FlashAttentionTriton(torch.autograd.Function):
    # The sweep winner on an RTX 4070 SUPER (bf16, D = 64): 4 warps x 16 rows = 64 rows,
    # and 64*64*2 + (3-1) * 2 * 32*64*2 = 24,576 bytes of shared memory per program.
    Q_TILE_SIZE = 64
    K_TILE_SIZE = 32
    NUM_WARPS = 4
    NUM_STAGES = 3

    @staticmethod
    def forward(ctx, Q, K, V, is_causal=False):
        ...
        grid = (triton.cdiv(N_q, FlashAttentionTriton.Q_TILE_SIZE), B)
        # The tuned config assumes bf16 rows of D = 64 (128 bytes). fp32 at D = 128 has
        # 512-byte rows: with 3 stages Triton asks for 107,520 bytes of shared memory,
        # more than the 101,376 an Ada block may have, so drop to 2 stages there.
        num_stages = FlashAttentionTriton.NUM_STAGES if D * Q.element_size() <= 256 else 2

        flash_fwd_kernel[grid](
            ...,
            Q_TILE_SIZE=FlashAttentionTriton.Q_TILE_SIZE,
            K_TILE_SIZE=FlashAttentionTriton.K_TILE_SIZE,
            is_causal=is_causal,
            num_warps=FlashAttentionTriton.NUM_WARPS,
            num_stages=num_stages,
        )

My first attempt did not have the num_stages line. I changed only the four class attributes, and the unchanged test file failed one case out of twenty:

FAILED test_flashattention_triton.py::test_forward_matches_sdpa[1-200-96-128-False-dtype0]
E   triton.runtime.errors.OutOfResources: out of resource: shared memory, Required: 107520, Hardware limit: 101376.

This failing case uses fp32 inputs with D = 128. My budget equation above assumed bf16 (2 bytes per value) and D = 64. Here each value is 4 bytes and each row is twice as long, so every tile takes four times as many bytes. The Q tile now takes 32 KB, and each K + V stage also takes 32 KB.

There is a second change. At D = 128, the 64-row query tile is shorter than the head dimension. So Triton places the warps side by side along the key axis instead of stacking them. It does this for bf16 at D = 128 too. As in Round 0, no warp then holds a whole row, and this adds two more buffers:

A buffer for P, 8 KB. With fp32 inputs, tl.dot uses TF32 instructions, so P stays in fp32 and takes 4 bytes per value.
1 KB of scratch space for the row sum. Each row is now split across warps, so their partial sums have to be combined through shared memory.

With 3 stages, the total is 98,304 + 8,192 + 1,024 = 107,520 bytes. That is 6 KB more than a program is allowed. With 2 stages, the total is 74,752 bytes, which fits.

The lesson is that a configuration tuned for one dtype and head dimension may not work for another. The test file covers both, so it caught the problem. The autotuned version in Round 3 handles this case on its own, without any special code.

The fix

The fix is one line in v1's forward. It lowers the pipeline depth from 3 stages to 2 when the rows are too wide:

num_stages = FlashAttentionTriton.NUM_STAGES if D * Q.element_size() <= 256 else 2

D * Q.element_size() is the number of bytes in one row of Q, K or V:

Inputs	D	Bytes per row	Stages used
bf16	64	64 × 2 = 128	3
bf16	128	128 × 2 = 256	3
fp32	64	64 × 4 = 256	3
fp32	128	128 × 4 = 512	2

Only fp32 at D = 128 gets 2 stages. With one fewer K + V stage, it needs 74,752 bytes of shared memory instead of 107,520. That is under the 101,376-byte limit, so the kernel launches and the failing test passes. Every other case keeps the tuned 3 stages, so the speed on the benchmark shape (bf16, D = 64) does not change.

This is a special case written by hand, and it only covers the one combination that failed. Round 3 removes it. With @triton.autotune, a configuration that does not fit in shared memory is skipped instead of crashing. For fp32 at D = 128, three of the four configurations are too big, so the autotuner uses the 2-stage one. In v5, the configuration (64, 32, 4, 2) is in the list only so that this case always has one that fits.

With that fixed, all tests pass, and the unchanged harness gives:

N	causal	baseline (ms)	v1 (ms)	v1 TFLOP/s	SDPA TFLOP/s	v1 % of SDPA
512	no	0.114	0.049	43.7	32.8	133%
1024	no	0.387	0.149	57.5	56.7	101%
2048	no	1.540	0.534	64.4	64.2	100%
4096	no	6.386	2.121	64.8	68.5	95%
512	yes	0.114	0.050	21.6	19.4	112%
1024	yes	0.385	0.153	28.1	35.7	79%
2048	yes	1.539	0.545	31.5	52.1	61%
4096	yes	6.383	2.162	31.8	61.7	52%

(These, and every harness table from here on, are medians of ten runs of the unchanged harness, with the versions interleaved and a 20-second pause before each run. "% of SDPA" is the median of each run's own ratio, since SDPA is measured in the same process, and in a version's own table SDPA's TFLOP/s come from that version's runs; the Results table pools all sixty. The medians matter: the short causal kernels at N = 512 vary by up to 22% between runs.)

Changing four numbers made the non-causal kernel three times faster: level with SDPA from N = 1024 to 2048, 5% behind at 4096, and ahead at N = 512, where SDPA itself only reaches 33 TFLOP/s. Two things hold SDPA back there. Its 128 × 128 tiles make only 128 programs at this size, 1.14 waves of the 112 that fit on the GPU at once, so Nsight Compute shows its SMs busy for only 78% of the kernel. And it loses more than this kernel does to the cache flush that do_bench runs before every call. Causal attention got the same 3× and is still half of SDPA at long sequences, because this kernel still computes every tile. Before fixing that, the profiler should confirm the new picture.

Round 2: profile the tuned kernel

I follow the same steps as in Round 0, this time on the tuned kernel. The script profile_all.sh profiles each kernel version and SDPA. It skips v2, because it runs the same kernel as v1 and only adds autotuning, and the causal report also skips the baseline. Each report's Summary page lists all the kernels side by side. To compare against the baseline, right-click the baseline's row and choose Add Baseline(s). Every number on the other pages then also shows how much it changed from the baseline, and every chart shows both kernels.

Step 1: Speed of Light, with the baseline

The tuned kernel (blue) takes 2.18 ms against the baseline's 6.27 ms (green), and its headline numbers are lower: Compute (SM) Throughput falls from 79.9% to 43.8%, Memory Throughput from 79.9% to 32.7%. Nsight Compute now adds a "Latency Issue" warning, because both are under 60%. Reading only these two bars may lead to the conclusion the tuned kernel is worse. However, round 0 already showed that the baseline's high percentages were because they were the LSU pipe doing work the faster kernel no longer needs.

Step 2: the breakdown, and a trap in it

The top compute row is now SM: Pipe Tensor Cycles Active at 43.8% (+92% against the baseline), and the LSU pipe has dropped to 25.5% (−68%). The tensor cores are the busiest unit, which is what a matrix-multiply kernel should look like.

But 43.8% is not the whole story. The roofline (next step) says the same kernel runs the tensor cores at 87.5% of their peak. For every kernel in this article the two numbers differ by a factor of two with rounding: 22.8% against 45.5% for the baseline, 43.8% against 87.5% for the tuned kernel, 44.6% against 89.2% for SDPA. A direct experiment shows where the factor of two comes from. Here is the same Triton matmul (8192³) with different input and accumulator types, with both counters:

Inputs → accumulator	Operations, % of that type's peak	Tensor pipe %
bf16 → fp32	97.6	48.8
fp16 → fp32	97.5	48.7
tf32 → fp32	93.0	46.5
fp16 → fp16	84.1	84.1
int8 → int32	64.0	64.0

For bf16, fp16 and TF32 inputs with fp32 accumulation, the pipe counter reads exactly half; with fp16 accumulation, and for int8 inputs, the two agree. A reply on Nvidia developer forums describes this as a documented defect of the counter on GeForce cards: for HMMA (fp16, bf16, TF32) with fp32 accumulation it tops out at 50%, "a defect in the metric that cannot be fixed". So the consequence is practical: for bf16 matrix multiplies with fp32 accumulation on a GeForce Ada card, about 50% on that row means the tensor cores are saturated, and the "Latency Issue" rule is a false alarm. The roofline divides by the right peak for each data type, so that is where I read how compute-bound a kernel is.

Step 3: the roofline, with the baseline

Every dot on the chart moved up, which means the tuned kernel does more work per second. The tuned kernel's dots now sit on the flat line at the top, so it is limited by the tensor cores, not by memory.

The dots also moved sideways. Remember that a dot's horizontal position is the number of operations divided by the number of bytes moved.

The L2 dot moved right, from 23.9 to 63.5 operations per byte. This is because the kernel now reads much less data from L2, i.e. 2.16 GB instead of 8.6 GB. That is about 4 times less, which is what we expected. Each program now handles 64 query rows instead of 16, so it reads K and V 4 times less often. The drop is 3.98 times rather than exactly 4, because the kernel still reads all of Q, and that part does not shrink.
The L1 and DRAM dots moved slightly left. The kernel no longer repeats the first multiply, so it does a third fewer operations. The bytes moved through L1 and DRAM fell by less than a third. Fewer operations divided by almost the same number of bytes gives a smaller number, so these dots moved left.

The table under the chart confirms the second prediction. The tensor cores did 137.4 G operations, 33.3% fewer than before. This is exactly the number the algorithm needs, so the repeated multiplies are gone.

The tuned kernel now runs the tensor cores at 87.5% of their maximum speed (63.1 TOP/s, at the 2.52 GHz clock that the profiler locks the GPU to). So for non-causal attention, there is not much speed left to gain.

Step 4: why the warps wait now

Math pipe throttle is now the longest stall, at 6.9 cycles per instruction. Warps are waiting for the tensor cores, the same signature as SDPA. Short scoreboard and MIO throttle, the shared-memory stalls that dominated the baseline, have almost disappeared. Barrier stalls remain at 2.1 cycles, and the Source page says where. The loop has two BAR.SYNCs per iteration, both generated for the K_j = tl.load(...) line, which guard the shared-memory buffers of the asynchronous K and V copies; they account for 99.5% of the barrier samples.

Side by side, with SDPA for reference:

N = 4096, non-causal	baseline (v0)	tuned (v1)	SDPA
Duration under Nsight Compute	6.27 ms	2.18 ms	2.14 ms
Tensor-core operations	206.2 G	137.4 G	137.4 G
Tensor-core % of peak (roofline)	45.5%	87.5%	89.2%
LSU pipe utilisation	79.9%	25.5%	12.4%
Instructions executed	1,606 M	394 M	220 M
Registers per thread	56	127	255
Shared memory per program	10.8 KB	24.6 KB	49.2 KB
Theoretical occupancy	66.7%	33.3%	16.7%
Bytes from L2 into the SMs	8.61 GB	2.16 GB	1.35 GB
Shared-memory wavefronts	545 M	97 M	40 M
Longest stall	short scoreboard	math pipe throttle	math pipe throttle

Occupancy went down by half and the kernel got three times faster. SDPA runs at half the tuned kernel's occupancy again. Occupancy is a means of hiding latency, not a goal, and a kernel that is waiting on its tensor cores does not need more warps to wait with.

Step 5: checking the claim about too many warps

Round 1 claimed that with 8 warps on a 64-row tile, Triton stacks the warps along the query axis and every multiply is done twice. One more profile, of the same 64 × 32 kernel with num_warps=8, settles the question. The tensor cores executed 274.9 G operations, exactly twice the algorithm's 137.4 G, at 89.3% of their peak, and the kernel took 4.26 ms instead of 2.18. The tensor cores were as busy as in the fast kernel. They were doing everything twice. "Busy" and "useful" are different questions, and the operation count tells them apart.

Step 6: the causal run

The same report for causal attention says the tuned kernel executed 137.4 G tensor operations there too, against the 68.7 G that causal attention needs (and that the harness divides by). Half of the tensor-core work is thrown away by the mask, which is why causal runs at the same speed as non-causal and half of SDPA. That is the largest thing left, and it is inside the inner loop, which is where Round 4 goes. First, though, the last item of the tuning job: autotuning.

Round 3: wrap the winner in `@triton.autotune`

Hard-coding one configuration has two weaknesses that the sweep and the test file have already shown. The best configuration can depend on the shape (4 stages at N = 512, a 16-key tile for long causal runs), and a configuration can be invalid for a dtype or head dimension it was not tuned on (fp32 at D = 128). Triton's answer is @triton.autotune, i.e. give it a list of configurations, and on the first call for each new combination of the key arguments it compiles and times every configuration on the real inputs, then keeps the fastest for that key.

kernels/v2_autotune lists the per-shape winners of the sweep and lets the grid depend on the chosen tile:

# The per-shape winners of the sweep (Q_TILE, K_TILE, num_warps, num_stages). They differ
# only in the pipeline depth and, for long causal runs, the key tile.
AUTOTUNE_CONFIGS = [
    triton.Config({"Q_TILE_SIZE": bq, "K_TILE_SIZE": bk}, num_warps=nw, num_stages=ns)
    for bq, bk, nw, ns in [(64, 32, 4, 3), (64, 32, 4, 4), (64, 32, 4, 2), (64, 16, 4, 5)]
]


# One tuning run per distinct (N_QUERIES, N_KEYS, D, is_causal) and input dtype; Triton adds
# the dtypes of the tensor arguments to the key on its own.
@triton.autotune(configs=AUTOTUNE_CONFIGS, key=["N_QUERIES", "N_KEYS", "D", "is_causal"])
@triton.jit
def flash_fwd_kernel(...):            # the body is unchanged
    ...


class FlashAttentionTriton(torch.autograd.Function):
    @staticmethod
    def forward(ctx, Q, K, V, is_causal=False):
        ...
        # The grid depends on the Q tile the autotuner picks, so it is a function of the config.
        grid = lambda meta: (triton.cdiv(N_q, meta["Q_TILE_SIZE"]), B)
        flash_fwd_kernel[grid](
            Q, K, V, O, L,
            ...,
            D=D,
            is_causal=is_causal,          # no tile sizes, num_warps or num_stages: the config supplies them
        )

Three details to watch out for:

Everything the configuration sets must be left out of the call. Passing Q_TILE_SIZE= as well raises Conflicting meta-parameters.
Changing a key value always means re-timing, but not always recompiling. Every time a value in the key changes, the autotuner times all the configurations again. Whether it also has to compile them again depends on the kind of value. N_QUERIES and N_KEYS are normal integers. Triton only checks a few simple things about them, i.e. whether they divide by 16, whether they equal 1, and whether they fit in 32 bits. So a new sequence length can reuse the kernels already compiled, and only the timing is repeated. D and is_causal are tl.constexpr, which means their values are built into the compiled code. A new value for either of them means compiling new kernels as well as timing them.
A configuration that does not fit is skipped, not fatal. The autotuner catches OutOfResources and times that configuration as infinitely slow. For fp32 at D = 128, three of the four configurations need more than 99 KB of shared memory, and the autotuner quietly picks the 2-stage one. The special case in v1's forward is no longer needed, and the unchanged test file passes.

The unchanged harness:

N	causal	v1, fixed (ms)	v2, autotuned (ms)	v2 TFLOP/s	v2 % of SDPA
512	no	0.049	0.048	44.6	131%
1024	no	0.149	0.150	57.1	101%
2048	no	0.534	0.537	64.0	100%
4096	no	2.121	2.129	64.6	95%
512	yes	0.050	0.046	23.3	120%
1024	yes	0.153	0.153	28.0	79%
2048	yes	0.545	0.546	31.5	61%
4096	yes	2.162	2.157	31.9	52%

As the sweep predicted, v2 runs at about the same speed as the hard-coded winner: within 1% from N = 1024 up, causal or not.

We can ask the autotuner about its choice With TRITON_PRINT_AUTOTUNING=1. The autotuner picked (64, 16, 4, 5) at N = 512 non-causal in 9 runs out of 10. That is the last candidate in the list. It is 6% slower than (64, 32, 4, 4) and 2.4% slower than v1's own configuration. This is not bad luck. After an idle period, this GPU stays at a lower clock (2.52 GHz) until it has been busy for 0.13 to 0.3 s (see "One warning about durations in a timeline" in Round 0). At the start of a harness run, GPU work is mixed with compilation. nvidia-smi updates only every half second, and it showed the clock at 2.52 GHz for 0.4 to 1.0 s in 58 of the 60 runs. The autotuner times the candidates in list order, so the first ones are timed on a slower GPU. In fresh processes after the same 20-second pause, it timed the first two candidates as 5% and 6% slower than their true speed, and it picked the last candidate 7 times out of 10. With one second of GPU work before tuning, it picked (64, 32, 4, 4) 10 times out of 10.

For the shapes that come later, the clock has settled, so differences in the choice are random noise, not bias. At N = 4096 causal, the ten harness runs picked three different configurations. When timed carefully, though, the four candidates are within 2% of each other there, so the choice barely matters. An autotuner is only as accurate as its timings, and the first timings after the GPU has been idle are the least accurate of all.

Round 4: inside the inner loop

There are two outstanding optimisations:

Skipping operations on the masked tiles of causal attention
Trimming the instructions that compete with the tensor cores for issue slots, primarily the exponential, i.e. switching tl.exp to tl.math.exp2 and folding the softmax scale and log₂e into one multiply-add [Note: An issue slot is one opportunity for a warp scheduler to send one instruction to an execution unit. Each Ada SM is split into four sub-partitions, and each sub-partition has its own warp scheduler. On every clock cycle, each scheduler looks at its resident warps, picks one that is ready, and issues one instruction from it. That single “issue one instruction this cycle” opportunity is an issue slot. An SM therefore has 4 issue slots per cycle. Over a kernel’s runtime, the total number of slots is 4 × SMs × cycles.]

Causal tile-skipping

In round 2, we found that for causal attention the tuned kernel executes 137.4 G tensor operations, double the expected count of 68.7 G. In the timeline below, two launches labelled v1_tiles causal=True are about as long as the two v1_tiles causal=False launches before them:

Why. The loop visits every key tile. For each one it computes Q Kᵀ, applies the mask, takes the exponential and multiplies by V. It does this even when the mask throws away every element of the tile.

Counting the tiles allow us to understand how many calculations are wasted. The tile sizes are Round 1's winner, which every comparison in this round keeps fixed, i.e. each program handles 64 query rows (Q_TILE_SIZE = 64) and steps through the keys 32 at a time (K_TILE_SIZE = 32). At N = 4096, one head therefore has:

4,096 / 64 = 64 query tiles, numbered i = 0 to 63, and
4,096 / 32 = 128 key tiles, numbered j = 0 to 127.

Every program meets every key tile, so one head has 64 × 128 = 8,192 (query tile, key tile) pairs to work through. All 32 heads look the same, so counting one head is enough.

Which key tiles does a query tile need? In causal attention, the query at position r may only look at keys 0 to r. Query tile i holds 64 rows in a row: tile 0 has rows 0–63, tile 1 has rows 64–127, and in general tile i has rows 64*i* to 64*i* + 63. The last row, 64*i* + 63, sees the most keys: keys 0 to 64*i* + 63, which is 64*i* + 64 keys. Every other row in the tile sees fewer. So this tile never needs a key after 64*i* + 63. At 32 keys per tile, those 64*i* + 64 keys fill exactly (64*i* + 64) / 32 = 2*i* + 2 key tiles. For example, query tile 3 holds rows 192–255, so it needs keys 0–255, which are key tiles 0 to 7: 2 × 3 + 2 = 8 tiles.

The figure draws N = 512 so the tiles are big enough to see. N = 4096 has the same staircase, eight times larger in each direction. Each row of the grid splits into three parts:

Fully visible (green). Key tiles 0 to 2*i* − 1 hold keys 0 to 64*i* − 1. Every one of these keys comes before the tile's first row, so every row sees every key. These tiles need no mask. There are 2*i* of them.
On the diagonal (orange). Key tiles 2*i* and 2*i* + 1 hold keys 64*i* to 64*i* + 63, the same positions as the tile's own query rows. Row 64*i* + t sees only the first t + 1 of these keys, so these two tiles need the mask. There are always exactly 2 of them, because a 64-row query tile is exactly two 32-key tiles wide.
Fully masked (grey). Every key tile after that comes after every row of the query tile. The mask would set all of it to zero, so it needs no work at all.

Adding these up over the 64 query tiles of one head:

	For query tile i	Over all 64 query tiles
Fully visible, no mask	2i	0 + 2 + 4 + … + 126 = 4,032
On the diagonal, masked	2	64 × 2 = 128
Needed (visible + diagonal)	2i + 2	4,160, which is 50.8% of 8,192
Fully masked, can be skipped	128 − (2i + 2)	8,192 − 4,160 = 4,032

So only about half of the 8,192 pairs need any work, and only 128 of them need a mask. (That 128 is not the same 128 as the number of key tiles, even though the numbers match. It is 2 diagonal tiles × 64 query tiles. The two are equal because a query tile is exactly two key tiles wide.) The kernel so far computes all 8,192 pairs.

The change to the Triton code. Split the work along the three groups in the figure.

Move the body of the loop into a helper function, _attend_key_tiles, which works through a range of key tiles.
Call it twice from the kernel. The first call covers the green tiles, which need no mask. The second call covers the two orange tiles on the diagonal, which do need it.
Never visit the grey tiles. Neither call's range reaches them.

The helper takes a flag, MASKED, that says which kind of tile it is working on. MASKED is a tl.constexpr, which means its value is fixed when Triton compiles the kernel, not when it runs. Triton therefore compiles two separate copies of the loop, one for each value. The copy for green tiles contains no masking code at all. It does not compare positions or replace scores with -1e6, and it does not check whether its loads run past the end of K or V. Here is the helper:

@triton.jit
def _attend_key_tiles(
    O_acc, l_acc, m_acc, Q_i, q_pos,
    K_block_ptr, V_block_ptr,
    start, stop, N_KEYS, scale,
    K_TILE_SIZE: tl.constexpr,
    MASKED: tl.constexpr,
    is_causal: tl.constexpr,
):
    # Fold the key tiles [start, stop) into the running state. MASKED is a
    # compile-time flag: the unmasked copy of this loop has no compare, no
    # select and no bounds checks on its loads.
    K_block_ptr = K_block_ptr.advance((start, 0))
    V_block_ptr = V_block_ptr.advance((start, 0))
    for k_start in range(start, stop, K_TILE_SIZE):
        if MASKED:
            K_j = tl.load(K_block_ptr, boundary_check=(0, 1), padding_option="zero")
            V_j = tl.load(V_block_ptr, boundary_check=(0, 1), padding_option="zero")
        else:
            K_j = tl.load(K_block_ptr)
            V_j = tl.load(V_block_ptr)

        S_ij = tl.dot(Q_i, tl.trans(K_j)) * scale          # (Q_TILE, K_TILE), fp32

        if MASKED:
            k_pos = (k_start + tl.arange(0, K_TILE_SIZE))[None, :]
            keep = k_pos < N_KEYS
            if is_causal:
                keep = keep & (k_pos <= q_pos)
            S_ij = tl.where(keep, S_ij, -1e6)

        # ... the online-softmax update and P @ V, unchanged ...

        K_block_ptr = K_block_ptr.advance((K_TILE_SIZE, 0))
        V_block_ptr = V_block_ptr.advance((K_TILE_SIZE, 0))
    return O_acc, l_acc, m_acc

What the helper does:

Its inputs. It receives the running state (O_acc, l_acc and m_acc), this program's query tile Q_i, and a range of keys to work through, from key start up to, but not including, key stop. Both are key positions, not tile numbers.
Jump to the first key. K_block_ptr and V_block_ptr point at key 0 when the helper is called. .advance((start, 0)) moves them down to key start.
The loop. k_start is the first key of the current tile: start, then start + 32, and so on, until it reaches stop. If start equals stop, the loop does not run and the state comes back unchanged. That happens, for example, for query tile 0, which has no green tiles.
Loading K and V. The masked copy loads with boundary_check. Any part of the tile that lies past the end of K or V is filled with zeros instead of being read from memory that does not belong to the tensor. The unmasked copy loads the tile directly, because it is only ever given tiles that lie completely inside the tensor.
The scores. S_ij holds one score for each pair of a query row and a key in this tile. It is 64 × 32 in this kernel.
The mask (masked copy only). k_pos lists the positions of the tile's 32 keys as a row, and q_pos lists the positions of the 64 query rows as a column. Comparing the two gives a 64 × 32 grid of true or false values called keep. A score is kept if its key exists (k_pos < N_KEYS, which removes the zero-filled keys past the end) and, for causal attention, if the key is not in the future (k_pos <= q_pos). Every score that is not kept becomes -1e6. Its exponential is then effectively zero, so it adds nothing to the output.
The rest of the loop. The softmax update and the multiply by V are the same as before. At the end of each pass, both pointers move down by one tile.
Its output. It returns the updated running state, which the next call continues from.

and in flash_fwd_kernel, after loading Q_i:

    q_start = query_tile_index * Q_TILE_SIZE
    q_pos = (q_start + tl.arange(0, Q_TILE_SIZE))[:, None]

    # Key tiles that end at or before full_stop lie entirely inside the tensor.
    full_stop = (N_KEYS // K_TILE_SIZE) * K_TILE_SIZE
    if is_causal:
        # Every key below the first query of this tile is visible to every row,
        # so those tiles need no mask. The tiles that straddle the diagonal get
        # the mask, and nothing beyond the last query of the tile is visited.
        unmasked_stop = tl.minimum((q_start // K_TILE_SIZE) * K_TILE_SIZE, full_stop)
        masked_stop = tl.minimum(q_start + Q_TILE_SIZE, N_KEYS)
    else:
        # Only a ragged last key tile needs the mask.
        unmasked_stop = full_stop
        masked_stop = N_KEYS

    O_acc, l_acc, m_acc = _attend_key_tiles(
        O_acc, l_acc, m_acc, Q_i, q_pos, K_block_ptr, V_block_ptr,
        0, unmasked_stop, N_KEYS, scale,
        K_TILE_SIZE=K_TILE_SIZE, MASKED=False, is_causal=is_causal,
    )
    O_acc, l_acc, m_acc = _attend_key_tiles(
        O_acc, l_acc, m_acc, Q_i, q_pos, K_block_ptr, V_block_ptr,
        unmasked_stop, masked_stop, N_KEYS, scale,
        K_TILE_SIZE=K_TILE_SIZE, MASKED=True, is_causal=is_causal,
    )

What the kernel does before and with the two calls:

q_start is the first query row of this program's tile: 64*i* for query tile i. q_pos lists the positions of all 64 rows as a column, which is what the mask compares against.
full_stop is N_KEYS rounded down to a whole number of key tiles. Every key tile that ends at or before it lies completely inside the tensor and can be loaded without bounds checks. When N_KEYS is a multiple of 32, full_stop is simply N_KEYS.
Causal: where the green tiles end. Every key before q_start comes before every row of the tile, so all of those keys are visible. unmasked_stop is q_start rounded down to the start of a key tile. With 64-row query tiles and 32-key tiles, q_start is already a multiple of 32, so the rounding changes nothing. It matters when the key tile is larger than the query tile. With 128-key tiles, for example, query tile 1 starts at key 64, in the middle of key tile 0, so that tile has to go through the masked loop. It is also capped at full_stop, so a ragged last key tile always goes through the masked copy, which has the bounds checks.
Causal: where the orange tiles end. masked_stop is one past the tile's last row, q_start + 64. No row of the tile can see a key at or beyond that position, so the grey tiles are never visited. It is capped at N_KEYS in case the tile reaches past the last key.
Non-causal. No key is in the future, so every complete key tile goes through the unmasked copy. Only a ragged last tile, if there is one, goes through the masked copy. When N_KEYS is a multiple of 32, the masked range is empty and the second call does nothing.
The two calls. The first call works through keys 0 to unmasked_stop without the mask. The second picks up the running state where the first left off and works through unmasked_stop to masked_stop with the mask. Both calls receive the same pointers, which point at key 0. Each call moves its own copy forward with advance, so the second call starts at the right key.

For example, take query tile i = 3 at N = 4096, causal, which holds rows 192–255. Then unmasked_stop = 192: the first call works through keys 0–191, which are key tiles 0–5, the six green tiles in the figure. masked_stop = 256: the second call works through keys 192–255, which are key tiles 6 and 7, the two orange tiles. Key tiles 8 to 127 are never touched.

For example, take query tile i = 3 at N = 512, causal, which holds rows 192–255. This is case 1 in the figure above. N = 512 gives 512 / 32 = 16 key tiles, numbered 0 to 15. Here unmasked_stop = 192, so the first call works through keys 0–191. Those are key tiles 0–5, the six green tiles. masked_stop = 256, so the second call works through keys 192–255. Those are key tiles 6 and 7, the two orange tiles. Key tiles 8 to 15 are never touched.

The unchanged test file, test_flashattention_triton.py, only tests causal attention with equal query and key lengths. Non-square causal attention is checked separately in check_nonsquare_causal.py against an explicit top-left mask, the convention SDPA's is_causal=True uses.

Measure. The test file passes, and the harness:

N	causal	v1 (ms)	v3 (ms)	v3 TFLOP/s	v3 % of SDPA
512	yes	0.050	0.038	28.2	146%
1024	yes	0.153	0.095	45.1	126%
2048	yes	0.545	0.303	56.7	109%
4096	yes	2.162	1.101	62.4	102%
4096	no	2.121	2.056	66.8	97%

Causal attention is twice as fast at N = 4096 (1.96×), and now level with or ahead of SDPA at every length. The non-causal path is 3% faster too, from dropping the mask on interior tiles.

Profile again. The same timeline shows the causal launches of v3 (last two) at half the width of everything else. In Nsight Compute, the causal kernel now executes 69.8 G tensor operations: the 68.7 G the algorithm needs plus 1.6% for the diagonal tiles, where half of each tile is computed and then masked. Instructions fell from 428 M to 184 M. SDPA executes 70.9 G here, a little more than v3, because its 128 × 128 tiles waste more on the diagonal.

One idea works only in the right form. In causal attention the last query tiles of each head have the most work, and the GPU starts them last, which can leave a tail of busy SMs at the end of the kernel. Reversing the order within each head (query_tile_index = tl.num_programs(0) - 1 - tl.program_id(0)) is a common fix, but here it gains only about 1% (1.099 ms to 1.087 ms at N = 4096 for v4, over ten interleaved rounds).

The likely reason is the launch order. The GPU starts programs roughly in order of program_id(0) first, then program_id(1), so the heads are launched one after another, and the longest tiles of the last heads still start last. Swapping the grid axes, so that program_id(0) is the head, and reversing the query tiles along the other axis starts every head's longest tiles first:

    batch_index = tl.program_id(0)
    query_tile_index = tl.num_programs(1) - 1 - tl.program_id(1)
    # launched with grid = (B, triton.cdiv(N_q, Q_TILE_SIZE))

That is 4 to 5% faster at N = 4096 and 7 to 10% at N = 2048, for 64 × 32, 64 × 64 and 64 × 128 tiles alike (its output matches SDPA's at both shapes, but it still needs the test file before it goes into a kernel). I found it after the measurements above, so it is not in the final kernel; it is the first thing I would add.

`exp2`

What the profile said. Nsight Compute's Source page shows the Python source next to the machine code, with a count of executed instructions for every line. For v3, one line of the loop accounts for 28.6% of every instruction the kernel executes:

That line is P_ij = tl.exp(S_ij - m_new). Selecting it highlights its SASS on the right, and the opcodes there is a mix of FFMA, FMUL, FSETP.GEU and MUFU.EX2, instead of simply MUFU.EX2. This is because tl.exp(x) compiles to:

mul.f32         %r485, %r453, 0f3FB8AA3B;   // x * log2(e), 0f3FB8AA3B is 1.4426950
ex2.approx.f32  %r486, %r485;               // 2^x

The GPU has no natural exponential instruction, only ex2 on the special-function unit (MUFU), so exp(x) becomes 2^(x · log₂e), which requires one hidden multiply per element. And ex2.approx.f32 without .ftz must handle results too small for a normal float (below 2⁻¹²⁶), so ptxas wraps every MUFU.EX2 in a range check and two fix-ups. Per element, with the registers renamed for readability (the real code interleaves several elements):

FSETP.GEU.AND  P0, PT, R1, -126, PT   // range check: is x below -126? (would the answer be too tiny?)
@!P0 FMUL      R1, R1, 0.5            // fix-up 1: if so, halve x
MUFU.EX2       R2, R1                 // compute 2^x (or 2^(x/2))
@!P0 FMUL      R2, R2, R2             // fix-up 2: if we halved, square the result

Add the softmax scale, which the kernel multiplies into every score on the line before (S_ij = tl.dot(...) * scale), and the subtraction of the row maximum, and each score pays for a scale, a subtraction, a multiply by log₂e, a compare, up to two fix-ups and the MUFU.EX2 itself.

The change. Fold the scale and log₂e into one constant, qk_scale = scale · log₂e, keep the running max m in log₂ units of the scaled scores, and call tl.math.exp2 directly. Because qk_scale is positive, the row maximum of the raw scores times qk_scale is the maximum of the scaled ones, so the scale is applied to Q_TILE_SIZE maxima instead of Q_TILE_SIZE × K_TILE_SIZE scores, and the subtraction fuses with the multiply:

        S_ij = tl.dot(Q_i, tl.trans(K_j))                   # raw scores, not yet scaled
        # ... mask as before ...

        # m is kept in log2 units of the scaled scores. qk_scale > 0, so the
        # row max of the raw scores times qk_scale is the max of the scaled ones,
        # and the scale is applied to Q_TILE values here instead of Q_TILE * K_TILE.
        m_new = tl.maximum(m_acc, tl.max(S_ij, axis=1, keep_dims=True) * qk_scale)
        # One FFMA and one ex2 per score: exp(scale*s - m) == 2^(s*qk_scale - m').
        P_ij = tl.math.exp2(S_ij * qk_scale - m_new)
        alpha = tl.math.exp2(m_acc - m_new)

with, outside the loop,

    # Fold log2(e) into the softmax scale once, so the inner loop can use exp2.
    qk_scale = scale * 1.4426950408889634
    ...
    # m and log2(l) are in log2 units; the backward pass expects natural-log L.
    L_i = tl.reshape((m_acc + tl.math.log2(l_acc)) * 0.6931471805599453, (Q_TILE_SIZE,))

The last line of code above matters because the test file checks L against torch.logsumexp, which is in natural-log units, so L has to be converted back (L = (m + log₂ l) · ln 2).

What the compiler did with it. tl.math.exp2 compiles to ex2.approx.ftz.f32, the range check and fix-ups disappear, and the scale, the subtraction and the multiply become one FFMA:

FFMA      R89, R17, UR21, -R3              // s * qk_scale - m
MUFU.EX2  R89, R89                         // 2^x

Across the whole compiled kernel (both loops), the number of FMUL instructions fell from 320 to 177, and all 40 range checks are gone. Those were the FSETP.GEU compares against −126, one for each MUFU.EX2. The four FSETP.GEUs that remain belong to the division and the logarithm after the loop, not to the exponential.

Measure. The test file passes (including L). The harness:

N	causal	v3 (ms)	v4 (ms)	v4 TFLOP/s	v4 % of SDPA
1024	no	0.147	0.144	59.5	105%
4096	no	2.056	2.027	67.8	99%
1024	yes	0.095	0.091	47.1	131%
4096	yes	1.101	1.083	63.5	103%

About 1.5 to 4% improvements on the above shapes.

Profile again.

The exp2 line is down to 12.7% of a smaller total. Executed instructions fell 25%, from 352 M to 264 M (non-causal), and issue-slot utilisation from 29% to 22%. And yet the kernel is only 1.5% faster. That is the lesson of this change: instruction count is not time. The kernel is limited by the tensor cores (the longest stall is still math pipe throttle, now 10.2 cycles), and the removed instructions ran on other pipes, mostly in parallel with the matrix multiplies. They only cost time where they delayed issuing the next HMMA.

The new profile also says the next most expensive line is O_acc = alpha * O_acc, at 12.7%. That rescale costs Q_TILE_SIZE × D multiplies per iteration no matter how many keys the iteration covers, so a bigger key tile spreads it over more scores.

Round 5: re-tune, then the final profile

Round 1's sweep measured the kernel as it was then. Round 4 changed what each loop iteration costs by shrinking the per-score work, so the per-iteration work (the alpha * O_acc rescale, the barrier around the asynchronous copies) is a bigger share. Causal programs now do half as much work each. Those are exactly the things the tile sizes trade against each other, so the old winner may no longer be the winner. The same sweep.py, pointed at kernels/v4_exp2, took 10.8 minutes:

N	causal	Best configuration (Q × K, warps, stages)	TFLOP/s	64 × 32, 4, 3 (Round 1's winner)
512	no	64 × 32, 4, 4	47.7	97.8% of best
1024	no	128 × 64, 4, 2	59.9	97.9%
2048	no	64 × 128, 4, 2	67.5	97.3%
4096	no	64 × 128, 4, 2	69.4	94.9%
512	yes	64 × 32, 4, 5	27.6	95.0%
1024	yes	64 × 32, 4, 3	46.6	100%
2048	yes	64 × 64, 4, 3	58.2	99.0%
4096	yes	64 × 64, 4, 2	63.5	98.3%

This time the best configuration depends on the shape and the mask. Long non-causal sequences now prefer 128-key tiles. Profiling the same kernel and mask with 64 × 32 tiles and 3 stages (Round 1's choice) and with 64 × 128 tiles and 2 stages (this round's) shows how much. Executed instructions fall by 35%, from 264 M to 171 M (26% or 31% if the stage count is held at 2 or 3). Of the drop:

the rescale of O_acc is 27%;
the K and V copies and their barriers are 22%;
the loop's own bookkeeping is 21%;
the per-row max, sum and alpha updates are 25%.

The lines that run once per iteration fall about 4×: the rescale, the bookkeeping, the alpha, m and l updates, and the warp shuffles inside tl.max and tl.sum. The copies fall about 2×, because their instructions grow with the tile. The per-score work barely changes: the exp2, the cast and the multiply by V not at all, and the reductions' compares and adds by 8%. Barrier stalls fall from 3.4 to 0.3 cycles per instruction, and the tensor cores go from 90.2% to about 92% of their peak. Removing the rescale outright, which gives wrong answers, makes the 32-key kernel only 2.7% faster, so the rescale is part of the story, not all of it. The price is 255 registers per thread (the non-causal kernel even spills two, though the spill code runs once per warp, not once per iteration) and 40,960 bytes of shared memory, which only pays off when there is plenty of work: short sequences and the causal diagonal (where a wide key tile means more masked waste) still prefer 32 or 64 keys. For 64-row tiles and larger, the warp rule from Round 1 still holds as a ceiling, never more warps than 16-row slices, but with the leaner loop, two slices per warp became competitive too: 128-row tiles now do best with 4 warps, the shape FlashAttention-2 itself uses.

So autotuning finally has something to choose. In the sweep, Round 1's winner is still within about 5% everywhere (97.5% of the best on average), but a short list can do better. Picking configurations greedily to raise the worst shape gives five that the sweep puts within 1% of the best on every shape, one of which is only there so that fp32 at D = 128 has a configuration that fits:

# From the second sweep (after tile-skipping and exp2), chosen so that every benchmark shape
# is within 1% of its best configuration in that sweep: 64x32 tiles with 3, 4 or 5 stages
# for short sequences, 64x128 with 2 stages for long non-causal ones, and 64x32 with 2
# stages, the only one that fits in 99 KB of shared memory for fp32 inputs at D = 128.
AUTOTUNE_CONFIGS = [
    triton.Config({"Q_TILE_SIZE": bq, "K_TILE_SIZE": bk}, num_warps=nw, num_stages=ns)
    for bq, bk, nw, ns in [(64, 32, 4, 3), (64, 32, 4, 4), (64, 32, 4, 5), (64, 128, 4, 2), (64, 32, 4, 2)]
]


@triton.autotune(configs=AUTOTUNE_CONFIGS, key=["N_QUERIES", "N_KEYS", "D", "is_causal"])
@triton.jit
def flash_fwd_kernel(...):            # the v4 body, unchanged
    ...

A single sweep measurement per configuration is noisy, though, and the list was chosen from the same data. Re-timed together with the sweep's five best configurations for each shape (14 configurations in all), five rounds each, the list is within 1% of the best on seven of the eight shapes and 2.3% behind on the eighth (N = 1024 causal, where 64 × 64 tiles with 3 stages win). Round 1's winner, re-timed the same way, is up to 7.3% behind (at N = 512 causal) and 2.8% behind on average.

That is kernels/v5_final. The unchanged test file passes. The unchanged harness, against the previous article's kernel and SDPA:

N	causal	baseline (ms)	final (ms)	final TFLOP/s	SDPA TFLOP/s	final % of SDPA
512	no	0.114	0.045	47.9	34.0	141%
1024	no	0.387	0.144	59.8	56.7	105%
2048	no	1.540	0.507	67.8	63.8	106%
4096	no	6.386	1.954	70.3	68.0	103%
512	yes	0.114	0.037	29.1	19.4	149%
1024	yes	0.385	0.091	47.1	35.8	132%
2048	yes	1.539	0.294	58.5	52.0	112%
4096	yes	6.383	1.077	63.8	61.5	104%

The final profile

Last loop of the method: profile the final kernel next to SDPA, with SDPA as the baseline.

Both kernels sit on the roof. In all ten harness runs the autotuner picked 64 × 128 tiles with 2 stages for this shape, so that is the configuration profiled here (profile_driver.py --config 64,128,4,2 leaves the autotuner only that one). It reaches 92.2% of the tensor-core peak against SDPA's 89.1%, with the same 137.4 G operations.The L2 dots differ: SDPA's 128-row tiles move 1.35 GB through L2 against my 2.16 GB, so SDPA's dot sits further right, but both are far enough right of the L2 ridge that it does not matter here.

N = 4096	final, non-causal	SDPA, non-causal	final, causal	SDPA, causal
Duration under Nsight Compute	2.06 ms	2.13 ms	1.15 ms	1.18 ms
Tensor-core operations	137.4 G	137.4 G	70.9 G	70.9 G
Tensor-core % of peak	92.2%	89.1%	85.4%	83.5%
Instructions executed	171 M	220 M	96 M	114 M
Registers per thread	255	255	255	255
Shared memory per program	41.0 KB	49.2 KB	41.0 KB	49.2 KB
Longest stall	math pipe throttle (9.1)	math pipe throttle (6.8)	math pipe throttle (8.1)	math pipe throttle (6.6)
Barrier stall	0.3	0.6	0.6	0.9

And the whole journey on one timeline, every version except v2 (whose kernel is v1's) at N = 4096, non-causal then causal, each launched twice:

What to look for, as a checklist

This is the order I now read profiles in, and what each step has to answer before moving to the next.

The benchmark

Use the same harness for every comparison, and report medians. On a power-limited consumer card, single runs are sometimes 10% or more slower than usual, so compare each kernel with a reference measured in the same run, run several times before believing a 3% difference, and interleave the versions.
Decide what "100%" means. For this kernel it is the tensor cores' bf16-with-fp32-accumulation peak: 512 operations per clock per SM, half the fp16-accumulation rate on GeForce Ada, checked against a large matmul (93 to 98% of it). In TFLOP/s it moves with the clock: 72 at Nsight Compute's 2.52 GHz, 75 to 77 at the 2.61 to 2.68 GHz the card runs the harness's N = 4096 kernels at. So compare Nsight Compute's "% of peak", not TFLOP/s against a fixed number.
Autotune on a warm GPU. After an idle period the clock stays low until the GPU has been kept busy for a few tenths of a second (much longer if the work comes in short bursts), so the configurations an autotuner times first look slower than they are.

The timeline (Nsight Systems)

Is the GPU row solid? Gaps suggests the CPU could be the bottleneck. A long launch latency in the kernel tooltip means the CPU is comfortably ahead.
How many kernels does one call launch, and which takes the time?
Compare the kernel with a reference doing the same work, on the same timeline, and do not quote single launches: the clock depends on what the GPU did just before. After an idle period the clock stays at 2.52 GHz until the GPU has been kept busy for a few tenths of a second. Once it has boosted, a pause of a few tens of milliseconds to 1.8 s lets the next tens of milliseconds of power-limited work run faster than it can sustain, and starting a capture pauses the program for 0.4 to 0.6 s. On my card the desktop's GUI tasks also interrupts a kernel now and then. nsys profile --gpu-metrics-devices=0 shows the clock, and any graphics work, next to each kernel.

The kernel (Nsight Compute)

Speed of Light: look beyond the two headline percentages. Open the breakdown and read the name of the unit at the top. In the baseline, the "80% compute, 80% memory" was the load/store unit moving data inside the SM.
On GeForce Ada, the tensor-pipe row tops out at 50% for matrix multiplies with bf16, fp16 or TF32 inputs and fp32 accumulation. With fp16 accumulation it agrees with the roofline's count. So the "Latency Issue" rule could fire on kernels that are compute-bound. Read compute-boundness from the tensor-core roofline's "Peak %" instead.
Roofline: compare "# Operations" with the algorithm's count. It found duplicated multiplies (1.5× in the baseline, 2× with too many warps) and wasted ones (2× for causal without tile-skipping).
Predict the bytes before reading them. L2 traffic was programs × bytes each program reads (8.6 GB, then 2.16 GB), within 1% of difference. When a prediction and a measurement disagree, one of the assumptions is wrong.
Warp state: math pipe throttle on top means the tensor cores are the limit, which is where a matrix-multiply kernel wants to be. Short scoreboard, MIO throttle and barrier on top mean shared-memory traffic and synchronisation. Long scoreboard means global-memory latency.
Barrier stalls are reported at an instruction after the BAR.SYNC (in this kernel, the first shared-memory load or MUFU), not at the BAR.SYNC itself. Count the BAR.SYNCs in the loop and give each one the samples that follow it.
Compute Workload Analysis and the Source page: which source lines produce the busiest pipe's instructions, and what share of all instructions each line costs. The Source page found tl.exp at 28.6% of the kernel.
Launch statistics and occupancy: which limit binds, registers or shared memory. Occupancy is a means, not a goal. The fastest kernels here run at a third or a sixth of the maximum. Occupancy also decides how a small grid splits into waves: 256 programs at 4 per SM is 1.14 waves, at 3 per SM 1.5.
Treat the "estimated speedup" suggestions as hypotheses. The top two for the baseline (occupancy and bank conflicts) were symptoms of the real problem.

The compiler (Triton)

compiled.n_regs, n_spills and metadata.shared for the configuration you are about to benchmark, and the budget arithmetic to predict them.
In the TTGIR: warpsPerCTA in the #mma layout (do the warps own whole rows?), any ttg.local_alloc or ttg.convert_layout in the loop (a round trip through shared memory?), and the memdesc shapes of the pipeline buffers.
The PTX or SASS of the hottest line. That is where ex2.approx.f32 and its range check turned up.

After every change

Profile again, and check the change did what you expected and nothing else. exp2 cut instructions by 25% and time by 1.5%, because the kernel was waiting on the tensor cores, not on the instructions it removed.

FlashAttention 2 from PyTorch to Triton

Lewis Won — Sun, 27 Sep 2026 02:27:24 +0000

Table of Contents

What I am building
Step 1: the tiled forward in PyTorch
A primer to tl.make_block_ptr
- base as the starting point in memory of a tensor
- shape tells Triton how big the tensor is
- strides tells Triton how far apart neighbours are in memory
- offsets tells Triton where the tile starts
- block_shape tells Triton the size of a tile
- order tells Triton which axis is side by side in memory
- .advance() tells Triton how to move the tile
- boundary_check tells Triton what happens past the edge
- Put it all together
- Cheat sheet
Step 2: the Triton forward kernel
Step 3: testing and bugs found
Step 4: the benchmark, and the baseline number
What is wrong with this kernel
Future articles

I wanted to learn kernel optimisation by starting with Triton. While I was able to follow basic tutorials like vector addition, the learning curve increased sharply when I moved towards more complex algorithms like Flash Attention 2.

After some exploration of various tutorials, I settled on following Stanford CS336 assignment 2 which provides the implementation of Flash Attention 2 (FA2) as an assignment. I liked it for the following reasons:

It began by asking for a PyTorch implementation of FA2, which helped to solidify my understanding of FA2 before introducing complexities associated with Triton.
It made assumptions to simplify the implementation of FA2 to help me build up my understanding of basic concepts in Triton first. These assumptions included (i) inputs of shape (batch, seq, d) with no head dimension; (ii) tile sizes of at least 16x16, which I hard-coded to 16x16; and (iii) sequence lengths and head dimensions that are powers of 2 and at least 16, so out-of-bounds accesses could be ignored.

Future articles will relax these assumptions.

I also retained some of the mistakes I made when implementing FA2, along with explanations of these mistakes and their corrections, because I learnt the most from mistakes.

All numbers in this series are from a single RTX 4070 Super (Ada, sm_89, 12 GB, roughly 100 KB of shared memory per SM) under Linux.

This article assumes familiarity with PyTorch and Einops, and no prior experience with Triton. This article was written with the assistance of AI.

If you spot any mistakes in this article, please let me know. I have also written articles on online softmax and FlashAttention by hand to build intuition for these two algorithms.

What I am building

The FlashAttention-2 (FA2) forward pass: for each query tile Q_i, iterate over key and value tiles K_j and V_j, rescaling the running output O_i on every step and dividing by l only once at the end (FA1 re-normalised at every step). The outer loop goes through the query tiles; on the GPU, each query tile gets its own program, and all the programs run in parallel.

The kernel computes standard scaled-dot-product attention without materialising the N×N score matrix. For a query tile Q_i of B_q rows and a key tile K_j of B_k rows, it computes the B_q × B_k score block, folds it into a running softmax, and moves to the next key tile. When all key tiles have been visited, the output rows for Q_i are complete and are written once.

Three running quantities per query tile make that possible, initialised on line 6:

m_i, the running maximum score for each row (line 10)
l_i, the running row sum of exp(S_i - m_i) (line 12)
O_i, the running sum of exp(S_i - m_i) times the value rows, kept unnormalised (line 13)

The superscript (j) marks the value after the j-th key tile.

On every step the previous l_i^(j-1) and O_i^(j-1) are rescaled by exp(m_i^(j-1) - m_i^(j)), which is 1 when the maximum is unchanged, before the new tile's contribution is added. This is the online softmax, which is applied one tile at a time.

At the end of each query tile, the kernel also writes a second output L, which is computed on line 16 and stored on line 18. L is the per-row logsumexp of the scaled scores L_i = m_i + log(l_i), which equals log Σ_k exp(S_ik) over the whole row. The backward pass will use it to recompute the softmax probabilities tile by tile as P = exp(S - L), since exp(S - m) / l = exp(S - m - log(l)) = exp(S - L). Storing L instead of P is what lets FA's backward avoid the NxN matrix too. The backward will be covered in a later article. For now L is computed, stored, and tested, but not used.

Step 1: the tiled forward in PyTorch

Before diving into Triton, I coded FA2 with PyTorch to familiarise myself with the algorithm. The PyTorch implementation is slow by design because it has a Python double loop over tiles, and launches a handful of small CUDA ops per iteration. The goal here is to be correct, not fast.

# flashattention_autograd_function_pytorch.py
import math
import torch
import einops


class FlashAttentionPytorch(torch.autograd.Function):

    @staticmethod
    def forward(ctx, Q, K, V, is_causal=False):
        # Note: Tile size is fixed at 16 as a simplifying assumption
        tile_size = 16
        # Split the sequence dimension into tiles: (..., T, B, d)
        # Leading dims are arbitrary (batch, heads, ...). The sequence axis N is split
        # into Tq tiles of Bq rows, so N must be a multiple of tile_size.
        Q_t = einops.rearrange(Q, "... (Tq Bq) d -> ... Tq Bq d", Bq=tile_size)
        K_t = einops.rearrange(K, "... (Tk Bk) d -> ... Tk Bk d", Bk=tile_size)
        V_t = einops.rearrange(V, "... (Tk Bv) d -> ... Tk Bv d", Bv=tile_size)

        O = torch.empty_like(Q)
        L = torch.empty(Q.shape[:-1], device=Q.device, dtype=Q.dtype)
        scale = 1.0 / math.sqrt(Q.shape[-1])

        for i in range(Q_t.shape[-3]):                      # outer loop: query tiles
            Q_i = Q_t[..., i, :, :]
            O_i = torch.zeros_like(Q_i)
            l_i = torch.zeros(Q_i.shape[:-1] + (1,), device=Q.device, dtype=Q.dtype)
            m_i = torch.full(Q_i.shape[:-1] + (1,), -torch.inf, device=Q.device, dtype=Q.dtype)

            for j in range(K_t.shape[-3]):                  # inner loop: key tiles
                K_j = K_t[..., j, :, :]
                V_j = V_t[..., j, :, :]
                S_ij = einops.einsum(Q_i, K_j, "... Bq d, ... Bk d -> ... Bq Bk") * scale

                m_new = torch.maximum(m_i, S_ij.amax(dim=-1, keepdim=True))
                P_ij = torch.exp(S_ij - m_new)
                alpha = torch.exp(m_i - m_new)               # rescale factor for the old state
                l_i = alpha * l_i + P_ij.sum(dim=-1, keepdim=True)
                O_i = alpha * O_i + einops.einsum(P_ij, V_j, "... Bq Bk, ... Bk d -> ... Bq d")
                m_i = m_new

            O[..., i * tile_size:(i + 1) * tile_size, :] = O_i / l_i
            L[..., i * tile_size:(i + 1) * tile_size] = (m_i + torch.log(l_i)).squeeze(-1)

        ctx.save_for_backward(Q, K, V, O, L)
        ctx.is_causal = is_causal
        return O, L

A few points to note:

m_i starts at negative infinity so that the first tile's row maximum "wins" in line 10 of the algorithm above. m_i starting at negative infinity also ensures on the first step alpha = exp(-inf - m_new) is exactly 0, not NaN, so that the rescale in lines 12 and 13 runs unchanged and there is no special case to handle for the first iteration.
O_i is the weighted sum of value rows, and the softmax needs that sum divided by the total weight l_i. There are two ways to book-keep: (i) keep the sum and divide once at the end (FA2); or (ii) keep the average and re-divide every time a new tile arrives (FA1). FA2's approach saves two elementwise passes over a B_q x d tile on every inner loop.
L is stored, not m and l separately. The backward pass recomputes each tile of S from Q and K, then recovers that tile's probabilities as exp(S - L). For context, assuming N = 4096, S with dimension NxN per head per batch element will have about 16.8 million entries, or about 33MB in bf16. L with dimension Nx1 per head per batch element contains 4096 entries, or about 16KB.
The is_causal flag is accepted and ignored. The PyTorch version does not implement causal masking.

A primer to `tl.make_block_ptr`

Before I introduce the implementation of FA2 in Triton, following Stanford CS336 assignment 2, I will introduce tl.make_block_ptr.

tl.make_block_ptr creates a block pointer. A block pointer is a tile sliding over a tensor. The code to create a block pointer in Triton is below. Using x_block_ptr, the kernel below adds up each row of x, with each column multiplied by a weight. The kernel is called weighted_sum_fwd. Instead of loading the entire row at once, the kernel slides across the columns of x one tile at a time, in a loop.

x_block_ptr = tl.make_block_ptr(
    x_ptr,
    shape=(NUM_ROWS, D),
    strides=(x_stride_row, x_stride_dim),
    offsets=(row_tile_idx * ROWS_TILE_SIZE, 0),
    block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),
    order=(1, 0),
)

# Walk across the columns of x, one tile per step.
# cdiv divides and rounds up, so no column is missed.
for i in range(tl.cdiv(D, D_TILE_SIZE)):
    row = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option="zero")
    ...  # multiply by weight, add to the totals
    x_block_ptr = x_block_ptr.advance((0, D_TILE_SIZE))   # move one tile right

There are 6 arguments in tl.make_block_ptr, illustrated in the image below. I will go through them one by one. The illustrations and HTML visualisations are drawn by Claude Opus 5.5.

The blue rectangle is the tile, i.e. the part of the tensor the block pointer points at. It sits over part of the tensor: .advance() slides it, and tl.load reads what is inside. Triton calls it a block, which is where make_block_ptr and block_shape get their names.

Why use a block pointer

A Triton kernel runs many programs in parallel. Each program works on a small tile of data. To read a tile, a program needs the memory address of every element in the tile. It also must not read past the edges of the tensor. There are two ways to do this.

The classical way with pointer maths

Before block pointers, programmers built the addresses directly, and many codebases, including the official Triton tutorials, still do. Below is the weighted_sum_fwd kernel written the classical way. I will explain how to read Triton code written the classical way as it is still a useful skill to master.

import triton
import triton.language as tl

@triton.jit
def weighted_sum_fwd(
    x_ptr, weight_ptr,              # inputs: where x and weight start in memory
    output_ptr,                     # output: where the results go
    x_stride_row, x_stride_dim,     # x is 2D: one stride per axis
    weight_stride_dim,              # weight is 1D: one stride (usually 1)
    output_stride_row,              # output is 1D: one stride (usually 1)
    NUM_ROWS, D,                    # the size of x
    ROWS_TILE_SIZE: tl.constexpr,   # tile sizes, fixed when the kernel is compiled
    D_TILE_SIZE: tl.constexpr,
):
    row_tile_idx = tl.program_id(0)

    # 1. Which rows this program owns
    rows = row_tile_idx * ROWS_TILE_SIZE + tl.arange(0, ROWS_TILE_SIZE)
    row_mask = rows < NUM_ROWS

    output = tl.zeros((ROWS_TILE_SIZE,), dtype=tl.float32)
    for i in range(tl.cdiv(D, D_TILE_SIZE)):
        # 2. Which columns this step covers
        cols = i * D_TILE_SIZE + tl.arange(0, D_TILE_SIZE)
        col_mask = cols < D

        # 3. A 2D grid of addresses, built by broadcasting
        x_ptrs = (x_ptr + rows[:, None] * x_stride_row
                        + cols[None, :] * x_stride_dim)
        w_ptrs = weight_ptr + cols * weight_stride_dim

        # 4. Masks for the edges, combined by hand
        row = tl.load(x_ptrs, mask=row_mask[:, None] & col_mask[None, :], other=0.0)
        weight = tl.load(w_ptrs, mask=col_mask, other=0.0)

        output += tl.sum(row * weight[None, :], axis=1)

    tl.store(output_ptr + rows * output_stride_row, output, mask=row_mask)

The code to start the kernel in Python is below, i.e. the launch code. grid says how many programs to run, and the other values fill the signature of weighted_sum_fwd in order. The block pointer version later in this section has the same signature, so the same launch code works for both.

import torch
import triton

x = torch.randn(5, 10, device="cuda")      # NUM_ROWS = 5, D = 10
weight = torch.randn(10, device="cuda")
output = torch.empty(5, device="cuda")

NUM_ROWS, D = x.shape
ROWS_TILE_SIZE, D_TILE_SIZE = 2, 4

# One program per ROWS_TILE_SIZE rows: 5 rows need cdiv(5, 2) = 3 programs
grid = (triton.cdiv(NUM_ROWS, ROWS_TILE_SIZE),)
weighted_sum_fwd[grid](
    x, weight, output,              # a tensor is passed as its start address
    x.stride(0), x.stride(1),       # 10, 1
    weight.stride(0),               # 1
    output.stride(0),               # 1
    NUM_ROWS, D,
    ROWS_TILE_SIZE=ROWS_TILE_SIZE, D_TILE_SIZE=D_TILE_SIZE,
)

What each line does

I follow the case from the launch code above. x has NUM_ROWS = 5 rows and D = 10 columns, stored row by row. So x_stride_row = 10 (one row down skips 10 elements) and x_stride_dim = 1 (one column right is the next element). The kernel also takes one stride for each 1D tensor, weight_stride_dim and output_stride_row. Both are 1, because the elements of each vector sit next to each other in memory. Tiles are ROWS_TILE_SIZE = 2 rows by D_TILE_SIZE = 4 columns. For the rest of this section, I will trace program 2, which handles the last row.

In one sentence: each program adds up x[r, c] * weight[c] over all columns c, for each of its rows r. It works 4 columns at a time and writes one total per row at the end.

Setup

row_tile_idx = tl.program_id(0)

Because there are 5 rows, and each program handles 2 rows, 5 rows need 3 programs (5 ÷ 2, rounded up), hence grid = 3. The programs are numbered 0, 1 and 2. When the kernel runs, the GPU gives each program its own number, and tl.program_id(0) reads it. The 0 asks for the first axis of this numbering. Here grid has only one entry, so there is only one axis. The program I am tracing gets 2, so it handles rows 4 and 5. It is the last program, and its tile is only half full: row 4 exists, but row 5 does not. Every line below runs separately inside each program, and the programs may run at the same time, in any order.

1. Which rows this program owns

rows = row_tile_idx * ROWS_TILE_SIZE + tl.arange(0, ROWS_TILE_SIZE)

tl.arange(0, 2) makes the small vector [0, 1]. It is a Triton tensor that holds both numbers at once. The line computes 2 * 2 = 4 and adds it to each element, so rows = [4, 5]. These are the rows this program handles. Triton needs to know the length of this vector when it compiles the kernel, which is why ROWS_TILE_SIZE is marked tl.constexpr, which means it is a value fixed at compile time.

row_mask = rows < NUM_ROWS

This compares each element with 5: [4 < 5, 5 < 5] gives [True, False]. This is a mask: one True or False per element, where True means "this spot exists, so it is safe to use". Row 5 does not exist, because x only has rows 0 to 4. This line runs once before the loop because the rows do not change.

output = tl.zeros((ROWS_TILE_SIZE,), dtype=tl.float32)

This makes one running total per row, starting at [0.0, 0.0]. It uses 32-bit floats even if x holds 16-bit numbers, because adding many low-precision numbers builds up rounding error. You can try the CS336 Assignment 2 experiments with floating points to understand why accumulators need to have higher precision.

for i in range(tl.cdiv(D, D_TILE_SIZE)):

tl.cdiv divides and rounds up: cdiv(10, 4) = 3, so i takes the values 0, 1 and 2. Rounding up matters because plain division would give 2 steps and skip columns 8 and 9. Unlike the rows and row_mask vectors above, this is a normal loop where the steps run one after another. The for loop is not parallelised because each step adds to the totals left by the step before.

2. Which columns this step covers

cols = i * D_TILE_SIZE + tl.arange(0, D_TILE_SIZE)

The same pattern as rows, but across the columns: [0, 1, 2, 3] when i = 0, then [4, 5, 6, 7] when i = 1, then [8, 9, 10, 11] when i = 2. I will walk through the case when i = 2, because that is where the edge of the tensor matters.

col_mask = cols < D

This gives [True, True, False, False]. Columns 10 and 11 do not exist, hence are masked as False.

3. A 2D grid of addresses

x_ptrs = (x_ptr + rows[:, None] * x_stride_row
                + cols[None, :] * x_stride_dim)

tl.load needs one address for each of the 8 elements in the tile. This line builds all 8 at once, from rows = [4, 5] (step 1) and cols = [8, 9, 10, 11] (step 2).

The strides turn row and column numbers into distances in memory. Each row of x holds 10 elements, so row 4 starts 4 × 10 = 40 elements from the start of x, and row 5 starts at 50.

Each column adds 1 more element (x_stride_dim = 1). rows[:, None] makes the rows a column, and cols[None, :] makes the columns a row. On their own, rows and cols are just flat lists of numbers:

rows = [4, 5]
cols = [8, 9, 10, 11]

I want one address for every combination of a row and a column: 2 rows × 4 columns = 8 addresses. Think of an addition table, with rows down the side, columns across the top, and each cell adding its row's number to its column's number.

rows[:, None] stands the list up as a column (2 rows, 1 column):

[[4],
[5]]

cols[None, :] lays the list down as a row (1 row, 4 columns):

[[8, 9, 10, 11]]

The : means "keep all the values," and None means "add a new direction here, with size 1." Where you put the None decides whether the list stands up or lies down.

When you add a column to a row, Triton copies the column across and the row down until both are 2 × 4. Then it adds cell by cell. That's the addition table:

[[40],   +   [[8, 9, 10, 11]]   →   [[48, 49, 50, 51],
 [50]]                               [58, 59, 60, 61]]

Broadcasting stretches both to a 2 × 4 shape and adds them cell by cell (panel A in the picture below):

	col 8	col 9	col 10	col 11
row 4 (starts at 40)	48	49	50	51
row 5 (starts at 50)	58	59	60	61

Each number counts elements from x_ptr, the address of x[0, 0]. So x_ptr + 48 is the address of x[4, 8]: 40 to reach row 4, plus 8 to reach column 8. Triton counts in elements, not bytes.

w_ptrs = weight_ptr + cols * weight_stride_dim

The weight tensor is a single list with one number per column of x, and every row of x uses the same weights. So only the column positions matter: four addresses, weight_ptr + [8, 9, 10, 11]. No rows are involved, so no grid is needed.

4. Loading safely at the edges

row = tl.load(x_ptrs, mask=row_mask[:, None] & col_mask[None, :], other=0.0)

First, the two masks are combined into a grid the same way as the addresses. row_mask becomes a column, col_mask becomes a row, and & ("and") makes a cell True only if both of its parts are True. This is panel B:

[[True,  True,  False, False],
 [False, False, False, False]]

tl.load then reads memory only at the addresses where the mask is True. Where it is False, it does not touch memory and uses the value given by other (here 0.0) instead. The result is the 2 × 4 tile in panel C: [[x[4,8], x[4,9], 0, 0], [0, 0, 0, 0]].

The masking matters because x has 5 × 10 = 50 elements, at addresses 0 to 49. Every False cell is at address 50 or more, past the end of x. Reading there returns garbage values or crashes the kernel.

Program 2 on its last loop step (i = 2). A: the 8 addresses, built from rows and cols. B: the mask, built the same way from row_mask and col_mask. C: the tile that tl.load returns.

weight = tl.load(w_ptrs, mask=col_mask, other=0.0)

This gives [weight[8], weight[9], 0, 0].

output += tl.sum(row * weight[None, :], axis=1)

weight[None, :] makes the weights a row, so they multiply every row of the tile. The product has dimension 2 × 4. tl.sum(..., axis=1) adds across the columns, which are axis 1, giving one number per row. Those numbers are added to the running totals in output.

The 0s at the edges add nothing, so the half-empty tile gives the right answer. Row 5 is all 0s, so its total stays 0.

After the loop

tl.store(output_ptr + rows * output_stride_row, output, mask=row_mask)

After the last step, output holds one total per row: [total for row 4, total for row 5]. There are no columns left, because tl.sum(..., axis=1) added them up on every step.

The output tensor is a single list with one number per row of x, so each total needs just one address. With output_stride_row = 1, rows * output_stride_row is [4, 5], so the addresses are output_ptr + [4, 5].

tl.store then matches up three lists of 2, position by position:

	first	second
address	`output_ptr + 4`	`output_ptr + 5`
value	total for row 4	total for row 5
`row_mask`	`True`	`False`

It writes row 4's total, which is the sum of x[4, c] * weight[c] over all 10 columns. It skips row 5, which would be past the end of output.

Unlike in step 3, rows is not turned into a column here. [:, None] and [None, :] are only needed to combine two lists into every combination, like rows × columns. Here there is only one list, so matching values one to one is exactly right.

Each program writes only its own rows, so no two programs ever write to the same place.

What you have to get right with pointer maths

With pointer maths, every detail is your job:

the index vectors, made with tl.arange,
the broadcasting, with [:, None] and [None, :] on the right vectors,
the strides, each multiplied into the right index,
one mask per axis, combined with &, plus an other value for the spots the mask leaves out,
new column indices on every loop step.

A slip in any of these gives wrong numbers, and there may not be any error message.

The block pointer way

Here is the same kernel written with block pointers. The signature is the same, so the launch code above works unchanged. Three block pointers are made once, before the loop. After that, three calls do the rest: tl.load reads a tile, .advance() moves it, and tl.store writes a tile.

import triton
import triton.language as tl

@triton.jit
def weighted_sum_fwd(
    x_ptr, weight_ptr,              # inputs: where x and weight start in memory
    output_ptr,                     # output: where the results go
    x_stride_row, x_stride_dim,     # x is 2D: one stride per axis
    weight_stride_dim,              # weight is 1D: one stride (usually 1)
    output_stride_row,              # output is 1D: one stride (usually 1)
    NUM_ROWS, D,                    # the size of x
    ROWS_TILE_SIZE: tl.constexpr,   # tile sizes, fixed when the kernel is compiled
    D_TILE_SIZE: tl.constexpr,
):
    row_tile_idx = tl.program_id(0)

    # 1. Describe each tensor, and where this program's tile starts
    x_block_ptr = tl.make_block_ptr(
        x_ptr,
        shape=(NUM_ROWS, D),
        strides=(x_stride_row, x_stride_dim),
        offsets=(row_tile_idx * ROWS_TILE_SIZE, 0),
        block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),
        order=(1, 0),
    )
    weight_block_ptr = tl.make_block_ptr(
        weight_ptr,
        shape=(D,),
        strides=(weight_stride_dim,),
        offsets=(0,),
        block_shape=(D_TILE_SIZE,),
        order=(0,),
    )
    output_block_ptr = tl.make_block_ptr(
        output_ptr,
        shape=(NUM_ROWS,),
        strides=(output_stride_row,),
        offsets=(row_tile_idx * ROWS_TILE_SIZE,),
        block_shape=(ROWS_TILE_SIZE,),
        order=(0,),
    )

    output = tl.zeros((ROWS_TILE_SIZE,), dtype=tl.float32)
    for i in range(tl.cdiv(D, D_TILE_SIZE)):
        # 2 to 4. Columns, addresses and edge masks: all handled inside tl.load
        row = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option="zero")
        weight = tl.load(weight_block_ptr, boundary_check=(0,), padding_option="zero")

        output += tl.sum(row * weight[None, :], axis=1)

        # Move both tiles one tile to the right for the next step
        x_block_ptr = x_block_ptr.advance((0, D_TILE_SIZE))
        weight_block_ptr = weight_block_ptr.advance((D_TILE_SIZE,))

    tl.store(output_block_ptr, output, boundary_check=(0,))

Triton does the index maths, the broadcasting and the masks for you. The code is not shorter here, because the make_block_ptr calls are long. However, block pointers help to abstract away the handling of memory addresses, which is error-prone.

Each piece of the hand-written version has a matching argument:

By hand	With a block pointer
`x_ptr`	`base`
`rows < NUM_ROWS`, `cols < D`	`shape` plus `boundary_check`
`* x_stride_row`, `* x_stride_dim`	`strides`
`row_tile_idx * ROWS_TILE_SIZE + ...`, `i * D_TILE_SIZE + ...`	`offsets`, then `.advance()`
`tl.arange(0, ROWS_TILE_SIZE)`, `tl.arange(0, D_TILE_SIZE)`	`block_shape`
`[:, None]` and `[None, :]`	done for you
`other=0.0`	`padding_option="zero"`
nothing	`order`, an extra hint about memory layout

Two lines are exactly the same in both versions: row_tile_idx = tl.program_id(0) at the start, and the maths, output += tl.sum(row * weight[None, :], axis=1). For the same inputs, both versions give the same result.

The kernel makes three block pointers: one for x (2D), one for weight (1D), and one for output (1D). The rest of this guide takes the arguments one at a time, using x since it is 2D and hence more general than 1D.

`base` as the starting point in memory of a tensor

GPU memory is a long line of slots and not a grid. A 2D tensor is stored on that line row by row, i.e. all of row 0, then all of row 1, and so on.

The number in each cell is its slot in memory. Row 0 fills slots 0 to 3, row 1 fills slots 4 to 7, and so on.

The first argument, base, is the address of the very first element, x[0, 0]. On its own, it says nothing about the size or layout of the tensor. The other arguments add that.

Pass the start of the whole tensor, not the start of your tile. You choose the tile later, with offsets.

In the kernel: x_ptr, weight_ptr and output_ptr.

`shape` tells Triton how big the tensor is

shape is the size of the full tensor, not the size of the tile. For x it is (NUM_ROWS, D). For weight it is (D,), and for output it is (NUM_ROWS,).

The thick outline is shape. The striped red cells are outside the tensor. The last tile in a row can hang over the edge, like this one.

Triton uses shape to know where the tensor ends. NUM_ROWS and D do not have to be multiples of the tile size, so a tile can stick out past the edge. When you load with boundary_check, Triton compares each spot in the tile with shape and fills the outside spots with a padding value, such as 0. The edges section below shows this.

shape can be a normal runtime number; it does not need to be a tl.constexpr (a value fixed at compile time).

`strides` tells Triton how far apart neighbours are in memory

strides has one number per axis. Each number says how many slots you jump in memory when you take one step along that axis.

address of x [r, c] = base + r \times row stride + c \times column stride

Try the three layouts. The grid is the same every time, but the elements land in different memory slots. The strides are what tell Triton which layout it is looking at.

Strides count elements, not bytes. PyTorch's x.stride() gives you exactly these numbers, so the Python side passes x.stride(0) and x.stride(1).

Do not assume stride_row = D. Transposed and sliced tensors break that rule. With the real strides, the kernel reads the right numbers.

`offsets` tells Triton where the tile starts

offsets is the (row, column) of the tile's top-left element. It counts elements.

Each program (one copy of the kernel running on the GPU) handles its own band of rows. tl.program_id(0) tells it which band. To turn a band number into a row number, multiply by the band height: row_tile_idx * ROWS_TILE_SIZE.

The column offset is 0 because every program starts at the left edge. The loop then walks it to the right.

The other two pointers follow the same idea. weight starts at (0,). output starts at (row_tile_idx * ROWS_TILE_SIZE,), so each program writes exactly the rows it read. Offsets can be computed at run time, and usually are.

`block_shape` tells Triton the size of a tile

block_shape is the size of the tile. tl.load always gives you a tensor of exactly this shape, even when part of the tile is past the edge.

Triton compiles the kernel for one fixed tile size, so these numbers must be known at compile time. That is why the kernel marks ROWS_TILE_SIZE and D_TILE_SIZE as tl.constexpr. Triton requires each size to be a power of 2. Triton splits a tile evenly across the GPU's threads, which come in groups sized by powers of 2, so this keeps the split even. Tensors can be any size, and tiles at the edges are padded.

Bigger tiles mean fewer loop steps, but each program needs more registers (the small, fast storage inside the GPU chip, which runs out quickly). Picking good sizes is an empirical question, answered by profiling.

`order` tells Triton which axis is side by side in memory

Walk through memory one slot at a time and watch each element's position. In a matrix stored row by row, the column number changes at every step, while the row number changes only once per row. So the column axis changes "fastest" and the row axis "slowest". order lists the axes from fastest to slowest, so here it is (1, 0).

Each stop shows the memory slot of that element. Solid lines connect neighbours in memory. Dashed lines jump to the next row (left) or the next column (right).

For a normal row-major matrix, elements in a row sit side by side, so axis 1 comes first: order=(1, 0). A 1D tensor has only one axis: order=(0,).

order is a layout hint. Triton uses it to pick fast ways to load and store.

`.advance()` tells Triton how to move the tile

ptr.advance((d_row, d_col)) moves the tile by that many elements along each axis. Only the offsets change. base, shape, strides, block_shape and order stay the same.

Blue is the starting tile. Moving by (0, D_TILE_SIZE) gives orange. Moving by (ROWS_TILE_SIZE, 0) gives purple.

advance does not change the pointer you call it on. It returns a new one. So always save the result: x_block_ptr = x_block_ptr.advance(...).

In weighted_sum_fwd, each loop step moves x one tile to the right with (0, D_TILE_SIZE), and moves weight with (D_TILE_SIZE,). Both move by the same amount, so the columns of x and the entries of weight stay lined up. The downward move (purple) is not in this kernel. Instead, the bands of rows are spread across parallel programs, which is managed through offsets.

# D = 8, D_TILE_SIZE = 4, so the loop runs cdiv(8, 4) = 2 times
i = 0   load x at (0, 0), then advance to (0, 4)
i = 1   load x at (0, 4), then advance to (0, 8)
# loop ends. (0, 8) is past the edge, but it is never loaded.

Moving the tile past the edge is fine; only loading or storing there needs care.

`boundary_check` tells Triton what happens past the edge

These two arguments belong to tl.load and tl.store, not to make_block_ptr. But they only work because you gave the block pointer a shape. boundary_check lists the axes to check, i.e. 0 means rows and 1 means columns. padding_option says what to put in the spots that are outside, e.g. "zero" for 0, or "nan" for NaN ("not a number", a special float value that turns any sum it touches into NaN).

I filled x with 0, 1, 2, 3, ... so you can see where each loaded number came from.

"zero" is right for this kernel because the outside spots become 0 in both row and weight, so row * weight adds nothing to the sum.

tl.store with boundary_check skips spots outside shape. The kernel stores output with boundary_check=(0,) because the last program may have rows past NUM_ROWS, and those must not be written.

If you check an axis but leave out padding_option, memory past the edge is not read, but the values in those spots are undefined.

Put it all together

This is the whole weighted_sum_fwd kernel. Each program sums a band of rows of x, weighted by weight, one tile at a time. Change the sizes, pick a program, and step through the loop.

Cheat sheet

Argument	In plain words	x	weight	output
`base`	Address of element 0	`x_ptr`	`weight_ptr`	`output_ptr`
`shape`	Size of the whole tensor	`(NUM_ROWS, D)`	`(D,)`	`(NUM_ROWS,)`
`strides`	Memory jump for one step on each axis	`(x_stride_row, x_stride_dim)`	`(weight_stride_dim,)`	`(output_stride_row,)`
`offsets`	Where the tile starts, in elements	`(row_tile_idx * ROWS_TILE_SIZE, 0)`	`(0,)`	`(row_tile_idx * ROWS_TILE_SIZE,)`
`block_shape`	Tile size, fixed at compile time	`(ROWS_TILE_SIZE, D_TILE_SIZE)`	`(D_TILE_SIZE,)`	`(ROWS_TILE_SIZE,)`
`order`	Axes from fastest to slowest in memory	`(1, 0)`	`(0,)`	`(0,)`

Recipe: a 2D row-major tensor

ptr = tl.make_block_ptr(
    base_ptr,
    shape=(M, N),
    strides=(stride_m, stride_n),   # t.stride(0), t.stride(1)
    offsets=(pid * BM, 0),          # in elements
    block_shape=(BM, BN),           # tl.constexpr, powers of 2
    order=(1, 0),                   # axis 1 has stride 1
)
tile = tl.load(ptr, boundary_check=(0, 1), padding_option="zero")
ptr = ptr.advance((0, BN))          # save the result

Recipe: a 1D tensor

ptr = tl.make_block_ptr(
    base_ptr,
    shape=(N,),          # note the comma: (N) is just N
    strides=(stride,),   # t.stride(0), usually 1
    offsets=(start,),    # in elements
    block_shape=(B,),    # tl.constexpr, a power of 2
    order=(0,),          # only one axis
)
tile = tl.load(ptr, boundary_check=(0,), padding_option="zero")
ptr = ptr.advance((B,))  # save the result; note the comma again

Common mistakes with block pointers

The examples use a 2D tensor x of shape (N, D), split into tiles of rows. After the first example, snippets show only the arguments that matter.

Passing the tile's address as base. Pass the start of the tensor and use offsets for the tile.

  row_tile_idx = tl.program_id(0)

  # Wrong: base already points at the tile
  tile_ptr = x_ptr + row_tile_idx * ROWS_TILE_SIZE * stride_row
  x_block_ptr = tl.make_block_ptr(
      base=tile_ptr,
      shape=(N, D),
      strides=(stride_row, stride_col),
      offsets=(0, 0),
      block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),
      order=(1, 0),
  )

  # Right: base is the tensor start, offsets locate the tile
  x_block_ptr = tl.make_block_ptr(
      base=x_ptr,
      shape=(N, D),
      strides=(stride_row, stride_col),
      offsets=(row_tile_idx * ROWS_TILE_SIZE, 0),
      block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),
      order=(1, 0),
  )

Putting a tile number in offsets. Use row_tile_idx * ROWS_TILE_SIZE, not row_tile_idx.

  # Wrong: tile 3 starts at row 3
  offsets=(row_tile_idx, 0),

  # Right: tile 3 starts at row 3 * ROWS_TILE_SIZE
  offsets=(row_tile_idx * ROWS_TILE_SIZE, 0),

Putting the tile size in shape. shape is the whole tensor. The tile size goes in block_shape.

  # Wrong: shape describes the tile
  shape=(ROWS_TILE_SIZE, D_TILE_SIZE),
  block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),

  # Right: shape describes the whole tensor
  shape=(N, D),
  block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),

Using a runtime value in block_shape. It must be tl.constexpr, and in practice a power of 2.

  # Wrong: tile sizes are runtime arguments
  @triton.jit
  def kernel(x_ptr, N, D, stride_row, stride_col, ROWS_TILE_SIZE, D_TILE_SIZE):
      ...

  # Right: tile sizes are constexpr
  @triton.jit
  def kernel(x_ptr, N, D, stride_row, stride_col,
             ROWS_TILE_SIZE: tl.constexpr, D_TILE_SIZE: tl.constexpr):
      ...

  # On the host, round up to a power of 2
  D_TILE_SIZE = triton.next_power_of_2(D)

Writing (D) for a 1D tuple. Python reads that as a plain number. Write (D,).

  # Wrong: every "tuple" here is a plain number
  shape=(D),
  strides=(stride),
  offsets=(tile_idx * TILE_SIZE),
  block_shape=(TILE_SIZE),
  order=(0),

  # Right
  shape=(D,),
  strides=(stride,),
  offsets=(tile_idx * TILE_SIZE,),
  block_shape=(TILE_SIZE,),
  order=(0,),

Giving the tuples different lengths. shape, strides, offsets, block_shape, and order need one entry per dimension.

  # Wrong: strides and offsets have one entry for a 2D tensor
  shape=(N, D),
  strides=(stride_row,),
  offsets=(row_tile_idx * ROWS_TILE_SIZE,),
  block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),
  order=(1, 0),

  # Right: two entries everywhere
  shape=(N, D),
  strides=(stride_row, stride_col),
  offsets=(row_tile_idx * ROWS_TILE_SIZE, 0),
  block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),
  order=(1, 0),

Assuming stride_row = D. Pass the real strides from x.stride().

  x = y.t()  # shape (N, D), but x.stride() == (1, N)

  # Wrong: assumes x is contiguous
  kernel[grid](x, N, D, D, 1, ROWS_TILE_SIZE=32, D_TILE_SIZE=64)

  # Right: works for transposed and sliced tensors too
  kernel[grid](x, N, D, x.stride(0), x.stride(1), ROWS_TILE_SIZE=32, D_TILE_SIZE=64)

Calling ptr.advance(...) without saving the result. It returns a new pointer, so write ptr = ptr.advance(...).

  # Wrong: the result is thrown away, so every iteration loads the first tile
  for _ in range(tl.cdiv(D, D_TILE_SIZE)):
      x = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option="zero")
      x_block_ptr.advance((0, D_TILE_SIZE))

  # Right
  for _ in range(tl.cdiv(D, D_TILE_SIZE)):
      x = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option="zero")
      x_block_ptr = x_block_ptr.advance((0, D_TILE_SIZE))

Putting tile counts in advance. Its offsets are in elements too. Use ptr.advance((0, D_TILE_SIZE)), not ptr.advance((0, 1)).

  # Wrong: moves one column
  x_block_ptr = x_block_ptr.advance((0, 1))

  # Right: moves one tile
  x_block_ptr = x_block_ptr.advance((0, D_TILE_SIZE))

Skipping boundary_check when the sizes do not divide evenly.

  # Wrong: no check, or booleans instead of dimension indices
  x = tl.load(x_block_ptr)
  x = tl.load(x_block_ptr, boundary_check=(True, True))
  tl.store(out_block_ptr, y)

  # Right: check dims 0 and 1 on both the load and the store
  x = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option="zero")
  tl.store(out_block_ptr, y, boundary_check=(0, 1))

Using padding_option="nan" for a sum. One NaN turns the whole sum into NaN.

  # Wrong: padding is undefined, or NaN
  x = tl.load(x_block_ptr, boundary_check=(0, 1))
  x = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option="nan")
  acc += tl.sum(x, axis=1)

  # Right
  x = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option="zero")
  acc += tl.sum(x, axis=1)

Step 2: the Triton forward kernel

The Triton version has the same shape as the PyTorch one, with the outer loop over query tiles replaced by the launch grid. Each program (Triton's name for a thread block) gets one query tile and one batch element, loads Q_i once, and loops over key tiles on its own.

Below is the kernel as it stands at the end of this article; the two lines marked FIX will be explained in Step 3.

# flashattention_autograd_function_triton.py
import math
import torch
import triton
import triton.language as tl


@triton.jit
def flash_fwd_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr, L_ptr,
    stride_qb, stride_qq, stride_qd,
    stride_kb, stride_kk, stride_kd,
    stride_vb, stride_vk, stride_vd,
    stride_ob, stride_oq, stride_od,
    stride_lb, stride_lq,
    N_QUERIES, N_KEYS,
    scale,
    D: tl.constexpr,
    Q_TILE_SIZE: tl.constexpr,
    K_TILE_SIZE: tl.constexpr,
    is_causal: tl.constexpr,
):
    query_tile_index = tl.program_id(0)
    batch_index = tl.program_id(1)

    # Block pointers: a (rows, D) tile into each tensor for this batch element.
    # Q and O tiles start at this program's query tile; K and V start at row 0
    # and are advanced inside the loop.
    Q_block_ptr = tl.make_block_ptr(
        Q_ptr + batch_index * stride_qb,
        shape=(N_QUERIES, D), strides=(stride_qq, stride_qd),
        offsets=(query_tile_index * Q_TILE_SIZE, 0),
        block_shape=(Q_TILE_SIZE, D), order=(1, 0),
    )
    K_block_ptr = tl.make_block_ptr(
        K_ptr + batch_index * stride_kb,
        shape=(N_KEYS, D), strides=(stride_kk, stride_kd),
        offsets=(0, 0), block_shape=(K_TILE_SIZE, D), order=(1, 0),
    )
    V_block_ptr = tl.make_block_ptr(
        V_ptr + batch_index * stride_vb,
        shape=(N_KEYS, D), strides=(stride_vk, stride_vd),
        offsets=(0, 0), block_shape=(K_TILE_SIZE, D), order=(1, 0),
    )
    O_block_ptr = tl.make_block_ptr(
        O_ptr + batch_index * stride_ob,
        shape=(N_QUERIES, D), strides=(stride_oq, stride_od),
        offsets=(query_tile_index * Q_TILE_SIZE, 0),
        block_shape=(Q_TILE_SIZE, D), order=(1, 0),
    )
    L_block_ptr = tl.make_block_ptr(
        L_ptr + batch_index * stride_lb,
        shape=(N_QUERIES,), strides=(stride_lq,),
        offsets=(query_tile_index * Q_TILE_SIZE,),
        block_shape=(Q_TILE_SIZE,), order=(0,),
    )

    # Running state, kept in fp32 regardless of input dtype.
    O_acc = tl.zeros((Q_TILE_SIZE, D), dtype=tl.float32)
    l_acc = tl.zeros((Q_TILE_SIZE, 1), dtype=tl.float32)
    m_acc = tl.full((Q_TILE_SIZE, 1), value=float("-inf"), dtype=tl.float32)

    Q_i = tl.load(Q_block_ptr, boundary_check=(0, 1), padding_option="zero")

    q_pos = (query_tile_index * Q_TILE_SIZE + tl.arange(0, Q_TILE_SIZE))[:, None]

    for j in range(tl.cdiv(N_KEYS, K_TILE_SIZE)):
        k_pos = (j * K_TILE_SIZE + tl.arange(0, K_TILE_SIZE))[None, :]

        K_j = tl.load(K_block_ptr, boundary_check=(0, 1), padding_option="zero")
        V_j = tl.load(V_block_ptr, boundary_check=(0, 1), padding_option="zero")

        S_ij = tl.dot(Q_i, tl.trans(K_j)) * scale          # (Q_TILE, K_TILE), fp32

        # FIX 2: zero-padded keys past N_KEYS score 0, not -inf. Mask them.
        keep = k_pos < N_KEYS
        if is_causal:
            keep = keep & (k_pos <= q_pos)
        S_ij = tl.where(keep, S_ij, -1e6)

        m_new = tl.maximum(m_acc, tl.max(S_ij, axis=1, keep_dims=True))
        P_ij = tl.exp(S_ij - m_new)
        alpha = tl.exp(m_acc - m_new)
        l_acc = alpha * l_acc + tl.sum(P_ij, axis=1, keep_dims=True)

        # FIX 1: the cast must be assigned. tl.dot needs both operands in the
        # same dtype; the fp32 accumulator is passed separately via acc=.
        P_ij = P_ij.to(V_j.dtype)
        O_acc = alpha * O_acc
        O_acc = tl.dot(P_ij, V_j, acc=O_acc)
        m_acc = m_new

        K_block_ptr = K_block_ptr.advance((K_TILE_SIZE, 0))
        V_block_ptr = V_block_ptr.advance((K_TILE_SIZE, 0))

    O_i = (O_acc / l_acc).to(O_block_ptr.type.element_ty)
    tl.store(O_block_ptr, O_i, boundary_check=(0, 1))

    L_i = tl.reshape(m_acc + tl.log(l_acc), (Q_TILE_SIZE,))
    tl.store(L_block_ptr, L_i, boundary_check=(0,))


class FlashAttentionTriton(torch.autograd.Function):
    Q_TILE_SIZE = 16
    K_TILE_SIZE = 16

    @staticmethod
    def forward(ctx, Q, K, V, is_causal=False):
        assert Q.ndim == 3, "expects (batch, seq, head_dim); flatten (B, H, N, D) to (B*H, N, D)"
        assert Q.stride(-1) == 1 and K.stride(-1) == 1 and V.stride(-1) == 1
        B, N_q, D = Q.shape
        N_k = K.shape[1]
        assert D in (16, 32, 64, 128), "block_shape dims must be powers of two"

        O = torch.empty_like(Q)
        L = torch.empty((B, N_q), device=Q.device, dtype=torch.float32)
        grid = (triton.cdiv(N_q, FlashAttentionTriton.Q_TILE_SIZE), B)

        flash_fwd_kernel[grid](
            Q, K, V, O, L,
            Q.stride(0), Q.stride(1), Q.stride(2),
            K.stride(0), K.stride(1), K.stride(2),
            V.stride(0), V.stride(1), V.stride(2),
            O.stride(0), O.stride(1), O.stride(2),
            L.stride(0), L.stride(1),
            N_q, N_k,
            1.0 / math.sqrt(D),
            D=D,
            Q_TILE_SIZE=FlashAttentionTriton.Q_TILE_SIZE,
            K_TILE_SIZE=FlashAttentionTriton.K_TILE_SIZE,
            is_causal=is_causal,
        )
        ctx.save_for_backward(Q, K, V, O, L)
        ctx.is_causal = is_causal
        return O

    @staticmethod
    def backward(ctx, dO):
        raise NotImplementedError("tiled backward is left for future articles")

Legend:

Term	Meaning
grid	How many programs to launch along each axis. Here, `(cdiv(N_q, 16), B*H)`.
`program_id(0)`	Which query tile this program handles: 0, 1, 2, …
`program_id(1)`	Which batch element this program handles. After flattening, that means one (sequence, head) pair.
query tile	A block of 16 consecutive query rows (tokens), handled by one program.
16	The query tile size: rows of Q per program.
`cdiv(a, b)`	`a ÷ b`, rounded up, so a partly filled last tile still gets a program.
`B`	Batch size: the number of sequences processed together.
`H`	The number of attention heads.
`N`	Sequence length: tokens per sequence.
`N_q`	The number of query tokens. In self-attention it equals `N`.
`D`	Head dimension: how many numbers each token has in one head.
`(B, H, N, D)`	The shape of Q, K and V: batch, heads, tokens, head dimension.
`(B*H, N, D)`	Batch and heads merged into one axis, so each (sequence, head) pair is treated as a separate batch element.

Reading it top to bottom:

The grid. program_id(0) indexes query tiles, program_id(1) indexes batch elements, so the launch has cdiv(N_q, 16) × B programs. Multi-head attention is handled by flattening (B, H, N, D) into (B*H, N, D) before the call. With the flattening, the kernel never knows about heads. At 16 rows per program and 4096 tokens × 32 heads, that is 8192 programs, which is plenty to fill a GPU. The problem, as I show in the last section, is that a 16x16 tile is too small to saturate the GPU.

GPUs usually start programs in order with axis 0 changing fastest, so putting query tiles on axis 0 makes programs that run together share the same K and V (better cache reuse), and on NVIDIA GPUs axis 0 is also the only one allowed more than 65,535 programs, which long sequences can need.

Block pointers. make_block_ptr describes a 2D tile of a strided tensor, i.e. the base pointer, the full logical shape, the strides, where the tile starts, and how big it is. order=(1, 0) says the last dimension is the fastest-varying one, i.e. row-major. The payoff is boundary_check and padding_option on load, which handle the ragged last tile for free, and .advance(), which slides the K and V tiles down by one tile per iteration without recomputing addresses.

Accumulators in fp32. O_acc, l_acc and m_acc are fp32 whatever the input dtype, and tl.dot returns fp32 for bf16 inputs. This is what makes it safe to run in bf16. The inputs and the P·V product are lower precision, but the sum across key tiles is higher precision.

One Q load, many K/V loads. Q_i is loaded once before the loop. K_j and V_j are loaded per iteration, and every program loads the entire K and V for its batch element. That is inherent to the algorithm: it is the trade-off FlashAttention-2 makes to avoid the N×N matrix, which is why K/V loads dominate the memory traffic and why tile size matters.

The inner loop is a line-for-line translation of the PyTorch version: score block, mask, new max, unnormalised probabilities, rescale factor, update l, update O, carry m forward. tl.dot(P_ij, V_j, acc=O_acc) fuses the matrix multiply and the accumulate.

The store. boundary_check on the store discards rows of the last tile that fall past N_QUERIES, so the padded rows computed with zeroed Q are never written.

Step 3: testing and bugs found

The testing harness ran one shape, (4, 128, 128, 64), in fp32, at a tolerance of 1e-2, and my kernel passed the test. However, the kernel was still wrong in two ways, and Claude found both bugs using the test file below, which was also generated by Claude.

# test_flashattention_triton.py
import math
import pytest
import torch
import torch.nn.functional as F
from flashattention_autograd_function_triton import FlashAttentionTriton


def reference(q, k, v, is_causal):
    o = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
    s = (q.float() @ k.float().transpose(-1, -2)) / math.sqrt(q.shape[-1])
    if is_causal:
        tril = torch.ones(s.shape[-2:], dtype=torch.bool, device=q.device).tril()
        s = s.masked_fill(~tril, float("-inf"))
    return o, torch.logsumexp(s, dim=-1)


# (batch, n_queries, n_keys, head_dim). Deliberately includes shapes that
# are not multiples of the 16-row tile and non-square attention.
SHAPES = [
    (4, 128, 128, 64),     # the assignment's shape
    (2, 1024, 1024, 64),
    (2, 100, 100, 64),     # ragged last tile on both axes
    (1, 37, 53, 32),       # ragged and non-square
    (1, 200, 96, 128),
]


@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
@pytest.mark.parametrize("is_causal", [False, True])
@pytest.mark.parametrize("B,Nq,Nk,D", SHAPES)
def test_forward_matches_sdpa(B, Nq, Nk, D, is_causal, dtype):
    if is_causal and Nq != Nk:
        pytest.skip("keep causal tests square so the mask convention is unambiguous")
    torch.manual_seed(0)
    q = torch.randn(B, Nq, D, device="cuda", dtype=dtype, requires_grad=True)
    k = torch.randn(B, Nk, D, device="cuda", dtype=dtype, requires_grad=True)
    v = torch.randn(B, Nk, D, device="cuda", dtype=dtype, requires_grad=True)

    o = FlashAttentionTriton.apply(q, k, v, is_causal)
    # L is not returned; pull it out of the saved tensors, as the assignment test does.
    L = [t for t in o.grad_fn.saved_tensors if t.shape == (B, Nq)][0]

    o_ref, L_ref = reference(q, k, v, is_causal)
    # fp32 tl.dot uses TF32 on tensor cores by default; bf16 rounds P before P@V.
    tol = dict(atol=2e-2, rtol=2e-2) if dtype == torch.bfloat16 else dict(atol=1e-2, rtol=1e-2)
    torch.testing.assert_close(o, o_ref, **tol)
    torch.testing.assert_close(L, L_ref, **tol)

Two details in the harness:

L is not part of the function's return value, so the test digs it out of o.grad_fn.saved_tensors, which is the same trick the assignment's test uses and requires the inputs to have requires_grad=True.
The fp32 tolerance is looser than you might expect because Triton's tl.dot on fp32 inputs uses TF32 on Ampere and later unless you pass input_precision="ieee".

TF32 keeps fp32's range but rounds each input to about 3 significant digits before multiplying, which makes matrix maths much faster on Tensor Cores but leaves errors around 1e-3. input_precision="ieee" forces true fp32 maths, which is slower but matches the reference, i.e. the plain PyTorch version the test compares against runs in true fp32.

Bug 1: the casts that never happened

In my original code which passed the harness test, I had the following three lines:

P_ij.to(V_j.type.element_ty)   # "cast P to bf16 before the dot"
...
O_i.to(O_block_ptr.type.element_ty)
...
L_i.to(L_block_ptr.type.element_ty)

In Triton, as in PyTorch, .to() returns a new tensor. .to() does not mutate. All three lines were no-ops, and the comment on the first line described a cast that was not taking place. The kernel still passed the harness test because the test ran in fp32, where casting fp32 to fp32 changed nothing.

The bf16 test found it at compile time rather than as a numerical error. tl.dot required both operands to share a dtype, and after the no-op P_ij was still fp32 while V_j was bf16. The fix: P_ij = P_ij.to(V_j.dtype). The O and L casts are fixed the same way, folded into the store lines.

The lesson I took from this is to check the dtype of every tl.dot operand when a kernel is first written.

Bug 2: padded keys that voted

tl.load with boundary_check and padding_option="zero" fills rows past the end of the tensor with zeros. For Q that is harmless because padded query rows produce garbage that the store then discards. For K it is problematic because a zero key row gives a score of exactly 0 for every query, so the padded keys enter the softmax with weight exp(0 - m) each. When N_KEYS isn't a multiple of the key tile (say 100 keys in tiles of 16), the last tile runs past the real keys, and its extra slots are "phantom" keys that do not exist. Unless they are masked out, the output becomes a weighted average over real keys plus these phantom ones.

The original harness test did not catch this because 128 is a multiple of 16. However, on a (2, 100, 100, 64) tensor, the kernel failed. The fix is the keep = k_pos < N_KEYS mask, applied whether or not attention is causal. It costs one compare and one select per score, which is nothing next to the two matrix multiplies, and it means the kernel is now correct for any sequence length, not just multiples of 16.

Step 4: the benchmark, and the baseline number

This section reports the kernel's speed as a percentage of PyTorch's scaled_dot_product_attention on the same inputs and the same NVIDIA RTX 4070 Super. SDPA is chosen because at head dim 64 in bf16 it dispatches to the FlashAttention-2 kernel that ships inside PyTorch, so "100% of SDPA" means "as fast as FA2". The harness is short and it never changes after this article.

# bench_flashattention.py
import torch
import torch.nn.functional as F
import triton
from flashattention_autograd_function_triton import FlashAttentionTriton


def attn_flops(batch, n, d, is_causal):
    # QK^T and PV: two matmuls of n*n*d multiply-adds each, 2 flops per MAC.
    f = 4 * batch * n * n * d
    return f / 2 if is_causal else f


def bench(B, H, N, D, is_causal, dtype=torch.bfloat16):
    q, k, v = (torch.randn(B * H, N, D, device="cuda", dtype=dtype) for _ in range(3))
    ms_ours = triton.testing.do_bench(
        lambda: FlashAttentionTriton.apply(q, k, v, is_causal), return_mode="median")

    q4, k4, v4 = (t.view(B, H, N, D) for t in (q, k, v))
    ms_sdpa = triton.testing.do_bench(
        lambda: F.scaled_dot_product_attention(q4, k4, v4, is_causal=is_causal),
        return_mode="median")

    flops = attn_flops(B * H, N, D, is_causal)
    tflops = lambda ms: flops / ms * 1e-9
    return ms_ours, ms_sdpa, tflops(ms_ours), tflops(ms_sdpa)


if __name__ == "__main__":
    print(f"{'N':>6} {'causal':>7} {'ours ms':>9} {'sdpa ms':>9} {'ours TF/s':>10} {'sdpa TF/s':>10} {'% sdpa':>7}")
    for N in (512, 1024, 2048, 4096):
        for causal in (False, True):
            a, b, c, d = bench(B=4, H=8, N=N, D=64, is_causal=causal)
            print(f"{N:>6} {str(causal):>7} {a:>9.3f} {b:>9.3f} {c:>10.1f} {d:>10.1f} {100 * b / a:>6.0f}%")

triton.testing.do_bench handles warm-up and flushes L2 between runs, and with return_mode="median" it reports the median, so a single call is enough. The FLOP count is the standard one for attention: two matmuls of N×N×d multiply-adds, halved for causal because the kernel is only supposed to do the lower triangle. Note the word supposed. This kernel computes every tile and masks half of them, so its causal FLOP rate is flattered by the accounting.

Results on the RTX 4070 Super, batch 4, 8 heads, head dim 64, bf16:

N	causal	mine (ms)	SDPA (ms)	mine TFLOP/s	SDPA TFLOP/s	% of SDPA
512	no	0.107	0.063	20.1	34.3	59%
512	yes	0.106	0.055	10.1	19.4	52%
1024	no	0.386	0.152	22.3	56.5	39%
1024	yes	0.385	0.121	11.2	35.5	31%
2048	no	1.534	0.539	22.4	63.7	35%
2048	yes	1.560	0.329	11.0	52.2	21%
4096	no	6.195	2.009	22.2	68.4	32%
4096	yes	6.155	1.120	11.2	61.3	18%

At N = 4096 without causal masking, the kernel runs at 32% of SDPA (22.2 vs 68.4 TFLOP/s). The gap widens as N grows: my kernel reaches about 22 TFLOP/s by N = 1024 and stays there, while SDPA keeps climbing from 34 to 68 TFLOP/s, so the ratio falls from 59% at N = 512 to 32% at N = 4096. Causal attention is worse again, at 18% of SDPA at N = 4096. My kernel's causal runtime is the same as its non-causal runtime at every N (6.16 ms vs 6.20 ms at N = 4096), because every tile above the diagonal is still computed and then masked, while SDPA's causal runtime is roughly half its non-causal runtime.

To check which SDPA backend you're actually comparing against, wrap the reference call in torch.nn.attention.sdpa_kernel([SDPBackend.FLASH_ATTENTION]): if it errors, you're not comparing against FA2 and the percentages mean something else.

What is wrong with this kernel

The kernel is correct and it is slow. Here is why, in the order I plan to fix it. Each item names the cause, not just the symptom, so that future articles can test it.

1. The tiles are too small. 16×16 tiles mean each tl.dot is a 16×64 by 64×16 product. Tensor cores want bigger operands than that to reach their throughput, and each program does so little work per K/V load that the kernel spends its time moving data rather than multiplying it. Every program also reloads all of K and V from L2 or HBM, and with a 16-row Q tile the ratio of loads to useful FLOPs is about eight times worse than with a 128-row tile. The right sizes depend on the card: Ada has roughly 100 KB of shared memory per SM against Hopper's 228 KB, so the 128×64 or 128×128 configurations in the FA2 paper are not automatically right here. In future articles I will dive into the shared-memory budget arithmetic, sweep tile sizes together with num_warps and num_stages, and wrap the winner in @triton.autotune.

2. Causal attention does all the work and throws half away. The mask is applied to every tile, including tiles that lie entirely above the diagonal, where every element is masked. Those tiles contribute nothing to the output and the kernel still loads K and V for them, runs both matmuls and the exponentials, then discards it all. The fix is to stop the key loop at the diagonal, which roughly halves causal runtime, and to only apply the mask on the one tile that actually straddles it.

3. exp instead of exp2. The GPU's fast exponential unit computes 2^x; tl.exp is exp2(x · log₂e), an extra multiply on every score. The standard trick is to fold scale · log₂e into the scaling of Q once per tile, then use tl.math.exp2 directly. It's a small win but it's in the inner loop.

4. The -1e6 mask sentinel. Masking with a large negative number rather than -inf sidesteps a real hazard: a row whose every score is -inf produces exp(-inf - (-inf)) = NaN. A fully masked tile is harmless once a row has seen a real key, because exp(-1e6 - m) underflows to exactly zero in fp32. The remaining risk is a row with no valid key at all: every score is then -1e6, the row maximum is -1e6, and each masked key gets weight exp(0) = 1, so the row silently becomes an average of values it should never see. In this kernel every row can see at least key 0, so it cannot happen yet, but kernels that allow such rows need to handle them explicitly.

5. No backward pass. backward raises NotImplementedError, and the PyTorch reference backward materialises the full N×N score matrix, so it isn't FlashAttention's backward at all. The tiled backward needs the L I have been carrying, a precomputed D = rowsum(dO ∘ O), and a kernel that iterates over query tiles for each key tile to build dK and dV, with dQ accumulated either by atomics or by a second pass.

6. The mask is computed unconditionally. The k_pos < N_KEYS compare and the tl.where run on every tile, including the full interior ones where nothing is masked. It is cheap, but in a kernel this small everything in the inner loop counts. Once tile-skipping is in place the mask can be restricted to boundary tiles.

7. No profile. Everything above is reasoning from first principles. None of it has been confirmed with Nsight Compute, and reasoning about GPU performance without a profiler is how people end up optimising the wrong thing. In future articles I will profile the forward and backward passes, compare these kernels against the GPU's roofline, and identify what is left between this kernel and FA2.

Future articles

Some ideas I have for future articles include:

Take the kernel above and change the tile sizes, num_warps and num_stages. Work through the shared-memory budget for an Ada SM, sweep the configuration space, show the heatmap, then wrap the winner in @triton.autotune and measure what that costs in compile time. The benchmark harness and the test file stay exactly as they are here, so the numbers are directly comparable to the ones above.
Go inside the inner loop for causal tile-skipping, exp2 and the P cast.
Backward pass of FA2.
Profiler and the roofline.
Rewrite the kernel in TileLang and CuTe DSL.

Batched inference by hand

Lewis Won — Sat, 22 Aug 2026 06:53:50 +0000

Table of Contents

Motivation
What is Batched Inference?
Setup
Scenario 1: UNBATCHED Inference (Sequential)
Scenario 2: BATCHED Inference (Parallel)
The Memory Bandwidth vs. Compute Trade-off
Code
Appendix A: The Problem of Uneven Sequence Lengths
Appendix B: Two Tokens Per User — Is Masking Needed?

Motivation

Batching is a foundational optimization in Large Language Model (LLM) inference. Working through the math helps with:

Understanding hardware utilization: GPUs are massively parallel calculators. Generating a single token for a single user leaves most of the GPU's compute units idle while waiting for the model weights to travel from HBM into the chip's compute units. Batching reduces this memory bandwidth bottleneck by amortizing the cost of loading model weights.
Debugging Performance: Understanding batching explains why your inference servers can handle 10 concurrent users at almost the same speed as 1 user, and why batch size is constrained by KV cache memory.

What is Batched Inference?

In a production environment, an LLM server receives multiple requests from different users at the same time.

Imagine User A asks: "What is the capital of France?"
Imagine User B asks: "Write a poem about a cat."

Note: Both prompts get processed in one shot (the prefill phase), and then the server enters the decode phase: generating one token at a time for each user, until they're done. Decode is where almost all the time goes, and it's where batching earns its keep — because at each step each user contributes only a single token, which is too little work to keep a GPU busy on its own. Everything below describes one decode step.

The model needs to predict the next token for User A and the next token for User B.

If we don't batch, the server processes User A's token, and then processes User B's token sequentially. To do this, it must load the massive model weight matrices into the GPU's compute units (SRAM) to process User A, and then load the exact same weights again to process User B.

Batched inference stacks the current token from User A and the current token from User B into a single mathematical matrix. By doing this, the model loads each weight matrix from HBM exactly once per forward pass instead of once per user, and multiplies it against both users' tokens simultaneously.

Note: "Once per forward pass" does not mean the weights get loaded once and stay put. On-chip SRAM holds tens of megabytes; the weights are tens of gigabytes. Each layer's weights are pulled in from HBM, used, and thrown away to make room for the next layer — so a single decoding step streams the whole model through the chip, and the next token does it all over again. Batching doesn't reduce that traffic. Batching just means more users get served per trip.

A note on memory hierachy

On an H100 (SXM variant) the tiers look like this:

Tier	Size	Bandwidth	What lives there
On-chip SRAM (L1 / shared memory, per SM)	256 KB × 132 SMs ≈ 33 MB	tens of TB/s, aggregate	the tile of weights and activations being multiplied right now
L2 cache	50 MB	well above HBM, well below SRAM	whatever was touched recently
HBM3	80 GB	3.35 TB/s	the model weights and every user's KV cache

The memory system has to feed about 990 TFLOPS of BF16 on the tensor cores in order to fully utilize the SM cores, or 990 TFLOPS / 3.35 TB/s ≈ 295 FLOPs per byte. Any kernel that does fewer than roughly 300 floating-point operations per byte it pulls from HBM is waiting on memory. Any kernel that does more is waiting on the SM compute. That crossover is the "ridge point" from memory-bound to compute-bound.

For a batch-size of 1 decode step: it loads 2 bytes (BF16), then do 2 FLOPs, one multiply and one add. That's 1 FLOP per byte, about 300× short of the ridge. The tensor cores sit idle roughly 99.7% of the time, and the step takes as long as streaming the weights takes. For an 8-billion-parameter model in BF16 that's 16 GB / 3.35 TB/s ≈ 4.8 ms per token, a ceiling of about 210 tokens per second for a single user.

Batching increases the numerator, without changing the denominator. With N users in the batch each weight byte still crosses the bus once but now feeds N multiply-adds, so the weight matmuls run at roughly N FLOPs per byte. At N = 10 that's 30× under the ridge, which is why ten users cost about the same per step as one. The math only catches up with the memory somewhere around N ≈ 300, which is the "second, softer limit" from the trade-off section with a number attached. (In practice the crossover comes earlier because the KV cache also has to be read from HBM every step, and that traffic grows with both batch size and context length. While it is possible to disable KV cache during inference, memory use will still grow with longer context.)

Because a H100 only has 33 MB of SRAM against 16 GB of weights, the chip can hold about a five-hundredth of the model at any moment. Hence, in the course of a single decode step the on-chip memory fills and empties roughly 500 times to let every weight through once.

Setup

We will walk through the calculation for predicting the next token for two different sequences.

Assume our embedding dimension is $d = 2$ . We have one weight matrix for our attention Queries: $W_{q}$ . (We'll just focus on projecting the Query, but the exact same logic applies to Keys, Values, and the dense Feed-Forward layers).

Let our Query weight matrix be:

W_{q} = (1011)

We have two users who have sent requests to our server. Both users are mid-conversation. Each has already had some tokens processed (their prompt, plus whatever has been generated so far), and the keys and values of those tokens sit in that user's KV cache. The scenarios below look only at the one new token per user that this step has to handle. The cached history comes back into play at attention, in the BONUS section of the code, where we'll assume it happens to be 3 tokens per user.

User A has a sequence that just produced a new token, $x_{A}$ .
User B has a completely unrelated sequence that just produced a new token, $x_{B}$ .

Let their current token embeddings be:

$x_{A} = [1, 0]$
$x_{B} = [0, 1]$

We need to calculate the new Query vector for both users so they can proceed with their respective attention mechanisms.

Scenario 1: UNBATCHED Inference (Sequential)

Without batching, the server processes the requests sequentially (a batch size of 1).

Step 1: Process User A

The GPU loads the weight matrix $W_{q}$ from its High Bandwidth Memory (HBM) into the small, fast on-chip SRAM that the compute cores read from. It performs a Vector-Matrix multiplication for User A.

q_{A} = x_{A} \times W_{q} = (10) \times (1011) = (11)

Step 2: Process User B

The GPU, having finished with User A, now moves to User B. It must load the weight matrix $W_{q}$ from HBM onto the chip all over again.

q_{B} = x_{B} \times W_{q} = (01) \times (1011) = (01)

We got our answers, but we loaded the weights twice to do very little math.

Scenario 2: BATCHED Inference (Parallel)

Now, let's process the requests together. We take User A's token and User B's token and stack them vertically into a single Input Batch Matrix, $X_{ba t c h}$ .

X_{ba t c h} = (1001)

(Row 1 is User A, Row 2 is User B)

Step 1: Process the Batch

The GPU loads the weight matrix $W_{q}$ from HBM once. It performs a Matrix-Matrix multiplication.

Q_{ba t c h} = X_{ba t c h} \times W_{q} = (1001) \times (1011) = (1011)

Notice the output:

Row 1 is $[1, 1]$ , which is exactly $q_{A}$ .
Row 2 is $[0, 1]$ , which is exactly $q_{B}$ .

We got the exact same answers for both users, but we only had to load the model weights from memory one time which is dominated by fetching W_q, not by the multiply itself. The extra row of math rides along for free in time the GPU was going to spend waiting anyway. That's why doubling the batch here costs almost nothing. A property that holds until the batch gets large enough to saturate the compute cores (see the trade-off section below).

The Memory Bandwidth vs. Compute Trade-off

In LLM inference (specifically during the decoding phase, when generating tokens one by one), the bottleneck is Memory Bandwidth, i.e. the time it takes to move the gigabytes of model weights from the GPU's HBM into the small, fast on-chip SRAM that feeds the processing cores, for every single token step.

Without Batching (Batch Size = 1): The GPU spends a significant portion of time waiting for weights to travel across the memory bus. Once the weights arrive, the processing cores do a tiny Vector-Matrix calculation instantly, throw the weights away, and wait for the next load. The compute cores are starving.
With Batching (Batch Size = N): You load the weights once, but now the compute cores do $N$ times as much math (Matrix-Matrix multiplication). You are using "free" compute time that would have otherwise been spent waiting for memory.

The Limit to Batching

While the model weights are shared across all users, every single user requires their own dedicated KV Cache to store their unique sequence history (as discussed in the KV Cache article). If you increase the batch size too much, the combined KV caches of all those users will exceed your GPU's total HBM, resulting in an Out Of Memory (OOM) error. Even without KV cache during inference, longer context will still demand more memory.

As the batch grows, you eventually give the compute cores more math than they can hide behind the weight loading. Once you cross that point, the GPU is compute-bound: doubling the batch roughly doubles the time of each decoding step, so total throughput flattens out while every individual user waits longer for each token. We can therefore cap batch size based on whichever comes first — hitting the HBM limit, or reaching the maximum per-token latency promised to users.

Code

Below is the PyTorch implementation showing how independent vectors are stacked into a batch. Note the use of torch.bmm (Batched Matrix Multiplication) to handle the attention step where every sequence in the batch has its own independent KV cache.

import torch

# 1. SETUP
# Weight matrix (2x2)
W_q = torch.tensor([[1.0, 1.0], [0.0, 1.0]])

# Individual tokens for User A and User B
x_A = torch.tensor([[1.0, 0.0]]) # Shape: 1x2
x_B = torch.tensor([[0.0, 1.0]]) # Shape: 1x2

print("=== SCENARIO 1: UNBATCHED ===")
# Streamed W_q from HBM onto the chip (Wait...)
q_A = x_A @ W_q
# Streamed W_q from HBM onto the chip AGAIN (Wait...)
q_B = x_B @ W_q

print(f"User A Query: {q_A}")
print(f"User B Query: {q_B}\n")

print("=== SCENARIO 2: BATCHED ===")
# Stack tokens into a batch
X_batch = torch.cat([x_A, x_B], dim=0) # Shape: 2x2
print(f"X_batch shape: {X_batch.shape}")

# Load W_q ONCE and multiply
Q_batch = X_batch @ W_q
print(f"Q_batch shape: {Q_batch.shape}")
print(f"Batched Queries:\n{Q_batch}\n")

# Verify mathematically
assert torch.allclose(q_A, Q_batch[0:1])
assert torch.allclose(q_B, Q_batch[1:2])
print("SUCCESS: Batched output perfectly matches sequential output!\n")

print("=== BATCHED ATTENTION (BONUS) ===")
# Each user has their own history, so the two K caches are stacked along a
# batch dimension into one 3D tensor: (Batch, Seq_Len, Dim). Stacking needs
# both histories to be the same length; assume 3 tokens each (Appendix A
# covers what happens when they aren't).
K_hist = torch.tensor([
    [[1., 0.], [1., 1.], [0., 1.]],  # User A's history (3x2)
    [[0., 1.], [2., 0.], [1., 1.]]   # User B's history (3x2)
])

# The new token's own key joins the cache before attention: a token attends
# to its history AND to itself. Take W_k = I so the key is just the embedding
# (the K projection is a per-row matmul exactly like W_q).
W_k = torch.eye(2)
K_new = (X_batch @ W_k).unsqueeze(1)             # (2, 2) -> (2, 1, 2)
K_cache = torch.cat([K_hist, K_new], dim=1)      # (2, 4, 2): 3 history + 1 new

Q_3D = Q_batch.unsqueeze(1)                      # (2, 2) -> (2, 1, 2)

# bmm multiplies User A's Q with User A's K and User B's Q with User B's K.
# K_cache is transposed on its last two dims: (2, 4, 2) -> (2, 2, 4).
# (Real implementations divide by sqrt(d), then softmax and multiply by V; omitted.)
attention_scores = torch.bmm(Q_3D, K_cache.transpose(1, 2))

print(f"Attention Scores Shape: {attention_scores.shape} -> (Batch, 1, Hist + 1)")
print(f"Scores:\n{attention_scores}")

Output:

=== SCENARIO 1: UNBATCHED ===
User A Query: tensor([[1., 1.]])
User B Query: tensor([[0., 1.]])

=== SCENARIO 2: BATCHED ===
X_batch shape: torch.Size([2, 2])
Q_batch shape: torch.Size([2, 2])
Batched Queries:
tensor([[1., 1.],
        [0., 1.]])

SUCCESS: Batched output perfectly matches sequential output!

=== BATCHED ATTENTION (BONUS) ===
Attention Scores Shape: torch.Size([2, 1, 4]) -> (Batch, 1, Hist + 1)
Scores:
tensor([[[1., 2., 1., 1.]],

        [[1., 0., 1., 1.]]])

Appendix A: The Problem of Uneven Sequence Lengths

The toy example above assumes a perfect world on two counts. User A and User B both produced a new token at the exact same moment, so they could share a decode step. And in the BONUS section, both happened to have exactly 3 tokens of history in their KV cache, so the two caches stacked into one rectangular tensor.

In reality, batched matrices must be perfectly rectangular. What happens if User A's prompt is 100 tokens long, but User B's prompt is only 5 tokens long?

The Padding Problem
Historically, the solution was padding. You would add 95 empty "pad" tokens to User B's cache just to make it the same shape as User A's cache so they could fit into a single (2, 100, d) 3D tensor, and specifically left-padding, with the pads in front, since a decoder generates from whatever sits in the last position and right-padding would leave the model staring at a pad token instead of User B's actual final token.

This is disastrous for efficiency. The GPU wastes precious FLOPs multiplying queries against padding tokens (which are then masked out and ignored anyway), and wastes precious HBM storing padding tokens in the KV cache. It is also a persistent source of bugs because the attention mask and the position IDs both have to be shifted to match the padding, and getting either wrong produces output that looks plausible but is subtly incorrect.

Three major innovations solved this:

Wasted FLOPs in the kernel — ragged / varlen attention (FlashAttention, xFormers): these solve padding inside a single forward pass. Instead of demanding a rectangular (Batch, Seq_Len, Dim) tensor, they take all tokens flattened into one long sequence plus an index of where each sequence starts, and never touch a padding token. (Concretely: the tokens are flattened into one long buffer, and the kernel is handed a list of cumulative offsets — [0, 100, 105] for a 100-token and a 5-token sequence, telling it where each sequence begins and ends.)
Wasted HBM in the cache — PagedAttention (vLLM): Padding's storage-side cousin: rather than reserving one contiguous max-length block of HBM per user, the cache is allocated in small fixed-size pages on demand, like virtual memory in an operating system. A 5-token user occupies 5 tokens' worth of pages, not 100, and the gaps between users' blocks stop fragmenting the rest. This is what makes the KV cache limit from the trade-off section a limit on tokens actually in flight rather than on worst-case reservations.
Wasted batch slots over time — continuous batching (iteration-level scheduling): Prompts aren't the only thing that varies in length; generations do too. In a static batch, when User B finishes after 10 tokens their seat sits empty (or worse, keeps generating pad tokens) while User A grinds through 500 more. Continuous batching ejects finished requests and admits new ones between token steps, so the batch stays full.

Appendix B: Two Tokens Per User — Is Masking Needed?

Everything so far has been one decode step: each user contributes exactly one new token. But there are common situations where a user contributes several tokens to the same forward pass — prefill (a prompt is many tokens), chunked prefill, and speculative decoding (verifying $k$ draft tokens at once). The smallest version of that situation: two users, two new tokens each.

Does batching still work? And now that each user's tokens sit next to each other in the same matrix, do we need a mask to keep them from interfering?

The projections need no mask at all. Attention needs one, but only the causal mask inside each user's own sequence. Keeping User A away from User B costs nothing extra as long as we keep the batch dimension.

Setup

Keep $W_{q}$ from before. User A's first new token is the same $[1, 0]$ as before and User B's is the same $[0, 1]$ ; each now has a second token behind it:

User A: $x_{A, 1} = [1, 0]$ , $x_{A, 2} = [1, 1]$
User B: $x_{B, 1} = [0, 1]$ , $x_{B, 2} = [2, 0]$

So each user contributes two new tokens, and the batch as a whole contains four. Both counts show up below: "two" whenever we're looking at one user (the 2 rows of a user's score block, the 2 new keys appended to a cache), "four" whenever we're looking at the whole batch (the 4 rows of the projection in Step 1, the 4 query rows of the grid in Step 2).

Step 1: The projection — no mask

Stack all four rows into one matrix and multiply, exactly as in Scenario 2:

Q = ⎝ ⎛ 11020110 ⎠ ⎞ \times (1011) = ⎝ ⎛ 11021212 ⎠ ⎞

Row 1 is $q_{A, 1} = [1, 1]$ — the same $q_{A}$ we have now computed three times. Row 2 is $q_{A, 2}$ , row 3 is $q_{B, 1}$ , row 4 is $q_{B, 2}$ .

A projection is a per-row operation. Row 2 is computed from $x_{A, 2}$ and $W_{q}$ and nothing else; it cannot see rows 1, 3, or 4. The matrix has no idea whether it holds four users with one token each, one user with four tokens, or two users with two. So the memory-bandwidth argument from Scenario 2 carries over unchanged — $W_{q}$ is loaded from HBM once and amortized over four rows instead of two — and there is nothing to mask.

The same is true of the K and V projections and every feed-forward layer. That's most of the model's weights, and none of them care who the rows belong to.

Notice that the history tokens don't appear in this step at all, and nothing is missing. Their projections were computed in the steps that produced them: their queries were used once and thrown away, and their keys and values went into the cache. This step projects only the four embeddings that are new, and the same multiplication with $W_{k}$ and $W_{v}$ yields the four new keys and values that get appended to the caches in Step 3. That is the division of labour in a decode step: the weight matrices only ever see new tokens, and the history reaches the computation exclusively through the cache at attention. It's why the projection cost of a step grows with the number of new tokens, not with the length of the conversation so far.

Step 2: Attention — where users and positions start to matter

Attention is the one place a token interacts with other tokens, so it's the one place where "which user?" and "which position?" matter. There are two things to prevent:

Cross-user leakage: $q_{A, 2}$ must not attend to any of User B's keys.
Future leakage: $q_{A, 1}$ must not attend to $k_{A, 2}$ , which comes after it in User A's sequence. A decoder only looks backward.

If we flattened the four new tokens into one long sequence and naively let every query score every key, we'd get a grid of scores, one per (query, key) pair: 4 rows for the new queries, and 10 columns for the keys, because each user's cache holds 3 history keys plus the 2 new ones. Not all of those pairs are legitimate. Each cell asks: may this query (row) look at this key (column)? Two rules decide. The key must belong to the same user, and it must not come later than the query in that user's sequence (a token may see itself). Apply both and the allowed pairs look like this:

              User A's keys           User B's keys
             hist      new           hist      new
            k0 k1 k2   k3 k4        k0 k1 k2   k3 k4
A1 (q3)   [  1  1  1    1  0    |    0  0  0    0  0  ]
A2 (q4)   [  1  1  1    1  1    |    0  0  0    0  0  ]
          ---------------------+----------------------
B1 (q3)   [  0  0  0    0  0    |    1  1  1    1  0  ]
B2 (q4)   [  0  0  0    0  0    |    1  1  1    1  1  ]

Read row 1: $q_{A, 1}$ , sitting at position 3, may see User A's three history keys (all in its past) and its own key, but not $k_{A, 2}$ at position 4, and nothing of User B's. Row 4: $q_{B, 2}$ may see all five of User B's keys and none of User A's.

The two rules leave different fingerprints. "Same user" zeroes the two off-diagonal blocks. "Not in the future" zeroes the upper triangle of the new-token columns inside each diagonal block; the history columns are untouched by it, because history is always in the past. The two rules also have different fixes.

It is not efficient to compute the entire 4×10 grid and masks it because the two off-diagonal 2×5 blocks are wasted work. There are two ways to avoid computing them, and they differ only in who enforces the block structure. Either way, what gets computed is exactly the two diagonal 2×5 blocks, and the top-left one is the score matrix we'll work through in Step 3.

The first is to keep the batch dimension. Instead of a (4, d) matrix, hold the new tokens as a 3D tensor of shape (Batch, Seq_New, Dim) = (2, 2, 2), and the KV caches as (2, 5, 2), since each user now has 3 history keys plus 2 new ones. torch.bmm multiplies User A's queries against User A's keys and User B's against User B's; the cross-user blocks are unreachable by construction. The cost is that everything must be rectangular, which is the padding problem of Appendix A.

Concretely: a tensor is a box, and every slice along a dimension has to be the same length. The middle dimension of (2, 5, 2) is one number, so it can only be 5 if both users have exactly 5 keys. Take one history token away from User B, so they have 2 + 2 = 4 keys while User A still has 3 + 2 = 5, and the stack refuses:

K_A = torch.tensor([[1., 0.], [1., 1.], [0., 1.], [1., 0.], [1., 1.]])  # 3 history + 2 new = 5 keys
K_B = torch.tensor([[0., 1.], [2., 0.], [0., 1.], [2., 0.]])            # 2 history + 2 new = 4 keys

torch.stack([K_A, K_B])
# RuntimeError: stack expects each tensor to be equal size,
#               but got [5, 2] at entry 0 and [4, 2] at entry 1

The only way into a (2, 5, 2) tensor is to give User B a fifth key that doesn't exist: a row of zeros, placed at the front (left-padding, for the reason given in Appendix A).

K_B_padded = torch.cat([torch.zeros(1, 2), K_B], dim=0)   # now 5 x 2
K_cache = torch.stack([K_A, K_B_padded])                  # (2, 5, 2) -- works
scores = torch.bmm(Q_batch, K_cache.transpose(1, 2))      # Q_batch: the (2, 2, 2) queries from Step 1

bmm runs now, but User B's scores have a column for a key that isn't real:

User B raw scores:
[[0., 1., 0., 1., 0.],
 [0., 2., 4., 2., 4.]]
   ^ pad column

That 0 is not "ignored". It is a score like any other, and softmax would hand the pad key a weight of $e^{0}$ divided by the row sum. So the pad forces a second kind of masking on top of the causal one, and it changes the mask's shape: User A's rows are causal-only, User B's are pad-plus-causal, so the mask becomes a per-user (2, 2, 5) rather than a single (2, 5) that broadcasts across the batch. The pad column also costs FLOPs to compute and HBM to store, for a key that is thrown away the moment it is scored. Multiply by real prompt lengths, 100 tokens versus 5, and you are looking at Appendix A's problem.

The second method is to flatten to a ragged layout and hand the kernel two offset lists: where each user's queries start in the flat query buffer, cu_seqlens_q = [0, 2, 4], and where each user's keys start in the flat KV buffer, cu_seqlens_k = [0, 5, 10]. Those are precisely the row groups and column groups of the grid above.

Either way, user isolation falls out of the layout. What does not fall out of the layout is the causal triangle inside each block: bmm (or the kernel) happily computes q_{A,1} · k_{A,2}.

The causal mask has to be explicit. Inside User A's block, bmm happily computes $q_{A, 1} \cdot k_{A, 2}$ , and nothing in the tensor layout says that pair is illegal. We have to say so.

Step 3: The numbers

Take $W_{k} = I$ so a new token's key is just its embedding (any $W_{k}$ works; this keeps the arithmetic small). User A's keys are the three history keys from the BONUS section plus the two new ones:

K_{A} = ⎝ ⎛ 1101101101 ⎠ ⎞, Q_{A} = (1112)

Scores for User A — one row per new query, one column per key:

Q_{A} K_{A T} = (1123121123)

The first three entries of row 1, $[1, 2, 1]$ , are exactly the scores from the BONUS section: same query, same history. The two new columns are the new keys. The top-right entry, $q_{A, 1} \cdot k_{A, 2} = 2$ , is the illegal one: token 1 peeking at token 2. Apply the causal mask:

(11231211 - \infty 3)

After softmax, $e^{- \infty} = 0$ , so the weight on that key is exactly zero. Every history column stays visible to both rows (history is always in the past), and the diagonal of the new block stays visible (a token may attend to itself). Because both users have the same amount of history, the mask is a single (2, 5) boolean that broadcasts across the batch dimension.

Code

import torch

# 1. SETUP
W_q = torch.tensor([[1.0, 1.0], [0.0, 1.0]])   # same W_q as before
W_k = torch.eye(2)                              # identity keeps the arithmetic small

# Each user now has TWO new tokens. Shape: (Batch, Seq_New, Dim) = (2, 2, 2)
X_batch = torch.tensor([
    [[1., 0.], [1., 1.]],   # User A: token 1, token 2
    [[0., 1.], [2., 0.]],   # User B: token 1, token 2
])

# 2. PROJECTIONS: no mask. A projection is a per-row operation.
Q_batch = X_batch @ W_q     # (2, 2, 2) — matmul broadcasts over the batch dim
K_new   = X_batch @ W_k     # (2, 2, 2)

# Flattening all four tokens into one (4, 2) matrix gives identical rows.
assert torch.allclose(Q_batch, (X_batch.reshape(4, 2) @ W_q).reshape(2, 2, 2))
print("Q_batch (Batch, Seq_New, Dim):")
print(Q_batch, end="\n\n")

# 3. KV CACHE: the same 3 tokens of history as before, extended by the 2 new keys.
K_hist = torch.tensor([
    [[1., 0.], [1., 1.], [0., 1.]],  # User A's history
    [[0., 1.], [2., 0.], [1., 1.]],  # User B's history
])
K_cache = torch.cat([K_hist, K_new], dim=1)   # (2, 5, 2)
print(f"K_cache shape: {K_cache.shape} -> (Batch, Hist + Seq_New, Dim)\n")

# 4. SCORES: bmm keeps users apart. User A's queries never meet User B's keys.
scores = torch.bmm(Q_batch, K_cache.transpose(1, 2))   # (2, 2, 5)
print("Raw scores (Batch, Seq_New, Hist + Seq_New):")
print(scores, end="\n\n")

# 5. CAUSAL MASK: the one place masking is needed.
# New token i may see all history and new tokens 0..i, but nothing after it.
hist_len, new_len = K_hist.shape[1], X_batch.shape[1]
mask = torch.ones(new_len, hist_len + new_len, dtype=torch.bool)
mask[:, hist_len:] = torch.tril(torch.ones(new_len, new_len, dtype=torch.bool))
print("Causal mask (True = may attend):")
print(mask, end="\n\n")

masked_scores = scores.masked_fill(~mask, float("-inf"))   # broadcasts over batch
print("Masked scores:")
print(masked_scores, end="\n\n")

weights = torch.softmax(masked_scores, dim=-1)
print("Attention weights (top-right of each user is exactly 0):")
print(weights)

Output:

Q_batch (Batch, Seq_New, Dim):
tensor([[[1., 1.],
         [1., 2.]],

        [[0., 1.],
         [2., 2.]]])

K_cache shape: torch.Size([2, 5, 2]) -> (Batch, Hist + Seq_New, Dim)

Raw scores (Batch, Seq_New, Hist + Seq_New):
tensor([[[1., 2., 1., 1., 2.],
         [1., 3., 2., 1., 3.]],

        [[1., 0., 1., 1., 0.],
         [2., 4., 4., 2., 4.]]])

Causal mask (True = may attend):
tensor([[ True,  True,  True,  True, False],
        [ True,  True,  True,  True,  True]])

Masked scores:
tensor([[[1., 2., 1., 1., -inf],
         [1., 3., 2., 1., 3.]],

        [[1., 0., 1., 1., -inf],
         [2., 4., 4., 2., 4.]]])

Attention weights (top-right of each user is exactly 0):
tensor([[[0.1749, 0.4754, 0.1749, 0.1749, 0.0000],
         [0.0513, 0.3790, 0.1394, 0.0513, 0.3790]],

        [[0.2969, 0.1092, 0.2969, 0.2969, 0.0000],
         [0.0414, 0.3057, 0.3057, 0.0414, 0.3057]]])

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}$ .

DEV Community: Lewis Won

How to profile and optimise kernels

Where the last article left off

The setup

The GPU

Software and files

The baseline, re-measured

How I profile

A small program to profile

The commands

What the compiler did

Round 0: profile the baseline

Step 1: the timeline

Step 2: the Speed of Light section

Step 3: which instructions keep the LSU busy

Step 4: why the warps wait

Step 5: the roofline

Step 6: ask the compiler why

Five facts to start from

What the compiler chose

How the 4 warps divide each tile

Problem 1: the first multiply is done twice

Problem 2: every row maximum and row sum goes through shared memory

Problem 3: P also goes through shared memory

STS: Store to Shared memory

BAR.SYNC: Barrier (synchronise)

LDSM: Load from Shared memory, as a Matrix

How they fit together: P's round trip in the baseline

Counting the barriers

What a bigger tile changes

The diagnosis

What resident warps are

Why more warps usually help

Why it would not help here

Why Nsight Compute suggested it anyway

Round 1: tile sizes, num_warps and num_stages

The shared-memory budget for an Ada SM

The two numbers

Using both limits

How to read the notation

Line 1: the Q tile

Line 2: the K buffers

Line 3: the V buffers

Significance of the three lines of code

Programs / SM (shared)

Registers / thread

Programs / SM (registers)

Putting them together

Exception 1: the warps sit side by side, so P needs its own buffer

Exception 2: one stage, but V still uses shared memory

The sweep

The heatmaps

The winner, through the unchanged harness and tests

Round 2: profile the tuned kernel

Step 1: Speed of Light, with the baseline

Step 2: the breakdown, and a trap in it

Step 3: the roofline, with the baseline

Step 4: why the warps wait now

Step 5: checking the claim about too many warps

Step 6: the causal run

Round 3: wrap the winner in @triton.autotune

Round 4: inside the inner loop

Causal tile-skipping

exp2

Round 5: re-tune, then the final profile

The final profile

What to look for, as a checklist

FlashAttention 2 from PyTorch to Triton

What I am building

Step 1: the tiled forward in PyTorch

A primer to tl.make_block_ptr

Why use a block pointer

The classical way with pointer maths

What each line does

Setup

1. Which rows this program owns

2. Which columns this step covers

3. A 2D grid of addresses

4. Loading safely at the edges

After the loop

Problem 3: `P` also goes through shared memory

Round 1: tile sizes, `num_warps` and `num_stages`

Exception 1: the warps sit side by side, so `P` needs its own buffer

Round 3: wrap the winner in `@triton.autotune`

`exp2`

A primer to `tl.make_block_ptr`

`base` as the starting point in memory of a tensor

`shape` tells Triton how big the tensor is

`strides` tells Triton how far apart neighbours are in memory

`offsets` tells Triton where the tile starts

`block_shape` tells Triton the size of a tile

`order` tells Triton which axis is side by side in memory

`.advance()` tells Triton how to move the tile

`boundary_check` tells Triton what happens past the edge

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

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