论文

Tangram: 解锁非均匀 KV 缓存压缩以实现高效的多轮 LLM 服务

Tangram: 解锁非均匀 KV 缓存压缩以实现高效的多轮 LLM 服务

多轮大语言模型(LLM)服务中,对话历史不断累积,导致键值缓存(KV cache)随轮次和用户数持续增长,容量很快超过模型权重本身。这使得内存(而非计算)成为吞吐量的瓶颈约束。非均匀 KV 压缩为不同注意力头分配不同预算,在精度上远优于均匀方案,但目前难以实际应用:现代服务栈假设各头的 KV 长度一致,因此非均匀性带来的内存碎片化、页面回收耗时(占预填充时间的 25%),以及 GPU 负载不均(解码延迟增加 1.7 倍或每步 15–20% 的重规划开销)等问题严重制约性能。 我们观察到这种非均匀性无需在运行时发现:头级保留遵循两级结构规律——输入无关的头部排序及窄边界下的每头比例——仅需 50 个样本即可离线校准。基于此,我们提出 Tangram,一个静态解决动态问题的服务框架: - 预算预留(Budget Reservation):在调度时固定每头压缩后的占用,消除页面回收; - 分页织补(Ragged Paging):将相似预算的头部聚类到独立页表,将碎片转化为可回收内存; - 提前负载均衡(Ahead-of-Time Load Balancing):预计算均衡的 GPU 分区,无需运行时规划。 通过在 vLLM 上实现,Tangram 可作为现有非均匀压缩方法的即插即用基座,在保持精度的同时,将端到端吞吐量相比全 KV 基线提升高达 2.6 倍。代码已开源。

论文精读

TL;DR Tangram 利用离线校准的注意力头重要性规律,静态分配异构 KV 缓存预算,消除运行时回收与负载均衡开销,使非均匀压缩在多轮 LLM 服务中落地,吞吐量最高提升 2.6 倍。

问题

多轮 LLM 推理的 KV 缓存 随对话回合指数增长,内存正取代算力成为吞吐瓶颈。为缓解这一压力,非均匀 KV 压缩(为不同注意力头分配异构预算)在精度保持上远超均匀方案,是业界前沿方向。然而,现有服务系统(如 vLLM)的 PagedAttention 强制所有注意力头使用相同长度的 KV 页,导致异构预算产生严重弊端:

  • 页碎片化:释放的细粒度内存被夹在页块中无法复用,形成大量碎片;
  • 高额回收开销:系统需持续扫描并回收散置的页,占用高达 25% 的预填充时间
  • GPU 负载失衡:各头计算量差异引起静态分区下解码延迟最高膨胀 1.7×,若动态重平衡则每步消耗 15–20% 时间进行调度规划。

这些工程痛点使非均匀压缩的精度优势几乎被运行时开销抵消,难以落地。问题的核心挑战在于:如何在保持非均匀压缩精度收益的同时,消除异构内存管理带来的碎片与调度惩罚? 这需要同时重构内存分配器、分页机制和负载均衡策略,与现有服务框架深度耦合。业界对高吞吐多轮 LLM 推理需求迫切(如客服对话、编程助手),任何内存效率改进都直接转化为可观的成本节约。类比 移动端推理 必须在有限 RAM 中精细管理模型状态,多轮 LLM 服务的 KV 缓存就像移动设备上的应用驻留,内存碎片化会严重拖慢整体响应,亟需一类“数据库表分区”式的静态管理技术来根治。

核心洞察

  • 注意力头的 KV 缓存重要性遵循输入无关的排名和狭窄的比率边界,可通过少量样本离线校准。此规律让 Tangram 免去运行时动态决策,直接在调度阶段静态分配各头预算,彻底消除传统非均匀压缩方法因运行时再规划导致的页面回收开销和负载不均衡,使非均匀压缩真正在服务系统中可用。
  • 将预算相似的注意力头聚类为独立页表(Ragged Paging),使压缩产生的自由内存从碎片变为可回收的连续空间。这解决了 PagedAttention 架构下非均匀压缩必然引入的页面碎片问题,避免了高达 25% 的预填充时间浪费在散页回收上,并配合 Ahead-of-Time Load Balancing 预计算均衡的 GPU 分区,实现了零运行时重规划的高效服务。

