Batch Matmul 算子特性分析（v0.5）

基于 issue#3（v0.3）扩展补全。所有分支的进入条件只出现 B/M/N/K 与芯片规格参数，每条条件附简要解释。
目标芯片：昇腾 950PR / DV100 同档（DAV_3510）。

算子功能与接口说明

算子功能：完成带 batch 的矩阵乘计算。

算子输入：

• 左矩阵（A 矩阵）：[BatchA, M, K]、数据类型 dtype、layout（典型 ND）
• 右矩阵（B 矩阵）：[BatchB, K, N]、数据类型 dtype、layout（典型 ND）
• 偏置 bias：维度固定 [B, 1, N]、数据类型 dtype、layout（固定 ND），可为空

算子输出：输出矩阵（C 矩阵）：[BatchC, M, N]、数据类型 dtype、layout（典型 ND）

计算公式：C = A @ B + bias

其中 A、B 输入维度典型为 3 维，最后两维做矩阵乘计算。例如 A、B 输入维度分别为 (B,M,K)、(B,K,N) 时，C 维度为 (B,M,N)。bias 是维度为 (B,1,N) 的向量。B 为矩阵乘法的 batch 数，M 为 A 的行数和 C 的行数，N 为 B 的列数和 C 的列数，K 为 A 的列数和 B 的行数。

符号与芯片参数约定

后文所有分支条件只使用 case 形状参数（B、M、N、K、dtype）与下列芯片规格参数：

以上经验常数均为该档芯片实测值；换芯片时逻辑结构不变，只换常数表。

BMM 算子最优实现分析

最优软件实现设计的系统推导

如何理解性能最优

NPU 上 BMM 的执行是核内多级硬件流水的并行——Cube 计算（MMAD）、GM/L2→L1 搬移（MTE2）、L1→L0 搬移（MTE1）、L0C 写出（Fixpipe），多核并行时各流水级时延可被双缓冲（double buffer）等机制相互掩盖，最终：

总时延 = 流水线最慢的一级。算子优化的关键就是对瓶颈流水级的优化。由此直接得出一个重要的设计自由度——瓶颈交换：当 MTE2（搬移）是瓶颈、MMAD（计算）不是瓶颈时，可以牺牲一定 MMAD 时延（例如冗余计算）换取 MTE2 性能提升；反之，当 MMAD 是瓶颈时，可以牺牲一定 MTE2 时延（例如重复搬移）换取计算效率提升。只要瓶颈级时延下降，总时延就下降。

后文会看到：MergeBatch 就是"牺牲算力换搬移效率"的典型，ASW_Basic 切 M/N 就是"牺牲搬移（重复读）换并行度"的典型——它们的存在正当性都来自这个 max 模型。

实现本质逻辑

BMM 在 NPU 上实现的本质是：把参与计算的数据分块（tile），由全部 AIC 核并行 + 串行地完成这些分块的计算，再组合成最终结果：

分块有 4 个维度：B、M、N、K。核间怎么分这 4 个维度，就是分支划分的第一性问题（核内分块是第二性问题，属于各分支内部的 tiling）。

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

不同切分差异的根本原因：

• B 维在数学上独立（BMM 语义就是逐 batch 独立矩阵乘），所以切 B 天然零重复、零依赖；
• K 维有 L0C 累加机制：核内切 K 时，切 K 后多轮 mmad 在 256KB 的 L0C 上原地累加（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 种是完备的。按切分特征分组：

15 种组合的代价结构只由两个布尔特征决定——是否含 K、是否含 M/N。

三个维度的"核间切分价格"严格排序：cost(切 B) = 0 < cost(切 M/N) ≪ cost(切 K)

• 切 B 免费：零重复读、零依赖、零中间写出；
• 切 M/N 廉价但有价：共享矩阵被重复读，但若共享部分能驻留 L2（128MB），重复读以 5.2TB/s 命中 L2 而非 1.6TB/s 的 GM，则代价大部分被 L2 吸收；配合 swizzle（执行顺序编排）压缩同时活跃的数据集，代价进一步压低；
• 切 K 昂贵：归约流量 $T_{Reduce} \propto grid_K \times$ 输出量，且引入核间同步——这是结构性代价，L2 吸收不掉。

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

BMM 算子软件分支推导

