从零开始构建PhotonInfer推理引擎2

本文详细解析PhotonInfer项目PR2的核心设计理念和技术实现，聚焦模型权重读取、mmap内存映射、tokenizer机制和张量抽象，带你深入理解现代推理引擎的数据加载和计算技术。

📋 文章导览

  • 项目背景：PR2的核心定位
  • 设计哲学：内存映射与零拷贝理念
  • 核心组件：权重加载、tokenizer、张量抽象、文件格式解析
  • 关键代码：mmap实现与数据结构详解
  • 性能优化：为什么mmap能大幅提升性能？
  • 教学价值：PR2教会了我们什么？

🎯 项目背景：PR2的核心使命

从PR1到PR2：基础设施到数据加载

PR1建立了现代C++20基础设施，PR2则实现了高效的数据加载系统：

1
2
3
4
5
6
// PR2完成的核心功能
✅ 内存映射权重加载（mmap + RAII）
✅ 多格式模型支持（FP32/INT8量化）
✅ TikTokenizer实现（Base64 + 贪心匹配）
✅ 张量抽象层（Eigen集成）
✅ 完整的文件格式解析

PR2的技术挑战

深度学习推理引擎面临的核心问题：

  • 大模型加载：7B参数模型需要14GB内存，如何高效加载？
  • 内存效率：避免数据拷贝，减少内存占用
  • 格式兼容：支持多种模型格式和量化方案
  • token处理：如何将文本转换为模型可理解的token序列？

🧠 设计哲学：内存映射与零拷贝

1. mmap：文件到内存的桥梁

传统文件读取的问题：

1
2
3
4
5
// 传统方式：多次拷贝
std::ifstream file("model.bin", std::ios::binary);
std::vector<char> buffer(1024 * 1024 * 1024); // 1GB
file.read(buffer.data(), buffer.size());     // 磁盘 -> 用户缓冲区
// 使用时还要拷贝到GPU...

PhotonInfer的mmap理念：

1
2
3
4
// mmap方式：零拷贝映射
void* data = mmap(nullptr, file_size, PROT_READ, MAP_PRIVATE, fd, 0);
// 文件内容直接映射到进程地址空间，无需拷贝
// 按需加载（page fault机制）

2. RAII包装的内存管理

1
2
3
4
5
6
7
8
9
10
11
class RawModelData {
public:
  virtual ~RawModelData() {
    if (data_ != nullptr && data_ != MAP_FAILED) {
      munmap(data_, file_size_);  // 自动释放映射
    }
    if (fd_ != -1) {
      close(fd_);  // 自动关闭文件
    }
  }
};

3. 类型安全的权重访问

1
2
3
4
5
6
7
8
// 编译期类型检查 + 运行时偏移计算
template <typename WeightType>
class RawModelDataTyped : public RawModelData {
public:
  const WeightType* weight(usize offset) const {
    return static_cast<const WeightType*>(weight_data_) + offset;
  }
};

🏗️ 核心架构：PR2的五大组件

组件1：模型加载器 (ModelLoader)

mmap文件映射的完整流程

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
Result<void> ModelLoader::mmap_file(
    const std::filesystem::path& path,
    i32& fd,
    void*& data,
    usize& file_size) {

  // 1. 打开文件
  fd = open(path.c_str(), O_RDONLY);
  if (fd == -1) {
    return Err<void>(ErrorCode::FileNotFound,
      "Failed to open model file: " + path.string());
  }

  // 2. 获取文件大小
  struct stat sb;
  if (fstat(fd, &sb) == -1) {
    close(fd);
    return Err<void>(ErrorCode::IOError,
      "Failed to get file size");
  }
  file_size = static_cast<usize>(sb.st_size);

  // 3. 内存映射文件
  data = mmap(nullptr, file_size, PROT_READ, MAP_PRIVATE, fd, 0);
  if (data == MAP_FAILED || data == nullptr) {
    close(fd);
    return Err<void>(ErrorCode::IOError,
      "Failed to mmap model file");
  }

  return Ok();
}

mmap关键参数解析：

  • PROT_READ：只读映射
  • MAP_PRIVATE：私有映射，修改不会写回文件
  • fd：文件描述符
  • offset=0：从文件开头映射

模型文件格式解析

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
Result<ModelLoader::LoadedModel> ModelLoader::load(...) {
  // 1. 读取配置头 (7个int32)
  ModelConfig config;
  std::fread(&config, sizeof(ModelConfig), 1, file);

  // 2. 读取量化参数 (可选)
  i32 group_size = 1;
  if (is_quantized) {
    std::fread(&group_size, sizeof(i32), 1, file);
  }

  // 3. mmap映射权重数据
  RawModelData* raw_data = is_quantized ?
    new RawModelDataInt8() : new RawModelDataFp32();

  mmap_file(model_path, fd, data, file_size);
  raw_data->weight_data_ = data + weight_offset(is_quantized);
}

