matmul-analysis/BMM/BMM_Theory/docs/02_分支理论/05_StreamK分支.md

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≈10.9	归约代价可接受：切 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 10.9

（Q16 为 Cube-only 口径 13.5T，issue#40；旧总算力口径 15.1875T 时 θ_c≈12。）代入数值：grid_K=2 → K>44；grid_K=4 → K>58；grid_K=8 → K>100。L2 读写（1.54 ps/元素）是主导项，AIV 求和仅占 5%。

访存 Bound 情形阈值恒低于计算 Bound（θ_m·[M^tN^t/(M^t+N^t)]/θ_c ≤ 128/270 ≈ 0.47 < 1），故汇总条件取计算 Bound 阈值（保守，覆盖两种情形）。

注意：θ_c 对 workspace 落点敏感——部分和落 GM 时读写带宽从 5.2TB/s 降到 ~0.64TB/s，θ_c 升至约 86。设计时应优先保证 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≈10.9，grid_K=32 时 K>360）远低于 8192，说明 8192 的绑定约束是 dValue（条件 2），不是归约代价（条件 3）。源码不动态算 grid_K，用固定阈值保守合并两条条件。本软件按理论动态计算 grid_K 与 θ_c，进入条件更精细。