跟着 Claude Code,我用 Python 实现了一个 轻量级的vLLM

这是一个普通开发者借助 Claude Code 协作，从零学习 LLM 推理引擎的记录。很多实现细节是在 Claude Code 的帮助下完成的，核心价值在于学习过程本身。

vLLM 是目前最流行的 LLM 推理框架之一，但它的代码量庞大、依赖复杂。当我想深入理解推理引擎内部原理时，我发现自己很难从头到尾读懂它的源码。

于是想：能不能用更少的代码，来学习 vLLM 的核心功能？ 很多架构设计和实现细节，都是在和 Claude Code 的对话中逐步明确的。我只是一个愿意动手尝试的学习者。

目前 支持 7 个模型(仅单卡推理)：LLaMA-3、Qwen2/2.5/3、Qwen3 MoE、Qwen 3.5（混合注意力）和 Gemma3。

目前不支持 分布式并行、量化、视觉语言大模型，多模态大模型

项目最终的组织方式是这样的：

mini-vllm/├── setup.py                        # CUDA 算子编译入口├── pyproject.toml                  # 项目元数据、依赖├── csrc/                           # CUDA/C++ 算子源码│   ├── bindings/│   │   └── ops_binding.cpp         # PyTorch 算子注册│   ├── include/mini_vllm/│   │   ├── utils.cuh               # CUDA 工具函数│   │   └── dispatch.h              # 类型分发宏│   └── kernels/                    # CUDA kernel 实现│       ├── activation/             # SiLU / GeLU 激活函数│       ├── embedding/              # 词嵌入│       ├── mlp/                    # Fused MLP（SwiGLU / GeGLU）│       ├── rmsnorm/                # RMSNorm（标准/plus_one/fused_add）│       ├── rope/                   # RoPE 旋转位置编码│       ├── sampling/               # 采样（greedy/top_k_top_p/softmax_multinomial）│       └── softmax/                # Softmax│├── mini_vllm/                      # Python 推理框架│   ├── engine/                     # 推理引擎│   │   ├── engine.py               # LLMEngine 核心引擎│   │   ├── scheduler.py            # 请求调度器│   │   ├── sequence.py             # 序列状态管理│   │   ├── kv_cache.py             # 分页 KV Cache 管理│   │   └── sampler.py              # 采样器（CUDA/PyTorch）│   ├── models/                     # 模型定义（PyTorch + CUDA）│   │   ├── loader.py               # 模型加载 + 自动检测 + CUDA/PyTorch 自动选择│   │   ├── model.py                # 通用模型基类（PyTorch）│   │   ├── model_cuda.py           # 通用模型基类（CUDA 加速）│   │   ├── standard_models.py      # LLaMA/Qwen2/2.5/3 子类│   │   ├── gemma3_model.py         # Gemma 3（PyTorch）│   │   ├── gemma3_cuda.py          # Gemma 3（CUDA）│   │   ├── qwen3_moe_model.py      # Qwen3 MoE（PyTorch）│   │   ├── qwen3_moe_cuda.py       # Qwen3 MoE（CUDA）│   │   └── qwen3_5_model.py        # Qwen 3.5 混合注意力（仅 PyTorch）│   ├── ops/                        # 算子封装层（自动 CUDA/PyTorch 切换）│   │   ├── _utils.py               # is_ext_loaded() 检测│   │   ├── activation.py           # SiLU / GeLU│   │   ├── embedding.py            # Embedding│   │   ├── mlp.py                  # Fused MLP（SwiGLU / GeGLU）│   │   ├── rmsnorm.py              # RMSNorm│   │   ├── rope.py                 # RoPE│   │   ├── sampling.py             # Greedy/Top-K+Top-P/Softmax+Multinomial│   │   └── softmax.py              # Softmax│   ├── server/                     # HTTP API 服务│   │   ├── api_server.py           # FastAPI（OpenAI 兼容接口）│   │   └── openai_types.py         # OpenAI 协议数据类型│   ├── tokenizer/│   │   └── tokenizer.py           │   ├── distributed/│   │   └── tensor_parallel.py      # 张量并行（预留）│   └── utils/│       ├── config.py           │       └── logger.py               # 日志│├── tests/                          # 测试│   ├── test_qwen3_moe.py           # Qwen3 MoE 测试│   ├── test_config.py              # 配置测试│   ├── test_cuda_rmsnorm.py        # CUDA RMSNorm 算子测试│   ├── test_engine.py              # 引擎测试│   ├── test_kv_cache.py            # KV Cache 测试│   ├── test_sampler.py             # 采样器测试│   ├── test_scheduler.py           # 调度器测试│   ├── test_sequence.py            # 序列状态测试│   ├── test_e2e_all_models.py      # 全模型端到端测试│├── examples/                       # 示例│   ├── offline_inference.py        # 离线推理（非流式 + 流式）│   └── openai_client.py            # OpenAI API 客户端│├── scripts/                        # 测试脚本│   ├── test_offline.py             # 离线推理测试│   └── test_online.py              # 在线服务测试

