Nanoflow

论文 | 幻灯片

NanoFlow 是一个面向高吞吐量的高性能大语言模型推理框架。与 vLLM、Deepspeed-FastGen 和 TensorRT-LLM 相比，NanoFlow 始终能够提供更出色的吞吐量。与 TensorRT-LLM 相比，NanoFlow 的吞吐量最高可提升 1.91 倍。 NanoFlow 的主要特性包括：

• 设备内并行化：通过利用纳米批处理和执行单元调度，将单个设备内的不同资源需求进行重叠，从而最大化硬件利用率。
• 异步 CPU 调度：通过采用异步控制流来管理 GPU 执行、CPU 批次生成以及 KV 缓存管理，实现高效的 CPU 调度。

新闻

• [2024/09] 🚀 Nanoflow 现已支持 Llama2 70B、Llama3 70B、Llama3.1 70B、Llama3 8B、Llama3.1 8B 以及 Qwen2 72B 模型。我们还发布了实验脚本，用于复现评估结果。

简介

NanoFlow 的核心理念在于：现有框架的传统流水线设计由于操作的顺序执行，导致硬件资源利用率低下。因此，NanoFlow 提出了设备内并行化（如以下动图所示），通过使用纳米批处理来调度计算密集型、内存密集型和网络密集型的操作，使其能够同时执行。这种重叠使得关键路径上的计算密集型操作得以保留，从而大幅提升资源利用率。

NanoFlow 的关键组件概览

设备内并行化的示意图

在 GPU 高度利用率的情况下，CPU 的开销——包括 KV 缓存管理、批次生成以及已结束请求的筛选——占到了推理时间的相当大比例（超过 10%）。因此，NanoFlow 采用了如下图所示的异步控制流。在任意迭代 i 中，NanoFlow 会在当前迭代结束之前就做出下一轮的批次决策，并为下一轮分配 KV 缓存条目。NanoFlow 不会等待迭代 i 中生成的序列结束标记（EOS）出现，而是直接启动迭代 i+1，并在迭代 i+2 时才对已完成的请求进行清理。

异步控制流调度的说明

为了避免重复计算并复用多轮对话中的 KV 缓存，NanoFlow 会主动将已完成请求的 KV 缓存卸载到 SSD 上。在一次迭代中，NanoFlow 会选取已退役请求的 KV 缓存，并以逐层的方式将其并行复制到主机端，同时进行实时推理操作。我们的测算表明，服务 LLaMA2-70B 仅需 5 GB/s 的卸载带宽，而单块 SSD 的带宽即可达到 3 GB/s。

综合以上各项技术，我们现在开源了 NanoFlow，包含基于 C++ 的后端和基于 Python 的演示前端，代码总量约为 4000 行。NanoFlow 集成了最先进的内核库，包括用于 GEMM 的 CUTLASS、用于注意力机制的 FlashInfer 以及用于网络通信的 MSCCL++。此外，该代码库还包含了环境搭建和实验复现所需的必要脚本。

基准测试

我们列出了一些主要的基准测试结果。更多细节请参阅我们的论文。我们在 A100 80GB SXM 上进行评估，并选择 vLLM v0.5.3、Deepspeed-FastGen v0.2.3 和 TensorRT-LLM v0.8.0 作为基准。请注意，所有框架均关闭了量化、推测解码、前缀缓存等特定优化功能。

离线吞吐量：Llama2-70B 在 8×A100（80GB）上

我们分别在两种场景下进行了离线吞吐量测试：一是基于收集到的轨迹的实际工作负载（Splitwise、LMSYS-Chat-1M、ShareGPT），二是输入输出长度恒定的情况。NanoFlow 始终优于所有基准。

离线吞吐量基准测试

在线延迟：Llama2-70B 在 8×A100（80GB）上

我们针对三种真实世界轨迹，测试了归一化延迟（即端到端请求延迟除以输出 token 数），并设置了不同的请求速率（每秒到达的请求数）。在所有数据集中，NanoFlow 均能以更低的延迟维持更高的请求速率，优于其他基准。

在线延迟基准测试

可行性：不同模型上的离线吞吐量

我们将 NanoFlow 移植到了 5 个具有代表性的模型上，以展示其灵活性。我们在输入长度固定为 1024、输出长度固定为 512 的情况下，评估了 NanoFlow 在这些大语言模型上的离线吞吐量（每 GPU 每秒的 token 数）。

NanoFlow 在不同模型上的离线吞吐量

代码库

摘要

大型语言模型（LLMs）的广泛应用导致了对行星级规模服务系统的需求激增，这类系统需要数以万计的GPU持续为数亿用户提供服务。因此，在合理的延迟约束下，吞吐量已成为衡量服务系统性能的关键指标。为了提升吞吐量，人们已探索了多种设备间并行化方法，如数据并行、张量并行和流水线并行等。然而，现有方法并未考虑在同一设备内部不同资源的使用重叠问题，从而导致资源利用率低下、性能无法达到最优。