1、问题规约（能降维就不在 BMM 本体里解决）

• 若 BatchA = 1 或 BatchB = 1，BMM 可通过维度折叠转化为普通 Matmul，直接复用 Matmul 的成熟优化体系（tiling、L2 切分、全载），不必在 BMM 框架内重新发明轮子 → 转 Matmul；
• 退化 case（K=0 无计算、K=1 无累加深度）让 Cube 完全或几乎无用，改走 AIV 向量通路——这属于"计算通路选择"，是正交于切分的前置判断。可称为特殊分支。

2、先买免费的 B 维度

• 切 B 可满足多轮都尽量填满核；
• 如果单 Batch 的 M×N 够大（能开出大 tile、L0C 利用率高、Cube 喂得饱）→ 每核逐个 batch 做完整 Matmul 就是最优，任何额外切分只增代价。可称为 IterBatch 分支；
• 如果单 Batch 的 M×N 很小，把多个 batch 在核内合并成大 tile 计算（[bM,K]@[K,bN]，取块对角线输出），用冗余算力换 Cube 利用率与搬移效率。可称为 MergeBatch 分支。

3、B 维度买不满核，加点价买廉价的 M/N 维度切分

• 优先切 B 后再切 M/N（或混合切），切 M、N 引入的共享矩阵重复读交给 L2 + swizzle 吸收。可称为 ASW_Basic 分支。

4、免费和廉价维度 B、M、N 都买不满核，才考虑加价买昂贵的 K 维度切分

• 用核间规约代价换并行度收益。可称为 StreamK 分支。

由此推导出 6 大分支：转Matmul、特殊分支、MergeBatch、IterBatch、ASW_Basic、StreamK。

各分支的进入条件给出的是"主场"；主场之间存在重叠区（如 B ≥ C 且 M×N 中等时 MergeBatch 与 IterBatch 都合法），重叠区的归属由端到端时延模型 $T_{total}$ 仲裁，分支体系保证候选集完备无冗余。

可转 Matmul 分支

进入分支条件

解释：单边 batch=1 的 BMM 与普通 Matmul 在数学上只差一个维度标签，无需在 BMM 框架内重新发明轮子。

实现方案

对于 BatchB = 1：左矩阵 [B, M, K] 的 batch 维与 M 维在 ND 布局下内存相邻、紧排，直接视图为 [B·M, K]；输出 [B·M, N] 与 [B, M, N] 的内存布局逐元素一致；无输入重排，无输出 Split，可直接转 Matmul 无任何代价。

对于 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。

此时当 A 矩阵比较小时（$MK \cdot \text{dtype} \le L1$），优先考虑将 A 矩阵加载到每个核的 L1，然后每个核均匀读取 B 矩阵完成全部计算，此时无需输入/输出重排的额外开销；

但 A 矩阵较大时，可将 BatchA==1 扩展到 BatchA==BatchB（广播），然后分别按 BMM 其他分支与"B 折叠转 Matmul + 重排"预估时延，择优选型。

MergeBatch 分支

每个核负责多个 Batch 的 Matmul 计算，核间无需同步或通信。假设单核每次要完成 b 个 Batch 的 Matmul 计算，单核计算时将 [b,M,K]@[b,K,N]=[b,M,N] 转换合并成 [bM,K]@[K,bN]=BlockTrace([bM,bN])=[b,M,N]。所谓 BlockTrace 指以 [M,N] 的 block 粒度将结果矩阵的块对角线取出作为输出；交叉项（不同 batch 的 A 与 B 的乘积）被算出但丢弃，浪费比例 (b−1)/b——进入该分支的 case 必然是访存 Bound（条件 5 保证），浪费的算力被搬移时延掩盖。

进入分支条件

进入 MergeBatch 分支需同时满足以下条件（b0 为单次合并数下限，b0 ≥ 2）：

1、batch 关系与每核份额

解释：无广播才能逐 batch 对应合并；每核至少分到 $2b_0$ 个 batch——$b_0$ 是"合并搬移有收益"的最小合并数（经验值 4，即 $b_{core} \ge 4$ 时至少能合并 2 组），2 组起步才能构成组间乒乓流水。

2、L0C 容量

