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
# （最大的位置嵌入长度，默认为2048），base（基数，默认为10000）和device（设备类型，例如CPU或GPU）。
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
super().__init__()
self.dim = dim # 将输入的dim参数保存到self.dim属性中。
# # 将输入的max_position_embeddings参数保存到self.max_position_embeddings属性中。
self.max_position_embeddings = max_position_embeddings
# 将输入的base参数保存到self.base属性中。
self.base = base
# 计算逆频率并保存到变量inv_freq中。逆频率是一种用于位置编码的技巧，
# 它可以帮助模型更好地捕捉位置信息。
inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
# 将inv_freq保存到模型的缓存中。register_buffer是PyTorch nn.Module的一个方法，
# 它用于保存一些不需要计算梯度的变量。
self.register_buffer("inv_freq", inv_freq, persistent=False)
# Build here to make `torch.jit.trace` work.
# 调用_set_cos_sin_cache方法，预先计算并保存正弦和余弦的缓存值。
self._set_cos_sin_cache(
seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
)
# 这是一个私有方法，接收三个参数：seq_len（序列长度），device（设备类型）和dtype（数据类型）
def _set_cos_sin_cache(self, seq_len, device, dtype):
# 将输入的seq_len参数保存到self.max_seq_len_cached属性中。
self.max_seq_len_cached = seq_len
# 生成一个长度为max_seq_len_cached的序列，并保存到变量t中。
t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
# 使用外积计算频率和t的乘积，结果保存到变量freqs中。
freqs = torch.einsum("i,j->ij", t, self.inv_freq)
# 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 上执行位置编码的旋转操作。
// <<<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 函数释放 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/fas... 这里的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/fas... 。
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 < 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/fas... 。
// 这段代码是一个用于从给定的 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) {
base[id] = (base[id] < 0 ? base[id] * config.repeat_penalty : base[id] / config.repeat_penalty);
}
}
// 计算温度的倒数 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 和 maxValue，maxValue 是 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后端实现采用的优化技巧。
作者：BBuf
文章来源：GiantPandaCV
推荐阅读
- 大模型部署框架 FastLLM 简要解析
- 中科大联合上海 AI Lab 发布 FreeDrag: 无需点跟踪即可稳定拖动语义内容!
- 编译入门那些事儿（1）：LLVM中的Pass和PassManager
- 冠绝榜单 | 百度联合上海AI实验室提出 CityTrack: 刷新城市规模多目标跟踪纪录！
- 华为诺亚实验室 | AIGC时代的ImageNet，百万生成图片助力AI生成图片检测器研发
更多嵌入式AI干货请关注嵌入式AI专栏。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
