添加《Batch Matmul算子特性分析》v1.0 与《BMM从分块计算到四大分支的逻辑推导》
- BatchMatmul算子特性分析.md/.html：基于 issue #1/#2 草稿完成的完整文档， 含功能接口、五大分支（转Matmul/MergeBatch/IterBatch/StreamK/ASW_Basic）的 系统推导、进入条件与实现方案、完备性审视（含转Matmul方向不对称等5处修正） - BMM从分块计算到四大分支的逻辑推导.md/.html：补齐分块公式与最优分支之间的 逻辑层——切分代价不对称性（切B<切M/N<切K）与15种组合的等价类坍缩
This commit is contained in:
DemosRight
363
BMM从分块计算到四大分支的逻辑推导.html
Normal file
363
BMM从分块计算到四大分支的逻辑推导.html
Normal file
@@ -0,0 +1,363 @@
|
|||||||
|
<!DOCTYPE html>
|
||||||
|
<html lang="zh-CN">
|
||||||
|
<head>
|
||||||
|
<meta charset="UTF-8">
|
||||||
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||||
|
<title>从分块计算到四大分支：BMM 最优实现的逻辑推导</title>
|
||||||
|
<script id="MathJax-script" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
|
||||||
|
<script>
|
||||||
|
MathJax = {
|
||||||
|
tex: {
|
||||||
|
inlineMath: [['$','$'], ['\\(','\\)']],
|
||||||
|
displayMath: [['$$','$$'], ['\\[','\\]']],
|
||||||
|
tags: 'ams'
|
||||||
|
}
|
||||||
|
};
|
||||||
|
</script>
|
||||||
|
<style>
|
||||||
|
:root{--bg:#fefefe;--panel:#1a2230;--ink:#1f2933;--muted:#5f6b7a;--accent:#0b6bcb;--accent2:#0e9f6e;--warn:#b45309;--line:#d9e2ec;--code-bg:#f4f6f9;--hl:#fff7e6}
|
||||||
|
*{box-sizing:border-box}
|
||||||
|
body{font-family:"PingFang SC","Microsoft YaHei","Helvetica Neue",Arial,sans-serif;color:var(--ink);background:#eef2f6;margin:0;line-height:1.8}
|
||||||
|
.page{max-width:1020px;margin:0 auto;padding:32px 44px 80px;background:#fff;box-shadow:0 0 24px rgba(0,0,0,.06)}
|
||||||
|
h1{font-size:28px;border-bottom:3px solid var(--accent);padding-bottom:12px;margin-top:8px}
|
||||||
|
h2{font-size:22px;margin-top:48px;border-left:6px solid var(--accent);padding-left:12px;color:#0b3d73}
|
||||||
|
h3{font-size:18px;margin-top:32px;color:#0b3d73;border-bottom:1px dashed var(--line);padding-bottom:6px}
|
||||||
|
table{border-collapse:collapse;width:100%;margin:14px 0;font-size:14px}
|
||||||
|
th,td{border:1px solid var(--line);padding:7px 10px;text-align:left;vertical-align:top}
|
||||||
|
th{background:#eaf2fb;color:#0b3d73;white-space:nowrap}
|
||||||
|
tr:nth-child(even) td{background:#f8fafc}
|
||||||
|
code,pre{font-family:"JetBrains Mono",Consolas,Menlo,monospace;font-size:13px}
|
||||||
|
code{background:var(--code-bg);padding:1px 5px;border-radius:4px;color:#9d2c5e}
|
||||||
|
pre{background:var(--code-bg);border:1px solid var(--line);border-radius:8px;padding:14px;overflow-x:auto;line-height:1.55}
|
||||||
|
pre code{background:none;color:#243447;padding:0}
|
||||||
|
blockquote{background:#eafaf3;border-left:5px solid var(--accent2);padding:10px 16px;border-radius:0 8px 8px 0;margin:14px 0;color:#243447}
|
||||||
|
blockquote p{margin:0}
|
||||||
|
.meta{color:var(--muted);font-size:13px;margin-bottom:24px}
|
||||||
|
</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<div class="page">
|
||||||
|
<h1>从分块计算到四大分支：BMM 最优实现的逻辑推导</h1>
|
||||||
|
<blockquote>
|
||||||
|
<p>昇腾 BatchMatMulV3 算子 | 目标芯片 950PR（DAV_3510，32 AIC / 1.65GHz）
|
||||||
|
上游文档：《BMM分块计算数学公式》（分块计算"是什么"） · 下游文档：《BMM最优软件实现方案设计》第六章（四大分支"怎么用"）
|
||||||
|
本文定位：补上中间缺失的一层——<strong>为什么从"分块计算 + 最优实现"出发，恰好推导出 MergeBatch / IterBatch / ASW_Basic / StreamK 这四大分支，不多也不少。</strong></p>
|
||||||
|
</blockquote>
|
||||||
|
<hr />
|
||||||
|
<h2>0. 问题陈述：缺失的那层逻辑是什么</h2>
|
||||||
|
<p>《BMM分块计算数学公式》告诉我们：BMM 的 4 重循环 $(b, m, n, k)$ 可以按硬件容量做嵌套分块，并用 swizzle 函数 $\sigma$ 重排块到核的映射。《BMM最优软件实现方案设计》第六章直接给出了四大分支。但中间有一跳没有论证：</p>
|
||||||
|
<blockquote>
|
||||||
|
<p>分块参数 $(B^t, M^t, N^t, K^t)$、核间切分方式、swizzle 函数 $\sigma$ 有无穷多种取法。凭什么最优实现的搜索空间恰好收敛到 <strong>4 个分支</strong>？</p>
|
||||||
|
</blockquote>
|
||||||
|
<p>本文的推导链条如下，每一节对应链条上的一环：</p>
|
||||||
|
<pre><code>分块计算公式（自由度）
|
||||||
|
→ 硬件强加的两条不变式（规则）
|
||||||
|
→ 四个维度核间切分的代价表（代价不对称）
|
||||||
|
→ "最优"的三个必要条件（目标函数展开）
|
||||||
|
→ 决策树：按代价从低到高购买并行度（推导）
|
||||||
|
→ 15 种切分组合坍缩为 4 个等价类（完备性 + 极小性证明）
|
||||||
|
→ 与分块公式的参数特化对应（闭环）
|
||||||
|
</code></pre>
|
||||||
|
<hr />
|
||||||
|
<h2>1. 起点：分块计算公式给出的自由度</h2>
|
||||||
|
<p>分块计算把 BMM 重组为两层结构（记号同《BMM分块计算数学公式》）：</p>
|
||||||
|
<p><strong>核间</strong>：$\tilde{B} \times \tilde{M} \times \tilde{N}$ 个基本块，经 swizzle $\sigma$ 映射到 $C$ 个核：</p>
|
||||||
|
<p>$$
|
||||||
|
(\beta, \mu, \nu) = \sigma^{-1}\big((c + rC) \bmod \tilde{B}\tilde{M}\tilde{N}\big), \qquad c \in [0, C)
|
||||||
|
$$</p>
|
||||||
|
<p><strong>核内</strong>：K 循环在 L0C 上原地累加，经 L1/L0 两级缓冲流水：</p>
|
||||||
|
<p>$$
|
||||||
|
C_{(\beta,\mu,\nu)}^{\text{L0C}} = \sum_{\kappa=0}^{\tilde{K}-1} \mathrm{mmad}\big(A[\beta,\mu,\kappa],\; B[\beta,\kappa,\nu]\big)
|
||||||
|
$$</p>
|
||||||
|
<p>这个数学结构里，<strong>自由决策只有三类</strong>：</p>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>自由度</th>
|
||||||
|
<th>数学对象</th>
|
||||||
|
<th>物理含义</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td>F1：核间切分</td>
|
||||||
|
<td>$\tilde B, \tilde M, \tilde N$ 如何对 $C$ 分解；$\kappa$ 是否跨核</td>
|
||||||
|
<td>哪些维度的块被分到<strong>不同核</strong>上</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>F2：核内组织</td>
|
||||||
|
<td>$B^t, M^t, N^t, K^t$ 及 L1/L0 驻留深度 $d_A, d_B$</td>
|
||||||
|
<td>单核内如何折叠/驻留/流水</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>F3：执行顺序</td>
|
||||||
|
<td>swizzle 函数 $\sigma$</td>
|
||||||
|
<td>时间上相邻的块在空间上的排布（L2 复用）</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<p><strong>"四大分支"问题，本质上就是 F1 和 F2 的选择问题</strong>（F3 是每个分支内部的二级优化）。所以推导的主线是：<strong>F1 有哪些本质上不同的选择，各自的代价是什么。</strong></p>
|
||||||
|
<hr />
|
||||||
|
<h2>2. 规则：硬件强加的两条不变式</h2>
|
||||||
|
<p>分块公式里有句话写得轻描淡写，却是整个分支体系的根：</p>
|
||||||
|
<blockquote>
|
||||||
|
<p>"<strong>K 循环 ($\kappa$) 在 L0C 上累加，不写回 GM。</strong>"</p>
|
||||||
|
</blockquote>
|
||||||
|
<p>把它和输出写回方式放在一起，就是硬件强加给所有实现的两条不变式：</p>
|
||||||
|
<p><strong>不变式 I1（L0C 累加不变式）</strong>：一个输出块 $C_{(\beta,\mu,\nu)}$ 的 $\tilde K$ 轮 mmad 结果驻留在 L0C（256KB，FP32）上原地累加，中间不产生任何 GM/L2 流量。</p>
|
||||||
|
<ul>
|
||||||
|
<li>推论 1：<strong>核内切 K 是免费的</strong>（$\tilde K$ 只是循环次数）。</li>
|
||||||
|
<li>推论 2：<strong>核间切 K 必然打破 I1</strong>——每个核只算了一段 K 的部分和，L0C 驻留不住"别人的 K"，部分和必须写出到 GM/L2 workspace，再由 AIV 做核间归约。分块公式第六章的 STREAM_K 行写的就是这个：</li>
|
||||||
|
</ul>
|
||||||
|
<p>$$
|
||||||
|
C_{(\beta,\mu,\nu)} = \sum_{c \in \text{group}} C_{(\beta,\mu,\nu)}^{(c)} \quad \text{（部分和经 workspace 归约）}
|
||||||
|
$$</p>
|
||||||
|
<p><strong>不变式 I2（输出专属不变式）</strong>：若输出块 $C_{(\beta,\mu,\nu)}$ 由一个核独占负责，则该核只写出最终结果，无核间同步；反之（切 K）则引入核间依赖与归约流量。</p>
|
||||||
|
<p>两条不变式合起来给出本文最重要的结构性事实：</p>
|
||||||
|
<blockquote>
|
||||||
|
<p><strong>K 维和其他三维在核间切分中的地位根本不对等</strong>：切 B/M/N 保持 I1、I2，切 K 同时打破 I1、I2。这就是为什么"是否切 K"天然是分支的第一分界线，StreamK 必然自成一支。</p>
|
||||||
|
</blockquote>
|
||||||
|
<hr />
|
||||||
|
<h2>3. 关键一步：四个维度核间切分的代价表</h2>
|
||||||
|
<p>对 F1，核间切分的本质选择是"4 个维度中切哪些"。逐维度分析切分代价（依据：batch 语义、数据通路、GM/L2 带宽差）：</p>
|
||||||
|
<h3>3.1 切 B —— 零代价并行</h3>
|
||||||
|
<p>不同 batch 的 $A[b], B[b], C[b]$ 在内存中天然不相交。因此：</p>
|
||||||
|
<ul>
|
||||||
|
<li><strong>读入</strong>：每个数据块被且仅被 1 个核读取，零重复；</li>
|
||||||
|
<li><strong>计算</strong>：每核独立产出最终结果，无核间依赖（保持 I1、I2）；</li>
|
||||||
|
<li><strong>写出</strong>：只有最终结果，无中间写出。</li>
|
||||||
|
</ul>
|
||||||
|
<p>代价 = 0。但并行度上限 = $B$，且单核分到的计算是**单个 batch 的 $[M,K]\times[K,N]$**——Cube 效率完全由单 batch 的 $M, N$ 决定（第 4 节会看到这是 MergeBatch 与 IterBatch 分家的原因）。</p>
|
||||||
|
<h3>3.2 切 M 或切 N —— 低代价并行（代价可被 L2 吸收）</h3>
|
||||||
|
<p>切 M：同一 batch 的右矩阵 $B[b,:,:]$（大小 $K \times N$）被负责不同 M 行块的所有核**重复读取**；切 N 对称（左矩阵 $M \times K$ 被重复读）。</p>
|
||||||
|
<ul>
|
||||||
|
<li><strong>读入</strong>：重复因子 = 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$（滑窗蛇形）进一步压缩同时活跃的工作集。</li>
|
||||||
|
<li><strong>计算</strong>：无核间依赖（保持 I1、I2）。</li>
|
||||||
|
<li><strong>写出</strong>：只有最终结果。</li>
|
||||||
|
</ul>
|
||||||
|
<p>代价 = 共享矩阵的重复读取，受 L2 容量约束，属于<strong>有上限的低代价</strong>。</p>
|
||||||
|
<h3>3.3 切 K —— 高代价并行（结构性代价，无法吸收）</h3>
|
||||||
|
<ul>
|
||||||
|
<li><strong>读入</strong>：无重复（各核读不同 K 段）；</li>
|
||||||
|
<li><strong>计算</strong>：<strong>打破 I1、I2</strong>——每个输出块由 $\text{grid}_K$ 个核共同产生，部分和写 workspace，再归约；</li>
|
||||||
|
<li><strong>写出</strong>：存在中间结果写出。额外时延项 $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，另有同步开销）。</li>
|
||||||
|
</ul>
|
||||||
|
<p>代价 = 归约流量 + 核间同步，是<strong>随 grid_K 线性增长的结构性代价</strong>，L2 吸收不掉。</p>
|
||||||
|
<h3>3.4 代价不对称性总结</h3>
|
||||||
|
<p>$$
|
||||||
|
\boxed{\;\text{cost}(\text{切}B) = 0 \;<\; \text{cost}(\text{切}M) \approx \text{cost}(\text{切}N) \;\ll\; \text{cost}(\text{切}K)\;}
|
||||||
|
$$</p>
|
||||||
|
<p>这个排序不是经验，是三条硬件事实的推论：① batch 维在数学上独立（BMM 语义）；② L0C 累加机制（I1）使切 K 产生独一份的归约开销；③ L2（5.2TB/s / 128MB）恰好能把 M/N 共享读取的代价吸收掉大半。<strong>整条决策树就是按这个代价排序"从便宜到贵"地购买并行度。</strong></p>
|
||||||
|
<hr />
|
||||||
|
<h2>4. 目标：最短端到端时延展开为三个必要条件</h2>
|
||||||
|
<p>最优实现的判据是 $T_{total} = \max(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{FIXPIPE} [, T_{REDUCE}])$ 最小。逐项拆开，$\min\max$ 等价于三个必要条件：</p>
|
||||||
|
<p><strong>R1 核占满（并行度条件）</strong>。$T_{MMAD}$ 与活跃核数成反比；GM 带宽也要求 ≥ 3/4 核（24 核）并发才能达到 90%+ 利用率。设不切 K 时的独立输出块数</p>
|
||||||
|
<p>$$
|
||||||
|
P \;=\; B \times \Big\lceil \tfrac{M}{16} \Big\rceil \times \Big\lceil \tfrac{N}{16} \Big\rceil
|
||||||
|
$$</p>
|
||||||
|
<p>（16 是 mmad 的 fractal 粒度，即输出块的最小可分单元）。$P \ge C = 32$ 是满核的必要条件；$P < 32$ 时必然有核闲置，闲置部分纯浪费。</p>
|
||||||
|
<p><strong>R2 搬运最快（复用条件）</strong>。MTE2 有效带宽由 L2 命中率决定：</p>
|
||||||
|
<p>$$
|
||||||
|
BW_{MTE2} \approx (1 - h) \times 1.6 + h \times 5.2 \ \text{TB/s}
|
||||||
|
$$</p>
|
||||||
|
<p>baseM = baseN = 256 时达成 Cube Bound 需要 $BW_{MTE2} \ge 2.64$ TB/s，即 $h \gtrsim 28.9%$。GM 流量必须最小化（每份数据 GM 只读一次），重复读取尽量落在 L2/L1。</p>
|
||||||
|
<p><strong>R3 Cube 喂饱（效率条件）</strong>。单核每次 mmad 序列的 tile 要足够大：L0C 利用率</p>
|
||||||
|
<p>$$
|
||||||
|
U = \frac{M^t \times N^t \times 4\text{B}}{256\text{KB}} = \frac{M^t N^t}{65536}
|
||||||
|
$$</p>
|
||||||
|
<p>越接近 1，Cube 空转越少；$M^t, N^t$ 需 16 对齐，且 MTE1 不拖后腿要求 baseM, baseN ≥ 80 左右。</p>
|
||||||
|
<p><strong>分支推导 = 在 R1/R2/R3 约束下，按 §3 的代价排序购买并行度与 Cube 效率。</strong> 下面逐分支看它们各自是哪一组约束矛盾的"最优解"。</p>
|
||||||
|
<hr />
|
||||||
|
<h2>5. 推导主链：四种典型矛盾 → 四个最优响应</h2>
|
||||||
|
<h3>5.1 情形一：$B \ge C$ 且单 batch $M \times N$ 大 → <strong>IterBatch</strong></h3>
|
||||||
|
<ul>
|
||||||
|
<li>R1：零代价并行度 $B \ge 32$ 已足够，无需切 M/N（省掉共享读取），更无需切 K。</li>
|
||||||
|
<li>R3：$M \times N$ 大（如 ≥ 128×128），单 batch 就能开出 baseM×baseN ≈ 256×256 的 tile，$U$ 接近 1，Cube 喂得饱。</li>
|
||||||
|
<li>R2：核内走标准三级流水（L1→L0→Cube→L0C），M/N 大 tile 保证复用。</li>
|
||||||
|
</ul>
|
||||||
|
<p>此时<strong>任何额外的切分都只会引入代价而无收益</strong>，最优解就是"每核分若干 batch，核内逐个 batch 做完整 Matmul"。这正是分块公式中 $B^t = 1$、batch 作为核间最外层维度的特化。</p>
|
||||||
|
<p><strong>判据</strong>：$B \ge C$ 且 $M N \gtrsim 128^2$（L0C 利用率 ≥ 25% 量级）。</p>
|
||||||
|
<h3>5.2 情形二：$B \ge C$ 但单 batch $M \times N$ 小 → <strong>MergeBatch</strong></h3>
|
||||||
|
<ul>
|
||||||
|
<li>R1：满足（$B \ge 32$）。</li>
|
||||||
|
<li>R3：<strong>不满足</strong>。$M = N = 64$ 时 $U = 64^2/256^2 = 6.25%$，Cube 大片闲置；$M = N = 16$ 时 $U = 0.4%$，mmad 几乎全在空转。并行度够，但每个核"吃不饱"。</li>
|
||||||
|
</ul>
|
||||||
|
<p>矛盾在 R3。唯一的解法：<strong>把"折叠维度"用在 batch 上</strong>——核内把 $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 块作为有效输出。</p>
|
||||||
|
<p>代价是交叉项算力浪费 $(b-1)/b$。为什么这个代价值得付？因为小 $M, N$ 意味着算术强度低：</p>
|
||||||
|
<p>$$
|
||||||
|
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}
|
||||||
|
$$</p>
|
||||||
|
<p>单核分摊的 GM 带宽约 50GB/s，Cube 要满转需要 AI ≈ 270 FLOP/B；小 $M,N$ 的 case 远低于此，<strong>瓶颈本来就在搬运，Cube 浪费的拍数被 MTE2 时延掩盖</strong>——浪费是免费的。这就是"为什么是 MergeBatch"的定量理由：它不是牺牲算力换效率，而是在算力本就用不完时回收闲置。</p>
|
||||||
|
<p><strong>判据</strong>：$B \ge 2C$（每核至少 2 batch）且 $MN < 128^2$ 且 $AI < 270$。</p>
|
||||||
|
<h3>5.3 情形三：$B < C$ 但 $B \times \lceil M/16 \rceil \times \lceil N/16 \rceil \ge C$ → <strong>ASW_Basic</strong></h3>
|
||||||
|
<ul>
|
||||||
|
<li>R1：零代价维度 B 买不够 32 核，但低代价维度 M/N 可以补齐（$P \ge 32$）。</li>
|
||||||
|
<li>R2：切 M/N 引入共享矩阵的重复读取——必须靠 L2 驻留 + swizzle 把代价压到最低，这正是 swizzle 函数 $\sigma$ 的主战场（滑窗 $W = \max{d \mid d \mid C, d \le \lfloor\sqrt C\rfloor}$ 使同窗口 A 行块与 B 列块的 L2 足迹最小）。</li>
|
||||||
|
</ul>
|
||||||
|
<p>矛盾是"R1 缺并行、只能靠低代价维度补"。最优解是一个<strong>通用框架</strong>：允许切 B/M/N 的任意组合（不切 K），核间切分维度与 swizzle 由 shape 决定——切 B 优先（零共享），B 不够切 M 或 N（共享一侧矩阵，L2 吸收），再不够混合切。这就是 ASW_Basic：它不是一个具体切法，而是"所有不切 K、含 M/N 共享切法"的总框架，$\sigma$ 取 ASW 滑窗蛇形。</p>
|
||||||
|
<p><strong>判据</strong>：$P \ge C$ 且不满足情形一/二（B 不够，或 M/N 大到不需要折叠 batch）。</p>
|
||||||
|
<blockquote>
|
||||||
|
<p>与 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）。</p>
|
||||||
|
</blockquote>
|
||||||
|
<h3>5.4 情形四：$P = B \lceil M/16 \rceil \lceil N/16 \rceil < C$ → <strong>StreamK</strong></h3>
|
||||||
|
<ul>
|
||||||
|
<li>R1：<strong>无法满足</strong>。B、M、N 三个便宜维度全部用尽（切到 16 粒度）仍凑不满 32 核。典型：$B = 1, M = N = 64$ 时 $P = 16 < 32$，一半核闲置。</li>
|
||||||
|
<li>此时<strong>唯一剩余的并行维度是 K</strong>。切 K 虽然代价高（打破 I1/I2，付 $T_{REDUCE}$），但不切的代价是核闲置——两害相权，当</li>
|
||||||
|
</ul>
|
||||||
|
<p>$$
|
||||||
|
T_{MMAD/\text{core}} \gg T_{REDUCE} \;\Longleftrightarrow\; K \;\gtrsim\; \text{grid}_K^2 \times 1690
|
||||||
|
$$</p>
|
||||||
|
<p>（推导：要求计算时延 ≥ 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 的组合）来降低归约组大小。</p>
|
||||||
|
<p>矛盾是"便宜维度用尽仍缺并行"。最优解是<strong>付归约代价买 K 维并行</strong>，且归约组能小则小。这就是 StreamK——分块公式中 $\tilde K$ 跨核拆分、部分和经 workspace 归约的特化。</p>
|
||||||
|
<p><strong>判据</strong>：$P < C$（或 P 虽够但核内 M/N 范围被压得太碎）且存在 grid_K 使 $K/\text{grid}_K \ge 256$ 且满足上式。</p>
|
||||||
|
<h3>5.5 一条链总结</h3>
|
||||||
|
<pre><code>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吸收代价）
|
||||||
|
</code></pre>
|
||||||
|
<p><strong>四大分支各自是唯一最优解的"矛盾区域"互不相同</strong>：IterBatch 解"零代价并行已够"的区域，MergeBatch 解"并行够但 Cube 饿"的区域，ASW_Basic 解"并行缺、低代价维度可补"的区域，StreamK 解"便宜维度用尽"的区域。区域不同，最优响应不同——这就是四个分支的存在性证明。</p>
|
||||||
|
<hr />
|
||||||
|
<h2>6. 完备性与极小性：为什么恰好四个，不多不少</h2>
|
||||||
|
<h3>6.1 完备性：15 种切分组合按代价特征坍缩为 4 个等价类</h3>
|
||||||
|
<p>核间切分的所有可能 = 四维 {B, M, N, K} 的非空子集，共 $2^4 - 1 = 15$ 种（《方案设计》§6.1 已枚举 C1~C15）。关键观察：<strong>一个切分组合的代价结构只由两个布尔特征决定</strong>——</p>
|
||||||
|
<ul>
|
||||||
|
<li><strong>是否含 K</strong>（决定要不要付归约代价，打破 I1/I2）；</li>
|
||||||
|
<li><strong>是否含 M 或 N</strong>（决定有没有共享矩阵的重复读取）。</li>
|
||||||
|
</ul>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>含 K？</th>
|
||||||
|
<th>含 M/N？</th>
|
||||||
|
<th>组合</th>
|
||||||
|
<th>代价结构</th>
|
||||||
|
<th>归入分支</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td>是</td>
|
||||||
|
<td>任意</td>
|
||||||
|
<td>C4, C7, C9, C10, C12~C15（8 个）</td>
|
||||||
|
<td>必有归约；grid_K×grid_B×grid_M×grid_N 只是参数差异</td>
|
||||||
|
<td><strong>StreamK</strong></td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>否</td>
|
||||||
|
<td>是</td>
|
||||||
|
<td>C2, C3, C5, C6, C8, C11（6 个）</td>
|
||||||
|
<td>必有共享读取；切哪几维只是 grid 参数</td>
|
||||||
|
<td><strong>ASW_Basic</strong></td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>否</td>
|
||||||
|
<td>否（纯 {B}）</td>
|
||||||
|
<td>C1（1 个）</td>
|
||||||
|
<td>零共享零归约</td>
|
||||||
|
<td>核内只有两种组织方式，见下</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<p>纯 {B} 的核内组织方式只有两种——<strong>把多个 batch 合并成一个大 tile 算（MergeBatch），或逐个 batch 算（IterBatch）</strong>——不存在第三种（要么利用 batch 间的 tile 级合并，要么不利用）。于是：</p>
|
||||||
|
<p>$$
|
||||||
|
15 \;\xrightarrow{\text{按 (含K, 含M/N) 归并}}\; 1 + 1 + 2 \;=\; \boxed{4}
|
||||||
|
$$</p>
|
||||||
|
<p>任何合法 case 的任意切法都落在这 4 个等价类之一 ⇒ <strong>完备</strong>。</p>
|
||||||
|
<blockquote>
|
||||||
|
<p>与源码对照：arch35 源码有 10 个策略，为什么这里只有 4 个？因为源码策略 = 4 个切分等价类 × 两个正交维度的笛卡尔积：① <strong>计算通路</strong>（Cube vs AIV：K=0 的 K_EQUAL_ZERO、K=1 的 TO_MUL 是 K 退化时的通路切换，发生在切分决策之前，属于预处理层）；② <strong>驻留策略</strong>（AL1/BL1_FULL_LOAD 是 ASW_Basic 内部 $M^t = M$ / $N^t = N$ 的 tiling 极限，ITER_BATCH_BROADCAST 是 IterBatch 在广播输入下的数据复用特化）。它们是分支<strong>内部</strong>的参数特化，不构成新的切分等价类。源码尾部还有 BASE=999 无条件兜底，也印证"特判 ⊂ 通用"的层级结构。</p>
|
||||||
|
</blockquote>
|
||||||
|
<h3>6.2 极小性：每个分支都有它是唯一最优的 shape 区域</h3>
|
||||||
|
<p>去掉任何一个分支，都存在 case 失去最优实现：</p>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>分支</th>
|
||||||
|
<th>独占最优的示例 case（BF16）</th>
|
||||||
|
<th>替代方案的劣势</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td>IterBatch</td>
|
||||||
|
<td>B=32, M=N=K=4096</td>
|
||||||
|
<td>ASW 切 M/N 引入无谓共享读取；MergeBatch 引入无谓浪费</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>MergeBatch</td>
|
||||||
|
<td>B=64, M=N=64, K=256</td>
|
||||||
|
<td>IterBatch 的 L0C 利用率仅 6.25%，Cube 闲置 ~94%</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>ASW_Basic</td>
|
||||||
|
<td>B=2, M=N=8192, K=1024</td>
|
||||||
|
<td>切 B 仅 2 核活跃；StreamK 付无谓归约</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>StreamK</td>
|
||||||
|
<td>B=1, M=N=64, K=65536</td>
|
||||||
|
<td>不切 K 只用 16/32 核（P=16），时延差数量级（文档 Case 3：19.5ms → 0.61μs/核量级）</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<p>⇒ 四者<strong>互为不可替代，构成极小完备集</strong>。</p>
|
||||||
|
<hr />
|
||||||
|
<h2>7. 边界是软的：竞争区域由时延模型仲裁</h2>
|
||||||
|
<p>上面的判据给出的是分支的"主场"，但相邻分支的边界不是硬切换。重叠区域（如 B ≥ 32 且 M×N 中等时 IterBatch 与 ASW 切 B 等效；M×N 在 64²~128² 之间时 MergeBatch 与 IterBatch 互有胜负）必须由统一点评仲裁：</p>
|
||||||
|
<p>$$
|
||||||
|
\text{branch}^* = \arg\min_{\text{cand} \in \bigcup \text{4 分支候选}} \max\big(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{FIXPIPE}, T_{REDUCE}\big)
|
||||||
|
$$</p>
|
||||||
|
<p>这就是《方案设计》§6.7 决策算法"全分支生成候选 → 时延模型评估 → 取 min"的合理性来源：<strong>分支体系保证候选集完备且无冗余（每个等价类只派一个代表框架），时延模型在等价类内部和边界上做精细仲裁。</strong> 两层结构缺一不可——只有时延模型没有分支体系，搜索空间是 15 种组合 × 全部 grid，爆炸且重复；只有分支体系没有时延模型，边界 case 会被错误硬切。</p>
|
||||||
|
<hr />
|
||||||
|
<h2>8. 闭环：与分块计算公式的参数特化对应</h2>
|
||||||
|
<p>最后把推导结果写回分块公式的语言，验证链条闭合（对应《BMM分块计算数学公式》第六章，但现在每行都有了"理由"）：</p>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>分支</th>
|
||||||
|
<th>分块参数特化</th>
|
||||||
|
<th>推导出处</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td>IterBatch</td>
|
||||||
|
<td>$B^t = 1$；$\tilde B \ge C$；batch 为核间最外层；核内标准 K 累加</td>
|
||||||
|
<td>§5.1：零代价并行够 + Cube 饱</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>MergeBatch</td>
|
||||||
|
<td>$B^t = b > 1$ 折叠进 $M^t_{eff} = bM^t, N^t_{eff} = bN^t$；L0C 二次方程求最优 $b$；BlockTrace 取对角</td>
|
||||||
|
<td>§5.2：Cube 饿 → 折叠 batch 换 $U \uparrow$，浪费被访存 Bound 掩盖</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>ASW_Basic</td>
|
||||||
|
<td>通用 $(\beta, \mu, \nu)$ 展开；$\sigma$ = ASW 滑窗蛇形（$W \le \lfloor\sqrt C\rfloor$ 的最大因子） | §5.3：低代价维度补并行，$\sigma$ 压缩 L2 足迹</td>
|
||||||
|
<td></td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>StreamK</td>
|
||||||
|
<td>$\tilde K$ 跨核拆分，$C_{(\beta,\mu,\nu)} = \sum_c C^{(c)}$ 部分和归约；grid_K×grid_B×grid_M×grid_N ≤ C</td>
|
||||||
|
<td>§5.4：便宜维度用尽，付归约代价买 K 维并行</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<hr />
|
||||||
|
<h2>9. 总结：一句话逻辑链</h2>
|
||||||
|
<blockquote>
|
||||||
|
<p><strong>分块计算给出自由度（4 维怎么切），硬件给出规则（L0C 累加使切 K 独贵、batch 独立使切 B 免费、L2 使切 M/N 廉价），最优目标给出约束（核要满、Cube 要饱、搬运要省）。按"最便宜的维度优先购买并行度、Cube 不够大就折叠 batch、便宜维度用尽才切 K"的原则展开，15 种切分组合在 (含K, 含M/N) 两个代价特征下恰好坍缩为 4 个等价类——这就是 MergeBatch / IterBatch / ASW_Basic / StreamK 四大分支，完备且极小；边界区域由端到端时延模型统一仲裁。</strong></p>
|
||||||
|
</blockquote>
|
||||||
|
<hr />
|
||||||
|
<p><em>文档版本：v1.0 | 上游：《BMM分块计算数学公式》v1.1+ | 下游：《BMM最优软件实现方案设计》v4.0 第六章</em></p>
|
||||||
|
|
||||||
|
</div>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
302
BMM从分块计算到四大分支的逻辑推导.md
Normal file
302
BMM从分块计算到四大分支的逻辑推导.md
Normal file
@@ -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<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**（决定要不要付归约代价，打破 I1/I2）；
|
||||||
|
- **是否含 M 或 N**（决定有没有共享矩阵的重复读取）。
|
||||||
|
|
||||||
|
| 含 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） | 替代方案的劣势 |
|
||||||
|
|---|---|---|
|
||||||
|
| 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 第六章*
|
||||||
754
BatchMatmul算子特性分析.html
Normal file
754
BatchMatmul算子特性分析.html
Normal file
@@ -0,0 +1,754 @@
|
|||||||
|
<!DOCTYPE html>
|
||||||
|
<html lang="zh-CN">
|
||||||
|
<head>
|
||||||
|
<meta charset="UTF-8">
|
||||||
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||||
|
<title>Batch Matmul 算子特性分析</title>
|
||||||
|
<script id="MathJax-script" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
|
||||||
|
<script>
|
||||||
|
MathJax = {
|
||||||
|
tex: {
|
||||||
|
inlineMath: [['$','$'], ['\\(','\\)']],
|
||||||
|
displayMath: [['$$','$$'], ['\\[','\\]']],
|
||||||
|
tags: 'ams'
|
||||||
|
}
|
||||||
|
};
|
||||||
|
</script>
|
||||||
|
<style>
|
||||||
|
:root{--bg:#fefefe;--panel:#1a2230;--ink:#1f2933;--muted:#5f6b7a;--accent:#0b6bcb;--accent2:#0e9f6e;--warn:#b45309;--line:#d9e2ec;--code-bg:#f4f6f9;--hl:#fff7e6}
|
||||||
|
*{box-sizing:border-box}
|
||||||
|
body{font-family:"PingFang SC","Microsoft YaHei","Helvetica Neue",Arial,sans-serif;color:var(--ink);background:#eef2f6;margin:0;line-height:1.8}
|
||||||
|
.page{max-width:1020px;margin:0 auto;padding:32px 44px 80px;background:#fff;box-shadow:0 0 24px rgba(0,0,0,.06)}
|
||||||
|
h1{font-size:28px;border-bottom:3px solid var(--accent);padding-bottom:12px;margin-top:8px}
|
||||||
|
h2{font-size:22px;margin-top:48px;border-left:6px solid var(--accent);padding-left:12px;color:#0b3d73}
|
||||||
|
h3{font-size:18px;margin-top:32px;color:#0b3d73;border-bottom:1px dashed var(--line);padding-bottom:6px}
|
||||||
|
table{border-collapse:collapse;width:100%;margin:14px 0;font-size:14px}
|
||||||
|
th,td{border:1px solid var(--line);padding:7px 10px;text-align:left;vertical-align:top}
|
||||||
|
th{background:#eaf2fb;color:#0b3d73;white-space:nowrap}
|
||||||
|
tr:nth-child(even) td{background:#f8fafc}
|
||||||
|
code,pre{font-family:"JetBrains Mono",Consolas,Menlo,monospace;font-size:13px}
|
||||||
|
code{background:var(--code-bg);padding:1px 5px;border-radius:4px;color:#9d2c5e}
|
||||||
|
pre{background:var(--code-bg);border:1px solid var(--line);border-radius:8px;padding:14px;overflow-x:auto;line-height:1.55}
|
||||||
|
pre code{background:none;color:#243447;padding:0}
|
||||||
|
blockquote{background:#eafaf3;border-left:5px solid var(--accent2);padding:10px 16px;border-radius:0 8px 8px 0;margin:14px 0;color:#243447}
|
||||||
|
blockquote p{margin:0}
|
||||||
|
.meta{color:var(--muted);font-size:13px;margin-bottom:24px}
|
||||||
|
</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<div class="page">
|
||||||
|
<h1>Batch Matmul 算子特性分析</h1>
|
||||||
|
<blockquote>
|
||||||
|
<p>版本：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 上的<strong>最优软件实现方案</strong>（端到端总时延最短），并论证分支体系的完备性与极小性。
|
||||||
|
阅读主线：第 1 章定义问题 → 第 2 章定义"最优" → 第 3 章推导分支划分 → 第 4 章逐分支给出进入条件与实现方案 → 第 5 章决策流程 → 第 6 章完备性审视（漏洞检查）。</p>
|
||||||
|
</blockquote>
|
||||||
|
<hr />
|
||||||
|
<h2>1. 算子功能与接口</h2>
|
||||||
|
<h3>1.1 算子功能</h3>
|
||||||
|
<p>Batch Matmul（BMM）完成<strong>带 batch 维的矩阵乘</strong>计算：对 batch 维的每个索引独立执行一次矩阵乘，再加可选偏置：</p>
|
||||||
|
<p>$$
|
||||||
|
C[b, m, n] = \sum_{k=0}^{K-1} A[b, m, k] \cdot B[b, k, n] + bias[b, 1, n]
|
||||||
|
$$</p>
|
||||||
|
<p>即 <code>C = A @ B + bias</code>。A、B 输入维度典型为 3 维（支持最多 4 级 batch 维 $b = b_0 b_1 b_2 b_3$ 展平），最后两维做矩阵乘。</p>
|
||||||
|
<h3>1.2 接口参数</h3>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>参数</th>
|
||||||
|
<th>形状</th>
|
||||||
|
<th>数据类型</th>
|
||||||
|
<th>layout</th>
|
||||||
|
<th>说明</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td>左矩阵 A</td>
|
||||||
|
<td><code>[BatchA, M, K]</code></td>
|
||||||
|
<td>dtype（FP16/BF16/FP8/…）</td>
|
||||||
|
<td>典型 ND</td>
|
||||||
|
<td>可带转置标记（isATrans）</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>右矩阵 B</td>
|
||||||
|
<td><code>[BatchB, K, N]</code></td>
|
||||||
|
<td>dtype</td>
|
||||||
|
<td>典型 ND</td>
|
||||||
|
<td>可带转置标记（isBTrans），推理场景可为 weightNz</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>偏置 bias</td>
|
||||||
|
<td><code>[B, 1, N]</code></td>
|
||||||
|
<td>dtype</td>
|
||||||
|
<td>固定 ND</td>
|
||||||
|
<td>可为空</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>输出 C</td>
|
||||||
|
<td><code>[BatchC, M, N]</code></td>
|
||||||
|
<td>dtype（可随路量化）</td>
|
||||||
|
<td>典型 ND</td>
|
||||||
|
<td>BatchC = broadcast(BatchA, BatchB)</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<p><strong>广播语义</strong>：batch 维兼容广播——每一级 batch 维上，两侧取值要么相等、要么为 1；输出该维取两者的最大值。注意两类广播的形态差异：</p>
|
||||||
|
<ul>
|
||||||
|
<li><strong>单边全广播</strong>：<code>BatchA = 1</code> 或 <code>BatchB = 1</code>（一侧整体只有 1 个 batch）；</li>
|
||||||
|
<li><strong>交叉广播</strong>：两侧均大于 1 但不同级为 1，例如 <code>BatchA = (6,1)</code>、<code>BatchB = (1,6)</code>，则 <code>BatchC = (6,6)</code> 共 36 个输出 batch。这类 case <strong>不属于</strong>单边广播，后面第 3、6 章会看到它只能落入通用分支。</li>
|
||||||
|
</ul>
|
||||||
|
<h3>1.3 case 空间：决定实现方式的特征量</h3>
|
||||||
|
<p>一个 BMM case 的实现方式完全由以下特征量决定：</p>
|
||||||
|
<p>$$
|
||||||
|
(\; B,\; M,\; N,\; K,\; \text{dtype},\; \text{layout/转置},\; \text{广播形态} \;)
|
||||||
|
$$</p>
|
||||||
|
<p>其中 $B$ = BatchC（展平后的输出 batch 数）。后续全部分析就是回答：<strong>给定这组特征量，最优实现是什么。</strong></p>
|
||||||
|
<hr />
|
||||||
|
<h2>2. 性能模型：什么叫"最优"</h2>
|
||||||
|
<h3>2.1 总时延 = 流水线最慢的一级</h3>
|
||||||
|
<p>NPU 上 BMM 的执行是核内多级硬件流水的并行——Cube 计算（MMAD）、GM/L2→L1 搬移（MTE2）、L1→L0 搬移（MTE1）、L0C 写出（Fixpipe），多核并行时各流水级时延可被双缓冲（double buffer）等机制相互掩盖，最终：</p>
|
||||||
|
<p>$$
|
||||||
|
\boxed{\;T_{total} = \max\big(T_{MMAD},\; T_{MTE2},\; T_{MTE1},\; T_{Fixpipe}\;[,\;T_{Reduce}]\big)\;}
|
||||||
|
$$</p>
|
||||||
|
<p><strong>算子优化的关键就是对瓶颈流水级的优化。</strong> 由此直接得出一个重要的设计自由度——<strong>瓶颈交换</strong>：</p>
|
||||||
|
<blockquote>
|
||||||
|
<p>当 MTE2（搬移）是瓶颈、MMAD（计算）不是瓶颈时，可以牺牲一定 MMAD 时延（例如冗余计算）换取 MTE2 性能提升；反之，当 MMAD 是瓶颈时，可以牺牲一定 MTE2 时延（例如重复搬移）换取计算效率提升。只要瓶颈级时延下降，总时延就下降。</p>
|
||||||
|
</blockquote>
|
||||||
|
<p>后文会看到：<strong>MergeBatch 就是"牺牲算力换搬移效率"的典型，ASW_Basic 切 M/N 就是"牺牲搬移（重复读）换并行度"的典型</strong>——它们的存在正当性都来自这个 max 模型。</p>
|
||||||
|
<h3>2.2 目标芯片关键规格与经验常数</h3>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>规格项</th>
|
||||||
|
<th>数值</th>
|
||||||
|
<th>对设计的意义</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td>AIC / AIV 核数</td>
|
||||||
|
<td>32 / 64（1:2）</td>
|
||||||
|
<td>核间并行度上限 $C = 32$；StreamK 归约可用 AIV</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>Cube 算力 FP16/BF16</td>
|
||||||
|
<td>≈486 TFLOPS（单核 ≈ 16×16×16 MAC/拍）</td>
|
||||||
|
<td>算存比分子</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>GM 带宽</td>
|
||||||
|
<td>1.6 TB/s（读写共享总线）</td>
|
||||||
|
<td>访存 Bound 的分母；读+写互相挤占</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>L2 Cache</td>
|
||||||
|
<td>128MB，5.2 TB/s（读写各独享）</td>
|
||||||
|
<td>重复读取的吸收层；Cube 输出写 L2 即算完成</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>L1 / L0A / L0B / L0C</td>
|
||||||
|
<td>512KB / 64KB / 64KB / 256KB（每 AIC）</td>
|
||||||
|
<td>核内分块的容量约束；L0C 按 FP32 累加计</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>UB</td>
|
||||||
|
<td>512KB（每 AIV）</td>
|
||||||
|
<td>向量通路/归约</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<p><strong>GM→L1 搬移效率（ND2NZ）达到 90%+ 带宽利用率的四条经验约束</strong>（DV100 实测，重要性递减）：</p>
|
||||||
|
<ol>
|
||||||
|
<li><strong>参与核数</strong>：建议 32 核并行搬移，至少 3/4 核（24 核）并发——核数不足时带宽利用率上限被压低；</li>
|
||||||
|
<li><strong>单核搬移数据总量</strong> ≥ 480KB（<code>min_DatamountPerCore</code>）；</li>
|
||||||
|
<li><strong>单次搬移 tile 大小</strong> ≥ 16KB（<code>min_TileSize</code>）；</li>
|
||||||
|
<li><strong>dValue</strong>（单次搬移的连续内轴字节数，如非转置 A 的 K 向、非转置 B 的 N 向）≥ 128B，建议 256B，最好 512B。</li>
|
||||||
|
</ol>
|
||||||
|
<blockquote>
|
||||||
|
<p>这些常数（480KB、16KB、256B/512B，以及后文的 $b_{thr}=4$、$k_{thr}$=32B/dtype、$minCoreNum \approx 0.8 \times 32$）都是<strong>该档芯片的实测经验值</strong>；换芯片时本文的所有逻辑结构不变，只需替换常数表。</p>
|
||||||
|
</blockquote>
|
||||||
|
<h3>2.3 算存比与 Bound 判定</h3>
|
||||||
|
<p><strong>约定</strong>：以"元素"为访存单位（1 元素 = dtype 字节），算存比单位为 FLOP/元素。</p>
|
||||||
|
<ul>
|
||||||
|
<li>case 的固有算存比（不计输出写出）：</li>
|
||||||
|
</ul>
|
||||||
|
<p>$$
|
||||||
|
AI = \frac{2MNK}{(MK + KN)} = \frac{2MN}{M + N} \quad [\text{FLOP/元素}]
|
||||||
|
$$</p>
|
||||||
|
<ul>
|
||||||
|
<li>计入输出写出（写 L2/GM）的完整算存比：</li>
|
||||||
|
</ul>
|
||||||
|
<p>$$
|
||||||
|
AI_{full} = \frac{2MNK}{MK + KN + MN} \quad [\text{FLOP/元素}]
|
||||||
|
$$</p>
|
||||||
|
<ul>
|
||||||
|
<li>芯片硬件平衡点（该 dtype 位宽下）：</li>
|
||||||
|
</ul>
|
||||||
|
<p>$$
|
||||||
|
R = \frac{\text{Cube 峰值算力}}{\text{GM 带宽} / \text{dtype 字节}} = \frac{486 \times 2}{1.6} \approx 607.5 \quad [\text{FLOP/元素，BF16/FP16}]
|
||||||
|
$$</p>
|
||||||
|
<p><strong>判定规则</strong>：$AI < R$ → 访存 Bound（瓶颈在 MTE2，优化重心是搬移效率与复用）；$AI > R$ → 计算 Bound（瓶颈在 MMAD，优化重心是 Cube 利用率与流水掩盖）。注意一个有用的事实：低位宽 dtype（FP8、MXFP4）算力同倍数翻倍而元素字节同倍数减半，$R$ 近似不变，因此该平衡点对该芯片各低位宽 dtype 通用。</p>
|
||||||
|
<p><strong>大白话</strong>：矩阵乘就像搬砖盖楼——Cube 是起重机（算力），GM 带宽是供货卡车（带宽）。算存比就是"每块砖要盖多少平米"。每块砖盖的平米数低于起重机的额定配比时，楼永远盖不快，问题出在供货而不在起重机——这时候优化要围绕"怎么少堵车、怎么一次多运点"，而不是换更大的起重机。</p>
|
||||||
|
<hr />
|
||||||
|
<h2>3. 分支划分的系统推导</h2>
|
||||||
|
<p>本章是全文逻辑主线：<strong>从分块计算的本质出发，不靠经验罗列，推导出五大分支的划分。</strong></p>
|
||||||
|
<h3>3.1 分块计算的本质与四个可切维度</h3>
|
||||||
|
<p>BMM 在 NPU 上实现的本质是：把参与计算的数据分块（tile），由 32 个 AIC 核<strong>并行 + 串行</strong>地完成这些分块的计算，再组合成最终结果：</p>
|
||||||
|
<p>$$
|
||||||
|
C[B, M, N] = \sum_{\kappa} A[B, M, K_\kappa] @ B[B, K_\kappa, N]
|
||||||
|
$$</p>
|
||||||
|
<p>分块有 4 个维度：<strong>B、M、N、K</strong>。<strong>核间怎么分这 4 个维度，就是分支划分的第一性问题</strong>（核内分块是第二性问题，属于各分支内部的 tiling）。</p>
|
||||||
|
<h3>3.2 四个维度的切分特征：一张决定性的表</h3>
|
||||||
|
<p>从"读入 / 计算 / 写出"三个视角考察每个维度的核间切分特征（这是后续一切推导的基石）：</p>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>切分维度</th>
|
||||||
|
<th>读入特征</th>
|
||||||
|
<th>计算特征</th>
|
||||||
|
<th>写出特征</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td><strong>切 B</strong></td>
|
||||||
|
<td>每个数据块只被固定的 1 个核读取，核间<strong>零重复读</strong></td>
|
||||||
|
<td>每个输出块由 1 个核独立完成，<strong>无核间依赖</strong></td>
|
||||||
|
<td>只写最终结果，无中间结果</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td><strong>切 M / 切 N</strong></td>
|
||||||
|
<td>切 M 则同一右矩阵块被多核<strong>重复读</strong>；切 N 则同一左矩阵块被多核重复读</td>
|
||||||
|
<td>每个输出块由 1 个核独立完成，无核间依赖</td>
|
||||||
|
<td>只写最终结果，无中间结果</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td><strong>切 K</strong></td>
|
||||||
|
<td>每个数据块只被固定的 1 个核读取，零重复读</td>
|
||||||
|
<td>每个输出块由<strong>多核共同</strong>完成，存在核间依赖</td>
|
||||||
|
<td><strong>有中间结果写出</strong>，需核间 Reduce 归约</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<p>为什么特征差异这么大？两条硬件层面的根本原因：</p>
|
||||||
|
<ul>
|
||||||
|
<li><strong>B 维在数学上独立</strong>（BMM 语义就是逐 batch 独立矩阵乘），所以切 B 天然零重复、零依赖；</li>
|
||||||
|
<li><strong>K 维有 L0C 累加机制</strong>：核内切 K 时，$\tilde K$ 轮 mmad 在 256KB 的 L0C 上原地累加（<code>cmatrixInitVal=false</code>），中间结果不出核。一旦把 K 切到<strong>核间</strong>，单核的 L0C 装不下"别的核算的 K 段"，部分和必须写出到 GM/L2 workspace，再由 AIV 归约——<strong>切 K 是唯一同时破坏"累加不出核"和"输出独占"两条性质的切法</strong>。</li>
|
||||||
|
</ul>
|
||||||
|
<h3>3.3 完备枚举：15 种核间切分组合</h3>
|
||||||
|
<p>4 个维度的任意非空子集都可作为一种核间切分方案，共 $2^4 - 1 = 15$ 种：</p>
|
||||||
|
<p>$$
|
||||||
|
\{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\}
|
||||||
|
$$</p>
|
||||||
|
<p>任何实现方案必属于其中一种 ⇒ 这 15 种是<strong>完备的</strong>。按 §3.2 的特征分组：</p>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>组</th>
|
||||||
|
<th>组合</th>
|
||||||
|
<th>共同特征</th>
|
||||||
|
<th>优化重心</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td>纯 B</td>
|
||||||
|
<td>{B}</td>
|
||||||
|
<td>零重复读，读写数据量固定</td>
|
||||||
|
<td>访存 Bound：搬移效率高；计算 Bound：计算效率高</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>含 M/N 不含 K</td>
|
||||||
|
<td>{M},{N},{M,N},{B,M},{B,N},{B,M,N}</td>
|
||||||
|
<td>可能有重复读</td>
|
||||||
|
<td>访存 Bound：重复读尽量少 + 搬移效率高；计算 Bound：计算效率高</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>含 K</td>
|
||||||
|
<td>{K},{B,K},{M,K},{K,N},{B,M,K},{B,K,N},{M,K,N},{B,M,K,N}</td>
|
||||||
|
<td>有中间结果写出 + 归约</td>
|
||||||
|
<td>计算效率高，且归约引入的额外 MTE2/Fixpipe 时延不能成为新瓶颈</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<h3>3.4 代价不对称性：切分维度的"价格表"</h3>
|
||||||
|
<p>15 种组合的<strong>代价结构只由两个布尔特征决定</strong>——是否含 K、是否含 M/N。三个维度的"核间切分价格"严格排序：</p>
|
||||||
|
<p>$$
|
||||||
|
\text{cost}(\text{切}B) = 0 \;<\; \text{cost}(\text{切}M/N) \;\ll\; \text{cost}(\text{切}K)
|
||||||
|
$$</p>
|
||||||
|
<ul>
|
||||||
|
<li><strong>切 B 免费</strong>：零重复读、零依赖、零中间写出；</li>
|
||||||
|
<li><strong>切 M/N 廉价但有价</strong>：共享矩阵被重复读，但若共享部分能驻留 128MB L2，重复读以 5.2TB/s 命中 L2 而非 1.6TB/s 的 GM——代价大部分被 L2 吸收；配合 swizzle（执行顺序编排）压缩同时活跃的工作集，代价进一步压低；</li>
|
||||||
|
<li><strong>切 K 昂贵</strong>：归约流量 $T_{Reduce} \propto grid_K \times$ 输出量，且引入核间同步——这是结构性代价，L2 吸收不掉。</li>
|
||||||
|
</ul>
|
||||||
|
<p><strong>"价格表"不是经验，是 BMM 语义 + L0C 累加机制 + L2/GM 带宽结构三条事实的推论。</strong> 整条分支决策树就是一句话：<strong>按价格从低到高购买并行度，买不够才加价。</strong></p>
|
||||||
|
<h3>3.5 推导主链：从价格表到五大分支</h3>
|
||||||
|
<p><strong>第 0 层：问题归约（能降维就不在 BMM 本体里解决）。</strong></p>
|
||||||
|
<ul>
|
||||||
|
<li>若 <code>BatchA = 1</code> 或 <code>BatchB = 1</code>，BMM 可通过维度折叠<strong>转化为普通 Matmul</strong>，直接复用 Matmul 的成熟优化体系（tiling、L2 切分、全载），不必在 BMM 框架内重新发明轮子 → <strong>分支一：转Matmul</strong>；</li>
|
||||||
|
<li>退化 case（K=0 无计算、K=1 无累加深度）让 Cube 完全或几乎无用，改走 AIV 向量通路——这属于"计算通路选择"，是正交于切分的前置判断（见 4.0）。</li>
|
||||||
|
</ul>
|
||||||
|
<p><strong>第 1 层：归约不掉的 case（BatchA = BatchB = B > 1，K ≥ 2），必须在 BMM 框架内组织 4 维分块。先买免费的 B 维。</strong></p>
|
||||||
|
<ul>
|
||||||
|
<li>$B \ge C = 32$ 时，切 B 就能填满核。此时分两种情形：
|
||||||
|
<ul>
|
||||||
|
<li><strong>单 batch 的 $M \times N$ 够大</strong>（能开出大 tile、L0C 利用率高、Cube 喂得饱）→ 每核逐个 batch 做完整 Matmul 就是最优，任何额外切分只增代价 → <strong>分支二：IterBatch</strong>；</li>
|
||||||
|
<li>**单 batch 的 $M \times N$ 太小**（如 64×64，L0C 利用率仅 6.25%，Cube 大片闲置）→ 把多个 batch 在核内合并成大 tile 计算（$[bM,K]@[K,bN]$，取块对角线输出），用冗余算力换 Cube 利用率与搬移效率。这类 case $AI$ 低、天然访存 Bound，浪费的算力被搬移时延掩盖（§2.1 瓶颈交换）→ <strong>分支三：MergeBatch</strong>。</li>
|
||||||
|
</ul>
|
||||||
|
</li>
|
||||||
|
</ul>
|
||||||
|
<p><strong>第 2 层：免费的 B 买不满 32 核，加价买廉价的 M/N。</strong></p>
|
||||||
|
<ul>
|
||||||
|
<li>$B < 32$ 但 $P = B \times \lceil M/16 \rceil \times \lceil N/16 \rceil \ge 32$（M/N 平面的分块补得上并行缺口）→ 切 M/N（或混合切），共享矩阵的重复读交给 L2 + swizzle 吸收。这是"所有不切 K、含共享读取切法"的通用框架 → <strong>分支四：ASW_Basic</strong>。</li>
|
||||||
|
</ul>
|
||||||
|
<p><strong>第 3 层：廉价维度也买不满，才买昂贵的 K。</strong></p>
|
||||||
|
<ul>
|
||||||
|
<li>$P < 32$（B、M、N 都小，通常 K 大）→ 唯一剩余的并行维度是 K。付归约代价换并行度，当且仅当并行收益 > 归约代价（定量条件见 4.4）→ <strong>分支五：StreamK</strong>。</li>
|
||||||
|
</ul>
|
||||||
|
<pre><code>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 买并行，付归约代价）◀──────────────────┘
|
||||||
|
</code></pre>
|
||||||
|
<h3>3.6 完备性与极小性：五大分支不多不少</h3>
|
||||||
|
<p><strong>完备性</strong>。对归约不掉的 case，15 种组合按"（是否含 K，是否含 M/N）"坍缩：</p>
|
||||||
|
<ul>
|
||||||
|
<li>含 K 的 8 种：代价结构相同（必付归约），grid_K × grid_B × grid_M × grid_N 只是参数差异 ⇒ 1 个分支（StreamK）统一覆盖；</li>
|
||||||
|
<li>不含 K 含 M/N 的 6 种：代价结构相同（共享重复读）⇒ 1 个分支（ASW_Basic）统一覆盖；</li>
|
||||||
|
<li>纯 {B} 的 1 种：核内组织只有两种本质不同的方式——<strong>合并多 batch 一起算</strong>（MergeBatch）或<strong>逐个 batch 算</strong>（IterBatch），不存在第三种 ⇒ 2 个分支。</li>
|
||||||
|
</ul>
|
||||||
|
<p>$1 + 1 + 2 = 4$，加上第 0 层的归约分支（转Matmul）= <strong>5</strong>。任何合法 case 必落其一，无空洞。</p>
|
||||||
|
<p><strong>极小性</strong>。每个分支都有它"唯一最优"的 shape 区域，去掉任何一个都会有 case 失去最优实现：</p>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>分支</th>
|
||||||
|
<th>独占最优的代表 case（BF16）</th>
|
||||||
|
<th>替代方案为何更差</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td>转Matmul</td>
|
||||||
|
<td>BatchB=1, B=128, M=N=K=2048</td>
|
||||||
|
<td>免费折叠后 Matmul 体系的 L2 切分/全载直接可用；BMM 分支内重做一遍无收益</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>IterBatch</td>
|
||||||
|
<td>B=32, M=N=K=4096</td>
|
||||||
|
<td>ASW 切 M/N 引入无谓共享读；MergeBatch 引入无谓冗余算力</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>MergeBatch</td>
|
||||||
|
<td>B=128, M=32, N=128, K=64</td>
|
||||||
|
<td>IterBatch 的 L0C 利用率仅 ~1.6%，Cube 空转，搬移 tile 碎（<16KB）带宽利用率崩</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>ASW_Basic</td>
|
||||||
|
<td>B=2, M=N=8192, K=1024</td>
|
||||||
|
<td>切 B 仅 2 核干活；StreamK 付无谓归约</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>StreamK</td>
|
||||||
|
<td>B=1, M=N=64, K=65536</td>
|
||||||
|
<td>不切 K 时 P=16 < 32，近半核闲置，时延差数量级</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<p>⇒ 五大分支构成<strong>极小完备集</strong>。</p>
|
||||||
|
<h3>3.7 大白话总结</h3>
|
||||||
|
<blockquote>
|
||||||
|
<p>把 BMM 想成给 32 个工人分一批"矩阵乘订单"：</p>
|
||||||
|
<ul>
|
||||||
|
<li><strong>能整单外包的就别自己干</strong>：一侧只有一个 batch，问题其实就是一次普通矩阵乘，直接交给成熟的 Matmul 产线（转Matmul）；</li>
|
||||||
|
<li><strong>按订单分（切 B）最省心</strong>：每人几单，互不干扰（零重复搬料、零协调）。订单多时，单大就一单接一单干（IterBatch）；单太小时一人同时干几单、拼成大活干，虽然多做了点无用功，反正瓶颈在等料不在干活（MergeBatch）；</li>
|
||||||
|
<li><strong>订单不够分，就按行/列拆单（切 M/N）</strong>：大家会重复领同一份料，但料放在近处仓库（L2），多跑几趟近仓库很便宜（ASW_Basic）；</li>
|
||||||
|
<li><strong>行/列也不够拆，只好沿深度 K 拆</strong>：几个人合做同一块输出，各算一段再汇总——汇总是要额外开会的（Reduce），所以只在实在分不满、且 K 足够长时才这么干（StreamK）。</li>
|
||||||
|
</ul>
|
||||||
|
</blockquote>
|
||||||
|
<hr />
|
||||||
|
<h2>4. 各分支详解：进入条件与实现方案</h2>
|
||||||
|
<blockquote>
|
||||||
|
<p>每节按统一结构展开：<strong>做什么（定义）→ 什么时候进（进入条件逐条 + 每条的理由）→ 怎么做（实现方案与参数求解）→ 与源码对照（批判性）</strong>。</p>
|
||||||
|
</blockquote>
|
||||||
|
<h3>4.0 前置通路层：K 退化 case</h3>
|
||||||
|
<p>在切分决策之前先做计算通路判断：</p>
|
||||||
|
<ul>
|
||||||
|
<li><strong>K = 0</strong>：无任何计算，C = bias 或 0，纯 AIV 写值；</li>
|
||||||
|
<li><strong>K = 1</strong>：退化为逐元素乘 <code>C = A ⊙ B</code>，无累加深度，Cube 的 16×16×16 粒度浪费 15/16，走 AIV 向量通路（GM→UB→Mul→GM）远优于 Cube 通路。</li>
|
||||||
|
</ul>
|
||||||
|
<p>这一层与"切分维度"正交：<strong>五大分支管理的都是 Cube 通路上的 case（K ≥ 2）</strong>。K 极小（如 K < 16，不足一个 fractal）时同理应优先考虑向量通路或尾部补齐。</p>
|
||||||
|
<h3>4.1 转Matmul（问题归约分支）</h3>
|
||||||
|
<p><strong>做什么</strong>：当 <code>BatchA = 1</code> 或 <code>BatchB = 1</code> 时，把 batch 维折叠进矩阵维度，转化为普通 Matmul，计算完成后按需恢复 batch 维。</p>
|
||||||
|
<p><strong>为什么这么做</strong>：单边 batch=1 的 BMM 与普通 Matmul 在数学上只差一个维度标签。Matmul 的优化体系（L2 切分、AL1/BL1 全载、swizzle）比 BMM 各分支在"单样本"情形下更成熟——<strong>站在巨人肩膀上，不重复造轮子</strong>。BMM 本体只研究相对 Matmul 有差异化的 case（双侧 batch > 1）。</p>
|
||||||
|
<p><strong>折叠规则与代价（关键的不对称性）</strong>：</p>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>情形</th>
|
||||||
|
<th>折叠方式</th>
|
||||||
|
<th>是否免费</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td><code>BatchB = 1</code></td>
|
||||||
|
<td>左矩阵 <code>[B, M, K]</code> 的 batch 维与 M 维在 ND 布局下<strong>内存相邻</strong>，直接视图为 <code>[B·M, K]</code>；输出 <code>[B·M, N]</code> 与 <code>[B, M, N]</code> 的内存布局逐元素一致</td>
|
||||||
|
<td><strong>完全免费</strong>：零输入重排、零输出 split</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td><code>BatchA = 1</code></td>
|
||||||
|
<td>需将右矩阵 <code>[B, K, N]</code> 折叠为 <code>[K, B·N]</code></td>
|
||||||
|
<td><strong>有代价</strong>：B 的 batch 维与 N 维在内存中不相邻（中间隔 K），折叠等价于一次 <code>[B,K,N]→[K,B,N]</code> 的转置重排（O(B·K·N) 读写），且输出 <code>[M, B·N]</code> 与目标 <code>[B, M, N]</code> 之间存在置换，需要随路 scatter</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<blockquote>
|
||||||
|
<p>这一不对称性是对"直接折叠"设想的必要修正（详见 6.2）。<code>BatchB = 1</code> 无脑走转Matmul；<code>BatchA = 1</code> 需比较"重排一次 B + Matmul 收益"与"广播友好的 IterBatch/ASW_Basic（A 驻留 L2/L1 天然共享，不重排任何数据）"——大 case 重排代价被摊薄，走转Matmul；小 case 直接走 BMM 分支。源码中的 <code>MergeBatchAndMAxis()</code> 只做了 batchB=1 方向的折叠，与"该方向免费"的判断互为印证。</p>
|
||||||
|
</blockquote>
|
||||||
|
<p><strong>进入条件</strong>：<code>BatchA = 1 || BatchB = 1</code>。其中 <code>BatchB = 1</code> 恒进；<code>BatchA = 1</code> 时按上述代价比较决定（决策规则见第 5 章）。</p>
|
||||||
|
<p><strong>BatchA = BatchB = 1</strong>：纯 Matmul，是转Matmul 的特例。</p>
|
||||||
|
<h3>4.2 MergeBatch（多 batch 合并计算）</h3>
|
||||||
|
<h4>4.2.1 做什么</h4>
|
||||||
|
<p>核间按 B 分核（每核负责 $b_{core}$ 个 batch），核间无同步无通信。核内<strong>将多个 batch 合并计算</strong>：</p>
|
||||||
|
<p>$$
|
||||||
|
[b, M, K] @ [b, K, N] \;\Rightarrow\; [bM, K] @ [K, bN] = [bM, bN] \;\xrightarrow{\text{BlockTrace}}\; [b, M, N]
|
||||||
|
$$</p>
|
||||||
|
<p>其中 <strong>BlockTrace</strong> 指以 $[M, N]$ 的 block 粒度取结果矩阵的块对角线作为各 batch 的有效输出：$C[i, m, n] = R[iM + m,; iN + n]$。交叉项（第 i 个 batch 的 A 乘第 j≠i 个 batch 的 B）被算出但丢弃——这就是"算力浪费"，浪费比例 $(b-1)/b$。</p>
|
||||||
|
<p><strong>为什么允许浪费</strong>：进入该分支的 case 必然是访存 Bound（条件三保证），瓶颈在 MTE2 不在 MMAD，浪费的算力被搬移时延掩盖（§2.1 瓶颈交换）——<strong>用本来闲置的 Cube 算力，换 tile 变大后的搬移效率与 Cube 利用率</strong>。</p>
|
||||||
|
<h4>4.2.2 进入条件（逐条 + 理由）</h4>
|
||||||
|
<p>设计原则四条：① 硬件时延可流水掩盖（double buffer 乒乓）；② 满足容量约束（每次计算 L1/L0A/L0B/L0C 放得下）；③ GM→L1 搬移高效（§2.2 四条经验约束）；④ 算力有浪费但计算不能成为瓶颈。</p>
|
||||||
|
<p>形式化后，进入 MergeBatch 需<strong>同时</strong>满足：</p>
|
||||||
|
<p><strong>条件 1（batch 够分且够合并）</strong>：<code>BatchA = BatchB</code>（无广播）且</p>
|
||||||
|
<p>$$
|
||||||
|
b_{core} = B / aicNum \ge b_{thr}
|
||||||
|
$$</p>
|
||||||
|
<p>$b_{thr}$ 是"多 batch 合并搬移能拿到效率收益"的最小合并数——$b$ 太小时合并的搬移收益抵不过实现复杂度，不如 IterBatch 且完全不浪费算力（功耗）。DV100 实测经验值 $b_{thr} = 4$。</p>
|
||||||
|
<p><strong>条件 2（合并数 b 的上下界）</strong>：设核内单次 Cube 计算合并 $b$ 个 batch（$b \le b_{core}$），要求 $b \ge b_{thr}$ 且 $b$ 由以下约束共同封顶：</p>
|
||||||
|
<ul>
|
||||||
|
<li><strong>算存比约束</strong>（计算不能变成瓶颈）：合并后单次计算的算存比 $AI(b) = \dfrac{2bMN}{M + N}$，保持访存 Bound 要求</li>
|
||||||
|
</ul>
|
||||||
|
<p>$$
|
||||||
|
AI(b) < R \;\Longleftrightarrow\; b < b_{AI} = \frac{R \cdot (M + N)}{2MN}
|
||||||
|
$$</p>
|
||||||
|
<ul>
|
||||||
|
<li><strong>L0C 容量约束</strong>：输出 $[bM, bN]$ 的 FP32 累加块须放入 256KB L0C；考虑乒乓（double buffer）则两份：</li>
|
||||||
|
</ul>
|
||||||
|
<p>$$
|
||||||
|
2 \cdot (bM)(bN) \cdot 4\text{B} \le 256\text{KB} \;\Rightarrow\; b \le b_{L0C}
|
||||||
|
$$</p>
|
||||||
|
<ul>
|
||||||
|
<li><strong>L0A/L0B 容量约束</strong>：$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}$，需联合求解。</li>
|
||||||
|
</ul>
|
||||||
|
<blockquote>
|
||||||
|
<p>乒乓取舍：若 $b_{core}$ 小于"不乒乓时 L0C 允许的合并数上限"，说明 batch 余量不足，只能不乒乓（牺牲流水掩盖换合并数）；否则乒乓。形式上 $b_{max} = \min(b_{AI}, b_{L0C}^{(\text{是否乒乓})}, b_{core})$。</p>
|
||||||
|
</blockquote>
|
||||||
|
<p><strong>条件 3（访存 Bound 性质）</strong>：</p>
|
||||||
|
<p>$$
|
||||||
|
\frac{2MN}{M + N} < \frac{R}{b}
|
||||||
|
$$</p>
|
||||||
|
<p>即 case 固有算存比显著低于平衡点（低一个合并倍数 $b$ 的量级）——这是"浪费可被掩盖"的定量保证，也是条件 2 中 $b_{AI} > b$ 的另一写法。</p>
|
||||||
|
<p><strong>条件 4（K 向搬移效率）</strong>：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 条）。</p>
|
||||||
|
<p><strong>条件 5（L1 驻留与 tile 效率）</strong>：L1 级 batch 驻留数</p>
|
||||||
|
<p>$$
|
||||||
|
b_{L1} = \frac{L1_{size}}{(M k_{L1} + k_{L1} N) \cdot \text{dtype}}, \qquad b_{L1} > b
|
||||||
|
$$</p>
|
||||||
|
<p>（L1 驻留的 batch 组必须大于单次计算的合并数，否则合并无从谈起），且搬移 tile 大小满足效率：左矩阵非转置时 $\max(b_{L1} M k_{L1} \cdot \text{dtype},; k_{L1} N \cdot \text{dtype}) > min_TileSize$（16KB）；左矩阵转置时相应调整。</p>
|
||||||
|
<p><strong>条件 6（单核搬移总量）</strong>：</p>
|
||||||
|
<p>$$
|
||||||
|
b_{core} \cdot (M k_{L1} + k_{L1} N) \cdot \text{dtype} \ge min\_DatamountPerCore = 480\text{KB}
|
||||||
|
$$</p>
|
||||||
|
<p>对应 §2.2 第 2 条：单核搬移数据总量不足时带宽利用率上限被压低。</p>
|
||||||
|
<p><strong>条件汇总逻辑</strong>：条件 1 定资格（batch 够多），条件 2/3 定上限（别算出瓶颈、别撑爆 L0），条件 4/5/6 定下限（搬移效率不能崩）。<strong>上限与下限之间必须有交集，交集为空则该 case 与 MergeBatch 无缘。</strong></p>
|
||||||
|
<h4>4.2.3 核内 b 与 baseK 的取值优化</h4>
|
||||||
|
<p>进入分支后，$b$ 与 baseK 的具体取值有一个重要性质：<strong>在访存 Bound 前提下，$b < b_{AI}$ 时 b 的取值不影响性能，baseK 只需满足分形约束（16 倍数）也不影响算存比</strong>——因为瓶颈是搬移，算力余量内怎么切都一样。据此：</p>
|
||||||
|
<ol>
|
||||||
|
<li>先由条件 2 求 $b_{max}$，实际 $b$ 在 $[b_{thr}, b_{max}]$ 内取，且<strong>尽量均匀</strong>（$b_{core}$ 整除 $b$）——每次计算的 $b$ 均匀一致对功耗更有利；</li>
|
||||||
|
<li>再由 $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 倍数。</li>
|
||||||
|
</ol>
|
||||||
|
<h4>4.2.4 数值例子（把条件过一遍）</h4>
|
||||||
|
<p>case：B=128, M=32, K=64, N=128, BF16，32 核。</p>
|
||||||
|
<ol>
|
||||||
|
<li><strong>资格</strong>：$b_{core} = 128/32 = 4 \ge b_{thr} = 4$ ✓；</li>
|
||||||
|
<li><strong>算存比</strong>：$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——算存比约束很宽；</li>
|
||||||
|
<li><strong>L0C 约束</strong>（乒乓）：$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$）——<strong>L0C 才是真正的紧约束</strong>；</li>
|
||||||
|
<li><strong>取 b = 2</strong>（= $b_{core}$ 的因子，均匀）：L0B 给出 $k_{L0} \le 64\text{KB} / (2 \times 2 \times 128 \times 2\text{B}) = 64$，$K = 64$ 一步到位，baseK = 64；</li>
|
||||||
|
<li><strong>搬移效率</strong>：$k_{L1} = 64 \ge k_{thr} = 16$ ✓；单核搬移量 $4 \times (32 \times 64 + 64 \times 128) \times 2\text{B} = 80\text{KB}$ —— <strong>不足 480KB</strong>，条件 6 不满足！</li>
|
||||||
|
</ol>
|
||||||
|
<p>结论：该 case 在"单核搬移总量"一项上不达标——这正是 MergeBatch 条件体系的用处：<strong>它提前告诉你瓶颈不在算力而在搬移效率</strong>，此 case 应尝试合并更多 K/提高 $b_{core}$ 数据量，或由时延模型与 IterBatch/ASW_Basic 比较后定夺。</p>
|
||||||
|
<h4>4.2.5 执行流程</h4>
|
||||||
|
<pre><code>核间：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
|
||||||
|
</code></pre>
|
||||||
|
<h4>4.2.6 与源码对照（批判性）</h4>
|
||||||
|
<p>源码 <code>batch_matmul_v3_mergebatch_basicapi_tiling.cpp</code> 的 IsCapable 条件：各级 batchA_i == batchB_i、<code>batchC ≥ 4 × aicNum</code>、<code>alignK ≥ 64</code>、<code>M ≤ N</code>、无 bias、非 NZ、拒绝非连续转置等。对照分析：</p>
|
||||||
|
<ul>
|
||||||
|
<li><code>batchC ≥ 4 × aicNum</code> 与草稿 $b_{core} \ge b_{thr} = 4$ <strong>一致</strong>（每核至少 4 batch）——合理；</li>
|
||||||
|
<li><code>alignK ≥ 64</code> 比草稿 $k_{thr}$=32B/dtype（BF16 为 16 元素）保守 4 倍——<strong>源码更粗</strong>，可能误杀 K∈[16,64) 的可获益 case；</li>
|
||||||
|
<li><code>M ≤ N</code> 是<strong>实现便利性约束而非原理约束</strong>——M > N 的镜像 case 在原理上同样可合并（交换合并方向），源码直接放弃，属于覆盖缺口；</li>
|
||||||
|
<li>源码<strong>没有显式的算存比/搬移效率判定</strong>（条件 3/4/6），把"是否最优"的责任推给了优先级顺序——本文的条件体系正是对此的补全。</li>
|
||||||
|
</ul>
|
||||||
|
<h3>4.3 IterBatch（逐 batch 计算）</h3>
|
||||||
|
<h4>4.3.1 做什么</h4>
|
||||||
|
<p>核间按 B 分核（每核 1 个或多个 batch），核间无同步无通信；核内<strong>逐个 batch 分别执行标准 Matmul 分块</strong>（L1→L0A/L0B→Cube→L0C→Fixpipe）并输出。无算力浪费、无跨 batch 依赖——是"切 B"最朴素的形态。</p>
|
||||||
|
<h4>4.3.2 进入条件（三大类 + 理由）</h4>
|
||||||
|
<p>满足以下<strong>任一</strong>：</p>
|
||||||
|
<p><strong>类 1（计算 Bound 型）</strong>：</p>
|
||||||
|
<p>$$
|
||||||
|
AI_{full} = \frac{2MKN}{MK + KN + MN} \ge R \quad \text{且} \quad BatchA = BatchB,\; b_{core} = B / aicNum \in \mathbb{Z}^+
|
||||||
|
$$</p>
|
||||||
|
<p>理由：连输出写出都计入仍是计算 Bound，则瓶颈恒在 Cube——只要负载均衡（$b_{core}$ 整除，各核同量），逐 batch 计算就是满算力实现，无需任何花哨。</p>
|
||||||
|
<p><strong>类 2（输出驻留 L2 的均衡型）</strong>：</p>
|
||||||
|
<p>$$
|
||||||
|
\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}^+
|
||||||
|
$$</p>
|
||||||
|
<p>理由：读侧访存不弱（$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 也不是瓶颈。三个条件合起来 = <strong>读、算、写三条路都不堵</strong>。</p>
|
||||||
|
<p><strong>类 3（访存 Bound 型）</strong>：$\dfrac{2MN}{M+N} < R$ 且同时满足：</p>
|
||||||
|
<ol>
|
||||||
|
<li>$BatchA = BatchB$ 且 $b_{core} = B / aicNum \ge 1$；</li>
|
||||||
|
<li><strong>负载均衡</strong>：访存 Bound 时核负载利用率建议 bAvg/bMax > 0.8（计算 Bound 时建议 100%）。操作化表述：B 整除核数，或尾波（B mod aicNum）活跃的核数 ≥ minCoreNum（DV100 取 0.8 × 32 ≈ 26）。注意：若只写 <code>B mod aicNum > minCoreNum</code> 会误杀整除 case（余数 0 恰恰是完全均衡），<strong>判据应以均衡比为准、余数规则为其近似</strong>；</li>
|
||||||
|
<li><strong>单核搬移不重复读</strong>（访存 Bound 下重复读就是纯损失），按 L1 容量分五种形态之一：
|
||||||
|
<ul>
|
||||||
|
<li>(a) $b_{core} = 1$ 且 $(MK + KN) \cdot \text{dtype} \le L1_{size}$：单 batch 左右矩阵同时驻留 L1，零重复读；</li>
|
||||||
|
<li>(b) $b_{core} > 1$ 且 $2(MK + KN) \cdot \text{dtype} \le L1_{size}$：L1 放下 2 个 batch 形成乒乓流水；</li>
|
||||||
|
<li>(c) 放不下的，$(MK + KN/Step) \cdot \text{dtype} \le L1_{size}$（或 M/Step 对称）：一矩阵不切、另一切分，分块大小仍须 > min_TileSize；</li>
|
||||||
|
<li>(d) $b_{core} > 1$ 时上一条的半容量版本（L1 双 batch 乒乓预算减半）；</li>
|
||||||
|
<li>(e) 左右都切 K：$(M \cdot K/Step + K/Step \cdot N) \cdot \text{dtype} \le L1_{size}$；</li>
|
||||||
|
<li>以上 (c)(d)(e) 切分后的分块均须满足搬移效率（tile ≥ 16KB、dValue ≥ 128B/256B）。</li>
|
||||||
|
</ul>
|
||||||
|
</li>
|
||||||
|
</ol>
|
||||||
|
<h4>4.3.3 设计原理（为什么是这些条件）</h4>
|
||||||
|
<ol>
|
||||||
|
<li>核间切 B 零共享零依赖，<strong>唯一的系统性风险是负载不均</strong>——所以均衡是第一条件；</li>
|
||||||
|
<li>输出须满足 L0C 容量：$MN \cdot 4\text{B} \le L0C$（否则核内还要切 M/N，那就不是纯 IterBatch 而是 ASW 行为——但注意 IterBatch 核内<strong>允许</strong>对单 batch 做标准 M/N/K tiling，此处的准确含义是"核内 tiling 不构成跨 batch 的耦合"）；</li>
|
||||||
|
<li>访存 Bound 时<strong>单核数据不得重复读</strong>——L1 装得下才不重复，装不下就按 Step 切分且切分后仍满足搬移效率下限；</li>
|
||||||
|
<li>与 MergeBatch 的分工：IterBatch 不浪费算力，但需要"单 batch 足够大"撑搬移效率与 Cube 利用率；MergeBatch 用浪费换效率，专治小 M×N。两者在 $M \times N$ 的中段重叠，由时延模型仲裁（第 5 章）。</li>
|
||||||
|
</ol>
|
||||||
|
<h4>4.3.4 执行流程</h4>
|
||||||
|
<pre><code>核间：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 回写可异步）
|
||||||
|
</code></pre>
|
||||||
|
<h3>4.4 StreamK（K 维核间切分）</h3>
|
||||||
|
<h4>4.4.1 做什么</h4>
|
||||||
|
<p>当 B、M、N 三个维度切到最碎仍填不满 32 核时，把 K 维切到核间：多核各算一段 K 的部分和，再归约：</p>
|
||||||
|
<p>$$
|
||||||
|
C_{(\beta,\mu,\nu)} = \sum_{c \in group} C^{(c)}_{(\beta,\mu,\nu)} \quad \text{（部分和写 workspace，AIV 归约或原子加）}
|
||||||
|
$$</p>
|
||||||
|
<h4>4.4.2 进入条件（两条缺一不可）</h4>
|
||||||
|
<p><strong>条件 1（并行缺口存在）</strong>：不切 K 时的独立输出块数</p>
|
||||||
|
<p>$$
|
||||||
|
P = B \times \lceil M / 16 \rceil \times \lceil N / 16 \rceil < C = 32
|
||||||
|
$$</p>
|
||||||
|
<p>（源码实现取更保守的 $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 的门槛降一档平方级）。</p>
|
||||||
|
<p><strong>条件 2（归约代价可接受）</strong>：每核计算时延须远大于归约时延（安全系数 α = 10）：</p>
|
||||||
|
<p>$$
|
||||||
|
T_{MMAD/core} \ge \alpha \cdot T_{Reduce}
|
||||||
|
\;\Longleftrightarrow\;
|
||||||
|
\frac{K}{grid_K} \;\gtrsim\; grid_K \times 1690
|
||||||
|
$$</p>
|
||||||
|
<p>即 $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 效率崩）。<strong>grid_K 越大对 K 的要求越苛刻——StreamK 内部的 grid 搜索自然淘汰归约过重的配置。</strong></p>
|
||||||
|
<h4>4.4.3 实现要点</h4>
|
||||||
|
<ul>
|
||||||
|
<li>部分和写出：L0C → workspace（GM 或 L2），归约由 AIV 执行（AIC:AIV = 1:2，2 个 AIV 伺候 1 个 AIC 的部分和流）；</li>
|
||||||
|
<li>原子加（AtomicAdd）变体省一遍读回，但破坏确定性——源码在确定性等级 > 1 时禁用 StreamK，工程上必须保留这个开关；</li>
|
||||||
|
<li>核间同步与任务编排建议直接复用 Matmul 高阶 API 的 StreamK 模板（避免手写 CrossCore flagId 冲突）。</li>
|
||||||
|
</ul>
|
||||||
|
<h4>4.4.4 与不切 K 分支的关系</h4>
|
||||||
|
<p>StreamK 从不"硬切换"进入：当不切 K 的分支候选已足够快（如已 Cube Bound），StreamK 候选的归约开销使其自然落败；只有当 B/M/N 并行度不足导致大量核闲置时，StreamK 才以数量级优势胜出（例：B=1、M=N=64、K=65536 时，不切 K 仅 16 核可用，StreamK 32 核满负荷）。</p>
|
||||||
|
<h3>4.5 ASW_Basic（通用切分框架）</h3>
|
||||||
|
<h4>4.5.1 做什么</h4>
|
||||||
|
<p>不切 K，允许切 B/M/N 的<strong>任意组合</strong>——它是"不含 K 且含共享读取"的 6 种切分组合（{M},{N},{M,N},{B,M},{B,N},{B,M,N}）的统一实现框架，也是 B < 32 但 P ≥ 32 时的最优归宿，同时兜住 B ≥ 32 但 IterBatch/MergeBatch 条件不满足的剩余 case。<strong>B 可以大、可以小、可以等于 1；ASW_Basic 是实践中最常命中的分支。</strong></p>
|
||||||
|
<h4>4.5.2 核心机制：swizzle + L2 管理</h4>
|
||||||
|
<p>切 M/N 的固有代价是共享矩阵的重复读，ASW_Basic 用两件武器把代价压到最低：</p>
|
||||||
|
<ol>
|
||||||
|
<li><strong>ASW 滑窗蛇形 swizzle</strong>：把 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 足迹越小，且因子性保证整窗被核数均分、窗口边界不碎；</li>
|
||||||
|
<li><strong>L2 切分</strong>：工作集超过 128MB 时，按 mL2TileNum × nL2TileNum 切分，每个 L2 块<strong>错位分核</strong>（对角线分配），避免多核同时抢同一地址的读读冲突，并优先选拖尾小的方案。</li>
|
||||||
|
</ol>
|
||||||
|
<h4>4.5.3 核间切分维度的选择顺序</h4>
|
||||||
|
<p>按共享代价从低到高：</p>
|
||||||
|
<ol>
|
||||||
|
<li><strong>切 B</strong>（B ≥ 核数）：零共享，永远先试；</li>
|
||||||
|
<li><strong>切 M</strong>（B 不够）：右矩阵 [K, N] 被共享，若 $KN \cdot \text{dtype} \le 128$MB 则驻留 L2，零 GM 重复读；</li>
|
||||||
|
<li><strong>切 N</strong>：对称；</li>
|
||||||
|
<li><strong>混合切</strong>（B×M、M×N、…）：双向共享，靠 swizzle + L2 切分管理；</li>
|
||||||
|
<li><strong>降核</strong>：P 远小于 32 且 K 也不够格走 StreamK 时，宁可部分核闲置（小 case 时延绝对值小，调度开销反而主导）。</li>
|
||||||
|
</ol>
|
||||||
|
<h4>4.5.4 内部特化（不是独立分支）</h4>
|
||||||
|
<ul>
|
||||||
|
<li><strong>AL1/BL1 全载</strong>：当单边无 batch 且该侧矩阵很小（如 batchA=1 且 M ≤ 256），把整个 A 常驻 L1，$\tilde M = 1$、A 的 GM→L1 搬运只发生一次——这是 ASW 内部"M 不切 + L1 驻留深度拉满"的参数极限，不是新分支；</li>
|
||||||
|
<li><strong>广播友好形态</strong>：交叉广播 case（§1.2）在 ASW 内通过对广播侧做 L1/L2 驻留处理，共享关系与切 M/N 同构。</li>
|
||||||
|
</ul>
|
||||||
|
<hr />
|
||||||
|
<h2>5. 分支仲裁与整体决策流程</h2>
|
||||||
|
<h3>5.1 为什么需要仲裁层</h3>
|
||||||
|
<p>第 4 章的进入条件给出的是各分支的"主场"，但主场之间有<strong>重叠区</strong>（例如 B ≥ 32 且 M×N 中等时，MergeBatch 与 IterBatch 都合法；B 较大时 IterBatch 与 ASW_Basic 切 B 等效）。重叠区的最优归属没有解析解，必须由统一的时延模型仲裁：</p>
|
||||||
|
<p>$$
|
||||||
|
\text{branch}^* = \arg\min_{cand \in \bigcup \text{各分支候选}} \max\big(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{Fixpipe}[, T_{Reduce}]\big)
|
||||||
|
$$</p>
|
||||||
|
<p>分支体系的价值在于<strong>候选集完备且无冗余</strong>：每个等价类只派一个代表框架生成候选，时延模型在等价类内部和边界上做精细仲裁。两层缺一不可——只有模型没有分支，搜索空间是 15 种组合 × 全部 grid 参数的爆炸；只有分支没有模型，边界 case 被硬阈值误杀。</p>
|
||||||
|
<h3>5.2 总决策流程</h3>
|
||||||
|
<pre><code>输入：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 切分方案 + 预估端到端时延
|
||||||
|
</code></pre>
|
||||||
|
<hr />
|
||||||
|
<h2>6. 五大分支设计的完备性审视（漏洞检查）</h2>
|
||||||
|
<p>按"转Matmul / MergeBatch / IterBatch / StreamK / ASW_Basic 五大分支 cover 全部 BMM case 最优实现"的设计目标，逐条拷问。</p>
|
||||||
|
<h3>6.1 覆盖矩阵：无空洞</h3>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>shape 区域</th>
|
||||||
|
<th>主分支</th>
|
||||||
|
<th>兜底</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td>K = 0 / K = 1</td>
|
||||||
|
<td>前置通路层（AIV）</td>
|
||||||
|
<td>—</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>BatchA=1 或 BatchB=1</td>
|
||||||
|
<td>转Matmul</td>
|
||||||
|
<td>ASW_Basic（广播友好形态）</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>BatchA=BatchB>1，B≥32，M×N 大</td>
|
||||||
|
<td>IterBatch</td>
|
||||||
|
<td>ASW_Basic 切 B</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>BatchA=BatchB>1，B≥32，M×N 小，访存 Bound</td>
|
||||||
|
<td>MergeBatch</td>
|
||||||
|
<td>IterBatch</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>上述但 B<32，P≥32</td>
|
||||||
|
<td>ASW_Basic</td>
|
||||||
|
<td>降核 ASW</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>P<32，K 大</td>
|
||||||
|
<td>StreamK</td>
|
||||||
|
<td>降核 ASW</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>P<32，K 也小</td>
|
||||||
|
<td>降核 ASW_Basic</td>
|
||||||
|
<td>（此时时延绝对值小，调度开销主导，分支选择不敏感）</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>交叉广播（BatchA≠BatchB 且均>1）</td>
|
||||||
|
<td>ASW_Basic</td>
|
||||||
|
<td>—</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<h3>6.2 审视发现的问题与修正</h3>
|
||||||
|
<p>设计骨架成立，但有五处需要修正或显式声明，否则会被反例拷问：</p>
|
||||||
|
<p><strong>问题 1：转Matmul 的方向不对称（原设想的最大漏洞）。</strong> "一侧 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 驻留共享，本身已接近最优）。</p>
|
||||||
|
<p><strong>问题 2：交叉广播 case 的归属必须显式声明。</strong> BatchA=(6,1)、BatchB=(1,6) 这类 case 既不进转Matmul（两侧都>1），也不满足 MergeBatch/IterBatch 的 BatchA=BatchB 前提——它们落入 ASW_Basic（对广播侧做驻留）。不声明就会显得"五分支有洞"。声明后完备性闭合。</p>
|
||||||
|
<p><strong>问题 3：K 退化不在五分支内，需前置通路层。</strong> K=0（无计算）、K=1（无累加深度）走 AIV 向量通路；这是"计算通路选择"维度，与"切分维度"正交，放在五分支之前判断。源码用 K_EQUAL_ZERO / TO_MUL 两个独立策略处理，印证了这一层的必要性——但它们是通路切换而非新切分等价类。</p>
|
||||||
|
<p><strong>问题 4：IterBatch 草稿条件 3(2) 的字面表述有误杀。</strong> <code>B mod aicNum > minCoreNum</code> 会拒绝 B 整除核数的完全均衡 case（余数 0）。正确表述应以负载均衡比为准：访存 Bound 要求 bAvg/bMax > 0.8，计算 Bound 要求 100%；余数规则只是它的近似操作化。</p>
|
||||||
|
<p><strong>问题 5：经验常数必须可标定，且源码硬编码条件偏粗。</strong> $b_{thr}=4$、$k_{thr}$=32B/dtype、480KB、16KB、dValue≥256B 都是 DV100 档芯片的实测经验值，文档中全部显式参数化（换芯片只换常数表）。对照之下源码把 <code>batchC ≥ 4×aicNum</code>、<code>alignK ≥ 64</code>、<code>M ≤ N</code> 等直接写死，其中 <code>alignK ≥ 64</code> 比草稿阈值保守 4 倍、<code>M ≤ N</code> 放弃了镜像 case——<strong>做最优实现时不应继承这些硬编码，而应按第 4 章的条件体系重新判定</strong>。</p>
|
||||||
|
<h3>6.3 与源码 10 策略的关系</h3>
|
||||||
|
<p>源码 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）与本文五分支不是同层概念：</p>
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>
|
||||||
|
<th>源码策略</th>
|
||||||
|
<th>本文归属</th>
|
||||||
|
</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
<tr>
|
||||||
|
<td>K_EQUAL_ZERO、TO_MUL</td>
|
||||||
|
<td>前置通路层（4.0），正交于切分</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>STREAM_K</td>
|
||||||
|
<td>StreamK</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>MERGE_BATCH</td>
|
||||||
|
<td>MergeBatch</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>ITER_BATCH、ITER_BATCH_BROADCAST</td>
|
||||||
|
<td>IterBatch（broadcast 变体是广播输入下的数据复用特化；单边广播在本文更多由转Matmul 吸收）</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>AL1_FULL_LOAD、BL1_FULL_LOAD</td>
|
||||||
|
<td>ASW_Basic 的 L1 驻留参数极限（4.5.4）</td>
|
||||||
|
</tr>
|
||||||
|
<tr>
|
||||||
|
<td>ASW_BASIC、BASE</td>
|
||||||
|
<td>ASW_Basic</td>
|
||||||
|
</tr>
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
<p><strong>结论</strong>：源码策略集 = 五分支 × 计算通路 × 驻留/广播特化的展开。五分支是"切分等价类"的最小完备骨架，源码的冗余策略可在新实现中收敛，源码遗漏的（条件 3/4/6 的显式判定、M>N 的 MergeBatch 镜像、转Matmul 的 BatchA=1 方向）正是优化的增量空间。</p>
|
||||||
|
<hr />
|
||||||
|
<h2>7. 总结</h2>
|
||||||
|
<ol>
|
||||||
|
<li><strong>最优的定义</strong>：$T_{total} = \max(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{Fixpipe}[, T_{Reduce}])$ 最小；瓶颈交换是合法且必要的手段。</li>
|
||||||
|
<li><strong>分支的推导</strong>：4 维可切 → 15 种组合完备 → 切分特征表（切 B 免费 / 切 M/N 廉价被 L2 吸收 / 切 K 昂贵需归约）→ 按价格从低到高购买并行度 → 坍缩为 MergeBatch、IterBatch、ASW_Basic、StreamK 四个等价类，加前置的转Matmul 问题归约层，共五大分支，<strong>完备且极小</strong>。</li>
|
||||||
|
<li><strong>各分支的条件</strong>都不是孤立经验，而是"资格（并行度够不够）→ 上限（容量/算存比封顶）→ 下限（搬移效率托底）"三层逻辑的实例化；所有经验常数可标定、可移植。</li>
|
||||||
|
<li><strong>边界不靠硬切</strong>：重叠区由端到端时延模型统一仲裁，分支体系负责候选集的完备无冗余。</li>
|
||||||
|
<li><strong>设计经受了漏洞拷问</strong>：修正了转Matmul 的方向不对称与 IterBatch 均衡条件的字面误杀，显式声明了交叉广播与 K 退化的归属；相对源码实现，本文条件体系更细、更真、覆盖更全。</li>
|
||||||
|
</ol>
|
||||||
|
<hr />
|
||||||
|
<p><em>参考文档：《BMM分块计算数学公式》《BMM最优软件实现方案设计》《BatchMatMulV3算子分支实现分析》《BMM从分块计算到四大分支的逻辑推导》；源码：<code>ops-nn/matmul/batch_mat_mul_v3</code>（arch35）；芯片资料：昇腾 950PR 架构白皮书与 CANN 9.0.0 性能建模文档。</em></p>
|
||||||
|
|
||||||
|
</div>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
618
BatchMatmul算子特性分析.md
Normal file
618
BatchMatmul算子特性分析.md
Normal file
@@ -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 性能建模文档。*
|
||||||
Reference in New Issue
Block a user