0x0. 前言

接著 大模型部署框架 FastLLM 簡(jiǎn)要解析 這篇文章首先梳理了一下FastLLM的調(diào)用鏈和關(guān)鍵的數(shù)據(jù)結(jié)構(gòu)，然后解析了 FastLLM 的一些實(shí)現(xiàn)細(xì)節(jié)和CPU/GPU后端實(shí)現(xiàn)采用的優(yōu)化技巧。

0x1. 調(diào)用鏈和數(shù)據(jù)結(jié)構(gòu)解析

以chatglm-6b的支持為例，函數(shù)入口在 https://github.com/ztxz16/fastllm/blob/master/src/models/chatglm.cpp#L626 ，這里的 input 就是輸入的 context（string類型）。然后 https://github.com/ztxz16/fastllm/blob/master/src/models/chatglm.cpp#L633 這行代碼對(duì) input 進(jìn)行 tokenizer encode并構(gòu)造好inputIds，再構(gòu)造好attentionMask之后就可以給Forward函數(shù)推理，拿到推理結(jié)果之后再使用tokenizer進(jìn)行decode得到輸出。

在這里，inputIds和attentionMask都是Data數(shù)據(jù)類型，類比于PyTorch的Tensor，來(lái)對(duì)輸入數(shù)據(jù)以及device，shape等信息進(jìn)行統(tǒng)一管理。下面的代碼展示了Data數(shù)據(jù)結(jié)構(gòu)的定義，源碼在：https://github.com/ztxz16/fastllm/blob/master/include/fastllm.h#L201-L286

classData{
public:
boollockInCPU=false;//如果lock在CPU上，那么不允許移動(dòng)到其余設(shè)備
WeightTypeweightType=WeightType::NONE;//權(quán)重類型，NONE代表非權(quán)重（或未知權(quán)重）

DataTypedataType=DataType::FLOAT32;//數(shù)據(jù)類型
intunitSize,unitSizeDiv=1;//單個(gè)元素的字節(jié)數(shù)=unitSIze/unitSizeDiv

std::vectordims;//數(shù)據(jù)形狀
std::vectorstrides;//跨度

uint64_texpansionSize=0;//擴(kuò)容后的尺寸
uint64_texpansionBytes=0;//擴(kuò)容后的字節(jié)數(shù)
std::vectorexpansionDims;//預(yù)擴(kuò)容的形狀
uint8_t*cpuData=nullptr;//數(shù)據(jù)指針

void*cudaData=nullptr;
std::vectorextraCudaData;

void*deviceData=nullptr;
std::vectorextraDeviceData;

DataDevicedataDevice=DataDevice::CPU;

//這兩個(gè)參數(shù)用于量化，對(duì)FLOAT數(shù)據(jù)不適用
intperChannelAxis=-1;//沿哪個(gè)軸分通道量化，-1代表沒有分通道
std::vectorperChannelsConfigs;//perChannelsConfigs[i]代表第i個(gè)通道的min,max;如果沒有分通道，perChannelsConfigs[0]代表全局min,max
std::vectorscales,mins;
std::vectorzeros;
std::vectorweightSum;//作為權(quán)重時(shí)，有時(shí)候需要存一些和加速計(jì)算

std::stringfileName;
longlongfilePos;
std::shared_ptrm_file;

Data(){};

Data(DataTypetype);

Data(DataTypetype,conststd::vector&dims);//構(gòu)造函數(shù)

//構(gòu)造函數(shù)，創(chuàng)建好之后從data復(fù)制數(shù)據(jù)
//data中是原始數(shù)據(jù)，如果type不是float那么需要量化
Data(DataTypetype,conststd::vector&dims,conststd::vector&data);

~Data();//析構(gòu)函數(shù)

Data(constData&ori);//深拷貝

voidCopyFrom(constData&ori);//復(fù)制

uint64_tGetBytes()const;//獲取總字節(jié)數(shù)

voidAllocate();//分配內(nèi)存

voidAllocate(floatv);//分配內(nèi)存并初始化

voidExpansion(conststd::vector&dims);//預(yù)擴(kuò)容到相應(yīng)尺寸

voidMallocSpace(uint64_tsize);//在設(shè)備上分配

voidFreeSpace();//回收設(shè)備上的內(nèi)存

voidUpdateUnitSize();//更新unitSize

voidResize(conststd::vector&dims);//更改尺寸

voidReshape(conststd::vector&dims);//更改尺寸,但不修改數(shù)據(jù)

uint64_tCount(inti)const;//dims[i]*strides[i]

voidPrintShape()const;//輸出形狀

voidPrint()const;//輸出

voidCalcWeightSum();//計(jì)算WeightSum

voidToDevice(DataDevicedevice);//移動(dòng)到指定device

voidToDevice(void*device);

voidset_file(std::shared_ptrfile){
m_file=file;
}
};

在Forward函數(shù)里面，以Data為核心載體，運(yùn)行chatglm-6b模型的流程，具體包含如下的一些算子：https://github.com/ztxz16/fastllm/blob/master/include/fastllm.h#L346-L408 。以Permute為例我們?yōu)g覽下它的實(shí)現(xiàn)：

voidPermute(constData&input,conststd::vector&axis,Data&output){
DataaxisData=Data(DataType::INT32PARAM,{(int)axis.size()});
axisData.Allocate();
for(inti=0;iRun("Permute",{
{"input",(Data*)&input},{"axis",&axisData},{"output",(Data*)&output}
},{},{});
}

這里的curExecutor負(fù)責(zé)根據(jù)FastLLM編譯開啟的后端選項(xiàng)把算子Dispatch到不同的device進(jìn)行執(zhí)行，{"input", (Data*)&input}, {"axis", &axisData}, {"output", (Data*)&output}} 這行代碼表示的是一個(gè)DataDict對(duì)象，也就是一個(gè)值為data的字典，原始定義為typedef std::map DataDict;。接著我們看一下curExecutor的定義和實(shí)現(xiàn)：

namespacefastllm{
classExecutor{
private:
std::vectordevices;
std::mapprofiler;

public:
Executor();//創(chuàng)建默認(rèn)的Executor

~Executor();//析構(gòu)

voidClearDevices();//清空devices

voidAddDevice(BaseDevice*device);//增加一個(gè)device

//運(yùn)行一個(gè)op
voidRun(conststd::string&opType,constfastllm::DataDict&datas,constfastllm::FloatDict&floatParams,
constfastllm::IntDict&intParams);

voidClearProfiler();

voidPrintProfiler();
};
}

從Executor類的定義我們可以判斷它負(fù)責(zé)了在設(shè)定的devices上根據(jù)opType和輸入數(shù)據(jù)等執(zhí)行Op的前向計(jì)算，也就是Run這個(gè)接口。由于Executor類是FastLLM的調(diào)度核心實(shí)現(xiàn)，所以我們來(lái)詳細(xì)解析一下它的實(shí)現(xiàn)。

