企业官网建设流程全解析

TGI PagedAttention 深度解析：分块 KV Cache、前缀复用与连续批处理的内存优化

【免费下载链接】text-generation-inferenceLarge Language Model Text Generation Inference项目地址: https://gitcode.com/GitHub_Trending/te/text-generation-inference

本指南围绕 TGI（Text Generation Inference）中 PagedAttention 的设计与实现展开：先厘清它解决的核心问题——解码阶段 KV Cache 的显存碎片与浪费，再拆解其"分块存储 + 查找表寻址 + 按需分配"的核心机制，随后结合本仓库源码深入讲解 TGI 如何复用 vLLM 的自定义 CUDA kernel、如何通过前缀缓存（Prefix Caching）与 Radix 树实现 KV 跨请求共享，以及这些机制如何与连续批处理（Continuous Batching）配合提升吞吐。读完本文，你将理解 PagedAttention 在 TGI 中的完整落地路径，并掌握与其相关的关键启动参数（--max_batch_total_tokens、PREFIX_CACHING、ATTENTION等）的配置方法。

一、问题背景：为什么大模型生成会"卡"在内存上

LLM（大语言模型）在生成过程中面临的主要瓶颈之一是内存限制。在生成（Generation）的解码阶段，模型需要将之前所有 token 产生的 attention 键值对（Key/Value）保存在 GPU 显存中，供后续 token 计算 attention 时重复使用。这部分缓存被称为KV cache（Key-Value Cache）。

KV cache 的大小随模型规模和序列长度快速增长：

• 大模型通常有数十乃至上百个 attention 头；
• 每个头都要保存键和值两组向量，维度等于head_size；
• 序列越长，需要缓存的 token 数越多。

对于大模型 + 长序列的组合，KV cache 会占据非常可观的显存。更关键的是，传统实现中 KV cache 需要连续的显存空间：每个请求提前按最大序列长度预留一整块内存。这带来两个问题：

1. 预留过度：实际生成长度往往远小于预留的max_total_tokens，大量显存被闲置浪费；
2. 碎片化：请求完成、内存释放后，不同请求占用的显存块大小不一，容易产生无法被有效利用的碎片。

这正是 TGI 引入 PagedAttention 的核心动机。

二、PagedAttention 核心思想：把 KV cache 切成"页"

PagedAttention 的优化思路是将 KV cache 划分为固定大小的块（Block），并通过一张**查找表（Lookup Table，即 block table）**来访问这些块，从而：

• KV cache 不再需要连续存储：逻辑上属于同一序列的 KV 数据，可以散落在显存的任意物理块中，由查找表记录"第几个逻辑块对应哪块物理内存"；
• 按需分配（allocate as needed）：请求开始时只分配当前需要的块，后续生成新 token 时再动态追加新块，用完即还，杜绝了"一次预留整条序列"的浪费。

内存利用率的提升直接反映在吞吐上：由于 PagedAttention 按需、精细地管理显存，内存受限（memory-bound）的推理负载可以获得更高的 GPU 利用率，从而支撑更大的推理批次（batch），即同一时刻并行处理更多请求。

三、KV 共享：并行采样与多代生成的基础

查找表机制带来的第二个好处是KV 跨多个生成共享（KV sharing）。

以并行采样（parallel sampling）为例：对同一个 prompt 同时生成多个输出时，所有生成在前缀阶段（prefill）产生的 KV 是完全相同的。在 PagedAttention 的块结构下，这些生成可以共享同一组已缓存的 KV 块，不必为每个输出各自复制一份前缀 KV，进一步节省显存、减少重复计算。

这一能力在 TGI 的 v3 后端中被进一步系统化，发展为前缀缓存（Prefix Caching）：任何两个请求只要共享一段 token 前缀，就可以复用对应的缓存 KV 块，避免重复 prefill（详见下文第五节）。

四、源码实现：TGI 如何落地 PagedAttention

4.1 复用 vLLM 的自定义 CUDA Kernel

TGI 的 PagedAttention 实现直接复用了 vLLM 项目开发的自定义 CUDA kernel（对应原文档所述，vLLM 即 vLLM Project 的自定义内核来源）。

在 CUDA 后端中，TGI 通过加载kernels-community/paged-attention内核模块来获得 PagedAttention 相关算子，见 server/text_generation_server/layers/attention/cuda.py：

if SYSTEM == "cuda": try: paged_attention_kernels = load_kernel( module="paged_attention", repo_id="kernels-community/paged-attention" ) except Exception as e: raise ImportError( f"Could not import attention kernels. Make sure your installation is correct. Complete error: {e}" )

