diff --git a/BMM从分块计算到四大分支的逻辑推导.html b/BMM从分块计算到四大分支的逻辑推导.html new file mode 100644 index 0000000..4cf47d1 --- /dev/null +++ b/BMM从分块计算到四大分支的逻辑推导.html @@ -0,0 +1,363 @@ + + + + + +从分块计算到四大分支:BMM 最优实现的逻辑推导 + + + + + +
+

从分块计算到四大分支:BMM 最优实现的逻辑推导

+
+

昇腾 BatchMatMulV3 算子 | 目标芯片 950PR(DAV_3510,32 AIC / 1.65GHz) +上游文档:《BMM分块计算数学公式》(分块计算"是什么") · 下游文档:《BMM最优软件实现方案设计》第六章(四大分支"怎么用") +本文定位:补上中间缺失的一层——为什么从"分块计算 + 最优实现"出发,恰好推导出 MergeBatch / IterBatch / ASW_Basic / StreamK 这四大分支,不多也不少。

+
+
+

0. 问题陈述:缺失的那层逻辑是什么

+

《BMM分块计算数学公式》告诉我们:BMM 的 4 重循环 $(b, m, n, k)$ 可以按硬件容量做嵌套分块,并用 swizzle 函数 $\sigma$ 重排块到核的映射。《BMM最优软件实现方案设计》第六章直接给出了四大分支。但中间有一跳没有论证:

+
+

分块参数 $(B^t, M^t, N^t, K^t)$、核间切分方式、swizzle 函数 $\sigma$ 有无穷多种取法。凭什么最优实现的搜索空间恰好收敛到 4 个分支

+
+

本文的推导链条如下,每一节对应链条上的一环:

+
分块计算公式(自由度)
+  → 硬件强加的两条不变式(规则)
+  → 四个维度核间切分的代价表(代价不对称)
+  → "最优"的三个必要条件(目标函数展开)
+  → 决策树:按代价从低到高购买并行度(推导)
+  → 15 种切分组合坍缩为 4 个等价类(完备性 + 极小性证明)
+  → 与分块公式的参数特化对应(闭环)
+
+
+

1. 起点:分块计算公式给出的自由度

+

分块计算把 BMM 重组为两层结构(记号同《BMM分块计算数学公式》):

+

核间:$\tilde{B} \times \tilde{M} \times \tilde{N}$ 个基本块,经 swizzle $\sigma$ 映射到 $C$ 个核:

+

$$ +(\beta, \mu, \nu) = \sigma^{-1}\big((c + rC) \bmod \tilde{B}\tilde{M}\tilde{N}\big), \qquad c \in [0, C) +$$

+

核内:K 循环在 L0C 上原地累加,经 L1/L0 两级缓冲流水:

+

$$ +C_{(\beta,\mu,\nu)}^{\text{L0C}} = \sum_{\kappa=0}^{\tilde{K}-1} \mathrm{mmad}\big(A[\beta,\mu,\kappa],\; B[\beta,\kappa,\nu]\big) +$$

+

这个数学结构里,自由决策只有三类

+ + + + + + + + + + + + + + + + + + + + + + + + + +
自由度数学对象物理含义
F1:核间切分$\tilde B, \tilde M, \tilde N$ 如何对 $C$ 分解;$\kappa$ 是否跨核哪些维度的块被分到不同核
F2:核内组织$B^t, M^t, N^t, K^t$ 及 L1/L0 驻留深度 $d_A, d_B$单核内如何折叠/驻留/流水
F3:执行顺序swizzle 函数 $\sigma$时间上相邻的块在空间上的排布(L2 复用)
+

"四大分支"问题,本质上就是 F1 和 F2 的选择问题(F3 是每个分支内部的二级优化)。所以推导的主线是:F1 有哪些本质上不同的选择,各自的代价是什么。

+
+

2. 规则:硬件强加的两条不变式

+

分块公式里有句话写得轻描淡写,却是整个分支体系的根:

+
+

"K 循环 ($\kappa$) 在 L0C 上累加,不写回 GM。"

+
+

把它和输出写回方式放在一起,就是硬件强加给所有实现的两条不变式:

+

不变式 I1(L0C 累加不变式):一个输出块 $C_{(\beta,\mu,\nu)}$ 的 $\tilde K$ 轮 mmad 结果驻留在 L0C(256KB,FP32)上原地累加,中间不产生任何 GM/L2 流量。

+ +

$$ +C_{(\beta,\mu,\nu)} = \sum_{c \in \text{group}} C_{(\beta,\mu,\nu)}^{(c)} \quad \text{(部分和经 workspace 归约)} +$$

+

不变式 I2(输出专属不变式):若输出块 $C_{(\beta,\mu,\nu)}$ 由一个核独占负责,则该核只写出最终结果,无核间同步;反之(切 K)则引入核间依赖与归约流量。

+

两条不变式合起来给出本文最重要的结构性事实:

+
+

K 维和其他三维在核间切分中的地位根本不对等:切 B/M/N 保持 I1、I2,切 K 同时打破 I1、I2。这就是为什么"是否切 K"天然是分支的第一分界线,StreamK 必然自成一支。

+
+
+

3. 关键一步:四个维度核间切分的代价表

+

对 F1,核间切分的本质选择是"4 个维度中切哪些"。逐维度分析切分代价(依据:batch 语义、数据通路、GM/L2 带宽差):

+

3.1 切 B —— 零代价并行

+

不同 batch 的 $A[b], B[b], C[b]$ 在内存中天然不相交。因此:

+ +

代价 = 0。但并行度上限 = $B$,且单核分到的计算是**单个 batch 的 $[M,K]\times[K,N]$**——Cube 效率完全由单 batch 的 $M, N$ 决定(第 4 节会看到这是 MergeBatch 与 IterBatch 分家的原因)。

+

3.2 切 M 或切 N —— 低代价并行(代价可被 L2 吸收)

+

切 M:同一 batch 的右矩阵 $B[b,:,:]$(大小 $K \times N$)被负责不同 M 行块的所有核**重复读取**;切 N 对称(左矩阵 $M \times K$ 被重复读)。

+ +

代价 = 共享矩阵的重复读取,受 L2 容量约束,属于有上限的低代价

+

3.3 切 K —— 高代价并行(结构性代价,无法吸收)

+ +

代价 = 归约流量 + 核间同步,是随 grid_K 线性增长的结构性代价,L2 吸收不掉。

+

3.4 代价不对称性总结

+

$$ +\boxed{\;\text{cost}(\text{切}B) = 0 \;<\; \text{cost}(\text{切}M) \approx \text{cost}(\text{切}N) \;\ll\; \text{cost}(\text{切}K)\;} +$$

+

这个排序不是经验,是三条硬件事实的推论:① batch 维在数学上独立(BMM 语义);② L0C 累加机制(I1)使切 K 产生独一份的归约开销;③ L2(5.2TB/s / 128MB)恰好能把 M/N 共享读取的代价吸收掉大半。整条决策树就是按这个代价排序"从便宜到贵"地购买并行度。

+
+

4. 目标:最短端到端时延展开为三个必要条件

+

最优实现的判据是 $T_{total} = \max(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{FIXPIPE} [, T_{REDUCE}])$ 最小。逐项拆开,$\min\max$ 等价于三个必要条件:

+

R1 核占满(并行度条件)。$T_{MMAD}$ 与活跃核数成反比;GM 带宽也要求 ≥ 3/4 核(24 核)并发才能达到 90%+ 利用率。设不切 K 时的独立输出块数

+

$$ +P \;=\; B \times \Big\lceil \tfrac{M}{16} \Big\rceil \times \Big\lceil \tfrac{N}{16} \Big\rceil +$$

+

(16 是 mmad 的 fractal 粒度,即输出块的最小可分单元)。$P \ge C = 32$ 是满核的必要条件;$P < 32$ 时必然有核闲置,闲置部分纯浪费。

+

R2 搬运最快(复用条件)。MTE2 有效带宽由 L2 命中率决定:

+

$$ +BW_{MTE2} \approx (1 - h) \times 1.6 + h \times 5.2 \ \text{TB/s} +$$

+

baseM = baseN = 256 时达成 Cube Bound 需要 $BW_{MTE2} \ge 2.64$ TB/s,即 $h \gtrsim 28.9%$。GM 流量必须最小化(每份数据 GM 只读一次),重复读取尽量落在 L2/L1。

+

R3 Cube 喂饱(效率条件)。单核每次 mmad 序列的 tile 要足够大:L0C 利用率

+

$$ +U = \frac{M^t \times N^t \times 4\text{B}}{256\text{KB}} = \frac{M^t N^t}{65536} +$$

+

越接近 1,Cube 空转越少;$M^t, N^t$ 需 16 对齐,且 MTE1 不拖后腿要求 baseM, baseN ≥ 80 左右。

+

分支推导 = 在 R1/R2/R3 约束下,按 §3 的代价排序购买并行度与 Cube 效率。 下面逐分支看它们各自是哪一组约束矛盾的"最优解"。

+
+

5. 推导主链:四种典型矛盾 → 四个最优响应

+

5.1 情形一:$B \ge C$ 且单 batch $M \times N$ 大 → IterBatch

+ +

此时任何额外的切分都只会引入代价而无收益,最优解就是"每核分若干 batch,核内逐个 batch 做完整 Matmul"。这正是分块公式中 $B^t = 1$、batch 作为核间最外层维度的特化。

+

判据:$B \ge C$ 且 $M N \gtrsim 128^2$(L0C 利用率 ≥ 25% 量级)。

+

5.2 情形二:$B \ge C$ 但单 batch $M \times N$ 小 → MergeBatch

+ +

矛盾在 R3。唯一的解法:把"折叠维度"用在 batch 上——核内把 $b$ 个 batch 的 A 沿 M 拼接、B 沿 N 拼接,等效大矩阵 $[bM, K] \times [K, bN] \to [bM, bN]$,L0C 利用率从 $MN/65536$ 提升到 $b^2MN/65536$(受 L0C 约束 $b^2MN \le 65536$,如 $M{=}N{=}64$ 时 $b \le 4$),最后用 BlockTrace 取 $b$ 个对角线 $M \times N 块作为有效输出。

+

代价是交叉项算力浪费 $(b-1)/b$。为什么这个代价值得付?因为小 $M, N$ 意味着算术强度低:

+

$$ +AI = \frac{b M N K \cdot 2}{(b M K + K b N + b^2 M N)\,\text{dtype}} < 270 \ \text{FLOP/B} \;\Rightarrow\; \text{访存 Bound} +$$

+

单核分摊的 GM 带宽约 50GB/s,Cube 要满转需要 AI ≈ 270 FLOP/B;小 $M,N$ 的 case 远低于此,瓶颈本来就在搬运,Cube 浪费的拍数被 MTE2 时延掩盖——浪费是免费的。这就是"为什么是 MergeBatch"的定量理由:它不是牺牲算力换效率,而是在算力本就用不完时回收闲置。

+

判据:$B \ge 2C$(每核至少 2 batch)且 $MN < 128^2$ 且 $AI < 270$。

+

5.3 情形三:$B < C$ 但 $B \times \lceil M/16 \rceil \times \lceil N/16 \rceil \ge C$ → ASW_Basic

+ +

