MLSYS推理服务

1. LLM基础与生成机制

1.1 语言模型 (LM) 定义

语言模型 (Language Model, LM) 的核心功能是对一系列词元 (token) 序列的概率分布进行建模

  • 词汇表 (A\mathcal{A}):这是一个包含所有可能词元的集合。
  • 语言模型的概率赋值:对于词汇表 A\mathcal{A} 中的任意一个词元序列 A1,A2,,ALAA_1, A_2, \dots, A_L \in \mathcal{A},语言模型会为其分配一个介于 0 和 1 之间的概率值 P(A1,A2,,AL)[0,1]P(A_1, A_2, \dots, A_L) \in [0,1]
  • 概率的直观含义:这个概率值直观地表示了一个词元序列的“合理性”或“流畅性”。概率越高,序列越符合语言习惯或训练数据中的模式。

示例: 假设词汇表为 A={the,man,walks,on,street}\mathcal{A}=\{\text{the}, \text{man}, \text{walks}, \text{on}, \text{street}\}。 语言模型可能会分配以下概率:

  • P(the,man,walks,on,street)=0.02P(\text{the}, \text{man}, \text{walks}, \text{on}, \text{street}) = 0.02
  • P(the,street,man,on,walks)=0.01P(\text{the}, \text{street}, \text{man}, \text{on}, \text{walks}) = 0.01
  • P(man,the,the,street,walks)=0.0001P(\text{man}, \text{the}, \text{the}, \text{street}, \text{walks}) = 0.0001 从上述示例可以看出,第一个序列的概率最高,因为它在语义和语法上更合理。

1.2 LLM的自回归特性

目前主流的大型语言模型 (LLM) 主要采用自回归 (autoregressive) 架构

  • 自回归生成过程:这意味着模型在生成序列中的下一个词元 x[i+1]x[i+1] 时,会依赖于它已经生成的所有前序词元 x[1:i]x[1:i]
    • 具体来说,给定已有的词元序列 x[1:i]x[1:i],大模型会计算下一个词元 x[i+1]x[i+1] 的概率分布,并从中选择一个词元。
    • 这个过程是递归进行的,每次生成一个新词元,直到生成终止符或达到最大长度。
  • 推理与训练的区别
    • 训练阶段:通常可以进行批处理 (batch processing),模型可以同时处理多个序列,并计算其损失。
    • 推理阶段:由于自回归的特性,推理过程是序列化、多次运行模型前向传播 (FP) 的过程。这意味着,为了生成一个完整的回答,模型需要逐个词元地进行计算和输出。
  • 概率输出机制:通常,在Transformer模型的最后一层、最后一个词元的输出特征上,通过 Softmax 函数计算出词汇表中每个词元的概率分布,从而进行下一个词元的预测。

1.3 Token选择策略

在自回归生成过程中,模型输出的是下一个词元的概率分布。如何从这个分布中“选择”一个实际的词元,是生成效果的关键。

1.3.1 最大概率搜索

这类方法的核心思想是倾向于选择概率最高的词元。

1.3.1.1 Greedy Search
  • 原理:在每一步生成时,始终选择当前概率最高的词元作为下一个词元。
  • 特点
    • 简单直接,计算开销小。
    • 局部最优:但由于只考虑当前步的最优选择,可能会错过全局最优的序列。生成的文本可能听起来合理,但缺乏多样性和创造性。
1.3.1.2 Beam Search
  • 原理
    • Beam Search 维护一个固定大小的“束” (beam size, kk),同时跟踪 kk 条当前最有可能的序列路径。
    • 在每一步生成时,它会扩展这 kk 条路径,并从所有可能的扩展中选择新的 kk 条概率最高的序列,作为下一轮的候选。
    • 最终,选择这 kk 条路径中累计概率最高的作为输出序列。
  • 特点
    • 比 Greedy Search 更能找到全局最优解,通常能生成更流畅、语法更正确的文本。
    • 计算开销大于 Greedy Search。
1.3.1.3 Search类最大概率方法的优劣

对话模型不喜欢 Beam Search 的核心原因

  • 会放大高概率词元,导致输出死板:Beam Search倾向于选择最“安全”的路径,即那些在训练数据中出现频率高、概率大的词元组合。
  • 表现为
    • 更容易出现重复:模型可能陷入重复的短语或模式中,难以跳出。
    • 更倾向模板式说法:生成的内容缺乏多样性,听起来像固定的模板。
    • 丢失创造性、想象力:对于需要开放式、创新性回答的对话任务,Beam Search 的输出会显得过于保守和缺乏新意。
    • 长文本结构容易塌陷:在生成长文本时,模型可能会过早地收敛到常见的结束语或总结性短语(如“in conclusion …”),导致文本结构僵硬,缺乏自然的过渡。

