 的选择与 Cube Bound 达成分析)
CANN ops-transformer FlashAttn 性能建模D256 下基本块 (M, N) 的选择与 Cube Bound 达成分析【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本文基于 CANN ops-transformer 仓库中 flash_attn_perf_analysis.md 的技术文档结合 FlashAttn 算子的 Host 侧 tiling 与 Kernel 侧源码推导 D256Q/K/V 头维 256非量化 fp16/bf16dV D条件下 Flash Attention 基本块 (M, N) 的选择依据。读者读完后将掌握如何以 Cube Bound 为建模目标推导性能下限与空间上限如何在 L0C/L1/UB 的容量约束与 MTE2/Fixpipe 的带宽约束之间找到可行域以及为什么 bmm1/bmm2 结果缓冲共用能从根本上扩大 D256 的可行解空间。1. 建模目标Cube BoundFlash Attention 在 Ascend 上由多条引擎流水组成CubeMmadbmm1Q·Kᵀ与 bmm2P·V两个矩阵乘MTEQ/K/V 从 GM 搬入 L1、L1 到 L0A/L0B 的分形搬运Fixpipebmm 结果从 L0C 搬出到 UBVectorsoftmax行 max/exp/sum与输出的在线累加。这些引擎中只有 Cube 计算不可压缩——矩阵乘的 FMA 数由问题规模唯一确定搬入、搬出与 softmax 都可以与其他阶段重叠执行。因此 kernel 的性能上界就是 Cube 满流水建模目标为Cube Bound$$ T_{cube} \ge T_{mte2}, \qquad T_{cube} \ge T_{fixpipe}, \qquad T_{cube} \ge T_{vec} $$即 Cube 永远有数据可算其余所有阶段的时间都被 Cube 时间完全遮挡。其中 Vector 侧 softmax/累加运行在与 Cube 配对的 Vector 核上工作量与 M(ND) 成正比在实用块大小范围内远小于 $T_{cube}$不构成额外约束本文重点分析最紧的MTE2搬入与Fixpipe搬出两个条件。从实现看FlashAttn 的 kernel 入口 flash_attn.cpp 以KERNEL_TYPE_MIX_AIC_1_2模式将同一份源码同时编译为 Cube 侧AIC与 Vector 侧AIV两个变体每个 AIC 与两个 AIV 配对CV_RATIO2——这与本文一个 Cube 核配两个 Vector 核、Vector 侧缓冲按 M/2 行计的模型设定完全一致。2. 模型设定2.1 基本块与迭代结构MQ 序列方向每次迭代、每个 Cube 核承担的块行数bmm 的 M 维NKV 序列方向的块大小bmm 的 N 维一次迭代处理 N 行 KVD头维bmm 的 K 维 / bmm2 的输出列数非量化时 dV D。单次外层迭代处理一个 KV 块BMM1: S_Q[M,D] × S_K[D,N] → scores[M,N] softmaxonline跨迭代维护行 max/sum BMM2: P[M,N] × V[N,D] → 部分输出[M,D]与历史值在线累加架构上一个 Cube 核与两个 Vector 核配对1:2Cube 一次算 M 行两个 Vector 核各负责 M/2 行的 softmax 与累加Vector 侧缓冲按 M/2 行计。2.2 硬件规格单核参数取值说明CoreNum32核数Frequency1.6 GHz主频L1512 KBQ/P/KV 分片驻留L0A / L0B64 KBbmm 左/右矩阵分形缓冲各 2×32K 子缓冲L0C256 KBbmm 结果累加/输出缓冲多份轮转UB248 KB可用bmm 结果落点与 Vector 侧全部缓冲L2128 MB / 5.2 TB/s折合单核 ~101 B/cycleHBM128 GB / 1.6 TB/s折合单核 ~31 B/cycle单核带宽换算总带宽 ÷ 核数 ÷ 主频。L25.2×10¹² ÷ (32 × 1.6×10⁹) ≈ 101 B/cycleHBM1.6×10¹² ÷ (32 × 1.6×10⁹) ≈ 31 B/cycle。Cube 单周期完成一次16×16×16mma吞吐FMA 4096 FMA/cycle模型假设。源码佐证容量数据硬件表中的 UB 可用 248 KB 与 Kernel 侧约束完全对应——flash_attn_block_vec_nd.h 中UbLayout结构体以static_assert(sizeof(UbLayout) 248KB)限制总占用注释明确说明UB 总 size 为 256KB但需要为 AscendC 和 VF 预留 8K业务只能使用 248K。L0C 的 4×64KB 轮转结构、L0A/L0B 各 2×32KB 双缓冲也都能在 flash_attn_block_cube_nd.h 的 buffer 分配表中找到一一对应的实现见 §4。3. Cube Bound 的达成条件性能侧约束3.1 时间模型按单次迭代总量计。两个 matmul 的 FMA 总量为2·M·N·D$$ T_{cube} \frac{2 M N D}{FMA} $$Q 块M·D仅在首次迭代搬入一次S2 足够大时被摊薄K/V 每次迭代搬入2·N·D$$ T_{mte2} \approx \frac{2 N D \cdot \text{sizeof}(kv)}{Mix_{bw}}, \qquad Mix_{bw} h \cdot L2_{bw} (1-h) \cdot HBM_{bw} $$其中 $h$ 为 L2 命中率$Mix_{bw} \in [31,\ 101]$ B/cycle。bmm1/bmm2 结果fp32由 L0C fixpipe 到 UB$$ T_{fixpipe} \frac{(M N M D) \cdot \text{sizeof}(fp32)}{UB_{bw}} $$3.2 条件一M 的下限对 MTE2 的遮挡$T_{cube} T_{mte2}$ 两侧约去 N、D与 D、N 无关$$ M \frac{2 \cdot FMA}{Mix_{bw}} \frac{8192}{Mix_{bw}} $$Mix_bwM 下限101L2 全命中M 81.1即M ≥ 8231全 HBM 缺失M 264.3即M ≥ 265即M 82 时无论 L2 命中多好都必然 mte2 boundM ∈ [82, 265) 时需 L2 命中率满足$$ h \ge \frac{8192/M - 31}{101 - 31} $$M 越大对 h 的要求越低M96 需 78%M128 需 47%M ≥ 265 起无要求。3.3 条件二N 的下限对 Fixpipe 的遮挡$T_{cube} T_{fixpipe}$ 约去 M与 M 无关$$ UB_{bw} 2 \cdot FMA \cdot \frac{ND}{N D} 8192 \cdot \left(\frac{1}{D} \frac{1}{N}\right) \quad\Longleftrightarrow\quad N \frac{8192}{UB_{bw} - 8192/D} $$D256 时UB_bw 32 8192/NN128 需 UB_bw 96N256 需 UB_bw 64。N 越大bmm2 输出M×D与 N 无关被越多 KV 行摊薄fixpipe 压力越小——该条件方向上偏好大 N。UB_bw 未列入硬件表见 §8 待补充若其与 L2 带宽同级~101则 N ≥ 119 才能满足N128 可以满足。源码佐证Fixpipe 压力方向D256 下 fixpipe 输出量大Kernel 侧在 flash_attn_block_cube_nd.h 中通过IterateBmm2l0Split将 BMM2 沿 dV 方向按 128 列分块nLoops循环、splitN逐 tile 占用轮转的 64KB L0C buffer 并独立 Fixpipe——单个 L0C buffer 无需容纳整块128×256×4B结果这正是 §4.1 中bmm2 结果可沿 dV 方向分块适配 buffer 容量的实现落点。3.4 数值示例M128, N128, D256项数值周期cube2×128×128×256 8.39M FMA2048KV 搬入131072 Bh100%: 1298 / h47%: 2048临界/ h0: 4228fixpipe 搬出(128²128×256)×4 196608 BUB_bw96: 2048临界/ 128: 1536注意D256 时 fixpipe 搬出量196.6K/迭代大于 KV 搬入量131.1K/迭代fixpipe 是仅次于 MTE2 的第二约束。对照 (M64, N256)T_cube同为 2048M·N 相同但 KV 搬入 262144 B即使 h100% 也需 2596 cycles 2048——M64 82必然 mte2 bound的直接体现。4. 空间约束容量侧约束各级存储需容纳基本块的工作集并以多缓冲轮转维持流水重叠。以下按bmm1 结果缓冲与 bmm2 结果缓冲各自独立分配的常规设计推导§7 将放宽该假设。4.1 L0Cbmm 结果 bufferL0C256K以 4 份轮转 buffer 容纳 bmm1 结果M·N·4与 bmm2 结果M·D·4每份 64Kbmm2 结果可沿 dV 方向分块适配 buffer 容量分块数 $t \lceil M D \cdot 4 / 64K \rceil$为保 Cube→Fixpipe 重叠每迭代消耗的 buffer 数1 t须小于总 buffer 数 4。由此$$ M N \cdot 4 \le 64K ;\Rightarrow; \mathbf{M N \le 16384}; \qquad t \le 2 ;\Rightarrow; M D \le 32K ;\Rightarrow; \mathbf{M \le 128}\ (D{}256) $$dV 分块不改变 Mmad 总量矩阵乘按分块累加代价是 L0C 轮转深度从 2.0t1降到 4/3t2、fixpipe 调用次数翻倍但总量不变。4.2 L1输入分片驻留分片缓冲大小Q 块双缓冲2 × M·D·2 4MDPsoftmax 结果回写供 bmm2三缓冲3 × M·N·2 6MNK/V分时复用同组 buffer至少双 buffer≥ 2 × 2 × N·D·2 4NDbuffer 份数越多流水越深$$ 4MD 6MN 4ND \le 512K $$源码佐证L1 分配策略flash_attn_block_cube_nd.h 的 L1 buffer 分配与此表逐项对应P 矩阵 3 个 buffer3×mBaseSize×s2BaseSize×sizeof(INPUT_T)Q 矩阵 2 个gS1 行内复用、isFirstS2Loop加载/isLastS2Loop释放K/V 4×64KBconfig5 即 s2BaseSize256 且 D128 时改用 2×128KB 大缓冲。P 3 个、Q 2 个、KV 至少 2 个的多缓冲轮转深度正是为了让 MTE2 与 Cube 在 PRELOAD_N2 流水flash_attn_kernel_dn.h 中创建任务与执行任务分离、执行慢于创建 2 轮下连续重叠。4.3 UB结果缓冲与 Vector 侧工作集buffer大小bmm1 结果fp32双缓冲2 × (M/2)·N·4 4MNbmm2 结果fp32双缓冲2 × (M/2)·D·4 4MDPbf16双缓冲1 行分形对齐 padding≈ 2 × (M/2)·N·2 2MN输出累加暂存fp32单缓冲(M/2)·D·4 2MD掩码bool双缓冲2 × (M/2)·N·1 MNsoftmax 标量max/sum/exp 各三缓冲 lse双缓冲 杂项≈ 54M 512$$ 7MN 6MD 54M 512 \le 248K\ (253952,B) $$4.4 L0A/L0B定性各 64K2×32K 子缓冲。bmm 输入按 16×16 分形驻留K 方向D256按 128 分块后M、N ≤ 128 时每个输入分片恰好 32K 填满一个子缓冲自然满足M 或 N 过大将迫使输入侧退化单缓冲或重复加载故 L0A/L0B 不单独给出不等式但支持 M、N ≤ 128 的上界。源码佐证分形与子缓冲Kernel 侧 L0A/L0B 由BufferManager各分配 2×32KB 双 buffer见 flash_attn_block_cube_nd.h 的l0aBufferManager_/l0bBufferManager_配合BuffersPolicyDB的 LOCK_UNLOCK 策略让 MTE1 搬运L1→L0与 PIPE_M 计算重叠。D72 复用 D128 的 config 变体时L1→L0 搬运宽度按 16 元素分形对齐72→80Mmad 按真实 D 精确累加见 README 基本块一节说明16 分形对齐是贯穿 D 各档位的实现约束。5. D256 代入(M, N) 的可行域将 D256 代入上述约束分离缓冲设计约束不等式D256N128N256条件一MTE2M 8192/Mix_bwM ≥ 82h100%~ 265h0同左条件二FixpipeN 8192/(UB_bw−32)需 UB_bw 96需 UB_bw 64L0CMN ≤ 16384M ≤ 128M ≤ 128M ≤ 64L14MD6MN4ND ≤ 512KM ≤ 219M ≤ 102UB分离缓冲7MN6MD54M512 ≤ 248KM ≤ 102M ≤ 74联合82 M ≤ 102M ≤ 64无 cube bound 窗口两点观察N256 被 L0C 锁死MN ≤ 16384使 M ≤ 64 82无论其他约束如何放松都到不了 cube bound 阈值。N128 的窗口极窄可行域 (82, 102] 紧贴阈值 82——即使取到上界 M102也要求 L2 命中率 ≥ 70%裕度极小。而块大小需按 16 分形对齐、实用档位取 2 的幂M ∈ {64, 128}M128 超出 UB 上限 102311K 248KM64 低于阈值 82——分离缓冲设计下D256 在常规档位上不存在 cube bound 可行解。退而求其次只能取 (64, 256)空间全部满足UB 212K、L1 416K但 M64 82注定 mte2 bound——这就是 D256 性能远差于 D128 的理论根源。源码佐证档位与 config 映射模型中的常规档位 M ∈ {64, 128}并非随意设定。Host 侧 fa_adjust_sinner_souter.h 的AdjustSinnerAndSouter对 D256 只有两档选择gSize × maxSeqQ ≥ 64时取 sOuter64 / sInner128否则取 sOuter32 / sInner256再结合mBaseSize sOuter × CV_RATIOCV_RATIO2即得 M ∈ {64, 128}、N ∈ {128, 256}。这两个档位在编译期模板 flash_attn_template_tiling_key.h 中编码为 config4sOuter64, sInner128, D256与 config5sOuter32, sInner256, D256并在 README 基本块表 中完整列出。也就是说M 取 64 还是 128在实现层面就等价于选 config 4 还是 config 5二者被 tiling key 的 Config 位段bits 25-28固化。6. 可行域内 M、N 对性能的影响上节给出的是不等式范围而非定值范围内如何取值取决于 M、N 各自的性能影响。6.1 M在可行域内越大越好算术强度每 KV 字节的 FMA 数为M/2M 是唯一的算术强度旋钮。M 越大条件一越宽松对 L2 命中率的要求越低M102 需 70%M128 需 47%cube bound 越稳Q 摊薄Q 搬入的摊薄负担相对 cube 时间与 M 无关两者同随 M 线性增长不构成偏好代价尾块浪费Q 侧总行数不整除 M 时末块空转与任务粒度变粗。Q 侧行数充足prefill / 长序列时代价可忽略Q 侧总行数不足 M 时大块大面积空转必须缩小 M——此时 M 反正低于 82cube bound 不可达属另一种取值逻辑见 §8。结论Q 侧行数充足时M 取可行域上界。源码佐证Q 侧行数充足的判据AdjustSinnerAndSouter中 D256 分支的判据是gSize × maxSeqQ ≥ 64——即Q 侧有效行数Q 序列长度 × GQA group 数是否够填满一个大块。这与本文Q 侧行数充足与否决定 M 取值逻辑的结论精确对应充足则选大 MsOuter64 → M128不足则降档sOuter32 → M64。此外 tiling 的 gS1 合轴gS1Size actSeqLensQ × gSize见 README正是把 Q 序列与 GQA group 两轴合并为同一任务维M 的分块实际作用于 gS1 轴。6.2 N双向影响且与 M 在 L0C 上耦合大 N 的收益条件二fixpipe更宽松KV 迭代次数S2/N减少循环与同步开销摊薄小 N 的收益UB/L0C 中的 MN 项、L1 中的 N·D 项都随 N 收缩K/V 单个 L1 分片更小N128 时 64K/片可 4 个 buffer 深轮转N256 时 128K/片仅 2 个 bufferMTE2 与 Cube 的重叠更深S2 不整除 N 时尾块浪费更小。关键耦合L0C 约束MN ≤ 16384把 M 与 N 绑在同一份预算里——N 每放大一档M 的上限就减半。条件一M ≥ 82是 cube bound 的唯一入口而条件二只是过线即可N128 在 UB_bw ≥ 96 时已满足。因此选择次序是先保 M 过阈值并尽量大N 在剩余预算内取满足条件二的值。D256 下即 N128、M 取上界。7. 缓冲冗余的发现bmm1/bmm2 结果共用后可行域扩大7.1 冗余性论证§4.3 中 bmm1 结果缓冲4MN与 bmm2 结果缓冲4MD各自独立双缓冲但两者的生命周期并不重叠bmm1(L) 结果 → softmax(L) 消费完毕 → 该 buffer 死亡 bmm2(L) 结果 → 累加(L) 消费完毕 → 该 buffer 死亡数据流是环形链bmm1 → softmax → bmm2 → 累加 → bmm1(下一轮)同一迭代内 softmax 先读完 bmm1 结果、bmm2 才开始产出bmm2 的结果要到下一轮 softmax 前才被累加读完。因此两组双 buffer 可以共用同一对轮转 bufferbuffer 容量取两者的较大值(M/2)·max(N,D)·4。D256 ≥ N 时bmm1 结果(M/2)·N·4完全装得进 bmm2 的 buffer——bmm1 的双 buffer4MN整体消失。源码佐证实现即共用这一理论上发现的冗余在 Kernel 中已是实际设计——flash_attn_block_cube_nd.h 与 flash_attn_block_vec_nd.h 中的ubMmResBuffers_明确注释BMM1/BMM2 fixpipe 结果分时复用buffer 大小取mBaseSize/2 × max(s2BaseSize, dVBaseSize) × sizeof(FP32)即本文的(M/2)·max(N,D)·4buffer 数按 D 分档UB_MM_RES_BUFCNT (dBaseSize 128) ? 2 : 4。读写两侧通过一 flag 一 buffer的核间同步接力AIV 侧CC_MM_0~3/ AIC 侧CROSSCORE_MM_0~3写端BMM1/BMM2fixpipe 完成后 Set结果就绪读端Vec1/Vec2Wait 后读取、读毕再 Setbuffer 释放同一 flag 沿「就绪 → 释放 → 就绪」交替接力。mmResBufId_单调递增计数器保证同一 ExecuteTask 内先后执行的 BMM1(L) 与 BMM2(L-2) 落在不同 buffer避免写写冲突——这正是 §7.3 所述按结果序而非算子序分配 buffer、以环形等待保证时序的工程实现。7.2 新的 UB 空间公式$$ \underbrace{4M \cdot \max(N,D)}_{\text{共用结果缓冲}} 2MN 2MD MN 54M 512 \le 248K $$D256 且 N ≤ 256 时 $\max(N,D)D$化简为$$ \mathbf{3MN 6MD 54M 512 \le 248K} \qquad (\text{原 } 7MN 6MD \dots) $$节省量 4MNN ≤ D 时。对 (M128, N128)节省 64KUB 占用从 311K超限降为247.25K ≤ 248K余 768B。7.3 扩大后的可行域与更好的选择约束N128N256UB共用缓冲M ≤ 128M ≤ 107L0C / L1M ≤ 128 / M ≤ 219M ≤ 64/ M ≤ 102联合82 M ≤ 128M ≤ 64仍无窗口与 §6 的结论呼应N128 的窗口从 (82, 102] 扩大到(82, 128]恰好把 M128 档位纳入对 L2 命中率的要求从 70% 降到 47%cube bound 从勉强贴线变为有裕度——这就是缓冲共用后能选出更好 (M, N) 的量化体现N256 侧不变被 L0C 的 MN ≤ 16K 锁死在 M ≤ 64进一步印证 §6.2 的判断L0C 耦合下 N 必须给 M 让路共用的前提是按结果序而非算子序分配 buffer并以环形等待保证bmm1 第 L 轮写入晚于bmm2 第 L−2 轮读完——即空间收益需要同步机制的配合不是免费的。源码佐证环形等待与流水Kernel 的 PRELOAD_N2 流水flash_attn_kernel_dn.h 的ExecuteTask保证同一时刻 AIC 上 BMM1(本轮) 与 BMM2(2 轮前) 并行、AIV 上 Vec1(本轮) 与 Vec2(2 轮前) 并行P 矩阵的 3 buffer 轮转AIV 写第 N 轮 P 与 AIC 读第 N-2 轮 P 并行恰好容纳写入中/待读取/已释放三态。这些都建立在结果生命周期错开这一前提之上与本文的环形链论证互相印证。8. 结论场景选择依据D256Q 侧行数充足prefillM128, N128落在共用 buffer 后的可行域 (82, 128] 上界L2 命中率 ≥ ~47% 即 cube boundbmm2 结果按 dV/128 分块适配 L0C bufferD256Q 侧总行数 128decode/小 batchM64, N256M 被工作量压低cube bound 反正不可达M 82取大 N 换取 KV 迭代次数减半且 (64,256) 恰好贴住MN 16K的 L0C/mask 界空间无浪费建模方法论小结性能条件给下限cube bound 等价于M 8192/Mix_bw82~265与 D/N 无关与N 8192/(UB_bw − 8192/D)空间约束给上限L0CMN ≤ 16384、M ≤ 128dV 分块后L14MD6MN4ND ≤ 512KUB 由结果缓冲与 Vector 工作集决定可行域内 M 越大越好唯一算术强度旋钮N 与 M 在 L0C 上共享预算、优先保 M缓冲复用扩大可行域bmm1/bmm2 结果生命周期错开、可共用 bufferUB 上界 M ≤ 102 → M ≤ 128使 cube bound 从临界变为可行——空间优化最终转化为性能选项的解放。待补充UB 带宽 UB_bw 实测值用于封闭条件二D256 下 fixpipe 搬出量大于 KV 搬入量值得单独 profiling 验证。附模型结论与实现的对应关系及验证途径本文的所有推导最终都落在 CANN ops-transformer 仓库 FlashAttn 算子的真实实现中可按下述对应关系在源码中逐一验证基本块档位M、N 的取值集合编译期 config 模板 flash_attn_template_tiling_key.hconfig 0~7 的 sOuter/sInner/D/DV 编码与运行期AdjustSinnerAndSouterfa_adjust_sinner_souter.h共同决定mBaseSize sOuter × 2与s2BaseSize sInnertiling key 编码5 个模板参数InOutLayoutType / KvLayoutType / HasAttenMask / TemplateId / Config按位拼接为 tiling keyHost tiling 的GenTilingKey生成、运行期 kernel 按 key 选变体细节见 README 的 TilingKey 章节缓冲与存储层级L0C 4×64KB、L0A/L0B 各 2×32KB、UBstatic_assert ≤ 248KB、L1 的 P/Q/KV 多缓冲见 flash_attn_block_cube_nd.h 与 flash_attn_block_vec_nd.h 的 buffer 分配表负载均衡与 FDsection 切分与 split-K 归约FlashDecode由flash_attn_metadataAICPU基于 SectionStreamK 算法产出 metadata主算子按 metadata 区间执行见 flash_attn_tiling_data.h 的 metadata 布局与 README 的负载均衡章节性能验证途径仓库提供 train/infer 各 220 性能用例tests/pytests/test_cases/performance_redline_train.py、performance_redline_infer.py可通过--perf_mode采集性能数据并配合 ASCEND DEBUG 日志中的GenTilingKey/kernel_nameFlashAttn_hash_tiling_key确认实际选用的档位——感兴趣的读者可在真机上用这些用例验证本文UB_bw 实测值待补充的开放项。此外FlashAttn 的对外接口flash_attn_metadataflash_attn两段式调用、参数说明、约束与调用示例见 torchapi_flash_attn.md可与本文的性能建模对照理解接口层的 D 支持范围64/72/128/256以及 MLA 形态的 QK192/V128与模型假设的 D256 场景一致接口文档中metadata 与主算子入参必须保持一致的约束则对应本文 §4 中tiling 与 metadata 共用AdjustSinnerAndSouter保证两侧切分口径一致的实现要求。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-transformer FlashAttn 性能建模：D=256 下基本块 (M, N) 的选择与 Cube Bound 达成分析
CANN ops-transformer FlashAttn 性能建模D256 下基本块 (M, N) 的选择与 Cube Bound 达成分析【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本文基于 CANN ops-transformer 仓库中 flash_attn_perf_analysis.md 的技术文档结合 FlashAttn 算子的 Host 侧 tiling 与 Kernel 侧源码推导 D256Q/K/V 头维 256非量化 fp16/bf16dV D条件下 Flash Attention 基本块 (M, N) 的选择依据。读者读完后将掌握如何以 Cube Bound 为建模目标推导性能下限与空间上限如何在 L0C/L1/UB 的容量约束与 MTE2/Fixpipe 的带宽约束之间找到可行域以及为什么 bmm1/bmm2 结果缓冲共用能从根本上扩大 D256 的可行解空间。1. 建模目标Cube BoundFlash Attention 在 Ascend 上由多条引擎流水组成CubeMmadbmm1Q·Kᵀ与 bmm2P·V两个矩阵乘MTEQ/K/V 从 GM 搬入 L1、L1 到 L0A/L0B 的分形搬运Fixpipebmm 结果从 L0C 搬出到 UBVectorsoftmax行 max/exp/sum与输出的在线累加。这些引擎中只有 Cube 计算不可压缩——矩阵乘的 FMA 数由问题规模唯一确定搬入、搬出与 softmax 都可以与其他阶段重叠执行。因此 kernel 的性能上界就是 Cube 满流水建模目标为Cube Bound$$ T_{cube} \ge T_{mte2}, \qquad T_{cube} \ge T_{fixpipe}, \qquad T_{cube} \ge T_{vec} $$即 Cube 永远有数据可算其余所有阶段的时间都被 Cube 时间完全遮挡。其中 Vector 侧 softmax/累加运行在与 Cube 配对的 Vector 核上工作量与 M(ND) 成正比在实用块大小范围内远小于 $T_{cube}$不构成额外约束本文重点分析最紧的MTE2搬入与Fixpipe搬出两个条件。从实现看FlashAttn 的 kernel 入口 flash_attn.cpp 以KERNEL_TYPE_MIX_AIC_1_2模式将同一份源码同时编译为 Cube 侧AIC与 Vector 侧AIV两个变体每个 AIC 与两个 AIV 配对CV_RATIO2——这与本文一个 Cube 核配两个 Vector 核、Vector 侧缓冲按 M/2 行计的模型设定完全一致。2. 模型设定2.1 基本块与迭代结构MQ 序列方向每次迭代、每个 Cube 核承担的块行数bmm 的 M 维NKV 序列方向的块大小bmm 的 N 维一次迭代处理 N 行 KVD头维bmm 的 K 维 / bmm2 的输出列数非量化时 dV D。单次外层迭代处理一个 KV 块BMM1: S_Q[M,D] × S_K[D,N] → scores[M,N] softmaxonline跨迭代维护行 max/sum BMM2: P[M,N] × V[N,D] → 部分输出[M,D]与历史值在线累加架构上一个 Cube 核与两个 Vector 核配对1:2Cube 一次算 M 行两个 Vector 核各负责 M/2 行的 softmax 与累加Vector 侧缓冲按 M/2 行计。2.2 硬件规格单核参数取值说明CoreNum32核数Frequency1.6 GHz主频L1512 KBQ/P/KV 分片驻留L0A / L0B64 KBbmm 左/右矩阵分形缓冲各 2×32K 子缓冲L0C256 KBbmm 结果累加/输出缓冲多份轮转UB248 KB可用bmm 结果落点与 Vector 侧全部缓冲L2128 MB / 5.2 TB/s折合单核 ~101 B/cycleHBM128 GB / 1.6 TB/s折合单核 ~31 B/cycle单核带宽换算总带宽 ÷ 核数 ÷ 主频。L25.2×10¹² ÷ (32 × 1.6×10⁹) ≈ 101 B/cycleHBM1.6×10¹² ÷ (32 × 1.6×10⁹) ≈ 31 B/cycle。Cube 单周期完成一次16×16×16mma吞吐FMA 4096 FMA/cycle模型假设。源码佐证容量数据硬件表中的 UB 可用 248 KB 与 Kernel 侧约束完全对应——flash_attn_block_vec_nd.h 中UbLayout结构体以static_assert(sizeof(UbLayout) 248KB)限制总占用注释明确说明UB 总 size 为 256KB但需要为 AscendC 和 VF 预留 8K业务只能使用 248K。L0C 的 4×64KB 轮转结构、L0A/L0B 各 2×32KB 双缓冲也都能在 flash_attn_block_cube_nd.h 的 buffer 分配表中找到一一对应的实现见 §4。3. Cube Bound 的达成条件性能侧约束3.1 时间模型按单次迭代总量计。两个 matmul 的 FMA 总量为2·M·N·D$$ T_{cube} \frac{2 M N D}{FMA} $$Q 块M·D仅在首次迭代搬入一次S2 足够大时被摊薄K/V 每次迭代搬入2·N·D$$ T_{mte2} \approx \frac{2 N D \cdot \text{sizeof}(kv)}{Mix_{bw}}, \qquad Mix_{bw} h \cdot L2_{bw} (1-h) \cdot HBM_{bw} $$其中 $h$ 为 L2 命中率$Mix_{bw} \in [31,\ 101]$ B/cycle。bmm1/bmm2 结果fp32由 L0C fixpipe 到 UB$$ T_{fixpipe} \frac{(M N M D) \cdot \text{sizeof}(fp32)}{UB_{bw}} $$3.2 条件一M 的下限对 MTE2 的遮挡$T_{cube} T_{mte2}$ 两侧约去 N、D与 D、N 无关$$ M \frac{2 \cdot FMA}{Mix_{bw}} \frac{8192}{Mix_{bw}} $$Mix_bwM 下限101L2 全命中M 81.1即M ≥ 8231全 HBM 缺失M 264.3即M ≥ 265即M 82 时无论 L2 命中多好都必然 mte2 boundM ∈ [82, 265) 时需 L2 命中率满足$$ h \ge \frac{8192/M - 31}{101 - 31} $$M 越大对 h 的要求越低M96 需 78%M128 需 47%M ≥ 265 起无要求。3.3 条件二N 的下限对 Fixpipe 的遮挡$T_{cube} T_{fixpipe}$ 约去 M与 M 无关$$ UB_{bw} 2 \cdot FMA \cdot \frac{ND}{N D} 8192 \cdot \left(\frac{1}{D} \frac{1}{N}\right) \quad\Longleftrightarrow\quad N \frac{8192}{UB_{bw} - 8192/D} $$D256 时UB_bw 32 8192/NN128 需 UB_bw 96N256 需 UB_bw 64。N 越大bmm2 输出M×D与 N 无关被越多 KV 行摊薄fixpipe 压力越小——该条件方向上偏好大 N。UB_bw 未列入硬件表见 §8 待补充若其与 L2 带宽同级~101则 N ≥ 119 才能满足N128 可以满足。源码佐证Fixpipe 压力方向D256 下 fixpipe 输出量大Kernel 侧在 flash_attn_block_cube_nd.h 中通过IterateBmm2l0Split将 BMM2 沿 dV 方向按 128 列分块nLoops循环、splitN逐 tile 占用轮转的 64KB L0C buffer 并独立 Fixpipe——单个 L0C buffer 无需容纳整块128×256×4B结果这正是 §4.1 中bmm2 结果可沿 dV 方向分块适配 buffer 容量的实现落点。3.4 数值示例M128, N128, D256项数值周期cube2×128×128×256 8.39M FMA2048KV 搬入131072 Bh100%: 1298 / h47%: 2048临界/ h0: 4228fixpipe 搬出(128²128×256)×4 196608 BUB_bw96: 2048临界/ 128: 1536注意D256 时 fixpipe 搬出量196.6K/迭代大于 KV 搬入量131.1K/迭代fixpipe 是仅次于 MTE2 的第二约束。对照 (M64, N256)T_cube同为 2048M·N 相同但 KV 搬入 262144 B即使 h100% 也需 2596 cycles 2048——M64 82必然 mte2 bound的直接体现。4. 空间约束容量侧约束各级存储需容纳基本块的工作集并以多缓冲轮转维持流水重叠。以下按bmm1 结果缓冲与 bmm2 结果缓冲各自独立分配的常规设计推导§7 将放宽该假设。4.1 L0Cbmm 结果 bufferL0C256K以 4 份轮转 buffer 容纳 bmm1 结果M·N·4与 bmm2 结果M·D·4每份 64Kbmm2 结果可沿 dV 方向分块适配 buffer 容量分块数 $t \lceil M D \cdot 4 / 64K \rceil$为保 Cube→Fixpipe 重叠每迭代消耗的 buffer 数1 t须小于总 buffer 数 4。由此$$ M N \cdot 4 \le 64K ;\Rightarrow; \mathbf{M N \le 16384}; \qquad t \le 2 ;\Rightarrow; M D \le 32K ;\Rightarrow; \mathbf{M \le 128}\ (D{}256) $$dV 分块不改变 Mmad 总量矩阵乘按分块累加代价是 L0C 轮转深度从 2.0t1降到 4/3t2、fixpipe 调用次数翻倍但总量不变。4.2 L1输入分片驻留分片缓冲大小Q 块双缓冲2 × M·D·2 4MDPsoftmax 结果回写供 bmm2三缓冲3 × M·N·2 6MNK/V分时复用同组 buffer至少双 buffer≥ 2 × 2 × N·D·2 4NDbuffer 份数越多流水越深$$ 4MD 6MN 4ND \le 512K $$源码佐证L1 分配策略flash_attn_block_cube_nd.h 的 L1 buffer 分配与此表逐项对应P 矩阵 3 个 buffer3×mBaseSize×s2BaseSize×sizeof(INPUT_T)Q 矩阵 2 个gS1 行内复用、isFirstS2Loop加载/isLastS2Loop释放K/V 4×64KBconfig5 即 s2BaseSize256 且 D128 时改用 2×128KB 大缓冲。P 3 个、Q 2 个、KV 至少 2 个的多缓冲轮转深度正是为了让 MTE2 与 Cube 在 PRELOAD_N2 流水flash_attn_kernel_dn.h 中创建任务与执行任务分离、执行慢于创建 2 轮下连续重叠。4.3 UB结果缓冲与 Vector 侧工作集buffer大小bmm1 结果fp32双缓冲2 × (M/2)·N·4 4MNbmm2 结果fp32双缓冲2 × (M/2)·D·4 4MDPbf16双缓冲1 行分形对齐 padding≈ 2 × (M/2)·N·2 2MN输出累加暂存fp32单缓冲(M/2)·D·4 2MD掩码bool双缓冲2 × (M/2)·N·1 MNsoftmax 标量max/sum/exp 各三缓冲 lse双缓冲 杂项≈ 54M 512$$ 7MN 6MD 54M 512 \le 248K\ (253952,B) $$4.4 L0A/L0B定性各 64K2×32K 子缓冲。bmm 输入按 16×16 分形驻留K 方向D256按 128 分块后M、N ≤ 128 时每个输入分片恰好 32K 填满一个子缓冲自然满足M 或 N 过大将迫使输入侧退化单缓冲或重复加载故 L0A/L0B 不单独给出不等式但支持 M、N ≤ 128 的上界。源码佐证分形与子缓冲Kernel 侧 L0A/L0B 由BufferManager各分配 2×32KB 双 buffer见 flash_attn_block_cube_nd.h 的l0aBufferManager_/l0bBufferManager_配合BuffersPolicyDB的 LOCK_UNLOCK 策略让 MTE1 搬运L1→L0与 PIPE_M 计算重叠。D72 复用 D128 的 config 变体时L1→L0 搬运宽度按 16 元素分形对齐72→80Mmad 按真实 D 精确累加见 README 基本块一节说明16 分形对齐是贯穿 D 各档位的实现约束。5. D256 代入(M, N) 的可行域将 D256 代入上述约束分离缓冲设计约束不等式D256N128N256条件一MTE2M 8192/Mix_bwM ≥ 82h100%~ 265h0同左条件二FixpipeN 8192/(UB_bw−32)需 UB_bw 96需 UB_bw 64L0CMN ≤ 16384M ≤ 128M ≤ 128M ≤ 64L14MD6MN4ND ≤ 512KM ≤ 219M ≤ 102UB分离缓冲7MN6MD54M512 ≤ 248KM ≤ 102M ≤ 74联合82 M ≤ 102M ≤ 64无 cube bound 窗口两点观察N256 被 L0C 锁死MN ≤ 16384使 M ≤ 64 82无论其他约束如何放松都到不了 cube bound 阈值。N128 的窗口极窄可行域 (82, 102] 紧贴阈值 82——即使取到上界 M102也要求 L2 命中率 ≥ 70%裕度极小。而块大小需按 16 分形对齐、实用档位取 2 的幂M ∈ {64, 128}M128 超出 UB 上限 102311K 248KM64 低于阈值 82——分离缓冲设计下D256 在常规档位上不存在 cube bound 可行解。退而求其次只能取 (64, 256)空间全部满足UB 212K、L1 416K但 M64 82注定 mte2 bound——这就是 D256 性能远差于 D128 的理论根源。源码佐证档位与 config 映射模型中的常规档位 M ∈ {64, 128}并非随意设定。Host 侧 fa_adjust_sinner_souter.h 的AdjustSinnerAndSouter对 D256 只有两档选择gSize × maxSeqQ ≥ 64时取 sOuter64 / sInner128否则取 sOuter32 / sInner256再结合mBaseSize sOuter × CV_RATIOCV_RATIO2即得 M ∈ {64, 128}、N ∈ {128, 256}。这两个档位在编译期模板 flash_attn_template_tiling_key.h 中编码为 config4sOuter64, sInner128, D256与 config5sOuter32, sInner256, D256并在 README 基本块表 中完整列出。也就是说M 取 64 还是 128在实现层面就等价于选 config 4 还是 config 5二者被 tiling key 的 Config 位段bits 25-28固化。6. 可行域内 M、N 对性能的影响上节给出的是不等式范围而非定值范围内如何取值取决于 M、N 各自的性能影响。6.1 M在可行域内越大越好算术强度每 KV 字节的 FMA 数为M/2M 是唯一的算术强度旋钮。M 越大条件一越宽松对 L2 命中率的要求越低M102 需 70%M128 需 47%cube bound 越稳Q 摊薄Q 搬入的摊薄负担相对 cube 时间与 M 无关两者同随 M 线性增长不构成偏好代价尾块浪费Q 侧总行数不整除 M 时末块空转与任务粒度变粗。Q 侧行数充足prefill / 长序列时代价可忽略Q 侧总行数不足 M 时大块大面积空转必须缩小 M——此时 M 反正低于 82cube bound 不可达属另一种取值逻辑见 §8。结论Q 侧行数充足时M 取可行域上界。源码佐证Q 侧行数充足的判据AdjustSinnerAndSouter中 D256 分支的判据是gSize × maxSeqQ ≥ 64——即Q 侧有效行数Q 序列长度 × GQA group 数是否够填满一个大块。这与本文Q 侧行数充足与否决定 M 取值逻辑的结论精确对应充足则选大 MsOuter64 → M128不足则降档sOuter32 → M64。此外 tiling 的 gS1 合轴gS1Size actSeqLensQ × gSize见 README正是把 Q 序列与 GQA group 两轴合并为同一任务维M 的分块实际作用于 gS1 轴。6.2 N双向影响且与 M 在 L0C 上耦合大 N 的收益条件二fixpipe更宽松KV 迭代次数S2/N减少循环与同步开销摊薄小 N 的收益UB/L0C 中的 MN 项、L1 中的 N·D 项都随 N 收缩K/V 单个 L1 分片更小N128 时 64K/片可 4 个 buffer 深轮转N256 时 128K/片仅 2 个 bufferMTE2 与 Cube 的重叠更深S2 不整除 N 时尾块浪费更小。关键耦合L0C 约束MN ≤ 16384把 M 与 N 绑在同一份预算里——N 每放大一档M 的上限就减半。条件一M ≥ 82是 cube bound 的唯一入口而条件二只是过线即可N128 在 UB_bw ≥ 96 时已满足。因此选择次序是先保 M 过阈值并尽量大N 在剩余预算内取满足条件二的值。D256 下即 N128、M 取上界。7. 缓冲冗余的发现bmm1/bmm2 结果共用后可行域扩大7.1 冗余性论证§4.3 中 bmm1 结果缓冲4MN与 bmm2 结果缓冲4MD各自独立双缓冲但两者的生命周期并不重叠bmm1(L) 结果 → softmax(L) 消费完毕 → 该 buffer 死亡 bmm2(L) 结果 → 累加(L) 消费完毕 → 该 buffer 死亡数据流是环形链bmm1 → softmax → bmm2 → 累加 → bmm1(下一轮)同一迭代内 softmax 先读完 bmm1 结果、bmm2 才开始产出bmm2 的结果要到下一轮 softmax 前才被累加读完。因此两组双 buffer 可以共用同一对轮转 bufferbuffer 容量取两者的较大值(M/2)·max(N,D)·4。D256 ≥ N 时bmm1 结果(M/2)·N·4完全装得进 bmm2 的 buffer——bmm1 的双 buffer4MN整体消失。源码佐证实现即共用这一理论上发现的冗余在 Kernel 中已是实际设计——flash_attn_block_cube_nd.h 与 flash_attn_block_vec_nd.h 中的ubMmResBuffers_明确注释BMM1/BMM2 fixpipe 结果分时复用buffer 大小取mBaseSize/2 × max(s2BaseSize, dVBaseSize) × sizeof(FP32)即本文的(M/2)·max(N,D)·4buffer 数按 D 分档UB_MM_RES_BUFCNT (dBaseSize 128) ? 2 : 4。读写两侧通过一 flag 一 buffer的核间同步接力AIV 侧CC_MM_0~3/ AIC 侧CROSSCORE_MM_0~3写端BMM1/BMM2fixpipe 完成后 Set结果就绪读端Vec1/Vec2Wait 后读取、读毕再 Setbuffer 释放同一 flag 沿「就绪 → 释放 → 就绪」交替接力。mmResBufId_单调递增计数器保证同一 ExecuteTask 内先后执行的 BMM1(L) 与 BMM2(L-2) 落在不同 buffer避免写写冲突——这正是 §7.3 所述按结果序而非算子序分配 buffer、以环形等待保证时序的工程实现。7.2 新的 UB 空间公式$$ \underbrace{4M \cdot \max(N,D)}_{\text{共用结果缓冲}} 2MN 2MD MN 54M 512 \le 248K $$D256 且 N ≤ 256 时 $\max(N,D)D$化简为$$ \mathbf{3MN 6MD 54M 512 \le 248K} \qquad (\text{原 } 7MN 6MD \dots) $$节省量 4MNN ≤ D 时。对 (M128, N128)节省 64KUB 占用从 311K超限降为247.25K ≤ 248K余 768B。7.3 扩大后的可行域与更好的选择约束N128N256UB共用缓冲M ≤ 128M ≤ 107L0C / L1M ≤ 128 / M ≤ 219M ≤ 64/ M ≤ 102联合82 M ≤ 128M ≤ 64仍无窗口与 §6 的结论呼应N128 的窗口从 (82, 102] 扩大到(82, 128]恰好把 M128 档位纳入对 L2 命中率的要求从 70% 降到 47%cube bound 从勉强贴线变为有裕度——这就是缓冲共用后能选出更好 (M, N) 的量化体现N256 侧不变被 L0C 的 MN ≤ 16K 锁死在 M ≤ 64进一步印证 §6.2 的判断L0C 耦合下 N 必须给 M 让路共用的前提是按结果序而非算子序分配 buffer并以环形等待保证bmm1 第 L 轮写入晚于bmm2 第 L−2 轮读完——即空间收益需要同步机制的配合不是免费的。源码佐证环形等待与流水Kernel 的 PRELOAD_N2 流水flash_attn_kernel_dn.h 的ExecuteTask保证同一时刻 AIC 上 BMM1(本轮) 与 BMM2(2 轮前) 并行、AIV 上 Vec1(本轮) 与 Vec2(2 轮前) 并行P 矩阵的 3 buffer 轮转AIV 写第 N 轮 P 与 AIC 读第 N-2 轮 P 并行恰好容纳写入中/待读取/已释放三态。这些都建立在结果生命周期错开这一前提之上与本文的环形链论证互相印证。8. 结论场景选择依据D256Q 侧行数充足prefillM128, N128落在共用 buffer 后的可行域 (82, 128] 上界L2 命中率 ≥ ~47% 即 cube boundbmm2 结果按 dV/128 分块适配 L0C bufferD256Q 侧总行数 128decode/小 batchM64, N256M 被工作量压低cube bound 反正不可达M 82取大 N 换取 KV 迭代次数减半且 (64,256) 恰好贴住MN 16K的 L0C/mask 界空间无浪费建模方法论小结性能条件给下限cube bound 等价于M 8192/Mix_bw82~265与 D/N 无关与N 8192/(UB_bw − 8192/D)空间约束给上限L0CMN ≤ 16384、M ≤ 128dV 分块后L14MD6MN4ND ≤ 512KUB 由结果缓冲与 Vector 工作集决定可行域内 M 越大越好唯一算术强度旋钮N 与 M 在 L0C 上共享预算、优先保 M缓冲复用扩大可行域bmm1/bmm2 结果生命周期错开、可共用 bufferUB 上界 M ≤ 102 → M ≤ 128使 cube bound 从临界变为可行——空间优化最终转化为性能选项的解放。待补充UB 带宽 UB_bw 实测值用于封闭条件二D256 下 fixpipe 搬出量大于 KV 搬入量值得单独 profiling 验证。附模型结论与实现的对应关系及验证途径本文的所有推导最终都落在 CANN ops-transformer 仓库 FlashAttn 算子的真实实现中可按下述对应关系在源码中逐一验证基本块档位M、N 的取值集合编译期 config 模板 flash_attn_template_tiling_key.hconfig 0~7 的 sOuter/sInner/D/DV 编码与运行期AdjustSinnerAndSouterfa_adjust_sinner_souter.h共同决定mBaseSize sOuter × 2与s2BaseSize sInnertiling key 编码5 个模板参数InOutLayoutType / KvLayoutType / HasAttenMask / TemplateId / Config按位拼接为 tiling keyHost tiling 的GenTilingKey生成、运行期 kernel 按 key 选变体细节见 README 的 TilingKey 章节缓冲与存储层级L0C 4×64KB、L0A/L0B 各 2×32KB、UBstatic_assert ≤ 248KB、L1 的 P/Q/KV 多缓冲见 flash_attn_block_cube_nd.h 与 flash_attn_block_vec_nd.h 的 buffer 分配表负载均衡与 FDsection 切分与 split-K 归约FlashDecode由flash_attn_metadataAICPU基于 SectionStreamK 算法产出 metadata主算子按 metadata 区间执行见 flash_attn_tiling_data.h 的 metadata 布局与 README 的负载均衡章节性能验证途径仓库提供 train/infer 各 220 性能用例tests/pytests/test_cases/performance_redline_train.py、performance_redline_infer.py可通过--perf_mode采集性能数据并配合 ASCEND DEBUG 日志中的GenTilingKey/kernel_nameFlashAttn_hash_tiling_key确认实际选用的档位——感兴趣的读者可在真机上用这些用例验证本文UB_bw 实测值待补充的开放项。此外FlashAttn 的对外接口flash_attn_metadataflash_attn两段式调用、参数说明、约束与调用示例见 torchapi_flash_attn.md可与本文的性能建模对照理解接口层的 D 支持范围64/72/128/256以及 MLA 形态的 QK192/V128与模型假设的 D256 场景一致接口文档中metadata 与主算子入参必须保持一致的约束则对应本文 §4 中tiling 与 metadata 共用AdjustSinnerAndSouter保证两侧切分口径一致的实现要求。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考