矛盾是"R1 缺并行、只能靠低代价维度补"。最优解是一个通用框架:允许切 B/M/N 的任意组合(不切 K),核间切分维度与 swizzle 由 shape 决定——切 B 优先(零共享),B 不够切 M 或 N(共享一侧矩阵,L2 吸收),再不够混合切。这就是 ASW_Basic:它不是一个具体切法,而是"所有不切 K、含 M/N 共享切法"的总框架,$\sigma$ 取 ASW 滑窗蛇形。

+

判据:$P \ge C$ 且不满足情形一/二(B 不够,或 M/N 大到不需要折叠 batch)。

+
+

与 IterBatch 的竞争边界:$B \ge 32$ 且 $M \times N$ 中等时两者都可行。分野在 L2:IterBatch 每核做完整 $M \times N \times K$,若单 batch 工作集 $MK + KN$ 超 L2 则 M/N 外循环反复挤兑 L2;ASW 切 M 时右矩阵 $KN \le 128$MB 可驻留 L2 供 32 核共享。谁的工作集能驻留 L2,谁优——这是时延模型比较,不是硬阈值(见 §7)。

+
+

5.4 情形四:$P = B \lceil M/16 \rceil \lceil N/16 \rceil < C$ → StreamK

+ +

$$ +T_{MMAD/\text{core}} \gg T_{REDUCE} \;\Longleftrightarrow\; K \;\gtrsim\; \text{grid}_K^2 \times 1690 +$$

+

(推导:要求计算时延 ≥ 10× 归约时延,代入 8192 FLOP/拍、1.65GHz、GM 1.6TB/s)时,切 K 净收益为正。K 越大,可承担的 grid_K 越大;K 不够大时,减少 grid_K(配合切 B/M/N,即 grid_K × grid_B × grid_M × grid_N ≤ 32 的组合)来降低归约组大小。

+

矛盾是"便宜维度用尽仍缺并行"。最优解是付归约代价买 K 维并行,且归约组能小则小。这就是 StreamK——分块公式中 $\tilde K$ 跨核拆分、部分和经 workspace 归约的特化。

+

判据:$P < C$(或 P 虽够但核内 M/N 范围被压得太碎)且存在 grid_K 使 $K/\text{grid}_K \ge 256$ 且满足上式。

+

5.5 一条链总结

+
R1 缺并行?────────────────────────────────────────────
+   │ 不缺                            │ 缺
+   ▼                                ▼
+R3 缺 Cube 效率?              便宜维度(B/M/N)用尽?
+   │ 不缺        │ 缺               │ 用尽
+   ▼            ▼                  ▼
+B 够切?     MergeBatch        StreamK(切K,付归约代价)
+ │B≥C  │B<C   (折叠batch
+ ▼      ▼     换Cube效率,
+Iter  ASW_Basic 浪费被访存Bound掩盖)
+Batch (切M/N共享,
+       L2+swizzle吸收代价)
+
+

四大分支各自是唯一最优解的"矛盾区域"互不相同:IterBatch 解"零代价并行已够"的区域,MergeBatch 解"并行够但 Cube 饿"的区域,ASW_Basic 解"并行缺、低代价维度可补"的区域,StreamK 解"便宜维度用尽"的区域。区域不同,最优响应不同——这就是四个分支的存在性证明。

+
+

6. 完备性与极小性:为什么恰好四个,不多不少

+

6.1 完备性:15 种切分组合按代价特征坍缩为 4 个等价类

+

核间切分的所有可能 = 四维 {B, M, N, K} 的非空子集,共 $2^4 - 1 = 15$ 种(《方案设计》§6.1 已枚举 C1~C15)。关键观察:一个切分组合的代价结构只由两个布尔特征决定——

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
含 K?含 M/N?组合代价结构归入分支
任意C4, C7, C9, C10, C12~C15(8 个)必有归约;grid_K×grid_B×grid_M×grid_N 只是参数差异StreamK
C2, C3, C5, C6, C8, C11(6 个)必有共享读取;切哪几维只是 grid 参数ASW_Basic
否(纯 {B})C1(1 个)零共享零归约核内只有两种组织方式,见下
+

纯 {B} 的核内组织方式只有两种——把多个 batch 合并成一个大 tile 算(MergeBatch),或逐个 batch 算(IterBatch)——不存在第三种(要么利用 batch 间的 tile 级合并,要么不利用)。于是:

+

$$ +15 \;\xrightarrow{\text{按 (含K, 含M/N) 归并}}\; 1 + 1 + 2 \;=\; \boxed{4} +$$

+

任何合法 case 的任意切法都落在这 4 个等价类之一 ⇒ 完备

+
+

与源码对照:arch35 源码有 10 个策略,为什么这里只有 4 个?因为源码策略 = 4 个切分等价类 × 两个正交维度的笛卡尔积:① 计算通路(Cube vs AIV:K=0 的 K_EQUAL_ZERO、K=1 的 TO_MUL 是 K 退化时的通路切换,发生在切分决策之前,属于预处理层);② 驻留策略(AL1/BL1_FULL_LOAD 是 ASW_Basic 内部 $M^t = M$ / $N^t = N$ 的 tiling 极限,ITER_BATCH_BROADCAST 是 IterBatch 在广播输入下的数据复用特化)。它们是分支内部的参数特化,不构成新的切分等价类。源码尾部还有 BASE=999 无条件兜底,也印证"特判 ⊂ 通用"的层级结构。

+
+

6.2 极小性:每个分支都有它是唯一最优的 shape 区域

+

去掉任何一个分支,都存在 case 失去最优实现:

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
分支独占最优的示例 case(BF16)替代方案的劣势
IterBatchB=32, M=N=K=4096ASW 切 M/N 引入无谓共享读取;MergeBatch 引入无谓浪费
MergeBatchB=64, M=N=64, K=256IterBatch 的 L0C 利用率仅 6.25%,Cube 闲置 ~94%
ASW_BasicB=2, M=N=8192, K=1024切 B 仅 2 核活跃;StreamK 付无谓归约
StreamKB=1, M=N=64, K=65536不切 K 只用 16/32 核(P=16),时延差数量级(文档 Case 3:19.5ms → 0.61μs/核量级)
+

⇒ 四者互为不可替代,构成极小完备集

+
+

7. 边界是软的:竞争区域由时延模型仲裁

+

上面的判据给出的是分支的"主场",但相邻分支的边界不是硬切换。重叠区域(如 B ≥ 32 且 M×N 中等时 IterBatch 与 ASW 切 B 等效;M×N 在 64²~128² 之间时 MergeBatch 与 IterBatch 互有胜负)必须由统一点评仲裁:

+

$$ +\text{branch}^* = \arg\min_{\text{cand} \in \bigcup \text{4 分支候选}} \max\big(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{FIXPIPE}, T_{REDUCE}\big) +$$

+

这就是《方案设计》§6.7 决策算法"全分支生成候选 → 时延模型评估 → 取 min"的合理性来源:分支体系保证候选集完备且无冗余(每个等价类只派一个代表框架),时延模型在等价类内部和边界上做精细仲裁。 两层结构缺一不可——只有时延模型没有分支体系,搜索空间是 15 种组合 × 全部 grid,爆炸且重复;只有分支体系没有时延模型,边界 case 会被错误硬切。

+
+

8. 闭环:与分块计算公式的参数特化对应

+

最后把推导结果写回分块公式的语言,验证链条闭合(对应《BMM分块计算数学公式》第六章,但现在每行都有了"理由"):

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
分支分块参数特化推导出处
IterBatch$B^t = 1$;$\tilde B \ge C$;batch 为核间最外层;核内标准 K 累加§5.1:零代价并行够 + Cube 饱
MergeBatch$B^t = b > 1$ 折叠进 $M^t_{eff} = bM^t, N^t_{eff} = bN^t$;L0C 二次方程求最优 $b$;BlockTrace 取对角§5.2:Cube 饿 → 折叠 batch 换 $U \uparrow$,浪费被访存 Bound 掩盖
ASW_Basic通用 $(\beta, \mu, \nu)$ 展开;$\sigma$ = ASW 滑窗蛇形($W \le \lfloor\sqrt C\rfloor$ 的最大因子) | §5.3:低代价维度补并行,$\sigma$ 压缩 L2 足迹
StreamK$\tilde K$ 跨核拆分,$C_{(\beta,\mu,\nu)} = \sum_c C^{(c)}$ 部分和归约;grid_K×grid_B×grid_M×grid_N ≤ C§5.4:便宜维度用尽,付归约代价买 K 维并行
+
+

9. 总结:一句话逻辑链

+
+

分块计算给出自由度(4 维怎么切),硬件给出规则(L0C 累加使切 K 独贵、batch 独立使切 B 免费、L2 使切 M/N 廉价),最优目标给出约束(核要满、Cube 要饱、搬运要省)。按"最便宜的维度优先购买并行度、Cube 不够大就折叠 batch、便宜维度用尽才切 K"的原则展开,15 种切分组合在 (含K, 含M/N) 两个代价特征下恰好坍缩为 4 个等价类——这就是 MergeBatch / IterBatch / ASW_Basic / StreamK 四大分支,完备且极小;边界区域由端到端时延模型统一仲裁。

+
+
+

文档版本:v1.0 | 上游:《BMM分块计算数学公式》v1.1+ | 下游:《BMM最优软件实现方案设计》v4.0 第六章

