当前位置： 首页 > news >正文

衡水企业网站建设,把手机网站做成app,内蒙古微信公众号114查,如何做快递api接口网站0x0. 前言 接着 大模型部署框架 FastLLM 简要解析 这篇文章首先梳理了一下FastLLM的调用链和关键的数据结构#xff0c;然后解析了 FastLLM 的一些实现细节和CPU/GPU后端实现采用的优化技巧。 0x1. 调用链和数据结构解析 以chatglm-6b的支持为例#xff0c;函数入口在 htt… 0x0. 前言 接着 大模型部署框架 FastLLM 简要解析 这篇文章首先梳理了一下FastLLM的调用链和关键的数据结构然后解析了 FastLLM 的一些实现细节和CPU/GPU后端实现采用的优化技巧。 0x1. 调用链和数据结构解析 以chatglm-6b的支持为例函数入口在 https://github.com/ztxz16/fastllm/blob/master/src/models/chatglm.cpp#L626 这里的 input 就是输入的 contextstring类型。然后 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来对输入数据以及deviceshape等信息进行统一管理。下面的代码展示了Data数据结构的定义源码在https://github.com/ztxz16/fastllm/blob/master/include/fastllm.h#L201-L286 class Data {public:bool lockInCPU false; // 如果lock在CPU上那么不允许移动到其余设备WeightType weightType WeightType::NONE; // 权重类型NONE代表非权重或未知权重DataType dataType DataType::FLOAT32; // 数据类型int unitSize, unitSizeDiv 1; // 单个元素的字节数 unitSIze / unitSizeDivstd::vector int dims; // 数据形状std::vector uint64_t strides; // 跨度uint64_t expansionSize 0; // 扩容后的尺寸uint64_t expansionBytes 0; // 扩容后的字节数std::vector int expansionDims; // 预扩容的形状uint8_t *cpuData nullptr; // 数据指针void cudaData nullptr;std::vector void extraCudaData;void deviceData nullptr;std::vector void extraDeviceData;DataDevice dataDevice DataDevice::CPU;// 这两个参数用于量化对FLOAT数据不适用int perChannelAxis -1; // 沿哪个轴分通道量化-1代表没有分通道std::vector LowBitConfig perChannelsConfigs; // perChannelsConfigs[i]代表第i个通道的min, max; 如果没有分通道perChannelsConfigs[0]代表全局min, maxstd::vector float scales, mins;std::vector int zeros;std::vector int weightSum; // 作为权重时有时候需要存一些和加速计算std::string fileName;long long filePos;std::shared_ptrFileMmap m_file;Data () {};Data (DataType type);Data (DataType type, const std::vector int dims); // 构造函数// 构造函数创建好之后从data复制数据// data中是原始数据如果type不是float那么需要量化Data (DataType type, const std::vector int dims, const std::vector float data);~Data(); // 析构函数Data (const Data ori); // 深拷贝void CopyFrom(const Data ori); // 复制uint64_t GetBytes() const; // 获取总字节数void Allocate(); // 分配内存void Allocate(float v); // 分配内存并初始化void Expansion(const std::vector int dims); // 预扩容到相应尺寸void MallocSpace(uint64_t size); // 在设备上分配void FreeSpace(); // 回收设备上的内存void UpdateUnitSize(); // 更新unitSizevoid Resize(const std::vector int dims); // 更改尺寸void Reshape(const std::vector int dims); // 更改尺寸,但不修改数据uint64_t Count(int i) const; // dims[i] * strides[i]void PrintShape() const; // 输出形状void Print() const; // 输出void CalcWeightSum(); // 计算WeightSumvoid ToDevice(DataDevice device); // 移动到指定devicevoid ToDevice(void device);void set_file(std::shared_ptrFileMmap file) {m_file file;}};在Forward函数里面以Data为核心载体运行chatglm-6b模型的流程具体包含如下的一些算子https://github.com/ztxz16/fastllm/blob/master/include/fastllm.h#L346-L408 。以Permute为例我们浏览下它的实现 void Permute(const Data input, const std::vectorint axis, Data output) {Data axisData Data(DataType::INT32PARAM, {(int)axis.size()});axisData.Allocate();for (int i 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 std::string, Data* DataDict;。接着我们看一下curExecutor的定义和实现 namespace fastllm {class Executor {private:std::vector BaseDevice* devices;std::map std::string, float profiler;public:Executor (); // 创建默认的Executor~Executor(); // 析构void ClearDevices(); // 清空 devicesvoid AddDevice(BaseDevice device); // 增加一个device// 运行一个opvoid Run(const std::string opType, const fastllm::DataDict datas, const fastllm::FloatDict floatParams,const fastllm::IntDict intParams);void ClearProfiler();void PrintProfiler();}; }从Executor类的定义我们可以判断它负责了在设定的devices上根据opType和输入数据等执行Op的前向计算也就是Run这个接口。由于Executor类是FastLLM的调度核心实现所以我们来详细解析一下它的实现。 namespace fastllm {Executor::Executor() {this-devices.clear(); #ifdef USE_CUDA// 将一个指向 CudaDevice 类对象的指针插入到 devices 向量的末尾。// 这里通过 new 运算符创建了一个 CudaDevice 对象并将返回的指针进行类型转换为 BaseDevice 类型。this-devices.push_back((BaseDevice) new CudaDevice()); #endifthis-devices.push_back((BaseDevice) new CpuDevice());}Executor::~Executor() {// 释放 devices 向量中的每个指针元素所占用的内存。for (int i 0; i devices.size(); i) {delete devices[i];}}void Executor::ClearDevices() {// this-devices 指的是当前对象的 devices 成员即指向 BaseDevice 类对象的指针向量。this-devices.clear();}// 该函数用于向 devices 向量中添加一个指向 BaseDevice 类对象的指针。void Executor::AddDevice(fastllm::BaseDevice *device) {this-devices.push_back(device);}void Executor::Run(const std::string opType, const fastllm::DataDict datas, const fastllm::FloatDict floatParams,const fastllm::IntDict intParams) {// 创建一个 st 变量用于记录函数开始执行的时间。auto st std::chrono::system_clock::now();// 创建一个布尔变量 lockInCPU用于记录是否将数据锁定在 CPU 上。bool lockInCPU false;// 在第一个 for 循环中遍历数据字典 datas查找是否有 ___batch 后缀的参数// 并根据情况设置 lockInCPU 的值。it.first 是数据字典中的键keyit.second // 是对应的值value。如果存在 ___batch 后缀的参数则将 lockInCPU 设置为// 对应数据的 lockInCPU 属性布尔值否则设置为当前数据的 lockInCPU 属性。for (auto it: datas) {if (intParams.find(it.first ___batch) ! intParams.end()) {int batch intParams.find(it.first ___batch)-second;for (int i 0; i batch; i) {lockInCPU | ((Data)it.second)[i]-lockInCPU;}} else {lockInCPU | it.second-lockInCPU;}}// 第二个 for 循环遍历 devices 向量中的所有设备指针 device。// 在循环中首先检查 lockInCPU 是否为真并且当前设备的类型不是 cpu// 如果是则跳过当前设备continue。这个检查是为了保证数据锁定在 CPU 上时只执行 CPU 设备上的操作。for (auto device: 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()) {int batch intParams.find(it.first ___batch)-second;for (int i 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 成员变量用于性能分析。float spend GetSpan(st, std::chrono::system_clock::now());profiler[opType] spend;}// 清除profile的信息void Executor::ClearProfiler() {profiler.clear();}// 打印profile信息也即输出每个层的运行时间和模型的总运行时间void Executor::PrintProfiler() {float sum 0.0;for (auto it : profiler) {printf(%s spend %f\n, it.first.c_str(), it.second);sum it.second;}printf(total spend %f\n, sum);} }自此前向计算就顺利完成了再把推理结果给 tokenizer 解码就结束了整体的调度执行流程是很简单明了的。 0x2. tokenizer 解析 接着我们来解析一下tokenizer的实现。先看一下tokenizer的定义https://github.com/ztxz16/fastllm/blob/master/include/fastllm.h#L287-L310 struct Tokenizer {struct TrieNode {int tokenId;std::map int, TrieNode next;TrieNode();};TrieNode root;std::unordered_map int, std::string tokenToStringDict;Tokenizer ();~Tokenizer();void Clear(); // 清空分词器void Insert(const std::string s, int tokenId); // 插入一个tokenData Encode(const std::string s); // 编码std::string Decode(const Data data); // 解码std::string DecodeTokens(const std::vector int tokens); // 解码};我们从实现来看tokenizer的细节 // 这是 Tokenizer 类的嵌套结构 TrieNode 的构造函数的实现。// 在构造函数中将 tokenId 成员变量的值初始化为 -999999。// 这个值在构造函数中被硬编码它是作为一个特殊标记来使用的。Tokenizer::TrieNode::TrieNode() {this-tokenId -999999;}// Tokenizer 类的构造函数的实现。// 在构造函数中通过 new 运算符创建一个新的 TrieNode 对象// 并将其指针赋值给 root 成员变量。这样构造函数创建了一个空的字典树// 并将其根节点指针存储在 root 中。Tokenizer::Tokenizer() {root new TrieNode();}// Tokenizer 类的析构函数的实现。// 在析构函数中首先调用 Clear() 函数用于释放动态分配的资源和清空数据。// 然后调用 delete 运算符释放通过 new 运算符创建的 root 对象的内存从而释放整个字典树的内存。Tokenizer::~Tokenizer() {Clear();delete root;}// 这是 Tokenizer 类的成员函数 Clear() 的定义用于清空分词器并释放动态分配的资源。void Tokenizer::Clear() {// 创建一个指向 TrieNode 的指针向量 q用于辅助遍历字典树。std::vector TrieNode q;// 将字典树的根节点 root 加入 q 向量作为遍历的起始点。q.push_back(root);// 开始遍历 q 向量中的节点这是一个广度优先搜索BFS的过程。for (int i 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 new TrieNode();// 清空 tokenToStringDict 映射表以确保所有 token 的映射被清空。tokenToStringDict.clear();}// 这是 Tokenizer 类的成员函数 Insert 的定义用于向分词器中插入一个 token。void Tokenizer::Insert(const std::string s, int tokenId) {// 创建一个指向 TrieNode 的指针 now并将其初始化为指向字典树的根节点 root。TrieNode *now this-root;// 开始遍历输入的字符串 s 中的每个字符。for (int i 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]] new TrieNode();}// 将 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 进行编码。Data Tokenizer::Encode(const std::string s) {// 创建一个浮点数向量 v用于存储编码结果。该向量将存储找到的 token 对应的 tokenId 值。std::vector float v;// 开始遍历输入的字符串 s 中的每个字符。for (int i 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 (int j 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(\n);// 遍历完成后v 向量中存储了输入字符串中所有找到的 token 对应的 tokenId 值。// 创建一个 Data 对象并返回表示编码的结果。这里 Data 是一个数据结构// 用于存储数据及其相关信息。编码结果是一个一维浮点数数组// 表示输入字符串中所有找到的 token 对应的 tokenId 值。return Data (DataType::FLOAT32, {1, (int)v.size()}, v);}// 这是 Tokenizer 类的成员函数 DecodeTokens 的定义// 用于对输入的 token 数组进行解码将 token 转换回原始的字符串。std::string Tokenizer::DecodeTokens(const std::vectorint tokens) {// 创建一个空字符串 ret用于存储解码结果。std::string ret ;// 开始遍历输入的 token 数组 tokens。for (int i 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() ) {int c 0;for (int i 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 n) {ret \n;} else if (s |tab|) {ret \t;} else {ret s;}}// 将特殊字符 \xE2\x96\x81UTF-8 编码替换为空格 这是用于表示空格的特殊字符。std::string blank ;blank 226, blank 150, blank 129;while (true) {std::string::size_type pos(0);if ((pos ret.find(blank)) ! std::string::npos)ret.replace(pos, blank.length(), );else break;}// 检查是否有 |blank数字 格式的特殊 token如果有将其解码成对应数量的空格字符。int pos 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 int tokens;for (int i 0; i data.Count(0); i) {tokens.push_back((int) ((float *) data.cpuData)[i]);}return DecodeTokens(tokens);} 上面的 if (pos ! -1) {int space_num atoi(ret.substr(8, ret.size() - 10).c_str());return std::string(space_num, );}这行代码应该是有bug假设 ret 的值为 “Hello|blank_4world!”那么在解码时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 上分配一块内存空间。bool CpuDevice::Malloc(void **ret, size_t size) {ret (void)new uint8_t [size];return true;}// 这是 CpuDevice 类的成员函数 Free 的定义用于在 CPU 上释放之前分配的内存。bool CpuDevice::Free(void ret) {delete[] (uint8_t)ret;return true;}// 这是 CpuDevice 类的成员函数 CopyDataFromCPU 的定义用于将数据从 CPU 拷贝到指定的设备上。// 这里什么都不做直接返回true。bool CpuDevice::CopyDataFromCPU(void *dst, void *src, size_t size) {return true;}// 这是 CpuDevice 类的成员函数 CopyDataToCPU 的定义用于将数据从指定的设备拷贝到 CPU 上。bool CpuDevice::CopyDataToCPU(void *dst, void *src, size_t size) {return true;}// 如果定义了 AVX 和 AVX2那么会启用第一个 DotU8U8 函数和 DotU4U8 函数。 // 如果只定义了 AVX但没有定义 AVX2那么会启用第二个 DotU8U8 函数和 DotU4U8 函数。#ifdef AVX #ifdef AVX2// 这是一段使用了 Intel AVX2 指令集Advanced Vector Extensions 2的代码// 用于计算两个8位无符号整数数组的点积。// 定义了一个函数 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();// 初始化两个变量i 用于循环计数ans 用于存储最后的结果。int i 0;int ans 0;// 等这几行代码初始化了一些常量向量const __m256i lowMask _mm256_set1_epi8(0xf);const __m256i ones _mm256_set1_epi16(1);const __m256i ones8 _mm256_set1_epi8(1);const m256i xors _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 中的每个元素都是 0xfones 中的每个元素都是 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指令细节读者想详细了解可以参考官方文档https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html 。 CpuEmbedding 算子解析 // CpuEmbedding 算子的形状推导函数这个函数接受四个参数 // 一个 std::string 类型的 opType两个字典类型的 datas 和 floatParams以及一个 intParams。 void CpuEmbedding::Reshape(const std::string opType, const fastllm::DataDict datas,const fastllm::FloatDict floatParams, const fastllm::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, Embeddings weights dim should be 2.\n);// 这行代码检查 weight 的数据类型是否为 FLOAT32 或 BFLOAT16。如果不是就会抛出一个错误。AssertInFastLLM(weight.dataType DataType::FLOAT32 ||weight.dataType DataType::BFLOAT16, Embeddings weights type should be float32 or bfloat16.\n);// 这行代码检查 input 的数据类型是否为 FLOAT32。如果不是就会抛出一个错误。AssertInFastLLM(input.dataType DataType::FLOAT32, Embeddings inputs type should be float32.\n);// 这行代码将 weight 的 weightType 属性设置为 EMBEDDING。weight.weightType WeightType::EMBEDDING;// 这行代码从 weight 的维度中提取词汇大小vocabSize和嵌入大小embSize。int vocabSize weight.dims[0], embSize weight.dims[1];// 这两行代码将 embSize 添加到 input 的维度中形成一个新的维度。std::vector int dims input.dims;dims.push_back(embSize);// 这两行代码将 output 的数据类型设置为 FLOAT32并重新调整其维度。output.dataType DataType::FLOAT32;output.Resize(dims);}// 这是一个名为 CpuEmbedding::Run 的函数它在某个名为 CpuEmbedding 的类中被定义。// 这个函数接受四个参数一个 std::string 类型的 opType// 两个字典类型的 datas 和 floatParams以及一个 intParams。// 这个函数的主要任务是执行嵌入层Embedding layer的运算。// 嵌入层通常用于将离散型特征例如词汇转换为连续的向量表示。// 具体的实现方法是对于每个输入的索引从权重矩阵中查找对应的行// 然后将其复制到输出矩阵的对应位置。void CpuEmbedding::Run(const std::string opType, const fastllm::DataDict datas,const fastllm::FloatDict floatParams, const fastllm::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。int vocabSize weight.dims[0], embSize weight.dims[1];// 这行代码计算 input 的长度。uint64_t inputLen 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 (int i 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); #elsefseek(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); #elsefseek(fi, (long long)token * embSize * sizeof(uint16_t) weight.filePos, 0); #endifint 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 解析 void CpuLayerNormOp::Run(const std::string opType, const fastllm::DataDict datas,const fastllm::FloatDict floatParams, const fastllm::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。在层归一化中轴通常是特征维度。int axis intParams.find(axis) ! intParams.end() ? intParams.find(axis)-second : -1;// 这两行代码计算 input 的维度数并将 axis 转换为非负数。// 这是为了处理负数的轴值因为在 Python 中轴可以是负数表示从后向前数的位置。int dimsLen input.dims.size();axis (axis % dimsLen dimsLen) % dimsLen;// 这三行代码计算 outer、channels 和 inner。// outer 是归一化操作的外部维度的元素总数channels 是归一化操作的轴的大小// inner 是归一化操作的内部维度的元素总数。int outer input.Count(0) / input.Count(axis);int channels input.dims[axis];int inner input.strides[axis];// 这行代码为 mean 和 var 分配内存它们用于存储每个归一化组的均值和方差。float *mean new float[inner], *var new float[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 (int i 0; i outer; i) {// 这行代码定义了三个浮点数变量分别用于存储均值、平方和和方差。float mean 0.f, s2 0.f, var 0.f;int j 0;// 这是一段条件编译的代码只有在目标平台为 ARM 架构时才会编译和执行。// 这段代码使用了 ARM 架构的 SIMD 指令来加速计算。 #ifdef aarch64float32x4_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 aarch64float32x4_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);} #endiffor (; j channels; j) {float a gammaData[j], b betaData[j];outputDataj / 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这个算子。 void CpuLinearOp::Run(const std::string opType, const fastllm::DataDict datas,const fastllm::FloatDict floatParams, const fastllm::IntDict intParams) { //auto st std::chrono::system_clock::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);int n input.Count(0) / input.dims.back();int m input.dims.back();int k 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 用来表示当前任务的起始位置。int threadNum GetThreads();int per k / threadNum;int cur 0;// 接着创建线程池通过 GetPool() 获取和用于保存线程任务的std::future数组。// 对于每个线程确定其需要处理的任务范围从 cur 到 end然后提交线程任务。// 线程任务是通过调用 FloatLinearPart 函数来执行的该函数需要输入数据、// 权重数据、偏置数据、输出数据、输入维度n、权重维度m、输出维度k// 以及任务范围从 cur 到 end作为参数。auto pool GetPool();std::vector std::future void futures;for (int i 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::future::get() // 方法来阻塞主线程直到对应的子线程完成任务。// 这样可以保证所有的线程任务都完成后主线程才继续执行。for (int i 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_ARITHMETICuint16_t *temp new uint16_t[n * m];for (int i 0; i n * m; i) {temp[i] float_to_half(inputData[i]);}inputData (float*)temp; #endifint threadNum GetThreads();int per k / threadNum;int cur 0;auto pool GetPool();std::vector std::future void futures;for (int i 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 (int i 0; i futures.size(); i) {futures[i].get();} #ifdef ARM_FEATURE_FP16_VECTOR_ARITHMETICdelete[] 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::vectorLowBitConfig对象// LowBitConfig是一个用于存储数据量化信息的类包括最小值、最大值、位宽和零点。// 这些信息是通过遍历输入数据获得的。std::vector LowBitConfig inputConfigs;for (int i 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::vectoruint8_t对象uinput并将其大小设置为输入数据的大小n * m。// uinput中的每个元素都是输入数据元素经过inputConfigs中对应配置信息量化后的结果。// 注意这里的量化过程可能会根据是否定义了AVX2__进行不同的处理。std::vector uint8_t uinput;uinput.resize(n * m);for (int i 0; i n * m; i) { #ifdef AVX2uinput[i] inputConfigs[i / m].quantization(inputData[i]);uinputi ^ 128; #elseuinput[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 AVX2inputSum uinput[i * m j] ^ 128; #elseinputSum 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; #endifvalue - 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::vector LowBitConfig inputConfigs;for (int i 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 uint8_t uinput;uinput.resize(n * m);for (int i 0; i n * m; i) {uinput[i] inputConfigs[i / m].quantization(inputData[i]);} #ifdef AVXuint8_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; #endifif (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 weights dataType.\n);} //float spend GetSpan(st, std::chrono::system_clock::now()); //float gops (float)n * m * k / spend / 1e9; // printf(n %d, m %d, k %d, spend %f s, gops %f\n, n, m, k, spend, gops);} 在上面的实现中MultiplyMultiThread完成了对量化输入的计算我们看一下它的实现细节 //a [n, m], b [k, m], c aT(b) [n, k]void MultiplyMultiThread(uint8_t *a, uint8_t *b, int32_t *c, int n, int m, int k, int threadNum) {int per k / threadNum;int cur 0;if (threadNum 1) {Multiply(a, b cur * m, c cur, n, m, k - cur, k);} else {auto pool GetPool();std::vectorstd::futurevoid futures;for (int i 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 (int i 0; i futures.size(); i) {futures[i].get();}}}可以看到这段代码仍然是在用线程池来启动多个线程完成计算核心部分是Multiply函数这个函数的实现细节 //a [n, m], b [k, m], c aT(b) [n, k]void Multiply(uint8_t *a, uint8_t *b, int32_t *c, int n, int m, int k, int kstride) { #ifdef ARM_FEATURE_DOTPRODint 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 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;}} #elseint 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)结果矩阵是 ©。矩阵 (a) 的形状是 ([n, m])矩阵 (b) 的形状是 ([k, m])所以矩阵 (c a^T b) 的形状是 ([n, k])。 在这段代码中使用了不同的方法进行矩阵乘法取决于系统是否支持特定的优化硬件指令。 如果系统支持 ARMv8.2 的点积指令ARM_FEATURE_DOTPROD那么会使用这个指令进行矩阵乘法。在这种情况下每次会同时处理32个元素这样可以加速计算。 如果系统支持 ARMv8aarch64但不支持 ARMv8.2 的点积指令那么会使用 NEON SIMD 指令进行矩阵乘法。在这种情况下每次会同时处理64个元素。 如果系统支持 AVXAVX那么会使用 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 。

