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)

A1 主要贡献

本文首先通过详细分析指出,尽管大型语言模型(LLM)服务包含内存密集型组件(如自注意力机制),但对于大多数常见工作负载和模型而言,端到端的LLM服务实际上是计算密集型(compute-bound)的。然而,现有服务引擎未能充分利用计算资源,因为计算、内存、网络等异构操作在单个设备内是顺序执行的,导致总计算利用率仅约40%。

针对这一问题,本文的核心贡献是提出了NanoFlow,一个旨在通过利用设备内并行性(intra-device parallelism)来最大化工作负载资源瓶颈利用率的服务框架。

主要创新点如下:
1. 提出设备内并行性框架NanoFlow:该框架将输入批次(batch)拆分为更小的纳米批次(nano-batch),并复制操作以独立处理每个部分,从而允许计算密集型、内存密集型和网络密集型等异构操作在单个设备内并行执行,实现计算、内存和网络资源的细粒度流水线化。
2. 设计自动搜索引擎(auto-search):为了应对确定纳米批次的数量、大小、顺序以及GPU资源分配等构成的巨大搜索空间,NanoFlow提出了一个自动搜索引擎。该引擎分两个阶段工作:
* 第一阶段:在不考虑内核间干扰的假设下,确定初始的流水线调度(纳米操作的数量、大小和顺序)。
* 第二阶段:通过分析实际的内核干扰来优化流水线,并重新规划资源分配,从而在效率和准确性之间取得平衡。
3. 实现端到端LLM服务运行时:NanoFlow提供了一个完整的运行时系统,用于执行优化后的流水线,能高效地形成纳米批次、分配GPU资源并管理KV缓存。
4. 通过全面的实验验证了有效性:在LLaMA-2-70B等多个流行模型上的评估表明,NanoFlow相比现有最先进的服务系统(如vLLM、DeepSpeed-FastGen、TensorRT-LLM)平均可实现1.91倍的吞吐量提升,达到了理论最优吞吐量的50%至72%。

A3 背景知识/关键Observation/设计原则

2 背景知识

2.1 LLM 推理工作流

最近的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 操作特性

根据推理过程中操作的特性,我们可以将其分为四类:


图1: Transformer架构。黄色框中的操作具有大批量大小,并在请求之间共享模型权重参数;因此,它们是计算密集型的。绿色框中的操作需要为每个请求加载一个唯一的KV缓存;因此,它们是内存密集型的。蓝色框代表执行操作之间同步的网络操作。

2.3 规模化服务LLM

对于大型模型,需要多个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】对不同的设备间并行范式进行了全面评估:

在实践中,这些并行范式是结合使用的。张量并行广泛用于节点内的GPU,以支持更大的批处理大小,而流水线并行则用于跨节点进一步扩展集群。

3 分析

3.1 影响服务吞吐量的关键因素

我们将吞吐量定义为服务的总吞吐量,即prefill和decode阶段每秒处理的token数量。吞吐量受硬件属性、模型配置、用户查询和批处理大小的影响。

3.2 LLM服务成本模型

在最大批处理大小的假设下,我们从所需内存、计算和网络资源的角度推导LLM服务迭代的延迟。

3.3 LLM服务工作负载分类

我们通过比较基于内存、计算和网络需求推导出的迭代延迟来对LLM服务工作负载进行分类。

3.4 成本模型验证

我们在8个NVIDIA A100 GPU上服务LLaMA-2 70B、密集批处理大小为2048个请求的场景下验证了成本模型。 我们计算了每个操作的GFLOPs、内存移动量和网络流量(见表2),并用这些数据结合硬件规格计算出$T_{compute}$、$T_{mem}$和$T_{net}$。最长的估算时间$T_{op} = \max(T_{compute}, T_{mem}, T_{net})$指出了最受限的资源。为了验证这些结果,我们测量了不同操作的实际运行时间。对于大多数操作,估算时间$T_{op}$与性能分析结果一致。一个例外是prefill attention,其实际测量时间更长,这是因为相关内核的启动开销主导了总时间。所有操作的$T_{compute}$、$T_{mem}$和$T_{net}$值之和表明,计算是资源最受限的部分,这与我们模型的发现一致。

表2: 成本模型估算与真实世界测量的操作运行时比较。

3.5 最优服务吞吐量