适合使用 Beam Search 的场景: Beam Search 更适用于任务需要强确定性输出的场景,例如:

  • 机器翻译:目标是生成与原文意思一致且语法正确的译文。
  • 代码自动补全:需要生成语法正确的代码片段。
  • 工业写作模板:如客服回复、合同草拟等,需要标准化和高准确性的文本。
  • 大规模离线批处理生成 (Large-scale batch offline generation):在不需要实时交互、对生成质量有较高要求的场景。

1.3.2 概率抽样

为了解决最大概率搜索方法(特别是 Beam Search)在对话场景中缺乏多样性的问题,引入了概率抽样方法。这些方法在生成过程中引入了一定的随机性。

1.3.2.1 Top-k 抽样
  • 原理
    • 模型首先根据概率分布,找出概率最高的 kk 个词元
    • 然后,从这 kk 个词元中进行随机抽样,而不是简单地选择概率最大的那一个。
  • 作用:在保留一定质量的同时,引入了多样性。
  • 对话模型常用 kk:通常选择 k[40,100]k \in [40, 100]
1.3.2.2 Top-p 抽样 (Nucleus Sampling)
  • 原理
    • 首先将词汇表中的词元按概率降序排列。
    • 然后,从累积概率刚好超过预设阈值 pp 的最小词元集合中进行随机抽样
    • 例如,如果 p=0.9p=0.9,模型会选择累积概率达到或超过 0.9 的词元子集,然后从这个子集中进行抽样。
  • 作用:动态地选择词元数量,避免了 Top-k 中固定 kk 值可能带来的问题(当分布平坦时,Top-k 可能包含太多不相关的词元;当分布尖锐时,Top-k 又过于集中)。它能够更灵活地适应不同的概率分布形状。
  • 对话模型常用 pp:通常选择 p[0.8,0.95]p \in [0.8, 0.95]

2. Transformer生成过程中的计算冗余与优化

2.1 Transformer生成过程分析

Transformer 模型在生成文本时,是自回归(autoregressive)的。这意味着模型是逐个 token 进行预测和生成的。

  • 生成过程迭代性: 为了生成下一个 token,模型需要考虑之前所有已生成的 token。在没有优化的情况下,每次生成新 token 时,整个序列(包括已生成的和新加入的 token)都需要重新进行前向传播 (Forward Pass) 计算,包括注意力机制和前馈网络。
  • 计算冗余: 这种逐个 token 生成的方式导致了大量的重复冗余计算。具体来说,当生成第 i+1i+1 个 token 时,模型需要重新计算从第 1 个 token 到第 ii 个 token 的所有中间激活值(包括 Key, Value, Query 等)。然而,这些激活值在生成第 ii 个 token 时就已经计算过了。
  • 核心问题: 这意味着每次生成一个新的 token,新的计算其实只针对于刚加入的 token 的前向传播过程,而之前 token 的计算结果却被重复计算,造成了巨大的性能浪费。
    • (若需查看原始图片详情,请参考原文中的“Transformer生成过程:token 1”和“Transformer生成过程:token 2”示意图,以及“GPT结构图”,它们通常展示了模型逐层、逐token处理数据的过程,以及这种处理方式如何导致重复计算的。)

2.2 KV Cache 优化

针对上述计算冗余问题,KV Cache 成为了一种关键的优化技术。

2.2.1 KV Cache 原理

  • 定义: Key-Value 缓存 (KV Cache) 是一种存储机制,用于缓存Transformer 模型中每个自注意力层在处理先前 tokens 时生成的 Key (K) 和 Value (V) 向量
  • 存储内容: 在 Transformer 的自注意力机制中,对于每个 token,都会计算出其对应的 Query (Q)、Key (K) 和 Value (V) 向量。KV Cache 就是用来存储 K 和 V 向量的。
  • 目的: 其核心思想是在后续的推理过程中重用这些已计算的 K 和 V,避免重复计算。

2.2.2 增量推理

  • 工作机制: 当模型接收到新 token 时,不再需要重新计算整个序列的 K 和 V 矩阵。
    1. 模型只计算新 token 自身的 Query (Q)
    2. 然后,将这个新 token 的 Query 与缓存中所有先前 token 的 Key (K) 矩阵以及新 token 自己的 Key进行注意力计算。
    3. 同时,将这个新 token 的 Value (V) 向量也添加到 KV Cache 中,以备后续 token 的生成使用。
  • 优势: 这种方法被称为增量推理 (Incremental Inference)。它大大减少了计算量,特别是对于长序列的生成,因为大部分 K 和 V 的计算都被缓存并重用了,从而显著提升了推理速度

