Featured image of post vLLM代码学习之一:概览

vLLM代码学习之一:概览

vllm作为一个实验目标,尝试实践验证系统地学习一个大项目的代码的方法。

方法本身很简单,类似宽度优先搜索,将项目的文件按照深度逐层分析,在不过度深究每个子文件的细节的情况下理解当前文件夹的内容和目的,同一层间做功能的划分,并逐层深入。最后,通过真实的执行顺序将这些功能串起来,这样能做到完整地掌握整个项目的内容,在需要修改的时候能有的放矢。

通过记录学习vllm的过程,希望能帮助我掌握快速上手这样大型项目的方法论。我按照文件系统结构分段,阅读或回头查阅可以直接看对应的章节。

NOTE: 本文基于f3f8d8fff4c5354d5214f0f6f29e4dc5c4e3a8ca

顶层结构

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
vllm
├── benchmarks
├── cmake
├── csrc
├── docs
├── examples
├── requirements
├── tests
├── tools
└── vllm

🔥 这几个文件夹的内容很重要:

  • csrc/ 包含项目的C++和CUDA源代码,是项目的kernel核心实现
  • vllm/ 包含项目的主要代码,是项目的主要实现

💡 这些模块重要性更低一些,但是在需要的时候可以查阅:

  • examples/ 包含项目的示例代码,提供了一些使用vllm的例子,可以辅助我们的理解
  • benchmark/ 提供各种基准测试和性能评估工具,包括纯文本和多模态场景,用于测试vllm的生成速度
  • tests/ 包含项目的测试代码,用于确保项目的正确性

❄️ 这些文件夹的内容对于学习vllm没有帮助,也不需要做修改:

  • cmake/ 包含项目的CMake构建文件,用于配置和构建项目
  • docs/ 包含项目的文档,自动化构建代码用的
  • requirements/ 包含项目的依赖项,列出项目运行所需的外部库和工具
  • tools/ 包含项目的工具,用于辅助项目的开发和测试
其他文件都不用细看,我们直接进`vllm`的核心代码
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
vllm
├── DCO
├── Dockerfile
├── Dockerfile.arm
├── Dockerfile.cpu
├── Dockerfile.hpu
├── Dockerfile.neuron
├── Dockerfile.ppc64le
├── Dockerfile.rocm
├── Dockerfile.rocm_base
├── Dockerfile.s390x
├── Dockerfile.tpu
├── Dockerfile.xpu
├── find_cuda_init.py
├── format.sh
├── LICENSE
├── MANIFEST.in
├── pyproject.toml
├── python_only_dev.py
├── README.md
├── RELEASE.md
├── requirements
├── SECURITY.md
├── setup.py
└── use_existing_torch.py

vllm

核心代码很多,这里用emoji符号区分一下目录和文件。
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
vllm/vllm/
├── 📁 adapter_commons/
├── 📁 assets/
├── 📁 attention/
├── 📁 benchmarks/
├── 📁 compilation/
├── 📁 core/
├── 📁 device_allocator/
├── 📁 distributed/
├── 📁 engine/
├── 📁 entrypoints/
├── 📁 executor/
├── 📁 inputs/
├── 📁 logging_utils/
├── 📁 lora/
├── 📁 model_executor/
├── 📁 multimodal/
├── 📁 platforms/
├── 📁 plugins/
├── 📁 profiler/
├── 📁 prompt_adapter/
├── 📁 reasoning/
├── 📁 spec_decode/
├── 📁 third_party/
├── 📁 transformers_utils/
├── 📁 triton_utils/
├── 📁 usage/
├── 📁 v1/
├── 📁 vllm_flash_attn/
├── 📁 worker/
├── 📄 beam_search.py
├── 📄 config.py
├── 📄 connections.py
├── 📄 _custom_ops.py
├── 📄 envs.py
├── 📄 forward_context.py
├── 📄 __init__.py
├── 📄 _ipex_ops.py
├── 📄 jsontree.py
├── 📄 logger.py
├── 📄 logits_process.py
├── 📄 outputs.py
├── 📄 pooling_params.py
├── 📄 py.typed
├── 📄 sampling_params.py
├── 📄 scalar_type.py
├── 📄 scripts.py
├── 📄 sequence.py
├── 📄 test_utils.py
├── 📄 tracing.py
├── 📄 utils.py
└── 📄 version.py

同时记录一下打印这样的目录结构的命令:

1
2
3
4
# 强制让目录以/结尾,以此区分目录和文件
tree vllm -F -L 1 --dirsfirst | sed -E \
  -e 's/([├└]── )([^ ]*\/)$/\1📁 \2/' \
  -e 's/([├└]── )([^ ]*)$/\1📄 \2/'

v1/目录比较特殊,是一个重构中的vllm版本,这里先暂时和主分支分开讨论。对于几乎没有内容的模块我也选择略去。

