The FlashAttention (FA) algorithm, as described in the original paper, isn’t performant for inference right out of the box. The technique to make it performant for inference was described in FlashDecoding. We will refer to this technique as “split-KV” and in this post, we will (1) look at why we need split-KV, (2) benchmark split-KV to see for ourselves that it is faster than vanilla FlashAttention for long sequences, and finally (3) implement the split-KV variant on top of the CuteDSL FA3 implementation for Hopper GPUs in the cutlass library.
Background: Parallelization in FlashAttention
FlashAttentionCovered in an earlier post parallelizes the work along the three dimensions of batch-size, number of heads and $Q$’s sequence length $L_Q$. In other words, the FlashAttention kernel generates $B \times H \times N_{Q}$ units of work where $B$ is the batch-size, $H$ is the number of heads and $N_{Q}$ is the number of blocks of size $M$ that the sequence length $L_Q$ is split into. Our first problem in inference is that $L_Q=1$ for decoding. Our second problem is that for longer sequences, the batch-size is typically small and can be as low as 1. This means fewer units of work to be done in parallel. And fewer things to parallelize means the GPU is underutilized and GPUs are expensive so that’s not great. For example, an H200 GPU has 132 Streaming Multiprocessors (SMs) and if we have less than 132 units of work, say, $B = 1$ and $H = 32$, then we would only use $32/132 \approx 24\%$ of the SMs at a time.
Let’s look at Figure $1$ to jog our memory of how FA works. It shows a single sequence of length $L_Q > 1$ where the resulting blocks $Q1$, $Q2$ and $Q3$ are independent units of work and can be processed in parallel by the GPU. Recall that the matrix $S$ is not actually materialized. See this post for a more detailed review of FA.
In Figure $2$, we depict the case of decoding with a single unit of work ($L_Q = M = 1$) for a given sequence.
In contrast to Figure $1$ which has a higher $L_Q$, we can see we have fewer units of work to parallelize. The solution to this conundrum is to find a way to increase the number of independent units of work per sequence for the GPU to process in parallel. One way to do that is to partition (split) the KV sequence, use FA to compute partial outputs for each of those splits in parallel and then combine the partial outputs to get our final output. This is the split-KV technique and we show it in Figure $3$.
Benchmarking split-KV against vanilla FA
We now turn to benchmarking the latency of split-KV against standard FATODO: add the bandwidth and SM utilization metrics from NCU. . To do this, we will use the FA3 kernel in the flash-attention repo which already implements split-KV and lets us toggle it on and off. We plot the latency of the FA3 kernel in Figure $4$. We see that split-KV is better at longer sequences for small batch-sizes (left subplot) and we see that the gains diminish as we increase the batch-size (right subplot).
The FA3 kernel uses a heuristic to determine the number of splits to use and is best summarized by this comment in the code:
// Find the number of splits that maximizes the occupancy. For example, if we have
// batch * n_heads = 48 and we have 108 SMs, having 2 splits (efficiency = 0.89) is
// better than having 3 splits (efficiency = 0.67). However, we also don't want too many
// splits as that would incur more HBM reads/writes.
// So we find the best efficiency, then find the smallest number of splits that gets 85%
// of the best efficiency.
To elaborate on the efficiency bit in the comment, we want to increase the units of work just enough to use all our SMs but not in a way that the efficiency (i.e. occupancy) goes down because of wave quantization. The $2$-split case gives an efficiency of $(2 \times 48)/108 = 0.89$ because it is just a single wave ($2 \times 48 < 108$). In the 3-split case, however, we have a total of $3 \times 48 = 144 > 108$ units of work which means the first wave will be 108 units (efficiency $= 108/108 = 1$) and the second will be 36 units (efficiency $= 36/108 = 0.33$) resulting in an average efficiency of $(1 + 0.33)/2 = 0.67$.
The other thing to note is that split-KV is not free and adds overhead for reading/writing the partial outputs from/for the combine kernel. So we can’t increase the number of splits arbitrarily. In fact, in the right subplot it Figure $1$, the number of KV splits is $1$ for batch_size >= 4 as there is enough parallel work available already without resorting to splitting.
Implementing split-KV for FA3 in CuteDSL
Now that we have convinced ourselves that split-KV is worth our time, let’s try to implement it. We will build on top of the CuteDSL FA3 example in the official cutlass repo.
Before we talk about split-KV related code changes, here is a very rough sketch of the standard FA3 implementation:
- It uses a grid size of a $N_Q \times H \times B$ where $N_Q$ is the number of blocks of size $M$ that the sequence length $L_Q$ is split into, $H$ is the number of heads and $B$ is the batch-size. (add side note this is for a non-persistent kernel. A persistent kernel would use a grid size equal to the number of SMs and then schedule $N_Q x H x B$ units of work on them)
- The kernel uses warp-specialization and has a dedicated warp-group for loading $Q, K, V$ (using TMA) and two warp-groups for attention computation (using WGMMA).
- $Q, K, V$ are loaded from the global memory (gmem) into shared memory (smem) with the pipelines staged in smem. For the pipeline, the “producer” is the TMA gmem -> smem load and the “consumer” is the WGMMA op.
- WGMMA writes the outputs to registers which the epilogue first copies to smem in smaller chunks and then TMA copies them from smem -> gmem.
Our plan then is to reuse as much of the existing FA3 kernel as possible We also keep our lives simpler by implementing the non-causal, non-windowed version. . The code for the split-KV implementation is here. The modified FA3 kernel is the fmha_splitkv.py file which we call from within the test_splitkv.py file. The latter also calls the flash_fwd_combine kernel to calculate the final output from the partial outputs we get from the split-KV kernel.
Here are the main changes we made to the FA3 kernel to implement split-KV:
- Included number of splits in the grid size in the “driver” script that calls the our modified split-KV FA3 kernel followed by the combine kernel. We do so by folding
ns(i.e. number of splits) into the existing batch dimensionbto make the new batch dimensionb * nsas follows:o_p_k = o_partial.view(b * ns, s_q, h, 1, d).permute(1, 4, 3, 2, 0) lse_p_k = lse_partial.view(b * ns, h, s_q).permute(2, 1, 0)[:, None, None, :, :] - Other than book keeping changes that are a result of including splits into the batch dim, the core thing we change in the kernel is updating the indexing logic when looping over the K/V blocks for the entire sequence. Whereas the standard kernel would loop over all the K/V blocks (ignoring masking for now), the split-KV kernel will only loop over the K/V blocks for the current split. These boundaries of where to start and end the K/V loop are implemented in the
get_split_kv_trip_infofunction.@cute.jit def get_split_kv_trip_info( mask_type, blk_coord: cute.Coord, tile_shape: cute.Shape, seqlen_q: Int32, seqlen_k: Int32, split_idx: Int32, num_splits: int, window_size_left: Optional[Int32] = None, window_size_right: Optional[Int32] = None, ) -> Tuple[Int32, Int32, Int32, Int32]: assert mask_type is fmha_utils.MaskEnum.RESIDUAL_MASK, ( "split-kv only supports RESIDUAL mask, not causal or windowed" ) full_start = fmha_utils.FusedMask.get_trip_start( mask_type, blk_coord, tile_shape, seqlen_q, seqlen_k, window_size_left ) full_count = fmha_utils.FusedMask.get_trip_count( mask_type, blk_coord, tile_shape, seqlen_q, seqlen_k, window_size_left, window_size_right ) base_split_size = full_count // num_splits rem = full_count % num_splits trip_count = base_split_size if split_idx < rem: trip_count = base_split_size + 1 trip_start = full_start + split_idx * base_split_size + min(split_idx, rem) ... return trip_start, trip_count, unmasked_count, trailing_count - We simplify the epilogue to directly write the partial outputs/accumulators from the registers to gmem. We want the partial results that will be read by the combine kernel to be FP32 and the accumulators are already in FP32 so we copy them directly without any down casting.
@cute.jit def epilogue_splitkv( self, acc_pv: cute.Tensor, pv_tiled_mma: cute.TiledMma, mO_qdl: cute.Tensor, wg_coord_partial: Tuple, tidx: cutlass.Int32, seqlen_q: cutlass.Int32, ): """ Partial output O is already in FP32. We copy from accumulators in the reg -> gmem directly. """ thr_mma = pv_tiled_mma.get_slice(tidx) gO = cute.local_tile( mO_qdl, self.pv_mma_tiler[:2], (wg_coord_partial[0], 0, wg_coord_partial[2]), ) tOgO = thr_mma.partition_C(gO) cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) tOcO = thr_mma.partition_C(cO) acc_mn = cute.make_tensor( acc_pv.iterator, self.layout_acc_mn(pv_tiled_mma, acc_pv.layout) ) tOgO_mn = cute.make_tensor( tOgO.iterator, self.layout_acc_mn(pv_tiled_mma, tOgO.layout) ) tOcO_mn = cute.make_tensor( tOcO.iterator, self.layout_acc_mn(pv_tiled_mma, tOcO.layout) ) m_offset = wg_coord_partial[0] * self.pv_mma_tiler[0] for i in cutlass.range_constexpr(cute.size(acc_mn, mode=[0])): if tOcO_mn[i][0] + m_offset < seqlen_q: for j in cutlass.range_constexpr(cute.size(acc_mn, mode=[1])): tOgO_mn[(i, j)] = acc_mn[i, j] - Finally, we re-useTODO: Write the combine kernel from scratch the flash_fwd_combine kernel from FA4 to combine the partial outputs (
o_partialandlse_partial) from the split-KV kernel above into the final output.
And that is it. In around ~200 lines of code and mostly bypassing the more involved parts of the kernel, we are done with our split-KV implementation.