DistServe

DistServe 通过将预填充和解码计算分离，从而提升大型语言模型（LLMs）推理服务的性能。现有的 LLM 推理系统通常将这两个阶段部署在一起，并对所有用户和请求的预填充与解码计算进行批处理。我们发现，这种策略不仅会导致预填充和解码之间产生严重的干扰，还会将两个阶段的资源分配和并行化方案紧密耦合起来。而在 DistServe 中，用户只需为这两个阶段分别设置并行度配置和调度策略，它就能像单个实例一样自动处理 KV 缓存通信和内存管理。

DistServe 使用高性能的 C++ Transformer 推理库 SwiftTransformer 作为执行后端，该库支持多种特性，如模型并行、流水线并行、FlashAttention、连续批处理以及分页注意力机制。

支持的模型包括：

• GPT-2 (gpt2, gpt2-xl, ...)
• OPT (facebook/opt-1.3b, facebook/opt-6.7b, ...)
• LLaMA2 (meta-llama/Llama-2-7b, meta-llama/Llama-2-13b, ...)

构建与安装

# 克隆项目
git clone https://github.com/LLMServe/DistServe.git && cd DistServe

# 设置 DistServe 的 Conda 环境
conda env create -f environment.yml && conda activate distserve

# 克隆并构建 SwiftTransformer 库
git clone https://github.com/LLMServe/SwiftTransformer.git && cd SwiftTransformer && git submodule update --init --recursive
cmake -B build && cmake --build build -j$(nproc)
cd ..

# 安装 DistServe
pip install -e .

启动

启动 Ray 集群

DistServe 依赖 Ray 来实现分布式工作节点。如果您未提前启动 Ray 运行时，它会自动在当前节点上初始化一个包含所有 GPU 的集群。若需使用多节点进行推理，则可能需要手动提前启动 Ray 运行时。

运行离线示例

DistServe 至少需要两块 GPU 才能运行。我们提供了一个离线推理示例，位于 examples/offline.py。

运行在线示例

要运行在线推理，您需要启动 DistServe API 服务器，请参阅 distserve/api_server/distserve_api_server.py 中的注释。

随后，运行 examples/online.py 中的客户端示例。

评估

要复现我们论文中的所有实验，请按照 指南 操作。

引用

如果您在研究中使用了 DistServe，请引用我们的论文：

@misc{zhong2024distserve,
      title={DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving}, 
      author={Yinmin Zhong and Shengyu Liu and Junda Chen and Jianbo Hu and Yibo Zhu and Xuanzhe Liu and Xin Jin and Hao Zhang},
      year={2024},
      eprint={2401.09670},
      archivePrefix={arXiv},
      primaryClass={cs.DC}
}

DistServe 快速上手指南

DistServe 通过解耦大语言模型（LLM）服务中的 Prefill（预填充） 和 Decoding（解码） 计算阶段，显著提升了服务性能。它允许用户独立配置两个阶段的并行策略与调度方案，并自动处理 KV-Cache 通信与内存管理。底层基于高性能 C++ 推理引擎 SwiftTransformer，支持 FlashAttention、PagedAttention 等特性。

环境准备

• 操作系统：Linux (推荐 Ubuntu)
• 硬件要求：至少 2 块 GPU（用于演示解耦效果）
• 软件依赖：
  • Git
  • Conda (推荐 Miniconda 或 Anaconda)
  • CMake
  • CUDA Toolkit (需与显卡驱动及 PyTorch 版本兼容)
  • Ray (分布式框架，安装脚本中会自动处理部分依赖)

提示：国内开发者建议在创建 Conda 环境和 pip 安装时使用清华源或阿里源以加速下载。

安装步骤

请依次执行以下命令完成项目克隆、环境构建及编译安装：

# 1. 克隆项目代码
git clone https://github.com/LLMServe/DistServe.git && cd DistServe

# 2. 创建并激活 Conda 环境
# 国内加速建议：在 environment.yml 中修改渠道源，或手动创建环境后安装依赖
conda env create -f environment.yml && conda activate distserve

# 3. 克隆并编译后端推理库 SwiftTransformer
git clone https://github.com/LLMServe/SwiftTransformer.git && cd SwiftTransformer
git submodule update --init --recursive

# 编译构建 (使用所有可用核心加速编译)
cmake -B build && cmake --build build -j$(nproc)
cd ..

# 4. 安装 DistServe 包
# 国内加速建议：pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install -e .

基本使用

DistServe 依赖 Ray 进行分布式任务调度。

1. 启动说明

• 单机多卡：若未预先启动 Ray 集群，DistServe 会自动检测当前节点的所有 GPU 并初始化集群。
• 多机部署：若需跨节点推理，请先手动启动 Ray 集群 (ray start --head 等)。

2. 离线推理示例 (Offline Inference)

这是最简单的测试方式，用于验证安装及基本功能。确保机器至少有 2 块可用 GPU。

# 运行离线示例脚本
python examples/offline.py

3. 在线服务示例 (Online Serving)

如需部署为 API 服务：

1. 启动 API 服务器： 查看并配置 distserve/api_server/distserve_api_server.py 中的参数，然后运行该脚本启动服务。
2. 启动客户端测试：

   python examples/online.py
   

支持模型

目前原生支持以下模型系列（可通过 HuggingFace ID 加载）：

• GPT-2: gpt2, gpt2-xl 等
• OPT: facebook/opt-1.3b, facebook/opt-6.7b 等
• LLaMA2: meta-llama/Llama-2-7b, meta-llama/Llama-2-13b 等