BMM 算子最优实现分析

1. 什么是 BMM 算子？

Batch MatMul：对 batch 维的每个索引独立做一次矩阵乘

输入：左矩阵 [BatchA, M, K]、右矩阵 [BatchB, K, N]

输出：[BatchC, M, N]，batch 维支持广播

核心挑战：B、M、N、K 四个维度都可能很大或很小，必须在32 个 AIC 核上高效并行完成计算

一个 case 由 (B, M, N, K, dtype, 转置, 广播形态) 完全决定

2. 硬件基础：昇腾 950PR

关键关系：L2（5.2TB/s）→ L1（512KB）→ L0A/B（64KB）→ Cube（16×16×16 一拍）→ L0C（256KB）→ Fixpipe → GM

GM 带宽 1.6TB/s 相对 L2 的 5.2TB/s 是瓶颈——数据复用就是一切

3. 什么叫"最优"？

总时延 = 流水线最慢的一级

核内多级流水并行：Cube 计算（MMAD）、GM→L1 搬移（MTE2）、L1→L0 搬移（MTE1）、L0C 写出（Fixpipe）

瓶颈交换：搬移是瓶颈时，可牺牲算力换搬移效率；计算是瓶颈时，可牺牲搬移换计算效率

算存比：AI = 2MN/(M+N) FLOP/元素

16bit 平衡点：R₁₆ = 486 TFLOPS / (1.6TB/s ÷ 2B) ≈ 607.5 FLOP/元素

AI < R₁₆ → 访存 Bound（瓶颈在搬移）｜AI > R₁₆ → 计算 Bound（瓶颈在 Cube）

4. 分块计算的本质：4 维可切

BMM 的实现 = 把数据按 B、M、N、K 切块，由 32 个 AIC 核并行 + 串行完成

核间怎么分这 4 个维度，就是分支划分的第一性问题

4 维任意非空子集共 15 种切分组合 → 完备枚举

为什么切 K 最贵？L0C 累加机制：核内切 K 时多轮 mmad 在 L0C 原地累加不出核；切到核间，部分和必须写出 workspace 再归约

5. 切分维度的"价格表"

切 B 免费：零重复读、零依赖、零中间写出——BMM 语义就是逐 batch 独立

切 M/N 廉价：共享矩阵被重复读，但 128MB L2（5.2TB/s）吸收大部分代价

切 K 昂贵：归约流量 ∝ 切 K 份数 × 输出量，且引入核间同步——结构性代价

整条分支决策树就是一句话：

按价格从低到高购买并行度，买不够才加价

6. 从价格表到 7 大分支

7 大路径：特殊(AIV)转MatmulIterBatchMergeBatchASW_BasicStreamK降核ASW

7. MergeBatch — 多 batch 合并计算

思路：单 batch M×N 太小时，把 b 个 batch 拼成 [bM,K]@[K,bN] 大矩阵，

算完取块对角线得各 batch 结果。交叉项被丢弃，浪费比例 (b−1)/b

为什么允许浪费？ 进这个分支的 case 必然是访存 Bound，瓶颈在搬移

不是在计算——浪费的算力被搬移时延掩盖，用闲置算力换搬移效率

进入条件精要：

• 每核 batch 数 ≥ 4（2 组 × 合并下限 b₀=2，组间乒乓）
• 合并块输出 ≤ L0C（b₀²·MN·4B·2 ≤ 256KB）
• 单核搬移总量 ≥ 480KB（保证 GM 带宽利用率）
• 搬移 tile ≥ 16KB（单次搬移效率下限）
• 2MN/(M+N) < R₁₆/b₀（访存 Bound，保证浪费被掩盖）

典型：B=128, M=N=64, K=512 → MergeBatch；K=256 → 单核搬移不够 480KB → IterBatch

8. IterBatch — 逐 batch 计算

思路：核间切 B，核内逐个 batch 做标准 Matmul。无浪费、无跨 batch 依赖

进入条件精要：B ≥ 32；负载均衡；L1 四形态之一满足

关键：能否进 IterBatch 不由算存比判定

即使计算 Bound，L1 放不下完整 M、N 维输入 → 单 batch 内部重复读 → 额外搬移可能把算子重新拖回访存 Bound

L1 四形态：

