vLLM 代码结构分析

1. 顶层目录一览

目录	用途
vllm/	主 Python 包：引擎、调度器、模型实现、Worker、采样器、API 入口等所有 Python 源码
csrc/	C++/CUDA 核心：PagedAttention 内核、KV cache 操作、MoE、量化 GEMM、自定义通信等
rust/	Rust 工作空间：推理 CLI 前端、推理解析器、分词器、聊天渲染、metrics 收集
tests/	测试套件：按功能区域（kernel、模型、量化、分布式、V1 等）组织
docs/	文档源文件：Sphinx/mkdocs 格式
benchmarks/	基准测试：延迟/吞吐量评估脚本
examples/	示例：OpenAI API、离线推理、LLM 类使用模式
docker/	Docker 部署：CUDA、ROCm、CPU 构建的 Dockerfile
cmake/	CMake 辅助模块：扩展构建宏、HIP 转换、外部项目集成
requirements/	pip 依赖：按平台分列（cuda.txt、rocm.txt、cpu.txt、common.txt）
scripts/	实用脚本：安装/构建/发布
tools/	开发者工具：安装 DeepGEMM、FlashInfer 等第三方编译依赖
.buildkite/	Buildkite CI：管道定义、硬件测试、发布管道
.github/	GitHub Actions：pre-commit、macOS smoke test、stale 标记

---

2. 构建系统

2.1 构建后端

pyproject.toml — 使用 setuptools.build_meta 作为构建后端，核心构建依赖为：

cmake >= 3.26.1
ninja
setuptools >= 77.0.3
setuptools-scm >= 8.0        (版本号从 git tag 派生)
setuptools-rust >= 1.9.0     (Rust 扩展构建)
torch == 2.11.0              (作为构建依赖引入，供 CMake 检测)

版本号由 setuptools-scm 自动从 git 标签生成（pyproject.toml:50-51）。CLI 入口点 vllm 指向 entrypoints/cli/main.py（pyproject.toml:44）。

2.2 setup.py：构建协调者

setup.py（~1207 行）是整个构建的神经中枢，负责：

a) 设备自动检测（setup.py:57-80）

自动设置 VLLM_TARGET_DEVICE 环境变量：

• 从 torch 报告检测：cuda → cuda、rocm → rocm、xpu → xpu、其他 → cpu
• 也支持 VLLM_TARGET_DEVICE=empty 用于纯 Python 安装

b) CMakeExtension：编译 C++/CUDA

setup.py:160-163 — CMakeExtension 继承 setuptools.Extension，表明这是一个通过 CMake 编译的 C/C++ 扩展。

setup.py:166-381 — cmake_build_ext 自定义 setuptools 命令：

cmake_build_ext:
  ├─ 调用 CMake 配置 + 构建 (make -j)
  ├─ 将生成的 .abi3.so 文件复制到 vllm/ 目录
  ├─ CPU 构建时捆绑 tcmalloc
  └─ 可编辑安装时复制 flash-attn/triton Python 文件

c) 扩展注册（setup.py:1037-1093）

根据目标设备注册不同的扩展模块：

扩展模块	行号	说明
vllm._C	L1091	主扩展：PagedAttention、KV cache、通信、多数量化内核
vllm._C_stable_libtorch	L1093	稳定 ABI 扩展：CUTLASS scaled MM、AWQ、GPTQ、GGUF 等
vllm._moe_C	L1040	MoE 扩展：top-k、moe_align、moe_wna16 等
vllm._rocm_C	L1049	ROCm 扩展：ROCm 专用内核（仅 HIP 平台）
vllm.vllm_flash_attn._vllm_fa2_C	L1052	FlashAttention v2 内核
vllm.vllm_flash_attn._vllm_fa3_C	L1053	FlashAttention v3 内核（需 CUDA 12.3+）
vllm._flashmla_C	L1069	FlashMLA 内核（需 CUDA 12.9+）
vllm._deep_gemm_C	L1078	DeepGEMM 内核（可选）
vllm.cumem_allocator	L1041	CUDA 虚拟内存分配器
vllm.spinloop	L1046	CPU spinloop 自旋锁

CPU 平台额外扩展：_C_AVX512、_C_AVX2（setup.py:1083-1088）。

d) Rust 扩展（setup.py:1149-1158）

RustExtension(
    target="vllm.vllm-rs",
    path="rust/src/cmd/Cargo.toml",
    args=["--bin", "vllm-rs"],
    binding=Binding.Exec,       # 独立二进制，非 Python 扩展
    optional=True,              # 可选安装
)

