FlashAttention Inner Loop Implementation

The Anatomy of a Tiled Attention Kernel

Before we dive into the inner loop, let's establish the overall structure of our CUDA kernel. The FlashAttention algorithm is organized around two nested loops. As the diagram from the last lesson showed, the outer loop iterates over blocks of the query matrix $Q$, while the inner loop iterates over blocks of the key $K$ and value $V$ matrices.

For this lesson, we will focus on the code that executes for a single block of $Q$, which is loaded into fast on-chip memory (registers) at the beginning of a thread block's execution. The "inner loop" of our kernel will then iterate through the tiles of $K$ and $V$ along the sequence length dimension.

The following blog post provides an excellent, from-the-ground-up implementation of a FlashAttention kernel. We will use its initial, unoptimized version as a clear, pedagogical guide.

Implementing the Inner Loop

The code you just reviewed provides a complete skeleton. Now, let's detail the operations inside that main for (int kv_idx = ...) loop. For each iteration, we process one tile of $K$ and $V$ to update our running output tile $O$.

1. Data Loading: From Global to On-Chip Memory

The first step inside the loop is to bring the current tiles of $K$ and $V$ from slow global HBM into fast on-chip memory. The process follows the standard pattern for high-performance kernels, which should be familiar from your work on GEMM optimization.

1. Global to Shared Memory (cp.async): A tile of $K$ (and later, $V$) is copied from global memory to shared memory. The blog post uses the cp.async PTX instruction, which is an asynchronous copy operation. This allows the GPU to initiate a memory transfer and proceed with other work, hiding memory latency.
2. Shared to Register Memory (ldmatrix): From shared memory, the data is loaded into registers just before it's needed for computation. The ldmatrix PTX instruction is specifically designed to load matrix fragments for use in Tensor Core operations.

The article you read previously explains this data movement flow in detail. You can refer back to the sections "Global to Shared memory data transfer" and "Shared memory to Register memory data transfer" in that resource if you need a refresher on how ldmatrix maps threads to matrix elements.

2. First MMA: Computing Scores ($S_{\text{block}}=Q_{\text{block}}{K}_{\text{block}}^T$)

Once a tile of $K$ is in registers (K_rmem), we can perform the first matrix multiplication. We multiply our persistent tile of $Q$ (held in Q_rmem) with the transpose of the current $K$ tile.

The attention_v1_kernel code uses a triple-nested loop to perform this mma.m16n8k16 operation, accumulating the results in the S_rmem register array. This array now holds a [WARP_Q, BLOCK_KV] tile of the attention scores, which we'll call $S_{\text{block}}$. This corresponds to one block of the full $Q{K}^T$ matrix.

3. Online Softmax and Output Update

This is where the magic happens. Instead of writing $S_{\text{block}}$ to memory, we immediately process it to update our running statistics. Let's trace the logic, which you saw in the previous lesson's numerical example.

For each row of our S_rmem tile, we must:

1. Find Local Max: Compute the maximum value in the current score tile, $m_{\text{block}}$.
2. Find New Global Max: Compare the local max with our running maximum from previous blocks, $m_{\text{old}}$, and find the new true maximum: $m_{\text{new}}=\max (m_{\text{old}},m_{\text{block}})$.
3. Rescale Previous Results: If the maximum has changed ($m_{\text{new}}>m_{\text{old}}$), we must rescale our previous running denominator sum $d_{\text{old}}$ and our previous output accumulator $o_{\text{old}}$. The update rules are: $$d_{\text{new}}=d_{\text{old}}\cdot e^{m_{\text{old}}-m_{\text{new}}}$$ $$o_{\text{new}}=o_{\text{old}}\cdot e^{m_{\text{old}}-m_{\text{new}}}$$
4. Compute $P_{\text{block}}$ and $d_{\text{block}}$: We compute the exponentiated values for the current score tile, normalized by the new maximum: $P_{\text{block}}=\exp (S_{\text{block}}-m_{\text{new}})$. The sum of these values gives us the local denominator contribution, $d_{\text{block}}$.
5. Update Global Denominator: Add the local contribution to the (potentially rescaled) running sum: $d_{\text{new}}=d_{\text{new}}+d_{\text{block}}$.
6. Second MMA ($P_{\text{block}}{V}_{\text{block}}$): Now, we load the corresponding tile of $V$ into registers and immediately perform the second matrix multiplication. This gives us the output contribution for this block.
7. Update Global Output: We add this new contribution to our (potentially rescaled) running output accumulator: $o_{\text{new}}=o_{\text{new}}+(P_{\text{block}}{V}_{\text{block}})$.

After this sequence, m_new, d_new, and o_new become m_old, d_old, and o_old for the next iteration of the kv_idx loop.

To see how this abstract algorithm is implemented in a real, albeit simplified, kernel, let's examine a pedagogical CUDA C++ implementation.

By tracing the code in FlashAttentionFirstAttempt.cu, you can see the direct, line-by-line implementation of the online softmax algorithm we've been discussing. The code doesn't use low-level PTX and is structured for clarity, making it an ideal bridge from the mathematical formulation to a working kernel.

Conclusion

In this lesson, you have dissected the inner loop of a FlashAttention kernel and seen how its abstract principles map to concrete CUDA C++ code. You are no longer just looking at a block diagram; you are now reading the implementation.

Key Takeaways:

• Kernel Structure: A FlashAttention kernel is typically structured with an outer loop over $Q$ tiles and an inner loop over $K$/$V$ tiles.
• Data Flow: Within the inner loop, tiles of $K$ and $V$ are loaded from global memory into shared memory and then registers for computation.
• Fused Computation: The core of the inner loop is a fused sequence of operations: QK^T matrix multiplication, online softmax update of running statistics ($m,d,o$), and the PV matrix multiplication.
• Online Update in Code: You have seen how if statements and rescaling factors in the code directly implement the update rules for the running maximum, denominator, and output accumulator.

You have now built a solid understanding of the most critical part of the FlashAttention algorithm's implementation. The next step is to put all the pieces together.

In the next lesson, we will expand our focus from the inner loop to the complete forward pass. You will learn how to initialize the loop, handle the final normalization after the loop concludes, and write the final output back to global memory. This will culminate in a simplified but complete FlashAttention kernel that you can verify against a standard implementation.