2.3 生成式推理 (Generative Inference)

由于用户输入 (prompt) 和模型输出 (answer) 的处理方式存在差异,生成式推理通常分为两个阶段。

  • 输入与输出的差异:
    • 人们输入的提示词 (prompt) 通常是一次性输入的,是一个完整的序列。
    • 模型的回答 (answer) token序列输出的,即逐个 token 生成。

因此,处理这两种类型的任务通常采用不同的策略:

2.3.1 预填充阶段 (Prefill phase)

  • 目的: 处理用户输入的完整 prompt。
  • 过程: 模型将整个提示(prompt)序列作为输入,进行一次完整的前向传播。在这个过程中,为每个 Transformer 层生成并填充键值缓存 (KV cache)。此时的计算模式与模型训练时的批处理(batch processing)相似。
  • 输出: 得到 prompt 中所有 token 的 KV Cache,以及第一个要生成的回答 token 的 logits (未归一化概率)。

2.3.2 解码阶段 (Decode phase)

  • 目的: 逐个生成回答 token。
  • 过程: 在预填充阶段完成后,模型进入解码阶段。在每一步解码过程中,模型利用增量推理(如 2.2 节所述),计算当前新生成的 token 的 Query,并结合更新后的 KV 缓存进行注意力计算,最终预测下一个 token。然后将新 token 的 Key 和 Value 加入 KV Cache,并重复此过程。
  • 特点: 这个阶段是自回归的、序列化的。

2.4 性能分析

理解 Prefill 和 Decode 阶段的计算特性对于优化模型性能至关重要。

2.4.1 Prefill与Decode阶段的计算特点

  • 预填充阶段 (Prefill phase):
    • 计算受限 (Computation bounded): 此阶段需要处理整个 prompt 序列,涉及大量的矩阵乘法和复杂的注意力计算。这些操作通常受限于 GPU 的计算能力(flops),即每秒浮点运算次数。
    • 原因: 需要处理较长的输入序列,并行计算量大。
  • 解码阶段 (Decode phase):
    • I/O 或带宽受限 (IO/bandwidth bounded): 在此阶段,每次生成一个新 token,计算量相对较小。主要的开销在于从内存中读取和写入 KV Cache,以及将模型参数加载到计算单元。这些操作受限于 GPU 的内存带宽,即每秒数据传输速率。
    • 原因: 每次只处理一个新 token,但需要频繁访问和更新可能非常大的 KV Cache。

2.4.2 内存占用组成

生成式推理过程中,主要的内存占用包括:

  • 模型参数: 存储大型语言模型(LLM)的所有权重和偏置。这部分内存占用通常是固定的,且非常大。
  • KV 缓存: 存储每个 Transformer 层中所有已处理 token 的 Key 和 Value 向量。这部分内存会随上下文长度的增加而正比增长,是导致推理内存占用过高的一个重要因素。上下文越长,KV Cache 越大。

3. 高级推理优化技术

本章将深入探讨两项关键的高级推理优化技术:KV Cache内存管理(PagedAttention)和投机采样(Speculative Decoding),它们旨在解决LLM推理过程中遇到的计算和内存瓶颈。

3.1 KV Cache 内存管理 (PagedAttention, SOSP’23)

KV Cache 是 Transformer 模型在自回归生成过程中存储历史 Key 和 Value 矩阵的机制。然而,随着序列长度的增加和多用户并发请求,KV Cache 会带来显著的内存管理挑战。

3.1.1 KV Cache 大小问题

  • 存储需求大:KV Cache 需要存储每个 Transformer 层中每个 token 的历史中间结果(Key 和 Value 向量)。对于大型语言模型,这会占用大量显存。
  • 训练与推理的差异:在训练阶段,模型通常通过批处理(batch training)和已知序列进行,不需要持久化KV Cache。但在推理阶段,尤其是自回归生成时,模型需要回顾之前的 token,因此 KV Cache 对于加速生成至关重要。

3.1.2 内存碎片问题

