1 代碼結(jié)構(gòu)&調(diào)用流程

1.1 代碼結(jié)構(gòu)

llama.cpp 的代碼結(jié)構(gòu)比較直觀，如下所示，為整體代碼結(jié)構(gòu)中的比較核心的部分的代碼結(jié)構(gòu)

|--example
||--main
||--main.cpp#推理llama2的主函數(shù)
|--ggml.c#ggml.c和.h文件定義一些框架的基礎(chǔ)數(shù)據(jù)結(jié)構(gòu)和函數(shù)等
|--ggml.h
|--ggml-alloc.c#內(nèi)存分配管理
|--ggml-alloc.h
|--llama.cpp#整個(gè)llama2的計(jì)算圖構(gòu)建和weight加載等
|--llama.h
|--ggml-cuda.cu#cuda版本的llama2中的kernel實(shí)現(xiàn)與調(diào)用
|--ggml-cuda.h
|--ggml-opencl.cpp#opencl版本的llama2中的kernel實(shí)現(xiàn)與調(diào)用
|--ggml-opencl.h
|--...#其他

1.2 調(diào)用流程

當(dāng)我們搭建完成環(huán)境，并對(duì)llama.cpp 進(jìn)行編譯后 在llama.cpp/build/bin/會(huì)生成一個(gè)main的可執(zhí)行文件，根據(jù)README.md給的相關(guān)命令即可進(jìn)行l(wèi)lama2的推理。大致梳理一下llama.cpp的調(diào)用執(zhí)行流程:

首先，main這個(gè)可執(zhí)行文件的源碼對(duì)應(yīng)于llama.cpp/examples/main/main.cpp，在main.cpp中會(huì)解析命令行的參數(shù)，如所用的模型文件，prompt信息等，之后進(jìn)行一系列操作后，用一個(gè)llama_token_bos()token并調(diào)用了一次llama_eval()函數(shù)來(lái)對(duì)模型進(jìn)行了一次warm up, 之后進(jìn)入一個(gè)while循環(huán)進(jìn)行模型的推理，期間會(huì)多次調(diào)用llama_eval()函數(shù)進(jìn)行推理，直到不滿足while條件。

llama_eval()函數(shù)的定義在llama.cpp/llama.cpp文件中,llama_eval()函數(shù)進(jìn)一步會(huì)去調(diào)用llama_eval_internal()函數(shù),llama_eval_internal()函數(shù)內(nèi)部會(huì)根據(jù)預(yù)先的宏定義走不同的推理模式，比如GGML_USE_MPI、GGML_USE_MPI和其他模式，因?yàn)楸疚氖且訡UDA推理模式進(jìn)行說(shuō)明的，所以我們主要看該模式下的函數(shù)調(diào)用：主要有兩個(gè) llama_build_graph()和ggml_graph_compute_helper() 。這兩個(gè)函數(shù)的功能分別是前者用于構(gòu)建推理計(jì)算圖 ，而后者則是根據(jù)計(jì)算圖調(diào)用對(duì)應(yīng)算子

llama_build_graph

llama_build_graph()函數(shù)接口，如下所示

staticstructggml_cgraph*llama_build_graph(
llama_context&lctx,//llamacontext存放著一些模型信息，包括模型文件、超參數(shù)等
constllama_token*tokens,//需要去處理的tokens
constfloat*embd,//embeddingsinput
intn_tokens,//numberoftokens
intn_past//已經(jīng)處理的tokens數(shù)量
);

這其中整個(gè)llama 2的模型結(jié)構(gòu)的推理計(jì)算圖全在該函數(shù)內(nèi)實(shí)現(xiàn)，代碼太長(zhǎng)了為了節(jié)省篇幅就不截取了，大家可以根據(jù)函數(shù)名找到對(duì)應(yīng)的函數(shù)實(shí)現(xiàn)

ggml_graph_compute_helper

ggml_graph_compute_helper()函數(shù)內(nèi)部主要會(huì)調(diào)用兩個(gè)函數(shù): ggml_graph_plan() 和 ggml_graph_compute() 。前者用于創(chuàng)建一個(gè)ggml_cplan結(jié)構(gòu)體cplan，同時(shí)根據(jù)之前l(fā)lama_build_graph()創(chuàng)建的計(jì)算圖，對(duì)圖中每個(gè)節(jié)點(diǎn)所對(duì)應(yīng)的算子OP確定cplan中的成員值，之后返回cplan。后者根據(jù)llama_build_graph()創(chuàng)建的計(jì)算圖和ggml_graph_plan()創(chuàng)建的cplan 進(jìn)一步調(diào)用ggml_graph_compute_thread,這個(gè)函數(shù)再根據(jù)當(dāng)前計(jì)算圖節(jié)點(diǎn)信息進(jìn)一步調(diào)用ggml_compute_forward

ggml_compute_forward

此函數(shù)會(huì)根據(jù)當(dāng)前節(jié)點(diǎn)的信息調(diào)用具體的算子。當(dāng)然根據(jù)不同的編譯選項(xiàng)會(huì)使得算子有不同的調(diào)用：當(dāng)定義了GGML_USE_CUBLAS，如果當(dāng)前節(jié)點(diǎn)所對(duì)應(yīng)的算子在CUDA平臺(tái)有具體的實(shí)現(xiàn)就會(huì)調(diào)用它，否則就會(huì)調(diào)用CPU端的實(shí)現(xiàn)

ggml_compute_forward

ggml_cuda_compute_forward 就會(huì)調(diào)用具體的CUDA節(jié)點(diǎn)

ggml_cuda_compute_forward

另外，整個(gè)llama.cpp中一個(gè)很重要的結(jié)構(gòu)體ggml_tensor，其定義如下