namespacefastllm{
Executor::Executor(){
this->devices.clear();
#ifdefUSE_CUDA
//將一個(gè)指向CudaDevice類對(duì)象的指針插入到devices向量的末尾。
//這里通過(guò)new運(yùn)算符創(chuàng)建了一個(gè)CudaDevice對(duì)象，并將返回的指針進(jìn)行類型轉(zhuǎn)換為BaseDevice*類型。
this->devices.push_back((BaseDevice*)newCudaDevice());
#endif
this->devices.push_back((BaseDevice*)newCpuDevice());
}

Executor::~Executor(){
//釋放devices向量中的每個(gè)指針元素所占用的內(nèi)存。
for(inti=0;idevices指的是當(dāng)前對(duì)象的devices成員，即指向BaseDevice類對(duì)象的指針向量。
this->devices.clear();
}

//該函數(shù)用于向devices向量中添加一個(gè)指向BaseDevice類對(duì)象的指針。
voidExecutor::AddDevice(fastllm::BaseDevice*device){
this->devices.push_back(device);
}

voidExecutor::Run(conststd::string&opType,constfastllm::DataDict&datas,constfastllm::FloatDict&floatParams,
constfastllm::IntDict&intParams){
//創(chuàng)建一個(gè)st變量，用于記錄函數(shù)開始執(zhí)行的時(shí)間。
autost=std::now();
//創(chuàng)建一個(gè)布爾變量lockInCPU，用于記錄是否將數(shù)據(jù)鎖定在CPU上。
boollockInCPU=false;
//在第一個(gè)for循環(huán)中，遍歷數(shù)據(jù)字典datas，查找是否有"___batch"后綴的參數(shù)，
//并根據(jù)情況設(shè)置lockInCPU的值。it.first是數(shù)據(jù)字典中的鍵（key），it.second
//是對(duì)應(yīng)的值（value）。如果存在"___batch"后綴的參數(shù)，則將lockInCPU設(shè)置為
//對(duì)應(yīng)數(shù)據(jù)的lockInCPU屬性（布爾值），否則設(shè)置為當(dāng)前數(shù)據(jù)的lockInCPU屬性。
for(auto&it:datas){
if(intParams.find(it.first+"___batch")!=intParams.end()){
intbatch=intParams.find(it.first+"___batch")->second;
for(inti=0;ilockInCPU;
}
}else{
lockInCPU|=it.second->lockInCPU;
}
}
//第二個(gè)for循環(huán)遍歷devices向量中的所有設(shè)備指針device。
//在循環(huán)中，首先檢查lockInCPU是否為真，并且當(dāng)前設(shè)備的類型不是"cpu"，
//如果是，則跳過(guò)當(dāng)前設(shè)備（continue）。這個(gè)檢查是為了保證數(shù)據(jù)鎖定在CPU上時(shí)，只執(zhí)行CPU設(shè)備上的操作。
for(autodevice:devices){
if(lockInCPU&&device->deviceType!="cpu"){
continue;
}
//然后，通過(guò)調(diào)用device->CanRun(opType,datas,floatParams,intParams)
//檢查當(dāng)前設(shè)備是否可以運(yùn)行指定的操作opType。如果可以運(yùn)行，則進(jìn)行以下操作：
if(device->CanRun(opType,datas,floatParams,intParams)){
//第三個(gè)for循環(huán)遍歷數(shù)據(jù)字典datas，如果存在"___batch"后綴的參數(shù)，
//則將對(duì)應(yīng)數(shù)據(jù)轉(zhuǎn)移到當(dāng)前設(shè)備上；否則，將當(dāng)前數(shù)據(jù)轉(zhuǎn)移到當(dāng)前設(shè)備上。
for(auto&it:datas){
if(intParams.find(it.first+"___batch")!=intParams.end()){
intbatch=intParams.find(it.first+"___batch")->second;
for(inti=0;iToDevice((void*)device);
}
}else{
it.second->ToDevice((void*)device);
}
}
//調(diào)用device->Reshape(opType,datas,floatParams,intParams)
//進(jìn)行形狀推導(dǎo)，device上的形狀推導(dǎo)調(diào)用了opType對(duì)應(yīng)的op的形狀推導(dǎo)，
//并且被各個(gè)不同的op重寫。
device->Reshape(opType,datas,floatParams,intParams);
//對(duì)opType對(duì)應(yīng)的這個(gè)算子進(jìn)行推理。
device->Run(opType,datas,floatParams,intParams);
break;
}
}
//最后，計(jì)算操作運(yùn)行時(shí)間，并將其加入profiler成員變量，用于性能分析。
floatspend=GetSpan(st,std::now());
profiler[opType]+=spend;
}

//清除profile的信息
voidExecutor::ClearProfiler(){
profiler.clear();
}

//打印profile信息，也即輸出每個(gè)層的運(yùn)行時(shí)間和模型的總運(yùn)行時(shí)間
voidExecutor::PrintProfiler(){
floatsum=0.0;
for(auto&it:profiler){
printf("%sspend%f
",it.first.c_str(),it.second);
sum+=it.second;
}
printf("totalspend%f
",sum);
}
}

自此，前向計(jì)算就順利完成了，再把推理結(jié)果給 tokenizer 解碼就結(jié)束了，整體的調(diào)度執(zhí)行流程是很簡(jiǎn)單明了的。

0x2. tokenizer 解析

接著，我們來(lái)解析一下tokenizer的實(shí)現(xiàn)。先看一下tokenizer的定義（https://github.com/ztxz16/fastllm/blob/master/include/fastllm.h#L287-L310）：

structTokenizer{
structTrieNode{
inttokenId;
std::mapnext;
TrieNode();
};
TrieNode*root;

std::unordered_maptokenToStringDict;

Tokenizer();

~Tokenizer();

voidClear();//清空分詞器

voidInsert(conststd::string&s,inttokenId);//插入一個(gè)token

DataEncode(conststd::string&s);//編碼

std::stringDecode(constData&data);//解碼

std::stringDecodeTokens(conststd::vector&tokens);//解碼
};

我們從實(shí)現(xiàn)來(lái)看tokenizer的細(xì)節(jié)：

//這是Tokenizer類的嵌套結(jié)構(gòu)TrieNode的構(gòu)造函數(shù)的實(shí)現(xiàn)。
//在構(gòu)造函數(shù)中，將tokenId成員變量的值初始化為-999999。
//這個(gè)值在構(gòu)造函數(shù)中被硬編碼，它是作為一個(gè)特殊標(biāo)記來(lái)使用的。
Tokenizer::TrieNode(){
this->tokenId=-999999;
}

//Tokenizer類的構(gòu)造函數(shù)的實(shí)現(xiàn)。
//在構(gòu)造函數(shù)中，通過(guò)new運(yùn)算符創(chuàng)建一個(gè)新的TrieNode對(duì)象，
//并將其指針賦值給root成員變量。這樣，構(gòu)造函數(shù)創(chuàng)建了一個(gè)空的字典樹，
//并將其根節(jié)點(diǎn)指針存儲(chǔ)在root中。
Tokenizer::Tokenizer(){
root=newTrieNode();
}

//Tokenizer類的析構(gòu)函數(shù)的實(shí)現(xiàn)。
//在析構(gòu)函數(shù)中，首先調(diào)用Clear()函數(shù)，用于釋放動(dòng)態(tài)分配的資源和清空數(shù)據(jù)。
//然后，調(diào)用delete運(yùn)算符釋放通過(guò)new運(yùn)算符創(chuàng)建的root對(duì)象的內(nèi)存，從而釋放整個(gè)字典樹的內(nèi)存。
Tokenizer::~Tokenizer(){
Clear();
deleteroot;
}

//這是Tokenizer類的成員函數(shù)Clear()的定義，用于清空分詞器并釋放動(dòng)態(tài)分配的資源。
voidTokenizer::Clear(){
//創(chuàng)建一個(gè)指向TrieNode的指針向量q，用于輔助遍歷字典樹。
std::vectorq;
//將字典樹的根節(jié)點(diǎn)root加入q向量，作為遍歷的起始點(diǎn)。
q.push_back(root);
//開始遍歷q向量中的節(jié)點(diǎn)，這是一個(gè)廣度優(yōu)先搜索（BFS）的過(guò)程。
for(inti=0;inext){
//將當(dāng)前節(jié)點(diǎn)now的子節(jié)點(diǎn)加入q向量中，以便繼續(xù)遍歷子節(jié)點(diǎn)的子節(jié)點(diǎn)。
q.push_back(it.second);
}
}
//當(dāng)遍歷完成后，q向量中包含了字典樹中的所有節(jié)點(diǎn)。
//創(chuàng)建一個(gè)新的TrieNode對(duì)象，并將其指針賦值給root成員變量，表示創(chuàng)建了一個(gè)空的字典樹。
root=newTrieNode();
//清空tokenToStringDict映射表，以確保所有token的映射被清空。
tokenToStringDict.clear();
}

//這是Tokenizer類的成員函數(shù)Insert的定義，用于向分詞器中插入一個(gè)token。
voidTokenizer::Insert(conststd::string&s,inttokenId){
//創(chuàng)建一個(gè)指向TrieNode的指針now，并將其初始化為指向字典樹的根節(jié)點(diǎn)root。
TrieNode*now=this->root;
//開始遍歷輸入的字符串s中的每個(gè)字符。
for(inti=0;inext中添加新的子節(jié)點(diǎn)，該子節(jié)點(diǎn)的鍵為當(dāng)前字符s[i]的編碼值，
//值為指向新創(chuàng)建的TrieNode對(duì)象的指針。這表示在字典樹中添加了一個(gè)新的字符節(jié)點(diǎn)。
if(now->next.find(s[i])==now->next.end()){
now->next[s[i]]=newTrieNode();
}
//將now移動(dòng)到下一個(gè)字符s[i]對(duì)應(yīng)的節(jié)點(diǎn)，以便繼續(xù)處理下一個(gè)字符。
now=now->next[s[i]];
}
//遍歷完成后，now將指向字典樹中最后一個(gè)字符的節(jié)點(diǎn)。
//設(shè)置當(dāng)前節(jié)點(diǎn)的tokenId成員變量，表示當(dāng)前節(jié)點(diǎn)代表一個(gè)token，
//并使用傳入的tokenId值來(lái)標(biāo)識(shí)該token。
now->tokenId=tokenId;
//將傳入的tokenId和對(duì)應(yīng)的字符串s添加到tokenToStringDict
//映射表中，用于后續(xù)的解碼過(guò)程。
tokenToStringDict[tokenId]=s;
}

//這是Tokenizer類的成員函數(shù)Encode的定義，用于對(duì)輸入的字符串s進(jìn)行編碼。
DataTokenizer::Encode(conststd::string&s){
//創(chuàng)建一個(gè)浮點(diǎn)數(shù)向量v，用于存儲(chǔ)編碼結(jié)果。該向量將存儲(chǔ)找到的token對(duì)應(yīng)的tokenId值。
std::vectorv;
//開始遍歷輸入的字符串s中的每個(gè)字符。
for(inti=0;iroot;
//從當(dāng)前字符s[i]開始繼續(xù)遍歷字符串s。
for(intj=i;jnext.find(s[j])!=now->next.end()){
//將now移動(dòng)到下一個(gè)字符s[j]對(duì)應(yīng)的節(jié)點(diǎn)。
now=now->next[s[j]];
//檢查當(dāng)前節(jié)點(diǎn)now是否代表一個(gè)token，即它的tokenId是否有效。
if(now->tokenId!=-999999){
//如果當(dāng)前節(jié)點(diǎn)代表一個(gè)token，將tokenId和當(dāng)前位置j存儲(chǔ)到
//tokenId和pos變量中，以便記錄找到的token的信息。
tokenId=now->tokenId;
pos=j;
}
}else{//如果當(dāng)前字符不再是token的一部分，退出內(nèi)層循環(huán)，繼續(xù)外層循環(huán)。
break;
}
}
//如果pos大于等于當(dāng)前位置i，表示找到了一個(gè)token。
//這里pos存儲(chǔ)了找到的token的結(jié)束位置，i移動(dòng)到pos處，以便繼續(xù)遍歷下一個(gè)字符。
if(pos>=i){
i=pos;
v.push_back(tokenId);
//printf("%d",tokenId);
}
}
//printf("
");
//遍歷完成后，v向量中存儲(chǔ)了輸入字符串中所有找到的token對(duì)應(yīng)的tokenId值。
//創(chuàng)建一個(gè)Data對(duì)象并返回，表示編碼的結(jié)果。這里Data是一個(gè)數(shù)據(jù)結(jié)構(gòu)，
//用于存儲(chǔ)數(shù)據(jù)及其相關(guān)信息。編碼結(jié)果是一個(gè)一維浮點(diǎn)數(shù)數(shù)組，
//表示輸入字符串中所有找到的token對(duì)應(yīng)的tokenId值。
returnData(DataType::FLOAT32,{1,(int)v.size()},v);
}

//這是Tokenizer類的成員函數(shù)DecodeTokens的定義，
//用于對(duì)輸入的token數(shù)組進(jìn)行解碼，將token轉(zhuǎn)換回原始的字符串。
std::stringTokenizer::DecodeTokens(conststd::vector&tokens){
//創(chuàng)建一個(gè)空字符串ret，用于存儲(chǔ)解碼結(jié)果。
std::stringret="";
//開始遍歷輸入的token數(shù)組tokens。
for(inti=0;i"格式的token（其中HH表示十六進(jìn)制數(shù)），
//則需要將其轉(zhuǎn)換為對(duì)應(yīng)的字符。首先，提取HH，然后將其轉(zhuǎn)換為對(duì)應(yīng)的字符，
//并用空格代替原始的token。
if(s.size()==6&&s.substr(0,3)=="<0x"?&&?s.back()?==?'>'){
intc=0;
for(inti=3;i='0'&&s[i]<=?'9')?{
????????????????????????c?+=?(s[i]?-?'0');
????????????????????}?else?{
????????????????????????c?+=?(s[i]?-?'A'?+?10);
????????????????????}
????????????????}

????????????????s?=?"?";
????????????????s[0]?=?c;
????????????}
????????????//?根據(jù)不同的?token?進(jìn)行解碼：
????????????if?(s?==?""){
ret+="
";
}elseif(s=="<|tab|>"){
ret+="	";
}else{
ret+=s;
}
}

//將特殊字符"xE2x96x81"（UTF-8編碼）替換為空格""，這是用于表示空格的特殊字符。
std::stringblank="";
blank+=226,blank+=150,blank+=129;
while(true){
std::string::size_typepos(0);
if((pos=ret.find(blank))!=std::string::npos)
ret.replace(pos,blank.length(),"");
elsebreak;
}
//檢查是否有"<|blank_數(shù)字>"格式的特殊token，如果有，將其解碼成對(duì)應(yīng)數(shù)量的空格字符。
intpos=ret.find("<|blank_");
????????if?(pos?!=?-1)?{
????????????int?space_num?=?atoi(ret.substr(8,?ret.size()?-?10).c_str());
????????????return?std::string(space_num,?'?');
????????}

????????return?ret;
????}