由于LLM推理过程中请求的输入和输出长度具有高度不确定性,KV Cache的内存管理面临以下挑战:

  • 预留 (Reservation):系统为了未来的使用需要提前保留资源,即使这些资源当前尚未被完全利用。这可能导致资源分配不灵活。
  • 内部碎片 (Internal Fragmentation):由于输出长度未知,系统为了避免频繁重新分配,可能会过度分配资源(例如,为某个序列分配一个超出实际所需大小的内存块)。这导致已分配空间的一部分未能被有效利用。
  • 外部碎片 (External Fragmentation):由于不同序列的长度不一致,在内存中释放和分配操作后,会出现大量不连续的小块空闲空间。尽管总的空闲内存可能足够大,但由于不连续性,无法满足大块内存分配的需求,这可能导致在明明还有显存的情况下出现**OOM(Out Of Memory)**错误,极大降低GPU利用率。

【若需查看原始图片详情,请参考原文中的“vLLM-pre.pdf”引用的统计图,其中展示了优化后真正用于保存KV Cache内容的部分从20%提升至96.3%】

3.1.3 PagedAttention 原理

PagedAttention 借鉴了操作系统中分页内存管理(Paging Memory Management)的思想来高效管理KV Cache。

3.1.3.1 KV Block 概念
  • 定义:一个KV Block 是一段固定大小、地址连续的内存块,用于存放特定数量 token 的 Key 和 Value 状态。
  • 类比:它在概念上类似于操作系统中的“内存页”(Memory Page)。通过将KV Cache划分为固定大小的块,可以更灵活、更细粒度地进行内存管理。
3.1.3.2 逻辑/物理Block映射
  • 逻辑Block:每个序列的KV Cache在逻辑上是连续的,由一系列逻辑Block组成。
  • 物理Block:这些逻辑Block被映射到物理内存中的非连续的物理Block上。
  • Block Table:系统维护一个Block Table(块表),记录了每个逻辑Block对应的物理Block地址,从而实现了逻辑地址到物理地址的映射查询。这种方式允许在物理内存中动态分配和重用KV Block,而无需保证物理连续性。
3.1.3.3 多请求存储与内部碎片优化
  • 多请求同时存储:PagedAttention允许来自不同请求的KV Cache Block在物理内存中交错存储,极大地提高了显存的利用率。
  • 内部碎片优化:通过使用固定大小的KV Block,每个序列浪费的 token 数被限制在一个Block的大小之内。例如,如果块大小是16或32个token,无论序列长度如何(0到1000个token),每个序列的最大浪费都远小于整个序列的长度,从而有效降低了内部碎片。
3.1.3.4 Copy-on-Write (写时复制) 机制
  • 应用场景:在Beam Search等场景下,模型会基于当前的KV Cache生成多条候选路径(形成树形结构的 token 展开)。传统的做法是为每条路径复制整个KV Cache,效率低下。
  • 机制:Copy-on-Write 机制允许不同的Beam共享相同的KV Cache Block,直到某个Beam需要修改该Block的内容(例如,当其生成一个新的 token 导致KV状态改变时)。只有在发生修改时,系统才会为该修改创建 Block 的副本,否则保持共享。这极大地减少了内存复制和占用,提高了Beam Search的效率。

3.2 投机采样 (Speculative Decoding, ICML’23)

投机采样(Speculative Decoding)是一种旨在加速LLM解码阶段(序列生成)过程的技术,因为这一阶段的序列生成是影响生成回答的重要瓶颈。

3.2.1 原理与加速效果

  • 核心思路:通过结合一个“轻量级模型预测”和“重模型验证”的方式实现并行化,以提高生成速度。
  • 轻量级模型预判:首先使用一个较小、速度更快的语言模型(通常称为“辅助模型”或“草稿模型”,Draft Model)快速生成多个候选输出 token 序列。
  • 大型模型验证:接着,将这些由轻量级模型生成的候选序列一次性提供给更强大但较慢的大型模型(例如 GPT-3/4,称为“验证模型”,Verifier Model)进行并行验证。
  • 加速效果
    • 如果轻量级模型预测正确,大模型直接接受这些 token,从而一次性“跳过”多个生成步骤。
    • 如果预测不正确,大模型会从第一个错误点重新生成并修正。
    • 通过这种方式,大模型不再需要从头逐个生成所有内容,而是大部分时间进行快速验证,显著减少了整体生成时间,降低了推理延迟。

3.2.2 投机采样算法流程

投机采样算法涉及两个模型:一个辅助小模型和一个大模型。

  1. 小模型生成序列:辅助小模型(Draft Model)根据当前已生成的 token 序列,快速生成一个包含kk个未来 token 的候选序列。
  2. 大模型并行验证:大模型(Verifier Model)在一个前向传播中并行检查小模型生成的这kk个 token 是否正确。它计算出这些 token 的真实概率分布。
  3. 修正与迭代
    • 大模型从候选序列的最早一个被验证为不正确的 token 开始进行修正。
    • 一旦修正完毕,辅助小模型会从大模型修正后的位置继续开始生成新的候选序列,重复上述过程。