structggml_tensor{
enumggml_typetype;//數(shù)據(jù)類型FP32INT8等
enumggml_backendbackend;//后端CPU/GPU
intn_dims;//幾維度
int64_tne[GGML_MAX_DIMS];//numberofelements，這就是tensor的shape,不過(guò)它的存放方式是倒著的
//比如[batch_size,multi-head,seq_len,head-dim]
//他的存儲(chǔ)方式是ne[0]=head-dim,ne[1]=seq_len,ne[2]=multi-head,ne[3]=batch_size
size_tnb[GGML_MAX_DIMS];//strideinbytes:
//nb[0]=sizeof(type)
//nb[1]=nb[0]*ne[0]+padding
//nb[i]=nb[i-1]*ne[i-1]
//computedata
enumggml_opop;
//opparams-allocatedasint32_tforalignment
int32_top_params[GGML_MAX_OP_PARAMS/sizeof(int32_t)];
boolis_param;
structggml_tensor*grad;
structggml_tensor*src[GGML_MAX_SRC];
//performance
intperf_runs;
int64_tperf_cycles;
int64_tperf_time_us;
void*data;
charname[GGML_MAX_NAME];
void*extra;//extrathingse.g.forggml-cuda.cu
charpadding[4];
};

至此，llama.cpp 在推理llama 2時(shí)的一個(gè)主要調(diào)用邏輯，就算說(shuō)完了，接下來(lái)，我們來(lái)看看本文的重點(diǎn)部分：llama 2 中每個(gè)Transformer Block的CUDA版本的算子調(diào)用及解析

2 逐算子解析

在之前Llama2 詳解中我們說(shuō)過(guò)，大模型的推理可以分為prompt和generation兩個(gè)階段，兩個(gè)階段在處理時(shí)的差異在于數(shù)據(jù)維度的差異，即prompt階段是多token輸入 input_tensor: [batch_size, seq_len, hidden_dim] ; 而generation階段的輸入則是 input_tensor: [batch_size, 1, hidden_dim] ，所以前者更多的計(jì)算算子是gemm，而后者更多的計(jì)算算子則是gemv 。

后文中為了方便說(shuō)明具體的參數(shù)信息，本文以Llama-2 7B模型 batch_size =1 為例來(lái)說(shuō)明llama.cpp 在推理時(shí)的tensor shape和其他參數(shù)信息

回顧一下，上圖為L(zhǎng)lama2 詳解中我畫的llama 2的模型結(jié)構(gòu)圖。根據(jù)模型結(jié)構(gòu)，我們來(lái)看看llama.cpp的推理流程。如下圖所示，為我通過(guò)Nsight Systems工具抓取的 在llama.cpp 選用CUDA作為推理后端時(shí)的算子調(diào)用和執(zhí)行情況，其中黃色框?yàn)橐淮蝫armup，綠色框就是prompting階段，紅色框的多個(gè)塊就是一次次的generation階段。

相信大家也不難發(fā)現(xiàn)，通過(guò)Nsight Systems所統(tǒng)計(jì)的執(zhí)行時(shí)間占比最大的kernel是dequantize_mul_mat_vec——簡(jiǎn)單解釋一下：這是一個(gè)反量化矩陣向量乘法，只會(huì)在generation階段調(diào)用。所以說(shuō)大模型的推理，generation階段占比更重，紅色框中的一個(gè)小塊即為生成一個(gè)token所調(diào)用kernel的時(shí)間，而隨著你需要生成的token的數(shù)量增多，紅色框的占比會(huì)越來(lái)越大。

那么接下來(lái)，我們就結(jié)合上述模型結(jié)構(gòu)圖和Nsight Systems 截圖來(lái)一起看看llama 2推理時(shí)都會(huì)調(diào)用哪些CUDA算子，以及l(fā)lama.cpp 對(duì)這些算子是如何實(shí)現(xiàn)的～

因?yàn)?strong>prompting和generation只是在tensor shape不一樣，而算子實(shí)現(xiàn)的算法功能都是一致的。那么我們可以根據(jù)llama 2的模型結(jié)構(gòu)圖中將一個(gè)Transformer Block拆分為的兩個(gè)塊（Attention Block和FeedForward Block），然后逐一比較這兩個(gè)塊在提示（prompting）和生成（generation）階段所調(diào)用的算子以及它們的實(shí)現(xiàn)。

2.1 Attention Block

通過(guò)對(duì)Nsight Systems Profile report文件放大后分析，可以得到Attention Block的上層算法流程以及其在prompting階段和generation階段所調(diào)用的CUDA算子，如下圖所示。根據(jù)這一對(duì)比示意圖，我們就來(lái)細(xì)看一下每個(gè)算子的功能以及具體的實(shí)現(xiàn)

2.1.1 rms_norm_f32

我們回憶一下RMS Norm的公式：

其中 為可學(xué)習(xí)的參數(shù)，推理時(shí)固定

#defineWARP_SIZE32
//callkernel
staticvoidrms_norm_f32_cuda(constfloat*x,float*dst,constintncols,constintnrows,constfloateps,cudaStream_tstream){
GGML_ASSERT(ncols%WARP_SIZE==0);
constdim3block_dims(WARP_SIZE,1,1);//(32,1,1)
//所以調(diào)用的cuda的gridDim=(nrows,1,1),blockDim=(32,1,1)
//也就是說(shuō)一個(gè)block處理一個(gè)row的數(shù)據(jù)，即每32個(gè)線程處理一行數(shù)據(jù)，共計(jì)nrows行
rms_norm_f32<<>>(x,dst,ncols,eps);
}

