0. 关键性结论
| # | 情况 | 策略 |
|---|---|---|
| 1 | N & NTile 大 M & Mtile 小 | Grid 分N block分M |
| 2 | N & NTile 小 M & Mtile 大 | Grid 分M block分N |
| 3 | K很小 | 增加M_tile 和 N_tile 大小 |
| 4 | K很大 | 双pingpong 计算,ping-pong 收益大,这个需要测试主要看 IO的耗时了 |
| 5 | 循环映射 M-N-K N-M-K 三层for循环 | 如果tilesize mtile <<< Ntile 建议使用N-M-K,同等cache命中数下,命中Ntile,Ntile更大数据更多。 这就要看 tile-size 和 LLC的大小 目前来说不知道 哪个是实现更优,通常情况下,命中数感觉都一样。 |
| 没有区别? | ||
| 6 | M/N 有余块 | 用 M/N 余块逻辑 |
1. 背景与目标
在深度学习大模型训练与推理中,矩阵乘法(GEMM)是核心计算模式,其性能直接影响端到端吞吐。本算子面向具备 私有 L1 缓存(容量 384KB) 和 异步数据搬运引擎(DTE) 的硬件设计,采用二维切分(grid × block)与双缓冲流水线技术,在有限 L1 资源下最大化数据复用,实现高效的 bf16 矩阵乘(C = A × B,A 为 bf16,B 为 bf16,C 为 fp32)。
新一代芯片是采用64Sip
LLC (Last Level Cache) 在计算卡中被定义为L2存储系统。LLC是分布式的结构,每个LLC都有唯一对应的地址空间。我粗浅理解就是有两个功能: 记录访存地址addr和缓存。先来先服务LRU机制。32thread 环境下LLC 容量16MB 维护一个table(key-value)。
2. 名词定义和解释
bf16: 目前用 bf16 作为类型计算。
矩阵乘法 计算公式 D = alpha *(A @ B) + C * Beta
其中矩阵A 为左矩阵, 矩阵B为右矩阵, C矩阵为偏置,D为结果矩阵。
A shape M K
B shape K N
D shape M N
Tile-size(M_L1 K_L1 N_L1),矩阵在L1分块的大小,分别是 M_L1 K_L1 N_L1 举例BF16 表示 32X32X128
Tile(M-tile N-tile K-tile): 在矩阵计算中,M K N维度上Tiling的块个数。
blockIdx&threadIdx和矩阵方向的映射,简称矩阵方向映射:矩阵涉及到两个方向:row方向 column方向。thread组织方式是有两种,blockIdx和threadIdx。那么就存在两种映射方式。
T-B模式(ThreadIdx- BlockIdx模式):由ThreadIdx去承担M-tile的计算任务,由BlockIdx去承担N-tile的计算任务。
B-T模式(BlockIdx-ThreadIdx模式):由BlockIdx去承担M-tile的计算任务,由ThreadIdx去承担N-tile的计算任务。
3. 约束
在做矩阵乘的算法设计的时候,提前是有许多硬件约束的。基于现有各个子系统的约束,从总体系统角度出发,近似逼近各个子系统约束的上限,才能近似达到总体计算能力的上限。全局视角去看,清楚约束,解决主要矛盾。
3.1 L1 约束
L1数据区的上限是384 KB。 这决定着单个 thread 能同时驻留多少 A/B/C/workspace,是 tile 尺寸上限
拿L1分配来说,针对于双buffer场景
| 名称 | 注释 | |
|---|---|---|
| A_Buffer_L1 | 2M_L1 \* K_L1 | |
| B_Buffer_L1 | 2K_L1 \* N_L1 | |
| out | K_L1 \* N_L1 | 分情况,可以使用,A_Buffer_L1空间替代。 |
| workspace |
3.2 LLC 约束
LLC(Last Level Cache)在计算卡中被定义为L2存储系统。单Die 16MB,32sip/thread。双Die32MB 64sip/thread。会有cache 命中,这样会增加数据搬运的效率。但是LLC对于程序员是无感知的。
3.3 L3约束
L3 总共144GB,64SIP
3.4 acore汇编 约束 BF6
acore层面,acore计算能力的下限 是 32x32x128,计算能力的上限是L1空间大小。
3.4 SIP 约束
64个sip 当前可用32sip测试
3.5 VACC的约束
vacc是1024个 ,一个vacc 512bit
3.6 算力和带宽
略
2. 对于Tile-Size的一些想法
以BF16为例子,32x32x256。
下限是acore的处理能力。上限是L1的大小。
充分利用汇编的计算能力。是一个比较推荐的使用方式。
4. 对于PingPong 缓冲的想法
结论:1. pingpong 不是原则,是手段。没有pingpong依然可以矩阵计算。
- 在K 维长,IO 延迟 >> 计算的情况下使用比较好
- L1 紧张需要更大 tile,IO延迟较短,数据基本在 LLC 命中。不使用比较好。
pingpong 不是原则,是手段。没有pingpong 依然可以完成矩阵计算。可以掩盖掉部分的IO时间,但是需要占用大量的L1空间。总体来说,用 L1 空间换 IO 延迟隐藏。
什么时候使用比较好。K 维长,IO 延迟 >> 计算
什么时候不用比较好,L1 紧张,需要更大 tile。IO延迟较短,数据基本在 LLC 命中。计算很短,同步开销可接受
无论哪种映射和循环,K 内层 ping-pong 结构一致。
Ping-pong流水图
cpp
K-tile: j=0 j=1 j=2 j=3
│ │ │ │
A slot0: [同步 load A0]──────────────[异步 load A2]──────────────
A slot1: [异步 load A1]──────────[异步 load A3]──────────
B slot0: [同步 load B0]──────────────[异步 load B2]──────────────
B slot1: [异步 load B1]──────────[异步 load B3]──────────
│ │ │ │
wait: skip wait A1/B1 wait A2/B2 wait A3/B3
prefetch: A1/B1 -> slot1 A2/B2 -> slot0 A3/B3 -> slot1 none
addmm: [addmm0] [addmm1] [addmm2] [addmm3]
用 slot0 用 slot1 用 slot0 用 slot1
store: store C
deslice: C -> L3
cpp
j=0: [同步 A0/B0] [异步预取 A1/B1] [addmm0]
j=1: [等待 A1/B1] [异步预取 A2/B2] [addmm1]
j=2: [等待 A2/B2] [异步预取 A3/B3] [addmm2]
j=3: [等待 A3/B3] [addmm3 + store + deslice]
5. 对于矩阵方向的映射的想法
尝试去讨论,针对于同一个矩阵数据形状,用不同的矩阵方向的映射。会有什么区别?从而得出针对于不同形状矩阵用什么典型的处理方式。
结论:
又因为左tile << 右tile 命中的数据量较大的一方,性能更好。
所以 建议在此 shape下 使用 用T-B模式,反之使用B-T模式。
实验假设
A、B 矩阵情况 M = 256 = 32*8 , N = 2048 = 256 *8
Tile-Size 情况 BF16 32x256x256 来说
假设<<<8,8 >>> kernel launch
实验1 用T-B模式(ThreadIdx- BlockIdx模式)计算
结果 左tile 8 miss 右tile 1miss 7 命中
| threadIdx | cache | Block | cache |
|---|---|---|---|
| 0 | miss | 0 | miss |
| 1 | miss | 1 | 命中 |
| 2 | miss | 2 | 命中 |
| 3 | miss | 3 | 命中 |
| 4 | miss | 4 | 命中 |
| 5 | miss | 5 | 命中 |
| 6 | miss | 6 | 命中 |
| 7 | miss | 7 | 命中 |
实验2 用B-T模式(BlockIdx-ThreadIdx模式)
结果 左tile 1miss 7 命中 右tile 8 miss
| threadIdx | cache | Block | cache |
|---|---|---|---|
| 0 | miss | 0 | miss |
| 1 | 命中 | 1 | miss |
| 2 | 命中 | 2 | miss |
| 3 | 命中 | 3 | miss |
| 4 | 命中 | 4 | miss |
| 5 | 命中 | 5 | miss |
| 6 | 命中 | 6 | miss |
| 7 | 命中 | 7 | miss |
又因为左tile << 右tile
所以 建议在此 shape下 使用 用T-B模式,反之使用B-T模式。
4. 对于循环层级的想法
刚才我们针对不同形状使用什么矩阵方向的映射进行了一个讨论。接下来我们对循环层级进行一个讨论。首先,限定矩阵方向的映射只用用T-B模式(ThreadIdx- BlockIdx模式)计算。这种情况下M-N-K和 N-M-K 对于计算的影响。
实验1 循环 M-N-K
情况一 假设 cache 无穷大 情况二 假设cache 只比 shape 大 1bit
| m | n | k | m | n | k | |
| 0 不命中 | 0 不 | 0 不 | 0 不 | |||
| 0 命中 | 1 不 | 0 命 | 1 不 | |||
| 1 不命中 | 0 命 | 1 不 | 0 不 | |||
| 1 命中 | 1 命 | 1 命 | 1 不 | |||
实验2 循环 N-M-K
情况一 假设 cache 无穷大 情况二 假设cache 只比 shape 大 1bit
| m | n | k | m | n | k | |
| 0 不命中 | 0 不 | 0 不 | 0 不 | |||
| 1 不命中 | 0 命 | 0 不 | 1 命 | |||
| 0 命中 | 1 不 | 1 不 | 0 不 | |||
| 1 命中 | 1 命 | 1 不 | 1 命 | |||
结论是, 看起来没有太大的区别,需要继续深入讨论
4. 对于LLC空间的使用的一些想法
略
6. 对于参数配置化的想法
jsx
template<bool M_TO_BLOCK, bool N_OUTER>
__global__ void matmul_flex(...) {
if constexpr (M_TO_BLOCK) {
// 模式 B:grid 切 M,thread 切 N
} else {
// 模式 A:grid 切 N,thread 切 M
}
if constexpr (N_OUTER) {
// NMK:N 外层,M 内层
} else {
// MNK:M 外层,N 内层
}
}
Host 侧 选择策略:
if (M > N) {
launch matmul_flex<true, true>; // 模式 B + NMK
} else {
launch matmul_flex<false, false>; // 模式 A + MNK
}
策略的枚举
cpp
enum class MatmulStrategy {
B_NMK, // 模式 B + NMK
B_MNK, // 模式 B + MNK
A_NMK, // 模式 A + NMK
A_MNK, // 模式 A + MNK
};
kernel 模板
cpp
template<int TILE_M, int TILE_N, int TILE_K,
bool M_TO_BLOCK, bool N_OUTER>
__global__ void matmul_flex(float *out,
__bf16 *L_ptr,
__bf16 *R_ptr,
__bf16 *bias_ptr,
int M, int K, int N) {}
策略选择
cpp
MatmulStrategy select_strategy(int M, int N, int K,
int TILE_M, int TILE_N, int TILE_K,
int max_grid, int max_block) {
}
lauch
cpp
void launch_matmul(float *out, __bf16 *L, __bf16 *R, __bf16 *bias,
int M, int K, int N) {
constexpr int TILE_M = 32;
constexpr int TILE_N = 256;
constexpr int TILE_K = 256;
MatmulStrategy strat = select_strategy(M, N, K, TILE_M, TILE_N, TILE_K,
max_grid, max_block);
dim3 grid, block;
switch (strat) {
case MatmulStrategy::B_NMK:
grid = dim3(ceil(M / TILE_M));
block = dim3(ceil(N / TILE_N));
matmul_flex<TILE_M, TILE_N, TILE_K, true, true><<<grid, block>>>(out, L, R, bias, M, K, N);
break;
case MatmulStrategy::B_MNK:
...
case MatmulStrategy::A_NMK:
...
case MatmulStrategy::A_MNK:
...
}
}