AICore约束-hyper-parallel-AtomGit

Open

AICore约束 #4

hedongdongAdd on Aug 18

Aug 18 Create

先给结论：这里没有统一的“切到多少就不会落HBM”。容量约束决定的是微块的最大尺寸；指令粒度、对齐、带宽利用率又决定它不能太小。

对当前910B/arch22和MegaMoE，可以看到三种不同的实际微块：

Dispatch内部搬运块：8 KiB
GMM Cube微块：      128 × 256 × 64
SwiGLU Vector微块： 19 × 512

它们分别由ACL SHMEM预留UB、L0A/B/C容量、SwiGLU活跃Buffer数量决定。

910B/arch22片上存储容量

本仓CANN依赖中的arch22平台信息是：

存储层	容量	范围/用途
AIV UB	192 KiB	每个AIV本地，Vector输入、输出、临时变量
AIC L1	512 KiB	每个AIC本地，Cube数据复用和中转
AIC L0A	64 KiB	Cube左矩阵
AIC L0B	64 KiB	Cube右矩阵
AIC L0C	128 KiB	Cube FP32累加结果
L2	32 MiB	芯片共享、硬件管理的GM Cache
ICache	32 KiB	每核代码Cache
DCache	32 KiB	每核Scalar数据Cache
BiasTable	1 KiB	Cube Bias

当前CANN平台测试信息可见 test_foreach_sigmoid_tiling.cpp。

不过需要注意：

UB/L1/L0是软件显式管理的本地Buffer，不是普通CPU意义上的Cache。
L2才是硬件管理的共享Cache。
算子能使用的UB通常小于192 KiB，因为还要扣除队列、临时变量、同步区以及通信库预留空间。
官方建议通过 GetCoreMemSize 获取目标设备实际容量，不要跨产品硬编码。Ascend C架构说明

对齐约束是：

UB：       32 Byte
L1：       32 Byte
L0A/L0B： 512 Byte
L0C：      64 Byte

GM单次搬运通常达到16 KiB以上才更容易发挥带宽，内轴最好按128/256/512 Byte对齐。Ascend C硬件约束

为什么GMM是128×256×64

当前：

baseM = 128
baseN = 256
baseK = 64
输入BF16/FP16：2 Byte
L0C累加FP32：  4 Byte

L0C约束

一个输出累加微块占用：

[
128 \times 256 \times 4
=131072\ \text{Bytes}
=128\ \text{KiB}
]

正好占满L0C：

C tile = 128×256 FP32 = 128 KiB
L0C容量                = 128 KiB

所以当前无法对L0C进行双缓冲。CANN的MatMul tiling也明确通过：

baseM * baseN * 4 * dbL0C <= l0CSize

判断能否双缓冲，见 batch_mat_mul_v3_base_tiling.cpp。

L0A约束

双缓冲A：

[
2 \times 128 \times 64 \times 2
=32768\ \text{Bytes}
=32\ \text{KiB}
]

低于64 KiB的L0A容量。

L0B约束

双缓冲B：

[
2 \times 64 \times 256 \times 2
=65536\ \text{Bytes}
=64\ \text{KiB}
]

正好占满L0B。

因此这组参数非常典型：

A双缓冲：32 KiB / 64 KiB
B双缓冲：64 KiB / 64 KiB
C单缓冲：128 KiB / 128 KiB

对于：

[10000,7168] × [7168,4096]

一个Cube微块实际上是：

A：128 × 64
B：64 × 256
C：128 × 256

K方向：

7168 / 64 = 112轮

这112轮的部分和一直累加在同一块L0C中：

K0=0..63       → 累加到L0C
K1=64..127     → 继续累加
...
K111           → 得到完整128×256结果

中间112次K分块不会把部分和写回GM。这就是“避免中间结果落HBM”真正发生的位置。

但是K全部规约完以后，128×256结果仍然会通过FixPipe写到GM。

SwiGLU实际不是128个Token一次装进UB

外层SwiGLU Task是：

128 token × 4096输入
→ 拆成左右两个128×2048

但当前生成的 swiglu_tiling.bin 显示：

rowLen     = 128
colLen     = 2048
baseRowLen = 19
baseColLen = 512
doubleBuffer = 1

也就是实际Vector微块：

19 token × 512对SwiGLU元素

一个128-token Task内部大约会拆成：

[
\lceil128/19\rceil \times \lceil2048/512\rceil
=7\times4
=28
]

个Vector微块。

为什么是 19×512？BF16 SwiGLU双缓冲时，每个配对元素需要：

