We have two new Colfax articles on optimizing FlashAttention-4, featuring work by myself (decode) and Jack Carlisle (bwd pass): https://research.colfax-intl.com/optimization-diaries-s-p-ping-pong-for-flashattention-4-decode/… https://research.colfax-intl.com/optimization-diaries-improving-flashattention-4-backward-for-head-dimension-64/… For decode, an S/P ping-pong is used to overlap the softmax for KV block i with QK^T for block i+1:

Optimization diaries: S/P ping-pong for FlashAttention-4 decode - Colfax Research

Source: https://research.colfax-intl.com/optimization-diaries-s-p-ping-pong-for-flashattention-4-decode/ LLM Inference is divided into a prefill phase and a decode phase. During prefill, the model processes a large number of input tokens and populates a key-value (KV) cache. During decode, it autoregressively generates one or a few new tokens at a time using the cached keys and values.

In this blog post, we discuss an optimization for FlashAttention-4 (FA4) decoding on NVIDIA Blackwell GPUs. Currently, in FA4 decoding, theQKTQK^Tmatmul forKVKVblocki+1i+1is issued only after the softmax forKVKVblockiihas completed, even though there is no mathematical dependency between the two. To instead overlap them, we can leverage spare tensor memory (TMEM) present in the decode path. The spare TMEM is used to ping-pongS/PS/Pbetween two slots so that while the softmax warps are writing the output of blockiito one slot, the MMA warp can issueQKTQK^Tto the other slot.

This change achieves up to a 16% performance gain on supported single and multi-token decoding configurations for head dimension 64 and 128. The code may be found inPR #2817on the FlashAttention repository.

Recap on the FA4 forward pass

LetQQ,KK, andVVbe the query, key, and value matrices. The forward pass of attention calculates the attention outputOOas follows:

S=1dQKT,P=softmax⁡(S),O=PVS = \frac{1}{\sqrt{d}}QK^T, \qquad P = \operatorname{softmax}(S), \qquad O = PV

where softmax is applied row-wise. In practice, FA4 does not materialize the fullSSorPPmatrices and instead processes them tile by tile. As such, softmax is computed “online” withOOrescaled when necessary.

To implement the forward pass, FA4 uses a web of overlapping pipelines across 5 different warp roles: load, MMA, softmax, correction, and epilogue. The load warp copies tiles ofQQ,KK, andVVfrom global memory (GMEM) to shared memory (SMEM). The MMA warp consumesQQandKKfrom load and issuesS=QKTS=QK^Tfor the softmax warps to consume. The softmax warps producePPand update the online softmax statistics. Based on the softmax statistics, the correction warps rescaleOOif necessary. The MMA warp then consumesPPwithVVto issuePVPV.

Moving forward, we suppress the transpose decoration onKK. For prefill, each CTA is assigned two 128-rowQQtiles, a highQQ-tileQHQ^Hand a lowQQ-tileQLQ^L. The purpose of this is to overlap softmax(SHS^H) withQLKQ^L K. In this scheme,SHS^His stored in TMEM columns 0-127 andSLS^Lis stored in columns 128-255. These two slots are reused to storePHP^HandPLP^L. Figure 1 is taken from the FA4 preprint and depicts this schedule.

**Figure 1.**Each CTA is assigned twoQQ-tiles. This figure shows how the two matmuls and softmax for each are scheduled relative to each other during prefill.For single and multi-token decode, however, there is typically only a single (usually padded) 128-rowQQtile. The overlap of operations in Figure 1 is not relevant in this case. Even so, TMEM columns 128-255 are still allocated and left unused. LetQK(i)QK(i)andS(i)S(i)denote theQKQKmatmul and corresponding scores forKVKVblockii. With only a singleQQ-tile,QK(i)QK(i),softmax⁡(S(i))\operatorname{softmax}(S(i)), andQK(i+1)QK(i+1)proceed serially. Thus, nothing hides the latency of softmax, even thoughQK(i+1)QK(i+1)does not depend onsoftmax⁡(S(i))\operatorname{softmax}(S(i)). We implement an alternative parallelism strategy to rectify this.