方法

Tangram 是一种面向多轮 LLM 服务的 KV 缓存压缩框架,其核心输入为堆积的对话历史及其在每轮、每用户下不断膨胀的 KV 缓存。方法的关键思路是:利用注意力头保留模式的结构化规律,将非均匀压缩中的动态决策转为静态规划,从而消除运行时内存管理开销。框架包含三个紧耦合模块:

  1. Budget Reservation (预算预留)
    在 prefill 阶段,根据离线校准得到的头级预算(head-wise budget),直接计算出每个头压缩后的精确 KV 尺寸,并一次性分配好页表空间。该预算来自对少量样本(如 50 条)的离线分析,表现出输入无关的头重要性排序每头保留比例边界狭窄的双层规律。相比动态回收,完全避免了 page reclamation 引发的高达 25% prefill 时间消耗。

  2. Ragged Paging (不规则分页)
    将预算相近的注意力头聚合成头组(head group),每个组管理独立的页表。这一设计把因异构预算导致的碎片内存转化为可回收的连续块,并通过预算感知聚类(budget-aware clustering) 平衡内存收益与管理开销。向量化的 block table 管理进一步降低了调度复杂度。

  3. Ahead-of-Time Load Balancing (提前负载均衡)
    基于静态划分策略,在编译期预计算平衡的 GPU 分区方案,运行时无需任何重规划步骤。该模块解决了两大难题:静态分区下因头组差异导致的straggler 效应(解码延迟膨胀可达 1.7×),以及动态重均衡带来的每步 15–20% 解码时间开销

最终输出是压缩后的 KV 缓存,其内存占用被精确控制,使吞吐不再受缓存容量限制。

与同类方法的差异:现有系统(如 vLLM 衍生方案)在运行时动态处理非均匀压缩,产生碎片回收和负载再均衡开销;Tangram 通过离线规律发现,将预算分配彻底静态化,在不牺牲精度的前提下,端到端吞吐较完整 KV 基线提升最高 2.6 倍。

实验

实验设计

在典型多轮对话服务场景下,将 Tangram 集成到 vLLM 推理框架,加载不同的非均匀 KV 压缩方法(如 FastGen、PyramidKV 等),在多轮对话数据集上评估准确率与系统性能。对比基线包括保留全 KV 缓存(full-KV)以及原有的 uniform 压缩系统,主要测量端到端吞吐量、prefill 延迟和 decode 延迟等指标。

关键发现

Tangram 在保持与非均匀压缩原始方法完全一致的精度的前提下,通过静态预算预留消除了原有系统中因页碎片回收造成的开销(占 prefill 时间高达 25%),并通过 Ragged PagingAhead-of-Time 负载均衡解决了 GPU 负载倾斜问题(原使 decode 延迟膨胀 1.7×,或每步重规划消耗 15–20% 时间),最终实现最高 2.6× 的端到端吞吐量提升。实验进一步验证了注意力头保留模式的二层次结构规律——头重要性排序与输入无关,每头保留比例的波动范围极窄——仅用约 50 个样本即可完成离线校准,保证了静态调度的可行性。

与基线对比深度解读

原有非均匀 KV 压缩在 vLLM 等系统中会因各头 KV 长度不一致产生大量内存碎片,运行时不得不频繁回收和重映射页表,严重抵消压缩收益。Tangram 的核心贡献在于将异构性从运行时转移到编译时:利用离线校准出的头保留规律,将每个头的压缩后大小固定为编译期常数,从而在调度阶段就能精确分配页面,完全消灭碎片回收与负载再规划。这一转变使得非均匀压缩首次能够将理论上的精度‑内存优势无损耗地转化为实际吞吐量,实质性地拓展了 LLM 多轮服务的可支撑序列长度与并发规模。

行业影响

落地场景