这个类是用来创建旋转位置编码Rotary Position Embedding的。

Llama模型引入了旋转位置编码以改进长序列处理的性能。

class LlamaRotaryEmbedding(torch.nn.Module):# 这是类的初始化方法接收四个参数dim嵌入的维度max_position_embeddings# 最大的位置嵌入长度默认为2048base基数默认为10000和device设备类型例如CPU或GPU。def init(self, dim, max_position_embeddings2048, base10000, deviceNone):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是PyTorch nn.Module的一个方法# 它用于保存一些不需要计算梯度的变量。self.register_buffer(inv_freq, inv_freq, persistentFalse)# Build here to make torch.jit.trace work.# 调用_set_cos_sin_cache方法预先计算并保存正弦和余弦的缓存值。self._set_cos_sin_cache(seq_lenmax_position_embeddings, deviceself.inv_freq.device, dtypetorch.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, devicedevice, dtypeself.inv_freq.dtype)# 使用外积计算频率和t的乘积结果保存到变量freqs中。freqs torch.einsum(i,j-ij, t, self.inv_freq)# Different from paper, but it uses a different permutation in order to obtain the same calculation# 将频率的两份副本拼接在一起结果保存到变量emb中。emb torch.cat((freqs, freqs), dim-1)# 计算emb的余弦值然后将结果保存到模型的缓存中。self.register_buffer(cos_cached, emb.cos()[None, None, :, :].to(dtype), persistentFalse)# 计算emb的正弦值然后将结果保存到模型的缓存中。self.register_buffer(sin_cached, emb.sin()[None, None, :, :].to(dtype), persistentFalse)# 这是模型的前向传播方法接收两个参数x输入数据和seq_len序列长度。def forward(self, x, seq_lenNone):# x: [bs, num_attention_heads, seq_len, head_size]# 如果输入的序列长度大于缓存的最大序列长度那么调用_set_cos_sin_cache方法更新缓存。if seq_len self.max_seq_len_cached:self._set_cos_sin_cache(seq_lenseq_len, devicex.device, dtypex.dtype)# 返回对应输入位置的正弦和余弦值。这些值将用于旋转位置编码。return (self.cos_cached[:, :, :seq_len, …].to(dtypex.dtype),self.sin_cached[:, :, :seq_len, …].to(dtypex.dtype),)def apply_rotary_pos_emb(q, k, cos, sin, position_ids):# The first two dimensions of cos and sin are always 1, so we can 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)return q_embed, k_embedCudaLlamaRotatePosition2DOp对应的就是上面的Python代码。 void CudaLlamaRotatePosition2DOp::Run(const std::string opType, const fastllm::DataDict datas,const fastllm::FloatDict floatParams, const fastllm::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);int rotaryDim intParams.find(rotaryDim) ! intParams.end() ? intParams.find(rotaryDim)-second : 128;FastllmCudaLlamaRotatePosition2D(data, positionIds, sinData, cosData, rotaryDim);}这里调用的是FastllmCudaLlamaRotatePosition2D这个函数它的实现和解析如下 // 这是一个在 GPU 上运行的 CUDA 函数用于执行 Llama 模型的位置编码旋转操作。 // data输入的数据这个数据将会被旋转。 // positionIds位置编码的数据。 // sinDatacosData用于旋转的 sin 和 cos 值。 // rotaryDim旋转的维度。 bool FastllmCudaLlamaRotatePosition2D(fastllm::Data data, const fastllm::Data positionIds,const fastllm::Data sinData, const fastllm::Data cosData, int rotaryDim) {// 使用 FastllmCudaPrepareInput 函数将输入的数据从 CPU 复制到 GPU。// 这个函数会返回一个指向 GPU 内存的指针。 float *cudaData (float *) FastllmCudaPrepareInput(data);float *cudaPositionIds (float *) FastllmCudaPrepareInput(positionIds);float *cudaSin (float *) FastllmCudaPrepareInput(sinData);float *cudaCos (float *) FastllmCudaPrepareInput(cosData);// 计算旋转操作需要的一些参数包括 outerspatialbslenn 和 m。// 这些参数是用于确定 CUDA 核函数的执行配置和一些数据操作的。int outer data.dims[0] * data.dims[1];int spatial data.Count(2);int bs data.dims[0], len data.dims[1];int n data.dims[2], m data.dims[3];// 调用 CUDA 核函数 FastllmLlamaRotatePosition2DKernel 来在 GPU 上执行位置编码的旋转操作。// outer * n, min(rotaryDim, m / 2) 是 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 函数释放 positionIdssinData 和 cosData 在 GPU 上的内存。// 这些数据在这个函数中不再需要。FastllmCudaFinishInput(positionIds, cudaPositionIds);FastllmCudaFinishInput(sinData, cudaSin);FastllmCudaFinishInput(cosData, cudaCos);// 使用 FastllmCudaFinishOutput 函数将旋转后的数据从 GPU 复制回 CPU。// 这个函数也会释放 data 在 GPU 上的内存。FastllmCudaFinishOutput(data, cudaData);return true; }最后再解析下这个cuda kernel。 // float *data输入数据大小为 [bs, len, n, m]其中 bs 是批量大小 // len 是序列长度n 是头的数量m 是每个头的维度。 // float *positionIds位置编码的索引大小为 [bs, len]。 // float *sin 和 float *cos预先计算的正弦和余弦值用于旋转编码。 // int len, int bs, int spatial, int n, int m输入数据的各个维度大小。 // int partStride 和 int sinCosStride用于索引 positionIds 和 sin/cos 的步长。 // int rotateDim旋转维度。 global void FastllmLlamaRotatePosition2DKernel(float *data, float *positionIds, float *sin, float *cos,int len, int bs, int spatial, int n, int m, int partStride, int sinCosStride, int rotateDim) {// 首先计算出当前线程应处理的位置 o长度 l 和批次 b。int o (blockIdx.x / n);int l o % len;int b o / len;int j threadIdx.x;// 然后根据 positionIds 获取对应的旋转角度的正弦值 curSin 和余弦值 curCos。int index (int) (positionIds[b * partStride l]);float curSin sin[index * sinCosStride j];float curCos cos[index * sinCosStride j];float *d (float *) data o * spatial j;int i blockIdx.x % n;// 接着获取输入数据对应位置的值 va 和 vb。float va 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个元素就是比较合理直观的。 void CpuLlamaRotatePosition2DOp::Run(const std::string opType, const fastllm::DataDict datas,const fastllm::FloatDict floatParams, const fastllm::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);int rotaryDim intParams.find(rotaryDim) ! intParams.end() ? intParams.find(rotaryDim)-second : 128;int bs data.dims[0], len data.dims[1];int spatial data.Count(2);int n data.dims[2], m data.dims[3];int stride (int)sinData.dims[1];for (int b 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 个值这里 K1所以是找到最大值。TopK(logits, topk, 1); topk.ToDevice(DataDevice::CPU);for (int b 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 参数的影响。 int LLMSampling(Data logits, int outerOffset,const GenerationConfig config, const LastTokensUnit tokens) {// 将 logits 数据从当前设备转移到 CPU。logits.ToDevice(DataDevice::CPU);// 从 logits 的维度中获取词汇量 vocabSize。int vocabSize 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 (int id : tokens.tokenSet) {baseid;}}// 计算温度的倒数 invTemp。float invTemp 1.0f / config.temperature;// 定义一个向量 v用于存储 logit值索引。std::vector std::pair float, int v;// 遍历每个 logit将其值乘以 invTemp并存入 v 中。for (int i 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 和 maxValuemaxValue 是 v 中最大的元素。float psum 0.0, maxValue -v.begin()-first;// 定义一个向量 ps用于存储处理后的概率。std::vector float ps;// 遍历 v 中的前 topk 个元素将其值取 exp 并减去 maxValue存入 ps同时更新 psum。for (int i 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。float rnd fastllmRandom.randP();curSum 0.0;// 遍历 ps 中的前 topk 个元素将其累加到 curSum// 当 curSum 大于 rnd 或者达到最后一个元素时// 返回对应 v[i].second也就是返回采样得到的 id。for (int i 0; i topk; i) {curSum ps[i];if (curSum rnd || i topk - 1) {return v[i].second;}}// 如果以上步骤都没有返回那么返回 -1。return -1;}LLMSampling实现了一种基于温度和惩罚的采样策略用于从给定的 logits 中选择一个 id。这种采样的方法可以控制输出文本的多样性。 0x6. 总结 接着 大模型部署框架 FastLLM 简要解析 这篇文章首先梳理了一下FastLLM的调用链和关键的数据结构然后解析了 FastLLM 的一些实现细节和CPU/GPU后端实现采用的优化技巧。