diff --git a/BMM/BatchMatmul算子特性分析.html b/BMM/BatchMatmul算子特性分析.html new file mode 100644 index 0000000..88fe6d2 --- /dev/null +++ b/BMM/BatchMatmul算子特性分析.html @@ -0,0 +1,754 @@ + + +
+ + +++版本:v1.0 | 目标芯片:DV100(与昇腾 950PR / DAV_3510 同档:32 AIC / 64 AIV,GM 1.6TB/s,L2 128MB / 5.2TB/s) +本文目标:任意 shape/dtype/layout 的 BMM case,都能系统地推导出其在 NPU 上的最优软件实现方案(端到端总时延最短),并论证分支体系的完备性与极小性。 +阅读主线:第 1 章定义问题 → 第 2 章定义"最优" → 第 3 章推导分支划分 → 第 4 章逐分支给出进入条件与实现方案 → 第 5 章决策流程 → 第 6 章完备性审视(漏洞检查)。
+
Batch Matmul(BMM)完成带 batch 维的矩阵乘计算:对 batch 维的每个索引独立执行一次矩阵乘,再加可选偏置:
+$$ +C[b, m, n] = \sum_{k=0}^{K-1} A[b, m, k] \cdot B[b, k, n] + bias[b, 1, n] +$$
+即 C = A @ B + bias。A、B 输入维度典型为 3 维(支持最多 4 级 batch 维 $b = b_0 b_1 b_2 b_3$ 展平),最后两维做矩阵乘。
| 参数 | +形状 | +数据类型 | +layout | +说明 | +
|---|---|---|---|---|
| 左矩阵 A | +[BatchA, M, K] |
+dtype(FP16/BF16/FP8/…) | +典型 ND | +可带转置标记(isATrans) | +
| 右矩阵 B | +[BatchB, K, N] |
+dtype | +典型 ND | +可带转置标记(isBTrans),推理场景可为 weightNz | +
| 偏置 bias | +[B, 1, N] |
+dtype | +固定 ND | +可为空 | +
| 输出 C | +[BatchC, M, N] |
+dtype(可随路量化) | +典型 ND | +BatchC = broadcast(BatchA, BatchB) | +
广播语义:batch 维兼容广播——每一级 batch 维上,两侧取值要么相等、要么为 1;输出该维取两者的最大值。注意两类广播的形态差异:
+BatchA = 1 或 BatchB = 1(一侧整体只有 1 个 batch);BatchA = (6,1)、BatchB = (1,6),则 BatchC = (6,6) 共 36 个输出 batch。这类 case 不属于单边广播,后面第 3、6 章会看到它只能落入通用分支。一个 BMM case 的实现方式完全由以下特征量决定:
+$$ +(\; B,\; M,\; N,\; K,\; \text{dtype},\; \text{layout/转置},\; \text{广播形态} \;) +$$
+其中 $B$ = BatchC(展平后的输出 batch 数)。后续全部分析就是回答:给定这组特征量,最优实现是什么。
+NPU 上 BMM 的执行是核内多级硬件流水的并行——Cube 计算(MMAD)、GM/L2→L1 搬移(MTE2)、L1→L0 搬移(MTE1)、L0C 写出(Fixpipe),多核并行时各流水级时延可被双缓冲(double buffer)等机制相互掩盖,最终:
+$$ +\boxed{\;T_{total} = \max\big(T_{MMAD},\; T_{MTE2},\; T_{MTE1},\; T_{Fixpipe}\;[,\;T_{Reduce}]\big)\;} +$$
+算子优化的关键就是对瓶颈流水级的优化。 由此直接得出一个重要的设计自由度——瓶颈交换:
+++当 MTE2(搬移)是瓶颈、MMAD(计算)不是瓶颈时,可以牺牲一定 MMAD 时延(例如冗余计算)换取 MTE2 性能提升;反之,当 MMAD 是瓶颈时,可以牺牲一定 MTE2 时延(例如重复搬移)换取计算效率提升。只要瓶颈级时延下降,总时延就下降。
+
后文会看到:MergeBatch 就是"牺牲算力换搬移效率"的典型,ASW_Basic 切 M/N 就是"牺牲搬移(重复读)换并行度"的典型——它们的存在正当性都来自这个 max 模型。
+| 规格项 | +数值 | +对设计的意义 | +
|---|---|---|
| AIC / AIV 核数 | +32 / 64(1:2) | +核间并行度上限 $C = 32$;StreamK 归约可用 AIV | +
| Cube 算力 FP16/BF16 | +≈486 TFLOPS(单核 ≈ 16×16×16 MAC/拍) | +算存比分子 | +
| GM 带宽 | +1.6 TB/s(读写共享总线) | +访存 Bound 的分母;读+写互相挤占 | +
| L2 Cache | +128MB,5.2 TB/s(读写各独享) | +重复读取的吸收层;Cube 输出写 L2 即算完成 | +
| L1 / L0A / L0B / L0C | +512KB / 64KB / 64KB / 256KB(每 AIC) | +核内分块的容量约束;L0C 按 FP32 累加计 | +
| UB | +512KB(每 AIV) | +向量通路/归约 | +
GM→L1 搬移效率(ND2NZ)达到 90%+ 带宽利用率的四条经验约束(DV100 实测,重要性递减):
+min_DatamountPerCore);min_TileSize);++这些常数(480KB、16KB、256B/512B,以及后文的 $b_{thr}=4$、$k_{thr}$=32B/dtype、$minCoreNum \approx 0.8 \times 32$)都是该档芯片的实测经验值;换芯片时本文的所有逻辑结构不变,只需替换常数表。
+
约定:以"元素"为访存单位(1 元素 = dtype 字节),算存比单位为 FLOP/元素。
+$$ +AI = \frac{2MNK}{(MK + KN)} = \frac{2MN}{M + N} \quad [\text{FLOP/元素}] +$$
+$$ +AI_{full} = \frac{2MNK}{MK + KN + MN} \quad [\text{FLOP/元素}] +$$
+$$ +R = \frac{\text{Cube 峰值算力}}{\text{GM 带宽} / \text{dtype 字节}} = \frac{486 \times 2}{1.6} \approx 607.5 \quad [\text{FLOP/元素,BF16/FP16}] +$$
+判定规则:$AI < R$ → 访存 Bound(瓶颈在 MTE2,优化重心是搬移效率与复用);$AI > R$ → 计算 Bound(瓶颈在 MMAD,优化重心是 Cube 利用率与流水掩盖)。注意一个有用的事实:低位宽 dtype(FP8、MXFP4)算力同倍数翻倍而元素字节同倍数减半,$R$ 近似不变,因此该平衡点对该芯片各低位宽 dtype 通用。
+大白话:矩阵乘就像搬砖盖楼——Cube 是起重机(算力),GM 带宽是供货卡车(带宽)。算存比就是"每块砖要盖多少平米"。每块砖盖的平米数低于起重机的额定配比时,楼永远盖不快,问题出在供货而不在起重机——这时候优化要围绕"怎么少堵车、怎么一次多运点",而不是换更大的起重机。
+本章是全文逻辑主线:从分块计算的本质出发,不靠经验罗列,推导出五大分支的划分。
+BMM 在 NPU 上实现的本质是:把参与计算的数据分块(tile),由 32 个 AIC 核并行 + 串行地完成这些分块的计算,再组合成最终结果:
+$$ +C[B, M, N] = \sum_{\kappa} A[B, M, K_\kappa] @ B[B, K_\kappa, N] +$$
+分块有 4 个维度:B、M、N、K。核间怎么分这 4 个维度,就是分支划分的第一性问题(核内分块是第二性问题,属于各分支内部的 tiling)。
+从"读入 / 计算 / 写出"三个视角考察每个维度的核间切分特征(这是后续一切推导的基石):
+| 切分维度 | +读入特征 | +计算特征 | +写出特征 | +
|---|---|---|---|
| 切 B | +每个数据块只被固定的 1 个核读取,核间零重复读 | +每个输出块由 1 个核独立完成,无核间依赖 | +只写最终结果,无中间结果 | +
| 切 M / 切 N | +切 M 则同一右矩阵块被多核重复读;切 N 则同一左矩阵块被多核重复读 | +每个输出块由 1 个核独立完成,无核间依赖 | +只写最终结果,无中间结果 | +
| 切 K | +每个数据块只被固定的 1 个核读取,零重复读 | +每个输出块由多核共同完成,存在核间依赖 | +有中间结果写出,需核间 Reduce 归约 | +
为什么特征差异这么大?两条硬件层面的根本原因:
+cmatrixInitVal=false),中间结果不出核。一旦把 K 切到核间,单核的 L0C 装不下"别的核算的 K 段",部分和必须写出到 GM/L2 workspace,再由 AIV 归约——切 K 是唯一同时破坏"累加不出核"和"输出独占"两条性质的切法。4 个维度的任意非空子集都可作为一种核间切分方案,共 $2^4 - 1 = 15$ 种:
+$$ +\{B\},\{M\},\{N\},\{K\},\{B,M\},\{B,K\},\{B,N\},\{M,K\},\{M,N\},\{K,N\},\{B,M,K\},\{B,M,N\},\{B,K,N\},\{M,K,N\},\{B,M,K,N\} +$$
+任何实现方案必属于其中一种 ⇒ 这 15 种是完备的。按 §3.2 的特征分组:
+| 组 | +组合 | +共同特征 | +优化重心 | +
|---|---|---|---|
| 纯 B | +{B} | +零重复读,读写数据量固定 | +访存 Bound:搬移效率高;计算 Bound:计算效率高 | +
| 含 M/N 不含 K | +{M},{N},{M,N},{B,M},{B,N},{B,M,N} | +可能有重复读 | +访存 Bound:重复读尽量少 + 搬移效率高;计算 Bound:计算效率高 | +
| 含 K | +{K},{B,K},{M,K},{K,N},{B,M,K},{B,K,N},{M,K,N},{B,M,K,N} | +有中间结果写出 + 归约 | +计算效率高,且归约引入的额外 MTE2/Fixpipe 时延不能成为新瓶颈 | +
15 种组合的代价结构只由两个布尔特征决定——是否含 K、是否含 M/N。三个维度的"核间切分价格"严格排序:
+$$ +\text{cost}(\text{切}B) = 0 \;<\; \text{cost}(\text{切}M/N) \;\ll\; \text{cost}(\text{切}K) +$$
+"价格表"不是经验,是 BMM 语义 + L0C 累加机制 + L2/GM 带宽结构三条事实的推论。 整条分支决策树就是一句话:按价格从低到高购买并行度,买不够才加价。
+第 0 层:问题归约(能降维就不在 BMM 本体里解决)。
+BatchA = 1 或 BatchB = 1,BMM 可通过维度折叠转化为普通 Matmul,直接复用 Matmul 的成熟优化体系(tiling、L2 切分、全载),不必在 BMM 框架内重新发明轮子 → 分支一:转Matmul;第 1 层:归约不掉的 case(BatchA = BatchB = B > 1,K ≥ 2),必须在 BMM 框架内组织 4 维分块。先买免费的 B 维。
+第 2 层:免费的 B 买不满 32 核,加价买廉价的 M/N。
+第 3 层:廉价维度也买不满,才买昂贵的 K。
+case (B, M, N, K, dtype, layout, 广播形态)
+ │
+ ▼
+[0] K=0? → AIV 清零;K=1? → AIV 逐元素乘 ← 前置通路层(正交于切分)
+ │ K ≥ 2
+ ▼
+[1] BatchA=1 或 BatchB=1? ──是──→ 转Matmul(问题归约,复用 Matmul 体系)
+ │ 否(BatchA=BatchB=B>1,或交叉广播)
+ ▼
+[2] 并行度账本:P = B·⌈M/16⌉·⌈N/16⌉ 与 C=32 比较
+ │
+ ├─ B ≥ 32(切 B 可满核)─────────────────────────┐
+ │ M×N 大,Cube 饱 → IterBatch │
+ │ M×N 小,Cube 饿,且访存 Bound → MergeBatch │
+ │ │
+ ├─ B < 32 且 P ≥ 32 → ASW_Basic(切 M/N 补并行, │
+ │ 共享读取由 L2+swizzle 吸收) │
+ │ │
+ └─ P < 32(B/M/N 用尽仍缺并行,K 大)→ StreamK │
+ (切 K 买并行,付归约代价)◀──────────────────┘
+
+完备性。对归约不掉的 case,15 种组合按"(是否含 K,是否含 M/N)"坍缩:
+$1 + 1 + 2 = 4$,加上第 0 层的归约分支(转Matmul)= 5。任何合法 case 必落其一,无空洞。
+极小性。每个分支都有它"唯一最优"的 shape 区域,去掉任何一个都会有 case 失去最优实现:
+| 分支 | +独占最优的代表 case(BF16) | +替代方案为何更差 | +
|---|---|---|
| 转Matmul | +BatchB=1, B=128, M=N=K=2048 | +免费折叠后 Matmul 体系的 L2 切分/全载直接可用;BMM 分支内重做一遍无收益 | +
| IterBatch | +B=32, M=N=K=4096 | +ASW 切 M/N 引入无谓共享读;MergeBatch 引入无谓冗余算力 | +
| MergeBatch | +B=128, M=32, N=128, K=64 | +IterBatch 的 L0C 利用率仅 ~1.6%,Cube 空转,搬移 tile 碎(<16KB)带宽利用率崩 | +
| ASW_Basic | +B=2, M=N=8192, K=1024 | +切 B 仅 2 核干活;StreamK 付无谓归约 | +
| StreamK | +B=1, M=N=64, K=65536 | +不切 K 时 P=16 < 32,近半核闲置,时延差数量级 | +
⇒ 五大分支构成极小完备集。
+++把 BMM 想成给 32 个工人分一批"矩阵乘订单":
++
+- 能整单外包的就别自己干:一侧只有一个 batch,问题其实就是一次普通矩阵乘,直接交给成熟的 Matmul 产线(转Matmul);
+- 按订单分(切 B)最省心:每人几单,互不干扰(零重复搬料、零协调)。订单多时,单大就一单接一单干(IterBatch);单太小时一人同时干几单、拼成大活干,虽然多做了点无用功,反正瓶颈在等料不在干活(MergeBatch);
+- 订单不够分,就按行/列拆单(切 M/N):大家会重复领同一份料,但料放在近处仓库(L2),多跑几趟近仓库很便宜(ASW_Basic);
+- 行/列也不够拆,只好沿深度 K 拆:几个人合做同一块输出,各算一段再汇总——汇总是要额外开会的(Reduce),所以只在实在分不满、且 K 足够长时才这么干(StreamK)。
+
++每节按统一结构展开:做什么(定义)→ 什么时候进(进入条件逐条 + 每条的理由)→ 怎么做(实现方案与参数求解)→ 与源码对照(批判性)。
+
在切分决策之前先做计算通路判断:
+C = A ⊙ B,无累加深度,Cube 的 16×16×16 粒度浪费 15/16,走 AIV 向量通路(GM→UB→Mul→GM)远优于 Cube 通路。这一层与"切分维度"正交:五大分支管理的都是 Cube 通路上的 case(K ≥ 2)。K 极小(如 K < 16,不足一个 fractal)时同理应优先考虑向量通路或尾部补齐。
+做什么:当 BatchA = 1 或 BatchB = 1 时,把 batch 维折叠进矩阵维度,转化为普通 Matmul,计算完成后按需恢复 batch 维。
为什么这么做:单边 batch=1 的 BMM 与普通 Matmul 在数学上只差一个维度标签。Matmul 的优化体系(L2 切分、AL1/BL1 全载、swizzle)比 BMM 各分支在"单样本"情形下更成熟——站在巨人肩膀上,不重复造轮子。BMM 本体只研究相对 Matmul 有差异化的 case(双侧 batch > 1)。
+折叠规则与代价(关键的不对称性):
+| 情形 | +折叠方式 | +是否免费 | +
|---|---|---|
BatchB = 1 |
+左矩阵 [B, M, K] 的 batch 维与 M 维在 ND 布局下内存相邻,直接视图为 [B·M, K];输出 [B·M, N] 与 [B, M, N] 的内存布局逐元素一致 |
+完全免费:零输入重排、零输出 split | +
BatchA = 1 |
+需将右矩阵 [B, K, N] 折叠为 [K, B·N] |
+有代价:B 的 batch 维与 N 维在内存中不相邻(中间隔 K),折叠等价于一次 [B,K,N]→[K,B,N] 的转置重排(O(B·K·N) 读写),且输出 [M, B·N] 与目标 [B, M, N] 之间存在置换,需要随路 scatter |
+
++这一不对称性是对"直接折叠"设想的必要修正(详见 6.2)。
+BatchB = 1无脑走转Matmul;BatchA = 1需比较"重排一次 B + Matmul 收益"与"广播友好的 IterBatch/ASW_Basic(A 驻留 L2/L1 天然共享,不重排任何数据)"——大 case 重排代价被摊薄,走转Matmul;小 case 直接走 BMM 分支。源码中的MergeBatchAndMAxis()只做了 batchB=1 方向的折叠,与"该方向免费"的判断互为印证。
进入条件:BatchA = 1 || BatchB = 1。其中 BatchB = 1 恒进;BatchA = 1 时按上述代价比较决定(决策规则见第 5 章)。
BatchA = BatchB = 1:纯 Matmul,是转Matmul 的特例。
+核间按 B 分核(每核负责 $b_{core}$ 个 batch),核间无同步无通信。核内将多个 batch 合并计算:
+$$ +[b, M, K] @ [b, K, N] \;\Rightarrow\; [bM, K] @ [K, bN] = [bM, bN] \;\xrightarrow{\text{BlockTrace}}\; [b, M, N] +$$
+其中 BlockTrace 指以 $[M, N]$ 的 block 粒度取结果矩阵的块对角线作为各 batch 的有效输出:$C[i, m, n] = R[iM + m,; iN + n]$。交叉项(第 i 个 batch 的 A 乘第 j≠i 个 batch 的 B)被算出但丢弃——这就是"算力浪费",浪费比例 $(b-1)/b$。
+为什么允许浪费:进入该分支的 case 必然是访存 Bound(条件三保证),瓶颈在 MTE2 不在 MMAD,浪费的算力被搬移时延掩盖(§2.1 瓶颈交换)——用本来闲置的 Cube 算力,换 tile 变大后的搬移效率与 Cube 利用率。
+设计原则四条:① 硬件时延可流水掩盖(double buffer 乒乓);② 满足容量约束(每次计算 L1/L0A/L0B/L0C 放得下);③ GM→L1 搬移高效(§2.2 四条经验约束);④ 算力有浪费但计算不能成为瓶颈。
+形式化后,进入 MergeBatch 需同时满足:
+条件 1(batch 够分且够合并):BatchA = BatchB(无广播)且
$$ +b_{core} = B / aicNum \ge b_{thr} +$$
+$b_{thr}$ 是"多 batch 合并搬移能拿到效率收益"的最小合并数——$b$ 太小时合并的搬移收益抵不过实现复杂度,不如 IterBatch 且完全不浪费算力(功耗)。DV100 实测经验值 $b_{thr} = 4$。
+条件 2(合并数 b 的上下界):设核内单次 Cube 计算合并 $b$ 个 batch($b \le b_{core}$),要求 $b \ge b_{thr}$ 且 $b$ 由以下约束共同封顶:
+$$ +AI(b) < R \;\Longleftrightarrow\; b < b_{AI} = \frac{R \cdot (M + N)}{2MN} +$$
+$$ +2 \cdot (bM)(bN) \cdot 4\text{B} \le 256\text{KB} \;\Rightarrow\; b \le b_{L0C} +$$
+++乒乓取舍:若 $b_{core}$ 小于"不乒乓时 L0C 允许的合并数上限",说明 batch 余量不足,只能不乒乓(牺牲流水掩盖换合并数);否则乒乓。形式上 $b_{max} = \min(b_{AI}, b_{L0C}^{(\text{是否乒乓})}, b_{core})$。
+
条件 3(访存 Bound 性质):
+$$ +\frac{2MN}{M + N} < \frac{R}{b} +$$
+即 case 固有算存比显著低于平衡点(低一个合并倍数 $b$ 的量级)——这是"浪费可被掩盖"的定量保证,也是条件 2 中 $b_{AI} > b$ 的另一写法。
+条件 4(K 向搬移效率):L1 级 K 切分 $k_{L1} = K / StepK$(StepK 为正整数),要求 $k_{L1} \ge k_{thr}$,DV100 下 $k_{thr} = 32\text{B}/\text{dtype}$(BF16 即 16 元素)。理由:$k_{L1}$ 决定 ND2NZ 的 dValue,过碎则搬移指令效率崩(§2.2 第 4 条)。
+条件 5(L1 驻留与 tile 效率):L1 级 batch 驻留数
+$$ +b_{L1} = \frac{L1_{size}}{(M k_{L1} + k_{L1} N) \cdot \text{dtype}}, \qquad b_{L1} > b +$$
+(L1 驻留的 batch 组必须大于单次计算的合并数,否则合并无从谈起),且搬移 tile 大小满足效率:左矩阵非转置时 $\max(b_{L1} M k_{L1} \cdot \text{dtype},; k_{L1} N \cdot \text{dtype}) > min_TileSize$(16KB);左矩阵转置时相应调整。
+条件 6(单核搬移总量):
+$$ +b_{core} \cdot (M k_{L1} + k_{L1} N) \cdot \text{dtype} \ge min\_DatamountPerCore = 480\text{KB} +$$
+对应 §2.2 第 2 条:单核搬移数据总量不足时带宽利用率上限被压低。
+条件汇总逻辑:条件 1 定资格(batch 够多),条件 2/3 定上限(别算出瓶颈、别撑爆 L0),条件 4/5/6 定下限(搬移效率不能崩)。上限与下限之间必须有交集,交集为空则该 case 与 MergeBatch 无缘。
+进入分支后,$b$ 与 baseK 的具体取值有一个重要性质:在访存 Bound 前提下,$b < b_{AI}$ 时 b 的取值不影响性能,baseK 只需满足分形约束(16 倍数)也不影响算存比——因为瓶颈是搬移,算力余量内怎么切都一样。据此:
+case:B=128, M=32, K=64, N=128, BF16,32 核。
+结论:该 case 在"单核搬移总量"一项上不达标——这正是 MergeBatch 条件体系的用处:它提前告诉你瓶颈不在算力而在搬移效率,此 case 应尝试合并更多 K/提高 $b_{core}$ 数据量,或由时延模型与 IterBatch/ASW_Basic 比较后定夺。
+核间:32 核,每核 b_core 个 batch
+核内(每核):
+ for k_l1 in range(0, K, kL1): # L1 级 K 切分
+ MTE2: A[bL1组, M, k_l1:k_l1+kL1]、B[bL1组, k_l1:k_l1+kL1, N] → L1(乒乓)
+ for b_start in range(0, bL1, b): # 按合并数 b 分组
+ for k_l0 in range(0, kL1, kL0): # L0 级 K 切分
+ MTE1: A[bM, kL0] → L0A;B[kL0, bN] → L0B
+ Cube: [bM, kL0] @ [kL0, bN] → L0C 累加
+ Fixpipe: BlockTrace 取 b 个 [M, N] 对角块 → L2/GM
+
+源码 batch_matmul_v3_mergebatch_basicapi_tiling.cpp 的 IsCapable 条件:各级 batchA_i == batchB_i、batchC ≥ 4 × aicNum、alignK ≥ 64、M ≤ N、无 bias、非 NZ、拒绝非连续转置等。对照分析:
batchC ≥ 4 × aicNum 与草稿 $b_{core} \ge b_{thr} = 4$ 一致(每核至少 4 batch)——合理;alignK ≥ 64 比草稿 $k_{thr}$=32B/dtype(BF16 为 16 元素)保守 4 倍——源码更粗,可能误杀 K∈[16,64) 的可获益 case;M ≤ N 是实现便利性约束而非原理约束——M > N 的镜像 case 在原理上同样可合并(交换合并方向),源码直接放弃,属于覆盖缺口;核间按 B 分核(每核 1 个或多个 batch),核间无同步无通信;核内逐个 batch 分别执行标准 Matmul 分块(L1→L0A/L0B→Cube→L0C→Fixpipe)并输出。无算力浪费、无跨 batch 依赖——是"切 B"最朴素的形态。
+满足以下任一:
+类 1(计算 Bound 型):
+$$ +AI_{full} = \frac{2MKN}{MK + KN + MN} \ge R \quad \text{且} \quad BatchA = BatchB,\; b_{core} = B / aicNum \in \mathbb{Z}^+ +$$
+理由:连输出写出都计入仍是计算 Bound,则瓶颈恒在 Cube——只要负载均衡($b_{core}$ 整除,各核同量),逐 batch 计算就是满算力实现,无需任何花哨。
+类 2(输出驻留 L2 的均衡型):
+$$ +\frac{2MN}{M+N} \ge R_{读GM} \;\;\text{且}\;\; 2K \ge R_{写L2} \;\;\text{且}\;\; MN \cdot \text{dtype} \le L2_{size} \;\;\text{且}\;\; BatchA = BatchB,\; b_{core} \in \mathbb{Z}^+ +$$
+理由:读侧访存不弱($AI \ge R_{读GM}$);单 batch 输出 $MN \cdot \text{dtype} \le 128$MB 可驻留 L2——Cube 输出写 L2 即算完成,写 GM 的流量被省掉;$2K$ 是"每写出一个元素对应的计算量",$2K \ge R_{写L2}$ 保证写 L2 的 5.2TB/s 也不是瓶颈。三个条件合起来 = 读、算、写三条路都不堵。
+类 3(访存 Bound 型):$\dfrac{2MN}{M+N} < R$ 且同时满足:
+B mod aicNum > minCoreNum 会误杀整除 case(余数 0 恰恰是完全均衡),判据应以均衡比为准、余数规则为其近似;核间:32 核分 batch(尽量整除,尾波核数 ≥ minCoreNum)
+核内(每核):
+ for batch in 本核的 b_core 个 batch: # 逐个 batch
+ for m_tile / n_tile(核内标准 tiling):
+ for k_tile in range(0, K, baseK):
+ MTE2 预取下一 k_tile → L1(双缓冲)
+ MTE1: L1 → L0A/L0B
+ Cube: mmad → L0C 原地累加 # K 循环不出核
+ Fixpipe: L0C → L2(写 L2 即完成,GM 回写可异步)
+
+当 B、M、N 三个维度切到最碎仍填不满 32 核时,把 K 维切到核间:多核各算一段 K 的部分和,再归约:
+$$ +C_{(\beta,\mu,\nu)} = \sum_{c \in group} C^{(c)}_{(\beta,\mu,\nu)} \quad \text{(部分和写 workspace,AIV 归约或原子加)} +$$
+条件 1(并行缺口存在):不切 K 时的独立输出块数
+$$ +P = B \times \lceil M / 16 \rceil \times \lceil N / 16 \rceil < C = 32 +$$
+(源码实现取更保守的 $B \cdot mCnt \cdot nCnt \le aicNum/2$,并附加 ND-only、无交叉广播、确定性等级 ≤ 1(原子加顺序不定)等工程限制)。注意**严格的 StreamK 不要求核间完全不切 B/M/N**——它是一般框架 `grid_K × grid_B × grid_M × grid_N ≤ C`,纯切 K 只是 grid_B=grid_M=grid_N=1 的特例;当 B/M/N 能提供部分并行度时,应该用组合 grid 把归约组 $grid_K$ 压到最小(grid_K 小一档,K 的门槛降一档平方级)。
+条件 2(归约代价可接受):每核计算时延须远大于归约时延(安全系数 α = 10):
+$$ +T_{MMAD/core} \ge \alpha \cdot T_{Reduce} +\;\Longleftrightarrow\; +\frac{K}{grid_K} \;\gtrsim\; grid_K \times 1690 +$$
+即 $K \gtrsim grid_K^2 \times 1690$:grid_K=2 → K ≥ 6.8K;grid_K=4 → K ≥ 27K;grid_K=8 → K ≥ 108K;grid_K=32 → K ≥ 1.7M(仅极端 case)。同时 $K / grid_K \ge 256$(单核 K 段过碎则 tiling 效率崩)。grid_K 越大对 K 的要求越苛刻——StreamK 内部的 grid 搜索自然淘汰归约过重的配置。
+StreamK 从不"硬切换"进入:当不切 K 的分支候选已足够快(如已 Cube Bound),StreamK 候选的归约开销使其自然落败;只有当 B/M/N 并行度不足导致大量核闲置时,StreamK 才以数量级优势胜出(例:B=1、M=N=64、K=65536 时,不切 K 仅 16 核可用,StreamK 32 核满负荷)。
+不切 K,允许切 B/M/N 的任意组合——它是"不含 K 且含共享读取"的 6 种切分组合({M},{N},{M,N},{B,M},{B,N},{B,M,N})的统一实现框架,也是 B < 32 但 P ≥ 32 时的最优归宿,同时兜住 B ≥ 32 但 IterBatch/MergeBatch 条件不满足的剩余 case。B 可以大、可以小、可以等于 1;ASW_Basic 是实践中最常命中的分支。
+切 M/N 的固有代价是共享矩阵的重复读,ASW_Basic 用两件武器把代价压到最低:
+按共享代价从低到高:
+第 4 章的进入条件给出的是各分支的"主场",但主场之间有重叠区(例如 B ≥ 32 且 M×N 中等时,MergeBatch 与 IterBatch 都合法;B 较大时 IterBatch 与 ASW_Basic 切 B 等效)。重叠区的最优归属没有解析解,必须由统一的时延模型仲裁:
+$$ +\text{branch}^* = \arg\min_{cand \in \bigcup \text{各分支候选}} \max\big(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{Fixpipe}[, T_{Reduce}]\big) +$$
+分支体系的价值在于候选集完备且无冗余:每个等价类只派一个代表框架生成候选,时延模型在等价类内部和边界上做精细仲裁。两层缺一不可——只有模型没有分支,搜索空间是 15 种组合 × 全部 grid 参数的爆炸;只有分支没有模型,边界 case 被硬阈值误杀。
+输入:B, M, N, K, dtype, layout, 广播形态, bias
+ │
+ ├─[0] K = 0 → AIV 清零;K = 1 → AIV 逐元素乘(TO_MUL 通路)
+ │
+ ├─[1] BatchB = 1 → 转Matmul(免费折叠 [B·M, K],必选)
+ │ BatchA = 1 → 比较"重排 B + Matmul" vs "广播友好 IterBatch/ASW",
+ │ 按时延模型选小者
+ │
+ ├─[2] BatchA = BatchB = B > 1:
+ │ 生成四类候选并逐一估算 T_total:
+ │ MergeBatch(条件 1~6 全过才生成,见 4.2.2)
+ │ IterBatch(三大类条件,见 4.3.2)
+ │ ASW_Basic(按 4.5.3 顺序试切 B/M/N/混合,含 swizzle 与 L2 切分)
+ │ StreamK(P < 32 且 K 满足 4.4.2 条件 2 时生成,grid 搜索)
+ │ → argmin T_total
+ │
+ └─[3] 输出:分支 + tiling 参数(baseM/baseN/baseK、b、bL1、kL1…)
+ + swizzle 方案 + L2 切分方案 + 预估端到端时延
+
+按"转Matmul / MergeBatch / IterBatch / StreamK / ASW_Basic 五大分支 cover 全部 BMM case 最优实现"的设计目标,逐条拷问。
+| shape 区域 | +主分支 | +兜底 | +
|---|---|---|
| K = 0 / K = 1 | +前置通路层(AIV) | +— | +
| BatchA=1 或 BatchB=1 | +转Matmul | +ASW_Basic(广播友好形态) | +
| BatchA=BatchB>1,B≥32,M×N 大 | +IterBatch | +ASW_Basic 切 B | +
| BatchA=BatchB>1,B≥32,M×N 小,访存 Bound | +MergeBatch | +IterBatch | +
| 上述但 B<32,P≥32 | +ASW_Basic | +降核 ASW | +
| P<32,K 大 | +StreamK | +降核 ASW | +
| P<32,K 也小 | +降核 ASW_Basic | +(此时时延绝对值小,调度开销主导,分支选择不敏感) | +
| 交叉广播(BatchA≠BatchB 且均>1) | +ASW_Basic | +— | +
设计骨架成立,但有五处需要修正或显式声明,否则会被反例拷问:
+问题 1:转Matmul 的方向不对称(原设想的最大漏洞)。 "一侧 batch=1 就折叠成 Matmul"在 BatchB=1 方向免费(batch 与 M 内存相邻,输出布局一致);但在 BatchA=1 方向,右矩阵 [B,K,N] 折叠成 [K,B·N] 需要一次真实的转置重排,且输出 [M,B·N] 与目标 [B,M,N] 之间存在置换——"只是输出后按 batch split"的设想在 ND 布局下不成立,split 实际是 scatter。修正:BatchA=1 时按"重排代价 vs Matmul 收益"走模型仲裁,小 case 留在 BMM 分支内(A 单侧数据天然可 L2/L1 驻留共享,本身已接近最优)。
+问题 2:交叉广播 case 的归属必须显式声明。 BatchA=(6,1)、BatchB=(1,6) 这类 case 既不进转Matmul(两侧都>1),也不满足 MergeBatch/IterBatch 的 BatchA=BatchB 前提——它们落入 ASW_Basic(对广播侧做驻留)。不声明就会显得"五分支有洞"。声明后完备性闭合。
+问题 3:K 退化不在五分支内,需前置通路层。 K=0(无计算)、K=1(无累加深度)走 AIV 向量通路;这是"计算通路选择"维度,与"切分维度"正交,放在五分支之前判断。源码用 K_EQUAL_ZERO / TO_MUL 两个独立策略处理,印证了这一层的必要性——但它们是通路切换而非新切分等价类。
+问题 4:IterBatch 草稿条件 3(2) 的字面表述有误杀。 B mod aicNum > minCoreNum 会拒绝 B 整除核数的完全均衡 case(余数 0)。正确表述应以负载均衡比为准:访存 Bound 要求 bAvg/bMax > 0.8,计算 Bound 要求 100%;余数规则只是它的近似操作化。
问题 5:经验常数必须可标定,且源码硬编码条件偏粗。 $b_{thr}=4$、$k_{thr}$=32B/dtype、480KB、16KB、dValue≥256B 都是 DV100 档芯片的实测经验值,文档中全部显式参数化(换芯片只换常数表)。对照之下源码把 batchC ≥ 4×aicNum、alignK ≥ 64、M ≤ N 等直接写死,其中 alignK ≥ 64 比草稿阈值保守 4 倍、M ≤ N 放弃了镜像 case——做最优实现时不应继承这些硬编码,而应按第 4 章的条件体系重新判定。
源码 arch35 的 10 个策略(K_EQUAL_ZERO / TO_MUL / STREAM_K / MERGE_BATCH / ITER_BATCH_BROADCAST / ITER_BATCH / AL1_FULL_LOAD / BL1_FULL_LOAD / ASW_BASIC / BASE)与本文五分支不是同层概念:
+| 源码策略 | +本文归属 | +
|---|---|
| K_EQUAL_ZERO、TO_MUL | +前置通路层(4.0),正交于切分 | +
| STREAM_K | +StreamK | +
| MERGE_BATCH | +MergeBatch | +
| ITER_BATCH、ITER_BATCH_BROADCAST | +IterBatch(broadcast 变体是广播输入下的数据复用特化;单边广播在本文更多由转Matmul 吸收) | +
| AL1_FULL_LOAD、BL1_FULL_LOAD | +ASW_Basic 的 L1 驻留参数极限(4.5.4) | +
| ASW_BASIC、BASE | +ASW_Basic | +
结论:源码策略集 = 五分支 × 计算通路 × 驻留/广播特化的展开。五分支是"切分等价类"的最小完备骨架,源码的冗余策略可在新实现中收敛,源码遗漏的(条件 3/4/6 的显式判定、M>N 的 MergeBatch 镜像、转Matmul 的 BatchA=1 方向)正是优化的增量空间。
+参考文档:《BMM分块计算数学公式》《BMM最优软件实现方案设计》《BatchMatMulV3算子分支实现分析》《BMM从分块计算到四大分支的逻辑推导》;源码:ops-nn/matmul/batch_mat_mul_v3(arch35);芯片资料:昇腾 950PR 架构白皮书与 CANN 9.0.0 性能建模文档。