大模型推理中 KV Cache为什么没有Q Cache

大模型推理中 KV Cache为什么没有Q Cache

flyfish

为什么没有Q Cache？因为Q是一次性的当前查询，只在当前步使用，而KV是持久的历史上下文，会被后续所有步骤复用。如果到这里已经明白了，下面的就不用看了。

KV Cache设计理念基于自回归解码的特性，通过复用历史K/V向量，将推理复杂度从 O(n^2) 降低到 O(n)，实现了模型的推理加速。既然可以缓存K和V向量，为什么不能同样缓存Query（Q）向量呢？

本文从Transformer注意力机制、自回归解码的工作原理来解析KV Cache设计的底层逻辑，并回答为什么没有Q Cache。

一、Transformer注意力机制回顾

1.1 注意力机制的核心公式

Transformer模型的核心是自注意力（Self-Attention）机制，其数学表达式为：

Attention

(

Q

,

K

,

V

)

=

softmax

(

Q

K

T

d

k

)

V

\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V

Attention(Q,K,V)=softmax(dk​

​QKT​)V

其中： Q（Query）：查询向量，代表当前位置需要获取什么样的信息 K（Key）：键向量，代表每个位置能够提供的信息标识 V（Value）：值向量，代表每个位置实际存储的信息内容

d

k

d_k

dk​：向量维度，用于缩放以防止梯度消失

这个机制的工作方式类似于信息检索系统：当需要查询某个位置的信息时，Q向量会与所有位置的K向量进行相似度计算，得到注意力权重，然后用这些权重对V向量进行加权求和，得到最终的输出。

1.2 Q、K、V的来源与关系

在Transformer的每一层中，输入向量首先通过三个独立的线性变换矩阵生成Q、K、V：

Q

=

X

W

Q

,

K

=

X

W

K

,

V

=

X

W

V

Q = XW_Q,\\quad K = XW_K,\\quad V = XW_V

Q=XWQ​,K=XWK​,V=XWV​

其中

W

Q

W_Q

WQ​、

W

K

W_K

WK​、

W

V

W_V

WV​是可训练的权重矩阵，

X

X

X是上一层的输出（或输入嵌入）。

这里有一个关键点需要理解：Q、K、V都来自同一个输入X的不同线性变换。这意味着在一次前向传播中，对于同一个输入序列，我们同时计算了所有位置上的Q、K、V向量。

二、Prefill与Decode阶段对引用文章的回顾

Prefill阶段：生成、写入初始的KV Cache Decode阶段：复用、追加KV Cache 大模型推理的两个阶段：Prefill（预填充）和 Decode（解码）

2.1 推理的两个阶段概述

大语言模型的推理过程分为两个截然不同的阶段：

Prefill阶段（预填充阶段）：这个阶段处理用户输入的完整prompt，生成初始的KV Cache。在这一阶段，模型一次性处理整个输入序列，计算每个token的Q、K、V向量，并将K和V缓存起来供后续使用。

Decode阶段（解码阶段）：这是自回归生成阶段，每次生成一个新的token。在每个解码步骤中，模型只需要计算新token的Q向量，然后与缓存的K、V向量进行注意力计算，从而避免重复计算历史token的K和V。

2.2 两个阶段的计算复杂度分析

在没有KV Cache的情况下，自回归解码的复杂度分析如下：

假设输入序列长度为

s

s

s，输出序列长度为

n

n

n：

Prefill阶段：处理

s

s

s个token，计算复杂度为

O

(

s

2

⋅

d

)

O(s^2 \\cdot d)

O(s2⋅d)，其中

d

d

d是模型维度 Decode阶段：每个新token都需要与所有历史token计算注意力，复杂度为

O

(

n

⋅

(

s

+

n

)

⋅

d

)

O(n \\cdot (s+n) \\cdot d)

O(n⋅(s+n)⋅d)

使用KV Cache后： Prefill阶段：同样需要

O

(

s

2

⋅

d

)

O(s^2 \\cdot d)

O(s2⋅d)计算，但存储了所有历史token的K、V Decode阶段：每个新token只需要计算自己的Q，然后与缓存的K、V进行注意力计算，复杂度降为

O

(

(

s

+

n

)

⋅

d

)

O((s+n) \\cdot d)

O((s+n)⋅d)

这就是KV Cache将复杂度从

O

(

n

2

)

O(n^2)

O(n2)降低到

O

(

n

)

O(n)

O(n)的数学原理。

三、为什么没有Q Cache

3.1 因果掩码（Causal Mask）带来的变化

要理解为什么只需要缓存KV而不需要缓存Q，我们需要深入理解因果掩码（Causal Mask）的作用。

在自回归语言模型中，为了保证生成的自回归性质（即生成第t个token时只能看到位置1到t-1的信息），我们需要在注意力计算中应用下三角掩码：

MaskedAttention

(

Q

,

K

,

V

)

=

softmax

(

Q

K

T

d

k

+

M

)

V

\\text{MaskedAttention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}} + M\\right)V

MaskedAttention(Q,K,V)=softmax(dk​

​QKT​+M)V

其中

M

M

M是掩码矩阵，

M

i

j

=

−

∞

M_{ij} = -\\infty

Mij​=−∞（如果

i

<

j

i < j

i<j），否则

M

i

j

=

0

M_{ij} = 0

Mij​=0。

这个掩码的物理意义是：每个位置只能关注其前面的位置，不能关注后面的位置。这导致了注意力矩阵下三角化的结构。

3.2 Q向量的一次性使用特性

在Decode阶段，每个新生成的token只需要计算一次自己的Q向量。具体流程如下：

这里的关键发现是：在每个解码步骤中，Q向量是位置相关的、时变的。第t步的Q向量

Q

t

Q_t

Qt​与第t+1步的Q向量

Q

t

+

1

Q_{t+1}

Qt+1​完全不同，因为它们对应不同的位置。

3.3 K、V向量的可复用特性

与Q向量不同，K和V向量具有时间不变性：

这正是因为因果掩码的存在：历史token不需要重新计算它们的K和V，因为它们的K和V的定义是固定的，只与它们自身的位置和内容有关。

3.4 数学公式的解释

在解码第t个token时：

不使用缓存：

Attention

(

Q

t

,

K

1

:

t

,

V

1

:

t

)

=

softmax

(

Q

t

K

1

:

t

T

d

k

)

V

1

:

t

\\text{Attention}(Q_t, K_{1:t}, V_{1:t}) = \\text{softmax}\\left(\\frac{Q_t K_{1:t}^T}{\\sqrt{d_k}}\\right) V_{1:t}

Attention(Qt​,K1:t​,V1:t​)=softmax(dk​

​Qt​K1:tT​​)V1:t​

使用KV Cache： 缓存了

K

1

:

t

−

1

K_{1:t-1}

K1:t−1​和

V

1

:

t

−

1

V_{1:t-1}

V1:t−1​ 只需要计算

Q

t

Q_t

Qt​和

K

t

,

V

t

K_t, V_t

Kt​,Vt​（如果还需要处理新的输入token） 然后执行注意力计算

如果尝试缓存Q： 缓存

Q

1

:

t

−

1

Q_{1:t-1}

Q1:t−1​意味着什么？ 当我们要生成第t个token时，

Q

t

Q_t

Qt​是全新的、从未计算过的 缓存

Q

1

:

t

−

1

Q_{1:t-1}

Q1:t−1​对计算

Q

t

Q_t

Qt​没有任何帮助，因为

Q

t

Q_t

Qt​必须重新计算