组件2：原始模型数据 (RawModelData)

模板方法模式的类型安全访问

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class RawModelData {
public:
  // 模板方法：子类实现具体类型访问
  virtual const void* weight(usize offset) const = 0;

protected:
  void* weight_data_ = nullptr;  // 权重数据起始地址
};

// FP32实现
class RawModelDataFp32 : public RawModelData {
public:
  const void* weight(usize offset) const override {
    return static_cast<const f32*>(weight_data_) + offset;
  }
};

// INT8实现
class RawModelDataInt8 : public RawModelData {
public:
  const void* weight(usize offset) const override {
    return static_cast<const i8*>(weight_data_) + offset;
  }
};

工厂模式的类型创建

1
2
3
4
5
6
7
8
9
10
11
12
template <typename T>
struct RawModelDataFactory;

template <>
struct RawModelDataFactory<f32> {
  using type = RawModelDataFp32;
};

template <>
struct RawModelDataFactory<i8> {
  using type = RawModelDataInt8;
};

组件3：TikTokenizer实现

Base64解码的编译期优化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// 编译期构建查找表
constexpr auto TikTokenizer::build_base64_table() noexcept -> std::array<i32, 256> {
  std::array<i32, 256> table{};
  constexpr std::string_view base64_chars =
      "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

  for (usize i = 0; i < base64_chars.size(); ++i) {
    table[static_cast<u8>(base64_chars[i])] = static_cast<i32>(i);
  }
  return table;
}

// 高效的Base64解码
auto TikTokenizer::base64_decode(std::string_view encoded) -> std::string {
  static constexpr auto lookup_table = build_base64_table();
  // 使用位操作进行解码...
}

贪心tokenization算法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
auto TikTokenizer::encode(std::string_view text) const -> std::vector<i32> {
  std::vector<i32> tokens;
  usize i = 0;

  while (i < text.length()) {
    // 贪心策略：优先匹配最长token
    const usize max_len = std::min(text.length() - i, usize{16});
    bool found = false;

    for (usize len = max_len; len > 0; --len) {
      const std::string_view substr = text.substr(i, len);
      const std::string key{substr};

      if (auto it = vocab_.find(key); it != vocab_.end()) {
        tokens.push_back(it->second);
        i += len;
        found = true;
        break;
      }
    }

    if (!found) {
      // 回退策略：字节级编码
      const std::string byte_str{text[i]};
      if (auto it = vocab_.find(byte_str); it != vocab_.end()) {
        tokens.push_back(it->second);
      }
      ++i;
    }
  }

  return tokens;
}

组件4：张量抽象层 (Tensor)

Eigen集成的零拷贝视图

1
2
3
4
5
6
7
8
9
10
11
class Tensor {
public:
  // 创建张量
  static Result<Tensor> create(std::vector<int32_t> dims, DataType dtype);

  // Eigen矩阵视图（零拷贝）
  template <typename T>
  auto matrix_map() -> Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>> {
    // 返回Eigen::Map，直接操作底层Buffer数据
  }
};

组件5：张量抽象层 (Tensor)

Eigen零拷贝集成的设计理念

张量是深度学习的核心数据结构，PR2实现了基于Eigen的张量抽象：

核心设计原则：

  • 存储vs计算分离：Tensor管理内存，Eigen::Map提供计算视图
  • 零拷贝访问：Eigen::Map直接操作底层Buffer数据
  • 类型安全：编译期类型检查，运行时类型转换
  • 设备抽象：统一的CPU/CUDA接口
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class Tensor {
public:
  // 创建多维张量
  static Result<Tensor> create(std::vector<int32_t> dims, DataType dtype);

  // 零拷贝Eigen视图
  template <typename T>
  Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>
  matrix_map();

  // 内存管理
  Buffer buffer_;                // 底层内存缓冲区
  std::vector<int32_t> dims_;    // 维度信息
  usize size_;                   // 元素总数
  DataType dtype_;               // 数据类型
};

内存布局与步长计算

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
std::vector<usize> Tensor::strides() const {
  if (dims_.empty()) {
    return {};
  }

  std::vector<usize> strides(dims_.size());
  strides.back() = 1;

  // 行优先(row-major)步长计算
  for (isize i = static_cast<isize>(dims_.size()) - 2; i >= 0; --i) {
    strides[i] = strides[i + 1] * dims_[i + 1];
  }

  return strides;
}

Eigen零拷贝矩阵视图

1
2
3
4
5
6
7
8
9
10
11
12
13
// 2D矩阵视图（行优先，与NumPy兼容）
template <typename T>
Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>
Tensor::matrix_map() {
  if (ndim() != 2) {
    return Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
        nullptr, 0, 0);
  }

  // 直接映射底层Buffer数据，无拷贝
  return Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
      ptr<T>(), dims_[0], dims_[1]);
}