????std::string?Tokenizer::Decode(const?Data?&data)?{
????????std::vector?tokens;
for(inti=0;i

	

	上面的：

	
if(pos!=-1){
intspace_num=atoi(ret.substr(8,ret.size()-10).c_str());
returnstd::string(space_num,'');
}


	

	這行代碼應(yīng)該是有bug，假設(shè) ret 的值為 "Hello<|blank_4>world!"，那么在解碼時(shí)，pos 將是 8，而 space_num 將是 4。然后，函數(shù)將返回 " "，即包含四個(gè)空格字符的字符串。在這種情況下，特殊 token "<|blank_4>" 被成功解碼成了四個(gè)空格字符，但是Hello和world!這部分被刪掉了。所以最終的解碼結(jié)果是不對(duì)的，需要修正一下。

	對(duì)tokenizer的解析可以發(fā)現(xiàn)，在c++中使用字典樹數(shù)據(jù)結(jié)構(gòu)來(lái)實(shí)現(xiàn)tokenizer是相對(duì)比較簡(jiǎn)單方便的。

	接下來(lái)，我們對(duì)CPU后端和GPU后端的算子實(shí)現(xiàn)進(jìn)行解析。

	0x3. CPU后端算子實(shí)現(xiàn)

	主要就是對(duì)這個(gè)文件進(jìn)行解析：https://github.com/ztxz16/fastllm/blob/master/src/devices/cpu/cpudevice.cpp 。

	輔助函數(shù)

	
//這是CpuDevice類的成員函數(shù)Malloc的定義，用于在CPU上分配一塊內(nèi)存空間。
boolCpuDevice::Malloc(void**ret,size_tsize){
*ret=(void*)newuint8_t[size];
returntrue;
}

//這是CpuDevice類的成員函數(shù)Free的定義，用于在CPU上釋放之前分配的內(nèi)存。
boolCpuDevice::Free(void*ret){
delete[](uint8_t*)ret;
returntrue;
}

//這是CpuDevice類的成員函數(shù)CopyDataFromCPU的定義，用于將數(shù)據(jù)從CPU拷貝到指定的設(shè)備上。
//這里什么都不做，直接返回true。
boolCpuDevice::CopyDataFromCPU(void*dst,void*src,size_tsize){
returntrue;
}

//這是CpuDevice類的成員函數(shù)CopyDataToCPU的定義，用于將數(shù)據(jù)從指定的設(shè)備拷貝到CPU上。
boolCpuDevice::CopyDataToCPU(void*dst,void*src,size_tsize){
returntrue;
}

//如果定義了__AVX__和__AVX2__，那么會(huì)啟用第一個(gè)DotU8U8函數(shù)和DotU4U8函數(shù)。
//如果只定義了__AVX__，但沒有定義__AVX2__，那么會(huì)啟用第二個(gè)DotU8U8函數(shù)和DotU4U8函數(shù)。

#ifdef__AVX__
#ifdef__AVX2__
//這是一段使用了IntelAVX2指令集（AdvancedVectorExtensions2）的代碼，
//用于計(jì)算兩個(gè)8位無(wú)符號(hào)整數(shù)數(shù)組的點(diǎn)積。
//定義了一個(gè)函數(shù)DotU8U8，它接受兩個(gè)指向8位無(wú)符號(hào)整數(shù)的指針a和b，
//以及一個(gè)整數(shù)n。這個(gè)函數(shù)的目的是計(jì)算數(shù)組a和b的點(diǎn)積，其中數(shù)組的長(zhǎng)度為n。
intDotU8U8(uint8_t*a,uint8_t*b,intn){
//初始化一個(gè)256位的整數(shù)向量acc，所有位都設(shè)置為零。這個(gè)向量用于存儲(chǔ)點(diǎn)積的累加值。
__m256iacc=_mm256_setzero_si256();
//初始化兩個(gè)變量，i用于循環(huán)計(jì)數(shù)，ans用于存儲(chǔ)最后的結(jié)果。
inti=0;
intans=0;
//等這幾行代碼初始化了一些常量向量
const__m256ilowMask=_mm256_set1_epi8(0xf);
const__m256iones=_mm256_set1_epi16(1);
const__m256iones8=_mm256_set1_epi8(1);
const__m256ixors=_mm256_set1_epi8(-128);
//這是一個(gè)循環(huán)，每次處理32個(gè)元素。這是因?yàn)锳VX2可以同時(shí)處理32個(gè)8位整數(shù)。
for(;i+31

	

	在啟用AVX2進(jìn)行點(diǎn)積計(jì)算時(shí)，有一個(gè)特殊的操作就是把b[i]轉(zhuǎn)換為有符號(hào)的整數(shù)并減掉128。我沒太懂這個(gè)操作的意義是什么，問了一下gpt4獲得了如下的回答：

	


	然后這里有個(gè)疑問是在DotU4U8的實(shí)現(xiàn)中調(diào)用的指令應(yīng)該是AVX2的指令集，但確是在AVX2宏關(guān)閉時(shí)調(diào)用的，不清楚這里是否會(huì)有bug。



	上述函數(shù)中涉及到大量的intel Intrinsics指令細(xì)節(jié)。

	

	CpuEmbedding 算子解析

	
//CpuEmbedding算子的形狀推導(dǎo)函數(shù)，這個(gè)函數(shù)接受四個(gè)參數(shù)：
//一個(gè)std::string類型的opType，兩個(gè)字典類型的datas和floatParams，以及一個(gè)intParams。
voidCpuEmbedding::Reshape(conststd::string&opType,constfastllm::DataDict&datas,
constfastllm::FloatDict&floatParams,constfastllm::IntDict&intParams){
//這三行代碼從datas字典中查找鍵為"input"、"output"和"weight"的元素，
//并將找到的元素的值賦給input、output和weight。
//這里的"input"、"output"和"weight"可以理解為嵌入層的輸入、輸出和權(quán)重。
Data&input=*(datas.find("input")->second);
Data&output=*(datas.find("output")->second);
Data&weight=*(datas.find("weight")->second);

//這行代碼檢查weight的維度數(shù)量是否為2。如果不是，就會(huì)拋出一個(gè)錯(cuò)誤。
AssertInFastLLM(weight.dims.size()==2,"Embedding'sweight'sdimshouldbe2.
");
//這行代碼檢查weight的數(shù)據(jù)類型是否為FLOAT32或BFLOAT16。如果不是，就會(huì)拋出一個(gè)錯(cuò)誤。
AssertInFastLLM(weight.dataType==DataType::FLOAT32||
weight.dataType==DataType::BFLOAT16,"Embedding'sweight'stypeshouldbefloat32orbfloat16.
");
//這行代碼檢查input的數(shù)據(jù)類型是否為FLOAT32。如果不是，就會(huì)拋出一個(gè)錯(cuò)誤。
AssertInFastLLM(input.dataType==DataType::FLOAT32,"Embedding'sinput'stypeshouldbefloat32.
");

//這行代碼將weight的weightType屬性設(shè)置為EMBEDDING。
weight.weightType=WeightType::EMBEDDING;
//這行代碼從weight的維度中提取詞匯大?。╲ocabSize）和嵌入大?。╡mbSize）。
intvocabSize=weight.dims[0],embSize=weight.dims[1];
//這兩行代碼將embSize添加到input的維度中，形成一個(gè)新的維度。
std::vectordims=input.dims;
dims.push_back(embSize);

//這兩行代碼將output的數(shù)據(jù)類型設(shè)置為FLOAT32，并重新調(diào)整其維度。
output.dataType=DataType::FLOAT32;
output.Resize(dims);
}

//這是一個(gè)名為CpuEmbedding::Run的函數(shù)，它在某個(gè)名為CpuEmbedding的類中被定義。
//這個(gè)函數(shù)接受四個(gè)參數(shù)：一個(gè)std::string類型的opType，
//兩個(gè)字典類型的datas和floatParams，以及一個(gè)intParams。
//這個(gè)函數(shù)的主要任務(wù)是執(zhí)行嵌入層（Embeddinglayer）的運(yùn)算。
//嵌入層通常用于將離散型特征（例如詞匯）轉(zhuǎn)換為連續(xù)的向量表示。
//具體的實(shí)現(xiàn)方法是，對(duì)于每個(gè)輸入的索引，從權(quán)重矩陣中查找對(duì)應(yīng)的行，
//然后將其復(fù)制到輸出矩陣的對(duì)應(yīng)位置。
voidCpuEmbedding::Run(conststd::string&opType,constfastllm::DataDict&datas,
constfastllm::FloatDict&floatParams,constfastllm::IntDict&intParams){
//這三行代碼從datas字典中查找鍵為"input"、"output"和"weight"的元素，
//并將找到的元素的值賦給input、output和weight。
//這里的"input"、"output"和"weight"可以理解為嵌入層的輸入、輸出和權(quán)重。
Data&input=*(datas.find("input")->second);
Data&output=*(datas.find("output")->second);
Data&weight=*(datas.find("weight")->second);;

output.Allocate();//這行代碼為output分配內(nèi)存。

//這行代碼從weight的維度中提取詞匯大小（vocabSize）和嵌入大?。╡mbSize）。
intvocabSize=weight.dims[0],embSize=weight.dims[1];
//這行代碼計(jì)算input的長(zhǎng)度。
uint64_tinputLen=input.Count(0);
//這行代碼獲取input的數(shù)據(jù)，并將其轉(zhuǎn)換為浮點(diǎn)數(shù)的指針。
float*inputData=(float*)input.cpuData;

//接下來(lái)的代碼根據(jù)內(nèi)存模式和權(quán)重的數(shù)據(jù)類型的不同，分別處理了四種情況。
//這四種情況可以歸納為兩個(gè)大類：內(nèi)存模式和權(quán)重的數(shù)據(jù)類型。
//內(nèi)存模式：如果GetLowMemMode()返回true，則表示處于低內(nèi)存模式。
//在這種模式下，權(quán)重?cái)?shù)據(jù)不會(huì)一次性全部加載到內(nèi)存中，而是每次只加載需要的部分。
//否則，權(quán)重?cái)?shù)據(jù)會(huì)全部加載到內(nèi)存中。
if(GetLowMemMode()){
FILE*fi=fopen(weight.fileName.c_str(),"rb");
//權(quán)重的數(shù)據(jù)類型：如果權(quán)重的數(shù)據(jù)類型為FLOAT32，則使用浮點(diǎn)數(shù)進(jìn)行計(jì)算。
//如果權(quán)重的數(shù)據(jù)類型為BFLOAT16，則使用16位浮點(diǎn)數(shù)進(jìn)行計(jì)算。
if(weight.dataType==DataType::FLOAT32){
float*outputData=(float*)output.cpuData;
for(inti=0;i

	

	CpuLayerNormOp 解析

	
voidCpuLayerNormOp::Run(conststd::string&opType,constfastllm::DataDict&datas,
constfastllm::FloatDict&floatParams,constfastllm::IntDict&intParams){
//這四行代碼從datas字典中查找鍵為"input"、"output"、"gamma"和"beta"的元素，
//并將找到的元素的值賦給input、output、gamma和beta。
//這里的"input"是層歸一化的輸入，"output"是輸出，
//"gamma"和"beta"是用于對(duì)歸一化后的結(jié)果進(jìn)行縮放和移位的可學(xué)習(xí)參數(shù)。
Data&input=*(datas.find("input")->second);
Data&output=*(datas.find("output")->second);
Data&gamma=*(datas.find("gamma")->second);
Data&beta=*(datas.find("beta")->second);

//這行代碼為output分配內(nèi)存。
output.Allocate();

//這行代碼從intParams字典中查找鍵為"axis"的元素。
//如果找到，則使用找到的值作為歸一化的軸；否則，使用默認(rèn)值-1。在層歸一化中，軸通常是特征維度。
intaxis=intParams.find("axis")!=intParams.end()?intParams.find("axis")->second:-1;
//這兩行代碼計(jì)算input的維度數(shù)，并將axis轉(zhuǎn)換為非負(fù)數(shù)。
//這是為了處理負(fù)數(shù)的軸值，因?yàn)樵?a href="http://hljzzgx.com/tags/python/" target="_blank">Python中，軸可以是負(fù)數(shù)，表示從后向前數(shù)的位置。
intdimsLen=input.dims.size();
axis=(axis%dimsLen+dimsLen)%dimsLen;

//這三行代碼計(jì)算outer、channels和inner。
//outer是歸一化操作的外部維度的元素總數(shù)，channels是歸一化操作的軸的大小，
//inner是歸一化操作的內(nèi)部維度的元素總數(shù)。
intouter=input.Count(0)/input.Count(axis);
intchannels=input.dims[axis];
intinner=input.strides[axis];

//這行代碼為mean和var分配內(nèi)存，它們用于存儲(chǔ)每個(gè)歸一化組的均值和方差。
float*mean=newfloat[inner],*var=newfloat[inner];
float*inputData=(float*)input.cpuData;
float*outputData=(float*)output.cpuData;
float*gammaData=(float*)gamma.cpuData;
float*betaData=(float*)beta.cpuData;

//在這個(gè)條件下，每個(gè)通道只有一個(gè)元素，所以可以并行地對(duì)每個(gè)通道進(jìn)行層歸一化。
if(inner==1){
//這是一個(gè)循環(huán)，對(duì)input中的每一個(gè)外部元素進(jìn)行處理。
for(inti=0;i

	

	CPULinearOp 解析

	最后簡(jiǎn)單讀一下CPULinearOp這個(gè)算子。

	
voidCpuLinearOp::Run(conststd::string&opType,constfastllm::DataDict&datas,
constfastllm::FloatDict&floatParams,constfastllm::IntDict&intParams){
//autost=std::now();
Data&input=*(datas.find("input")->second);
Data&output=*(datas.find("output")->second);
Data&weight=*(datas.find("weight")->second);
Data&bias=*(datas.find("bias")->second);

output.Allocate(0.0f);
intn=input.Count(0)/input.dims.back();
intm=input.dims.back();
intk=output.dims.back();

//這段代碼處理權(quán)重?cái)?shù)據(jù)類型為FLOAT32的情況。首先，它將輸入、權(quán)重、輸出和
//偏置數(shù)據(jù)的指針分別轉(zhuǎn)換為float*類型的指針。對(duì)于偏置數(shù)據(jù)，如果其維度長(zhǎng)度大于0，
//則獲取其數(shù)據(jù)指針，否則設(shè)為nullptr。
if(weight.dataType==DataType::FLOAT32){
float*inputData=(float*)input.cpuData;
float*weightData=(float*)weight.cpuData;
float*outputData=(float*)output.cpuData;
float*biasData=bias.dims.size()>0?(float*)bias.cpuData:nullptr;

//接下來(lái)，計(jì)算需要的線程數(shù)（threadNum）。這里用的是用戶設(shè)定的線程數(shù)
//（通過(guò)GetThreads()獲得）。然后，每個(gè)線程負(fù)責(zé)的任務(wù)數(shù)（per）
//為k（輸出數(shù)據(jù)的最后一個(gè)維度）除以線程數(shù)。cur用來(lái)表示當(dāng)前任務(wù)的起始位置。
intthreadNum=GetThreads();
intper=k/threadNum;
intcur=0;
//接著，創(chuàng)建線程池（通過(guò)GetPool()獲?。┖陀糜诒４婢€程任務(wù)的std::future數(shù)組。
//對(duì)于每個(gè)線程，確定其需要處理的任務(wù)范圍（從cur到end），然后提交線程任務(wù)。
//線程任務(wù)是通過(guò)調(diào)用FloatLinearPart函數(shù)來(lái)執(zhí)行的，該函數(shù)需要輸入數(shù)據(jù)、
//權(quán)重?cái)?shù)據(jù)、偏置數(shù)據(jù)、輸出數(shù)據(jù)、輸入維度（n）、權(quán)重維度（m）、輸出維度（k）
//以及任務(wù)范圍（從cur到end）作為參數(shù)。
autopool=GetPool();
std::vector>futures;
for(inti=0;iSubmit(FloatLinearPart,inputData,weightData,biasData,outputData,
n,m,k,cur,end));
cur=end;
}

//然后，主線程也執(zhí)行一部分任務(wù)，處理范圍為從cur到k。
FloatLinearPart(inputData,weightData,biasData,outputData,n,m,k,cur,k);
//最后，主線程等待所有子線程完成工作。通過(guò)調(diào)用std::get()
//方法來(lái)阻塞主線程，直到對(duì)應(yīng)的子線程完成任務(wù)。
//這樣，可以保證所有的線程任務(wù)都完成后，主線程才繼續(xù)執(zhí)行。
for(inti=0;i0?(float*)bias.cpuData:nullptr;
#ifdef__ARM_FEATURE_FP16_VECTOR_ARITHMETIC
uint16_t*temp=newuint16_t[n*m];
for(inti=0;i>futures;
for(inti=0;iSubmit(Float16LinearPart,inputData,weightData,biasData,outputData,
n,m,k,cur,end));
cur=end;
}

Float16LinearPart(inputData,weightData,biasData,outputData,n,m,k,cur,k);
for(inti=0;i0?(float*)bias.cpuData:nullptr;
weight.CalcWeightSum();

//之后，代碼創(chuàng)建一個(gè)std::vector對(duì)象，
//LowBitConfig是一個(gè)用于存儲(chǔ)數(shù)據(jù)量化信息的類，包括最小值、最大值、位寬和零點(diǎn)。
//這些信息是通過(guò)遍歷輸入數(shù)據(jù)獲得的。
std::vectorinputConfigs;
for(inti=0;i對(duì)象uinput，并將其大小設(shè)置為輸入數(shù)據(jù)的大?。╪*m）。
//uinput中的每個(gè)元素都是輸入數(shù)據(jù)元素經(jīng)過(guò)inputConfigs中對(duì)應(yīng)配置信息量化后的結(jié)果。
//注意這里的量化過(guò)程可能會(huì)根據(jù)是否定義了__AVX2__進(jìn)行不同的處理。
std::vectoruinput;
uinput.resize(n*m);
for(inti=0;i0?(float*)bias.cpuData:nullptr;
weight.CalcWeightSum();

std::vectorinputConfigs;
for(inti=0;iuinput;
uinput.resize(n*m);
for(inti=0;i

	

	在上面的實(shí)現(xiàn)中，MultiplyMultiThread完成了對(duì)量化輸入的計(jì)算，我們看一下它的實(shí)現(xiàn)細(xì)節(jié)：

	
//a=[n,m],b=[k,m],c=aT(b')=[n,k]
voidMultiplyMultiThread(uint8_t*a,uint8_t*b,int32_t*c,intn,intm,intk,intthreadNum){
intper=k/threadNum;
intcur=0;
if(threadNum==1){
Multiply(a,b+cur*m,c+cur,n,m,k-cur,k);
}else{
autopool=GetPool();
std::vector>futures;
for(inti=0;iSubmit(Multiply,a,b+cur*m,c+cur,n,m,end-cur,k));
cur=end;
}
for(inti=0;i

	

	可以看到這段代碼仍然是在用線程池來(lái)啟動(dòng)多個(gè)線程完成計(jì)算，核心部分是Multiply函數(shù)，這個(gè)函數(shù)的實(shí)現(xiàn)細(xì)節(jié)：

	
//a=[n,m],b=[k,m],c=aT(b')=[n,k]
voidMultiply(uint8_t*a,uint8_t*b,int32_t*c,intn,intm,intk,intkstride){
#ifdef__ARM_FEATURE_DOTPROD
intblock=0;
for(;block

	

	這段代碼實(shí)現(xiàn)了兩個(gè)矩陣的乘法。輸入的兩個(gè)矩陣是 (a) 和 (b)，結(jié)果矩陣是 (c)。矩陣 (a) 的形狀是 ([n, m])，矩陣 (b) 的形狀是 ([k, m])，所以矩陣 (c = a^T b) 的形狀是 ([n, k])。

	在這段代碼中，使用了不同的方法進(jìn)行矩陣乘法，取決于系統(tǒng)是否支持特定的優(yōu)化硬件指令。

	如果系統(tǒng)支持 ARMv8.2 的點(diǎn)積指令（__ARM_FEATURE_DOTPROD），那么會(huì)使用這個(gè)指令進(jìn)行矩陣乘法。在這種情況下，每次會(huì)同時(shí)處理32個(gè)元素，這樣可以加速計(jì)算。

	如果系統(tǒng)支持 ARMv8（__aarch64__），但不支持 ARMv8.2 的點(diǎn)積指令，那么會(huì)使用 NEON SIMD 指令進(jìn)行矩陣乘法。在這種情況下，每次會(huì)同時(shí)處理64個(gè)元素。

	如果系統(tǒng)支持 AVX（__AVX__），那么會(huì)使用 AVX 指令進(jìn)行矩陣乘法。在這種情況下，會(huì)使用 DotU8U8 函數(shù)來(lái)計(jì)算向量的點(diǎn)積。

	如果系統(tǒng)不支持上述任何一種優(yōu)化指令，那么會(huì)使用基礎(chǔ)的方法進(jìn)行矩陣乘法。在這種情況下，每次只處理一個(gè)元素。

	這段代碼的優(yōu)化部分主要利用了 SIMD（單指令多數(shù)據(jù)）的并行化特性，通過(guò)同時(shí)處理多個(gè)元素來(lái)加速計(jì)算。而選擇使用哪種優(yōu)化方法，取決于系統(tǒng)支持哪種硬件指令。

	CPU后端的算子解析就暫時(shí)講到這里，我們發(fā)現(xiàn)CPU的算子實(shí)現(xiàn)不僅考慮了Intel CPU也考慮了Arm端的優(yōu)化，這也是FastLLM可以在Arm邊緣端部署大模型的原因。

	0x4. GPU后端算子實(shí)現(xiàn)

	GPU后端算子實(shí)現(xiàn)在 https://github.com/ztxz16/fastllm/blob/master/src/devices/cuda/cudadevice.cpp 和 https://github.com/ztxz16/fastllm/blob/master/src/devices/cuda/fastllm-cuda.cu 。我們還是挑幾個(gè)算子來(lái)講解。

	CudaLlamaRotatePosition2DOp

	LLama的ROPE實(shí)現(xiàn)在：https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L92-L126 。

	
#這個(gè)類是用來(lái)創(chuàng)建旋轉(zhuǎn)位置編碼（RotaryPositionEmbedding）的。
#Llama模型引入了旋轉(zhuǎn)位置編碼，以改進(jìn)長(zhǎng)序列處理的性能。
classLlamaRotaryEmbedding(torch.nn.Module):
#這是類的初始化方法，接收四個(gè)參數(shù)：dim（嵌入的維度），max_position_embeddings
#（最大的位置嵌入長(zhǎng)度，默認(rèn)為2048），base（基數(shù)，默認(rèn)為10000）和device（設(shè)備類型，例如CPU或GPU）。
def__init__(self,dim,max_position_embeddings=2048,base=10000,device=None):
super().__init__()
self.dim=dim#將輸入的dim參數(shù)保存到self.dim屬性中。
##將輸入的max_position_embeddings參數(shù)保存到self.max_position_embeddings屬性中。
self.max_position_embeddings=max_position_embeddings
#將輸入的base參數(shù)保存到self.base屬性中。
self.base=base
#計(jì)算逆頻率并保存到變量inv_freq中。逆頻率是一種用于位置編碼的技巧，
#它可以幫助模型更好地捕捉位置信息。
inv_freq=1.0/(self.base**(torch.arange(0,self.dim,2).float().to(device)/self.dim))
#將inv_freq保存到模型的緩存中。register_buffer是PyTorchnn.Module的一個(gè)方法，
#它用于保存一些不需要計(jì)算梯度的變量。
self.register_buffer("inv_freq",inv_freq,persistent=False)

#Buildheretomake`torch.jit.trace`work.
#調(diào)用_set_cos_sin_cache方法，預(yù)先計(jì)算并保存正弦和余弦的緩存值。
self._set_cos_sin_cache(
seq_len=max_position_embeddings,device=self.inv_freq.device,dtype=torch.get_default_dtype()
)

#這是一個(gè)私有方法，接收三個(gè)參數(shù)：seq_len（序列長(zhǎng)度），device（設(shè)備類型）和dtype（數(shù)據(jù)類型）
def_set_cos_sin_cache(self,seq_len,device,dtype):
#將輸入的seq_len參數(shù)保存到self.max_seq_len_cached屬性中。
self.max_seq_len_cached=seq_len
#生成一個(gè)長(zhǎng)度為max_seq_len_cached的序列，并保存到變量t中。
t=torch.arange(self.max_seq_len_cached,device=device,dtype=self.inv_freq.dtype)

#使用外積計(jì)算頻率和t的乘積，結(jié)果保存到變量freqs中。
freqs=torch.einsum("i,j->ij",t,self.inv_freq)
#Differentfrompaper,butitusesadifferentpermutationinordertoobtainthesamecalculation
#將頻率的兩份副本拼接在一起，結(jié)果保存到變量emb中。
emb=torch.cat((freqs,freqs),dim=-1)
#計(jì)算emb的余弦值，然后將結(jié)果保存到模型的緩存中。
self.register_buffer("cos_cached",emb.cos()[None,None,:,:].to(dtype),persistent=False)
#計(jì)算emb的正弦值，然后將結(jié)果保存到模型的緩存中。
self.register_buffer("sin_cached",emb.sin()[None,None,:,:].to(dtype),persistent=False)

#這是模型的前向傳播方法，接收兩個(gè)參數(shù)：x（輸入數(shù)據(jù)）和seq_len（序列長(zhǎng)度）。
defforward(self,x,seq_len=None):
#x:[bs,num_attention_heads,seq_len,head_size]
#如果輸入的序列長(zhǎng)度大于緩存的最大序列長(zhǎng)度，那么調(diào)用_set_cos_sin_cache方法，更新緩存。
ifseq_len>self.max_seq_len_cached:
self._set_cos_sin_cache(seq_len=seq_len,device=x.device,dtype=x.dtype)

#返回對(duì)應(yīng)輸入位置的正弦和余弦值。這些值將用于旋轉(zhuǎn)位置編碼。
return(
self.cos_cached[:,:,:seq_len,...].to(dtype=x.dtype),
self.sin_cached[:,:,:seq_len,...].to(dtype=x.dtype),
)

defapply_rotary_pos_emb(q,k,cos,sin,position_ids):
#Thefirsttwodimensionsofcosandsinarealways1,sowecan`squeeze`them.
cos=cos.squeeze(1).squeeze(0)#[seq_len,dim]
sin=sin.squeeze(1).squeeze(0)#[seq_len,dim]
cos=cos[position_ids].unsqueeze(1)#[bs,1,seq_len,dim]
sin=sin[position_ids].unsqueeze(1)#[bs,1,seq_len,dim]
q_embed=(q*cos)+(rotate_half(q)*sin)
k_embed=(k*cos)+(rotate_half(k)*sin)
returnq_embed,k_embed


	

	CudaLlamaRotatePosition2DOp對(duì)應(yīng)的就是上面的Python代碼。

	
voidCudaLlamaRotatePosition2DOp::Run(conststd::string&opType,constfastllm::DataDict&datas,
constfastllm::FloatDict&floatParams,constfastllm::IntDict&intParams){
Data&data=*(datas.find("input")->second);
Data&positionIds=*(datas.find("positionIds")->second);
Data&sinData=*(datas.find("sin")->second);
Data&cosData=*(datas.find("cos")->second);
introtaryDim=intParams.find("rotaryDim")!=intParams.end()?intParams.find("rotaryDim")->second:128;

FastllmCudaLlamaRotatePosition2D(data,positionIds,sinData,cosData,rotaryDim);
}


	

	這里調(diào)用的是FastllmCudaLlamaRotatePosition2D這個(gè)函數(shù)，它的實(shí)現(xiàn)和解析如下：

	
//這是一個(gè)在GPU上運(yùn)行的CUDA函數(shù)，用于執(zhí)行Llama模型的位置編碼旋轉(zhuǎn)操作。
//data：輸入的數(shù)據(jù)，這個(gè)數(shù)據(jù)將會(huì)被旋轉(zhuǎn)。
//positionIds：位置編碼的數(shù)據(jù)。
//sinData，cosData：用于旋轉(zhuǎn)的sin和cos值。
//rotaryDim：旋轉(zhuǎn)的維度。
boolFastllmCudaLlamaRotatePosition2D(fastllm::Data&data,constfastllm::Data&positionIds,
constfastllm::Data&sinData,constfastllm::Data&cosData,introtaryDim){
//使用FastllmCudaPrepareInput函數(shù)將輸入的數(shù)據(jù)從CPU復(fù)制到GPU。
//這個(gè)函數(shù)會(huì)返回一個(gè)指向GPU內(nèi)存的指針。
float*cudaData=(float*)FastllmCudaPrepareInput(data);
float*cudaPositionIds=(float*)FastllmCudaPrepareInput(positionIds);
float*cudaSin=(float*)FastllmCudaPrepareInput(sinData);
float*cudaCos=(float*)FastllmCudaPrepareInput(cosData);

//計(jì)算旋轉(zhuǎn)操作需要的一些參數(shù)，包括outer，spatial，bs，len，n和m。
//這些參數(shù)是用于確定CUDA核函數(shù)的執(zhí)行配置和一些數(shù)據(jù)操作的。
intouter=data.dims[0]*data.dims[1];
intspatial=data.Count(2);
intbs=data.dims[0],len=data.dims[1];
intn=data.dims[2],m=data.dims[3];
//調(diào)用CUDA核函數(shù)FastllmLlamaRotatePosition2DKernel來(lái)在GPU上執(zhí)行位置編碼的旋轉(zhuǎn)操作。
//<<>>是CUDA中定義并行線程塊和線程的語(yǔ)法，
//outer*n是線程塊的數(shù)量，min(rotaryDim,m/2)是每個(gè)線程塊中的線程數(shù)量。
//核函數(shù)的參數(shù)包括之前準(zhǔn)備的數(shù)據(jù)和一些計(jì)算參數(shù)。
FastllmLlamaRotatePosition2DKernel<<>>(cudaData,cudaPositionIds,cudaSin,cudaCos,
len,bs,spatial,n,m,
(int)positionIds.dims.back(),(int)sinData.dims[1],rotaryDim);

//使用FastllmCudaFinishInput函數(shù)釋放positionIds，sinData和cosData在GPU上的內(nèi)存。
//這些數(shù)據(jù)在這個(gè)函數(shù)中不再需要。
FastllmCudaFinishInput(positionIds,cudaPositionIds);
FastllmCudaFinishInput(sinData,cudaSin);
FastllmCudaFinishInput(cosData,cudaCos);
//使用FastllmCudaFinishOutput函數(shù)將旋轉(zhuǎn)后的數(shù)據(jù)從GPU復(fù)制回CPU。
//這個(gè)函數(shù)也會(huì)釋放data在GPU上的內(nèi)存。
FastllmCudaFinishOutput(data,cudaData);
returntrue;
}


	

	最后再解析下這個(gè)cuda kernel。

	
//float*data：輸入數(shù)據(jù)，大小為[bs,len,n,m]，其中bs是批量大小，
//len是序列長(zhǎng)度，n是頭的數(shù)量，m是每個(gè)頭的維度。
//float*positionIds：位置編碼的索引，大小為[bs,len]。
//float*sin和float*cos：預(yù)先計(jì)算的正弦和余弦值，用于旋轉(zhuǎn)編碼。
//intlen,intbs,intspatial,intn,intm：輸入數(shù)據(jù)的各個(gè)維度大小。
//intpartStride和intsinCosStride：用于索引positionIds和sin/cos的步長(zhǎng)。
//introtateDim：旋轉(zhuǎn)維度。
__global__voidFastllmLlamaRotatePosition2DKernel(float*data,float*positionIds,float*sin,float*cos,
intlen,intbs,intspatial,intn,intm,intpartStride,intsinCosStride,introtateDim){
//首先，計(jì)算出當(dāng)前線程應(yīng)處理的位置o，長(zhǎng)度l和批次b。
into=(blockIdx.x/n);
intl=o%len;
intb=o/len;
intj=threadIdx.x;
//然后，根據(jù)positionIds獲取對(duì)應(yīng)的旋轉(zhuǎn)角度的正弦值curSin和余弦值curCos。
intindex=(int)(positionIds[b*partStride+l]);

floatcurSin=sin[index*sinCosStride+j];
floatcurCos=cos[index*sinCosStride+j];
float*d=(float*)data+o*spatial+j;
inti=blockIdx.x%n;
//接著，獲取輸入數(shù)據(jù)對(duì)應(yīng)位置的值va和vb。
floatva=d[i*m],vb=d[i*m+m/2];
//最后，根據(jù)旋轉(zhuǎn)矩陣的公式，計(jì)算旋轉(zhuǎn)后的值，并將結(jié)果寫回輸入數(shù)據(jù)中。
d[i*m]=va*curCos-vb*curSin;
d[i*m+m/2]=va*curSin+vb*curCos;
}



	

	直接看這個(gè)cuda kernel可能比較難理解，可以結(jié)合https://github.com/ztxz16/fastllm/blob/master/src/devices/cpu/cpudevice.cpp#L2204-L2233 這里的cpu實(shí)現(xiàn)來(lái)看，這樣來(lái)看設(shè)置batch * seq_length * n個(gè)block，每個(gè)block處理m個(gè)元素就是比較合理直觀的。

	
voidCpuLlamaRotatePosition2DOp::Run(conststd::string&opType,constfastllm::DataDict&datas,
constfastllm::FloatDict&floatParams,constfastllm::IntDict&intParams){
Data&data=*(datas.find("input")->second);
Data&positionIds=*(datas.find("positionIds")->second);
Data&sinData=*(datas.find("sin")->second);
Data&cosData=*(datas.find("cos")->second);
introtaryDim=intParams.find("rotaryDim")!=intParams.end()?intParams.find("rotaryDim")->second:128;

intbs=data.dims[0],len=data.dims[1];
intspatial=data.Count(2);
intn=data.dims[2],m=data.dims[3];
intstride=(int)sinData.dims[1];
for(intb=0;b

	

	FastLLM在cuda上的實(shí)現(xiàn)不算高校，不過(guò)優(yōu)點(diǎn)在于它支持了完整的int8和int4量化的計(jì)算，有興趣的讀者可以自行研究這部分kernel實(shí)現(xiàn)。

	0x5. LLMSamping解析

	在 chatglm-6b 的實(shí)現(xiàn)中，在前向推理完成后以及tokenizer解碼之前有一個(gè)根據(jù)logits取label的過(guò)程：https://github.com/ztxz16/fastllm/blob/master/src/models/chatglm.cpp#L267-L279 。

	
if(generationConfig.IsSimpleGreedy()){
//對(duì)logits進(jìn)行TopK操作，將結(jié)果存儲(chǔ)在topk中。
//這里的TopK操作是找到logits中最大的K個(gè)值，這里K=1，所以是找到最大值。
TopK(logits,topk,1);
topk.ToDevice(DataDevice::CPU);
for(intb=0;b

	

	LLMSampling是一種常見的在序列生成任務(wù)中，根據(jù)不同的需求，使用不同的策略生成序列的方法。我們這里來(lái)研究一下它的實(shí)現(xiàn)。它的實(shí)現(xiàn)在：https://github.com/ztxz16/fastllm/blob/master/src/fastllm.cpp#L874-L916 。

	
//這段代碼是一個(gè)用于從給定的logits（通常表示預(yù)測(cè)的概率分布）進(jìn)行采樣的函數(shù)，
//采樣策略主要受GenerationConfig和LastTokensUnit參數(shù)的影響。
intLLMSampling(Data&logits,intouterOffset,
constGenerationConfig&config,constLastTokensUnit&tokens){
//將logits數(shù)據(jù)從當(dāng)前設(shè)備轉(zhuǎn)移到CPU。
logits.ToDevice(DataDevice::CPU);
//從logits的維度中獲取詞匯量vocabSize。
intvocabSize=logits.dims.back();
//計(jì)算base指針，指向要處理的logits的開始位置。
float*base=((float*)logits.cpuData)+outerOffset*vocabSize;

//判斷config.repeat_penalty是否不等于1，如果不等于1，
//則對(duì)tokens.tokenSet中每個(gè)id對(duì)應(yīng)的base[id]值進(jìn)行修改。
if(fabs(config.repeat_penalty-1.0)>1e-6){
for(intid:tokens.tokenSet){
base[id]=(base[id]。
std::vector>v;
//遍歷每個(gè)logit，將其值乘以invTemp，并存入v中。
for(inti=0;ifirst;
//定義一個(gè)向量ps，用于存儲(chǔ)處理后的概率。
std::vectorps;
//遍歷v中的前topk個(gè)元素，將其值取exp并減去maxValue，存入ps，同時(shí)更新psum。
for(inti=0;iconfig.top_p){
topk=i+1;
break;
}
}
//生成一個(gè)隨機(jī)數(shù)rnd。
floatrnd=fastllmRandom.randP();
curSum=0.0;
//遍歷ps中的前topk個(gè)元素，將其累加到curSum，
//當(dāng)curSum大于rnd或者達(dá)到最后一個(gè)元素時(shí)，
//返回對(duì)應(yīng)v[i].second，也就是返回采樣得到的id。
for(inti=0;irnd||i==topk-1){
returnv[i].second;
}
}
//如果以上步驟都沒有返回，那么返回-1。
return-1;
}


	

	LLMSampling實(shí)現(xiàn)了一種基于溫度和懲罰的采樣策略，用于從給定的 logits 中選擇一個(gè) id。這種采樣的方法可以控制輸出文本的多樣性。

	0x6. 總結(jié)

	接著 大模型部署框架 FastLLM 簡(jiǎn)要解析 這篇文章首先梳理了一下FastLLM的調(diào)用鏈和關(guān)鍵的數(shù)據(jù)結(jié)構(gòu)，然后解析了 FastLLM 的一些實(shí)現(xiàn)細(xì)節(jié)和CPU/GPU后端實(shí)現(xiàn)采用的優(yōu)化技巧。

	審核編輯：湯梓紅