大模型部署框架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;i< axisData.Count(0); i++) {
            ((int32_t*)axisData.cpuData)[i] = axis[i];
        }
        curExecutor->Run("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;i< devices.size(); i++) {
            delete devices[i];
        }
    }

    void Executor::ClearDevices() {
        // this->devices指的是当前对象的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;i< batch; i++) {
                    lockInCPU |= ((Data**)it.second)[i]->lockInCPU;
}
}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;i< batch; i++) {
                            ((Data**)it.second)[i]->ToDevice((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;i< q.size(); i++) {
            // 取出当前遍历到的节点 now。
            TrieNode *now = q[i];
            // 对当前节点 now 的所有子节点进行遍历。
            for (auto it : now->next){
//将当前节点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;i< s.size(); i++) {
            // 检查当前字符 s[i] 是否已经存在于当前节点 now 的 next 映射表中。
            // 如果当前字符 s[i] 不存在于当前节点 now 的子节点中，
            // 在 now->next中添加新的子节点，该子节点的键为当前字符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;i< s.size(); i++) {
            // 创建两个整数变量 tokenId 和 pos，
            // 用于记录找到的 token 的 tokenId 值和 token 的结束位置。
            int tokenId = -999999, pos = i - 1;
            // 创建一个指向 TrieNode 的指针 now，并将其初始化为指向字典树的根节点 root。
            TrieNode *now = this->root;
//从当前字符s[i]开始继续遍历字符串s。
for(intj=i;j< s.size(); j++) {
              // 检查当前字符 s[j] 是否存在于当前节点 now 的 next 映射表中。
              // 如果存在，表示当前字符构成了一个 token 的一部分，继续遍历子节点。
                if (now->next.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< tokens.size(); i++) {
            // 获取当前 token 对应的原始字符串 s，通过查询 tokenToStringDict 映射表，
            // 将 tokens[i] 转换回字符串。
            std::string s = tokenToStringDict[tokens[i]];
            // 判断当前 token 是否需要特殊处理：
            // 如果 s 是类似 "<0xHH>"格式的token（其中HH表示十六进制数），
//则需要将其转换为对应的字符。首先，提取HH，然后将其转换为对应的字符，
//并用空格代替原始的token。
if(s.size()==6&&s.substr(0,3)=="<0x" && s.back() == '>'){
intc=0;
for(inti=3;i< 5; i++) {
                    c *= 16;
                    if (s[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< data.Count(0); i++) {
            tokens.push_back((int) ((float *) data.cpuData)[i]);
        }
        return DecodeTokens(tokens);
    }

上面的：

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< n; i += 32) {
            // 这两行代码从数组 a 和 b 中加载数据到 256 位的向量 bx 和 by。
            __m256i bx = _mm256_loadu_si256((const __m256i *) (a + i));
            __m256i by = _mm256_loadu_si256((const __m256i *) (b + i));
            
            // 这行代码将 by 中的每个元素减去 128，这对应于上面表达式中的 ((int)b[i] - 128)。
            by = _mm256_xor_si256(by, xors);
            // 这行代码对于那些原本是 0 的元素（在减去 128 后变为 -128 的元素）加 1，
            // 以避免后续乘法操作时的溢出。
            by = _mm256_add_epi8(by, _mm256_and_si256(_mm256_cmpeq_epi8(by, xors), ones8));
            
            //  这行代码将 bx 中的符号应用到 by 中，对应于上面表达式中的 ((int8_t*)a)[i]。
            by = _mm256_sign_epi8(by, bx);
            // 这行代码将 bx 中的所有非零元素变为 1，这是为了在后续的乘法操作中保持 by 中元素的原值。
            bx = _mm256_sign_epi8(bx, bx);
            
            // 这行代码先对 bx 和 by 进行乘法运算（这对应于上面表达式中的 * 操作），
            // 然后再与 acc 进行加法操作（这对应于上面表达式中的 += 操作）。
            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(_mm256_maddubs_epi16(bx, by), ones));
        }
        // 这是另一个循环，用于处理数组中剩余的元素（数量小于 32）。
        // 这些元素通过常规的方式计算点积，然后累加到 ans 中。
        for (; i < n; i++) {
            ans += ((int8_t*)a)[i] * ((int)b[i] - 128);
        }
        
        // 最后，将 acc 中的所有元素相加，然后再加上 ans，返回最终的结果。
        return ans + I32sum(acc);
    };
#else
    // 定义了一个函数 DotU8U8，它接受两个指向 8 位无符号整数的指针 a 和 b，
    // 以及一个整数 n。这个函数的目的是计算数组 a 和 b 的点积，其中数组的长度为 n。
    int DotU8U8(uint8_t *a, uint8_t *b, int n) {
        // 初始化一个 256 位的整数向量 acc，所有位都设置为零。这个向量用于存储点积的累加值。
        __m256i acc = _mm256_setzero_si256();

        int i = 0;
        int ans = 0;
        // 这是一个循环，每次处理 32 个元素。这是因为 AVX 可以同时处理 32 个 8 位整数。
        for (; i + 31 < n; i += 32) {
            // 这两行代码从数组 a 和 b 中加载数据到 256 位的向量 bx 和 by。
            __m256i bx = _mm256_loadu_si256((const __m256i *) (a + i));
            __m256i by = _mm256_loadu_si256((const __m256i *) (b + i));
            
            // 接下来的四行代码将 bx 和 by 中的 8 位整数扩展为 16 位整数。
            // 这是因为在后续的乘法和累加操作中，如果仍然使用 8 位整数，可能会发生溢出。
            __m256i mx0 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(bx, 0));
            __m256i mx1 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(bx, 1));

            __m256i my0 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(by, 0));
            __m256i my1 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(by, 1));
            
            // 这两行代码首先对 mx0 和 my0，以及 mx1 和 my1 进行乘法累加操作，
            // 然后再与 acc 进行加法操作，结果存储在 acc 中。
            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(mx0, my0));
            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(mx1, my1));
        }
        //  这是另一个循环，用于处理数组中剩余的元素（数量小于 32）。
        // 这些元素通过常规的方式计算点积，然后累加到 ans 中。
        for (; i < n; i++) {
            ans += a[i] * b[i];
        }
        
        // 最后，将 acc 中的所有元素相加，然后再加上 ans，返回最终的结果。
        return ans + I32sum(acc);
    };
