大模型部署框架FastLLM實現細節解析

0x0. 前言

接著 大模型部署框架 FastLLM 簡要解析 這篇文章首先梳理了一下FastLLM的調用鏈和關鍵的數據結構，然后解析了 FastLLM 的一些實現細節和CPU/GPU后端實現采用的優化技巧。

0x1. 調用鏈和數據結構解析

以chatglm-6b的支持為例，函數入口在 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 這行代碼對 input 進行 tokenizer encode并構造好inputIds，再構造好attentionMask之后就可以給Forward函數推理，拿到推理結果之后再使用tokenizer進行decode得到輸出。

在這里，inputIds和attentionMask都是Data數據類型，類比于PyTorch的Tensor，來對輸入數據以及device，shape等信息進行統一管理。下面的代碼展示了Data數據結構的定義，源碼在：https://github.com/ztxz16/fastllm/blob/master/include/fastllm.h#L201-L286

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

DataTypedataType=DataType::FLOAT32;//數據類型
intunitSize,unitSizeDiv=1;//單個元素的字節數=unitSIze/unitSizeDiv

std::vectordims;//數據形狀
std::vectorstrides;//跨度

uint64_texpansionSize=0;//擴容后的尺寸
uint64_texpansionBytes=0;//擴容后的字節數
std::vectorexpansionDims;//預擴容的形狀
uint8_t*cpuData=nullptr;//數據指針

void*cudaData=nullptr;
std::vectorextraCudaData;

void*deviceData=nullptr;
std::vectorextraDeviceData;

DataDevicedataDevice=DataDevice::CPU;

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

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

Data(){};

Data(DataTypetype);

Data(DataTypetype,conststd::vector&dims);//構造函數

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

~Data();//析構函數

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

voidCopyFrom(constData&ori);//復制

uint64_tGetBytes()const;//獲取總字節數

voidAllocate();//分配內存

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

voidExpansion(conststd::vector&dims);//預擴容到相應尺寸

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

voidFreeSpace();//回收設備上的內存

voidUpdateUnitSize();//更新unitSize

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

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

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

voidPrintShape()const;//輸出形狀

voidPrint()const;//輸出

voidCalcWeightSum();//計算WeightSum

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

voidToDevice(void*device);

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

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

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負責根據FastLLM編譯開啟的后端選項把算子Dispatch到不同的device進行執行，{"input", (Data*)&input}, {"axis", &axisData}, {"output", (Data*)&output}} 這行代碼表示的是一個DataDict對象，也就是一個值為data的字典，原始定義為typedef std::map DataDict;。接著我們看一下curExecutor的定義和實現：

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

public:
Executor();//創建默認的Executor

~Executor();//析構

voidClearDevices();//清空devices

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

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

voidClearProfiler();

voidPrintProfiler();
};
}

從Executor類的定義我們可以判斷它負責了在設定的devices上根據opType和輸入數據等執行Op的前向計算，也就是Run這個接口。由于Executor類是FastLLM的調度核心實現，所以我們來詳細解析一下它的實現。

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

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

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

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

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

//打印profile信息，也即輸出每個層的運行時間和模型的總運行時間
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);
}
}

自此，前向計算就順利完成了，再把推理結果給 tokenizer 解碼就結束了，整體的調度執行流程是很簡單明了的。

0x2. tokenizer 解析

接著，我們來解析一下tokenizer的實現。先看一下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);//插入一個token

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

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

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

我們從實現來看tokenizer的細節：

//這是Tokenizer類的嵌套結構TrieNode的構造函數的實現。
//在構造函數中，將tokenId成員變量的值初始化為-999999。
//這個值在構造函數中被硬編碼，它是作為一個特殊標記來使用的。
Tokenizer::TrieNode(){
this->tokenId=-999999;
}

//Tokenizer類的構造函數的實現。
//在構造函數中，通過new運算符創建一個新的TrieNode對象，
//并將其指針賦值給root成員變量。這樣，構造函數創建了一個空的字典樹，
//并將其根節點指針存儲在root中。
Tokenizer::Tokenizer(){
root=newTrieNode();
}