The S/P Ping-Pong

For simplicity, we refer to TMEM columns 0-127 as slot 0 and columns 128-255 as slot 1. Figure 2 shows which operands are resident in each slot during the prologue and the first few iterations of the main loop for the first work for the base path.

**Figure 2.**The operands live in slots 0 and 1 for the base path.SStiles are in blue,PPtiles are in yellow, and grey means empty.Ping-pong overlapsQK(i+1)withsoftmax⁡(S(i))by utilizing the unused TMEM in slot 1. In particular, where the MMA warp writesSSand the softmax warps writePPwill ping-pong between the two buffers as shown in Figure 3:

**Figure 3.**The operands live in slots 0 and 1 for the S/P ping-pong path.where the new issue order is as follows:

**Figure 4.**MMA/softmax schedule for the S/P ping-pong path.The cell widths in Figure 2 are not proportional to execution time.

Figure 5 summarizes some of the synchronizations among the MMA, load, softmax, and correction warps with respect to the ping-pong. Dashed arrows indicate a TMEM value being consumed. TheOOaccumulator uses a single buffer, but is drawn twice for visual clarity. While the correction warps are depicted issuing “rescale” each iteration, whether or not rescaling is actually done ultimately depends on the extent to which the row-max has changed.

Figure 5. Overview of pipelines and synchronizations for theS/PS/Pping-pong. A dotted arrow depicts an operand in TMEM being consumed. The twoOO‘s represent the same accumulation buffer.## Implementation

To implement the ping-pong, careful coordination is necessary to ensure warps are waiting on or consuming from the appropriate slot. The original code only needed a single phase bit since, with a singleSSslot, there is a single barrier whose phase flips once per block. With twoSSslots, each slot’s barrier flips once every other block. So we must keep track of which barrier to use and how many times that barrier has already flipped. Thus, we use a pair of barriers and two bits: bit 0 and bit 1 of the global count ofPVPVmatmuls issued (mma\_pv\_count). Bit 0 selects the barrier and bit 1 tracks (mod 2) how many times the given slot’s barrier has flipped.

The prologue proceeds as follows:

1. Load warp copiesQ0Q0andK0K0to SMEM.
2. MMA warp issuesQK(0)QK(0).
3. Load warp copiesK1K1to SMEM.

The important change is theKVKVload order. Previously, the load warp produced tiles ofKKandVVin the following order:K0,V0,K1,V1,K2,V2,…,K0, V0, K1, V1, K2, V2, \dots ,. The ping-pong path path producesK0,K1,V0,K2,V1,K3,V2,…K0, K1, V0, K2, V1, K3, V2, \dotsbefore appending the final tile ofVVat the end. The secondKKtile being loaded is to enable the MMA warp to issueQK(1)QK(1)immediately upon entering the main loop. Here is the code for the prologue:

if const_expr(self.use_s_ping_pong):
   # Wait for Q
   pipeline_q.consumer_wait_w_index_phase(0, mma_q_consumer_phase)
   # Wait for K(0)
   pipeline_kv.consumer_wait(mma_kv_consumer_state)
   Ki_index, Ki_phase = mma_kv_consumer_state.index, mma_kv_consumer_state.phase
   sK_cur = sK[None, None, None, Ki_index]
   if const_expr(self.uneven_kv_smem):
       sK_cur = self.offset_kv_smem(sK_cur, Ki_index, Ki_phase)
   # Issue QK(0).
   if (mma_pv_count & 1) == 0:
       gemm_Si[0](smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK_cur.iterator))
       pipeline_s_p_o.producer_commit_w_index(0)
   else:
       gemm_Si[1](smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK_cur.iterator))
       pipeline_s_p_o.producer_commit_w_index(1)
   mma_q_consumer_phase ^= 1
   # Release K(0)
   pipeline_kv.consumer_release(mma_kv_consumer_state)
   # Advance to K(1)
   mma_kv_consumer_state.advance() 
   O_should_accumulate = False