目录结构一目了然：

• • csrc/ 放 CUDA 算子源码
• • mini_vllm/engine/ 是推理引擎核心
• • mini_vllm/models/ 是模型定义，每个模型都有 PyTorch 和 CUDA 两个版本
• • mini_vllm/ops/ 是算子封装层，负责运行时自动切换后端
• • mini_vllm/server/ 提供 OpenAI 兼容的 HTTP API

三层架构

mini-vllm 采用了经典的三层架构：

┌─────────────────────────────────────────┐│           HTTP API 层 (FastAPI)          ││   /v1/chat/completions  /v1/completions  ││   支持 streaming 和 non-streaming         │└──────────────────┬──────────────────────┘                   ▼┌─────────────────────────────────────────┐│           推理引擎层 (LLMEngine)          ││   Scheduler + KV Cache + Sampler         ││   threading.Lock 线程安全                 │└──────────────────┬──────────────────────┘                   ▼┌─────────────────────────────────────────┐│     模型层 + 算子层 (PyTorch / CUDA)      ││   自动检测 CUDA 扩展，无缝切换后端          │└─────────────────────────────────────────┘

推理引擎层：三大核心组件

LLMEngine 是整个系统的核心，负责协调三个关键组件：

• • Scheduler：管理 waiting/running 队列，实现连续批处理
• • KVCacheManager：分页管理 GPU 显存，按 block 分配和释放
• • Sampler：支持 greedy、temperature、top-k、top-p 采样策略

引擎通过 threading.Lock 保证线程安全。generate() 方法内部自动持锁；generate_stream() 则要求调用方持锁，因为需要跨 yield 保持状态一致性。

模型层与算子层：双版本自动切换

每个算子都有 PyTorch 和 CUDA 两个版本，运行时自动切换。

模型加载时，loader.py 会检查 mini_vllm._C 扩展模块是否可用：

• • 可用 → 选择 *_cuda.py 模型文件，使用 CUDA kernel
• • 不可用 → 回退到 *.py 模型文件，使用 PyTorch 原生实现

这样在没有 GPU 的机器上也能开发调试，编译了 CUDA 算子后又能获得加速。

核心机制一：分页 KV Cache

LLM 推理中，KV Cache 是最大的显存消耗者。每生成一个 token，都需要把当前层的 Key 和 Value 向量缓存下来，供后续 token 做注意力计算时使用。传统做法是为每个序列预分配一大块连续显存——按最大序列长度分配。但实际生成的长度往往远小于最大长度，造成严重的显存浪费。

举个例子：如果模型的最大序列长度是 4096，你就要为每个请求预留 4096 个 token 的 KV Cache 空间。但如果这个请求只生成了 100 个 token，那剩下 3996 个 token 的空间就白白浪费了。

vLLM 论文中的 PagedAttention 是借鉴操作系统的虚拟内存思想，把 KV Cache 切分成固定大小的 block，按需分配。

具体来说：

• • 每个 block 存储固定数量的 token 的 K/V 向量（比如 16 个 token）
• • 序列生成时，每需要更多空间就动态申请一个新 block
• • 序列结束后释放所有 block（free() ，在 finally 块中无条件调用，确保不会泄漏）

这种设计的好处是：

1. 1. 消除内部碎片：不需要按最大长度预分配，按实际需要申请
2. 2. 减少外部碎片：block 大小固定，不存在内存碎片问题
3. 3. 天然支持不同长度的序列混合批处理：每个序列的 block 数量可以不同

核心机制二：连续批处理

传统的 LLM 推理批处理是静态的：凑够 N 个请求一起处理，等最长的那个生成完才释放。这在请求长度差异大时效率很低——短请求早就生成完了，但必须等长请求一起结束。

