I rebuilt FlashAttention in Triton to understand the performance archaeology

⏲️ Estimated reading time ~45min.

Flash Attention has become one of the most impactful optimizations in modern deep learning. Since the original paper was published in 2022, we’ve seen four major versions—each squeezing more performance out of increasingly powerful hardware. But here’s the thing: reading papers is one thing, understanding why these optimizations were made is another entirely.

My goal here is simple: start from first principles, implement FlashAttention v1 exactly as described in the paper, profile it, find the bottlenecks, and see how far we can push it. We’ll build intuition by iterating on the algorithm, discovering through profiling exactly why v2, v3, and v4 were necessary. Think of this as archaeology—digging through the performance layers to understand what each version was really solving.

So let’s put ourselves in the shoes of that mythical Stanford grad student. You’ve finally finished configuring your neovim and archlinux setup (a multi-year endeavor, naturally). You open up a fresh LLaMA model to peek under the hood. Text goes in, gets tokenized, embedded, then flows through a stack of transformer blocks. Standard stuff. But then you look closer at the attention mechanism—three projections and then… there it is:

scores = torch.matmul(q, k.transpose(3, 2)) / math.sqrt(self.head_dim)
scores = F.softmax(scores, dim=-1)
output = torch.matmul(scores, v)  # (B, N_h, T, D_h)

This monstrosity of a code is staring you right in the face. For those who don’t immediately see the problem with these 4 lines, let me add some annotations on how this would normally execute in PyTorch (without compilation).

# 1. We load q `(B, N_h, S, D_h)` and k `(B, N_h, S, D_h)`
# 2. We compute Q.Kt and write it back to HBM. Note score is `(B, N_h, S, S)`.
scores = torch.matmul(q, k.transpose(3, 2)) / math.sqrt(self.head_dim)
# 3. Reload the scores tensor to compute the softmax and write it back to HBM
scores = F.softmax(scores, dim=-1)
# 4. Load v `(B, N_h, S, D_h)`, load the scores from HBM
# 5. Compute scores@v and write it back to HBM.
output = torch.matmul(scores, v)  # `(B, N_h, T, D_h)`

Do you see it now? Well, we have three tensors q, k, v each of dimension (B, N_h, T, D_h)¹. The output tensor of the attention mechanism is (B, N_h, T, D_h), and somehow in the middle we had to materialize a (B, N_h, S, S) tensor for funsies. The attention mechanism has a critical bottleneck: quadratic memory complexity. Let’s take a standard training sequence length S=8192. Computing attention naively requires O(S²) memory to store the full attention matrix, which means consuming several gigabytes of GPU memory. Crucially, in modern transformers, S >> D_h (sequence length is much larger than head dimension) - we typically have S=8192 or more while D_h=64 or 128. This massive asymmetry is what makes Flash Attention algorithm possible. The second big issue here is the back and forth to HBM. Modern GPUs have compute throughput vastly exceeding memory bandwidth. Repeatedly reading Q, K, V, and scores from slow High Bandwidth Memory (HBM) is going to greatly impact performance (more on this later).

The whole idea of Flash Attention is to bypass these intermediate steps—i.e., go from tensors q, k, v to the output tensor directly and compute the attention in one go, with minimal memory footprint (materializing only the tensors we need) and minimal back and forth to HBM, and hopefully getting to O(S) memory complexity.

The Plan: Start with FlashAttention v1 from the original paper. Implement it faithfully in Triton, profile with NVIDIA’s tooling, identify the bottlenecks, then iterate. Each optimization will teach us something about the GPU memory hierarchy and why the subsequent versions (v2, v3, v4) introduced the changes they did. Four versions represents an enormous amount of engineering work—let’s see how much of that journey we can reconstruct by just following the profiler’s breadcrumbs.

Setup and Hardware

First things first, I’ll use triton to implement the attention kernels. I wanted to implement them in CUDA but I thought this was a good opportunity to learn triton. I’m always skeptical about DSLs and abstractions that promise “write once, run fast everywhere,” but like my very wise friend Chris Fleetwood said, you can’t go wrong learning something shipped inside PyTorch. Running triton is also extremely straightforward and reduces the mess of boilerplate and C++ code that you have to write when implementing CUDA kernels, especially if you want to call them from Python. I actually went back and reimplemented Flash Attention in CUDA to have a bit more control, but that will come later (a little teasing).

I’ll compare these kernels directly with a reference PyTorch implementation, first to make sure the kernel is doing what it’s supposed to, and second to profile the kernel against the baseline.

Why Triton? Block-level Programming Without Thread Hell