The variablemma\_pv\_countis a global count ofPVPVmatmuls issued. Callinggemm\_Si\[0\]instructsQK(0QK(0) to write to slot 0. Thepipeline\_s\_p\_o\.producer\_commit\_w\_index\(0\)call signals to the softmax warpgroup to wait for theQK(0)QK(0)matmul to finish before giving the go-ahead to consumeS(0)S(0)from slot 0. Similarly,gemm\_Si\[1\]andpipeline\_s\_p\_o\.producer\_commit\_w\_index\(1\)do the same except for slot 1.

The ping-pong main loop proceeds as follows:

for i in cutlass.range(block_iter_count - 1, unroll=1):
   # Wait for K(i+1)
   pipeline_kv.consumer_wait(mma_kv_consumer_state)
   Ki_index, Ki_phase = mma_kv_consumer_state.index, mma_kv_consumer_state.phase
   sK_cur = sK[None, None, None, Ki_index]
   if const_expr(self.uneven_kv_smem):
       sK_cur = self.offset_kv_smem(sK_cur, Ki_index, Ki_phase)
   # Issue QK(i+1). Even global block count writes to TMEM slot 0 (0-127) and odd writes to TMEM slot 1 (128-255)
   if ((mma_pv_count + 1) & 1) == 0:
       gemm_Si[0](smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK_cur.iterator))
       pipeline_s_p_o.producer_commit_w_index(0)
   else:
       gemm_Si[1](smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK_cur.iterator))
       pipeline_s_p_o.producer_commit_w_index(1)
   # Release K(i+1)
   pipeline_kv.consumer_release(mma_kv_consumer_state)
   # Advance to V(i)
   mma_kv_consumer_state.advance()
   # Wait for V(i)
   pipeline_kv.consumer_wait(mma_kv_consumer_state)
   Vi_index, Vi_phase = mma_kv_consumer_state.index, mma_kv_consumer_state.phase
   tOrVi = tOrV[None, None, None, Vi_index]
   sV_cur = sV[None, None, None, Vi_index]
   if const_expr(self.uneven_kv_smem):
       sV_cur = self.offset_kv_smem(sV_cur, Vi_index, Vi_phase)
   # Phase for this block's slot. Each slot is reused every second block, and its barrier phase flips on each reuse.
   pv_phase = (mma_pv_count >> 1) & 1
   # Issue PV(i)
   if (mma_pv_count & 1) == 0:
       pipeline_s_p_o.producer_acquire_w_index_phase(0, pv_phase)
       gemm_Pi[0](
           tCrB=tOrVi,
           sB=sV_cur,
           zero_init=not O_should_accumulate,
           mbar_ptr=pipeline_p_lastsplit.sync_object_full.get_barrier(0) if self.split_P_arrive > 0 else None,
           mbar_phase=pv_phase,
       )
   else:
       pipeline_s_p_o.producer_acquire_w_index_phase(1, pv_phase)
       gemm_Pi[1](
           tCrB=tOrVi,
           sB=sV_cur,
           zero_init=not O_should_accumulate,
           mbar_ptr=pipeline_p_lastsplit.sync_object_full.get_barrier(1) if self.split_P_arrive > 0 else None,
           mbar_phase=pv_phase,
       )
   pipeline_o_acc.producer_commit_w_index(mma_pv_count & 1)
   mma_pv_count += 1
   # Release V(i)
   pipeline_kv.consumer_release(mma_kv_consumer_state)
   # Advance to K(i+2)
   mma_kv_consumer_state.advance() 
   O_should_accumulate = True

