文章标题: 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 操作特性

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

  • 密集操作(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向量,因此是内存密集型的。
  • 网络操作(Network operations):包括AllGather (AG) 和 AllReduce (AR) 在内的集体通信操作是网络密集型的,需要节点间的高带宽互连,如NVLink【索引编号25, Nvlink and nvlink switch, 2024】。
  • 其他操作(Other operations):我们将Transformer架构中剩余的操作(如层归一化、位置嵌入等)归为其他操作,因为它们与密集或注意力操作相比运行时间很短。


图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】对不同的设备间并行范式进行了全面评估:

  • 张量并行(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,以支持更大的批处理大小,而流水线并行则用于跨节点进一步扩展集群。

3 分析

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

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

  • 硬件规格
    • $N_{GPU}$:GPU数量
    • $MemBW$ (GB/s):总GPU内存带宽
    • $MemSize$ (GB):总GPU内存容量
    • $Compute$ (GFLOP/s):总GPU计算能力
    • $NetBW$ (GB/s):总GPU互连带宽
  • 模型配置
    • $D_{model}$:隐藏层维度大小
    • $L$:层数
    • $P_{Model}$:参数数量
    • $R_{GQA}$:GQA的组大小
    • $S_{type}$ (Bytes):模型参数数据类型的字节数(例如,FP16为2)
  • 用户查询统计
    • $p$:待prefill的提示中平均token数
    • $d$:待解码的输出中平均token数
      因此,一个用户请求平均对应$p$个prefill token和$d$个decode token,总共$p+d$个token。使用用户查询统计数据,总吞吐量可以轻松转换为其他吞吐量指标,例如,解码吞吐量为 $\frac{d}{p+d} \times Throughput_{total}$,每秒请求数为 $\frac{1}{p+d} \times Throughput_{total}$。
  • 批处理大小:一个以吞吐量为导向的优化服务系统需要处理大批次。对于计算密集型的密集操作(如GEMM),更大的批次可以分摊权重加载开销并增加执行单元的占用率,从而提供更高的计算利用率。对于内存和网络密集型操作,所有请求共享相同的预热成本(如内核启动开销、流水线设置开销、网络同步开销等),因此更大的批处理大小能带来更高的吞吐量。因此,在考虑最优系统吞吐量时,我们假设系统始终在总可用内存能容纳模型权重和所有请求的KV缓存的最大批处理大小下运行。

3.2 LLM服务成本模型

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

  • 内存:从内存角度看,延迟为:

    这是因为在使用最大批处理大小时,每次迭代需要将整个设备内存内容加载到GPU缓存和寄存器中一次(这在现代服务系统中很典型【索引编号17, Efficient memory management for large language model serving with pagedattention, 2023】)。尽管模型权重在迭代间共享,但由于重用距离长,缓存模型权重以避免从设备内存加载是不可行的。
  • 计算:在LLM服务中,所有操作都贡献总计算量,而密集操作产生绝大部分计算(如§3.4验证)。因此我们对密集操作的计算延迟进行建模。对于密集操作中的每个GEMM,需要在权重和激活值上执行 $2B_{Dense}N_wK_w$ 次计算,其中 $N_w$ 和 $K_w$ 是权重矩阵的维度,$B_{Dense}$ 是密集操作的批处理大小。因此,密集操作的总计算量为 $2B_{Dense}\sum N_wK_w$。$\sum N_wK_w$ 项是所有层中权重元素的数量,可以用模型参数量 $P_{Model}$ 近似,因此表达式简化为 $2B_{Dense}P_{Model}$。所以,从计算角度看,延迟为:
  • 网络:对于张量并行,在多个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: 计算时间与内存时间的比较。颜色越接近黄色,工作负载计算密集程度越高;越接近绿色,则内存密集程度越高。
  • 跨加速器的普适性:尽管分析主要基于NVIDIA A100,但结果对其他加速器也类似。如表1所示,不同厂商和代际的加速器其计算/内存比和计算/网络比保持稳定,意味着工作负载特性也基本不变。
    总之,我们的成本模型表明,现代LLM服务工作负载在计算密集型区域运行。
    表1: 各厂商和发布年份的加速器型号特性。

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)进行性能分析。

  • 确定最大密集批处理大小:给定特定的模型架构和用户输入统计数据(如§3.1所述),NanoFlow计算出能够放入GPU内存的最大密集批处理大小(例如,对于8×A100服务LLaMA-2 70B,为2048)。这个密集批处理大小作为自动搜索的输入,为性能分析中的形状设定了上限。
  • 分析无干扰的内核:NanoFlow接着确定内核在无干扰情况下的执行时间,即每个操作独占使用GPU。NanoFlow以128的倍数,从128到密集批处理大小,对离散输入批处理大小的内核进行性能分析,因为128是GEMM分块的硬件友好形状。NanoFlow探索所有可能的内核实现,包括改变线程块数、warp数和GEMM、GEMV及网络内核的tile大小,以找出执行时间最短的实现。分析结果输出一个从(内核,批处理大小)到其最佳实现的映射。
  • 分析有干扰的内核:理想情况下,我们希望独立地为每个纳米操作分配GPU计算、内存和网络带宽的一部分。我们称这个分数为$R_{physical}$。然而,当两个或多个内核在设备上并行执行时,它们会比单独运行时慢。我们称之为内核干扰,其原因是内核竞争GPU硬件资源,如执行单元、缓存和内存带宽【索引编号43, Orion: Interference-aware, fine-grained gpu sharing for ml applications, 2024】。不幸的是,内核干扰是不可预测的,因为NVIDIA GPU不提供对计算、内存和网络带宽的明确控制。因此,我们无法直接控制$R_{physical}$。
  • 使用GEMM性能作为代理并建模干扰:我们使用GEMM性能作为$R_{physical}$的代理$R$,并对成对的内核干扰进行建模。我们以GEMM为中心定义$R$,因为计算密集型操作占主导地位,而优化计算是NanoFlow的目标。具体来说,我们通过测量计算内核A与另一内核B重叠时的实际性能来分析成对的内核干扰模式。例如,如果重叠时,内核A的性能达到其单独运行最佳性能的40%,我们就说$R_A=0.4$。然后我们假设重叠的内核获得了剩余的资源,所以我们设置$R_B = 1 - R_A = 0.6$。接着,我们通过测量当分配给内核A的资源为$R_A$时,非计算内核B的性能损失来建模干扰。当A和B重叠时,我们将A和B的性能归一化到它们单独运行时的峰值性能。我们称这个值为$P$。对于计算内核A,$P_A$根据定义等同于$R_A$,而$P_B$则捕捉了当$R_A$资源分配给计算内核时,内存或网络的利用率。
    直观地说,这个过程建立了一个$R$计算利用率到$P$内存或网络利用率的“汇率”。整个内核干扰模型可以被捕捉在一个表中,如表3,该表为每种内核类型将$R$映射到$P$。

  • 减少分析空间:尽管如此,由于可能的内核组合数量庞大,分析$R$和$P$之间的映射仍然具有挑战性。NanoFlow通过将GEMV和网络内核的线程块数限制在8到128之间(步长为8)来减少分析空间,并排除效率较低的GEMM内核。NanoFlow专注于成对干扰分析(计算-内存和计算-网络),而不是同时检查三种内核类型,并假设成对干扰分析出的$R$到$P$的映射在三种内核重叠时仍然成立。

图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问题受以下约束:

  • 优化目标:MILP的主要目标是消除计算操作的流水线气泡(如图4中的wasted所示),以最小化执行时间。
  • 纳米操作数量的约束:为了重叠资源,每个操作需要至少被分割成两个纳米操作。由于分割会降低批处理效果,纳米操作的数量应最小化。因此,自动搜索从将所有操作分为两个纳米操作开始,并构建MILP问题来确定批处理大小、执行时间和顺序。如果临时最优解存在计算气泡,自动搜索会增加气泡附近操作的纳米操作数量,直到MILP无法产生更好的解。
  • 批处理大小和执行时间的约束:自动搜索强制批处理大小从128到密集批处理大小中选择,并使用§4.1.1中计算的给定批处理大小下纳米操作的无干扰执行时间。
  • 依赖关系的约束:纳米操作的依赖关系由它们的父操作(存在于基线实现中,如pytorch)及其输入批次决定。两个纳米操作是依赖的,当且仅当它们的父操作是依赖的且它们的输入批次相交。
  • 重叠的约束:重叠受相同资源(如计算)约束的内核对提高利用率无益,所以自动搜索限制流水线只重叠受不同资源约束的操作。
  • 操作变换的约束:一些网络集体通信有具有不同性能特征和依赖关系的等效变换。例如,一个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探索所有这些网络纳米操作的替代变换,并搜索最短的流水线设计。

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

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

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

  • GPU资源利用率的约束:如§4.1.1所述,并发内核应竞争总计1.0的GPU资源利用率R。因此,自动搜索强制任何给定时间的R之和小于或等于1.0。
  • 执行时间的约束:给定纳米操作的资源利用率R,自动搜索使用 $D_{best}/P$ 计算操作的执行时间,其中$D_{best}$代表无干扰时的最佳执行时间,P使用表3得出。

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

4.1.4 示例流水线
  • 70B流水线:我们将自动搜索应用于70B规模的模型,包括LLaMA-2 70B、LLaMA-3 70B、Qwen2.5-72B和Deepseek-67B。与LLaMA-2 70B相比,LLaMA-3 70B的词汇量更大,增加了采样时间;Deepseek-67B的层数和隐藏维度略有不同;Qwen2.5-72B在KQV生成中引入了偏置。然而,这些调整并未显著改变性能特性,导致调度相似。图6中我们提供了为LLaMA-2 70B生成的NanoFlow流水线作为示例。由于在解码层的开始部分有三种资源(计算、内存、网络)重叠,自动搜索最终为这部分使用了4个纳米操作以减少流水线气泡。在这里,解码注意力操作的资源利用率为0.4,达到了最大解码注意力性能的80%。在流水线的其余部分,GEMM操作被优先考虑,只使用两个纳米操作。

    图6: LLaMA-2 70B的执行流水线,由NanoFlow自动生成。实心背景和阴影背景分别代表输入批次0-768和768-2048。R代表资源利用率。通过重叠计算密集、内存密集和网络密集的操作,NanoFlow提高了计算利用率并改善了服务吞吐量。
  • 8B流水线:8B模型(如Llama3-8B)不需要网络操作,因为它们可以装入单个GPU。因此,自动搜索将所有操作分成两个纳米操作,其中解码注意力与Up Gate Down投影重叠。
  • MoE流水线:与70B模型相比,MoE模型具有不同的隐藏维度和层数。由于专家之间的不平衡【索引编号16, Mistral 7b, 2023】,NanoFlow对MoE使用张量并行,其中FFN使用分组GEMM实现。FFN还涉及一个额外的门控路由操作。自动搜索对MoE模型的工作方式类似,并自动生成一个高效的流水线(§6)。

4.2 NanoFlow 运行时

4.2.1 请求调度
  • 批次形成: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请求进行分块,以精确填满预选的最佳性能密集批次的剩余容量。这有效地减少了尾部延迟,因为密集操作在迭代中以一致的批处理大小执行。
  • 内存管理与预测:为了优化GPU内存使用并避免内存耗尽,NanoFlow根据用户请求的状态预测未来的峰值内存使用情况。NanoFlow跟踪每个请求已解码的token数,使用平均解码长度估算完成时间,并且仅在预测的内存使用量保持在GPU限制内时才预填充新请求。如果GPU内存耗尽,NanoFlow会将一个请求卸载到CPU,并在内存可用时重新加载,无需重新计算。
  • 异步调度:批次形成过程,包括估算内存使用、调度新请求、移除完成的请求以及调整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%)。
