MergeBatch vs IterBatch:BMM 小 tile 场景的两种实现路径分析

版本:v1.1(新增 §5.3 IterBatch L0C fallback 合理性公式推导)。目标芯片:昇腾 950PR(DAV_3510)。理论分析基于《BMM 算子优化分析 v0.98》,源码参考 [cann-ops-nn](https://gitcode.com/cann/ops-nn/tree/master/matmul/batch_mat_mul_v3) 实现。

摘要

BMM 中当 B ≥ C 时核间切 B 免费,核内有 IterBatch(逐 batch)和 MergeBatch(合并 $b_0$ 个 batch)两种路径。本文从 GM→L1 搬移粒度差异出发,证明 MergeBatch 的核心优势是合并搬移减少 DMA 命令数(每次搬移有固定开销 $T_{cmd}$),核心劣势是 **drain 暴露放大 $b_0$ 倍。分界条件:MergeBatch 仅在 K 截断**($k_{L1} = K$,小 M/N)且 $b_{core}$ 足够大时优于 IterBatch;**L1 绑定情形($k_{L1} < K$)MergeBatch 恒劣于 IterBatch**。与源码对比发现:源码 MergeBatch 进入条件偏宽(缺搬移量/tile/算存比约束),IterBatch 只覆盖 L1 形态 b(缺 a/c/d)。


一、问题定义

BMM(BatchMatMul)中,当 B ≥ C(batch 数 ≥ AIC 核数)时,核间切 B 是免费的——每核独立处理若干 batch,无共享无依赖。核内的两种实现路径:

核心问题:什么条件下 MergeBatch 优于 IterBatch?


二、IterBatch 理论最优实现

2.1 执行模型

核间切 B,每核 $b_{core} = B/C$ 个 batch。核内逐 batch 做标准 Matmul:

对每 batch i ∈ [0, b_core):
    GM→L1: 搬入 A_i[M, K] + B_i[K, N](或按 kL1 分块)
    L1→L0→Cube: 计算 C_i[M, N]
    Fixpipe: L0C→GM 写回 C_i

2.2 L1 驻留形态

IterBatch 按 L1 容量分四种形态(v0.96 §六):

形态条件特征
a) 单 batch 全驻留$b_{core}=1$,$(MK+KN) \cdot dtype \le L1$零重复读
b) 双 batch 乒乓$b_{core}>1$,$2(MK+KN) \cdot dtype \le L1$batch 间流水掩盖
c) 一侧驻留+对侧切 K$MK \cdot dtype \le L1/2$(驻留侧)驻留侧零重复读
d) 两侧都切 K兜底两侧都有 K 段级重复读

2.3 实现方案


三、MergeBatch 理论最优实现

3.1 执行模型

核间切 B(同 IterBatch),核内把 $b_0$ 个 batch 合并计算:

对每合并组 g ∈ [0, b_core/b0):
    GM→L1: 搬入 A'[b0·M, K] + B'[K, b0·N](单次 DMA,ndNum=1)
    L1→L0→Cube: 计算 C'[b0·M, b0·N]
    Fixpipe: L0C→GM 写回(只取对角块,丢弃交叉项)

交叉项 $(b_0^2 - b_0) \cdot MN$ 被算出但丢弃——浪费比例 $(b_0-1)/b_0$。进入该分支的 case 必然访存 Bound(条件 5 保证),浪费的算力被搬移时延掩盖。

3.2 进入条件

同时满足($b_0$ = 最小合并数,取 2):

  1. $BatchA = BatchB$ 且 $b_{core} = B/C \ge 2b_0$
  2. $2(b_0 M)(b_0 N) \cdot 4B \le L0C$(合并后输出放得下 L0C 双缓冲)
  3. $b_{core}(MK+KN) \cdot dtype \ge min\_DatamountPerCore$(单核搬移量足够)
  4. $\max(MK, KN) \cdot dtype \ge min\_TileSize$(搬移 tile 不太小)
  5. $2MN/(M+N) < R_{16}/b_0$(合并后仍访存 Bound)

3.3 实现方案

Step 1:合并数 b₀(L0C + 算存比双上限)

$$ b_0 = \min\Big(\sqrt{\frac{L0C}{2 \cdot MN \cdot 4B}},\; \frac{R_{16}(M+N)}{2MN},\; b_{core}\Big) $$