#endif
    // 它接受两个指向 8 位无符号整数的指针 a 和 b，以及一个整数 n。
    // 这个函数的目的是计算数组 a 和 b 的点积，其中数组的长度为 n。
    int DotU4U8(uint8_t *a, uint8_t *b, int n) {
        // 初始化一个 256 位的整数向量 acc，所有位都设置为零。这个向量用于存储点积的累加值。
        __m256i acc = _mm256_setzero_si256();

        int i = 0;
        int ans = 0;
        // 初始化两个常量向量，lowMask 中的每个元素都是 0xf，ones 中的每个元素都是 1。
        const __m256i lowMask = _mm256_set1_epi8(0xf);
        const __m256i ones = _mm256_set1_epi16(1);
        for (; i + 31 < n; i += 32) {
            // 从数组 a 中加载 16 个元素到 128 位的向量 orix 中。
            // 这里 i / 2 的原因是每个元素实际上只有 4 位。
            __m128i orix = _mm_loadu_si128((const __m128i *) (a + i / 2));
            // 将 orix 中的元素分成高 4 位和低 4 位，然后将它们合并成一个 256 位的向量 bytex。
            __m256i bytex = _mm256_set_m128i(_mm_srli_epi16(orix, 4), orix);
            // 使用按位与操作，取 bytex 中的每个元素的低 4 位，结果存储在 bx 中。
            __m256i bx = _mm256_and_si256(lowMask, bytex);
            // 从数组 b 中加载数据到 256 位的向量 by。
            __m256i by = _mm256_loadu_si256((const __m256i *) (b + i));
            // 这行代码首先进行了两个向量的乘法累加操作，然后再与 acc 进行加法操作，结果存储在 acc 中。
            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(_mm256_maddubs_epi16(by, bx), ones));
        }
        for (; i < n; i++) {
            ans += a[i] * b[i];
        }

        return ans + I32sum(acc);
    };
#endif