4.2.2 KV缓存管理
  • 同步卸载:为了高效服务多轮应用,NanoFlow将运行中请求的KV缓存卸载到由主机CPU内存和SSD组成的分层缓存中,并在新一轮会话到达时加载前一轮的KV缓存。NanoFlow在每个Transformer层的KQV生成后直接卸载token的KV缓存,而不是等待请求完成。由于NanoFlow的密集批次保持稳定,卸载的数据大小在迭代间是平衡的。为最小化卸载开销,NanoFlow在FFN的计算密集型操作期间执行设备到主机的拷贝——这是一个使用少量GPU资源的GPU发起的拷贝操作。此外,NanoFlow通过NUMA感知的线程绑定来减少卸载时间。
  • 主机KV缓存管理:NanoFlow使用LRU策略来管理CPU内存和SSD的分层缓存。例如,当CPU达到其内存限制时,KV缓存被驱逐到SSD,并在请求的下一轮对话到来时从CPU或SSD中检索以启动加载过程。
  • KV缓存加载与分散:NanoFlow使用PagedAttention【索引编号17, Efficient memory management for large language model serving with pagedattention, 2023】,因此KV缓存的页面在GPU内存中是碎片化存放的。为了避免拷贝到碎片化的GPU页面目的地,NanoFlow首先将数据拷贝到GPU上的一个连续空间,然后将页面分散到它们的目的地。这使得主机到设备的拷贝带宽提高了7-10倍。

