admin/matmul-analysis
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
00c4a01a2d5078f9ffe397ef9783a4430c0d243b
matmul-analysis/BMM最优软件实现方案设计.md

39 KiB
Raw Blame History

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.850.92；≤8 核并发、单核数据量小 → η ≈ 0.40.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+