解码阶段（每步只生成 1 个 token）调用paged_attention()函数，其中根据序列长度与并行度在PagedAttention V1 / V2之间做启发式选择（cuda.py）：

• 若分区数（partition）为 1，或序列数 × 头数足够大（num_seqs * num_heads > 512），使用V1，避免 V2 的规约（reduction）开销；
• 否则使用V2，通过_PARTITION_SIZE = 512把长序列分成多个 partition 并行计算后再合并，降低单次 kernel 的寄存器与共享内存压力。

KV 写入侧同样依赖 vLLM 内核：paged_reshape_and_cache负责把新计算的 key/value 按slots索引写入块缓存，并支持 FP8 KV cache 的量化写入（见 server/text_generation_server/layers/attention/kv_cache.py）。在 ROCm 与 Intel IPEX 后端，则分别调用 vLLM 自定义算子ops.reshape_and_cache与 IPEX 的PagedAttention.reshape_and_cache。

4.2 三种注意力后端与 BLOCK_SIZE

PagedAttention 并不是 TGI 唯一的注意力实现。TGI 通过环境变量ATTENTION选择后端（server/text_generation_server/models/globals.py），合法取值包括：

BLOCK_SIZE（即每个块容纳的 token 数）是 PagedAttention 的关键超参：块越小，显存碎片越少、按需分配越精细，但 block table 与调度开销越大；块越大则相反。paged后端默认 16，与 vLLM 默认值一致。

KV cache 的物理布局也因后端而异：paged后端采用[num_blocks, num_heads, head_size, BLOCK_SIZE]的块布局（见 kv_cache.py），即"按块连续存放 BLOCK_SIZE 个 token 的同一头"，保证访问缓存块时内存连续、利于 GPU 并行读取。

4.3 块表（block table）与按需分配

解码时，每个请求维护一张block table：记录其逻辑 KV 序列由哪些物理块组成。TGI 在 CUDA 后端将其作为block_tables张量直接传给 vLLM 内核（cuda.py），内核据此在解码时通过查找表访问散落各处的 KV 块。

分配与释放的"按需"语义体现在 v3 后端的SimpleAllocator（backends/v3/src/block_allocator.rs）：

• 初始化时把所有块放入free_blocks空闲列表（块 0 保留给健康检查）；
• 分配时按tokens.div_ceil(block_size)向上取整计算所需块数，从空闲列表末尾一次性取出；
• 生成结束（BlockAllocation被 Drop）时通过free()把块归还空闲列表；
• 对 HPU（Habana）设备还额外多申请 1 个 slot 用于 ping-pong 优化。

4.4 解码吞吐的关键：连续批处理

PagedAttention 的按需分配与连续批处理（Continuous Batching）是天生的一对。TGI v3 后端在后台batching_task中持续组批（backends/v3/src/backend.rs），每个请求在入队时获得自己的BlockAllocation（backends/v3/src/queue.rs），因此：

• 不同长度的请求可以共享同一个 batch，无需按最长序列 padding；
• 一个请求生成完成后立即释放其 KV 块，腾出的块马上可被新请求使用；
• 块预算的上限由max_batch_total_tokens决定，batch 内所有请求的 token 总和不得超过该值（launcher/src/main.rs）。

正是因为 PagedAttention 把 KV cache 细化为可独立分配/释放的块，连续批处理才能把"批次中任何时刻的总 token 数"当作唯一约束，从而最大化显存利用率。

五、前缀缓存：Radix 树实现的 KV 共享

5.1 两级分配器架构

在 v3 后端，块分配由一个后台block_allocator_task统一管理（block_allocator.rs）：当启用前缀缓存时实例化RadixAllocator，否则实例化SimpleAllocator。二者都实现同一Allocatortrait（allocate/free）。

5.2 Radix 树（基数树）的设计

RadixAllocator的核心是一棵radix trie（backends/v3/src/radix.rs），其设计受 SGLang 的 RadixAttention 启发，专门为前缀缓存优化：

• key = token 序列，value = 物理块序列，且键值等长——插入前缀abc → 块xyz后，a、ab也能查到对应块（因为节点按块粒度共享）；
• 每个节点记录ref_count（引用计数）与last_accessed（最近访问时间）：
  • ref_count > 0表示该前缀正被某个活跃请求使用，不可被驱逐；
  • 访问节点会更新其访问时间，用于 LRU 驱逐决策；
