Files
matmul-analysis/BMM/BMM算子优化分析_Release/ASW_Basic分支分析/BMM尾轮处理策略对比分析_v1.2.md

36 KiB
Raw Permalink Blame History

BMM 尾轮处理策略对比分析:不重切 / 尾轮重切 / 整轮均匀重切

目标芯片:昇腾 950PRDAV_3510。本文自包含——不依赖其他分析文档全部推导在文内完成仅引用外部公开资料昇腾白皮书、CANN 文档、开源算子源码仓)。 版本v1.22026-08-31§7.2 重写:损失公式逐步推导链、方形近似可靠性审视、主流 prefill/decode 场景实际损失分析) 历史v1.1 新增 §七 B1 无脑版门限与损失分析v1.0 首版(三策略两两对比与 B/M/K/N 判定流程) 历史v1.0 首版(三策略两两对比与 B/M/K/N 判定流程)


摘要

BMMBatchMatMul算子在核间切分 M/N 后,总输出块数 N_{blk} = B \cdot mCnt \cdot nCnt 一般不是核数 C 的整数倍,尾轮只有 r = N_{blk} \bmod C 个核工作、其余核空转。本文对四种尾轮处理策略做严格的两两对比:A0不重切、A1a尾轮整数倍切分、A1b尾轮 tile 重选凑满核、B整轮均匀重切,在计算 Bound 与访存 Bound 两类场景下给出完整的时延公式与适用条件分界。

核心结论:

  1. A0 从来不是最优r > 0 时 A1 或 B 严格优);
  2. A1a 是 A1b 的真子集枚举空间包含关系A1b 恒不劣于 A1a
  3. 计算 BoundA1b 与 B 理论时延严格相等(总计算量/C 守恒A1b 搬移周长和恒 ≤ B均值不等式——时延打平、按工程权衡选择
  4. 访存 Bound$\rho \ge \rho_{dv}$(尾轮占比 ≥ dValue 平方阈值)时 A1b 严格优;\rho < \rho_{dv} 时 A1b 被 dValue 卡死B 严格优当且仅当 $\sqrt{\rho_{dv}} > n_{wave}/\sqrt{g} - (n_{wave}-1)$——B 反超的唯一通道是 dValueA1b 的尾轮 tile 缩放 \sqrt{\rho} 恒狠于 B 的 $1/\sqrt{g}$\rho < 1/g 恒成立),卡死更深时 B 胜出,该区间窄且临界实测打平;
  5. 给出基于 B/M/K/N/dtype 的闭式判定流程:五步前置计算 + 决策表,无需逐项仿真;
  6. v1.1 新增 §七B 的简化形态 B1(无脑版:满足门限就直接整数轮满核切,门限闭式 $r>0 \land (K \ge K^* \lor g \le (sN \cdot dtype/256\text{B})^2)$)与 B0/B1 相对 A1b 的损失分析
  7. v1.2 重写 §7.2:损失公式逐步推导链(T_B/T_{A1b} = \sqrt{w(w-1+\rho)}/(w-1+\sqrt{\rho}) 的四步推导与 5.4%/16.6%/30% 的数值来源);方形近似可靠性审视(非方形惩罚因子 $p(\lambda)$、k_{L1} 跨 regime——理论值是损失下界极端 case 实际损失可达 ~30%主流 prefill/decode 场景实际损失M,N 为 256 倍数时 r=0 无尾轮;M 奇异或 decode 场景 w 大,损失 ≤2.5%~30% 仅出现在主流形状不会命中的边角 case——主流场景 B1 是安全的工程选择

一、问题背景与执行模型

1.1 BMM 的核间切分与数据流

BatchMatMulV3 在昇腾 NPU 上的兜底分支ASW_Basic按输出块并行。适用范围说明ASW_Basic 是兜底分支——能进入更靠前特殊分支K=0、转 Matmul/Mul、StreamK、MergeBatch、IterBatch、AL1/BL1 全载等)的 case 会被优先截胡、不会进入 ASW_Basic因此本文只讨论 ASW_Basic 场景(核间切 M/N的尾轮处理那些在特殊分支下被承接的 case 不在本文对比范围。输出平面 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_____________↑

核间不切 KsingleCoreK = 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} GMHBM总带宽 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的两级效率阈值

  1. dValue ≥ 256BDMA 硬件突发下限dValue 是 ND 排布中连续维的字节数——B 矩阵非转置 [K, N] 时连续维为 NdValue = $sN \cdot dtype$B 转置 [N, K] 时连续维为 KdValue = $k_{L1} \cdot dtype$A 非转置 [M, K] 的 dValue = $k_{L1} \cdot dtype$A 的 dValue 由 K 向保证,不约束 $sM$B 非转置时 sN 有 256B 下限——约束非对称
  2. 单次搬移量 ≥ 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或 Ms^* 份,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/2s^* = \lfloor C/r \rfloor = 1 完全失效A1b 允许非整数比例与凑满 C 核(无 r \le C/2 限制)。下文 A1 一律指 A1bA1a 视为退化形态)。

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^*$(计算 BoundA1b 凑满 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 重复读少、掩盖余量更大。

结论(计算 BoundA1b 与 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$A1as^* = \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%

例 2r \mid C 完美点)M=1536、N=2048、K=4096、BF16。$N_{blk}=48$、$r=16$。A1as^* = 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 恒优于 A0dValue 允许时)。

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 恒优于 A0dValue 约束 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$A1as^* = 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$。

  • A1bs_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 最优策略