Triton is a Python-based DSL for writing GPU kernels. The pitch is simple: write Python-like code at the block level, and the compiler handles individual thread management, memory coalescing, and other low-level optimizations. What makes Triton fundamentally different from CUDA is the abstraction level - in CUDA, you explicitly program individual threads of a block. In Triton, you write your kernel thinking about blocks/tiles of data a little bit how you’d write torch code, and the compiler figures out how to map that to threads. It handles threads data access pattern, synchronization barriers and work splitting. My skepticism comes from the fact that I don’t really trust these “magic” compilers for running parallel code. Usually, the promised performance falls apart quickly. But Triton has a few things going for it:

1. It’s maintained by OpenAI and ships with PyTorch 2.0+
2. The generated PTX is easily inspectable, so you can see what it’s actually doing
3. It handles the tedious bits (pointer arithmetic, bounds checking) while still giving you control over the algorithm. You write your kernel at block level and the compilers issues the correct code to load your tensors in.

Setup

I’m running this on my personal machine:

• OS: Arch Linux (btw)
• GPU: NVIDIA GeForce RTX 2070 (8 GB VRAM)
• Driver: 580.95.05
• CUDA: 13.0
• Triton: 3.5

Here are the GPU specs from cudaGetDeviceProperties. These numbers will be useful once we start profiling our kernel to detect performance issues:

Profiling Tools

For profiling, I use three tools at different granularities:

1. torch.profiler: Quick and dirty. Good for seeing wall-clock time and basic GPU utilization. I use this for initial sanity checks.

2. NVIDIA Nsight Systems (nsys): System-wide profiler. Shows CPU/GPU timeline, kernel launches, memory transfers. Great for spotting gaps where the GPU is idle.

3. NVIDIA Nsight Compute (ncu): The heavy hitter. Gives you everything: occupancy, memory throughput, warp stalls, instruction mix, bank conflicts. This is what I’ll use for deep-diving into kernel performance. You can run it with:

sudo ncu --set full --kernel-name "attn_kernel" -o profile_output -f python script.py

Then open the .ncu-rep file with ncu-ui for the full analysis.

The Flash Attention Algorithm

For modern GPUs, compute throughput vastly exceeds memory bandwidth - an A100 can do ~300 TFLOPs but only has ~2 TB/s of memory bandwidth. This massive compute-to-memory ratio means that naive algorithms spending most of their time waiting for memory, not computing. Flash Attention restructures the computation to maximize arithmetic intensity (FLOPs per byte transferred).

Let’s build the intuition behind the core ideas from walking back from the solution. The goal is to one shot (single kernel) build the output tensor.

We are in GPU programming land. If you want to minimize transfers to main memory—like any good GPU engineer who has read Simon’s matmul blog post (stop reading this post and go read it. I am serious. It’s the best matmul writeup on the internet.)—you’d immediately think about computing blocks of output in parallel and writing each block to memory.

This is commonly known as tiling, and every GPU kernel out there uses this trick. The idea is simple but powerful: instead of streaming data directly from slow global memory for each operation, you load a tile (a small block of data) into fast shared memory once, reuse it across multiple computations within a thread block, and only write the final results back. This dramatically reduces memory bandwidth bottlenecks by exploiting data locality and the GPU’s memory hierarchy—exactly what you need when global memory access is orders of magnitude slower than compute.

Let’s first look at how the output tensor is computed based on the attention formula. Let’s focus on a single batch and head (B, N_h) to clearly see the shapes of the matmuls and the elements. Also, we’ll omit the causal masking here (although it’s not really complicated to add it to the attention score before softmax).

If you have been following along, you’re probably scratching your head wondering why we are sharding the V tensor row-wise. Sorry for pulling a quick one here, but this formulation of the output (using the rows of V instead of the columns) will actually help us build the flash attention algorithm. For each row $O_i$Oi​ of the output:

$$O_i = \frac{\sum_j e^{S_{ij}} V_j}{\sum_j e^{S_{ij}}} \quad\text{where}\quad S_{ij} = Q_i \cdot K_j^{\top}$$Oi​=∑j​eSij​∑j​eSij​Vj​​whereSij​=Qi​⋅Kj⊤​

Still, why split $V_j$Vj​ row-wise and not column-wise? First of all, sharding on the D dim makes less. Like I pointed out earlier, D is usually 64 or 128 whereas S sequence length could be 8192 or more depending on the context window. If we tried to split $V$V column-wise (along D), the math would look like this:

$$\begin{bmatrix} \dots & P_{ij} & \dots \end{bmatrix} \times \begin{bmatrix} V_{\text{left}} & V_{\text{right}} \end{bmatrix} = \begin{bmatrix} P \cdot V_{\text{left}} & P \cdot V_{\text{right}} \end{bmatrix}$$[…​Pij​​…​]×[Vleft​​Vright​​]=[P⋅Vleft​​P⋅Vright​​]