• find操作用来查找请求 prefill token 与缓存的最大公共前缀，把命中的块写入分配结果，同时提升沿途节点的访问时间；
• insert操作把新 prefill 的 token→块映射写入 trie；遇到部分重叠的前缀时通过split_node把节点一分为二，保证只有真正共享的部分才复用（radix.rs）；
• evict操作按访问时间从最旧到最新驱逐叶节点（ref_count 为 0），回收块供新分配使用（radix.rs）。

5.3 一次分配的生命周期

RadixAllocator::allocate的完整流程（radix.rs）：

1. 在 trie 中查找 prefill tokens 的最长公共前缀，得到命中的物理块与prefix_len；
2. 对命中节点incref（即使后续分配失败也要保持引用，防止前缀被驱逐）；
3. 计算需要新分配的 suffix 块数，从空闲列表或驱逐缓存块中补齐；
4. 生成BlockAllocation，其中prefix_len表示无需重算 KV 的部分——解码/后续 prefill 可以直接跳过这些 token 的 attention 计算；
5. free时对前缀节点decref，并把超出已缓存前缀的新 token→块映射insert回 trie，供后续请求复用。

5.4 测试与基准验证

仓库为 Radix 分配器提供了详尽的单元测试（backends/v3/src/radix.rs 内嵌#[cfg(test)] mod tests），覆盖：

• 前缀复用：同一 prefill 第二次分配时prefix_len从 0 变为完整长度（如allocator_reuses_prefixes）；
• 块对齐：block_size=2时只按块边界复用前缀；
• 完全/部分重叠 prefill 的释放与内存回收；
• LRU 驱逐顺序：先回收"更旧"的分配；
• 随机压力测试：10 万次随机 allocate/free 后校验"无重复块、无块泄漏、前缀引用计数正确"等不变量。

此外，backends/v3/benches/prefix_cache.rs 提供了基于 Criterion 的分配器基准，用于评估长随机 prefill 场景下 Radix 分配/释放的性能。

5.5 前缀缓存的启用条件

值得注意，前缀缓存在 v3 后端默认启用（launcher 未显式设置PREFIX_CACHING时默认true，见 launcher/src/main.rs），但要求注意力后端为flashinfer——globals.py中明确校验：仅当ATTENTION为flashinfer/flashdecoding/flashdecoding-ipex时才允许前缀缓存（globals.py），否则直接抛错。同时，启用 LoRA 适配器时会自动关闭前缀缓存（launcher/src/main.rs），因为不同适配器下的 KV 语义不同，不能跨适配器共享。

这也解释了默认BLOCK_SIZE的差异：flashinfer后端BLOCK_SIZE=1，即每个块只装 1 个 token，前缀复用的粒度最细；而paged后端块大小为 16。

六、实战配置：如何调优 PagedAttention 相关参数

6.1 关键启动参数

以下是 TGI launcher 中与 PagedAttention / 批处理强相关的参数（launcher/src/main.rs）：

6.2 环境变量

6.3 调优建议（基于源码推断）

• 优先让max_batch_total_tokens尽可能大：TGI 会自动推断该值以尽量利用剩余显存；手工设置时建议以"模型加载后剩余显存"为上限，结合TGI_WIGGLE_ROOM保留余量，避免 OOM。
• 块大小与碎片权衡：paged后端块大小为 16，flashinfer为 1。块越小越省显存但调度开销越高；若追求前缀复用精度，可考虑flashinfer后端。
• 长 prompt / 高并发场景开启前缀缓存：大量请求共享 system prompt 或文档前缀时，前缀缓存可显著减少重复 prefill；确认后端为flashinfer且未使用 LoRA 即可生效。
• prefill 与 decode 的节奏平衡：通过--max_batch_prefill_tokens限制单次 prefill 规模、用--max_waiting_tokens控制抢占频率，避免 decode 被频繁打断导致吞吐下降。

七、总结

PagedAttention 是 TGI 解决 KV cache 显存瓶颈的核心机制：通过固定大小分块 + 查找表寻址 + 按需分配，消除连续内存预留带来的浪费与碎片，使内存受限负载的 GPU 利用率与批次容量显著提升；同时，查找表结构天然支持KV 跨请求共享，为并行采样与更通用的前缀缓存（Radix 树实现）奠定基础。在仓库层面，TGI 复用了 vLLM 的自定义 CUDA kernel 完成块访问与写入，并在 v3 后端用 Rust 实现了块分配器与 Radix 前缀缓存，配合连续批处理与max_batch_total_tokens等参数，共同支撑起高吞吐的文本生成服务。

【免费下载链接】text-generation-inferenceLarge Language Model Text Generation Inference项目地址: https://gitcode.com/GitHub_Trending/te/text-generation-inference