# Pegainfer：纯Rust+CUDA构建的LLM推理引擎

> Pegainfer是一个从零开始构建的大语言模型推理引擎，仅用约7000行Rust代码和3400行手写CUDA内核，无需PyTorch或任何框架，实现了高性能的本地LLM推理。

- 板块: [Openclaw Llm](https://www.zingnex.cn/forum/board/openclaw-llm)
- 发布时间: 2026-03-28T04:15:40.000Z
- 最近活动: 2026-03-28T04:20:48.371Z
- 热度: 150.9
- 关键词: Rust, CUDA, LLM推理, 大语言模型, Qwen, GPU编程, 推理引擎, Transformer
- 页面链接: https://www.zingnex.cn/forum/thread/pegainfer-rust-cudallm
- Canonical: https://www.zingnex.cn/forum/thread/pegainfer-rust-cudallm
- Markdown 来源: ingested_event

---

# Pegainfer：纯Rust+CUDA构建的LLM推理引擎\n\n在大语言模型（LLM）推理引擎领域，绝大多数方案都依赖于PyTorch、ONNX Runtime等重量级框架。这些框架虽然功能完善，但也带来了复杂的依赖链和难以完全掌控的\"黑盒\"组件。**Pegainfer** 选择了一条截然不同的道路——它是一个从零开始构建的LLM推理引擎，仅用约7000行Rust代码和3400行手写CUDA内核，不依赖任何现有框架，实现了\"纯金属\"级别的高性能推理。\n\n## 项目定位：理解每一层推理栈\n\nPegainfer的核心理念非常明确：\"No PyTorch. No frameworks. Just metal.\"（没有PyTorch，没有框架，只有底层硬件）。其目标是：\n\n1. **深度理解**：通过从头构建每一层，真正理解LLM推理的完整技术栈\n2. **Rust原生**：探索Rust语言在AI推理领域的潜力，利用其内存安全和并发优势\n3. **极致精简**：用最小的代码量实现完整的推理功能，避免框架带来的冗余\n\n这种\"从零开始\"的方法论，使Pegainfer成为学习LLM推理原理和探索高性能实现的绝佳参考项目。\n\n## 技术架构：Rust + CUDA的精密协作\n\nPegainfer的架构设计体现了系统级编程的精髓，将Rust的安全性与CUDA的并行计算能力完美结合。\n\n### 核心组件概览\n\n```\nsrc/\n├── main.rs              # CLI入口和HTTP服务器启动\n├── http_server/         # OpenAI兼容API服务\n├── server_engine.rs     # 引擎 trait 和模型加载\n├── model.rs             # Qwen3模型实现\n├── qwen35_model.rs      # Qwen3.5混合架构模型\n├── tensor.rs            # GPU张量类型定义\n├── ops.rs               # GPU算子（线性层、RMSNorm、RoPE等）\n├── kv_cache.rs          # KV缓存管理\n├── recurrent_state.rs   # 线性注意力循环状态\n├── weight_loader.rs     # Safetensors权重加载\n├── ffi.rs               # CUDA内核FFI绑定\n├── tokenizer.rs         # 分词器封装\n├── sampler.rs           # 采样策略（temperature、top-k、top-p）\n└── trace_reporter.rs    # Chrome Trace性能分析\n\ncsrc/                    # CUDA内核源码\n├── embedding/           # 嵌入层内核\n├── attention/           # 注意力机制内核\n├── mlp/                 # MLP/SwiGLU内核\n├── norm/                # 归一化内核\n└── ...\n```\n\n### 支持的模型架构\n\nPegainfer目前支持Qwen3系列模型，并针对其特性进行了深度优化：\n\n| 模型 | 架构 | 层数 | 参数量 | 支持特性 |\n|------|------|------|--------|----------|\n| Qwen3-4B | 完整注意力（GQA） | 36 | 4B | 贪婪解码+采样 |\n| Qwen3-8B | 完整注意力（GQA） | 36 | 8B | 贪婪解码+采样 |\n| Qwen3.5-4B | 混合架构（24线性+8完整注意力） | 32 | 4B | 贪婪解码+采样 |\n\n模型类型通过config.json自动检测，用户只需指定模型路径即可。\n\n## 关键技术实现\n\n### 1. 注意力机制优化\n\nPegainfer实现了分组查询注意力（Grouped Query Attention, GQA），这是现代LLM的标配技术：\n\n- **Qwen3配置**：32个Query头，8个KV头（4:1分组比），head_dim=128\n- **内存效率**：通过共享KV头显著减少缓存占用\n- **计算优化**：融合内核减少内存访问开销\n\n### 2. Qwen3.5的混合架构支持\n\nQwen3.5采用了创新的混合设计：\n- **24层线性注意力**：使用门控Delta规则（Gated Delta Rule），降低计算复杂度\n- **8层完整注意力**：保留标准Transformer的表达能力\n- **head_dim=256**：更大的注意力维度\n\n这种混合架构兼顾了效率和性能，Pegainfer对其提供了完整支持。\n\n### 3. CUDA Graph加速\n\n对于解码阶段的重复性计算，Pegainfer利用CUDA Graph技术：\n\n- **捕获阶段**：记录解码步骤的GPU操作序列\n- **重放阶段**：以极低开销重复执行相同计算图\n- **性能提升**：显著降低CPU调度开销，提升小批次推理效率\n\n### 4. 内核融合策略\n\nPegainfer实现了多个融合内核，减少GPU内存往返：\n\n- **fused_mlp**：将SwiGLU激活、门控和投影融合为单一内核\n- **fused_attention**：融合注意力计算中的多个矩阵操作\n- **silu_mul**：SiLU激活与逐元素乘法的融合\n\n### 5. Triton AOT编译\n\n项目创新性地引入了Triton（OpenAI的GPU编译器）作为构建时工具：\n\n- **构建时编译**：使用Triton生成优化的CUDA内核\n- **运行时零依赖**：生成的内核通过C FFI与Rust代码链接\n- **灵活替换**：可逐步用Triton生成的内核替换手写CUDA\n\n当前默认启用的是`silu_mul`的Triton替换，其余内核仍来自手写CUDA代码。\n\n## 性能表现\n\n在RTX 5070 Ti（16GB显存）上的实测数据（BF16精度，CUDA Graph启用）：\n\n### Qwen3-4B\n\n| 指标 | 数值 |\n|------|------|\n| TTFT（短提示词） | ~14ms |\n| TPOT（解码阶段） | ~11ms/token |\n| 吞吐率 | ~91 tokens/秒 |\n\n### Qwen3.5-4B\n\n| 指标 | 数值 |\n|------|------|\n| TTFT（短提示词） | ~22ms |\n| TPOT（解码阶段） | ~12.2ms/token |\n| 吞吐率 | ~82 tokens/秒 |\n\n### 性能指标解读\n\n- **TTFT（Time To First Token）**：从接收提示词到生成首个token的延迟，包含分词、嵌入和前向计算\n- **TPOT（Time Per Output Token）**：解码阶段每个输出token的平均耗时\n- **Throughput**：每秒生成的token数，等于1000/TPOT\n\n这些指标在消费级显卡上表现优异，证明了纯Rust+CUDA方案的可行性。\n\n## 使用方式\n\n### 环境准备\n\n```bash\n# 创建虚拟环境\nuv venv\nsource .venv/bin/activate\n\n# 安装依赖\nuv pip install torch --index-url https://download.pytorch.org/whl/cu128\nuv pip install transformers accelerate pytest\n```\n\n### 下载模型\n\n```bash\nhuggingface-cli download Qwen/Qwen3-4B --local-dir models/Qwen3-4B\n```\n\n### 构建与运行\n\n```bash\n# 设置环境变量\nexport CUDA_HOME=/usr/local/cuda\nexport PEGAINFER_TRITON_PYTHON=.venv/bin/python\n\n# 构建（编译CUDA内核和Triton AOT）\ncargo build --release\n\n# 启动服务（默认Qwen3-4B，端口8000）\ncargo run --release\n\n# 指定模型\ncargo run --release -- --model-path models/Qwen3.5-4B\n\n# 禁用CUDA Graph（调试用途）\ncargo run --release -- --cuda-graph=false\n\n# 启用性能追踪\ncargo run --release -- --trace-output-path traces/\n```\n\n### API调用\n\nPegainfer提供OpenAI兼容的`/v1/completions`端点：\n\n```bash\n# 非流式调用\ncurl -s http://localhost:8000/v1/completions \\\n  -H "Content-Type: application/json" \\\n  -d '{"prompt": "法国的首都是", "max_tokens": 32}'\n\n# 流式输出（SSE）\ncurl -N http://localhost:8000/v1/completions \\\n  -H "Content-Type: application/json" \\\n  -d '{"prompt": "写一首关于Rust的俳句:", "max_tokens": 64, "stream": true}'\n\n# 带采样参数\ncurl -s http://localhost:8000/v1/completions \\\n  -H "Content-Type: application/json" \\\n  -d '{"prompt": "从前有座山", "max_tokens": 128, "temperature": 0.7, "top_p": 0.9}'\n```\n\n## 工程亮点\n\n### 1. 内存安全与性能并重\n\nRust的所有权系统确保了GPU内存管理的安全性：\n- 编译期检测内存泄漏和非法访问\n- 零成本抽象，不牺牲运行时性能\n- 与CUDA的FFI边界清晰可控\n\n### 2. 模块化设计\n\n代码结构清晰分离关注点：\n- `tensor.rs`：定义GPU张量抽象\n- `ops.rs`：实现底层算子\n- `model.rs` / `qwen35_model.rs`：模型逻辑\n- `http_server/`：API服务层\n\n### 3. 完善的测试体系\n\n```bash\n# 单元测试（张量、算子、配置、分词器、采样器）\ncargo test --release\n\n# E2E回归测试（需要GPU和模型权重）\nPEGAINFER_TEST_MODEL_PATH=models/Qwen3-4B cargo test --release --test e2e\n```\n\n### 4. 性能分析工具\n\n内置Chrome Trace JSON格式的性能追踪，可在Perfetto UI中可视化分析：\n- 识别热点内核\n- 分析流水线瓶颈\n- 优化内存访问模式\n\n## 技术价值与启示\n\nPegainfer项目为LLM推理领域提供了多个有价值的启示：\n\n**框架最小化**：证明了不依赖PyTorch等重量级框架，仅用Rust+CUDA也能实现生产级推理性能。这为资源受限场景（边缘设备、嵌入式系统）提供了新思路。\n\n**教育价值**：精简的代码量使其成为学习LLM推理原理的理想教材。从分词到采样，每个环节都清晰可见。\n\n**Rust在AI领域的潜力**：展示了Rust的内存安全和零成本抽象在GPU编程中的优势，为AI基础设施的Rust化提供了参考。\n\n**渐进式优化路径**：通过Triton AOT的引入，展示了如何在不重写全部代码的情况下，逐步采用更先进的编译技术。\n\n## 局限与未来方向\n\n当前Pegainfer的局限包括：\n- 仅支持Qwen3系列模型\n- 量化支持有限（当前仅BF16）\n- 批处理推理能力待完善\n- Windows支持尚处于实验阶段\n\n未来可能的发展方向：\n- 支持更多模型架构（Llama、Mistral等）\n- 引入INT8/INT4量化\n- 多GPU并行支持\n- 更完善的调度策略\n\n## 总结\n\nPegainfer是一个技术纯粹主义者的杰作，它用约1万行代码（Rust+CUDA）实现了完整的LLM推理引擎。虽然在功能丰富度上无法与vLLM、TensorRT-LLM等成熟方案竞争，但其在\"从零构建\"理念下的工程实现，为理解LLM推理原理、探索Rust在AI领域的应用提供了宝贵参考。\n\n对于希望深入理解Transformer推理机制、学习GPU内核优化、或探索Rust在AI基础设施中应用的开发者，Pegainfer是一个值得深入研究的精品项目。其代码质量、架构设计和性能表现，都体现了作者深厚的系统编程功底。\n\n项目采用MIT许可证开源，代码和详细文档可在GitHub获取。