This would force us to load the entire score matrix $P=softmax(Q.K^t)$P=softmax(Q.Kt) just to compute the left half of the output. But we can’t store the entire $P$P matrix—that’s the exact problem FlashAttention solves! By splitting row-wise:

$$\left[ \begin{array}{c|c} P_{i0} & P_{i1} \end{array} \right] \times \begin{bmatrix} V_0 \\ \hline V_1 \end{bmatrix} = (P_{i0} \cdot V_0) + (P_{i1} \cdot V_1)$$[Pi0​​Pi1​​]×[V0​V1​​​]=(Pi0​⋅V0​)+(Pi1​⋅V1​)

Now, the genius piece of flash attention is this: if we found a way to build $O_i$Oi​ iteratively for a given $i$i, we could calculate a small chunk $P_{ij}$Pij​, multiply by $V_j$Vj​, add it to the sum, and discard $P_{ij}$Pij​. This means we are only allocating a small block of memory!

But why does this matter anyway? Isn’t it still mathematically equivalent if we compute the whole P_i row? Well, if you live in some mythical mathematical idea world, yes. But this code runs on hardware, and when and how we access memory does matter a LOT!

Quick Detour into GPU Memory Land

Understanding GPU memory is crucial, especially because GPUs are SIMT (Single Instruction, Multiple Thread) machines. If you’re familiar with SIMD programming on the CPU side, you usually need to use explicit vector instructions to pack and align vectors before executing instructions. SIMT, on the other hand, lets you write scalar thread code that the GPU transparently executes as a lockstep vector, coalescing loads automatically when memory is well-aligned.

GPUs execute thousands of lightweight threads in lockstep (called warps). Multiple warps form a block, and blocks are scheduled into SMs. Looking back at the Setup section, we can see that my RTX 2070 has 36 SMs (streaming multiprocessors), and each SM can handle 1024 threads (32 warps) concurrently.

GPU cores are simple and rely on this massive parallelism to hide latency—whenever one warp stalls waiting on memory, the scheduler swaps in another warp that already has its data “in flight.” This only works if kernels expose enough parallelism and if memory accesses are structured to take advantage of locality². A core part of the CUDA programming model is that it exposes the memory hierarchy directly, forcing the programmer (or compiler) to reason explicitly about where data lives and how far it is from the compute units. Memory proximity matters enormously:

• DRAM/HBM (High Bandwidth Memory): I don’t have HBM on my RTX 2070, but even on high-end cards like the A100 with large capacity (~80GB), this memory is relatively far from the SMs. Despite impressive raw bandwidth (~1–2 TB/s), latency is high, so naive repeated reads kill performance.

• L2 Cache: A chip-wide cache that helps buffer global memory traffic. Still far slower than on-SM memories. L2 cache behavior is typically leveraged implicitly through access patterns rather than explicit optimization.

• L1 / Shared Memory (SRAM): On-SM, extremely fast, and explicitly managed in CUDA. This is where SRAM truly shines compared to HBM. First, the physics: SRAM sits directly on the SM, mere micrometers from the compute units (vs millimeters for HBM - that’s 1000x closer!). It uses 6-transistor flip-flop circuits that hold state without refresh cycles, unlike HBM’s capacitor-based cells. This means SRAM access takes only ~20-30 cycles vs ~200-600 for HBM.

  The bandwidth story is even more compelling: each SM gets its own ~164KB of SRAM (48 Kb for my poor RTX 2070) with ~1-2 TB/s bandwidth. With 108 SMs on an A100, that’s theoretically ~100+ TB/s aggregate SRAM bandwidth across the chip, compared to “only” ~2 TB/s to HBM. But here’s the real beauty - SRAM is explicitly programmer-controlled! Unlike HBM which goes through L2 cache, L1 cache, and complex replacement policies you can’t control, with SRAM you orchestrate the exact choreography of data movement with computation. No cache thrashing, no surprise evictions, just deterministic high-speed access exactly when you need it.

• Registers: The closest memory to the compute units. Tiny (~256KB per SM) but insanely fast (100+ TB/s effective). The compiler allocates these, and register pressure directly impacts occupancy (how many warps can run concurrently, which in turn affects latency hiding).

Look at this diagram of memory hierarchy with bandwidth and memory size taken from Flash Attention v1 very closely, and try to internalize it. It is extremely important to take I/O into account when designing highly performant algorithms.

If you’re still following along with me, the main insight is to restructure attention so that all intermediate activations fit within these fast memories—registers and shared memory—minimizing slow HBM traffic and maximizing warp-level parallelism. This is why building the output block by block using the same size $P_{ij}$Pij​ is a core pillar of fast attention. But it is still unclear at this stage how we could build the output incrementally.

Online Softmax Mathematics

Before we go deeper into how we could build these blocks, we need one last quick detour into something I omitted until now for simplification reasons. The last lego block we need in our toolbox before building flash attention is softmax numerical stability.

