NanoFlow: Towards Optimal Large Language Model Serving Throughput
文章标题: NanoFlow: 迈向最优的大型语言模型服务吞吐量 作者/机构: Kan Zhu (University of Washington), Yufei Gao (University of Washington, Tsinghua University), Yilong Zhao (University of Washington, UC Berkeley), Gefei Zuo (University of Michigan), Liangyu Zhao (University of Washington), Dedong Xie (University of Washington), Yile Gu (University of Washington), Tian Tang (University of Washington, Tsinghua University), Qinyu Xu (University of Washington, Tsinghua University), Zihao Ye (University of Washington), Keisuke Kamahori (University of Washington), Chien-Yu Lin (University of Washington), Ziren Wang (University of Washington, Tsinghua University), Stephanie Wang (University of Washington), Arvind Krishnamurthy (University of Washington), Baris Kasikci (University of Washington)
最近的LLM(如GPT-4, LLaMA, Mistral【索引编号16, Mistral 7b, 2023】【索引编号29, Gpt-4 technical report, 2024】【索引编号49, Llama 2: Open foundation and fine-tuned chat models, 2023】)基于仅解码器(decoder-only)的Transformer架构【索引编号51, Attention is all you need, 2017】。其推理过程包括两个阶段【索引编号53, Orca: A distributed serving system for {Transformer-Based} generative models, 2022】:(1) prefill阶段,一次性处理输入提示;(2) decode阶段,自回归地逐个生成输出token。Prefill阶段会初始化KV缓存(KV-cache)【索引编号34, Efficiently scaling transformer inference, 2023】,用于加速decode阶段。
两个阶段都使用图1所示的相同推理流程。输入在每个解码器层中,首先进入注意力阶段,通过与权重矩阵$W_Q$, $W_K$, $W_V$相乘形成Query、Key和Value,其中Key和Value会拼接到已有的KV缓存中。然后Query与所有已有的Key相乘以评估当前token与之前所有token的相似度,该相似度用于计算Value的加权平均值,以聚合上下文信息。结果通过一个使用$W_O$的线性变换(O-projection)后,再进行层归一化操作【索引编号6, Layer normalization, 2016】。接着,在全连接网络(feed forward network)阶段,激活值分别与$W_{up}$和$W_{gate}$相乘生成$o_u$和$o_g$。之后,$o_g$通过一个激活函数(如SiLU【索引编号10, Sigmoidweighted linear units for neural network function approximation in reinforcement learning, 2017】),并与$o_u$进行逐元素相乘。最后,应用$W_{down}$作为Down-projection生成该层的输出,并作为下一层的输入。
2.2 操作特性
根据推理过程中操作的特性,我们可以将其分为四类:
密集操作(Dense operations):包括KQV生成、O projection、Up、Gate和Down projection,这些操作计算激活值与模型权重的乘法。在prefill阶段,所有新请求的token形成一个大批次;在decode阶段,批处理解码(batched decoding)会聚合所有解码请求的激活值【索引编号53, Orca: A distributed serving system for {Transformer-Based} generative models, 2022】。此外,由于prefill和decode阶段共享权重矩阵,两者的激活值可以组合起来进一步利用批处理效应,分摊权重加载成本【索引编号3, Sarathi: Efficient llm inference by piggybacking decodes with chunked prefills, 2023】。因此,密集操作是计算密集型的。
注意力操作(Attention operations):Transformer通过注意力机制捕获token之间的关系【索引编号51, Attention is all you need, 2017】。Prefill阶段同时处理所有输入token,是计算密集型的;而decode阶段每次迭代生成一个token,需要从KV缓存中加载缓存的K和V向量,因此是内存密集型的。
对于大型模型,需要多个GPU来提供足够的内存和计算资源以实现高效的LLM服务【索引编号18, Alpaserve: Statistical multiplexing with model parallelism for deep learning serving, 2023】。现有工作【索引编号12, Fastermoe: Modeling and optimizing training of large-scale dynamic pre-trained models, 2022】【索引编号38, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020】【索引编号57, Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning, 2022】对不同的设备间并行范式进行了全面评估:
张量并行(Tensor Parallelism)【索引编号38, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020】【索引编号57, Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning, 2022】将每个权重矩阵切分到不同的GPU上,并集体执行每个操作,避免了权重复制,从而扩展了计算能力。然而,如图1所示,操作后需要频繁的集体通信来同步每个操作的结果。
流水线并行(Pipeline Parallelism)【索引编号14, Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2019】按层将模型切分为多个阶段,每个GPU只需持有部分权重。与所有设备在同一批数据上执行的张量并行不同,流水线并行的设备在按阶段划分的不同微批次数据上操作。
内存:从内存角度看,延迟为:
这是因为在使用最大批处理大小时,每次迭代需要将整个设备内存内容加载到GPU缓存和寄存器中一次(这在现代服务系统中很典型【索引编号17, Efficient memory management for large language model serving with pagedattention, 2023】)。尽管模型权重在迭代间共享,但由于重用距离长,缓存模型权重以避免从设备内存加载是不可行的。
网络:对于张量并行,在多个GPU上对不同数据分片执行操作后,需要集体网络通信原语来同步结果。集体通信传输的数据大小等于密集操作的输入大小,其维度为 $[B_{dense}, D_{model}]$。为支持张量并行,需要两个AG和一个AR(或两个AR)【索引编号39, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020】。一个AR需要收集其他GPU的输出,执行本地求和并广播结果,因此需要传输两次激活值,而一个AG传输一次。因此,两种情况下,我们可以估算一个GPU的总数据移动量(字节)为 $4 \cdot B_{Dense}D_{model}S_{type} \cdot L$。因此,从网络角度看,延迟为:
3.3 LLM服务工作负载分类
我们通过比较基于内存、计算和网络需求推导出的迭代延迟来对LLM服务工作负载进行分类。
网络时间与计算时间的比较:首先,我们将公式3(网络时间)与公式2(计算时间)进行比较,其比率为:
注意到 $P_{model} \approx 12D_{model}^2 L$,因此 $\frac{2D_{model}L}{P_{model}} \approx \frac{1}{6D_{model}}$。对于需要多GPU的模型,$D_{model}$ 通常大于4096【索引编号9, Alibaba cloud’s qwen2 with enhanced capabilities tops llm leaderboard, 2024】【索引编号16, Mistral 7b, 2023】【索引编号22, Build the future of ai with meta llama 3, 2024】【索引编号48, Llama: Open and efficient foundation language models, 2023】,使得该项小于 $1 \times 10^{-5}$。由于数据中心GPU之间的高带宽互连(如NVLink【索引编号25, Nvlink and nvlink switch, 2024】、Infinity Fabric【索引编号54, Forestcoll: Efficient collective communications on heterogeneous network fabrics, 2024】),$\frac{N_{GPU}Compute}{NetBW/S}$ 的范围在 $1 \times 10^4$ 到 $1 \times 10^5$ 之间。因此,这个比率通常小于1,表明网络不是瓶颈。我们在图2中说明了这一比率,对于大型模型和数据中心GPU,计算时间比网络时间更长。
图2: 网络时间与计算时间的比较。颜色越接近黄色,工作负载计算密集程度越高;越接近蓝色,表示工作负载网络密集程度越高。
内存时间与计算时间的比较:接下来,我们比较公式1(内存时间)和公式2(计算时间)来最终确定工作负载的特性,即以下比值:
现代模型广泛采用分组查询注意力(GQA),如LLaMA-3【索引编号48, Llama: Open and efficient foundation language models, 2023】、NeMo【索引编号4, Mistral nemo, 2024】和Qwen2【索引编号9, Alibaba cloud’s qwen2 with enhanced capabilities tops llm leaderboard, 2024】。在GQA中,多个注意力头共享一个KV缓存,这意味着在相同的内存容量下,批次可以包含更多请求,因此 $B_{dense}$ 往往更大。例如,在8×A100 GPU上服务LLaMA-2 70B时,解码请求的最大批处理大小可达1024。结合批次中的prefill token,$B_{dense}$ 可以达到2048,而同样大小的非GQA模型 $B_{dense}$ 只能为256。因此,比值 $T_R$ 小于1,计算成为主要瓶颈。此外,现代模型 Pmodel 越来越大,进一步导致 $T_R$ 降至1以下。图3显示了五个代表性模型的 $T_R$ 值,表明许多工作负载是计算密集型的,除了在LLaMA-3 8B模型上的长解码工作负载,其 $T_R$ 约为1。
图3: 计算时间与内存时间的比较。颜色越接近黄色,工作负载计算密集程度越高;越接近绿色,则内存密集程度越高。
现有服务系统远未达到最优吞吞吐量的主要原因是GPU资源利用率低下。 图4展示了使用现有服务框架(如SGLang【索引编号58, Sglang: Efficient execution of structured language model programs, 2024】、vLLM【索引编号17, Efficient memory management for large language model serving with pagedattention, 2023】、DeepSpeed-FastGen【索引编号13, Deepspeed-fastgen: High-throughput text generation for llms via mii and deepspeed-inference, 2024】)时,GPU内Transformer操作的执行流程。这些框架在单个设备上按顺序对一个大批次执行计算密集型操作、解码注意力和网络操作。这种顺序执行流程形成了多个流水线气泡(在图4中表示为"WASTED"),导致最受限的资源——计算资源——未被充分利用。
图4: 现有系统的执行流水线。绿色、黄色和蓝色操作分别对应于内存密集型、计算密集型和网络密集型操作。前一层和后一层的操作用虚线边框表示。“WASTED”显示了流水线中受限最严重的计算资源未被充分利用的阶段。为简化起见,省略了小型操作(如层归一化、激活等)。
操作变换的约束:一些网络集体通信有具有不同性能特征和依赖关系的等效变换。例如,一个AG可以转换为一个AR,采用不同的权重分区方法【索引编号39, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020】【索引编号57, Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning, 2022】。NanoFlow探索所有这些网络纳米操作的替代变换,并搜索最短的流水线设计。
批次形成:NanoFlow假设自动扩缩、工作负载均衡和优先级感知路由由外部控制平面管理【索引编号15, Intelligent router for llm workloads: Improving performance through workload-aware load balancing, 2025】【索引编号33, Hierarchical autoscaling for large language model serving with chiron, 2025】【索引编号45, Aibrix: Towards scalable, cost-effective large language model inference infrastructure, 2025】。一个NanoFlow实例假设有充足的请求,并平等对待每个请求。当请求不充足时,控制平面应减少NanoFlow实例数量,以维持足够大的单实例批处理大小。值得注意的是,虽然$B_{dense}$(即GEMM的token批处理大小)在数千级别,但这代表了prefill和decode token的组合。因此,请求批处理大小(例如,1个prefill和256个decode请求)在数百级别,这在现实世界的大规模服务中很容易达到。在形成批次时,NanoFlow优先处理未完成的decode请求,并遵循SarathiServe【索引编号3, Sarathi: Efficient llm inference by piggybacking decodes with chunked prefills, 2023】的方法,以token粒度对prefill请求进行分块,以精确填满预选的最佳性能密集批次的剩余容量。这有效地减少了尾部延迟,因为密集操作在迭代中以一致的批处理大小执行。
异步调度:批次形成过程,包括估算内存使用、调度新请求、移除完成的请求以及调整PagedAttention【索引编号17, Efficient memory management for large language model serving with pagedattention, 2023】的页表,会消耗不可忽略的CPU时间【索引编号42, MLSys @ WukLab - Can Scheduling Overhead Dominate LLM Inference Performance? A Study of CPU Scheduling Overhead on Two Popular LLM Inference Systems, 2024】。为了避免在此期间GPU被闲置,NanoFlow将批次形成与GPU执行并行进行异步调度。在任何迭代$i$,NanoFlow在当前迭代结束前形成下一次迭代的批次。这意味着NanoFlow无法检测到在迭代$i$生成的EOS token。在启动迭代$i+1$后,NanoFlow形成迭代$i+2$的批次,检测来自迭代$i$的EOS token,并移除完成的请求。幸运的是,由于典型工作负载的平均解码长度超过100(见表4),考虑到隐藏批次形成开销的好处,一个额外解码token的开销可以忽略不计(<1%)。
KV缓存加载与分散:NanoFlow使用PagedAttention【索引编号17, Efficient memory management for large language model serving with pagedattention, 2023】,因此KV缓存的页面在GPU内存中是碎片化存放的。为了避免拷贝到碎片化的GPU页面目的地,NanoFlow首先将数据拷贝到GPU上的一个连续空间,然后将页面分散到它们的目的地。这使得主机到设备的拷贝带宽提高了7-10倍。
[42] MLSys @ WukLab - Can Scheduling Overhead Dominate LLM Inference Performance? A Study of CPU Scheduling Overhead on Two Popular LLM Inference Systems (2024)
引用段落: §4.2.1 请求调度
原文描述: 引用来说明批次形成过程会消耗不可忽略的CPU时间。
[43] Orion: Interference-aware, fine-grained gpu sharing for ml applications (2024)
引用段落: §4.1.1 内核性能分析与干扰建模
原文描述: 引用来说明内核并行执行时会因竞争GPU硬件资源而产生性能干扰。
[46] CUTLASS (2023)
引用段落: §3.5 最优服务吞吐量
原文描述: 引用作为测量GPU峰值计算能力的工具库。
[49] Llama 2: Open foundation and fine-tuned chat models (2023)
引用段落: §2.1 LLM 推理工作流
原文描述: 作为基于decoder-only Transformer架构的LLM例子。
[51] Attention is all you need (2017)
引用段落: §2.1 LLM 推理工作流, §2.2 操作特性
原文描述: 作为Transformer架构的基础论文被引用;以及注意力机制的来源。
[53] Orca: A distributed serving system for {Transformer-Based} generative models (2022)