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. 核心功能流程

1. 模型加载：

   • 创建 Graph 对象
   • 构建模型结构 (construct_llm)
   • 加载模型权重

2. 推理执行：

   • 准备输入数据
   • 执行计算图 (Graph::execute)
   • 模块按顺序执行 (OprModuleBase::execute)
   • 获取输出结果

3. 内存管理：

   • 按需分配内存 (Tensor::prepare_data)
   • 使用完毕释放内存 (Tensor::recall_data)
   • 通过引用计数管理内存生命周期

总结

InferLLM 的 core 模块提供了一个灵活、高效的框架，用于构建和执行 LLM 模型。它通过抽象设备、张量、计算图等概念，实现了对不同硬件平台的支持和高效的内存管理。核心模块的设计清晰，各组件职责明确，便于扩展和维护。