Numerically stable softmax

Until now I wrote the softmax function for a vector $z = (z_1, \dots, z_n)$z=(z1​,…,zn​) as:

$$\operatorname{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{n} e^{z_j}}$$softmax(zi​)=∑j=1n​ezj​ezi​​

But what happens if input values $z$z are large (e.g., $z_i = 1000$zi​=1000)? Computing $e^{1000}$e1000 will exceed the maximum value a floating-point number can hold, resulting in inf (infinity) or NaN (Not a Number) errors. Conversely, underflow can happen if values are very negative—the denominator might vanish to zero, leading to division by zero errors. The trick is to numerically stabilize by subtracting the maximum logit.

The key property is that Softmax is invariant to adding or subtracting a constant ( c ):

$$\operatorname{softmax}(z_i + c) = \operatorname{softmax}(z_i)$$softmax(zi​+c)=softmax(zi​)

For numerical stability, we usually subtract the maximum value ($m = \max_j z_j$m=maxj​zj​). Setting ( c = -m ) ensures that the exponentials are at most 1, preventing overflow :

$$\operatorname{softmax}(z_i) = \frac{e^{z_i - m}}{\sum_{j=1}^{n} e^{z_j - m}}$$softmax(zi​)=∑j=1n​ezj​−mezi​−m​

So let’s go back to incrementally computing a row of the output $O_i$Oi​. Let’s start slow to build the intuition. We take row i=0 and j=0. This means we are focusing on the first block of $P_{00}=Q_0@K_0^t$P00​=Q0​@K0t​ and the first row of the value tensor $V_0$V0​.

If you already see an issue with the current softmax, bear with me a minute. Remember that each element of the score matrix is computed as $P_{ij} = \frac{e^{Q_i \cdot K_j^T - m_i}}{\sum_j e^{Q_i \cdot K_j^T- m_i}}$Pij​=∑j​eQi​⋅KjT​−mi​eQi​⋅KjT​−mi​​. So the denominator is clearly broken (we need to compute the whole row to compute this sum).

For now we only have $Q_0\cdot K_0^t$Q0​⋅K0t​ in the row, so the first value we computed is clearly broken. Similarly, $m_i=m_0$mi​=m0​ here represents the maximum value of the row. For now we’ve only computed one value, so the max is also broken. Let’s continue to the second value of the output row and see how we can fix what is broken:

When we see a new block, we can’t just throw away our previous work. Instead, we need to rescale what we’ve already computed. How though? If we simplify the problem to only having two values in this row (D=2) and we have computed the first one, the correct output value should be :

$$O_1 = \frac{e^{Q_1 \cdot K_1^T - m_1} \cdot V_1 + e^{Q_1 \cdot K_2^T - m_1} \cdot V_2}{e^{Q_1 \cdot K_1^T - m_1} + e^{Q_1 \cdot K_2^T - m_1}} \quad \text{where} \quad m_1 = \max(Q_1 \cdot K_1^T, Q_1 \cdot K_2^T)$$O1​=eQ1​⋅K1T​−m1​+eQ1​⋅K2T​−m1​eQ1​⋅K1T​−m1​⋅V1​+eQ1​⋅K2T​−m1​⋅V2​​wherem1​=max(Q1​⋅K1T​,Q1​⋅K2T​)

Thanks to the exponential multiplicative property we can easily build these from the previous work and currently computed block. This is done using running correction factors:

$$\begin{gathered} m_1 = Q_1 \cdot K_1^T, \quad m_2 = Q_1 \cdot K_2^T, \quad m_{\text{new}} = \max(m_1, m_2) \quad \alpha = e^{m_1 - m_{\text{new}}}, \quad \beta = e^{m_2 - m_{\text{new}}} \end{gathered}$$m1​=Q1​⋅K1T​,m2​=Q1​⋅K2T​,mnew​=max(m1​,m2​)α=em1​−mnew​,β=em2​−mnew​​

These $\alpha$α and $\beta$β factors let us rescale our previous exponentials to the new maximum. Think of it as “adjusting the baseline” - when we find a new maximum or a new sum, we scale down the previous values accordingly.

$$l_{\text{new}} = \alpha l_1 + \beta l_2 = e^{m_1 - m_{\text{new}}} l_1 + e^{m_2 - m_{\text{new}}} l_2$$lnew​=αl1​+βl2​=em1​−mnew​l1​+em2​−mnew​l2​

$$\boxed{l_{\text{new}} = e^{Q_1 \cdot K_1^T - m_{\text{new}}} + e^{Q_1 \cdot K_2^T - m_{\text{new}}}}$$lnew​=eQ1​⋅K1T​−mnew​+eQ1​⋅K2T​−mnew​​

