机器心脏报告

机器之心编辑部

提高GPU利用率，就这么简单。

人工智能的快速发展伴随着大量的计算量。这自然引出了一个问题：如何降低AI的计算需求，提高现有AI的计算效率。

为了回答这个问题，斯坦福大学的研究人员在博客《GPUs Go Brrr》中给出了答案。

博客地址：https://hazyresearch.stanford.edu/blog/2024-05-12-tk

文章主要关注两个问题：第一，真正需要的硬件是什么？二、如何满足硬件要求？

文章讨论了如何让GPU 运行得更快，并发布了一个库ThunderKittens，它允许用户在CUDA 上轻松编写快速的深度学习内核。它具有以下特点：

简单，ThunderKittens 很容易写。可扩展性，如果用户需要ThunderKittens无法提供的功能，可以扩展功能。高速。 GitHub 链接：https://github.com/HazyResearch/ThunderKittens

ThunderKittens 使一些困难的事情变得非常简单，允许在现代硬件上实现非常高的利用率。在项目中，作者使用ThunderKittens为RTX 4090编写了一个简单的FlashAttention-2内核。代码总共有58行代码（不包括空格）。结果显示，ThunderKittens 在RTX 4090 上实现了大约122 TFLOP（理论最大值的74%）。另外，当内核程序只有100行时，ThunderKittens性能比H100上的FlashAttention-2高出30%左右。

Nvidia H100 有一些怪癖

本研究重点关注NVIDIA H100，但所提供的信息也适用于其他GPU。

H100 SXM GPU 包括：

80 GB HBM3，带宽3 TB/s（实际上更少）； 50 MB L2 缓存，带宽12 TB/s，在GPU 上分为两个25MB 部分，通过crossbar 连接； 132 流式多处理（SM，流式多处理器）。除了上述之外，H100 SXM GPU还有很多需要注意的地方，比如内存控制器、指令缓存等。

研究人员表示，保持张量核心平稳运行并不容易。他们在AI 硬件中发现了一些怪癖，其中许多也适用于非H100 GPU，但H100 特别棘手。 （相比之下，RTX 4090 非常易于使用。）这些怪癖包括：

需要WGMMA指令，但用起来也很烦人；共享内存实际上并没有那么快，需要非常小心；地址生成成本高昂；占用仍然有帮助，寄存器通常是关键资源。本文进一步描述了这些GPU 怪癖的具体情况。

WGMMA 指令很烦人

H100 有一组新的指令，称为“扭曲组矩阵乘法累加，WGMMA”（PTX 中的wgmma.mma_async，或SASS 中的HGMMA/IGMMA/QGMMA/BGMMA）。以前的GPU 上可用的张量核心指令是wmma.mma.sync 和mma.sync 。通过这些指令，SM 单个象限上的32 个线程将同步将其数据块馈送到张量核心并等待结果。

与wgmma.mma_async 指令不同，128 个连续线程（分布在SM 的所有象限上）协作直接从共享内存（以及可选的寄存器）异步同步和启动矩阵乘法。

在基准测试中，研究团队发现这些指令对于提取H100 的完整计算是必要的。如果没有它们，GPU 似乎只能达到峰值利用率的63% 左右。

共享内存

单次访问共享内存的延迟约为30 个周期，听起来可能不多，但在这段时间内，SM 的张量核心可以完成几乎两个完整的32x32 矩阵乘法运算。

处理共享内存有点棘手，因为它存储在32 个独立的内存存储中。如果你不小心，这可能会导致所谓的存储体冲突，即要求同一个存储体同时提供多个不同的内存段，导致请求被序列化，这可能会不成比例地减慢内核的速度- 而wgmma 和mma 指令所需的寄存器布局受这些存储体冲突的影响。解决方案是使用各种交错模式重新排列共享内存，以避免这些冲突。

地址生成

H100 的特点之一是张量核心和内存都足够快，仅生成用于获取数据的内存地址就占用了芯片的很大一部分资源。

NVIDIA 似乎已经意识到了这一点，因为他们为GPU 提供了张量内存加速器（或他们所说的TMA）。 TMA允许用户在全局和共享内存中指定多维张量布局，这节省了所有地址生成成本，也使构建管道变得更容易。

研究团队还发现，TMA 与wgmma.mma_async 一样，对于发挥H100 的全部潜力是完全不可或缺的。

占据

在某些方面，H100 比前几代硬件对占用率的依赖程度更低。 NVIDIA 在设计GPU 时确实考虑到了占用情况。而对于H100来说，占用率只能用有用来形容，但并没有多大作用。研究人员发现，它在A100 和RTX 4090 上变得越来越重要。

雷霆小猫

那么，如何才能更轻松地编写内核，同时仍然拥有硬件的全部功能呢？

研究团队设计了一个嵌入CUDA 的DSL，命名为ThunderKittens。

