人工智能系统性能工程 (Chinese Edition)

第18章 高级预填充 -解码与键值缓存调优

本作品已使用人工智能进行翻译。欢迎您提供反馈和意见：[e-mail]

本章在第17章基础上，深入探讨推理预填充与解码阶段的高级优化策略。我们将基于高阶扩展策略，涵盖低级技术手段，包括单解码“巨型内核”、智能KV缓存调优与GPU间共享、快速GPU间 prompt 状态传输、自适应资源调度，以及预填充与解码工作进程间的动态路由。

同时重点介绍硬件与软件创新成果，这些突破将带来性能与效率的新高度。通过应用这些技术，您可显著降低解码延迟，提升单GPU吞吐量，并在大规模部署中满足严格的延迟服务水平目标（SLO）。

优化解码内核

迄今为止，我们主要关注 的高层系统与集群优化策略。在提升超大规模推理性能时，另一组值得关注的技术是底层内核与内存管理调优——尤其针对解码阶段。

解码阶段采用分布式处理且常受内存有界限制，这促使研究者和实践者致力于实现解码阶段的极致加速并针对特定硬件进行优化。该领域两大突破性成果是FlashMLA（DeepSeek）、ThunderMLA（斯坦福大学）和FlexDecoding（PyTorch），它们专门针对LLM工作负载中常见的变长序列场景，优化了变压器模型解码阶段的多头注意力效率。接下来我们将逐一探讨这些方案。

FlashMLA（DeepSeek）

闪存多潜伏注意力（FlashMLA） 是DeepSeek推出的优化解码内核。其核心聚焦于单令牌解码步骤——本质上是生成下一个令牌的Transformer层前向传播过程。通过融合运算与优化GPU内存层次结构，FlashMLA显著提升了解码速度。

FlashMLA（解码）之于推理，正如FlashAttention（预填充）之于训练。它显著降低内存访问开销与延迟。相较标准内核，FlashMLA可为解码阶段带来大幅延迟缩减。

FlashMLA通过将多个注意力操作融合为单次运算提升算术密集度。这种方式可在单次融合内核调用中处理多个注意力头和多个时间步，即使在小批量处理时也能保持数学运算单元持续工作，从而提高解码过程中的GPU利用率。图18-1展示了在Hopper H100 GPU上，相较于分组查询注意力（GQA）和多查询注意力（MQA）等其他注意力实现方案，MLA在算术强度方面的提升。（注：Blackwell架构通过更高TFLOPs和HBM带宽使两条性能曲线均上移。）

图18-1. MLA接近 计算有界状态（基于NVIDIA Hopper H100架构测量）

FlashMLA的引入意义重大，它证明即使在次优GPU硬件上，也能有效降低解码阶段的瓶颈——包括内存带宽和内核启动开销。该方案通过减少独立GPU内核启动次数并优化内存访问模式，在受限硬件上为解码任务榨取了最大性能。

DeepSeek开源的FlashMLA实现方案已广泛应用。SGLang和vLLM均提供对DeepSeek模型的原生支持。因此，您应评估FlashMLA方案，在无需修改高层架构的前提下提升单令牌解码吞吐量。

由于DeepSeek开源的FlashMLA已集成至现代推理服务系统，您应将其作为提升每个解码工作进程吞吐量（或降低单令牌延迟）的解决方案，且无需进行任何高层架构变更。

ThunderMLA（斯坦福大学）

基于FlashMLA，斯坦福大学的 ...