我们提出了NanoFlow这一新型服务框架，通过操作协同调度的方式，充分挖掘设备内部的并行性，实现计算、内存和网络等资源的重叠利用。为实现设备内并行，NanoFlow引入了两项关键创新：首先，NanoFlow提出纳米批处理技术，将请求按操作粒度进行拆分，打破LLM推理中顺序操作的依赖关系，从而实现操作间的重叠执行；其次，为了充分利用这种重叠效应，NanoFlow采用了一种基于执行单元调度的设备级流水线机制，该机制将设备的功能单元进行划分，并在每个单元中同时执行不同的操作。NanoFlow还借助参数搜索算法自动完成流水线的配置，使得NanoFlow能够轻松适配不同模型。我们在NVIDIA GPU上实现了NanoFlow，并在LLaMA-2-70B、Mixtral 8×7B、LLaMA-3-8B等多款主流模型上评估了端到端的服务吞吐量。实验结果表明，NanoFlow可达到最优吞吐量的68.5%。在实际工作负载下，与当前最先进的服务系统相比，NanoFlow的吞吐量提升了1.91倍，而后者在移植后的模型上仅能达到最优吞吐量的59%至72%。

安装

Docker环境搭建

mkdir -p ~/framework-test
docker run --gpus all --net=host --privileged -v /dev/shm:/dev/shm --name nanoflow -v ~/framework-test:/code -it nvcr.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04

apt update
apt install pybind11-dev
apt install liburing-dev
apt install libopenmpi-dev
sysctl -w kernel.io_uring_disabled=0
sysctl -w vm.nr_hugepages=65536

Gurobi许可证配置（适用于Docker）

请按照以下步骤获取Gurobi许可证并进行配置，以便您的Docker容器能够使用该许可证。

1. 申请Gurobi许可证

1. 访问Gurobi官网并注册账号（https://www.gurobi.com/）。
2. 登录后，进入My Gurobi → Get License页面。
3. 选择“WLS Academic”许可证类型，并填写相关必要信息。
4. Gurobi会向您发送一封包含许可证文件gurobi.lic的电子邮件（或提供许可证密钥字符串）。

2. 将许可证放置于宿主机

mkdir -p ~/gurobi/license
mv /path/to/downloaded/gurobi.lic ~/gurobi/license/
ls ~/gurobi/license
echo "export GRB_LICENSE_FILE=$(pwd)/gurobi.lic" >> ~/.bashrc

安装依赖

git clone git@github.com:efeslab/Nanoflow.git
cd Nanoflow
chmod +x ./installAnaconda.sh
./installAnaconda.sh
# 重启终端
source ~/.bashrc

# 登录Hugging Face
mkdir -p hf
echo "export HF_HOME=$(pwd)/hf" >> ~/.bashrc
source ~/.bashrc

huggingface-cli login

# 配置环境
cd Nanoflow-python
yes | bash setup.sh

构建

cd pybind
mkdir -p build
cmake ..
make -j 256

端到端测试

cd entry
python run_llama3.py -load_hf_weight=True

引用

如果您在研究中使用NanoFlow，请引用我们的论文（arXiv预印本）：

@misc{zhu2024nanoflowoptimallargelanguage,
      title={NanoFlow: Towards Optimal Large Language Model Serving Throughput}, 
      author={Kan Zhu and Yilong Zhao and Liangyu Zhao and Gefei Zuo and Yile Gu and Dedong Xie and Yufei Gao and Qinyu Xu and Tian Tang and Zihao Ye and Keisuke Kamahori and Chien-Yu Lin and Stephanie Wang and Arvind Krishnamurthy and Baris Kasikci},
      year={2024},
      eprint={2408.12757},
      archivePrefix={arXiv},
      primaryClass={cs.DC},
      url={https://arxiv.org/abs/2408.12757}, 
}

致谢

NanoFlow的灵感来源于并复用了以下项目的代码：CUTLASS、FlashInfer、MSCCL++以及Punica。此外，GoogleTest、NVBench和spdlog等工具也为NanoFlow的开发提供了极大便利。我们特别感谢陈思琴在NanoFlow标志设计方面的帮助。

NanoFlow 快速上手指南

环境准备

• 系统：Ubuntu 22.04（推荐）
• GPU：NVIDIA A100 80 GB（官方测试环境）
• CUDA：12.8+
• Docker：已安装 nvidia-docker
• 网络：可访问 Hugging Face 与 Gurobi 官网

安装步骤

1. 启动容器

mkdir -p ~/framework-test
docker run --gpus all --net=host --privileged -v /dev/shm:/dev/shm --name nanoflow -v ~/framework-test:/code -it nvcr.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04

2. 安装系统依赖

apt update
apt install -y pybind11-dev liburing-dev libopenmpi-dev
sysctl -w kernel.io_uring_disabled=0
sysctl -w vm.nr_hugepages=65536

3. 配置 Gurobi 许可证（学术用户）

1. 访问 https://www.gurobi.com 注册账号
2. 在「My Gurobi → Get License」选择 WLS Academic 并下载 gurobi.lic
3. 将许可证挂载到容器：

mkdir -p ~/gurobi/license
# 把下载的 gurobi.lic 放到 ~/gurobi/license/
echo "export GRB_LICENSE_FILE=/gurobi/license/gurobi.lic" >> ~/.bashrc

4. 克隆与安装

git clone https://github.com/efeslab/Nanoflow.git
cd Nanoflow
chmod +x ./installAnaconda.sh && ./installAnaconda.sh
source ~/.bashrc

# 登录 Hugging Face（国内可配镜像）
mkdir -p hf
echo "export HF_HOME=$(pwd)/hf" >> ~/.bashrc
source ~/.bashrc
huggingface-cli login

# 一键环境
cd Nanoflow-python
yes | bash setup.sh

5. 编译

cd pybind
mkdir -p build && cd build
cmake .. && make -j$(nproc)

基本使用

cd entry
python run_llama3.py -load_hf_weight=True

首次运行会自动下载 Llama-3-8B 权重并启动本地推理服务，默认端口 8000。