第 13 讲 · Pallas 与 Splash Attention · Pallas Kernels and Splash Attention

对应 Maps to：CMU 11-868 / 11-968 Large Language Model Systems（Spring 2026，Lei Li），Lecture 13 · Wed 2/25 · Guest Lecture by Srinath Mandalapu（Google）: Pallas and Splash Attention（slides：llmsys-13-pallas_splash_attention.pdf，111 页）。 阅读材料：Pallas 文档（docs.jax.dev 的 “Pallas: a JAX kernel language”、“Pallas TPU: pipelining”、“Pallas TPU: matmul”）；Dao et al. (2022) FlashAttention；Dao (2023) FlashAttention-2；Shah et al. (2024) FlashAttention-3；Tokamax 仓库里的 Splash Attention 实现。 依据说明：本讲按嘉宾 slides 的四部分组织：动机（XLA 融合的 HBM 瓶颈与 VPU 瓶颈、Pallas 在训练 / 推理中的用途）、Pallas 基础（存储空间、抽象、矩阵加法、VMEM OOM、流水线、BlockSpec 与 grid、pallas_call 的循环语义、输出别名、两级分块、MLIR、VPU 流水线、分块矩阵乘的 1D / 2D / 3D 网格、算术强度与 roofline、块大小调优的实测、编译路径、融合激活、手动 DMA 与显式流水线）、Flash Attention（Tokamax、注意力 FLOPs、XLA 融合在大规模下的实测、在线 softmax、延迟归一化、MHA / GQA、网格与 MFU、Pacchetto 的 trace 分析、块大小与微分块、MXU 的零填充与 MXU / VALU 流水线、Max Logit Estimate）、Splash Attention（掩码与稀疏、MaskInfo、segment ID 与联合掩码、Splash 的 pallas_call、索引映射、残差输出、内核结构、soft cap、DeepSeek-V3 的配置、优化总结与练习）。Pallas 风格的模拟器、流水线时间线模型、对实测的反推、Flash / Splash 内核的模拟实现与正确性检验、掩码的块分类与 L = 32768 的统计、块大小的玩具模型、MLE 安全范围的实验，都是本讲义自己设计的，不是 Pallas / Tokamax 的真实代码；本机没有 JAX / TPU，所有”实测”数字都来自 slides。 配套代码 Code：code/lecture13/（5 个脚本 + pallas_sim.py；都在 2 秒以内） 预计阅读时间：3–4 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲，你应该能够：

1. 说出什么时候需要绕过编译器、手写内核：XLA 的融合在注意力上仍要把 L × L 的矩阵写回 HBM、让 MXU 等着 VPU。
2. 用 Pallas 的三件套描述一个内核：grid（迭代空间）、BlockSpec（每个网格点把哪一块搬进 VMEM）、内核（只看见 VMEM 里的块）；知道 pallas_call 等价于一组嵌套循环，并会估算 VMEM 用量、避免 OOM。
3. 解释双缓冲的流水线为什么能把搬运藏进计算里，什么时候藏不住；知道手写 DMA、更深的流水线解决什么问题。
4. 为分块矩阵乘法选网格和块大小：3D 网格、K 放在最后一维、输出块驻留 VMEM；用”每步算术强度 b/2”和”每步固定开销”解释块大小的影响。
5. 推导在线 softmax 与两段合并的公式，写出 Flash Attention 的内核主体（m、l、o 三个累加器、α 缩放、延迟归一化），估算它相对 XLA 写法省下的 HBM 流量；说出 Max Logit Estimate 的前提。
6. 把一个注意力掩码预处理成 Splash Attention 的稀疏执行表（block_mask、active_rows/cols、partial_mask_blocks、mask_next），并分析因果、局部窗口、文档掩码及其组合能跳过多少计算。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  动机：XLA 的融合还差在哪里；Pallas 内核在训练、推理里的用处
§2  Pallas 基础：存储空间、pallas_call、grid 与 BlockSpec、别名、两级分块           ← 代码 01
§3  流水线：双缓冲；手写 DMA；计算受限与访存受限                                  ← 代码 02
§4  分块矩阵乘法：1D / 2D / 3D 网格；算术强度；块大小调优；融合激活；编译路径       ← 代码 03
§5  Flash Attention：在线 softmax；延迟归一化；块大小与 VMEM；trace 分析；MLE         ← 代码 04
§6  Splash Attention：稀疏执行表；segment ID 与联合掩码；块大小的取舍               ← 代码 05
§7 小结   §8 常见误区   §9 练习   §10 延伸阅读   §11 下一讲预告

1. 动机：为什么要手写内核 · Motivation

1.1 XLA 的融合还差在哪里

第 12 讲看到，XLA 把注意力编译成三个融合块，而且懂得重算 exp。slides 第 5 页指出它在规模变大时的问题：

• HBM 瓶颈：softmax 要整行的信息，标准的融合仍然要把大的分数矩阵写回 HBM 再读出来。
• 向量单元的瓶颈：exp、求和在 VPU 上跑的时候，MXU 闲着。
• 而 Pallas 让你直接控制迭代空间和 HBM ↔ VMEM 的数据搬运，可以用”在线 softmax”让数据一直留在片上；你管”什么时候算什么”（时间上的调度），后端管”数据怎样摆”（T(8,128) 的空间布局）。Pallas 内核以自定义调用的形式绕过 HLO，经 MLIR 直接降到 LLO，编译器不会把你手工安排的顺序重新打乱。