最优吞吐量(以token/s/GPU为单位)是在最有限的资源(即计算资源)被充分利用时达到的。 应用公式2,系统的最优吞吐量为:


在计算密集型场景下,最优吞吞吐量仅取决于GPU针对特定数据类型的总计算能力和模型中的参数数量。值得注意的是,其他因素,包括GPU内存大小、带宽、模型数据类型或prefill和decode阶段的长度,对实现最优吞吐量没有显著影响。我们通过分析最先进的GEMM供应商库CUTLASS【索引编号46, CUTLASS, 2023】来测量GPU的计算能力。例如,在8×A100 GPU上服务LLaMA-2 70B时,分析出的FP16峰值计算能力为280 TFLOPS,$P_{Model}$为70B。将这些值代入公式5,得到的最优吞吐量为1857 tokens/s/GPU。

3.6 与最优吞吐量的差距

如图7所示,现有的服务系统远未达到最优吞吐量。 在服务离线请求时,vLLM、DeepSpeed-FastGen和TensorRT-LLM分别只达到了最优吞吐量的22.0%、22.9%和37.8%。

现有服务系统远未达到最优吞吞吐量的主要原因是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”显示了流水线中受限最严重的计算资源未被充分利用的阶段。为简化起见,省略了小型操作(如层归一化、激活等)。

3.7 设备内并行

表面上看,现有服务框架在GPU中使用单个大批次的优势在于,相比于多个小批次及其相应操作,它们需要加载模型权重的次数更少。 然而,我们本节前面概述的成本模型及其验证提供了一个有价值的见解:由于现代模型在很大程度上是计算密集型的(如图3和表2所示,系数超过2),创建和处理更小的批次可能是合理的,特别是如果能提高计算利用率。

这些观察启发了通过纳米批处理(nano-batching)实现设备内并行,这也是NanoFlow设计的基础。 NanoFlow创建了多个纳米操作,这些操作与原始操作相同,但在我们称之为纳米批次的更细粒度批次上运行。例如,对于原始的在批处理大小为2048上操作的Up projection,NanoFlow可以创建两个纳米操作UP1和UP2,分别在0-768和768-204t8范围内的批次上操作,同时执行完全相同的投影。不同的纳米批次之间没有数据依赖关系。因此,使用不同资源——计算、内存、网络——的纳米操作可以在一个设备内并行地对不同的纳米批次进行操作,以最大化最受限资源——计算——的利用率。尽管需要重复加载模型权重,我们证明了NanoFlow提供的吞吐量显著高于现有框架(平均高出1.91倍),进一步缩小了与最优服务吞吐量的差距(§6)。

A2 方法细节

4 设计

4.1 自动化流水线搜索

由于模型和硬件的多样性,为每个纳米操作决定资源分配、启动内核和调度设备内并行极具挑战。 为此,NanoFlow提出了自动搜索(auto-search),它使用混合整数线性规划来获取纳米批次的重叠策略,并通过两阶段近似来减少搜索空间。我们首先概述自动搜索的先决条件,包括内核性能分析和干扰建模(§4.1.1),然后详细介绍其两个阶段:流水线结构搜索(§4.1.2)和GPU资源分配(§4.1.3)。最后,我们为流行模型提供示例流水线(§4.1.4)。

4.1.1 内核性能分析与干扰建模

理解不同操作的性能特性是自动化流水线搜索的关键。 为此,NanoFlow对GEMM内核(用于计算密集型操作,如KQV生成)、GEMV内核(用于内存密集型操作,如解码注意力)和网络内核(用于网络密集型操作,如AR和AG)进行性能分析。

图5展示了NanoFlow对A100 GPU上GEMM-GEMV内核对(约100对)的分析结果。这些配对按GEMM性能降序排列。例如,归一化性能为0.8意味着内核的执行时间是其最快实现的1/0.8倍。对GEMM干扰更大但GEMV性能更差的配对(图5中灰色部分)被丢弃。此过程在不同的GEMM-GEMV性能权衡点上识别出性能最佳的内核组合。