+ +
+ + diff --git a/BMM从分块计算到四大分支的逻辑推导.md b/BMM从分块计算到四大分支的逻辑推导.md new file mode 100644 index 0000000..c418efa --- /dev/null +++ b/BMM从分块计算到四大分支的逻辑推导.md @@ -0,0 +1,302 @@ +# 从分块计算到四大分支:BMM 最优实现的逻辑推导 + +> 昇腾 BatchMatMulV3 算子 | 目标芯片 950PR(DAV_3510,32 AIC / 1.65GHz) +> 上游文档:《BMM分块计算数学公式》(分块计算"是什么") · 下游文档:《BMM最优软件实现方案设计》第六章(四大分支"怎么用") +> 本文定位:补上中间缺失的一层——**为什么从"分块计算 + 最优实现"出发,恰好推导出 MergeBatch / IterBatch / ASW_Basic / StreamK 这四大分支,不多也不少。** + +--- + +## 0. 问题陈述:缺失的那层逻辑是什么 + +《BMM分块计算数学公式》告诉我们:BMM 的 4 重循环 $(b, m, n, k)$ 可以按硬件容量做嵌套分块,并用 swizzle 函数 $\sigma$ 重排块到核的映射。《BMM最优软件实现方案设计》第六章直接给出了四大分支。但中间有一跳没有论证: + +> 分块参数 $(B^t, M^t, N^t, K^t)$、核间切分方式、swizzle 函数 $\sigma$ 有无穷多种取法。凭什么最优实现的搜索空间恰好收敛到 **4 个分支**? + +本文的推导链条如下,每一节对应链条上的一环: + +``` +分块计算公式(自由度) + → 硬件强加的两条不变式(规则) + → 四个维度核间切分的代价表(代价不对称) + → "最优"的三个必要条件(目标函数展开) + → 决策树:按代价从低到高购买并行度(推导) + → 15 种切分组合坍缩为 4 个等价类(完备性 + 极小性证明) + → 与分块公式的参数特化对应(闭环) +``` + +--- + +## 1. 起点:分块计算公式给出的自由度 + +分块计算把 BMM 重组为两层结构(记号同《BMM分块计算数学公式》): + +**核间**:$\tilde{B} \times \tilde{M} \times \tilde{N}$ 个基本块,经 swizzle $\sigma$ 映射到 $C$ 个核: + +$$ +(\beta, \mu, \nu) = \sigma^{-1}\big((c + rC) \bmod \tilde{B}\tilde{M}\tilde{N}\big), \qquad c \in [0, C) +$$ + +**核内**:K 循环在 L0C 上原地累加,经 L1/L0 两级缓冲流水: + +$$ +C_{(\beta,\mu,\nu)}^{\text{L0C}} = \sum_{\kappa=0}^{\tilde{K}-1} \mathrm{mmad}\big(A[\beta,\mu,\kappa],\; B[\beta,\kappa,\nu]\big) +$$ + +这个数学结构里,**自由决策只有三类**: + +| 自由度 | 数学对象 | 物理含义 | +|---|---|---| +| F1:核间切分 | $\tilde B, \tilde M, \tilde N$ 如何对 $C$ 分解;$\kappa$ 是否跨核 | 哪些维度的块被分到**不同核**上 | +| F2:核内组织 | $B^t, M^t, N^t, K^t$ 及 L1/L0 驻留深度 $d_A, d_B$ | 单核内如何折叠/驻留/流水 | +| F3:执行顺序 | swizzle 函数 $\sigma$ | 时间上相邻的块在空间上的排布(L2 复用) | + +**"四大分支"问题,本质上就是 F1 和 F2 的选择问题**(F3 是每个分支内部的二级优化)。所以推导的主线是:**F1 有哪些本质上不同的选择,各自的代价是什么。** + +--- + +## 2. 规则:硬件强加的两条不变式 + +分块公式里有句话写得轻描淡写,却是整个分支体系的根: + +> "**K 循环 ($\kappa$) 在 L0C 上累加,不写回 GM。**" + +把它和输出写回方式放在一起,就是硬件强加给所有实现的两条不变式: + +**不变式 I1(L0C 累加不变式)**:一个输出块 $C_{(\beta,\mu,\nu)}$ 的 $\tilde K$ 轮 mmad 结果驻留在 L0C(256KB,FP32)上原地累加,中间不产生任何 GM/L2 流量。 +- 推论 1:**核内切 K 是免费的**($\tilde K$ 只是循环次数)。 +- 推论 2:**核间切 K 必然打破 I1**——每个核只算了一段 K 的部分和,L0C 驻留不住"别人的 K",部分和必须写出到 GM/L2 workspace,再由 AIV 做核间归约。分块公式第六章的 STREAM_K 行写的就是这个: + +$$ +C_{(\beta,\mu,\nu)} = \sum_{c \in \text{group}} C_{(\beta,\mu,\nu)}^{(c)} \quad \text{(部分和经 workspace 归约)} +$$ + +**不变式 I2(输出专属不变式)**:若输出块 $C_{(\beta,\mu,\nu)}$ 由一个核独占负责,则该核只写出最终结果,无核间同步;反之(切 K)则引入核间依赖与归约流量。 + +两条不变式合起来给出本文最重要的结构性事实: + +> **K 维和其他三维在核间切分中的地位根本不对等**:切 B/M/N 保持 I1、I2,切 K 同时打破 I1、I2。这就是为什么"是否切 K"天然是分支的第一分界线,StreamK 必然自成一支。 + +--- + +## 3. 关键一步:四个维度核间切分的代价表 + +对 F1,核间切分的本质选择是"4 个维度中切哪些"。逐维度分析切分代价(依据:batch 语义、数据通路、GM/L2 带宽差): + +### 3.1 切 B —— 零代价并行 + +不同 batch 的 $A[b], B[b], C[b]$ 在内存中天然不相交。因此: + +- **读入**:每个数据块被且仅被 1 个核读取,零重复; +- **计算**:每核独立产出最终结果,无核间依赖(保持 I1、I2); +- **写出**:只有最终结果,无中间写出。 + +代价 = 0。但并行度上限 = $B$,且单核分到的计算是**单个 batch 的 $[M,K]\times[K,N]$**——Cube 效率完全由单 batch 的 $M, N$ 决定(第 4 节会看到这是 MergeBatch 与 IterBatch 分家的原因)。 + +### 3.2 切 M 或切 N —— 低代价并行(代价可被 L2 吸收) + +切 M:同一 batch 的右矩阵 $B[b,:,:]$(大小 $K \times N$)被负责不同 M 行块的所有核**重复读取**;切 N 对称(左矩阵 $M \times K$ 被重复读)。 + +- **读入**:重复因子 = M(或 N)方向的核间切分数 $g_M$($g_N$)。若被共享矩阵能驻留 L2($K N \cdot \text{dtype} \le 128\text{MB}$),则重复读取以 5.2TB/s 命中 L2 而非 1.6TB/s 的 GM——代价大部分被 L2 吸收;swizzle $\sigma$(滑窗蛇形)进一步压缩同时活跃的工作集。 +- **计算**:无核间依赖(保持 I1、I2)。 +- **写出**:只有最终结果。 + +代价 = 共享矩阵的重复读取,受 L2 容量约束,属于**有上限的低代价**。 + +### 3.3 切 K —— 高代价并行(结构性代价,无法吸收) + +- **读入**:无重复(各核读不同 K 段); +- **计算**:**打破 I1、I2**——每个输出块由 $\text{grid}_K$ 个核共同产生,部分和写 workspace,再归约; +- **写出**:存在中间结果写出。额外时延项 $T_{\text{REDUCE}} \propto \text{grid}_K \times B_c M_c N_c \times 4\text{B} / \text{BW}$(AtomicAdd 到 L2 约 5.2TB/s、到 GM 约 1.6TB/s,另有同步开销)。 + +代价 = 归约流量 + 核间同步,是**随 grid_K 线性增长的结构性代价**,L2 吸收不掉。 + +### 3.4 代价不对称性总结 + +$$ +\boxed{\;\text{cost}(\text{切}B) = 0 \;<\; \text{cost}(\text{切}M) \approx \text{cost}(\text{切}N) \;\ll\; \text{cost}(\text{切}K)\;} +$$ + +这个排序不是经验,是三条硬件事实的推论:① batch 维在数学上独立(BMM 语义);② L0C 累加机制(I1)使切 K 产生独一份的归约开销;③ L2(5.2TB/s / 128MB)恰好能把 M/N 共享读取的代价吸收掉大半。**整条决策树就是按这个代价排序"从便宜到贵"地购买并行度。** + +--- + +## 4. 目标:最短端到端时延展开为三个必要条件 + +最优实现的判据是 $T_{total} = \max(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{FIXPIPE} [, T_{REDUCE}])$ 最小。逐项拆开,$\min\max$ 等价于三个必要条件: + +**R1 核占满(并行度条件)**。$T_{MMAD}$ 与活跃核数成反比;GM 带宽也要求 ≥ 3/4 核(24 核)并发才能达到 90%+ 利用率。设不切 K 时的独立输出块数 + +$$ +P \;=\; B \times \Big\lceil \tfrac{M}{16} \Big\rceil \times \Big\lceil \tfrac{N}{16} \Big\rceil +$$ + +(16 是 mmad 的 fractal 粒度,即输出块的最小可分单元)。$P \ge C = 32$ 是满核的必要条件;$P < 32$ 时必然有核闲置,闲置部分纯浪费。 + +**R2 搬运最快(复用条件)**。MTE2 有效带宽由 L2 命中率决定: + +$$ +BW_{MTE2} \approx (1 - h) \times 1.6 + h \times 5.2 \ \text{TB/s} +$$ + +baseM = baseN = 256 时达成 Cube Bound 需要 $BW_{MTE2} \ge 2.64$ TB/s,即 $h \gtrsim 28.9\%$。GM 流量必须最小化(每份数据 GM 只读一次),重复读取尽量落在 L2/L1。 + +**R3 Cube 喂饱(效率条件)**。单核每次 mmad 序列的 tile 要足够大:L0C 利用率 + +$$ +U = \frac{M^t \times N^t \times 4\text{B}}{256\text{KB}} = \frac{M^t N^t}{65536} +$$ + +越接近 1,Cube 空转越少;$M^t, N^t$ 需 16 对齐,且 MTE1 不拖后腿要求 baseM, baseN ≥ 80 左右。 + +**分支推导 = 在 R1/R2/R3 约束下,按 §3 的代价排序购买并行度与 Cube 效率。** 下面逐分支看它们各自是哪一组约束矛盾的"最优解"。 + +--- + +## 5. 推导主链:四种典型矛盾 → 四个最优响应 + +### 5.1 情形一:$B \ge C$ 且单 batch $M \times N$ 大 → **IterBatch** + +- R1:零代价并行度 $B \ge 32$ 已足够,无需切 M/N(省掉共享读取),更无需切 K。 +- R3:$M \times N$ 大(如 ≥ 128×128),单 batch 就能开出 baseM×baseN ≈ 256×256 的 tile,$U$ 接近 1,Cube 喂得饱。 +- R2:核内走标准三级流水(L1→L0→Cube→L0C),M/N 大 tile 保证复用。 + +此时**任何额外的切分都只会引入代价而无收益**,最优解就是"每核分若干 batch,核内逐个 batch 做完整 Matmul"。这正是分块公式中 $B^t = 1$、batch 作为核间最外层维度的特化。 + +**判据**:$B \ge C$ 且 $M N \gtrsim 128^2$(L0C 利用率 ≥ 25% 量级)。 + +### 5.2 情形二:$B \ge C$ 但单 batch $M \times N$ 小 → **MergeBatch** + +- R1:满足($B \ge 32$)。 +- R3:**不满足**。$M = N = 64$ 时 $U = 64^2/256^2 = 6.25\%$,Cube 大片闲置;$M = N = 16$ 时 $U = 0.4\%$,mmad 几乎全在空转。并行度够,但每个核"吃不饱"。 + +矛盾在 R3。唯一的解法:**把"折叠维度"用在 batch 上**——核内把 $b$ 个 batch 的 A 沿 M 拼接、B 沿 N 拼接,等效大矩阵 $[bM, K] \times [K, bN] \to [bM, bN]$,L0C 利用率从 $MN/65536$ 提升到 $b^2MN/65536$(受 L0C 约束 $b^2MN \le 65536$,如 $M{=}N{=}64$ 时 $b \le 4$),最后用 BlockTrace 取 $b$ 个对角线 $M \times N 块作为有效输出。 + +代价是交叉项算力浪费 $(b-1)/b$。为什么这个代价值得付?因为小 $M, N$ 意味着算术强度低: + +$$ +AI = \frac{b M N K \cdot 2}{(b M K + K b N + b^2 M N)\,\text{dtype}} < 270 \ \text{FLOP/B} \;\Rightarrow\; \text{访存 Bound} +$$ + +单核分摊的 GM 带宽约 50GB/s,Cube 要满转需要 AI ≈ 270 FLOP/B;小 $M,N$ 的 case 远低于此,**瓶颈本来就在搬运,Cube 浪费的拍数被 MTE2 时延掩盖**——浪费是免费的。这就是"为什么是 MergeBatch"的定量理由:它不是牺牲算力换效率,而是在算力本就用不完时回收闲置。 + +**判据**:$B \ge 2C$(每核至少 2 batch)且 $MN < 128^2$ 且 $AI < 270$。 + +### 5.3 情形三:$B < C$ 但 $B \times \lceil M/16 \rceil \times \lceil N/16 \rceil \ge C$ → **ASW_Basic** + +- R1:零代价维度 B 买不够 32 核,但低代价维度 M/N 可以补齐($P \ge 32$)。 +- R2:切 M/N 引入共享矩阵的重复读取——必须靠 L2 驻留 + swizzle 把代价压到最低,这正是 swizzle 函数 $\sigma$ 的主战场(滑窗 $W = \max\{d \mid d \mid C, d \le \lfloor\sqrt C\rfloor\}$ 使同窗口 A 行块与 B 列块的 L2 足迹最小)。 + +矛盾是"R1 缺并行、只能靠低代价维度补"。最优解是一个**通用框架**:允许切 B/M/N 的任意组合(不切 K),核间切分维度与 swizzle 由 shape 决定——切 B 优先(零共享),B 不够切 M 或 N(共享一侧矩阵,L2 吸收),再不够混合切。这就是 ASW_Basic:它不是一个具体切法,而是"所有不切 K、含 M/N 共享切法"的总框架,$\sigma$ 取 ASW 滑窗蛇形。 + +**判据**:$P \ge C$ 且不满足情形一/二(B 不够,或 M/N 大到不需要折叠 batch)。 + +> 与 IterBatch 的竞争边界:$B \ge 32$ 且 $M \times N$ 中等时两者都可行。分野在 L2:IterBatch 每核做完整 $M \times N \times K$,若单 batch 工作集 $MK + KN$ 超 L2 则 M/N 外循环反复挤兑 L2;ASW 切 M 时右矩阵 $KN \le 128$MB 可驻留 L2 供 32 核共享。谁的工作集能驻留 L2,谁优——这是时延模型比较,不是硬阈值(见 §7)。 + +### 5.4 情形四:$P = B \lceil M/16 \rceil \lceil N/16 \rceil < C$ → **StreamK** + +- R1:**无法满足**。B、M、N 三个便宜维度全部用尽(切到 16 粒度)仍凑不满 32 核。典型:$B = 1, M = N = 64$ 时 $P = 16 < 32$,一半核闲置。 +- 此时**唯一剩余的并行维度是 K**。切 K 虽然代价高(打破 I1/I2,付 $T_{REDUCE}$),但不切的代价是核闲置——两害相权,当 + +$$ +T_{MMAD/\text{core}} \gg T_{REDUCE} \;\Longleftrightarrow\; K \;\gtrsim\; \text{grid}_K^2 \times 1690 +$$ + +(推导:要求计算时延 ≥ 10× 归约时延,代入 8192 FLOP/拍、1.65GHz、GM 1.6TB/s)时,切 K 净收益为正。K 越大,可承担的 grid_K 越大;K 不够大时,减少 grid_K(配合切 B/M/N,即 grid_K × grid_B × grid_M × grid_N ≤ 32 的组合)来降低归约组大小。 + +矛盾是"便宜维度用尽仍缺并行"。最优解是**付归约代价买 K 维并行**,且归约组能小则小。这就是 StreamK——分块公式中 $\tilde K$ 跨核拆分、部分和经 workspace 归约的特化。 + +**判据**:$P < C$(或 P 虽够但核内 M/N 范围被压得太碎)且存在 grid_K 使 $K/\text{grid}_K \ge 256$ 且满足上式。 + +### 5.5 一条链总结 + +``` +R1 缺并行?──────────────────────────────────────────── + │ 不缺 │ 缺 + ▼ ▼ +R3 缺 Cube 效率? 便宜维度(B/M/N)用尽? + │ 不缺 │ 缺 │ 用尽 + ▼ ▼ ▼ +B 够切? MergeBatch StreamK(切K,付归约代价) + │B≥C │B 与源码对照:arch35 源码有 10 个策略,为什么这里只有 4 个?因为源码策略 = 4 个切分等价类 × 两个正交维度的笛卡尔积:① **计算通路**(Cube vs AIV:K=0 的 K_EQUAL_ZERO、K=1 的 TO_MUL 是 K 退化时的通路切换,发生在切分决策之前,属于预处理层);② **驻留策略**(AL1/BL1_FULL_LOAD 是 ASW_Basic 内部 $M^t = M$ / $N^t = N$ 的 tiling 极限,ITER_BATCH_BROADCAST 是 IterBatch 在广播输入下的数据复用特化)。它们是分支**内部**的参数特化,不构成新的切分等价类。源码尾部还有 BASE=999 无条件兜底,也印证"特判 ⊂ 通用"的层级结构。 + +### 6.2 极小性:每个分支都有它是唯一最优的 shape 区域 + +去掉任何一个分支,都存在 case 失去最优实现: + +| 分支 | 独占最优的示例 case(BF16) | 替代方案的劣势 | +|---|---|---| +| IterBatch | B=32, M=N=K=4096 | ASW 切 M/N 引入无谓共享读取;MergeBatch 引入无谓浪费 | +| MergeBatch | B=64, M=N=64, K=256 | IterBatch 的 L0C 利用率仅 6.25%,Cube 闲置 ~94% | +| ASW_Basic | B=2, M=N=8192, K=1024 | 切 B 仅 2 核活跃;StreamK 付无谓归约 | +| StreamK | B=1, M=N=64, K=65536 | 不切 K 只用 16/32 核(P=16),时延差数量级(文档 Case 3:19.5ms → 0.61μs/核量级) | + +⇒ 四者**互为不可替代,构成极小完备集**。 + +--- + +## 7. 边界是软的:竞争区域由时延模型仲裁 + +上面的判据给出的是分支的"主场",但相邻分支的边界不是硬切换。重叠区域(如 B ≥ 32 且 M×N 中等时 IterBatch 与 ASW 切 B 等效;M×N 在 64²~128² 之间时 MergeBatch 与 IterBatch 互有胜负)必须由统一点评仲裁: + +$$ +\text{branch}^* = \arg\min_{\text{cand} \in \bigcup \text{4 分支候选}} \max\big(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{FIXPIPE}, T_{REDUCE}\big) +$$ + +这就是《方案设计》§6.7 决策算法"全分支生成候选 → 时延模型评估 → 取 min"的合理性来源:**分支体系保证候选集完备且无冗余(每个等价类只派一个代表框架),时延模型在等价类内部和边界上做精细仲裁。** 两层结构缺一不可——只有时延模型没有分支体系,搜索空间是 15 种组合 × 全部 grid,爆炸且重复;只有分支体系没有时延模型,边界 case 会被错误硬切。 + +--- + +## 8. 闭环:与分块计算公式的参数特化对应 + +最后把推导结果写回分块公式的语言,验证链条闭合(对应《BMM分块计算数学公式》第六章,但现在每行都有了"理由"): + +| 分支 | 分块参数特化 | 推导出处 | +|---|---|---| +| IterBatch | $B^t = 1$;$\tilde B \ge C$;batch 为核间最外层;核内标准 K 累加 | §5.1:零代价并行够 + Cube 饱 | +| MergeBatch | $B^t = b > 1$ 折叠进 $M^t_{eff} = bM^t, N^t_{eff} = bN^t$;L0C 二次方程求最优 $b$;BlockTrace 取对角 | §5.2:Cube 饿 → 折叠 batch 换 $U \uparrow$,浪费被访存 Bound 掩盖 | +| ASW_Basic | 通用 $(\beta, \mu, \nu)$ 展开;$\sigma$ = ASW 滑窗蛇形($W \le \lfloor\sqrt C\rfloor$ 的最大因子) | §5.3:低代价维度补并行,$\sigma$ 压缩 L2 足迹 | +| StreamK | $\tilde K$ 跨核拆分,$C_{(\beta,\mu,\nu)} = \sum_c C^{(c)}$ 部分和归约;grid_K×grid_B×grid_M×grid_N ≤ C | §5.4:便宜维度用尽,付归约代价买 K 维并行 | + +--- + +## 9. 总结:一句话逻辑链 + +> **分块计算给出自由度(4 维怎么切),硬件给出规则(L0C 累加使切 K 独贵、batch 独立使切 B 免费、L2 使切 M/N 廉价),最优目标给出约束(核要满、Cube 要饱、搬运要省)。按"最便宜的维度优先购买并行度、Cube 不够大就折叠 batch、便宜维度用尽才切 K"的原则展开,15 种切分组合在 (含K, 含M/N) 两个代价特征下恰好坍缩为 4 个等价类——这就是 MergeBatch / IterBatch / ASW_Basic / StreamK 四大分支,完备且极小;边界区域由端到端时延模型统一仲裁。** + +--- + +*文档版本:v1.0 | 上游:《BMM分块计算数学公式》v1.1+ | 下游:《BMM最优软件实现方案设计》v4.0 第六章* diff --git a/BatchMatmul算子特性分析.html b/BatchMatmul算子特性分析.html new file mode 100644 index 0000000..88fe6d2 --- /dev/null +++ b/BatchMatmul算子特性分析.html @@ -0,0 +1,754 @@ + + + + + +Batch Matmul 算子特性分析 + + + + + +
+

