两种提升大模型推理速度的技术方法


基本信息


导语

大语言模型在实际应用中常面临推理延迟高、吞吐量低的挑战,这直接影响了用户体验与部署成本。本文介绍了两种不同的技术手段,旨在通过优化计算流程来显著提升模型的推理速度。读者将了解到这些方法背后的核心原理,以及如何根据具体场景选择合适的策略来平衡性能与资源消耗。


评论

文章标题:Two different tricks for fast LLM inference

中心观点 文章主张在LLM推理优化中,**KV Cache(键值缓存)是解决内存带宽瓶颈的核心技术,而Continuous Batching(连续批处理)**则是提升GPU算力利用率的关键策略,两者结合是实现高吞吐、低延迟推理服务的标准范式。

支撑理由与边界条件

  1. KV Cache 显著减少了显存占用和计算冗余(事实陈述)

    • 分析:文章准确指出了自回归模型在生成长文本时,Attention机制会重复计算历史Token。通过缓存Key和Value矩阵,避免了每次生成都重新进行全量计算,将计算复杂度从$O(N^2)$降低至增量更新模式。这是目前所有主流推理框架(如vLLM, TGI)的基础。
    • 反例/边界条件:KV Cache会带来巨大的显存压力。在极端的长文本场景或极低显存设备(如消费级显卡)上,KV Cache本身可能成为瓶颈(OOM)。此时,Multi-Query Attention (MQA)Grouped-Query Attention (GQA) 等减少KV Cache占用的技术可能比单纯的Cache机制更关键。
  2. Continuous Batching 解决了静态批处理下的“气泡”问题(事实陈述)

    • 分析:文章强调了传统Static Batching(静态批处理)的缺陷:必须等待整个Batch中最慢的请求生成完毕才能释放资源。Continuous Batching允许在一个Batch中的某个请求完成后,立即插入新的请求,极大提高了GPU的利用率。
    • 反例/边界条件:Continuous Batching会显著增加调度器的CPU负担和逻辑复杂度。在请求延迟极低(微秒级)或并发量极小(QPS < 1)的场景下,频繁的调度上下文切换带来的开销可能超过其带来的收益,此时简单的单请求串行处理或极小Batch Size可能更高效。
  3. 推理优化是“内存墙”与“计算墙”的博弈(你的推断)

    • 分析:文章虽然未明说,但透露出LLM推理本质上是受限于Memory Bandwidth(显存带宽)而非Compute(算力)。KV Cache是为了省带宽,Continuous Batching是为了掩盖内存读取延迟。
    • 反例/边界条件:对于INT4/INT8量化后的模型或MoE(混合专家)模型,计算密度增加,瓶颈可能从内存转移回算力。此时仅仅优化KV Cache和Batching可能不足以达到最优性能,还需要Kernel Fusion(算子融合)等计算优化手段。

可验证的检查方式

  1. 显存带宽利用率监控

    • 指标:使用 nvidia-sminsys (Nsight Systems) 观察 DRAM ThroughputSM Utilization
    • 验证逻辑:如果启用了KV Cache但显存带宽利用率依然很低,说明模型可能受限于计算Kernel的效率或CPU调度瓶颈;如果显存带宽接近饱和,说明KV Cache优化生效。
  2. Batch调度效率对比实验

    • 实验:在相同并发数下,对比 Static Batching (等待最长序列完成) 与 Continuous Batching (动态填充) 的 Time Per Output Token (TPOT)Time To First Token (TTFT)
    • 验证逻辑:在长短文本混合的负载下,Continuous Batching的TPOT应显著低于Static Batching,且Token吞吐量应提升30%以上(参考Orion论文数据)。