ThunderKittens 的目标是尽可能简单，包括四种模板类型：

在寄存器文件中注册tile—— 2D 张量。寄存器文件中的寄存器向量—— 1D 张量。共享内存中的共享tile—— 2D张量。 SharedVector —— 共享内存中的一维张量。平铺块通过高度、宽度和布局进行参数化，寄存器向量通过长度和布局进行参数化，共享向量仅通过长度进行参数化。这样您通常就不会遭受银行冲突的困扰。

研究团队还提供了一些必要的操作：

初始化，例如将共享向量清零

一元运算，例如exp 二元运算，例如mul 行/列运算，例如row_sum 这项研究给出了一个用ThunderKittens 编写的用于RTX 4090 的简单前向闪存注意内核：

#define NUM_WORKERS 16 //该内核每个块并行使用16 个工作线程，以帮助更快地发出指令。

使用命名空间小猫； //为了简单起见，该内核仅处理headdim=64。另外，这里n 应该是256 的倍数。

__global__ void attend_ker64(int n, const bf16* __restrict__ __q__, const bf16* __restrict__ __k__, const bf16* __restrict__ __v__, bf16* __o__) {

自动warpid=kittens:warpid();

自动block_start=blockIdx.x*(n*64);

const bf16 *_q=__q__ + block_start，*_k=__k__ + block_start，*_v=__v__ + block_start；

bf16 *_o=__o__ + block_start;

外部__shared__alignment_dummy __shm[]; //这是CUDA 共享内存

共享分配器al((int*)__shm[0]);

//K 和V 位于共享内存中——这就是所有适合的内容。

st_bf_1x4 (k_smem)[NUM_WORKERS]=al.分配， NUM_WORKERS();

st_bf_1x4 (v_smem)[NUM_WORKERS]=al.分配， NUM_WORKERS();

//初始化所有寄存器块。

rt_bf_1x4 q_reg、k_reg、v_reg； //v_reg需要交换成col_l

rt_fl_1x1 att_block；

rt_bf_1x1 att_block_mma；

rt_fl_1x4 o_reg;

rt_fl_1x1:col_vec max_vec_last，max_vec； //这些是注意力块的列向量

rt_fl_1x1:col_vecnorm_vec_last，norm_vec； //这些是注意力块的列向量

int qo_blocks=n/(q_reg.rows*NUM_WORKERS), kv_blocks=n/(q_reg.rows*NUM_WORKERS);

for(自动q_blk=0; q_blk qo_blocks; q_blk++) {

//每个扭曲加载自己的16x64 Q 块，然后乘以1/sqrt(d)

加载(q_reg, _q + (q_blk*NUM_WORKERS + warpid)*q_reg.num_elements, q_reg.cols);

mul(q_reg, q_reg, __float2bfloat16(0.125f)); //温度调节

//零闪存注意L、M 和O 寄存器。

neg_infty(max_vec); //Q 块的零寄存器

零（norm_vec）；

零（o_reg）；

//针对已加载的这些q 迭代k、v

for(自动kv_idx=0; kv_idx kv_blocks; kv_idx++) {

//每个warp 将自己的k、v 块加载到共享内存中

加载(v_smem[warpid], _v + (kv_idx*NUM_WORKERS + warpid)*q_reg.num_elements, q_reg.cols);

加载(k_smem[warpid], _k + (kv_idx*NUM_WORKERS + warpid)*q_reg.num_elements, q_reg.cols);

__syncthreads(); //我们需要确保在开始计算阶段之前加载所有内存

//现在每个warp 都会遍历所有子图块，加载它们，然后执行flash 注意内部alg。

for(int subtile=0; subtile NUM_WORKERS; subtile++) {

加载（k_reg，k_smem [subtile]）； //将k 从共享加载到寄存器中

零（att_block）； //零16x16 注意力图块

mma_ABt(att_block, q_reg, k_reg, att_block); //Q@K.T

复制（norm_vec_last，norm_vec）；

复制（max_vec_last，max_vec）；

row_max(max_vec, att_block, max_vec); //累加到max_vec

sub_row(att_block, att_block, max_vec); //从注意力中减去最大值——现在全部=0

exp(att_block, att_block); //就地对块求幂。

sub(max_vec_last, max_vec_last, max_vec); //从旧的最大值中减去新的最大值以找到新的标准化。

exp(max_vec_last, max_vec_last); //对该向量求幂——这就是我们需要标准化的。

mul(norm_vec,norm_vec,max_vec_last); //范数vec 现在已标准化。

row_sum(norm_vec, att_block,norm_vec); //将新的注意力块累积到现在重新缩放的norm_vec上

p_row（att_block，att_block，norm_vec）； //现在注意力块已正确标准化

mul(norm_vec_last,norm_vec_last,max_vec_last); //根据新的最大值标准化先前的范数vec

p（norm_vec_last，norm_vec_last，norm_vec）； //根据新范数对先前范数vec 进行归一化

复制（att_block_mma，att_block）； //mma_AB 转换为bf16

加载（v_reg，v_smem [subtile]）； //将v 从共享加载到寄存器中。

rt_bf_1x4 v_reg_col=swap_layout_inplace(v_reg); //这是一个引用，调用使v_reg 无效

mul_row（o_reg，o_reg，norm_vec_last）； //在mma_AB 之前标准化o_reg

mma_AB(o_reg, att_block_mma, v_reg_col, o_reg); //使用局部注意力@V matmul 将mfma 转移到o_reg 上。

}

__syncthreads(); //我们需要确保所有扭曲都完成，然后才能开始加载下一个kv 块

}

store(_o + (q_blk*NUM_WORKERS + warpid)*q_reg.num_elements, o_reg, q_reg.cols); //写出o.如果将d 设置为constexpr q_reg.rows :/，编译器会出现寄存器使用问题

}

}

