FlashAttention 的分块大小到底怎么选？一文拆透 Tiling 策略*

徐安安ye

86人浏览 · 2026-05-20 21:41:41

徐安安ye · 2026-05-20 21:41:41 发布

所有讲 FlashAttention 的文章都会提“分块计算”四个字，但从来没人说清楚：块到底要多大？为什么是 128 不是 64？为什么 NVIDIA 的 block size 和昇腾 NPU 的不一样？今天把这个问题彻底拆透。

一、先建立直觉：分块就像“炒菜分批次”

回到那个炒菜的比喻。
标准注意力是：所有菜切好，全放冰箱（显存），要用了再全拿出来。——冰箱往返次数太多，灶台（算力）闲着。
FlashAttention 的思路是：一次只炒一小锅，炒完立刻装盘（输出），中间不进冰箱。
这个“一小锅”有多大，就是 block size。

block size 太大 → 锅太大，灶台放不下（UB 溢出）
block size 太小 → 锅太小，要炒很多次，开销（kernel launch）太大

目标：让锅刚好占满灶台，又不溢出。

二、Block Size 受什么约束？

FlashAttention 的 block size 不是随便设的，受三个硬件约束：

2.1 约束一：UB（片上存储）大小
这是硬约束，直接决定 block size 的上限。
以昇腾 NPU（Atlas 800）为例：

每个 AI Core 的 UB = 1-2MB
block size = 128 时，Q/K/V 三个矩阵各占： $head_dim × 2 字节 = 128 × 128 × 2 = 32 KB 128 \times \text{head\_dim} \times 2\text{字节} = 128 \times 128 \times 2 = 32\text{KB}$
三个加起来 ~96KB，远小于 1MB ——所以 128 是安全的

如果 head_dim = 512（某些大模型用大 head_dim）：

单个矩阵： $128 \times 512 \times 2 = 128\text{KB}$
三个加起来 ~384KB，还是小于 1MB

但如果 block size = 256，head_dim = 512：

单个矩阵： $256 \times 512 \times 2 = 256\text{KB}$
三个加起来 ~768KB，加上 Softmax 的中间变量，接近 1MB 上限——可能有风险

2.2 约束二：计算粒度（Thread Block / AI Core 利用率）
block size 太小，会导致并行度不够。

以 NVIDIA A100 为例：
- SM 数量：108 个
- 每个 SM 可以同时跑多个 thread block
- 如果 block size = 32（太小），需要很多 block 才能喂饱所有 SM，kernel launch 开销变大。
以昇腾 NPU 为例：
- AI Core 数量：30 个
- 每个 AI Core 可以同时跑多个 block
- block size = 64 时，30 个 AI Core 可能跑不满；block size = 128 或 256 更能喂饱 AI Core。

2.3 约束三：内存访问效率（Memory Access Pattern）
block size 会影响显存访问的合并度（coalescing）。

如果 block size 是 2 的幂次（64/128/256），显存访问更容易合并，有效带宽更高。
如果 block size = 100（不是 2 的幂），显存访问会碎片化，实际带宽打折扣。
所以几乎所有实现都用 64/128/256 这几个值。

三、为什么 NVIDIA 和昇腾 NPU 的 block size 不一样？

这是很多人困惑的地方：同一套算法，为什么不同硬件上 block size 不一样？
答案：UB 大小不一样，计算粒度不一样。

3.1 NVIDIA GPU（以 A100 为例）

L1 缓存（SRAM）：256KB / SM
实际给 FlashAttention 用的：~100-150KB（因为还要存其他东西）
block size 上限（head_dim=128）：128-256
通常选 block size = 128，因为：
1. 128 足够喂饱 SM（108 个 SM，每个跑多个 block）。
2. 256 会让 SRAM 压力大，可能 spill 到 L2（更慢）。
3. 128 是 NVIDIA 官方实现的标准值（FlashAttention 原论文用 128）。

3.2 昇腾 NPU（Atlas 800）

UB（Unified Buffer）：1-2MB / AI Core（比 GPU 的 L1 大 4-8 倍）
block size 上限（head_dim=128）：256-512
通常选 block size = 256，因为：
1. UB 够大，256 不会溢出。
2. 256 能更好喂饱 AI Core（只有 30 个，比 GPU 的 108 个少很多）。
3. 更大的 block 减少 kernel launch 次数，开销更低。