3.2.3 概率分布一致性证明

投机采样的重要特性在于,它能够在显著加速生成速度的同时,保持与大模型单独生成时相同的概率分布。这意味着投机采样生成的文本在质量和多样性上与大模型直接生成的文本是等效的。

原论文证明了该方法生成的序列概率分布与大模型本身生成序列的概率相同。其证明核心基于接收-拒绝采样(Accept-Reject Sampling)机制:

  • 为了从目标概率分布 P(x)P(x) 中采样,我们不直接采样,而是从一个提议分布 Q(x)Q(x) 中采样。
  • 接受条件:如果从 Q(x)Q(x) 采样的样本 xx 满足 Q(x)P(x)Q(x) \le P(x),则接受该样本。
  • 拒绝条件与调整:如果 Q(x)>P(x)Q(x) > P(x),则以概率 1P(x)Q(x)1 - \frac{P(x)}{Q(x)} 拒绝该样本。被拒绝的样本不会直接丢弃,而是从一个调整后的分布 P(x)P'(x) 中重新采样: P(x)=norm(max(0,P(x)Q(x)))P'(x) = norm(\max(0, P(x) - Q(x))) 其中 normnorm 表示归一化,确保调整后的概率总和为1。通过这种严谨的概率调整和重采样机制,投机采样保证了最终输出序列的概率分布与大模型本身的概率分布一致。

总结:投机采样通过“并行预测+串行验证”的策略,显著降低了LLM的推理延迟,同时确保了生成质量与大模型原生生成的一致性。

总结

PagedAttention 通过操作系统分页内存管理的概念,解决了KV Cache带来的内存碎片和高内存占用问题,提升了多用户并发下的GPU利用率。投机采样则通过轻量级模型预判和大模型验证相结合的方式,在不牺牲生成质量的前提下,显著加速了LLM的解码过程。这两种技术都是当前LLM推理服务中至关重要的优化手段。

4. 推理系统 (Serving System) 优化

4.1 推理系统概述

推理系统(Serving System)负责将经过训练的机器学习模型部署到生产环境中,并为终端用户提供模型推理服务。其主要功能和目标包括:

4.1.1 模型部署与推理API

  • 模型部署:将训练好的机器学习模型打包,部署到服务器或云平台上。这包括模型的序列化、加载、环境配置等步骤,目标是确保模型能够在生产环境中稳定、高效地运行。
  • 推理API:提供一个接口,接收用户或应用程序的输入数据(请求数据),调用部署好的模型进行推理计算,并返回推理结果。推理API需要确保服务具备低延迟(快速响应用户请求)、高吞吐(在单位时间内处理大量请求)和可扩展性(能够根据负载变化动态调整资源)。

4.1.2 分布式部署

对于大规模模型(特别是大语言模型LLM),单台设备往往难以满足其计算或内存需求,或者难以应对高并发的推理请求。在这种情况下,通常采用分布式部署的方式:

  • 负载均衡:利用负载均衡技术将用户的推理请求分发到不同的模型服务实例上,避免单个实例过载。
  • 提高整体性能:通过多个实例并行处理请求,提高整体的吞吐能力可用性

4.2 生成式推理中的并行策略

在生成式推理中,由于模型通常较大且生成过程是自回归的,对并行策略的设计尤为重要。

  • Naïve Generative Inference (朴素生成式推理)

    • 通常是逐个token序列推理。这意味着每生成一个token,模型都需要运行一次前向传播,效率较低。
    • 在之前的讨论中(如KV Cache),我们已经看到了这种逐个token生成带来的计算冗余。
  • 重新思考训练阶段的并行方式

    • 训练阶段常用的并行策略,如数据并行(DP)、模型并行(TP)和管道并行(PP),也可以应用于推理阶段,但需要根据推理的特点进行调整和优化。
    • 数据并行 (Data Parallelism, DP)
      • 在推理中,DP主要用于增加运行实例,从而同时服务多个用户请求。每个实例加载完整的模型,处理一部分用户的请求。
    • 张量并行 (Tensor Parallelism, TP)
      • TP将模型的单个层(或张量)拆分到多个设备上,用于加速单个推理请求的计算。特别适用于单个大模型无法完全载入单个GPU内存的情况,或需要降低单个token生成延迟的场景。
    • 管道并行 (Pipeline Parallelism, PP)
      • PP将模型的不同层分配到不同的设备上,形成一个处理“管道”。一个请求的数据会依次流经这些设备。
      • 潜在优势:尽管管道并行可能会增加单个请求的处理延迟(因为数据需要在设备间传输,产生通讯开销),但它能够很好地实现GPU复用。通过将模型的不同部分分摊到多个GPU上,可以提高整体设备的利用率,让所有设备尽可能同时工作。