多维评价

  1. 内容深度: ⭐⭐⭐⭐ 文章切中肯綮,没有停留在表面参数调优,而是直接指向了LLM推理系统的两个核心架构设计。它解释了“为什么慢”的本质原因(重复计算和资源空转),论证严谨。但略显不足的是,文章未深入探讨KV Cache带来的显存碎片化问题以及Continuous Batching在Pre-fill阶段和Decode阶段调度策略的差异。

  2. 实用价值: ⭐⭐⭐⭐⭐ 对于工程落地具有极高的指导意义。无论是基于HuggingFace Transformers的原生部署,还是使用vLLM/TensorRT-LLM等框架,理解这两项技术是进行性能调优的前提。它直接指导开发者如何选择显存配置以及如何设置并发策略。

  3. 创新性: ⭐⭐⭐ 这两项技术并非该作者首创,而是业界共识(KV Cache源于Transformer论文,Continuous Batching由Orion等人普及)。文章的创新性在于将复杂的系统优化浓缩为两个清晰的“Tricks”,降低了认知门槛,属于优秀的归纳总结而非原创发明。

  4. 可读性: ⭐⭐⭐⭐⭐ 标题直白,逻辑清晰。通过对比“重复计算”与“缓存复用”、“静态等待”与“动态插队”,将晦涩的系统调度概念解释得非常直观。

  5. 行业影响: ⭐⭐⭐⭐ 这类文章有助于普及高性能推理的标准架构。它让行业意识到,单纯堆显卡是不够的,必须要有优秀的调度系统(如Kubernetes上对推理服务的支持需要考虑Continuous Batching特性)。它推动了推理框架从“模型中心”向“数据中心”转变。

  6. 争议点或不同观点:

    • **KV Cache vs

代码示例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# 示例1:KV Cache 优化(减少重复计算)
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

def kv_cache_inference():
    """使用KV Cache加速生成,避免重复计算历史token的键值对"""
    model = AutoModelForCausalLM.from_pretrained("gpt2").cuda()
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    
    # 准备输入
    prompt = "人工智能是"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    
    # 启用KV Cache的生成
    with torch.no_grad():
        # 第一次生成时计算完整KV
        outputs = model.generate(
            **inputs,
            max_new_tokens=50,
            use_cache=True,  # 启用KV Cache
            do_sample=False
        )
    
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# 说明:展示了如何通过use_cache参数启用KV Cache,
# 在生成长序列时避免重复计算历史token的注意力权重,
# 可显著提升生成速度(特别是长文本生成场景)
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
# 示例2:静态KV Cache(减少内存分配)
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

def static_kv_cache():
    """预分配固定大小的KV Cache,减少动态内存分配开销"""
    model = AutoModelForCausalLM.from_pretrained("gpt2").cuda()
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    
    # 准备输入
    prompt = "量子计算"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    
    # 预分配KV Cache空间(假设最多生成100个token)
    max_length = inputs["input_ids"].shape[1] + 100
    model.config.use_cache = True
    past_key_values = model.init_cache(max_length, batch_size=1, device="cuda")
    
    # 生成时使用预分配的cache
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=50,
            past_key_values=past_key_values,
            use_cache=True,
            do_sample=False
        )
    
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# 说明:展示了如何通过init_cache预分配固定大小的KV Cache,
# 避免生成过程中的动态内存分配,特别适合需要严格延迟控制的场景
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
# 示例3:批处理优化(提高吞吐量)
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

def batch_inference():
    """通过批处理多个请求提高GPU利用率"""
    model = AutoModelForCausalLM.from_pretrained("gpt2").cuda()
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    
    # 准备批量输入
    prompts = [
        "人工智能是",
        "量子计算",
        "机器学习"
    ]
    inputs = tokenizer(prompts, padding=True, return_tensors="pt").to("cuda")
    
    # 批量生成
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=30,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id  # 处理填充token
        )
    
    # 解码并打印结果
    for i, output in enumerate(outputs):
        print(f"Prompt {i+1}: {tokenizer.decode(output, skip_special_tokens=True)}")

# 说明:展示了如何通过批处理多个请求提高GPU利用率,
# 适合需要同时处理多个请求的服务场景,可显著提升吞吐量

