Tokamak：从零构建的极简LLM推理引擎

原作者/维护者：WeichunK 来源平台：github 原始链接：https://github.com/WeichunK/tokamak 发布时间：2026-07-12T18:21:40Z

Tokamak是用PyTorch从零实现的轻量级LLM推理引擎，完整涵盖KV缓存、分页注意力、连续批处理、投机解码等核心优化技术，并与vLLM进行性能对标测试。项目以极简代码展示推理机制底层原理，适合学习、原型验证等场景。

大型语言模型（LLM）推理优化是AI基础设施领域活跃方向，模型规模扩大带来硬件资源高效利用的挑战。现有框架（如vLLM、TensorRT-LLM）代码复杂，学习曲线陡峭。Tokamak应运而生，以极简实现展示核心技术，名称取自核聚变装置“托卡马克”，寓意可控释放大模型能量。

1. KV缓存：缓存之前token的Key/Value向量，避免重复计算，生成新token时间复杂度从O(n²)降为O(n)。
2. 分页注意力：借鉴vLLM技术，将KV缓存划分为固定大小页，类似虚拟内存管理，提高GPU内存利用率，支持更多并发请求。
3. 连续批处理：动态添加/移除请求，GPU利用率最大化，无需等待整个批次完成。
4. 投机解码：用轻量级草稿模型生成候选token，大模型并行验证，一次接受多个token，提升吞吐率。

Tokamak提供与vLLM的性能对比基准测试，在相同硬件和模型参数下比较吞吐量、延迟、内存占用等指标。这种开源透明的对比帮助用户了解实际表现，为框架优化提供数据支撑，也便于研究者复现工业级推理功能。

• 选择PyTorch框架，代码可读性高、易修改调试，适合算法实验和教学。
• 代码结构清晰，核心组件（KV缓存管理、注意力计算等）模块化封装，可单独研究或组合使用。
• 极简设计，代码量远小于生产级框架，是学习LLM推理优化的理想起点。

• 教学学习：为学生和研究者提供精简完整的参考实现。
• 算法原型验证：快速实验新推理优化算法，验证后迁移到生产框架。
• 边缘部署：轻量级特性适合资源受限的边缘设备部署小模型。
• 定制化需求：作为二次开发基础，满足现有框架无法覆盖的特定需求。

Tokamak以极简代码实现现代LLM推理核心技术，为社区提供学习资源和实验平台。它证明轻量级实现可涵盖关键优化技术，降低推理优化领域的进入门槛。随着大模型技术发展，推理优化仍有创新空间，Tokamak助力更多人参与底层技术研究改进。