揭秘vLLM：从KV Cache计算到GQA原理的深度之旅

揭秘vLLM：从KV Cache计算到GQA原理的深度之旅

发布日期: 2025年7月11日

如果你正在与大语言模型（LLM）打交道，那么你一定对“显存”这个词深感敬畏。在模型推理过程中，除了模型权重本身，KV-Cache 是另一个巨大的显存消耗者。高效管理KV-Cache是提升LLM服务吞吐量、降低成本的关键。

vLLM 作为当前最流行的LLM推理和服务框架之一，其核心优势就在于通过PagedAttention等技术对KV-Cache进行了极致的优化。但vLLM在启动时，究竟是如何确定需要为KV-Cache分配多少显存的呢？而现代模型又是如何从架构上“节流”，减少KV-Cache大小的呢？

这篇博客将带你进行一次深度之旅，我们将一起：

1. 通过一段核心代码，精确计算一个大模型所需的KV-Cache初始化显存。
2. 深入剖析**GQA（分组查询注意力）**这一关键优化技术，理解它如何“在不牺牲质量的前提下，大幅节省显存”。

让我们开始吧！

Part 1：精确计算KV-Cache初始化显存

在vLLM中，系统会预先分配一块足以容纳模型所支持的**最大序列长度（max_model_len）**的KV-Cache。我们通过分析其核心代码，可以推导出精确的计算过程。

源码逻辑

# vLLM中计算单层最大显存占用的逻辑
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
    max_model_len = vllm_config.model_config.max_model_len
    # 计算需要多少个块，并乘以每个块的大小
    return cdiv(max_model_len, self.block_size) * self.page_size_bytes

其中，每个块（Page）的大小由以下公式决定：

# 计算单个物理块(Page)大小的逻辑
page_size_bytes = 2 * self.block_size * self.num_kv_heads * self.head_size * get_dtype_size(self.dtype)

通过将两个公式结合并进行数学简化，我们可以得出一个更直观的计算公式。

最终计算公式

单层最大显存 (Bytes) = max_model_len × 2 × num_kv_heads × head_size × dtype_size

这个公式清晰地告诉我们，单层的KV-Cache大小由以下几个核心参数决定：

• max_model_len: 模型能处理的最大Token数量。
• 2: 代表 K(Key) 和 V(Value) 两个向量。
• num_kv_heads: K/V头的数量，这是GQA等优化的关键。
• head_size: 每个注意力头的维度。
• dtype_size: 数据类型占用的字节数（如FP16/BF16为2字节）。

实战演练

让我们用之前讨论过的一个具体模型参数来进行计算：

# 模型配置参数
max_model_len = 131072
num_hidden_layers = 28
num_attention_heads = 28
num_kv_heads = 4       # 注意！这里的K/V头数远小于总注意力头数
hidden_size = 3584
dtype_size = 2         # BF16
# ------------------------------------
# 推导参数
head_size = hidden_size / num_attention_heads  # 3584 / 28 = 128

计算步骤：

1. 计算单层总KV-Cache大小：

   单层总大小 = 131072 * 2 * 4 * 128 * 2
                = 268,435,456 字节
                = 256 MB
   

2. 计算模型总KV-Cache大小：

   模型总大小 = 单层总大小 × 总层数
                = 256 MB/层 × 28 层
                = 7168 MB
                = 7.00 GB
   

通过这个计算，我们清晰地得知，vLLM为了保证能处理最长131k的序列，会为这个模型初始化分配整整7.00 GB的显存专门用于KV-Cache。

看到这里你可能会问，为什么num_kv_heads是4而不是28？如果用28来算，结果将是惊人的49GB！这背后就是GQA的功劳。

Part 2：GQA——节省显存的“秘密武器”

GQA（Grouped-Query Attention，分组查询注意力） 是现代LLM架构中一项革命性的优化。要理解它，我们先用一个“专家顾问团”的比喻来解释三种不同的注意力机制。

注意力机制	核心思想	比喻：专家顾问团	优缺点
MHA (多头注意力)	每个Q头都有独立的K/V头	一人一本专属笔记	性能最好，但最慢、最耗显存
MQA (多查询注意力)	所有Q头共享唯一一个K/V头	所有人共用一本笔记	最快、最省显存，但性能损失风险大
GQA (分组查询注意力)	将Q头分组，组内共享K/V头	一个小组共用一本笔记	性能接近MHA，速度和显存接近MQA

GQA完美地解决了MHA成本过高和MQA性能下降的问题，成为了一个黄金般的平衡点。

GQA的核心：共享的K/V，独特的Q

一个关键问题：如果一个小组内的所有注意力头（Heads）都共享相同的K和V，那它们之间还有什么区别？它们如何产生不同的结果？

答案就在于：小组内的每一个Head，都拥有一个独一无二的Q（Query）向量。

让我们回到“专家顾问团”的比喻：

• 共享的K/V向量：这就像一个小组的专家们共同阅读的同一份“背景参考资料”。它定义了需要被关注的“是什么”。
• 独立的Q向量：这就像每个专家被赋予的独特“提问角度”或“分析任务”。它定义了“要关注什么”。

工作流程可以简化为：

1. 共享资料：一个小组的所有专家（Heads）都拿到同一份参考资料（K/V）。
2. 独立提问：每位专家根据自己独特的任务（Q向量），对这份资料提出不同的问题。例如，专家A问“这份资料有何财务风险？”，专家B问“这份资料有何市场机会？”。
3. 得出不同结论：因为提问的角度不同，每位专家从同一份资料中得出的“相关性”和“结论”也完全不同。

总结一下：

组件	在GQA小组内	作用
K (Key) & V (Value)	共享	作为“参考资料”，定义了信息的“内容”
Q (Query)	独立且唯一	作为“提问角度”，定义了信息的“关注点”

GQA的智慧在于，它意识到注意力的多样性主要来源于“提问视角Q”的多样性。通过让成本高昂的“参考资料K/V”被小组共享，它极大地压缩了KV-Cache的大小，同时保留了模型从多角度审视信息的核心能力。

结论

通过这次深度之旅，我们不仅学会了如何根据模型config.json文件精确计算其在vLLM中的KV-Cache显存占用，更重要的是，我们揭开了GQA这一关键优化的神秘面纱。

下一次当你配置LLM服务时，你将能更好地理解这些参数背后的意义：

• 显存计算为你提供了硬件规划的依据。
• GQA原理则让你深刻体会到，正是这些优雅的架构设计，才使得越来越强大的大语言模型能够更高效地运行在我们的硬件之上。