<aside>

本文由人类编写

</aside>

0. 引言

在即将推出的 FA4 中，FA 的作者在其 main 分支上，基于 cutlass 的示例代码，实现了一版高性能的 FMHA 算子。代码可见具体链接。这里我们带着上文的理解，开始对比 FA4 的实现中，做了哪些改动，实现了哪些新的特性？概括地说，主要包含了以下的改动和特性。

<aside>

  1. 功能覆盖对比
  2. 流水线与同步
  3. 调度器
  4. 张量布局与读写

总得来说，由于 FA 是一个完整的库，所以不仅支持了 MHA，还支持了 GQA、MQA 两种模式，还有 Deepseek 需要的变体，以及一些推理优化的特性，比如支持 page-attention，滑动窗口、learnable_sink等特性。下面会以 DeepSeek-V3 系列的 MLA 的计算逻辑进行分析。并按照存和算两大部分进行拆解，并最后将其合在一起，看看是如何进行排布的。

1 存储分析

第一步，我们先不看细节性的代码，先看看代码里，是如何分配 SMEM，TMEM，RMEM 的？

SMEM：

@cute.struct
class SharedStorage:
# m_barriers for pipelines
mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mbar_total] # 15 * 8B = 120B
# Tmem holding buffer
tmem_holding_buf: Int32
# Smem tensors
# store row max and row sum
sScale: cute.struct.MemRange[Float32, self.q_stage * self.m_block_size * 2] # row_max 和 row_sum ，128 * 2 * 2 * 4B = 512B = 0.5 KB
sO: cute.struct.Align[
cute.struct.MemRange[self.o_dtype, sO_size], # 大小为 0 ，O 放在 TMEM 上
self.buffer_align_bytes, # 1024b = 128B
]
sQ: cute.struct.Align[
cute.struct.MemRange[self.q_dtype, cute.cosize(sQ_layout)], # 128 * 192  * 2B * 2 = 96 KB
self.buffer_align_bytes,
]
sK: cute.struct.Align[
# cute.cosize(sK_layout) is correct even in the case of self.uneven_kv_smem
cute.struct.MemRange[self.k_dtype, cute.cosize(sK_layout)], # 128 * 192 * 2B * 2 + 128 * 128 * 2B = 128 KB
self.buffer_align_bytes,
]

可以看到，SMEM 已经用的比较满了，共计用了 224.5 KB 的大小，很接近 228KB 的上限了。这也算是一个 SMEM 瓶颈。SMEM主要包含：两片 Q 的切片，三份 KV 的切片（因为大小受限，这里是 2 块 192 \* 128, 一块 128 \* 128）。且需要注意，这里 Large+Small+Large 的拆分，中间的 Small并不是一成不变，因为不是一直都是 Large+Small+Large 的排列，而是会随着流水线的进行，变成 Small+Large+Small 的组合。下面 load 阶段会着重分析。

TMEM：

self.tmem_s_offset = [0, self.n_block_size] # e.g., 0, 128， S 分两块，每块 128 列
self.tmem_o_offset = [
self.tmem_s_offset[-1] + self.n_block_size + i * self.head_dim_v_padded
for i in range(self.q_stage)
] # e.g., 256, 384，O 也分两块，每块 128 列
self.tmem_total = self.tmem_o_offset[-1] + self.head_dim_v_padded # 计算 O的最后一个分片的尾部
assert self.tmem_total <= SM100_TMEM_CAPACITY_COLUMNS
self.tmem_s_to_p_offset = self.n_block_size // 2
self.tmem_p_offset = [
self.tmem_s_offset[i] + self.tmem_s_to_p_offset for i in range(2)
] # 0, 128

可以看到，这里 TMEM 的分配就回到了 C 风格，主要还是 TMEM 是一个新出的硬件，NV 官方支持不够，需要用手动偏移值来区分分块。可以看到，这里 TMEM 也用满了。S 占两片\[128 \* 128], O 占两片 \[128 * head_dim_v_padded]。在 Deepseek 情况下，这里已经占满了全部 512 列的存储，没有任何空间浪费。同时这里的 P 会复用 S 的空间，且复用是按 S 的一半来复用的。

RMEM：

if self.head_dim_padded < 96:
	self.num_regs_softmax = 200
	self.num_regs_correction = 64
	self.num_regs_other = 48
else:
	# self.num_regs_softmax = 192 if self.is_causal or self.is_local else 184
	self.num_regs_softmax = 200
	# self.num_regs_softmax = 176
	# self.num_regs_correction = 96
	# self.num_regs_correction = 80
	# self.num_regs_correction = 64 if self.is_causal or self.is_local else 80
	self.num_regs_correction = 64
	# self.num_regs_other = 32
	# self.num_regs_other = 64
	# self.num_regs_other = 80
	self.num_regs_other = 48
	# self.num_regs_other = 96 if self.is_causal or self.is_local else 80
	# self.num_regs_other = 64 if self.is_causal or self.is_local else 80
	self.num_regs_empty = 24

可以看到，在 Deepseek 例子下，会走 else 分支，softmax 会使用 200 个寄存器，correction 使用 64 个，other 占用 48，empty 占用 24。比如 Softmax，这里分配的是一个线程占有 200 个寄存器，这 200 是怎么来的呢？前文说到了，一个 softmax 处理一个中间结果 S，S 的shape 是 \[128, 128]， 那么这里 4 个warp 处理，正好一个线程处理一行，也就是 128 个数，占据 128 个寄存器；同时，P 的写回是S 的一半，所以寄存器需要 64 个。加起来是 192，同时还需要一些辅助中间值，比如 row_max 等，故而申请到 200 个。同时比如 correction 这个 WG，启动 4 个 warp，来处理一个 O 的结果，所以一个线程处理 128 个数，且 O 是 BF16 格式，那么就使用 64 个寄存器就 OK。同样地，我们来分析 others，主要压力来自于 MMA，每次 UMMA，处理一个 \[128,32] * \[32, 128]的矩阵乘，所以分配到每个线程来看，矩阵 A 需要16 个，B 需要 16 个，总体需要 32 个寄存器，加上一些其他变量，48 是比较安全的分配值。一点疑惑：empty 为啥要分 24，从代码上没看出来为啥。

2 Warp 分析

结合上文的存储分析后，我们对数据流的来处和去处都有了一个比较清晰的认识了：