//kernelcode
static__global__voidrms_norm_f32(constfloat*x,float*dst,constintncols,constfloateps){
constintrow=blockIdx.x*blockDim.y+threadIdx.y;
constinttid=threadIdx.x;

floattmp=0.0f;//partialsumforthreadinwarp
//一個(gè)線程求和(ncols/WARP_SIZE)個(gè)數(shù)據(jù)的x^2
for(intcol=tid;col0;mask>>=1){
tmp+=__shfl_xor_sync(0xffffffff,tmp,mask,32);
}

constfloatmean=tmp/ncols;//mean(x^2)
constfloatscale=rsqrtf(mean+eps);//1/根號(hào)mean
//算完之后寫回原數(shù)組
for(intcol=tid;col

	

	所以，rms_norm_f32這個(gè)kernel就是在計(jì)算RMS Norm的前一部分，之后再通過(guò)如下kernel mul_f32乘上，就得到了完整的RMS Norm。rms_norm_f32這個(gè)kernel 在prompting階段處理的tensor shape 是[1, seq_len , 4096] ，nrows 等于seq_len，在generation階段處理的tensor shape 則是[1, 1 , 4096] 。

	batch_size =1 ,7B模型的 hidden_dim = 4096

	
#defineCUDA_MUL_BLOCK_SIZE256
//callkernel
staticvoidmul_f32_cuda(constfloat*x,constfloat*y,float*dst,constintkx,constintky,cudaStream_tstream){
constintnum_blocks=(kx+CUDA_MUL_BLOCK_SIZE-1)/CUDA_MUL_BLOCK_SIZE;
mul_f32<<>>(x,y,dst,kx,ky);
}

//kernel
static__global__voidmul_f32(constfloat*x,constfloat*y,float*dst,constintkx,constintky){
constinti=blockDim.x*blockIdx.x+threadIdx.x;

if(i>=kx){
return;
}
dst[i]=x[i]*y[i%ky];
}


	

	mul_f32算子比較簡(jiǎn)單就是挨個(gè)元素乘上對(duì)應(yīng)的，其實(shí)rms_norm_f32和mul_f32 可以合并成一個(gè)kernel，后續(xù)文章會(huì)推出一些kernel優(yōu)化的解析會(huì)講到。

	2.1.2 Linear Layer

	對(duì)比Llama 2模型結(jié)構(gòu)可以發(fā)現(xiàn)，一個(gè)Transformer Block中總共有7個(gè)Linear層：生成Q、K、V的三個(gè)Linear、Attention Block中最后一個(gè)Linear和FeedForward Block中的3個(gè)Linear。雖然這些Linear層處理的tensor 的shape是不同的，但是在相同的階段調(diào)用的算子都是同一個(gè)，所以可以舉一反三。

	此外，llama.cpp中涉及到量化推理的主要就是Linear層，前文提過(guò)本文先導(dǎo)知識(shí)之一就是模型量化，所謂模型量化就是將模型中的weight數(shù)據(jù)和input-tensor數(shù)據(jù)，通過(guò)量化算法將原始FP32類型邏輯等價(jià)地轉(zhuǎn)換為int8以及更低bit數(shù)據(jù)，這樣做的好處就是在對(duì)模型進(jìn)行推理時(shí)能節(jié)省內(nèi)存和計(jì)算加速的好處。模型量化算法有很多種，以常見(jiàn)的對(duì)稱均勻量化為例，模型量化時(shí)都會(huì)對(duì)原始FP32數(shù)據(jù)在pre-tensor/pre-channel域計(jì)算得到一個(gè)scale，然后通過(guò)量化公式：將數(shù)據(jù)由FP32量化為INT-8(或更低bit)數(shù)據(jù) 。

	這里解釋一下：模型量化后計(jì)算速度的加快的主要原因在于：在同等帶寬的情況下能一次向量化的load更多數(shù)據(jù)(比如原始load 1個(gè)FP32的時(shí)間 現(xiàn)在能load 4個(gè)int8的數(shù)據(jù))

	以llama.cpp 提供的LLaMA 2 7B chat 8bit模型為例，Llama 2中Linear層的weight數(shù)據(jù)就是int-8類型，更具體的說(shuō)，Linear層中的weight數(shù)據(jù)是以如下結(jié)構(gòu)體的形式保存的，其中d為前文中提到的量化算法中的scale，int8_t qs[QK8_0] 即為量化后的INT-8數(shù)據(jù)

	
#defineQK8_032
typedefstruct{
halfd;//delta量化的scale
int8_tqs[QK8_0];//quants量化的weight數(shù)據(jù)
}block_q8_0;


	

	這里的量化scale既不是以pre-tensor為單位也不是以pre-channel為單位，而是以32為單位，主要原因是因?yàn)镃UDA 一個(gè)warp就是32個(gè)線程

	Linear層在進(jìn)行量化推理時(shí)可以選用兩種方式：

	反量化int-8的weight，之后將fp32的input-tensor與fp32的weight進(jìn)行Linear層的運(yùn)算

	量化input-tensor，之后將int-8的input-tensor與int-8的weight進(jìn)行Linear層的運(yùn)算

	接下來(lái)，我們就來(lái)看看llama.cpp對(duì)于這兩種方式是如何實(shí)現(xiàn)的

	dequantize Linear

	首先，對(duì)于Linear層prompting階段和generation階段都會(huì)被調(diào)用，但是因?yàn)樘幚淼膖ensor的shape不一樣，所以在不同階段執(zhí)行時(shí)調(diào)用的kernel不一樣，而在同一個(gè)階段調(diào)用的kernel又是一樣的，所以為了節(jié)省篇幅，后文在講解時(shí)會(huì)挑選每個(gè)階段的其中一個(gè)Linear層來(lái)進(jìn)行說(shuō)明。

	以反量化的dequantize Linear推理時(shí)，在prompting階段是使用dequantize_block+cublasSgemm實(shí)現(xiàn)的，其中前者是一個(gè)反量化kernel，將weight反量化為FP32 ，后者就是直接調(diào)庫(kù)的gemm，沒(méi)啥好說(shuō)的。所以我們主要來(lái)看看generation階段的實(shí)現(xiàn)。前面說(shuō)過(guò)生成階段處理的tensor shape相對(duì)于prompting階段不同，在generation階段 以生成Q K V的Linear層為例，其input tensor 's shape 是[batch_size, 1 , hidden_dim] ,而weight還是一個(gè)矩陣數(shù)據(jù)，所以此時(shí)llama.cpp會(huì)去調(diào)用矩陣向量乘法的kernel，具體而言就是 dequantize_mul_mat_vec,  如下圖為Nsight 抓取的generation階段生成QKV的三個(gè)Linear的調(diào)用情況

	

	dequantize-linear

	具體的核函數(shù)調(diào)用函數(shù)和kernel代碼如下所示。這三個(gè)Linear 調(diào)用的都是dequantize_mul_mat_vec算子，girdDim=(1,4096,1) ，blockDim=(32,1,1) ，輸入tensor shape是[1, 1, 4096]，而weight shape則是[1, 4096 , 4096]，所以CUDA kernel用了4096個(gè)block來(lái)處理整個(gè)gemv，每個(gè)block處理一行的weight * input tensor。

	input tensor shape和weight shape依然是以Llama-2 7B模型，batch_size = 1 為例說(shuō)明

	又因?yàn)閣eight是以結(jié)構(gòu)體的方式存儲(chǔ)的32個(gè)weight數(shù)據(jù)共享一個(gè)scale數(shù)據(jù)，所以結(jié)構(gòu)體的數(shù)量為(4096 * 4096 / 32)。kernel的具體寫法在以下代碼塊中做了詳細(xì)注釋。

	dequantize_mul_mat_vec的gridDim = {1,4096,1} , blockDim = {32,1,1}，也就說(shuō)對(duì)于[1,4096] * [4096,4096]的矩陣向量乘，一共用了4096個(gè)block來(lái)處理，每個(gè)block中32個(gè)線程，每個(gè)block處理中的線程處理一個(gè)[1,4096] * [1,4096]的逐元素乘后累和。

	
