Tangram:解锁非均匀KV缓存压缩以实现高效多轮大语言模型服务 Tangram: Unlocking Non-Uniform KV Cache Compression for Efficient Multi-turn LLM Serving
通过预算保留、锯齿分页和AOT负载均衡三大技术,将非均匀KV缓存压缩的理论优势转化为实际系统性能,实现高达2.6倍的吞吐量提升。
前置知识
KV缓存
在Transformer模型的自注意力机制中,每个token的Key和Value向量需要被缓存以供后续token生成时使用。对于L层H头的模型,KV缓存存储了所有层和头中每个token的K和V张量,其大小随上下文长度线性增长。在多轮对话中,历史对话会被保留以维持上下文连贯性,导致KV缓存随着轮次和并发用户数快速膨胀,常常超过模型参数本身的大小,成为内存瓶颈而非计算瓶颈。
本文的核心问题就是多轮对话中KV缓存无限增长导致的内存压力,理解KV缓存如何工作、为什么需要压缩、以及压缩对模型精度的影响,是理解本文贡献的基础。
PagedAttention
PagedAttention是vLLM等现代服务框架采用的内存管理技术,借鉴操作系统的分页思想,将KV缓存分割为固定大小的块。每个块可以独立分配和释放,消除了跨请求的外部内存碎片。PagedAttention使用一个块表来记录每个请求的物理页面地址,在调度步骤中预先计算并分配所需页面。其关键约束是单体页面结构:一个物理块同时跨越所有层和所有注意力头,这导致无法进行细粒度的头级压缩,因为一个页面只有在所有共享它的头都释放时才能被回收。
PagedAttention的单体页面结构是导致非均匀压缩在现有系统中无法实用的根本原因之一。本文的Ragged Paging技术正是为了解决这个核心设计约束。
非均匀KV缓存压缩
与均匀压缩为每个注意力头分配相同数量的token预算不同,非均匀压缩在层级别强制执行全局保留率。它将所有H头的重要性分数展平,然后应用全局token预算top选择,产生高度不规则的隐式per-head保留率——有些头保留完整历史,而其他头被大量修剪。这种不均匀的per-head预算自然地反映了注意力的异构集中模式,关键的长期信息集中在少量检索头中,而其他头仅关注局部信息。
非均匀压缩从算法上更优,能够保持远高于均匀压缩的多轮对话精度。但本文的核心贡献就是解决了非均匀压缩在系统层面的三大障碍,使其真正可用于生产环境。
连续批处理和块化预填充
连续批处理允许在调度步骤中动态地添加和移除请求,使得长时间运行的预填充请求不会垄断一个迭代。块化预填充将长上下文预填充请求分割为固定大小的token块,在连续的调度步骤中处理,并与其他请求的解码步骤交错执行。这种组合确保了每个请求有界的时间到首个令牌,是现代高吞吐量服务的基础。
本文的核心创新之一是将压缩完全融合到块化预填充和连续批处理中,而不是作为一个单独的后期处理阶段。这意味着每个prefill块在生成后立即被压缩,保持与现有服务流程的完全兼容。
注意力头专业化
研究表明,Transformer模型中的注意力头具有模型内在的、专门化的角色,这些角色与输入无关而持续存在。不同头关注不同类型的信息:某些头专注于检索长程上下文中的关键信息,其注意力权重集中在少量token上;其他头则更多地关注局部语法或近邻信息,注意力分布较广。这种专业化意味着在KV缓存压缩中,不同头对历史信息的保留需求存在本质差异。
本文的关键观察是这种头专业化直接反映在KV保留配置中,展现出两层结构规律性:按保留需求排序的头排名基本与输入无关,且每个头的绝对保留率变化范围很小。这使作者能够离线校准而非运行时动态发现压缩策略。
研究动机
多轮LLM服务中,KV缓存会随着对话轮次和并发用户数增长而线性增长。对于Qwen2.5-32B模型,仅16个并发请求,累积的KV缓存在十轮对话内就超过了模型权重本身的大小,并且此后无限制地继续增长。在多轮服务中,内存容量而非计算能力成为限制批处理量从而限制吞吐量的约束条件。KV缓存压缩是标准解决方案,但均匀压缩为所有头保留相同数量的token,忽略了注意力的异构集中模式,导致在多轮设置下精度急剧下降。非均匀压缩虽然在算法上更优,但在现代服务栈中仍然不实用。现有系统的三大支柱都隐式假设所有注意力头保持统一长度的KV缓存,导致三个严重开销:单体页面结构导致非均匀压缩释放的内存无法回收,作为不可恢复的页面碎片被困住;调度器必须过度分配然后跟踪、回收和重新映射分散的页面,这个控制平面混乱消耗多达25%的prefill执行时间;异构的per-head KV长度使GPU SM上的工作负载倾斜,静态内核分区导致stragglers使解码注意力延迟增加高达1.7倍,而动态per-layer重新规划则在每个解码步骤上消耗15-20%的CPU时间。简而言之,非均匀压缩的理论内存节省在服务系统中蒸发殆尽,甚至可能降低端到端吞吐量。
本文的目标是本文的目标是构建一个服务框架Tangram,将非均匀KV缓存压缩的理论效率与高吞吐量服务的实际约束相协调。具体来说,作者希望解决三个核心问题:消除页面碎片化,使非均匀压缩释放的内存能够真正被回收;消除页面管理开销,避免运行时的过度分配和回收过程;消除负载不平衡,使GPU SM的利用率保持均衡而无需运行时规划开销。最终目标是在保持非均匀压缩方法精度的同时,显著提高多轮LLM服务的端到端吞吐量。作者希望证明,通过利用头级KV保留的模型内在规律性,可以将所有运行时负担转化为离线可解析的属性,从而将非均匀压缩从算法承诺转化为实际系统性能。
与已有工作不同的是,本文的独特切入点在于发现并利用了一个关键观察:头级KV保留展现出模型内在的、而非输入驱动的两层结构规律性。具体来说,按保留需求排序的注意力头排名基本与输入无关,且每个头的绝对保留率仅在狭窄、可估计的范围内变化。这个观察改变了整个问题格局:KV缓存异构性不需要在运行时动态发现;它可以一次性离线校准,并被视为静态蓝图。这使所有三个运行时负担——页面簿记、回收和内核规划——都成为可以在执行前确定性地解决的决策。与此前的异构内存系统不同,它们要么在更粗的层粒度上管理异构性,要么独立地处理每个头,但都没有利用这种跨头结构规律性。Tangram通过聚类具有相似、模型内在保留的头到锯齿页面中,使页面边界与实际保留分布对齐,使头级异构内存管理对高吞吐量服务变得实用。
核心方法
Tangram是一个整体服务框架,旨在将非均匀KV缓存压缩的理论优势转化为实际的高吞吐量服务。方法整体思路基于一个核心洞察:每个注意力头的保留遵循模型内在的模式,局限于狭窄的、与输入无关的范围。Tangram利用这个规律性在执行前固定每个头的预算,用静态计划的内存占用替代动态的、输入依赖的压缩。技术路线分为三个组件:预算保留、离线固定每个头的保留率,让调度器在执行前精确分配所需的页面数,消除过度分配和整个压缩回收路径;锯齿分页,将相似预算的头聚类到独立的页面表中,使碎片化成为可回收的内存;AOT负载均衡,从保留的预算配置文件预计算均衡的GPU分区,实现零运行时规划的均衡SM利用率。这三个组件共同作用,将非均匀压缩从算法创新转化为系统实践。
核心创新点是利用头级KV保留的两层结构规律性——输入不变的头排名和狭窄的per-head比例变化范围——将所有运行时负担转化为离线可解析的属性。与现有系统在运行时动态发现KV异构性不同,Tangram意识到这个异构性是模型内在的确定性结构,可以一次性离线校准。这个观察使作者能够做三个关键的静态设计决策:固定每个头的保留率而非为每个输入重新计算,使调度器能够在调度时间精确分配所需页面;聚类具有相似稳定保留水平的头一次,使相似的预算共享一个页面表,实现真正的物理内存回收;基于固定的头组形状预计算最优GPU工作负载分区,实现均衡SM利用率而无需运行时规划。这三个静态决策分别对应解决三个运行时瓶颈:页面簿记、回收和内核规划,将它们从关键路径上完全移除。
方法步骤详情
方法步骤的完整描述如下:第一步是预算校准,在少量样本上运行非均匀压缩,记录每个头从全局top选择中接收的隐式per-head保留率,总结为其均值和标准差。第二步是预算固定,用受控安全边缘固定每个头的per-head预算比率,其中是安全边缘系数,实验中设为2。第三步是预算感知聚类,对于每一层,按校准预算对所有头排序,将个连续头组成一个头组。第四步是预算保留,当请求到达时,Tangram在调度时间分配精确满足所需页面数,每个prefill块生成后立即被压缩,使用per-head top选择而非全局top选择。第五步是锯齿分页,每个头组由一个独立的页面表管理,页面大小为个头的键值对,只保留该组内头实际需要的条目。第六步是AOT负载均衡,预计算静态工作负载分割映射,使用CUDAMaxOccupancy API确定总CTA数量,然后为每个头组分配与其聚合预算成正比的分区因子。第七步是运行时执行,解码时Tangram检索预计算的表来配置注意力内核,实现零运行时规划成本。
技术新颖性
技术新颖性体现在三个方面。首先,预算保留的静态规划方法与动态压缩形成鲜明对比。现有系统在运行时发现每个头的KV缓存将被压缩多少以及需要多少页面,因此必须过度分配然后运行昂贵的压缩和回收过程,这消耗多达25%的prefill时间。Tangram通过固定每个头的预算使请求的压缩后占用在执行前完全确定,调度器在准入请求时精确分配所需页面,每个页面以压缩后状态进入缓存,消除了过度供应和整个页面回收路径。其次,锯齿分页的单体页面解耦是新颖的。PagedAttention的单体页面结构导致一个页面只有在所有共享它的头都释放时才能被回收,所以一个长保留头阻塞每个页面。Ragged Paging通过相似预算的头聚类解决了这个问题,短保留头被一起打包,释放更多页面。最后,AOT负载均衡的静态分区避免了运行时规划成本。FlashInfer在每个解码步骤前重新计算每层的独特分区,但Tangram的静态计划保持有效,因为它从预算比率而非绝对长度计算,每个头组的保留长度与相同上下文长度线性缩放,所以组之间的相对工作量——唯一决定平衡的数量——对上下文长度不变。
实验结果
核心发现在准确性和性能两方面。在准确性方面,Tangram在与三种非均匀压缩方法和五个模型的组合中,紧密跟踪每种方法的原始无Tangram实现,在目标全局保留率下匹配、有时甚至超过原始精度。例如,在Qwen3-4B模型上,当全局保留率为50%时,w/ Tangram和w/o Tangram的精度几乎完全一致。这表明Tangram充当了非均匀压缩的忠实系统基板而非改变压缩本身。在性能方面,Tangram成功将非均匀KV缓存压缩转化为实际系统级增益,在长上下文任务上实现高达2.6倍的吞吐量提升。增量消融实验确认每个组件提供累加的吞吐量增益:从仅预算保留到添加锯齿分页再到完整的AOT负载均衡。在服务质量方面,在75%驱逐下,当请求率增长时Tangram维持低首token时间,而vLLM的首token时间急剧上升。在内存回收方面,动态压缩在75%驱逐率下消耗多达25%的prefill执行时间来跟踪和回收分散的页面,而Tangram几乎没有额外成本。锯齿分页是关键机制,通过head group级别的独立页面回收将非均匀压缩转化为实际内存节省。AOT负载均衡始终实现最低的解码注意力延迟,避免了FlashDecoding的straggler效应和FlashInfer的显著规划开销。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Long Context Serving (SCBench, >100K tokens) | Throughput (requests per second) | 0.009 (Qwen3-4B at 25% retention) | 0.005 (vLLM Full-KV) | 2.6x |
| Mid Context Serving (SCBench, 20K-100K tokens) | Throughput (requests per second) | 0.018 (Qwen3-4B at 25% retention) | 0.007 (vLLM Full-KV) | 2.3x |
| Short Context Serving (SCBench, <20K tokens) | Throughput (requests per second) | 0.450 (Qwen3-4B at 25% retention) | 0.300 (vLLM Full-KV) | 1.5x |
| Decode Attention Latency | Latency reduction vs FlashDecoding | +24.9% (Qwen3-4B) | FlashDecoding | lower latency |
| Page Reclamation Overhead | Prefill time spent on reclamation | ~0% | 25% (dynamic allocation) | eliminated |
| Memory Reclaimed | Fraction of full KV cache reclaimed | 12-25% additional over adjacent grouping | without Budget-Aware Clustering | more efficient |
局限与改进
局限性分析包括作者承认的和观察到的。首先,预算校准依赖于每个头保留率的狭窄变化范围这一假设,虽然有50个样本的实验支持,但在极端输入分布下可能失效。其次,安全边缘系数的选择需要权衡:较大的提供更大的安全裕度但消耗更多预算,而较小的可能不足以覆盖per-input偏差。作者在实验中使用等于2,但没有提供系统的方法来选择这个值。第三,的选择引入基本权衡:小的最大化内存回收但增加管理开销,大的减少管理开销但导致剩余碎片化。作者建议等于4-8作为最佳平衡点,但这个选择可能因硬件和模型而异。第四,预算感知聚类假设头排名是模型内在且输入不变的,虽然作者在三个模型族和三个任务上验证了这个观察,但可能在某些专业模型或极端任务上不成立。第五,AOT负载均衡假设预算比率而非绝对长度决定工作负载平衡,这对于上下文长度线性缩放的情况成立,但在某些非线性场景下可能需要重新评估。最后,作者主要在单个GPU设置下评估,多GPU设置的扩展性可能需要额外的协调机制。
独立分析的弱点
独立分析的弱点包括:预算校准依赖于离线样本,如果服务数据分布与校准样本差异很大,可能导致子优性能。改进方向可以是自适应校准机制,周期性地用实际服务数据更新预算配置文件。的选择需要手动调整,缺乏自动调优机制。改进方向可以是基于实时性能指标的动态选择算法。安全边缘系数是固定的,无法适应不同头的偏差特性。改进方向可以是per-head的自适应,根据观察到的方差动态调整。AOT负载均衡假设预算比率决定工作负载平衡,但在某些情况下可能需要运行时微调。改进方向可以是混合方法:离线计算基准分区,运行时进行轻微调整。当前实现主要针对单GPU,多GPU场景下的扩展性可能需要额外的协调。改进方向可以是在多GPU设置下的分布式负载均衡策略。作者主要在合成基准上评估,实际生产环境中的噪声、抖动和故障模式可能影响性能。改进方向可以是在生产级工作负载下的鲁棒性测试和故障恢复机制。
未来方向
未来研究方向包括作者提出的和基于成果可延伸的。作者没有明确列出未来工作,但基于成果可以延伸:扩展到更多压缩方法,当前评估三种非均匀压缩方法,可以集成更多方法并评估兼容性。自适应校准,开发周期性更新预算配置文件的机制,使系统能够适应数据分布漂移。跨模型迁移,研究校准的预算配置文件是否可以在模型间共享或迁移,减少校准成本。多GPU扩展,将Tangram扩展到多GPU设置,研究分布式场景下的负载均衡和内存管理。硬件感知优化,根据不同GPU架构的SM数量和内存层次结构优化选择和分区策略。端到端优化,将Tangram与其他优化技术如量化、稀疏化、模型蒸馏结合,探索组合效应。理论分析,更深入地分析头级KV保留规律性的理论基础,例如为什么头排名是输入不变的,什么因素影响per-head方差。
复现评估
复现评估:作者声明他们的实现公开发布在https://github.com/aiha-lab/TANGRAM,基于vLLM框架。实验在专用服务器节点上进行,配备Intel Xeon Gold 6326 CPU和四个NVIDIA A100 GPU。评估五个模型,都支持超过100K tokens的上下文窗口。使用SCBench基准,评估长上下文能力。预算校准使用50个样本,安全系数设为2。报告的吞吐量、延迟和碎片化率是在该硬件上实际运行执行的测量结果,而非分析估计或模拟。由于代码公开、硬件配置明确、数据集标准、评估指标清晰,复现难度中等。主要挑战可能是需要配备多个A100 GPU的服务器以及处理长时间运行的实验。
论文图表
Figure 1(a)展示了KV缓存大小随着对话轮次或并发请求数的增长,在Qwen2.5-32B模型上,16个并发请求下不到十轮对话,KV缓存就超过了模型权重大小,并且此后无限制继续增长。Figure 1(b)对比了均匀压缩和非均匀压缩在50%全局保留率下的策略,可以看出非均匀压缩根据不同头的需要分配不同数量的tokens。Figure 1(c)比较了均匀和非均匀压缩在长期对话QA基准上的准确性,显示非均匀压缩在激进压缩下保持远高于均匀压缩的精度。
这张图对理解论文至关重要,因为它直观地展示了问题的严重性、非均匀压缩的算法优势以及本文试图解决的核心挑战。Figure 1(a)建立了问题动机,Figure 1(b)展示了均匀vs非均匀压缩的本质区别,Figure 1(c)提供了精度证据说明为什么非均匀压缩值得在系统层面优化。
这张图展示了非均匀压缩下每个头的KV保留率,在三个模型族和三个SCBench任务上。x轴上的头按保留预算排序,每个box聚合50个输入样本。虚线标记Tangram用安全系数等于2保留的per-head预算。关键观察:每个头的box很窄,表示per-head方差低;不同任务的marker沿单调排序上升,直接可视化多样化工作负载诱导相同的头排名;Tangram的静态预算覆盖每个头的per-input偏差,同时只花费边际额外预算。
这张图对理解论文至关重要,因为它提供了关键观察的视觉证据:头级KV保留展现出两层结构规律性——输入不变的头排名和狭窄的per-head比例变化范围。这个观察是Tangram整个设计的基础,使作者能够离线校准而非运行时动态发现压缩策略。
这个表格展示了负载均衡开销。比例表示总解码推理步骤时间中花在负载均衡阶段的部分。在四个模型上,负载均衡时间从5.64毫秒到17.13毫秒,占总解码时间的15.48%到20.20%。关键观察:per-layer规划开销消耗总解码迭代时间的15-20%,否定动态平衡本应提供的GPU利用率增益。
这个表格对理解论文至关重要,因为它量化了动态负载均衡如FlashInfer的开销。15-20%的解码时间是一个显著的成本,表明动态规划虽然改善了GPU利用率,但CPU规划成本抵消了这些收益。