Batch Matmul 算子特性分析

+
+

版本: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 章完备性审视(漏洞检查)。

+
+
+

1. 算子功能与接口

+

1.1 算子功能

+

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$ 展平),最后两维做矩阵乘。

+

1.2 接口参数

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数形状数据类型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(可随路量化)典型 NDBatchC = broadcast(BatchA, BatchB)
+

广播语义:batch 维兼容广播——每一级 batch 维上,两侧取值要么相等、要么为 1;输出该维取两者的最大值。注意两类广播的形态差异:

+ +

1.3 case 空间:决定实现方式的特征量

+

一个 BMM case 的实现方式完全由以下特征量决定:

+

$$ +(\; B,\; M,\; N,\; K,\; \text{dtype},\; \text{layout/转置},\; \text{广播形态} \;) +$$

+

其中 $B$ = BatchC(展平后的输出 batch 数)。后续全部分析就是回答:给定这组特征量,最优实现是什么。

+
+

2. 性能模型:什么叫"最优"

+

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

+

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 模型。

+

2.2 目标芯片关键规格与经验常数

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
规格项数值对设计的意义
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 Cache128MB,5.2 TB/s(读写各独享)重复读取的吸收层;Cube 输出写 L2 即算完成
L1 / L0A / L0B / L0C512KB / 64KB / 64KB / 256KB(每 AIC)核内分块的容量约束;L0C 按 FP32 累加计
UB512KB(每 AIV)向量通路/归约
+

GM→L1 搬移效率(ND2NZ)达到 90%+ 带宽利用率的四条经验约束(DV100 实测,重要性递减):

+
    +
  1. 参与核数:建议 32 核并行搬移,至少 3/4 核(24 核)并发——核数不足时带宽利用率上限被压低;
  2. +
  3. 单核搬移数据总量 ≥ 480KB(min_DatamountPerCore);
  4. +
  5. 单次搬移 tile 大小 ≥ 16KB(min_TileSize);
  6. +
  7. dValue(单次搬移的连续内轴字节数,如非转置 A 的 K 向、非转置 B 的 N 向)≥ 128B,建议 256B,最好 512B。
  8. +
+
+

