Bendi新闻
>
只需百行代码,让H100提速30%,斯坦福开源全新AI加速框架

只需百行代码,让H100提速30%,斯坦福开源全新AI加速框架

5月前
机器之心报道
机器之心编辑部
提高 GPU 利用率,就是这么简单。
AI 的快速发展,伴随而来的是大计算量。这就自然而然的引出了一个问题:如何减少 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%。

英伟达 H100 有些小怪癖

该研究重点关注 NVIDIA H100,不过所介绍的内容也适用于其他 GPU。

H100 SXM GPU 包含:

  • 80 GB HBM3,带宽为 3 TB/s(实际上带宽会少一些);
  • 50 MB 二级缓存,带宽 12 TB/s,在 GPU 上分成两个 25MB 的部分,通过 crossbar 连接;
  • 132 个流多处理器 (SM,streaming multiprocessors)。

除了上述这些,H100 SXM GPU 还有很多可关注的东西,例如内存控制器、指令缓存等。

研究者表示保持张量核心的运行流畅并不容易。他们发现了一些 AI 硬件上的怪癖,这些怪癖中的很多内容也适用于非 H100 GPU,但 H100 尤其棘手。(相比之下,RTX 4090 则非常容易使用),这些怪癖包括:

  • WGMMA 指令是必需的,但使用起来也非常令人恼火;
  • 共享内存实际上并没有那么快,并且需要非常小心;
  • 地址生成成本很高;
  • 占用率仍然有帮助,寄存器通常是关键资源。

文章进一步描述了 GPU 这些怪癖的具体内容。

WGMMA 指令令人恼火

H100 有一组新指令,称为「warp group matrix multiply accumulate,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 矩阵乘法运算。

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

地址生成

H100 其中一个特点是张量核心和内存都足够快,以至于仅仅生成用于获取数据的内存地址就占据了芯片资源的相当一部分。

NVIDIA 似乎已经意识到了这一点,因为他们赋予了 GPU 张量内存加速器(或称之为 TMA)。TMA 允许用户在全局和共享内存中指定多维张量布局,这节省了所有的地址生成成本,并且还使得构建 pipeline 更加容易。

研究团队还发现 TMA 和 wgmma.mma_async 一样,在实现 H100 的全部潜力方面是完全不可或缺的。

占用

在某些方面,与前几代硬件相比,H100 对占用率的依赖程度较低。NVIDIA 确实在设计 GPU 时考虑了占用率。虽然对于 H100 来说,占用率只能说有用,但作用不大。研究者发现在 A100 和 RTX 4090 上它变得越来越重要。

ThunderKittens

那么,如何才能更轻松地编写内核,同时仍兼具硬件的全部功能?

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

ThunderKittens 旨在尽可能简单,并包含四种模板类型:

  • 寄存器 tile—— 寄存器文件中的 2D 张量。
  • 寄存器向量 —— 寄存器文件中的 1D 张量。
  • 共享 tile—— 共享内存中的 2D 张量。
  • 共享向量 —— 共享内存中的 1D 张量。

tile 通过高度、宽度和布局进行参数化,寄存器向量由长度和布局参数化,共享向量仅由长度参数化。这样通常不会遭受 bank 冲突的困扰。

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

初始化,如将共享向量清零

  • 一元运算,如 exp
  • 二元运算,如 mul
  • 行 / 列操作,如 row_sum

该研究给出了一个用 ThunderKittens 编写的,用于 RTX 4090 的简单前向 flash attention 内核:

#define NUM_WORKERS 16 // This kernel uses 16 workers in parallel per block, to help issue instructions more quickly.using namespace kittens; // this kernel only handles headdim=64 for simplicity. Also n should be a multiple of 256 here.__global__ void attend_ker64(int n, const bf16* __restrict__ __q__, const bf16* __restrict__ __k__, const bf16* __restrict__ __v__, bf16* __o__) {
auto warpid = kittens::warpid(); auto 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;
extern __shared__ alignment_dummy __shm[]; // this is the CUDA shared memory shared_allocator al((int*)&__shm[0]); // K and V live in shared memory -- this is about all that will fit. st_bf_1x4<ducks::st_layout::swizzle> (&k_smem)[NUM_WORKERS] = al.allocate<st_bf_1x4<ducks::st_layout::swizzle>, NUM_WORKERS>(); st_bf_1x4<ducks::st_layout::swizzle> (&v_smem)[NUM_WORKERS] = al.allocate<st_bf_1x4<ducks::st_layout::swizzle>, NUM_WORKERS>();
// Initialize all of the register tiles. rt_bf_1x4<> q_reg, k_reg, v_reg; // v_reg need to be swapped into 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; // these are column vectors for the attention block rt_fl_1x1<>::col_vec norm_vec_last, norm_vec; // these are column vectors for the attention block int qo_blocks = n / (q_reg.rows*NUM_WORKERS), kv_blocks = n / (q_reg.rows*NUM_WORKERS);
for(auto q_blk = 0; q_blk < qo_blocks; q_blk++) {
// each warp loads its own Q tile of 16x64, and then multiplies by 1/sqrt(d) load(q_reg, _q + (q_blk*NUM_WORKERS + warpid)*q_reg.num_elements, q_reg.cols); mul(q_reg, q_reg, __float2bfloat16(0.125f)); // temperature adjustment
// zero flash attention L, M, and O registers. neg_infty(max_vec); // zero registers for the Q chunk zero(norm_vec); zero(o_reg);
// iterate over k, v for these q's that have been loaded for(auto kv_idx = 0; kv_idx < kv_blocks; kv_idx++) {
// each warp loads its own chunk of k, v into shared memory load(v_smem[warpid], _v + (kv_idx*NUM_WORKERS + warpid)*q_reg.num_elements, q_reg.cols); load(k_smem[warpid], _k + (kv_idx*NUM_WORKERS + warpid)*q_reg.num_elements, q_reg.cols); __syncthreads(); // we need to make sure all memory is loaded before we can begin the compute phase
// now each warp goes through all of the subtiles, loads them, and then does the flash attention internal alg. for(int subtile = 0; subtile < NUM_WORKERS; subtile++) {
load(k_reg, k_smem[subtile]); // load k from shared into registers
zero(att_block); // zero 16x16 attention tile mma_ABt(att_block, q_reg, k_reg, att_block); // [email protected]
copy(norm_vec_last, norm_vec); copy(max_vec_last, max_vec);
row_max(max_vec, att_block, max_vec); // accumulate onto the max_vec sub_row(att_block, att_block, max_vec); // subtract max from attention -- now all <=0 exp(att_block, att_block); // exponentiate the block in-place.
sub(max_vec_last, max_vec_last, max_vec); // subtract new max from old max to find the new normalization. exp(max_vec_last, max_vec_last); // exponentiate this vector -- this is what we need to normalize by. mul(norm_vec, norm_vec, max_vec_last); // and the norm vec is now normalized.
row_sum(norm_vec, att_block, norm_vec); // accumulate the new attention block onto the now-rescaled norm_vec div_row(att_block, att_block, norm_vec); // now the attention block is correctly normalized
mul(norm_vec_last, norm_vec_last, max_vec_last); // normalize the previous norm vec according to the new max div(norm_vec_last, norm_vec_last, norm_vec); // normalize the previous norm vec according to the new norm
copy(att_block_mma, att_block); // convert to bf16 for mma_AB
load(v_reg, v_smem[subtile]); // load v from shared into registers. rt_bf_1x4<ducks::rt_layout::col> &v_reg_col = swap_layout_inplace(v_reg); // this is a reference and the call has invalidated v_reg
mul_row(o_reg, o_reg, norm_vec_last); // normalize o_reg in advance of mma_AB'ing onto it mma_AB(o_reg, att_block_mma, v_reg_col, o_reg); // mfma onto o_reg with the local attention@V matmul. } __syncthreads(); // we need to make sure all warps are done before we can start loading the next kv chunk }
store(_o + (q_blk*NUM_WORKERS + warpid)*q_reg.num_elements, o_reg, q_reg.cols); // write out o. compiler has an issue with register usage if d is made constexpr q_reg.rows :/ }}

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

TMA、WGMMA、swizzling 模式和描述符的复杂度又如何呢?如下是用 ThunderKittens 编写的, H100 的 FlashAttention-2 前向传递:

template<int D>__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) {    extern __shared__ int __shm[]; // this is the CUDA shared memory    tma_swizzle_allocator al((int*)&__shm[0]);
constexpr int tile_width = fwd_attend_ker_tile_dims<D>::tile_width; // constants constexpr int qo_height = fwd_attend_ker_tile_dims<D>::qo_height; constexpr int kv_height = fwd_attend_ker_tile_dims<D>::kv_height;
st_bf<qo_height, tile_width, layout_q> (&q_smem) [NUM_WARPGROUPS] = al.allocate<st_bf<qo_height, tile_width, layout_q>, NUM_WARPGROUPS>(); st_bf<kv_height, tile_width, layout_k> (&k_smem)[2][NUM_WORKERS_KV] = al.allocate<st_bf<kv_height, tile_width, layout_k>, 2, NUM_WORKERS_KV>(); st_bf<kv_height, tile_width, layout_v> (&v_smem)[2][NUM_WORKERS_KV] = al.allocate<st_bf<kv_height, tile_width, layout_v>, 2, NUM_WORKERS_KV>();
int tic = 0, toc = 1; rt_fl<1, kv_height> att_block; rt_bf<1, kv_height> att_block_mma; rt_fl<1, qo_height> o_prev; col_vec<rt_fl<1, kv_height>> max_vec_last, max_vec; col_vec<rt_fl<1, kv_height>> norm_vec_last, norm_vec;
int warpid = kittens::warpid(); int warpgroupid = warpid/kittens::WARPGROUP_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<st_bf<qo_height, tile_width, layout_q>, NUM_WARPGROUPS>(qsmem_barrier, 1); tma::init_barrier<st_bf<kv_height, tile_width, layout_k>, NUM_WORKERS_KV*2>(kvsmem_barrier, 1); }
if (warpid == 0) { for (int wg = 0; wg < NUM_WORKERS/kittens::WARPGROUP_WARPS; wg++) { // load q int tile_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++) { // load k, v int tile_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); // zero registers for the Q chunk zero(norm_vec); zero(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)); } else { warpgroup::mul(q_smem[warpgroupid], q_smem[warpgroupid], __float2bfloat16(0.08838834764f)); }
for (auto 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(); if (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++) { int tile_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); warpgroup::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 div_row(att_block, att_block, norm_vec);
mul(norm_vec_last, norm_vec_last, max_vec_last); div(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<st_bf<qo_height, tile_width, layout_o>(*)>(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 也将很快推出。




© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

微信扫码关注该文公众号作者

来源:机器之心

相关新闻

传微软组建新团队开发更小、更便宜AI模型;Pika联手北大斯坦福开源文生图框架;传和硕独家拿下Ai Pin代工订单丨AIGC日报扩散模型更懂复杂提示词!Pika北大斯坦福开源新框架,利用LLM提升理解力北大快手攻克复杂视频生成难题!新框架轻松组合各种细节,代码将开源小洞不补,大洞吃苦:西交、麦马开源全新「拖动式编辑」框架&数据集红杉资本入局,马斯克的AI公司接近达成60亿美元融资;苹果发布基于开源训练和推理框架的语言模型OpenELM丨AIGC日报字节豆包全新图像Tokenizer:生成图像最低只需32个token,最高提速410倍可「自主进化」的Agent?首个端到端智能体符号化训练框架开源了开源框架中的责任链模式实践腾讯:终于补齐了Muse系列数字人开源框架,感谢阿里!阿里通义实验室薄列峰:从兵马俑跳“科目三”到照片唱歌,四大框架让AI生成的人物活起来丨GenAICon 2024微软TaskWeaver开源框架:携手数据分析与行业定制,打造顶级Agent解决方案中科大/华为诺亚出手!芯片性能≠布局评分,EDA物理设计框架全面开源现场Live震撼!OmAgent框架强势开源!行业应用已全面开花苹果、AMD和高通GPU被爆存在漏洞!只需十行代码即可窃取数据,数百万台苹果设备或将受到影响ICLR 2024 | 机器人领域首个开源视觉-语言操作大模型!RoboFlamingo框架激发开源VLMs更大潜能机器人领域首个开源视觉-语言操作大模型,RoboFlamingo框架激发开源VLMs更大潜能ICLR 2024 | Adobe提出DMV3D:3D生成只需30秒!让文本、图像都动起来的新方法!滴滴开源Flutter混合开发框架Unify小米大模型提效新框架:训练最高提速34%,推理最高提速52%!Kaldi之父合作出品复刻Sora的通用视频生成能力,开源多智能体框架Mora来了5秒完成3D生成,真香合成数据集已开源,上交港中文新框架超越Instant3D开放开源!蚂蚁集团浙江大学联合发布开源大模型知识抽取框架OneKE寡姐带货国风Polo衫,马斯克穿牛仔走红毯!虚拟试衣新框架火了,只需两张图30秒即生成模块化重构LLaVA,替换组件只需添加1-2个文件,开源TinyLLaVA Factory来了
logo
联系我们隐私协议©2024 bendi.news
Bendi新闻
Bendi.news刊载任何文章,不代表同意其说法或描述,仅为提供更多信息,也不构成任何建议。文章信息的合法性及真实性由其作者负责,与Bendi.news及其运营公司无关。欢迎投稿,如发现稿件侵权,或作者不愿在本网发表文章,请版权拥有者通知本网处理。