//callkernel
#defineQK8_032
#defineQR8_01
staticvoiddequantize_mul_mat_vec_q8_0_cuda(constvoid*vx,constdfloat*y,float*dst,constintncols,constintnrows,cudaStream_tstream){
GGML_ASSERT(ncols%GGML_CUDA_DMMV_X==0);
constintblock_num_y=(nrows+GGML_CUDA_MMV_Y-1)/GGML_CUDA_MMV_Y;
constdim3block_nums(1,block_num_y,1);//(1,4096,1)
constdim3block_dims(WARP_SIZE,GGML_CUDA_MMV_Y,1);//(32,1,1)
dequantize_mul_mat_vec
<<>>(vx,y,dst,ncols,nrows);
}

//kernelcode
static__device____forceinline__voiddequantize_q8_0(constvoid*vx,constintib,constintiqs,dfloat2&v){
//因?yàn)榇藭r(shí)的int8量化是采用的是均勻?qū)ΨQ量化
//根據(jù)量化公式，反量化就是int8*scale
constblock_q8_0*x=(constblock_q8_0*)vx;

constdfloatd=x[ib].d;//scale

v.x=x[ib].qs[iqs+0];//int8weight
v.y=x[ib].qs[iqs+1];//int8weight

#ifdefGGML_CUDA_F16
//FP16的情況
v=__hmul2(v,{d,d});
#else
//FP32的情況
v.x*=d;//反量化
v.y*=d;//反量化
#endif//GGML_CUDA_F16
}

//gridDim={1,4096,1},blockDim={32,1,1}

template
static__global__voiddequantize_mul_mat_vec(constvoid*__restrict__vx,constdfloat*__restrict__y,float*__restrict__dst,constintncols,constintnrows){
//qk=quantizedweightsperxblock
//qr=numberofquantizedweightsperdatavalueinxblock

constintrow=blockIdx.y*blockDim.y+threadIdx.y;//0-4095

if(row>=nrows){
return;
}

constinttid=threadIdx.x;//0-31

constintiter_stride=2*GGML_CUDA_DMMV_X;//2*32
constintvals_per_iter=iter_stride/WARP_SIZE;//2numquantizedvalsperthreadandiiter
//單個(gè)線程，for循環(huán)的一次迭代所處理的數(shù)據(jù)量--2
constinty_offset=qr==1?1:qk/2;//1

//partialsumforeachthread
#ifdefGGML_CUDA_F16
half2tmp={0.0f,0.0f};//twosumsforf16totakeadvantageofhalf2intrinsics
#else
floattmp=0.0f;
#endif//GGML_CUDA_F16
//32個(gè)線程需要處理4096組數(shù)據(jù)的乘加

for(inti=0;i2valuesperiiterisfasterforfastGPUs
//前面說(shuō)過(guò)一個(gè)線程每次迭代都處理兩個(gè)數(shù)據(jù)，向量化存取有效利用量化所節(jié)省的帶寬
#pragmaunroll
for(intj=0;j0;mask>>=1){
tmp+=__shfl_xor_sync(0xffffffff,tmp,mask,32);
}
//當(dāng)tid=0時(shí)再把每個(gè)block的結(jié)果寫會(huì)結(jié)果
if(tid==0){
#ifdefGGML_CUDA_F16
dst[row]=tmp.x+tmp.y;
#else
dst[row]=tmp;
#endif//GGML_CUDA_F16
}
}


	

	quantize Linear

	Linear層的第二種調(diào)用方式就是對(duì)輸入tensor做量化，之后再與int8的weight做int8的運(yùn)算。同樣，我們以generation階段的生成QKV的三個(gè)Linear之一為例的實(shí)現(xiàn)進(jìn)行說(shuō)明。通過(guò)Nsight 抓取的kernel 調(diào)用情況可以發(fā)現(xiàn)，每個(gè)mul_mat_vec_q在被調(diào)用之前，都會(huì)有一個(gè)quantize_q8_1 ，quantize_q8_1用于對(duì)輸入tensor進(jìn)行量化，mul_mat_vec_q則是進(jìn)行int8的矩陣向量乘法。前面說(shuō)過(guò)此時(shí)的輸入tensor shape是[1, 1, 4096]，weight shape是[1, 4096 , 4096]，其中weight的數(shù)據(jù)依然是采用block_q8_0這種結(jié)構(gòu)體的方式存儲(chǔ)。

	

	quantize-linear

	具體的kernel實(shí)現(xiàn)和調(diào)用如下

	量化函數(shù) ——quantize_q8_1

	quantize_q8_1 kernel用于對(duì)數(shù)據(jù)進(jìn)行int8的對(duì)稱均勻量化，具體而言，此時(shí)就是對(duì)輸入shape為[1,1,4096]的fp32數(shù)據(jù)進(jìn)行量化，girdDim ={16,1,1} , blockDim={256 , 1, 1}

	16 * 256 = 4096，實(shí)際就是一個(gè)線程量化一個(gè)數(shù)據(jù)

	
#defineQK8_132
typedefstruct{
half2ds;//ds.x=delta,ds.y=sum，這里的ds存了兩個(gè)half第一個(gè)half是scale，第二half是sum
int8_tqs[QK8_0];//quants，int8數(shù)據(jù)
}block_q8_1;
static__global__voidquantize_q8_1(constfloat*__restrict__x,void*__restrict__vy,constintkx,constintkx_padded){
constintix=blockDim.x*blockIdx.x+threadIdx.x;//0-4096

if(ix>=kx_padded){
return;
}
constintiy=blockDim.y*blockIdx.y+threadIdx.y;//0
constinti_padded=iy*kx_padded+ix;//ix
block_q8_1*y=(block_q8_1*)vy;

constintib=i_padded/QK8_1;//blockindex因?yàn)榻Y(jié)構(gòu)體數(shù)據(jù)是以32為一組，所以ib計(jì)算得到當(dāng)前數(shù)據(jù)所在結(jié)構(gòu)體block的index
constintiqs=i_padded%QK8_1;//quantindexiqs計(jì)算的就是當(dāng)前數(shù)據(jù)所在結(jié)構(gòu)體內(nèi)部的index

constfloatxi=ix0;mask>>=1){
amax=fmaxf(amax,__shfl_xor_sync(0xffffffff,amax,mask,32));
sum+=__shfl_xor_sync(0xffffffff,sum,mask,32);
}
//套用均勻?qū)ΨQ量化的量化公式
//q=round(clip(r_i/scale,Q_{min},Q_{max}))
//scale=fmax-fmin/qmax-qmin
constfloatd=amax/127;
constint8_tq=amax==0.0f?0:roundf(xi/d);
//存儲(chǔ)量化后的值
y[ib].qs[iqs]=q;

if(iqs>0){
return;
}
//只用iqs==0的線程將scale和sum寫回
y[ib].ds.x=d;
y[ib].ds.y=sum;
}


	

	量化的矩陣向量乘——mul_mat_vec_q

	對(duì)input_tensor調(diào)用了量化函數(shù)之后，再調(diào)用mul_mat_vec_q執(zhí)行int8的輸入與int8的weight之間的矩陣向量乘法運(yùn)算，mul_mat_vec_q的gridDim = {1,4096,1} blockDim ={32,1,1}，所以同樣是每個(gè)block處理一行的weight與input_tensor做乘加運(yùn)算，一共4096行。每個(gè)block內(nèi)部32個(gè)線程又處理4096個(gè)數(shù)據(jù)的乘加運(yùn)算

	
