从零开始构建LLM推理引擎：深入理解Transformer推理与系统优化

项目概述

本开源项目由ashwinvijayakumar24开发（GitHub repo: llm_inference_engine, 发布时间：2026年6月5日），旨在从零构建LLM推理引擎，深入解析生产级推理系统的核心原理。项目涵盖Transformer前向传播实现、KV缓存机制、连续批处理、分页注意力（PagedAttention）及CUDA内核优化等关键技术，目标是在NVIDIA A100/H100/H200 GPU上实现高效的Llama3.2 1B模型推理，并与HuggingFace Transformers、llama.cpp进行基准对比，为开发者提供实践指南。

背景与动机

当前LLM开发中，多数开发者依赖model.generate()或成熟框架（如vLLM、llama.cpp）完成推理，但对底层系统细节（Transformer前向传播、KV缓存管理、批处理调度、CUDA优化）缺乏深入理解。本项目通过从零构建完整推理引擎，帮助开发者掌握生产级推理的核心原理。项目采用Python进行高层编排，结合CUDA C++编写高性能内核，聚焦Llama3.2 1B模型在NVIDIA GPU上的高效推理。

技术栈与架构

• 允许使用的底层库：NumPy/PyTorch张量存储、cuBLAS（GEMM运算）、HuggingFace Tokenizer（Rust后端）、safetensors权重解析。
• 自主实现组件：Transformer各层逻辑（RoPE位置编码、GQA分组查询注意力、RMSNorm归一化、SwiGLU前馈网络）、KV缓存管理（朴素连续缓冲→分页式块管理）、连续批处理调度器、量化路径（int8/int4）、自定义CUDA注意力内核。
• 开发环境：Apple M4 MacBook Air（正确性验证、Python引擎开发）；NVIDIA PACE集群（CUDA内核开发、性能基准测试，使用A100/H100/H200 GPU及Slurm调度）。

核心实现细节

1. Transformer前向传播：手动实现完整流程（词嵌入→多层Transformer→LM头→采样），关键细节包括RoPE位置编码、GQA分组查询注意力、RMSNorm归一化、SwiGLU前馈网络。
2. KV缓存管理：从朴素连续缓冲区（预分配最大长度，内存浪费）演进到分页式块管理（PagedAttention风格，消除碎片，支持变长序列）。
3. 连续批处理：允许新请求加入、完成请求移出当前批次，提升GPU利用率；建议先实现单线程调度器验证正确性，再并行优化。
4. 量化优化：计划实现int8（内存减半）和int4（内存减至1/4）权重量化，评估指标包括内存占用、困惑度、吞吐量。
5. CUDA注意力内核：开发流程为NumPy参考→CUDA代码→Python绑定→性能剖析→对比PyTorch基线；重点优化解码阶段（序列长度1）的内存访问模式。

开发计划与基准测试

• 分阶段开发：
  1. 环境搭建（CMake+nanobind、权重解析、Tokenizer验证）；
  2. 正确性验证（NumPy实现前向传播，与HF对比隐藏状态）；
  3. 可用性完善（朴素KV缓存、采样策略、CLI/HTTP端点）；
  4. 基线基准（性能测试框架，HF/llama.cpp对比）；
  5. 差异化优化（量化→连续批处理→分页KV→CUDA内核→推测解码）；
  6. 项目打磨（文档完善、性能图表）。
• 基准指标：吞吐量（prefill/decode）、TTFT（首token时间）、ITL（token间延迟）、峰值显存、吞吐量vs批次大小、显存vs上下文长度；对比对象为HF Transformers和llama.cpp。

潜在风险与应对

1. Logits与HF不匹配：逐层对比隐藏状态，排查RoPE参数、RMSNorm epsilon位置、注意力缩放因子等问题。
2. PACE Slurm队列等待：在Mac完成Python/正确性工作，仅CUDA开发时使用PACE。
3. CUDA内核正确性：先用NumPy参考版本验证输出，再进行性能测试。
4. 范围蔓延：专注2-3高质量功能，避免半成品。

项目价值与未来方向

• 价值：教育意义（理解推理内在机制）、工程实践（模块化系统设计）、面试准备（简历亮点）、全栈技术深度（算法到硬件）。
• 未来方向：多GPU/张量并行推理、Flash Attention融合内核、LoRA适配器热切换、完整推测解码实现。