计算 Boundr \le C/2r \mid C A1a 优 B 优 A1a = B选 A1a搬移少 A1a
计算 Boundr \le C/2r \nmid C A1a 优 B 优 B 略优(取整损失) A1a/B 皆可
计算 Boundr > 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 的直接判定流程

五步闭式前置计算后查表即得最优策略,无需逐项建模仿真:

  1. 首轮切分:按搬入时延最小化枚举得 $mCnt, nCnt, sM, sN, k_{L1}$
  2. 尾轮参数$N_{blk} = B \cdot mCnt \cdot nCnt$$n_{wave} = \lceil N_{blk}/C \rceil$$r = N_{blk} \bmod C$$\rho = r/C$
  3. 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
  4. 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$
  5. 单块主导项$T_{block} = \max(T_{MMAD}, T_{load})$。

决策表

# 条件 最优策略 端到端时延
1 r = 0 A0无尾轮 n_{wave} \cdot T_{block}
2 计算 Bound0 < r \le C/2 A1as^* = \lfloor C/r \rfloor T_{MMAD}(n_{wave} - 1 + 1/s^*)
3 计算 Boundr > 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 访存 BoundB 也不可行($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 ? ──┬─ 是 ─▶ A1as* = ⌊C/r⌋
    │                    └─ 否 ─▶ A1bs_t = sM√ρ ↓16可非方形枚举或 B
    │                              (时延相等;工程简洁选 B
    │
    └─ 否(访存 Bound── ρρ_dv ? ──┬─ 是 ─▶ A1b周长和恒 ≤ B
                                       └─ 否 ─▶ 分界公式判定:
                                            √ρ_dv > n_wave/√g (n_wave1) ? ── 是 ─▶ B
                                                                             └─ 否 ─▶ A1b/B 打平

七、B 方案的简化形态 B1 与 B0/B1 相对 A1b 的损失分析

前文方案 B整轮均匀重切与 A1b 的分界判定需要逐 case 计算 $\rho$、$\rho_{dv}$、分界公式——本节回答两个工程问题:①能否给 B 一个无脑版门限(满足条件就直接整数轮满核切,不做 A1b/B 分界)?②若完全采用 B无论无脑版 B1 还是完整版 B0相对理论最优的 A1b 会损失多少、最大损失多少?

7.1 B0 与 B1 的定义

  • B0完整版 B:§三~§五的方案 B——与 A1b 做分界判定后选择(访存 Bound 且 \rho \ge \rho_{dv} 时应选 A1b 而非 B
  • B1无脑版 B:不做 A1b/B 分界,只要门限满足就直接按整数轮满核切(无论计算 Bound 还是访存 Bound

B1 门限推导B1 要成立需两个条件——

条件 1有尾轮$r = N_{blk} \bmod C > 0$(否则无尾轮可处理);

条件 2可行 + 值得)B0 恒优于 A0§4.2/§5.2 已证),故"可行即值得"。可行性只在访存 Bound 下受 dValue 约束($g \le g_{dv}$);计算 Bound 下无条件dValue 放宽为搬移掩盖,且 §4.3 已证 B 与 A1b 时延严格相等)。

合起来,B1 的判定门限(闭式,仅依赖 B/M/K/N/dtype


\boxed{\; r > 0 \;\land\; \Big( K \ge K^* \;\lor\; g \le \Big(\frac{sN \cdot dtype}{256\text{B}}\Big)^2 \Big) \;}

即:计算 Bound 时无条件直接整数轮满核切;访存 Bound 时满足 dValue 门限($g \le g_{dv}$)就直接整数轮满核切。门限的全部输入($r$、$g$、$K^*$、$sN$都来自首轮切分结果host 端一次计算即可。

7.2 B0/B1 相对 A1b 的损失分析

计算 Bound损失严格为 0。§4.3 已证 $T_{A1b} = T_B = T_{MMAD}(n_{wave}-1+\rho)$(总计算量/C 守恒——B1 在计算 Bound 下零损失。唯一差异是搬移量(周长和 A1b ≤ B但计算 Bound 下搬移被掩盖,不体现为时延。

访存 Bound损失集中在区间 I$\rho \ge \rho_{dv}$A1b 可凑满)。以下分三步:先给出损失公式的逐步推导,再审视其"方形近似"假设的可靠性,最后落到主流 prefill/decode 场景的实际损失。

7.2.1 损失公式的逐步推导

第一步:两策略的时延表达式(访存 Bound主导项为搬入 $T_{load}$)。

A1b主轮 n_{wave}-1 轮整块 + 尾轮 tile 缩 \sqrt{\rho} 凑满 C 核:


T_{A1b} = (n_{wave}-1) \cdot T_{load} + \sqrt{\rho} \cdot T_{load} = T_{load}\big(n_{wave} - 1 + \sqrt{\rho}\big)

BN_{blk}' = n_{wave} C 块均匀重切,单块搬移缩 \sqrt{g} 倍($g = n_{wave} C/N_{blk}$n_{wave} 轮满载:


T_B = n_{wave} \cdot \frac{T_{load}}{\sqrt{g}} = n_{wave} \cdot T_{load} \cdot \sqrt{\frac{N_{blk}}{n_{wave} C}} = T_{load}\sqrt{\frac{n_{wave} \cdot N_{blk}}{C}}

第二步:把 N_{blk}(n_{wave}, \rho) 表示。由 N_{blk} = (n_{wave}-1)C + r\rho = r/C 得 $N_{blk}/C = n_{wave} - 1 + \rho$,代入:


T_B = T_{load}\sqrt{n_{wave}\big(n_{wave} - 1 + \rho\big)}

第三步:作比得损失公式


\frac{T_B}{T_{A1b}} = \frac{\sqrt{n_{wave}\big(n_{wave}-1+\rho\big)}}{n_{wave} - 1 + \sqrt{\rho}} \triangleq f(\rho, n_{wave})

第四步:数值代入

  • 5.4% 的来源$sN=256$、BF16$\rho_{dv} = (256/(256 \times 2))^2 = 0.25$A1b 可行区间为 $\rho \ge 0.25$,损失在该区间单调减(见下),最大值在 $\rho = \rho_{dv} = 0.25$、n_{wave}=2 处(对应 $r=8$、$N_{blk}=40$、$g = 64/40 = 1.6$

T_B = 2 \cdot T_{load}/\sqrt{1.6} = 1.581 \cdot T_{load},\qquad T_{A1b} = (1 + \sqrt{0.25}) \cdot T_{load} = 1.5 \cdot T_{load}

f = \frac{1.581}{1.5} = 1.054 \;\Rightarrow\; \text{损失 } 5.4\%
  • 16.6% 的来源$sN=512$、BF16$\rho_{dv} = (256/1024)^2 = 0.0625$,最大值在 $\rho = 0.0625$$r=2$、$N_{blk}=34$、$g = 64/34 = 1.882$)、n_{wave}=2 处:

T_B = 2/\sqrt{1.882} \cdot T_{load} = 1.458 \cdot T_{load},\qquad T_{A1b} = (1 + 0.25) \cdot T_{load} = 1.25 \cdot T_{load}

f = \frac{1.458}{1.25} = 1.166 \;\Rightarrow\; \text{损失 } 16.6\%
  • 30% 极端例的来源M=16896、N=512、K=140、BF16理论公式给 16.6%,但实际时延计算为 $T_{A1b} = 6.7\mu s$、$T_B = 8.7\mu s$,比值 $1.30$。超出的 14 个百分点来自理论公式的两个近似同时失效(详见 7.2.2①B 重切后周长缩小使 L1 约束放松,k_{L1} 从 128 升至 140K 全载),每块搬移量反升 140/128 = 1.094 倍;②N_{blk}' = 64 只能分解为 $(64,1)$sN'=512 被 dValue 卡死不能更小tile 长宽比 264:512 非方形,周长和比方形基准高 5.6%。两因子相乘 $1.094 \times 1.056 = 1.155$1.166 \times 1.155 \approx 1.35 与实际 1.30 基本吻合(残差来自 A1b 尾轮 k_{L1} 同升等细节)。

单调性证明(最大损失位置):对 f 关于 \rho 求导(固定 $n_{wave}$\partial f/\partial\rho 的符号由 \big(n_{wave}-1+\sqrt{\rho}\big) - \big(2\sqrt{\rho}(n_{wave}-1+\rho)/\sqrt{n_{wave}(n_{wave}-1+\rho)}\big) 决定;数值上 f\rho 单调减、随 n_{wave} 单调减(\rho \to 1n_{wave} \to \infty 时 $f \to 1$)。故最大损失在 $\rho = \rho_{min} = \max(1/C, \rho_{dv})$、n_{wave} = 2取得。

7.2.2 方形近似假设的可靠性审视

损失公式建立在三个近似上①tile 近方形(缩放按 $\sqrt{g}$);②k_{L1} 不随重切变化③16 对齐取整忽略。逐一审视:

1非方形分解——惩罚因子与方向性

$N_{blk}'$(或尾轮区域)的整数分解 (mCnt', nCnt') 偏离方形比例 M/N 时,周长和高于方形基准。设分解长宽比 $\lambda = sM'/sN'$,非方形惩罚因子:


p(\lambda) = \frac{\sqrt{\lambda} + 1/\sqrt{\lambda}}{2} \ge 1 \quad \text{(等号当 } \lambda = 1 \text{ 方形)}

非方形修正后的损失公式:


\frac{T_B}{T_{A1b}} = \frac{n_{wave} \cdot p_B / \sqrt{g}}{(n_{wave}-1) + \sqrt{\rho} \cdot p_t}

其中 $p_B$、p_t 分别为 B 全局分解与 A1b 尾轮分解的惩罚因子A1b 主轮 tile 不变,惩罚为 1回答"非方形是否一定有损失"

  • 损失的存在性($T_B > T_{A1b}$)由均值不等式的结构决定——只要两策略在同一基准 tile 形状上做相同的面积→周长缩放,不等式方向不变;非方形惩罚同时抬高两者,方向保持
  • 但损失的大小不再等于理论值:p_B > p_t 时损失放大(极端例的 264:512 分解 $p_B = 1.057$p_B < p_t 时损失缩小甚至 B 反超。无绝对保证,需按实际分解计算;
  • 主流 caseM, N 大且分解自由度高)近方形可达$p \approx 1$),理论值可靠;极端长宽比 case(如 N=512 使 $nCnt'=1$B 只能 (64,1) 分解)p_B 显著 $> 1$,实际损失可超理论值近一倍。

2k_{L1} 跨 regime

缩放律假设 k_{L1} 不变,但 $k_{L1} = \min(K, L1/(2(sM{+}sN) \cdot dtype))$:重切使周长缩小 → L1 约束放松 → k_{L1} 可能从 L1 主导值升至 $K$K 全载)。k_{L1} 上升使每块搬移量反升(搬更多 K 数据B 的总搬移量进一步增大——k_{L1} 跨 regime 只会放大 B 的损失,不会缩小A1b 主轮 tile 不变、不受影响)。极端例中该效应贡献 140/128 = 1.094 倍放大。

316 对齐取整tile 小时($s' \sim 128$)取整误差占比 $\sim 16/128 = 12.5%$,大 tile 时 <6%。对两策略平等,不改变方向。

小结:方形近似在主流 case分解自由度高、k_{L1} 同 regime下可靠误差 <5%);在"分解被 dValue/整数约束强制非方形 + k_{L1} 跨 regime"的边角 case 下会低估损失(极端例 16.6% → 30%)。损失公式的理论值应视为损失下界

7.2.3 主流 prefill/decode 场景的实际损失

大模型推理的 BMM 典型形状用户给定prefill 场景 B∈[2,128]、M∈[4k,128k]、K/N∈[128,10240]decode 场景 B∈[2,128]、M∈[1,128]、K/N∈[128,10240]。按本文模型逐一分析:

1M, N 为 256 倍数的主流形状:r=0 恒成立,损失为 0

sM=sN=256 时 $mCnt = M/256$、$nCnt = N/256$。主流模型维度4096/8192/14336 等)均为 256 倍数,mCnt, nCnt 含充足 2 的幂因子,N_{blk} = B \cdot mCnt \cdot nCntB \ge 2 几乎恒为 32 倍数 → 无尾轮。prefill 的典型 case如 B=8、M=16384、N=8192$N_{blk}=4096$、$w=128$、$r=0$)根本不触发尾轮问题。

2M 奇异的 prefill caseseq_len 非 256 倍数):w 大 → 损失 <2%

B=2、M=4224$mCnt=17$、N=4224$nCnt=17$$N_{blk}=578$、$w=19$、$r=2$。即使访存 BoundK 小),f = \sqrt{19 \times 18.06}/(18+0.25) = 1.015损失 1.5%w 越大损失越小——prefill 的 M 大使 N_{blk} 大、w 大,损失天然被稀释。

3decode case$M \le 128$):访存 Bound 居多,但 w 大或落入打平区

  • M 小 → T_{MMAD} 小 → 访存 Bound 为主(如 M=64、sN=256、K=512$T_{MMAD}=1.1\mu s < T_{load}=5.1\mu s$
  • $N_{blk} = B \cdot nCnt$$mCnt=1$$B$、nCnt 的乘积通常使 w 大(如 B=100、N=3328$N_{blk}=1300$、$w=41$、$r=20$、$\rho=0.625 \ge \rho_{dv}=0.25$f = \sqrt{41 \times 40.625}/(40+0.79) \approx 1.0005损失 ≈0.05%
  • w 小的 decode caseB 小且 N 小):如 B=17、M=64、N=2048$nCnt=8$$N_{blk}=136$、$w=5$、$r=8$、$\rho=0.25=\rho_{dv}$f = 1.024损失 2.4%
  • N_{blk} < C 的 case如 B=8、N=512$N_{blk}=16$$w=1$——核都填不满,属降核/并行度不足场景(不是尾轮问题),不在本文策略讨论范围。

4主流场景损失结论

场景 尾轮出现条件 典型 w B0/B1 相对 A1b 损失
prefillM,N 为 256 倍数 不出现($r=0$ 0
prefillM 奇异 B \cdot mCnt \cdot nCnt \bmod 32 \neq 0 \ge 9 <2%
decodew \ge 5 同上 \ge 5 ≤2.4%
decodeN_{blk} < C w=1 降核场景,不在本文范围
极端边角(M/N \gg 1 极端长宽比 + K \in (k_{L1}, K^*) 窄区间 + w=2 + r 小) 罕见 2 理论 16.6%,实际可达 ~30%

主流 prefill/decode case 的 B1 损失 ≤ 2.5%——B1无脑整数轮满核切在主流场景下是安全的工程选择~30% 的损失只出现在主流形状不会命中的边角 case极端长宽比 + 窄 K 区间 + 两轮尾轮)。

7.3 损失汇总与工程建议

场景 B1 相对 A1b 的损失 说明
计算 Bound任意 $r$ 0 时延严格相等§4.3
访存 Bound$\rho \ge \rho_{dv}$sN=256 5.4% 典型场景(方形 tile、$w=2$、r=8 时取到)
访存 Bound$\rho \ge \rho_{dv}$sN=512 ≤ 16.6%(理论)/ ~30%(实际极端例) 需极端长宽比 + 窄 K 区间,罕见
访存 Bound\rho < \rho_{dv} ≤ 0B 反超或打平) §5.3(iii) 分界公式
sN \ge 608 0该场景恒计算 Bound K^* 区间为空
主流 prefill/decode case§7.2.3 ≤2.5% 256 倍数形状 $r=0$M 奇异/decode 场景 w 大稀释损失