这些常数(480KB、16KB、256B/512B,以及后文的 $b_{thr}=4$、$k_{thr}$=32B/dtype、$minCoreNum \approx 0.8 \times 32$)都是该档芯片的实测经验值;换芯片时本文的所有逻辑结构不变,只需替换常数表。

+
+

2.3 算存比与 Bound 判定

+

约定:以"元素"为访存单位(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 带宽是供货卡车(带宽)。算存比就是"每块砖要盖多少平米"。每块砖盖的平米数低于起重机的额定配比时,楼永远盖不快,问题出在供货而不在起重机——这时候优化要围绕"怎么少堵车、怎么一次多运点",而不是换更大的起重机。

+
+

3. 分支划分的系统推导

+

本章是全文逻辑主线:从分块计算的本质出发,不靠经验罗列,推导出五大分支的划分。

+

3.1 分块计算的本质与四个可切维度

+

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)。

+

3.2 四个维度的切分特征:一张决定性的表

+

从"读入 / 计算 / 写出"三个视角考察每个维度的核间切分特征(这是后续一切推导的基石):

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
切分维度读入特征计算特征写出特征
切 B每个数据块只被固定的 1 个核读取,核间零重复读每个输出块由 1 个核独立完成,无核间依赖只写最终结果,无中间结果
切 M / 切 N切 M 则同一右矩阵块被多核重复读;切 N 则同一左矩阵块被多核重复读每个输出块由 1 个核独立完成,无核间依赖只写最终结果,无中间结果
切 K每个数据块只被固定的 1 个核读取,零重复读每个输出块由多核共同完成,存在核间依赖有中间结果写出,需核间 Reduce 归约
+

为什么特征差异这么大?两条硬件层面的根本原因:

+ +

3.3 完备枚举:15 种核间切分组合

+

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 时延不能成为新瓶颈
+

3.4 代价不对称性:切分维度的"价格表"

+

15 种组合的代价结构只由两个布尔特征决定——是否含 K、是否含 M/N。三个维度的"核间切分价格"严格排序:

+

$$ +\text{cost}(\text{切}B) = 0 \;<\; \text{cost}(\text{切}M/N) \;\ll\; \text{cost}(\text{切}K) +$$

+ +

"价格表"不是经验,是 BMM 语义 + L0C 累加机制 + L2/GM 带宽结构三条事实的推论。 整条分支决策树就是一句话:按价格从低到高购买并行度,买不够才加价。

+

3.5 推导主链:从价格表到五大分支

+

第 0 层:问题归约(能降维就不在 BMM 本体里解决)。

+ +

第 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 买并行,付归约代价)◀──────────────────┘
+
+

3.6 完备性与极小性:五大分支不多不少

+

完备性。对归约不掉的 case,15 种组合按"(是否含 K,是否含 M/N)"坍缩:

+ +

$1 + 1 + 2 = 4$,加上第 0 层的归约分支(转Matmul)= 5。任何合法 case 必落其一,无空洞。

+

极小性。每个分支都有它"唯一最优"的 shape 区域,去掉任何一个都会有 case 失去最优实现:

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
分支独占最优的代表 case(BF16)替代方案为何更差
转MatmulBatchB=1, B=128, M=N=K=2048免费折叠后 Matmul 体系的 L2 切分/全载直接可用;BMM 分支内重做一遍无收益
IterBatchB=32, M=N=K=4096ASW 切 M/N 引入无谓共享读;MergeBatch 引入无谓冗余算力
MergeBatchB=128, M=32, N=128, K=64IterBatch 的 L0C 利用率仅 ~1.6%,Cube 空转,搬移 tile 碎(<16KB)带宽利用率崩
ASW_BasicB=2, M=N=8192, K=1024切 B 仅 2 核干活;StreamK 付无谓归约
StreamKB=1, M=N=64, K=65536不切 K 时 P=16 < 32,近半核闲置,时延差数量级
+

⇒ 五大分支构成极小完备集

+

3.7 大白话总结

+
+

把 BMM 想成给 32 个工人分一批"矩阵乘订单":

+
    +
  • 能整单外包的就别自己干:一侧只有一个 batch,问题其实就是一次普通矩阵乘,直接交给成熟的 Matmul 产线(转Matmul);
  • +
  • 按订单分(切 B)最省心:每人几单,互不干扰(零重复搬料、零协调)。订单多时,单大就一单接一单干(IterBatch);单太小时一人同时干几单、拼成大活干,虽然多做了点无用功,反正瓶颈在等料不在干活(MergeBatch);
  • +
  • 订单不够分,就按行/列拆单(切 M/N):大家会重复领同一份料,但料放在近处仓库(L2),多跑几趟近仓库很便宜(ASW_Basic);
  • +
  • 行/列也不够拆,只好沿深度 K 拆:几个人合做同一块输出,各算一段再汇总——汇总是要额外开会的(Reduce),所以只在实在分不满、且 K 足够长时才这么干(StreamK)。
  • +
+
+
+

4. 各分支详解:进入条件与实现方案

+
+

每节按统一结构展开:做什么(定义)→ 什么时候进(进入条件逐条 + 每条的理由)→ 怎么做(实现方案与参数求解)→ 与源码对照(批判性)

+
+

4.0 前置通路层:K 退化 case

+

在切分决策之前先做计算通路判断:

+ +

这一层与"切分维度"正交:五大分支管理的都是 Cube 通路上的 case(K ≥ 2)。K 极小(如 K < 16,不足一个 fractal)时同理应优先考虑向量通路或尾部补齐。

+

4.1 转Matmul(问题归约分支)

+

做什么:当 BatchA = 1BatchB = 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 的特例。

+

4.2 MergeBatch(多 batch 合并计算)

+

4.2.1 做什么

+

核间按 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 利用率

+

4.2.2 进入条件(逐条 + 理由)

+

设计原则四条:① 硬件时延可流水掩盖(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 无缘。

+

4.2.3 核内 b 与 baseK 的取值优化

+

进入分支后,$b$ 与 baseK 的具体取值有一个重要性质:在访存 Bound 前提下,$b < b_{AI}$ 时 b 的取值不影响性能,baseK 只需满足分形约束(16 倍数)也不影响算存比——因为瓶颈是搬移,算力余量内怎么切都一样。据此:

+
    +
  1. 先由条件 2 求 $b_{max}$,实际 $b$ 在 $[b_{thr}, b_{max}]$ 内取,且尽量均匀($b_{core}$ 整除 $b$)——每次计算的 $b$ 均匀一致对功耗更有利;
  2. +
  3. 再由 $b$ 反查 L0AB 允许的最大 baseK:$baseK_{max} = \min\left(\dfrac{64\text{KB}}{2 \cdot bM \cdot \text{dtype}},; \dfrac{64\text{KB}}{2 \cdot bN \cdot \text{dtype}},; \dfrac{C0_{size}}{\text{dtype}}\right)$,向下取 16 倍数。
  4. +
+

4.2.4 数值例子(把条件过一遍)

+

case:B=128, M=32, K=64, N=128, BF16,32 核。

+
    +
  1. 资格:$b_{core} = 128/32 = 4 \ge b_{thr} = 4$ ✓;
  2. +
  3. 算存比:$AI = 2 \times 32 \times 128 / (32 + 128) = 51.2$ FLOP/元素 $< R = 607.5$ ✓ 访存 Bound;$b_{AI} = 607.5 / 51.2 \approx 11.86$,即 $b \le 11$ 都不会进入算力 Bound——算存比约束很宽;
  4. +
  5. L0C 约束(乒乓):$2(b \cdot 32)(b \cdot 128) \times 4\text{B} \le 256\text{KB} \Rightarrow b^2 \le 8 \Rightarrow b \le 2$(不乒乓则 $b \le 4$)——L0C 才是真正的紧约束
  6. +
  7. 取 b = 2(= $b_{core}$ 的因子,均匀):L0B 给出 $k_{L0} \le 64\text{KB} / (2 \times 2 \times 128 \times 2\text{B}) = 64$,$K = 64$ 一步到位,baseK = 64;
  8. +
  9. 搬移效率:$k_{L1} = 64 \ge k_{thr} = 16$ ✓;单核搬移量 $4 \times (32 \times 64 + 64 \times 128) \times 2\text{B} = 80\text{KB}$ —— 不足 480KB,条件 6 不满足!
  10. +
+

结论:该 case 在"单核搬移总量"一项上不达标——这正是 MergeBatch 条件体系的用处:它提前告诉你瓶颈不在算力而在搬移效率,此 case 应尝试合并更多 K/提高 $b_{core}$ 数据量,或由时延模型与 IterBatch/ASW_Basic 比较后定夺。

+

4.2.5 执行流程

+
核间: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
+
+

4.2.6 与源码对照(批判性)

+

源码 batch_matmul_v3_mergebatch_basicapi_tiling.cpp 的 IsCapable 条件:各级 batchA_i == batchB_i、batchC ≥ 4 × aicNumalignK ≥ 64M ≤ N、无 bias、非 NZ、拒绝非连续转置等。对照分析:

+ +

4.3 IterBatch(逐 batch 计算)

+

4.3.1 做什么

+

核间按 B 分核(每核 1 个或多个 batch),核间无同步无通信;核内逐个 batch 分别执行标准 Matmul 分块(L1→L0A/L0B→Cube→L0C→Fixpipe)并输出。无算力浪费、无跨 batch 依赖——是"切 B"最朴素的形态。

+

4.3.2 进入条件(三大类 + 理由)

+

满足以下任一

+

类 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$ 且同时满足:

+
    +
  1. $BatchA = BatchB$ 且 $b_{core} = B / aicNum \ge 1$;
  2. +
  3. 负载均衡:访存 Bound 时核负载利用率建议 bAvg/bMax > 0.8(计算 Bound 时建议 100%)。操作化表述:B 整除核数,或尾波(B mod aicNum)活跃的核数 ≥ minCoreNum(DV100 取 0.8 × 32 ≈ 26)。注意:若只写 B mod aicNum > minCoreNum 会误杀整除 case(余数 0 恰恰是完全均衡),判据应以均衡比为准、余数规则为其近似
  4. +
  5. 单核搬移不重复读(访存 Bound 下重复读就是纯损失),按 L1 容量分五种形态之一: +
      +
    • (a) $b_{core} = 1$ 且 $(MK + KN) \cdot \text{dtype} \le L1_{size}$:单 batch 左右矩阵同时驻留 L1,零重复读;
    • +
    • (b) $b_{core} > 1$ 且 $2(MK + KN) \cdot \text{dtype} \le L1_{size}$:L1 放下 2 个 batch 形成乒乓流水;
    • +
    • (c) 放不下的,$(MK + KN/Step) \cdot \text{dtype} \le L1_{size}$(或 M/Step 对称):一矩阵不切、另一切分,分块大小仍须 > min_TileSize;
    • +
    • (d) $b_{core} > 1$ 时上一条的半容量版本(L1 双 batch 乒乓预算减半);
    • +
    • (e) 左右都切 K:$(M \cdot K/Step + K/Step \cdot N) \cdot \text{dtype} \le L1_{size}$;
    • +
    • 以上 (c)(d)(e) 切分后的分块均须满足搬移效率(tile ≥ 16KB、dValue ≥ 128B/256B)。
    • +
    +
  6. +
+

4.3.3 设计原理(为什么是这些条件)

+
    +
  1. 核间切 B 零共享零依赖,唯一的系统性风险是负载不均——所以均衡是第一条件;
  2. +
  3. 输出须满足 L0C 容量:$MN \cdot 4\text{B} \le L0C$(否则核内还要切 M/N,那就不是纯 IterBatch 而是 ASW 行为——但注意 IterBatch 核内允许对单 batch 做标准 M/N/K tiling,此处的准确含义是"核内 tiling 不构成跨 batch 的耦合");
  4. +
  5. 访存 Bound 时单核数据不得重复读——L1 装得下才不重复,装不下就按 Step 切分且切分后仍满足搬移效率下限;
  6. +
  7. 与 MergeBatch 的分工:IterBatch 不浪费算力,但需要"单 batch 足够大"撑搬移效率与 Cube 利用率;MergeBatch 用浪费换效率,专治小 M×N。两者在 $M \times N$ 的中段重叠,由时延模型仲裁(第 5 章)。
  8. +
+

4.3.4 执行流程

+
核间: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 回写可异步)
+
+