Here are a few notes to supplement the comments in the code block. The slot theQKQKmatmul writes to depends on the parity ofmma\_pv\_count \+ 1sinceQKQKruns one ahead ofPVPV. The linepv\_phase = \(mma\_pv\_count \>\> 1\) & 1extracts bit 1 frommma\_pv\_count, which is the parity of the reuse count of slotmma\_pv\_count & 1. Figure 6 summarizes how the slot PV writes to andpv\_phasechange withmma\_pv\_countthrough the first several iterations.

**Figure 6.**PV slot and pv_phase for the givenmma\_pv\_count.Thepipeline\_s\_p\_o\.producer\_acquire\_w\_index\_phase\(0/1, pv\_phase\)call serves two purpose:

1. It waits on the softmax warps to finish producingP(i)P(i)and release the slot.
2. It waits on the correction warps to finish any requiredOOrescale and release the slot.

Thepipeline\_o\_acc\.producer\_commit\_w\_index\(mma\_pv\_count & 1\)call signals that the MMA warp is done accumulating this iteration’sPVPVinto theOOaccumulation buffer so it is safe for correction warps to consume. Note that this is not required for the original path, since the serialization guaranteed thatPV(i)PV(i)would complete before the potential correction rescale forS(i+1)S(i+1)would arrive.

IKET Profiling

NVIDIA’s In-Kernel Event Tracing (IKET) enables us to examine the activity of individual warps over the lifetime of a kernel. While our changes in theory allow the MMA warp to issueQK(i+1)QK(i+1)while the softmax warps are still processingS(i)S(i), IKET can verify that this scheduling change actually occurs. To that end, we provide two traces from IKET: base vs ping-pong. Here is the trace from the base path:

Figure 7. Order of instructions issued by load, MMA, and softmax warps in the FA4 base path during decode.The relevant sequence is:

1. QKQKmatmul (mma\_issue\_QK)
2. softmax(SS) (sm\_compute)
3. PVPVmatmul (mma\_issue\_PV)
4. QKQKmatmul (mma\_issue\_QK)

There is little to no overlap between the bars for these operations, so they are effectively serialized. For comparison, here is the trace from the ping-pong path:

Figure 8. Order of instructions issued by load, MMA, and softmax warps in the FA4 S/P ping-pong path during decode.Notice how the bluemma\_issue\_QKbars overlap significantly with the purplesm\_computebars. This is concrete evidence that the ping-pong path is issuingQKQKmatmuls and softmax concurrently.

Performance

Decoding is memory-bound so we report achieved memory bandwidth as the performance metric. Figure 9 displays benchmark results measured on an NVIDIA B200 Blackwell GPU for a subset of shapes where the ping-pong path is taken, including grouped-query attention (GQA) with ratiosHq:HkvH_q:H_{kv}= 16:1 and 16:2. The benefit is small at shorterKVKVsequence lengths, but grows substantially as the sequences become longer. This is consistent with the change targeting steady-state, where longer sequences expose more iterations in which theQKQKand softmax overlap can be benefited from. For head dimension 64, the increase in achieved bandwidth reaches 15.6% for GQA ratio16:216:2and 16.0% for GQA ratio 16:1. For head dimension 128, the increases are up to 9.3% for 16:1 and 14.9% for 16:2. Figure 10 compares both single and multi-token decode at a fixed sequence length of 128k. The improvements are consistent across all query lengths considered.

**Figure 9.**Single token FA4 decode benchmark results for base path (in blue) vs ping-pong path (in orange) measured on an NVIDIA B200 Blackwell GPU.**Figure 10.**Single and multi-token FA4 decode benchmark results for base path (in blue) vs ping-pong path (in orange) measured on an NVIDIA B200 Blackwell GPU. Sequence length is fixed to 128k.## Conclusion

In this blog post, we discussed an optimization for FA4 decode that removes an unnecessary serialization betweenQKQKand softmax by ping-pongingS/PS/Pacross two TMEM slots, one of which was previously idle. We illustrated how parallelism strategies for prefill do not always apply to decode, covered implementation details of the ping-pong, and generated IKET traces to confirm the intended overlap. Benchmarks reported gains in performance of up to 16%.