slides 第 58 页给了大规模下的实测（128 个头、序列 4096、q/k 头维度 192、v 128）：XLA 的三个融合块各要 3.34、3.29、3.68 ms，HBM 带宽利用率 76%、71%、69%，而 FLOP 利用率只有 24%、0.19%、22%——典型的访存受限。中间的分数矩阵 [128, 4096, 4096] 是 f32，8 GiB，被写一次、读两次。

1.2 Pallas 内核用在哪里

slides 第 10–11 页：

• 训练：Splash Attention（融合 + 跳过掩码里的空块，§6）；Ring Attention（把 K、V 块在设备之间”传圈”，处理一台机器放不下的超长序列）；MoE 的 ragged all-to-all（不同设备之间数量不等的 token 交换）和分组矩阵乘 GMM（一次算多个专家，不用填充、不用丢 token；第 17 讲）。
• 推理：Paged Attention 把 KV 缓存切成固定大小的物理页（第 24 讲），逻辑上连续的序列在物理上可以不连续。XLA 处理不了这种”查页表再取数”的间接访问，要靠 Pallas 内核手动做指针间接。

2. Pallas 基础 · Pallas fundamentals

2.1 存储空间与最简单的内核

slides 第 16 页的存储空间：ANY（通常就是 HBM，大张量的默认位置）、VMEM（向量单元旁边的片上 SRAM，小而快）、SMEM（标量单元的片上 SRAM，放标量和控制流的值）、SEMAPHORE（很小的同步用存储）。

最简单的内核（slides 第 18 页）：

def add_matrices_kernel(x_vmem_ref, y_vmem_ref, z_vmem_ref):
    z_vmem_ref[:, :] = x_vmem_ref[:, :] + y_vmem_ref[:, :]   # 从 VMEM 读进寄存器、相加、写回 VMEM

z = pl.pallas_call(add_matrices_kernel, out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype))(x, y)

pallas_call 先在 VMEM 里给 x、y、z 各分配一块、把 x、y 从 HBM 拷进来，调用内核，再把 z 拷回 HBM。内核本身只看见 VMEM。问题是 VMEM 很小：01 号脚本第 1 部分，两个 2048 × 2048 的 float32（各 16 MiB）加上输出共 48 MiB，超过了每个核心 32 MiB 的默认上限——和 slides 第 19 页的真实报错 “Scoped allocation with size 48.01M and limit 32.00M” 一样。处理 LLM 规模的张量必须分块。

2.2 grid、BlockSpec 与 pallas_call

slides 第 22–25 页的四个部件：

pallas_call 的语义就是一组嵌套循环（slides 第 25 页）：

for i in range(grid[0]):
  for j in range(grid[1]):
    in_refs  = 按 in_specs 的 index_map(i, j) 把输入块从 HBM 搬进 VMEM
    out_refs = 输出块在 VMEM 里的缓冲区
    kernel(*in_refs, *out_refs)
    按 out_specs 的 index_map(i, j) 把输出块写回 HBM

而且它会自动把下一块的搬运和这一块的计算重叠（§3）。01 号脚本第 2 部分（f32[2048, 2048] 相加）：

HBM 的读写总量和不分块时一样（每个元素读一次、写一次），但 VMEM 只需放几块。内核代码一个字都没改——“怎么切、怎么搬”全写在 BlockSpec 里。

index_map 决定搬什么（01 号脚本第 3 部分）：pallas_call 遵循两条规则——输入块的下标和上一步相同就不重新搬；输出块的下标和上一步相同就一直留在 VMEM 里、不写回。例如给 f32[1024, 1024] 加一个按行广播的 bias（BlockSpec 的 index_map 是 (0, j)）：列块 j 在内层循环时，bias 每一步都要搬，DMA 共 192 次；把 j 放到外层，同一块 bias 被连续 8 步共用，DMA 降到 136 次。矩阵乘法把归约维 K 放在网格的最后一维，靠的就是第二条规则（§4）。

输出别名（slides 第 26 页）：input_output_aliases={2: 0} 让输出直接复用输入的 HBM 缓冲区，内核原地更新（如 out[...] = out[...] + x[...]），不用新分配显存、也省掉一次拷贝。01 号脚本第 4 部分验证了输出和输入是同一块内存、原来的值被覆盖。优化器更新参数、KV 缓存追加 token 都是这种模式。

2.3 两级分块

slides 第 27–29 页：

• 宏分块（Pallas / 用户）：管 HBM ↔ VMEM。块要大到能藏住 DMA 的延迟、摊薄固定开销，又要小到放得进 VMEM。
• 微分块（编译器 / 硬件）：管 VMEM ↔ 向量寄存器。编译器把宏块拆成 8 × 128 的硬件块，生成一串 vector load、ALU、store 指令。256 × 256 的 float32 块就是 64 个向量寄存器块（bf16 是 32 个）。