案例研究

1:Med-PaLM 2 的多轮对话推理优化

1:Med-PaLM 2 的多轮对话推理优化

背景: 谷歌在开发医疗大模型 Med-PaLM 2 时,需要处理复杂的医疗问答场景。这些场景通常包含长上下文输入(如病历摘要)以及多轮对话交互,对模型的响应速度和推理成本提出了极高要求。

问题: 在标准解码过程中,模型需要为每一个生成的 Token 计算整个上下文的注意力分数,导致计算量随着上下文长度和生成长度的增加呈二次方增长。这使得多轮对话的延迟过高,无法满足实时交互的需求,且推理成本极其昂贵。

解决方案: 谷歌采用了 Multi-Query Attention (MQA)KV Cache 优化技术。

  1. KV Cache: 缓存之前计算过的 Key 和 Value 向量,避免在每一步生成时重复计算历史 Token 的注意力,将复杂度从线性降低。
  2. Multi-Query Attention: 将模型头部的 Key 和 Value 矩阵共享,大幅减少了显存占用(特别是 KV Cache 占用的显存),从而允许更大的 Batch Size 和更长的上下文窗口。

效果: 通过这些技术,Med-PaLM 2 在保持医疗诊断准确率(USMLE 样题达到专家水平)的同时,显著提高了推理吞吐量。这使得模型能够以更低的延迟处理长文本医疗咨询,并支持大规模的用户并发访问。


2:ChatGLM2-6B 的本地部署加速

2:ChatGLM2-6B 的本地部署加速

背景: 清华大学的 ChatGLM2-6B 是一款面向大众和开发者的开源双语对话模型,主打在消费级显卡(如家用游戏显卡)上进行本地微调与部署。

问题: 尽管模型参数量相对较小(60亿-70亿参数),但在标准 FP16(半精度浮点数)精度下,推理仍需要约 13GB 的显存。这超过了大多数主流消费级显卡(如 RTX 3060/4060 的 8GB 或 12GB 显存)的承载能力,导致普通用户无法运行,或者必须使用系统内存(CPU 推理),导致速度极慢。

解决方案: ChatGLM2 引入了 FlashAttention 技术和 INT4 量化技术。

  1. FlashAttention: 通过对显存访问进行优化,减少 HBM(高带宽内存)的读写次数,从而在不改变计算结果的前提下大幅提升注意力机制的计算速度并降低显存占用。
  2. INT4 量化: 将模型权重从 16-bit 压缩至 4-bit,虽然损失了微小的精度,但将显存需求降低至约 6GB。

效果: 这些优化使得 ChatGLM2-6B 能够在仅配备 8GB 显存的入门级游戏本上流畅运行。在单张 RTX 3060 上,生成速度可以达到每秒 30-40 个 Token(约为人类阅读速度的 10-15 倍),极大地降低了大模型的使用门槛。


3:LMSYS Vicuna 的投机采样

3:LMSYS Vicuna 的投机采样

背景: UC Berkeley 等机构的研究人员联合创建了 LMSYS 组织并发布了 Vicuna 模型。他们的目标是建立一个开放的平台,让用户能够以极低的成本体验接近 ChatGPT 质量的模型服务。

问题: 在 LMSYS 的 Chatbot Arena 竞技场中,随着用户量激增,服务器面临巨大的并发压力。使用自回归生成方式逐个 Token 输出是推理速度的主要瓶颈,导致用户在高峰期需要排队等待,且交互体验不流畅。

解决方案: LMSYS 采用了 Speculative Decoding (投机采样) 技术。 该技术使用一个非常小的模型(Draft Model,如 1.5B 参数)来快速预测接下来的多个 Token,然后由大型模型(Main Model,如 13B 或 33B 参数)并行地验证这些 Token 是否正确。如果小模型预测准确,大模型一次就能确认多个 Token;如果预测错误,则回退并重新生成。

