最硬核仓颉AI框架实现：从零构建Transformer推理引擎全指南
项目地址: https://gitcode.com/abcd1234-wyj/learning_ml_cj 你是否还在为大模型推理引擎的复杂实现望而却步？是否想亲手用仓颉(Cangjie)语言构建高效的Transformer推理系统？本文将带你深入学习_ml_cj项目的核心架构与实现细节，掌握从张量操作到完整推理链路的全流程开发。读完本文你将获得：
- 仓颉语言实现高性能AI算子的实战经验
- Transformer核心组件(RoPE/Self-Attention/MLP)的底层实现
- 大模型推理引擎的完整架构设计方法论
- 从零到一部署可运行的Llama模型推理系统
项目背景与核心价值
learning_ml_cj项目脱胎于"InfiniTensor：人工智能编译器与大模型系统训练营"，旨在复现Python中Transformer库的核心功能并移植至仓颉社区。作为东北大学仓颉社区软件工程课程的实践项目，它不仅实现了基础的大模型推理能力，更为仓颉语言在AI领域的应用提供了宝贵的技术参考。
项目定位与特性
该项目是一个支持safetensors格式的大模型推理引擎，主要特点包括：
| 核心特性 | 技术指标 | 应用场景 |
|---|---|---|
| 纯仓颉实现 | 零外部依赖，原生仓颉代码 | 端侧部署、教育学习 |
| 完整Transformer架构 | 包含Self-Attention、MLP等核心模块 | 大模型推理研究 |
| 支持安全张量格式 | 兼容safetensors模型文件 | 模型安全加载 |
| 轻量级设计 | 核心代码不足2000行 | 嵌入式系统集成 |
项目目前已实现基础的文本生成功能，自带一个656K参数的TinyStories示例模型，可直接运行推理测试。
系统架构深度解析
整体架构设计
learning_ml_cj采用模块化设计，整体架构分为存储层和推理层两大核心部分：
核心工作流程：输入张量经过模型结构处理，调用算子库中的计算函数，利用参数管理模块加载的模型权重，通过键值缓存优化推理效率，最终输出推理结果。
源码目录结构
.
├── src # 仓颉源代码
│ ├── kvcache # KV缓存实现
│ ├── model # Transformer核心结构
│ │ ├── mlp.cj # 多层感知机
│ │ ├── model.cj # Llama模型主类
│ │ └── self_attention.cj # 自注意力机制
│ ├── operators # 算子实现
│ │ ├── rope.cj # 旋转位置编码
│ │ ├── masked_softmax.cj # 带掩码的softmax
│ │ └── matmul_transb.cj # 矩阵转置乘法
│ ├── params # 模型参数管理
│ └── tensor # 张量实现
└── models # 示例模型文件
└── story # TinyStories模型
核心代码集中在src目录，按功能划分为5个主要模块，每个模块职责单一清晰，便于维护和扩展。
核心模块实现详解
1. 张量系统(Tensor)
张量是整个系统的基础数据结构，用于存储和操作多维数组。tensor.cj实现了一个基础但高效的张量类，支持多种操作：
public class Tensor {
// 构造函数，支持指定数据和形状
public init(dataInner: Array<Float32>, shapeInner: Array<UInt>) { ... }
// 创建指定形状的零张量
public static func default(shapeInner: Array<UInt>): Tensor { ... }
// 重塑张量形状
public func reshape(new_shape: Array<UInt>): Tensor { ... }
// 切片操作
public func slice(start: UInt, new_shape: Array<UInt>): Tensor { ... }
// 数据访问接口
public func data(): Array<Float32> { ... }
public func shape(): Array<UInt> { ... }
}
张量实现采用扁平化存储方式，所有数据存储在一维数组中，通过形状(shape)信息进行维度管理。这种设计兼顾了存储效率和访问灵活性，特别适合小模型推理场景。
2. 旋转位置编码(RoPE)
RoPE(Rotary Position Embedding)是当前大模型中广泛使用的位置编码技术，rope.cj实现了这一核心算子：
public func rope(y: Tensor, start_pos: UInt, theta: Float32) {
let shape = y.shape();
let seq_len = Int64(shape[0]);
let n_heads = Int64(shape[1]);
let d = Int64(shape[2]);
let data = y.data();
for (tok in 0..seq_len) {
let pos = Int64(start_pos) + tok;
for (head in 0..n_heads) {
for (i in 0..d / 2) {
// 计算频率
let freq = Float32(pos) / Float32(pow(theta, Float32(i * 2) / Float32(d)));
let (sin, cos) = (sin(freq), cos(freq));
// 复数旋转操作
let a = data[tok * n_heads * d + head * d + i];
let b = data[tok * n_heads * d + head * d + i + d / 2];
data[tok * n_heads * d + head * d + i] = a * cos - b * sin;
data[tok * n_heads * d + head * d + i + d / 2] = b * cos + a * sin;
}
}
}
}
该实现通过复数旋转的方式将位置信息编码到注意力的Query和Key中，使模型能够感知序列中词元的位置关系。代码中通过三重循环遍历序列长度、注意力头数和维度，对每个元素执行旋转计算。
3. 自注意力机制(Self-Attention)
自注意力是Transformer的核心组件，self_attention.cj实现了完整的多头注意力机制：
实现代码中，首先对Q、K、V进行线性变换和RoPE编码，然后计算注意力分数并应用掩码，最后通过Softmax归一化并与V相乘得到输出：
public func self_attention(
output: Tensor,
att_scores: Tensor,
q: Tensor,
k: Tensor,
v: Tensor,
n_kv_heads: UInt,
n_groups: UInt,
seq_len: UInt,
total_seq_len: UInt,
dqkv: UInt
) {
// 1. 计算注意力分数 (Q*K^T / sqrt(d))
dot(att_scores, q, k, n_kv_heads, n_groups, seq_len, total_seq_len, dqkv);
// 2. 应用掩码和Softmax
masked_softmax(att_scores, seq_len, total_seq_len, n_kv_heads, n_groups);
// 3. 计算输出 (注意力权重 * V)
matmul_transb(output, 0.0, att_scores, v, 1.0 / Float32(dqkv).sqrt());
}
该实现支持多组注意力(Grouped Attention)，通过n_groups参数控制，可在保持性能的同时减少计算量。
4. Llama模型主类
model.cj中的Llama类是整个推理系统的核心，封装了完整的前向传播过程：
public class Llama {
// 模型配置参数
public var vocab: UInt; // 词表大小
public var n_layers: UInt; // 层数
public var n_q_h: UInt; // 查询头数
public var n_kv_h: UInt; // 键值头数
public var d: UInt; // 隐藏层维度
public var dqkv: UInt; // QKV维度
public var di: UInt; // 中间层维度
public var eps: Float32; // 归一化epsilon
public var rope_theta: Float32;// RoPE参数
public var max_seq_len: UInt; // 最大序列长度
// 参数对象
public var params: LLamaParams;
// 构造函数
public static func from_example_files(model_path: Option<Path>): Llama { ... }
// 前向传播
public func forward(input: Tensor, cache: KVCache): Tensor {
let seq_len = input.size();
let past_seq_len = cache.len();
cache.increment(seq_len);
let total_seq_len = past_seq_len + seq_len;
let n_groups = this.n_q_h / this.n_kv_h;
// 初始化缓冲区
let residual = Tensor.default([seq_len, this.d]);
var hidden_states = Tensor.default([seq_len, this.d]);
// 词嵌入查找
gather(residual, input, this.params.embedding_table);
// 逐层计算
for (layer in 0..this.n_layers) {
// 自注意力模块
rms_norm(hidden_states, residual, this.params.rms_att_w[layer], this.eps);
// ... 注意力计算过程 ...
// MLP模块
mlp(residual, hidden_states, ...);
}
// 输出层计算
rms_norm(_hidden_states, _residual, this.params.rms_out_w, this.eps);
matmul_transb(logits, 0.0, _hidden_states, this.params.lm_head, 1.0);
return logits;
}
// 生成函数
public func generate(
token_ids: Array<UInt32>,
max_len: UInt32,
top_p: Float32,
top_k: UInt32,
temperature: Float32
): ArrayList<UInt32> { ... }
}
前向传播过程遵循标准Transformer架构：输入通过词嵌入层后，依次经过多个Transformer块(每个包含自注意力和MLP模块)，最后通过输出层生成下一个token的概率分布。
generate函数实现了文本生成逻辑，支持top-p、top-k等采样参数，通过循环调用forward函数实现多轮推理。
快速上手实战教程
环境准备与安装
前置条件
- 仓颉编译器(cjc) v0.55.3或更高版本
- Git版本控制工具
- 基本开发环境(GCC等)
安装步骤
# 克隆仓库
git clone https://gitcode.com/abcd1234-wyj/learning_ml_cj
cd learning_ml_cj
# 编译项目
cjpm run
项目使用仓颉包管理器(cjpm)构建，编译成功后将自动运行示例推理程序。
核心功能示例
基础文本生成
以下是使用自带示例模型进行文本生成的完整代码：
import std.fs.*
import learning_ml_cj.model.*
import learning_ml_cj.translator.*
main(args: Array<String>): Int64 {
// 输入token ids: "Once upon a time,"
let input: Array<UInt32> = [1, 80, 147, 201, 282, 57];
// 加载模型和翻译器
let model = Llama.from_example_files(Option<Path>.None);
let translator = Translator.from_example_files(Option<Path>.None);
// 生成文本(最大长度200，top_p=0.9，top_k=40)
let output = model.generate(input, 200, 0.9, 40, -1.0);
// 解码并输出结果
let user_in = translator.translate(input);
let result = translator.translate(output);
println(user_in + result);
return 0;
}
执行后将输出类似以下内容：
Once upon a time, a little girl named Lily lived in a small house with her mom, dad, and her dog, Spot. Spot loved to play all day. One day, Lily saw a small bird on the ground. She picked it up and tried to help it fly again...
自定义模型加载
若要加载自定义模型，只需将模型路径传递给from_example_files方法：
// 加载自定义模型
let model_path = Path("/path/to/your/model");
let model = Llama.from_example_files(Option<Path>.Some(model_path));
自定义模型需包含以下文件：
- config.json: 模型配置
- model.safetensors: 模型权重
- tokenizer.json: 分词器配置
性能优化与高级特性
推理效率优化
learning_ml_cj实现了多项推理优化技术，主要包括：
-
KV缓存机制：避免重复计算已处理序列的键值对
public func new_cache(): KVCache { return KVCache( this.n_layers, // 层数 this.max_seq_len, // 最大序列长度 this.n_kv_h * this.dqkv, // 每个头的维度 0 // 初始长度 ); } -
算子融合：将多个操作合并为单一函数调用，减少内存访问
-
预分配缓冲区：在前向传播开始时分配所有所需缓冲区，避免运行时内存分配
这些优化使模型在低配置设备上也能流畅运行，以656K参数模型为例，在普通CPU上生成100个token仅需约2秒。
扩展性设计
项目架构考虑了未来扩展需求，主要扩展点包括：
- 算子扩展：operators目录下可直接添加新算子，只需实现标准接口
- 模型扩展：通过继承Llama类可实现新模型结构
- 数据类型扩展：Tensor类设计支持未来添加float16/int8等量化类型
项目 roadmap 与贡献指南
未来发展计划
| 版本 | 计划发布时间 | 主要功能 |
|---|---|---|
| v1.0.0 | 2024.11.01 | 基础推理功能，示例模型 |
| v1.0.1 | 2025.02.01 | 模型参数转换工具，翻译模块优化 |
| v1.1.0 | 2025.Q2 | 支持量化模型，性能优化 |
| v2.0.0 | 2025.Q4 | 支持更大模型，添加训练功能 |
如何贡献代码
- Fork本仓库并创建分支
- 提交PR前确保所有测试通过
- 新增功能需包含对应的测试用例
- 代码风格遵循仓颉社区规范
项目采用MIT开源协议，欢迎任何形式的贡献，包括但不限于代码、文档、测试用例等。
总结与学习资源
learning_ml_cj项目为理解大模型推理引擎提供了绝佳的实践案例，通过纯仓颉实现，展示了如何从零构建一个完整的Transformer推理系统。无论是AI初学者还是有经验的开发者，都能从中学习到大模型的核心原理和实现细节。
关键知识点回顾
- Transformer推理的完整流程与核心组件
- 仓颉语言在AI领域的应用实践
- 张量操作与算子实现的底层细节
- 大模型优化技术(KV缓存、RoPE等)的工程实现
扩展学习资源
- 官方文档：项目doc目录下的设计文档和API说明
- 示例模型：models/story目录包含完整的模型文件
- 测试用例：各模块下的*_test.cj文件提供了功能验证代码
通过深入学习和实践该项目，你将获得构建高性能AI系统的核心能力，为进一步探索大模型技术奠定坚实基础。现在就动手尝试修改代码、优化性能或添加新功能，开启你的大模型系统开发之旅吧！
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
