添加 BMM最优软件实现方案设计.md
This commit is contained in:
882
BMM最优软件实现方案设计.md
Normal file
882
BMM最优软件实现方案设计.md
Normal file
@@ -0,0 +1,882 @@
|
||||
# BMM 最优软件实现方案设计
|
||||
|
||||
> 目标芯片：昇腾 950PR（32 AIC / 64 AIV，1.65GHz，GM 128GB/1.6TB/s，L2 128MB/5.2TB/s）
|
||||
> 算子接口：BMM (Batch Matrix Multiply)，接口形态与 `BatchMatMulV3` 对齐，支持任意 shape/dtype 自动决策
|
||||
> 设计目标：任意 shape/dtype 的 BMM case 都能得到在 950PR 上**最短端到端算子时延**的软件实现方案
|
||||
> 架构版本：NPU 架构版本 351x（DAV_3510）
|
||||
|
||||
---
|
||||
|
||||
## 1. 设计方法论
|
||||
|
||||
### 1.1 核心思路
|
||||
|
||||
BMM 的计算本质是：给定 batch 维度，每个 batch 执行一次矩阵乘 `[M, K] × [K, N] → [M, N]`。算子端到端时延由以下流水线中最慢的一级决定：
|
||||
|
||||
```
|
||||
T_total = max(T_MMAD, T_MTE2, T_MTE1, T_FIXPIPE [, T_REDUCE])
|
||||
```
|
||||
|
||||
要达成最短时延，需要同时决策三个维度：
|
||||
|
||||
1. **软件实现分支选择**（FoldBatchToMatMul / MergeBatch / IterBatch / StreamK / ASW_Basic）——决定了核间切分和数据流组织方式
|
||||
2. **Tiling 参数选择**（baseM, baseN, baseK, L1 tile 尺寸）——决定了单核内的分块策略
|
||||
3. **Swizzle 编排**（分块执行顺序、核分配、L2 切分策略）——决定了数据复用和 L2 命中率
|
||||
|
||||
### 1.2 950PR 关键硬件参数
|
||||
|
||||
| 参数 | 数值 | 说明 |
|
||||
|---|---|---|
|
||||
| AIC 核数 / 频率 | 32 / 1.65GHz | 源码 `aicNum = ascendcPlatform.GetCoreNumAic()` |
|
||||
| GM 带宽 | 1.6TB/s（读写共享） | 读写共享同一总线，至少 3/4 核（24 核）并发才能达到 90%+ 带宽利用率 |
|
||||
| L2 Cache 容量 | 128MB（源码 `l2Size` 字段） | 读写数据均占用 L2 容量；L2 满时发生替换：L0C 写出的 dirty 数据需先写回 GM（占用 GM 写带宽），GM 读入的 clean 数据可直接替换 |
|
||||
| L2 读带宽 | 5.2TB/s | 仅 L2→L1 方向 |
|
||||
| L2 写带宽 | 5.2TB/s | 仅 L0C→Fixpipe→L2 方向，与 L2 读带宽独立 |
|
||||
| L1 Buffer | 512KB / AIC | 源码 `l1Size` 字段；L1 容量直接影响数据重复读取量——L1 越大，每个 M/N tile 可容纳的 K 分片越大，内外循环重复读取次数越少 |
|
||||
| L0A/L0B/L0C | 64KB / 64KB / 256KB | 源码 `l0ASize/l0BSize/l0CSize`；L0C 容量约束 baseM × baseN |
|
||||
| UB | 512KB / AIV | 源码 `ubSize` |
|
||||
| MTE1 | 256 BPC ≈ 422.4GB/s（单 AIC） | L1→L0A/L0B |
|
||||
| Cube 峰值 | 13.5 TFLOPS BF16（单 AIC） | 16×16×16 MAC/拍 × 1.65GHz |
|
||||
|
||||
### 1.3 数据通路关键约束
|
||||
|
||||
**GM→L1 通路**：数据可从 GM 搬运至 L1（DataCopy 指令，走 MTE2 流水），搬运过程中若 L2 配置为 Cache 模式，数据会**随路驻留**到 L2。同一数据块再次被 AIC 需要时，若仍在 L2 中则命中，以 L2 带宽（5.2TB/s）读取，不再走 GM。
|
||||
|
||||
**L0C 写出通路**：L0C 可直接写出至 L2（Fixpipe，L0C→L2），也可直接写出至 GM（L0C→GM）。写出至 L2 时使用 5.2TB/s 独享写带宽，不与 L2 读竞争。
|
||||
|
||||
**L2 替换机制**：L2 容量有限（128MB）。当 L2 被占满后，新的 GM 读入或 L0C 写出数据需要腾出空间：
|
||||
- 对于 L2 中已有的 L0C 写出数据（dirty）：需先写回 GM（占用 GM 写带宽 1.6TB/s），再腾出空间
|
||||
- 对于 L2 中已有的 GM 读入数据（clean）：直接替换，无需写回
|
||||
|
||||
**GM 访存效率**（影响 MTE2 实测带宽的关键因素，优先级依次递减）：
|
||||
1. **参与核数**：GM 带宽由所有核共享并发读取才能充分利用。至少 3/4 核（24 核）并发才能达到 90% 以上带宽利用率
|
||||
2. **单核读取数据量**：单核读取数据量越大，越容易达到高带宽利用率
|
||||
3. **单次搬移数据块大小**：单次搬移数据块越大，越容易达到高带宽利用率
|
||||
4. **ND2NZ 指令的 dValue**：dValue 应 ≥128Byte，推荐 256Byte 或 512Byte。例如 GM 中以 ND 排布的非转置左矩阵的 K 维度、非转置右矩阵的 N 维度对应的搬移数据量即为 dValue
|
||||
|
||||
---
|
||||
|
||||
## 2. 软件实现分支体系
|
||||
|
||||
BMM 算子有五个主要软件实现分支：
|
||||
|
||||
| 分支 | 核间切分方式 | 核内计算方式 | 适用场景 |
|
||||
|---|---|---|---|
|
||||
| **FoldBatchToMatMul** | 将 batch 维折叠到 M/N，转为普通 MatMul | 标准 Matmul tiling（非 BMM 本体的研究范围） | 左矩阵或右矩阵的 Batch 维为 1 |
|
||||
| **MergeBatch** | 切 B | 核内分级合并 batch（L1/L0/Cube 三级），切 K 累加后 BlockTrace 取对角线 | B 大，单 batch 的 M/N 小 |
|
||||
| **IterBatch** | 切 B | 核内逐个 batch 分别执行 Matmul | B 大，单 batch 的 M/N 大 |
|
||||
| **StreamK** | 切 K | 核间沿 K 维动态分配，原子累加输出 | B/M/N 切分无法有效利用核数，K 足够大 |
|
||||
| **ASW_Basic** | 切 B / M / N（不切 K） | 标准 Matmul tiling（L1→L0A/L0B→Cube→L0C→Fixpipe） | 通用 case，B 可大可小 |
|
||||
|
||||
源码中 950PR（DAV_3510）的策略优先级为：
|
||||
|
||||
```
|
||||
K_EQUAL_ZERO → TO_MUL → STREAM_K → MERGE_BATCH_BASICAPI →
|
||||
ITER_BATCH_BROADCAST_BASICAPI → ITER_BATCH_BASICAPI → ITER_BATCH →
|
||||
AL1_FULL_LOAD_BASIC → BL1_FULL_LOAD_BASIC → ASW_BASIC → BASE
|
||||
```
|
||||
|
||||
> 注意：源码中 `TO_MUL`（策略 1）是 K==1 时用 AIV 做向量乘，与本文的 FoldBatchToMatMul（折叠 batch 到 M/N 转为普通 MatMul）含义完全不同。本文不使用源码的 TO_MUL 命名，以避免混淆。源码优先级也非针对所有 shape 都最优，本文以批判性视角重新设计决策框架。
|
||||
|
||||
### 2.1 FoldBatchToMatMul ——桥接分支
|
||||
|
||||
**定位**：这不是一个 BMM 内部的优化分支，而是一个**桥接分支**——当 `BatchA == 1` 或 `BatchB == 1` 时，BMM 可以通过维度折叠转化为普通 MatMul。BMM 算子的研究范围只关注 BatchA > 1 **且** BatchB > 1 的差异化 case。对于可折叠的 case，直接借助 Matmul 优化体系（tiling、L2 切分、全载等）的已有成果，站在巨人肩膀上高效实现。
|
||||
|
||||
**折叠规则**：
|
||||
- BatchA=1：左矩阵 `[M, K]` 不变，右矩阵 `[B, K, N]` 沿 N 维拼接为 `[K, B×N]`。计算 `[M, K] × [K, B×N] → [M, B×N]`，输出后 Split 恢复为 `[B, M, N]`
|
||||
- BatchB=1：左矩阵 `[B, M, K]` 沿 M 维拼接为 `[B×M, K]`，右矩阵 `[K, N]` 不变。计算 `[B×M, K] × [K, N] → [B×M, N]`，输出后 Split
|
||||
- 两个均为 1：纯 MatMul
|
||||
|
||||
**适用条件**：`BatchA == 1 || BatchB == 1`。满足即走此分支，无需进入 BMM 后续决策。
|
||||
|
||||
### 2.2 MergeBatch
|
||||
|
||||
**原理**：核间按 Batch 分核，每个核负责 `B_core` 个 batch。核内受 L1 和 L0 容量约束，将 batch 分级合并——L1 级合并 `bL1` 个 batch，L0 级合并 `bL0` 个 batch。核内可切 K（`kL0` 粒度），Cube 一次计算 `[bL0×M, kL0] × [kL0, bL0×N]` 并累加到 L0C，跨 K 完成全部累加后通过 BlockTrace 提取对角线 `[M, N]` 块作为各 batch 的有效输出。交叉项（不同 batch 的 A 与 B 的乘积）被丢弃。
|
||||
|
||||
**优点**：核内 Cube 利用率高（等效大矩阵的 M、N 维度大），K 维复用次数高。
|
||||
|
||||
**缺点**：存在算力浪费——交叉项被丢弃，浪费比例 `(bL0-1)/bL0`。但 MergeBatch 场景下通常是访存 Bound，浪费可被掩盖。
|
||||
|
||||
### 2.3 IterBatch
|
||||
|
||||
**原理**：核间按 Batch 分核，每个核负责若干个 batch。核内逐个 batch 分别调用 Matmul 完成计算。无算力浪费，无跨 batch 的数据依赖。
|
||||
|
||||
**子变体**：`ITER_BATCH_BROADCAST_BASICAPI`（B 矩阵跨 batch 相同时广播）、`ITER_BATCH_BASICAPI`（Tensor API）、`ITER_BATCH`（基础 API）。
|
||||
|
||||
### 2.4 StreamK
|
||||
|
||||
**原理**：切 K 维到核间，核间沿 K 维动态分配分片，每个核处理一段 K 后通过原子累加（AtomicAdd）更新输出。核间**仅切 K，不切 B/M/N**——因为切 K 本身已经提供了足够的核间并行度，再切 B/M/N 只会减少单核工作量而不会增加并行度，反而增加复杂度。
|
||||
|
||||
**优点**：在 B/M/N 维度不足时，利用 K 维提供并行度；动态分配天然负载均衡。
|
||||
|
||||
**缺点**：存在核间 Reduce 的额外开销（原子累加）。
|
||||
|
||||
### 2.5 ASW_Basic
|
||||
|
||||
**原理**：不切 K，切 B/M/N 维度。核间按 B/M/N 分配，每核独立完成自己负责的输出区域。**这是最通用的分支——B 可以 == 1（纯 MatMul），也可以 < 32（切 M/N），也可以 ≥ 32（切 B，此时与 IterBatch 等效但核内不逐个 batch 迭代而是通过 tiling 自然覆盖）。**
|
||||
|
||||
**核间切分模式**：
|
||||
- 切 B：每个核负责若干个 batch，核间无共享。要求 B ≥ num_cores
|
||||
- 切 M：每个核负责若干 M 行，右矩阵 [K, N] 被所有核共享
|
||||
- 切 N：每个核负责若干 N 列，左矩阵 [M, K] 被所有核共享
|
||||
- 切 B×M / B×N / M×N：混合切分
|
||||
|
||||
---
|
||||
|
||||
## 3. MatmulConfig 与 Tiling 参数体系
|
||||
|
||||
### 3.1 核心 Tiling 参数
|
||||
|
||||
| 层级 | 参数 | 含义 | 约束 |
|
||||
|---|---|---|---|
|
||||
| L2 切分 | `mL2TileNum`, `nL2TileNum` | M/N 方向 L2 切分块数 | 单次 L2 切分的数据量 ≤ L2 容量 |
|
||||
| L1 切分 | `l1M`, `l1N`, `l1K` | L1 Buffer 中一次处理的 M/N/K 范围 | `l1M × l1K + l1K × l1N` × sizeof ≤ L1/(2×db) |
|
||||
| L0 切分 | `baseM`, `baseN`, `baseK` | Cube 单次计算的 M/N/K 分块 | L0A: `baseM×baseK`×sizeof ≤ 64KB；L0B: `baseK×baseN`×sizeof ≤ 64KB；L0C: `baseM×baseN`×4 ≤ 256KB |
|
||||
|
||||
**L0C 容量对 M/N 切分的影响**：L0C = 256KB，以 FP32 累加计，`baseM × baseN ≤ 256×256 = 65536`。在 MergeBatch 中，合并 bL0 个 batch 后，L0C 约束变为 `(bL0×M) × (bL0×N) × 4 ≤ 256KB`。
|
||||
|
||||
### 3.2 Tiling Key 模板参数
|
||||
|
||||
| 参数 | 含义 | 取值 |
|
||||
|---|---|---|
|
||||
| `MULTIBATCHL1FULLLOAD` | 多 batch 时 L1 全载 | 0/1 |
|
||||
| `MULTIBATCH` | 是否为多 batch 模式 | 0/1 |
|
||||
| `LOADMODE` | 加载模式 | BASE_FULLLOAD / AL1_FULLLOAD / BL1_FULLLOAD / VECTOR_FULLLOAD |
|
||||
| `ISMULTIBATCHOUT` | 输出是否为多 batch 格式 | 0/1 |
|
||||
| `MIXND2NZ` | 是否混合 ND 和 NZ 格式 | 0/1 |
|
||||
| `KERNEL_TYPE` | 核类型 | AIC_ONLY / MIX_AIC_1_2 |
|
||||
|
||||
### 3.3 Tiling 参数选择算法
|
||||
|
||||
```
|
||||
输入: M_core, N_core, K_core, dtype, L1_size, L0A_size, L0B_size, L0C_size
|
||||
输出: baseM, baseN, baseK, l1M, l1N, l1K
|
||||
|
||||
Step 1: 确定 baseM/baseN 候选集
|
||||
// L0C 约束
|
||||
for baseM in [16, 32, ..., min(M_core, 256)] (step 16):
|
||||
for baseN in [16, 32, ..., min(N_core, 256)] (step 16):
|
||||
if baseM * baseN * 4 > L0C_size: continue
|
||||
|
||||
// L0A/L0B 约束
|
||||
maxK_A = L0A_size / (baseM * sizeof(dtype))
|
||||
maxK_B = L0B_size / (baseN * sizeof(dtype))
|
||||
baseK_max = min(maxK_A, maxK_B)
|
||||
|
||||
// L1 双缓冲约束
|
||||
while baseK_max >= 16:
|
||||
l1_usage = l1BufferNum × (baseM × baseK_max + baseK_max × baseN) × sizeof(dtype)
|
||||
if l1_usage ≤ L1_size: break
|
||||
baseK_max -= 16
|
||||
|
||||
if baseK_max >= 16:
|
||||
baseK = min(baseK_max, K_core)
|
||||
score = baseM * baseN * baseK
|
||||
candidates.append((score, baseM, baseN, baseK))
|
||||
|
||||
Step 2: 选择最优
|
||||
candidates.sort(descending by score)
|
||||
优先选择 baseM/baseN 整除 M_core/N_core 的，减少尾块开销
|
||||
|
||||
Step 3: 确定 L1 tile
|
||||
l1M = baseM; l1N = baseN; l1K = baseK
|
||||
若 L1 容量充裕，可增大 l1K 以减少 M/N 外循环的重复读取次数
|
||||
```
|
||||
|
||||
### 3.4 典型 Tiling 配置（950PR，FP16/BF16）
|
||||
|
||||
| 场景 | baseM | baseN | baseK | L0C 利用率 | 适用条件 |
|
||||
|---|---|---|---|---|---|
|
||||
| 大 M×N（compute-bound） | 256 | 256 | 128 | 100% | M≥256, N≥256, K≥128 |
|
||||
| 瘦 M（M 小 N 大） | 64 | 256 | 256 | 25% | M=64, N≥256 |
|
||||
| 矮 N（M 大 N 小） | 256 | 64 | 256 | 25% | M≥256, N=64 |
|
||||
| 方阵中等 | 128 | 128 | 256 | 25% | M≈N≈128~1024 |
|
||||
| 极小 shape | 16 | 16 | 256 | 0.4% | M,N 极小 |
|
||||
|
||||
---
|
||||
|
||||
## 4. Swizzle 与 L2 管理策略
|
||||
|
||||
### 4.1 Swizzle 核心原则
|
||||
|
||||
1. **外循环维度 = 重复读取代价最小的维度**：哪个维度的输入矩阵在 L2 中驻留成本低，哪个维度做外循环
|
||||
2. **K 始终是内循环**：K 维累加在 L0C 中完成（除非 StreamK 分支）
|
||||
3. **L2 友好**：L2 容量足够时数据只从 GM 读一次，后续重复读取均从 L2 命中
|
||||
4. **最小化尾块**：选择 baseM/baseN 尽可能整除 M_core/N_core
|
||||
|
||||
### 4.2 L2 切分策略
|
||||
|
||||
当单次计算的数据量超过 L2 容量（128MB）时使能 L2 Cache 切分：
|
||||
1. 计算 M 方向 L2 切分数 `mL2TileNum` 和 N 方向 L2 切分数 `nL2TileNum`
|
||||
2. 确保每次 L2 切分的数据量 ≤ L2 阈值
|
||||
3. **错位分核**：将每个 L2 切分块按对角线分配给不同核，避免多核同时访问相同地址导致读读冲突
|
||||
4. 优先选择拖尾较小的切分方案
|
||||
|
||||
### 4.3 各分支的 Swizzle 模式
|
||||
|
||||
#### MergeBatch（核内三级合并，切 K）
|
||||
|
||||
```
|
||||
// 核间：按 Batch 分核，每核负责 B_core 个 batch
|
||||
// 核内：分级合并，切 K
|
||||
|
||||
for k_l1 in range(0, K, kL1): # L1 级 K 切分
|
||||
MTE2: A[bL1×M, k_l1:k_l1+kL1] → L1_ping
|
||||
MTE2: B[k_l1:k_l1+kL1, bL1×N] → L1_ping
|
||||
|
||||
for b_start in range(0, bL1, bL0): # 遍历 batch 组
|
||||
for k_l0 in range(0, kL1, kL0): # L0 级 K 切分
|
||||
MTE1: A[b_start×M:(b_start+bL0)×M, k_l0:k_l0+kL0] → L0A
|
||||
MTE1: B[k_l0:k_l0+kL0, b_start×N:(b_start+bL0)×N] → L0B
|
||||
Cube: [bL0×M, kL0] × [kL0, bL0×N] → 累加到 L0C
|
||||
|
||||
// 完成所有 kL0 子循环后，L0C 中有 [bL0×M, bL0×N]
|
||||
BlockTrace: 提取 bL0 个对角线 [M, N] 块 → Fixpipe → L2/GM
|
||||
```
|
||||
|
||||
#### IterBatch
|
||||
|
||||
```
|
||||
for each core_batch in batches:
|
||||
for each batch in core_batches:
|
||||
for m_tile in M_tiles:
|
||||
for n_tile in N_tiles:
|
||||
for k_tile in K_tiles:
|
||||
MTE2: 预取 A[k_next], B[k_next] → L1
|
||||
MTE1: L1 → L0A, L1 → L0B
|
||||
Cube: L0A × L0B → L0C (累加)
|
||||
Fixpipe: L0C → L2
|
||||
```
|
||||
|
||||
#### StreamK
|
||||
|
||||
```
|
||||
for each k_tile in dynamic_K_tiles: # 核间动态分配 K 分片
|
||||
MTE2: 预取 A[k_tile], B[k_tile] → L1
|
||||
MTE1: L1 → L0A, L1 → L0B
|
||||
Cube: L0A × L0B → L0C (累加)
|
||||
AtomicAdd: L0C 部分和 → GM/L2 中的输出 C
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 5. 端到端时延评估模型
|
||||
|
||||
### 5.1 通用评估框架
|
||||
|
||||
```
|
||||
T_total = max(T_MMAD, T_MTE2, T_MTE1, T_FIXPIPE [, T_REDUCE])
|
||||
```
|
||||
|
||||
### 5.2 MMAD 计算时间
|
||||
|
||||
```
|
||||
FP16/BF16: T_MMAD = (B_core × M_core × N_core × K_core) / (8192 × 1.65e9)
|
||||
FP8/MXFP8: T_MMAD = (B_core × M_core × N_core × K_core) / (16384 × 1.65e9)
|
||||
```
|
||||
|
||||
其中 `B_core, M_core, N_core, K_core` 为单核负责的计算空间。**MergeBatch 中 Cube 实际执行的计算量为 `bL0×M × bL0×N × K × 2`，有效计算量为 `B_core × M × N × K × 2`，差值为交叉项。**
|
||||
|
||||
### 5.3 MTE2 搬运时间
|
||||
|
||||
```
|
||||
Size_A_MTE2 = B_core × (N_core/baseN) × M_core × K_core × sizeof(A)
|
||||
Size_B_MTE2 = B_core × (M_core/baseM) × N_core × K_core × sizeof(B)
|
||||
|
||||
Size_A_GM = B_core × M_core × K_core × sizeof(A) // 首次从 GM
|
||||
Size_A_L2 = Size_A_MTE2 - Size_A_GM // 重复从 L2
|
||||
|
||||
BW_MTE2_A = Size_A_MTE2 / (Size_A_GM / 1.6e12 + Size_A_L2 / 5.2e12)
|
||||
T_MTE2 = max(Size_A_MTE2 / BW_MTE2_A, Size_B_MTE2 / BW_MTE2_B)
|
||||
```
|
||||
|
||||
**GM 访存效率修正**：实际有效 GM 带宽需乘以效率因子 η。`BW_GM_eff = 1.6e12 × η`。保守估计：24+ 核并发、单核数据量 ≥ 1MB、dValue ≥ 128B → η ≈ 0.85~0.92；≤8 核并发、单核数据量小 → η ≈ 0.4~0.6。
|
||||
|
||||
### 5.4 MTE1 搬运时间
|
||||
|
||||
```
|
||||
Size_MTE1_per_tile = baseK × (baseM + baseN) × sizeof(dtype)
|
||||
T_MTE1 = B_core × (M_core/baseM) × (N_core/baseN) × (K_core/baseK) × Size_MTE1_per_tile / 422.4e9
|
||||
```
|
||||
|
||||
### 5.5 FIXPIPE 搬出时间
|
||||
|
||||
```
|
||||
Size_C = B_core × M_core × N_core × sizeof(dtype_C)
|
||||
T_FIXPIPE = Size_C / 5.2e12 + Size_C / 1.6e12 // L2 写 + 写回 GM
|
||||
```
|
||||
|
||||
### 5.6 各分支特有项
|
||||
|
||||
**MergeBatch 额外项**：浪费比例 `(bL0-1)/bL0`。合并后有效 AI = `(bL0×M×N×K×2) / ((bL0×M×K + K×bL0×N + bL0²×M×N) × sizeof)`。若 AI < 270 则浪费被掩盖。
|
||||
|
||||
**StreamK 额外项**：`T_REDUCE = (num_cores × B × M × N × 4) / BW_reduce`。部分和在 L2 中原子累加时 `BW_reduce ≈ 5.2e12`（受原子操作开销影响），在 GM 中时 `BW_reduce ≈ 1.6e12`。
|
||||
|
||||
---
|
||||
|
||||
## 6. 策略自动决策框架
|
||||
|
||||
### 6.1 从完备性出发：所有可能的核间切分组合
|
||||
|
||||
BMM 的迭代空间为 4D：`(b, m, n, k)`。核间切分的本质是将这 4D 空间分配到 32 个 AIC 核上。完备性角度，所有可能的切分组合是四个维度 {B, M, N, K} 的任意非空子集，共 2⁴-1 = 15 种：
|
||||
|
||||
| 组合 | 切分维度 | 特征 |
|
||||
|---|---|---|
|
||||
| C1 | {B} | 纯切 B |
|
||||
| C2 | {M} | 纯切 M |
|
||||
| C3 | {N} | 纯切 N |
|
||||
| C4 | {K} | 纯切 K |
|
||||
| C5 | {B, M} | 切 B 和 M |
|
||||
| C6 | {B, N} | 切 B 和 N |
|
||||
| C7 | {B, K} | 切 B 和 K |
|
||||
| C8 | {M, N} | 切 M 和 N |
|
||||
| C9 | {M, K} | 切 M 和 K |
|
||||
| C10 | {N, K} | 切 N 和 K |
|
||||
| C11 | {B, M, N} | 切 B、M、N |
|
||||
| C12 | {B, M, K} | 切 B、M、K |
|
||||
| C13 | {B, N, K} | 切 B、N、K |
|
||||
| C14 | {M, N, K} | 切 M、N、K |
|
||||
| C15 | {B, M, N, K} | 切全部四维 |
|
||||
|
||||
### 6.2 从完备性到性能最优：筛选逻辑
|
||||
|
||||
性能最优的分支集合必须是完备性组合的子集。排除某个组合的唯一充分理由是：**对于所有可能的 shape，该组合的时延都严格大于等于另一个组合**。若某个组合在至少一种 shape 下能达到最优，则必须保留。
|
||||
|
||||
**筛选规则 1：含 K 的组合全部保留，统一归入 StreamK 分支。**
|
||||
|
||||
与上一版文档的错误论证不同，这里的正确分析如下：
|
||||
|
||||
切 K 到核间后，再切 B/M/N 并不会"减少每核的 K 范围"——恰恰相反。以 B=4, K=4096, 32 核为例：
|
||||
|
||||
- 纯切 K（C4）：grid_K=32, grid_B=1 → K_core=128, B_core=4。Reduce 组大小为 32。
|
||||
- 切 K+B（C7）：grid_K=8, grid_B=4 → K_core=512, B_core=1。Reduce 组大小为 8。
|
||||
|
||||
**K_core 从 128 增大到 512**，同时 Reduce 组从 32 缩小到 8。K_core 增大意味着每核有更多的 K 维计算量，部分和写出次数更少；Reduce 组缩小意味着核间累加的数据量和同步开销都更低。C7 可能优于 C4。
|
||||
|
||||
同理，C9（{M,K}）、C10（{N,K}）以及更高维组合 C12~C15 都可能在某些 shape 下优于纯切 K。参考知识库中 `cann-samples/matmul_story` 文档的描述：
|
||||
|
||||
> "切K模板（StreamK）在普通模板多核并行切分M轴、N轴的基础上，进一步增加了K轴的切分"
|
||||
|
||||
StreamK 本身就是 M/N 切分 + K 切分的组合框架。因此，全部 8 个含 K 的组合（C4, C7, C9, C10, C12, C13, C14, C15）统一归入 StreamK 分支，通过 grid 配置（grid_K × grid_B × grid_M × grid_N ≤ 32）来覆盖。StreamK 内部的 grid 选择问题见 6.6 节。
|
||||
|
||||
**筛选规则 2：不切 K 的 {B} 拆分为 MergeBatch 和 IterBatch。**
|
||||
|
||||
论证：纯 {B} 的核内处理只有两种方式——合并多个 batch 一起算（MergeBatch）或逐个 batch 算（IterBatch）。不存在第三种方式。因此将 {B} 拆分为 MergeBatch（当 M×N 小，合并提升 Cube 利用率）和 IterBatch（当 M×N 大，单 batch 已能有效利用 Cube）。
|
||||
|
||||
**筛选规则 3：不切 K 的其余 6 个组合归入 ASW_Basic 分支。**
|
||||
|
||||
论证：{M}, {N}, {B,M}, {B,N}, {M,N}, {B,M,N} 的共同特征是不切 K。它们的差异仅在于 B/M/N 中的哪些维度切、各切多少——这是 ASW_Basic 内部的核间切分维度选择问题。ASW_Basic 内部根据 shape 特征自动选择最优的切分维度组合。
|
||||
|
||||
**最终性能最优分支集合（4 个）：**
|
||||
|
||||
| 分支 | 覆盖的完备性组合 | 特征 |
|
||||
|---|---|---|
|
||||
| **MergeBatch** | {B}（核内合并） | 不切 K，核内多 batch 合并 |
|
||||
| **IterBatch** | {B}（核内逐个） | 不切 K，核内逐个 batch |
|
||||
| **ASW_Basic** | {M}, {N}, {B,M}, {B,N}, {M,N}, {B,M,N} | 不切 K，灵活切 B/M/N |
|
||||
| **StreamK** | {K}, {B,K}, {M,K}, {N,K}, {B,M,K}, {B,N,K}, {M,N,K}, {B,M,N,K} | 切 K，含 B/M/N 组合切分 |
|
||||
|
||||
**FoldBatchToMatMul 的处理**：当 BatchA=1 或 BatchB=1 时，BMM 退化为可折叠的 case。此时不进入上述 4 个分支的决策，而是直接借助 Matmul 优化体系。BMM 的研究范围限定为 BatchA > 1 且 BatchB > 1 的差异化 case。
|
||||
|
||||
### 6.3 MergeBatch 分支
|
||||
|
||||
#### 6.3.1 核间切分
|
||||
|
||||
核间按 Batch 分核：`num_cores = min(32, max(8, B/2))`，每个核负责 `B_core = ceil(B / num_cores) ≥ 2` 个 batch。
|
||||
|
||||
#### 6.3.2 核内三级 Batch 合并与 K 切分
|
||||
|
||||
单核负责 `B_core` 个 batch，受 L1 和 L0 容量约束，分级合并：
|
||||
|
||||
```
|
||||
B_core → [ L1 级: bL1 个 batch 合并 ] → [ L0 级: bL0 个 batch 合并 ] → [ Cube: bL0 个 batch ]
|
||||
↑ ↑ ↑
|
||||
(bL1 ≤ B_core) L1 容量约束 L0 容量约束
|
||||
```
|
||||
|
||||
**L0 级约束**：
|
||||
```
|
||||
L0A: bL0 × M × kL0 × sizeof(dtype) ≤ 64KB
|
||||
L0B: kL0 × bL0 × N × sizeof(dtype) ≤ 64KB
|
||||
L0C: (bL0 × M) × (bL0 × N) × 4 ≤ 256KB
|
||||
```
|
||||
|
||||
**L1 级约束**：
|
||||
```
|
||||
2 × (bL1 × M × kL1 + kL1 × bL1 × N) × sizeof(dtype) ≤ 512KB [双缓冲]
|
||||
```
|
||||
|
||||
**K 切分与累加流程**：
|
||||
|
||||
```
|
||||
for k_l1 in range(0, K, kL1): # L1 级 K 切分
|
||||
MTE2: A[bL1×M, k_l1:k_l1+kL1], B[k_l1:k_l1+kL1, bL1×N] → L1
|
||||
|
||||
for b_start in range(0, bL1, bL0): # 遍历 batch 组
|
||||
for k_l0 in range(0, kL1, kL0): # L0 级 K 切分
|
||||
MTE1: A[b_start×M:(b_start+bL0)×M, k_l0:k_l0+kL0] → L0A
|
||||
MTE1: B[k_l0:k_l0+kL0, b_start×N:(b_start+bL0)×N] → L0B
|
||||
Cube: [bL0×M, kL0] × [kL0, bL0×N] → 累加到 L0C
|
||||
|
||||
// K 循环完成，L0C 中有完整的 [bL0×M, bL0×N]
|
||||
BlockTrace: 提取 bL0 个对角线 [M, N] 块 → Fixpipe → L2/GM
|
||||
```
|
||||
|
||||
BlockTrace：`C[i, m, n] = Result[i×M + m, i×N + n]`，`i ∈ [0, bL0)`。
|
||||
|
||||
#### 6.3.3 适用条件（逐条定量）
|
||||
|
||||
**条件 1：每核至少 2 个 batch。** `B_core ≥ 2`，即 `num_cores ≤ B/2`。
|
||||
|
||||
**条件 2：单 batch 的 M×N 足够小，Merge 后 Cube 利用率提升显著。**
|
||||
```
|
||||
M × N < 128 × 128 = 16384（L0C 利用率 < 25%）
|
||||
```
|
||||
若 M×N 更大，单 batch 已能较有效利用 Cube，合并收益不足以抵消浪费。
|
||||
|
||||
**条件 3：合并 L0 级 batch 后结果矩阵能放入 L0C。**
|
||||
```
|
||||
(bL0 × M) × (bL0 × N) × 4 ≤ 256KB
|
||||
```
|
||||
例如 M=N=64：bL0≤4（256×256×4=256KB）；M=N=128：bL0≤2。
|
||||
|
||||
**条件 4：算力浪费可被访存时延掩盖。**
|
||||
```
|
||||
有效 AI = (bL0×M×N×K×2) / ((bL0×M×K + K×bL0×N + bL0²×M×N) × sizeof)
|
||||
若 AI < 270（GM 带宽比）→ 访存 Bound → 浪费被掩盖
|
||||
```
|
||||
|
||||
#### 6.3.4 候选生成
|
||||
|
||||
```
|
||||
def candidate_mergebatch(B, M, N, K, dtype):
|
||||
if M * N >= 128 * 128: return [] # 条件2
|
||||
candidates = []
|
||||
for num_cores in [32, 16, 8]:
|
||||
B_core = ceil(B / num_cores)
|
||||
if B_core < 2: continue # 条件1
|
||||
for bL0 in range(min(B_core, 4), 1, -1):
|
||||
if (bL0*M)*(bL0*N)*4 > 256*1024: continue # 条件3
|
||||
kL0 = min(64*1024//(bL0*M*sizeof(dtype)), 64*1024//(bL0*N*sizeof(dtype)))
|
||||
for bL1 in range(min(B_core, 8), bL0-1, -1):
|
||||
if 2*(bL1*M*kL0 + kL0*bL1*N)*sizeof(dtype) <= 512*1024: break
|
||||
AI = (bL0*M*N*K*2)/((bL0*M*K+K*bL0*N+bL0*bL0*M*N)*sizeof(dtype))
|
||||
candidates.append({"branch":"MergeBatch","num_cores":num_cores,"B_core":B_core,
|
||||
"bL0":bL0,"bL1":bL1,"kL0":kL0,"waste":(bL0-1)/bL0,"waste_hidden":AI<270})
|
||||
break
|
||||
return candidates
|
||||
```
|
||||
|
||||
### 6.4 IterBatch 分支
|
||||
|
||||
#### 6.4.1 核间切分与核内计算
|
||||
|
||||
核间切 B：`num_cores = min(32, B)`，每核 `B_core = ceil(B / num_cores)` 个 batch。核内逐个 batch 独立执行标准 Matmul tiling。
|
||||
|
||||
#### 6.4.2 适用条件
|
||||
|
||||
**条件 1：每核至少 1 个 batch。** `B_core ≥ 1`，即 `num_cores ≤ B`。
|
||||
|
||||
**条件 2：单 batch 的 M×N 足够大，独立 Matmul 能有效利用 Cube。**
|
||||
```
|
||||
M × N ≥ 64 × 64 = 4096
|
||||
```
|
||||
若 M×N < 4096，单 batch 的 L0C 利用率 < 6.25%，应优先考虑 MergeBatch。
|
||||
|
||||
#### 6.4.3 候选生成
|
||||
|
||||
```
|
||||
def candidate_iterbatch(B, M, N, K, dtype):
|
||||
candidates = []
|
||||
for num_cores in [min(32,B), min(32,B)//2, min(32,B)//4]:
|
||||
if num_cores < 1: continue
|
||||
B_core = ceil(B / num_cores)
|
||||
baseM, baseN, baseK = select_tiling(M, N, K, dtype)
|
||||
candidates.append({"branch":"IterBatch","num_cores":num_cores,"B_core":B_core,
|
||||
"M_core":M,"N_core":N,"K_core":K,"baseM":baseM,"baseN":baseN,"baseK":baseK})
|
||||
return candidates
|
||||
```
|
||||
|
||||
### 6.5 ASW_Basic 分支
|
||||
|
||||
#### 6.5.1 定位
|
||||
|
||||
ASW_Basic 是**不切 K 分支的通用框架**。它覆盖了完备性分析中 {M}、{N}、{B,M}、{B,N}、{M,N}、{B,M,N} 六种组合，以及 {B} 中不能归入 MergeBatch/IterBatch 的剩余情况。**ASW_Basic 不要求 B=1——B 可以很大，也可以很小，也可以等于 1。**
|
||||
|
||||
#### 6.5.2 ASW_Basic 与 IterBatch 的竞争关系
|
||||
|
||||
当 B ≥ 32 时，IterBatch（切 B）和 ASW_Basic（切 B 或切 M/N）都是可行的。需要比较：
|
||||
|
||||
- **IterBatch（切 B=32）**：每核 1 个 batch，核内做完整 M×N 的 Matmul。若 M×N 很大，核内 M/N tiling 可能导致 L2 反复替换——同一个 batch 的 A 和 B 在 M/N 循环间需要反复从 GM 读取。
|
||||
- **ASW_Basic（切 M=32 或切 N=32）**：每核负责 M/32 行或 N/32 列。若切 M，右矩阵 [K, N] 被所有核共享。若 K×N×sizeof ≤ 128MB，右矩阵可驻留 L2，所有核以 5.2TB/s 共享读取，无 GM 带宽压力。
|
||||
|
||||
**决策依据**：当单 batch 的 M×N 很大，使得核内 M/N tiling 时的 L2 替换开销超过切 M/N 时的输入共享开销时，ASW_Basic 优于 IterBatch。定量判断：
|
||||
|
||||
```
|
||||
IterBatch 的 L2 压力：单核负责 M×N×K 的完整 Matmul，M/baseM 次 M 循环 + N/baseN 次 N 循环
|
||||
若 M_core×K_core > L2_SIZE，A 矩阵在 M 循环间无法驻留 L2 → 每次 M 循环都要从 GM 重读 A
|
||||
|
||||
ASW_Basic 切 M 的 L2 压力：每核负责 (M/32)×N×K，右矩阵 [K,N] 被 32 核共享
|
||||
若 K×N×sizeof ≤ L2_SIZE，右矩阵驻留 L2 → 零 GM 读取
|
||||
若 K×N×sizeof > L2_SIZE，右矩阵需从 GM 重复读取 → 与 IterBatch 的 L2 替换类似
|
||||
|
||||
决策规则：
|
||||
if K×N×sizeof ≤ L2_SIZE 且 M_core×K_core > L2_SIZE:
|
||||
ASW_Basic 切 M 优于 IterBatch
|
||||
elif K×M×sizeof ≤ L2_SIZE 且 N_core×K_core > L2_SIZE:
|
||||
ASW_Basic 切 N 优于 IterBatch
|
||||
else:
|
||||
IterBatch 优于 ASW_Basic（切 B 零共享，无额外开销）
|
||||
```
|
||||
|
||||
#### 6.5.3 核间切分维度选择
|
||||
|
||||
按代价从低到高：
|
||||
|
||||
1. **切 B**：`grid_B = num_cores`，零共享。要求 `B ≥ num_cores`
|
||||
2. **切 M**：`grid_M = num_cores`，右矩阵共享。要求 `B=1 或 B < num_cores`
|
||||
3. **切 N**：`grid_N = num_cores`，左矩阵共享。要求 `B=1 或 B < num_cores`
|
||||
4. **切 B×M / B×N / M×N**：混合切分，双向共享
|
||||
|
||||
#### 6.5.4 候选生成
|
||||
|
||||
```
|
||||
def candidate_asw_basic(B, M, N, K, dtype):
|
||||
candidates = []
|
||||
# 切B
|
||||
if B >= 32:
|
||||
for nc in [32, 16, 8]:
|
||||
baseM, baseN, baseK = select_tiling(M, N, K, dtype)
|
||||
candidates.append({"branch":"ASW_Basic","split":"B","num_cores":nc,
|
||||
"B_core":B//nc,"M_core":M,"N_core":N,"baseM":baseM,"baseN":baseN,"baseK":baseK})
|
||||
# 切M
|
||||
if M >= 32 and B < 32:
|
||||
for grid_M in [32, 16, 8]:
|
||||
M_core = M // grid_M
|
||||
B_in_L2 = (K * N * sizeof(dtype) <= 128*1024*1024)
|
||||
baseM, baseN, baseK = select_tiling(M_core, N, K, dtype)
|
||||
candidates.append({"branch":"ASW_Basic","split":"M","num_cores":grid_M,
|
||||
"M_core":M_core,"N_core":N,"shared_B":True,"B_in_L2":B_in_L2,
|
||||
"baseM":baseM,"baseN":baseN,"baseK":baseK})
|
||||
# 切N（对称）
|
||||
# 切M×N（B=1时）
|
||||
if B == 1 and M * N >= 32:
|
||||
for grid_M in [8, 4, 2]:
|
||||
grid_N = 32 // grid_M
|
||||
M_core = M // grid_M; N_core = N // grid_N
|
||||
baseM, baseN, baseK = select_tiling(M_core, N_core, K, dtype)
|
||||
candidates.append({"branch":"ASW_Basic","split":"MxN","num_cores":32,
|
||||
"grid_M":grid_M,"grid_N":grid_N,"M_core":M_core,"N_core":N_core,
|
||||
"baseM":baseM,"baseN":baseN,"baseK":baseK})
|
||||
# 降核
|
||||
if B * M * N < 32:
|
||||
nc = max(1, B * (M//256) * (N//256))
|
||||
# ...
|
||||
return candidates
|
||||
```
|
||||
|
||||
### 6.6 StreamK 分支
|
||||
|
||||
#### 6.6.1 定位
|
||||
|
||||
StreamK 是**切 K 分支的统一框架**，覆盖全部 8 个含 K 的完备性组合（{K}, {B,K}, {M,K}, {N,K}, {B,M,K}, {B,N,K}, {M,N,K}, {B,M,N,K}）。如知识库中 `cann-samples/matmul_story` 所述：
|
||||
|
||||
> "切K模板（StreamK）在普通模板多核并行切分M轴、N轴的基础上，进一步增加了K轴的切分"
|
||||
|
||||
StreamK 的核间切分通过 grid 配置来描述：`grid_K × grid_B × grid_M × grid_N ≤ 32`。其中 `grid_K` 为 K 维切分数，`grid_B/grid_M/grid_N` 为并行维度的切分数。纯切 K 只是 `grid_B=grid_M=grid_N=1` 的特例；切 K+B 是 `grid_K>1, grid_B>1, grid_M=grid_N=1` 的特例，以此类推。
|
||||
|
||||
#### 6.6.2 为什么含 K 的组合都可能最优
|
||||
|
||||
以 B=4, K=4096, M=N=64, 32 核为例，对比纯切 K 与切 K+B：
|
||||
|
||||
| 配置 | grid_K | grid_B | K_core | B_core | Reduce 组 |
|
||||
|---|---|---|---|---|---|
|
||||
| 纯切 K | 32 | 1 | 128 | 4 | 32 |
|
||||
| 切 K+B | 8 | 4 | 512 | 1 | 8 |
|
||||
|
||||
切 K+B 时 K_core 从 128 增大到 512（4×），Reduce 组从 32 缩小到 8（1/4）。每核的 K 范围更大意味着部分和写出的次数更少，Reduce 组更小意味着核间累加的数据量和同步开销更低。当 B≥4 且 K 不是极大时，切 K+B 可能优于纯切 K。
|
||||
|
||||
同理，切 K+M 或切 K+N 在 M 或 N 较大时也可能优于纯切 K。所有 8 个含 K 的组合都应作为候选参与评估，由时延模型选出最优。
|
||||
|
||||
#### 6.6.3 进入条件
|
||||
|
||||
StreamK 的进入条件不是"是否应该切 K"，而是"**是否需要在候选集中加入含 K 的组合**"。
|
||||
|
||||
**条件 1：不切 K 时并行度不足，或切 K 组合可能提供更优的 MTE2/L2 利用。** 即以下两者之一成立：
|
||||
- `total_tiles = B × ceil(M/16) × ceil(N/16) < 32`（不切 K 填不满核）
|
||||
- 或者：不切 K 时核内 M/N tiling 的 L2 压力大，切 K 可以减少每核的 M/N 范围
|
||||
|
||||
**条件 2：K 足够大，使得 Reduce 开销可接受。** 切 K 后每核的 MMAD 时延应远大于 Reduce 时延。设 `α = 10`（安全系数）：
|
||||
|
||||
```
|
||||
T_MMAD_per_core = B_core × M_core × N_core × K_core × 2 / (8192 × 1.65e9)
|
||||
T_REDUCE = grid_K × B_core × M_core × N_core × 4 / 1.6e12
|
||||
|
||||
要求 T_MMAD_per_core ≥ α × T_REDUCE
|
||||
→ K_core ≥ α × grid_K × 4 × 8192 × 1.65e9 / (2 × 1.6e12)
|
||||
= 10 × grid_K × 4 × 8192 × 1.65 / (2 × 1.6)
|
||||
≈ grid_K × 1690
|
||||
|
||||
简化：K / grid_K ≥ grid_K × 1690 → K ≥ grid_K² × 1690
|
||||
|
||||
对于 grid_K=32: K ≥ 1.73M（极严格，仅极大 K 满足）
|
||||
对于 grid_K=16: K ≥ 432K
|
||||
对于 grid_K=8: K ≥ 108K
|
||||
对于 grid_K=4: K ≥ 27K
|
||||
对于 grid_K=2: K ≥ 6.8K
|
||||
```
|
||||
|
||||
这说明：grid_K 越大，对 K 的要求越高。StreamK 内部的 grid 搜索自然会淘汰那些 Reduce 开销过大的配置。
|
||||
|
||||
**条件 3：存在至少一个 grid 配置使 K_core ≥ 256。** 单核 K 范围太小则 tiling 效率低。
|
||||
|
||||
```
|
||||
K / grid_K ≥ 256 → grid_K ≤ K / 256
|
||||
```
|
||||
|
||||
#### 6.6.4 候选生成（含 grid 搜索）
|
||||
|
||||
```
|
||||
def candidate_streamk(B, M, N, K, dtype):
|
||||
candidates = []
|
||||
# 尝试所有可行的 grid 配置
|
||||
for grid_K in [32, 16, 8, 4, 2, 1]:
|
||||
if grid_K > K / 256: continue # 条件3: K_core≥256
|
||||
remaining = 32 // grid_K
|
||||
# 在 remaining 个核中分配 B/M/N
|
||||
for grid_B in factors_of(min(remaining, B)):
|
||||
rem2 = remaining // grid_B
|
||||
for grid_M in factors_of(min(rem2, M/16)):
|
||||
grid_N = rem2 // grid_M
|
||||
if grid_N > N/16: continue
|
||||
|
||||
B_core = ceil(B / grid_B)
|
||||
M_core = ceil(M / grid_M)
|
||||
N_core = ceil(N / grid_N)
|
||||
K_core = K / grid_K
|
||||
|
||||
baseM, baseN, baseK = select_tiling(M_core, N_core, K_core, dtype)
|
||||
|
||||
# 条件2: 评估 Reduce 是否可接受
|
||||
T_MMAD = eval_mmad_local(B_core, M_core, N_core, K_core, baseM, baseN, baseK, dtype)
|
||||
T_REDUCE = grid_K * B_core * M_core * N_core * 4 / 1.6e12
|
||||
if T_MMAD < 10 * T_REDUCE: continue # Reduce 占比过大，淘汰
|
||||
|
||||
candidates.append({
|
||||
"branch": "StreamK",
|
||||
"grid_K": grid_K, "grid_B": grid_B,
|
||||
"grid_M": grid_M, "grid_N": grid_N,
|
||||
"B_core": B_core, "M_core": M_core,
|
||||
"N_core": N_core, "K_core": K_core,
|
||||
"baseM": baseM, "baseN": baseN, "baseK": baseK,
|
||||
"reduce_size": grid_K * B_core * M_core * N_core * 4
|
||||
})
|
||||
return candidates
|
||||
```
|
||||
|
||||
#### 6.6.5 与不切 K 分支的关系
|
||||
|
||||
当不切 K 的分支（MergeBatch/IterBatch/ASW_Basic）的候选 T_total 已经足够小（例如已经 Cube Bound），且 StreamK 候选的 Reduce 开销使其无法超越时，StreamK 自然被淘汰。只有当不切 K 的分支因并行度不足或 L2 压力大而表现不佳时，StreamK 才会胜出。
|
||||
|
||||
### 6.7 分支选择决策算法
|
||||
|
||||
```
|
||||
def decide(B, M, N, K, dtype, BatchA, BatchB):
|
||||
# Step 1: 桥接分支
|
||||
if BatchA == 1 or BatchB == 1:
|
||||
return fold_to_matmul(B, M, N, K, dtype, BatchA, BatchB)
|
||||
|
||||
# Step 2: 收集所有候选
|
||||
all_candidates = []
|
||||
all_candidates.extend(candidate_mergebatch(B, M, N, K, dtype))
|
||||
all_candidates.extend(candidate_iterbatch(B, M, N, K, dtype))
|
||||
all_candidates.extend(candidate_asw_basic(B, M, N, K, dtype))
|
||||
all_candidates.extend(candidate_streamk(B, M, N, K, dtype))
|
||||
|
||||
# Step 3: 评估时延
|
||||
for cand in all_candidates:
|
||||
cand["T_total"] = max(
|
||||
eval_mmad(cand), eval_mte2(cand), eval_mte1(cand),
|
||||
eval_fixpipe(cand), eval_reduce(cand)
|
||||
)
|
||||
|
||||
# Step 4: 选择最优
|
||||
all_candidates.sort(key=lambda c: c["T_total"])
|
||||
return all_candidates[0]
|
||||
```
|
||||
|
||||
### 6.8 决策示例
|
||||
|
||||
**Case A: B=32, BatchA=BatchB=32, M=N=4096, K=4096, BF16**
|
||||
|
||||
```
|
||||
Step 1: BatchA=32≠1, BatchB=32≠1 → 不走桥接
|
||||
|
||||
MergeBatch: M×N=16M > 128×128 → 条件2不满足 → 无候选
|
||||
|
||||
IterBatch: num_cores=32, B_core=1, 每核1 batch
|
||||
T_MMAD = 4096³/(8192×1.65e9) = 5.10ms, Cube Bound ✓
|
||||
|
||||
ASW_Basic: B=32 → 切B, num_cores=32, 与 IterBatch 等效
|
||||
|
||||
StreamK: total_tiles=32×256×256 >> 32 → 条件1不满足 → 无候选
|
||||
|
||||
→ 选择 IterBatch（切B=32），T_total = 5.10ms
|
||||
```
|
||||
|
||||
**Case B: B=64, BatchA=BatchB=64, M=N=64, K=256, BF16**
|
||||
|
||||
```
|
||||
Step 1: BatchA=64≠1, BatchB=64≠1 → 不走桥接
|
||||
|
||||
MergeBatch: M×N=4096 < 128×128, B_core=2, bL0=2, kL0=256
|
||||
waste=50%, AI≈0.33 → 访存Bound → waste_hidden
|
||||
T_MMAD = 128×128×256/(8192×1.65e9) = 0.31ms
|
||||
|
||||
IterBatch: M×N=4096 < 64×64 → 条件2不满足
|
||||
|
||||
ASW_Basic: B=64 → 切B, 与 IterBatch 等效
|
||||
|
||||
StreamK: total_tiles=64×4×4=1024 ≥ 32 → 条件1不满足
|
||||
|
||||
→ 选择 MergeBatch(num_cores=32, B_core=2, bL0=2)，T_total ≈ 0.31ms
|
||||
```
|
||||
|
||||
**Case C: B=2, BatchA=1, BatchB=2, M=2048, N=2048, K=2048, BF16**
|
||||
|
||||
```
|
||||
Step 1: BatchA=1 → 走桥接 FoldBatchToMatMul
|
||||
等效 MatMul: [2048, 2048] × [2048, 4096] → [2048, 4096]
|
||||
→ Matmul 决策体系（切 N 到 32 核）→ T_total ≈ 5.0ms
|
||||
+ Split 开销 ≈ 3μs（可忽略）
|
||||
|
||||
若不走桥接：B=2 → ASW_Basic 切 B=2，只用 2 核 → T_total ≈ 80ms
|
||||
→ 桥接分支加速 16×
|
||||
```
|
||||
|
||||
**Case D: B=32, BatchA=BatchB=32, M=N=2048, K=65536, BF16**
|
||||
|
||||
```
|
||||
Step 1: BatchA=32≠1, BatchB=32≠1 → 不走桥接
|
||||
|
||||
MergeBatch: M×N=4M > 128×128 → 条件2不满足
|
||||
|
||||
IterBatch: num_cores=32, B_core=1, 每核1 batch
|
||||
M_core=2048, N_core=2048, K_core=65536
|
||||
M/baseM=8, N/baseN=8 → L2 压力大
|
||||
T_MMAD = 2048²×65536/(8192×1.65e9) = 20.4ms
|
||||
|
||||
ASW_Basic: 切 M, grid_M=32, M_core=64
|
||||
M_core=64, N_core=2048, K_core=65536
|
||||
右矩阵 [65536, 2048]×2 = 256MB > 128MB → 不驻留 L2
|
||||
→ 共享代价高，不如 IterBatch
|
||||
|
||||
StreamK: total_tiles=32×128×128=524288 ≥ 32 → 条件1不满足
|
||||
|
||||
→ 选择 IterBatch（切B=32），T_total ≈ 20.4ms
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 7. 典型 Case 时延评估
|
||||
|
||||
### 7.1 Case 1: B=32, M=N=K=4096, BF16（LLM FFN 层）
|
||||
|
||||
**最优分支：IterBatch（切 B=32）**
|
||||
|
||||
```
|
||||
每核: B_core=1, M_core=4096, N_core=4096, K_core=4096
|
||||
Tiling: baseM=256, baseN=256, baseK=128
|
||||
|
||||
T_MMAD = 4096³ / (8192 × 1.65e9) = 5.10ms
|
||||
|
||||
Size_A_MTE2 = 1 × (4096/256) × 4096 × 4096 × 2 = 2.15GB
|
||||
Size_A_GM = 1 × 4096 × 4096 × 2 = 33.6MB (首次，1.6TB/s)
|
||||
Size_A_L2 = 2.15GB - 33.6MB ≈ 2.12GB (重复，5.2TB/s)
|
||||
BW_MTE2_A = 2.15GB / (33.6MB/1.6TB/s + 2.12GB/5.2TB/s) = 5.0TB/s
|
||||
T_MTE2_A = 2.15GB / 5.0TB/s = 0.43ms
|
||||
|
||||
Size_B_MTE2 = 1 × (4096/256) × 4096 × 4096 × 2 = 2.15GB
|
||||
T_MTE2_B = 0.43ms
|
||||
|
||||
T_MTE2 = max(0.43, 0.43) = 0.43ms
|
||||
|
||||
T_FIXPIPE = 4096² × 2 / 5.2e12 + 4096² × 2 / 1.6e12 = 0.006 + 0.021 = 0.027ms
|
||||
|
||||
T_total = max(5.10, 0.43, ~0.3, 0.027) = 5.10ms → Cube Bound ✓
|
||||
```
|
||||
|
||||
### 7.2 Case 2: B=64, M=N=64, K=256, BF16
|
||||
|
||||
**最优分支：MergeBatch（num_cores=32, B_core=2, bL0=2）**
|
||||
|
||||
```
|
||||
每核: B_core=2, M=64, N=64, K=256
|
||||
MergeBatch: bL0=2, bM=128, bN=128, kL0=256
|
||||
|
||||
Cube 实际计算: [128, 256] × [256, 128] → [128, 128]
|
||||
有效输出: 2 × [64, 64]（对角线 BlockTrace）
|
||||
算力浪费: (128×128 - 2×64×64) / (128×128) = 50%
|
||||
|
||||
T_MMAD = 128 × 128 × 256 / (8192 × 1.65e9) = 0.31ms
|
||||
|
||||
Size_A_MTE2 = 128 × 256 × 2 × 1 = 64KB; Size_B_MTE2 = 64KB
|
||||
T_MTE2 ≈ 128KB / 5.2TB/s ≈ 0.02μs
|
||||
|
||||
T_total ≈ 0.31ms → Cube Bound
|
||||
```
|
||||
|
||||
### 7.3 Case 3: B=1, M=64, N=64, K=65536, BF16
|
||||
|
||||
**最优分支：StreamK（切 K=32）**
|
||||
|
||||
```
|
||||
每核: K_core=2048, M_core=64, N_core=64
|
||||
Tiling: baseM=64, baseN=64, baseK=256
|
||||
|
||||
T_MMAD = 64 × 64 × 2048 / (8192 × 1.65e9) = 0.61μs
|
||||
|
||||
Size_A_MTE2 = 64 × 2048 × 2 = 256KB; Size_B_MTE2 = 64 × 2048 × 2 = 256KB
|
||||
T_MTE2 = 256KB × 2 / 5.2TB/s = 0.1μs
|
||||
|
||||
T_REDUCE = 32 × 64 × 64 × 4 / 1.6e12 = 0.33μs
|
||||
|
||||
T_total ≈ 0.61μs → Cube Bound
|
||||
```
|
||||
|
||||
> 若用 ASW_Basic 切 M（B=1 无法切 B），M=64 仅能切 1 核 → T_total = 19.5ms，StreamK 加速 32000×。
|
||||
|
||||
---
|
||||
|
||||
## 8. 实现路线图
|
||||
|
||||
### 8.1 推荐技术栈
|
||||
|
||||
- **模板库**：catlass v1.6.0（已支持 Ascend 950PR，含 BlockMmad pingpong mutex、EVG 声明式后处理、ascend950 Tile 层组件）
|
||||
- **算子仓参考**：`ops-nn/matmul/batch_mat_mul_v3`（950PR 的 arch35 目录）
|
||||
- **编程接口**：AscendC Matmul 高阶 API（Tensor API 优先，基础 API 兜底）
|
||||
- **编译**：CANN 9.0.0（bisheng 编译器，C++ 17）
|
||||
|
||||
### 8.2 实现步骤
|
||||
|
||||
1. **Shape 分类器**：实现第 6 章的决策框架，输入 shape/dtype，输出推荐分支和 tiling 参数
|
||||
2. **Tiling 求解器**：实现第 3.3 节的 tiling 选择算法，输出 baseM/baseN/baseK 和 L1 切分
|
||||
3. **L2 切分决策**：判断是否需要 L2 切分，若需要则计算 `mL2TileNum`/`nL2TileNum` 和错位分核方案
|
||||
4. **Kernel 实现**：基于 catlass 的 `BlockMmad` ping-pong 模板，按分支实现参数化 Kernel
|
||||
5. **Host 侧 Tiling 逻辑**：实现核间切分映射、L2 切分策略、Swizzle 编排
|
||||
6. **时延评估器**：实现第 5 章的时延评估模型，支持两种模式
|
||||
7. **性能验证**：使用 msProf 采集各流水线实际耗时，与理论模型对比，迭代修正
|
||||
|
||||
### 8.3 参考实现
|
||||
|
||||
| 参考 | 路径 | 说明 |
|
||||
|---|---|---|
|
||||
| batch_mat_mul_v3 源码 | `代码仓/ops-nn/matmul/batch_mat_mul_v3` | 各分支的完整实现 |
|
||||
| catlass 基础 Matmul | `代码仓/catlass/examples/43_ascend950_basic_matmul` | 950PR 基本 matmul 模板 |
|
||||
| catlass StreamK | `代码仓/catlass/examples/66_ascend950_streamk_matmul` | K 维动态切分 |
|
||||
| cann-samples matmul_story | `代码仓/cann-samples/Samples/2_Performance/matmul_story` | 性能建模与调优 |
|
||||
| L2 Cache 切分案例 | CANN 9.0.0 算子实践参考 | L2 切分算法与错位分核 |
|
||||
|
||||
---
|
||||
|
||||
*文档版本：v4.0 | 适用芯片：昇腾 950PR（DAV_3510） | CANN 版本：9.0.0+*
|
||||
Reference in New Issue
Block a user