• a) 单 batch 全驻留（b_core=1 且 MK+KN ≤ L1）
• b) 双 batch 乒乓（b_core>1 且 2(MK+KN) ≤ L1）
• c) 一侧驻留+对侧切 K（预算 L1 或 L1/2，b_core≥2 时半区预取下一 batch 驻留侧）
• d) 两侧切 K（无缝兜底，驻留侧超 L1/2 时接管）

9. StreamK — K 维核间切分

思路：B/M/N 都填不满核时，把 K 切给多个核各算一段，再归约

进入条件：P = B·MN·4B/L0C < 16（=C/2）；K/grid_K ≥ 128 元素（dValue 256B）

阈值取 C/2：grid_K≥2 时每块需 2 核，P×2 ≤ C 才放得下

归约代价（按实现流程推导）：

归约时延 = 写部分和 + AIV 读回 + AIV 求和 + 写回，四段之和

要求计算时延 ≥ 归约时延 × 10 → K ≥ grid_K² × 109

grid_K=2→K≥0.5K；4→1.8K；8→7.0K；32→112K

源码固定门槛 8192 ≈ 本模型 grid_K≤8 的要求，互为印证

10. ASW_Basic — 通用框架（swizzle）

思路：不切 K，切 B/M/N 任意组合。最常命中的分支（35.2%）

swizzle 滑窗蛇形：编排输出块的执行顺序，压缩每一波核的活跃工作集

窗口怎么取：M 向每 W 块划窗，W = C 的 ≤ √C 的最大因子（32 核 → W=4）

一波核的 L2 足迹 ≈ (W·M^t·K + C/W·K·N^t)·dtype，W=√C 时最小

蛇形只在窗口行边界：奇数窗口行 N 方向反向，使上一窗口末尾的 B 列带延续到下一窗口（LRU 热线），跨窗切换几乎零增量。窗内不蛇形——窗内 W 个 A 行块全程驻留 L2，扫序无关。

对比：朴素行优先（Ñ=16 时）一波足迹 18 单位；滑窗 12 单位 → 缩小 1/3

11. ASW_Basic — L2 切分 & 写出策略

切的是输出平面：把 M×N 切成大矩形块，每块的输入工作集 ≤ L2 可用空间

L2 是读写共用的：输出驻留 L2 会压缩读入空间 → 写与读要联合决策

写出 Bound 先判定：BW_out = C·Q₁₆·outB/(2K)，K=256 时已超 GM 总线

三场景：

• A：S_in+S_out ≤ L2 → 全驻留，输出走 L2 异步回写（削峰）
• B：S_in ≤ L2 < S_in+S_out → 输入驻留保住 r_in=1，输出直写 GM（不占 L2）
• C：S_in > L2 → 必须切分，输出直写 GM，块内错位分核（对角线，同地址冲突 4→2）

r_in = 1 含义：每个输入从 GM 只读一遍，后续复用全在 L2 命中

12. case 遍历验证：20736 个 case 全覆盖

B∈[1,2048]、M,N,K∈[1,10240] 对数采样，按进入条件严格分类

7 个分支全部有真实 case 命中，无覆盖空洞

降核 ASW 是 P<C 且 K 小区域的理性归宿：并行度凑不满，碎切反而不如少用核

这份分析靠谱吗？

✅ 根基扎实

• 源码级验证：分支条件与 CANN ops-nn batch_mat_mul_v3 源码逐条对照（arch35 路径），MergeBatch/IterBatch/StreamK/ASW 的 IsCapable 条件均能在源码中找到对应，差异处已标注并给出理由
• 硬件规格准确：950PR 架构白皮书（表3-1/表4-2）+ CANN 9.0 AscendC 编程文档（054/056/192篇等）为所有参数（L1 512KB、L0C 256KB、dValue 256B、GM 1.6TB/s）提供依据
• 数学推导闭环：15 种切分组合完备枚举 → 价格表（免费/廉价/昂贵）→ 7 分支体系，极小性有反例表格验证

⚠️ 需要注意

• 经验常数（min_TileSize=16KB、480KB、均衡率 0.8、CUBE_BOUND_RATIO 0.85 等）为该芯片档实测值，换芯片需重新标定
• 重叠区（MergeBatch vs IterBatch 在 M×N 中段）的归属由端到端时延模型 $T_{total}$ 精确仲裁，文档条件给出的是"主场"
• StreamK θ=109 基于 L2 5.2TB/s 推导，若 workspace 落 GM 则退回 θ≈1700

整体评价：逻辑体系完整，推导自洽，可指导工程实现