使用模式：

1
2
3
4
5
6
7
8
9
10
11
// 创建张量
auto tensor_result = Tensor::create({2, 3}, DataType::Float32);
Tensor tensor = std::move(tensor_result.value());

// 零拷贝Eigen操作
auto mat = tensor.matrix_map<float>();
mat = mat * 2.0f + 1.0f;  // 原地操作，直接修改tensor数据

// 向量视图
auto vec = tensor.vector_map<float>();
vec = vec.array().sqrt();  // 逐元素开方

张量创建工厂方法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
// 从std::vector创建1D张量
template <typename T>
Result<Tensor> Tensor::from_vector(const std::vector<T>& data, DeviceType device) {
  DataType dtype = cpp_type_to_data_type_v<T>;
  std::vector<int32_t> dims = {static_cast<int32_t>(data.size())};

  auto result = create(dims, dtype, device);
  if (!result) return result;

  Tensor tensor = std::move(result.value());

  // 直接内存拷贝
  if (device == DeviceType::CPU) {
    std::memcpy(tensor.ptr<T>(), data.data(), data.size() * sizeof(T));
  }

  return Ok(std::move(tensor));
}

Move-only语义与RAII

1
2
3
4
5
6
7
8
9
10
class Tensor {
public:
  // 禁止拷贝（独占内存所有权）
  Tensor(const Tensor&) = delete;
  Tensor& operator=(const Tensor&) = delete;

  // 只允许移动
  Tensor(Tensor&&) noexcept = default;
  Tensor& operator=(Tensor&&) noexcept = default;
};

📊 性能特性分析

mmap的性能优势

1. 内存效率对比

1
2
3
4
5
传统加载:
磁盘 -> 内核缓冲区 -> 用户缓冲区 -> GPU (多次拷贝)

mmap加载:
磁盘 -> 虚拟内存 -> 进程地址空间 (零拷贝，按需加载)

2. 启动时间优化

1
2
3
4
5
// mmap的优势
auto start = std::chrono::high_resolution_clock::now();
auto result = ModelLoader::load("llama-7b.bin");
// mmap几乎瞬时完成（只建立映射关系）
auto end = std::chrono::high_resolution_clock::now();

3. 内存占用优化

  • 传统方式：需要额外缓冲区存储整个文件
  • mmap方式：只占用虚拟内存地址空间，物理内存按需分配

量化模型支持

1
2
3
4
5
6
7
// INT8量化权重加载
auto loaded = ModelLoader::load_typed<i8>("model_quant.bin");
auto weights = loaded.raw_data->weight(offset);  // int8指针

// FP32权重加载
auto loaded_fp32 = ModelLoader::load_typed<f32>("model.bin");
auto weights_fp32 = loaded_fp32.raw_data->weight(offset);  // float指针

🎬 实战示例：完整的模型加载流程

示例1：加载Llama模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
#include "photon/model/model_loader.hpp"
#include "photon/model/tokenizer.hpp"

void load_and_use_model() {
  // 1. 加载tokenizer
  auto tokenizer_result = TikTokenizer::load("tokenizer.model");
  if (!tokenizer_result) {
    std::cerr << "Failed to load tokenizer" << std::endl;
    return;
  }
  TikTokenizer tokenizer = std::move(tokenizer_result.value());

  // 2. 加载模型（自动检测量化）
  auto model_result = ModelLoader::load("llama-7b.bin", false);
  if (!model_result) {
    std::cerr << "Failed to load model" << std::endl;
    return;
  }
  auto& loaded = model_result.value();

  // 3. 编码输入文本
  std::string prompt = "Hello, how are you?";
  auto tokens = tokenizer.encode(prompt);

  // 4. 添加特殊token
  tokens.insert(tokens.begin(), tokenizer.bos_id());
  tokens.push_back(tokenizer.eos_id());

  std::cout << "Encoded " << prompt.size() << " chars into "
            << tokens.size() << " tokens" << std::endl;
}

示例2：量化模型处理

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
// 加载INT8量化模型
auto quant_result = ModelLoader::load_typed<i8>("llama-7b-quant.bin");
if (!quant_result) {
  return;
}

auto& loaded = quant_result.value();

// 访问量化权重
const i8* attention_weights = static_cast<const i8*>(
  loaded.raw_data->weight(attention_offset)
);

// 反量化计算
float scale = 1.0f / loaded.group_size;
for (size_t i = 0; i < weight_count; ++i) {
  float dequantized = attention_weights[i] * scale;
  // 使用反量化后的权重进行计算
}

示例3：张量操作与Eigen集成

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
#include "photon/core/tensor.hpp"