5 实现

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

A4 实验环境

  • 硬件配置:
    • GPU: 8块 NVIDIA A100 80GB SXM GPU,通过NVLink互连。
  • 模型配置:
    • 主要评估模型: LLaMA-2-70B。
    • 其他评估模型: LLaMA-3-70B, LLaMA-3-8B, QWen2-72B, Deepseek-67B, Mixtral 8×7B。
    • 数据类型: 所有模型均使用FP16权重和激活。
  • 软件配置与基线:
    • vLLM: 一个先进的高吞吐量服务系统,实现了PagedAttention和分块prefill。
    • DeepSpeed-FastGen: 微软开发的服务框架,动态组合prefill和decode请求。
    • TensorRT-LLM: 基于NVIDIA TensorRT SDK的高性能LLM推理引擎,启用了paged KV-cache和动态批处理优化。
  • 数据集:
    • Splitwise: 微软生产环境收集的对话追踪数据,约20000个请求。
    • LMSYS-Chat-1M: 包含100万个真实世界对话的大规模数据集。
    • ShareGPT: 从ShareGPT API收集的对话数据集。
    • 采样: 从ShareGPT和LMSYS-Chat-1M中随机采样50,000个请求用于评估。数据集的输入输出长度统计见表4。

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

A4 实验结果