e) 预编译扩展支持（setup.py:412-861）

当无法本地编译时，自动从 wheels.vllm.ai 下载预编译的 .abi3.so 文件。支持的预编译文件（setup.py:716-733）：

vllm/_C.abi3.so
vllm/_C_stable_libtorch.abi3.so
vllm/_moe_C.abi3.so
vllm/_flashmla_C.abi3.so
vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so
vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so
vllm/cumem_allocator.abi3.so
vllm/spinloop.abi3.so

2.3 CMakeLists.txt：C++/CUDA 编译

CMakeLists.txt（~1312 行）是主要的 C++/CUDA 构建文件：

• CMake 项目名：vllm_extensions，C++20 标准（L14）
• 目标设备：由 -DVLLM_TARGET_DEVICE 控制（cuda/rocm/cpu）（L25）
• Python 发现：通过 VLLM_PYTHON_EXECUTABLE 精确匹配 Python 环境（L71-76）
• Torch 集成：find_package(Torch REQUIRED)（L95）
• CUDA 架构控制：从 torch 的编译标志中提取架构信息，按文件粒度设置 CUDA 架构以减少编译时间（L180-337）
• CUTLASS 集成：通过 FetchContent 从 GitHub 获取 CUTLASS v4.4.2（L329-357）
• 外部项目（L1298-1311）：

外部项目	平台	说明
triton_kernels.cmake	CUDA + HIP	Triton kernel 编译
deepgemm.cmake	CUDA only	DeepSeek FP8 GEMM 库
flashmla.cmake	CUDA only	Flash MLA attention
qutlass.cmake	CUDA only	Quantized CUTLASS 集成
vllm_flash_attn.cmake	CUDA only	vLLM 定制的 FlashAttention

2.4 define_extension_target 宏

cmake/utils.cmake:549-612 — 定义 CMake 编译目标的中心宏：

define_extension_target(target_name):
  ├─ Python_add_library(target MODULE)           # 创建 Python 可加载的 .so
  ├─ 若语言=HIP: 执行 hipify 预处理 (.cu → .hip)
  ├─ 设置 include dirs, compile flags, arch
  ├─ 链接 torch, CUDA::cudart
  └─ install(TARGETS target DESTINATION vllm/)    # 安装到 vllm/ 包目录

---

3. C++/CUDA 代码（csrc/）

3.1 目录结构

csrc/
├── torch_bindings.cpp        ← 主算子注册（~50 个算子）
├── cache_kernels.cu          ← KV cache 写入/读写内核
├── cache_kernels_fused.cu    ← 融合 KV cache 操作
├── custom_all_reduce.cu/.cuh ← 自定义全规约通信
├── custom_quickreduce.cu     ← 快速规约 (ROCm)
├── cumem_allocator.cpp       ← CUDA 虚拟内存分配器
├── spinloop.cpp              ← CPU 自旋锁
├── sampler.cu / topk.cu      ← 采样/top-k 内核
├── ops.h / cache.h           ← 算子/缓存公共头文件
├── cuda_utils.h / dispatch_utils.h  ← CUDA 工具/分发
│
├── attention/                ← PagedAttention 内核
│   ├── paged_attention_v1.cu
│   ├── paged_attention_v2.cu
│   ├── attention_kernels.cuh  ← 核心 attention 算法
│   ├── attention_utils.cuh    ← Qk_dot 等辅助结构
│   └── merge_attn_states.cu   ← Cascade attention 合并
│
├── core/                     ← 基础设施
│   ├── registration.h        ← REGISTER_EXTENSION 宏
│   ├── scalar_type.hpp
│   └── batch_invariant.hpp
│
├── cpu/                      ← CPU 内核
│   ├── cpu_attn.cpp          ← CPU attention
│   └── 融合 MoE、MLA、norm、RoPE 等
│
├── cutlass_extensions/       ← CUTLASS 扩展
│   └── scaled_mm_epilogues   ← 缩放 MM epilogue
│
├── libtorch_stable/          ← 稳定 ABI 算子
│   ├── torch_bindings.cpp    ← STABLE_TORCH_LIBRARY_FRAGMENT
│   └── cutlass_scaled_mm、awq、gptq、fp4、w8a8 等
│
├── mamba/                    ← Mamba SSM 内核
│   └── selective_scan_fwd.cu
│
├── moe/                      ← MoE 内核
│   ├── torch_bindings.cpp    ← MoE 算子注册
│   ├── topk_softmax、moe_align、moe_wna16
│   ├── marlin_moe_wna16      ← Marlin 量化 MoE
│   ├── dsv3_router_gemm      ← DeepSeek V3 路由器
│   └── mxfp8_moe             ← SM100+ MXFP8 分组 GEMM
│
├── quantization/             ← 量化内核
│   ├── marlin/               ← GPTQ/AWQ/FP8/NVFP4 Marlin GEMM
│   ├── machete/              ← Hopper W4A16 GEMM
│   ├── gguf/                 ← GGUF 去量化
│   ├── w8a8/                 ← FP8/int8 缩放 MM
│   └── fused_kernels/        ← 融合 SiLU+multiply+block quant
│
├── quickreduce/              ← ROCm QuickReduce
└── rocm/                     ← ROCm 专用内核 + torch_bindings