#defineVDR_Q8_0_Q8_1_MMVQ2
staticvoidmul_mat_vec_q8_0_q8_1_cuda(constvoid*vx,constvoid*vy,float*dst,constintncols,constintnrows,cudaStream_tstream){
GGML_ASSERT(ncols%QK8_0==0);
constintblock_num_y=(nrows+GGML_CUDA_MMV_Y-1)/GGML_CUDA_MMV_Y;
constdim3block_nums(1,block_num_y,1);
constdim3block_dims(WARP_SIZE,GGML_CUDA_MMV_Y,1);
//QK8_0=32,QI8_0=8
mul_mat_vec_q
<<>>(vx,vy,dst,ncols,nrows);
}
templatestatic__device____forceinline__floatvec_dot_q8_0_q8_1_impl(
constint*v,constint*u,constfloat&d8_0,constfloat&d8_1){
#if__CUDA_ARCH__>=MIN_CC_DP4A//lowestcomputecapabilityforintegerintrinsics
intsumi=0;
#pragmaunroll
for(inti=0;i=MIN_CC_DP4A
}
static__device____forceinline__floatvec_dot_q8_0_q8_1(
constvoid*__restrict__vbq,constblock_q8_1*__restrict__bq8_1,constint&iqs){

constblock_q8_0*bq8_0=(constblock_q8_0*)vbq;

intv[VDR_Q8_0_Q8_1_MMVQ];
intu[VDR_Q8_0_Q8_1_MMVQ];

#pragmaunroll
for(inti=0;iqs,iqs+i);//？
u[i]=get_int_from_int8_aligned(bq8_1->qs,iqs+i);
}

returnvec_dot_q8_0_q8_1_impl(v,u,bq8_0->d,bq8_1->ds.x);
}

//gridDim={1,4096,1}blockDim={32,1,1}
//qk=32，qi=8，block_q_t=block_q8_0，vdr=2，vec_dot_q_cuda=vec_dot_q8_0_q8_1
template
static__global__voidmul_mat_vec_q(constvoid*__restrict__vx,constvoid*__restrict__vy,float*__restrict__dst,constintncols,constintnrows){
constintrow=blockIdx.y*blockDim.y+threadIdx.y;//theindexofrow

if(row>=nrows){
return;
}
constintblocks_per_row=ncols/qk;//4096/32=128一行數(shù)據(jù)共有128個(gè)block_q8_0的結(jié)構(gòu)體
constintblocks_per_warp=vdr*WARP_SIZE/qi;//2*32/8=8一個(gè)warp一次處理8個(gè)block_q8_0的結(jié)構(gòu)體
//所以，每輪迭代4個(gè)線程處理一個(gè)block_q8_0的結(jié)構(gòu)體
//所以，每輪迭代每個(gè)線程處理8個(gè)int8數(shù)據(jù)

//partialsumforeachthread
floattmp=0.0f;
//block_q_t=block_q8_0
constblock_q_t*x=(constblock_q_t*)vx;//weight所對(duì)應(yīng)的結(jié)構(gòu)體指針數(shù)據(jù)
constblock_q8_1*y=(constblock_q8_1*)vy;//量化input_tensor的結(jié)構(gòu)體指針
//根據(jù)之前的關(guān)系，所以每個(gè)線程需要迭代128/8=16次
for(inti=0;i0;mask>>=1){
tmp+=__shfl_xor_sync(0xffffffff,tmp,mask,32);
}
//結(jié)果寫回
if(threadIdx.x==0){
dst[row]=tmp;
}
}


	

	2.2.3 rope_f32

	前文Llama 2詳解 對(duì)rope 這種位置編碼的方式說(shuō)明過(guò)，我們這里在回憶一下rope的公式

	

	對(duì)于rope的處理，只會(huì)對(duì)Q和K進(jìn)行位置編碼，通過(guò)Nsight 抓取的kernel調(diào)用也能發(fā)現(xiàn)

	

	rope

	
