機(jī)器心臟報(bào)告

機(jī)器之心編輯部

提高GPU利用率，就這么簡單。

人工智能的快速發(fā)展伴隨著大量的計(jì)算量。這自然引出了一個(gè)問題：如何降低AI的計(jì)算需求，提高現(xiàn)有AI的計(jì)算效率。

為了回答這個(gè)問題，斯坦福大學(xué)的研究人員在博客《GPUs Go Brrr》中給出了答案。

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

文章主要關(guān)注兩個(gè)問題：第一，真正需要的硬件是什么？二、如何滿足硬件要求？

文章討論了如何讓GPU 運(yùn)行得更快，并發(fā)布了一個(gè)庫ThunderKittens，它允許用戶在CUDA 上輕松編寫快速的深度學(xué)習(xí)內(nèi)核。它具有以下特點(diǎn)：

簡單，ThunderKittens 很容易寫?？蓴U(kuò)展性，如果用戶需要ThunderKittens無法提供的功能，可以擴(kuò)展功能。高速。 GitHub 鏈接：https://github.com/HazyResearch/ThunderKittens

ThunderKittens 使一些困難的事情變得非常簡單，允許在現(xiàn)代硬件上實(shí)現(xiàn)非常高的利用率。在項(xiàng)目中，作者使用ThunderKittens為RTX 4090編寫了一個(gè)簡單的FlashAttention-2內(nèi)核。代碼總共有58行代碼（不包括空格）。結(jié)果顯示，ThunderKittens 在RTX 4090 上實(shí)現(xiàn)了大約122 TFLOP（理論最大值的74%）。另外，當(dāng)內(nèi)核程序只有100行時(shí)，ThunderKittens性能比H100上的FlashAttention-2高出30%左右。

Nvidia H100 有一些怪癖

本研究重點(diǎn)關(guān)注NVIDIA H100，但所提供的信息也適用于其他GPU。

H100 SXM GPU 包括：

80 GB HBM3，帶寬3 TB/s（實(shí)際上更少）； 50 MB L2 緩存，帶寬12 TB/s，在GPU 上分為兩個(gè)25MB 部分，通過crossbar 連接； 132 流式多處理（SM，流式多處理器）。除了上述之外，H100 SXM GPU還有很多需要注意的地方，比如內(nèi)存控制器、指令緩存等。

研究人員表示，保持張量核心平穩(wěn)運(yùn)行并不容易。他們在AI 硬件中發(fā)現(xiàn)了一些怪癖，其中許多也適用于非H100 GPU，但H100 特別棘手。 （相比之下，RTX 4090 非常易于使用。）這些怪癖包括：

需要WGMMA指令，但用起來也很煩人；共享內(nèi)存實(shí)際上并沒有那么快，需要非常小心；地址生成成本高昂；占用仍然有幫助，寄存器通常是關(guān)鍵資源。本文進(jìn)一步描述了這些GPU 怪癖的具體情況。

WGMMA 指令很煩人

H100 有一組新的指令，稱為“扭曲組矩陣乘法累加，WGMMA”（PTX 中的wgmma.mma_async，或SASS 中的HGMMA/IGMMA/QGMMA/BGMMA）。以前的GPU 上可用的張量核心指令是wmma.mma.sync 和mma.sync 。通過這些指令，SM 單個(gè)象限上的32 個(gè)線程將同步將其數(shù)據(jù)塊饋送到張量核心并等待結(jié)果。

與wgmma.mma_async 指令不同，128 個(gè)連續(xù)線程（分布在SM 的所有象限上）協(xié)作直接從共享內(nèi)存（以及可選的寄存器）異步同步和啟動矩陣乘法。

在基準(zhǔn)測試中，研究團(tuán)隊(duì)發(fā)現(xiàn)這些指令對于提取H100 的完整計(jì)算是必要的。如果沒有它們，GPU 似乎只能達(dá)到峰值利用率的63% 左右。

共享內(nèi)存

單次訪問共享內(nèi)存的延遲約為30 個(gè)周期，聽起來可能不多，但在這段時(shí)間內(nèi)，SM 的張量核心可以完成幾乎兩個(gè)完整的32x32 矩陣乘法運(yùn)算。

處理共享內(nèi)存有點(diǎn)棘手，因?yàn)樗鎯υ?2 個(gè)獨(dú)立的內(nèi)存存儲中。如果你不小心，這可能會導(dǎo)致所謂的存儲體沖突，即要求同一個(gè)存儲體同時(shí)提供多個(gè)不同的內(nèi)存段，導(dǎo)致請求被序列化，這可能會不成比例地減慢內(nèi)核的速度- 而wgmma 和mma 指令所需的寄存器布局受這些存儲體沖突的影響。解決方案是使用各種交錯(cuò)模式重新排列共享內(nèi)存，以避免這些沖突。