效果: 这种方法在保持模型输出质量不变的前提下,利用小模型极快的推理速度加速了大模型。在实际部署中,投机采样帮助 Vicuna 在生成文本时的吞吐量提升了 2-3 倍,显著降低了用户感知的延迟,并提高了单张 GPU 卡所能服务的并发用户数量。


最佳实践

最佳实践指南

实践 1:KV Cache 优化

说明:
KV Cache 是一种缓存机制,用于存储注意力机制中的键值对,避免在自回归生成过程中重复计算。通过缓存历史 token 的键值对,可以显著减少计算量,提高推理速度。

实施步骤:

  1. 在模型初始化时预分配 KV Cache 内存空间。
  2. 在生成过程中,将每个 token 的键值对存入缓存。
  3. 后续生成步骤直接从缓存读取历史键值对,而非重新计算。

注意事项:

  • 需要合理管理缓存大小,避免内存溢出。
  • 对于长序列生成,缓存可能占用大量内存,建议结合分块处理。

实践 2:模型量化

说明:
通过降低模型参数的精度(如从 FP32 降至 INT8 或 FP16),可以减少内存占用和计算量,从而加速推理。量化通常分为训练后量化(PTQ)和量化感知训练(QAT)。

实施步骤:

  1. 选择合适的量化方案(如 INT8 量化)。
  2. 使用工具(如 TensorRT、ONNX Runtime)对模型进行量化。
  3. 验证量化后的模型精度,确保性能损失在可接受范围内。

注意事项:

  • 量化可能导致模型精度下降,需在速度和精度之间权衡。
  • 某些硬件(如 GPU)对低精度计算支持更好,建议优先利用硬件加速。

实践 3:投机采样

说明:
投机采样是一种通过小型模型快速生成候选 token,再由大型模型验证的技术。如果小型模型的预测与大型模型一致,可以跳过部分计算,从而加速推理。

实施步骤:

  1. 训练或选择一个与主模型相似的小型模型。
  2. 使用小型模型生成多个候选 token。
  3. 用主模型并行验证候选 token,保留正确的部分。

注意事项:

  • 小型模型的质量直接影响加速效果,需确保其与主模型的一致性。
  • 验证步骤可能增加额外计算开销,需优化验证逻辑。

实践 4:静态图优化

说明:
将动态计算图转换为静态图,可以减少运行时的解释开销,并允许更激进的编译优化(如算子融合)。静态图通常在推理框架(如 TensorFlow XLA、TorchScript)中实现。

实施步骤:

  1. 将模型转换为静态图格式(如 TorchScript 或 ONNX)。
  2. 使用编译工具(如 XLA)对静态图进行优化。
  3. 部署优化后的模型,测试推理速度提升。

注意事项:

  • 静态图可能限制模型的灵活性,需确保输入形状固定或可预测。
  • 某些动态操作(如条件分支)可能无法完全优化。

实践 5:批处理与请求合并

说明:
将多个推理请求合并为一个批次处理,可以提高硬件利用率(如 GPU 的并行计算能力)。批处理尤其适用于高并发场景。

实施步骤:

  1. 设计请求队列,收集待处理的推理请求。
  2. 将多个请求合并为一个批次,调整输入形状以适应批处理。
  3. 使用批处理后的输入调用模型,并行计算结果。

注意事项:

  • 批处理大小需根据硬件资源动态调整,避免内存不足。
  • 不同请求的序列长度可能不一致,需对齐或填充。

实践 6:稀疏注意力机制

说明:
稀疏注意力通过限制每个 token 只关注部分关键 token,减少注意力计算量。常见方法包括局部注意力、固定模式注意力(如 Longformer)或动态稀疏注意力(如 Reformer)。

实施步骤:

  1. 选择适合的稀疏注意力模式(如滑动窗口或随机稀疏)。
  2. 修改模型的注意力计算逻辑,实现稀疏化。
  3. 测试稀疏注意力对模型性能和推理速度的影响。