工程建议

  1. 计算 Bound 为主力的场景直接上 B1——零损失、实现最简host 一次枚举、kernel 无尾轮分支);
  2. 访存 Bound 且 tile 小($sN \le 256$)时 B1 损失 ≤5.4%,实现简洁性通常值得这点损失;
  3. 访存 Bound 且 $sN \approx 512$、$n_{wave}=2$、r 小的 case 损失可达 ~30%(极端例),此类 case 应落到 A1b或按 §5.3 分界公式精确判定,即 B0
  4. B1 的判定门限全部闭式B/M/K/N/dtype → $r$、$g$、$K^*$、$sN$host 端零成本。

八、边界说明

  1. 带宽模型敏感性:访存 Bound 结论依赖"每核 MTE2 带宽上限 $BW_{pc}$"假设。若 HBM 为全局共享池(尾轮 r 核可吃满 $W_{GM}$A0 尾轮搬移时延已是 \rho \cdot T_{load} 接近理想A1b/B 的搬移增量无带宽补偿,结论反转。昇腾 950PR 的 MTE2 为每核独立 DMA 引擎、带宽按核数配平,采用固定份额结论;临界 case 建议实测复核。
  2. A1b 的工程代价:需两套 tile 参数(主轮大 tile + 尾轮小 tile与尾轮区域的边界处理r 个原块的并一般为 L 形按矩形分解重切host 端多一次枚举NPU 侧 kernel 需支持尾轮 tile 尺寸切换。这些复杂度不改变时延结论,但影响实现成本。
  3. 方形基准的适用条件§2.5 已述——K 全载时方形严格最优、L1 主导时仅面积相关。非方形分解的实际收益由同一目标函数评估,例 5 的非方形惩罚即此体现。
  4. B 不整除 $N_{blk}'$B \nmid N_{blk}' 时按 batch 分组切分、部分 batch 多一块,收益略降,判定不变。

参考文献

  1. 昇腾 950 NPU 架构白皮书华为技术有限公司2026
  2. cann-ops-nn 源码仓BatchMatMulV3kernel 侧尾轮线性映射见 op_kernel/arch35/batch_mat_mul_v3_asw_block_advanced.hUpdateBasicIndexindex = newBlockIdx + roundIdx × usedCoreNumif (index < totalCnt) 跳过空转核——即本文的策略 A0