Tangram 直接面向多轮对话 LLM 服务的高吞吐场景,如客服系统、AI 助手、教育答疑、医疗问诊、金融投顾等。这些产品中,单用户对话越长、并发越高,KV 缓存膨胀越快,极易触发内存墙。借助 Tangram 的非均匀 KV 压缩与静态内存规划,服务端可在不损失回复质量的前提下,将单 GPU 能承载的会话数提升 2.6×,尤其适合需要长上下文理解的文档分析、代码审查、合同解读等场景。

商业价值

核心降本路径是提高 GPU 利用率与吞吐,直接减少每请求的硬件成本。在同等 SLA 下,企业可以用更少的 GPU 实例支撑相同甚至更高的并发量,显著降低云端推理开支。同时,Tangram 无需运行时重规划,消除了因页回收和负载倾斜导致的解码延迟抖动(原系统有 1.7× 延迟膨胀),使得用户体验更稳定,对于面向终端用户的 SaaS 产品,体验提升可直接转化为续费与口碑。

与现有产品/工作流的接口

Tangram 直接以插件形式嵌入 vLLM,与现有 serving stack 无缝衔接。工程上只需在模型部署配置中启用 Tangram 的算子替换,并提供少量离线校准样本(~50 条),即可将任意支持 vLLM 的 Transformer 模型升级为非均匀 KV 压缩的持续批处理服务。它与现有 PagedAttention 兼容,通过预算聚类形成独立页表,不破坏上层调度逻辑。对于已使用 KV 缓存压缩的团队,Tangram 作为统一 substrate,可替换原 uniform 策略并集成更先进的 head-wise 压缩算法,实现更优的精度-吞吐权衡。

具体落地用例

  • 电商智能客服:大促期间客服机器人需同时处理海量咨询,每段对话可能持续数十轮,KV 缓存极易占满 GPU 显存。Tangram 允许为不同注意力头分配非均匀预算,在保证回答准确率相当的情况下,将单 GPU 的并发对话数提升 2–3 倍,降低扩容成本。
  • 医疗影像报告生成:放射科助手需根据连续的多轮描述(患者历史、前片对比)生成结构化报告。使用 Tangram 后,服务端可维持更长的对话窗口而不触发 OOM,同时避免因动态回收页造成的生成延迟尖刺,确保报告生成的实时性与可靠性。

局限

  • **静态预算分配的泛化风险**:Tangram 依赖离线校准获得的输入无关头排序和窄界比率,假设头重要性不随输入剧烈变化。然而,在多领域对话或长尾任务中,注意力头的关键性可能发生偏移,导致预分配的压缩预算与实际需求错配,造成精度折损。论文仅展示少量样本(50)校准后的平均准确度,未系统评估领域迁移或对抗性输入下的鲁棒性,这在实际多租户、开放域服务中可能成为瓶颈。
  • **模型与规模覆盖不足**:评估仅在单个模型系列(如 Llama 类)上进行,未展示跨架构(如非自回归模型、混合专家模型)或更大规模(如 70B+)的适用性。不同模型的头数量、隐藏维度差异可能影响 ragged paging 的聚类效率与内存收益,而论文未讨论相关可扩展性。此外,实验中的压缩方法仅为现有非均匀技术的接入,Tangram 本身不改进压缩率,其吞吐提升高度依赖所选压缩方法的效果,但论文未对比多种压缩策略在 Tangram 上的表现差异。
  • **工程复杂度与生态绑定**:Tangram 深度耦合 vLLM 的 PagedAttention 和 continuous batching 机制,迁移到其他推理框架(如 TensorRT-LLM、LMDeploy)需要重大适配。ragged paging 引入的多粒度页表管理与向量化块表操作虽提高了可回收内存,但也增加了调度器实现复杂度和潜在的缓存未命中问题。对于已有高度优化的推理栈,引入 Tangram 的改造成本可能抵消部分收益,而论文未分析在非 vLLM 环境下的部署开销与兼容性。
论文Hyungmin Kim2026-06-15原文

相关内容