Pallas 经由 MLIR（TPU 上的方言叫 Mosaic）降到硬件指令；数据按 T(8,128) 布局；不整齐的边缘自动补 0，保证每一块大小相同、内核里不需要”最后一块”的分支（01 号脚本第 5 部分：f32[1000, 1000] 用 256 的块，计算了 1,048,576 个元素，有效的占 95.4%，结果正确）。

VPU 内部同样是流水线（slides 第 29 页）：取下一块（VLOAD）、算这一块（VALU）、存上一块（VSTORE）同时进行。逐元素加法的算术强度只有 1 FLOP / 12 字节，trace 里 VLOAD / VSTORE 100% 忙碌、VALU 只有 30%–50%——它受限于 VMEM 的带宽。

3. 流水线：让搬运和计算重叠 · Pipelining

3.1 双缓冲

slides 第 20–21 页的例子：给 x₀、y₀ 和 x₁、y₁ 各分配一份 VMEM；把 x₀、y₀ 拷进来；开始拷 x₁、y₁ 的同时，把 x₀、y₀ 读进寄存器、相加、存回；开始把 z₀ 拷回 HBM 的同时，等 x₁、y₁ 拷完、再算……02 号脚本第 1 部分模拟了”每块搬 1、算 2”的 6 块（字符甘特图）：

单缓冲（总时间 18）：
  DMA  |000      111      222      333      444      555      |
  计算 |   000000   111111   222222   333333   444444   555555|
双缓冲（总时间 13）：
  DMA  |00001111    22222    3333    4444    55555            |
  计算 |    00000000111111111222222223333333344444444455555555|

双缓冲时只有第 0 块的搬运露在外面：总时间 = 1 + 6 × 2。一般地，双缓冲时每块的时间 ≈ max(搬运, 计算)——流水线只能藏住较短的那个。

3.2 手写 DMA

pallas_call 自动做双缓冲；要更精细的控制，就自己写 DMA（slides 第 48–52 页）：pallas_call 把输入声明成 HBM 引用（memory_space=ANY）、在 scratch_shapes 里分配两份 VMEM 缓冲区、用 SMEM 里的标量记录当前用哪一份；内核里用 pl.make_async_copy(...) 发起拷贝、.start() / .wait() 控制时机，用 pl.ds(start, size) 手动切片；pltpu.emit_pipeline 则把流水线的编排搬进内核内部。好处：可以更早发起下一块的拷贝、用三个或四个槽位的环形缓冲区、精确控制块在 VMEM 里的驻留。

slides 第 50 页的反面教材：copy.start() 后面紧跟 copy.wait()——没有重叠。在 (4096, 7168) @ (7168, 18432)、1024³ 的块上，MXU 利用率只有 23.4%。

02 号脚本第 2 部分用理想模型算这个形状：每步搬 x、y 两块 bf16[1024, 1024]（4 MiB），按峰值带宽约 1.14 μs；计算 2·1024³ FLOP 约 2.09 μs。理想模型下单缓冲也有 64.7%，双缓冲 99.9%——和实测的 23.4% 差很远。反推一下：23.4% 意味着每步的搬运实际要 6.83 μs（有效带宽约 614 GB/s，峰值的 17%）。真实的单次 DMA 有启动延迟、同步开销，远达不到峰值带宽；这时搬运比计算还慢，单引擎模型下双缓冲的上限只有 30.5%。（slides 用 pallas_call 自动流水线实测到 62.70%，说明真实系统里同时有多个 DMA 在飞，单引擎的模型太保守。）理想模型告诉你上限在哪里，实测与上限的差距告诉你该去 trace 里找什么。

3.3 访存受限时、搬运有波动时

02 号脚本第 3 部分：矩阵乘 1024³ 的块计算受限（每块搬运 1.14 μs < 计算 2.09 μs），双缓冲后每块 2.1 μs；256³ 的块和逐元素加法都是访存受限的，双缓冲只能把计算藏进搬运里。对访存受限的运算，唯一的办法是少搬：融合进别的内核，或用更大的块提高复用。

第 4 部分：每块计算 1.0、搬运平均 0.8 但随机波动（标准差 1.0）时，2、3、4 个缓冲区的利用率是 76.0%、86.7%、91.6%——多一个缓冲区就能提前多搬一块，用平时的富余去填偶尔的慢。代价是 VMEM。

4. 分块矩阵乘法 · Blockwise matmul in Pallas

4.1 三种网格

slides 第 30–37 页。NumPy 版的分块矩阵乘是三重循环：外层 M 块、中层 N 块、内层 K 块，z[m, n] += x[m, k] @ y[k, n]；关键是”当前的 Z 块在整个累加阶段留在快速存储里”。Pallas 有三种写法：

03 号脚本第 1 部分：三种写法在小矩阵上结果都对；放大到 slides 的形状 (4096, 7168) @ (7168, 18432)（bf16 输入、f32 累加、块边长 1024），1D 要 664 MiB、2D 要 64 MiB，都 OOM，3D 只要 16 MiB。

3D 网格的内核（slides 第 44 页）：

def matmul_kernel(x_ref, y_ref, z_ref):
    @pl.when(pl.program_id(2) == 0)
    def _():
        z_ref[...] = jnp.zeros_like(z_ref)
    z_ref[...] += x_ref[...] @ y_ref[...]