3.2 算子注册机制

vLLM 使用 两组并行 的 PyTorch 算子注册路径：

路径 A：TORCH_LIBRARY_EXPAND（不稳定 ABI）

用于 _C 和 _moe_C 扩展。

csrc/torch_bindings.cpp:18 — 主算子注册：

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("paged_attention_v1(Tensor out, Tensor query, ...) -> ()");
  ops.impl("paged_attention_v1", torch::kCUDA, &paged_attention_v1);
  // ...约 50 个算子
}

算子分布在四个命名空间： | 命名空间 | 功能 | |———-|——| | ops | PagedAttention、position encoding、norm、量化 GEMM | | _cache_ops | KV cache 读写（reshape_and_cache 等） | | _cuda_utils | CUDA 设备属性查询 | | _custom_ar | 自定义全规约通信 |

csrc/moe/torch_bindings.cpp — MoE 算子注册：

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
  m.def("topk_softmax", ...);
  m.def("moe_align_block_size", ...);
  m.def("moe_wna16_gemm", ...);
}

csrc/rocm/torch_bindings.cpp — ROCm 算子注册（类似模式）。

路径 B：STABLE_TORCH_LIBRARY_FRAGMENT（稳定 ABI）

用于 _C_stable_libtorch 扩展，兼容 PyTorch >= 2.10 的稳定 ABI。

csrc/libtorch_stable/torch_bindings.cpp:9：

STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
  ops.def("cutlass_scaled_mm", ...);
  ops.def("cutlass_fp4_moe_mm", ...);
  // AWQ、GPTQ、GGUF、FP4 等算子
}

这里使用 #include <torch/csrc/stable/library.h>，生成的 .abi3.so 在不同 PyTorch 版本间兼容。

扩展加载入口

csrc/core/registration.h:22-27：

#define REGISTER_EXTENSION(name)
extern "C" PYBIND11_EXPORT PyObject* PyInit_##name() {
  // 初始化 torch 库，返回 PyModuleDef
}

TORCH_EXTENSION_NAME 由 CMake 定义——_C 生成 PyInit__C()，_moe_C 生成 PyInit__moe_C()。这使得 Python 可以通过 import vllm._C 直接加载扩展。

3.3 Python 侧调用方式

平台层加载（platforms/interface.py:242-250）：

class Platform:
    def import_kernels(self):
        import vllm._C       # 加载主扩展 → 注册 torch.ops._C.*
        import vllm._moe_C    # 加载 MoE 扩展 → 注册 torch.ops._moe_C.*

各平台覆盖此方法：

• CUDA (cuda.py)：默认行为，加载 vllm._C + vllm._moe_C
• ROCm (rocm.py)：额外加载 vllm._rocm_C
• CPU (cpu.py)：加载 vllm._C（映射到 _C_AVX2/_C_AVX512）

算子包装（_custom_ops.py，3900+ 行）—— 所有 Python 代码通过此文件调用编译后的内核：

# _custom_ops.py:20
current_platform.import_kernels()

# _custom_ops.py:135 — PagedAttention V1
def paged_attention_v1(out, query, key_cache, value_cache, ...):
    torch.ops._C.paged_attention_v1(out, query, key_cache, ...)

# _custom_ops.py:182 — PagedAttention V2
def paged_attention_v2(out, exp_sum, max_logits, tmp, query, ...):
    torch.ops._C.paged_attention_v2(out, exp_sum, max_logits, ...)

# MoE 算子通过 torch.ops._moe_C 调用
# ROCm 算子通过 torch.ops._rocm_C 调用