6.2 吞吐量

  • 实验内容: 评估NanoFlow在离线场景下的总吞吐量(输入和输出token总数/执行时间),并与vLLM、DeepSpeed-FastGen、TensorRT-LLM以及理论最优吞吐量进行比较。工作负载的输入输出长度从真实数据集(Splitwise, LMSYS, ShareGPT)采样,并辅以固定长度的实验。
  • 实验结果:
    • NanoFlow在所有设置下均表现出最高的吞-吐量,在最佳情况下达到了理论最优值的68.5%(如图7)。
    • 对于固定长度,NanoFlow的吞吐量平均比vLLM高2.62倍,比DeepSpeed-FastGen高2.78倍,比TensorRT-LLM高1.73倍。
    • 对于从数据集中抽取的长度,NanoFlow的吞吐量平均比vLLM高4.18倍,比DeepSpeed-FastGen高3.45倍,比TensorRT-LLM高1.91倍。
  • 分析结论: NanoFlow通过设备内并行显著提升了吞吐量,大幅超越了现有的SOTA系统,并有效逼近了理论性能上限。


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

6.3 延迟

  • 实验内容: 在请求到达间隔服从指数分布的在线服务场景下,评估NanoFlow的平均归一化延迟(端到端请求延迟/输出token数)。逐步增加请求率,测量在200ms的服务水平目标(SLO)内系统能维持的最大请求率。
  • 实验结果:
    • 在低请求率下,NanoFlow的延迟与基线相当但略高,因为它为吞吐量场景优化,使用了较大的批处理大小。
    • 在所有数据集中,NanoFlow均能在200ms延迟SLO内维持比所有基线更高的请求率。例如,在LMSys-Chat-1M数据集上,NanoFlow可支持的请求率比TensorRT-LLM高1.64倍(如图8)。
    • NanoFlow的P99延迟仅为平均延迟的1.07倍(在接近最大吞吐量时),显示出良好的尾部延迟性能,这得益于其恒定的密集批处理大小。
  • 分析结论: NanoFlow不仅在吞吐量上领先,在满足延迟SLO的前提下,也能处理更高的请求负载,表现出优越的在线服务能力。


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

