tiny-llm - 一周学会大模型推理服务

这是一门面向系统工程师的课程，使用 MLX 实现大模型推理服务。整个代码库几乎完全基于 MLX 的数组/矩阵 API，不依赖任何高级神经网络框架，以便我们能够从零开始构建模型推理基础设施，并深入探索各种优化技术。

我们的目标是掌握高效部署和运行大型语言模型（例如 Qwen2 系列）的核心技术。

在第 1 周，你将用纯 Python 实现必要的组件，以利用 Qwen2 模型生成响应（如注意力机制、RoPE 等）。第 2 周，你将实现一个类似于 vLLM 的推理系统，但更加简单（如 KV 缓存、连续批处理、Flash Attention 等）。第 3 周，我们将探讨更高级的主题，以及模型如何与外部世界交互。

为什么选择 MLX？因为如今搭建基于 macOS 的本地开发环境比配置 NVIDIA GPU 更加容易。

为什么选择 Qwen2？这是我最早接触的大模型之一——它也是 vLLM 文档中的典型示例。我曾仔细研究过 vLLM 的源码，并积累了一些相关知识。

书籍

tiny-llm 的配套书籍已发布在 https://skyzh.github.io/tiny-llm/。你可以按照指南逐步构建自己的项目。

社区

欢迎加入 skyzh 的 Discord 社区，与 tiny-llm 社区一起学习交流。

路线图

目前第 1 周和第 2 周的内容已完成，第 3 周正在推进中。

周次 + 章节	主题	代码	测试	文档
1.1	注意力机制	✅	✅	✅
1.2	RoPE 位置编码	✅	✅	✅
1.3	分组查询注意力机制	✅	✅	✅
1.4	RMSNorm 归一化与 MLP 层	✅	✅	✅
1.5	加载模型	✅	✅	✅
1.6	生成响应（即解码过程）	✅	✅	✅
1.7	采样策略	✅	✅	✅
2.1	键值缓存	✅	✅	✅
2.2	CPU 上的量化矩阵乘法与线性变换	✅	✅	✅
2.3	GPU 上的量化矩阵乘法与线性变换	✅	✅	✅
2.4	CPU 上的 Flash Attention 2	✅	✅	✅
2.5	GPU 上的 Flash Attention 2	✅	✅	✅
2.6	连续批处理	✅	✅	✅
2.7	分块预填充	✅	✅	✅
3.1	分页注意力机制 - 第一部分	🚧	🚧	🚧
3.2	分页注意力机制 - 第二部分	🚧	🚧	🚧
3.3	MoE（专家混合模型）	🚧	🚧	🚧
3.4	推测解码	🚧	✅	🚧
3.5	RAG 管道	🚧	🚧	🚧
3.6	AI 助手 / 工具调用	🚧	🚧	🚧
3.7	长上下文支持	🚧	🚧	🚧

其他未涵盖的主题：量化/压缩的 KV 缓存、前缀/提示缓存；采样、微调；更小的核函数（softmax、silu 等）。

星标历史

tiny-llm 快速上手指南

tiny-llm 是一个面向系统工程师的 LLM 服务课程项目。它仅基于 MLX 数组/矩阵 API（不使用高级神经网络框架），旨在帮助开发者从零构建模型服务基础设施，深入理解注意力机制、KV Cache、连续批处理等核心优化技术。

环境准备

系统要求

• 操作系统：macOS (推荐 Apple Silicon M1/M2/M3 芯片)
  • 注：该项目主要基于 Apple MLX 框架设计，旨在利用 macOS 本地开发环境替代复杂的 NVIDIA GPU 配置。
• 硬件：具备统一内存架构的 Mac 电脑（内存建议 8GB 以上，运行大模型需更大内存）

前置依赖

• Python: 3.9 或更高版本
• 包管理工具: pip 或 conda
• MLX 框架: 项目核心依赖

安装步骤

1. 克隆项目代码

   git clone https://github.com/skyzh/tiny-llm.git
   cd tiny-llm
   

2. 创建虚拟环境并安装依赖

   python3 -m venv venv
   source venv/bin/activate
   pip install -r requirements.txt
   

   如果 requirements.txt 未自动包含 MLX，请手动安装：

   pip install mlx
   

3. 下载模型权重 本项目默认使用 Qwen2 系列模型。你需要从 Hugging Face 下载对应的权重文件并放置于项目指定目录（通常为 weights/ 或根据代码指引配置）。

   国内加速方案：由于 Hugging Face 在国内访问受限，建议使用镜像站下载：

   # 使用 huggingface.co 镜像 (例如 hf-mirror.com)
   export HF_ENDPOINT=https://hf-mirror.com
   
   # 使用 huggingface-cli 下载 Qwen2 模型 (以 Qwen2-0.5B-Instruct 为例)
   huggingface-cli download --resume-download Qwen/Qwen2-0.5B-Instruct --local-dir weights/qwen2-0.5b
   

基本使用

tiny-llm 是一个分周进度的学习项目，目前 Week 1 和 Week 2 的核心功能已完成。以下是最基础的模型加载与文本生成示例（基于 Week 1 完成后的状态）。

运行推理脚本

在项目根目录下，运行 Python 脚本加载模型并生成回复。假设你已经完成了 Week 1 的代码实现（或使用了作者提供的参考实现）：

python generate.py \
    --model-path weights/qwen2-0.5b \
    --prompt "你好，请介绍一下你自己。" \
    --max-tokens 100

代码逻辑简述

该工具不依赖 transformers 等高层库，而是直接使用 MLX 原生 API 构建计算图。核心流程如下：

1. Load Model: 读取 safetensors 权重并映射为 MLX arrays。
2. Preprocessing: 对输入 prompt 进行 Tokenization。
3. Forward Pass: 依次执行 Embedding -> RMSNorm -> Attention (RoPE) -> MLP。
4. Sampling: 从输出概率分布中采样下一个 Token。
5. Decoding: 循环上述过程直到生成结束符或达到最大长度。

进阶功能测试 (Week 2)

若已完成后两周内容，可体验带 KV Cache 和连续批处理的高效推理：

python serve.py \
    --model-path weights/qwen2-0.5b \
    --batch-size 4 \
    --enable-flash-attn

提示：详细的教学文档、每章的代码实现细节及原理解析，请访问官方文档：https://skyzh.github.io/tiny-llm/