matmul 总体设计

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 × BA 为 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依然可以矩阵计算。

  1. 在K 维长,IO 延迟 >> 计算的情况下使用比较好
  2. 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:
	   ...
  }
}
相关推荐
Dream Cosmos1 小时前
C++ 中 static 关键字的作用与使用规则
linux·c++
蒸蒸yyyyzwd4 小时前
cpp 选手秋招学习笔记 day33
c++·笔记·求职招聘
hansang_IR10 小时前
【题解】P4460 [CQOI2018] 解锁屏幕
c++·算法
ChampaignWolf11 小时前
在 SAP S/4HANA 内部基于 ICF Handler 从零实现 ABAP MCP Server:架构、代码与踩坑实录
c++·架构·sap·abap·mcp
尘客-追梦12 小时前
Day 01:插件化之前,先看清 ABI 这堵墙
开发语言·c++·qt
有点。13 小时前
*C++哈夫曼树与哈夫曼编码
java·c++·servlet
啦啦啦啦啦zzzz13 小时前
利用RAII机制进行日志的采集
linux·服务器·c++·网络编程·c++20
nLif13 小时前
基于GTK的简易MFC对话框兼容层 -- MfcToGTK
linux·c++·mfc·gtk
小小、码农15 小时前
C++11右值引用与移动语义 —— 从拷贝资源到转移资源
开发语言·网络·c++·网络协议