阅
[InferLLM大模型推理框架项目](03)InferLLM核心模块core(src/core)综合介绍
InferLLM核心模块core(src/core)结构与功能
InferLLM 的 core(src/core) 模块包含了框架的核心组件，负责模型的基础数据结构、内存管理、计算图构建和执行等功能。下面对 core 模块下的主要文件进行分析：
1. 核心数据结构
1.1 Tensor (tensor.h/cpp)
Tensor 是框架的基础数据结构，用于表示多维数组数据：
- 内存管理：支持三种内存来源
- 自己分配的内存
- 外部共享的内存
- 从文件映射的内存（权重）
- 数据类型：支持多种数据类型（Float32, Int4, Int8等）
- 内存状态：通过
TensorState管理（Own, OutSide） - 用户计数：通过引用计数管理内存生命周期
- 数据预处理：支持权重数据的预处理
关键方法：
TensorState prepare_data(); // 准备数据，如果数据不存在则分配或从文件读取
TensorState recall_data(); // 回收数据，释放不需要的内存
void* ptr(); // 获取数据指针
1.2 WorkSpace (tensor.h)
WorkSpace 是一个简单的内存容器，用于临时存储计算过程中的数据：
class WorkSpace {
public:
void* ptr();
template <typename T> T* ptr();
size_t length();
void set_memory(void* data, size_t length);
private:
void* m_data = nullptr;
size_t m_length = 0;
};
2. 设备抽象
2.1 Device (device.h/cpp)
Device 是对计算设备的抽象，提供内存分配和数据传输功能：
- 内存管理：分配和释放设备内存
- 数据传输：在主机和设备之间传输数据
- 内核管理：持有对应设备的计算内核
派生类：
CPUDevice：CPU设备实现GPUDevice：GPU设备实现（条件编译）
关键方法：
virtual void* allocate(size_t len) = 0;
virtual void free_device(void* ptr) = 0;
virtual void host2device_copy(void* device, const void* host, size_t size, bool async = false) = 0;
3. 计算图构建
3.1 Graph (graph.h/cpp)
Graph 是模型的骨架，定义了计算流程和模型结构：
- 模块管理：管理所有计算模块
- 权重管理：管理模型权重
- 执行流程：控制模型的前向计算
- 工作空间：管理计算过程中的临时内存
关键方法：
void execute(std::vector<int32_t> in_token, std::vector<float>& logist, uint32_t nr_past, bool prefill = false);
virtual void construct_llm() = 0; // 构建LLM模型结构，由派生类实现
3.2 OprModuleBase (graph.h)
OprModuleBase 是计算图中的基本单元，封装了一组相关的算子：
- 算子管理：管理一组相关的算子
- 输入输出：管理模块的输入输出张量
- 执行控制：控制模块的执行
派生类：
AttentionModule：注意力模块LlamaFFNModule：LLaMA前馈网络模块GlmFFNModule：GLM前馈网络模块HeadModule：输出头模块EmbdModule：嵌入模块
4. 键值存储
4.1 KvStorage (kvstorage.h/cpp)
KvStorage 继承自 Tensor，专门用于存储和管理键值对：
- 索引管理：管理存储的键值对索引
- 内存管理：管理键值对的内存
- 数据访问：提供对当前键值对的访问
关键方法：
void* get_current_data(); // 获取当前数据指针
size_t add_id(uint32_t id); // 添加索引
void reset_id(); // 重置索引
5. 模型接口
5.1 Model (model.cpp)
Model 提供了对外的模型接口，封装了内部实现：
- 解码接口：提供文本解码功能
- 模型管理：管理模型的生命周期
关键方法：
std::string decode(const std::string& user_input, int& token);
std::string decode_iter(int& token);
6. 线程池
6.1 ThreadPool (thread_pool.h/cpp)
ThreadPool 提供了线程管理和任务调度功能：
- 线程管理：创建和管理工作线程
- 任务调度：分配任务给工作线程
- 同步机制：提供任务同步机制
7. 核心功能流程
-
模型加载：
- 创建
Graph对象 - 构建模型结构 (
construct_llm) - 加载模型权重
- 创建
-
推理执行：
- 准备输入数据
- 执行计算图 (
Graph::execute) - 模块按顺序执行 (
OprModuleBase::execute) - 获取输出结果
-
内存管理：
- 按需分配内存 (
Tensor::prepare_data) - 使用完毕释放内存 (
Tensor::recall_data) - 通过引用计数管理内存生命周期
- 按需分配内存 (
总结
InferLLM 的 core 模块提供了一个灵活、高效的框架，用于构建和执行 LLM 模型。它通过抽象设备、张量、计算图等概念，实现了对不同硬件平台的支持和高效的内存管理。核心模块的设计清晰，各组件职责明确，便于扩展和维护。
赏
