跳到主内容
@wquguru
精选88Avi Chawla技巧与观点

vLLM KV缓存压缩无效原因及NVIDIA TriAttention方案

A tricky LLM interview question:

原文
发到 X
推荐理由

深入剖析了vLLM PagedAttention在KV缓存压缩中的底层陷阱,给出了可落地的优化思路与指标建议,做推理部署的同学必看。

A tricky LLM interview question:

一个棘手的 LLM 面试问题:

You're serving a reasoning model on vLLM, and long traces keep pushing it out of GPU memory.

你在 vLLM 上部署了一个推理模型,而长长的推理轨迹不断将其挤出 GPU 内存。

You add KV cache compression and evict 90% of the cached tokens.

你添加了 KV 缓存压缩,并驱逐了 90% 的已缓存 token。

VRAM usage doesn't move, and the GPU still runs out of memory.

显存(VRAM)使用量没有变化,GPU 仍然内存不足。

Why?

为什么?

(answer below)

(答案见下文)

Evicting 90% of the KV cache can free close to zero bytes on a paged server.

在基于分页(paged)的服务器上,驱逐 90% 的 KV 缓存可能几乎释放不了任何字节。

During decoding, every new token appends its key and value vectors at every layer. vLLM holds all of them until the request finishes.

在解码过程中,每个新 token 都会在每一层追加其键(key)和值(value)向量。vLLM 会一直保留它们,直到请求结束。

For reasoning models with long chains of thought, this cache is the main memory cost.

对于具有长思维链的推理模型,这个缓存是主要的内存开销。

A Qwen3-32B with 4-bit weights on a 24GB GPU needs KV for every token of a 32K-token trace, and it hits OOM around 24K tokens.

在 24GB 的 GPU 上运行权重为 4-bit 的 Qwen3-32B,需要为 32K token 轨迹中的每个 token 分配 KV 缓存,它在大约 24K token 处就会触发 OOM(内存溢出)。

Attention over long contexts is sparse, so you can keep the high-importance tokens and drop the rest with little accuracy loss.

对长上下文的注意力机制是稀疏的,因此你可以保留高重要性的 token 并丢弃其余部分,而精度损失很小。

That reduces the number of cached tokens but it doesn't necessarily reduce allocated memory.

这会减少缓存 token 的数量,但并不一定会减少已分配的内存。

vLLM manages the cache with PagedAttention. It splits GPU memory into fixed physical blocks of 16 token slots each.

vLLM 通过 PagedAttention 管理缓存。它将 GPU 内存划分为固定的物理块,每块包含 16 个 token 槽位。

The allocator reclaims a block only when all 16 slots are empty.

分配器仅当所有 16 个槽位都为空时才会回收该块。

Importance-based eviction keeps tokens scattered across the full sequence, so the retained tokens are spread over nearly every block.

基于重要性的驱逐会使 token 分散在整个序列中,因此保留下来的 token 分布在几乎所有的块中。

Evict 14K of 16K tokens across 1,000 blocks, and almost every block still holds at least one token.

如果在 1,000 个块中驱逐了 16K 个 token 中的 14K,那么几乎每个块仍然至少持有一个 token。

So the allocator reclaims close to nothing.

因此,分配器几乎无法回收任何东西。

You could write new tokens into the empty slots. If token 16,001 goes into the slot token 40 held, the block now reads positions 38, 16,001, 41.

你可以将新 token 写入空槽位。如果 token 16,001 被放入 token 40 所在的槽位,该块现在读取的位置将是 38、16,001、41。

The cache is no longer in position order. Attention still gives the correct result, but each slot now needs a stored position index, and every kernel has to read it.

缓存不再按位置顺序排列。注意力机制仍能给出正确结果,但每个槽位现在都需要存储一个位置索引,且每个内核都必须读取它。

An in-order layout doesn't need that metadata.

有序布局不需要这些元数据。

So the cache shrinks 90% logically and 0% physically. Many compression papers miss this because they benchmark on pre-allocated contiguous tensors instead of a paged server.

因此,缓存逻辑上缩小了 90%,但物理上缩小了 0%。许多压缩论文忽略了这一点,因为它们在预分配的连续张量上进行基准测试,而不是在分页服务器上。

The second problem is the eviction signal.

第二个问题是驱逐信号。

Most eviction methods rank tokens by their attention scores.

大多数驱逐方法根据注意力分数对 token 进行排名。

FlashAttention never materializes those scores. It computes attention tile by tile in on-chip SRAM and discards each score tile after use, which is where its speed comes from.

FlashAttention 从不物化这些分数。它以片(tile)为单位在片上 SRAM 中计算注意力,并在每次使用后丢弃每个分数片,这正是其速度的来源。

To get the scores, you have to fall back to eager attention and build the full score matrix in HBM, which removes the speedup FlashAttention provides.

要获取分数,你必须退回到急切注意力机制,并在 HBM 中构建完整的分数矩阵,这会消除 FlashAttention 提供的加速效果。

NVIDIA's TriAttention handles both problems.

NVIDIA 的 TriAttention 解决了这两个问题。

It ranks tokens without attention scores. It uses the geometry of the query and key vectors before RoPE is applied, where they form stable clusters.

它在没有注意力分数的情况下对 token 进行排序。它使用应用 RoPE 之前查询向量和键向量的几何结构,此时它们形成稳定的聚类。

Every 128 decoded tokens, it runs a compaction pass. Retained tokens move forward to fill the evicted slots, so trailing blocks empty out completely and return to the allocator, and the cache stays in position order.

每解码 128 个 token,它就会执行一次压缩传递。保留的 token 向前移动以填充被驱逐的槽位,因此尾部块完全清空并返回给分配器,且缓存保持位置顺序。

On long reasoning traces, it matches full-attention accuracy with 2.5x faster decoding and 10.7x less KV memory.

在长推理轨迹上,它以 2.5 倍的解码速度和 10.7 倍的 KV 内存节省匹配了全注意力的准确性。

For KV cache compression on a real server, the metric to track is blocks freed, not tokens evicted.

对于真实服务器上的 KV 缓存压缩,需要跟踪的指标是释放的块数,而不是被驱逐的 token 数。

NVIDIA's write-up: https://research.nvidia.com/labs/eai/blogs/kv-cache-compression-and-its-infra-problems/

NVIDIA 的文章:https://research.nvidia.com/labs/eai/blogs/kv-cache-compression-and-its-infra-problems/

I also wrote a first-principles breakdown of the KV cache. It covers why the model caches keys and values, why the cache grows with every token, and how generation speed compares with and without caching.

我还撰写了一篇关于 KV 缓存的第一性原理分解文章。它涵盖了模型为何缓存键和值、缓存为何随每个 token 增长,以及有无缓存时的生成速度对比。

Read it below.

请在下方阅读。

更进一步:量化金融体系

看懂新闻只是起点——沿量化金融路径,把它变成能交付的工程能力

进入量化体系 →

相似阅读

关联信息,但可能不是同一事件