6.4 消融研究

  • 实验内容: 为了量化NanoFlow中各项技术(纳米批处理、纳米操作重叠、KV缓存卸载)的贡献,将其与两个变体基线进行比较:一个是不使用纳米批次的顺序执行基线,另一个是使用纳米批次但顺序执行的基线。
  • 实验结果:
    • 纳米批处理开销: 仅引入纳米批处理而不重叠,会导致性能下降13.2%,这证实了批处理效应减弱带来的开销(如图9)。
    • 操作重叠收益:
      • 在prefill-only工作负载中,通过重叠网络密集型和计算密集型内核,NanoFlow获得了1.07倍的加速。
      • 在decode-heavy工作负载中,通过同时重叠网络和内存密集型内核,NanoFlow获得了1.17倍的加速。
    • KV缓存卸载开销与收益: 启用KV缓存卸载会因内核干扰导致流水线性能下降3.0%。然而,对于多轮对话的LMSYS-Chat工作负载,卸载可以减少3.02倍的重计算。
  • 分析结论: 尽管纳米批处理会带来开销,但通过操作重叠获得的收益远大于此开销,是NanoFlow性能提升的关键。KV缓存卸载则是在轻微性能损失下换取显著计算节省的有效权衡。


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

6.5 资源使用

  • 实验内容: 对比NanoFlow与非重叠基线在推理单个层时的计算、内存和网络资源利用率。
  • 实验结果: 如图10所示,非重叠基线在任一时刻主要只利用一种资源。而NanoFlow能够并发利用多种资源,实现了68.5%的平均计算利用率。
  • 分析结论: NanoFlow通过设备内并行有效地将流水线气泡(资源空闲时间)转化为生产力,显著提高了瓶颈资源(计算)的利用率,这是其高吞吐量的根本原因。由于内核干扰,计算利用率仍低于理论最优值。


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

6.6 在其他LLM上的性能

  • 实验内容: 评估NanoFlow在LLaMA-3-70B, LLaMA-3-8B, QWen2-72B, Deepseek-67B, Mixtral 8×7B等多个流行LLM上的吞吐量表现,并与vLLM进行比较。
  • 实验结果: 如图11所示,对于这些模型,NanoFlow将吞吐量提升至理论最优值的50%至72%之间,显著超过vLLM。
  • 分析结论: NanoFlow的自动搜索机制能够很好地适应不同的模型架构(包括MoE),并自动生成高效的流水线,证明了其方法的通用性和有效性。


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

A5 结论

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

