添加 BMM算子优化分析_v0.4.html
This commit is contained in:
659
BMM算子优化分析_v0.4.html
Normal file
659
BMM算子优化分析_v0.4.html
Normal file
@@ -0,0 +1,659 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>BMM 算子优化分析 v0.4 — 昇腾 950PR</title>
|
||||
<script>
|
||||
MathJax = {
|
||||
tex: {
|
||||
inlineMath: [['$','$']],
|
||||
displayMath: [['$$','$$']],
|
||||
tags: 'ams',
|
||||
processEscapes: true
|
||||
}
|
||||
};
|
||||
</script>
|
||||
<script id="MathJax-script" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
|
||||
<style>
|
||||
:root{--ink:#1f2933;--muted:#5f6b7a;--accent:#0b6bcb;--accent2:#0e9f6e;--warn:#b45309;--line:#d9e2ec;--code-bg:#f4f6f9}
|
||||
*{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}
|
||||
h4{font-size:16px;margin-top:20px;color:#123}
|
||||
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}
|
||||
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}
|
||||
.math{background:#fafbfc;border:1.5px solid #c3d6ee;border-radius:10px;padding:10px 22px;margin:14px 0;overflow-x:auto}
|
||||
ul.tight li,ol.tight li{margin:3px 0}
|
||||
hr{border:none;border-top:1px solid var(--line);margin:24px 0}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="page">
|
||||
<h1>BMM 算子优化分析</h1>
|
||||
<blockquote>版本:v0.4 | 目标芯片:昇腾 950PR(DAV_3510,32 AIC / 64 AIV,GM 1.6TB/s,L2 128MB/5.2TB/s)<br>基于 issue#3 (v0.3) 完善:补全 StreamK/ASW_Basic 实现方案,精化全部分支的进入条件与实现步骤,加入源码对照与硬件依据</blockquote>
|
||||
<hr>
|
||||
<h2>一、算子功能与接口说明</h2>
|
||||
<h3>1.1 算子功能</h3>
|
||||
<p>Batch Matmul(BMM)完成<b>带 batch 维的矩阵乘</b>:对 batch 维的每个索引独立执行一次矩阵乘,再加可选偏置:</p>
|
||||
<div class="math">$$
|
||||
C[b, m, n] = \sum_{k=0}^{K-1} A[b, m, k] \cdot B[b, k, n] + bias[b, 1, n]
|
||||
$$</div>
|
||||
<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><tr><th>参数</th><th>形状</th><th>数据类型</th><th>layout</th><th>说明</th></tr>
|
||||
<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></table>
|
||||
<p><b>广播语义</b>:batch 维兼容广播——每一级 batch 维上,两侧取值要么相等、要么为 1;输出该维取两者的最大值。两类广播的形态差异:</p>
|
||||
<ul class="tight">
|
||||
<li><b>单边全广播</b>:<code>BatchA = 1</code> 或 <code>BatchB = 1</code>(一侧整体只有 1 个 batch);</li>
|
||||
<li><b>交叉广播</b>:两侧均大于 1 但不同级为 1,例如 <code>BatchA = (6,1)</code>、<code>BatchB = (1,6)</code>,则 <code>BatchC = (6,6)</code> 共 36 个输出 batch。这类 case <b>不属于</b>单边广播,只能落入通用分支。</li>
|
||||
</ul>
|
||||
<h3>1.3 case 空间</h3>
|
||||
<p>一个 BMM case 的实现方式完全由以下特征量决定:</p>
|
||||
<div class="math">$$
|
||||
(\; B,\; M,\; N,\; K,\; \text{dtype},\; \text{layout/转置},\; \text{广播形态} \;)
|
||||
$$</div>
|
||||
<p>其中 $B$ = BatchC(展平后的输出 batch 数)。后续全部分析就是回答:<b>给定这组特征量,最优实现是什么。</b></p>
|
||||
<hr>
|
||||
<h2>二、性能模型:什么叫"最优"</h2>
|
||||
<h3>2.1 总时延 = 流水线最慢的一级</h3>
|
||||
<p>NPU 上 BMM 的执行是核内多级硬件流水的并行——Cube 计算(MMAD)、GM/L2→L1 搬移(MTE2)、L1→L0 搬移(MTE1)、L0C 写出(Fixpipe)。多核并行时各流水级时延可被双缓冲(double buffer)等机制相互掩盖,最终:</p>
|
||||
<div class="math">$$
|
||||
\boxed{\;T_{total} = \max\big(T_{MMAD},\; T_{MTE2},\; T_{MTE1},\; T_{Fixpipe}\;[,\;T_{Reduce}]\big)\;}
|
||||
$$</div>
|
||||
<p><b>总时延 = 流水线最慢的一级。</b> 算子优化的关键就是对瓶颈流水级的优化。由此得出一个重要的设计自由度——<b>瓶颈交换</b>:</p>
|
||||
<blockquote>当 MTE2(搬移)是瓶颈、MMAD(计算)不是瓶颈时,可以牺牲一定 MMAD 时延(例如冗余计算)换取 MTE2 性能提升;反之,当 MMAD 是瓶颈时,可以牺牲一定 MTE2 时延(例如重复搬移)换取计算效率提升。只要瓶颈级时延下降,总时延就下降。</blockquote>
|
||||
<p>后文会看到:<b>MergeBatch 就是"牺牲算力换搬移效率"的典型,ASW_Basic 切 M/N 就是"牺牲搬移(重复读)换并行度"的典型</b>——它们的存在正当性都来自这个 max 模型。</p>
|
||||
<h3>2.2 目标芯片关键规格与经验常数</h3>
|
||||
<table><tr><th>规格项</th><th>数值</th><th>对设计的意义</th></tr>
|
||||
<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/拍 × 1.65GHz)</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></table>
|
||||
<p><b>GM→L1 搬移效率(ND2NZ)达到 90%+ 带宽利用率的四条经验约束</b>(重要性递减):</p>
|
||||
<ol class="tight">
|
||||
<li><b>参与核数</b>:建议 32 核并行搬移,至少 3/4 核(24 核)并发——核数不足时带宽利用率上限被压低;</li>
|
||||
<li><b>单核搬移数据总量</b> ≥ 480KB(<code>min_DatamountPerCore</code>);</li>
|
||||
<li><b>单次搬移 tile 大小</b> ≥ 16KB(<code>min_TileSize</code>);</li>
|
||||
<li><b>dValue</b>(单次搬移的连续内轴字节数,如非转置 A 的 K 向、非转置 B 的 N 向)≥ 128B,建议 256B,最好 512B。</li>
|
||||
</ol>
|
||||
<blockquote>这些常数(480KB、16KB、256B/512B,以及后文的 $b_{thr}=4$、$k_{thr}$=32B/dtype、$minCoreNum \approx 0.8 \times 32$)都是<b>该档芯片的实测经验值</b>;换芯片时所有逻辑结构不变,只需替换常数表。</blockquote>
|
||||
<h3>2.3 算存比与 Bound 判定</h3>
|
||||
<p><b>约定</b>:以"元素"为访存单位(1 元素 = dtype 字节),算存比单位为 FLOP/元素。</p>
|
||||
<ul class="tight">
|
||||
<li>case 的固有算存比(不计输出写出):</li>
|
||||
</ul>
|
||||
<div class="math">$$
|
||||
AI = \frac{2MNK}{(MK + KN)} = \frac{2MN}{M + N} \quad [\text{FLOP/元素}]
|
||||
$$</div>
|
||||
<ul class="tight">
|
||||
<li>计入输出写出(写 L2/GM)的完整算存比:</li>
|
||||
</ul>
|
||||
<div class="math">$$
|
||||
AI_{full} = \frac{2MNK}{MK + KN + MN} \quad [\text{FLOP/元素}]
|
||||
$$</div>
|
||||
<ul class="tight">
|
||||
<li>芯片硬件平衡点(该 dtype 位宽下):</li>
|
||||
</ul>
|
||||
<div class="math">$$
|
||||
R = \frac{\text{Cube 峰值算力}}{\text{GM 带宽} / \text{dtype 字节}} = \frac{486 \times 2}{1.6} \approx 607.5 \quad [\text{FLOP/元素,BF16/FP16}]
|
||||
$$</div>
|
||||
<p><b>判定规则</b>:$AI < R$ → <b>访存 Bound</b>(瓶颈在 MTE2,优化重心是搬移效率与复用);$AI > R$ → <b>计算 Bound</b>(瓶颈在 MMAD,优化重心是 Cube 利用率与流水掩盖)。</p>
|
||||
<blockquote><b>大白话</b>:矩阵乘就像搬砖盖楼——Cube 是起重机(算力),GM 带宽是供货卡车(带宽)。算存比就是"每块砖要盖多少平米"。每块砖盖的平米数低于起重机的额定配比时,楼永远盖不快,问题出在供货而不在起重机——这时候优化要围绕"怎么少堵车、怎么一次多运点",而不是换更大的起重机。</blockquote>
|
||||
<hr>
|
||||
<h2>三、分支划分的系统推导</h2>
|
||||
<p>本章是全文逻辑主线:<b>从分块计算的本质出发,不靠经验罗列,推导出六大分支的划分。</b></p>
|
||||
<h3>3.1 分块计算的本质与四个可切维度</h3>
|
||||
<p>BMM 在 NPU 上实现的本质是:把参与计算的数据分块(tile),由 32 个 AIC 核<b>并行 + 串行</b>地完成这些分块的计算,再组合成最终结果:</p>
|
||||
<div class="math">$$
|
||||
C[B, M, N] = \sum_{\kappa} A[B, M, K_\kappa] \cdot B[B, K_\kappa, N]
|
||||
$$</div>
|
||||
<p>分块有 4 个维度:<b>B、M、N、K</b>。<b>核间怎么分这 4 个维度,就是分支划分的第一性问题</b>(核内分块是第二性问题,属于各分支内部的 tiling)。</p>
|
||||
<h3>3.2 四个维度的切分特征:一张决定性的表</h3>
|
||||
<p>从"读入 / 计算 / 写出"三个视角考察每个维度的核间切分特征(这是后续一切推导的基石):</p>
|
||||
<table><tr><th>切分维度</th><th>读入特征</th><th>计算特征</th><th>写出特征</th></tr>
|
||||
<tr><td><b>切 B</b></td><td>每个数据块只被固定的 1 个核读取,核间<b>零重复读</b></td><td>每个输出块由 1 个核独立完成,<b>无核间依赖</b></td><td>只写最终结果,无中间结果</td></tr>
|
||||
<tr><td><b>切 M / 切 N</b></td><td>切 M 则同一右矩阵块被多核<b>重复读</b>;切 N 则同一左矩阵块被多核重复读</td><td>每个输出块由 1 个核独立完成,无核间依赖</td><td>只写最终结果,无中间结果</td></tr>
|
||||
<tr><td><b>切 K</b></td><td>每个数据块只被固定的 1 个核读取,零重复读</td><td>每个输出块由<b>多核共同</b>完成,存在核间依赖</td><td><b>有中间结果写出</b>,需核间 Reduce 归约</td></tr></table>
|
||||
<p>为什么特征差异这么大?两条硬件层面的根本原因:</p>
|
||||
<ul class="tight">
|
||||
<li><b>B 维在数学上独立</b>(BMM 语义就是逐 batch 独立矩阵乘),所以切 B 天然零重复、零依赖;</li>
|
||||
<li><b>K 维有 L0C 累加机制</b>:核内切 K 时,$\tilde{K}$ 轮 mmad 在 256KB 的 L0C 上原地累加(<code>cmatrixInitVal=false</code>),中间结果不出核。一旦把 K 切到<b>核间</b>,单核的 L0C 装不下"别的核算的 K 段",部分和必须写出到 GM/L2 workspace,再由 AIV 归约——<b>切 K 是唯一同时破坏"累加不出核"和"输出独占"两条性质的切法</b>。</li>
|
||||
</ul>
|
||||
<h3>3.3 完备枚举:15 种核间切分组合</h3>
|
||||
<p>4 个维度的任意非空子集都可作为一种核间切分方案,共 $2^4 - 1 = 15$ 种:</p>
|
||||
<div class="math">$$
|
||||
\{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\}
|
||||
$$</div>
|
||||
<p>任何实现方案必属于其中一种 ⇒ 这 15 种是<b>完备的</b>。按 §3.2 的特征分组:</p>
|
||||
<table><tr><th>组</th><th>组合</th><th>共同特征</th><th>优化重心</th></tr>
|
||||
<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></table>
|
||||
<h3>3.4 代价不对称性:切分维度的"价格表"</h3>
|
||||
<p>15 种组合的<b>代价结构只由两个布尔特征决定</b>——是否含 K、是否含 M/N。三个维度的"核间切分价格"严格排序:</p>
|
||||
<div class="math">$$
|
||||
\text{cost}(\text{切}B) = 0 \;<\; \text{cost}(\text{切}M/N) \;\ll\; \text{cost}(\text{切}K)
|
||||
$$</div>
|
||||
<ul class="tight">
|
||||
<li><b>切 B 免费</b>:零重复读、零依赖、零中间写出;</li>
|
||||
<li><b>切 M/N 廉价但有价</b>:共享矩阵被重复读,但若共享部分能驻留 128MB L2,重复读以 5.2TB/s 命中 L2 而非 1.6TB/s 的 GM——代价大部分被 L2 吸收;配合 swizzle(执行顺序编排)压缩同时活跃的工作集,代价进一步压低;</li>
|
||||
<li><b>切 K 昂贵</b>:归约流量 $T_{Reduce} \propto grid_K \times$ 输出量,且引入核间同步——这是结构性代价,L2 吸收不掉。</li>
|
||||
</ul>
|
||||
<p><b>"价格表"不是经验,是 BMM 语义 + L0C 累加机制 + L2/GM 带宽结构三条事实的推论。</b> 整条分支决策树就是一句话:<b>按价格从低到高购买并行度,买不够才加价。</b></p>
|
||||
<h3>3.5 推导主链:从价格表到六大分支</h3>
|
||||
<p><b>第 0 层:问题归约与通路选择。</b></p>
|
||||
<ul class="tight">
|
||||
<li><b>K = 0</b>:无任何计算,C = bias 或 0,纯 AIV 写值 → <b>分支一:特殊分支(K=0)</b>;</li>
|
||||
<li><b>K = 1</b>:退化为逐元素乘 $C = A \odot B$,无累加深度,Cube 的 16×16×16 粒度浪费 15/16 → <b>分支二:特殊分支(K=1)</b>;</li>
|
||||
<li><b>BatchA = 1 或 BatchB = 1</b>:BMM 可通过维度折叠<b>转化为普通 Matmul</b>,复用 Matmul 的成熟优化体系(tiling、L2 切分、全载)→ <b>分支三:转Matmul</b>。</li>
|
||||
</ul>
|
||||
<p><b>第 1 层:归约不掉的 case(BatchA = BatchB = B > 1,K ≥ 2),必须在 BMM 框架内组织 4 维分块。先买免费的 B 维。</b></p>
|
||||
<ul class="tight">
|
||||
<li>$B \ge C = 32$ 时,切 B 就能填满核。此时分两种情形:</li>
|
||||
<ul class="tight">
|
||||
<li>**单 batch 的 $M \times N$ 够大<b>(能开出大 tile、L0C 利用率高、Cube 喂得饱)→ 每核逐个 batch 做完整 Matmul 就是最优 → </b>分支四:IterBatch**;</li>
|
||||
<li>**单 batch 的 $M \times N$ 太小**(如 64×64,L0C 利用率仅 6.25%,Cube 大片闲置)→ 把多个 batch 在核内合并成大 tile 计算($[bM,K]@[K,bN]$,取块对角线输出),用冗余算力换 Cube 利用率与搬移效率 → <b>分支五:MergeBatch</b>。</li>
|
||||
</ul>
|
||||
</ul>
|
||||
<p><b>第 2 层:免费的 B 买不满 32 核,加价买廉价的 M/N。</b></p>
|
||||
<ul class="tight">
|
||||
<li>$B < 32$ 但 $P = B \times \lceil M/16 \rceil \times \lceil N/16 \rceil \ge 32$(M/N 平面的分块补得上并行缺口)→ 切 M/N(或混合切),共享矩阵的重复读交给 L2 + swizzle 吸收 → <b>分支六:ASW_Basic</b>。</li>
|
||||
</ul>
|
||||
<p><b>第 3 层:廉价维度也买不满,才买昂贵的 K。</b></p>
|
||||
<ul class="tight">
|
||||
<li>$P < 32$(B、M、N 都小,通常 K 大)→ 唯一剩余的并行维度是 K。付归约代价换并行度 → <b>分支七:StreamK</b>。</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(问题归约)
|
||||
│ 否(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>
|
||||
<p>由此推导出<b>六大 BMM 本体分支</b>(转Matmul、MergeBatch、IterBatch、ASW_Basic、StreamK)加一个<b>前置通路层</b>(K 退化),共七个决策路径。</p>
|
||||
<h3>3.6 完备性与极小性:不多不少</h3>
|
||||
<p><b>完备性</b>。对归约不掉的 case,15 种组合按"(是否含 K,是否含 M/N)"坍缩:</p>
|
||||
<ul class="tight">
|
||||
<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 种:核内组织只有两种本质不同的方式——<b>合并多 batch 一起算</b>(MergeBatch)或<b>逐个 batch 算</b>(IterBatch),不存在第三种 ⇒ 2 个分支。</li>
|
||||
</ul>
|
||||
<p>$1 + 1 + 2 = 4$,加上前置的归约分支(转Matmul)与通路层(K 退化)= <b>7</b>。任何合法 case 必落其一,无空洞。</p>
|
||||
<p><b>极小性</b>。每个分支都有它"唯一最优"的 shape 区域,去掉任何一个都会有 case 失去最优实现:</p>
|
||||
<table><tr><th>分支</th><th>独占最优的代表 case(BF16)</th><th>替代方案为何更差</th></tr>
|
||||
<tr><td>特殊(K=0/1)</td><td>K=0 或 K=1</td><td>Cube 通路完全或几乎无用</td></tr>
|
||||
<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></table>
|
||||
<p>⇒ 七大路径构成<b>极小完备集</b>。</p>
|
||||
<h3>3.7 大白话总结</h3>
|
||||
<blockquote>把 BMM 想成给 32 个工人分一批"矩阵乘订单":<br>- <b>K=0/1 的订单太简单,不需要起重机</b>:Cube 派不上用场,直接让搬运工(AIV)干完(特殊分支);<br>- <b>能整单外包的就别自己干</b>:一侧只有一个 batch,问题其实就是一次普通矩阵乘,直接交给成熟的 Matmul 产线(转Matmul);<br>- <b>按订单分(切 B)最省心</b>:每人几单,互不干扰(零重复搬料、零协调)。订单多时,单大就一单接一单干(IterBatch);单太小时一人同时干几单、拼成大活干,虽然多做了点无用功,反正瓶颈在等料不在干活(MergeBatch);<br>- <b>订单不够分,就按行/列拆单(切 M/N)</b>:大家会重复领同一份料,但料放在近处仓库(L2),多跑几趟近仓库很便宜(ASW_Basic);<br>- <b>行/列也不够拆,只好沿深度 K 拆</b>:几个人合做同一块输出,各算一段再汇总——汇总是要额外开会的(Reduce),所以只在实在分不满、且 K 足够长时才这么干(StreamK)。</blockquote>
|
||||
<hr>
|
||||
<h2>四、各分支详解:进入条件与实现方案</h2>
|
||||
<blockquote>每节按统一结构展开:<b>做什么(定义)→ 什么时候进(进入条件逐条 + 每条的理由)→ 怎么做(实现方案与参数求解)→ 与源码对照(批判性)</b>。</blockquote>
|
||||
<h3>4.0 前置通路层:K 退化 case</h3>
|
||||
<p>在切分决策之前先做计算通路判断:</p>
|
||||
<table><tr><th>K 值</th><th>数学本质</th><th>计算通路</th><th>数据通路</th><th>核类型</th></tr>
|
||||
<tr><td><b>K = 0</b></td><td>无任何计算,C = bias 或 0</td><td>无 mmad</td><td>GM → AIV → GM</td><td>AIV_ONLY</td></tr>
|
||||
<tr><td><b>K = 1</b></td><td>退化为逐元素乘 $C = A \odot B$,无累加深度</td><td>无 mmad,向量 Mul</td><td>GM → UB → Mul → GM</td><td>AIV_ONLY</td></tr></table>
|
||||
<p><b>为什么是 AIV 而非 AIC</b>:Cube 阵列一拍完成 16×16×16(fp16)的 fractal 乘加。K=1 时 K 维只有 1/16 被利用——阵列 15/16 的 MAC 空转,还要付出 GM→L1→L0A/L0B 的分形搬运、L0C 累加与 fixpipe 写出全链路开销。而 AIV 每拍处理 256B 连续数据,逐元素乘 + UB 多 batch pingpong 能把 MTE 带宽吃满;64 个 AIV 的聚合向量算力(FP16 54 TFLOPS)对这个计算密度绰绰有余。</p>
|
||||
<p><b>源码对照</b>:K_EQUAL_ZERO(策略 0)与 TO_MUL(策略 1)两个独立策略处理,tiling 直接下发 AIV_ONLY kernel。触发条件 TO_MUL 要求 <code>batchC ≥ 128</code>(= aivNum×2,开 UB pingpong)、N 不在 (32B/dtype, 256B/dtype] 区间(中间区让给其它分支)、UB 容量放得下。</p>
|
||||
<h3>4.1 转Matmul(问题归约分支)</h3>
|
||||
<p><b>做什么</b>:当 <code>BatchA = 1</code> 或 <code>BatchB = 1</code> 时,把 batch 维折叠进矩阵维度,转化为普通 Matmul,计算完成后按需恢复 batch 维。</p>
|
||||
<p><b>为什么这么做</b>:单边 batch=1 的 BMM 与普通 Matmul 在数学上只差一个维度标签。Matmul 的优化体系(L2 切分、AL1/BL1 全载、swizzle)比 BMM 各分支在"单样本"情形下更成熟——<b>站在巨人肩膀上,不重复造轮子</b>。</p>
|
||||
<p><b>折叠规则与代价(关键的不对称性)</b>:</p>
|
||||
<table><tr><th>情形</th><th>折叠方式</th><th>是否免费</th><th>代价分析</th></tr>
|
||||
<tr><td><code>BatchB = 1</code></td><td>左矩阵 <code>[B, M, K]</code> 的 batch 维与 M 维在 ND 布局下<b>内存相邻</b>,直接视图为 <code>[B·M, K]</code>;输出 <code>[B·M, N]</code> 与 <code>[B, M, N]</code> 的内存布局逐元素一致</td><td><b>完全免费</b></td><td>零输入重排、零输出 split,直接转 Matmul 无任何代价</td></tr>
|
||||
<tr><td><code>BatchA = 1</code></td><td>需将右矩阵 <code>[B, K, N]</code> 折叠为 <code>[K, B·N]</code></td><td><b>有代价</b></td><td>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></table>
|
||||
<p><b>BatchA = 1 时的决策规则</b>:</p>
|
||||
<ul class="tight">
|
||||
<li><b>A 矩阵较小</b>($MK \cdot \text{dtype} \le L1_{size}$,即 A 可全载 L1):优先考虑将 A 常驻 L1,各核均匀读取 B 完成全部计算,无需输入/输出重排的额外开销——留在 BMM 分支内(广播友好的 IterBatch/ASW_Basic);</li>
|
||||
<li><b>A 矩阵较大</b>:将 <code>BatchA=1</code> 扩展为 <code>BatchA=BatchB</code>(广播语义),然后比较"B 折叠转 Matmul + 重排"与"BMM 分支"的预估时延,择优选择。</li>
|
||||
</ul>
|
||||
<p><b>进入条件</b>:<code>BatchA = 1 || BatchB = 1</code>。其中 <code>BatchB = 1</code> 恒进(免费);<code>BatchA = 1</code> 时按上述代价比较决定。</p>
|
||||
<p><b>源码对照</b>:源码中的 <code>MergeBatchAndMAxis()</code> 只做了 <code>batchB=1</code> 方向的折叠(<code>args_.mValue = batchA * mValue</code>),与"该方向免费"的判断互为印证。<code>batchA=1</code> 方向源码未做折叠,而是走 AL1_FULL_LOAD 或 ASW_Basic 的广播处理。</p>
|
||||
<h3>4.2 MergeBatch(多 batch 合并计算)</h3>
|
||||
<h4>4.2.1 做什么</h4>
|
||||
<p>核间按 B 分核(每核负责 $b_{core}$ 个 batch),核间无同步无通信。核内<b>将多个 batch 合并计算</b>:</p>
|
||||
<div class="math">$$
|
||||
[b, M, K] @ [b, K, N] \;\Rightarrow\; [bM, K] @ [K, bN] = [bM, bN] \;\xrightarrow{\text{BlockTrace}}\; [b, M, N]
|
||||
$$</div>
|
||||
<p>其中 <b>BlockTrace</b> 指以 $[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><b>为什么允许浪费</b>:进入该分支的 case 必然是访存 Bound(条件三保证),瓶颈在 MTE2 不在 MMAD,浪费的算力被搬移时延掩盖(§2.1 瓶颈交换)——<b>用本来闲置的 Cube 算力,换 tile 变大后的搬移效率与 Cube 利用率</b>。</p>
|
||||
<h4>4.2.2 进入条件(逐条 + 理由)</h4>
|
||||
<p>设计原则四条:① 硬件时延可流水掩盖(double buffer 乒乓);② 满足容量约束(每次计算 L1/L0A/L0B/L0C 放得下);③ GM→L1 搬移高效(§2.2 四条经验约束);④ 算力有浪费但计算不能成为瓶颈。</p>
|
||||
<p>形式化后,进入 MergeBatch 需<b>同时</b>满足:</p>
|
||||
<p><b>条件 1(batch 够分且够合并)</b>:<code>BatchA = BatchB</code>(无广播)且</p>
|
||||
<div class="math">$$
|
||||
b_{core} = \frac{B}{C} \ge b_{thr}
|
||||
$$</div>
|
||||
<p>$b_{thr}$ 是"多 batch 合并搬移能拿到效率收益"的最小合并数——$b$ 太小时合并的搬移收益抵不过实现复杂度,不如 IterBatch 且完全不浪费算力(功耗)。经验值 $b_{thr} = 4$(源码 <code>MIN_BATCH_L0 = 4</code>,注释 "each aic should process at least 4 batchs")。</p>
|
||||
<p><b>条件 2(合并数 b 的上下界)</b>:设核内单次 Cube 计算合并 $b$ 个 batch($b \le b_{core}$),要求 $b \ge b_{thr}$ 且 $b$ 由以下约束共同封顶:</p>
|
||||
<ul class="tight">
|
||||
<li><b>算存比约束</b>(计算不能变成瓶颈):合并后单次计算的算存比 $AI(b) = \dfrac{2bMN}{M + N}$,保持访存 Bound 要求</li>
|
||||
</ul>
|
||||
<div class="math">$$
|
||||
AI(b) < R \;\Longleftrightarrow\; b < b_{AI} = \frac{R \cdot (M + N)}{2MN}
|
||||
$$</div>
|
||||
<ul class="tight">
|
||||
<li><b>L0C 容量约束</b>:输出 $[bM, bN]$ 的 FP32 累加块须放入 256KB L0C;考虑乒乓(double buffer)则两份:</li>
|
||||
</ul>
|
||||
<div class="math">$$
|
||||
2 \cdot (bM)(bN) \cdot 4\text{B} \le 256\text{KB} \;\Rightarrow\; b \le b_{L0C} = \sqrt{\frac{256\text{KB}}{2 \cdot MN \cdot 4\text{B}}}
|
||||
$$</div>
|
||||
<ul class="tight">
|
||||
<li><b>L0A/L0B 容量约束</b>:$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)。</li>
|
||||
</ul>
|
||||
<blockquote><b>乒乓取舍</b>:若 $b_{core}$ 小于"不乒乓时 L0C 允许的合并数上限",说明 batch 余量不足,只能不乒乓(牺牲流水掩盖换合并数);否则乒乓。形式上 $b_{max} = \min(b_{AI},\; b_{L0C}^{(\text{是否乒乓})},\; b_{core})$。</blockquote>
|
||||
<p><b>条件 3(访存 Bound 性质)</b>:</p>
|
||||
<div class="math">$$
|
||||
\frac{2MN}{M + N} < \frac{R}{b}
|
||||
$$</div>
|
||||
<p>即 case 固有算存比显著低于平衡点(低一个合并倍数 $b$ 的量级)——这是"浪费可被掩盖"的定量保证。</p>
|
||||
<p><b>条件 4(K 向搬移效率)</b>:L1 级 K 切分 $k_{L1} = K / StepK$(StepK 为正整数),要求 $k_{L1} \ge k_{thr}$,$k_{thr} = 32\text{B}/\text{dtype}$(BF16 即 16 元素)。理由:$k_{L1}$ 决定 ND2NZ 的 dValue,过碎则搬移指令效率崩(§2.2 第 4 条)。</p>
|
||||
<p><b>条件 5(L1 驻留与 tile 效率)</b>:L1 级 batch 驻留数</p>
|
||||
<div class="math">$$
|
||||
b_{L1} = \frac{L1_{size}}{(M k_{L1} + k_{L1} N) \cdot \text{dtype}}, \qquad b_{L1} > b
|
||||
$$</div>
|
||||
<p>(L1 驻留的 batch 组必须大于单次计算的合并数,否则合并无从谈起),且搬移 tile 大小满足效率:左矩阵非转置时 $\max(b_{L1} M k_{L1} \cdot \text{dtype},\; k_{L1} N \cdot \text{dtype}) > min\_TileSize$(16KB);左矩阵转置时相应调整。</p>
|
||||
<p><b>条件 6(单核搬移总量)</b>:</p>
|
||||
<div class="math">$$
|
||||
b_{core} \cdot (M k_{L1} + k_{L1} N) \cdot \text{dtype} \ge min\_DatamountPerCore = 480\text{KB}
|
||||
$$</div>
|
||||
<p>对应 §2.2 第 2 条:单核搬移数据总量不足时带宽利用率上限被压低。</p>
|
||||
<p><b>条件汇总逻辑</b>:条件 1 定资格(batch 够多),条件 2/3 定上限(别算出瓶颈、别撑爆 L0),条件 4/5/6 定下限(搬移效率不能崩)。<b>上限与下限之间必须有交集,交集为空则该 case 与 MergeBatch 无缘。</b></p>
|
||||
<h4>4.2.3 核内参数求解</h4>
|
||||
<p>进入分支后,$b$ 与 baseK 的具体取值有一个重要性质:**在访存 Bound 前提下,$b < b_{AI}$ 时 b 的取值不影响性能,baseK 只需满足分形约束(16 倍数)也不影响算存比**——因为瓶颈是搬移,算力余量内怎么切都一样。据此:</p>
|
||||
<ol class="tight">
|
||||
<li>先由条件 2 求 $b_{max}$,实际 $b$ 在 $[b_{thr}, b_{max}]$ 内取,且<b>尽量均匀</b>($b_{core}$ 整除 $b$)——每次计算的 $b$ 均匀一致对功耗更有利;</li>
|
||||
<li>再由 $b$ 反查 L0AB 允许的最大 baseK:</li>
|
||||
</ol>
|
||||
<div class="math">$$
|
||||
baseK_{max} = \min\left(\frac{64\text{KB}}{2 \cdot bM \cdot \text{dtype}},\; \frac{64\text{KB}}{2 \cdot bN \cdot \text{dtype}},\; \frac{C0_{size}}{\text{dtype}}\right)
|
||||
$$</div>
|
||||
<p>向下取 16 倍数。</p>
|
||||
<ol class="tight">
|
||||
<li>L1 级参数:$k_{L1} \ge \min(k_{L0\_max},\; 128\text{B}/\text{dtype})$,$b_{L1} = \min(b_{L1\_max},\; b_{core})$。</li>
|
||||
</ol>
|
||||
<h4>4.2.4 数值例子</h4>
|
||||
<p>case:B=128, M=32, K=64, N=128, BF16,32 核。</p>
|
||||
<ol class="tight">
|
||||
<li><b>资格</b>:$b_{core} = 128/32 = 4 \ge b_{thr} = 4$ ✓;</li>
|
||||
<li><b>算存比</b>:$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><b>L0C 约束</b>(乒乓):$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$)——<b>L0C 才是真正的紧约束</b>;</li>
|
||||
<li><b>取 b = 2</b>(= $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><b>搬移效率</b>:$k_{L1} = 64 \ge k_{thr} = 16$ ✓;单核搬移量 $4 \times (32 \times 64 + 64 \times 128) \times 2\text{B} = 80\text{KB}$ —— <b>不足 480KB</b>,条件 6 不满足!</li>
|
||||
</ol>
|
||||
<p>结论:该 case 在"单核搬移总量"一项上不达标——这正是 MergeBatch 条件体系的用处:<b>它提前告诉你瓶颈不在算力而在搬移效率</b>。</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>
|
||||
<table><tr><th>源码条件</th><th>本文对应</th><th>差异</th></tr>
|
||||
<tr><td><code>batchC ≥ 4 × aicNum</code></td><td>条件 1:$b_{core} \ge b_{thr} = 4$</td><td><b>一致</b>(每核至少 4 batch)——合理</td></tr>
|
||||
<tr><td><code>alignK ≥ 64</code></td><td>条件 4:$k_{L1} \ge k_{thr}$=32B/dtype(BF16 为 16 元素)</td><td>源码保守 <b>4 倍</b>,可能误杀 K∈[16,64) 的可获益 case</td></tr>
|
||||
<tr><td><code>M ≤ N</code></td><td>无此限制(M > N 的镜像 case 原理上同样可合并,交换合并方向即可)</td><td>源码放弃镜像 case,属于<b>覆盖缺口</b></td></tr>
|
||||
<tr><td>无显式算存比/搬移效率判定</td><td>条件 3/4/6</td><td>源码把"是否最优"推给优先级顺序,本文条件体系是补全</td></tr></table>
|
||||
<h3>4.3 IterBatch(逐 batch 计算)</h3>
|
||||
<h4>4.3.1 做什么</h4>
|
||||
<p>核间按 B 分核(每核 1 个或多个 batch),核间无同步无通信;核内<b>逐个 batch 分别执行标准 Matmul 分块</b>(L1→L0A/L0B→Cube→L0C→Fixpipe)并输出。无算力浪费、无跨 batch 依赖——是"切 B"最朴素的形态。</p>
|
||||
<h4>4.3.2 进入条件(三大类 + 理由)</h4>
|
||||
<p>满足以下<b>任一</b>:</p>
|
||||
<p><b>类 1(计算 Bound 型)</b>:</p>
|
||||
<div class="math">$$
|
||||
AI_{full} = \frac{2MKN}{MK + KN + MN} \ge R \quad \text{且} \quad BatchA = BatchB,\; b_{core} = \frac{B}{C} \in \mathbb{Z}^+
|
||||
$$</div>
|
||||
<p>理由:连输出写出都计入仍是计算 Bound,则瓶颈恒在 Cube——只要负载均衡($b_{core}$ 整除,各核同量),逐 batch 计算就是满算力实现,无需任何花哨。</p>
|
||||
<p><b>类 2(输出驻留 L2 的均衡型)</b>:</p>
|
||||
<div class="math">$$
|
||||
\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}^+
|
||||
$$</div>
|
||||
<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 也不是瓶颈。三个条件合起来 = <b>读、算、写三条路都不堵</b>。</p>
|
||||
<p><b>类 3(访存 Bound 型)</b>:$\dfrac{2MN}{M+N} < R$ 且同时满足:</p>
|
||||
<ol class="tight">
|
||||
<li>$BatchA = BatchB$ 且 $b_{core} = B / C \ge 1$;</li>
|
||||
<li><b>负载均衡</b>:访存 Bound 时核负载利用率建议 $b_{Avg}/b_{Max} > 0.8$(计算 Bound 时建议 100%)。操作化表述:B 整除核数,或尾波(B mod C)活跃的核数 $\ge minCoreNum$(取 $0.8 \times 32 \approx 26$)。注意:若只写 <code>B mod C > minCoreNum</code> 会误杀整除 case(余数 0 恰恰是完全均衡),<b>判据应以均衡比为准、余数规则为其近似</b>;</li>
|
||||
<li><b>单核搬移不重复读</b>(访存 Bound 下重复读就是纯损失),按 L1 容量分五种形态之一:</li>
|
||||
<ul class="tight">
|
||||
<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>
|
||||
</ol>
|
||||
<h4>4.3.3 设计原理</h4>
|
||||
<ol class="tight">
|
||||
<li>核间切 B 零共享零依赖,<b>唯一的系统性风险是负载不均</b>——所以均衡是第一条件;</li>
|
||||
<li>输出须满足 L0C 容量:$MN \cdot 4\text{B} \le L0C$(否则核内还要切 M/N,那就不是纯 IterBatch 而是 ASW 行为——但 IterBatch 核内<b>允许</b>对单 batch 做标准 M/N/K tiling,此处的准确含义是"核内 tiling 不构成跨 batch 的耦合");</li>
|
||||
<li>访存 Bound 时<b>单核数据不得重复读</b>——L1 装得下才不重复,装不下就按 Step 切分且切分后仍满足搬移效率下限;</li>
|
||||
<li>与 MergeBatch 的分工:IterBatch 不浪费算力,但需要"单 batch 足够大"撑搬移效率与 Cube 利用率;MergeBatch 用浪费换效率,专治小 M×N。两者在 $M \times N$ 的中段重叠,由时延模型仲裁(第五章)。</li>
|
||||
</ol>
|
||||
<h4>4.3.4 核内参数求解</h4>
|
||||
<p>**情形 (a):$b_{core} = 1$ 且 L1 可放完整单 batch**</p>
|
||||
<pre><code>if (L0C_Size >= M * N * L0C_ElementSize):
|
||||
BaseM = M; BaseN = N
|
||||
BaseK = min(align(L0A / BaseM, 16), align(L0B / BaseN, 16))
|
||||
else:
|
||||
if M < N:
|
||||
BaseM = align(M, 16)
|
||||
BaseN = floor(L0C / BaseM)
|
||||
BaseK = min(floor_align(L0A / BaseM, 16), floor_align(L0B / BaseN, 16))
|
||||
else:
|
||||
BaseN = align(N, 16)
|
||||
BaseM = floor(L0C / BaseN)
|
||||
BaseK = min(floor_align(L0A / BaseM, 16), floor_align(L0B / BaseN, 16))</code></pre>
|
||||
<p>**情形 (b):$b_{core} > 1$ 且 L1 可放 2 batch 乒乓**</p>
|
||||
<pre><code>if (L0C_Size >= 2 * M * N * L0C_ElementSize):
|
||||
BaseM = M; BaseN = N
|
||||
BaseK = min(floor_align(L0A / BaseM, 16), floor_align(L0B / BaseN, 16))
|
||||
else:
|
||||
if M < N:
|
||||
BaseM = align(M, 16)
|
||||
BaseK = min(floor_align(L0A / 2 / BaseM, 16), K)
|
||||
BaseN = max(floor_align(L0C / 2 / BaseM, 16), floor_align(L0B / 2 / BaseK, 16))
|
||||
else:
|
||||
BaseN = align(N, 16)
|
||||
BaseK = min(floor_align(L0B / 2 / BaseN, 16), K)
|
||||
BaseM = max(floor_align(L0C / 2 / BaseN, 16), floor_align(L0A / 2 / BaseK, 16))</code></pre>
|
||||
<p><b>情形 (c)/(d)/(e):L1 放不下完整 batch</b></p>
|
||||
<p>一矩阵不切、另一切分:假设 L1 放完整左矩阵和部分右矩阵,则右矩阵应搬入的 K 向长度:</p>
|
||||
<div class="math">$$
|
||||
k_{L1\_b} = \min\!\left(\frac{L1_{size} - M \cdot K \cdot \text{dtype}}{N \cdot \text{dtype}},\; K\right)
|
||||
$$</div>
|
||||
<p>且 $k_{L1\_b} \ge k_{thr}$,切分后 tile ≥ 16KB。</p>
|
||||
<h4>4.3.5 执行流程</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>
|
||||
<h4>4.3.6 与源码对照</h4>
|
||||
<p>源码有三个 iterbatch 变体:<code>ITER_BATCH_BROADCAST_BASICAPI</code>(单边广播)、<code>ITER_BATCH_BASICAPI</code>(基础 API)、<code>ITER_BATCH</code>(高阶 API)。关键差异:</p>
|
||||
<table><tr><th>维度</th><th>源码 basicapi</th><th>源码高阶 API</th><th>本文</th></tr>
|
||||
<tr><td>L1 容量计算</td><td>$(sizeA + sizeB + bias) \times 2 \le l1Size$(除 DB)</td><td>$iterBatch = l1Size / inputSizeOneBatch$(不除 DB)</td><td>按 L1 形态 (a)~(e) 分类</td></tr>
|
||||
<tr><td>截断</td><td><code>mmadCount=8</code>(issue queue)、<code>fullCopySize=64KB</code></td><td><code>iterBatch ≤ 4</code> 时 singleCoreK 减半</td><td>统一为搬移效率 + 均衡率约束</td></tr>
|
||||
<tr><td>均衡率</td><td>0.8</td><td>0.8</td><td>0.8(访存 Bound)/ 1.0(计算 Bound)</td></tr>
|
||||
<tr><td>广播</td><td>独立分支(单边单轴)</td><td>不支持</td><td>由转Matmul 或 ASW_Basic 吸收</td></tr></table>
|
||||
<h3>4.4 StreamK(K 维核间切分)</h3>
|
||||
<h4>4.4.1 做什么</h4>
|
||||
<p>当 B、M、N 三个维度切到最碎仍填不满 32 核时,把 K 维切到核间:多核各算一段 K 的部分和,再归约:</p>
|
||||
<div class="math">$$
|
||||
C_{(\beta,\mu,\nu)} = \sum_{c \in \text{group}} C^{(c)}_{(\beta,\mu,\nu)} \quad \text{(部分和写 workspace,AIV 归约或原子加)}
|
||||
$$</div>
|
||||
<h4>4.4.2 进入条件(两条缺一不可)</h4>
|
||||
<p><b>条件 1(并行缺口存在)</b>:不切 K 时的独立输出块数</p>
|
||||
<div class="math">$$
|
||||
P = B \times \lceil M / 16 \rceil \times \lceil N / 16 \rceil < C = 32
|
||||
$$</div>
|
||||
<p>(源码取更保守的 $B \cdot mCnt \cdot nCnt \le aicNum/2$,并附加 ND-only、无交叉广播、确定性等级 ≤ 1 等工程限制)。注意<b>严格的 StreamK 不要求核间完全不切 B/M/N</b>——它是一般框架 <code>grid_K × grid_B × grid_M × grid_N ≤ C</code>,纯切 K 只是 grid_B=grid_M=grid_N=1 的特例;当 B/M/N 能提供部分并行度时,应该用组合 grid 把归约组 $grid_K$ 压到最小(grid_K 小一档,K 的门槛降一档平方级)。</p>
|
||||
<p><b>条件 2(归约代价可接受)</b>:每核计算时延须远大于归约时延(安全系数 $\alpha = 10$):</p>
|
||||
<div class="math">$$
|
||||
T_{MMAD/core} \ge \alpha \cdot T_{Reduce}
|
||||
\;\Longleftrightarrow\;
|
||||
\frac{K}{grid_K} \;\gtrsim\; grid_K \times 1690
|
||||
$$</div>
|
||||
<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 效率崩)。<b>grid_K 越大对 K 的要求越苛刻——StreamK 内部的 grid 搜索自然淘汰归约过重的配置。</b></p>
|
||||
<h4>4.4.3 实现方案</h4>
|
||||
<p><b>核间组织</b>:grid_K 个核组成一个归约组,共享同一个输出块 $(\beta, \mu, \nu)$ 的计算。核间分配:</p>
|
||||
<div class="math">$$
|
||||
\text{核 } c \text{ 的 K 段:} \left[\frac{c \cdot K}{grid_K},\; \frac{(c+1) \cdot K}{grid_K}\right)
|
||||
$$</div>
|
||||
<p><b>核内流水</b>(每核):</p>
|
||||
<pre><code>输入:K 段 [k_start, k_end),输出块 (β, μ, ν)
|
||||
for m_tile / n_tile(核内标准 tiling):
|
||||
for k_tile in range(k_start, k_end, baseK):
|
||||
MTE2: A[β, m_tile, k_tile:k_tile+baseK] → L1
|
||||
MTE1: L1 → L0A/L0B
|
||||
Cube: mmad → L0C 原地累加 # 本核的 K 段内累加
|
||||
Fixpipe: L0C → workspace(GM 或 L2) # 写出部分和</code></pre>
|
||||
<p><b>归约阶段</b>:</p>
|
||||
<ul class="tight">
|
||||
<li><b>方式 A:workspace + AIV 归约</b>(确定性)。部分和写入 GM workspace(每核一块 256×256 fp32 缓冲),归约由 AIV 执行(AIC:AIV = 1:2,2 个 AIV 伺候 1 个 AIC 的部分和流)。workspace 大小 = $C \times 256 \times 256 \times 4\text{B} + 20\text{MB}$(RPC 区);</li>
|
||||
<li><b>方式 B:原子加(AtomicAdd)</b>(非确定性)。部分和直接原子累加到输出 GM,省一遍读回,但归约顺序不定——源码在确定性等级 > 1 时禁用此变体。</li>
|
||||
</ul>
|
||||
<p><b>fixpipe 优化</b>:当 N 不对齐且输出块足够大时($n > 64$ 且 $n \% 16 \ne 0$ 且 $m > 2$ 且 $m \times n \ge 256$),启用 ND_FIXPIPE_1_2(1 AIC : 2 AIV 的 ND fixpipe 通路),由 AIV 分担搬出。</p>
|
||||
<p><b>grid 搜索</b>:$grid_K$ 从 2 开始递增(2, 4, 8, ...),每档检查条件 2 是否满足;同时检查 $K / grid_K \ge 256$。取满足条件的最小 $grid_K$(归约代价最小)。</p>
|
||||
<p><b>核间同步</b>:复用 Matmul 高阶 API 的 StreamK 模板(避免手写 CrossCore flagId 冲突,087篇/054篇)。</p>
|
||||
<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>
|
||||
<h4>4.4.5 与源码对照</h4>
|
||||
<p>源码 <code>batch_matmul_v3_basic_streamk_tiling.cpp</code> 的条件:</p>
|
||||
<table><tr><th>源码条件</th><th>本文对应</th><th>差异</th></tr>
|
||||
<tr><td><code>batchC × mCnt × nCnt ≤ aicNum/2</code></td><td>条件 1:$P < C$</td><td>源码取一半核数为阈值,更保守;本文用精确并行缺口</td></tr>
|
||||
<tr><td><code>CeilAlign(k,256) ≥ max(8192, aicNum×256B/dtype)</code></td><td>条件 2:$K \gtrsim grid_K^2 \times 1690$ 且 $K/grid_K \ge 256$</td><td>源码用固定阈值,本文用 grid 搜索动态判定</td></tr>
|
||||
<tr><td>fp32 非 hf32 时 K ≤ 200 万</td><td>无(本文不限)</td><td>源码因 binary accumulation 精度限制 fp32 超长 K;本文不限 dtype,通过 grid 搜索控制归约深度</td></tr>
|
||||
<tr><td><code>aivNum == 2 × aicNum</code></td><td>归约方式 A 的硬件前提</td><td>一致</td></tr>
|
||||
<tr><td>workspace = aicNum×256×256×4B + 20MB</td><td>归约方式 A 的 workspace</td><td>一致</td></tr></table>
|
||||
<h3>4.5 ASW_Basic(通用切分框架)</h3>
|
||||
<h4>4.5.1 做什么</h4>
|
||||
<p>不切 K,允许切 B/M/N 的<b>任意组合</b>——它是"不含 K 且含共享读取"的 6 种切分组合({M},{N},{M,N},{B,M},{B,N},{B,M,N})的统一实现框架,也是 B < 32 但 P ≥ 32 时的最优归宿,同时兜住 B ≥ 32 但 IterBatch/MergeBatch 条件不满足的剩余 case。<b>B 可以大、可以小、可以等于 1;ASW_Basic 是实践中最常命中的分支。</b></p>
|
||||
<h4>4.5.2 核心机制:swizzle + L2 管理</h4>
|
||||
<p>切 M/N 的固有代价是共享矩阵的重复读,ASW_Basic 用两件武器把代价压到最低:</p>
|
||||
<p><b>武器 1:ASW 滑窗蛇形 swizzle</b></p>
|
||||
<p>把 M 向按窗口 $W$ 分组,窗口内蛇形遍历 N 向:</p>
|
||||
<div class="math">$$
|
||||
W = \max\{\,d \mid d \mid C,\; d \le \lfloor\sqrt{C}\rfloor \,\}
|
||||
$$</div>
|
||||
<p>32 核取 $W = 4$。数学效果:同一时刻 32 个核活跃的工作集被压缩到"$W$ 个 A 行块 + 一条 B 列块带",L2 足迹最小 ⇒ 共享读取基本命中 5.2TB/s 的 L2 而非 1.6TB/s 的 GM。</p>
|
||||
<p>**为什么窗口取 $\lfloor\sqrt{C}\rfloor$ 的最大因子**:窗口越接近正方形,A 行块 + B 列块的 L2 足迹越小(方形窗的二维足迹是 $W \cdot K + K \cdot N$,随 W 偏离 $\sqrt{C}$ 而增大),且因子性保证整窗被核数均分、窗口边界不碎。</p>
|
||||
<p>窗口内的遍历顺序:</p>
|
||||
<div class="math">$$
|
||||
\sigma(\mu, \nu) = \big(\mu_{row} \cdot W + \mu_{col}\big) \cdot \tilde{N} + \nu'
|
||||
$$</div>
|
||||
<p>其中 $\mu_{row} = \lfloor \mu / W \rfloor$,$\mu_{col} = \mu \bmod W$,$\nu'$ 由蛇形决定:</p>
|
||||
<div class="math">$$
|
||||
\nu' = \begin{cases} \nu & \text{若 } \mu_{row} \text{ 为偶数(正向)} \\ \tilde{N} - 1 - \nu & \text{若 } \mu_{row} \text{ 为奇数(反向)} \end{cases}
|
||||
$$</div>
|
||||
<p>核号与轮次:</p>
|
||||
<div class="math">$$
|
||||
c = \mathrm{idx}(\beta, \mu, \nu) \bmod C,\qquad r = \lfloor \mathrm{idx}(\beta, \mu, \nu) / C \rfloor
|
||||
$$</div>
|
||||
<p><b>武器 2:L2 切分 + 错位分核</b></p>
|
||||
<p>工作集超过 128MB 时,按 mL2TileNum × nL2TileNum 切分。每个 L2 块<b>错位分核</b>(对角线分配):</p>
|
||||
<div class="math">$$
|
||||
\sigma_{diag}(\mu, \nu) = \mu \cdot \tilde{N} + \left(\nu + \left\lfloor \frac{\mu \cdot C}{\mathrm{lcm}(\tilde{M}, \tilde{N})} \right\rfloor \right) \bmod \tilde{N}
|
||||
$$</div>
|
||||
<p>避免多核同时抢同一地址的读读冲突,并优先选拖尾小的方案。</p>
|
||||
<h4>4.5.3 核间切分维度的选择顺序</h4>
|
||||
<p>按共享代价从低到高:</p>
|
||||
<table><tr><th>优先级</th><th>切分方式</th><th>共享矩阵</th><th>条件</th><th>代价</th></tr>
|
||||
<tr><td>1</td><td><b>切 B</b>(B ≥ 核数)</td><td>无</td><td>$B \ge C$</td><td>零共享,永远先试</td></tr>
|
||||
<tr><td>2</td><td><b>切 M</b>(B 不够)</td><td>右矩阵 [K, N]</td><td>$KN \cdot \text{dtype} \le 128$MB 则驻留 L2</td><td>右矩阵重复读,L2 吸收</td></tr>
|
||||
<tr><td>3</td><td><b>切 N</b></td><td>左矩阵 [M, K]</td><td>$MK \cdot \text{dtype} \le 128$MB 则驻留 L2</td><td>左矩阵重复读,L2 吸收</td></tr>
|
||||
<tr><td>4</td><td><b>混合切</b>(B×M、M×N、…)</td><td>双向共享</td><td>靠 swizzle + L2 切分管理</td><td>代价最高,最后试</td></tr>
|
||||
<tr><td>5</td><td><b>降核</b></td><td>—</td><td>P 远小于 32 且 K 也不够格走 StreamK</td><td>宁可部分核闲置(小 case 时延绝对值小)</td></tr></table>
|
||||
<h4>4.5.4 基本块寻优(cubeBound 模型)</h4>
|
||||
<p>ASW_Basic 的 baseM/baseN 由 cubeBound 解析模型求解。平台指标:hbmBW/l2BW 由频率×核数×平台速率算出,computePower = 单核算力 × aicNum。</p>
|
||||
<p>cubeBoundEdge 公式:</p>
|
||||
<div class="math">$$
|
||||
\text{cubeBoundEdge} = \frac{l2BW}{\text{computePower}} + l2CacheUsage \cdot \left(1 - \frac{l2BW}{hbmBW}\right) \cdot \text{cmr} - \frac{1 + l2BW/hbmBW}{kValue}
|
||||
$$</div>
|
||||
<p>其中 $\text{cmr} = (m+n)/(m \cdot n)$(每单位输出元素对应的输入行/列搬运代价),$l2CacheUsage = \max(batch \cdot (m+n) \cdot k \cdot \text{dtype} / l2Size, 1.0)$。</p>
|
||||
<p>逐项含义:</p>
|
||||
<ul class="tight">
|
||||
<li><b>第 1 项</b>:L2 供数速率 ÷ cube 耗数速率——cube bound 的理论阈值;</li>
|
||||
<li><b>第 2 项</b>:L2 装不下工作集时的访存惩罚,按 HBM 与 L2 带宽差加权;</li>
|
||||
<li><b>第 3 项</b>:K 向流水复用修正——K 越大,A/B 复用越充分。</li>
|
||||
</ul>
|
||||
<p>求解流程:在 baseM/baseN 候选解空间内枚举,按"计算访存比 cubeBoundParam = 1/baseM + 1/baseN 与尾块负载均衡率 balanceRate"综合评分:</p>
|
||||
<ul class="tight">
|
||||
<li>剪枝 ①:balanceRate ≥ 0.9 且候选 cubeBoundParam 既差于当前最优又大于 edge → 跳过;</li>
|
||||
<li>剪枝 ②:fp32 cubeBound 且多轮时 baseM/N < 64 → 跳过;</li>
|
||||
<li>评分:cubeBoundCond(满足 cubebound 且更均衡)优先;否则以 cubeBoundParam/balanceRate 为综合分。</li>
|
||||
</ul>
|
||||
<p>进入枚举前 edge 先乘 <b>CUBE_BOUND_RATIO = 0.85</b> 预留 15% 余量。</p>
|
||||
<h4>4.5.5 L1 步进与 buffer 管理</h4>
|
||||
<div class="math">$$
|
||||
\max StepK = \min\left(\left\lceil\frac{K}{baseK}\right\rceil,\; 8\right)
|
||||
$$</div>
|
||||
<p>stepK ≤ 8 是 issue queue 深度约束(源码注释:"Shape约束 && issue queue约束")。容量约束:</p>
|
||||
<div class="math">$$
|
||||
(aL1 + bL1) \times 2 \le L1_{size} - bias - scale
|
||||
$$</div>
|
||||
<p>且单边不超过 L1 一半:</p>
|
||||
<div class="math">$$
|
||||
\max(aL1, bL1) \times 4 \le L1_{size}
|
||||
$$</div>
|
||||
<p>stepK 枚举优先级:首个可行解 > 256B 对齐 > 单次搬运 ≥ 48KB。</p>
|
||||
<h4>4.5.6 内部特化(不是独立分支)</h4>
|
||||
<ul class="tight">
|
||||
<li><b>AL1/BL1 全载</b>:当单边无 batch 且该侧矩阵很小(如 batchA=1 且 M ≤ 256),把整个 A 常驻 L1,$\tilde{M} = 1$、A 的 GM→L1 搬运只发生一次——这是 ASW 内部"M 不切 + L1 驻留深度拉满"的参数极限,不是新分支。源码条件:<code>batchA ≤ 1</code> 且 <code>m ≤ 256</code> 且 <code>alignMatASize × 2 ≤ l1Size</code> 且 B 侧"值得全载"(总量 ≥ L1×aicNum 或 每核循环 ≥ 4 轮);</li>
|
||||
<li><b>广播友好形态</b>:交叉广播 case(§1.2)在 ASW 内通过对广播侧做 L1/L2 驻留处理,共享关系与切 M/N 同构。</li>
|
||||
</ul>
|
||||
<h4>4.5.7 执行流程</h4>
|
||||
<pre><code>核间:32 核,按 swizzle 滑窗分配 (β, μ, ν) 块
|
||||
核内(每核):
|
||||
for round in range(R): # R = ⌈B̃·M̃·Ñ / C⌉
|
||||
(β, μ, ν) = swizzle_inv(round × C + core_id)
|
||||
for k_tile in range(0, K, baseK):
|
||||
MTE2: A[β, μ, k_tile:k_tile+baseK] → L1(双缓冲/4buffer)
|
||||
MTE1: L1 → L0A/L0B
|
||||
Cube: mmad → L0C 原地累加 # K 循环不出核
|
||||
Fixpipe: L0C → L2/GM</code></pre>
|
||||
<h4>4.5.8 与源码对照</h4>
|
||||
<p>源码 <code>batch_matmul_v3_asw_basic_tiling.cpp</code> 的 IsCapable:A、B 的非连续转置状态必须一致;batch 完全相等;batchBias ≤ 1;dtype 限 fp16/bf16 系。DoOpTiling:ResetBase → GetRebalanceBlock → CalL1Tiling。</p>
|
||||
<table><tr><th>维度</th><th>源码</th><th>本文</th></tr>
|
||||
<tr><td>基本块默认值</td><td>baseM=256, baseN=256, baseK=128B/dtype</td><td>由 cubeBound 模型求解,不固定</td></tr>
|
||||
<tr><td>寻优模型</td><td>GetRebalanceBlock(cubeBound 解析模型)</td><td>同(§4.5.4),但源码的经验常数(0.85/0.9/48KB)需注意移植</td></tr>
|
||||
<tr><td>L1 buffer</td><td>4 buffer 判定:<code>abL1TensorSize × 4 ≤ l1Size</code></td><td>同(§4.5.5)</td></tr>
|
||||
<tr><td>swizzle</td><td>ASW 滑窗(窗长取 √C 的最大因子)</td><td>同(§4.5.2)</td></tr>
|
||||
<tr><td>fp32 splitK</td><td>K > 1024(或 8192 当 K > 2^28)时启用</td><td>由 StreamK grid 搜索统一管理</td></tr></table>
|
||||
<hr>
|
||||
<h2>五、分支仲裁与整体决策流程</h2>
|
||||
<h3>5.1 为什么需要仲裁层</h3>
|
||||
<p>第四章的进入条件给出的是各分支的"主场",但主场之间有<b>重叠区</b>(例如 B ≥ 32 且 M×N 中等时,MergeBatch 与 IterBatch 都合法;B 较大时 IterBatch 与 ASW_Basic 切 B 等效)。重叠区的最优归属没有解析解,必须由统一的时延模型仲裁:</p>
|
||||
<div class="math">$$
|
||||
\text{branch}^* = \arg\min_{cand \in \bigcup \text{各分支候选}} \max\big(T_{MMAD},\; T_{MTE2},\; T_{MTE1},\; T_{Fixpipe}\;[,\;T_{Reduce}]\big)
|
||||
$$</div>
|
||||
<p>分支体系的价值在于<b>候选集完备且无冗余</b>:每个等价类只派一个代表框架生成候选,时延模型在等价类内部和边界上做精细仲裁。两层缺一不可——只有模型没有分支,搜索空间是 15 种组合 × 全部 grid 参数的爆炸;只有分支没有模型,边界 case 被硬阈值误杀。</p>
|
||||
<h3>5.2 总决策流程</h3>
|
||||
<pre><code>输入:B, M, N, K, dtype, layout, 广播形态, bias
|
||||
│
|
||||
├─[0] K = 0 → AIV 清零;K = 1 → AIV 逐元素乘(特殊分支)
|
||||
│
|
||||
├─[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>六、完备性审视(漏洞检查)</h2>
|
||||
<p>按"转Matmul / MergeBatch / IterBatch / StreamK / ASW_Basic + 前置通路层 cover 全部 BMM case 最优实现"的设计目标,逐条拷问。</p>
|
||||
<h3>6.1 覆盖矩阵:无空洞</h3>
|
||||
<table><tr><th>shape 区域</th><th>主分支</th><th>兜底</th></tr>
|
||||
<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></table>
|
||||
<h3>6.2 审视发现的问题与修正</h3>
|
||||
<p>设计骨架成立,但有五处需要修正或显式声明,否则会被反例拷问:</p>
|
||||
<p><b>问题 1:转Matmul 的方向不对称(原设想的最大漏洞)。</b> "一侧 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 分支内。</p>
|
||||
<p><b>问题 2:交叉广播 case 的归属必须显式声明。</b> BatchA=(6,1)、BatchB=(1,6) 这类 case 既不进转Matmul(两侧都>1),也不满足 MergeBatch/IterBatch 的 BatchA=BatchB 前提——它们落入 ASW_Basic(对广播侧做驻留)。不声明就会显得"五分支有洞"。声明后完备性闭合。</p>
|
||||
<p><b>问题 3:K 退化不在五分支内,需前置通路层。</b> K=0(无计算)、K=1(无累加深度)走 AIV 向量通路;这是"计算通路选择"维度,与"切分维度"正交,放在五分支之前判断。源码用 K_EQUAL_ZERO / TO_MUL 两个独立策略处理,印证了这一层的必要性。</p>
|
||||
<p><b>问题 4:IterBatch 草稿条件的字面表述有误杀。</b> <code>B mod C > minCoreNum</code> 会拒绝 B 整除核数的完全均衡 case(余数 0)。正确表述应以负载均衡比为准:访存 Bound 要求 $b_{Avg}/b_{Max} > 0.8$,计算 Bound 要求 100%;余数规则只是它的近似操作化。</p>
|
||||
<p><b>问题 5:经验常数必须可标定,且源码硬编码条件偏粗。</b> $b_{thr}=4$、$k_{thr}$=32B/dtype、480KB、16KB、dValue≥256B 都是该档芯片的实测经验值,文档中全部显式参数化(换芯片只换常数表)。对照之下源码把 <code>batchC ≥ 4×aicNum</code>、<code>alignK ≥ 64</code>、<code>M ≤ N</code> 等直接写死,其中 <code>alignK ≥ 64</code> 比草稿阈值保守 4 倍、<code>M ≤ N</code> 放弃了镜像 case——<b>做最优实现时不应继承这些硬编码,而应按第四章的条件体系重新判定</b>。</p>
|
||||
<h3>6.3 与源码策略的关系</h3>
|
||||
<p>源码 arch35 的 11 个策略(K_EQUAL_ZERO / TO_MUL / STREAM_K / MERGE_BATCH / ITER_BATCH_BROADCAST / ITER_BATCH_BASICAPI / ITER_BATCH / AL1_FULL_LOAD / BL1_FULL_LOAD / ASW_BASIC / BASE)与本文分支不是同层概念:</p>
|
||||
<table><tr><th>源码策略</th><th>本文归属</th></tr>
|
||||
<tr><td>K_EQUAL_ZERO、TO_MUL</td><td>前置通路层(§4.0),正交于切分</td></tr>
|
||||
<tr><td>STREAM_K</td><td>StreamK(§4.4)</td></tr>
|
||||
<tr><td>MERGE_BATCH</td><td>MergeBatch(§4.2)</td></tr>
|
||||
<tr><td>ITER_BATCH、ITER_BATCH_BASICAPI、ITER_BATCH_BROADCAST</td><td>IterBatch(§4.3;broadcast 变体是广播输入下的数据复用特化;单边广播在本文更多由转Matmul 吸收)</td></tr>
|
||||
<tr><td>AL1_FULL_LOAD、BL1_FULL_LOAD</td><td>ASW_Basic 的 L1 驻留参数极限(§4.5.6)</td></tr>
|
||||
<tr><td>ASW_BASIC、BASE</td><td>ASW_Basic(§4.5)</td></tr></table>
|
||||
<p><b>结论</b>:源码策略集 = 本文分支 × 计算通路 × 驻留/广播特化的展开。本文的分支是"切分等价类"的最小完备骨架,源码的冗余策略可在新实现中收敛,源码遗漏的(条件 3/4/6 的显式判定、M>N 的 MergeBatch 镜像、转Matmul 的 BatchA=1 方向)正是优化的增量空间。</p>
|
||||
<hr>
|
||||
<h2>七、总结</h2>
|
||||
<ol class="tight">
|
||||
<li><b>最优的定义</b>:$T_{total} = \max(T_{MMAD}, T_{MTE2}, T_{MTE1}, T_{Fixpipe}[, T_{Reduce}])$ 最小;瓶颈交换是合法且必要的手段。</li>
|
||||
<li><b>分支的推导</b>:4 维可切 → 15 种组合完备 → 切分特征表(切 B 免费 / 切 M/N 廉价被 L2 吸收 / 切 K 昂贵需归约)→ 按价格从低到高购买并行度 → 坍缩为 MergeBatch、IterBatch、ASW_Basic、StreamK 四个等价类,加前置的转Matmul 问题归约层与 K 退化通路层,共七大路径,<b>完备且极小</b>。</li>
|
||||
<li><b>各分支的条件</b>都不是孤立经验,而是"资格(并行度够不够)→ 上限(容量/算存比封顶)→ 下限(搬移效率托底)"三层逻辑的实例化;所有经验常数可标定、可移植。</li>
|
||||
<li><b>边界不靠硬切</b>:重叠区由端到端时延模型统一仲裁,分支体系负责候选集的完备无冗余。</li>
|
||||
<li><b>设计经受了漏洞拷问</b>:修正了转Matmul 的方向不对称与 IterBatch 均衡条件的字面误杀,显式声明了交叉广播与 K 退化的归属;相对源码实现,本文条件体系更细、更真、覆盖更全。</li>
|
||||
</ol>
|
||||
<hr>
|
||||
<h2>附录 A:经验常数表(可标定)</h2>
|
||||
<table><tr><th>常数</th><th>数值</th><th>出处/依据</th><th>说明</th></tr>
|
||||
<tr><td>$b_{thr}$</td><td>4</td><td>源码 <code>MIN_BATCH_L0</code></td><td>合并搬移的最小有效合并数</td></tr>
|
||||
<tr><td>$k_{thr}$</td><td>32B/dtype</td><td>§2.2 第 4 条</td><td>ND2NZ dValue 下限(BF16 为 16 元素)</td></tr>
|
||||
<tr><td>$min\_DatamountPerCore$</td><td>480KB</td><td>§2.2 第 2 条</td><td>单核搬移总量下限</td></tr>
|
||||
<tr><td>$min\_TileSize$</td><td>16KB</td><td>§2.2 第 3 条</td><td>单次搬移 tile 下限(168篇官方为 16KB)</td></tr>
|
||||
<tr><td>$minCoreNum$</td><td>26(≈0.8×32)</td><td>§2.2 第 1 条</td><td>尾波活跃核数下限</td></tr>
|
||||
<tr><td>$CUBE\_BOUND\_RATIO$</td><td>0.85</td><td>源码 GetRebalanceBlock</td><td>cubeBound edge 余量</td></tr>
|
||||
<tr><td>$balanceRateEdge$</td><td>0.9</td><td>源码 GetRebalanceBlock</td><td>负载均衡剪枝阈值</td></tr>
|
||||
<tr><td>均衡率(访存 Bound)</td><td>0.8</td><td>源码 iterbatch</td><td>$b_{Avg}/b_{Max}$ 下限</td></tr>
|
||||
<tr><td>$\alpha$(StreamK 安全系数)</td><td>10</td><td>本文设定</td><td>$T_{MMAD} \ge \alpha \cdot T_{Reduce}$</td></tr>
|
||||
<tr><td>$R$(BF16/FP16)</td><td>607.5 FLOP/元素</td><td>§2.3</td><td>芯片算存比平衡点</td></tr></table>
|
||||
<h2>附录 B:与 issue#3 (v0.3) 的差异</h2>
|
||||
<table><tr><th>章节</th><th>v0.3 状态</th><th>v0.4 补全内容</th></tr>
|
||||
<tr><td>StreamK 实现方案</td><td>空</td><td>§4.4.3 完整实现:核间组织、核内流水、归约两方式、fixpipe 优化、grid 搜索、核间同步</td></tr>
|
||||
<tr><td>ASW_Basic 实现方案</td><td>空</td><td>§4.5.2~4.5.7 完整实现:swizzle 公式、L2 切分、切分顺序、cubeBound 模型、L1 步进、执行流程</td></tr>
|
||||
<tr><td>MergeBatch 进入条件</td><td>5 条(简略)</td><td>6 条(含公式推导),补充算存比约束 $b_{AI}$、L0C 解析式 $b_{L0C}$、乒乓取舍</td></tr>
|
||||
<tr><td>IterBatch 实现方案</td><td>部分伪代码</td><td>补全 (a)~(e) 五种情形的完整伪代码 + 设计原理 + 源码对照</td></tr>
|
||||
<tr><td>转Matmul 实现方案</td><td>简略</td><td>补充 BatchA=1 的决策规则(A 大小判断 + 广播扩展 + 时延比较)</td></tr>
|
||||
<tr><td>特殊分支</td><td>简述</td><td>补充 AIV vs AIC 的定量分析、源码触发条件</td></tr>
|
||||
<tr><td>源码对照</td><td>无</td><td>每个分支新增"与源码对照"小节,指出源码的保守/缺口/一致项</td></tr>
|
||||
<tr><td>完备性论证</td><td>无</td><td>§3.6 + 第六章:15 种组合 → 7 路径的完备性与极小性论证</td></tr></table>
|
||||
<hr>
|
||||
<p>*参考文档:《BMM分块计算数学公式》《BMM最优软件实现方案设计》《BatchMatMulV3算子分支实现分析》《BatchMatmul算子特性分析》;源码:<code>ops-nn/matmul/batch_mat_mul_v3</code>(arch35);芯片资料:昇腾 950PR 架构白皮书与 CANN 9.0.0 性能建模文档。*</p>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
Reference in New Issue
Block a user