pl.pallas_call(matmul_kernel, grid=(m // bm, n // bn, k // bk),
    in_specs=[pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)), pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))],
    out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),
    compiler_params=pltpu.TPUCompilerParams(dimension_semantics=("parallel", "parallel", "arbitrary")))

输出块的下标 (i, j) 在整个 K 循环里不变，所以它一直留在 VMEM 里累加：k == 0 时清零，K 循环结束才写回 HBM 一次。M、N 两维互不依赖，标为 parallel，可以分给不同的核心（slides 第 36 页的 MegaCore）；K 维有累加依赖，标为 arbitrary。编译路径（slides 第 44 页）：JAX + Pallas 内核 → HLO 里的 tpu_custom_call（内核体以序列化的 Mosaic MLIR 嵌在 backend_config 里）→ Mosaic 方言 → LLO。

4.2 算术强度与块大小

slides 第 38–39 页：方阵乘法 [M, M] @ [M, M] 的 FLOPs 是 2M³，最少要搬 3M² 个数，float32 时强度是 M/6 FLOP/字节；Ironwood 每个张量核心的平衡点是 1028.75 TFLOP/s ÷ 3685.5 GB/s ≈ 279，所以 M > 1674 才计算受限（bf16 时 M > 837）。

但分块以后，每个输入块要被搬很多次，真正的强度由块大小决定：每一步算 2b³ FLOP、搬两块 bf16（4b² 字节），强度 b/2。03 号脚本第 2 部分在 slides 的形状上扫了块大小（总计算量 1.082 TFLOP）：

块边长翻倍，网格步数变成 1/8、每步强度翻倍：512³ 的块（256 < 279）即使完美流水线也访存受限，1024³ 变成计算受限，2.75 倍的提速大半来自这里。把实测与理想时间之差摊到每一步上，分别约 0.86 和 1.25 μs——同一个量级，说明每次内核调用都有约 1 μs 的固定成本。slides 的结论：“调优的首要目标是把网格做得尽量小”，即在 VMEM 放得下的前提下找最大的块。作为参照，XLA 自己的 jnp.dot 在这个形状上是 1.45 ms、72.75%（slides 第 40 页）——手写内核要调好块大小才追得上编译器。

4.3 融合激活函数

slides 第 46–47 页：在最后一个 K 步里，对 VMEM 中累加好的输出块原地做 ReLU，再写回——激活函数不需要额外的 HBM 往返。03 号脚本第 4 部分：小例子里 HBM 读写从 2304 + 512 KiB 降到 2048 + 256 KiB；在 slides 的形状上省下一整个 bf16[4096, 18432] 的读 + 写（302 MB，约 82 μs，相对 1.05 ms 的矩阵乘约 8%）。XLA 会自动做这种输出融合，手写内核要自己写进去。

5. Flash Attention 的 Pallas 写法 · Flash Attention

5.1 在线 softmax

slides 第 60–64 页。Flash Attention 的核心是不让 L × L 的矩阵出现在 HBM 里：把 Q、K、V 切成块，对每个 Q 块依次扫过所有 K、V 块，块的计算全在 VMEM 里完成。难点在 softmax：它要整行的最大值和指数和。在线 softmax 只维护两个量：到目前为止的最大值 m、以 m 为基准的指数和 l。每来一块：

$m_{\text{new}} = \max(m, \max(\text{块})),\quad l \leftarrow l\cdot e^{m - m_{\text{new}}} + \sum e^{\text{块} - m_{\text{new}}}$

04 号脚本第 1 部分把 x = (1, 3, 2, 5, 0, 4) 每两个一块扫一遍：(m, l) 依次是 (3, 1.1353)、(5, 1.2034)、(5, 1.5781)，最后的 softmax 与直接计算的差为 0。更一般地，两段各自算出 (m, l, o) 后可以合并：

$m = \max(m_1, m_2),\ l = l_1 e^{m_1-m} + l_2 e^{m_2-m},\ o = \frac{o_1 e^{m_1-m} + o_2 e^{m_2-m}}{l}$

（脚本里与整行一起算的差 2.2e-16。）合并规则与顺序无关：一行可以按块顺序扫（Flash Attention），也可以分给不同的设备各算一段再合并（Ring Attention、推理时的 split-KV）。

5.2 内核主体与延迟归一化

slides 第 66、105 页的内核主体（一个网格点 = 一个头、一个 Q 块、一个 KV 块）：

qk = lax.dot_general(q, k, ...)                 # 1. 分数 [bq, bkv]
qk = apply_mask_and_soft_cap(qk)                # 2. 掩码、软截断
m_next = jnp.maximum(m_prev, qk.max(axis=-1))   # 3. 更新最大值
s_curr = jnp.exp(qk - m_next)                   # 4. 分子
alpha  = jnp.exp(m_prev - m_next)               # 5. 旧结果的缩放
l_next = s_curr.sum(axis=-1) + alpha * l_prev
o_scratch[:] = alpha * o_scratch[:] + s_curr @ v   # 6–7. 加权和，缩放旧的累加值
# 最后一个 KV 块：o = o_scratch / l（延迟归一化，只除一次）