//Tokenizer類的析構函數的實現。
//在析構函數中，首先調用Clear()函數，用于釋放動態分配的資源和清空數據。
//然后，調用delete運算符釋放通過new運算符創建的root對象的內存，從而釋放整個字典樹的內存。
Tokenizer::~Tokenizer(){
Clear();
deleteroot;
}

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

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

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

//這是Tokenizer類的成員函數DecodeTokens的定義，
//用于對輸入的token數組進行解碼，將token轉換回原始的字符串。
std::stringTokenizer::DecodeTokens(conststd::vector&tokens){
//創建一個空字符串ret，用于存儲解碼結果。
std::stringret="";
//開始遍歷輸入的token數組tokens。
for(inti=0;i"格式的token（其中HH表示十六進制數），
//則需要將其轉換為對應的字符。首先，提取HH，然后將其轉換為對應的字符，
//并用空格代替原始的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;
????????????}
????????????//?根據不同的?token?進行解碼：
????????????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_數字>"格式的特殊token，如果有，將其解碼成對應數量的空格字符。
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,'');
}


	

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

	對tokenizer的解析可以發現，在c++中使用字典樹數據結構來實現tokenizer是相對比較簡單方便的。

	接下來，我們對CPU后端和GPU后端的算子實現進行解析。

	0x3. CPU后端算子實現

	主要就是對這個文件進行解析：https://github.com/ztxz16/fastllm/blob/master/src/devices/cpu/cpudevice.cpp 。

	輔助函數

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

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

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

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

//如果定義了__AVX__和__AVX2__，那么會啟用第一個DotU8U8函數和DotU4U8函數。
//如果只定義了__AVX__，但沒有定義__AVX2__，那么會啟用第二個DotU8U8函數和DotU4U8函數。

