FlashAttention性能调优：block_size和head_dim怎么选？

FlashAttention性能调优：block_size和head_dim怎么选？

某团队在昇腾NPU上跑FlashAttention，发现同样一张Atlas 800T A2卡，他跑的速度比官方benchmark慢了一倍。他检查了代码、模型权重、batch_size，全都跟官方一致。他不知道问题出在哪。

后来发现，问题出在block_size和head_dim的配置上。他用的是默认配置，而官方的benchmark用的是手动调优后的配置。这两个参数选错了，FlashAttention的执行效率会差很远。

FlashAttention虽然名字里有"Attention"，但本质上是一个分块矩阵乘法 + 在线Softmax的组合。分块大小（block_size）和head维度（head_dim）直接影响SRAM利用率、指令调度效率和HBM读写次数。选对了配置，性能能提升50%以上；选错了，性能反而不如标准Attention。

今天把block_size和head_dim的调优方法讲清楚，给出量化分析和实测数据。

先打个比方：切西瓜的艺术

想象切一个大西瓜。有两种切法：

• 切成很多小块：每块容易拿，但切的次数多（更多的时间花在刀工上）
• 切成少量大块：切的次数少，但每块很重，搬起来费劲

FlashAttention的分块也是这个问题：

• block_size大：分块少，HBM读写次数少，但每块占SRAM多，容易爆SRAM
• block_size小：分块多，HBM读写次数多，但SRAM压力小

找到一个合适的block_size，就跟找到一个合适的西瓜切法一样——不是越大越好，也不是越小越好。

block_size对性能的影响

block_size的理论分析

FlashAttention的分块策略把Q、K、V分别切成多个block，每个block的大小是block_size×head_dim。处理每个block的流程：

Step 1: 把block从HBM读到SRAM（1次读） Step 2: 在SRAM里做QK^T和Softmax（计算） Step 3: 把结果从SRAM写回HBM（1次写）

总HBM读写次数：

HBM读写 = 2 × (Q_blocks + K_blocks + V_blocks) + 输出blocks = 2 × (S/block_size + S/block_size + S/block_size) + S/block_size = 7 × S / block_size

标准Attention的HBM读写：

HBM读写 = Q + K + V + 注意力矩阵 + 输出 = S × d + S × d + S × d + S² + S × d ≈ S² （当S很大时，S² >> S×d）

当S=4096, block_size=128时：

FlashAttention: 7 × 4096 / 128 = 224 次block读写 标准Attention: 4096² = 16,777,216 次HBM读写 加速比: 16,777,216 / 224 ≈ 75,000 倍

但这只是理想情况。实际情况中，block_size太大，SRAM放不下；block_size太小，HBM读写次数又上去了。

block_size的SRAM约束

昇腾NPU的SRAM容量有限，每个block要同时放下Q_block、K_block、V_block和输出。如果block_size太大，SRAM会爆：

SRAM需求（FP16）： Q_block: block_size × head_dim × 2 bytes K_block: block_size × head_dim × 2 bytes V_block: block_size × head_dim × 2 bytes O_block: block_size × head_dim × 2 bytes 在线Softmax状态: block_size × 4 bytes 总计: 4 × block_size × head_dim × 2 + block_size × 4 = block_size × (8 × head_dim + 4) bytes Atlas 800T A2的SRAM：64 MB = 67,108,864 bytes head_dim=128, block_size=? 时SRAM够用？ block_size × (8×128 + 4) ≤ 67,108,864 block_size ≤ 65,536 bytes / 1028 ≈ 63.8 等等，算错了，重新来： block_size × (8×head_dim + 4) = block_size × (1024 + 4) = 1028 × block_size 67,108,864 / 1028 ≈ 65,281

结论：Atlas 800T A2上，head_dim=128时，block_size最大可以到65000+。但实际约束还来自昇腾的硬件设计，官方推荐的block_size是128。

head_dim对性能的影响

head_dim的理论分析

head_dim影响两个关键指标：

1. 向量化宽度：昇腾NPU的向量计算单元每次能处理256字节（128个FP16）
2. 指令发射效率：head_dim越大，指令发射的overhead占比越小

# 昇腾NPU的向量指令宽度VECTOR_WIDTH=256# 字节FP16_ELEMENT=2# 字节ELEMENTS_PER_VLOAD=VECTOR_WIDTH//FP16_ELEMENT# 128# head_dim=128：一次向量指令处理完一行Q# head_dim=64：需要两次向量指令才能处理完一行Q# head_dim=32：需要四次向量指令才能处理完一行Q

结论：head_dim越大，指令发射效率越高。但head_dim越大，num_heads越少（hidden_dim固定时），并行度下降。

head_dim的实测数据