文件按功能管理： | 功能区域 | 行号范围 | 算子示例 | |———-|———|———| | Attention | L135-290 | paged_attention_v1/v2、merge_attn_states | | Position Encoding | L299-360 | rotary_embedding | | Normalization | L316-328 | rms_norm、fused_add_rms_norm | | 量化 GEMM | L600-1600 | awq_gemm、gptq_gemm、marlin_gemm、cutlass_scaled_mm | | FP4 | L1600-1900 | scaled_fp4_quant、cutlass_fp4_moe_mm | | FP8 | L1905+ | scaled_fp8_quant | | MoE | L2242+ | topk_softmax、moe_align_block_size、dsv3_router_gemm | | 自定义 AllReduce | — | torch.ops._C.custom_ar.* |

---

4. Rust 代码（rust/）

4.1 工作空间

rust/Cargo.toml — Rust 工作空间定义：

包	路径	用途
cmd	rust/src/cmd/	vllm-rs 二进制：vllm serve、vllm complete 等 CLI 入口
chat	rust/src/chat/	Jinja2 聊天模板渲染引擎
engine-core-client	rust/src/engine-core-client/	vLLM EngineCore 的 gRPC 客户端
llm	rust/src/llm/	LLM 高层接口：请求、输出、统计
managed-engine	rust/src/managed-engine/	托管（远程）引擎客户端
metrics	rust/src/metrics/	Prometheus metrics 的 Rust 收集端（请求、调度、API）
reasoning-parser	rust/src/reasoning-parser/	推理内容解析：DeepSeek R1、Kimi、Qwen3、Gemma 4 等
server	rust/src/server/	前端构建脚本
text	rust/src/text/	文本处理、输出解码
tokenizer	rust/src/tokenizer/	基于 HuggingFace tokenizers (fastoken) 的增量分词
tool-parser	rust/src/tool-parser/	工具调用格式解析

4.2 构建

由 build_rust.sh 驱动：

• 安装 rustup
• 读取 rust-toolchain.toml（通道 1.95）
• 构建 vllm-rs 二进制文件
• 复制到 vllm/vllm-rs

通过 setuptools-rust 集成到 Python 包：

RustExtension(
    target="vllm.vllm-rs",
    path="rust/src/cmd/Cargo.toml",
    args=["--bin", "vllm-rs"],
    binding=Binding.Exec,  # 独立二进制，非 Python 扩展 .so
    optional=True,
)

4.3 Python 如何调用 Rust

• CLI 入口：vllm-rs 二进制作为独立的 CLI 替代入口，pyproject.toml 的 [project.scripts] 指向 vllm = vllm.entrypoints.cli.main:main，该 Python 入口在 Rust 前端可用时委托给它
• 推理解析器：reasoning-parser 包通过 FFI（ctypes/cffi）从 Python 调用，提供高性能的推理内容解析能力
• gRPC 客户端：engine-core-client 提供 Rust 侧的 gRPC 客户端与 EngineCore 通信

---

5. 完整构建链路