m、l、o 三个累加器放在 VMEM 的 scratch 里，在整个 KV 循环中驻留；除法只在最后做一次（延迟归一化），每一块只用乘法和加法。04 号脚本第 2 部分在模拟器里按 grid = (头, Q 块, KV 块) 运行这个内核：与 float64 参考的差 2.0e-07（无掩码）和 5.4e-07（因果；掩码用全局位置在内核里现算）。

省了多少流量（同一部分）：在 slides 的规模上，XLA 的写法要把 f32 分数矩阵写一次、读两次，共 24.8 GiB，按 HBM 带宽至少 7.21 ms（slides 实测三块合计 10.31 ms）；Flash Attention 只读写 Q、K、V、O（K、V 对每个 Q 块读一遍，br = 1024 时读 4 遍），1.69 GiB，下限 0.49 ms，而计算量 1.3744 TFLOP 按峰值要 1.34 ms——从访存受限变成了计算受限。（slides 第 71 页把总量写成”1373.4 TFLOPs”，按公式应为 1373.4 GFLOP；它在 4 ms 内跑完，343 TFLOP/s，MFU 33%。）

slides 第 67–68 页还提到，多头注意力的各个头互相独立，是网格的第一维；GQA 让多个查询头共用一组 K、V，减少 KV 缓存和带宽（第 7 讲）。

5.3 块大小、VMEM 与 trace 分析

slides 第 72–83 页用 Pacchetto（把 LLO 指令包画成 Perfetto 时间线的工具）分析这个内核：MXU 时有空闲，向量寄存器 100% 占满，频繁出现 VSTORE.SPILL（寄存器放不下、溢出到 VMEM），VLOAD / VSTORE 断断续续，softmax 的向量运算没能藏在 MXU 的时间里。

调大块（slides 第 77、85 页，Q / K [128, 4096, 192]、V [128, 4096, 128]）：

04 号脚本第 3 部分估算了 VMEM 里都放了什么：输入块（双缓冲）只有几 MiB，大头是 [br, bc] 的 f32 分数和它的 exp：1024 × 2048 时 16 MiB，2048 × 2048 时 32 MiB，4096 × 4096 时 128 MiB——超过 VMEM 的容量。增大 br 让每个 Q 块对 K、V 的复用翻倍、网格步数减半，所以 1024 → 2048 快了 19%（slides：说明之前被 Q 块的大小限制住了）。

几个更细的手段：

• 微分块（bkv_compute，slides 第 82–83 页）：KV 块仍按 bc 搬进来，但在内核里分成更窄的几段算分数，中间结果从 [br, bc] 缩到 [br, bkv_compute]（4096 × 4096 时从 128 降到 32 MiB），目标是留在寄存器里、消除 spill。
• MXU 的零填充（slides 第 79 页）：头维度 192 不是 256 的倍数，LLO 把它补到 256：装 Kᵀ 只需约 48 次显式的 vmatpush，剩下四分之一由 vmatprep 隐式清零，省下 16 次推送。两个 256 × 256 的 MXU 按 Q 的行各分一半（上下 512 行）。
• MXU / VALU 流水线（slides 第 80 页）：VALU 直接消费 MXU 弹出的分数块算 max、exp，算好的 s 打包后直接推给 MXU 做 s·V——寄存器之间交接，不经过 VMEM。slides 第 81 页的 named scope 显示，如果中间的 qkᵀ 矩阵太大、必须整块写进 VMEM 再读回来，它就成了一道隐式的栅栏，三个阶段无法重叠。

5.4 Max Logit Estimate

slides 第 84–85 页：用一个事先确定的常数代替每块算出的最大值 m_next。省掉了 vmax 归约和 jnp.maximum，α 恒为 1、旧的累加值不用重新缩放——2048 × 2048 的块从 3.36 ms 降到 2.64 ms（50.27%）。为什么这样不会算错？因为

$\frac{e^{x_i - c}}{\sum_j e^{x_j - c}}$

里的 c 被约掉了：c 不必是真正的最大值，只要 exp 既不溢出、也不全部下溢。04 号脚本第 4 部分（因果注意力，分数的真实最大值 25.1）：

float32 的 exp 在 88.7 以上溢出、在 −87.3 以下进入次正规数。所以”正确的常数”要满足：c 不能比真实最大值小 88 以上，也不能比某一行的最大值大出约 87 以上。怎样保证？一个办法是 soft cap（slides 第 106 页）：logits = cap · tanh(qk / cap)，把分数限制在 (−cap, cap) 里，取 c = cap 就永远不会溢出。

6. Splash Attention：稀疏的 Flash Attention · Splash Attention

6.1 掩码就是稀疏

slides 第 89 页：因果掩码、局部窗口掩码让注意力变成稀疏的运算——把 [L, L] 按块切开，有的块全空（不用算）、有的全满（不用加掩码）、有的部分（要逐元素加掩码）。05 号脚本第 1 部分画出了 slides 的小例子（长度 8、块 2 × 2）：

P · · ·
F P · ·
F F P ·
F F F P

“Splash = Sparse + Flash”：先把掩码预处理成稀疏执行表，内核只在活跃块上运行。