地址生成

H100 的特點(diǎn)之一是張量核心和內(nèi)存都足夠快，僅生成用于獲取數(shù)據(jù)的內(nèi)存地址就占用了芯片的很大一部分資源。

NVIDIA 似乎已經(jīng)意識到了這一點(diǎn)，因?yàn)樗麄優(yōu)镚PU 提供了張量內(nèi)存加速器（或他們所說的TMA）。 TMA允許用戶在全局和共享內(nèi)存中指定多維張量布局，這節(jié)省了所有地址生成成本，也使構(gòu)建管道變得更容易。

研究團(tuán)隊(duì)還發(fā)現(xiàn)，TMA 與wgmma.mma_async 一樣，對于發(fā)揮H100 的全部潛力是完全不可或缺的。

占據(jù)

在某些方面，H100 比前幾代硬件對占用率的依賴程度更低。 NVIDIA 在設(shè)計(jì)GPU 時(shí)確實(shí)考慮到了占用情況。而對于H100來說，占用率只能用有用來形容，但并沒有多大作用。研究人員發(fā)現(xiàn)，它在A100 和RTX 4090 上變得越來越重要。

雷霆小貓

那么，如何才能更輕松地編寫內(nèi)核，同時(shí)仍然擁有硬件的全部功能呢？

研究團(tuán)隊(duì)設(shè)計(jì)了一個(gè)嵌入CUDA 的DSL，命名為ThunderKittens。

ThunderKittens 的目標(biāo)是盡可能簡單，包括四種模板類型：

在寄存器文件中注冊tile—— 2D 張量。寄存器文件中的寄存器向量—— 1D 張量。共享內(nèi)存中的共享tile—— 2D張量。 SharedVector —— 共享內(nèi)存中的一維張量。平鋪塊通過高度、寬度和布局進(jìn)行參數(shù)化，寄存器向量通過長度和布局進(jìn)行參數(shù)化，共享向量僅通過長度進(jìn)行參數(shù)化。這樣您通常就不會遭受銀行沖突的困擾。

研究團(tuán)隊(duì)還提供了一些必要的操作：

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

一元運(yùn)算，例如exp 二元運(yùn)算，例如mul 行/列運(yùn)算，例如row_sum 這項(xiàng)研究給出了一個(gè)用ThunderKittens 編寫的用于RTX 4090 的簡單前向閃存注意內(nèi)核：

#define NUM_WORKERS 16 //該內(nèi)核每個(gè)塊并行使用16 個(gè)工作線程，以幫助更快地發(fā)出指令。

使用命名空間小貓； //為了簡單起見，該內(nèi)核僅處理headdim=64。另外，這里n 應(yīng)該是256 的倍數(shù)。

__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 共享內(nèi)存

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

//K 和V 位于共享內(nèi)存中——這就是所有適合的內(nèi)容。

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++) {

//每個(gè)扭曲加載自己的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)); //溫度調(diào)節(jié)

//零閃存注意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++) {

//每個(gè)warp 將自己的k、v 塊加載到共享內(nèi)存中

加載(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(); //我們需要確保在開始計(jì)算階段之前加載所有內(nèi)存

//現(xiàn)在每個(gè)warp 都會遍歷所有子圖塊，加載它們，然后執(zhí)行flash 注意內(nèi)部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); //[email protected]

復(fù)制（norm_vec_last，norm_vec）；

復(fù)制（max_vec_last，max_vec）；

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

sub_row(att_block, att_block, max_vec); //從注意力中減去最大值——現(xiàn)在全部=0

exp(att_block, att_block); //就地對塊求冪。

sub(max_vec_last, max_vec_last, max_vec); //從舊的最大值中減去新的最大值以找到新的標(biāo)準(zhǔn)化。

exp(max_vec_last, max_vec_last); //對該向量求冪——這就是我們需要標(biāo)準(zhǔn)化的。

mul(norm_vec,norm_vec,max_vec_last); //范數(shù)vec 現(xiàn)在已標(biāo)準(zhǔn)化。

row_sum(norm_vec, att_block,norm_vec); //將新的注意力塊累積到現(xiàn)在重新縮放的norm_vec上

p_row（att_block，att_block，norm_vec）； //現(xiàn)在注意力塊已正確標(biāo)準(zhǔn)化

mul(norm_vec_last,norm_vec_last,max_vec_last); //根據(jù)新的最大值標(biāo)準(zhǔn)化先前的范數(shù)vec

p（norm_vec_last，norm_vec_last，norm_vec）； //根據(jù)新范數(shù)對先前范數(shù)vec 進(jìn)行歸一化

復(fù)制（att_block_mma，att_block）； //mma_AB 轉(zhuǎn)換為bf16

加載（v_reg，v_smem [subtile]）； //將v 從共享加載到寄存器中。

rt_bf_1x4 v_reg_col=swap_layout_inplace(v_reg); //這是一個(gè)引用，調(diào)用使v_reg 無效

mul_row（o_reg，o_reg，norm_vec_last）； //在mma_AB 之前標(biāo)準(zhǔn)化o_reg

mma_AB(o_reg, att_block_mma, v_reg_col, o_reg); //使用局部注意力@V matmul 將mfma 轉(zhuǎn)移到o_reg 上。

}

__syncthreads(); //我們需要確保所有扭曲都完成，然后才能開始加載下一個(gè)kv 塊

}

