From 74d455e65b37530902da1e855ca2935ef4a362e6 Mon Sep 17 00:00:00 2001
From: tlopex <820958424@qq.com>
Date: Thu, 6 Aug 2026 01:41:31 -0400
Subject: [PATCH] Clarify FA4 warpgroup terminology

---
 chapter_flash_attention/index.md    | 4 +---
 zh/chapter_flash_attention/index.md | 4 +---
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/chapter_flash_attention/index.md b/chapter_flash_attention/index.md
index 1e43cba7..2ecb3f2f 100644
--- a/chapter_flash_attention/index.md
+++ b/chapter_flash_attention/index.md
@@ -17,8 +17,6 @@ Here, `QKᵀ` gives the attention scores between queries and keys, and $d$ is th
 
 FlashAttention divides the computation into blocks and keeps only the current tiles and per-row softmax state on chip, avoiding the full score matrix while preserving the result of standard attention. Successive versions differ mainly in how this algorithm maps to the GPU. FlashAttention-2 improved work partitioning across thread blocks and warps. FlashAttention-3 used TMA, WGMMA, and warp specialization on Hopper to interleave data movement, the two MMAs, and softmax. FA4 targets Blackwell and reorganizes the pipeline around `tcgen05` and TMEM.
 
-The [FA4 paper](https://arxiv.org/abs/2603.05451) notes that, from Hopper to Blackwell, Tensor Core matrix-multiply throughput grew faster than exponential throughput, general-purpose floating-point throughput, and on-chip data movement. As the QKᵀ and PV MMAs become shorter, softmax exponentials, `O` rescaling, and TMEM/SMEM traffic account for a larger share of execution time. FA4 therefore cannot be obtained by simply replacing FA3's WGMMA instructions with `tcgen05.mma`; it must also reorganize where intermediate results live, which warpgroups execute each stage, and how those stages overlap.
-
 The preceding GEMM chapters introduced these Blackwell hardware paths: TMA moves tiles, `tcgen05` executes MMA, and TMEM holds accumulators. FA4 connects them into a different computation chain: a QKᵀ MMA computes the score tile `S = QKᵀ`, CUDA cores turn `S` into the unnormalized attention-weight tile `P`, and a PV MMA uses `P` and `V` to update the output accumulator `O`. Following the terminology in the FA4 paper, this chapter calls these operations the QKᵀ MMA and the PV MMA. Whenever softmax changes its exponent reference, the existing `O` in TMEM must first be converted to the new scale.
 
 This chapter is organized around three questions: how TMEM connects the two MMAs with softmax, how conditional rescaling reduces the number of `O` rescaling operations, and how multiple floating-point execution paths share exponential evaluation. We first derive the mathematical dependencies, then examine the TMEM layouts of `S`, `P`, and `O`, the division of work among warpgroups, and the barriers that hand off data and storage resources.
@@ -895,7 +893,7 @@ FA4 reuses the TMA, `tcgen05`, TMEM, and barrier machinery developed for the GEM
 2. Trace these four paths separately: Q/K in SMEM → S in TMEM, S in TMEM → P in TMEM, P in TMEM + V in SMEM → O in TMEM, and O in TMEM → O in GMEM. For each path, identify the executing role, source and destination storage, tile primitive, and hardware path. Which paths do not exist in the preceding GEMM kernel?
 3. A column $c$ in the fp16 view maps to physical 32-bit column $\lfloor c/2\rfloor$. Use this relation to derive the physical column ranges of `S0`, `S1`, `P0`, `P1`, `O0`, and `O1`. Which regions overlap, and which waits or barriers prevent an overlapping region from being read or overwritten too early?
 4. Trace one K/V block through `s_ready`, `p_o_rescale`, `p_ready_2`, and `o_ready`. For each barrier, identify who waits, who contributes arrivals, and which tile becomes safe to consume. Why does `p_o_rescale` expect 256 arrivals, and what overlap is gained by handing `P` to the PV MMA as 96 columns followed by 32 columns?
-5. The driver warpgroup reduces its register ceiling to 48 registers per thread, the two softmax warpgroups raise theirs to 200, and WG2 uses 64. Compute the total register budget for the four 128-thread warpgroups, then compare it with assigning 200 registers to every thread in the CTA. Why do the softmax roles need the largest allocation, and how does reducing WG3's ceiling make that allocation possible?
+5. WG3, which issues the TMA and MMA instructions, reduces its register ceiling to 48 registers per thread. The two softmax warpgroups, WG0 and WG1, raise theirs to 200, while WG2 uses 64. Compute the total register budget for the four 128-thread warpgroups, then compare it with assigning 200 registers to every thread in the CTA. Why do the softmax roles need the largest allocation, and how does reducing WG3's ceiling make that allocation possible?
 6. The kernel already rewrites the natural exponential as base-2 `exp2`. Why can the hardware exponential path still bottleneck softmax? Explain how splitting the elements between hardware `exp2` and the FMA-based cubic approximation changes execution-unit utilization, and identify which online-softmax equations remain unchanged.
 7. Let `SEQ_LEN_Q=6` and `SEQ_LEN_KV=8` with a bottom-right-aligned causal mask. What is the largest key index visible to query positions 0 and 5? With `BLK_N=4`, classify the K/V blocks for each query as fully valid, partially valid, or skipped. How does this difference affect causal task cost and scheduling order?
 8. Let `num_qo_heads=32`, `num_kv_heads=8`, and `BLK_M=128`. Compute `GQA_RATIO` and `SEQ_Q_PER_TILE`. For `kv_head_idx=3`, map packed rows 0, 5, and 127 to `(sequence offset, query head)`, and explain why all 128 rows can share one K/V tile.
diff --git a/zh/chapter_flash_attention/index.md b/zh/chapter_flash_attention/index.md
index 21b52d9a..f44da2a5 100644
--- a/zh/chapter_flash_attention/index.md
+++ b/zh/chapter_flash_attention/index.md
@@ -17,8 +17,6 @@ $$O = \text{softmax}(QK^{\top} / \sqrt{d})V$$
 
 FlashAttention 的核心做法是将计算分块，只在片上保留当前 tiles 和逐行 softmax 状态，从而避免保存完整的 score matrix，计算结果仍与标准 attention 相同。各版本的主要区别在于如何把这套算法映射到当代 GPU。FlashAttention-2 改进了 thread blocks 和 warps 之间的任务划分。FlashAttention-3 在 Hopper 上使用 TMA、WGMMA 和 warp specialization，将数据搬运、两次 MMA 与 softmax 交错执行。FA4 则面向 Blackwell，围绕 `tcgen05` 和 TMEM 重新组织这条 pipeline。
 
-[FA4 论文](https://arxiv.org/abs/2603.05451)指出，从 Hopper 到 Blackwell，Tensor Core 的矩阵乘吞吐量提升得比 exponential、普通浮点运算和片上数据搬运能力更快。QKᵀ MMA 和 PV MMA 所需的时间缩短后，softmax 中的指数计算、`O` 的重缩放以及 TMEM/SMEM 数据移动会占据更明显的执行时间。因此，FA4 不能只把 FA3 的 WGMMA 换成 `tcgen05.mma`，还需要重新安排中间结果的位置、warpgroup 的角色和各阶段的执行顺序。
-
 前面的 GEMM kernel 已经介绍了这些 Blackwell 硬件路径：TMA 搬运 tiles，`tcgen05` 执行 MMA，accumulator 保存在 TMEM 中。FA4 将它们连接成一条新的计算链：QKᵀ MMA 先计算 score tile `S = QK^T`，CUDA cores 再将 `S` 转换为尚未归一化的权重 tile `P`，PV MMA 最后用 `P` 和 `V` 更新 output accumulator `O`。本章沿用 FA4 论文的写法，将这两次操作分别称为 QKᵀ MMA 和 PV MMA。当 softmax 使用的指数参考值发生变化时，TMEM 中已有的 `O` 还需要先转换到新的尺度。
 
 本章围绕三个问题展开：TMEM 如何连接两次 MMA 与 softmax，conditional rescaling 如何减少 `O` 的重缩放次数，以及不同浮点执行路径如何共同承担指数计算。下面先推导这些操作的数学关系，再说明 `S`、`P` 和 `O` 的 TMEM layout、各个 warpgroups 的分工，以及 barriers 如何交接数据和存储资源。
@@ -897,7 +895,7 @@ FA4 复用了 GEMM kernel 中的 TMA、`tcgen05`、TMEM 和 barrier 机制，但
 2. 分别追踪以下四段数据路径：Q/K SMEM → S TMEM、S TMEM → P TMEM、P TMEM + V SMEM → O TMEM，以及 O TMEM → O GMEM。对每一段列出执行角色、源和目标存储位置、tile primitive 与硬件路径，并指出其中哪些步骤在前面的 GEMM kernel 中不存在。
 3. 根据 fp16 view 中的 column $c$ 对应物理 32-bit column $\lfloor c/2\rfloor$，推导 `S0`、`S1`、`P0`、`P1`、`O0` 和 `O1` 的物理 column ranges。哪些 regions 会发生重叠？哪些 waits 或 barriers 能防止重叠区域被过早读取或覆盖？
 4. 追踪一个 K/V block 依次经过 `s_ready`、`p_o_rescale`、`p_ready_2` 和 `o_ready` 的过程。对每个 barrier，说明谁执行 wait、谁贡献 arrivals，以及随后哪块 tile 可以安全使用。为什么 `p_o_rescale` 需要等待 256 次 arrivals？将 `P` 按 96 columns 和 32 columns 分两段交给 PV MMA，又获得了什么重叠机会？
-5. Driver warpgroup 将每个 thread 的 register 上限降到 48，两个 softmax warpgroups 将上限提高到 200，WG2 则使用 64。计算四个 128-thread warpgroups 的 register 总预算，再与 CTA 中所有 threads 都使用 200 个 registers 的情况比较。Softmax 角色为什么需要最大的配额？降低 WG3 的上限又如何使这项分配成为可能？
+5. 负责发起 TMA 和 MMA 指令的 WG3 将每个 thread 的 register 上限降到 48，两个 softmax warpgroups WG0/WG1 将上限提高到 200，WG2 则使用 64。计算四个 128-thread warpgroups 的 register 总预算，再与 CTA 中所有 threads 都使用 200 个 registers 的情况比较。Softmax 角色为什么需要最大的配额？降低 WG3 的上限又如何使这项分配成为可能？
 6. Kernel 已经将自然指数改写为 base-2 `exp2`，为什么 hardware exponential path 仍可能成为 softmax 的瓶颈？说明将元素分配给硬件 `exp2` 和基于 FMA 的三次多项式近似后，执行单元的利用方式发生了什么变化，以及哪些 online-softmax 公式保持不变。
 7. 设 `SEQ_LEN_Q=6`、`SEQ_LEN_KV=8`，并采用右对齐 causal mask。Query positions 0 和 5 分别可以访问到哪个最大 key index？若 `BLK_N=4`，它们各自需要处理哪些完整、部分有效或完全跳过的 K/V blocks？这会怎样影响 causal tasks 的工作量和调度顺序？
 8. 设 `num_qo_heads=32`、`num_kv_heads=8`、`BLK_M=128`。求 `GQA_RATIO` 和 `SEQ_Q_PER_TILE`；当 `kv_head_idx=3` 时，分别将 packed rows 0、5 和 127 映射到 `(sequence offset, query head)`，并说明为什么这 128 行可以共享同一份 K/V tile。