#ifdef__AVX__
#ifdef__AVX2__
//這是一段使用了IntelAVX2指令集（AdvancedVectorExtensions2）的代碼，
//用于計算兩個8位無符號整數數組的點積。
//定義了一個函數DotU8U8，它接受兩個指向8位無符號整數的指針a和b，
//以及一個整數n。這個函數的目的是計算數組a和b的點積，其中數組的長度為n。
intDotU8U8(uint8_t*a,uint8_t*b,intn){
//初始化一個256位的整數向量acc，所有位都設置為零。這個向量用于存儲點積的累加值。
__m256iacc=_mm256_setzero_si256();
//初始化兩個變量，i用于循環計數，ans用于存儲最后的結果。
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);
//這是一個循環，每次處理32個元素。這是因為AVX2可以同時處理32個8位整數。
for(;i+31

	

	在啟用AVX2進行點積計算時，有一個特殊的操作就是把b[i]轉換為有符號的整數并減掉128。我沒太懂這個操作的意義是什么，問了一下gpt4獲得了如下的回答：

	


	然后這里有個疑問是在DotU4U8的實現中調用的指令應該是AVX2的指令集，但確是在AVX2宏關閉時調用的，不清楚這里是否會有bug。



	上述函數中涉及到大量的intel Intrinsics指令細節。

	

	CpuEmbedding 算子解析

	
//CpuEmbedding算子的形狀推導函數，這個函數接受四個參數：
//一個std::string類型的opType，兩個字典類型的datas和floatParams，以及一個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"可以理解為嵌入層的輸入、輸出和權重。
Data&input=*(datas.find("input")->second);
Data&output=*(datas.find("output")->second);
Data&weight=*(datas.find("weight")->second);

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

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

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

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

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

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

//接下來的代碼根據內存模式和權重的數據類型的不同，分別處理了四種情況。
//這四種情況可以歸納為兩個大類：內存模式和權重的數據類型。
//內存模式：如果GetLowMemMode()返回true，則表示處于低內存模式。
//在這種模式下，權重數據不會一次性全部加載到內存中，而是每次只加載需要的部分。
//否則，權重數據會全部加載到內存中。
if(GetLowMemMode()){
FILE*fi=fopen(weight.fileName.c_str(),"rb");
//權重的數據類型：如果權重的數據類型為FLOAT32，則使用浮點數進行計算。
//如果權重的數據類型為BFLOAT16，則使用16位浮點數進行計算。
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"是用于對歸一化后的結果進行縮放和移位的可學習參數。
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分配內存。
output.Allocate();

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

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

//這行代碼為mean和var分配內存，它們用于存儲每個歸一化組的均值和方差。
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;

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

	

	CPULinearOp 解析

	最后簡單讀一下CPULinearOp這個算子。

	
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();

//這段代碼處理權重數據類型為FLOAT32的情況。首先，它將輸入、權重、輸出和
//偏置數據的指針分別轉換為float*類型的指針。對于偏置數據，如果其維度長度大于0，
//則獲取其數據指針，否則設為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;

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

//然后，主線程也執行一部分任務，處理范圍為從cur到k。
FloatLinearPart(inputData,weightData,biasData,outputData,n,m,k,cur,k);
//最后，主線程等待所有子線程完成工作。通過調用std::get()
//方法來阻塞主線程，直到對應的子線程完成任務。
//這樣，可以保證所有的線程任務都完成后，主線程才繼續執行。
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();

//之后，代碼創建一個std::vector對象，
//LowBitConfig是一個用于存儲數據量化信息的類，包括最小值、最大值、位寬和零點。
//這些信息是通過遍歷輸入數據獲得的。
std::vectorinputConfigs;
for(inti=0;i對象uinput，并將其大小設置為輸入數據的大?。╪*m）。
//uinput中的每個元素都是輸入數據元素經過inputConfigs中對應配置信息量化后的結果。
//注意這里的量化過程可能會根據是否定義了__AVX2__進行不同的處理。
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

	

	在上面的實現中，MultiplyMultiThread完成了對量化輸入的計算，我們看一下它的實現細節：

	
//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

	

	可以看到這段代碼仍然是在用線程池來啟動多個線程完成計算，核心部分是Multiply函數，這個函數的實現細節：

	
//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

	

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

	在這段代碼中，使用了不同的方法進行矩陣乘法，取決于系統是否支持特定的優化硬件指令。

	如果系統支持 ARMv8.2 的點積指令（__ARM_FEATURE_DOTPROD），那么會使用這個指令進行矩陣乘法。在這種情況下，每次會同時處理32個元素，這樣可以加速計算。

	如果系統支持 ARMv8（__aarch64__），但不支持 ARMv8.2 的點積指令，那么會使用 NEON SIMD 指令進行矩陣乘法。在這種情況下，每次會同時處理64個元素。

	如果系統支持 AVX（__AVX__），那么會使用 AVX 指令進行矩陣乘法。在這種情況下，會使用 DotU8U8 函數來計算向量的點積。

	如果系統不支持上述任何一種優化指令，那么會使用基礎的方法進行矩陣乘法。在這種情況下，每次只處理一個元素。

	這段代碼的優化部分主要利用了 SIMD（單指令多數據）的并行化特性，通過同時處理多個元素來加速計算。而選擇使用哪種優化方法，取決于系統支持哪種硬件指令。

	CPU后端的算子解析就暫時講到這里，我們發現CPU的算子實現不僅考慮了Intel CPU也考慮了Arm端的優化，這也是FastLLM可以在Arm邊緣端部署大模型的原因。

	0x4. GPU后端算子實現

	GPU后端算子實現在 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 。我們還是挑幾個算子來講解。

	CudaLlamaRotatePosition2DOp

	LLama的ROPE實現在：https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L92-L126 。

	
#這個類是用來創建旋轉位置編碼（RotaryPositionEmbedding）的。
#Llama模型引入了旋轉位置編碼，以改進長序列處理的性能。
classLlamaRotaryEmbedding(torch.nn.Module):
#這是類的初始化方法，接收四個參數：dim（嵌入的維度），max_position_embeddings
#（最大的位置嵌入長度，默認為2048），base（基數，默認為10000）和device（設備類型，例如CPU或GPU）。
def__init__(self,dim,max_position_embeddings=2048,base=10000,device=None):
super().__init__()
self.dim=dim#將輸入的dim參數保存到self.dim屬性中。
##將輸入的max_position_embeddings參數保存到self.max_position_embeddings屬性中。
self.max_position_embeddings=max_position_embeddings
#將輸入的base參數保存到self.base屬性中。
self.base=base
#計算逆頻率并保存到變量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的一個方法，
#它用于保存一些不需要計算梯度的變量。
self.register_buffer("inv_freq",inv_freq,persistent=False)

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

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

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

#這是模型的前向傳播方法，接收兩個參數：x（輸入數據）和seq_len（序列長度）。
defforward(self,x,seq_len=None):
#x:[bs,num_attention_heads,seq_len,head_size]
#如果輸入的序列長度大于緩存的最大序列長度，那么調用_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)

#返回對應輸入位置的正弦和余弦值。這些值將用于旋轉位置編碼。
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對應的就是上面的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);
}


	

	這里調用的是FastllmCudaLlamaRotatePosition2D這個函數，它的實現和解析如下：

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

//計算旋轉操作需要的一些參數，包括outer，spatial，bs，len，n和m。
//這些參數是用于確定CUDA核函數的執行配置和一些數據操作的。
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];
//調用CUDA核函數FastllmLlamaRotatePosition2DKernel來在GPU上執行位置編碼的旋轉操作。
//<<>>是CUDA中定義并行線程塊和線程的語法，
//outer*n是線程塊的數量，min(rotaryDim,m/2)是每個線程塊中的線程數量。
//核函數的參數包括之前準備的數據和一些計算參數。
FastllmLlamaRotatePosition2DKernel<<>>(cudaData,cudaPositionIds,cudaSin,cudaCos,
len,bs,spatial,n,m,
(int)positionIds.dims.back(),(int)sinData.dims[1],rotaryDim);

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


	

	最后再解析下這個cuda kernel。

	
