MiniInfer 是一个学习性质的大语言模型推理引擎,从零实现了现代 LLM 推理系统的核心组件。本项目旨在帮助开发者深入理解 LLM 推理的底层机制,包括注意力机制、KV 缓存、批处理调度等关键技术。
- ⚡ 性能优化:实现主流推理优化技术
- 🔧 易于扩展:模块化设计,便于添加新功能
- 📚 完整实现:从基础算子到完整推理流程
-
✅ 注意力机制
- Multi-Head Attention (MHA)
- Grouped Query Attention (GQA)
- 支持因果掩码 (Causal Mask)
-
✅ 位置编码
- Rotary Position Embedding (RoPE)
- 支持传统和非传统模式(Qwen2 使用 non-tranditional)
-
✅ 激活函数
-
✅ 归一化层
- RMSNorm
-
✅ KV Cache
- 单请求 KV Cache
- 批处理 KV Cache (尾部对齐)
- 支持动态请求管理
-
✅ 批处理推理
- Continuous Batching
- 动态请求调度
- 支持多请求并发
-
✅ CUDA 扩展
- C++/CUDA 自定义算子
- PyTorch C++ 扩展框架
- 向量加法示例(可扩展更多算子)
- ✅ Qwen2 系列模型 (0.5B, 1.5B, 7B)
💀 7B 未进行测试
- 完整的模型实现
- 权重加载与转换
- 量化支持 (FP16)
🔮 计划功能 (详见 ROADMAP.md)
- 🚧 项目代码重构中(重构目标类似 nano-vllm)
- 🔄 张量并行 (Tensor Parallelism)
- 🔄 流水线并行 (Pipeline Parallelism)
- 🔄 Flash Attention 集成
- 🔄 PagedAttention
- 🔄 推测解码 (Speculative Decoding)
- 🔄 量化支持 (INT8/INT4)
- 🔄 更多模型架构
💡 新手推荐: 查看详细的 快速入门指南 获取完整教程和常见问题解答。
- Python 3.10-3.12
- CUDA 11.8+ (用于 GPU 加速)
- 克隆仓库
git clone https://github.com/tom-jerr/MiniInfer .git
cd MiniInfer- 安装依赖
使用 PDM:
pip install pdm
pdm install- 构建 CUDA 扩展 (可选)
pdm run build-ext
pdm run build-ext-test # 测试扩展# 使用简单生成 (无 KV Cache)
pdm main
# 或使用带 KV Cache 版本
pdm main-with-kvcachepdm batch-main# 运行所有测试
pdm test
# 运行特定测试
pdm test --fun gqa # 会测试 unittests/test_gqa.py🚧 正在施工
🚧 正在施工
🚧 正在施工
项目包含完整的测试套件:
# 运行所有测试
pdm run test
# 运行特定测试
# 运行特定测试
pdm test --fun gqa # 会测试 unittests/test_gqa.py查看 ROADMAP.md 了解详细的开发计划和进度。
- Flash Attention 集成
- PagedAttention 实现
- 张量并行基础支持
- 完整的 Tensor Parallelism
- Pipeline Parallelism
- INT8 量化支持
- 快速入门指南 - 5 分钟上手教程,包含常见问题解答
- 开发路线图 - 详细的功能规划和开发进度
- 贡献指南 - 如何为项目做贡献
- 更新日志 - 版本历史和变更记录
- 激活函数接口文档 - 激活函数使用说明
欢迎贡献代码、报告问题或提出建议!
- Fork 项目
- 创建特性分支 (
git checkout -b feature/AmazingFeature) - 提交更改 (
git commit -m 'Add some AmazingFeature') - 推送到分支 (
git push origin feature/AmazingFeature) - 开启 Pull Request
本项目采用 MIT 许可证 - 详见 LICENSE 文件
本项目在开发过程中参考了以下优秀项目和资源:
- tiny-llm - 最初的项目框架代码参考
- nano-vllm - 重构后项目框架代码参考
- vLLM - 高性能 LLM 推理引擎,PagedAttention 和 Continuous Batching 的创新实现给了我们很大启发
- Attention Is All You Need - Transformer 架构的奠基论文
- GQA: Training Generalized Multi-Query Transformer Models - Grouped Query Attention 的设计思想
- Efficient Memory Management for Large Language Model Serving with PagedAttention - vLLM 的核心技术论文
感谢所有为开源 LLM 生态系统做出贡献的开发者和研究者。本项目作为学习性质的实现,旨在帮助更多人理解现代 LLM 推理系统的工作原理。
如果本项目的代码中有任何参考未明确标注来源,请联系我们,我们会立即补充说明。
- 作者: lzy
- Email: tomlzy213@gmail.com
- GitHub: @tom-jerr
如果这个项目对你有帮助,请给个 ⭐️ Star!
Made with ❤️ by tom-jerr