**Step 2:L0 级 K 粒度 $k_{L0}$(由 b₀ 和 L0A/L0B 决定)**

$$ k_{L0} = \min\Big(\frac{L0A}{2 b_0 M \cdot dtype},\; \frac{L0B}{2 b_0 N \cdot dtype}\Big) \text{ 向下 16 对齐} $$

**Step 3:L1 级 $k_{L1}$(先反推再截断)**

先令 $b_{L1}^* = b_0$,由 L1 容量反推 $k_{L1}^*$:

$$ k_{L1}^* = \frac{L1}{2 \cdot b_0 \cdot (M+N) \cdot dtype} $$

然后取 $k_{L1} = \min(k_{L1}^*,\; K,\; 512B/dtype)$——不超过 K(K 截断),也不超过 dValue 推荐值 512B。

**Step 4:$b_{L1}$ 最大化**

$$ b_{L1} = \min\Big(\frac{L1}{2 \cdot k_{L1} \cdot (M+N) \cdot dtype},\; b_{core}\Big) $$

若 $k_{L1}^* > 512B/dtype$,$k_{L1}$ 被 512B 截断,省出的 L1 空间可容纳更多 batch($b_{L1} > b_0$),提升 batch 间流水深度。


四、MergeBatch vs IterBatch 分界分析

4.1 执行模型与符号定义

两分支的核间切分相同——B 维切分到 C 核,每核 $b_{core} = B/C$ 个 batch;核内不切 M/N。差异在核内 K 维处理和 GM→L1 搬移粒度:

*符号表*:

符号含义表达式/取值
$C$AIC 核数32
$b_{core}$每核 batch 数$B/C$
$b_0$合并因子(单次合并的 batch 数)由 L0C + 算存比双上限确定
$k_{L1}$L1 级 K 分块粒度$\min(K,\; L1/(2(M{+}N) \cdot dtype))$
$n_K$K 分块数$\lceil K/k_{L1} \rceil$
$T_{load}$每 K 分块搬移时延$k_{L1}(M{+}N) \cdot dtype / BW_{pc}$
$T_{comp}$每 K 分块计算时延$2MN \cdot k_{L1} / Q_{16}$
$T_{write}$单 batch 输出写回时延$MN \cdot outB / W_{GM}$
$T_{cmd}$单次 GM→L1 DMA 搬移固定开销描述符配置 + 地址生成 + 突发启动
$BW_{pc}$单核 GM 带宽份额$W_{GM}/C = 50$ GB/s
$Q_{16}$单核 Cube BF16 峰值算力486/32 ≈ 15.2 TFLOPS
$W_{GM}$GM 带宽1.6 TB/s
$outB$输出元素字节数BF16 → 2B

4.2 端到端时延建模

L0C 双缓冲使 fixpipe 与 Cube 完全交叠,batch 边界无额外同步开销。两分支的差异仅来自 GM→L1 搬移次数和 drain 暴露。

IterBatch 时延(每核处理 $b_{core}$ 个 batch,每 batch $n_K$ 个 K 分块):

$$ T_{iter} = \underbrace{b_{core} \cdot n_K \cdot (T_{load} + T_{cmd})}_{\text{搬移(逐 batch,每次含 } T_{cmd}\text{)}} + \underbrace{T_{comp} + T_{write}}_{\text{末 batch drain}} $$

MergeBatch 时延(每核处理 $b_{core}/b_0$ 个合并 batch,每合并 batch $n_K^m$ 个 K 分块):

$$ T_{mb} = \underbrace{\frac{b_{core}}{b_0} \cdot n_K^m \cdot (T_{load}^m + T_{cmd})}_{\text{搬移(合并,每次含 } T_{cmd}\text{)}} + \underbrace{b_0 \cdot (T_{comp} + T_{write})}_{\text{末合并 batch drain}} $$

其中 $T_{load}^m$ 是合并后每 K 分块的搬移时延,$n_K^m$ 是合并后的 K 分块数——两者的取值取决于 $k_{L1}$ 是否截断(见下)。

合并后 K 分块粒度:MergeBatch 合并 $b_0$ 个 batch 后,L1 中的 tile 变为 $[b_0 M, k_{L1}^m] + [k_{L1}^m, b_0 N]$。L1 容量约束:

$$ 2 \cdot k_{L1}^m \cdot b_0 \cdot (M+N) \cdot dtype \le L1 \Rightarrow k_{L1}^m = \min\Big(K,\; \frac{k_{L1}}{b_0}\Big) $$

4.3 K 截断情形($k_{L1} = K$)

$n_K = n_K^m = 1$,$T_{load}^m = b_0 \cdot T_{load}$(合并后单次搬移数据量 $b_0$ 倍):

$$ T_{iter} = b_{core}(T_{load} + T_{cmd}) + T_{comp} + T_{write} $$
$$ T_{mb} = \frac{b_{core}}{b_0}(b_0 \cdot T_{load} + T_{cmd}) + b_0(T_{comp} + T_{write}) $$

两式相减:

$$ \Delta_K = T_{mb} - T_{iter} = \underbrace{(b_0-1)(T_{comp} + T_{write})}_{\text{drain 惩罚}} - \underbrace{b_{core}\Big(1-\frac{1}{b_0}\Big) T_{cmd}}_{\text{搬移命令节省}} $$

分界条件(MergeBatch 优于 IterBatch 当且仅当 $\Delta_K < 0$):

$$ b_{core} > \frac{b_0 \cdot (T_{comp} + T_{write})}{T_{cmd}} $$

小 MN → $T_{comp}$ 小 → 惩罚小 → MergeBatch 更容易赢;大 B → $b_{core}$ 大 → 搬移节省多 → MergeBatch 更容易赢。

4.4 L1 绑定情形($k_{L1} < K$)

$k_{L1}^m = k_{L1}/b_0$,$n_K^m = b_0 \cdot n_K$,$T_{load}^m = T_{load}$(每块搬移量相同,因为 $k_{L1}^m \cdot b_0 = k_{L1}$):

$$ T_{iter} = b_{core} \cdot n_K \cdot (T_{load} + T_{cmd}) + T_{comp} + T_{write} $$
$$ T_{mb} = \frac{b_{core}}{b_0} \cdot b_0 n_K \cdot (T_{load} + T_{cmd}) + b_0(T_{comp} + T_{write}) = b_{core} \cdot n_K \cdot (T_{load} + T_{cmd}) + b_0(T_{comp} + T_{write}) $$

搬移次数相同($b_{core} \cdot n_K$),每次搬移量相同——MergeBatch 无搬移优势

$$ \Delta_L = T_{mb} - T_{iter} = (b_0-1)(T_{comp} + T_{write}) > 0 $$

L1 绑定情形 MergeBatch 恒劣于 IterBatch——合并只放大了 drain 暴露,没有换来搬移命令节省。

4.5 统一分界条件

$$ \text{MergeBatch 最优} \iff \underbrace{k_{L1} = K}_{\text{K 截断}} \;\land\; \underbrace{b_{core} > \frac{b_0 \cdot (T_{comp} + T_{write})}{T_{cmd}}}_{\text{搬移节省} > \text{drain 惩罚}} $$

K 截断的物理含义:$L1/(2(M{+}N) \cdot dtype) \ge K$——小 M/N 使 L1 一次装下整个 K 维。

维度IterBatchMergeBatch差异
稳态搬移吞吐相同相同总搬移量相同
GM→L1 搬移命令数$b_{core}$ 次$b_{core}/b_0$ 次(K 截断时)**MergeBatch 少 $b_0$ 倍** ← 核心优势
drain 暴露$T_{comp} + T_{write}$$b_0(T_{comp} + T_{write})$IterBatch 少 $b_0$ 倍 ← 核心劣势
L0C 利用率访存 Bound 下不影响时延

4.6 $T_{cmd}$ 的物理成因

每次 GM→L1 DMA 搬移的固定开销包括:Nd2Nz 描述符配置(ndNum/nValue/dValue/srcStride 等 7 个字段写入 DMA 寄存器)、地址生成与突发启动、与 L1 buffer 的同步握手(SetFlag<MTE2_MTE1> / WaitFlag<MTE1_MTE2>)。量级估计为数十 ns,精确值需实测标定。


五、与源码实现的对比