Now here’s where the magic happens. We also need to update our output $O_1$O1​. Remember, we had computed $O_1^{\text{prev}}$O1prev​ using only the first block, but now we need to incorporate the second block. The update formula rescales the old output and adds the contribution from the new block:

$$O_1^{\text{new}} = \frac{\alpha \cdot l_1 \times O_1^{\text{prev}} + \beta \cdot e^{Q_1 \cdot K_2^T - m_2} \cdot V_2}{l_{\text{new}}}$$O1new​=lnew​α⋅l1​×O1prev​+β⋅eQ1​⋅K2T​−m2​⋅V2​​

Let’s expand this to see what’s really happening. The old output was $O_1^{\text{prev}} = \frac{e^{Q_1 K_1^T - m_1}}{l_1} \cdot V_1$O1prev​=l1​eQ1​K1T​−m1​​⋅V1​, so:

$$= \frac{1}{l_{\text{new}}} \left( e^{m_1 - m_{\text{new}}} \cdot \cancel{l_1} \cdot \frac{e^{Q_1 K_1^T - m_1}}{\cancel{l_1}} \cdot V_1 + e^{m_2 - m_{\text{new}}} \cdot e^{Q_1 \cdot K_2^T - m_2} \cdot V_2 \right)$$=lnew​1​(em1​−mnew​⋅l1​​⋅l1​​eQ1​K1T​−m1​​⋅V1​+em2​−mnew​⋅eQ1​⋅K2T​−m2​⋅V2​)

Simplifying by combining the exponentials:

$$= \frac{1}{l_{\text{new}}} \left( e^{Q_1 \cdot K_1^T - m_{\text{new}}} \cdot V_1 + e^{Q_1 K_2^T - m_{\text{new}}} \cdot V_2 \right)$$=lnew​1​(eQ1​⋅K1T​−mnew​⋅V1​+eQ1​K2T​−mnew​⋅V2​)

And voilà! We get exactly what we’d expect - the proper softmax formula:

$$= \frac{\sum_{j=1}^{2} e^{Q_1 \cdot K_j^T - m_{\text{new}}} V_j}{\sum_{j=1}^{2} e^{Q_1 \cdot K_j^T - m_{\text{new}}}}$$=∑j=12​eQ1​⋅KjT​−mnew​∑j=12​eQ1​⋅KjT​−mnew​Vj​​

The beauty of this approach is that we never had to store the full attention matrix. We just kept updating our running statistics ($m$m and $l$l) and our output incrementally!

This is the core insight of Flash Attention’s online softmax - we can compute exact attention without materializing the entire attention matrix in memory.

Core Idea: Tiling and Online Softmax

TLDR, Flash Attention solves the memory problem through two key insights:

1. Block-wise Computation: Instead of computing the full attention matrix, process Q, K, V in small blocks that fit in fast on-chip SRAM. We have been working with a single row of the output matrix, but we could easily generalize to computing a chunk of n rows of the matrix. For a single chunk n of the output, we would need n rows of the query matrix Q and all j rows of K and V. More concretely, the algorithm divides the sequence into blocks of size Bc and processes them iteratively.
2. Online Softmax: Compute softmax incrementally without materializing the full attention matrix. This is usually the unintuitive part of the flash attention algorithm. Hopefully, if you have been following along, you can see the core idea behind the online softmax. The key is keeping track of running statistics (max value m and sum l) to compute the correct softmax normalization.

Let’s roll up our sleeves and implement the algorithm!

Implementing Flash Attention v1

Here is the triton implementation of the FA1 algorithm³ straight out of the paper. I tried to follow the algorithm as closely as possible to have a good baseline that we can work from:

$$\small \begin{array}{l} \hline \textbf{Algorithm 1 } \text{FlashAttention} \\ \hline \\ \textbf{Require: } \text{Matrices } Q, K, V \in \mathbb{R}^{N \times d} \text{ in HBM, on-chip SRAM of size } M. \\ \\ 1. \quad \text{Set block sizes } B_c = \lceil \frac{M}{4d} \rceil, \ B_r = \min(\lceil \frac{M}{4d} \rceil, d). \\ 2. \quad \text{Initialize } O = \mathbf{0}_{N \times d}, \ \ell = \mathbf{0}_N, \ m = (-\infty)_N. \\ 3. \quad \text{Divide } Q \text{ into } T_r \text{ blocks } Q_i, \text{ and } K, V \text{ into } T_c \text{ blocks } K_j, V_j. \\ 4. \quad \text{Divide } O, \ell, m \text{ into blocks } O_i, \ell_i, m_i. \\ 5. \quad \textbf{for } 1 \le j \le T_c \textbf{ do} \\ 6. \qquad \text{Load } K_j, V_j \text{ from HBM to SRAM.} \\ 7. \qquad \textbf{for } 1 \le i \le T_r \textbf{ do} \\ 8. \qquad \quad \text{Load } Q_i, O_i, \ell_i, m_i \text{ from HBM to SRAM.} \\ 9. \qquad \quad \text{Compute } S_{ij} = Q_i K_j^\top \in \mathbb{R}^{B_r \times B_c}. \\ 10.\qquad \quad \tilde{m}_{ij} = \mathrm{rowmax}(S_{ij}), \ \tilde{P}_{ij} = \exp(S_{ij} - \tilde{m}_{ij}), \ \tilde{\ell}_{ij} = \mathrm{rowsum}(\tilde{P}_{ij}). \\ 11.\qquad \quad m_i^{\text{new}} = \max(m_i, \tilde{m}_{ij}), \ \ell_i^{\text{new}} = e^{m_i - m_i^{\text{new}}} \ell_i + e^{\tilde{m}_{ij} - m_i^{\text{new}}} \tilde{\ell}_{ij}. \\ 12.\qquad \quad O_i \leftarrow \mathrm{diag}(\ell_i^{\text{new}})^{-1} \left( \mathrm{diag}(\ell_i) e^{m_i - m_i^{\text{new}}} O_i + e^{\tilde{m}_{ij} - m_i^{\text{new}}} \tilde{P}_{ij} V_j \right). \\ 13.\qquad \quad \ell_i \leftarrow \ell_i^{\text{new}}, \ m_i \leftarrow m_i^{\text{new}}. \\ 14.\qquad \textbf{end for} \\ 15.\quad \textbf{end for} \\ 16.\quad \textbf{return } O. \\ \hline \end{array}$$Algorithm 1 FlashAttentionRequire: Matrices Q,K,V∈RN×d in HBM, on-chip SRAM of size M.1.Set block sizes Bc​=⌈4dM​⌉, Br​=min(⌈4dM​⌉,d).2.Initialize O=0N×d​, ℓ=0N​, m=(−∞)N​.3.Divide Q into Tr​ blocks Qi​, and K,V into Tc​ blocks Kj​,Vj​.4.Divide O,ℓ,m into blocks Oi​,ℓi​,mi​.5.for 1≤j≤Tc​ do6.Load Kj​,Vj​ from HBM to SRAM.7.for 1≤i≤Tr​ do8.Load Qi​,Oi​,ℓi​,mi​ from HBM to SRAM.9.Compute Sij​=Qi​Kj⊤​∈RBr​×Bc​.10.m~ij​=rowmax(Sij​), P~ij​=exp(Sij​−m~ij​), ℓ~ij​=rowsum(P~ij​).11.minew​=max(mi​,m~ij​), ℓinew​=emi​−minew​ℓi​+em~ij​−minew​ℓ~ij​.12.Oi​←diag(ℓinew​)−1(diag(ℓi​)emi​−minew​Oi​+em~ij​−minew​P~ij​Vj​).13.ℓi​←ℓinew​, mi​←minew​.14.end for15.end for16.return O.​​

We recognize the usual suspects from the math we derived earlier. I will only focus on the forward pass without causal masking to keep focus on the core algorithm. Let’s start with the simple reference implementation in pytorch:

# q,k,v are of shape: `(B, N_h, S, D_h)`
# B: batch_size
# N_h: num heads
# S: sequence length
# D_h: head dim
def simple_attn(q, k, v):
    att = q @ k.transpose(-2, -1) * (1.0 / math.sqrt(k.size(-1)))
    att = F.softmax(att, dim=-1)
    y = att @ v
    return y

My first implementation kernels/triton_flash_att.py follows the original Flash Attention paper very closely:

@triton.jit
def attn_kernel(Q, K, V, O, S, D, Tc, Tr, Bc, Br, softmax_scale, l, m):
    # ... [SETUP offsets here]
    # Outer loop over K, V blocks
    for j in range(0, Tc):
        # Load K_j, V_j from HBM to SRAM
        kj = tl.load(k_ptr + offset_j)  # (Bc, D)
        vj = tl.load(v_ptr + offset_j)  # (Bc, D)

        # Inner loop over Q blocks
        for i in range(0, Tr):
            # Load Q_i and previous O_i, l_i, m_i
            qi = tl.load(q_ptr + offset_i)  # `(Bc, D)`
            prev_oi = tl.load(o_ptr + offset_i)
            prev_li = tl.load(l_ptr + S_i_offset)
            prev_mi = tl.load(m_ptr + S_i_offset)

            # Compute attention scores for this block
            Sij = tl.dot(qi, tl.trans(kj)) * softmax_scale  # (Bc, Bc)

            # Online softmax update
            mij = tl.max(Sij, 1)  # Row-wise max
            pij = tl.exp(Sij - mij[:, None])
            lij = tl.sum(pij, 1)  # Row-wise sum

            # Update running statistics
            mi_new = tl.maximum(prev_mi, mij)
            alpha = tl.exp(prev_mi - mi_new)
            beta = tl.exp(mij - mi_new)
            li_new = prev_li * alpha + lij * beta

            # Update output
            oi_new = (alpha[:, None] * prev_li[:, None] * prev_oi
                      + beta[:, None] * tl.dot(pij, vj)) / li_new[:, None]

            # Write back to HBM
            tl.store(o_ptr + offset_i, oi_new)
            tl.store(m_ptr + S_i_offset, mi_new)
            tl.store(l_ptr + S_i_offset, li_new)