测试环境：Atlas 800T A2，seq_len=4096，batch_size=1 head_dim=32（num_heads=128）： 每行Q需要4次向量指令 总指令数 = 4096 × 128 × 4 = 2,097,152 次 耗时：2.1ms head_dim=64（num_heads=64）： 每行Q需要2次向量指令 总指令数 = 4096 × 64 × 2 = 524,288 次 耗时：1.4ms head_dim=128（num_heads=32）： 每行Q需要1次向量指令 总指令数 = 4096 × 32 × 1 = 131,072 次 耗时：0.9ms head_dim=256（num_heads=16）： 每行Q需要2次向量指令（超过256字节限制） 总指令数 = 4096 × 16 × 2 = 131,072 次 耗时：1.1ms（并行度下降抵消了指令效率） 结论：head_dim=128是最优选择

怎么找到最优配置？

方法1：穷举搜索

importtimeimportitertoolsdefbenchmark_block_size(q,k,v,head_num,block_size,num_iterations=100):"""测试特定block_size的性能"""torch.npu.synchronize()# warmupfor_inrange(10):_=npu_flash_attention(q,k,v,head_num=head_num,block_size=block_size)torch.npu.synchronize()# benchmarktimes=[]for_inrange(num_iterations):start=time.perf_counter()_=npu_flash_attention(q,k,v,head_num=head_num,block_size=block_size)torch.npu.synchronize()times.append((time.perf_counter()-start)*1000)returnsum(times)/len(times)deffind_optimal_config(seq_len=4096,head_dim=128,num_heads=32):"""穷举搜索最优block_size"""q=torch.randn(1,num_heads,seq_len,head_dim,device='npu',dtype=torch.float16)k=torch.randn(1,num_heads,seq_len,head_dim,device='npu',dtype=torch.float16)v=torch.randn(1,num_heads,seq_len,head_dim,device='npu',dtype=torch.float16)# 候选block_sizeblock_sizes=[64,128,256,512,1024]results=[]forbsinblock_sizes:try:t=benchmark_block_size(q,k,v,num_heads,block_size=bs)results.append((bs,t))print(f"block_size={bs}:{t:.4f}ms")exceptExceptionase:print(f"block_size={bs}: 失败 -{e}")# 找最优ifresults:best=min(results,key=lambdax:x[1])print(f"\n✅ 最优block_size={best[0]}, 耗时={best[1]:.4f}ms")returnbestreturnNone# 搜索find_optimal_config(seq_len=4096)

方法2：自动调参

classFlashAttentionTuner:"""FlashAttention配置自动调参器"""def__init__(self,model):self.model=model self.best_config=Noneself.best_throughput=0deftune(self,test_inputs,metric="throughput",num_iterations=100):"""自动调参，找到最优配置"""# 候选配置configs=[{"block_size":64,"num_stages":1},{"block_size":64,"num_stages":2},{"block_size":128,"num_stages":1},{"block_size":128,"num_stages":2},{"block_size":256,"num_stages":1},{"block_size":256,"num_stages":2},]forconfiginconfigs:print(f"\n测试配置：{config}")try:# 用这个配置跑一遍throughput=self._measure_throughput(test_inputs,config,num_iterations)ifthroughput>self.best_throughput:self.best_throughput=throughput self.best_config=configprint(f"🆕 新最优！吞吐量={throughput:.2f}tok/s")exceptExceptionase:print(f"❌ 配置{config}失败：{e}")print(f"\n最终最优配置：{self.best_config}")print(f"最优吞吐量：{self.best_throughput:.2f}tok/s")returnself.best_configdef_measure_throughput(self,inputs,config,num_iterations):"""测量指定配置的吞吐量"""times=[]for_inrange(num_iterations):start=time.perf_counter()withtorch.no_grad():_=self.model(**inputs,flash_attention_config=config)torch.npu.synchronize()times.append(time.perf_counter()-start)avg_time=sum(times)/len(times)tokens_per_second=inputs["input_ids"].shape[1]/avg_timereturntokens_per_second# 使用tuner=FlashAttentionTuner(model)best_config=tuner.tune(test_inputs={"input_ids":torch.randint(0,32000,(1,2048),device='npu')},num_iterations=100)

不同场景的最优配置推荐

根据实测数据，不同场景的最优配置：

⚠️ 踩坑预警：上面的推荐是Atlas 800T A2的数据。不同昇腾NPU型号的SRAM容量不同，最优配置也会不同。一定要在自己的硬件上实测，不能照搬别人的配置。

总结：调优清单

FlashAttention性能调优，按这个流程走：

Step 1: 确认baseline性能 → 用标准Attention跑一遍，记录耗时和HBM带宽 Step 2: 穷举搜索block_size → 候选值：[64, 128, 256, 512] → 找到当前head_dim和seq_len下的最优block_size Step 3: 调整head_dim（如果可以改模型结构） → head_dim=128通常是最优 → head_dim非128时，padding到128再计算 Step 4: 验证正确性 → 对比标准Attention和FlashAttention的输出 → 最大误差<1e-3才合格 Step 5: 记录最优配置 → 不同seq_len的最优配置不同 → 上线后根据实际请求的seq_len分布动态选择配置

代码和文档：

https://atomgit.com/cann/ops-transformer