利用干扰分析,自动搜索生成一个资源映射表(表3)。该表量化了GEMV和网络内核性能(P)作为资源利用率(R)的函数。例如,如图5中红色虚线所示,要达到0.3的GEMV性能,需要牺牲0.2的GEMM性能(从1降至0.8),这意味着$R_{GEMV}=0.2$对应于$P_{GEMV}=0.3$。对所有评估的LLM模型中的GEMM形状和64种批处理大小组合进行的敏感性分析显示,$R$到$P$的映射是一致的,标准差在均值的5%以内。因此,表3可用于后续自动搜索分析中所有GEMV和网络内核实现的$R$到$P$的映射。


图5: GEMM和GEMV内核之间的干扰特性。x轴上的点对应唯一的GEMM-GEMV实现对。y轴表示GEMM和GEMV内核的归一化性能P。

表3: 不同资源利用率R下GEMV和网络内核的性能P。

4.1.2 自动搜索第一阶段:流水线结构搜索

在第一阶段,自动搜索以密集批处理大小、操作依赖关系和无干扰内核的性能分析为输入,然后使用混合整数线性规划(MILP)输出每个纳米操作的数量、批处理大小和顺序。 这一阶段不考虑纳米操作之间的干扰以简化问题,这部分由第二阶段处理。MILP问题受以下约束:

由于搜索空间巨大,搜索最优流水线可能需要数小时甚至数天。然而,通过优先搜索可行解而非证明最优解,可以在大约10分钟内找到一个实用的流水线。

4.1.3 自动搜索第二阶段:优化流水线

第一阶段得出的流水线没有考虑内核干扰,这在实际部署中是不切实际的。 在第二阶段,自动搜索通过考虑内核因干扰而减速的情况来优化此流水线。此阶段构建一个独立的MILP问题,其中纳米操作的数量、批处理大小和排序约束保持不变(即第一阶段的输出)。自动搜索探索对单个纳米操作的各种资源利用率(R)分配,并使用干扰分析中得出的表3将R映射到P。一个特定的资源利用率值R直接对应于§4.1.1中确定的内核实现。第二阶段的目标同样是最小化流水线执行时间。

自动搜索仅在模型架构或工作负载(输入长度、输出长度)发生重大变化时执行。与长时间运行的部署时间相比,搜索时间可以忽略不计。

4.1.4 示例流水线

4.2 NanoFlow 运行时

4.2.1 请求调度
4.2.2 KV缓存管理

5 实现

我们在NVIDIA GPU上实现了NanoFlow,包含约10K行CUDA代码和约6K行Python代码。 根据自动搜索确定的GPU资源需求,NanoFlow基于其性能分析结果知道为每个纳米批次启动哪些内核以达到给定的利用率水平。NanoFlow尊重自动搜索的排序约束来启动纳米操作,并使用CUDA事件来强制执行排序依赖关系。

A4 实验环境

表4: 采样数据集中输入和输出长度的平均值和标准差。

A4 实验结果

6.2 吞吐量


图7: 离线吞吐量比较。NanoFlow在所有工作负载设置中均优于所有基线。TP表示用于张量并行的GPU数量。

6.3 延迟


图8: 延迟比较。x轴显示每秒传入的请求数,y轴显示归一化延迟。NanoFlow在200ms SLO约束内能处理更高的请求率。

6.4 消融研究


图9: NanoFlow的消融研究结果。纳米批处理和重叠提升了NanoFlow的性能。

6.5 资源使用


图10: 非重叠流水线与NanoFlow在单层推理期间的资源使用情况。NanoFlow通过同时利用内存和网络带宽,在整个流水线中实现了高计算使用率。

6.6 在其他LLM上的性能


图11: NanoFlow实例在其他模型上的性能(以每秒每GPU的token数计)。与vLLM相比,NanoFlow显著提高了吞吐量。NanoFlow最高可达最优吞吐量的78.5%。

A5 结论

本文通过分析揭示了现代LLM服务工作负载本质上是计算密集型的,并指出当前系统由于顺序执行计算、内存和网络密集型操作,导致计算资源利用率低下,吞吐量未达最优。为解决此问题,本文提出了NanoFlow,一个新颖的端到端LLM服务系统,它通过设备内并行来重叠具有不同资源需求的操作。NanoFlow的核心是一个能够适应各种LLM模型的自动搜索(auto-search)引擎,该引擎自动构建一个重叠纳米操作的流水线。实验证明,与最先进的系统相比,NanoFlow取得了1.91倍的吞吐量提升。

方法细节中的引用汇总