4.4 StreamK(K 维核间切分)

+

4.4.1 做什么

+

当 B、M、N 三个维度切到最碎仍填不满 32 核时,把 K 维切到核间:多核各算一段 K 的部分和,再归约:

+

$$ +C_{(\beta,\mu,\nu)} = \sum_{c \in group} C^{(c)}_{(\beta,\mu,\nu)} \quad \text{(部分和写 workspace,AIV 归约或原子加)} +$$

+

4.4.2 进入条件(两条缺一不可)

+

条件 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 搜索自然淘汰归约过重的配置。

+

4.4.3 实现要点

+ +

4.4.4 与不切 K 分支的关系

+

StreamK 从不"硬切换"进入:当不切 K 的分支候选已足够快(如已 Cube Bound),StreamK 候选的归约开销使其自然落败;只有当 B/M/N 并行度不足导致大量核闲置时,StreamK 才以数量级优势胜出(例:B=1、M=N=64、K=65536 时,不切 K 仅 16 核可用,StreamK 32 核满负荷)。

+

4.5 ASW_Basic(通用切分框架)

+

4.5.1 做什么

+

不切 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 是实践中最常命中的分支。

+

4.5.2 核心机制:swizzle + L2 管理

+

切 M/N 的固有代价是共享矩阵的重复读,ASW_Basic 用两件武器把代价压到最低:

+
    +
  1. ASW 滑窗蛇形 swizzle:把 M 向按窗口 $W$ 分组($W = \max{d : d \mid C,\ d \le \lfloor\sqrt{C}\rfloor}$,32 核取 W=4),窗口内蛇形遍历 N 向。数学效果:同一时刻 32 个核活跃的工作集被压缩到"W 个 A 行块 + 一条 B 列块带",L2 足迹最小 ⇒ 共享读取基本命中 5.2TB/s 的 L2 而非 1.6TB/s 的 GM。窗口取 $\lfloor\sqrt C\rfloor$ 的最大因子,是因为窗口越接近正方形,A 行块 + B 列块的 L2 足迹越小,且因子性保证整窗被核数均分、窗口边界不碎;
  2. +
  3. L2 切分:工作集超过 128MB 时,按 mL2TileNum × nL2TileNum 切分,每个 L2 块错位分核(对角线分配),避免多核同时抢同一地址的读读冲突,并优先选拖尾小的方案。
  4. +
+

4.5.3 核间切分维度的选择顺序

+

按共享代价从低到高:

+
    +
  1. 切 B(B ≥ 核数):零共享,永远先试;
  2. +
  3. 切 M(B 不够):右矩阵 [K, N] 被共享,若 $KN \cdot \text{dtype} \le 128$MB 则驻留 L2,零 GM 重复读;
  4. +
  5. 切 N:对称;
  6. +
  7. 混合切(B×M、M×N、…):双向共享,靠 swizzle + L2 切分管理;
  8. +
  9. 降核:P 远小于 32 且 K 也不够格走 StreamK 时,宁可部分核闲置(小 case 时延绝对值小,调度开销反而主导)。
  10. +
+

4.5.4 内部特化(不是独立分支)

+ +
+

5. 分支仲裁与整体决策流程

+

5.1 为什么需要仲裁层

+

第 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 被硬阈值误杀。

+

5.2 总决策流程

+
输入: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 切分方案 + 预估端到端时延
+
+
+

6. 五大分支设计的完备性审视(漏洞检查)

+

按"转Matmul / MergeBatch / IterBatch / StreamK / ASW_Basic 五大分支 cover 全部 BMM case 最优实现"的设计目标,逐条拷问。

+

6.1 覆盖矩阵:无空洞

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
shape 区域主分支兜底
K = 0 / K = 1前置通路层(AIV)
BatchA=1 或 BatchB=1转MatmulASW_Basic(广播友好形态)
BatchA=BatchB>1,B≥32,M×N 大IterBatchASW_Basic 切 B
BatchA=BatchB>1,B≥32,M×N 小,访存 BoundMergeBatchIterBatch
上述但 B<32,P≥32ASW_Basic降核 ASW
P<32,K 大StreamK降核 ASW
P<32,K 也小降核 ASW_Basic(此时时延绝对值小,调度开销主导,分支选择不敏感)
交叉广播(BatchA≠BatchB 且均>1)ASW_Basic
+

6.2 审视发现的问题与修正

+

设计骨架成立,但有五处需要修正或显式声明,否则会被反例拷问:

+

问题 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×aicNumalignK ≥ 64M ≤ N 等直接写死,其中 alignK ≥ 64 比草稿阈值保守 4 倍、M ≤ N 放弃了镜像 case——做最优实现时不应继承这些硬编码,而应按第 4 章的条件体系重新判定

+

6.3 与源码 10 策略的关系

+

源码 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_KStreamK
MERGE_BATCHMergeBatch
ITER_BATCH、ITER_BATCH_BROADCASTIterBatch(broadcast 变体是广播输入下的数据复用特化;单边广播在本文更多由转Matmul 吸收)
AL1_FULL_LOAD、BL1_FULL_LOADASW_Basic 的 L1 驻留参数极限(4.5.4)
ASW_BASIC、BASEASW_Basic
+

结论:源码策略集 = 五分支 × 计算通路 × 驻留/广播特化的展开。五分支是"切分等价类"的最小完备骨架,源码的冗余策略可在新实现中收敛,源码遗漏的(条件 3/4/6 的显式判定、M>N 的 MergeBatch 镜像、转Matmul 的 BatchA=1 方向)正是优化的增量空间。

+
+

7. 总结

+
    +
  1. 最优的定义:$T_{total} = \max(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{Fixpipe}[, T_{Reduce}])$ 最小;瓶颈交换是合法且必要的手段。
  2. +
  3. 分支的推导:4 维可切 → 15 种组合完备 → 切分特征表(切 B 免费 / 切 M/N 廉价被 L2 吸收 / 切 K 昂贵需归约)→ 按价格从低到高购买并行度 → 坍缩为 MergeBatch、IterBatch、ASW_Basic、StreamK 四个等价类,加前置的转Matmul 问题归约层,共五大分支,完备且极小
  4. +
  5. 各分支的条件都不是孤立经验,而是"资格(并行度够不够)→ 上限(容量/算存比封顶)→ 下限(搬移效率托底)"三层逻辑的实例化;所有经验常数可标定、可移植。
  6. +
  7. 边界不靠硬切:重叠区由端到端时延模型统一仲裁,分支体系负责候选集的完备无冗余。
  8. +
  9. 设计经受了漏洞拷问:修正了转Matmul 的方向不对称与 IterBatch 均衡条件的字面误杀,显式声明了交叉广播与 K 退化的归属;相对源码实现,本文条件体系更细、更真、覆盖更全。
  10. +
+
+

参考文档:《BMM分块计算数学公式》《BMM最优软件实现方案设计》《BatchMatMulV3算子分支实现分析》《BMM从分块计算到四大分支的逻辑推导》;源码:ops-nn/matmul/batch_mat_mul_v3(arch35);芯片资料:昇腾 950PR 架构白皮书与 CANN 9.0.0 性能建模文档。