5.1 MergeBatch 源码(batch_matmul_v3_mergebatch_basicapi_tiling.cpp

IsCapable 条件

// 源码:https://gitcode.com/cann/ops-nn/tree/master/matmul/batch_mat_mul_v3
if (batchC < MIN_BATCH_L0 * aicNum) return false;  // B < 128,即 b_core < 4
if (alignK < 64 || M > N) return false;             // K ≥ 64,M ≤ N
if (L0 buffer check with b0=4) return false;        // L0C 容量

对比分析

维度理论源码差异
合并数 b₀min(L0C 上限, 算存比上限, b_core)MIN_BATCH_L0 = 4(固定),实际由多项式动态求解源码 b₀=4 是硬编码下限,实际值由 L0C 容量动态计算
b_core 下限≥ 2b₀ = 4≥ 4(batchC ≥ 128)一致
L0C 约束2(b₀M)(b₀N)·4B ≤ L0CtempAlignM × tempAlignN × 4B × 2 ≤ L0C一致
搬移量约束b_core(MK+KN)·dtype ≥ 480KB源码缺少
tile 大小约束max(MK,KN)·dtype ≥ 16KB源码缺少
算存比约束2MN/(M+N) < R₁₆/b₀源码缺少
额外约束K ≥ 64,M ≤ N源码独有

结论:源码的 MergeBatch 进入条件偏宽——缺少搬移量、tile 大小、算存比三条约束。可能导致不应进入的 case 被误捕获(如计算 Bound 的 case 进入 MergeBatch 后冗余计算成为瓶颈)。

5.2 IterBatch 源码(batch_matmul_v3_iterbatch_tiling.cpp

IsCapable 条件

// 源码:https://gitcode.com/cann/ops-nn/tree/master/matmul/batch_mat_mul_v3
if (batchC <= aicNum) return false;  // B ≤ 32 不进 IterBatch
if ((alignM*alignK + alignK*alignN)*dtype*2 > L1) return false;  // L1 双缓冲
if (!l0CanLoadBatch_ && balanceRate < 0.8) return false;  // 负载均衡

对比分析

维度理论源码差异
L1 形态四种(a/b/c/d)只覆盖形态 b源码缺 a/c/d
B 下限b_core ≥ 1B > 32(b_core ≥ 2)源码排除了 b_core=1 的形态 a
负载均衡B mod C ≥ minCoreNumbalanceRate ≥ 0.8实现方式不同,目的一致

结论:源码的 IterBatch 只覆盖形态 b(双 batch 乒乓),且要求 B > C。形态 a/c/d 的 case 由 ASW/AL1/BL1 分支承接。

5.3 L0C 放不下 2 Batch 时的 fallback 合理性分析

源码在 L0C 双缓冲放不下 1 个 batch 时($L0C < 2 \cdot MN \cdot 4\text{B}$)不直接拒绝,而是检查 formulate balance rate ≥ 0.8。这个软门限的合理性需要回答两个问题:①此时 fixpipe 与 Cube 还能流水掩盖吗?②为什么主要风险是负载不均而非流水串行?

5.3.1 fallback case 的特征推导

fallback 进入的条件是 $L0C < 2 \cdot MN \cdot 4\text{B}$,同时 IterBatch 的 L1 硬约束是 $2(MK{+}KN) \cdot dtype \le L1$。两者联合限定了 K 的范围:

$$ MN > \frac{L0C}{8\text{B}} = 32768,\qquad K(M+N) \le \frac{L1}{2 \cdot dtype} = 131072 $$

由均值不等式 $M + N \ge 2\sqrt{MN} > 2\sqrt{32768} \approx 362$,代入 L1 约束:

$$ K < \frac{131072}{M+N} \le \frac{131072}{362} \approx 362 $$

结论:所有 fallback case 的 K < 362。 例:M=N=256(MN=65536>32768)时 K ≤ 131072/512 = 256。

5.3.2 写出 Bound 判定与串行损失量化

单 batch 计算时延 $T_{comp} = 2MNK/Q_{16}$,写出时延 $T_{write} = MN \cdot outB/W_{pc}$($W_{pc}$ 为单核 GM 写带宽份额 50GB/s,$Q_{16}$ 为单核算力 15.19 TFLOPS):

$$ \frac{T_{comp}}{T_{write}} = \frac{2K \cdot W_{pc}}{Q_{16} \cdot outB} = \frac{K}{304} $$

K < 304 时写出 Bound(Fixpipe 是瓶颈);K ∈ [304, 362] 时接近均衡。

L0C 双缓冲时每 batch 时延 = $\max(T_{comp}, T_{write})$;dbL0c 退化为 1 后串行 = $T_{comp} + T_{write}$,额外损失:

$$ \Delta = \min(T_{comp}, T_{write}) $$

写出 Bound(K < 304)时 $\Delta = T_{comp}$,占串行总时延的比例:

$$ \frac{\Delta}{T_{comp} + T_{write}} = \frac{T_{comp}}{T_{comp} + T_{write}} = \frac{K}{K + 304} $$

K=64 → 17%;K=128 → 30%;K=256 → 46%;K=362(均衡区上界)→ 约 54%。K 越小(写出 Bound 越严重),串行损失占比越小——直观理解:写出 Bound 时瓶颈本来就在 Fixpipe,Cube 交叠与否不影响 Fixpipe 持续工作,L0C 单缓冲只是让 Cube 在 Fixpipe 工作时空闲,损失的是本就不是瓶颈的计算时间。

5.3.3 为什么软门限防的是负载不均而非串行

串行损失 $\Delta$ 是确定性的、有界的(≤ 54%,且随 K 减小单调下降)。更重要的是,fallback case 里 L0C 双缓冲方案根本不存在($MN > 32768 \Rightarrow 2 \cdot MN \cdot 4\text{B} > L0C$)——真正的替代方案是 ASW 切 M/N(L0C 正常流水但引入核间重复读)。两者权衡下,L0C 串行的确定性损失通常优于切 M/N 的重复读,所以源码不因串行而拒绝。

真正不可控的风险在 batch 负载不均:batch 按 iterBatchL1 粒度分组后轮转分配到 C 核,尾轮不满时部分核提前空闲。balance rate 量化这一浪费:

$$ rate = \frac{avgIterBatch}{actualMaxIterBatch} = \frac{B/C}{\lceil \lceil B/iterBatchL1 \rceil / C \rceil \times iterBatchL1} $$

rate < 0.8 ⟺ 尾轮 padding 浪费超过 25% 核时。例:B=64、C=32、iterBatchL1=8 → 总组数 8、仅 8 核有活、rate = (64/32)/(1×8) = 0.25 < 0.8 → 拒绝,合理(32 核中 24 核闲置)。例:B=256、C=32、iterBatchL1=8 → 32 组满核、rate = 1 ✓。

结论链:L0C 放不下 ⟹ K < 362 ⟹ 写出 Bound 或接近均衡 ⟹ L0C 串行的确定性损失有界且随 K 减小而减小 ⟹ 主要风险在 batch 负载不均(尾轮拖长)⟹ 源码用 rate ≥ 0.8 软门限控制——这一设计与其 case 特征自洽。

5.4 综合评价

理论分析的价值在于:

  1. 分界条件——源码无显式 MergeBatch vs IterBatch 分界,理论给出了 $k_{L1} = K \land b_{core} > b_0(T_{comp}+T_{write})/T_{cmd}$ 的闭式条件
  2. L1 绑定情形——理论证明 MergeBatch 在 L1 绑定时恒劣于 IterBatch,源码未区分
  3. T_cmd 的物理意义——GM→L1 DMA 命令固定开销,是 MergeBatch 优势的本质来源

源码实现的价值在于:

  1. 工程完备性——负载均衡、格式检查、bias 处理等
  2. cubeBound 模型——经实测调优的 baseM/baseN 寻优
  3. L2 Cache 管理——100MB 阈值 + transConflict 控制

参考文献

  1. [昇腾 950 NPU 架构白皮书](https://public-download.obs.cn-east-2.myhuaweicloud.com/ascend/%E6%98%87%E8%85%BE950%20NPU%E6%9E%B6%E6%9E%84%E7%99%BD%E7%9A%AE%E4%B9%A6.pdf),华为技术有限公司,2026
  2. [CANN 9.2.0-beta.1 Ascend C API 参考](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/920beta1/API/ascendcopapi/atlasascendc_api_07_0003.html)——UnitFlag 机制详见"矩阵计算(ISASI)/Mmad 计算关键特性说明/UnitFlag"章节
  3. [cann-ops-nn 源码仓](https://gitcode.com/cann/ops-nn/tree/master/matmul/batch_mat_mul_v3)