阅
[InferLLM大模型推理框架项目](26)kern中ARM优化模块的kernel代码分析(src/kern/optimized/arm/kernel_gpu.h+.cpp)
InferLLM 框架中 ARM 优化模块的 kernel.h 和 kernel.cpp 分析
InferLLM 框架中的 ARM 优化模块主要通过 kernel.h 和 kernel.cpp 文件实现了针对 ARM 架构的优化计算函数。这两个文件共同构成了 ARM 平台上大语言模型推理的核心计算部分。
1. kernel.h 文件分析
kernel.h 文件主要声明了各种计算函数和注册内核的宏。
1.1 头文件包含
#pragma once
#include "kern/naive/naive.h"
#include "math.h"
#include "string.h"
这里包含了必要的头文件，其中 "kern/naive/naive.h" 包含了朴素实现的函数，用于在 ARM 优化不可用时作为回退方案。
1.2 函数声明
kernel.h 文件声明了一系列计算函数，这些函数都返回 TaskSet 类型，用于多线程并行计算：
namespace inferllm {
namespace opt {
// 嵌入层计算
TaskSet llm_embedding_get_int4_float(
const void* weights, const uint32_t* index, float* dst, uint32_t len_seq,
uint32_t embd);
// 元素级计算
TaskSet llm_elemwise_compute_float(
InData<float> srcs, float* dst, size_t len, ElemMode mode);
// 广播计算
TaskSet llm_elemwise_broadcast_dim0_src1_compute_float(
const float* src0, const float* src1, float* dst, uint32_t len0, uint32_t len1,
ElemMode mode);
// RMS 归一化
TaskSet llm_rms_norm_compute_float(
const float* src, float* dst, uint32_t seq_len, uint32_t embd, float eps);
// Softmax 计算
TaskSet llm_softmax_compute_float(
const float* src, float* dst, uint32_t len_row, uint32_t col);
// 量化矩阵乘法
TaskSet llm_matmul_compute_int4_float(
float* dst, const void* src0, const float* bias, const float* src1, uint32_t M,
uint32_t N, uint32_t K, void* workspace, uint32_t size);
// 打包的量化矩阵乘法
TaskSet llm_matmul_compute_int4_float_packed(
float* dst, const void* src0, const float* bias, const float* src1, uint32_t M,
uint32_t N, uint32_t K, void* workspace, uint32_t size);
// 获取矩阵乘法所需的工作空间大小
size_t llm_matmul_get_workspace_float(
uint32_t nr_thread, uint32_t M, uint32_t N, uint32_t K);
// 多头注意力中的 Q 和 K 的矩阵乘法
TaskSet llm_matmul_compute_with_head_stride_float(
float* dst, const float* srck, const float* srcq, uint32_t seqlen,
uint32_t embd, uint32_t head, uint32_t nr_past);
// 多头注意力中的 QK 和 V 的矩阵乘法
TaskSet llm_head_batched_matmul_compute_float(
float* dst, const float* v, const float* qk, uint32_t seqlen, uint32_t embd,
uint32_t head, uint32_t nr_past);
1.3 内核注册
kernel.h 文件使用 PartialImplementKernel 和 PartialImplementSpace 宏注册内核：
// 注册计算函数
PartialImplementKernel(ElemwiseFloat, llm_elemwise_compute_float);
PartialImplementKernel(
ElemwiseBroadcastDim0Src1Float, llm_elemwise_broadcast_dim0_src1_compute_float);
PartialImplementKernel(RmsNormFloat, llm_rms_norm_compute_float);
PartialImplementKernel(EmbeddingGetInt4Float, llm_embedding_get_int4_float);
PartialImplementKernel(MatmulInt4Float, llm_matmul_compute_int4_float);
PartialImplementKernel(
MatmulWithHeadStrideFloat, llm_matmul_compute_with_head_stride_float);
PartialImplementKernel(HeadBatchedMatmulFloat, llm_head_batched_matmul_compute_float);
// 注册工作空间计算函数
PartialImplementSpace(MatmulInt4Float, llm_matmul_get_workspace_float);
这些宏将内核 ID 与具体的实现函数关联起来，使得框架可以在运行时根据内核 ID 选择合适的实现。
2. kernel.cpp 文件分析
kernel.cpp 文件实现了 kernel.h 中声明的各种计算函数。
2.1 头文件包含
#include <assert.h>
#include "math.h"
#include "string.h"
#include "utils.h"
#include "core/tensor.h"
#include "kernel.h"
#include "optimized.h"
#include "quantize.h"
这里包含了必要的头文件，其中 "optimized.h" 包含了基础向量运算函数，"quantize.h" 包含了量化和反量化操作。
2.2 嵌入层计算
TaskSet llm_embedding_get_int4_float(
const void* weights, const uint32_t* index, float* dst, uint32_t len_seq,
uint32_t embd) {
auto task = [=](const TaskId& id) {
for (uint32_t i = id.start; i < id.end; ++i) {
const int row = index[i];
const int weight_stride =
embd * dtype_in_byte(DType::Int4) / dtype_block_size(DType::Int4);
dequantize_row_q4_0(
(static_cast<const char*>(weights) + row * weight_stride),
dst + i * embd, embd);
}
};
return TaskSet{{task, len_seq}};
}
这个函数实现了嵌入层的计算，将4位整数量化的权重反量化为浮点数。每个任务处理一个或多个序列位置，通过 dequantize_row_q4_0 函数将量化权重反量化为浮点数。
2.3 元素级计算
TaskSet llm_elemwise_compute_float(
InData<float> srcs, float* dst, size_t length, ElemMode mode) {
MultiThreadingTask task;
switch (mode) {
case ElemMode::Add: {
task = [=](const TaskId& id) {
uint32_t offset = id.start;
uint32_t len = id.end - id.start;
elemwise_vector_add(
len, srcs[0] + offset, srcs[1] + offset, dst + offset);
};
break;
}
case ElemMode::Mul: {
task = [=](const TaskId& id) {
uint32_t offset = id.start;
uint32_t len = id.end - id.start;
elemwise_vector_mul(
len, srcs[0] + offset, srcs[1] + offset, dst + offset);
};
break;
}
case ElemMode::Silu: {
task = [=](const TaskId& id) {
uint32_t offset = id.start;
uint32_t len = id.end - id.start;
return elemwise_vector_silu(len, srcs[0] + offset, dst + offset);
};
break;
}
case ElemMode::Gelu: {
task = [=](const TaskId& id) {
uint32_t offset = id.start;
uint32_t len = id.end - id.start;
return elemwise_vector_gelu(len, srcs[0] + offset, dst + offset);
};
break;
}
default:
INFER_ASSERT(0, "Not supported.");
}
return TaskSet{{task, length}};
}
这个函数实现了元素级计算，支持加法、乘法、SiLU 激活函数和 GELU 激活函数。每个任务处理一段连续的数据，通过 elemwise_vector_add、elemwise_vector_mul、elemwise_vector_silu 和 elemwise_vector_gelu 函数实现具体的计算。
2.4 广播计算
TaskSet llm_elemwise_broadcast_dim0_src1_compute_float(
const float* src0, const float* src1, float* dst, uint32_t len0, uint32_t len1,
ElemMode mode) {
MultiThreadingTask task;
switch (mode) {
case ElemMode::Add: {
task = [=](const TaskId& id) {
for (size_t i = id.start; i < id.end; i++) {
const float* p_src = src0 + i * len1;
float* p_dst = dst + i * len1;
elemwise_vector_add(len1, p_src, src1, p_dst);
}
};
break;
}
case ElemMode::Mul: {
task = [=](const TaskId& id) {
for (size_t i = id.start; i < id.end; i++) {
auto p_src = src0 + i * len1;
auto p_dst = dst + i * len1;
elemwise_vector_mul(len1, p_src, src1, p_dst);
}
};
break;
}
default:
INFER_ASSERT(0, "Not supported.");
}
return TaskSet{{task, len0}};
}
这个函数实现了广播计算，将 src1 广播到 src0 的第0维，然后进行元素级计算。每个任务处理 src0 的一个或多个行，通过 elemwise_vector_add 和 elemwise_vector_mul 函数实现具体的计算。
2.5 RMS 归一化
TaskSet llm_rms_norm_compute_float(
const float* src, float* dst, uint32_t seq_len, uint32_t embd, float eps) {
auto task = [=](const TaskId& id) {
for (uint32_t i = id.start; i < id.end; i++) {
const float* row = src + i * embd;
float* out = dst + i * embd;
float mean = reduce_square_sum(embd, row) / embd;
const float scale = 1.0 / sqrt(mean + eps);
elemwise_vec_scale(embd, row, scale, out);
}
};
return TaskSet{{task, seq_len}};
}
这个函数实现了 RMS 归一化，每个任务处理一个或多个序列位置。首先计算平方和的均值，然后计算缩放因子，最后将输入向量乘以缩放因子得到归一化结果。
2.6 Softmax 计算
TaskSet llm_softmax_compute_float(
const float* src, float* dst, uint32_t len_row, uint32_t col) {
auto task = [=](const TaskId& id) {
for (uint32_t row = id.start; row < id.end; row++) {
const float* psrc = src + row * col;
float* pdst = dst + row * col;
float max = reduce_max(col, psrc);
float sum = select_sub_max_and_reduce_sum(col, psrc, pdst, max);
sum = 1.0 / sum;
elemwise_vec_scale(col, pdst, sum, pdst);
}
};
return TaskSet{{task, len_row}};
}
这个函数实现了 Softmax 计算，每个任务处理一个或多个行。首先找到最大值，然后减去最大值并计算指数和，最后将每个元素除以指数和得到 Softmax 结果。
2.7 量化矩阵乘法
TaskSet llm_matmul_compute_int4_float(
float* dst, const void* src0, const float* bias, const float* src1, uint32_t M,
uint32_t N, uint32_t K, void* workspace, uint32_t size) {
// src0 是量化权重，每32个数据为一个块，一个块共享相同的缩放因子
// src1 是特征图，src0 布局为 {N, K}，src1 布局为 {M, K}，dst 布局为 {M, N}
INFER_ASSERT(sizeof(float) * K <= size, "workspace is not enough.");
uint32_t weight_q40_stride =
K * dtype_in_byte(DType::Int4) / dtype_block_size(DType::Int4);
uint32_t weight_q80_stride =
K * dtype_in_byte(DType::Int8) / dtype_block_size(DType::Int8);
// 第一阶段：量化输入，并存储在工作空间中
// 因为输入比权重小，量化输入可以减少内存流量
auto task1 = [=](const TaskId& id) {
for (uint32_t m = id.start; m < id.end; m++) {
BlockQ80* q_src1 = (BlockQ80*)(static_cast<uint8_t*>(workspace) +
m * weight_q80_stride);
quantize_row_q8_0(src1 + m * K, q_src1, K);
}
};
// 第二阶段：计算矩阵乘法
int8_t* q_src = static_cast<int8_t*>(workspace);
auto task2 = [=](const TaskId& id) {
uint32_t N_len = id.end - id.start;
uint32_t n_block_4 = N_len / 4;
uint32_t n_block_4_left = N_len - n_block_4 * 4;
// 处理4列一组的部分
for (uint32_t block4 = 0; block4 < n_block_4; block4++) {
uint32_t n = block4 * 4 + id.start;
float b0 = bias ? bias[n] : 0;
float b1 = bias ? bias[n + 1] : 0;
float b2 = bias ? bias[n + 2] : 0;
float b3 = bias ? bias[n + 3] : 0;
// 加载权重
const void* q_weight0 = static_cast<const char*>(src0) + n * weight_q40_stride;
const void* q_weight1 = static_cast<const char*>(src0) + (n + 1) * weight_q40_stride;
const void* q_weight2 = static_cast<const char*>(src0) + (n + 2) * weight_q40_stride;
const void* q_weight3 = static_cast<const char*>(src0) + (n + 3) * weight_q40_stride;
// 计算矩阵乘法
for (uint32_t m = 0; m < M; m++) {
int8_t* src = q_src + m * weight_q80_stride;
dst[m * N + n] = vec_vec_dot_q40_with_q80(K, q_weight0, src) + b0;
dst[m * N + n + 1] = vec_vec_dot_q40_with_q80(K, q_weight1, src) + b1;
dst[m * N + n + 2] = vec_vec_dot_q40_with_q80(K, q_weight2, src) + b2;
dst[m * N + n + 3] = vec_vec_dot_q40_with_q80(K, q_weight3, src) + b3;
}
}
// 处理剩余列
for (uint32_t left = 0; left < n_block_4_left; left++) {
uint32_t n = n_block_4 * 4 + left + id.start;
float b0 = bias ? bias[n] : 0;
const void* q_weight = static_cast<const char*>(src0) + n * weight_q40_stride;
for (uint32_t m = 0; m < M; m++) {
int8_t* src = q_src + m * weight_q80_stride;
dst[m * N + n] = vec_vec_dot_q40_with_q80(K, q_weight, src) + b0;
}
}
};
return TaskSet{{task1, M}, {task2, N}};
}
这个函数实现了4位整数权重与浮点数激活值的矩阵乘法。它分为两个阶段：第一阶段将输入量化为8位整数，第二阶段计算矩阵乘法。每个阶段都使用多线程并行计算，第一阶段按行分解，第二阶段按列分解。
2.8 打包的量化矩阵乘法
TaskSet llm_matmul_compute_int4_float_packed(
float* dst, const void* src0, const float* bias, const float* src1, uint32_t M,
uint32_t N, uint32_t K, void* workspace, uint32_t size) {
// 与 llm_matmul_compute_int4_float 类似，但针对打包的权重进行了优化
// 打包的权重是指将多个权重矩阵打包在一起，以提高缓存命中率
// ... 实现细节 ...
}
这个函数实现了打包的量化矩阵乘法，针对打包的权重进行了优化，以提高缓存命中率。
2.9 获取矩阵乘法所需的工作空间大小
size_t llm_matmul_get_workspace_float(
uint32_t nr_thread, uint32_t M, uint32_t N, uint32_t K) {
// 计算矩阵乘法所需的工作空间大小
// 工作空间用于存储量化的输入
uint32_t weight_q80_stride =
K * dtype_in_byte(DType::Int8) / dtype_block_size(DType::Int8);
return weight_q80_stride * M;
}
这个函数计算矩阵乘法所需的工作空间大小，工作空间用于存储量化的输入。
2.10 多头注意力计算
TaskSet llm_matmul_compute_with_head_stride_float(
float* dst, const float* srck, const float* srcq, uint32_t seqlen,
uint32_t embd, uint32_t head, uint32_t nr_past) {
uint32_t sub_embd = embd / head;
uint32_t length = nr_past + seqlen;
uint32_t line_stride = embd;
auto task = [=](const TaskId& id) {
for (uint32_t h = id.start; h < id.end; h++) {
auto dst_head = dst + h * seqlen * (nr_past + seqlen);
auto srck_head = srck + h * sub_embd;
auto srcq_head = srcq + h * sub_embd;
compute_src_offset_embd_matmul(
srcq_head, embd, srck_head, embd, dst_head, seqlen, length,
sub_embd);
}
};
return TaskSet{{task, head}};
}
TaskSet llm_head_batched_matmul_compute_float(
float* dst, const float* v, const float* qk, uint32_t seqlen, uint32_t embd,
uint32_t head, uint32_t nr_past) {
uint32_t sub_embd = embd / head;
uint32_t length = nr_past + seqlen;
uint32_t line_stride = embd;
auto task = [=](const TaskId& id) {
for (uint32_t h = id.start; h < id.end; h++) {
float* dst_head = dst + h * sub_embd;
const float* v_head = v + h * sub_embd;
const float* qk_head = qk + h * seqlen * length;
comput_matmul_with_dst_uncontinue(
dst_head, embd, v_head, embd, qk_head, seqlen, length, sub_embd);
}
};
return TaskSet{{task, head}};
}
这两个函数实现了多头注意力计算中的矩阵乘法。llm_matmul_compute_with_head_stride_float 计算 Q 和 K 的矩阵乘法，llm_head_batched_matmul_compute_float 计算 QK 和 V 的矩阵乘法。每个任务处理一个注意力头，通过 compute_src_offset_embd_matmul 和 comput_matmul_with_dst_uncontinue 函数实现具体的计算。
3. 内核注册机制分析
InferLLM 框架使用内核注册机制，将函数与内核 ID 关联起来。在 kernel.h 文件中，使用 PartialImplementKernel 和 PartialImplementSpace 宏注册内核：
PartialImplementKernel(ElemwiseFloat, llm_elemwise_compute_float);
PartialImplementKernel(
ElemwiseBroadcastDim0Src1Float, llm_elemwise_broadcast_dim0_src1_compute_float);
PartialImplementKernel(RmsNormFloat, llm_rms_norm_compute_float);
PartialImplementKernel(EmbeddingGetInt4Float, llm_embedding_get_int4_float);
PartialImplementKernel(MatmulInt4Float, llm_matmul_compute_int4_float);
PartialImplementKernel(
MatmulWithHeadStrideFloat, llm_matmul_compute_with_head_stride_float);
PartialImplementKernel(HeadBatchedMatmulFloat, llm_head_batched_matmul_compute_float);
PartialImplementSpace(MatmulInt4Float, llm_matmul_get_workspace_float);
这些宏的定义可能在其他头文件中，它们的作用是将内核 ID 与具体的实现函数关联起来，使得框架可以在运行时根据内核 ID 选择合适的实现。
与 ImplementKernel 和 ImplementSpace 宏不同，PartialImplementKernel 和 PartialImplementSpace 宏可能表示这些实现是部分实现，可能只支持某些特定的参数组合或者只在某些特定的条件下可用。
4. 多线程并行策略分析
InferLLM 框架使用 TaskSet 实现多线程并行，每个计算函数都返回一个 TaskSet，包含一个或多个任务及其子任务数量：
TaskSet llm_rms_norm_compute_float(
const float* src, float* dst, uint32_t seq_len, uint32_t embd, float eps) {
auto task = [=](const TaskId& id) {
for (uint32_t i = id.start; i < id.end; i++) {
// ... 计算逻辑 ...
}
};
return TaskSet{{task, seq_len}};
}
不同的计算函数使用不同的任务分解策略：
- 按序列长度分解：如
llm_rms_norm_compute_float，每个任务处理一个或多个序列位置 - 按头数分解：如
llm_matmul_compute_with_head_stride_float，每个任务处理一个或多个注意力头 - 按矩阵行列分解：如
llm_matmul_compute_int4_float，使用两个任务集，一个按行分解，一个按列分解
这种设计使得计算任务可以在多线程环境中高效执行，通过将大型计算任务分解为多个子任务，并分配给不同的线程处理。
5. 优化策略分析
5.1 NEON 指令集优化
ARM 优化模块使用 NEON 指令集进行向量化计算，提高计算效率。这些优化主要在 optimized.h 和 quantize.h 文件中实现，kernel.cpp 文件中的函数调用这些优化的基础向量运算函数。
5.2 分块处理策略
ARM 优化模块使用分块处理策略，将数据分成多个块进行处理：
// 矩阵乘法中的分块处理
auto task2 = [=](const TaskId& id) {
uint32_t N_len = id.end - id.start;
uint32_t n_block_4 = N_len / 4;
uint32_t n_block_4_left = N_len - n_block_4 * 4;
// 处理4列一组的部分
for (uint32_t block4 = 0; block4 < n_block_4; block4++) {
// ... 处理4列 ...
}
// 处理剩余列
for (uint32_t left = 0; left < n_block_4_left; left++) {
// ... 处理1列 ...
}
};
这种策略可以最大化利用 NEON 指令集的并行性，同时处理所有数据。
5.3 量化计算优化
ARM 优化模块使用量化计算减少内存占用和计算量：
TaskSet llm_matmul_compute_int4_float(
float* dst, const void* src0, const float* bias, const float* src1, uint32_t M,
uint32_t N, uint32_t K, void* workspace, uint32_t size) {
// 第一阶段：量化输入
auto task1 = [=](const TaskId& id) {
for (uint32_t m = id.start; m < id.end; m++) {
BlockQ80* q_src1 = (BlockQ80*)(static_cast<uint8_t*>(workspace) +
m * weight_q80_stride);
quantize_row_q8_0(src1 + m * K, q_src1, K);
}
};
// 第二阶段：使用量化数据计算
auto task2 = [=](const TaskId& id) {
// ... 使用 vec_vec_dot_q40_with_q80 计算点积 ...
};
return TaskSet{{task1, M}, {task2, N}};
}
这种设计减少了内存占用和内存带宽需求，提高了计算效率。特别是对于大型矩阵乘法，量化计算可以显著提高性能。
5.4 内存访问优化
ARM 优化模块优化了内存访问模式，减少内存访问开销：
// 矩阵乘法中的内存访问优化
for (uint32_t block4 = 0; block4 < n_block_4; block4++) {
uint32_t n = block4 * 4 + id.start;
// ... 加载偏置和权重 ...
// 计算矩阵乘法
for (uint32_t m = 0; m < M; m++) {
int8_t* src = q_src + m * weight_q80_stride;
dst[m * N + n] = vec_vec_dot_q40_with_q80(K, q_weight0, src) + b0;
dst[m * N + n + 1] = vec_vec_dot_q40_with_q80(K, q_weight1, src) + b1;
dst[m * N + n + 2] = vec_vec_dot_q40_with_q80(K, q_weight2, src) + b2;
dst[m * N + n + 3] = vec_vec_dot_q40_with_q80(K, q_weight3, src) + b3;
}
}
这种设计提高了缓存命中率，减少了内存访问开销。通过一次加载输入数据，然后计算多个输出元素，可以减少内存带宽需求。
6. 与 naive 实现的比较
ARM 优化模块与 naive 模块的主要区别：
- 向量化实现：ARM 优化模块使用 NEON 指令集进行向量化计算，naive 模块使用标量实现
- 分块处理：ARM 优化模块使用分块处理策略，naive 模块使用简单的循环
- 量化计算：ARM 优化模块使用量化计算减少内存占用和计算量，naive 模块使用浮点计算
- 多线程并行：ARM 优化模块使用 TaskSet 实现多线程并行，naive 模块也使用 TaskSet，但任务分解策略不同
在实际应用中，ARM 优化模块的性能明显优于 naive 模块，特别是在支持 NEON 指令集的 ARM 处理器上。
7. 内核注册机制的实现
InferLLM 框架使用 PartialImplementKernel 和 PartialImplementSpace 宏注册内核。这些宏的定义可能如下：
#define PartialImplementKernel(kernel_id, fun) \
template <> \
struct PartialComp<KernelID::kernel_id, KernelPlatform::Optimized> { \
template <typename... Args> \
static TaskSet exec(Args... args) { \
return fun(std::forward<Args>(args)...); \
} \
};
#define PartialImplementSpace(kernel_id, fun) \
template <> \
struct PartialSpace<KernelID::kernel_id, KernelPlatform::Optimized> { \
template <typename... Args> \
static size_t get(Args... args) { \
return fun(std::forward<Args>(args)...); \
} \
};
这些宏将内核 ID 与具体的实现函数关联起来，使得框架可以在运行时根据内核 ID 选择合适的实现。与 ImplementKernel 和 ImplementSpace 宏不同，PartialImplementKernel 和 PartialImplementSpace 宏表示这些实现是部分实现，可能只支持某些特定的参数组合或者只在某些特定的条件下可用。
在运行时，框架会首先检查是否有完整实现，如果没有，则检查是否有部分实现，如果都没有，则使用 naive 实现作为回退方案。
8. TaskSet 的实现
TaskSet 是 InferLLM 框架中的一个重要概念，用于多线程并行计算。它的定义可能如下：
struct TaskId {
uint32_t start;
uint32_t end;
};
using MultiThreadingTask = std::function<void(const TaskId&)>;
struct TaskItem {
MultiThreadingTask task;
uint32_t nr_task;
};
using TaskSet = std::vector<TaskItem>;
每个 TaskSet 包含一个或多个 TaskItem，每个 TaskItem 包含一个任务函数和子任务数量。在运行时，框架会将每个 TaskItem 分解为多个子任务，并分配给不同的线程处理。
9. 未来优化方向
基于当前实现，可以考虑以下优化方向：
9.1 更多 ARM 指令集支持
- 支持 ARMv8.2-A 的 FP16 指令，使用半精度浮点数进行计算
- 支持 ARMv8.2-A 的 DotProd 指令，加速点积计算
- 支持 ARMv8.6-A 的 BFloat16 指令，使用 BFloat16 进行计算
9.2 更高效的算法
- 使用 Winograd 算法优化矩阵乘法
- 使用 Flash Attention 算法优化注意力计算
- 使用混合精度计算提高性能
9.3 更多量化方法
- 支持 3 位、2 位甚至 1 位量化
- 支持非对称量化
- 支持组量化
9.4 更高级的并行策略
- 使用流水线并行减少内存占用
- 使用张量并行和模型并行处理大型模型
- 使用异步计算提高计算效率
总结
InferLLM 框架中的 ARM 优化模块通过 kernel.h 和 kernel.cpp 文件实现了针对 ARM 架构的优化计算函数。这些函数使用 NEON 指令集进行向量化计算，使用分块处理策略和量化计算减少内存占用和计算量，使用 TaskSet 实现多线程并行，提高大语言模型推理的性能。
与 naive 模块相比，ARM 优化模块的性能明显更高，特别是在支持 NEON 指令集的 ARM 处理器上。未来可以考虑支持更多 ARM 指令集、使用更高效的算法、支持更多量化方法和使用更高级的并行策略等方向进行优化，进一步提高大语言模型在 ARM 平台上的推理性能。
赏