解释：合并 $b_0$ 个 batch 的输出块 $[b_0 M, b_0 N]$（FP32 累加、双缓冲两份）必须放得下 256KB L0C。连最小合并都放不下，合并无从谈起。

3、单核搬移总量

解释：单核搬移数据总量不足时 GM 带宽利用率上限被压低（经验约束）。

4、搬移 tile 大小

解释：单 batch 单矩阵的最大连续搬移块须达到单次搬移效率下限；合并只是在此之上进一步放大。

5、访存 Bound（算力浪费可被掩盖）

解释：合并把单次计算的算存比放大 $b_0$ 倍后仍须低于芯片平衡点 R，保证瓶颈留在搬移侧——冗余算力被搬移时延掩盖，而不是反过来成为新瓶颈。

实现方案

确定 Tiling 参数，重点是确定 $b_{L0}$、$k_{L1}$、$b_{L1}$。

Step 1：基于 L0C 容量 + 访存 Bound 约束，先确定每次计算的合并数 b

b 尽量取 $b_{core}$ 的因子（每次计算的合并数均匀一致，负载与功耗更优）。

**Step 2：基于 $b_{L0}$ 和 L0A/L0B 确定 $k_{L0}$**

解释：L0A/L0B 各 64KB、双缓冲两份，装入合并后 $bM$ 行（或 $bN$ 列）× $k_{L0}$ 的 fractal；第三项保证 K 向内轴 dValue ≥ 32B 的搬移下限。

**Step 3：基于 $k_{L1}$ 和 L1 容量确定 $b_{L1}$**

解释：$k_{L1}$ 是 GM→L1 的 K 向粒度，须满足 128B 对齐（dValue 效率）；L1 双缓冲两份，每份驻留 $b_{L1}$ 个 batch 的 A、B 各一块；要求 $b_{L1} \ge b$（L1 驻留组不小于单次合并数，否则合并断供）。

IterBatch 分支

核间按 B 分核（每核 1 个或多个 batch），核内逐个 batch 分别执行标准 Matmul 分块计算。无算力浪费、无跨 batch 依赖，是"切 B"最朴素的形态。

进入分支条件

进入 IterBatch 分支需同时满足以下条件：

1、batch 关系与每核份额

2、负载均衡

解释：切 B 零共享零依赖，唯一系统性风险是负载不均。B 整除核数时完全均衡；不整除时要求尾波中活跃的核数不少于 minCoreNum（≈0.8C=26），保证尾波也有足够核并发以维持 GM 带宽利用率。

3、L1 容量约束（五选一）——核心要求：单核数据不重复读

注意：IterBatch 能否进入不取决于算存比判定。即使 case 本身是计算 Bound，只要 L1 放不下单核要处理的完整输入（M、N 维度的 A/B 数据），就会产生跨 batch 的重复读，额外搬移可能把算子重新拖回访存 Bound。所以进入条件必须直接由 L1 容量刻画：

a) 单核单 batch 完整驻留：

解释：每核 1 个 batch，左右矩阵同时驻留 L1，零重复读。

b) 单核多 batch 乒乓：

解释：L1 同时放下 2 个 batch 的输入，构成 batch 间双缓冲流水。

c) 单 batch 放不下：一侧完整驻留、另一侧按 Step 切分：

解释：Step 为大于 1 的正整数。左（或右）矩阵完整驻留 L1、只搬一次，右（或左）矩阵按 K（或 M）切 Step 段流水搬入——驻留侧零重复读，切分侧重复读 M（或 N）维但每段仍整块连续。

d) 多 batch 时 (c) 的乒乓版本（容量预算减半）：

e) 左右都按 K 切分：

解释：两侧矩阵都装不下时，统一按 K 切 Step 段，A/B 成段配套搬入。

4、Step 切分后的带宽效率

(c)(d)(e) 切分后的每个搬移分块须满足：

解释：切分是把粒度切小，必须守住搬移效率下限，否则切分本身把带宽打崩。

实现方案

情形 (a)：每核 1 batch 直接搬入 L1。L1→L0 时先看 L0C 能否放下完整单 batch 输出 M×N：

if (L0C >= M*N*4B):                      # L0C 放得下完整输出
    BaseM = M;  BaseN = N
    BaseK = min(align(L0A/M, 16), align(L0B/N, 16))     # 只切 K