store(_o + (q_blk*NUM_WORKERS + warpid)*q_reg.num_elements, o_reg, q_reg.cols); //寫出o.如果將d 設(shè)置為constexpr q_reg.rows :/，編譯器會出現(xiàn)寄存器使用問題

}

}

CUDA代碼總共約60行，硬件利用率為75%。雖然非常密集，但大部分復(fù)雜性在于算法，而不是混合模式或寄存器布局。

TMA、WGMMA、混合模式和描述符的復(fù)雜性如何？以下是使用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 共享內(nèi)存

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(); } 這個(gè)內(nèi)核只有 100 行代碼，它在 H100 上的性能比 FlashAttention-2 高出約 30%。ThunderKittens 負(fù)責(zé) wrap up 布局和指令，并提供一個(gè)可以在 GPU 上使用的 mini-pytorch。H100 SXM 上各種配置的 FlashAttention-2（Pytorch）與 ThunderKittens 的比較。 此外，研究團(tuán)隊(duì)還發(fā)布了基于線性注意力的內(nèi)核和其他架構(gòu)。基于線性注意力內(nèi)核的運(yùn)行速度為 215 TFLOP（如果考慮算法中固有的重計(jì)算，則運(yùn)行速度超過 300 TFLOP）。 雖然理論上線性注意力更高效，但從實(shí)踐經(jīng)驗(yàn)來看，線性注意力在硬件上的效率大大降低。因此，ThunderKittens 有望開辟廣泛的高吞吐量應(yīng)用。使用 ThunderKittens 可以非?？斓貙?shí)現(xiàn)線性注意力。 tile 看起來是個(gè)好點(diǎn)子 在研究團(tuán)隊(duì)看來，ThunderKittens 之所以運(yùn)行良好，是因?yàn)樗粫噲D做所有事情。CUDA 確實(shí)比 ThunderKittens 更有表現(xiàn)力，而 ThunderKittens 又小又簡單。 不過，ThunderKittens 具有很好的抽象能力，它具有小的 tile，這與 AI 和硬件的發(fā)展相匹配。ThunderKittens 不支持任何少于 16 的維數(shù)。但在研究團(tuán)隊(duì)看來，這一點(diǎn)并不重要，尤其對于硬件而言。如果你的矩陣乘法小于 16x16，你確定自己做的還是 AI 嗎？ 從哲學(xué)的視角來看，研究團(tuán)隊(duì)認(rèn)為框架遷移是合理的。「寄存器」當(dāng)然不應(yīng)該像舊 CPU 那樣的 32 位。CUDA 使用的 1024 位寬向量寄存器無疑朝著正確方向邁出了一步。但對研究團(tuán)隊(duì)而言，「寄存器」是 16x16 的數(shù)據(jù) tile。他們認(rèn)為 AI 想要這樣，它仍然只是矩陣乘法、規(guī)約和重塑。當(dāng)然硬件也想要這樣，小的矩陣乘法尋求硬件支持，而不僅僅是 systolic mma。 實(shí)際上，從更廣泛的視角來看，研究團(tuán)隊(duì)認(rèn)為應(yīng)該圍繞硬件的良好映射來重新調(diào)整 AI 思路。比如，循環(huán)狀態(tài)應(yīng)該有多大？SM 能夠容納多大尺寸？計(jì)算密度是多少？這些都不亞于硬件的要求。 研究團(tuán)隊(duì)表示，這項(xiàng)工作未來的一個(gè)重要方向是利用他們對硬件的了解來幫助設(shè)計(jì)與硬件相匹配的 AI。 最后，AMD 硬件上適配的 ThunderKittens 也將很快推出。