3组BF16队列 × 2 Byte × 双缓冲 = 12 Byte
2组FP32临时Buffer × 4 Byte     =  8 Byte
总计                           = 20 Byte/元素

那么：

[
192\text{ KiB}/20
=9830\text{个元素左右}
]

实际选择：

[
19\times512=9728
]

占用：

[
9728\times20
=194560\text{ Bytes}
=190\text{ KiB}
]

剩余约2 KiB处理对齐和其他开销。

对应的UB计算逻辑见 swi_glu_tiling.cpp。

所以：

128 token：调度Task粒度
19×512：  真正装入UB的Vector计算微块

Dispatch当前只使用16 KiB通信UB

虽然物理UB是192 KiB，但ACL SHMEM默认只给MTE通信预留：

DEFAULT_MTE_UB_SIZE = 16 * 1024;

见 shmem_init.cpp。

Dispatch使用ping-pong，因此每块：

16 KiB / 2 = 8 KiB

对于一个128-token、hidden=7168、BF16的Task：

总数据 = 128 × 7168 × 2
       = 1,835,008 Byte
       = 1.75 MiB

内部搬运次数：

[
1835008/8192=224
]

也就是：

一个128-token Dispatch Task
    ↓
224个8 KiB GM→UB→远端GM搬运块
    ↓
全部完成后发一次signal

这里的8 KiB只是通信暂存块，不是新的调度Task。

“不落HBM”要区分两种情况

同一计算核内部

例如GMM的112个K分块：

A/B：GM→L1→L0A/L0B
C部分和：始终留在L0C
K归约完成后：L0C→GM

这里确实避免了每个K分块都写HBM。

满足条件主要是：

[
db_A M_bK_bs_A \le L0A
]

[
db_B K_bN_bs_B \le L0B
]

[
db_C M_bN_bs_C \le L0C
]

AIC到AIV之间

910B/arch22是AIC、AIV分离架构，官方架构描述中两者通过Global Memory交换数据。NPU架构版本220x

所以当前：

GMM：
L0C → FixPipe → GM

SwiGLU：
GM → UB → Vector

当前GMM和SwiGLU之间并没有把结果直接从L0C交给AIV的UB。严格从编程模型看，中间结果已经写入GM。

但是“写入GM地址”不一定意味着每个字节都立即到物理HBM，因为L2可能命中：

AIC写GM → L2
AIV很快读GM → L2命中

问题是L2只有32 MiB，而且是共享、硬件管理的，无法仅凭tile尺寸保证驻留。

例如一个专家：

10000 × 4096 × 2 Byte
= 78.125 MiB

已经远大于32 MiB L2。如果当前必须等整个专家GMM完成后再启动SwiGLU，那么前面的输出大概率无法全部留在L2。

如果改成128-token粒度及时交给SwiGLU：

128 × 4096 × 2 Byte = 1 MiB

这1 MiB切片更可能在被驱逐前由AIV消费，从而减少真实HBM流量。但仍然是“GM地址上的L2复用”，不是严格的L0C→UB片上直连。

最终受什么限制

选取内部计算块时需要同时满足：

容量：所有活跃输入、输出、临时变量、双缓冲必须装入UB/L0/L1。
数据类型：BF16/FP16输入与FP32累加占用不同。
指令粒度：Cube基本矩阵粒度通常按16对齐。
搬运对齐：UB 32B、L0A/B 512B、L0C 64B。
带宽效率：GM单次搬运过小无法充分利用带宽。
双缓冲：块过大会导致无法ping-pong，计算与搬运不能重叠。
Bank冲突：即使容量装得下，UB地址排布不合理仍可能串行。
并行度：块太大导致tile数量少、部分核心空闲；块太小导致调度和Scalar开销过高。
跨核数据通路：AIC和AIV分离时，本地Buffer无法自然跨核直接共享。
事件粒度：消费者启动太晚，即便单个切片能放L2，也可能已被驱逐。

因此所谓“最小块”其实是两个边界之间的折中：

块不能太大：否则装不进片上Buffer
块不能太小：否则搬运、指令和调度效率下降

对当前MegaMoE最直观的三层答案就是：

调度层：Dispatch/SwiGLU按128 token生成Task
算子层：SwiGLU继续切成19×512；Dispatch继续切成8 KiB
Cube层：GMM按128×256×64计算，K部分和留在L0C
跨AIC/AIV：仍需经过GM，主要争取L2命中，并非严格不落GM/HBM