else:                                     # L0C 放不下，按较小维切
    if M < N:
        BaseM = align(M, 16)
        BaseN = floor(L0C/4B / BaseM)
        BaseK = min(floor_align(L0A/BaseM, 16), floor_align(L0B/BaseN, 16))
    else:
        BaseN = align(N, 16)
        BaseM = floor(L0C/4B / BaseN)
        BaseK = min(floor_align(L0A/BaseM, 16), floor_align(L0B/BaseN, 16))

情形 (b)：L1 放下 2 个 batch 则按 batch 做 double buffer，核内 GM→L1→L0→Cube→L0C→GM/L2 流水。L1→L0 分块：

if (L0C >= 2*M*N*4B):                    # 双 buffer 放得下两份完整输出
    BaseM = M;  BaseN = N
    BaseK = min(floor_align(L0A/M, 16), floor_align(L0B/N, 16))
else:
    if M < N:
        BaseM = align(M, 16)
        BaseK = min(floor_align(L0A/2/BaseM, 16), K)
        BaseN = max(floor_align(L0C/2/4B/BaseM, 16), floor_align(L0B/2/BaseK, 16))
    else:
        BaseN = align(N, 16)
        BaseK = min(floor_align(L0B/2/BaseN, 16), K)
        BaseM = max(floor_align(L0C/2/4B/BaseN, 16), floor_align(L0A/2/BaseK, 16))

情形 (c)：每核 1 batch，L1 放完整一侧矩阵 + 另一侧的一部分。假设 L1 放完整左矩阵和部分右矩阵，关键是确定右矩阵搬入的 K 向长度：

fixpipe 应开 unitflag（单侧驻留场景下输出通路按单元化组织）。情形 (d)(e) 同理，按各自容量预算计算 Step 与分块。

StreamK 分支

进入分支条件

1、并行缺口存在（B/M/N 用尽仍填不满核）：

解释：以 Cube 最小分形（16×16）为粒度，B×M×N 三维最多切出 P 个互不依赖的输出块；P < 32 意味着不切 K 必有核闲置。

2、K 足够长（归约代价可接受），设核间切 K 的份数为 $grid_K$：

解释：第一式保证单核 K 段不太碎、核内 tiling 效率不崩。第二式来自 $T_{MMAD/core} \ge \alpha \cdot T_{Reduce}$（安全系数 α≈10）：每核计算时延 ∝ $MNK/grid_K$（除以单核 Cube 速率），归约时延 ∝ $grid_K \cdot MN \cdot 4\text{B}$（除以 GM 带宽）——归约流量随 $grid_K$ 线性涨、每核收益也随 $grid_K$ 线性涨，相抵后对 K 的要求随 $grid_K$ 平方增长。数值上：$grid_K$=2 → K≥6.8K；4 → K≥27K；8 → K≥108K；32 → K≥1.7M（仅极端 case）。

实现方案

核间组织：先按 B/M/N 切出 P 个输出块，再把剩余核预算折成 K 向份数。对每个 batch：

其中 mCnt、nCnt 收拢为 blocksPerBatch 的因子（避免碎核尾块）；由条件 1 知 $mCnt \cdot nCnt \le blocksPerBatch/2$，故 $grid_K \ge 2$。每个输出块由 $grid_K$ 个核组成的归约组共同计算，组内核 c 负责 K 段 $[c \cdot K/grid_K,\; (c{+}1) \cdot K/grid_K)$，$singleCoreK = K / grid_K$。

核内流水：每核对自己的 K 段做标准分块流水（MTE2→L1→L0→mmad），段内多轮在 L0C 原地累加；段算完后部分和经 Fixpipe 写出。

归约（两种方式）：

• workspace + AIV 归约（确定性）：部分和写入 GM workspace（每核一块 256×256×4B 缓冲，另加 20MB 核间通信区），由 AIV 读出各段并累加输出——AIC:AIV=1:2 配比下 2 个 AIV 伺候 1 个 AIC 的部分和流；
• AtomicAdd（非确定性）：部分和直接原子累加到输出 GM，省一遍读回，但归约顺序不定。确定性等级 > 1 的场景禁用。