6.2 稀疏执行表（MaskInfo）

slides 第 90–91 页：[4096, 4096] 的因果掩码、块 [1024, 2048] → 4 × 2 的块网格。05 号脚本逐项复现了 slides 的数组：

部分块有 4 个，图样只有 2 种：(0, 0) 和 (2, 1) 的 Q、KV 起点之差相同，图样一样；(1, 0) 和 (3, 1) 也是。这些表算一次、放在 SMEM 里，内核开始前预取（slides 第 98 页：num_scalar_prefetch=6）。

slides 第 104 页还有两个数组：bounds_start（这一块是不是所在 Q 行的第一块，要初始化累加器）= [T T T F T F]、bounds_end（是不是最后一块，要归一化并写回）= [T T F T F T]。

6.3 Splash 的 pallas_call

slides 第 97–103 页：

• 网格 = (查询头数, num_active_blocks)：第二维不再是”KV 块号”，而是”第几个活跃块”。
• 索引映射：index_map 收到头号 h、线性的网格下标 g 和预取的数组，用 rows[g]、cols[g] 查出 Q 块、KV 块的位置——Q 只用行号，K、V 只用列号；部分块的掩码用 mask_next[g] 选图样。
• 输入（VMEM）：q、k、v、segment ID、partial_mask_blocks、q_sequence（Q 块里每个 token 的全局位置，用于现算掩码）。
• 输出：注意力输出 O，以及反向要用的两个残差——每行的 logsumexp（L）和最大值（M）。
• 内核结构：bounds_start 时初始化 O、L、M 的 scratch；累加循环里按 block_mask 决定要不要加掩码；bounds_end 时归一化、写回。

05 号脚本第 2 部分在模拟器里运行这个内核（2 个头、L = 512、块 64 × 64）：

与 float64 的稠密掩码注意力的差都在 4e-07 左右。

6.4 segment ID 与联合掩码

slides 第 95–96 页：训练时常把多个短样本打包成一条长序列，用 segment ID（每个 token 属于第几个样本）防止样本之间互相看到——只允许 q_id == kv_id，形成块对角的稀疏结构。多个掩码用逻辑 AND 组合（因果 ∧ 局部 ∧ 文档），约束越多、空块越多，最后烘焙进同一张 MaskInfo。

05 号脚本第 3 部分在 L = 32768 上统计（文档长度随机、平均约 1820）：

因果最多省一半；局部窗口和文档掩码把有效区域压成对角线附近的一条带，能省 6–12 倍的计算。块越小越贴合掩码，但部分块和不同的图样也越多。

6.5 块大小的取舍

slides 第 94 页：大块（如 2048）才能喂饱 MXU，但会在因果掩码里算很多被掩掉的”填充”；小块（如 512）跳过得更彻底，却可能让 MFU 崩掉。05 号脚本第 4 部分用一个玩具模型（每个活跃块 = 固定开销 + 计算量 ÷（峰值 × 效率），效率 = min(1, b / 512)）扫了一遍：每步固定开销 1 μs（§4.2 反推的量级）时，L = 2048 和 8192 都是 2048 的块最快；固定开销 0.1 μs 时，L = 2048 的最佳块变成 1024；为 0 时变成 512。最佳点取决于每一步的固定成本和硬件的真实效率曲线，只能靠实测——这正是 slides 说”较短序列上的因果掩码”难以两全的原因。

6.6 其他优化与练习

slides 第 107–110 页：

• DeepSeek-V3 的配置（128 个查询头、q/k 头维度 192、v 128、序列 4096）用 bq = bkv = 2048，block_mask 只有 4 项。
• 向量对齐：m、l 这两个”每行一个数”的累加器被补成 (bq, 128)，和 128 个 lane 对齐，把许多次慢的标量读变成一次向量读。
• trace（bq = bkv = bkv_compute = 1024）仍显示 MXU 有空隙、寄存器压力导致 spill。
• slides 的练习：bq = bkv = 2048 配 bkv_compute = 512；用 max_logit_const（如 70）开启 MLE；use_base2_exp=True 用 exp2 代替 exp（第 12 讲 §4.6）；不存 partial_mask_blocks、在内核里用位置现算掩码；在 16K、32K、128K 的序列上验证扩展性。

05 号脚本第 5 部分算了最后那条练习的账（L = 16384、块 1024）：因果掩码的部分块图样只有 1 种（1 MiB），局部窗口 2 种（2 MiB），文档掩码 28 种（28 MiB，几乎每个部分块都是新图样）。现算只需用 q_sequence、segment ID 做逐元素比较——不占 HBM 流量和 VMEM，在 VPU 上很便宜。又一次”用便宜的计算换昂贵的搬运”。

7. 小结 · Summary