//float*data：輸入數據，大小為[bs,len,n,m]，其中bs是批量大小，
//len是序列長度，n是頭的數量，m是每個頭的維度。
//float*positionIds：位置編碼的索引，大小為[bs,len]。
//float*sin和float*cos：預先計算的正弦和余弦值，用于旋轉編碼。
//intlen,intbs,intspatial,intn,intm：輸入數據的各個維度大小。
//intpartStride和intsinCosStride：用于索引positionIds和sin/cos的步長。
//introtateDim：旋轉維度。
__global__voidFastllmLlamaRotatePosition2DKernel(float*data,float*positionIds,float*sin,float*cos,
intlen,intbs,intspatial,intn,intm,intpartStride,intsinCosStride,introtateDim){
//首先，計算出當前線程應處理的位置o，長度l和批次b。
into=(blockIdx.x/n);
intl=o%len;
intb=o/len;
intj=threadIdx.x;
//然后，根據positionIds獲取對應的旋轉角度的正弦值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;
//接著，獲取輸入數據對應位置的值va和vb。
floatva=d[i*m],vb=d[i*m+m/2];
//最后，根據旋轉矩陣的公式，計算旋轉后的值，并將結果寫回輸入數據中。
d[i*m]=va*curCos-vb*curSin;
d[i*m+m/2]=va*curSin+vb*curCos;
}



	

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

	
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上的實現不算高校，不過優點在于它支持了完整的int8和int4量化的計算，有興趣的讀者可以自行研究這部分kernel實現。

	0x5. LLMSamping解析

	在 chatglm-6b 的實現中，在前向推理完成后以及tokenizer解碼之前有一個根據logits取label的過程：https://github.com/ztxz16/fastllm/blob/master/src/models/chatglm.cpp#L267-L279 。

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

	

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

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

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


	

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

	0x6. 總結

	接著 大模型部署框架 FastLLM 簡要解析 這篇文章首先梳理了一下FastLLM的調用鏈和關鍵的數據結構，然后解析了 FastLLM 的一些實現細節和CPU/GPU后端實現采用的優化技巧。

	審核編輯：湯梓紅