 的選擇與 Cube Bound 達(dá)成分析)
CANN ops-transformer FlashAttn 性能建模D256 下基本塊 (M, N) 的選擇與 Cube Bound 達(dá)成分析【免費(fèi)下載鏈接】ops-transformer本項(xiàng)目是CANN提供的transformer類大模型算子庫實(shí)現(xiàn)網(wǎng)絡(luò)在NPU上加速計(jì)算。項(xiàng)目地址: https://gitcode.com/cann/ops-transformer本文基于 CANN ops-transformer 倉庫中 flash_attn_perf_analysis.md 的技術(shù)文檔結(jié)合 FlashAttn 算子的 Host 側(cè) tiling 與 Kernel 側(cè)源碼推導(dǎo) D256Q/K/V 頭維 256非量化 fp16/bf16dV D條件下 Flash Attention 基本塊 (M, N) 的選擇依據(jù)。讀者讀完后將掌握如何以 Cube Bound 為建模目標(biāo)推導(dǎo)性能下限與空間上限如何在 L0C/L1/UB 的容量約束與 MTE2/Fixpipe 的帶寬約束之間找到可行域以及為什么 bmm1/bmm2 結(jié)果緩沖共用能從根本上擴(kuò)大 D256 的可行解空間。1. 建模目標(biāo)Cube BoundFlash Attention 在 Ascend 上由多條引擎流水組成CubeMmadbmm1Q·K?與 bmm2P·V兩個(gè)矩陣乘MTEQ/K/V 從 GM 搬入 L1、L1 到 L0A/L0B 的分形搬運(yùn)Fixpipebmm 結(jié)果從 L0C 搬出到 UBVectorsoftmax行 max/exp/sum與輸出的在線累加。這些引擎中只有 Cube 計(jì)算不可壓縮——矩陣乘的 FMA 數(shù)由問題規(guī)模唯一確定搬入、搬出與 softmax 都可以與其他階段重疊執(zhí)行。因此 kernel 的性能上界就是 Cube 滿流水建模目標(biāo)為Cube Bound$$ T_{cube} \ge T_{mte2}, \qquad T_{cube} \ge T_{fixpipe}, \qquad T_{cube} \ge T_{vec} $$即 Cube 永遠(yuǎn)有數(shù)據(jù)可算其余所有階段的時(shí)間都被 Cube 時(shí)間完全遮擋。其中 Vector 側(cè) softmax/累加運(yùn)行在與 Cube 配對(duì)的 Vector 核上工作量與 M(ND) 成正比在實(shí)用塊大小范圍內(nèi)遠(yuǎn)小于 $T_{cube}$不構(gòu)成額外約束本文重點(diǎn)分析最緊的MTE2搬入與Fixpipe搬出兩個(gè)條件。從實(shí)現(xiàn)看FlashAttn 的 kernel 入口 flash_attn.cpp 以KERNEL_TYPE_MIX_AIC_1_2模式將同一份源碼同時(shí)編譯為 Cube 側(cè)AIC與 Vector 側(cè)AIV兩個(gè)變體每個(gè) AIC 與兩個(gè) AIV 配對(duì)CV_RATIO2——這與本文一個(gè) Cube 核配兩個(gè) Vector 核、Vector 側(cè)緩沖按 M/2 行計(jì)的模型設(shè)定完全一致。2. 模型設(shè)定2.1 基本塊與迭代結(jié)構(gòu)MQ 序列方向每次迭代、每個(gè) Cube 核承擔(dān)的塊行數(shù)bmm 的 M 維NKV 序列方向的塊大小bmm 的 N 維一次迭代處理 N 行 KVD頭維bmm 的 K 維 / bmm2 的輸出列數(shù)非量化時(shí) dV D。單次外層迭代處理一個(gè) KV 塊BMM1: S_Q[M,D] × S_K[D,N] → scores[M,N] softmaxonline跨迭代維護(hù)行 max/sum BMM2: P[M,N] × V[N,D] → 部分輸出[M,D]與歷史值在線累加架構(gòu)上一個(gè) Cube 核與兩個(gè) Vector 核配對(duì)1:2Cube 一次算 M 行兩個(gè) Vector 核各負(fù)責(zé) M/2 行的 softmax 與累加Vector 側(cè)緩沖按 M/2 行計(jì)。2.2 硬件規(guī)格單核參數(shù)取值說明CoreNum32核數(shù)Frequency1.6 GHz主頻L1512 KBQ/P/KV 分片駐留L0A / L0B64 KBbmm 左/右矩陣分形緩沖各 2×32K 子緩沖L0C256 KBbmm 結(jié)果累加/輸出緩沖多份輪轉(zhuǎn)UB248 KB可用bmm 結(jié)果落點(diǎn)與 Vector 側(cè)全部緩沖L2128 MB / 5.2 TB/s折合單核 ~101 B/cycleHBM128 GB / 1.6 TB/s折合單核 ~31 B/cycle單核帶寬換算總帶寬 ÷ 核數(shù) ÷ 主頻。L25.2×1012 ÷ (32 × 1.6×10?) ≈ 101 B/cycleHBM1.6×1012 ÷ (32 × 1.6×10?) ≈ 31 B/cycle。Cube 單周期完成一次16×16×16mma吞吐FMA 4096 FMA/cycle模型假設(shè)。源碼佐證容量數(shù)據(jù)硬件表中的 UB 可用 248 KB 與 Kernel 側(cè)約束完全對(duì)應(yīng)——flash_attn_block_vec_nd.h 中UbLayout結(jié)構(gòu)體以static_assert(sizeof(UbLayout) 248KB)限制總占用注釋明確說明UB 總 size 為 256KB但需要為 AscendC 和 VF 預(yù)留 8K業(yè)務(wù)只能使用 248K。L0C 的 4×64KB 輪轉(zhuǎn)結(jié)構(gòu)、L0A/L0B 各 2×32KB 雙緩沖也都能在 flash_attn_block_cube_nd.h 的 buffer 分配表中找到一一對(duì)應(yīng)的實(shí)現(xiàn)見 §4。3. Cube Bound 的達(dá)成條件性能側(cè)約束3.1 時(shí)間模型按單次迭代總量計(jì)。兩個(gè) matmul 的 FMA 總量為2·M·N·D$$ T_{cube} \frac{2 M N D}{FMA} $$Q 塊M·D僅在首次迭代搬入一次S2 足夠大時(shí)被攤薄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 結(jié)果fp32由 L0C fixpipe 到 UB$$ T_{fixpipe} \frac{(M N M D) \cdot \text{sizeof}(fp32)}{UB_{bw}} $$3.2 條件一M 的下限對(duì) MTE2 的遮擋$T_{cube} T_{mte2}$ 兩側(cè)約去 N、D與 D、N 無關(guā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 時(shí)無論 L2 命中多好都必然 mte2 boundM ∈ [82, 265) 時(shí)需 L2 命中率滿足$$ h \ge \frac{8192/M - 31}{101 - 31} $$M 越大對(duì) h 的要求越低M96 需 78%M128 需 47%M ≥ 265 起無要求。3.3 條件二N 的下限對(duì) Fixpipe 的遮擋$T_{cube} T_{fixpipe}$ 約去 M與 M 無關(guān)$$ 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 時(shí)UB_bw 32 8192/NN128 需 UB_bw 96N256 需 UB_bw 64。N 越大bmm2 輸出M×D與 N 無關(guān)被越多 KV 行攤薄fixpipe 壓力越小——該條件方向上偏好大 N。UB_bw 未列入硬件表見 §8 待補(bǔ)充若其與 L2 帶寬同級(jí)~101則 N ≥ 119 才能滿足N128 可以滿足。源碼佐證Fixpipe 壓力方向D256 下 fixpipe 輸出量大Kernel 側(cè)在 flash_attn_block_cube_nd.h 中通過IterateBmm2l0Split將 BMM2 沿 dV 方向按 128 列分塊nLoops循環(huán)、splitN逐 tile 占用輪轉(zhuǎn)的 64KB L0C buffer 并獨(dú)立 Fixpipe——單個(gè) L0C buffer 無需容納整塊128×256×4B結(jié)果這正是 §4.1 中bmm2 結(jié)果可沿 dV 方向分塊適配 buffer 容量的實(shí)現(xiàn)落點(diǎn)。3.4 數(shù)值示例M128, N128, D256項(xiàng)數(shù)值周期cube2×128×128×256 8.39M FMA2048KV 搬入131072 Bh100%: 1298 / h47%: 2048臨界/ h0: 4228fixpipe 搬出(1282128×256)×4 196608 BUB_bw96: 2048臨界/ 128: 1536注意D256 時(shí) fixpipe 搬出量196.6K/迭代大于 KV 搬入量131.1K/迭代fixpipe 是僅次于 MTE2 的第二約束。對(duì)照 (M64, N256)T_cube同為 2048M·N 相同但 KV 搬入 262144 B即使 h100% 也需 2596 cycles 2048——M64 82必然 mte2 bound的直接體現(xiàn)。4. 空間約束容量側(cè)約束各級(jí)存儲(chǔ)需容納基本塊的工作集并以多緩沖輪轉(zhuǎn)維持流水重疊。以下按bmm1 結(jié)果緩沖與 bmm2 結(jié)果緩沖各自獨(dú)立分配的常規(guī)設(shè)計(jì)推導(dǎo)§7 將放寬該假設(shè)。4.1 L0Cbmm 結(jié)果 bufferL0C256K以 4 份輪轉(zhuǎn) buffer 容納 bmm1 結(jié)果M·N·4與 bmm2 結(jié)果M·D·4每份 64Kbmm2 結(jié)果可沿 dV 方向分塊適配 buffer 容量分塊數(shù) $t \lceil M D \cdot 4 / 64K \rceil$為保 Cube→Fixpipe 重疊每迭代消耗的 buffer 數(shù)1 t須小于總 buffer 數(shù) 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 總量矩陣乘按分塊累加代價(jià)是 L0C 輪轉(zhuǎn)深度從 2.0t1降到 4/3t2、fixpipe 調(diào)用次數(shù)翻倍但總量不變。4.2 L1輸入分片駐留分片緩沖大小Q 塊雙緩沖2 × M·D·2 4MDPsoftmax 結(jié)果回寫供 bmm2三緩沖3 × M·N·2 6MNK/V分時(shí)復(fù)用同組 buffer至少雙 buffer≥ 2 × 2 × N·D·2 4NDbuffer 份數(shù)越多流水越深$$ 4MD 6MN 4ND \le 512K $$源碼佐證L1 分配策略flash_attn_block_cube_nd.h 的 L1 buffer 分配與此表逐項(xiàng)對(duì)應(yīng)P 矩陣 3 個(gè) buffer3×mBaseSize×s2BaseSize×sizeof(INPUT_T)Q 矩陣 2 個(gè)gS1 行內(nèi)復(fù)用、isFirstS2Loop加載/isLastS2Loop釋放K/V 4×64KBconfig5 即 s2BaseSize256 且 D128 時(shí)改用 2×128KB 大緩沖。P 3 個(gè)、Q 2 個(gè)、KV 至少 2 個(gè)的多緩沖輪轉(zhuǎn)深度正是為了讓 MTE2 與 Cube 在 PRELOAD_N2 流水flash_attn_kernel_dn.h 中創(chuàng)建任務(wù)與執(zhí)行任務(wù)分離、執(zhí)行慢于創(chuàng)建 2 輪下連續(xù)重疊。4.3 UB結(jié)果緩沖與 Vector 側(cè)工作集buffer大小bmm1 結(jié)果fp32雙緩沖2 × (M/2)·N·4 4MNbmm2 結(jié)果fp32雙緩沖2 × (M/2)·D·4 4MDPbf16雙緩沖1 行分形對(duì)齊 padding≈ 2 × (M/2)·N·2 2MN輸出累加暫存fp32單緩沖(M/2)·D·4 2MD掩碼bool雙緩沖2 × (M/2)·N·1 MNsoftmax 標(biāo)量max/sum/exp 各三緩沖 lse雙緩沖 雜項(xiàng)≈ 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 時(shí)每個(gè)輸入分片恰好 32K 填滿一個(gè)子緩沖自然滿足M 或 N 過大將迫使輸入側(cè)退化單緩沖或重復(fù)加載故 L0A/L0B 不單獨(dú)給出不等式但支持 M、N ≤ 128 的上界。源碼佐證分形與子緩沖Kernel 側(cè) L0A/L0B 由BufferManager各分配 2×32KB 雙 buffer見 flash_attn_block_cube_nd.h 的l0aBufferManager_/l0bBufferManager_配合BuffersPolicyDB的 LOCK_UNLOCK 策略讓 MTE1 搬運(yùn)L1→L0與 PIPE_M 計(jì)算重疊。D72 復(fù)用 D128 的 config 變體時(shí)L1→L0 搬運(yùn)寬度按 16 元素分形對(duì)齊72→80Mmad 按真實(shí) D 精確累加見 README 基本塊一節(jié)說明16 分形對(duì)齊是貫穿 D 各檔位的實(shí)現(xiàn)約束。5. D256 代入(M, N) 的可行域?qū)?D256 代入上述約束分離緩沖設(shè)計(jì)約束不等式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聯(lián)合82 M ≤ 102M ≤ 64無 cube bound 窗口兩點(diǎn)觀察N256 被 L0C 鎖死MN ≤ 16384使 M ≤ 64 82無論其他約束如何放松都到不了 cube bound 閾值。N128 的窗口極窄可行域 (82, 102] 緊貼閾值 82——即使取到上界 M102也要求 L2 命中率 ≥ 70%裕度極小。而塊大小需按 16 分形對(duì)齊、實(shí)用檔位取 2 的冪M ∈ {64, 128}M128 超出 UB 上限 102311K 248KM64 低于閾值 82——分離緩沖設(shè)計(jì)下D256 在常規(guī)檔位上不存在 cube bound 可行解。退而求其次只能取 (64, 256)空間全部滿足UB 212K、L1 416K但 M64 82注定 mte2 bound——這就是 D256 性能遠(yuǎn)差于 D128 的理論根源。源碼佐證檔位與 config 映射模型中的常規(guī)檔位 M ∈ {64, 128}并非隨意設(shè)定。Host 側(cè) fa_adjust_sinner_souter.h 的AdjustSinnerAndSouter對(duì) D256 只有兩檔選擇gSize × maxSeqQ ≥ 64時(shí)取 sOuter64 / sInner128否則取 sOuter32 / sInner256再結(jié)合mBaseSize sOuter × CV_RATIOCV_RATIO2即得 M ∈ {64, 128}、N ∈ {128, 256}。這兩個(gè)檔位在編譯期模板 flash_attn_template_tiling_key.h 中編碼為 config4sOuter64, sInner128, D256與 config5sOuter32, sInner256, D256并在 README 基本塊表 中完整列出。也就是說M 取 64 還是 128在實(shí)現(xiàn)層面就等價(jià)于選 config 4 還是 config 5二者被 tiling key 的 Config 位段bits 25-28固化。6. 可行域內(nèi) M、N 對(duì)性能的影響上節(jié)給出的是不等式范圍而非定值范圍內(nèi)如何取值取決于 M、N 各自的性能影響。6.1 M在可行域內(nèi)越大越好算術(shù)強(qiáng)度每 KV 字節(jié)的 FMA 數(shù)為M/2M 是唯一的算術(shù)強(qiáng)度旋鈕。M 越大條件一越寬松對(duì) L2 命中率的要求越低M102 需 70%M128 需 47%cube bound 越穩(wěn)Q 攤薄Q 搬入的攤薄負(fù)擔(dān)相對(duì) cube 時(shí)間與 M 無關(guān)兩者同隨 M 線性增長(zhǎng)不構(gòu)成偏好代價(jià)尾塊浪費(fèi)Q 側(cè)總行數(shù)不整除 M 時(shí)末塊空轉(zhuǎn)與任務(wù)粒度變粗。Q 側(cè)行數(shù)充足prefill / 長(zhǎng)序列時(shí)代價(jià)可忽略Q 側(cè)總行數(shù)不足 M 時(shí)大塊大面積空轉(zhuǎn)必須縮小 M——此時(shí) M 反正低于 82cube bound 不可達(dá)屬另一種取值邏輯見 §8。結(jié)論Q 側(cè)行數(shù)充足時(shí)M 取可行域上界。源碼佐證Q 側(cè)行數(shù)充足的判據(jù)AdjustSinnerAndSouter中 D256 分支的判據(jù)是gSize × maxSeqQ ≥ 64——即Q 側(cè)有效行數(shù)Q 序列長(zhǎng)度 × GQA group 數(shù)是否夠填滿一個(gè)大塊。這與本文Q 側(cè)行數(shù)充足與否決定 M 取值邏輯的結(jié)論精確對(duì)應(yīng)充足則選大 MsOuter64 → M128不足則降檔sOuter32 → M64。此外 tiling 的 gS1 合軸gS1Size actSeqLensQ × gSize見 README正是把 Q 序列與 GQA group 兩軸合并為同一任務(wù)維M 的分塊實(shí)際作用于 gS1 軸。6.2 N雙向影響且與 M 在 L0C 上耦合大 N 的收益條件二fixpipe更寬松KV 迭代次數(shù)S2/N減少循環(huán)與同步開銷攤薄小 N 的收益UB/L0C 中的 MN 項(xiàng)、L1 中的 N·D 項(xiàng)都隨 N 收縮K/V 單個(gè) L1 分片更小N128 時(shí) 64K/片可 4 個(gè) buffer 深輪轉(zhuǎn)N256 時(shí) 128K/片僅 2 個(gè) bufferMTE2 與 Cube 的重疊更深S2 不整除 N 時(shí)尾塊浪費(fèi)更小。關(guān)鍵耦合L0C 約束MN ≤ 16384把 M 與 N 綁在同一份預(yù)算里——N 每放大一檔M 的上限就減半。條件一M ≥ 82是 cube bound 的唯一入口而條件二只是過線即可N128 在 UB_bw ≥ 96 時(shí)已滿足。因此選擇次序是先保 M 過閾值并盡量大N 在剩余預(yù)算內(nèi)取滿足條件二的值。D256 下即 N128、M 取上界。7. 緩沖冗余的發(fā)現(xiàn)bmm1/bmm2 結(jié)果共用后可行域擴(kuò)大7.1 冗余性論證§4.3 中 bmm1 結(jié)果緩沖4MN與 bmm2 結(jié)果緩沖4MD各自獨(dú)立雙緩沖但兩者的生命周期并不重疊bmm1(L) 結(jié)果 → softmax(L) 消費(fèi)完畢 → 該 buffer 死亡 bmm2(L) 結(jié)果 → 累加(L) 消費(fèi)完畢 → 該 buffer 死亡數(shù)據(jù)流是環(huán)形鏈bmm1 → softmax → bmm2 → 累加 → bmm1(下一輪)同一迭代內(nèi) softmax 先讀完 bmm1 結(jié)果、bmm2 才開始產(chǎn)出bmm2 的結(jié)果要到下一輪 softmax 前才被累加讀完。因此兩組雙 buffer 可以共用同一對(duì)輪轉(zhuǎn) bufferbuffer 容量取兩者的較大值(M/2)·max(N,D)·4。D256 ≥ N 時(shí)bmm1 結(jié)果(M/2)·N·4完全裝得進(jìn) bmm2 的 buffer——bmm1 的雙 buffer4MN整體消失。源碼佐證實(shí)現(xiàn)即共用這一理論上發(fā)現(xiàn)的冗余在 Kernel 中已是實(shí)際設(shè)計(jì)——flash_attn_block_cube_nd.h 與 flash_attn_block_vec_nd.h 中的ubMmResBuffers_明確注釋BMM1/BMM2 fixpipe 結(jié)果分時(shí)復(fù)用buffer 大小取mBaseSize/2 × max(s2BaseSize, dVBaseSize) × sizeof(FP32)即本文的(M/2)·max(N,D)·4buffer 數(shù)按 D 分檔UB_MM_RES_BUFCNT (dBaseSize 128) ? 2 : 4。讀寫兩側(cè)通過一 flag 一 buffer的核間同步接力AIV 側(cè)CC_MM_0~3/ AIC 側(cè)CROSSCORE_MM_0~3寫端BMM1/BMM2fixpipe 完成后 Set結(jié)果就緒讀端Vec1/Vec2Wait 后讀取、讀畢再 Setbuffer 釋放同一 flag 沿「就緒 → 釋放 → 就緒」交替接力。mmResBufId_單調(diào)遞增計(jì)數(shù)器保證同一 ExecuteTask 內(nèi)先后執(zhí)行的 BMM1(L) 與 BMM2(L-2) 落在不同 buffer避免寫寫沖突——這正是 §7.3 所述按結(jié)果序而非算子序分配 buffer、以環(huán)形等待保證時(shí)序的工程實(shí)現(xiàn)。7.2 新的 UB 空間公式$$ \underbrace{4M \cdot \max(N,D)}_{\text{共用結(jié)果緩沖}} 2MN 2MD MN 54M 512 \le 248K $$D256 且 N ≤ 256 時(shí) $\max(N,D)D$化簡(jiǎn)為$$ \mathbf{3MN 6MD 54M 512 \le 248K} \qquad (\text{原 } 7MN 6MD \dots) $$節(jié)省量 4MNN ≤ D 時(shí)。對(duì) (M128, N128)節(jié)省 64KUB 占用從 311K超限降為247.25K ≤ 248K余 768B。7.3 擴(kuò)大后的可行域與更好的選擇約束N128N256UB共用緩沖M ≤ 128M ≤ 107L0C / L1M ≤ 128 / M ≤ 219M ≤ 64/ M ≤ 102聯(lián)合82 M ≤ 128M ≤ 64仍無窗口與 §6 的結(jié)論呼應(yīng)N128 的窗口從 (82, 102] 擴(kuò)大到(82, 128]恰好把 M128 檔位納入對(duì) L2 命中率的要求從 70% 降到 47%cube bound 從勉強(qiáng)貼線變?yōu)橛性６取@就是緩沖共用后能選出更好 (M, N) 的量化體現(xiàn)N256 側(cè)不變被 L0C 的 MN ≤ 16K 鎖死在 M ≤ 64進(jìn)一步印證 §6.2 的判斷L0C 耦合下 N 必須給 M 讓路共用的前提是按結(jié)果序而非算子序分配 buffer并以環(huán)形等待保證bmm1 第 L 輪寫入晚于bmm2 第 L?2 輪讀完——即空間收益需要同步機(jī)制的配合不是免費(fèi)的。源碼佐證環(huán)形等待與流水Kernel 的 PRELOAD_N2 流水flash_attn_kernel_dn.h 的ExecuteTask保證同一時(shí)刻 AIC 上 BMM1(本輪) 與 BMM2(2 輪前) 并行、AIV 上 Vec1(本輪) 與 Vec2(2 輪前) 并行P 矩陣的 3 buffer 輪轉(zhuǎn)AIV 寫第 N 輪 P 與 AIC 讀第 N-2 輪 P 并行恰好容納寫入中/待讀取/已釋放三態(tài)。這些都建立在結(jié)果生命周期錯(cuò)開這一前提之上與本文的環(huán)形鏈論證互相印證。8. 結(jié)論場(chǎng)景選擇依據(jù)D256Q 側(cè)行數(shù)充足prefillM128, N128落在共用 buffer 后的可行域 (82, 128] 上界L2 命中率 ≥ ~47% 即 cube boundbmm2 結(jié)果按 dV/128 分塊適配 L0C bufferD256Q 側(cè)總行數(shù) 128decode/小 batchM64, N256M 被工作量壓低cube bound 反正不可達(dá)M 82取大 N 換取 KV 迭代次數(shù)減半且 (64,256) 恰好貼住MN 16K的 L0C/mask 界空間無浪費(fèi)建模方法論小結(jié)性能條件給下限cube bound 等價(jià)于M 8192/Mix_bw82~265與 D/N 無關(guān)與N 8192/(UB_bw ? 8192/D)空間約束給上限L0CMN ≤ 16384、M ≤ 128dV 分塊后L14MD6MN4ND ≤ 512KUB 由結(jié)果緩沖與 Vector 工作集決定可行域內(nèi) M 越大越好唯一算術(shù)強(qiáng)度旋鈕N 與 M 在 L0C 上共享預(yù)算、優(yōu)先保 M緩沖復(fù)用擴(kuò)大可行域bmm1/bmm2 結(jié)果生命周期錯(cuò)開、可共用 bufferUB 上界 M ≤ 102 → M ≤ 128使 cube bound 從臨界變?yōu)榭尚小臻g優(yōu)化最終轉(zhuǎn)化為性能選項(xiàng)的解放。待補(bǔ)充UB 帶寬 UB_bw 實(shí)測(cè)值用于封閉條件二D256 下 fixpipe 搬出量大于 KV 搬入量值得單獨(dú) profiling 驗(yàn)證。附模型結(jié)論與實(shí)現(xiàn)的對(duì)應(yīng)關(guān)系及驗(yàn)證途徑本文的所有推導(dǎo)最終都落在 CANN ops-transformer 倉庫 FlashAttn 算子的真實(shí)實(shí)現(xiàn)中可按下述對(duì)應(yīng)關(guān)系在源碼中逐一驗(yàn)證基本塊檔位M、N 的取值集合編譯期 config 模板 flash_attn_template_tiling_key.hconfig 0~7 的 sOuter/sInner/D/DV 編碼與運(yùn)行期AdjustSinnerAndSouterfa_adjust_sinner_souter.h共同決定mBaseSize sOuter × 2與s2BaseSize sInnertiling key 編碼5 個(gè)模板參數(shù)InOutLayoutType / KvLayoutType / HasAttenMask / TemplateId / Config按位拼接為 tiling keyHost tiling 的GenTilingKey生成、運(yùn)行期 kernel 按 key 選變體細(xì)節(jié)見 README 的 TilingKey 章節(jié)緩沖與存儲(chǔ)層級(jí)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 分配表負(fù)載均衡與 FDsection 切分與 split-K 歸約FlashDecode由flash_attn_metadataAICPU基于 SectionStreamK 算法產(chǎn)出 metadata主算子按 metadata 區(qū)間執(zhí)行見 flash_attn_tiling_data.h 的 metadata 布局與 README 的負(fù)載均衡章節(jié)性能驗(yàn)證途徑倉庫提供 train/infer 各 220 性能用例tests/pytests/test_cases/performance_redline_train.py、performance_redline_infer.py可通過--perf_mode采集性能數(shù)據(jù)并配合 ASCEND DEBUG 日志中的GenTilingKey/kernel_nameFlashAttn_hash_tiling_key確認(rèn)實(shí)際選用的檔位——感興趣的讀者可在真機(jī)上用這些用例驗(yàn)證本文UB_bw 實(shí)測(cè)值待補(bǔ)充的開放項(xiàng)。此外FlashAttn 的對(duì)外接口flash_attn_metadataflash_attn兩段式調(diào)用、參數(shù)說明、約束與調(diào)用示例見 torchapi_flash_attn.md可與本文的性能建模對(duì)照理解接口層的 D 支持范圍64/72/128/256以及 MLA 形態(tài)的 QK192/V128與模型假設(shè)的 D256 場(chǎng)景一致接口文檔中metadata 與主算子入?yún)⒈仨毐３忠恢碌募s束則對(duì)應(yīng)本文 §4 中tiling 與 metadata 共用AdjustSinnerAndSouter保證兩側(cè)切分口徑一致的實(shí)現(xiàn)要求?！久赓M(fèi)下載鏈接】ops-transformer本項(xiàng)目是CANN提供的transformer類大模型算子庫實(shí)現(xiàn)網(wǎng)絡(luò)在NPU上加速計(jì)算。項(xiàng)目地址: https://gitcode.com/cann/ops-transformer創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 CANN ops-transformer FlashAttn 性能建模：D=256 下基本塊 (M, N) 的選擇與 Cube Bound 達(dá)成分析
CANN ops-transformer FlashAttn 性能建模D256 下基本塊 (M, N) 的選擇與 Cube Bound 達(dá)成分析【免費(fèi)下載鏈接】ops-transformer本項(xiàng)目是CANN提供的transformer類大模型算子庫實(shí)現(xiàn)網(wǎng)絡(luò)在NPU上加速計(jì)算。項(xiàng)目地址: https://gitcode.com/cann/ops-transformer本文基于 CANN ops-transformer 倉庫中 flash_attn_perf_analysis.md 的技術(shù)文檔結(jié)合 FlashAttn 算子的 Host 側(cè) tiling 與 Kernel 側(cè)源碼推導(dǎo) D256Q/K/V 頭維 256非量化 fp16/bf16dV D條件下 Flash Attention 基本塊 (M, N) 的選擇依據(jù)。讀者讀完后將掌握如何以 Cube Bound 為建模目標(biāo)推導(dǎo)性能下限與空間上限如何在 L0C/L1/UB 的容量約束與 MTE2/Fixpipe 的帶寬約束之間找到可行域以及為什么 bmm1/bmm2 結(jié)果緩沖共用能從根本上擴(kuò)大 D256 的可行解空間。1. 建模目標(biāo)Cube BoundFlash Attention 在 Ascend 上由多條引擎流水組成CubeMmadbmm1Q·K?與 bmm2P·V兩個(gè)矩陣乘MTEQ/K/V 從 GM 搬入 L1、L1 到 L0A/L0B 的分形搬運(yùn)Fixpipebmm 結(jié)果從 L0C 搬出到 UBVectorsoftmax行 max/exp/sum與輸出的在線累加。這些引擎中只有 Cube 計(jì)算不可壓縮——矩陣乘的 FMA 數(shù)由問題規(guī)模唯一確定搬入、搬出與 softmax 都可以與其他階段重疊執(zhí)行。因此 kernel 的性能上界就是 Cube 滿流水建模目標(biāo)為Cube Bound$$ T_{cube} \ge T_{mte2}, \qquad T_{cube} \ge T_{fixpipe}, \qquad T_{cube} \ge T_{vec} $$即 Cube 永遠(yuǎn)有數(shù)據(jù)可算其余所有階段的時(shí)間都被 Cube 時(shí)間完全遮擋。其中 Vector 側(cè) softmax/累加運(yùn)行在與 Cube 配對(duì)的 Vector 核上工作量與 M(ND) 成正比在實(shí)用塊大小范圍內(nèi)遠(yuǎn)小于 $T_{cube}$不構(gòu)成額外約束本文重點(diǎn)分析最緊的MTE2搬入與Fixpipe搬出兩個(gè)條件。從實(shí)現(xiàn)看FlashAttn 的 kernel 入口 flash_attn.cpp 以KERNEL_TYPE_MIX_AIC_1_2模式將同一份源碼同時(shí)編譯為 Cube 側(cè)AIC與 Vector 側(cè)AIV兩個(gè)變體每個(gè) AIC 與兩個(gè) AIV 配對(duì)CV_RATIO2——這與本文一個(gè) Cube 核配兩個(gè) Vector 核、Vector 側(cè)緩沖按 M/2 行計(jì)的模型設(shè)定完全一致。2. 模型設(shè)定2.1 基本塊與迭代結(jié)構(gòu)MQ 序列方向每次迭代、每個(gè) Cube 核承擔(dān)的塊行數(shù)bmm 的 M 維NKV 序列方向的塊大小bmm 的 N 維一次迭代處理 N 行 KVD頭維bmm 的 K 維 / bmm2 的輸出列數(shù)非量化時(shí) dV D。單次外層迭代處理一個(gè) KV 塊BMM1: S_Q[M,D] × S_K[D,N] → scores[M,N] softmaxonline跨迭代維護(hù)行 max/sum BMM2: P[M,N] × V[N,D] → 部分輸出[M,D]與歷史值在線累加架構(gòu)上一個(gè) Cube 核與兩個(gè) Vector 核配對(duì)1:2Cube 一次算 M 行兩個(gè) Vector 核各負(fù)責(zé) M/2 行的 softmax 與累加Vector 側(cè)緩沖按 M/2 行計(jì)。2.2 硬件規(guī)格單核參數(shù)取值說明CoreNum32核數(shù)Frequency1.6 GHz主頻L1512 KBQ/P/KV 分片駐留L0A / L0B64 KBbmm 左/右矩陣分形緩沖各 2×32K 子緩沖L0C256 KBbmm 結(jié)果累加/輸出緩沖多份輪轉(zhuǎn)UB248 KB可用bmm 結(jié)果落點(diǎn)與 Vector 側(cè)全部緩沖L2128 MB / 5.2 TB/s折合單核 ~101 B/cycleHBM128 GB / 1.6 TB/s折合單核 ~31 B/cycle單核帶寬換算總帶寬 ÷ 核數(shù) ÷ 主頻。L25.2×1012 ÷ (32 × 1.6×10?) ≈ 101 B/cycleHBM1.6×1012 ÷ (32 × 1.6×10?) ≈ 31 B/cycle。Cube 單周期完成一次16×16×16mma吞吐FMA 4096 FMA/cycle模型假設(shè)。源碼佐證容量數(shù)據(jù)硬件表中的 UB 可用 248 KB 與 Kernel 側(cè)約束完全對(duì)應(yīng)——flash_attn_block_vec_nd.h 中UbLayout結(jié)構(gòu)體以static_assert(sizeof(UbLayout) 248KB)限制總占用注釋明確說明UB 總 size 為 256KB但需要為 AscendC 和 VF 預(yù)留 8K業(yè)務(wù)只能使用 248K。L0C 的 4×64KB 輪轉(zhuǎn)結(jié)構(gòu)、L0A/L0B 各 2×32KB 雙緩沖也都能在 flash_attn_block_cube_nd.h 的 buffer 分配表中找到一一對(duì)應(yīng)的實(shí)現(xiàn)見 §4。3. Cube Bound 的達(dá)成條件性能側(cè)約束3.1 時(shí)間模型按單次迭代總量計(jì)。兩個(gè) matmul 的 FMA 總量為2·M·N·D$$ T_{cube} \frac{2 M N D}{FMA} $$Q 塊M·D僅在首次迭代搬入一次S2 足夠大時(shí)被攤薄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 結(jié)果fp32由 L0C fixpipe 到 UB$$ T_{fixpipe} \frac{(M N M D) \cdot \text{sizeof}(fp32)}{UB_{bw}} $$3.2 條件一M 的下限對(duì) MTE2 的遮擋$T_{cube} T_{mte2}$ 兩側(cè)約去 N、D與 D、N 無關(guā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 時(shí)無論 L2 命中多好都必然 mte2 boundM ∈ [82, 265) 時(shí)需 L2 命中率滿足$$ h \ge \frac{8192/M - 31}{101 - 31} $$M 越大對(duì) h 的要求越低M96 需 78%M128 需 47%M ≥ 265 起無要求。3.3 條件二N 的下限對(duì) Fixpipe 的遮擋$T_{cube} T_{fixpipe}$ 約去 M與 M 無關(guān)$$ 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 時(shí)UB_bw 32 8192/NN128 需 UB_bw 96N256 需 UB_bw 64。N 越大bmm2 輸出M×D與 N 無關(guān)被越多 KV 行攤薄fixpipe 壓力越小——該條件方向上偏好大 N。UB_bw 未列入硬件表見 §8 待補(bǔ)充若其與 L2 帶寬同級(jí)~101則 N ≥ 119 才能滿足N128 可以滿足。源碼佐證Fixpipe 壓力方向D256 下 fixpipe 輸出量大Kernel 側(cè)在 flash_attn_block_cube_nd.h 中通過IterateBmm2l0Split將 BMM2 沿 dV 方向按 128 列分塊nLoops循環(huán)、splitN逐 tile 占用輪轉(zhuǎn)的 64KB L0C buffer 并獨(dú)立 Fixpipe——單個(gè) L0C buffer 無需容納整塊128×256×4B結(jié)果這正是 §4.1 中bmm2 結(jié)果可沿 dV 方向分塊適配 buffer 容量的實(shí)現(xiàn)落點(diǎn)。3.4 數(shù)值示例M128, N128, D256項(xiàng)數(shù)值周期cube2×128×128×256 8.39M FMA2048KV 搬入131072 Bh100%: 1298 / h47%: 2048臨界/ h0: 4228fixpipe 搬出(1282128×256)×4 196608 BUB_bw96: 2048臨界/ 128: 1536注意D256 時(shí) fixpipe 搬出量196.6K/迭代大于 KV 搬入量131.1K/迭代fixpipe 是僅次于 MTE2 的第二約束。對(duì)照 (M64, N256)T_cube同為 2048M·N 相同但 KV 搬入 262144 B即使 h100% 也需 2596 cycles 2048——M64 82必然 mte2 bound的直接體現(xiàn)。4. 空間約束容量側(cè)約束各級(jí)存儲(chǔ)需容納基本塊的工作集并以多緩沖輪轉(zhuǎn)維持流水重疊。以下按bmm1 結(jié)果緩沖與 bmm2 結(jié)果緩沖各自獨(dú)立分配的常規(guī)設(shè)計(jì)推導(dǎo)§7 將放寬該假設(shè)。4.1 L0Cbmm 結(jié)果 bufferL0C256K以 4 份輪轉(zhuǎn) buffer 容納 bmm1 結(jié)果M·N·4與 bmm2 結(jié)果M·D·4每份 64Kbmm2 結(jié)果可沿 dV 方向分塊適配 buffer 容量分塊數(shù) $t \lceil M D \cdot 4 / 64K \rceil$為保 Cube→Fixpipe 重疊每迭代消耗的 buffer 數(shù)1 t須小于總 buffer 數(shù) 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 總量矩陣乘按分塊累加代價(jià)是 L0C 輪轉(zhuǎn)深度從 2.0t1降到 4/3t2、fixpipe 調(diào)用次數(shù)翻倍但總量不變。4.2 L1輸入分片駐留分片緩沖大小Q 塊雙緩沖2 × M·D·2 4MDPsoftmax 結(jié)果回寫供 bmm2三緩沖3 × M·N·2 6MNK/V分時(shí)復(fù)用同組 buffer至少雙 buffer≥ 2 × 2 × N·D·2 4NDbuffer 份數(shù)越多流水越深$$ 4MD 6MN 4ND \le 512K $$源碼佐證L1 分配策略flash_attn_block_cube_nd.h 的 L1 buffer 分配與此表逐項(xiàng)對(duì)應(yīng)P 矩陣 3 個(gè) buffer3×mBaseSize×s2BaseSize×sizeof(INPUT_T)Q 矩陣 2 個(gè)gS1 行內(nèi)復(fù)用、isFirstS2Loop加載/isLastS2Loop釋放K/V 4×64KBconfig5 即 s2BaseSize256 且 D128 時(shí)改用 2×128KB 大緩沖。P 3 個(gè)、Q 2 個(gè)、KV 至少 2 個(gè)的多緩沖輪轉(zhuǎn)深度正是為了讓 MTE2 與 Cube 在 PRELOAD_N2 流水flash_attn_kernel_dn.h 中創(chuàng)建任務(wù)與執(zhí)行任務(wù)分離、執(zhí)行慢于創(chuàng)建 2 輪下連續(xù)重疊。4.3 UB結(jié)果緩沖與 Vector 側(cè)工作集buffer大小bmm1 結(jié)果fp32雙緩沖2 × (M/2)·N·4 4MNbmm2 結(jié)果fp32雙緩沖2 × (M/2)·D·4 4MDPbf16雙緩沖1 行分形對(duì)齊 padding≈ 2 × (M/2)·N·2 2MN輸出累加暫存fp32單緩沖(M/2)·D·4 2MD掩碼bool雙緩沖2 × (M/2)·N·1 MNsoftmax 標(biāo)量max/sum/exp 各三緩沖 lse雙緩沖 雜項(xiàng)≈ 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 時(shí)每個(gè)輸入分片恰好 32K 填滿一個(gè)子緩沖自然滿足M 或 N 過大將迫使輸入側(cè)退化單緩沖或重復(fù)加載故 L0A/L0B 不單獨(dú)給出不等式但支持 M、N ≤ 128 的上界。源碼佐證分形與子緩沖Kernel 側(cè) L0A/L0B 由BufferManager各分配 2×32KB 雙 buffer見 flash_attn_block_cube_nd.h 的l0aBufferManager_/l0bBufferManager_配合BuffersPolicyDB的 LOCK_UNLOCK 策略讓 MTE1 搬運(yùn)L1→L0與 PIPE_M 計(jì)算重疊。D72 復(fù)用 D128 的 config 變體時(shí)L1→L0 搬運(yùn)寬度按 16 元素分形對(duì)齊72→80Mmad 按真實(shí) D 精確累加見 README 基本塊一節(jié)說明16 分形對(duì)齊是貫穿 D 各檔位的實(shí)現(xiàn)約束。5. D256 代入(M, N) 的可行域?qū)?D256 代入上述約束分離緩沖設(shè)計(jì)約束不等式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聯(lián)合82 M ≤ 102M ≤ 64無 cube bound 窗口兩點(diǎn)觀察N256 被 L0C 鎖死MN ≤ 16384使 M ≤ 64 82無論其他約束如何放松都到不了 cube bound 閾值。N128 的窗口極窄可行域 (82, 102] 緊貼閾值 82——即使取到上界 M102也要求 L2 命中率 ≥ 70%裕度極小。而塊大小需按 16 分形對(duì)齊、實(shí)用檔位取 2 的冪M ∈ {64, 128}M128 超出 UB 上限 102311K 248KM64 低于閾值 82——分離緩沖設(shè)計(jì)下D256 在常規(guī)檔位上不存在 cube bound 可行解。退而求其次只能取 (64, 256)空間全部滿足UB 212K、L1 416K但 M64 82注定 mte2 bound——這就是 D256 性能遠(yuǎn)差于 D128 的理論根源。源碼佐證檔位與 config 映射模型中的常規(guī)檔位 M ∈ {64, 128}并非隨意設(shè)定。Host 側(cè) fa_adjust_sinner_souter.h 的AdjustSinnerAndSouter對(duì) D256 只有兩檔選擇gSize × maxSeqQ ≥ 64時(shí)取 sOuter64 / sInner128否則取 sOuter32 / sInner256再結(jié)合mBaseSize sOuter × CV_RATIOCV_RATIO2即得 M ∈ {64, 128}、N ∈ {128, 256}。這兩個(gè)檔位在編譯期模板 flash_attn_template_tiling_key.h 中編碼為 config4sOuter64, sInner128, D256與 config5sOuter32, sInner256, D256并在 README 基本塊表 中完整列出。也就是說M 取 64 還是 128在實(shí)現(xiàn)層面就等價(jià)于選 config 4 還是 config 5二者被 tiling key 的 Config 位段bits 25-28固化。6. 可行域內(nèi) M、N 對(duì)性能的影響上節(jié)給出的是不等式范圍而非定值范圍內(nèi)如何取值取決于 M、N 各自的性能影響。6.1 M在可行域內(nèi)越大越好算術(shù)強(qiáng)度每 KV 字節(jié)的 FMA 數(shù)為M/2M 是唯一的算術(shù)強(qiáng)度旋鈕。M 越大條件一越寬松對(duì) L2 命中率的要求越低M102 需 70%M128 需 47%cube bound 越穩(wěn)Q 攤薄Q 搬入的攤薄負(fù)擔(dān)相對(duì) cube 時(shí)間與 M 無關(guān)兩者同隨 M 線性增長(zhǎng)不構(gòu)成偏好代價(jià)尾塊浪費(fèi)Q 側(cè)總行數(shù)不整除 M 時(shí)末塊空轉(zhuǎn)與任務(wù)粒度變粗。Q 側(cè)行數(shù)充足prefill / 長(zhǎng)序列時(shí)代價(jià)可忽略Q 側(cè)總行數(shù)不足 M 時(shí)大塊大面積空轉(zhuǎn)必須縮小 M——此時(shí) M 反正低于 82cube bound 不可達(dá)屬另一種取值邏輯見 §8。結(jié)論Q 側(cè)行數(shù)充足時(shí)M 取可行域上界。源碼佐證Q 側(cè)行數(shù)充足的判據(jù)AdjustSinnerAndSouter中 D256 分支的判據(jù)是gSize × maxSeqQ ≥ 64——即Q 側(cè)有效行數(shù)Q 序列長(zhǎng)度 × GQA group 數(shù)是否夠填滿一個(gè)大塊。這與本文Q 側(cè)行數(shù)充足與否決定 M 取值邏輯的結(jié)論精確對(duì)應(yīng)充足則選大 MsOuter64 → M128不足則降檔sOuter32 → M64。此外 tiling 的 gS1 合軸gS1Size actSeqLensQ × gSize見 README正是把 Q 序列與 GQA group 兩軸合并為同一任務(wù)維M 的分塊實(shí)際作用于 gS1 軸。6.2 N雙向影響且與 M 在 L0C 上耦合大 N 的收益條件二fixpipe更寬松KV 迭代次數(shù)S2/N減少循環(huán)與同步開銷攤薄小 N 的收益UB/L0C 中的 MN 項(xiàng)、L1 中的 N·D 項(xiàng)都隨 N 收縮K/V 單個(gè) L1 分片更小N128 時(shí) 64K/片可 4 個(gè) buffer 深輪轉(zhuǎn)N256 時(shí) 128K/片僅 2 個(gè) bufferMTE2 與 Cube 的重疊更深S2 不整除 N 時(shí)尾塊浪費(fèi)更小。關(guān)鍵耦合L0C 約束MN ≤ 16384把 M 與 N 綁在同一份預(yù)算里——N 每放大一檔M 的上限就減半。條件一M ≥ 82是 cube bound 的唯一入口而條件二只是過線即可N128 在 UB_bw ≥ 96 時(shí)已滿足。因此選擇次序是先保 M 過閾值并盡量大N 在剩余預(yù)算內(nèi)取滿足條件二的值。D256 下即 N128、M 取上界。7. 緩沖冗余的發(fā)現(xiàn)bmm1/bmm2 結(jié)果共用后可行域擴(kuò)大7.1 冗余性論證§4.3 中 bmm1 結(jié)果緩沖4MN與 bmm2 結(jié)果緩沖4MD各自獨(dú)立雙緩沖但兩者的生命周期并不重疊bmm1(L) 結(jié)果 → softmax(L) 消費(fèi)完畢 → 該 buffer 死亡 bmm2(L) 結(jié)果 → 累加(L) 消費(fèi)完畢 → 該 buffer 死亡數(shù)據(jù)流是環(huán)形鏈bmm1 → softmax → bmm2 → 累加 → bmm1(下一輪)同一迭代內(nèi) softmax 先讀完 bmm1 結(jié)果、bmm2 才開始產(chǎn)出bmm2 的結(jié)果要到下一輪 softmax 前才被累加讀完。因此兩組雙 buffer 可以共用同一對(duì)輪轉(zhuǎn) bufferbuffer 容量取兩者的較大值(M/2)·max(N,D)·4。D256 ≥ N 時(shí)bmm1 結(jié)果(M/2)·N·4完全裝得進(jìn) bmm2 的 buffer——bmm1 的雙 buffer4MN整體消失。源碼佐證實(shí)現(xiàn)即共用這一理論上發(fā)現(xiàn)的冗余在 Kernel 中已是實(shí)際設(shè)計(jì)——flash_attn_block_cube_nd.h 與 flash_attn_block_vec_nd.h 中的ubMmResBuffers_明確注釋BMM1/BMM2 fixpipe 結(jié)果分時(shí)復(fù)用buffer 大小取mBaseSize/2 × max(s2BaseSize, dVBaseSize) × sizeof(FP32)即本文的(M/2)·max(N,D)·4buffer 數(shù)按 D 分檔UB_MM_RES_BUFCNT (dBaseSize 128) ? 2 : 4。讀寫兩側(cè)通過一 flag 一 buffer的核間同步接力AIV 側(cè)CC_MM_0~3/ AIC 側(cè)CROSSCORE_MM_0~3寫端BMM1/BMM2fixpipe 完成后 Set結(jié)果就緒讀端Vec1/Vec2Wait 后讀取、讀畢再 Setbuffer 釋放同一 flag 沿「就緒 → 釋放 → 就緒」交替接力。mmResBufId_單調(diào)遞增計(jì)數(shù)器保證同一 ExecuteTask 內(nèi)先后執(zhí)行的 BMM1(L) 與 BMM2(L-2) 落在不同 buffer避免寫寫沖突——這正是 §7.3 所述按結(jié)果序而非算子序分配 buffer、以環(huán)形等待保證時(shí)序的工程實(shí)現(xiàn)。7.2 新的 UB 空間公式$$ \underbrace{4M \cdot \max(N,D)}_{\text{共用結(jié)果緩沖}} 2MN 2MD MN 54M 512 \le 248K $$D256 且 N ≤ 256 時(shí) $\max(N,D)D$化簡(jiǎn)為$$ \mathbf{3MN 6MD 54M 512 \le 248K} \qquad (\text{原 } 7MN 6MD \dots) $$節(jié)省量 4MNN ≤ D 時(shí)。對(duì) (M128, N128)節(jié)省 64KUB 占用從 311K超限降為247.25K ≤ 248K余 768B。7.3 擴(kuò)大后的可行域與更好的選擇約束N128N256UB共用緩沖M ≤ 128M ≤ 107L0C / L1M ≤ 128 / M ≤ 219M ≤ 64/ M ≤ 102聯(lián)合82 M ≤ 128M ≤ 64仍無窗口與 §6 的結(jié)論呼應(yīng)N128 的窗口從 (82, 102] 擴(kuò)大到(82, 128]恰好把 M128 檔位納入對(duì) L2 命中率的要求從 70% 降到 47%cube bound 從勉強(qiáng)貼線變?yōu)橛性６取@就是緩沖共用后能選出更好 (M, N) 的量化體現(xiàn)N256 側(cè)不變被 L0C 的 MN ≤ 16K 鎖死在 M ≤ 64進(jìn)一步印證 §6.2 的判斷L0C 耦合下 N 必須給 M 讓路共用的前提是按結(jié)果序而非算子序分配 buffer并以環(huán)形等待保證bmm1 第 L 輪寫入晚于bmm2 第 L?2 輪讀完——即空間收益需要同步機(jī)制的配合不是免費(fèi)的。源碼佐證環(huán)形等待與流水Kernel 的 PRELOAD_N2 流水flash_attn_kernel_dn.h 的ExecuteTask保證同一時(shí)刻 AIC 上 BMM1(本輪) 與 BMM2(2 輪前) 并行、AIV 上 Vec1(本輪) 與 Vec2(2 輪前) 并行P 矩陣的 3 buffer 輪轉(zhuǎn)AIV 寫第 N 輪 P 與 AIC 讀第 N-2 輪 P 并行恰好容納寫入中/待讀取/已釋放三態(tài)。這些都建立在結(jié)果生命周期錯(cuò)開這一前提之上與本文的環(huán)形鏈論證互相印證。8. 結(jié)論場(chǎng)景選擇依據(jù)D256Q 側(cè)行數(shù)充足prefillM128, N128落在共用 buffer 后的可行域 (82, 128] 上界L2 命中率 ≥ ~47% 即 cube boundbmm2 結(jié)果按 dV/128 分塊適配 L0C bufferD256Q 側(cè)總行數(shù) 128decode/小 batchM64, N256M 被工作量壓低cube bound 反正不可達(dá)M 82取大 N 換取 KV 迭代次數(shù)減半且 (64,256) 恰好貼住MN 16K的 L0C/mask 界空間無浪費(fèi)建模方法論小結(jié)性能條件給下限cube bound 等價(jià)于M 8192/Mix_bw82~265與 D/N 無關(guān)與N 8192/(UB_bw ? 8192/D)空間約束給上限L0CMN ≤ 16384、M ≤ 128dV 分塊后L14MD6MN4ND ≤ 512KUB 由結(jié)果緩沖與 Vector 工作集決定可行域內(nèi) M 越大越好唯一算術(shù)強(qiáng)度旋鈕N 與 M 在 L0C 上共享預(yù)算、優(yōu)先保 M緩沖復(fù)用擴(kuò)大可行域bmm1/bmm2 結(jié)果生命周期錯(cuò)開、可共用 bufferUB 上界 M ≤ 102 → M ≤ 128使 cube bound 從臨界變?yōu)榭尚小臻g優(yōu)化最終轉(zhuǎn)化為性能選項(xiàng)的解放。待補(bǔ)充UB 帶寬 UB_bw 實(shí)測(cè)值用于封閉條件二D256 下 fixpipe 搬出量大于 KV 搬入量值得單獨(dú) profiling 驗(yàn)證。附模型結(jié)論與實(shí)現(xiàn)的對(duì)應(yīng)關(guān)系及驗(yàn)證途徑本文的所有推導(dǎo)最終都落在 CANN ops-transformer 倉庫 FlashAttn 算子的真實(shí)實(shí)現(xiàn)中可按下述對(duì)應(yīng)關(guān)系在源碼中逐一驗(yàn)證基本塊檔位M、N 的取值集合編譯期 config 模板 flash_attn_template_tiling_key.hconfig 0~7 的 sOuter/sInner/D/DV 編碼與運(yùn)行期AdjustSinnerAndSouterfa_adjust_sinner_souter.h共同決定mBaseSize sOuter × 2與s2BaseSize sInnertiling key 編碼5 個(gè)模板參數(shù)InOutLayoutType / KvLayoutType / HasAttenMask / TemplateId / Config按位拼接為 tiling keyHost tiling 的GenTilingKey生成、運(yùn)行期 kernel 按 key 選變體細(xì)節(jié)見 README 的 TilingKey 章節(jié)緩沖與存儲(chǔ)層級(jí)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 分配表負(fù)載均衡與 FDsection 切分與 split-K 歸約FlashDecode由flash_attn_metadataAICPU基于 SectionStreamK 算法產(chǎn)出 metadata主算子按 metadata 區(qū)間執(zhí)行見 flash_attn_tiling_data.h 的 metadata 布局與 README 的負(fù)載均衡章節(jié)性能驗(yàn)證途徑倉庫提供 train/infer 各 220 性能用例tests/pytests/test_cases/performance_redline_train.py、performance_redline_infer.py可通過--perf_mode采集性能數(shù)據(jù)并配合 ASCEND DEBUG 日志中的GenTilingKey/kernel_nameFlashAttn_hash_tiling_key確認(rèn)實(shí)際選用的檔位——感興趣的讀者可在真機(jī)上用這些用例驗(yàn)證本文UB_bw 實(shí)測(cè)值待補(bǔ)充的開放項(xiàng)。此外FlashAttn 的對(duì)外接口flash_attn_metadataflash_attn兩段式調(diào)用、參數(shù)說明、約束與調(diào)用示例見 torchapi_flash_attn.md可與本文的性能建模對(duì)照理解接口層的 D 支持范圍64/72/128/256以及 MLA 形態(tài)的 QK192/V128與模型假設(shè)的 D256 場(chǎng)景一致接口文檔中metadata 與主算子入?yún)⒈仨毐３忠恢碌募s束則對(duì)應(yīng)本文 §4 中tiling 與 metadata 共用AdjustSinnerAndSouter保證兩側(cè)切分口徑一致的實(shí)現(xiàn)要求?！久赓M(fèi)下載鏈接】ops-transformer本項(xiàng)目是CANN提供的transformer類大模型算子庫實(shí)現(xiàn)網(wǎng)絡(luò)在NPU上加速計(jì)算。項(xiàng)目地址: https://gitcode.com/cann/ops-transformer創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考