//(1,32,1)(512,1,1)
static__global__voidrope_f32(constfloat*x,float*dst,constintncols,constfloatp0,
constfloatp_delta,constintp_delta_rows,constfloattheta_scale){

constintcol=2*(blockDim.x*blockIdx.x+threadIdx.x);//0-1022
//ncols=128
if(col>=ncols){
return;
}
//做了截?cái)嗨詂ol的值域?yàn)閧0，2，4...126}
//其實(shí)這里不是太懂為什么要512個(gè)線程處理，然后又做截?cái)?，?shí)際每個(gè)block只有64個(gè)線程進(jìn)行了后續(xù)運(yùn)算
constintrow=blockDim.y*blockIdx.y+threadIdx.y;//0-31
constinti=row*ncols+col;//數(shù)據(jù)的索引
//p0=n_past就是在生成當(dāng)前token之前已處理的token長(zhǎng)度
//p_delta=1.0
//theta_scale=0.865964
//p_delta_rows=32
constfloattheta=(p0+p_delta*(row/p_delta_rows))*powf(theta_scale,col/2);
constfloatsin_theta=sinf(theta);
constfloatcos_theta=cosf(theta);

constfloatx0=x[i+0];
constfloatx1=x[i+1];
//這里用了32個(gè)block處理32個(gè)頭的rope計(jì)算
//其中每個(gè)block中又只有64
dst[i+0]=x0*cos_theta-x1*sin_theta;
dst[i+1]=x0*sin_theta+x1*cos_theta;
}


	

	2.1.4 Copy Kernel

	我們之前說(shuō)過(guò)，無(wú)論在prompting階段還是generation階段，生成的K和V都是要緩存下來(lái)的，區(qū)別在于prompting階段是將提示token對(duì)應(yīng)的KV直接寫入，而generation階段則是將單個(gè)token對(duì)應(yīng)的KV追加至KV cache。在對(duì)K V緩存時(shí)，可以直接對(duì)FP32數(shù)據(jù)進(jìn)行緩存，也可以通過(guò)將FP32數(shù)據(jù)轉(zhuǎn)換為FP16之后再進(jìn)行緩存，后者雖然會(huì)損失一定的精度，但是節(jié)省了顯存。如下兩個(gè)kernel用于對(duì)數(shù)據(jù)類型進(jìn)轉(zhuǎn)換

	llama.cpp 中會(huì)提前為KV cache分配顯存空間，然后prompting階段和generation階段生成的KV都會(huì)寫入。如當(dāng)最大context大小設(shè)置為512時(shí)，以FP32為例，每一個(gè)Transformer Block會(huì)分別給K cache 和 V cache分配512 * 4096 * 4 = 8MB 的存儲(chǔ)空間，KV cache一共16MB，那么32個(gè)Transformer Block一共512 MB的 KV cache空間。如果使用FP16緩存，則KV cache空間減半

	
//fp32->fp32
static__device__voidcpy_1_f32_f32(constchar*cxi,char*cdsti){
constfloat*xi=(constfloat*)cxi;
float*dsti=(float*)cdsti;

*dsti=*xi;
}
//fp32->fp16
static__device__voidcpy_1_f32_f16(constchar*cxi,char*cdsti){
constfloat*xi=(constfloat*)cxi;
half*dsti=(half*)cdsti;

*dsti=__float2half(*xi);//通過(guò)內(nèi)置函數(shù)將數(shù)據(jù)從fp32轉(zhuǎn)換為fp16
}



	

	cpy_f32_f16就是實(shí)際的顯存拷貝核函數(shù)，通過(guò)上述的兩個(gè)kernel的調(diào)用將數(shù)據(jù)拷貝至fp32或fp16。

	
template
static__global__voidcpy_f32_f16(constchar*cx,char*cdst,constintne,
constintne00,constintne01,constintnb00,constintnb01,constintnb02,
constintne10,constintne11,constintnb10,constintnb11,constintnb12){
constinti=blockDim.x*blockIdx.x+threadIdx.x;

if(i>=ne){
return;
}
//determineindicesi02/i12,i01/i11,i00/i10asafunctionofindexiofflattenedtensor
//thencombinethoseindiceswiththecorrespondingbyteoffsetstogetthetotaloffsets
//結(jié)合之前的ggml_tensor的ne和nb的定義
//nb[i]=nb[i-1]*ne[i-1],nb[0]=sizeof(type)

constinti02=i/(ne00*ne01);//theindexofne02
constinti01=(i-i02*ne01*ne00)/ne00;//theindexofne01
constinti00=i-i02*ne01*ne00-i01*ne00;//theindexofne00
constintx_offset=i00*nb00+i01*nb01+i02*nb02;//計(jì)算偏移

constinti12=i/(ne10*ne11);//dst同上
constinti11=(i-i12*ne10*ne11)/ne10;
constinti10=i-i12*ne10*ne11-i11*ne10;
constintdst_offset=i10*nb10+i11*nb11+i12*nb12;

cpy_1(cx+x_offset,cdst+dst_offset);//將cx[x_offset]轉(zhuǎn)換為fp16寫到cdst[dst_offset]
}


	

	另外，注意一點(diǎn)，一旦對(duì)tensor做了view或者reshape之類的操作使得內(nèi)存排布不在連續(xù)，nb[i] = nb[i-1] * ne[i-1]這個(gè)條件可能就不滿足了

	2.1.5 Multi-Head-Attention

	MHA 是整個(gè)Transformer Block中最核心的Kernel了，Attention的計(jì)算公式如下

	

	在前文中我們也說(shuō)過(guò)Llama 2 會(huì)采用一種分組共享KV cache的Attention計(jì)算GQA，但是因?yàn)槲覀兪且?B模型為例進(jìn)行說(shuō)明的，7B模型并沒(méi)有采用GQA，依然是采用的MHA，可參考Llama repo。所以本文還是以MHA為例進(jìn)行說(shuō)明此處的kernel調(diào)用情況

	對(duì)于Q*K 和 Attention Score * V這兩個(gè)乘法操作，在prompting階段和generation階段所調(diào)用的算子并不一樣，在prompting階段因?yàn)镼KV三個(gè)都是多個(gè)單詞所對(duì)應(yīng)的tensor，即shape為[1, 32, seq_len , 128] ，所以prompting階段在處理時(shí)依然是直接調(diào)用的cublasSgemm實(shí)現(xiàn)。所以我們還是主要來(lái)看看generation階段所調(diào)用的算子，在generation階段 Q的shape為[1, 32, 1 , 128] ，Q需要與新生成的K V，以及之前推理緩存下來(lái)的KV cache一起做self-attention運(yùn)算

	

	上圖為generation階段Multi-Head-Attention的算子調(diào)用，主要包括5個(gè)kernel

	Q*K 的算子——mul_mat_p021_f16_f32

	generation階段此時(shí)Q的shape是[1, 32,1,128] , K cache的shape為[1,32,seq_len,128]，此處的seq_len就是當(dāng)前處理的tokens的數(shù)量，會(huì)隨著generation階段，逐漸加1。如下kernel

	這里說(shuō)一句，這里說(shuō)的Q和K的shape只是方便理解Attention的計(jì)算過(guò)程，實(shí)際Q和K在物理內(nèi)存上可能不是按shape排列的方式存儲(chǔ)的，比如這里的K cache，內(nèi)存存放的次序還是[1, seq_len, 4096]

	以下kernel調(diào)用時(shí)的gridDim = {1,seq_len,32} ,blockDim = {32,1,1}，也就是說(shuō)一個(gè)block處理一個(gè)頭中的Q與K cache中的一行K進(jìn)行乘加運(yùn)算

	