1. 编译器的融合在注意力这类运算上仍有极限：L × L 的中间结果要回 HBM，VPU 的工作和 MXU 的工作不能重叠。Pallas 让你直接控制迭代空间和 HBM ↔ VMEM 的搬运。
2. Pallas = grid + BlockSpec + 内核；pallas_call 等价于嵌套循环：按 index_map 搬入块、调用内核、写回块；相同的输入块不重复搬，相同的输出块驻留 VMEM。VMEM 很小（scoped 默认 32 MiB），块大小要按 VMEM 估算。
3. 双缓冲让每块的时间 ≈ max(搬运, 计算)；访存受限时只能少搬。真实 DMA 远达不到峰值带宽，需要更深的流水线或更大的块。
4. 分块矩阵乘用 3D 网格、K 在最后一维、输出块驻留累加；每步算术强度 b/2；每次内核调用约 1 μs 的固定成本——在 VMEM 放得下的前提下块越大越好。
5. Flash Attention = 分块 + 在线 softmax（m、l、o 三个累加器、α 缩放）+ 延迟归一化；在 slides 的规模上把 HBM 流量从 24.8 GiB 降到 1.69 GiB，从访存受限变成计算受限。VMEM 的大头是分数块的中间结果，微分块缓解寄存器压力；MLE 用常数代替最大值，前提是常数在安全范围内（soft cap 可以保证）。
6. Splash Attention 把掩码预处理成稀疏执行表，网格只遍历活跃块；因果省一半，局部窗口、文档掩码及其组合能省更多；块大小在”贴合掩码”和”喂饱硬件”之间取舍；部分块的掩码可以现算而不必存储。

8. 常见误区 · Common pitfalls

• “pallas_call 会自动处理任何大小的数组。” 不给 BlockSpec 时整个数组都要放进 VMEM，2048 × 2048 的 float32 相加就会 OOM。
• “分块会增加 HBM 流量。” 逐元素运算分块后每个元素仍只读写一次；矩阵乘法分块后输入会被重复搬运，块越小重复越多。
• “网格的顺序无所谓。” 顺序决定哪些块能被相邻的步骤复用（bias 的例子：192 → 136 次 DMA），也决定输出块能不能驻留累加（K 必须在最后一维）。
• “双缓冲总能把搬运藏起来。” 只在计算比搬运慢时成立；访存受限时藏住的是计算。
• “块越小越灵活、越好。” 每次调用都有固定成本，块小时算术强度也低；在 VMEM 允许的范围内通常越大越好。
• “在线 softmax 是近似的。” 它和一次性计算在数学上完全相同，只是计算顺序不同。
• “Max Logit Estimate 需要知道真正的最大值。” 任何常数都行，只要 exp 不溢出、不全部下溢；soft cap 可以让常数总是安全的。
• “因果掩码能省一半，所以 Splash 只比 Flash 快一倍。” 局部窗口、文档掩码及其组合能跳过 90% 以上的块。
• “部分块的掩码必须存下来。” 可以用位置和 segment ID 在内核里现算，省掉存储和搬运。

9. 练习 · Exercises

1. 用 pallas_call 计算 z = x + y，x、y 是 f32[4096, 4096]。不分块需要多少 VMEM？用 (512, 512) 的块、双缓冲需要多少？网格多大？
2. 一个 bias 加法的输入 x: [M, N]、bias: [1, N]，块 (bm, bn)。网格写成 (M/bm, N/bn) 和 (N/bn, M/bm) 时，bias 各被搬运多少次？
3. 每块搬运时间 d、计算时间 c，共 n 块。写出单缓冲和双缓冲的总时间。d = 3、c = 1、n = 100 时各是多少？
4. 为什么 3D 网格矩阵乘要把 K 放在网格的最后一维？如果把 K 放在第一维，输出块会怎样？
5. 块边长 b 的 3D 矩阵乘（bf16 输入），每一步的算术强度是多少？Ironwood 的平衡点 279，b 至少多大才计算受限？
6. 实测 512³ 的块 4.63 ms、4032 步，理想时间 1.15 ms。如果每步固定开销是 0.86 μs，1024³ 的块（504 步、理想 1.05 ms）预计多少？和实测 1.68 ms 比如何？
7. 用在线 softmax 处理 x = (2, 0, 4, 1)，每块 2 个。写出每一步的 (m, l)，以及最后的 softmax。
8. 证明两段在线 softmax 的合并公式与整行一起算的结果相同。
9. Flash Attention 的 HBM 读写：H 个头、长度 L、头维度 d（q、k）和 dv（v），bf16 输入、f32 输出，Q 块 br。写出读写字节数的公式；H = 128、L = 4096、d = 192、dv = 128、br = 2048 时是多少 GiB？
10. 分数的最大值是 40、最小值是 −30。用 Max Logit Estimate 时，常数 c 的安全范围大约是多少（float32）？
11. 因果掩码 [4096, 4096]、块 [1024, 1024]：有多少个活跃块？其中多少全满、多少部分？部分块有几种图样？
12. 打包了 3 个文档（长度 1000、2000、1096）的序列，块 1024 × 1024、只用文档掩码（不加因果）：哪些块是全空、全满、部分？

参考答案

1. 不分块：3 × 64 MiB = 192 MiB。块 (512, 512)：每块 1 MiB，三个数组各两份 = 6 MiB；网格 (8, 8)。

2. (M/bm, N/bn)（列块在内层）：bias 每一步都换块，搬 (M/bm)·(N/bn) 次；(N/bn, M/bm)（列块在外层）：只在列块变化时搬，N/bn 次。