Some notes about this implementation :

• Double loop: Outer loop over K/V blocks, inner loop over Q blocks. Notice we are reloading block qi = tl.load(q_ptr + offset_i) # (Bc, D) in the inner loop which is pretty inefficient.
• We dispatch this kernel on a grid of B x N_h blocks. This means that each block has to find the correct matrix offset in Q,K,V and compute a full of the output O for the batch and attention head. I immediately thought that we are wasting compute here as each row O can be computed separately from each other. But don’t worry I’ll come back to this in our V2.
• Another wasteful thing I threw in is allocating l and m as l = torch.zeros(B, N_h, S).cuda() ; m = torch.full((B, N_h, S), float("-inf")).cuda(). If you recall the memory access bit earlier, we are effectively loading rows using S_i_offset from HBM which isn’t great. What we would want is for each block to allocate these accumulators locally in SRAM and avoid allocating them entirely on HBM.
• Notice also that I threw in another simplification Br = Bc = 32. This means that we are splitting Q and K and V into same size chunks. This isn’t the greatest idea because we might need to have different values for optimization reason. It is just another simplification I threw in for the implementation to be straightforward.

Before looking at the profiling of this v1 naive implementation let’s first look at another requirement I brushed over quickly but that is really important: SRAM limit. Because we are living in the real world we can’t have infinite SRAM on GPUs (sigh! ..), we need to have a ballpark estimation of how much SRAM we need to allocate and if this fits into my device SRAM.

Quick back-of-the-enveloppe calculation, the kernel at minimum needs to store in SRAM simultaneously:

• Query block Q_i: Bc × D floats
• Key block K_j: Bc × D floats
• Value block V_j: Bc × D floats
• Attention scores S_ij: Bc × Bc floats
• Running maximum and running sum: 2 × Bc floats

For B=10; Bc = 32; N_h = 64; S = 64; D_h = 32, Total SRAM: 2 × Bc + 3 × Bc × D + Bc² floats = 2 × 32 + 3 × 32 × 64 + 32² = 7,232 floats = 7,232 × 4 bytes ≈ 28KB. This fits comfortably in the ~48KB of shared memory available per thread block. It does mean, though, that we are limited to ~2 blocks per SM. If you’re confused about this limit, don’t worry, I’ll explain this promptly once we get into profiling.

Profiling the v1 Implementation

As I mentioned in the Setup and hardware section, I’ll be using the Nsight Compute ncu tool from the start. It’s one of those tools that comes with a LOT of information, and it can be pretty overwhelming. I recommend watching this lecture to get a global overview of all the ncu-ui sections. I also highly recommend reading Modal’s GPU glossary performance section front to back to at least familiarize yourself with the key terms you need to understand GPU program performance, as well as gain a general understanding of how GPUs execute code. If you don’t know what a warp scheduler is, please go read that resource before continuing to the next section.

I ran the ncu profiler using this command to get the full set of timer and metrics:

#!/bin/bash

sudo CUDA_HOME=/opt/cuda /usr/local/NVIDIA-Nsight-Compute/ncu \
    --set full \
    --kernel-name "attn_kernel" \
    --import-source yes \
    -o profile_flash_attn_v1 \
    -f .venv/bin/python3 kernels/triton_flash_att.py

Once the profiling is done, we can open it using ncu-ui, here is an annotated view of the profile once we open it:

Great, let’s drill down on the key metrics. First, we can see that our kernel took 166.47ms to execute. Our kernel used a grid size of (10, 64, 1) matchs our triton grid of (B, N_h). Each block has a size of (128, 1, 1). Because we are using triton to write our kernel, there is no way to specify the block size directly (although there is a num_warps param that we could use). This means that triton chose to use 4 warps per block to run the kernel.

Next, the profiler actually picked up of a very interesting problem that we identified earlier:

The 2.00 theoretical warps per scheduler this kernel can issue according to its occupancy are below the hardware maximum of 8. This kernel’s theoretical occupancy (25.0%) is limited by the required amount of shared memory.

To understand this occupancy problem a little bit more, ncu provides a quite handy occupancy calculator tool. Let’s look at what the tool says :

