From 8664dd5e57532985f78d89de12ca0783d0bf92cf Mon Sep 17 00:00:00 2001 From: admin Date: Thu, 3 Sep 2026 09:25:49 +0000 Subject: [PATCH] =?UTF-8?q?Add=20BMM=5FTheory:=20docs/02=5F=E5=88=86?= =?UTF-8?q?=E6=94=AF=E7=90=86=E8=AE=BA/05=5FStreamK=E5=88=86=E6=94=AF.md?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../docs/02_分支理论/05_StreamK分支.md | 73 +++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 BMM/BMM_Theory/docs/02_分支理论/05_StreamK分支.md diff --git a/BMM/BMM_Theory/docs/02_分支理论/05_StreamK分支.md b/BMM/BMM_Theory/docs/02_分支理论/05_StreamK分支.md new file mode 100644 index 0000000..68c3f36 --- /dev/null +++ b/BMM/BMM_Theory/docs/02_分支理论/05_StreamK分支.md @@ -0,0 +1,73 @@ +# StreamK 分支理论 + +> 整理自《BMM算子优化分析 v0.98》§七. 对应软件实现 `bmm_theory/branches/stream_k.py`. + +## 1. 定位:最后才买的并行维度 + +决策树的末端分支——当 **B/M/N 三维的并行度加起来都填不满核** 时,才动用最贵的切 K。回顾切分价格:`cost(切B)=0 < cost(切M/N) ≪ cost(切K)`。切 K 贵的根因:K 维有 L0C 累加机制,核内切 K 时多轮 mmad 在 L0C 原地累加、中间结果不出核;一旦切到**核间**,部分和必须写出 workspace 再归约——**切 K 是唯一同时破坏"累加不出核"和"输出独占"的切法**。 + +适用 case 画像是"细长"形:B 小、M/N 小、K 大(如 B=4, M=N=128, K=10240)。 + +## 2. 进入条件(4 条同时满足) + +记 `P = B·M·N·4B / L0C`——以"L0C 满载的输出基本块"为粒度估计不切 K 的最大并行度。 + +| # | 条件 | 含义 | +|---|---|---| +| 1 | `P ≤ C/2` | **并行缺口**:不切 K 时至多一半核有事做,K 是唯一剩余的并行维度。阈值取 C/2 而非 C:StreamK 定义即 grid_K≥2(至少 2 路切 K),grid_K=2 时每块需 2 核,总核数需求 = ⌈P⌉×2 ≤ C | +| 2 | `K/grid_K ≥ 256B/dtype` | **单核 K 段下限**:每核 K 段内轴连续长度不小于 dValue 推荐值,保证段内搬移效率不崩。grid_K = ⌊C/⌈P⌉⌋ | +| 3 | `K > grid_K²/(grid_K-1)·θ_c`,θ_c≈12 | **归约代价可接受**:切 K 后总时延须小于不切 K(降核 ASW),推导见 §3 | +| 4 | 确定性等级 ≤ 1 且 ND 格式 | 归约顺序不定引入浮点非确定性,确定性等级 2/3 的业务禁用 | + +## 3. 归约代价与 θ_c 推导(条件 3 的来源) + +**收益判据**:切 K 后芯片级总时延须小于不切 K。 + +- **不切 K(降核 ASW)**:P 个 tile 各占一个核,总时延 = 单 tile 流水时延 `T_pipe = max(T_MMAD^t, T_MTE2^t)`; +- **切 K(StreamK)**:同样 P 个 tile,每 tile 由 grid_K 个核共同完成(各算 K/grid_K 段),流水时延缩为 `T_pipe/grid_K`,但串行追加归约时延 `T_Reduce^t`。 + +$$T_{SK} = \frac{T_{pipe}}{grid_K} + T_{Reduce}^t < T_{pipe} \iff T_{Reduce}^t < T_{pipe}\Big(1-\frac{1}{grid_K}\Big)$$ + +**归约时延构成**(每 tile,部分和驻留 L2、AIV 归约;M^t·N^t = L0C/4B): + +$$T_{Reduce}^t = \underbrace{\frac{grid_K\cdot M^t N^t\cdot 4B}{W_{L2}}}_{\text{AIC 写部分和}} + \underbrace{\frac{grid_K\cdot M^t N^t\cdot 4B}{W_{L2}}}_{\text{AIV 读回}} + \underbrace{\frac{grid_K\cdot M^t N^t}{Q_{AIV}}}_{\text{AIV 求和}} + \underbrace{\frac{M^t N^t\cdot outB}{W_{L2}}}_{\text{写回}}$$ + +**计算 Bound 情形**(T_pipe = T_MMAD^t)代入判据,两边除以 L0C/4B 后 **tile 尺寸消去**——K 阈值不依赖 M/N 具体值: + +$$K > \frac{grid_K^2}{grid_K-1}\cdot\theta_c,\qquad \theta_c = \frac{Q_{16}}{2}\Big(\frac{8B}{W_{L2}}+\frac{1}{Q_{AIV}}\Big) \approx 12$$ + +代入数值:grid_K=2 → K>49;grid_K=4 → K>66;grid_K=8 → K>112。L2 读写(1.54 ps/元素)是主导项,AIV 求和仅占 5%。 + +访存 Bound 情形阈值恒低于计算 Bound(θ_m·[M^tN^t/(M^t+N^t)]/θ_c ≤ 128/304 ≈ 0.42 < 1),故**汇总条件取计算 Bound 阈值**(保守,覆盖两种情形)。 + +> **注意**:θ_c 对 workspace 落点敏感——部分和落 GM 时读写带宽从 5.2TB/s 降到 ~0.64TB/s,θ_c 升至约 97。**设计时应优先保证 workspace 驻留 L2**。 + +## 4. 实现方案 + +**核间组织**:先按 B/M/N 切出输出块,剩余核预算折成 K 向份数: + +$$blocksPerBatch = \Big\lfloor\frac{C}{B}\Big\rfloor,\qquad grid_K = \frac{blocksPerBatch}{mCnt\cdot nCnt}$$ + +mCnt、nCnt 收拢为 blocksPerBatch 的因子(避免碎核尾块)。归约组内核 c 负责 K 段 `[cK/grid_K, (c+1)K/grid_K)`。 + +**核内流水**:对自己的 K 段做标准分块流水(MTE2→L1→L0→mmad),段内多轮在 L0C 原地累加;段完部分和经 Fixpipe 写出。 + +**写出 dtype 关键点**(用户特别强调):StreamK 的中间部分和为**防精度丢失**,写出 dtype = L0C dtype(4B),**不随 C 矩阵的 fp16/fp8 转换**——因为部分和还要再次读入 AIV 完成 K 维度累加,中途转 fp16 会丢精度。只有最终归约结果才按 C dtype 写出。 + +**归约**:grid_K 个核的部分和经归约(Fixpipe 模式或 AIV)求和为最终结果,workspace 驻留 L2。 + +## 5. 时延模型 + +$$T_{SK} = \max\Big(\frac{T_{MMAD}}{grid_K},\; \frac{T_{MTE2}}{grid_K}\Big) + T_{Reduce} + T_{drain}$$ + +其中 T_MMAD、T_MTE2 为不切 K 的单 tile 流水时延,切 K 后并行缩 grid_K 倍;T_Reduce 为串行追加的归约尾。 + +## 6. 与源码的差异 + +源码 `batch_matmul_v3_basic_streamk_tiling.cpp` 的 K 固定门槛为 `CeilAlign(K,256) ≥ max(8192, aicNum×256B/dtype)`: + +- `aicNum×256B/dtype` = 4096(BF16)= dValue 下限在最大 grid_K=C 下的保障(对应条件 2); +- `8192` = C×512B = dValue 推荐值在最大 grid_K=C 下的保障(同为条件 2,取推荐值); +- max 取更严的 8192。 + +修正后的归约阈值(θ_c≈12,grid_K=32 时 K>396)远低于 8192,说明 **8192 的绑定约束是 dValue(条件 2),不是归约代价(条件 3)**。源码不动态算 grid_K,用固定阈值保守合并两条条件。本软件按理论动态计算 grid_K 与 θ_c,进入条件更精细。