//gridDim={1,seq_len,32},blockDim={32,1,1}
static__global__voidmul_mat_p021_f16_f32(
constvoid*__restrict__vx,constfloat*__restrict__y,float*__restrict__dst,
constintncols_x,constintnrows_x,constintnchannels_x,constintnchannels_y){

consthalf*x=(consthalf*)vx;//vx就是Kcache

constintrow_x=blockDim.y*blockIdx.y+threadIdx.y;//這個(gè)維度是seq_len的索引，[0,..,seq_len-1]
constintchannel=blockDim.z*blockIdx.z+threadIdx.z;//這個(gè)維度是multihead的索引[0,1,2..,31]
constintchannel_x=channel/(nchannels_y/nchannels_x);//這個(gè)是對(duì)于GQA的時(shí)候用的，就是Q分組共享Kcache
//此處我們是以7B模型為例，依然是MHA

constintnrows_y=ncols_x;//128
constintnrows_dst=nrows_x;//seq_len
constintrow_dst=row_x;//[0,..,seq_len-1]

floattmp=0.0f;
//因?yàn)橐粋€(gè)block(32個(gè)線程)處理128個(gè)數(shù)據(jù)，所以每個(gè)線程for循環(huán)迭代次數(shù)為128/32
for(intcol_x0=0;col_x0=ncols_x){
break;
}

//xistransposedandpermuted
//計(jì)算Kcache的index
//前面說(shuō)過(guò)Kcache在內(nèi)存存的次序還是[seq_len,multihead,head_dim]
//所以這里的index的計(jì)算方式理解一下
constintix=row_x*nchannels_x*ncols_x+channel_x*ncols_x+col_x;
//Kcache不是為了節(jié)省內(nèi)存用的FP16存著嘛，所以用一個(gè)__half2float內(nèi)置函數(shù)將FP16轉(zhuǎn)換為FP32
constfloatxi=__half2float(x[ix]);
//Kcache的列索引等于Q的列索引
//名字叫row_y但還是列索引，因?yàn)镼的內(nèi)存排布還是[32,128]
constintrow_y=col_x;

//yisnottransposedbutpermuted
constintiy=channel*nrows_y+row_y;//計(jì)算Q的全局index

tmp+=xi*y[iy];//乘后累和到tmp
}

//dstisnottransposedandnotpermuted
//dst的shape為[32,1,seq_len],所以內(nèi)存排布為[32,seq_len]
//所以dst的index計(jì)算方式如下
constintidst=channel*nrows_dst+row_dst;

//sumuppartialsumsandwritebackresult
//又是熟悉的block內(nèi)求和
#pragmaunroll
for(intmask=16;mask>0;mask>>=1){
tmp+=__shfl_xor_sync(0xffffffff,tmp,mask,32);
}

if(threadIdx.x==0){
dst[idst]=tmp;//寫回dst
}
}


	

	除以 —— scale_f32

	Attention(Q,K,V)公式中的Q乘K之后的除以，這個(gè)kernel沒(méi)啥好說(shuō)的，按元素乘

	
static__global__voidscale_f32(constfloat*x,float*dst,constfloatscale,constintk){
constinti=blockDim.x*blockIdx.x+threadIdx.x;

if(i>=k){
return;
}

dst[i]=scale*x[i];
}


	

	attention mask —— diag_mask_inf_f32

	對(duì)于Attention中的mask操作，我們看下面這個(gè)圖，一目了然，在原生Transformer Decode階段，加入mask的是為了防止前面token的Q與后面token的K計(jì)算得到一個(gè)較高的Attention Score，所以通過(guò)一個(gè)上三角(且上三角元素全為-INF)矩陣，來(lái)保證句子中單詞之間的時(shí)序性。

	

	attention-mask

	如下kernel就是實(shí)現(xiàn)mask操作的，也是逐元素，根據(jù)其坐標(biāo)來(lái)判斷是否需要mask。不過(guò)這里多說(shuō)一句，在通過(guò)Nsight抓取的MHA部分的kernel調(diào)用情況的截圖中可以看到，generation階段也調(diào)用了diag_mask_inf_f32這個(gè)kernel，實(shí)際是不需要調(diào)用的。因?yàn)樯呻A段生成的Q就是最新的單詞所對(duì)應(yīng)的Q，他與KV cache中的每個(gè)KV 計(jì)算的Attention Score都不會(huì)mask， mask的操作只需要存在于prompting階段中，想來(lái)這里也是因?yàn)閘lama.cpp的作者為了省事～

	這里不太明白為什么generation階段不需要mask的可以移步至B站CodeLearner

	