pyproject.toml (setuptools.build_meta)
  │
  └─ setup.py
       │
       ├─ [setuptools] ── RustExtension ──→ rust/ (cargo build)
       │                                      └─ vllm/vllm-rs (独立二进制)
       │
       ├─ [cmake_build_ext] ──→ CMakeLists.txt
       │    │
       │    ├─ define_extension_target(_C)
       │    │   ├─ csrc/torch_bindings.cpp (TORCH_LIBRARY_EXPAND)
       │    │   ├─ csrc/attention/*.cu        — PagedAttention
       │    │   ├─ csrc/cache_kernels.cu      — KV cache 操作
       │    │   ├─ csrc/custom_all_reduce.cu  — 自定义通信
       │    │   ├─ csrc/quantization/marlin/  — Marlin GEMM
       │    │   ├─ csrc/quantization/machete/ — Machete GEMM
       │    │   ├─ csrc/moe/                  — MoE（部分）
       │    │   └─ → vllm/_C.abi3.so
       │    │
       │    ├─ define_extension_target(_C_stable_libtorch)
       │    │   ├─ csrc/libtorch_stable/torch_bindings.cpp
       │    │   ├─ cutlass_scaled_mm, fp4, awq, gptq, gguf, w8a8
       │    │   └─ → vllm/_C_stable_libtorch.abi3.so
       │    │
       │    ├─ define_extension_target(_moe_C)
       │    │   ├─ csrc/moe/torch_bindings.cpp
       │    │   ├─ topk_softmax, moe_align, moe_wna16
       │    │   └─ → vllm/_moe_C.abi3.so
       │    │
       │    ├─ define_extension_target(_rocm_C) [HIP only]
       │    │   ├─ csrc/rocm/torch_bindings.cpp
       │    │   └─ → vllm/_rocm_C.abi3.so
       │    │
       │    ├─ FetchContent: vllm_flash_attn
       │    │   └─ → vllm/vllm_flash_attn/_vllm_fa{2,3}_C.abi3.so
       │    │
       │    ├─ FetchContent: flashmla
       │    │   └─ → vllm/_flashmla_C.abi3.so
       │    │
       │    ├─ FetchContent: deepgemm
       │    │   └─ → vllm/_deep_gemm_C.abi3.so
       │    │
       │    ├─ define_extension_target(cumem_allocator)
       │    │   └─ → vllm/cumem_allocator.abi3.so
       │    │
       │    └─ define_extension_target(spinloop)
       │         └─ → vllm/spinloop.abi3.so
       │
       └─ [setuptools] ── Python 纯代码
            ├─ vllm/ (Python 推理引擎)
            ├─ 包名发现: vllm* 通配符
            └─ 入口点: vllm = vllm.entrypoints.cli.main:main

5.1 运行时调用链路

Python 推理代码
  │
  ├─ current_platform.import_kernels()
  │   └─ import vllm._C          ─→ TORCH_LIBRARY_EXPAND 初始化
  │      import vllm._moe_C       ─→ MoE 算子注册
  │      (import vllm._rocm_C)    ─→ ROCm 算子注册 (conditionally)
  │
  ├─ vllm/_custom_ops.py
  │   ├─ paged_attention_v1()  ─→ torch.ops._C.paged_attention_v1()
  │   ├─ awq_gemm()            ─→ torch.ops._C.awq_gemm()
  │   ├─ topk_softmax()        ─→ torch.ops._moe_C.topk_softmax()
  │   └─ cutlass_scaled_mm()   ─→ torch.ops._C_stable_libtorch.cutlass_scaled_mm()
  │       (通过 _C_stable_libtorch 或 _C，取决于注册路径)
  │
  └─ vllm-rs (可选 Rust 前端)
       ├─ CLI 入口替代
       ├─ 推理解析器 FFI (ctypes)
       └─ gRPC EngineCore 客户端

---

6. 关键第三方编译依赖

6.1 通过 CMake FetchContent 源码编译的

依赖	来源	用途
CUTLASS v4.4.2	GitHub	NVIDIA GEMM 模板库；支撑 scaled MM、MoE、Marlin、Machete
Flash Attention	vllm-project/flash-attention	vLLM 定制的 FlashAttention v2/v3
DeepGEMM	deepseek-ai/DeepGEMM	DeepSeek FP8 高吞吐 GEMM
FlashMLA	GitHub	Flash MLA attention 内核（CUDA 12.9+）
quack-kernels / nvidia-cutlass-dsl	GitHub	FA4 的 CuteDSL 支持

6.2 通过 pip 安装的

依赖	用途
flashinfer	FlashInfer attention 内核
triton	Triton 语言编译器（GPU kernel）
numba	JIT 编译器（部分 kernel）
fastsafetensors	加速权重加载
torch >= 2.11.0	PyTorch 主框架

---

7. CI 与测试

7.1 测试组织

tests/ 按功能分区：

目录	焦点
kernels/	GPU kernel 单元测试
models/	端到端模型正确性测试（按模型分子目录）
quantization/	量化 kernel 测试
distributed/	多 GPU 分布式测试
v1/	V1 引擎测试
compile/	torch.compile 集成
spec_decode/	推测解码
multimodal/	多模态（图/音/视频）
lora/	LoRA 适配器
engine/	引擎核心逻辑
entrypoints/	API 端点
samplers/	采样策略
weight_loading/	权重加载
basic_correctness/	基本正确性
config/	配置解析

7.2 CI 系统

系统	配置位置	内容
Buildkite	.buildkite/	主 CI：test-pipeline.yaml + 30 个 test area YAML 文件（attention.yaml、models_basic.yaml、distributed.yaml、rust_frontend.yaml 等）
GitHub Actions	.github/workflows/	pre-commit linting、macOS smoke test、stale issue/PR 管理

pytest 标记定义在 pyproject.toml:103-114：slow_test、core_model、distributed、cpu_test 等，用于按类型筛选测试。