4.3 AlpaServe:模型并行与统计复用 (OSDI’23)

  • 论文名称:Statistical Multiplexing with Model Parallelism for Deep Learning Serving. Zhuohan Li, et al. OSDI 2023.
  • 核心观察
    • 在深度学习推理服务中,即使一个模型可以完全放入一个设备(GPU)中,如果能够将其拆分进行管道并行(PP),并结合统计复用(Statistical Multiplexing),也可能带来性能提升。
    • 传统观念认为PP会增加延迟,但在多用户、动态负载的服务场景下,通过精巧的调度,PP可以显著提高整体资源的利用率。
  • 设计原则
    • 最大化GPU复用:目标是让集群中的所有GPU尽可能同时工作。
    • 统计复用:通过将多个并发请求的计算在不同管道阶段交错执行,来隐藏设备间的通信延迟和计算空闲时间。即使单个请求的延迟可能略有增加,系统的总吞吐量和资源利用率可以大幅提高。
  • 目的让所有设备同时工作,从而提升整个推理服务的吞吐量和效率。

4.4 Orca:连续批处理 (Continuous Batching) (OSDI’22)

  • 核心问题:为了充分利用计算资源,推理过程中通常会将多个用户请求组织成一个批次(batch)进行模型推理。然而,传统批处理方式存在以下问题:

    • 批次组织和完成:经典方法通常是等待请求队列中的多个请求积累到一定数量后,将它们组织成一个批次,并一起完成推理。
    • GPU利用率低:由于不同请求的回答长度往往不一致(例如,一个请求只需要生成几个token,另一个需要生成几十个token),在一个批次中,完成较早的请求会因为等待批次中其他请求完成而导致其所占用的GPU资源被空闲,从而造成GPU利用率低下
  • Orca的解决方案:Continuous Batching (连续批处理)

    • Batch一直维护:Orca采用一种动态的批处理机制,批次不再是固定大小或固定生命周期的。只要批次中存在空闲位置,新的请求就可以即时加入
    • Token-level batching (Token级批处理):而非传统的请求级批处理。这意味着Orca不是等待整个请求(直到生成所有token)完成后才将其从批次中移除,而是在每个生成步(每次生成一个token)后,动态地管理批次中的token。当一个请求生成完它的当前token后,它可以继续留在批次中生成下一个token,或者如果它已经完成,则可以离开批次。这种细粒度的管理使得GPU可以更有效地被利用。
    • 效果:通过连续批处理,Orca显著提高了GPU的利用率,减少了因请求长度不一导致的资源浪费。
  • 实现挑战

    • “What behind the elegant idea may be … labor work” (一个优雅的想法背后可能是大量的工程工作)。
    • 为了实现continuous batching,特别是在注意力(attention)机制的计算中,难以保持简单高效的批处理矩阵运算。传统的Attention计算通常假定批次中的所有序列长度一致,但在动态的Token级批处理中,序列长度是变化的,这需要复杂的内存管理和调度机制来高效地执行矩阵运算。

