大模型部署框架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
class data { public: bool lockincpu = false; // 如果lock在cpu上，那么不允许移动到其余设备 weighttype weighttype = weighttype::none; // 权重类型，none代表非权重（或未知权重） datatype datatype = datatype::float32; // 数据类型 int unitsize, unitsizediv = 1; // 单个元素的字节数 = unitsize / unitsizediv std::vector dims; // 数据形状 std::vector strides; // 跨度 uint64_t expansionsize = 0; // 扩容后的尺寸 uint64_t expansionbytes = 0; // 扩容后的字节数 std::vector expansiondims; // 预扩容的形状 uint8_t *cpudata = nullptr; // 数据指针 void *cudadata = nullptr; std::vector extracudadata; void *devicedata = nullptr; std::vector extradevicedata; datadevice datadevice = datadevice::cpu; // 这两个参数用于量化，对float数据不适用 int perchannelaxis = -1; // 沿哪个轴分通道量化，-1代表没有分通道 std::vector perchannelsconfigs; // perchannelsconfigs[i]代表第i个通道的min, max; 如果没有分通道，perchannelsconfigs[0]代表全局min, max std::vector scales, mins; std::vector zeros; std::vector weightsum; // 作为权重时，有时候需要存一些和加速计算 std::string filename; long long filepos; std::shared_ptr m_file; data () {}; data (datatype type); data (datatype type, const std::vector &dims); // 构造函数 // 构造函数，创建好之后从data复制数据 // data中是原始数据，如果type不是float那么需要量化 data (datatype type, const std::vector &dims, const std::vector &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 &dims); // 预扩容到相应尺寸 void mallocspace(uint64_t size); // 在设备上分配 void freespace(); // 回收设备上的内存 void updateunitsize(); // 更新unitsize void resize(const std::vector &dims); // 更改尺寸 void reshape(const std::vector &dims); // 更改尺寸,但不修改数据 uint64_t count(int i) const; // dims[i] * strides[i] void printshape() const; // 输出形状 void print() const; // 输出 void calcweightsum(); // 计算weightsum void todevice(datadevice device); // 移动到指定device void todevice(void *device); void set_file(std::shared_ptr 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::vector &axis, data &output) { data axisdata = data(datatype::int32param, {(int)axis.size()}); axisdata.allocate(); for (int i = 0; i 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的定义和实现：
namespace fastllm { class executor { private: std::vector devices; std::map profiler; public: executor (); // 创建默认的executor ~executor(); // 析构 void cleardevices(); // 清空 devices void adddevice(basedevice *device); // 增加一个device // 运行一个op void 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());#endif this->devices.push_back((basedevice*) new cpudevice()); } executor::~executor() { // 释放 devices 向量中的每个指针元素所占用的内存。 for (int i = 0; i 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::now(); // 创建一个布尔变量 lockincpu，用于记录是否将数据锁定在 cpu 上。 bool lockincpu = 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()) { int batch = intparams.find(it.first + ___batch)->second; for (int i = 0; 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 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::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, it.first.c_str(), it.second); sum += it.second; } printf(total spend %f, 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 next; trienode(); }; trienode *root; std::unordered_map tokentostringdict; tokenizer (); ~tokenizer(); void clear(); // 清空分词器 void insert(const std::string &s, int tokenid); // 插入一个token data encode(const std::string &s); // 编码 std::string decode(const data &data); // 解码 std::string decodetokens(const std::vector &tokens); // 解码 };
我们从实现来看tokenizer的细节：
// 这是 tokenizer 类的嵌套结构 trienode 的构造函数的实现。 // 在构造函数中，将 tokenid 成员变量的值初始化为 -999999。 // 这个值在构造函数中被硬编码，它是作为一个特殊标记来使用的。 tokenizer::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 q; // 将字典树的根节点 root 加入 q 向量，作为遍历的起始点。 q.push_back(root); // 开始遍历 q 向量中的节点，这是一个广度优先搜索（bfs）的过程。 for (int i = 0; i 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 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 v; // 开始遍历输入的字符串 s 中的每个字符。 for (int i = 0; i root; // 从当前字符 s[i] 开始继续遍历字符串 s。 for (int j = i; j 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 值。 return data (datatype::float32, {1, (int)v.size()}, v); } // 这是 tokenizer 类的成员函数 decodetokens 的定义， // 用于对输入的 token 数组进行解码，将 token 转换回原始的字符串。 std::string tokenizer::decodetokens(const std::vector &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 是类似 格式的 token（其中 hh 表示十六进制数）， // 则需要将其转换为对应的字符。首先，提取 hh，然后将其转换为对应的字符， // 并用空格代替原始的 token。 if (s.size() == 6 && s.substr(0, 3) == ') { int c = 0; for (int i = 3; i = '0' && s[i] <= '9') { c += (s[i] - '0'); } else { c += (s[i] - 'a' + 10); } } s = ; s[0] = c; } // 根据不同的 token 进行解码： if (s == ) { ret += ; } else if (s == ) { ret += ; } else { ret += s; } } // 将特殊字符 xe2x96x81（utf-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; } // 检查是否有 格式的特殊 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 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 的值为 helloworld!，那么在解码时，pos 将是 8，而 space_num 将是 4。然后，函数将返回 ，即包含四个空格字符的字符串。在这种情况下，特殊 token 被成功解码成了四个空格字符，但是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 中的每个元素都是 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 second); data &output = *(datas.find(output)->second); data &weight = *(datas.find(weight)->second); // 这行代码检查 weight 的维度数量是否为 2。如果不是，就会抛出一个错误。 assertinfastllm(weight.dims.size() == 2, embedding's weight's dim should be 2.); // 这行代码检查 weight 的数据类型是否为 float32 或 bfloat16。如果不是，就会抛出一个错误。 assertinfastllm(weight.datatype == datatype::float32 || weight.datatype == datatype::bfloat16, embedding's weight's type should be float32 or bfloat16.); // 这行代码检查 input 的数据类型是否为 float32。如果不是，就会抛出一个错误。 assertinfastllm(input.datatype == datatype::float32, embedding's input's type should be float32.); // 这行代码将 weight 的 weighttype 属性设置为 embedding。 weight.weighttype = weighttype::embedding; // 这行代码从 weight 的维度中提取词汇大小（vocabsize）和嵌入大小（embsize）。 int vocabsize = weight.dims[0], embsize = weight.dims[1]; // 这两行代码将 embsize 添加到 input 的维度中，形成一个新的维度。 std::vector 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);#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 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 __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 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 futures; for (int i = 0; i < threadnum - 1; i++) { int end = cur + per + (cur + per * (threadnum - i) 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 (int i = 0; i 0 ? (float *) bias.cpudata : nullptr;#ifdef __arm_feature_fp16_vector_arithmetic uint16_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;#endif int threadnum = getthreads(); int per = k / threadnum; int cur = 0; auto pool = getpool(); std::vector futures; for (int i = 0; i < threadnum - 1; i++) { int end = cur + per + (cur + per * (threadnum - i) 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 0 ? (float *) bias.cpudata : nullptr; weight.calcweightsum(); // 之后，代码创建一个std::vector对象， // lowbitconfig是一个用于存储数据量化信息的类，包括最小值、最大值、位宽和零点。 // 这些信息是通过遍历输入数据获得的。 std::vector 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对象uinput，并将其大小设置为输入数据的大小（n * m）。 // uinput中的每个元素都是输入数据元素经过inputconfigs中对应配置信息量化后的结果。 // 注意这里的量化过程可能会根据是否定义了__avx2__进行不同的处理。 std::vector uinput; uinput.resize(n * m); for (int i = 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 0 ? (float *) bias.cpudata : nullptr; weight.calcweightsum(); std::vector 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 uinput; uinput.resize(n * m); for (int i = 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] 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::vector futures; for (int i = 0; i < threadnum; i++) { int end = cur + per + (cur + per * (threadnum - i) 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_dotprod 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 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 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), persistent=false) # 计算emb的正弦值，然后将结果保存到模型的缓存中。 self.register_buffer(sin_cached, emb.sin()[none, none, :, :].to(dtype), persistent=false) # 这是模型的前向传播方法，接收两个参数：x（输入数据）和seq_len（序列长度）。 def forward(self, x, seq_len=none): # 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_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), )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_embed
cudallamarotateposition2dop对应的就是上面的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：位置编码的数据。// sindata，cosdata：用于旋转的 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); // 计算旋转操作需要的一些参数，包括 outer，spatial，bs，len，n 和 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 上执行位置编码的旋转操作。 // <> 是 cuda 中定义并行线程块和线程的语法， // outer * n 是线程块的数量，min(rotarydim, m / 2) 是每个线程块中的线程数量。 // 核函数的参数包括之前准备的数据和一些计算参数。 fastllmllamarotateposition2dkernel <> (cudadata, cudapositionids, cudasin, cudacos, len, bs, spatial, n, m, (int)positionids.dims.back(), (int)sindata.dims[1], rotarydim); // 使用 fastllmcudafinishinput 函数释放 positionids，sindata 和 cosdata 在 gpu 上的内存。 // 这些数据在这个函数中不再需要。 fastllmcudafinishinput(positionids, cudapositionids); fastllmcudafinishinput(sindata, cudasin); fastllmcudafinishinput(cosdata, cudacos); // 使用 fastllmcudafinishoutput 函数将旋转后的数据从 gpu 复制回 cpu。 // 这个函数也会释放 data 在 gpu 上的内存。 fastllmcudafinishoutput(data, cudadata); 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 个值，这里 k=1，所以是找到最大值。 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 1e-6) { for (int id : 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 (int i = 0; i first; // 定义一个向量 ps，用于存储处理后的概率。 std::vector 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 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 rnd || i == topk - 1) { return v[i].second; } } // 如果以上步骤都没有返回，那么返回 -1。 return -1; }
llmsampling实现了一种基于温度和惩罚的采样策略，用于从给定的 logits 中选择一个 id。这种采样的方法可以控制输出文本的多样性。
0x6. 总结
接着 大模型部署框架 fastllm 简要解析 这篇文章首先梳理了一下fastllm的调用链和关键的数据结构，然后解析了 fastllm 的一些实现细节和cpu/gpu后端实现采用的优化技巧。
关于机器视觉与芯片之间的联系和分析
高通和中兴通讯利用5G网络实现了语音通话
区块链加密算法RSA加密的原理解析
荣耀路由Pro2搭载自主研发的凌霄芯片
石英玻璃纤维可以极大延长数据传输距离
大模型部署框架FastLLM实现细节解析
低功耗蓝牙模块实现低功耗的原理是怎样的
怎么用摇表测量确定三相电机的好坏
PCB线路板有哪一些散热方式
Verizon 45W USB-PD快充充电器采用英飞凌 EZ-PD™ PAG1 AC-DC电源解决方案
5G将会如何影响物联网的安全
一看便知：晶振如何匹配电容
借助磁控微流控芯片,建立埃博拉病毒核酸适配体的高效筛选平台
可编程任意电源的功能说明
发光二极管工作电压和电流是多少
技术新动向微生物细菌驱动的电池！
西门子PLC通信不上经验分享
三星A90 5G版曝光 预计售价4500元左右
万用表和钳形表检测电流的原理及区别
如何保证LED显示屏的质量并延长其使用寿命？
接着 大模型部署框架 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
class data { public: bool lockincpu = false; // 如果lock在cpu上，那么不允许移动到其余设备 weighttype weighttype = weighttype::none; // 权重类型，none代表非权重（或未知权重） datatype datatype = datatype::float32; // 数据类型 int unitsize, unitsizediv = 1; // 单个元素的字节数 = unitsize / unitsizediv std::vector dims; // 数据形状 std::vector strides; // 跨度 uint64_t expansionsize = 0; // 扩容后的尺寸 uint64_t expansionbytes = 0; // 扩容后的字节数 std::vector expansiondims; // 预扩容的形状 uint8_t *cpudata = nullptr; // 数据指针 void *cudadata = nullptr; std::vector extracudadata; void *devicedata = nullptr; std::vector extradevicedata; datadevice datadevice = datadevice::cpu; // 这两个参数用于量化，对float数据不适用 int perchannelaxis = -1; // 沿哪个轴分通道量化，-1代表没有分通道 std::vector perchannelsconfigs; // perchannelsconfigs[i]代表第i个通道的min, max; 如果没有分通道，perchannelsconfigs[0]代表全局min, max std::vector scales, mins; std::vector zeros; std::vector weightsum; // 作为权重时，有时候需要存一些和加速计算 std::string filename; long long filepos; std::shared_ptr m_file; data () {}; data (datatype type); data (datatype type, const std::vector &dims); // 构造函数 // 构造函数，创建好之后从data复制数据 // data中是原始数据，如果type不是float那么需要量化 data (datatype type, const std::vector &dims, const std::vector &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 &dims); // 预扩容到相应尺寸 void mallocspace(uint64_t size); // 在设备上分配 void freespace(); // 回收设备上的内存 void updateunitsize(); // 更新unitsize void resize(const std::vector &dims); // 更改尺寸 void reshape(const std::vector &dims); // 更改尺寸,但不修改数据 uint64_t count(int i) const; // dims[i] * strides[i] void printshape() const; // 输出形状 void print() const; // 输出 void calcweightsum(); // 计算weightsum void todevice(datadevice device); // 移动到指定device void todevice(void *device); void set_file(std::shared_ptr 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::vector &axis, data &output) { data axisdata = data(datatype::int32param, {(int)axis.size()}); axisdata.allocate(); for (int i = 0; i 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的定义和实现：
namespace fastllm { class executor { private: std::vector devices; std::map profiler; public: executor (); // 创建默认的executor ~executor(); // 析构 void cleardevices(); // 清空 devices void adddevice(basedevice *device); // 增加一个device // 运行一个op void 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());#endif this->devices.push_back((basedevice*) new cpudevice()); } executor::~executor() { // 释放 devices 向量中的每个指针元素所占用的内存。 for (int i = 0; i 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::now(); // 创建一个布尔变量 lockincpu，用于记录是否将数据锁定在 cpu 上。 bool lockincpu = 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()) { int batch = intparams.find(it.first + ___batch)->second; for (int i = 0; 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 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::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, it.first.c_str(), it.second); sum += it.second; } printf(total spend %f, 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 next; trienode(); }; trienode *root; std::unordered_map tokentostringdict; tokenizer (); ~tokenizer(); void clear(); // 清空分词器 void insert(const std::string &s, int tokenid); // 插入一个token data encode(const std::string &s); // 编码 std::string decode(const data &data); // 解码 std::string decodetokens(const std::vector &tokens); // 解码 };
我们从实现来看tokenizer的细节：
// 这是 tokenizer 类的嵌套结构 trienode 的构造函数的实现。 // 在构造函数中，将 tokenid 成员变量的值初始化为 -999999。 // 这个值在构造函数中被硬编码，它是作为一个特殊标记来使用的。 tokenizer::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 q; // 将字典树的根节点 root 加入 q 向量，作为遍历的起始点。 q.push_back(root); // 开始遍历 q 向量中的节点，这是一个广度优先搜索（bfs）的过程。 for (int i = 0; i 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 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 v; // 开始遍历输入的字符串 s 中的每个字符。 for (int i = 0; i root; // 从当前字符 s[i] 开始继续遍历字符串 s。 for (int j = i; j 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 值。 return data (datatype::float32, {1, (int)v.size()}, v); } // 这是 tokenizer 类的成员函数 decodetokens 的定义， // 用于对输入的 token 数组进行解码，将 token 转换回原始的字符串。 std::string tokenizer::decodetokens(const std::vector &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 是类似 格式的 token（其中 hh 表示十六进制数）， // 则需要将其转换为对应的字符。首先，提取 hh，然后将其转换为对应的字符， // 并用空格代替原始的 token。 if (s.size() == 6 && s.substr(0, 3) == ') { int c = 0; for (int i = 3; i = '0' && s[i] <= '9') { c += (s[i] - '0'); } else { c += (s[i] - 'a' + 10); } } s = ; s[0] = c; } // 根据不同的 token 进行解码： if (s == ) { ret += ; } else if (s == ) { ret += ; } else { ret += s; } } // 将特殊字符 xe2x96x81（utf-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; } // 检查是否有 格式的特殊 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 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 的值为 helloworld!，那么在解码时，pos 将是 8，而 space_num 将是 4。然后，函数将返回 ，即包含四个空格字符的字符串。在这种情况下，特殊 token 被成功解码成了四个空格字符，但是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 中的每个元素都是 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 second); data &output = *(datas.find(output)->second); data &weight = *(datas.find(weight)->second); // 这行代码检查 weight 的维度数量是否为 2。如果不是，就会抛出一个错误。 assertinfastllm(weight.dims.size() == 2, embedding's weight's dim should be 2.); // 这行代码检查 weight 的数据类型是否为 float32 或 bfloat16。如果不是，就会抛出一个错误。 assertinfastllm(weight.datatype == datatype::float32 || weight.datatype == datatype::bfloat16, embedding's weight's type should be float32 or bfloat16.); // 这行代码检查 input 的数据类型是否为 float32。如果不是，就会抛出一个错误。 assertinfastllm(input.datatype == datatype::float32, embedding's input's type should be float32.); // 这行代码将 weight 的 weighttype 属性设置为 embedding。 weight.weighttype = weighttype::embedding; // 这行代码从 weight 的维度中提取词汇大小（vocabsize）和嵌入大小（embsize）。 int vocabsize = weight.dims[0], embsize = weight.dims[1]; // 这两行代码将 embsize 添加到 input 的维度中，形成一个新的维度。 std::vector 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);#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 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 __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 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 futures; for (int i = 0; i < threadnum - 1; i++) { int end = cur + per + (cur + per * (threadnum - i) 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 (int i = 0; i 0 ? (float *) bias.cpudata : nullptr;#ifdef __arm_feature_fp16_vector_arithmetic uint16_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;#endif int threadnum = getthreads(); int per = k / threadnum; int cur = 0; auto pool = getpool(); std::vector futures; for (int i = 0; i < threadnum - 1; i++) { int end = cur + per + (cur + per * (threadnum - i) 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 0 ? (float *) bias.cpudata : nullptr; weight.calcweightsum(); // 之后，代码创建一个std::vector对象， // lowbitconfig是一个用于存储数据量化信息的类，包括最小值、最大值、位宽和零点。 // 这些信息是通过遍历输入数据获得的。 std::vector 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对象uinput，并将其大小设置为输入数据的大小（n * m）。 // uinput中的每个元素都是输入数据元素经过inputconfigs中对应配置信息量化后的结果。 // 注意这里的量化过程可能会根据是否定义了__avx2__进行不同的处理。 std::vector uinput; uinput.resize(n * m); for (int i = 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 0 ? (float *) bias.cpudata : nullptr; weight.calcweightsum(); std::vector 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 uinput; uinput.resize(n * m); for (int i = 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] 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::vector futures; for (int i = 0; i < threadnum; i++) { int end = cur + per + (cur + per * (threadnum - i) 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_dotprod 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 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 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), persistent=false) # 计算emb的正弦值，然后将结果保存到模型的缓存中。 self.register_buffer(sin_cached, emb.sin()[none, none, :, :].to(dtype), persistent=false) # 这是模型的前向传播方法，接收两个参数：x（输入数据）和seq_len（序列长度）。 def forward(self, x, seq_len=none): # 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_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), )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_embed
cudallamarotateposition2dop对应的就是上面的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：位置编码的数据。// sindata，cosdata：用于旋转的 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); // 计算旋转操作需要的一些参数，包括 outer，spatial，bs，len，n 和 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 上执行位置编码的旋转操作。 // <> 是 cuda 中定义并行线程块和线程的语法， // outer * n 是线程块的数量，min(rotarydim, m / 2) 是每个线程块中的线程数量。 // 核函数的参数包括之前准备的数据和一些计算参数。 fastllmllamarotateposition2dkernel <> (cudadata, cudapositionids, cudasin, cudacos, len, bs, spatial, n, m, (int)positionids.dims.back(), (int)sindata.dims[1], rotarydim); // 使用 fastllmcudafinishinput 函数释放 positionids，sindata 和 cosdata 在 gpu 上的内存。 // 这些数据在这个函数中不再需要。 fastllmcudafinishinput(positionids, cudapositionids); fastllmcudafinishinput(sindata, cudasin); fastllmcudafinishinput(cosdata, cudacos); // 使用 fastllmcudafinishoutput 函数将旋转后的数据从 gpu 复制回 cpu。 // 这个函数也会释放 data 在 gpu 上的内存。 fastllmcudafinishoutput(data, cudadata); 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 个值，这里 k=1，所以是找到最大值。 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 1e-6) { for (int id : 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 (int i = 0; i first; // 定义一个向量 ps，用于存储处理后的概率。 std::vector 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 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 rnd || i == topk - 1) { return v[i].second; } } // 如果以上步骤都没有返回，那么返回 -1。 return -1; }
llmsampling实现了一种基于温度和惩罚的采样策略，用于从给定的 logits 中选择一个 id。这种采样的方法可以控制输出文本的多样性。
0x6. 总结
接着 大模型部署框架 fastllm 简要解析 这篇文章首先梳理了一下fastllm的调用链和关键的数据结构，然后解析了 fastllm 的一些实现细节和cpu/gpu后端实现采用的优化技巧。
关于机器视觉与芯片之间的联系和分析
高通和中兴通讯利用5G网络实现了语音通话
区块链加密算法RSA加密的原理解析
荣耀路由Pro2搭载自主研发的凌霄芯片
石英玻璃纤维可以极大延长数据传输距离
大模型部署框架FastLLM实现细节解析
低功耗蓝牙模块实现低功耗的原理是怎样的
怎么用摇表测量确定三相电机的好坏
PCB线路板有哪一些散热方式
Verizon 45W USB-PD快充充电器采用英飞凌 EZ-PD™ PAG1 AC-DC电源解决方案
5G将会如何影响物联网的安全
一看便知：晶振如何匹配电容
借助磁控微流控芯片,建立埃博拉病毒核酸适配体的高效筛选平台
可编程任意电源的功能说明
发光二极管工作电压和电流是多少
技术新动向微生物细菌驱动的电池！
西门子PLC通信不上经验分享
三星A90 5G版曝光 预计售价4500元左右
万用表和钳形表检测电流的原理及区别
如何保证LED显示屏的质量并延长其使用寿命？