3. 单缓冲 n(d + c)；双缓冲 d + Σ max(·) ≈ d + (n − 1)·max(d, c) + c。d = 3、c = 1、n = 100：单缓冲 400，双缓冲 3 + 99 × 3 + 1 = 301（访存受限，藏住的是计算）。

4. 输出块的下标是 (i, j)。K 在最后一维时，同一个 (i, j) 在连续的 K 步里不变，输出块留在 VMEM 里累加、最后写回一次。K 在第一维时，相邻两步的 (i, j) 不同，输出块每一步都要写回、下次再读回来（而 pallas_call 默认不会把它读回来——结果就错了，或者得用别名手动处理），HBM 流量暴增。

5. 每步 2b³ FLOP，搬 2b² 个 bf16 = 4b² 字节，强度 b/2；b/2 > 279 ⇒ b ≥ 560（实际取 2 的幂 1024）。

6. 1.05 ms + 504 × 0.86 μs ≈ 1.48 ms，实测 1.68 ms：每步的固定开销并不完全相同（1024³ 的块每步多 1.25 μs），大块的每一步本身也更长、流水线的启动 / 收尾占比不同。

7. 块 1 (2, 0)：m = 2，l = e⁰ + e⁻² = 1.1353；块 2 (4, 1)：m = 4，l = 1.1353·e⁻² + e⁰ + e⁻³ = 0.1536 + 1 + 0.0498 = 1.2034。softmax = exp(x − 4) / 1.2034 = (0.1125, 0.0152, 0.8310, 0.0414)。

8. 设整行的最大值 m = max(m₁, m₂)。l₁e^(m₁−m) = Σ_{段1} e^(x−m₁)·e^(m₁−m) = Σ_{段1} e^(x−m)，同理段 2，相加就是整行的 Σ e^(x−m)；o 同理（每项乘 vᵢ）。所以合并结果与整行一起算相同。

9. 读 Q 一次：2HLd；K、V 每个 Q 块读一遍：(L/br)·2HL(d + dv)；写 O：4HL·dv。代入：Q 192 MiB、K 192 MiB × 2、V 128 MiB × 2、O 256 MiB，共 192 + 384 + 256 + 256 = 1088 MiB ≈ 1.06 GiB（br = 1024 时 1.69 GiB）。

10. 不溢出：40 − c < 88.7 ⇒ c > −48.7；不让最大值较小的行全部下溢：粗略地 c − 行最大值 < 87。若每行的最大值都不小于 −30，则 c < 57。安全范围约 (−48, 57)，取 c = 40 附近最稳妥。

11. 4 × 4 = 16 个块，对角线以上 6 个全空 → 10 个活跃块：对角线上 4 个部分、对角线以下 6 个全满；4 个对角块的 Q、KV 起点之差都是 0，只有 1 种图样。

12. 文档边界在 1000、3000。块 0 = [0, 1024) 含文档 1、2；块 1 = [1024, 2048) 只含文档 2；块 2 = [2048, 3072) 含文档 2、3；块 3 = [3072, 4096) 只含文档 3。只要两个块有共同的文档就不是全空；两边都只含同一个文档才是全满：

              KV 块 0  1  2  3
    Q 块 0         P  P  P  ·
    Q 块 1         P  F  P  ·
    Q 块 2         P  P  P  P
    Q 块 3         ·  ·  P  F

    全空 4 个（文档 1 / 2 与文档 3 互不相见），全满 2 个，部分 10 个。

10. 延伸阅读 · Further reading

• Pallas 文档（docs.jax.dev）：“Pallas: a JAX kernel language”、“Pallas TPU: pipelining”、“Pallas TPU: matmul”、“Pallas TPU: distributed computing”。
• Dao et al. (2022) FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness；Dao (2023) FlashAttention-2；Shah et al. (2024) FlashAttention-3（第 21 讲）。
• Milakov & Gimelshein (2018) Online normalizer calculation for softmax：在线 softmax 的来源。
• Tokamax（github.com/openxla/tokamax）：Splash Attention 等 Pallas 内核的实现；vLLM 的 tpu-inference 里的 ragged paged attention 内核。
• Liu et al. (2023) Ring Attention with Blockwise Transformers for Near-Infinite Context。
• Gemma 2 技术报告：注意力 logits 的 soft cap。

11. 下一讲预告 · Next lecture

第 10–13 讲都在榨干一块芯片。可是第 1 讲的那道总题目——训练一个千亿参数的模型——一块芯片无论如何也放不下、算不完。第 14 讲进入分布式训练：为什么要多卡、并行有哪几个维度（数据、流水线、张量、专家）；GPU 之间怎样通信（NCCL；broadcast、reduce、all-reduce、reduce-scatter、all-gather 这些集合通信原语，第 12 讲 §5 已经见过）；为什么环形 all-reduce 的通信量是 2(N − 1)/N、和设备数几乎无关；参数服务器和 all-reduce 的对比；以及最常用的数据并行：每张卡算一部分数据的梯度、all-reduce 求平均——它和单卡上的大批量训练为什么等价。我们会在一台机器上模拟多张卡，逐步追踪环形 all-reduce，并验证数据并行与单卡训练逐位相同。