4.5 DistServe:解耦推理 (Disaggregated Inference) (OSDI’24)

  • 核心观察:生成式大语言模型(LLM)的推理过程通常包含两个截然不同的阶段:Prefill(预填充)Decoding(解码)。这两个阶段具有不同的计算特性和资源瓶颈。

  • 两个任务的特点不同

    • Prefill 阶段
      • 模型将提示(prompt)序列作为输入,并计算出所有prompt token的KV Cache。
      • 这一阶段是计算能力限制(Computation bounded)。因为它涉及到对整个输入prompt序列的一次性处理,计算量大。
      • 批处理带来的收益很小:由于prompt长度可能各异,且每个prompt需要完全计算一次,批处理带来的效率提升不如解码阶段显著。
    • Decoding 阶段
      • 模型在每一步解码过程中,基于前一个生成的token和已有的KV Cache生成下一个token。
      • 这一阶段是内存带宽限制(Memory bandwidth bounded)。每次生成一个token时,需要频繁地从内存中读取和写入KV Cache,内存访问成为主要瓶颈。
      • 批处理能显著提升吞吐量:在解码阶段,将多个请求(每个请求生成一个token)打包成一个批次,可以更高效地利用GPU的计算单元,并减少内存访问的开销,从而大幅提高吞吐量。
  • 两个任务会相互影响

    • 当一个新的请求的Prefill阶段(计算密集型)被插入到正在进行Decoding的批处理中时,可能会导致**“卡顿”现象**。例如,在使用流式推理API时,用户可能会感觉到某个输出token突然停顿,这是因为Prefill任务抢占了资源,影响了Decoding任务的流畅性。
  • DistServe的解决方案:将Prefill和Decoding任务分离

    • 专用资源:DistServe建议在不同的GPU集群上分别处理Prefill和Decoding任务。
      • 某些GPU专门用于Prefill计算
      • 其他一些GPU专门用于Decoding计算
    • 不同的并行配置:由于Prefill和Decoding阶段的特点不同,它们可以采用不同的并行配置。例如,Prefill集群可能更侧重于计算能力,而Decoding集群可能更侧重于高内存带宽和批处理效率。
    • 动态资源比率:Prefill和Decoding服务器实例的集群资源比率可以进行动态配置,以适应不同的工作负载模式。
    • 主要开销:这种分离式架构的主要开销在于KV Cache的通信。由于Prefill阶段生成的KV Cache需要传输到Decoding集群进行后续使用,这会引入额外的数据传输延迟。

4.6 LoRA推理优化

LoRA (Low-Rank Adaptation) 是一种高效的微调技术,它通过引入少量可训练参数来适应大型预训练模型,而无需修改原始模型的全部参数。LoRA在服务系统中变得极其重要,原因如下:

  • 大规模LoRA模型部署需求:例如,Hugging Face 提供了超过 9,000 个版本的微调 BERT 模型。这意味着在实际生产环境中,可能需要同时部署和服务成千上万个基于LoRA微调的同一个基础模型。

4.6.1 朴素LoRA推理服务的问题

若按照传统方式,将每个LoRA模型独立部署到GPU上,会面临显著的效率问题:

  • GPU利用率低:如果一张或多张GPU仅仅服务一个LoRA模型,那么它们将未能充分利用基础LLM的共享权重。因为LoRA只添加了少量参数,大部分计算仍然在共享的基础模型上进行。
    • 【若需查看原始图片详情,请参考原文中的“LoRA推理服务示意图”】,该图展示了每个GPU独立服务一个LoRA模型(LoRA 1, LoRA 2, LoRA 3分别对应GPU 0, GPU 1, GPU 2),导致资源利用不佳。
  • 批处理潜力未充分利用:LLM具有很强的批处理效应,即同时处理多个请求可以显著提高吞吐量。
    • 错误做法示例:如果仅仅把使用同一个LoRA的请求放入一个批次,会导致批次大小波动(例如:b=2, b=1, b=3)。当不同LoRA的请求混合到来时,如果严格按照LoRA模型来区分批次,将无法形成大的批次。
    • 正确做法示例:理想情况是能够将所有请求(无论属于哪个LoRA)都放入一个大批次(例如:b=6),这样才能充分发挥LLM的批处理优势。

4.6.2 Punica & S-LoRA:批处理不同LoRA请求

  • 核心思想:由于不同的LoRA推理都使用相同的基础模型部分(这部分通常占模型参数的99%以上),因此存在将不同LoRA的请求批处理的可能性。
  • 朴素做法的局限:仅仅将属于同一LoRA模型的输入连续地归为一组进行批处理,无法充分利用LLM的批处理能力,因为它忽略了不同LoRA请求之间共享的基础模型计算。
  • Punica (Proceedings of Machine Learning and Systems, 2024) & S-LoRA (arXiv preprint, 2023)
    • 这些工作提出了创新的方法来批处理来自不同LoRA适配器的请求
    • 原理:LoRA的计算可以表示为 Y=XW+X(BTAT)=XW+X(ALBLT)Y = XW + X(B^T A^T) = XW + X(A_L B_L^T)。其中 XWXW 是基础模型的计算,X(ALBLT)X(A_L B_L^T) 是LoRA适配器的增量计算。
      • 这两个优化方案的核心思想是:对于基础模型计算 XWXW,所有不同LoRA的请求可以共享一个大的批次,因为它们都使用相同的权重 WW
      • 对于LoRA的增量计算 X(ALBLT)X(A_L B_L^T),这些计算是独立于各个LoRA的,但可以在批处理中高效地执行。
      • 这意味着,可以通过精巧的调度和内存管理,将来自多个不同LoRA适配器的请求在同一个批次中一起处理,共享基础模型的前向传播计算,然后只在最后的LoRA层进行区分和合并,从而大幅提高GPU利用率和吞吐量。

