Add BMM_Theory: docs/02_分支理论/05_StreamK分支.md

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，进入条件更精细。