第三部分:单机篇 —— 极致优化显存的高性能引擎¶
在第二部分中,我们剖析了 LLM 推理的物理与数学瓶颈:KV Cache 引发的显存海啸,以及 Prefill 与 Decode 之间的核心不对称性。这些瓶颈制约了大模型在生产环境中的并发能力和响应速度。
为了解决这些问题,系统工程师和算法科学家们提出了多种软硬件协同的优化方案。这一部分将深入探讨现代推理引擎(如 vLLM 和 SGLang)以及模型架构本身是如何解决上述瓶颈的。我们将从两个核心战场展开:
- 突破显存瓶颈:通过 GQA(模型架构)、PagedAttention(分页管理)和 RadixAttention(前缀缓存),将显存利用率提升数倍。
- 优化不对称性:通过连续批处理(Continuous Batching)和分块预填充(Chunked Prefill),消灭 Padding 浪费,让 GPU 的算力与带宽时刻保持饱和。
[!IMPORTANT] 需要注意的是,连续批处理和分块预填充虽然极大地提升了单机效率,但它们属于“战术级”的局部优化,并未从根本上解决 Prefill 和 Decode 的硬件错配与资源竞争。我们将在第四部分的集群篇中,解析更深入的“战略级”演进——分离式推理(Disaggregated Serving)。
下面,让我们首先探讨第一个优化方向——从模型架构层面减少显存占用。
第九章:模型架构层面的显存优化:GQA¶
第一节:从 MHA 到 GQA 的演进¶
在 Transformer 的早期设计中,采用的是 MHA(Multi-Head Attention,多头注意力)。
- MHA:如果有 \(H\) 个 Query 头,就必须对应有 \(H\) 个 Key 头和 \(H\) 个 Value 头。正如我们在第六章第三节所讨论的,KV Cache 的空间复杂度为 \(O(L \cdot N \cdot d)\) 。在 MHA 中,虽然隐藏层维度 \(d\) 被切分为了 \(H\) 个头(即 \(d = H \times d_k\) ),但因为 KV 头数与 Query 头数完全相等,所以总的 KV Cache 大小依然由完整的维度 \(d\) 决定。对于千亿参数的模型,为了保证表达能力, \(d\) 通常设定得非常大(例如 Llama 3 405B 的 \(d = 16384\) ),而在长上下文场景下,序列长度 \(N\) 也会变得非常大,这两者的叠加直接导致了 KV Cache 极其庞大。
为了缩减 KV Cache,研究人员提出了 MQA(Multi-Query Attention):
- MQA:所有的 Query 头共享同一组 Key 和 Value 头。这样 KV Cache 的大小直接缩减到了原来的 \(1/H\) 。显存压力大幅下降,但代价是模型的表达能力有所下降。
GQA(Grouped Query Attention,分组查询注意力) 则是两者的有效平衡,目前已被广泛应用于各类开源和闭源大模型中(例如 Llama 3 405B 也采用了该方案):
- GQA:将 Query 头进行分组(例如 8 个一组),每一组共享一组 Key 和 Value 头。
- 收益:既大幅削减了 KV Cache 的显存占用, 提升了系统的 Throughput(能容纳更多并发) ,同时因为减少了 Decode 阶段的数据搬运量,也间接提升了单请求的 TPS,且几乎没有损失模型的性能。
实战对比:基于 Llama 3 405B 的 KV Cache 计算
为了让大家对这三种机制的显存瘦身效果有更直观的认识,我们再次以 Llama 3 405B 的参数规格为例(层数 \(L=126\) ,隐藏层维度 \(d=16384\) ,Query 头数 \(H=128\) ,每个头维度 \(d_k=128\) ,FP16 格式),计算在 100 万 Token 上下文时,不同机制下的 KV Cache 大小:
- 标准 MHA 模式(假设每个 Query 头都有独立的 KV 头):
- 每个 Token 每层大小: \(2 \times 128 \times 128 \times 2 \text{ 字节} = 64 \text{ KB}\)
- 总大小: \(64 \text{ KB} \times 126 \text{ 层} \times 1,000,000 \approx \mathbf{8.06 \text{ TB}}\)
- MQA 模式(所有 Query 头共享 1 组 KV 头):
- 每个 Token 每层大小: \(2 \times 1 \times 128 \times 2 \text{ 字节} = 0.5 \text{ KB}\)
- 总大小: \(0.5 \text{ KB} \times 126 \text{ 层} \times 1,000,000 \approx \mathbf{63 \text{ GB}}\)
- GQA 模式(实际 Llama 3 采用,8 组 KV 头):
- 每个 Token 每层大小: \(2 \times 8 \times 128 \times 2 \text{ 字节} = 4 \text{ KB}\)
- 总大小: \(4 \text{ KB} \times 126 \text{ 层} \times 1,000,000 \approx \mathbf{504 \text{ GB}}\)
从 \(8 \text{ TB}\) 降到 \(504 \text{ GB}\) ,GQA 实现了 16 倍的显存压缩,直接让单卡或小规模集群处理长上下文成为可能。
第二节:深层思考 —— 为什么只砍 \(K\) 和 \(V\) 行得通?¶
这是一个非常深刻的问题:为什么只削减 \(K\) 和 \(V\) 的头数,而不削减 \(Q\) 的头数,模型的表达能力仍能保持?
要理解这背后的含义,我们需要从 \(Q\) 、 \(K\) 、 \(V\) 的角色分工、信息冗余 以及 大模型的知识存储 三个层面来剖析。
1. 角色不对称:提问者( \(Q\) )与被查者( \(K\) 、 \(V\) ) 在 Attention 机制中, \(Q\) 、 \(K\) 、 \(V\) 扮演着完全不同的角色:
- Query( \(Q\) ) 代表 “意图”或“问题”(我现在要找什么?)。它是动态的,随着模型生成每一个新词,意图都在变。
- Key( \(K\) ) 代表 “索引”或“标签”(我这里有什么?)。
- Value( \(V\) ) 代表 “内容”或“实质”(我这里真正的信息是什么?)。
这背后的物理含义是:同一个客观事实( \(K\) 和 \(V\) ),可以回答很多个不同的问题( \(Q\) )。
[!NOTE] 比喻:图书馆里的研究员 想象一个场景:图书馆里有 128 个研究员(代表 128 个 \(Q\) 头),他们每个人都在研究不同的课题,提出不同的问题。
- 在普通的 MHA 中:系统非常奢侈。为了服务这 128 个研究员,图书馆不仅配备了 128 个研究员,还为每个人配了一套专属的“百科全书”(128 个 \(K\) 和 \(V\) 头)。虽然这些书的本质都是记录相同的历史事实,内容有极大的重叠,但它们被编成了 128 个略有不同的“版本”。每个研究员只翻看自己桌上的那一套,这显然造成了极大的空间浪费。
- 在 GQA 中:系统进行了优化。图书馆里依然有 128 个研究员,但现在只买了 8 套百科全书(8 个 \(K\) 和 \(V\) 头)。每 16 个研究员共享一套书。
虽然 16 个研究员问的问题千奇百怪( \(Q\) 不同),但他们要查找的 历史背景和客观事实( \(K\) 和 \(V\) )是完全相同的。一套书足够回答他们所有人的问题。这就是为什么 \(Q\) 头不能少(保证意图的多样性),而 \(K\) 、 \(V\) 头可以少(共享知识库)。
2. MHA 中存在严重的“信息冗余” 在标准的 MHA 中,研究人员通过可视化分析发现:很多不同的 Key 头和 Value 头学到的特征是高度重复的。 比如,可能有 5 个头都在关注“句子的主语是谁”,另外 4 个头都在关注“代词指代的是谁”。GQA 的本质是减少参数冗余:既然这些头提取的是类似的信息,它们就可以共享同一组 Key 和 Value。这种共享迫使模型在训练时更加高效地利用参数,把冗余的信息压缩,从而在不损失表达能力的情况下,大幅削减了显存占用。
3. 大模型的“硬知识”根本不在 KV 里 根本原因在于:大模型的绝大多数知识,其实都存在 FFN(前馈网络)中,而不是 Attention 里。 我们在前面的章节算过账:在 Transformer 的每一层中,FFN 的参数量占了大约 82%,而 Attention 只占了不到 20%。
- FFN 像是一个巨大的“知识硬盘”,存储了海量的客观规律和事实。
- Attention 则像是一个“调度器”和“路由器”,负责在上下文之间搬运信息、建立关联。 既然 Attention 只是负责搬运和关联上下文的,那么我们把它的 KV 缓存压缩一下,并不会破坏模型本身在 FFN 里存储的千亿级知识储备。模型依然知道“法国的首都是巴黎”,它只是在推理时,用更少的“内存指针”(KV)去指向这个知识而已。
第三节:百花齐放:压缩 KV 的其他前沿进展¶
除了 GQA 这种在主流模型中被广泛采用的方案外,为了进一步对抗“显存海啸”,学术界和工业界最近涌现出了许多令人兴奋的新进展。虽然我们不在此深挖细节,但了解这些方向对于把握大模型的技术走向至关重要:
- MLA (Multi-head Latent Attention):
- 原理:由 DeepSeek 团队在 DeepSeek-V2 中提出。它通过低秩联合压缩(Low-Rank Joint Compression)技术,将 Key 和 Value 投影到一个低维的潜空间(Latent Space)中。在推理时,只需要缓存这个极小的潜向量,从而实现比 GQA 还要显著的 KV Cache 压缩比(最高可达数倍)。
- 参考:DeepSeek-V2 Paper
- 滑动窗口注意力 (Sliding Window Attention):
- 原理:模型在计算注意力时,不再关注“从头到尾”的所有历史 Token,而是只维护一个固定大小的滑动窗口(例如只看最近的 4096 个 Token)。超出窗口的 KV Cache 会被直接丢弃,从而将 KV Cache 的空间复杂度从 \(O(N)\) 降到了 \(O(1)\) 。
- 参考:Mistral 7B Paper
- 交错局部/全局注意力 (Interleaved Local/Global Attention):
- 原理:结合了滑动窗口和全局注意力的优点。在模型的某些层使用滑动窗口注意力以节省显存,而在另一些层保留全局注意力以捕捉长距离依赖(例如 Gemma 2 和 Mistral 的部分模型采用了类似策略)。
- 参考:可以参考相关模型的官方技术报告。
- 压缩内存注意力 (Infini-Attention):
- 原理:由 Google 提出,通过在标准的点积注意力中引入“压缩内存”机制,将旧的 KV 状态存储在固定大小的内存中。它结合了局部掩码注意力和长期线性注意力,使得模型能够理论上处理无限长的上下文,而不会导致 KV Cache 爆炸。
- 参考:Leave No Context Behind Paper
这些进展向我们揭示了一个清晰的趋势:通过模型的架构优化,进一步减少 K 和 V 的大小,依然是当前大模型优化的核心战场之一。
GQA 和上述的架构优化都是模型本身的改进,需要模型在训练时就采用这种架构。然而,除了从架构层面砍掉 KV 头数,我们还可以在不改变架构的情况下,从精度层面对 KV Cache 进行压缩。
第十章:精度降维:KV Cache 量化(FP8/INT8)¶
如果说 GQA 是在空间结构上做到了极致(减少数据量),那么 KV Cache 量化则是在数据密度上进行了极致压缩。
第一节:以计算换带宽的权衡策略¶
在标准的推理中,K 和 V 都是以 16 位(FP16 或 BF16)的精度存储的,每个元素占 2 字节。而 KV Cache 量化 的核心思想,就是将它们压缩为 8 位(FP8 或 INT8),使每个元素仅占 1 字节,从而将显存占用直接减半。
这确实引入了额外的计算开销:
- 写入时:当 Decode 阶段生成一个新的 Token 时,它的 K 和 V 向量必须先经过缩放(Scaling)和截断,从 16 位转化为 8 位,然后才能存入显存。
- 读取时:在计算 Attention 时,GPU 从显存中读取这 8 位的 K 和 V,必须先将它们“反量化”回 16 位(或者直接在支持低精度的硬件上进行计算)。
这本质上是用额外的计算换取存储空间的压缩。但这在 Decode 阶段是一种非常有效的权衡。
第二节:为什么这很划算?¶
我们在第八章深入讨论过,Decode 阶段是典型的内存带宽密集型(Memory-Bound)。GPU 的计算核心绝大多数时间都处于闲置状态,等待数据从显存传输。
- 不进行量化:数据体积大,传输慢,GPU 核心处于闲置状态。
- 进行量化:虽然多了几步量化转换的计算,但需要搬运的数据量直接减半。显存带宽的压力减轻了一半,数据能更快地供应给 GPU 核心。
在 NVIDIA H100 等现代 GPU 上,硬件已经原生支持 FP8 的张量计算,这种量化转换的计算开销几乎可以忽略不计。因此,用极小的计算代价,换取显存占用减半、搬运速度翻倍,成为了现代高性能推理引擎的标配。更重要的是,将 KV Cache 压缩到 8-bit(FP8/INT8)带来的精度损失通常不到 0.5%,几乎无损,却能让系统承载的并发用户量直接翻倍。
第三节:INT8 与 FP8:截然不同的量化范式¶
在 KV Cache 8位量化的实践中,存在两种截然不同的技术路线:
- INT8(8位整数):这是一种“浮点到整数”的映射。它通过缩放(Scale)和取整(Round),把连续的浮点数映射到 256 个离散的整数区间中。在计算时,必须先反量化回 FP16。这种方式在数学上是明显有损的,但通过细粒度量化(按 Token 或 Channel 独立计算缩放系数),可以把对模型精度的影响降到极低。
- FP8(8位浮点数):这依然是浮点数,保留了符号位、指数位和尾数位。它与 FP16 的区别仅仅是“变短了”,精度和表示范围降低了。但它依然代表“同一个数”,更符合直觉。在 NVIDIA H100 等支持 FP8 张量核心的硬件上,FP8 的计算速度极快,且不需要显式的反量化过程,正在成为服务端的主流。
第四节:动态与静态:与模型权重量化的对比¶
除了量化 KV Cache,模型本身也可以量化(如 GPTQ、AWQ),且应用更广。但两者在效果和难度上有着本质的区别:
- 模型权重量化(静态):模型的参数在训练好后是固定不变的。我们可以有充足的时间离线分析参数分布,甚至用少量数据进行校准,因此哪怕压到 4-bit,大模型依然能保持极高的性能。
- KV Cache 量化(动态):它是根据用户输入动态生成的。你无法提前预知会出现什么极端值(Outliers)。如果把 KV Cache 压到 4-bit,极端值会破坏数据的动态范围,导致模型输出质量严重下降。因此 KV Cache 通常只敢用 8-bit。
在实际部署中,两者往往组合使用:
- W4A16(仅量化权重):主攻“显存容量”,让大模型能塞进小显卡。
- W8A8 / FP8(全量化):主攻“推理速度”,是云端大厂高并发服务的首选。
虽然通过 GQA 和量化技术,KV Cache 的体积被大幅压缩,但随着上下文的增长,它在显存中依然会引发严重的碎片化问题。这就是为什么我们仍然需要 PagedAttention。
第十一章:引擎层面的显存管理:PagedAttention¶
第一节:碎片化危机:静态分配的弊端¶
在大模型生成文本时,由于无法预知用户最终生成的序列长度(可能仅有几个 Token,也可能达到模型的最大上下文长度),传统的显存管理方式采用了静态连续分配策略。系统必须根据模型的最大上下文长度(例如 8000 Token),为每个请求预先分配一块足够大的连续显存空间。
这导致了严重的内存浪费:
- 内部碎片(Internal Fragmentation)与预留浪费:系统必须按最大上下文长度为每个请求预先分配一整块连续显存。在请求处理期间,这块显存被完全独占。这意味着,无论请求最终是长是短,那些尚未使用的空间(预留给未来 Token)以及请求提前结束而永远用不上的空间,在请求被最终释放前都无法被其他请求复用。这种独占机制造成了严重的显存闲置。
- 外部碎片(External Fragmentation):即使显存中仍有足够的总空闲空间,但如果这些空间在物理上是不连续的,系统也无法将其分配给需要大块连续空间的新请求。
根据 vLLM 论文的统计,在传统静态连续分配方式下,由于碎片化问题,真正用来存储有效 KV Cache 的显存往往不到 20%,多达 80% 的内存未被有效利用。
第二节:OS 的灵感:虚拟内存分页¶
面对这种显存闲置问题,加州大学伯克利分校的研究人员(vLLM 的作者们)分析指出:这与早期计算机内存不足时遇到的问题类似。
在计算机操作系统中,为了解决内存碎片化问题,早已发明了虚拟内存分页机制(Paging)。操作系统将物理内存切分成固定大小的“页(Pages)”,程序在逻辑上看到的是连续的内存,但在物理上可以分散在内存的任何角落。
vLLM 的核心思想就是将操作系统的虚拟内存分页机制(Paging)移植到 GPU 显存管理上:
- 块化管理(Blocks):不再为单个请求分配巨型连续显存,而是将物理显存划分为固定大小的物理块(Physical Blocks)。例如,每个块固定存放 16 个 Token 的 K 和 V 矩阵。
- 非连续物理分配:在逻辑上,一个请求的 Token 序列是连续的;但在物理显存中,这些 Token 对应的块可以离散地分布在显存的任何位置,无需物理连续。
第三节:块表(Block Tables):近乎零的内存浪费¶
既然物理位置是打散的,模型在计算注意力时,怎么找得到前面所有的词呢?
为了在非连续的物理空间中高效进行注意力计算,PagedAttention 引入了块表(Block Tables)。块表负责维护逻辑块与物理块之间的映射关系,类似于操作系统中的页表(Page Table)。在计算 \(Q \cdot K^T\) 时,注意力机制通过查询块表,动态地从离散的物理块中读取对应的 K 和 V 向量,完成计算。
解析(直观理解映射机制):
为了更直观地理解 BlockTable 是如何维护 Request -> Token Block -> Memory Block 的映射的,我们来看它的张量维度设计:
block_table 在底层是一个二维张量,形状为 [max_num_reqs, max_num_blocks_per_req]。
- 第一维(行):对应不同的 Request(请求),用
req_idx索引。 - 第二维(列):对应某个请求的第几个 Token Block(逻辑块),用
logical_block_idx索引。 - 存储的值:就是对应的 Memory Block(物理显存块 ID)。
也就是说,查找公式为:physical_block_id = block_table[req_idx][logical_block_idx]。
PagedAttention 带来的显著收益:
显存浪费几乎降为零:因为是按需分配(填满一个 16 座的小包厢,才开下一个),显存的浪费被严格限制在了最后一个没填满的 Block 里(最多浪费 15 个 Token 的位置)。整体显存的碎片浪费率从 80% 大幅下降到 4% 以下。
然而,PagedAttention 主要解决的是显存碎片化和单请求内的浪费问题。在默认情况下,不同 Request 之间的 KV Cache 仍然是严格分开存储的。如果两个独立的请求在不同时间段到达,即使它们包含了完全相同的前缀(例如相同的背景文档),PagedAttention 也无法自动识别并重用之前算好的物理块。
这种更高级的、跨请求的动态“缓存重用”,正是我们下一章要讨论的核心。
第十二章:基于基数树的前缀缓存机制 (RadixAttention)¶
在第八章中,我们提到了多用户共享系统提示词的场景。然而,在实际应用中,尤其是多轮对话和 RAG(检索增强生成)场景下,情况要复杂得多。本章将介绍如何利用 基数树(Radix Tree) 实现更高级的缓存复用——前缀缓存(Prefix Caching)。
第一节:RAG 与多轮对话的困境¶
在大模型的实际应用中,我们输入的 Prompt 往往包含大量的背景资料。
场景一:多轮对话
- 第一轮:你问“什么是苹果?”(模型计算并生成回答)。
- 第二轮:你接着问“它好吃吗?”。此时输入给模型的实际 Prompt 是:[第一轮你的问题 + 第一轮AI的回答 + 第二轮你的问题]。
这里涉及到一个非常基础但容易被忽略的工程细节:HTTP 协议是无状态的。 这意味着,从大模型服务引擎(如 vLLM)的角度来看,第二轮对话是一个完全独立的、全新的 Request,会被系统分配一个全新的 Request ID。
在传统的 PagedAttention 中,显存中的块表(Block Table)是与 Request ID 强绑定的。虽然第二轮请求的 Prompt 里包含了第一轮的内容,但引擎仅通过 Request ID 进行识别。由于 ID 不同,引擎无法自动识别并复用第一轮已经算好的 KV Cache。
因此,在没有前缀缓存的时代,大模型只能重复进行相同的计算,把第一轮已经算过的内容重新计算一遍 QKV!随着对话轮数的增加,Prompt 越来越长,首字延迟(TTFT)就会快速增加。这种跨请求的显存孤岛,直接推动了前缀缓存(Prefix Caching)技术的发展。
场景二:RAG(知识库问答) 你上传了一本 10 万字的手册,然后不断向它提问。每次提问,这 10 万字都会作为前置背景。如果没有优化,每次提问都要重新处理这 10 万字, 首字延迟(TTFT) 会显著增加,影响用户体验。
场景三:固定系统提示词与 Few-Shot 示例 在企业级应用或 Agent 中,通常会在每次请求前塞入一段冗长且固定的 System Prompt 或 Few-Shot 示例。如果没有优化,哪怕有 1000 个用户并发访问,系统也会为这 1000 个独立的请求重复计算 1000 遍相同的 KV Cache,造成算力和显存的严重浪费。
场景四:并行采样与 Beam Search 在代码生成(要求模型输出多个候选方案)或使用 Beam Search(束搜索)时,系统需要为同一个 Prompt 生成多个不同的后续。如果没有优化,系统需要为每个分支复制并重复计算 Prompt 的 KV Cache。
第二节:基数树(Radix Trees):共享物理内存¶
为了解决这个问题,SGLang 等框架(以及后来的 vLLM)引入了 RadixAttention,利用 基数树(Radix Tree) 数据结构来管理 KV Cache。
在 Transformer 中,由于位置编码和自注意力的特性,只要前面的词汇序列完全一样,它们计算出的 KV Cache 是完全相同的。
基数树的运作方式如下:
- 根节点:空序列。
- 边:代表一串连续的 Token 序列(例如 16 个 Token 的 Block)。
- 节点:指向这些 Token 对应的物理块。
当一个新请求进来时,系统从根节点出发,拿请求的 Token 序列去和树的边做最长前缀匹配。 如果匹配成功,直接复用该节点指向的物理块,跳过这些 Token 的矩阵乘法计算。对于后面不匹配的部分(比如用户的新问题),再分配新的物理块并在树上长出新的叶子节点。
通过这种方式,多轮对话的历史记录、RAG 的公共文档,都可以被直接复用,首字弹出的速度从几秒钟大幅下降到几十毫秒。
[!NOTE] 深入细节:Token 如何索引与引擎差异
- Token 是如何索引的?:基数树上的索引绝非文本比较。大模型在处理文本时,早已将文本转化为了 Token ID(整数)。树的边存储的就是这些整数序列。在匹配时,系统进行的是高效的整数序列比对,或者对 Token 序列进行 Hash 计算(如 vLLM 依靠哈希值来快速锚定 Block)。
- vLLM 与 SGLang 的实现差异:虽然两者都利用了基数树来实现前缀缓存,但它们的颗粒度不同。SGLang 的 RadixAttention 是 Token 级别 的,匹配非常灵活(边可以代表任意长度的 Token 序列);而 vLLM 的 APC(Automatic Prefix Caching)则继承了 PagedAttention 的基因,是 Block 级别 的(通常以 16 个 Token 为一个固定大小的块进行管理和哈希匹配)。
- 基数树与 Block Table 的关系:基数树并没有替换 Block Table,它们最终都指向相同的物理块(Physical Blocks),只是索引的维度不同:
- Block Table:是由 CPU 端的调度器(Block Manager) 维护的
Request -> 逻辑块 -> 物理块的映射。在推理执行时,它会以张量形式平铺在 GPU 端,供 Attention Kernel 查找物理显存块。- 基数树 是基于
去重的 Token 序列前缀 -> 物理块的索引。它是全局的,驻留在 CPU 端用于跨请求的缓存共享和 LRU 淘汰。- 协同工作流:CPU 调度器利用基数树找到可复用的物理块,再加上新分配的块,拼装成某个请求的 Block Table 传给 GPU。GPU 的查找逻辑完全不需要改变。
为了更直观地展示它们的关系,我们可以用下面这张图来表示:
graph TD subgraph CPU ["CPU 管理面"] RadixTree["🌲 基数树 (Radix Tree)<br>索引: Token 序列 ➔ 物理块 ID"] BlockManager["⚙️ Block Manager<br>(维护映射与分配)"] end subgraph GPU_Mem ["GPU 显存 (数据面)"] BlockTableTensor["📋 Block Table Tensor<br>(供 GPU 执行查找)"] subgraph KV_Cache ["物理块池 (Physical Blocks)"] B10["📦 物理块 10<br>缓存: 'System Prompt...'"] B11["📦 物理块 11<br>缓存: 'User Question...'"] end end RadixTree -->|映射到| B10 RequestA["👤 请求 A (命中前缀)"] -->|1. 查树| RadixTree RadixTree -->|2. 返回匹配块| BlockManager BlockManager -->|3. 组装| BT_A["📋 请求 A 的 Block Table: [10, 11]"] BT_A -->|4. 传给 GPU| BlockTableTensor BlockTableTensor -->|5. 指向| B10 BlockTableTensor -->|5. 指向| B11 GPU_Kernel["⚡ GPU Attention Kernel"] -->|6. 读取| BlockTableTensor
第十三章:动态调度机制:连续批处理与分块预填充¶
在第七章中,我们看到了“静态批处理”的缺陷:为了迁就长句子,短句子被迫填充大量的 Padding,浪费了算力和显存。本章将介绍现代推理引擎是如何通过连续批处理和分块预填充有效解决这个问题的。
第一节:连续批处理(Continuous Batching):旋转门机制¶
为了打破静态批处理中“短请求必须等待长请求完成”的木桶效应,Orca 论文提出并由 vLLM 等引擎广泛应用的 连续批处理(Continuous Batching,也叫 In-flight Batching) 应运而生。
比喻:旋转门与永不停歇的高铁 想象一台永远在运行的高铁,每一站(每一次模型前向传播,耗时几十毫秒)都有人上车,有人下车。
- 动态进出:系统不再等待一整个 Batch 的请求全部生成完毕。在每一次生成 Token 的间隙,调度器都会检查:哪个请求遇到了结束符(EOS)?立刻将其移出 Batch(下车);队列里有没有新请求在排队?立刻把它塞进 Batch(上车)。
- 消灭填充:在底层,vLLM 借助了 FlashAttention 等高级算子,将不同请求的 Token 展平成一个一维的连续数据流送给 GPU。通过传入每个请求的“边界路标”(cu_seqlens 数组),GPU 能够物理隔离不同请求的计算,消灭了 Padding。
图解:矩阵 vs. 向量
1. 静态批处理 (Static Batching) — 平成“矩阵”
| 列 1 | 列 2 | 列 3 | 列 4 | |
|---|---|---|---|---|
| 行 1 (🔵) | 🔵 T1 | 🔵 T2 | 🔵 T3 | ❌ PAD |
| 行 2 (🟢) | 🟢 T1 | 🟢 T2 | ❌ PAD | ❌ PAD |
| 行 3 (🟡) | 🟡 T1 | 🟡 T2 | 🟡 T3 | 🟡 T4 |
2. 连续批处理 (Continuous Batching) — 接成“向量”
| 列 1 | 列 2 | 列 3 | 列 4 | 列 5 | 列 6 | 列 7 | 列 8 | 列 9 | |
|---|---|---|---|---|---|---|---|---|---|
| 行 1 | 🔵 T1 | 🔵 T2 | 🔵 T3 | 🟢 T1 | 🟢 T2 | 🟡 T1 | 🟡 T2 | 🟡 T3 | 🟡 T4 |
[!NOTE]
cu_seqlens = [0, 3, 5, 9](通过偏移量物理隔离不同请求的计算)
这种 Iteration 级别的调度,让 GPU 算力时刻保持饱满, Throughput(吞吐量) 相比静态批处理提升了数倍,同时保证了每个请求的 TBT(吐字间隔) 相对稳定,避免了短请求被长请求阻塞的问题。
连续批处理的底层运转流程与三大数据结构¶
为了让“永不停歇的高铁”能够高效运转,并在 GPU 内部精准识别不同请求的 Token,推理引擎在底层依赖三个关键的“账本”(数据结构)。这解释了 GPU 到底是如何在没有复杂指针和动态查找的情况下,高效准确地定位数据的:
cu_seqlens(累积序列长度数组):- 作用:负责隔离 Batch 中不同请求的 Token 边界。
-
原理:在连续批处理中,所有请求的输入 Token 被展平成一个一维的连续流送给 GPU。
cu_seqlens记录了每个请求在该数据流中的起止边界(例如[0, 3, 4]表示前 3 个属于请求 A,第 4 个属于请求 B)。这在 Batch 中同时包含 Prefill(多 Token)和 Decode(单 Token)的混合场景下尤为关键,Attention Kernel 看到它,就知道在计算时绝不能“跨界”去读邻居请求的数据。 -
Block Table(块表): - 作用:负责历史 KV 追溯(不仅用于 Decode 阶段寻找历史,在 Chunked Prefill 读取前半部分长 Prompt 以及前缀缓存读取共享前缀时,GPU 同样需要通过它来查找历史 KV)。
-
原理:这是一个平铺的二维数组,行对应 Batch 中的请求 index,列对应物理块 ID。当 GPU 拿到 Batch 中某个请求的当前新 Token 时,它不会去按指针遍历历史,而是直接用该请求在 Batch 中的 index(比如第 3 个)去查
BlockTable[3],直接拿到该请求所有的物理块 ID 列表,然后依次从显存中读取历史 KV 数据。 -
slot_mapping(槽位映射): - 作用:负责所有阶段的精准写入。
- 原理:这是一个长度等于当前 Batch 总 Token 数的一维数组。它由 CPU 调度器提前算好,直接告诉 GPU 当前 Batch 里的每一个新 Token 算完 KV 后,应该写入到物理显存的哪一个绝对槽位(Slot)中。GPU 只需要执行高效的
kv_cache[slot_mapping[i]] = new_kv即可,完全避免了在 GPU 内部做复杂的物理地址计算。
这种“CPU 准备映射关系,GPU 通过张量索引直接访问显存地址”的设计,是实现高并发、低延迟的关键机制。
特殊情况:队头阻塞与长 Prompt 上述的连续批处理机制在处理 Decode 请求(每次只吐 1 个 Token,访存密集)时非常完美。但是,当队列中突然塞进一个包含几万字的长 Prompt 请求时,如果系统在一次迭代中处理完整段 Prompt,会占用 GPU 较长的时间,导致队列中其他正在 Decode 的请求被迫中断。这种“队头阻塞”问题,直接催生了我们下一节要讨论的——分期付款式的分块预填充。
第二节:分块预填充(Chunked Prefill):完美互补¶
虽然连续批处理解决了 Decode 阶段的卡顿,但在 Prefill 阶段(处理用户输入的 Prompt)又遇到了新的问题:队头阻塞。
假设来了一个包含 1000 个 Token 的长 Prompt 请求。如果系统在一次迭代中处理完整段 Prompt,会占用 GPU 较长的时间,导致其他正在 Decode 的老请求被迫“卡住”不吐字,引发首字延迟的剧烈波动。
优化方案:分块处理 为了解决这个问题,业界引入了 Chunked Prefill(分块预填充):
- 系统会设定一个单次迭代的最大 Token 预算(例如 256)。
- 那 1000 个 Token 的长 Prompt 会被切成小块(比如
[256, 256, 256, 232]),分在多次迭代中“分期付款”式地处理。
终极合体:算力与带宽的有效协同 此外,系统可以将 “1 个被切小的 Prefill 块” 和 “几十个正在 Decode 的老请求” 打包在同一个 Batch 里发送给 GPU!
- Decode 请求每次只生成 1 个 Token,是访存密集型的,大量闲置了 GPU 的 Tensor Core 算力,但占满了显存带宽。
- Prefill 块包含几百个 Token,是计算密集型的,刚好填补了 Decode 请求闲置出来的 GPU 算力!
这种“拼车”模式让 GPU 的算力和带宽同时达到饱和,实现了高效的资源利用率。
第十四章:显存饱和时的抢占与调度机制¶
即使有了 PagedAttention 的精细化管理,在极端高并发或超长文本的压力下, GPU 显存依然可能被完全占满。当显存饱和但仍有请求挂起时,调度器该如何做出抉择?
第一节:调度器的困境¶
想象一下,你的 GPU 显存已经 100% 满负荷运转,连一个多余的 Block 都开辟不出来了。此时系统面临两类不同的内存需求:
- 非活跃数据:之前交互产生的、缓存在基数树中的前缀 KV Cache(引用计数为 0)。
- 活跃数据:当前正在“车上”处理的请求,在下一步生成新 Token 时需要申请新的物理块。
如果不做处理,系统就会直接发生 OOM(显存溢出) 崩溃。作为系统的核心调度组件,调度器必须启动不同的应对机制。
第二节:非活跃缓存的淘汰与分层卸载(Tiered Offloading)¶
当显存爆满时,系统首先会尝试清理那些暂时没人在用的缓存数据。
系统采用了类似于操作系统的 LRU(最近最少使用) 淘汰策略:
- 系统会维护每个节点的最后访问时间戳与引用计数。
- 引用计数大于 0 的节点(说明正有请求在使用)不能被回收。
- 系统会扫描那些引用计数为 0 的叶子节点,找出最久未被访问的进行回收。
从“二选一”到“多级存储”:Tiered Offloading
在传统的缓存管理中,满员意味着“丢弃”。但为了不直接丢弃宝贵的计算结果,现代引擎(如支持大上下文的先进系统)引入了 Tiered KV Cache Offloading(分层 KV Cache 卸载) 机制。
它将计算机经典的多级存储架构(Memory Hierarchy)应用到了 KV Cache 管理中:
- 热数据(GPU HBM):极高带宽,极低延迟,存放当前最活跃的 KV Cache。
- 温数据(CPU RAM):通过 PCIe 总线连接,带宽较 GPU 降了一个数量级,但容量大且便宜。最久未使用的物理块会被 卸载(Offload) 到这里。
- 冷数据(NVMe SSD):容量近乎无限,但访问速度最慢。在极长的文本或海量历史对话场景下,数据可进一步下沉到 SSD。
当该用户再次提问并命中这些历史缓存时,系统会异步将数据从 SSD/CPU 内存拉回到 GPU 显存。这种“空间换时间”的精细化管理,极大增加了系统的“短期记忆”容量。
第三节:活跃请求的抢占:Swap vs. Recompute¶
如果清理了缓存后显存依然不够,调度器就必须对正在运行的请求执行抢占了,即启动 抢占(Preemption) 机制:暂停一部分请求,释放它们的显存,优先保证其他请求顺利完成。
抢占请求的选择策略
在决定启动抢占时,调度器面临的第一个问题是:在当前“车上”的所有活跃请求中,应该选择哪一个请求进行抢占?
与非活跃缓存使用 LRU(最近最少使用)这种“看过去”的策略不同,对于活跃请求,现代推理引擎(如 vLLM)通常遵循 “优先保护存量请求,抢占新请求” 的原则,采用 逆向先来先服务(Reverse FCFS) 策略:
- 沉没成本(Wasted Work)最小化:老请求已经完成了庞大的 Prefill 计算,并生成了较多 Token,抢占它们会造成巨大的算力浪费。而新请求刚开始,损失最小。
- 保证最终完成(避免饥饿):如果随机抢占或优先抢占老请求,长文本请求可能会因为不断被中断而永远无法完成。
- 优先级机制:如果系统支持请求优先级,则低优先级的活跃请求会被优先选择。
在决定了抢占哪个请求之后,接下来的问题是:如何处理它已经算好的 KV Cache?vLLM 提供了两种经典的权衡策略:
策略 A:换出与换入(Swapping)
- 做法:把被暂停请求的 KV Cache 从昂贵的 GPU 显存,通过 PCIe 总线拷贝到便宜的 CPU 内存(Host RAM) 中。等 GPU 显存宽裕了,再从 CPU 内存拷回来(Swap In)继续算。
- 优缺点:它节省了 GPU 算力(不需要重算),但极其消耗 PCIe 带宽。在大并发下,海量数据的频繁搬运很容易导致 PCIe 通道发生严重拥堵,导致系统吞吐量急剧下降。
策略 B:丢弃并重计算(Recomputation)
- 做法:直接将被暂停请求在 GPU 里的 KV Cache 删除。等轮到它再次执行时,把它的历史输入重新走一遍 Prefill 阶段,把丢掉的 KV Cache 重新算出来。
- 优缺点:虽然这种方法增加了计算量,但在 A100/H100 等顶级显卡上,GPU 的计算能力是严重过剩的,而显存和带宽才是真正的瓶颈。重新计算的代价,在很多时候比通过 PCIe 搬运几十 GB 数据还要快!
vLLM 默认会优先尝试 Swap,但在显存和带宽极度紧张的极限场景下,Recomputation 往往是保证系统稳定运行的关键策略。
第四节:SGLang 的树状管理:将抢占与淘汰融为一体¶
在讨论完传统的 Swap 和 Recompute 策略后,我们来看看以 SGLang 为代表的、原生基于 Radix Tree(基数树) 的推理引擎是如何处理显存饱满情况的。
SGLang 的核心思想是将缓存共享、淘汰与活跃请求的抢占完全统一到了一棵树的拓扑结构中:
-
抢占即解引用: 在 SGLang 中,所有请求的 KV Cache 都是树上的分支。当一个活跃请求因为显存不足需要被抢占(暂停)时,系统不需要做任何跨介质的数据搬运(如 Swap),也不需要立即删除数据(如 Recompute)。它只需要将该请求暂停,其在 Radix Tree 上对应节点的引用计数归零即可。
-
最佳努力缓存(Best-effort Retention): 这些引用计数归零的节点依然保留在显存中,退化为“非活跃缓存”。如果接下来的显存压力得到缓解,且这些节点在 LRU(最近最少使用)淘汰机制中幸存下来,那么当该请求恢复时,系统会直接命中缓存,实现无缝恢复。
-
高度的简洁性: 这种设计避免了显式的 Swap 状态机和复杂的跨设备内存调度。如果显存实在不够用,树的 LRU 机制会自然地剔除最久未使用的叶子节点,被剔除的请求在恢复时则自动退化为 Recompute。这种简化的设计,在处理多轮对话和智能体分支等复杂场景时非常有效。
[!NOTE] SGLang 的演进:从纯 Recompute 到 Tiered Offload 需要指出的是,上述设计描述的是早期 SGLang 的核心逻辑(被剔除即丢弃并重算)。而在最新的发展中,SGLang 也引入了分层缓存卸载机制(如 HiCache 架构),允许将 Radix Tree 中的非活跃节点卸载到 CPU 内存或 SSD 中,从而结合了树状管理的灵活性与多级存储的容量优势。详情可参考 LMSYS Blog: SGLang HiCache。
第十五章:用“闲置算力”换取“极致延迟”:投机采样(Speculative Decoding)¶
在大模型推理的场景中,我们已经通过 PagedAttention 解决了显存碎片的空间浪费,通过连续批处理消灭了 Padding 带来的无效计算。然而,在自回归的 Decode 阶段,我们依然面临着显著的物理限制:显存带宽瓶颈(Memory-Bound)。
正如我们在第八章所描述的,Decode 阶段就像是“开着重型卡车去送一颗螺丝钉”。为了生成仅仅 1 个 Token,GPU 必须把几百 GB 的模型权重从显存完整搬运到计算核心一次。GPU 强大的 Tensor Core 算力绝大多数时间都处于闲置状态等待数据。
既然硬件的显存带宽已经锁死,大模型在每次搬运权重后,能否增加单次处理的计算量?系统工程师和算法科学家们联合推出了一项巧妙的技术——投机采样(Speculative Decoding)。它的核心思想是:它并不试图去改变显存带宽这一物理限制,而是通过“以计算换时间”的策略,充分利用 GPU 闲置的算力,提高单次访存的计算强度(Arithmetic Intensity),从而大幅降低低并发下的生成延迟,显著减小 TBT(提升单用户 TPS)。但需要注意的是,它在极高并发场景下可能会争抢算力,反而对系统整体的 Throughput 产生负面影响。
第一节:教授与助手 —— 投机采样的核心逻辑¶
投机采样的工作原理可以用一个 “教授与助手” 的比喻来理解:
- 目标模型(Target Model):知识渊博的老教授(大模型)。学识顶级,但写字极慢,不过一眼就能看出别人写得对不对。
- 草稿模型(Draft Model):年轻手快的助手(超小模型或外挂头)。学识一般,但手速极快,虽然偶尔会犯错。
如果让老教授自己一个字一个字地写,耗时极长。投机采样的做法是:
- 助手打草稿(Drafting):助手凭感觉快速生成 \(K\) 个词(比如 5 个词)。因为助手模型很小,这 5 个词生成得极快。
- 教授批改(Verification):大模型一次性把这 5 个词全部吃进去。注意:这里大模型不是一个字一个字地去读,而是利用 GPU 的并行能力(类似 Prefill),一次性计算出这 5 个位置上自己想说的词。
- 接受与纠正:如果老教授发现前 3 个词助手猜对了,第 4 个词用词不当。教授会把第 4 个词改成正确的,把第 5 个字扔掉(即回滚 KV Cache)。 这一轮下来,老教授只进行了一次“前向传播”,但我们实际上得到了 4 个完全由老教授认可的高质量词!这比老教授自己单独生成 4 次耗时显著减少。
第二节:计算强度的“逆向交易”¶
你可能会敏锐地发现:小模型算了一遍,大模型又验证了一遍,总计算量(FLOPs)不是增加了吗?
总计算量确实增加了,但在 Memory-Bound 场景下,这种权衡通常是有效的。
我们在 第八章 算过,Decode 阶段的计算强度极低(比如 1.9 FLOPs/Byte),远远低于硬件的拐点。GPU 的计算核心处于严重闲置状态。 投机采样虽然增加了计算量,但它让大模型一次性处理多个 Token。大模型验证 5 个词所需要读取的权重,和生成 1 个词所需要读取的权重是完全一样的(都是几百 GB)。我们只从显存中读取了一次模型权重,却同时验证了 5 个 Token 的生成。 这虽然没有物理性地改变显存带宽的限制,但它显著提高了单次访存的计算强度,让原本饥饿的 GPU 核心得以更充分利用,用闲置的算力换取了更短的耗时。
第三节:从双模型到外挂头:架构的演进¶
如何优雅地“打草稿”,是近年来技术演进的核心。这里经历了从“双模型”到“单模型外挂”的演进。
1. 经典的双模型方案(Dual-Model)
- 原理:这是投机采样的鼻祖。它在显存中同时加载两个独立的模型:一个参数量巨大的目标模型(如 Llama-3 70B)和一个参数量极小的草稿模型(如 Llama-3 8B,甚至更小的蒸馏模型)。
- 工作流:小模型老老实实做自回归 Decode,生成 \(K\) 个词;大模型并行验证这 \(K\) 个词。
- 痛点:工程实现较为复杂。你需要同时 serve 两个模型,管理两套独立的 KV Cache。如果小模型独立性太强,猜中率难以保证,且小模型本身也会占用宝贵的显存。
2. Medusa(美杜莎):多头并行预测 为了解决双模型带来的复杂性,Medusa 提出了一个新的思路:不使用独立的小模型。 它直接在大模型的最后一层特征(Hidden State)上,接了几个并列的、超轻量级的外挂头(Heads)。
- 原理:Head 1 预测 \(t+1\) 的词;Head 2 直接盲猜 \(t+2\) 的词;Head 3 直接盲猜 \(t+3\) 的词。
- 局限:这是一种非自回归的“盲猜”。Head 2 在不知道 \(t+1\) 是什么的情况下硬猜 \(t+2\) ,随着步数增加,准确率会急剧下降。
3. Eagle(老鹰):特征级自回归 为了解决 Medusa 的上述局限性,Eagle 引入了更符合语言链式法则的设计。
- 原理:它在特征层面引入了一个极小的单层 Transformer。它把大模型的特征 \(h_t\) 和预测的 Token 结合,用小网络自回归地推导出 \(h'_{t+1}\) ,再基于它预测 \(t+2\) 。这既保持了极高的预测准确率,又避免了大模型几十层网络带来的巨大耗时,是一种更有效的“打草稿”方式。
[!NOTE] 核心思考:这些“外挂”是谁训练的?为什么大模型不自带?又是在哪里实现的?
- 谁训练的?:通常是开源社区或企业针对特定场景(如代码生成、法律文档)训练的。因为猜中率极度依赖业务场景的词汇分布,通用的头往往不如场景定制的头有效。
- 为什么不自带?:Base Model 的创作者(如 Meta)专注于打造最聪明的“大脑”,而投机采样属于“系统级加速”。解耦可以让大模型保持纯粹,而让下游根据需求定制加速插件。
- 在哪里实现?:所有的调度、双模型通信、以及最难的 树状 KV Cache 动态修剪(回滚),全都是在 推理框架(Inference Framework,如 vLLM, SGLang) 中实现的。大模型本身只负责做矩阵乘法,而框架负责全局的调度与控制。
第四节:树状注意力与生产环境的取舍¶
1. 树状注意力(Tree Attention) 无论是 Medusa 还是 Eagle,为了提高猜中率,都会采用生成多个候选分支的策略:在每一步都分叉给出 Top-K 个候选词。这在自回归的过程中,形成了 “草稿树” 的拓扑结构。 推理框架(如 vLLM)会把这棵树一次性喂给大模型。大模型的 KV Cache 会在验证时扩展为树状结构,大模型用 Tree Attention 找出那条最正确的路径,然后框架会把其他分支的缓存 修剪(Pruning,即 KV Cache 回滚) 掉,恢复为线性序列。
需要指出的是,虽然有大模型(老教授)的验证机制托底,最终输出的正确率不会受到任何影响,但在 GPU 上进行这种 KV Cache 的树状修剪和物理回滚也伴随着相应的开销。如果草稿模型(助手)的“命中率”太低,系统频繁地在“生成-验证-丢弃-回滚”的循环中消耗资源,不仅无法加速,反而会因为频繁的显存指针操作和管理开销拖慢整体的推理速度。这进一步解释了为什么投机采样在生产环境中需要根据并发情况和命中率进行“动态取舍”。
2. 生产环境的动态取舍 在生产环境中,投机采样通常需要根据场景动态开启:
- 低并发/追求极致延迟时(如 Batch Size = 1,实时对话):开启投机采样。GPU 算力空闲,用多余的算力换取更快的吐字速度。
- 高并发/追求极致吞吐时(高峰期):关闭或降低投机采样。因为此时 GPU 的计算核心已经被海量的 Batch 塞满,再做投机采样就是纯粹的浪费,反而会导致排队时间变长。此时,多 Batch 几个请求比多猜几个词更划算。
第三部分我们探讨了单机推理优化的核心技术,分析了 GQA、PagedAttention、连续批处理以及投机采样等优化方法如何协同工作,将单张显卡的性能发挥到极致。这些优化成功地让大模型走出了实验室,具备了服务海量用户的能力。
然而,当模型的参数量走向千亿、万亿,当上下文窗口走向百万级别,单台机器的物理极限终究会被打破。我们该如何让成百上千张显卡协同工作?如何实现更深入的“战略级”资源隔离?请随我进入第四部分,我们将跳出单机的束缚,俯瞰分布式推理与集群调度的整体架构。