+ +
+ + diff --git a/BatchMatmul算子特性分析.md b/BatchMatmul算子特性分析.md new file mode 100644 index 0000000..eb1ac46 --- /dev/null +++ b/BatchMatmul算子特性分析.md @@ -0,0 +1,618 @@ +# Batch Matmul 算子特性分析 + +> 版本: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 章完备性审视(漏洞检查)。 + +--- + +## 1. 算子功能与接口 + +### 1.1 算子功能 + +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$ 展平),最后两维做矩阵乘。 + +### 1.2 接口参数 + +| 参数 | 形状 | 数据类型 | 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); +- **交叉广播**:两侧均大于 1 但不同级为 1,例如 `BatchA = (6,1)`、`BatchB = (1,6)`,则 `BatchC = (6,6)` 共 36 个输出 batch。这类 case **不属于**单边广播,后面第 3、6 章会看到它只能落入通用分支。 + +### 1.3 case 空间:决定实现方式的特征量 + +一个 BMM case 的实现方式完全由以下特征量决定: + +$$ +(\; B,\; M,\; N,\; K,\; \text{dtype},\; \text{layout/转置},\; \text{广播形态} \;) +$$ + +其中 $B$ = BatchC(展平后的输出 batch 数)。后续全部分析就是回答:**给定这组特征量,最优实现是什么。** + +--- + +## 2. 性能模型:什么叫"最优" + +### 2.1 总时延 = 流水线最慢的一级 + +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 模型。 + +### 2.2 目标芯片关键规格与经验常数 + +| 规格项 | 数值 | 对设计的意义 | +|---|---|---| +| 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 实测,重要性递减): + +1. **参与核数**:建议 32 核并行搬移,至少 3/4 核(24 核)并发——核数不足时带宽利用率上限被压低; +2. **单核搬移数据总量** ≥ 480KB(`min_DatamountPerCore`); +3. **单次搬移 tile 大小** ≥ 16KB(`min_TileSize`); +4. **dValue**(单次搬移的连续内轴字节数,如非转置 A 的 K 向、非转置 B 的 N 向)≥ 128B,建议 256B,最好 512B。 + +> 这些常数(480KB、16KB、256B/512B,以及后文的 $b_{thr}=4$、$k_{thr}$=32B/dtype、$minCoreNum \approx 0.8 \times 32$)都是**该档芯片的实测经验值**;换芯片时本文的所有逻辑结构不变,只需替换常数表。 + +### 2.3 算存比与 Bound 判定 + +**约定**:以"元素"为访存单位(1 元素 = dtype 字节),算存比单位为 FLOP/元素。 + +- case 的固有算存比(不计输出写出): + +$$ +AI = \frac{2MNK}{(MK + KN)} = \frac{2MN}{M + N} \quad [\text{FLOP/元素}] +$$ + +- 计入输出写出(写 L2/GM)的完整算存比: + +$$ +AI_{full} = \frac{2MNK}{MK + KN + MN} \quad [\text{FLOP/元素}] +$$ + +- 芯片硬件平衡点(该 dtype 位宽下): + +$$ +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 带宽是供货卡车(带宽)。算存比就是"每块砖要盖多少平米"。每块砖盖的平米数低于起重机的额定配比时,楼永远盖不快,问题出在供货而不在起重机——这时候优化要围绕"怎么少堵车、怎么一次多运点",而不是换更大的起重机。 + +--- + +## 3. 分支划分的系统推导 + +本章是全文逻辑主线:**从分块计算的本质出发,不靠经验罗列,推导出五大分支的划分。** + +### 3.1 分块计算的本质与四个可切维度 + +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)。 + +### 3.2 四个维度的切分特征:一张决定性的表 + +从"读入 / 计算 / 写出"三个视角考察每个维度的核间切分特征(这是后续一切推导的基石): + +| 切分维度 | 读入特征 | 计算特征 | 写出特征 | +|---|---|---|---| +| **切 B** | 每个数据块只被固定的 1 个核读取,核间**零重复读** | 每个输出块由 1 个核独立完成,**无核间依赖** | 只写最终结果,无中间结果 | +| **切 M / 切 N** | 切 M 则同一右矩阵块被多核**重复读**;切 N 则同一左矩阵块被多核重复读 | 每个输出块由 1 个核独立完成,无核间依赖 | 只写最终结果,无中间结果 | +| **切 K** | 每个数据块只被固定的 1 个核读取,零重复读 | 每个输出块由**多核共同**完成,存在核间依赖 | **有中间结果写出**,需核间 Reduce 归约 | + +为什么特征差异这么大?两条硬件层面的根本原因: + +- **B 维在数学上独立**(BMM 语义就是逐 batch 独立矩阵乘),所以切 B 天然零重复、零依赖; +- **K 维有 L0C 累加机制**:核内切 K 时,$\tilde K$ 轮 mmad 在 256KB 的 L0C 上原地累加(`cmatrixInitVal=false`),中间结果不出核。一旦把 K 切到**核间**,单核的 L0C 装不下"别的核算的 K 段",部分和必须写出到 GM/L2 workspace,再由 AIV 归约——**切 K 是唯一同时破坏"累加不出核"和"输出独占"两条性质的切法**。 + +### 3.3 完备枚举:15 种核间切分组合 + +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 时延不能成为新瓶颈 | + +### 3.4 代价不对称性:切分维度的"价格表" + +15 种组合的**代价结构只由两个布尔特征决定**——是否含 K、是否含 M/N。三个维度的"核间切分价格"严格排序: + +$$ +\text{cost}(\text{切}B) = 0 \;<\; \text{cost}(\text{切}M/N) \;\ll\; \text{cost}(\text{切}K) +$$ + +- **切 B 免费**:零重复读、零依赖、零中间写出; +- **切 M/N 廉价但有价**:共享矩阵被重复读,但若共享部分能驻留 128MB L2,重复读以 5.2TB/s 命中 L2 而非 1.6TB/s 的 GM——代价大部分被 L2 吸收;配合 swizzle(执行顺序编排)压缩同时活跃的工作集,代价进一步压低; +- **切 K 昂贵**:归约流量 $T_{Reduce} \propto grid_K \times$ 输出量,且引入核间同步——这是结构性代价,L2 吸收不掉。 + +**"价格表"不是经验,是 BMM 语义 + L0C 累加机制 + L2/GM 带宽结构三条事实的推论。** 整条分支决策树就是一句话:**按价格从低到高购买并行度,买不够才加价。** + +### 3.5 推导主链:从价格表到五大分支 + +**第 0 层:问题归约(能降维就不在 BMM 本体里解决)。** + +- 若 `BatchA = 1` 或 `BatchB = 1`,BMM 可通过维度折叠**转化为普通 Matmul**,直接复用 Matmul 的成熟优化体系(tiling、L2 切分、全载),不必在 BMM 框架内重新发明轮子 → **分支一:转Matmul**; +- 退化 case(K=0 无计算、K=1 无累加深度)让 Cube 完全或几乎无用,改走 AIV 向量通路——这属于"计算通路选择",是正交于切分的前置判断(见 4.0)。 + +**第 1 层:归约不掉的 case(BatchA = BatchB = B > 1,K ≥ 2),必须在 BMM 框架内组织 4 维分块。先买免费的 B 维。** + +- $B \ge C = 32$ 时,切 B 就能填满核。此时分两种情形: + - **单 batch 的 $M \times N$ 够大**(能开出大 tile、L0C 利用率高、Cube 喂得饱)→ 每核逐个 batch 做完整 Matmul 就是最优,任何额外切分只增代价 → **分支二:IterBatch**; + - **单 batch 的 $M \times N$ 太小**(如 64×64,L0C 利用率仅 6.25%,Cube 大片闲置)→ 把多个 batch 在核内合并成大 tile 计算($[bM,K]@[K,bN]$,取块对角线输出),用冗余算力换 Cube 利用率与搬移效率。这类 case $AI$ 低、天然访存 Bound,浪费的算力被搬移时延掩盖(§2.1 瓶颈交换)→ **分支三:MergeBatch**。 + +**第 2 层:免费的 B 买不满 32 核,加价买廉价的 M/N。** + +- $B < 32$ 但 $P = B \times \lceil M/16 \rceil \times \lceil N/16 \rceil \ge 32$(M/N 平面的分块补得上并行缺口)→ 切 M/N(或混合切),共享矩阵的重复读交给 L2 + swizzle 吸收。这是"所有不切 K、含共享读取切法"的通用框架 → **分支四:ASW_Basic**。 + +**第 3 层:廉价维度也买不满,才买昂贵的 K。** + +- $P < 32$(B、M、N 都小,通常 K 大)→ 唯一剩余的并行维度是 K。付归约代价换并行度,当且仅当并行收益 > 归约代价(定量条件见 4.4)→ **分支五:StreamK**。 + +``` +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 买并行,付归约代价)◀──────────────────┘ +``` + +### 3.6 完备性与极小性:五大分支不多不少 + +**完备性**。对归约不掉的 case,15 种组合按"(是否含 K,是否含 M/N)"坍缩: + +- 含 K 的 8 种:代价结构相同(必付归约),grid_K × grid_B × grid_M × grid_N 只是参数差异 ⇒ 1 个分支(StreamK)统一覆盖; +- 不含 K 含 M/N 的 6 种:代价结构相同(共享重复读)⇒ 1 个分支(ASW_Basic)统一覆盖; +- 纯 {B} 的 1 种:核内组织只有两种本质不同的方式——**合并多 batch 一起算**(MergeBatch)或**逐个 batch 算**(IterBatch),不存在第三种 ⇒ 2 个分支。 + +$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,近半核闲置,时延差数量级 | + +⇒ 五大分支构成**极小完备集**。 + +### 3.7 大白话总结 + +> 把 BMM 想成给 32 个工人分一批"矩阵乘订单": +> - **能整单外包的就别自己干**:一侧只有一个 batch,问题其实就是一次普通矩阵乘,直接交给成熟的 Matmul 产线(转Matmul); +> - **按订单分(切 B)最省心**:每人几单,互不干扰(零重复搬料、零协调)。订单多时,单大就一单接一单干(IterBatch);单太小时一人同时干几单、拼成大活干,虽然多做了点无用功,反正瓶颈在等料不在干活(MergeBatch); +> - **订单不够分,就按行/列拆单(切 M/N)**:大家会重复领同一份料,但料放在近处仓库(L2),多跑几趟近仓库很便宜(ASW_Basic); +> - **行/列也不够拆,只好沿深度 K 拆**:几个人合做同一块输出,各算一段再汇总——汇总是要额外开会的(Reduce),所以只在实在分不满、且 K 足够长时才这么干(StreamK)。 + +--- + +## 4. 各分支详解:进入条件与实现方案 + +> 每节按统一结构展开:**做什么(定义)→ 什么时候进(进入条件逐条 + 每条的理由)→ 怎么做(实现方案与参数求解)→ 与源码对照(批判性)**。 + +### 4.0 前置通路层:K 退化 case + +在切分决策之前先做计算通路判断: + +- **K = 0**:无任何计算,C = bias 或 0,纯 AIV 写值; +- **K = 1**:退化为逐元素乘 `C = A ⊙ B`,无累加深度,Cube 的 16×16×16 粒度浪费 15/16,走 AIV 向量通路(GM→UB→Mul→GM)远优于 Cube 通路。 + +这一层与"切分维度"正交:**五大分支管理的都是 Cube 通路上的 case(K ≥ 2)**。K 极小(如 K < 16,不足一个 fractal)时同理应优先考虑向量通路或尾部补齐。 + +### 4.1 转Matmul(问题归约分支) + +**做什么**:当 `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 的特例。 + +### 4.2 MergeBatch(多 batch 合并计算) + +#### 4.2.1 做什么 + +核间按 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 利用率**。 + +#### 4.2.2 进入条件(逐条 + 理由) + +设计原则四条:① 硬件时延可流水掩盖(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) = \dfrac{2bMN}{M + N}$,保持访存 Bound 要求 + +$$ +AI(b) < R \;\Longleftrightarrow\; b < b_{AI} = \frac{R \cdot (M + N)}{2MN} +$$ + +- **L0C 容量约束**:输出 $[bM, bN]$ 的 FP32 累加块须放入 256KB L0C;考虑乒乓(double buffer)则两份: + +$$ +2 \cdot (bM)(bN) \cdot 4\text{B} \le 256\text{KB} \;\Rightarrow\; b \le b_{L0C} +$$ + +- **L0A/L0B 容量约束**:$2 \cdot bM \cdot k_{L0} \cdot \text{dtype} \le 64\text{KB}$ 且 $2 \cdot k_{L0} \cdot bN \cdot \text{dtype} \le 64\text{KB}$(乒乓两份),同时 $k_{L0}$ 不得低于 Cube 分形下限(16bit 位宽下典型 min baseK = 16)。注意 L0AB 约束耦合 $k_{L0}$,需联合求解。 + +> 乒乓取舍:若 $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 无缘。** + +#### 4.2.3 核内 b 与 baseK 的取值优化 + +进入分支后,$b$ 与 baseK 的具体取值有一个重要性质:**在访存 Bound 前提下,$b < b_{AI}$ 时 b 的取值不影响性能,baseK 只需满足分形约束(16 倍数)也不影响算存比**——因为瓶颈是搬移,算力余量内怎么切都一样。据此: + +1. 先由条件 2 求 $b_{max}$,实际 $b$ 在 $[b_{thr}, b_{max}]$ 内取,且**尽量均匀**($b_{core}$ 整除 $b$)——每次计算的 $b$ 均匀一致对功耗更有利; +2. 再由 $b$ 反查 L0AB 允许的最大 baseK:$baseK_{max} = \min\left(\dfrac{64\text{KB}}{2 \cdot bM \cdot \text{dtype}},\; \dfrac{64\text{KB}}{2 \cdot bN \cdot \text{dtype}},\; \dfrac{C0_{size}}{\text{dtype}}\right)$,向下取 16 倍数。 + +#### 4.2.4 数值例子(把条件过一遍) + +case:B=128, M=32, K=64, N=128, BF16,32 核。 + +1. **资格**:$b_{core} = 128/32 = 4 \ge b_{thr} = 4$ ✓; +2. **算存比**:$AI = 2 \times 32 \times 128 / (32 + 128) = 51.2$ FLOP/元素 $< R = 607.5$ ✓ 访存 Bound;$b_{AI} = 607.5 / 51.2 \approx 11.86$,即 $b \le 11$ 都不会进入算力 Bound——算存比约束很宽; +3. **L0C 约束**(乒乓):$2(b \cdot 32)(b \cdot 128) \times 4\text{B} \le 256\text{KB} \Rightarrow b^2 \le 8 \Rightarrow b \le 2$(不乒乓则 $b \le 4$)——**L0C 才是真正的紧约束**; +4. **取 b = 2**(= $b_{core}$ 的因子,均匀):L0B 给出 $k_{L0} \le 64\text{KB} / (2 \times 2 \times 128 \times 2\text{B}) = 64$,$K = 64$ 一步到位,baseK = 64; +5. **搬移效率**:$k_{L1} = 64 \ge k_{thr} = 16$ ✓;单核搬移量 $4 \times (32 \times 64 + 64 \times 128) \times 2\text{B} = 80\text{KB}$ —— **不足 480KB**,条件 6 不满足! + +结论:该 case 在"单核搬移总量"一项上不达标——这正是 MergeBatch 条件体系的用处:**它提前告诉你瓶颈不在算力而在搬移效率**,此 case 应尝试合并更多 K/提高 $b_{core}$ 数据量,或由时延模型与 IterBatch/ASW_Basic 比较后定夺。 + +#### 4.2.5 执行流程 + +``` +核间: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 +``` + +#### 4.2.6 与源码对照(批判性) + +源码 `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 在原理上同样可合并(交换合并方向),源码直接放弃,属于覆盖缺口; +- 源码**没有显式的算存比/搬移效率判定**(条件 3/4/6),把"是否最优"的责任推给了优先级顺序——本文的条件体系正是对此的补全。 + +### 4.3 IterBatch(逐 batch 计算) + +#### 4.3.1 做什么 + +核间按 B 分核(每核 1 个或多个 batch),核间无同步无通信;核内**逐个 batch 分别执行标准 Matmul 分块**(L1→L0A/L0B→Cube→L0C→Fixpipe)并输出。无算力浪费、无跨 batch 依赖——是"切 B"最朴素的形态。 + +#### 4.3.2 进入条件(三大类 + 理由) + +满足以下**任一**: + +**类 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$ 且同时满足: + +1. $BatchA = BatchB$ 且 $b_{core} = B / aicNum \ge 1$; +2. **负载均衡**:访存 Bound 时核负载利用率建议 bAvg/bMax > 0.8(计算 Bound 时建议 100%)。操作化表述:B 整除核数,或尾波(B mod aicNum)活跃的核数 ≥ minCoreNum(DV100 取 0.8 × 32 ≈ 26)。注意:若只写 `B mod aicNum > minCoreNum` 会误杀整除 case(余数 0 恰恰是完全均衡),**判据应以均衡比为准、余数规则为其近似**; +3. **单核搬移不重复读**(访存 Bound 下重复读就是纯损失),按 L1 容量分五种形态之一: + - (a) $b_{core} = 1$ 且 $(MK + KN) \cdot \text{dtype} \le L1_{size}$:单 batch 左右矩阵同时驻留 L1,零重复读; + - (b) $b_{core} > 1$ 且 $2(MK + KN) \cdot \text{dtype} \le L1_{size}$:L1 放下 2 个 batch 形成乒乓流水; + - (c) 放不下的,$(MK + KN/Step) \cdot \text{dtype} \le L1_{size}$(或 M/Step 对称):一矩阵不切、另一切分,分块大小仍须 > min_TileSize; + - (d) $b_{core} > 1$ 时上一条的半容量版本(L1 双 batch 乒乓预算减半); + - (e) 左右都切 K:$(M \cdot K/Step + K/Step \cdot N) \cdot \text{dtype} \le L1_{size}$; + - 以上 (c)(d)(e) 切分后的分块均须满足搬移效率(tile ≥ 16KB、dValue ≥ 128B/256B)。 + +#### 4.3.3 设计原理(为什么是这些条件) + +1. 核间切 B 零共享零依赖,**唯一的系统性风险是负载不均**——所以均衡是第一条件; +2. 输出须满足 L0C 容量:$MN \cdot 4\text{B} \le L0C$(否则核内还要切 M/N,那就不是纯 IterBatch 而是 ASW 行为——但注意 IterBatch 核内**允许**对单 batch 做标准 M/N/K tiling,此处的准确含义是"核内 tiling 不构成跨 batch 的耦合"); +3. 访存 Bound 时**单核数据不得重复读**——L1 装得下才不重复,装不下就按 Step 切分且切分后仍满足搬移效率下限; +4. 与 MergeBatch 的分工:IterBatch 不浪费算力,但需要"单 batch 足够大"撑搬移效率与 Cube 利用率;MergeBatch 用浪费换效率,专治小 M×N。两者在 $M \times N$ 的中段重叠,由时延模型仲裁(第 5 章)。 + +#### 4.3.4 执行流程 + +``` +核间: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 回写可异步) +``` + +### 4.4 StreamK(K 维核间切分) + +#### 4.4.1 做什么 + +当 B、M、N 三个维度切到最碎仍填不满 32 核时,把 K 维切到核间:多核各算一段 K 的部分和,再归约: + +$$ +C_{(\beta,\mu,\nu)} = \sum_{c \in group} C^{(c)}_{(\beta,\mu,\nu)} \quad \text{(部分和写 workspace,AIV 归约或原子加)} +$$ + +#### 4.4.2 进入条件(两条缺一不可) + +**条件 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 搜索自然淘汰归约过重的配置。** + +#### 4.4.3 实现要点 + +- 部分和写出:L0C → workspace(GM 或 L2),归约由 AIV 执行(AIC:AIV = 1:2,2 个 AIV 伺候 1 个 AIC 的部分和流); +- 原子加(AtomicAdd)变体省一遍读回,但破坏确定性——源码在确定性等级 > 1 时禁用 StreamK,工程上必须保留这个开关; +- 核间同步与任务编排建议直接复用 Matmul 高阶 API 的 StreamK 模板(避免手写 CrossCore flagId 冲突)。 + +#### 4.4.4 与不切 K 分支的关系 + +StreamK 从不"硬切换"进入:当不切 K 的分支候选已足够快(如已 Cube Bound),StreamK 候选的归约开销使其自然落败;只有当 B/M/N 并行度不足导致大量核闲置时,StreamK 才以数量级优势胜出(例:B=1、M=N=64、K=65536 时,不切 K 仅 16 核可用,StreamK 32 核满负荷)。 + +### 4.5 ASW_Basic(通用切分框架) + +#### 4.5.1 做什么 + +不切 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 是实践中最常命中的分支。** + +#### 4.5.2 核心机制:swizzle + L2 管理 + +切 M/N 的固有代价是共享矩阵的重复读,ASW_Basic 用两件武器把代价压到最低: + +1. **ASW 滑窗蛇形 swizzle**:把 M 向按窗口 $W$ 分组($W = \max\{d : d \mid C,\ d \le \lfloor\sqrt{C}\rfloor\}$,32 核取 W=4),窗口内蛇形遍历 N 向。数学效果:同一时刻 32 个核活跃的工作集被压缩到"W 个 A 行块 + 一条 B 列块带",L2 足迹最小 ⇒ 共享读取基本命中 5.2TB/s 的 L2 而非 1.6TB/s 的 GM。窗口取 $\lfloor\sqrt C\rfloor$ 的最大因子,是因为窗口越接近正方形,A 行块 + B 列块的 L2 足迹越小,且因子性保证整窗被核数均分、窗口边界不碎; +2. **L2 切分**:工作集超过 128MB 时,按 mL2TileNum × nL2TileNum 切分,每个 L2 块**错位分核**(对角线分配),避免多核同时抢同一地址的读读冲突,并优先选拖尾小的方案。 + +#### 4.5.3 核间切分维度的选择顺序 + +按共享代价从低到高: + +1. **切 B**(B ≥ 核数):零共享,永远先试; +2. **切 M**(B 不够):右矩阵 [K, N] 被共享,若 $KN \cdot \text{dtype} \le 128$MB 则驻留 L2,零 GM 重复读; +3. **切 N**:对称; +4. **混合切**(B×M、M×N、…):双向共享,靠 swizzle + L2 切分管理; +5. **降核**:P 远小于 32 且 K 也不够格走 StreamK 时,宁可部分核闲置(小 case 时延绝对值小,调度开销反而主导)。 + +#### 4.5.4 内部特化(不是独立分支) + +- **AL1/BL1 全载**:当单边无 batch 且该侧矩阵很小(如 batchA=1 且 M ≤ 256),把整个 A 常驻 L1,$\tilde M = 1$、A 的 GM→L1 搬运只发生一次——这是 ASW 内部"M 不切 + L1 驻留深度拉满"的参数极限,不是新分支; +- **广播友好形态**:交叉广播 case(§1.2)在 ASW 内通过对广播侧做 L1/L2 驻留处理,共享关系与切 M/N 同构。 + +--- + +## 5. 分支仲裁与整体决策流程 + +### 5.1 为什么需要仲裁层 + +第 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 被硬阈值误杀。 + +### 5.2 总决策流程 + +``` +输入: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 切分方案 + 预估端到端时延 +``` + +--- + +## 6. 五大分支设计的完备性审视(漏洞检查) + +按"转Matmul / MergeBatch / IterBatch / StreamK / ASW_Basic 五大分支 cover 全部 BMM case 最优实现"的设计目标,逐条拷问。 + +### 6.1 覆盖矩阵:无空洞 + +| 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 | — | + +### 6.2 审视发现的问题与修正 + +设计骨架成立,但有五处需要修正或显式声明,否则会被反例拷问: + +**问题 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 章的条件体系重新判定**。 + +### 6.3 与源码 10 策略的关系 + +源码 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 方向)正是优化的增量空间。 + +--- + +## 7. 总结 + +1. **最优的定义**:$T_{total} = \max(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{Fixpipe}[, T_{Reduce}])$ 最小;瓶颈交换是合法且必要的手段。 +2. **分支的推导**:4 维可切 → 15 种组合完备 → 切分特征表(切 B 免费 / 切 M/N 廉价被 L2 吸收 / 切 K 昂贵需归约)→ 按价格从低到高购买并行度 → 坍缩为 MergeBatch、IterBatch、ASW_Basic、StreamK 四个等价类,加前置的转Matmul 问题归约层,共五大分支,**完备且极小**。 +3. **各分支的条件**都不是孤立经验,而是"资格(并行度够不够)→ 上限(容量/算存比封顶)→ 下限(搬移效率托底)"三层逻辑的实例化;所有经验常数可标定、可移植。 +4. **边界不靠硬切**:重叠区由端到端时延模型统一仲裁,分支体系负责候选集的完备无冗余。 +5. **设计经受了漏洞拷问**:修正了转Matmul 的方向不对称与 IterBatch 均衡条件的字面误杀,显式声明了交叉广播与 K 退化的归属;相对源码实现,本文条件体系更细、更真、覆盖更全。 + +--- + +*参考文档:《BMM分块计算数学公式》《BMM最优软件实现方案设计》《BatchMatMulV3算子分支实现分析》《BMM从分块计算到四大分支的逻辑推导》;源码:`ops-nn/matmul/batch_mat_mul_v3`(arch35);芯片资料:昇腾 950PR 架构白皮书与 CANN 9.0.0 性能建模文档。*