CUDA代码总共约60行，硬件利用率为75%。虽然非常密集，但大部分复杂性在于算法，而不是混合模式或寄存器布局。

TMA、WGMMA、混合模式和描述符的复杂性如何？以下是使用ThunderKittens、H100的FlashAttention-2前向传递编写的：

模板

__global__ __launch_bounds__((NUM_WORKERS)*kittens:WARP_THREADS, 2)

void fwd_attend_ker_dim(int N, const CUtensorMap* tma_q, const CUtensorMap* tma_k, const CUtensorMap* tma_v, CUtensorMap* tma_o) {

外部__shared__ int __shm[]; //这是CUDA 共享内存

tma_swizzle_allocator al((int*)__shm[0]);

constexpr inttile_width=fwd_attend_ker_tile_dims:tile_width; //常量

constexpr int qo_height=fwd_attend_ker_tile_dims:qo_height;

constexpr int kv_height=fwd_attend_ker_tile_dims:kv_height;

st_bf (q_smem) [NUM_WARPGROUPS]=al.allocate, NUM_WARRPGROUPS();

st_bf (k_smem)[2][NUM_WORKERS_KV]=al.allocate, 2, NUM_WORKERS_KV();

st_bf (v_smem)[2][NUM_WORKERS_KV]=al.allocate, 2, NUM_WORKERS_KV();

积分tic=0，toc=1；

rt_fl1，kv_height att_block；

rt_bf1，kv_height att_block_mma；

rt_fl1，qo_height o_prev；

col_vec max_vec_last, max_vec;

col_vecnorm_vec_last，norm_vec；

int warpid=kittens:warpid();

int warpgroupid=warpid/kittens:WARRPGROUP_WARPS;

int kv_blocks=N/(NUM_WORKERS_KV*k_smem[0][0].rows);

__shared__ uint64_t qsmem_barrier, kvsmem_barrier;//, vsmem_barrier;

int q_phasebit=0;

int kv_phasebit=0;

if (threadIdx.x==0) {

tma:init_barrier, NUM_WARRPGROUPS(qsmem_barrier, 1);

tma:init_barrier, NUM_WORKERS_KV*2(kvsmem_barrier, 1);

}

如果（warpid==0）{

for (int wg=0; wg NUM_WORKERS/kittens:WARRPGROUP_WARPS; wg++) { //加载q

inttile_idx=(blockIdx.y * NUM_WARPGROUPS * gridDim.x) + (blockIdx.x * NUM_WARPGROUPS) + wg;

tma:load_async（（q_smem [wg]），tma_q，qsmem_barrier，tile_idx）；

}

for (int w=0; w NUM_WORKERS_KV; w++) { //加载k, v

inttile_idx=(blockIdx.y * NUM_WORKERS_KV * kv_blocks) + (0 * NUM_WORKERS_KV) + w;

tma:load_async（（k_smem [tic] [w]），tma_k，kvsmem_barrier，tile_idx）；

tma:load_async（（v_smem [tic] [w]），tma_v，kvsmem_barrier，tile_idx）；

}

}

neg_infty(max_vec); //Q 块的零寄存器

零（norm_vec）；

零（o_prev）；

__syncthreads();

tma:arrive_and_wait(qsmem_barrier, q_phasebit);

q_phasebit ^=1;

if constexpr (D==64) { warpgroup:mul(q_smem[warpgroupid], q_smem[warpgroupid], __float2bfloat16(0.125f)); }

否则{ warpgroup:mul(q_smem[warpgroupid], q_smem[warpgroupid], __float2bfloat16(0.08838834764f)); }

for (自动kv_idx=0; kv_idx kv_blocks; kv_idx++, tic ^=1, toc ^=1) {

tma:arrive_and_wait(kvsmem_barrier, kv_phasebit);

kv_phasebit ^=1;

__syncthreads();

如果（warpid==0）{

tma:set_bytes(kvsmem_barrier, 2 * NUM_WORKERS_KV * k_smem[0][0].num_elements * sizeof(bf16));

if (kv_idx + 1 kv_blocks) {

for (int w=0; w NUM_WORKERS_KV; w++) {

inttile_idx=(blockIdx.y * NUM_WORKERS_KV * kv_blocks) + ((kv_idx + 1) * NUM_WORKERS_KV) + w;

tma:load_async（（k_smem [toc] [w]），tma_k，kvsmem_barrier，tile_idx）；

tma:load_async（（v_smem [toc] [w]），tma_v，kvsmem_barrier，tile_idx）；

}

}

}

warpgroup:mma_fence(att_block);

扭曲组：mm_ABt(att_

block, q_smem[warpgroupid], k_smem[tic][0]); warpgroup::mma_commit_group(); copy(norm_vec_last, norm_vec); copy(max_vec_last, max_vec); warpgroup::mma_async_wait(); row_max(max_vec, att_block, max_vec); // accumulate onto the max_vec sub_row(att_block, att_block, max_vec); exp(att_block, att_block); sub(max_vec_last, max_vec_last, max_vec); exp(max_vec_last, max_vec_last); mul(norm_vec, norm_vec, max_vec_last); row_sum(norm_vec, att_block, norm_vec); // accumulate onto the norm_vec p_row(att_block, att_block, norm_vec); mul(norm_vec_last, norm_vec_last, max_vec_last); p(norm_vec_last, norm_vec_last, norm_vec); copy(att_block_mma, att_block); // convert to bf16 for mma mul_row(o_prev, o_prev, norm_vec_last); // normalize o_prev in advance of mma'ing onto it warpgroup::mma_fence(o_prev); warpgroup::mma_AB(o_prev, att_block_mma, v_smem[tic][0]); warpgroup::mma_commit_group(); } auto (*o_smem) = reinterpret_cast(*)>(q_smem); // reuse q memory warpgroup::store(o_smem[warpgroupid], o_prev); __syncthreads(); if (warpid % 4 == 0) { // store o int tile_idx = (blockIdx.y * NUM_WARPGROUPS * gridDim.x) + (blockIdx.x * NUM_WARPGROUPS) + warpgroupid; tma::store_async(tma_o, (o_smem[warpgroupid]), tile_idx); tma::store_commit_group(); } tma::store_async_wait(); } 这个内核只有 100 行代码，它在 H100 上的性能比 FlashAttention-2 高出约 30%。ThunderKittens 负责 wrap up 布局和指令，并提供一个可以在 GPU 上使用的 mini-pytorch。H100 SXM 上各种配置的 FlashAttention-2（Pytorch）与 ThunderKittens 的比较。 此外，研究团队还发布了基于线性注意力的内核和其他架构。基于线性注意力内核的运行速度为 215 TFLOP（如果考虑算法中固有的重计算，则运行速度超过 300 TFLOP）。 虽然理论上线性注意力更高效，但从实践经验来看，线性注意力在硬件上的效率大大降低。因此，ThunderKittens 有望开辟广泛的高吞吐量应用。使用 ThunderKittens 可以非常快地实现线性注意力。 tile 看起来是个好点子 在研究团队看来，ThunderKittens 之所以运行良好，是因为它不会试图做所有事情。CUDA 确实比 ThunderKittens 更有表现力，而 ThunderKittens 又小又简单。 不过，ThunderKittens 具有很好的抽象能力，它具有小的 tile，这与 AI 和硬件的发展相匹配。ThunderKittens 不支持任何少于 16 的维数。但在研究团队看来，这一点并不重要，尤其对于硬件而言。如果你的矩阵乘法小于 16x16，你确定自己做的还是 AI 吗？ 从哲学的视角来看，研究团队认为框架迁移是合理的。「寄存器」当然不应该像旧 CPU 那样的 32 位。CUDA 使用的 1024 位宽向量寄存器无疑朝着正确方向迈出了一步。但对研究团队而言，「寄存器」是 16x16 的数据 tile。他们认为 AI 想要这样，它仍然只是矩阵乘法、规约和重塑。当然硬件也想要这样，小的矩阵乘法寻求硬件支持，而不仅仅是 systolic mma。 实际上，从更广泛的视角来看，研究团队认为应该围绕硬件的良好映射来重新调整 AI 思路。比如，循环状态应该有多大？SM 能够容纳多大尺寸？计算密度是多少？这些都不亚于硬件的要求。 研究团队表示，这项工作未来的一个重要方向是利用他们对硬件的了解来帮助设计与硬件相匹配的 AI。 最后，AMD 硬件上适配的 ThunderKittens 也将很快推出。