static__global__voiddiag_mask_inf_f32(constfloat*x,float*dst,constintncols,constintrows_per_channel,constintn_past){
constintcol=blockDim.x*blockIdx.x+threadIdx.x;
constintrow=blockDim.y*blockIdx.y+threadIdx.y;

if(col>=ncols){
return;
}

constinti=row*ncols+col;
//dst[i]=col>n_past+row?-INFINITY:x[i];
dst[i]=x[i]-(col>n_past+row%rows_per_channel)*INT_MAX;//equivalentwithinroundingerrorbutslightlyfasteronGPU
}


	

	Attention Score * V 的算子 mul_mat_vec_nc_f16_f32

	generation階段Attention Score 的shape為[1, 32 , 1, seq_len]，V的shape為[1, 32 ,seq_len,128] . mul_mat_vec_nc_f16_f32算子調(diào)用的gridDim={1,128,32} ，blockDim={32,1,1} ，所以 blockDim.z維度對(duì)應(yīng)于multihead=32 維度，blockDim.y維度對(duì)應(yīng)于head_dim=128維度，然后每個(gè)block中32個(gè)線程用來(lái)處理每個(gè)seq_len長(zhǎng)度序列的乘加。

	
//gridDim={1,128,32},blockDim={32,1,1}
static__global__voidmul_mat_vec_nc_f16_f32(//nc==non-contiguous
constvoid*__restrict__vx,constfloat*__restrict__y,float*__restrict__dst,constintncols_x,constintnrows_x,
constintrow_stride_x,constintchannel_stride_x,constintchannel_x_divisor){
//ncols_x=seq_len,nrows_x=128,row_stride_x=512,channel_stride_x=65536,channel_x_divisor=1

consthalf*x=(consthalf*)vx;//Vcache存儲(chǔ)時(shí)使用的FP16

constintrow_x=blockDim.y*blockIdx.y+threadIdx.y;//indexofhead_dim->0-127
constintchannel=blockDim.z*blockIdx.z+threadIdx.z;//indexofmulti-head->0-31
constintchannel_x=channel/channel_x_divisor;//channel/1

constintnrows_y=ncols_x;//seq_len
constintnrows_dst=nrows_x;//128
constintrow_dst=row_x;//indexofhead_dim->0-127

//AttentionScore*V最終的shape為[1,32,1,128]
//所以idst=（indexofmulti-head）*(128)+(indexofhead_dim)
constintidst=channel*nrows_dst+row_dst;

floattmp=0.0f;
//循環(huán)處理seq_len序列，每個(gè)線程處理seq_len/blockDim.x個(gè)數(shù)
for(intcol_x0=0;col_x0=ncols_x){
break;
}
//Vcache的index
constintix=channel_x*channel_stride_x+row_x*row_stride_x+col_x;
//fp16轉(zhuǎn)fp32
constfloatxi=__half2float(x[ix]);
//AttentionScoreindex
constintrow_y=col_x;
constintiy=channel*nrows_y+row_y;

tmp+=xi*y[iy];//乘加
}

//sumuppartialsumsandwritebackresult
//還是熟悉的block內(nèi)部求和
#pragmaunroll
for(intmask=16;mask>0;mask>>=1){
tmp+=__shfl_xor_sync(0xffffffff,tmp,mask,32);
}
//結(jié)果寫回
if(threadIdx.x==0){
dst[idst]=tmp;
}
}


	

	2.1.6 add_f32

	add_f32 用于殘差連接一下輸入tensor 與Attention Block的輸出，kernel的實(shí)現(xiàn)沒(méi)啥好說(shuō)的，就是最簡(jiǎn)單向量相加

	
static__global__voidadd_f32(constfloat*x,constfloat*y,float*dst,constintkx,constintky){
constinti=blockDim.x*blockIdx.x+threadIdx.x;

if(i>=kx){
return;
}
dst[i]=x[i]+y[i%ky];
}


	

	輸入tensor shape為[1,1,4096] ,Attention Block的輸出為[1,1,4096] 。在FeedForward Block最后也是同樣會(huì)調(diào)用add_f32 將FeedForward Block的輸入連殘差連接到輸出，所調(diào)用的kernel為同一個(gè)

	2.2 FeedForward Block

	FeedForward Block 上層算法流程以及其在prompting階段和generation階段所調(diào)用的CUDA算子，如下圖所示。整個(gè)過(guò)程中主要的就是幾個(gè)Linear層，在前面的2.1.2節(jié)中詳細(xì)介紹過(guò)了，所以這里就不過(guò)多贅述了～

	

	2.2.1 silu_f32

	FeedForward Block中在之前沒(méi)有出現(xiàn)過(guò)的kenrel就是silu_f32這個(gè)激活函數(shù)kernel。同樣，我們先回顧一下SiLU函數(shù)的公式

	

	
static__global__voidsilu_f32(constfloat*x,float*dst,constintk){
constinti=blockDim.x*blockIdx.x+threadIdx.x;

if(i>=k){
return;
}
//silu公式
dst[i]=x[i]/(1.0f+expf(-x[i]));
}


	

	整個(gè)Kernel在調(diào)用時(shí)blocksize 在prompting階段和generation階段的值不一樣，因?yàn)镕eedForward Block中前兩個(gè)Linear層的輸出尺寸是11008 ，所以在prompting階段需要prompting_length * 11008 個(gè)線程來(lái)處理prompting_length * 11008 個(gè)數(shù)據(jù)，而在generation階段則需要11008 個(gè)線程來(lái)處理11008 個(gè)數(shù)據(jù)。所以如下圖所示，為generation階段的調(diào)用silu的Nsight截圖，該kernel用了43個(gè)block，每個(gè)block 256個(gè)線程，= 43 * 256。prompting階段類比。

	

	silu

	至此就算把Llama 2 中完整的單個(gè)Tranformer Block中的所有l(wèi)lama.cpp調(diào)用的CUDA Kernel 說(shuō)明完啦～

	






	審核編輯：劉清