参数搜索：$grid_K$ 从 2 起按 2 的幂递增，取同时满足条件 2 两式的最小值（归约代价最小）；若到 $grid_K = blocksPerBatch$ 仍不满足条件 2，则该 case 不宜 StreamK，退而接受降核的部分闲置。

ASW_Basic 分支

进入分支条件

1、并行度补齐（B 买不满核时，M/N 平面能补上）：

解释：基本块粒度 $M^t \times N^t$ 受 L0C 容量封顶（最小 16×16）；按此粒度 B×M×N 能切出至少 C 个独立输出块，则切 M/N 的并行度够用。

2、典型进入路径：$B < C$（切 B 买不满核），或 $B \ge C$ 但 IterBatch/MergeBatch 条件不满足（L1 驻留失败、负载不均、MergeBatch 上下限无交集）时的兜底。无 batch 结构限制——交叉广播 case 也落入本分支。

解释：切 M/N 的固有代价是共享矩阵被多核重复读，但共享部分若能驻留 128MB L2，重复读以 5.2TB/s 命中 L2 而非 1.6TB/s 的 GM，代价大部分被吸收；配合 swizzle 压缩同时活跃的工作集，代价进一步压低。

实现方案

1、核间切分维度选择（按共享代价从低到高）：

• 切 B（$B \ge C$）：零共享，永远先试；
• 切 M：右矩阵 [K,N] 被所有核共享，若 $KN \cdot \text{dtype} \le L2$ 则驻留 L2，零 GM 重复读；
• 切 N：对称，左矩阵 [M,K] 共享驻留；
• 混合切（B×M、M×N 等）：双向共享，靠 swizzle + L2 切分管理；
• 降核：P 远小于 C 且 K 也不满足 StreamK 条件时，宁可部分核闲置（小 case 时延绝对值小，调度开销反而主导）。

2、swizzle：ASW 滑窗蛇形。M 向按窗口 W 分组，窗内 N 向蛇形遍历（偶数窗行正向、奇数窗行反向）：

数学效果：同一时刻 C 个核的活跃工作集被压缩到"W 个 A 行块 + 一条 B 列块带"，L2 足迹最小，共享读取基本命中 L2。窗口取 $\lfloor\sqrt C\rfloor$ 的最大因子：窗越接近方形、A/B 两侧足迹之和越小，且因子性保证整窗被核数均分、窗口边界不碎。

3、L2 切分：工作集超过 128MB 时按 mL2TileNum × nL2TileNum 切块，每个 L2 块内错位分核（对角线分配），避免多核同时抢同一地址的读读冲突，并优先选拖尾小的方案。

4、核内 tiling：baseM/baseN 由 L0C 容量（$M^t N^t \cdot 4\text{B} \cdot DB \le L0C$）与计算访存比（$1/M^t + 1/N^t$ 越小越接近 Cube Bound）联合选取，优先整除 M、N 减少尾块；baseK 由 L0A/L0B 反推（$M^t K^t \cdot \text{dtype} \cdot 2 \le L0A$，$K^t N^t \cdot \text{dtype} \cdot 2 \le L0B$），内轴对齐 128B/256B；L1 按容量开双缓冲，余量充足时开 4 buffer。

5、内部特化（参数极限，不是独立分支）：

• L1 全载：单边无 batch 且该侧矩阵小（如 BatchA=1 且 $M \le 256$、$MK \cdot \text{dtype} \cdot 2 \le L1$），且对侧循环轮数足够（$\ge 4$/核）时，把小侧矩阵整个常驻 L1、只搬一次——搬运次数从"每基本块一次"降为"整个 kernel 一次"；
• 广播驻留：交叉广播时对广播侧做 L1/L2 驻留，共享关系与切 M/N 同构。

特殊分支

K = 0：无任何计算，C = bias 或 0，纯 AIV 写值。

K = 1：退化为逐元素乘 C = A ⊙ B，无累加深度，Cube 的 16×16×16 粒度浪费 15/16，走 AIV 向量通路（GM→UB→Mul→GM）优于 Cube 通路。

解释：这一层与"切分维度"正交，是计算通路选择的前置判断——K 退化时 Cube 完全或几乎无用，AIV（64 核、每拍 256B）纯向量流水更优。触发需 batch 足够多（≥ 2×AIV 核数，开 UB 乒乓）且单 batch 输入输出能驻留 UB。