3.3 一句话总结

硬件	UB/SRAM 大小	AI Core/SM 数量	推荐 block size
NVIDIA A100	~256KB/SM	108	128
昇腾 NPU Atlas 800	~1-2MB/Core	30	256

GPU 用更小的 block 来喂饱更多 SM；NPU 用更大的 block 来减少 kernel launch 开销（因为 Core 少）。

四、Head Dimension 的影响：为什么 128 和 512 的 block size 不一样？

head_dim（每个注意力头的维度）是另一个重要变量。
FlashAttention 的 block size 指的是序列维度的分块大小（N 维度），但 head_dim 会影响每个 block 占多少显存。

4.1 head_dim = 128（Llama2-7B/70B 用这个）

Q 的一个 block： $128 \times 128 \times 2 = 32\text{KB}$
三个矩阵（Q/K/V）：~96KB
UB 占用很低，可以放心用 block_size = 256。

4.2 head_dim = 512（某些大模型用更大的 head）

Q 的一个 block： $128 \times 512 \times 2 = 128\text{KB}$
三个矩阵：~384KB
UB 占用中等，block_size = 256 可能有风险（~768KB + Softmax 中间变量）。

4.3 head_dim = 1024（超大规模模型）

Q 的一个 block： $128 \times 1024 \times 2 = 256\text{KB}$
三个矩阵：~768KB
UB 占用很高，block_size 必须降到 64 甚至 32。

规律：head_dim 越大，block_size 要越小——否则 UB 放不下。

五、CANN 8.0 的自适应 Tiling 策略

手动调 block size 太麻烦了。CANN 8.0 的 FlashAttention 实现（在 ops-transformer 仓库里）用了自适应 tiling：

运行时检测 head_dim、序列长度、UB 大小。
自动选最优 block size（64/128/256 之一）。
不需要用户手动配置。

但如果你想手动调（比如做性能实验），可以通过环境变量覆盖：

# 强制 block size = 256
export ASCEND_FA_BLOCK_SIZE=256

# 强制 block size = 128
export ASCEND_FA_BLOCK_SIZE=128

踩坑：block size 设太大导致 UB 溢出，会报 UB overflow 错误。设太小导致性能差，不会报错但会慢。

六、怎么验证你的 block size 是不是最优的？

有一个实用方法：测不同 block size 的吞吐，画图。

import torch
import time
import os

def benchmark_fa(model, seq_len, block_size):
    # 设置环境变量
    os.environ["ASCEND_FA_BLOCK_SIZE"] = str(block_size)
    torch.npu.reset_peak_memory_stats()
    start = time.time()
    # 跑推理 (伪代码，需替换为实际模型调用)
    # outputs = model.generate(input_ids, max_new_tokens=50)
    # 这里仅示意时间计算
    elapsed = time.time() - start
    return elapsed

# 试不同的 block size
for bs in [64, 128, 256]:
    t = benchmark_fa(model, seq_len=4096, block_size=bs)
    print(f"block_size={bs}, time={t:.2f}s")

在昇腾 NPU 上，通常的规律是：

head_dim=128，seq=2048：128 和 256 差不多，256 略好。
head_dim=128，seq=8192：256 明显好于 128（因为长序列下 kernel launch 开销占比更大）。
head_dim=512，seq=4096：128 好于 256（因为 UB 压力更大）。

七、ops-transformer 里的实现细节

如果你想深入看实现，ops-transformer 仓库里的关键文件：

ops-transformer/
  ├── fa_tiling.py          # Tiling 策略（block size 计算逻辑）
  ├── fa_kernel.cpp         # Kernel 实现（分块计算核心）
  ├── fa_pipeline.py        # 流水线调度（Cube/Vector 协同）
  └── test_fa_tiling.py    # Tiling 策略单元测试

重点关注 fa_tiling.py 里的 calculate_block_size() 函数——它实现了上面说的所有约束（UB 大小、head_dim、AI Core 数量）。

仓库地址（纯文本，直接粘浏览器打开）：
https://atomgit.com/cann/ops-transformer