方法细节中的引用汇总

  • [3] Sarathi: Efficient llm inference by piggybacking decodes with chunked prefills (2023)
    • 引用段落: §2.2 操作特性, §4.2.1 请求调度
    • 原文描述: 引用来说明prefill和decode阶段可以组合以利用批处理效应,分摊权重加载成本。在方法中,NanoFlow遵循其以token粒度对prefill请求进行分块的思想。
  • [4] Mistral nemo (2024)
    • 引用段落: §3.3 LLM服务工作负载分类
    • 原文描述: 引用作为现代模型广泛采用GQA的一个例子。
  • [5] GQA: Training generalized multi-query transformer models from multi-head checkpoints (2023)
    • 引用段落: 摘要
    • 原文描述: 引用来说明分组查询注意力(GQA)等优化可以减少内存负载。
  • [6] Layer normalization (2016)
    • 引用段落: §2.1 LLM 推理工作流
    • 原文描述: 引用来说明在注意力阶段后会进行层归一化操作。
  • [9] Alibaba cloud’s qwen2 with enhanced capabilities tops llm leaderboard (2024)
    • 引用段落: §3.3 LLM服务工作负载分类
    • 原文描述: 引用作为现代模型广泛采用GQA的例子,并指出其$D_{model}$通常大于4096。
  • [10] Sigmoidweighted linear units for neural network function approximation in reinforcement learning (2017)
    • 引用段落: §2.1 LLM 推理工作流
    • 原文描述: 引用来说明全连接网络(FFN)阶段可能使用的激活函数SiLU。
  • [12] Fastermoe: Modeling and optimizing training of large-scale dynamic pre-trained models (2022)
    • 引用段落: §2.3 规模化服务LLM
    • 原文描述: 作为评估不同设备间并行范式的现有工作之一被引用。
  • [13] Deepspeed-fastgen: High-throughput text generation for llms via mii and deepspeed-inference (2024)
    • 引用段落: §3.6 与最优吞吐量的差距
    • 原文描述: 引用作为现有的一个服务框架,其实际吞吐量远低于理论最优值。
  • [14] Gpipe: Efficient training of giant neural networks using pipeline parallelism (2019)
    • 引用段落: §2.3 规模化服务LLM
    • 原文描述: 引用来介绍流水线并行的概念。
  • [15, 33, 45]
    • 引用段落: §4.2.1 请求调度
    • 原文描述: 引用来说明NanoFlow假设自动扩缩、负载均衡等由外部控制平面管理。
  • [16] Mistral 7b (2023)
    • 引用段落: §2.1 LLM 推理工作流, §3.3 LLM服务工作负载分类, §4.1.4 示例流水线
    • 原文描述: 作为基于decoder-only Transformer架构的LLM例子;其$D_{model}$大于4096的例子;以及在MoE流水线中,指出其专家间存在不平衡。
  • [17] Efficient memory management for large language model serving with pagedattention (2023)
    • 引用段落: §3.2 LLM服务成本模型, §3.6 与最优吞吐量的差距, §4.2.1 请求调度, §4.2.2 KV缓存管理
    • 原文描述: 引用来说明现代服务系统在最大批处理大小时需要加载整个设备内存;作为现有服务框架吞吐量未达最优的例子;NanoFlow的异步调度需要调整其页表;NanoFlow使用其PagedAttention技术来管理KV缓存。
  • [18] Alpaserve: Statistical multiplexing with model parallelism for deep learning serving (2023)
    • 引用段落: §2.3 规模化服务LLM
    • 原文描述: 引用来说明大型模型服务需要多个GPU。
  • [22, 48]
    • 引用段落: §3.3 LLM服务工作负载分类
    • 原文描述: 作为现代模型$D_{model}$通常大于4096的例子。
  • [25] Nvlink and nvlink switch (2024)
    • 引用段落: §2.2 操作特性, §3.3 LLM服务工作负载分类
    • 原文描述: 作为网络密集型操作所需的高带宽互连技术例子;以及数据中心GPU间的高带宽互连技术例子。
  • [34] Efficiently scaling transformer inference (2023)
    • 引用段落: §2.1 LLM 推理工作流
    • 原文描述: 引用来介绍KV缓存的概念,它可以加速decode阶段。
  • [38, 57]
    • 引用段落: §2.3 规模化服务LLM
    • 原文描述: 作为评估不同设备间并行范式的现有工作被引用,并具体介绍张量并行的概念。
  • [39, 57]
    • 引用段落: §3.2 LLM服务成本模型, §4.1.2 自动搜索第一阶段
    • 原文描述: 引用来说明支持张量并行所需的集体通信操作数量;以及说明网络集体通信操作(如AG)可以有等效变换(如AR)。
  • [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)
    • 引用段落: §2.1 LLM 推理工作流, §2.2 操作特性
    • 原文描述: 引用来说明LLM推理的两个阶段(prefill和decode);以及说明批处理解码会聚合所有解码请求的激活值。
  • [54] Forestcoll: Efficient collective communications on heterogeneous network fabrics (2024)
    • 引用段落: §3.3 LLM服务工作负载分类
    • 原文描述: 作为数据中心GPU间的高带宽互连技术(Infinity Fabric)的例子。
  • [58] Sglang: Efficient execution of structured language model programs (2024)
    • 引用段落: §3.6 与最优吞吐量的差距
    • 原文描述: 引用作为现有服务框架的一个例子,其执行流程是顺序的。