🔥 以下的模块十分关键,实现了关键的基本能力:

  • attention/提供不同的注意力机制的实现,包含了算子、注意力算法、注意力层三个层面的实现。
  • core/包含vllm的块管理实现。
  • device_allocator/实现了对CUDA tensor的内存管理。
  • distributed/实现了分布式状态(TP、PP等)的管理、设备通信和KV传输。
  • sequence.py实现了序列类,保存了一个序列的完整信息,如分配的Block等。

🔥 以下的模块十分关键,实现了主要的结构体,但是不用专门看:

  • inputs/实现了输入类和一些预处理逻辑。
  • outputs.py实现了各种输出类。
  • config.py实现了各种配置类。
  • envs.py包含了默认的配置。
  • sampling_params.py实现了采样参数类。

🔥 以下模块尤其关键,且相互依赖,共同完成了vllm的主要数据流程,从最上层到最底层分别为:

  • entrypoints/提供了vllm的一系列入口类,包括离线生成的接口LLM、在线生成的接口API Server(基于fastapi)和OpenAI 兼容接口,是vllm暴露的最上层的接口。
  • engine/实现了LLMEngine,这是vllm的核心引擎类,负责整体协调和管理推理过程。它是最上层的组件,通过ExecutorWorker来实际执行模型推理。
  • executor/实现了Executor,是LLMEngineWorker的中间层,负责将上层LLMEngine的调度翻译成下层Worker的具体执行行为。
  • worker/实现了Worker,只负责在单个 GPU 上执行模型的实际计算。
  • model_executor/在最底层,包含不同的模型的具体实现和参数加载方式等,另外还包括了Guided decoding(e.g., JSON mode)。
以上几个模块的关系示意图如下。

请求执行:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
LLMEngine (顶层)
    ├─── 管理整体流程
    │    - 请求调度
    │    - 资源管理
    │    - 输出处理
    v
Executor (执行层)
    ├─── 协调执行
    │    - 分布式通信
    │    - 任务分发
    │    - 结果聚合
    v
Worker (工作层)
    ├─── 实际执行
    │    - GPU 计算
    │    - 内存管理
    │    - 模型加载
    v
ModelExecutor (模型层)
    ├─── 模型实现
    │    - 模型架构
    │    - 前向传播
    │    - 权重管理

接口:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
                 entrypoints
           ┌────────┴────────┐
           │                 │
        LLM 类          API Server
           │                 │
           │                 │
     LLMEngine ◄────────────┘
           ├─── Executor
           │        │
           │    Worker(s)
           └─── 其他组件
                (调度器、缓存等)

💡 以下的模块值得一读,但是不是掌握vllm所必须:

  • compilation/提供不同后端引擎支持下的编译加速,对计算图做优化,默认是TorchInductor
  • lora/实现了LoRA adapter的相关逻辑,包括层实现和应用、取消LoRA等功能。
  • prompt_adapter/实现了Prompt Adapter的相关逻辑。
  • multimodal/实现了多模态(语音、图像和视频)相关的逻辑。
  • platforms/用于支持不同硬件平台,其中Platform类定义了一个通用的平台接口。
  • reasoning/支持了推理模型相关的功能,例如划分推理部分和回答部分。
  • spec_decode/支持了投机采样相关的功能。
  • transformers_utils/有一些对HuggingFace的transformers库的封装。

❄️ 以下的模块重要性更低,不用做深度的阅读:

  • adapter_commons/提供一个统一的接口和工具集来支持不同类型的模型适配器(如 LoRA、P-Tuning 等)。
  • assets/提供对图像、视频等的支持。
  • benchmarks/提供对基准测试的支持,包括起服务器等。
  • logging_utils/提供日志功能。
  • profiler/提供测速相关的工具。
  • third_party/包了第三方的工具,目前只有pynvml
  • connections.py包含网络连接工具。

vllm/v1

v1的目录结构和主分支大体相似但是更简洁,所有模块都需要细读。
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
vllm/v1/
├── 📁 attention/
├── 📁 core/
├── 📁 engine/
├── 📁 executor/
├── 📁 metrics/
├── 📁 sample/
├── 📁 spec_decode/
├── 📁 stats/
├── 📁 structured_output/
├── 📁 worker/
├── 📄 __init__.py
├── 📄 kv_cache_interface.py
├── 📄 outputs.py
├── 📄 request.py
├── 📄 serial_utils.py
└── 📄 utils.py

11 directories, 6 files
  • attention/实现了三种注意力机制后端和MLA。
  • core/中包含了块和KV cache的实现。
  • executor/中包含了v1版本的Executor,在其上层的LLMEngine是和主分支共用的。
  • worker/中包含了v1版本的Worker
  • sample/中包含了各种采样的实现和算子,例如拒绝采样、topk、topp等。
使用 Hugo 构建
主题 StackJimmy 设计