AI 研究助理

基于站内最新资料,回答 AI 前沿问题

回答带出处引用,支持多轮追问,当前页面内容会作为上下文

快速响应
AI 前沿跟踪
今日首页重点内容GitHub 趋势榜
领域
学术研究
大模型
具身智能
AI 基础设施
智能体 · Skill · MCP
机器人
多模态
自动驾驶
数据 / MLOps
AI 硬件
安全对齐
综述与报告
论文AI 硬件KV cache长上下文NAS推理优化

MOSAIC: Multi-Objective Search for Allocating Inference Cache Across LLM Decoder Layers

查看原文(PDF)
AI 解读已生成 · AI 生成,数字以原文为准
AI 全文总结值得深入基于论文全文 3.4 万字 · 完整阅读

MOSAIC 用多目标搜索为每层 KV cache 分配预算,146/154 设置胜 uniform

长上下文LLM推理的KV cache瓶颈下,SnapKV、H2O、AdaKV、L2Norm等逐token淘汰方法给每层相同预算,忽视层间需求差异。MOSAIC用多目标神经架构搜索联合学习每层预算:在无预算约束的幂次网格上做分类器引导贝叶斯优化得到Pareto前沿,再在全量数据上选winner形状并缩放到任意目标预算,可选做预算约束细化。在Llama-3-8B-Instruct与Mistral-7B-Instruct-v0.2上、LongBench各任务类别和RULER检索子任务中,MOSAIC在146/154个(task, method, budget)设置里超过均匀分配,RULER最高提升**+48分、LongBench的LCC最高提升+11.9分,达到均匀分配同分最多省4×**缓存。学习到的分配一致偏向少量中间层,形状更具任务特异性而非方法特异性;相对EvolKV搜索空间大得多且评估更省。局限是仅评测上述稠密解码器模型,Mistral的Summarization部分仅完成Stage 1,未来扩展至MoE、混合架构及KeyDiff、QUOKA。

  1. 1146/154 设置胜 uniform,RULER 最高 +47.9 点
  2. 2同分下 KV cache 最多省 4×,即小 75%
  3. 3预算集中在少数中间层,一次搜索可迁移全预算

为什么重要 不换淘汰算法即可在同等显存下提分,或用最多 1/4 缓存达到原分数,可做成离线预算表适合 做长上下文 LLM 推理服务的工程与性能优化同学

胜率146/154 设置
RULER 提升+2.3 到 +47.9 点
显存节省最多 4× KV cache
搜索成本9–40 GPU-hour
对比 EvolKV5/6 设置更高
来源
arxiv.org
发布
2026/10/20

中文(AI 翻译)

Key-Value(KV)缓存是长上下文 LLM 推理的主要内存瓶颈。SnapKV、H2O、AdaKV 和 L2Norm 等 token 驱逐方法可压缩缓存,但它们给每个 Transformer 层相同预算,而各层需要的 token 并不相同。我们提出 MOSAIC,一种多目标神经架构搜索,可为任意驱逐方法学习逐层 KV 缓存预算。

一次搜索可刻画平均预算与任务分数之间的权衡;其最佳分配无需进一步搜索即可迁移到每个目标预算,可选的预算约束搜索可进一步细化。

在 Llama-3-8B-Instruct 与 Mistral-7B-Instruct-v0.2 上、四种驱逐方法中,MOSAIC 在 154 个(任务、方法、预算)设置中有 146 个优于均匀分配,在 RULER 上最高提升 48 分,在 LongBench(LCC)上提升 11.9 分。在两个模型上,它以最高 4 倍更少的 KV 缓存达到均匀分配的分数(7……(摘要截断)

原文

The Key-Value (KV) cache is the main memory bottleneck of long-context LLM inference.

Token eviction methods such as SnapKV, H2O, AdaKV and L2Norm shrink it, but they give every transformer layer the same budget, although layers differ in which tokens they need.

We present MOSAIC, a multi-objective neural architecture search that learns per-layer KV-cache budgets for any eviction method.

One search maps the trade-off between average budget and task score; its best allocation transfers to every target budget without further search, and an optional budget-constrained search refines it.

Across four eviction methods on Llama-3-8B-Instruct and Mistral-7B-Instruct-v0.2, MOSAIC beats uniform allocation in 146 of 154 (task, method, budget) settings, by up to +48 points on RULER and +11.9 on LongBench (LCC).

On both models, it reaches the score of uniform allocation with up to 4x less KV cache (75% smaller).

Against EvolKV, MOSAIC explores a ~1040 x larger search space with 4x fewer evaluations and 18% fewer GPU-hours per budget, and scores higher in 5 of 6 settings.

The learned allocations consistently favour a small set of middle layers.

Llama-3-8B-Instruct 上 MOSAIC 相对 uniform 的 RULER 各子任务/各预算增益,红色格子表示更高的任务,多数为正,最难的检索子任务增益最大
Llama-3-8B-Instruct 上 LongBench 各数据集按类别分组的增益热图,Code 类全为正,其中 LCC 达到 +11.9
对比项之前 / 基线现在 / 本方案
逐层预算 vs uniform每层同预算146/154 设置更高4 种淘汰方法×2 模型
RULER 分数uniform 分配最高 +47.9 点最低增益 +2.3 点
同分所需 KV cache1×最多 1/4(小 75%)17 个 Llama 设置只需一半或更少
搜索空间EvolKV $\leq 10^{29}$MOSAIC 约大 $10^{40}$ 倍全层联合 vs 8 层分组
搜索开销EvolKV 每组搜索评估次数 4× 更少,GPU 小时少 18%每个预算口径
测试模型
Llama-3-8B-Instruct, Mistral-7B-Instruct-v0.2
层数
Llama-3-8B 为 32 层
每层预算范围
bmin​=64, bmax​=4096 token
搜索评估上限
Stage 1 ≤200 个,Stage 4 ≤150 个
基准
RULER 11 个子任务,LongBench 4 类

深度总结

1 / 6
180 GB
问题与动机· 要点

KV cache 是长上下文推理的主要显存瓶颈:30B 模型、batch 128、序列长度 1024 已需 180 GB KV cache

  • KV cache 是长上下文推理的主要显存瓶颈:30B 模型、batch 128、序列长度 1024 已需 180 GB KV cache
  • SnapKV、H2O、AdaKV、L2Norm 只决定丢掉哪些 token,却给每个 transformer 层相同的预算
  • 层与层需要的 token 不同,uniform 分配把缓存给了不需要的层,导致真正需要更多缓存的层得不到足够资源
  • 已有逐层方案要么预设形状(PyramidKV),要么依赖注意力统计(CAKE),要么分组逐组搜索(EvolKV)