在启用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的维度中提取词汇大小（vocabSize）和嵌入大小（embSize）。
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的维度中提取词汇大小（vocabSize）和嵌入大小（embSize）。
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< inputLen; i++) {
                    // 这行代码从 inputData 中取出第 i 个元素，并将其四舍五入到最近的整数。
                    int token = (int) (inputData[i] + 1e-9);
                    // 这两行代码将文件指针移动到第 token 行的开始位置。
#if defined(_WIN32) or defined(_WIN64)
                    _fseeki64(fi, (long long)token * embSize * sizeof(float) + weight.filePos, 0);
#else
                    fseek(fi, (long long)token * embSize * sizeof(float) + weight.filePos, 0);
#endif
                    // 这行代码从文件中读取 embSize 个浮点数，并将它们存储在 outputData 的对应位置。
                    int ret = fread(outputData + i * embSize, sizeof(float), embSize, fi);
                }
            } else {
                // 如果权重的数据类型为 BFLOAT16，则使用 16 位浮点数进行计算。
                // 这部分代码的逻辑与 FLOAT32 部分的逻辑类似，只是多了一个步骤：
                // 将 16 位的浮点数转换为 32 位的浮点数。
                uint16_t *outputData = (uint16_t *) output.cpuData;
                uint16_t *weightData = new uint16_t[embSize];
                for (int i = 0; i < inputLen; i++) {
                    int token = (int) (inputData[i] + 1e-9);
#if defined(_WIN32) or defined(_WIN64)
                    _fseeki64(fi, (long long)token * embSize * sizeof(uint16_t) + weight.filePos, 0);
#else
                    fseek(fi, (long long)token * embSize * sizeof(uint16_t) + weight.filePos, 0);
#endif
                    int ret = fread(weightData, sizeof(uint16_t), embSize, fi);
                    for (int j = 0; j < embSize; j++) {
                        outputData[i * embSize * 2 + j * 2] = 0;
                        outputData[i * embSize * 2 + j * 2 + 1] = weightData[j];
                    }
                }
                delete[] weightData;
            }
            // 最后，fclose(fi); 这行代码关闭了文件。
            fclose(fi);
        } else {
            if (weight.dataType == DataType::FLOAT32) {
                // 这两行代码获取 output 和 weight 的数据，并将它们转换为浮点数的指针。
                float *outputData = (float *) output.cpuData;
                float *weightData = (float *) weight.cpuData;
                for (int i = 0; i < inputLen; i++) {
                    int token = (int) (inputData[i] + 1e-9);
                    // 这行代码从 weightData 中复制 embSize 个浮点数到 outputData 的对应位置。
                    // 这里的 token 是索引，embSize 是嵌入向量的长度。
                    memcpy(outputData + i * embSize, weightData + token * embSize, embSize * sizeof(float));
                }
            } else {
                uint16_t *outputData = (uint16_t *) output.cpuData;
                uint16_t *weightData = (uint16_t *) weight.cpuData;
                for (int i = 0; i < inputLen; i++) {
                    int token = (int) (inputData[i] + 1e-9);
                    for (int j = 0; j < embSize; j++) {
                        outputData[i * embSize * 2 + j * 2] = 0;
                        outputData[i * embSize * 2 + j * 2 + 1] = weightData[token * embSize + j];
                    }
                }
            }
        }
    }

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< outer; i++) {
                // 这行代码定义了三个浮点数变量，分别用于存储均值、平方和和方差。
                float mean = 0.f, s2 = 0.f, var = 0.f;
                int j = 0;
                // 这是一段条件编译的代码，只有在目标平台为 ARM 架构时才会编译和执行。
                // 这段代码使用了 ARM 架构的 SIMD 指令来加速计算。
#ifdef __aarch64__
                float32x4_t sums = vdupq_n_f32(0.0);
                    float32x4_t sums2 = vdupq_n_f32(0.0);
                    for (; j + 3 < channels; j += 4) {
                        float32x4_t vi = vld1q_f32(inputData + j);
                        sums = vaddq_f32(sums, vi);
                        sums2 = vaddq_f32(sums2, vmulq_f32(vi, vi));
                    }
                    mean = sums[0] + sums[1] + sums[2] + sums[3];
                    s2 = sums2[0] + sums2[1] + sums2[2] + sums2[3];
#endif
#ifdef __AVX2__
                // 这是另一段条件编译的代码，只有在目标平台支持 AVX2 指令集时才会编译和执行。
                // 这段代码使用了 AVX2 的 SIMD 指令来加速计算。
                __m256 sum_vec = _mm256_setzero_ps();
                __m256 squared_sum_vec = _mm256_setzero_ps();

                for (; j < channels - 7; j += 8) {
                    __m256 data_vec = _mm256_loadu_ps(inputData + j);
                    sum_vec = _mm256_add_ps(sum_vec, data_vec);

                    __m256 squared_data_vec = _mm256_mul_ps(data_vec, data_vec);
                    squared_sum_vec = _mm256_add_ps(squared_sum_vec, squared_data_vec);
                }

                float sum_array[8];
                _mm256_storeu_ps(sum_array, sum_vec);
                mean = sum_array[0] + sum_array[1] + sum_array[2] + sum_array[3] +
                            sum_array[4] + sum_array[5] + sum_array[6] + sum_array[7];

                float squared_sum_array[8];
                _mm256_storeu_ps(squared_sum_array, squared_sum_vec);
                s2 = squared_sum_array[0] + squared_sum_array[1] +
                                    squared_sum_array[2] + squared_sum_array[3] +
                                    squared_sum_array[4] + squared_sum_array[5] +
                                    squared_sum_array[6] + squared_sum_array[7];
#endif
                // 这是一个循环，对 input 中剩余的每一个通道进行处理。
                for (; j < channels; j++) {
                    mean += inputData[j];
                    s2 += inputData[j] * inputData[j];
                }
                // 这两行代码计算了均值和方差。
                mean /= channels;
                var = sqrt(s2 / channels - mean*mean + 1e-10);
                // 接下来是对output的每一个通道进行并行处理
                j = 0;
#ifdef __aarch64__
                float32x4_t means = vdupq_n_f32(mean);
                    float32x4_t vars = vdupq_n_f32(1.0 / var);
                    for (; j + 3 < channels; j += 4) {
                        float32x4_t va = vld1q_f32(gammaData + j), vb = vld1q_f32(betaData + j);
                        float32x4_t vi = vld1q_f32(inputData + j);
                        float32x4_t vo = vaddq_f32(vmulq_f32(vmulq_f32(vsubq_f32(vi, means), vars), va), vb);
                        vst1q_f32(outputData + j, vo);
                    }
#endif
                for (; j < channels; j++) {
                    float a = gammaData[j], b = betaData[j];
                    outputData[j] = (inputData[j] - mean) / var * a + b;
                }
                
                // 这两行代码更新了 inputData 和 outputData 的指针位置，
                // 以便在下一轮循环中处理下一个外部元素。
                inputData += channels;
                outputData += channels;
            }
            return;
        } else {
            // 这段代码同样是执行层归一化（Layer Normalization）操作，但这次的操作更为通用，
            // 能处理 inner 不等于 1 的情况，即每个通道有多个元素的情况。
            // 这是一个循环，对 input 中的每一个外部元素进行处理。
            for (int i = 0; i < outer; i++) {
                // 这两行代码将 mean 和 var 数组的所有元素初始化为 0。
                std::fill(mean, mean + inner, 0.f);
                std::fill(var, var + inner, 0.f);
                // 这行代码定义了一个指针 inputWalk，指向 inputData。
                float *inputWalk = inputData;
                // 这是一个循环，对每个通道进行处理。
                for (int j = 0; j < channels; j++) {
                   // 这是一个嵌套循环，对每个通道内的每个元素进行处理。
                    for (int k = 0; k < inner; k++) {
                        // 这行代码将当前元素的值加到对应的 mean 中，然后 inputWalk 指针向后移动。
                        mean[k] += *inputWalk++; 
                    }
                }
                // 这是另一个循环，计算每个通道的均值。
                for (int k = 0; k < inner; k++) {
                    mean[k] /= channels;
                }
                // 方差类似
                inputWalk = inputData;
                for (int j = 0; j < channels; j++) {
                    for (int k = 0; k < inner; k++) {
                        float x = (*inputWalk++) - mean[k];
                        var[k] += x * x;
                    }
                }
                for (int k = 0; k < inner; k++) {
                    var[k] = sqrt(var[k] / channels + 1e-5);
                }
                
                // 计算输出也是类似
                inputWalk = inputData;
                float *outputWalk = outputData;
                for (int j = 0; j < channels; j++) {
                    float a = gammaData[j], b = betaData[j];
                    for (int k = 0; k < inner; k++) {
                        *outputWalk++ = ((*inputWalk++) - mean[k]) / var[k] * a + b;
                    }
                }

                inputData += channels * inner;
                outputData += channels * inner;
            }
            delete[] mean;
            delete[] var;
        }
    }

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;i< threadNum - 1; i++) {
                int end = cur + per + (cur + per * (threadNum - i) < k);
                futures.push_back(pool->Submit(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;i< futures.size(); i++) {
                futures[i].get();
            }
        } else if (weight.dataType == DataType::FLOAT16) {
            float *inputData = (float *) input.cpuData;
            uint16_t *weightData = (uint16_t *) weight.cpuData;
            float *outputData = (float *) output.cpuData;
            float *biasData = bias.dims.size() >0?(float*)bias.cpuData:nullptr;
#ifdef__ARM_FEATURE_FP16_VECTOR_ARITHMETIC
uint16_t*temp=newuint16_t[n*m];
for(inti=0;i< n * m; i++) {
                temp[i] = float_to_half(inputData[i]);
            }
            inputData = (float*)temp;
#endif
            int threadNum = GetThreads();
            int per = k / threadNum;
            int cur = 0;
            auto pool = GetPool();
            std::vector >futures;
for(inti=0;i< threadNum - 1; i++) {
                int end = cur + per + (cur + per * (threadNum - i) < k);
                futures.push_back(pool->Submit(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;i< futures.size(); i++) {
                futures[i].get();
            }
#ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
            delete[] temp;
#endif
        } else if (weight.dataType == DataType::INT8) { // 这段代码处理权重数据类型为 INT8 的情况。
            // 这段代码首先对输入、权重、输出和偏置数据的指针进行类型转换，
            // 并根据偏置数据的维度是否大于0来决定是否获取偏置数据的指针。然后，它计算了权重数据的总和。
            float *inputData = (float *) input.cpuData;
            uint8_t *weightData = (uint8_t *) weight.cpuData;
            float *outputData = (float *) output.cpuData;
            float *biasData = bias.dims.size() >0?(float*)bias.cpuData:nullptr;
weight.CalcWeightSum();

//之后，代码创建一个std::vector对象，
//LowBitConfig是一个用于存储数据量化信息的类，包括最小值、最大值、位宽和零点。
//这些信息是通过遍历输入数据获得的。
std::vectorinputConfigs;
for(inti=0;i< n; i++) {
                float minValue = 1e9, maxValue = -1e9;
                for (int j = 0; j < m; j++) {
                    minValue = std::min(minValue, inputData[i * m + j]);
                    maxValue = std::max(maxValue, inputData[i * m + j]);
                }
                inputConfigs.push_back(LowBitConfig(minValue, maxValue, 8, 0));
            }
            // 接着，创建一个std::vector对象uinput，并将其大小设置为输入数据的大小（n*m）。
//uinput中的每个元素都是输入数据元素经过inputConfigs中对应配置信息量化后的结果。
//注意这里的量化过程可能会根据是否定义了__AVX2__进行不同的处理。
std::vectoruinput;
uinput.resize(n*m);
for(inti=0;i< n * m; i++) {
#ifdef __AVX2__
                uinput[i] = inputConfigs[i / m].quantization(inputData[i]);
                uinput[i] = (uinput[i] + !uinput[i]) ^ 128;
#else
                uinput[i] = inputConfigs[i / m].quantization(inputData[i]);
#endif
            }
            
            // 随后，调用MultiplyMultiThread函数，使用多线程并行计算uinput和weightData的乘积，
            // 并将结果存储在outputData中。
            MultiplyMultiThread(uinput.data(), weightData, (int32_t*)outputData, n, m, k, GetThreads());
            // 这段代码的目的是把在使用INT8进行量化计算时由于量化造成的误差进行修正，
            // 使得结果更接近于使用浮点数进行计算的结果。也就是反量化过程。
            for (int i = 0; i < n; i++) {
                // 这一步中，对于每一个输入向量（i从0到n），代码首先初始化inputSum为0，
                // 然后遍历输入向量的每个元素（j从0到m），将元素值加到inputSum上。
                // 如果定义了__AVX2__，则在加到inputSum之前，元素值会先与128进行异或操作。
                uint32_t inputSum = 0;
                for (int j = 0; j < m; j++) {
#ifdef __AVX2__
                    inputSum += uinput[i * m + j] ^ 128;
#else
                    inputSum += uinput[i * m + j];
#endif
                }
                
                // 接下来，代码遍历每个输出元素（j从0到k），并按照以下步骤进行调整和缩放：
                for (int j = 0; j < k; j++) {
                    // 首先，获取输出元素的原始值value。
                    int value = ((int32_t*)outputData)[i * k + j];
#ifdef __AVX2__
                    // 如果定义了__AVX2__，则value会增加128 * weight.weightSum[j]、
                    // 128 * inputSum，并减去m * 128 * 128。
                    value += (128 * weight.weightSum[j]);
                    value += (128 * inputSum);
                    value -= m * 128 * 128;
#endif
                    value -= weight.weightSum[j] * inputConfigs[i].zeroPoint;
                    value -= inputSum * weight.perChannelsConfigs[j].zeroPoint;
                    value += (int)inputConfigs[i].zeroPoint * weight.perChannelsConfigs[j].zeroPoint * m;
                    outputData[i * k + j] = weight.perChannelsConfigs[j].scale * inputConfigs[i].scale * value +
                                            (biasData == nullptr ? 0.0 : biasData[j]);
                }
            }
        } else if (weight.dataType == DataType::INT4 || weight.dataType == DataType::INT4_NOZERO) {
            float *inputData = (float *) input.cpuData;
            uint8_t *weightData = (uint8_t *) weight.cpuData;
            float *outputData = (float *) output.cpuData;
            float *biasData = bias.dims.size() >0?(float*)bias.cpuData:nullptr;
weight.CalcWeightSum();

std::vectorinputConfigs;
for(inti=0;i< n; i++) {
                float minValue = 1e9, maxValue = -1e9;
                for (int j = 0; j < m; j++) {
                    minValue = std::min(minValue, inputData[i * m + j]);
                    maxValue = std::max(maxValue, inputData[i * m + j]);
                }
                inputConfigs.push_back(LowBitConfig(minValue, maxValue, 8, 0));
            }
            std::vector uinput;
uinput.resize(n*m);
for(inti=0;i< n * m; i++) {
                uinput[i] = inputConfigs[i / m].quantization(inputData[i]);
            }
#ifdef __AVX__
            uint8_t *temp = new uint8_t[32];
            for (int i = 0; i < n; i++) {
                for (int j = 0; j + 31 < m; j += 32) {
                    memcpy(temp, uinput.data() + i * m + j, 32);
                    for (int k = 0; k < 16; k++) {
                        uinput[i * m + j + k] = temp[k * 2 + 1];
                        uinput[i * m + j + k + 16] = temp[k * 2];
                    }
                }
            }
            delete[] temp;
#endif
            if (weight.dataType == DataType::INT4) {
                MultiplyInt4MultiThread(uinput.data(), weightData, (int32_t *) outputData, n, m, k,
                                        weight.weightSum.data(), weight.zeros.data(), weight.scales.data(), biasData,
                                        inputConfigs, GetThreads());
            } else {
                MultiplyInt4NoZeroMultiThread(uinput.data(), weightData, (int32_t *) outputData, n, m, k,
                                        weight.weightSum.data(), weight.mins.data(), weight.scales.data(), biasData,
                                        inputConfigs, GetThreads());
            }
        } else {
            ErrorInFastLLM("Linear error: unsupport weight's dataType.
");
        }
//float spend = GetSpan(st, std::now());
//float gops = (float)n * m * k / spend / 1e9;
// printf("n = %d, m = %d, k = %d, spend %f s, gops = %f
", n, m, k, spend, gops);
    }

在上面的实现中，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;i< threadNum; i++) {
                int end = cur + per + (cur + per * (threadNum - i) < k);
                if (i == threadNum - 1) {
                    end = k;
                }
                futures.push_back(pool->Submit(Multiply,a,b+cur*m,c+cur,n,m,end-cur,k));
cur=end;
}
for(inti=0;i< futures.size(); i++) {
                futures[i].get();
            }
        }
    }