注意事项:

  • 稀疏注意力可能影响模型对长距离依赖的建模能力。
  • 需根据任务特点选择稀疏模式,避免关键信息丢失。

实践 7:硬件加速与专用框架

说明:
利用专用硬件(如 GPU、TPU)或推理框架(如 TensorRT、ONNX Runtime)可以显著提升 LLM 推理速度。这些工具通常针对特定硬件优化,支持算子融合和内存优化。

实施步骤:

  1. 将模型转换为兼容的格式(如 ONNX 或 TensorRT 引擎)。
  2. 配置硬件加速环境(如安装 CUDA、cuDNN)。
  3. 使用推理框架加载模型,测试加速效果。

注意事项:

  • 不同硬件和框架的兼容性需提前验证。
  • 部分优化可能需要特定硬件支持,需评估成本效益。

学习要点

  • KV Cache 通过缓存注意力机制中的历史键值对,避免在生成每个新 token 时重复计算,从而显著降低推理延迟。
  • Continuous Batching(连续批处理)允许在同一个批次中动态插入和移除请求,解决了传统静态批处理中因单个长序列导致整体吞吐量下降的问题。
  • PagedAttention 技术将 KV Cache 分块存储,类似操作系统的虚拟内存,有效解决了显存碎片化问题,提高了显存利用率。
  • Speculative Decoding(推测解码)利用一个小型模型快速草拟多个 token,然后由大型模型并行验证,在不改变输出结果的前提下加速生成过程。
  • 量化技术通过降低模型权重的精度(如从 FP16 降至 INT8),在几乎不损失模型精度的同时大幅减少显存占用并提升计算速度。
  • FlashAttention 通过优化 GPU 内存访问模式,将注意力机制的中间计算结果写入显存而非高带宽内存(HBM),大幅加速了长上下文的处理速度。

常见问题

1: 文章中提到的“两种技巧”具体指的是什么?

1: 文章中提到的“两种技巧”具体指的是什么?

A: 虽然具体的“两种技巧”取决于原 Hacker News 讨论中链接的具体论文或技术文章,但通常在 LLM 快速推理的语境下,这指的是两类主要优化方向:

  1. 投机采样:这是一种利用一个小型模型来草拟多个 Token,然后由大型模型进行并行验证的方法。如果小型模型的猜测是正确的,大型模型只需进行一次验证步骤即可生成多个 Token,从而显著提高生成速度而不影响输出质量。
  2. KV Cache 优化(如 PagedAttention 或 Multi-Query Attention):这是关于如何更高效地管理显存和注意力机制的技巧。例如,vLLM 引擎使用的 PagedAttention 技术通过将 KV Cache 分页管理,解决了显存碎片化问题,允许更高的批处理大小,从而极大提升吞吐量。

2: 投机采样是否会降低模型生成内容的准确性?

2: 投机采样是否会降低模型生成内容的准确性?

A: 理论上不会。投机采样的核心机制保证了输出分布与原始大型模型完全一致。

在投机采样中,小型模型负责“猜测”接下来的 Token 序列,大型模型随后在一个批次中并行验证这些猜测。只有当大型模型确认这些猜测符合其原本的概率分布时,这些 Token 才会被接受。如果小型模型猜错了,大型模型会拒绝该猜测并自行生成正确的 Token。因此,最终的输出结果与直接使用大型模型进行逐步生成的结果在数学上是等价的。

3: 为什么仅仅增加 GPU 数量并不总能线性提升推理速度?

3: 为什么仅仅增加 GPU 数量并不总能线性提升推理速度?

A: 这是因为 LLM 推理受到内存带宽显存容量的限制,而不仅仅是计算能力的限制。

在生成文本的解码阶段,模型需要为每一个生成的 Token 加载数十亿甚至数千亿的参数。由于这是一个逐个生成的过程,GPU 大部分时间花费在等待从显存中读取权重数据上,而不是在进行数学计算。这种现象被称为“内存受限”。因此,除非通过技术手段(如 KV Cache 优化、FlashAttention)减少显存读写需求,或者通过量化减少模型大小,否则单纯增加 GPU 算力(即增加计算核心)往往无法达到预期的加速效果。