Because our shared memory requirement per block is quite high, we can only have 2 active blocks at a time per SM! Low occupancy really hurts performance when there aren’t enough warps to hide the latency. But once occupancy is sufficient for latency hiding, increasing it further can degrade performance. Higher occupancy reduces resources per thread, potentially bottlenecking the kernel on registers or reducing the arithmetic intensity.

To fix this we need to change our SRAM requirement: SRAM only depends on Bc and D_h. So we can either have smaller chunks (although this could impact memory bandwidth), or have more attention heads to reduce each D (if we control model architecture). Either way, we can go back to tweaking this later. Let’s look at the detail section of the profile report to see what other surprises await us!

The whole idea behind Flash Attention is to avoid going to main memory (HBM) for intermediate steps. The 11.58 GB reads and 5.54 GB writes confirm a bottleneck. Quick math: In this naive implementation, we iterate over columns (K, V) in the outer loop and rows (Q, O) in the inner loop. This loop order forces us to reload the Accumulator ($O$O) and Query ($Q$Q) from HBM for every single chunk of keys. With $S/Bc = 64$S/Bc=64 chunks, we are reading and writing the entire output matrix 64 times! Math: Size of $O \approx 83 \text{ MB}$O≈83 MB. Reads $\approx 64 \times (Q + O) \approx 10.6 \text{ GB}$≈64×(Q+O)≈10.6 GB. Writes $\approx 64 \times O \approx 5.3 \text{ GB}$≈64×O≈5.3 GB.

  for j in range(0, Tc):
      # Load K_j, V_j from HBM
      # ....
      for i in range(0, Tr):
          # ....
          # -> Reading previous O_i (Bc,D)
          prev_oi = tl.load(o_ptr + offset_i)
          # -> Writing previous O_i (Bc,D)
          tl.store(o_ptr + offset_i, oi_new)

We are treating HBM like a register, which explains the massive bandwidth usage!

The last thing to look at before starting to look at the potential solutions we can implement is the source from the profiler.

ncu marks problematic kernel code lines with a warning emoji ⚠️, which is quite helpful! The div operation is usually costly on CPUs, and it’s basically the same story for GPUs. CUDA implements floating-point division on GPUs via reciprocal + multiply ops: using MUFU (Multi-Function Unit) followed by one or two FFMA or FMUL instructions to refine precision. If we look at the SASS disassembly for this line:

     MUFU.RCP R8, R8
     ...
@P5  FMUL R29, R29, 0.25
@P5  FMUL R15, R15, 0.25
@!P4 FMUL R10, R10, 16777216
@!P6 FMUL R29, R29, 16777216
@!P4 FMUL R23, R23, 16777216
@!P6 FMUL R15, R15, 16777216
     FMUL R10, R9, R10
     FMUL R29, R8, R29
     FMUL R9, R9, R23
     FMUL R8, R8, R15

Because we are doing the division inside the hot loop we are doing a lot more work here and we can start thinking about how we can avoid this division all together.

  oi_new = (
      alpha[:, None] * prev_li[:, None] * prev_oi
      + beta[:, None] * tl.dot(pij, vj)
  ) / li_new[:, None]

Next Implementation Plan

Based on the ncu profiling data and our analysis of the memory traffic, we have identified three critical bottlenecks to address in the next iteration (V2).

1. Invert the Loop Order: The massive 11.58 GB of main memory reads and 5.54 GB of writes is our biggest performance killer. It stems from treating HBM as a temporary register for our output accumulator. Current loop forces us to repeatedly read/write the accumulator $O$O and reload $Q$Q for every block of $K$K. We must invert the loops: Parallelize the kernel over Queries (rows) so that each thread block handles a tile of $Q$Q:

   • Stationary Data: Rows of $O$O and $Q$Q remain in SRAM/registers for the entire kernel.
   • Streaming Data: $K$K and $V$V stream from HBM. Since all blocks share the same $K/V$K/V, the L2 cache will absorb most of the traffic (coalesced access).

2. Defer Normalization: Before describing the fix, let’s recall the true attention output: $O_i = \frac{\sum_j e^{S_{ij}} V_j}{\sum_j e^{S_{ij}}}$Oi​=∑j​eSij​∑j​eSij​Vj​​. Inside the kernel, we currently divide during every iteration of the loop: $O_{\text{new}} = \frac{\dots}{l_{\text{new}}}$Onew​=lnew​…​. This normalization step does not need to happen each time. We will store the unnormalized accumulated numerator and denominator in registers throughout the loop. We will perform the normalization once, at the end of the kernel, just before writing the result to HBM.

3. Tuning Block Sizes: Theoretical occupancy is capped at 25% due to shared-memory pressure. Smaller B_c may reduce register/SRAM pressure and increase the number of active warps. But too small a tile can fragment memory access and reduce bandwidth efficiency. I’ll leave tweaking this value for later.