可以看到这段代码仍然是在用线程池来启动多个线程完成计算，核心部分是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< n; block++) {
            uint8_t *weightWalk = b;
            uint8_t *inputStart = a + block * m;

            for (int i = 0; i < k; i++) {
                int value = 0;
                uint8_t *inputWalk = inputStart;
                int j = 0;
                uint32x4_t sum0 = {0, 0, 0, 0};
                for (; j + 31 < m; j += 32) {
                    uint8x16_t vi = vld1q_u8(inputWalk);
                    uint8x16_t vi0 = vld1q_u8(inputWalk + 16);
                    uint8x16_t vw = vld1q_u8(weightWalk);
                    uint8x16_t vw0 = vld1q_u8(weightWalk + 16);
                    sum0 = vdotq_u32(sum0, vi, vw);
                    sum0 = vdotq_u32(sum0, vi0, vw0);
                    inputWalk += 32;
                    weightWalk += 32;
                }

                value += sum0[0] + sum0[1] + sum0[2] + sum0[3];
                for (; j < m; j++) {
        value += (int)(*(weightWalk++)) * (*(inputWalk++));
       }
                c[block * kstride + i] = value;
            }
        }
#elif defined(__aarch64__)
        int block = 0;
        for (; block < n; block++) {
            uint8_t *weightWalk = b;
            uint8_t *inputStart = a + block * m;

            for (int i = 0; i < k; i++) {
                int value = 0;
                uint8_t *inputWalk = inputStart;

                int per = 64;
                int cnt = m / per;
                int sur = m % per;

                uint32x4_t sum = {0};
                uint16x8_t temp = {0};
                uint16x8_t temp1 = {0};
                uint16x8_t temp2 = {0};
                uint16x8_t temp3 = {0};
                uint16x8_t temp4 = {0};
                uint16x8_t temp5 = {0};
                uint16x8_t temp6 = {0};
                uint16x8_t temp7 = {0};

                while (cnt--) {
                    temp = vmull_u8(vld1_u8(inputWalk), vld1_u8(weightWalk));
                    temp1 = vmull_u8(vld1_u8(inputWalk + 8), vld1_u8(weightWalk + 8));
                    temp2 = vmull_u8(vld1_u8(inputWalk + 16), vld1_u8(weightWalk + 16));
                    temp3 = vmull_u8(vld1_u8(inputWalk + 24), vld1_u8(weightWalk + 24));
                    temp4 = vmull_u8(vld1_u8(inputWalk + 32), vld1_u8(weightWalk + 32));
                    temp5 = vmull_u8(vld1_u8(inputWalk + 40), vld1_u8(weightWalk + 40));
                    temp6 = vmull_u8(vld1_u8(inputWalk + 48), vld1_u8(weightWalk + 48));
                    temp7 = vmull_u8(vld1_u8(inputWalk + 56), vld1_u8(weightWalk + 56));

                    sum = vpadalq_u16(sum, temp);
                    sum = vpadalq_u16(sum, temp1);
                    sum = vpadalq_u16(sum, temp2);
                    sum = vpadalq_u16(sum, temp3);
                    sum = vpadalq_u16(sum, temp4);
                    sum = vpadalq_u16(sum, temp5);
                    sum = vpadalq_u16(sum, temp6);
                    sum = vpadalq_u16(sum, temp7);

                    inputWalk += per;
                    weightWalk += per;
                }

                value += (sum[0] + sum[1] + sum[2] + sum[3]);
                while (sur--) {
                    value += (int)(*(weightWalk++)) * (*(inputWalk++));
                }

                c[block * kstride + i] = value;
            }
        }