void tensor_operations() {
  // 1. 创建多维张量
  auto tensor_result = Tensor::create({2, 3}, DataType::Float32);
  if (!tensor_result) {
    std::cerr << "Failed to create tensor" << std::endl;
    return;
  }
  Tensor tensor = std::move(tensor_result.value());

  // 2. 初始化数据
  auto mat = tensor.matrix_map<float>();
  mat << 1.0f, 2.0f, 3.0f,
         4.0f, 5.0f, 6.0f;

  // 3. Eigen原地操作（零拷贝）
  mat = mat * 2.0f + 1.0f;  // [3, 5, 7; 9, 11, 13]

  // 4. 向量操作
  auto vec = tensor.vector_map<float>();
  vec = vec.array().sqrt();  // 逐元素开方

  // 5. 张量信息查询
  std::cout << "Shape: [" << tensor.dim(0) << ", " << tensor.dim(1) << "]" << std::endl;
  std::cout << "Size: " << tensor.size() << " elements" << std::endl;
  std::cout << "Data type: " << static_cast<int>(tensor.dtype()) << std::endl;

  // 6. 从向量创建张量
  std::vector<float> data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
  auto vec_tensor_result = Tensor::from_vector(data);
  if (vec_tensor_result) {
    Tensor vec_tensor = std::move(vec_tensor_result.value());
    std::cout << "Created 1D tensor with " << vec_tensor.size() << " elements" << std::endl;
  }
}

📊 文件格式详解

模型文件结构

1
2
3
4
5
┌─────────────────┬─────────────────┬─────────────────┐
│   ModelConfig   │   Group Size    │   Weight Data   │
│   (7 * int32)   │   (int32)       │   (float/int8)  │
│                 │   [可选]        │                 │
└─────────────────┴─────────────────┴─────────────────┘

ModelConfig结构

1
2
3
4
5
6
7
8
9
struct ModelConfig {
  int32_t dim = 0;         // 模型维度 (4096 for Llama-7B)
  int32_t hidden_dim = 0;  // FFN隐藏维度 (11008 for Llama-7B)
  int32_t layer_num = 0;   // Transformer层数 (32 for Llama-7B)
  int32_t head_num = 0;    // 注意力头数 (32 for Llama-7B)
  int32_t kv_head_num = 0; // KV头数 (32 for Llama-7B, 支持GQA)
  int32_t vocab_size = 0;  // 词汇表大小 (32000 for Llama)
  int32_t seq_len = 0;     // 最大序列长度 (4096)
};

Tokenizer文件格式

1
2
3
4
5
# tokenizer.model 文件格式
base64_encoded_token token_id
aGVsbG8= 12345          # "hello" -> 12345
d29ybGQ= 12346          # "world" -> 12346
...

Base64编码原理

1
2
3
4
5
// 原始字符串 -> Base64编码 -> 存储
"hello" -> base64_encode("hello") -> "aGVsbG8="

// 读取时解码
base64_decode("aGVsbG8=") -> "hello"

🎓 教学价值：PR2教会了我们什么？

1. 系统级编程思维

  • 内存映射：理解虚拟内存和按需加载
  • RAII资源管理：确保系统资源正确释放
  • 零拷贝设计：性能优化的核心思想

2. 深度学习系统设计

  • 文件格式设计：如何设计高效的二进制格式
  • 量化支持：INT8量化的数据流和计算
  • tokenizer机制：文本到token的转换算法
  • 张量抽象：内存管理和计算视图的分离设计

3. 现代C++特性应用

  • 模板元编程：编译期类型计算
  • constexpr函数：编译期常量计算
  • std::filesystem：跨平台文件操作

4. 性能优化思维

  • 避免内存拷贝：mmap和零拷贝视图
  • 按需加载：page fault机制的应用
  • 缓存友好：数据结构设计的影响
  • 存储vs计算分离：Tensor内存管理，Eigen计算加速

🚀 PR2的意义：数据加载的艺术

PR2完成了推理引擎最核心的数据加载基础设施：

  1. 高效权重加载：mmap + RAII实现零拷贝大模型加载
  2. 多格式支持：FP32/INT8量化模型的统一接口
  3. tokenizer集成：完整的文本预处理流水线
  4. 张量抽象：Eigen零拷贝集成的多维数组管理
  5. 类型安全：编译期检查的模板元编程

这些设计为后续的Transformer推理实现奠定了坚实基础。

📚 延伸阅读

结语：PR2展示了如何用系统级编程技巧构建高性能的深度学习基础设施。通过mmap、RAII和类型安全的结合，我们不仅获得了高效的数据加载能力，更重要的是学会了一种性能导向的系统设计思维。

这正是PhotonInfer项目的魅力所在：用严谨的工程实践，支撑AI推理的无限可能。

关注我的B站账号，获取更多深度学习系统优化的教学内容！

#PhotonInfer #C++20 #mmap #深度学习 #推理引擎 #从零构建 #TikTokenizer