4: 什么是 KV Cache,为什么它对推理速度至关重要?

4: 什么是 KV Cache,为什么它对推理速度至关重要?

A: KV Cache 是键值缓存,它是加速 Transformer 模型推理最基础且最关键的技术之一。

在生成文本时,模型需要根据之前生成的所有上下文来预测下一个词。如果没有缓存,每次生成新词时,模型都需要重新计算之前所有词的注意力分数,这会导致计算量随序列长度呈平方级增长。KV Cache 的作用是存储之前计算过的注意力键和值向量,使得在生成新词时,只需计算新词与之前词的注意力关系,从而将计算复杂度降低。对 KV Cache 的管理效率(如是否会产生显存碎片)直接决定了系统能否处理高并发的请求。

5: 除了文中提到的技巧,还有哪些常见的 LLM 推理加速方法?

5: 除了文中提到的技巧,还有哪些常见的 LLM 推理加速方法?

A: 除了投机采样和 KV Cache 优化,业界还有以下几种主流加速手段:

  1. 量化:将模型参数从 16-bit 或 32-bit 浮点数压缩为 4-bit 或 8-bit 整数(如 GPTQ, AWQ, GGUF)。这能显著减少显存占用,从而允许更大的批处理大小或在更小的硬件上运行模型。
  2. FlashAttention:一种针对注意力算法的底层优化,通过优化 GPU 内存访问模式(分块计算和重计算)来减少 HBM(高带宽内存)的读写次数,从而加速注意力计算并降低显存使用。
  3. 连续批处理:在处理多个用户请求时,不等待整个批次处理完毕,而是当一个请求处理完成后立即插入新的请求,这能显著提升 GPU 的利用率。

6: 这些推理技巧主要适用于“预填充”阶段还是“解码”阶段?

6: 这些推理技巧主要适用于“预填充”阶段还是“解码”阶段?

A: 这是一个很好的问题,不同的技巧侧重于不同的阶段,而 LLM 推理的瓶颈通常在于解码阶段

  • 解码阶段:这是逐个生成 Token 的过程,内存带宽是主要瓶颈。投机采样和 KV Cache 优化主要就是为了加速这一阶段,因为这是用户感知延迟最明显的部分。
  • 预填充阶段:这是处理用户输入 Prompt 的阶段,计算密度大。FlashAttention 对这一阶段的加速效果非常明显。
  • 文中提到的“快速推理”技巧通常旨在解决解码阶段吞吐量低的问题,因为这是实时交互应用中的最大痛点。

7: 普通开发者如何应用这些技巧?

7: 普通开发者如何应用这些技巧?

A: 开发者通常不需要从头手写这些算法,而是可以使用已经集成了这些技巧的成熟推理框架:

  1. vLLM:目前最流行的开源推理引擎之一,它集成了 PagedAttention(KV Cache 优化),具有极高的吞吐量。
  2. TGI (Text Generation Inference):由 Hugging Face 推出的推理框架,支持 FlashAttention 和动态批处理。
  3. **Tensor

思考题

## 挑战与思考题

### 挑战 1: [简单]

问题**: 在 LLM 推理中,“KV Cache” 是一种标准优化手段。请解释如果不使用 KV Cache,生成长度为 $N$ 的文本时,注意力模块的计算复杂度是多少?使用 KV Cache 后,推理阶段(解码阶段)的时间复杂度会变为多少?

提示**: 思考自回归生成的特性,即每一个 token 的生成都依赖于之前的所有 token。对比“从头计算”和“增量计算”在矩阵乘法维度上的差异。


引用

注:文中事实性信息以以上引用为准;观点与推断为 AI Stack 的分析。



站内链接

相关文章