#elif defined(__AVX__)
        int block = 0;
        for (; block < n; block++) {
            uint8_t *weightWalk = b;
            uint8_t *inputStart = a + block * m;

            for (int i = 0; i < k; i++) {
                uint8_t *inputWalk = inputStart;

                c[block * kstride + i] = DotU8U8(inputWalk, weightWalk, m);
                weightWalk += m;
            }
        }
#else
        int block = 0;
     for (; block < n; block++) {
      uint8_t *weightWalk = b;
      uint8_t *inputStart = a + block * m;

      for (int i = 0; i < k; i++) {
       int value = 0;
       uint8_t *inputWalk = inputStart;
       for (int j = 0; j < m; j++) {
        value += (int)(*(weightWalk++)) * (*(inputWalk++));
       }

       c[block * kstride + i] = value;
      }
     }
#endif
    }

这段代码实现了两个矩阵的乘法。输入的两个矩阵是 (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<<< outer * n, min(rotaryDim, m / 2) >>>(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< bs; b++) {
            for (int l = 0; l < len; l++) {
                int index = (int) ((float *) positionIds.cpuData)[b * positionIds.dims.back() + l];
                float *sin = ((float *) sinData.cpuData) + stride * index;
                float *cos = ((float *) cosData.cpuData) + stride * index;
                float *d = (float *) data.cpuData + (b * len + l) * spatial;
                for (int i = 0; i < n; i++) {
                    for (int j = 0; j < rotaryDim && j < m / 2; j++) {
                        float a = d[j], b = d[j + m / 2];
                        d[j] = a * cos[j] - b * sin[j];
                        d[j + m / 2] = a * sin[j] + b * cos[j];
                    }

                    d += m;
                }
            }
        }
    }

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< batch; b++) {
                int base = (maxLen - 1) * batch + b; // 计算基础索引值 base。
                // 将 topk 中对应索引的值取整并添加到 lastRet 中。
                lastRet.push_back((int) (((float *) topk.cpuData)[base * 2] + 1e-3));
            }
        } else {
            for (int b = 0; b < batch; b++) {
                int base = (maxLen - 1) * batch + b; // 计算基础索引值 base。
                // 使用 LLMSampling 方法进行抽样，将结果添加到 lastRet 中。
                lastRet.push_back(LLMSampling(logits, base, generationConfig, lastTokens.units[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]< 0 ? base[id] * config.repeat_penalty : base[id] / config.repeat_penalty);
            }
        }
        // 计算温度的倒数 invTemp。
        float invTemp = 1.0f / config.temperature;
        // 定义一个向量 v，用于存储 。
std::vector>v;
//遍历每个logit，将其值乘以invTemp，并存入v中。
for(inti=0;i< vocabSize; i++) {
            v.push_back(std::make_pair(-base[i] * invTemp, i));
        }
        // 计算 topk，它是词汇量 vocabSize 和 config.top_k 中的较小值。
        int topk = std::min(vocabSize, config.top_k);
        // 对 v 中的前 topk 个元素进行排序。
        std::partial_sort(v.begin(), v.begin() + topk, v.end());
        // 初始化 psum 和 maxValue，maxValue 是 v 中最大的元素。
        float psum = 0.0, maxValue = -v.begin()->first;
//定义一个向量ps，用于存储处理后的概率。
std::vectorps;
//遍历v中的前topk个元素，将其值取exp并减去maxValue，存入ps，同时更新psum。
for(inti=0;i< topk; i++) {
            ps.push_back(expf(-v[i].first - maxValue));
            psum += ps.back();
        }
        float curSum = 0.0;
        // 遍历 ps，将其每个元素除以 psum 并更新 curSum，
        // 当 curSum 大于 config.top_p 时，更新 topk 并退出循环。
        for (int i = 0; i < topk; i++) {
            ps[i] /= psum;
            curSum += ps[i];
            if (curSum >config.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;i< topk; i++) {
            curSum += ps[i];
            if (curSum >rnd||i==topk-1){
returnv[i].second;
}
}
//如果以上步骤都没有返回，那么返回-1。
return-1;
}

LLMSampling实现了一种基于温度和惩罚的采样策略，用于从给定的 logits 中选择一个 id。这种采样的方法可以控制输出文本的多样性。

0x6. 总结

接着 大模型部署框架 FastLLM 简要解析 这篇文章首先梳理了一下FastLLM的调用链和关键的数据结构，然后解析了 FastLLM 的一些实现细节和CPU/GPU后端实现采用的优化技巧。

审核编辑：汤梓红