【若需查看原始图片详情,请参考原文中的“Punica & S-LoRA示意图”和“LLM具有强大的批处理效应图示”】

5. LoRA推理优化

5.1 LoRA在服务系统中的重要性

LoRA(Low-Rank Adaptation)是一种高效的参数微调技术,它在**大型语言模型(LLM)**的服务系统中扮演着至关重要的角色。其重要性体现在:

  • 模型定制化需求激增:随着LLM的广泛应用,用户对模型特定任务或领域知识的定制化需求爆炸式增长。例如,Hugging Face平台就提供了超过9,000个不同版本的微调BERT模型。
  • 高效微调:LoRA允许在不修改LLM大部分参数的情况下,通过引入少量可训练的低秩矩阵来适应新任务,这大大降低了微调成本和存储需求。
  • 服务复杂性:尽管LoRA微调本身很高效,但在生产环境中同时服务成百上千个定制化的LoRA模型却带来了新的挑战,尤其是如何高效地利用计算资源。

5.2 朴素LoRA推理服务的问题

传统的或“朴素”的LoRA推理服务方法,通常将每个LoRA模型作为一个独立的实体进行部署和管理。这种做法存在显著的效率问题:

  • GPU利用率低
    • 如果为每个LoRA适配器分配一张或多张GPU(如图所示:LoRA 1 在 GPU 0,LoRA 2 在 GPU 1,LoRA 3 在 GPU 2),会导致未能充分利用基座LLM的共享权重
    • 由于LoRA适配器只占模型总参数的一小部分(例如,LoRA参数可能只占整个模型参数的0.01%到1%),而基座模型参数是所有LoRA共用的,为每个LoRA实例单独分配GPU会使得GPU大部分计算能力闲置。这最终导致GPU利用率低整体性能效率下降
  • 批处理潜力未充分利用
    • LLM推理具有显著的批处理效应(Batching Effect),即将多个请求打包成一个批次并行处理可以显著提高吞吐量。
    • 然而,朴素做法仅仅是将使用同一个LoRA的请求放入一个批次
    • 例如,如果请求队列中有A、A、B、C、C、C这六个请求,朴素方法可能会分为三个批次:一个批次包含两个A,一个批次包含一个B,一个批次包含三个C。
    • 这种做法不能充分利用批处理的潜力,因为它忽略了不同LoRA请求之间共享基座模型的计算。理想情况下,如果能够批处理所有请求,即使它们对应不同的LoRA,也可以形成一个更大的批次(如示例中的包含六个请求的批次),从而更好地发挥批处理的优势。

5.3 Punica & S-LoRA:批处理不同LoRA请求

针对朴素LoRA推理服务的问题,Punica和S-LoRA等系统提出了创新的解决方案,核心思想是批处理来自不同LoRA适配器的推理请求

  • 核心思想与原理
    • 参数共享:不同LoRA的推理请求都要使用相同的基础模型部分。基座模型(如GPT-3、LLaMA等)的参数占据了模型总参数的绝大部分(通常高达99%)。
    • 计算分解:一个经过LoRA微调的模型 WW' 的输出 YY 可以表示为: Y=WX=(W+ΔW)X=WX+ΔWXY = W'X = (W + \Delta W)X = WX + \Delta W X 其中,WW 是基座模型的原始权重,ΔW\Delta W 是LoRA引入的低秩更新矩阵,XX 是输入。 关键在于WXWX 部分的计算对于所有基于同一基座模型的LoRA请求是通用且共享的。ΔWX\Delta W X 部分是LoRA特有的,但其计算量相对较小。
    • 批处理可能性:正因为 WXWX 部分的共享性,存在将不同LoRA的请求批处理的可能性。这意味着可以在一个大的批次中,首先计算所有请求共享的基座模型输出 WXWX,然后再为每个请求独立地计算并叠加其特定的LoRA更新 ΔWX\Delta W X
  • 实现方式
    • Punica和S-LoRA通过巧妙的系统设计,实现了对不同LoRA请求的动态批处理(Multi-tenant LoRA Serving)。
    • 它们不再将每个LoRA视为一个独立的模型实例,而是将基座模型权重多个LoRA适配器权重在内存中高效管理,并允许在一个批次中同时处理来自不同LoRA的请求。
    • 在推理过程中,模型可以执行一个大的批处理操作来处理所有请求的基座模型计算,然后根据每个请求所属的LoRA,在内存中快速切换并应用相应的低秩更新矩阵。这种方式极大地提高了GPU的利用率和推理吞吐量。