连续批处理（Continuous Batching）的核心想法是：prefill 和 decode 阶段可以在同一个 batch 中混合执行。

调度器维护两个队列：

• • waiting 队列：新请求等待 prefill（处理完整 prompt）
• • running 队列：正在 decode（逐 token 生成）的请求

每一步调度器都可以：

• • 把 waiting 队列中的新请求插入到当前 batch（做 prefill）
• • 把某个序列生成完成后将其从 running 队列移出
• • 把 waiting 队列中的请求提升到 running 队列

这样 GPU 始终保持忙碌，不会因为某个长序列而阻塞整个 batch。短请求生成完就立即释放资源，长请求可以继续生成。

核心机制三：CUDA/PyTorch 自动切换

整个项目有多个CUDA 算子，覆盖七大类：

CUDA kernel 的编写是 Claude Code 帮助最多的部分。我负责理解算子的数学原理，Claude Code 负责生成 kernel 代码和调试。

核心机制四：权重自动检测

加载 HuggingFace 模型时，config.json 中的信息有时不完整或不准确。比如某些模型的 config 写着 attention_bias: false，但实际权重里有 bias。

我们采用了直接检查权重文件的 key 来推断架构特性：

# 检测 attention biasq_bias_key = "model.layers.0.self_attn.q_proj.bias"if q_bias_key in state_dict:    model_config["attention_bias"] = True# 检测 QK-norm（Qwen3、Gemma 3 等）q_norm_key = "model.layers.0.self_attn.q_norm.weight"if q_norm_key in state_dict:    model_config["_qk_norm"] = True# 检测 feedforward layernorms（Gemma 3）pre_ff_key = "model.layers.0.pre_feedforward_layernorm.weight"if pre_ff_key in state_dict:    model_config["_has_feedforward_layernorms"] = True

这种方法比单纯依赖 config 更可靠，因为权重文件是模型的真实状态，不会骗人。

核心机制五：Qwen3.5 混合注意力

Qwen 3.5 的混合架构是我遇到的最大挑战。

它采用了一种混合设计：部分层使用标准 Full Attention（配合 KV Cache），部分层使用 Gated Delta Net 线性注意力（不用 KV Cache，用循环状态）。这两种注意力的计算方式完全不同，需要在同一个模型里同时支持。

这部分的实现几乎完全依赖 Claude Code。我负责理解混合注意力的设计思路，Claude Code 负责将论文中的数学公式转化为可运行的代码。

每个 layer 根据 _layer_types 路由到不同的计算路径：

• • Full Attention 层：标准的 QKV 投影 + RoPE + KV Cache + SDPA
• • Linear Attention 层：QKV 投影 + Conv1d + chunked delta rule（prefill）/ 循环状态更新（decode）

线性注意力层的状态管理包括三个部分：

每次 generate() 调用前自动重置线性注意力缓存，确保不同请求之间的状态隔离。

快速上手

安装

# 安装依赖pip install torch>=2.1.0 safetensors>=0.4.0 tokenizers>=0.15.0pip install fastapi>=0.104.0 uvicorn>=0.24.0 pydantic>=2.0.0 transformers>=4.36.0# 开发模式安装（不编译 CUDA）pip install -e . --no-build-isolation# 编译 CUDA 算子python setup.py build_ext --inplace

离线推理

from mini_vllm.engine.engine import LLMEnginefrom mini_vllm.utils.config import EngineConfig, SamplingParamsengine = LLMEngine(EngineConfig(model_path="/path/to/model"))# 非流式results = engine.generate(    [{"role": "user", "content": "你是谁"}],    SamplingParams(temperature=0.7, max_tokens=512),)print(results[0].output_text)# 流式for idx, token_text, finished in engine.generate_stream(    [{"role": "user", "content": "用Python写一个快速排序"}],    SamplingParams(temperature=0.3, max_tokens=512),):    print(token_text, end="", flush=True)

启动 API 服务

python -m mini_vllm.server.api_server --model-path /path/to/model

然后用 OpenAI SDK 连接：

from openai import OpenAIclient = OpenAI(base_url="http://localhost:8000/v1", api_key="")response = client.chat.completions.create(    model="your-model",    messages=[{"role": "user", "content": "你是谁"}],)print(response.choices[0].message.content)