23 KiB
BMM 尾轮处理策略对比分析:不重切 / 尾轮重切 / 整轮均匀重切
目标芯片:昇腾 950PR(DAV_3510)。本文自包含——不依赖其他分析文档,全部推导在文内完成;仅引用外部公开资料(昇腾白皮书、CANN 文档、开源算子源码仓)。 版本:v1.0,2026-08-31
摘要
BMM(BatchMatMul)算子在核间切分 M/N 后,总输出块数 N_{blk} = B \cdot mCnt \cdot nCnt 一般不是核数 C 的整数倍,尾轮只有 r = N_{blk} \bmod C 个核工作、其余核空转。本文对四种尾轮处理策略做严格的两两对比:A0(不重切)、A1a(尾轮整数倍切分)、A1b(尾轮 tile 重选凑满核)、B(整轮均匀重切),在计算 Bound 与访存 Bound 两类场景下给出完整的时延公式与适用条件分界。
核心结论:
- A0 从来不是最优(
r > 0时 A1 或 B 严格优); - A1a 是 A1b 的真子集(枚举空间包含关系),A1b 恒不劣于 A1a;
- 计算 Bound:A1b 与 B 理论时延严格相等(总计算量/C 守恒),A1b 搬移周长和恒 ≤ B(均值不等式)——时延打平、按工程权衡选择;
- 访存 Bound:$\rho \ge \rho_{dv}$(尾轮占比 ≥ dValue 平方阈值)时 A1b 严格优;
\rho < \rho_{dv}时 A1b 被 dValue 卡死,B 严格优当且仅当 $\sqrt{\rho_{dv}} > n_{wave}/\sqrt{g} - (n_{wave}-1)$——B 反超的唯一通道是 dValue:A1b 的尾轮 tile 缩放\sqrt{\rho}恒狠于 B 的 $1/\sqrt{g}$(\rho < 1/g恒成立),卡死更深时 B 胜出,该区间窄且临界实测打平; - 给出基于 B/M/K/N/dtype 的闭式判定流程:五步前置计算 + 决策表,无需逐项仿真。
一、问题背景与执行模型
1.1 BMM 的核间切分与数据流
BatchMatMulV3 在昇腾 NPU 上的兜底分支(ASW_Basic)按输出块并行:输出平面 M \times N 被切为 mCnt \times nCnt 个块,乘上 batch 维共 N_{blk} = B \cdot mCnt \cdot nCnt 个独立输出块,按序分配到 C 个 AIC 核。每个核内执行标准 Cube 数据流:
GM ──MTE2──> L1 ──MTE1──> L0A/L0B ──MMAD──> L0C ──Fixpipe──> GM
↑_____________ L2 Cache(读 5.2TB/s)_____________↑
核间不切 K(singleCoreK = K),每个输出块独立累加完成后经 Fixpipe 写出 GM。块与块之间无依赖——这是尾轮可以独立重切的前提。
1.2 尾轮问题
块按序分配:核 i 处理块 $i, i+C, i+2C, \ldots$。总块数 N_{blk} 不能被 C 整除时,最后一轮(尾轮)只有 r = N_{blk} \bmod C 个核有块可算,C - r 个核空转一个整块时间。
问题:给定 B、M、K、N、dtype 与首轮切分(singleCoreM/N),如何处理尾轮使端到端时延最小?
1.3 硬件规格与符号定义
昇腾 950PR 关键规格(昇腾 950 NPU 架构白皮书):
| 符号 | 含义 | 取值 |
|---|---|---|
C |
AIC 核数 | 32 |
Q_{16} |
单核 Cube BF16 峰值算力 | 486/32 ≈ 15.2 TFLOPS |
W_{GM} |
GM(HBM)总带宽 | 1.6 TB/s |
BW_{pc} |
单核 GM 带宽份额 | W_{GM}/C = 50 GB/s |
L1 |
L1 容量/核 | 512 KB |
L0C |
L0C 容量/核 | 256 KB |
dtype / outB |
输入/输出元素字节数 | BF16 → 2B |
Tiling 符号:sM, sN = 单核输出 tile 的 M/N 维度(singleCoreM/N 的简写);k_{L1} = GM→L1 的 K 向粒度;$mCnt = \lceil M/sM \rceil$、$nCnt = \lceil N/sN \rceil$。
尾轮符号:$N_{blk} = B \cdot mCnt \cdot nCnt$(总块数);$n_{wave} = \lceil N_{blk}/C \rceil$(总轮次);$r = N_{blk} \bmod C$(尾轮块数);$\rho = r/C$(尾轮占比)。
二、基础模型
2.1 单块时延三项
单个输出块 $[sM, sN]$(K 维全量累加)的三段时延:
T_{MMAD} = \frac{2 \cdot sM \cdot sN \cdot K}{Q_{16}},\qquad T_{MTE2} = \frac{(sM + sN) \cdot k_{L1} \cdot dtype}{BW_{pc}},\qquad T_{FIX} = \frac{sM \cdot sN \cdot outB}{BW_{pc}}
其中 $k_{L1} = \min\big(K,; \lfloor L1/(2(sM{+}sN) \cdot dtype) \rfloor_{16}\big)$(L1 双缓冲容量约束)。单块时延取主导项:$T_{block} = \max(T_{MMAD}, T_{MTE2}, T_{FIX})$(三段流水掩盖,瓶颈项决定)。
Bound 判据(由 T_{MMAD} = T_{MTE2} 解出):
K^* = \frac{k_{L1} \cdot dtype \cdot Q_{16}}{2 \cdot BW_{pc}}\Big(\frac{1}{sM} + \frac{1}{sN}\Big)
K \ge K^* 为计算 Bound,否则访存 Bound(例:$sM=sN=256$、$k_{L1}=256$、BF16 时 $K^* \approx 608$)。
2.2 块几何缩放律(全文分析的枢纽)
把块数放大 g 倍(面积缩 g 倍、线性尺寸缩 \sqrt{g} 倍,近方形比例)时:
| 时延项 | 依赖 | 缩放律 |
|---|---|---|
T_{MMAD} |
\propto 面积 sM \cdot sN |
\propto g^{-1} |
T_{MTE2} |
\propto 周长 (sM + sN) |
\propto g^{-1/2} |
T_{FIX} |
\propto 面积 |
$\propto g^{-1}$(总量与切分无关) |
搬移随周长缩放是关键:切分越多,计算缩 g 倍而搬移只缩 \sqrt{g} 倍——"切分越多重复读越多"在每块粒度上的体现(输出是面积、输入搬运是周长:A 行带 sM \times K + B 列带 $K \times sN$)。
2.3 搬移效率约束
GM→L1 搬移(Nd2Nz DMA)的两级效率阈值:
- dValue ≥ 256B(DMA 硬件突发下限):dValue 是 ND 排布中连续维的字节数——B 矩阵非转置
[K, N]时连续维为 N,dValue = $sN \cdot dtype$;B 转置[N, K]时连续维为 K,dValue = $k_{L1} \cdot dtype$;A 非转置[M, K]的 dValue = $k_{L1} \cdot dtype$(A 的 dValue 由 K 向保证,不约束 $sM$;B 非转置时sN有 256B 下限——约束非对称); - 单次搬移量 ≥ min_TileSize(推荐 16KB):如 $k_{L1} \cdot sN \cdot dtype \ge 16$KB。
2.4 带宽模型声明
昇腾 950PR 每核 MTE2 为独立 DMA 引擎、带宽按核数配平,建模为每核带宽上限 $BW_{pc} = W_{GM}/C$。该假设下尾轮 r 核聚合带宽仅 $r \cdot BW_{pc} < W_{GM}$,尾轮搬移不加速——r 个整块用 r 核、每核 1 块,时延仍为一个整块搬移时间 $T_{load} = T_{MTE2}$。HBM 全局共享池模型的敏感性见 §六边界说明。
2.5 首轮切分基准(近方形的来源与适用条件)
尾轮策略建立在首轮切分 (mCnt, nCnt, sM, sN) 之上。首轮由搬入时延最小化确定:$T_{MTE2}^{total} \propto (1/sM + 1/sN) \cdot k_{L1}$,在块数约束下枚举取最优。其连续极限为方形($sM = sN$):$k_{L1} = K$(K 全载)时,面积固定使周长最小化 → 均值不等式给出方形严格最优;k_{L1} 被 L1 容量压低(K 大)时,$(1/sM+1/sN) \cdot L1/(2(sM{+}sN)dtype) = L1/(2 \cdot sM \cdot sN \cdot dtype)$——只与面积有关、与长宽比无关,形状自由度释放。本文的缩放律以近方形为基准形态(线性尺寸 \sqrt{g} 缩放);非方形分解的实际收益由同一目标函数另行评估(§四的"非方形惩罚"即此体现)。
三、四种策略定义与实现方式
3.1 策略定义
| 策略 | 做法 | 块大小 |
|---|---|---|
| A0:不重切 | 尾轮 r 核各处理 1 个整块,C-r 核空转 |
主轮尾轮同大小 |
| A1a:尾轮整数倍切分 | 尾轮每块沿 N(或 M)切 s^* 份,r \cdot s^* 个小块分给 C 核 |
主轮整块 + 尾轮 1/s^* 小块 |
| A1b:尾轮 tile 重选凑满核 | 尾轮 r 个原块覆盖的区域(面积 $r \cdot sM \cdot sN$)用更小 tile (sM_t, sN_t) 重新切分,凑满 C 核 |
主轮整块 + 尾轮小 tile |
| B:整轮均匀重切 | 总块数向上取整到 $N_{blk}' = n_{wave} \cdot C$,全局重新枚举 tile 使每轮每核恰好一个同样大小的块 | 全部块同大小 |
实现方式对比:
| 策略 | host 侧 tiling | kernel 侧执行 | tile 参数套数 |
|---|---|---|---|
| A0 | 一套参数 | 线性映射 index = blockIdx + round × usedCoreNum,尾轮不满载核跳过 |
1 |
| A1a | 主参数 + 尾轮切分参数($s^*$) | 最后一轮切换小块尺寸 | 2 |
| A1b | 主参数 + 尾轮 tile($sM_t, sN_t$) | 尾轮区域独立子网格映射 | 2 |
| B | 一套重切后参数 | 全程统一 tile,无尾轮分支 | 1 |
对齐约束的平等性:16 对齐是 Cube 计算粒度(16×16×16 基本块)的硬件要求,四种策略的 tile——无论重切前后——都必须是 16 的倍数,无一例外。差异只在枚举空间(A1b 对尾轮子区域、B 对全局)与 tile 套数。
3.2 A1b 的 tile 确定与 A1a ⊆ A1b 证明
A1b 的理想 tile(方形同步缩小):
s_t^* = sM \cdot \sqrt{\frac{r}{C}} \quad (sM = sN \text{ 时})
再按 16 对齐调整,使重切块数 $\lceil r \cdot sM \cdot sN / (sM_t \cdot sN_t) \rceil \le C$;也允许非方形枚举((sM_t, sN_t) 独立按 16 步进)选最贴合的组合。
A1a 是 A1b 的真子集:A1a 的尾轮 tile 为 (sM, sN/s^*) 或 $(sM/s^*, sN)$(s^* \in \{2, \ldots, \lfloor C/r \rfloor\} 整数),恰为 A1b 枚举空间中 $sM_t = sM$、sN_t = sN/s^* 的特定组合——其 16 对齐、dValue、块数约束在 A1b 枚举中同样检查。记 A1b 枚举空间为 $\Omega_{A1b}$,A1a 解空间 $\Omega_{A1a} \subseteq \Omega_{A1b}$,枚举取最优:
T_{A1b} = \min_{\Omega_{A1b}} T \;\le\; \min_{\Omega_{A1a} \subseteq \Omega_{A1b}} T = T_{A1a}
A1b 恒不劣于 A1a,且通常严格优:A1a 受"s^* 整数"与"每块切同样份数"双重限制(r > C/2 时 s^* = \lfloor C/r \rfloor = 1 完全失效),A1b 允许非整数比例与凑满 C 核(无 r \le C/2 限制)。下文 A1 一律指 A1b(A1a 视为退化形态)。
3.3 三策略通用时延式
T_{A0} = n_{wave} \cdot T_{block},\qquad T_{A1} = (n_{wave}-1) \cdot T_{block} + T_{tail},\qquad T_B = n_{wave} \cdot T_{block}'
A1a 时 $T_{tail} = T_{block}/s^*$(计算 Bound);A1b 凑满 C 核时尾轮小块面积为 $\rho \cdot sM \cdot sN$,计算 Bound 下 $T_{tail}^{A1b} = \rho \cdot T_{block}$。
四、计算 Bound 完整推导($T_{MMAD} > T_{MTE2}$)
4.1 A0 vs A1
A1a 可行时($r \le C/2$):$\Delta_{A0 \to A1a} = T_{MMAD}(1 - 1/s^*)$。r > C/2 时 A1a 失效,但 A1b 不退化:
\Delta_{A0 \to A1b}^{calc} = T_{MMAD}\Big(1 - \frac{r}{C}\Big)
A1 恒优于 A0($r > 0$),A0 在计算 Bound 下永不最优。
4.2 A0 vs B
代入 T_{block}' = T_{MMAD}/g 与 $n_{wave}/g = N_{blk}/C = n_{wave}-1+\rho$:
\Delta_{A0 \to B}^{calc} = n_{wave} T_{MMAD} - T_{MMAD}\Big(n_{wave} - 1 + \frac{r}{C}\Big) = T_{MMAD} \cdot \frac{C - r}{C} > 0
B 恒优于 A0(可行性校验:搬移掩盖 $\sqrt{g} \le T_{MMAD}/T_{MTE2}$;分解对齐)。
4.3 A1 vs B
A1b 凑满核时 $T_{A1b} = T_{MMAD}(n_{wave} - 1 + \rho)$,与 T_B = T_{MMAD}(n_{wave} - 1 + r/C) 理论时延严格相等——两者都是"总计算量/C"(计算 Bound 下时延与切分方式无关,只要轮轮满载)。
结构性差异在搬移量(周长和):A1b 主轮保持大 tile、只缩尾轮;B 全局均匀缩小。设方形 tile 边长 $s$:
S_{A1b} = 2sC\big(n_{wave} - 1 + \sqrt{\rho}\big),\qquad S_B = 2sC\sqrt{n_{wave}\big(n_{wave} - 1 + \rho\big)}
\big(n_{wave}-1+\sqrt{\rho}\big)^2 \le n_{wave}(n_{wave}-1+\rho) \iff 2\sqrt{\rho} \le 1 + \rho \iff (\sqrt{\rho}-1)^2 \ge 0 \quad \checkmark
均值不等式:A1b 周长和恒 ≤ B(等号当 \rho = 1 即无尾轮)。搬移少意味着 L2 重复读少、掩盖余量更大。
结论(计算 Bound):A1b 与 B 理论时延严格相等;离散 16 对齐后互有胜负(数值依赖、无系统性优劣——对齐约束对两者平等);A1b 搬移周长和恒 ≤ B(结构性)。工程简洁选 B(一套 tile);追求搬移下限选 A1b。
4.4 数值实例
例 1($r > C/2$):B=1、M=N=1792、K=4096、BF16、C=32。首轮 $mCnt=nCnt=7$($sM=sN=256$),$N_{blk}=49$、$n_{wave}=2$、$r=17$、$\rho=0.53$。$T_{MMAD}=35.3\mu s$。
- A0:$70.6\mu s$;A1a:
s^* = \lfloor 32/17 \rfloor = 1失效; - A1b 非方形枚举 $(sM_t, sN_t) = (224, 160)$:尾轮
\lceil 17 \times 256^2/35840 \rceil = 32块恰好凑满,尾轮时延 $19.3\mu s$,$T_{A1b} = 54.6\mu s$; - B:$N_{blk}'=64=8\times8$、$sM'=sN'=224$,$T_B = 54.1\mu s$。
A1b 与 B 打平(54.6 vs 54.1,差 <1%);周长和此例恰好持平(28672 = 28672)。相对 A0 均省约 23%。
例 2(r \mid C 完美点):M=1536、N=2048、K=4096、BF16。$N_{blk}=48$、$r=16$。A1a:s^* = 2 = C/r 完美,$T_{A1a} = 53.0\mu s$;B:$T_B = 53.0\mu s$。A1a = B($1/s^* = r/C$),选 A1a(搬移增量小)。
五、访存 Bound 完整推导($T_{MTE2} \ge T_{MMAD}$)
主导项 $T_{block} = T_{load} = (sM + sN) \cdot k_{L1} \cdot dtype / BW_{pc}$。
5.1 A0 vs A1
A1a($r \le C/2$):尾轮沿 N 切 s^* 份,小块搬移 $(sM + sN/s^) k_{L1} dtype$(A 行带 [sM, k_{L1}] 每个小块都要完整搬一次,不随 s^* 缩小——结构性弱点),方形下 $T_{tail} = T_{load}(1+1/s^)/2$:
\Delta_{A0 \to A1a}^{mem} = \frac{T_{load}}{2}\Big(1 - \frac{1}{s^*}\Big)
受 dValue 硬约束 $s^* \le sN \cdot dtype/256\text{B}$(BF16、sN=256 时 $s^* \le 2$)。
A1b(任意 $r$):尾轮区域用 s_t = s\sqrt{\rho} tile 重切凑满 C 核,尾轮总搬移 $= C \cdot 2s\sqrt{\rho} \cdot k_{L1}$,C 核满载聚合带宽 $C \cdot BW_{pc}$:
T_{tail}^{A1b} = \sqrt{\rho} \cdot T_{load},\qquad \Delta_{A0 \to A1b}^{mem} = T_{load}\big(1 - \sqrt{\rho}\big) > 0
A1 恒优于 A0(dValue 允许时)。
5.2 A0 vs B
\Delta_{A0 \to B}^{mem} = n_{wave} T_{load}\Big(1 - \frac{1}{\sqrt{g}}\Big) > 0 \quad (g > 1)
B 恒优于 A0(dValue 约束 g \le (sN \cdot dtype/256\text{B})^2 满足时)。物理解释:尾轮 r 核聚合带宽仅 $r \cdot BW_{pc}$,搬移不加速;B 让所有轮次满核满带宽。
5.3 A1 vs B——完整分界推导
(i)无约束 regime($\rho \ge \rho_{dv}$,A1b 可凑满):总时延正比于周长和(满载轮聚合带宽相同),由 §4.3 的均值不等式:
\frac{T_{A1b}^{mem}}{T_B^{mem}} = \frac{n_{wave} - 1 + \sqrt{\rho}}{\sqrt{n_{wave}(n_{wave}-1+\rho)}} \le 1
A1b 恒不劣于 B,\rho 小时优势大($\rho=0.1$、n_{wave}=2 时优 11%)。
(ii)关键结构事实:A1b 的 dValue 约束恒比 B 更严。A1b 的尾轮 tile 缩放因子为 $\sqrt{\rho}$,B 的全局 tile 缩放因子为 $1/\sqrt{g} = \sqrt{(n_{wave}-1+\rho)/n_{wave}}$:
\rho < \frac{1}{g} \iff \rho \cdot n_{wave} < n_{wave} - 1 + \rho \iff \rho(n_{wave} - 1) < n_{wave} - 1 \iff \rho < 1 \quad \checkmark \text{(} r < C \text{ 恒成立)}
物理含义:A1b 只缩尾轮 r 块的区域去凑满 C 核,tile 必须缩得比 B 的全局缩放更狠——所以 A1b 的尾轮 tile 更早跌破 dValue 下限。这是 B 可能反超的唯一通道。
(iii)三区间判定($\rho_{dv} = (256\text{B}/(sN \cdot dtype))^2$):
| 区间 | 条件 | 结论 |
|---|---|---|
| I | $\rho \ge \rho_{dv}$(A1b 凑满可行;此时 1/g > \rho \ge \rho_{dv} 故 B 也可行) |
A1b 恒优(均值不等式) |
| II | $\rho < \rho_{dv} \le 1/g$(A1b 卡死、B 可行) | 分界公式判定(下) |
| III | $\rho_{dv} > 1/g > \rho$(都卡死) | A0/A1a 兜底 |
区间 II 中 A1b 退化为 $s_t = 256\text{B}/dtype$(dValue 下限),尾轮块数 $= r/\rho_{dv} < C$(凑不满),尾轮时延 $= \sqrt{\rho_{dv}} \cdot T_{load}$(小块搬移按 s_t/sN = \sqrt{\rho_{dv}} 缩放):
T_{A1b}^{dv} = T_{load}\big(n_{wave} - 1 + \sqrt{\rho_{dv}}\big),\qquad T_B = \frac{n_{wave}}{\sqrt{g}} \cdot T_{load}
B 严格优当且仅当:
\sqrt{\rho_{dv}} > n_{wave} \cdot \sqrt{\frac{n_{wave} - 1 + \rho}{n_{wave}}} - (n_{wave} - 1)
物理解读:A1b 卡死后尾轮时延被锁在 $\sqrt{\rho_{dv}} \cdot T_{load}$(dValue 下限决定),不再随 \rho 减小;而 B 的 n_{wave}/\sqrt{g} 随 \rho 减小而降低(全局重切幅度减小)。\rho 足够小(卡死足够深)时 B 反超。
5.4 数值实例
例 3(\rho 小、dValue 卡死,临界打平):M=N=1536、K=256、BF16。$mCnt=nCnt=6$、$N_{blk}=36$、$n_{wave}=2$、$r=4$、$\rho=0.125$、$\rho_{dv}=0.25$。$T_{load} = 5.24\mu s$。
- A0:$10.5\mu s$;A1a:$s^* = 2$(dValue 卡死),$T_{A1a} = 9.2\mu s$;
- A1b:$s_t^* = 256\sqrt{0.125} = 90.5 < 128$(dValue 下限)卡死,只能 $s_t = 128$,尾轮 16 块(半满载),$T_{A1b} = 7.9\mu s$(理论无约束值
7.1\mu s达不到); - B:$N_{blk}'=64$、$sM'=sN'=192$,$T_B = 7.9\mu s$。
分界公式核验:\sqrt{\rho_{dv}} = 0.5 vs $2/\sqrt{1.78} - 1 = 0.5$——两侧精确相等,实测 A1b = B = 7.9μs 打平 ✓(公式精确捕捉临界点)。
例 4($\rho \ge \rho_{dv}$,A1b 优):M=N=2304、K=256、BF16。$N_{blk}=81$、$n_{wave}=3$、$r=17$、$\rho=0.53$。\sqrt{\rho} \cdot sN \cdot dtype = 373\text{B} \ge 256\text{B} ✓。
- A0:$15.7\mu s$;A1a:
s^* = 1失效; - A1b:$s_t = 192$,尾轮 32 块,$T_{A1b} = 2 \times 5.24 + 3.93 = 14.4\mu s$;
- B:$N_{blk}'=96$、
(8,12)分解($sM'=288, sN'=192$),$T_B = 14.7\mu s$。
A1b 优 2.2%(周长和 $45056 < 46080$)。
例 5(\rho 更小、临界区):M=1280、N=1792、K=256、BF16。$mCnt=5, nCnt=7$、$N_{blk}=35$、$n_{wave}=2$、$r=3$、$\rho=0.094$。
- A1b:
s_t^* = 78.4 < 128卡死,$s_t=128$、尾轮 12 块,$T_{A1b} = 5.24 + 2.62 = 7.9\mu s$; - B:$N_{blk}'=64$、
(8,8)分解($sM'=160, sN'=224$),$T_B = 7.9\mu s$。
分界公式:\sqrt{\rho_{dv}} = 0.5 vs $2\sqrt{1.094/2} - 1 = 0.479$——公式判定 B 微优(理论差 0.11μs),实际 B 的非方形分解惩罚(周长 384 vs 方形 379)抵消理论优势,实测打平。B 优的区间存在但窄,且实际分解的非方形惩罚会进一步压缩。
六、决策总表与基于 B/M/K/N 的判定流程
6.1 决策总表
| 场景 | A0 vs A1 | A0 vs B | A1 vs B | 最优策略 |
|---|---|---|---|---|
计算 Bound,r \le C/2 且 r \mid C |
A1a 优 | B 优 | A1a = B(选 A1a,搬移少) | A1a |
计算 Bound,r \le C/2 且 r \nmid C |
A1a 优 | B 优 | B 略优(取整损失) | A1a/B 皆可 |
计算 Bound,r > C/2 |
A1b 优(不退化) | B 优 | 时延理论相等,离散打平(A1b 周长和 ≤ B) | A1b 或 B(工程简洁选 B,搬移下限选 A1b) |
访存 Bound,\rho \ge \rho_{dv} |
A1b 优 | B 优 | A1b 恒优(均值不等式) | A1b |
访存 Bound,\rho < \rho_{dv} 且分界公式成立 |
A1b 部分凑满 | B 优 | B 严格优 | B |
访存 Bound,\rho < \rho_{dv} 且分界公式不成立 |
A1b 部分凑满 | B 优 | 打平 | A1b/B 皆可 |
| dValue 全面卡死(B 也不可行) | A1a($s^* \ge 2$)或 A0 | B 不可行 | — | A1a/A0 |
6.2 基于 B/M/K/N/dtype 的直接判定流程
五步闭式前置计算后查表即得最优策略,无需逐项建模仿真:
- 首轮切分:按搬入时延最小化枚举得 $mCnt, nCnt, sM, sN, k_{L1}$;
- 尾轮参数:$N_{blk} = B \cdot mCnt \cdot nCnt$,$n_{wave} = \lceil N_{blk}/C \rceil$,$r = N_{blk} \bmod C$,$\rho = r/C$;
- Bound 判定:$K^* = \dfrac{k_{L1} \cdot dtype \cdot Q_{16}}{2 BW_{pc}}\Big(\dfrac{1}{sM} + \dfrac{1}{sN}\Big)$,
K \ge K^*→ 计算 Bound; - dValue 可行性:$\rho_{dv} = (256\text{B}/(sN \cdot dtype))^2$;$g = n_{wave}C/N_{blk}$;$g_{dv} = (sN \cdot dtype/256\text{B})^2$;
- 单块主导项:$T_{block} = \max(T_{MMAD}, T_{load})$。
决策表:
| # | 条件 | 最优策略 | 端到端时延 |
|---|---|---|---|
| 1 | r = 0 |
A0(无尾轮) | n_{wave} \cdot T_{block} |
| 2 | 计算 Bound,0 < r \le C/2 |
A1a,s^* = \lfloor C/r \rfloor |
T_{MMAD}(n_{wave} - 1 + 1/s^*) |
| 3 | 计算 Bound,r > C/2 |
A1b 或 B——时延理论相等 | T_{MMAD}(n_{wave} - 1 + \rho) |
| 4 | 访存 Bound,\rho \ge \rho_{dv} |
A1b | T_{load}(n_{wave} - 1 + \sqrt{\rho}) |
| 5 | 访存 Bound,\rho < \rho_{dv} 且 \sqrt{\rho_{dv}} > n_{wave}/\sqrt{g} - (n_{wave}-1) 且 g \le g_{dv} |
B | T_{load} \cdot n_{wave}/\sqrt{g} |
| 6 | 访存 Bound,\rho < \rho_{dv} 但分界公式不满足 |
A1b(部分凑满)或 B,打平 | T_{load}(n_{wave}-1+\sqrt{\rho_{dv}}) \approx T_B |
| 7 | 访存 Bound,B 也不可行($g > g_{dv}$) | A1a($s^* \ge 2$)或 A0 | T_{load}(n_{wave}-1) + T_{load}(1+1/s^*)/2 |
判定流程图:
[B, M, K, N, dtype]
│
▼
<首轮枚举 → mCnt, nCnt, sM, sN, kL1>
│
▼
<N_blk, n_wave, r, ρ = r/C>
│
├─ r = 0 ────────────────▶ A0(无尾轮)
▼
<K ≥ K* ?(计算 Bound)>
│
├─ 是 ── r ≤ C/2 ? ──┬─ 是 ─▶ A1a(s* = ⌊C/r⌋)
│ └─ 否 ─▶ A1b(s_t = sM√ρ ↓16,可非方形枚举)或 B
│ (时延相等;工程简洁选 B)
│
└─ 否(访存 Bound)── ρ ≥ ρ_dv ? ──┬─ 是 ─▶ A1b(周长和恒 ≤ B)
└─ 否 ─▶ 分界公式判定:
√ρ_dv > n_wave/√g −(n_wave−1) ? ── 是 ─▶ B
└─ 否 ─▶ A1b/B 打平
七、边界说明
- 带宽模型敏感性:访存 Bound 结论依赖"每核 MTE2 带宽上限 $BW_{pc}$"假设。若 HBM 为全局共享池(尾轮
r核可吃满 $W_{GM}$),A0 尾轮搬移时延已是\rho \cdot T_{load}接近理想,A1b/B 的搬移增量无带宽补偿,结论反转。昇腾 950PR 的 MTE2 为每核独立 DMA 引擎、带宽按核数配平,采用固定份额结论;临界 case 建议实测复核。 - A1b 的工程代价:需两套 tile 参数(主轮大 tile + 尾轮小 tile)与尾轮区域的边界处理(
r个原块的并一般为 L 形,按矩形分解重切);host 端多一次枚举,NPU 侧 kernel 需支持尾轮 tile 尺寸切换。这些复杂度不改变时延结论,但影响实现成本。 - 方形基准的适用条件:§2.5 已述——K 全载时方形严格最优、L1 主导时仅面积相关。非方形分解的实际收益由同一目标函数评估,例 5 的非方形惩罚即此体现。
- B 不整除 $N_{blk}'$:
B \nmid N_{blk}'时按 batch 分组切分、部分 batch 多一块,收益略降,判定不变。
参考文献
- 昇腾 950 NPU 架构白皮书,华为技术有限公司,2026
- cann-ops-nn 源码仓(BatchMatMulV3):kernel 侧尾轮线性映射见
op_kernel/arch35/batch_mat_mul_v3_asw_block_advanced.h(UpdateBasicIndex:index = newBlockIdx + roundIdx × usedCoreNum,if (index < totalCnt)跳过空转核——即本文的策略 A0)