{"slug": "optimization-diaries-s-p-ping-pong-for-flashattention-4-decode", "title": "Optimization diaries: S/P ping-pong for FlashAttention-4 decode", "summary": "A new optimization for FlashAttention-4 (FA4) decoding on NVIDIA Blackwell GPUs overlaps softmax and matrix multiplication operations by using spare tensor memory (TMEM) in a ping-pong scheme, achieving up to a 16% performance gain on supported single and multi-token decoding configurations for head dimensions 64 and 128. The code is available in PR #2817 on the FlashAttention repository.", "body_md": "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.\n\nIn this blog post, we discuss an optimization for FlashAttention-4 (FA4) decoding on NVIDIA Blackwell GPUs. Currently, in FA4 decoding, the matmul for block is issued only after the softmax for block has 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-pong between two slots so that while the softmax warps are writing the output of block to one slot, the MMA warp can issue to the other slot.\n\nThis 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 in [PR #2817](https://github.com/Dao-AILab/flash-attention/pull/2817) on the FlashAttention repository.\n\n## Recap on the FA4 forward pass\n\nLet , , and be the query, key, and value matrices. The forward pass of attention calculates the attention output as follows:\n\nwhere softmax is applied row-wise. In practice, FA4 does not materialize the full or matrices and instead processes them tile by tile. As such, softmax is computed “online” with rescaled when necessary.\n\nTo 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 of , , and from global memory (GMEM) to shared memory (SMEM). The MMA warp consumes and from load and issues for the softmax warps to consume. The softmax warps produce and update the online softmax statistics. Based on the softmax statistics, the correction warps rescale if necessary. The MMA warp then consumes with to issue .\n\nMoving forward, we suppress the transpose decoration on . For prefill, each CTA is assigned two 128-row tiles, a high -tile and a low -tile . The purpose of this is to overlap softmax() with . In this scheme, is stored in TMEM columns 0-127 and is stored in columns 128-255. These two slots are reused to store and . Figure 1 is taken from the FA4 preprint and depicts this schedule.\n\nFor single and multi-token decode, however, there is typically only a single (usually padded) 128-row tile. 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. Let and denote the matmul and corresponding scores for block . With only a single -tile, , , and proceed serially. Thus, nothing hides the latency of softmax, even though does not depend on . We implement an alternative parallelism strategy to rectify this.\n\n## The S/P Ping-Pong\n\nFor 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.\n\nPing-pong overlaps with by utilizing the unused TMEM in slot 1. In particular, where the MMA warp writes and the softmax warps write will ping-pong between the two buffers as shown in Figure 3:\n\nwhere the new issue order is as follows:\n\nThe cell widths in Figure 2 are not proportional to execution time.\n\nFigure 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. The accumulator 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.\n\n## Implementation\n\nTo 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 single  slot, there is a single barrier whose phase flips once per block. With two  slots, 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 of  matmuls 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.\n\nThe prologue proceeds as follows:\n\n1. Load warp copies and to SMEM.\n2. MMA warp issues .\n3. Load warp copies to SMEM.\n\nThe important change is the load order. Previously, the load warp produced tiles of and in the following order: . The ping-pong path path produces before appending the final tile of at the end. The second tile being loaded is to enable the MMA warp to issue immediately upon entering the main loop. Here is the code for the prologue:\n\n```\nif const_expr(self.use_s_ping_pong):\n   # Wait for Q\n   pipeline_q.consumer_wait_w_index_phase(0, mma_q_consumer_phase)\n   # Wait for K(0)\n   pipeline_kv.consumer_wait(mma_kv_consumer_state)\n   Ki_index, Ki_phase = mma_kv_consumer_state.index, mma_kv_consumer_state.phase\n   sK_cur = sK[None, None, None, Ki_index]\n   if const_expr(self.uneven_kv_smem):\n       sK_cur = self.offset_kv_smem(sK_cur, Ki_index, Ki_phase)\n   # Issue QK(0).\n   if (mma_pv_count & 1) == 0:\n       gemm_Si[0](smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK_cur.iterator))\n       pipeline_s_p_o.producer_commit_w_index(0)\n   else:\n       gemm_Si[1](smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK_cur.iterator))\n       pipeline_s_p_o.producer_commit_w_index(1)\n   mma_q_consumer_phase ^= 1\n   # Release K(0)\n   pipeline_kv.consumer_release(mma_kv_consumer_state)\n   # Advance to K(1)\n   mma_kv_consumer_state.advance() \n   O_should_accumulate = False\n```\n\nThe variable `mma_pv_count` is a global count of  matmuls issued. Calling `gemm_Si[0]` instructs ) to write to slot 0. The `pipeline_s_p_o.producer_commit_w_index(0)` call signals to the softmax warpgroup to wait for the  matmul to finish before giving the go-ahead to consume  from slot 0. Similarly, `gemm_Si[1]` and `pipeline_s_p_o.producer_commit_w_index(1)` do the same except for slot 1.\n\nThe ping-pong main loop proceeds as follows:\n\n```\nfor i in cutlass.range(block_iter_count - 1, unroll=1):\n   # Wait for K(i+1)\n   pipeline_kv.consumer_wait(mma_kv_consumer_state)\n   Ki_index, Ki_phase = mma_kv_consumer_state.index, mma_kv_consumer_state.phase\n   sK_cur = sK[None, None, None, Ki_index]\n   if const_expr(self.uneven_kv_smem):\n       sK_cur = self.offset_kv_smem(sK_cur, Ki_index, Ki_phase)\n   # Issue QK(i+1). Even global block count writes to TMEM slot 0 (0-127) and odd writes to TMEM slot 1 (128-255)\n   if ((mma_pv_count + 1) & 1) == 0:\n       gemm_Si[0](smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK_cur.iterator))\n       pipeline_s_p_o.producer_commit_w_index(0)\n   else:\n       gemm_Si[1](smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK_cur.iterator))\n       pipeline_s_p_o.producer_commit_w_index(1)\n   # Release K(i+1)\n   pipeline_kv.consumer_release(mma_kv_consumer_state)\n   # Advance to V(i)\n   mma_kv_consumer_state.advance()\n   # Wait for V(i)\n   pipeline_kv.consumer_wait(mma_kv_consumer_state)\n   Vi_index, Vi_phase = mma_kv_consumer_state.index, mma_kv_consumer_state.phase\n   tOrVi = tOrV[None, None, None, Vi_index]\n   sV_cur = sV[None, None, None, Vi_index]\n   if const_expr(self.uneven_kv_smem):\n       sV_cur = self.offset_kv_smem(sV_cur, Vi_index, Vi_phase)\n   # Phase for this block's slot. Each slot is reused every second block, and its barrier phase flips on each reuse.\n   pv_phase = (mma_pv_count >> 1) & 1\n   # Issue PV(i)\n   if (mma_pv_count & 1) == 0:\n       pipeline_s_p_o.producer_acquire_w_index_phase(0, pv_phase)\n       gemm_Pi[0](\n           tCrB=tOrVi,\n           sB=sV_cur,\n           zero_init=not O_should_accumulate,\n           mbar_ptr=pipeline_p_lastsplit.sync_object_full.get_barrier(0) if self.split_P_arrive > 0 else None,\n           mbar_phase=pv_phase,\n       )\n   else:\n       pipeline_s_p_o.producer_acquire_w_index_phase(1, pv_phase)\n       gemm_Pi[1](\n           tCrB=tOrVi,\n           sB=sV_cur,\n           zero_init=not O_should_accumulate,\n           mbar_ptr=pipeline_p_lastsplit.sync_object_full.get_barrier(1) if self.split_P_arrive > 0 else None,\n           mbar_phase=pv_phase,\n       )\n   pipeline_o_acc.producer_commit_w_index(mma_pv_count & 1)\n   mma_pv_count += 1\n   # Release V(i)\n   pipeline_kv.consumer_release(mma_kv_consumer_state)\n   # Advance to K(i+2)\n   mma_kv_consumer_state.advance() \n   O_should_accumulate = True\n```\n\nHere are a few notes to supplement the comments in the code block. The slot the  matmul writes to depends on the parity of `mma_pv_count + 1` since  runs one ahead of . The line `pv_phase = (mma_pv_count >> 1) & 1` extracts bit 1 from `mma_pv_count`, which is the parity of the reuse count of slot `mma_pv_count & 1`. Figure 6 summarizes how the slot PV writes to and `pv_phase` change with `mma_pv_count` through the first several iterations.\n\nThe `pipeline_s_p_o.producer_acquire_w_index_phase(0/1, pv_phase)` call serves two purpose:\n\n1. It waits on the softmax warps to finish producing and release the slot.\n2. It waits on the correction warps to finish any required rescale and release the slot.\n\nThe `pipeline_o_acc.producer_commit_w_index(mma_pv_count & 1)` call signals that the MMA warp is done accumulating this iteration’s  into the  accumulation buffer so it is safe for correction warps to consume. Note that this is not required for the original path, since the serialization guaranteed that  would complete before the potential correction rescale for  would arrive.\n\n## IKET Profiling\n\nNVIDIA’s In-Kernel Event Tracing ([IKET](https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/guides/iket_profiling.html))  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 issue  while the softmax warps are still processing , 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:\n\nThe relevant sequence is:\n\n1.  matmul (`mma_issue_QK` )\n2. softmax() (`sm_compute` )\n3.  matmul (`mma_issue_PV` )\n4.  matmul (`mma_issue_QK` )\n\nThere 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:\n\nNotice how the blue `mma_issue_QK` bars overlap significantly with the purple `sm_compute` bars. This is concrete evidence that the ping-pong path is issuing  matmuls and softmax concurrently. \n\nPerformance\n\nDecoding 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 ratios = 16:1 and 16:2. The benefit is small at shorter sequence 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 the and softmax overlap can be benefited from. For head dimension 64, the increase in achieved bandwidth reaches 15.6% for GQA ratio and 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.\n\n## Conclusion\n\nIn this blog post, we discussed an optimization for FA4 decode that removes an unnecessary serialization between and softmax by ping-ponging across 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%.", "url": "https://wpnews.pro/news/optimization-diaries-s-p-ping-pong-for-flashattention-4-decode", "canonical_source": "https://research.colfax-intl.com/optimization-diaries-s-p-ping-pong-for-flashattention-4-decode/", "published_at": "2026-09-07 17:59:28+00:00", "updated_at": "2026-09-07 18:31:09.235968+00:00", "lang": "en", "topics": ["machine-learning", "ai-research", "ai-infrastructure"], "entities": ["NVIDIA", "FlashAttention-4", "Dao-AILab"], "alternates": {"html": "https://wpnews.pro/news/optimization-diaries-s-p-ping-pong-for-flashattention-4-decode", "markdown": "https://wpnews.pro/news/optimization-diaries-s-p-ping-pong-for-flashattention-4-decode.md", "text": "https://wpnews.pro/news/optimization-diaries-s-p-ping-pong-for-flashattention-4-decode.txt", "jsonld": "https://wpnews.pro/news/optimization-diaries-s-p-ping-pong-for-flashattention-4-decode.jsonld"}}