作者: ZxZhao,复旦大学集成电路与系统设计博士
https://zhuanlan.zhihu.com/p/1988221694591128167
大模型部署技术的演进太快了,相信有很多朋友两三年没有关注这个领域,对相关的名词已经是一头雾水了。这篇文章以杂谈的形式,和大家聊聊这里的几个关键技术节点。
先拎出来个时间线,再展开补充技术细节。
最早,人们是沿袭着深度学习时期的思路来进行LLM的部署,简单来说就是简单粗暴的padding不同请求然后通过Static Batching进行推理部署。
之后,针对模型计算角度的技术不断被提出,Flash Attention是核心技术之一,解决了Attention部分的tile-based处理问题。
之后Continuous Batching引入更细粒度的调度和PagedAttention引入了更细粒度的显存管理,以及到Chunked Prefill解决了流水线阻塞问题。慢慢PD分离的架构逐渐成型且成为主流。
再到今年,为了更进一步利用模型各部分的计算特性,达到更高的系统级效率,AF分离逐渐成为新的趋势。总的来说,整个技术演进是通过软硬件协同设计,不断优化整个分布式系统在大模型推理上的综合表现的过程。
Background
方便无大模型经验的读者理解后面的故事,先补充一点最基本的大模型算法知识。
Transformer模型的推理过程并非单一维度的计算负载,而是分为两个截然不同的阶段:Prefill与Decode。Prefill阶段处理输入的Prompt,计算模式为矩阵-矩阵乘法(GEMM),具有较高的算术强度(Arithmetic Intensity),能够有效利用GPU的Tensor Core,通常属于Compute-bound场景。
进入Decode阶段,模型转为自回归生成,每生成一个Token,都需要读取之前所有Token的KV Cache。这是一个典型的矩阵-向量乘法(GEMV)过程。属于Memory-bound场景。
随着生成序列变长,KV Cache的体积线性增长。对于一个70B的模型、128k的上下文,KV Cache可能达到数GB甚至数十GB。每生成一个Token,就需要搬运这数十GB的数据穿过显存带宽,而计算量却相对极小。
另外从模型本身角度,可以分为两个部分Attention部分和FFN部分(有时也叫MLP,后文MoE也是指这部分)。Attention负责捕捉Token之间的依赖关系。其计算复杂度与序列长度呈二次方 O(N^2)(在Flash Attention之前)或线性关系。它是KV Cache的主要驻留地,也是长文本推理的显存瓶颈所在。
FFN部分,在标准Transformer中,它占据了模型参数量的2/3以上(MoE网络中会更多,90%以上)。对于MoE(混合专家)模型,这部分被替换为多个专家网络。它是模型知识的主要存储地,也是模型参数量的主要来源。
另外一个关键点,为了提高部署效率,Batching 在多用户场景下是必须的。
Flash Attention
在讨论系统级调度之前,必须先谈论一个关键技术Flash Attention,FA甚至不能被当作一个简单的关键技术,个人认为它甚至是大模型部署的基石。简单来讲,Flash Attention可以理解为,在attention的模型结构中,实现了softmax的tile-based处理,并且和原本的模型结构在计算上是等价的。
在2022年之前,Transformer的注意力机制 O(N^2) 的显存占用和访问开销限制了上下文长度的扩展,32k甚至8k的上下文在当时被认为是极难部署的。Tile-based的attention,Flash attention就成为了一个救星。
至于Flash attention版本的不断迭代,其实是对最新Nvidia GPU架构的不断微调的过程。FA1 首先通过 Tiling 和重计算策略将 O(N^2) 的 HBM 访问转化为 SRAM 操作,大幅缓解了带宽瓶颈;
FA2 针对 Ampere 架构特性,引入了sequence维度的并行,并将工作负载划分细粒度至 Warp 级,优化了 Thread Block 间的同步开销以提升 Occupancy;
FA3 为 Hopper 架构定制,利用 TMA 引擎实现数据的异步搬运,配合 WGMMA 指令集,在 Warp Group 层面实现了 GEMM 计算与 Softmax 的深度流水线overlap,进一步提升效率。
Static batching
为了应对这种workload,在 2021-2022 年左右,如在生产环境高性能部署 GPT-3 或 BLOOM 这样的模型,当时的方案是使用 NVIDIA 的 Triton 推理服务器作为前端,后端挂载 FasterTransformer 引擎。注意这里的 Triton 和后面 OpenAI 的 triton 语言完全是两个东西。
Nvidia Triton 服务器。维护一个请求队列。它会等待积攒到预设的 Batch Size(比如 8),或者达到超时时间(比如 50ms),然后将这组请求打包成一个巨大的 Tensor 发送给 GPU 上的模型。最开始会采用这种技术的主要原因是,当时在 Transformer 架构刚刚接管 NLP 任务,推理系统几乎完全照搬了深度学习训练阶段的 Static Batching。
这样的方案在现在看来是非常低效的,因为LLM推理具有显著的“变长”特性,不同请求的Prompt长度参差不齐,生成的Output长度更是不可预测。
在静态批处理的框架下,为了利用GPU的并行计算能力,系统会受限于最大序列长度,将一个Batch内的所有请求按照最长序列进行Padding dummy数据。这种做法在工程上比较粗糙,Padding操作不仅浪费了宝贵的显存容量,更导致GPU大量的算力被消耗在无效的“零计算”上。
另外,请求间的相互阻塞也是很大的一个问题,整个Batch的端到端延迟(Latency)完全取决于生成长度最长的那个请求。即便某个短请求仅需生成10个Token,它也必须在显存中“陪跑”,直到Batch中那个需要生成1000个Token的长请求结束。
这种粗放的调度机制,显存碎片化严重,throughput被锁死在低位,单位Token的生成成本居高不下 。
Continuous Batching
当然这一情况并没有持续很久,很快Continuous Batching成为了主流。简单来说,调度器不再等待整个 Batch 结束。每完成一个 Token 的生成迭代,系统就会检查有没有请求结束了。
如果有,立即释放其占用的显存槽位,然后检查等待队列,拉取一个新的请求插入到刚才空出的槽位中。这样,GPU 的 Batch Size 始终保持在硬件允许的上限,流水线被填满。
支撑这一调度的底层技术是PagedAttention。 在传统 PyTorch 实现中,KV Cache 要求显存必须连续分配。这导致了严重的内存碎片。
PagedAttention 借用了操作系统虚拟内存的思想,将 KV Cache 切分为固定大小的Block,比如每块存 16 个 Token。逻辑上连续的 KV Cache,在物理显存中可以是非连续的,通过维护一张Block Table来记录逻辑块到物理块的映射。
这一时期的代表是Orca和vLLM。
Chunked Prefill
然而,随着应用场景向长文档分析拓展,上下文长度(Context Length)开始呈指数级增长,新的系统瓶颈又出现了。
前文也提到了大模型推理会分为Prefill和Decode两个阶段。在早期的部署模式中,这两个阶段被放在同一个GPU上执行。
当一个长Context(如32k)的请求进入Prefill阶段时,会瞬间占满GPU的计算资源,导致正在进行Decode任务的其他请求被迫暂停。用户感知就是在逐token输出的过程中,会突然有一个卡顿,很影响实际的使用体验。
为了缓解这一问题,Chunked Prefills(分块预填充)技术被提出。它不再一次性处理完长Prompt,而是将其切分成多个小块(Chunk),利用Decode阶段的计算空隙执行。
在每个调度周期,GPU 处理一个 Chunk 的 Prefill,加上其他请求的 Decode。巨大的 Prefill 延时分摊到了多个decode过程中,这种方式平滑了系统整体的响应延迟
相关技术里以Sarathi系统为代表。
PD 分离
再往后,随着模型越来越大,Prefill和Decode两个阶段放在同一节点,其内在计算特性不同导致的对物理硬件需求的冲突越来越突显。强行放在同一张卡上,要么 Prefill 跑不满 Tensor Core,要么 Decode 跑不满带宽。
PD分离的架构逐渐成为主流。需要注意一下,这里的PD分离并不是指算法上分为Prefill和Decode两个阶段(这个早已是业界共识了),这里更多是指部署Prefill和Decode阶段的物理节点不同以及对两种场景下的硬件差异化趋势。
在DistServe等框架的推动下,推理集群被划分为两个独立的资源池:Prefill Instance和Decode Instance。用专用的Prefill Instances,也就是专用的Prefill节点群,只负责吞吐 Prompt,生成 KV Cache。专用的Decode节点群,只负责接收 KV Cache,自回归生成 Token。
Prefill节点可以配置算力强大的GPU(如NVIDIA H800),专注于快速吞噬Prompt并生成KV Cache;Decode节点则可以配置显存带宽较高或成本更低的GPU,专注于高效的Token自回归生成。当Prefill节点完成计算后,生成的KV Cache通过高速RDMA传输给Decode节点。
这种物理上的解耦,给整个系统的throughput和效率带来了巨大的提升。但新的瓶颈又变成了,KV cache的传输,这给系统带来了很大的带宽压力。在大Batch、长Context场景下,KV Cache的数据量极为庞大(例如70B模型、128k Context下的KV Cache可达数GB)
这一阶段,各种各样的并行策略是探索的一个重要主题。比如在Prefill阶段使用Pipeline Parallelism(流水线并行)来扩展上下文长度,在Decode阶段使用Tensor Parallelism(张量并行)来降低延迟。
值得一提的是,面对 PD 分离带来的带宽压力,Mooncake,也就是Kimi 背后的架构,提出了一种以 KV Cache 为中心的部署方式。这个工作值得仔细研读。
总的来说,它利用 GPU 集群中闲置的 CPU 内存、DRAM 甚至本地 NVMe SSD,构建了一个全局分布式的 KV Cache 池。在Mooncake中,直接利用RDMA和GPUDirect技术,在不同机器的 GPU 显存、CPU 内存之间进行的高速搬运。热数据在 HBM,温数据在 CPU DRAM,冷数据在 SSD。系统像管理数据库一样管理 KV Cache。
AF 分离
随着模型MoE(Mix of Expert)架构的模型逐渐称为主流,新的系统级挑战出现了。MoE的一个显著特点是,其MLP/FFN部分会有极大规模的权重,但在单次inference过程中,只会激活其中的一部分权重,以专家为单位,计算视角来看,专家本质上也就是包含几个离线训练好的FC层的sub-module。但是专家的激活是per-token的,这意味着在大batch情况下,根据大数定律,绝大部分的专家都会被激活。
这个传统的MoE部署中,Attention和Expert通常部署在同一个GPU上。Attention层需要大量的显存来存储KV Cache,而Expert层则需要大量的显存来存储海量的专家权重。
注意这里AF分离,我们关注点是模型内部的两个部分计算特点的不同。前文PD分离是整个模型在两个推理阶段的计算特点的不同。本质上来讲,AF分离是在PD分离基础上的,进一步对系统资源的细粒度调配优化。
AF分离架构将Attention层和Expert层物理拆分,部署在不同的GPU集群上。Attention Cluster专注于维护KV Cache和计算Attention,Expert Cluster则专注于存储专家权重和计算FFN。
这种拆分带来了极高的灵活性:系统可以根据Attention和FFN不同的计算量比例,独立配置两类集群的资源。例如,对于长Context任务,可以增加Attention节点以容纳更多KV Cache;对于复杂推理任务(激活更多专家),可以增加Expert节点以提升FFN计算能力。
这带来的代价也是很明显的,AF分离引入了高频的跨节点通信,每一层的Attention输出都需要发送给Expert集群,Expert计算完后再发回。这种通信模式是M-to-N的,且数据包极小、频率极高。
最大的问题变成了,每一层都要在 Attention 机器和 Expert 机器之间来回倒腾 Token。代表性的工作比如Janus, 设计了Adaptive Two-Phase Communication,先在 Attention 节点内部,利用超高带宽的NVLink,把发往同一个 Expert 机器的 Token 聚合起来,打成一个大包。
再通过InfiniBand网络,把大包发给 Expert 节点。Attention 和 Expert 分离了,调度器就可以在 Expert 集群中动态选择负载最低的 Expert 副本(Replica)来处理请求,实现了微秒级的负载均衡。
总结
总的来说,现在大模型推理优化,已经变成了一个高度复杂化的系统工程。越来越需要端到端的系统级优化,软硬件的协同优化,学科交叉人才的需求也越来越大。