TGI PagedAttention 深度解析分块 KV Cache、前缀复用与连续批处理的内存优化【免费下载链接】text-generation-inferenceLarge Language Model Text Generation Inference项目地址: https://gitcode.com/GitHub_Trending/te/text-generation-inference本指南围绕 TGIText 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 cacheKey-Value Cache。KV cache 的大小随模型规模和序列长度快速增长大模型通常有数十乃至上百个 attention 头每个头都要保存键和值两组向量维度等于head_size序列越长需要缓存的 token 数越多。对于大模型 长序列的组合KV cache 会占据非常可观的显存。更关键的是传统实现中 KV cache 需要连续的显存空间每个请求提前按最大序列长度预留一整块内存。这带来两个问题预留过度实际生成长度往往远小于预留的max_total_tokens大量显存被闲置浪费碎片化请求完成、内存释放后不同请求占用的显存块大小不一容易产生无法被有效利用的碎片。这正是 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 如何落地 PagedAttention4.1 复用 vLLM 的自定义 CUDA KernelTGI 的 PagedAttention 实现直接复用了 vLLM 项目开发的自定义 CUDA kernel对应原文档所述vLLM 即 vLLM Project 的自定义内核来源。在 CUDA 后端中TGI 通过加载kernels-community/paged-attention内核模块来获得 PagedAttention 相关算子见 server/text_generation_server/layers/attention/cuda.pyif SYSTEM cuda: try: paged_attention_kernels load_kernel( modulepaged_attention, repo_idkernels-community/paged-attention ) except Exception as e: raise ImportError( fCould 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_SIZEPagedAttention 并不是 TGI 唯一的注意力实现。TGI 通过环境变量ATTENTION选择后端server/text_generation_server/models/globals.py合法取值包括ATTENTION 取值说明默认 BLOCK_SIZEpagedPagedAttentionvLLM 内核16flashdecodingFlash DecodingFlash Attention v2 内核256flashinferFlashInfer 内核支持前缀缓存1flashdecoding-ipexIntel IPEX 上的 Flash Decoding64BLOCK_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 后端的SimpleAllocatorbackends/v3/src/block_allocator.rs初始化时把所有块放入free_blocks空闲列表块 0 保留给健康检查分配时按tokens.div_ceil(block_size)向上取整计算所需块数从空闲列表末尾一次性取出生成结束BlockAllocation被 Drop时通过free()把块归还空闲列表对 HPUHabana设备还额外多申请 1 个 slot 用于 ping-pong 优化。4.4 解码吞吐的关键连续批处理PagedAttention 的按需分配与连续批处理Continuous Batching是天生的一对。TGI v3 后端在后台batching_task中持续组批backends/v3/src/backend.rs每个请求在入队时获得自己的BlockAllocationbackends/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。二者都实现同一Allocatortraitallocate/free。5.2 Radix 树基数树的设计RadixAllocator的核心是一棵radix triebackends/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.rsevict操作按访问时间从最旧到最新驱逐叶节点ref_count 为 0回收块供新分配使用radix.rs。5.3 一次分配的生命周期RadixAllocator::allocate的完整流程radix.rs在 trie 中查找 prefill tokens 的最长公共前缀得到命中的物理块与prefix_len对命中节点incref即使后续分配失败也要保持引用防止前缀被驱逐计算需要新分配的 suffix 块数从空闲列表或驱逐缓存块中补齐生成BlockAllocation其中prefix_len表示无需重算 KV 的部分——解码/后续 prefill 可以直接跳过这些 token 的 attention 计算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_size2时只按块边界复用前缀完全/部分重叠 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_SIZE1即每个块只装 1 个 token前缀复用的粒度最细而paged后端块大小为 16。六、实战配置如何调优 PagedAttention 相关参数6.1 关键启动参数以下是 TGI launcher 中与 PagedAttention / 批处理强相关的参数launcher/src/main.rs参数默认值作用--max_batch_total_tokens自动推断批次内所有请求的 token 总数上限直接决定 KV 块预算与批次容量应在模型加载后尽可能大的前提下设置以充分利用剩余显存--max_batch_prefill_tokensmax_input_tokens 50限制单次 prefill 操作的 token 数prefill 计算密集、显存占用高限制其规模可避免峰值内存超限--max_waiting_tokens20允许等待队列中的请求抢占运行批次前累积的 token 数过小会导致频繁 prefill 打断 decode过大则等待请求延迟过高--waiting_served_ratio0.3等待请求数 / 运行请求数的比例阈值达到该比例且批次有空间时才触发抢占式组批--max_batch_size无强制限制单批次请求数主要面向不支持无 padding 推理的硬件目标6.2 环境变量环境变量取值作用ATTENTIONpaged/flashdecoding/flashinfer/flashdecoding-ipex选择注意力内核后端paged即本文主题的 PagedAttentionPREFIX_CACHING0/1/true/false是否启用前缀缓存仅flashinfer等后端支持LoRA 场景自动关闭PREFILL_CHUNKING默认1启用是否允许把超长 prefill 拆成多个 chunk与max_batch_prefill_tokens配合避免峰值内存见 globals.pyTGI_WIGGLE_ROOM默认0.90显存预算的余量系数用于估算 KV cache 可分配块数时留出安全空间globals.py6.3 调优建议基于源码推断优先让max_batch_total_tokens尽可能大TGI 会自动推断该值以尽量利用剩余显存手工设置时建议以模型加载后剩余显存为上限结合TGI_WIGGLE_ROOM保留余量避免 OOM。块大小与碎片权衡paged后端块大小为 16flashinfer为 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考