KV Cache 笔记

题图

When voice from the memories call
当旧忆传来声声呼唤
Unresolved the riddles stood ever long
昔日的谜犹然未解
Remember the colours of evergreen
莫忘那恒碧长青
The cherished answer bound in us all
正是你我珍存于心的解答

—— 寻诀

一块只读张量总算混成了对象

题图:pixiv 147576297

引子

时值transformers v5把Cache拆成逐层对象、连续批处理调度器搬进主库之际,KV cache已经从”推理小技巧”变成显存账单上的主角——参数可以不涨,上下文要涨,一个128K的请求能把70B的权重比下去。回望2023年的llama2.c,cache是三行指针算术加一整块malloc出来的扁平数组,那就是它的全部定义;今天它是list[CacheLayerMixin],同一样东西的两种活法。管cache是很有必要的。

本文以transformers 5.16.1的源码为基线(所有行号来自本机/usr/lib/python3.14/site-packages/transformers,装机版不pin commit,所以只保证这版对得上),从”为什么只缓存K和V”讲到”缓存怎么分配、怎么驱逐、怎么量化、怎么分页”。顺手把两件事摆在一起看:GQA/MLA/滑窗/CSA/线性状态这些结构各自要什么形状的cache,eager/sdpa/flash/flex这些kernel各自怎么读cache。作为读代码比读论文习惯的我,本文不谈训练细节也不谈采样策略,只谈这一堆张量的生命周期。

基线:从embedding走到KV cache

先把一条完整通路摆出来,标清楚cache挂在哪一步。以llama2-7B为例(dd=4096、HH=32、HkvH_{kv}=32、dhd_h=128、LL=32):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
input_ids [B,S]                        整数,词表下标
↓ Embedding 查表 (32000 → 4096)
x [B,S,4096] 层入口
↓ RMSNorm
├─ q = x·Wq [B,S,4096] → view+transpose → Q [B,32,S,128]
├─ k = x·Wk [B,S,4096] → → K [B,32,S,128] ← 头数是 H_kv
├─ v = x·Wv [B,S,4096] → → V [B,32,S,128]
↓ RoPE(Q, K, pos) 只旋转 Q 与 K,V 不转
★ 写入 cache:update(K, V, layer_idx) → K,V 长成 [B,32,S_kv,128]
↓ Attention(Q, K_cache, V_cache) → [B,32,S,128] → 拼回 + o_proj → [B,S,4096]
⊕ 残差
↓ RMSNorm → SwiGLU FFN (4096 → 11008 → 4096) → ⊕ 残差
↓ × 32 层,每层各自一份 cache
↓ RMSNorm → lm_head (4096 → 32000)
logits [B,S,32000] 取最后一行采样

三个观察:

  1. cache的挂点非常具体:在RoPE之后、attention kernel之前,而且只挂在K/V两条支路上——Q那条线直接进kernel,不经过任何存储
  2. 整条管线里Embedding是查表(不产生随序列增长的中间量),FFN是逐token的matmul(算完即弃),残差流的hidden每层重算——唯一”当前token要用到历史全部token产物”的地方就是K和V。KV cache不是”一种优化”,它是这个结构里唯一可缓存的东西
  3. 每层一份、互不相干,所以LL是显存公式的乘数:32层就是32份cache,浅层和深层的长度永远一样(跨层共享是后面的事)

投影矩阵的形状直接决定cache的形状:

WqRd×Hdh,Wk,WvRd×HkvdhW^q\in\mathbb{R}^{d\times H d_h},\qquad W^k,W^v\in\mathbb{R}^{d\times H_{kv} d_h}

其中 d=4096d=4096 为hidden、H=32H=32 为Q头数、HkvH_{kv} 为KV头数、dh=128d_h=128 为头维。llama2-7B是MHA所以 Hkv=HH_{kv}=H;70B把 HkvH_{kv} 压到8,投影矩阵变窄、cache同步变小——这是”结构决定cache形状”的第一次现身,全部展开在结构那一节。

Attention这一格展开。上面管线里那个Attention(Q, K_cache, V_cache)到底算了什么,单头写全:

Attn(Q,K,V)=softmax ⁣(QKdh+M)V\text{Attn}(Q,K,V)=\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_h}}+M\right)V

其中 Q,K,VRS×dhQ,K,V\in\mathbb{R}^{S\times d_h}单个头、单个样本的投影结果(批量与头在外面循环);M{0,}S×SM\in\{0,-\infty\}^{S\times S} 为因果掩膜,Mij=0M_{ij}=0jij\le i、否则 -\infty1/dh1/\sqrt{d_h} 是方差补偿:若 q,kq,k 分量近独立等方差,Var(qk)=dhσ4\mathrm{Var}(q^\top k)=d_h\sigma^4,不除就要把 softmax 推进饱和区。

分头之后到拼回去,一共五步:

1
2
3
4
5
6
Q [B,H,S,d_h]   K [B,H,S,d_h]   V [B,H,S,d_h]
① 打分 z = QKᵀ/√d_h + M → [B,H,S,S] ← 唯一的 S² 项
② 归一化 p = softmax(z, dim=-1) → [B,H,S,S] ← 每行和为 1
③ 汇总 o = pV → [B,H,S,d_h]
④ 拼头 transpose(1,2)+reshape → [B,S,H·d_h] = [B,S,4096]
⑤ 投影 o = o·W^o → [B,S,4096] ← 唯一跨头混合的一步

其中 ①②③ 逐头完全独立(这既是”多头”的并行性来源,也是 cache 可以按头切第2维的理由),只有 ④⑤ 把 32 个头拼起来混合;llama2-7B 满足 Hdh=d=4096H\cdot d_h=d=4096,所以 o_proj 是个方阵。逐元素看:

pij=exp(zij)miexp(zim),oi=jipijvjp_{ij}=\frac{\exp(z_{ij})}{\sum_{m\le i}\exp(z_{im})},\qquad o_i=\sum_{j\le i}p_{ij}v_j

其中 pijp_{ij} 为第 ii 个 token 分给第 jj 个历史 token 的权重,每行和为 1 —— 所以 oio_i 是历史 value 的凸组合,量级跟 VV 同阶。这个”每行和为1”就是后面 sink、温度、稀疏化能撬动一切的支点:多开一列不参与取值的”垃圾桶”不往输出里加任何东西,但它进分母,整行的概率就被压下去。

三个工程细节:

  1. ②处必先减行最大值。前篇 llama2.c 里那段expf(x[i] - max_val)就是这个,数学不变、fp32 不溢出;transformers 的 eager 写成softmax(..., dtype=torch.float32).to(query.dtype)是同一个动机的另一半
  2. ①处的MM是加性而非布尔[B,1,S,S]广播到所有头,加进 zz 后由 softmax 自动把未来位置压成0——所以 eager 里那个 S×SS\times S 掩膜是实打实的内存,而 kernel 版宁可传is_causal=True也不开它
  3. 每步的成本不同阶。①②③ 是 O(S2dh)O(S^2 d_h)/头,④⑤ 是 O(Sd2)O(Sd^2);prefill 两者都要付。到 decode,QQ 只有一行,①③ 降成 O(Sdh)O(S d_h)、⑤ 仍是 O(d2)O(d^2) —— 单步看是权重主导,但权重是”每层一次、全 batch 共享”,而 O(Sdh)O(S d_h) 那份要每层每请求各读一份:这就是上一节”写 O(1)O(1)、读 O(S)O(S)“的出处,也是并发一高 cache 就先撞显存的原因

对齐一下挂点:上面①③两步要的 K,VK,V 全体历史行,就是update返回的东西;而①里出现的 QQ 只用当前那一行,用完即丢。

再看两阶段。同一条管线,差别只在入口的 SS

1
2
prefill:  q[1..S] × K[1..S]ᵀ → 下三角 → 一次算完 S 行
decode : q[S+1] × K[1..S+1]ᵀ → 只多 1 行 → 长度 S+1 的向量
输入 Q cache里的K/V 输出
prefill input_ids [B,S] [B,H,S,d_h] [B,H,S,d_h] logits [B,S,32000]
decode input_ids [B,1] [B,H,1,d_h] [B,H,S+1,d_h] logits [B,1,32000]

其中decode行只有Q的那一维塌成1,KV那一维反而在涨——q_len与kv_len不对称是decode的全部特征,sdpa里那句is_causal = q_length > 1 and ...读的就是它。

换成逐 token 的增量写法(decode 每步的活就这么多):

zj(t)=qtkjdk,ot=jtexp(zj(t))mtexp(zm(t))vjz^{(t)}_j=\frac{q_t^\top k_j}{\sqrt{d_k}},\qquad o_t=\sum_{j\le t}\frac{\exp(z^{(t)}_j)}{\sum_{m\le t}\exp(z^{(t)}_m)}\,v_j

其中 qt,kt,vtRdhq_t,k_t,v_t\in\mathbb{R}^{d_h} 为单头向量(就是上面管线里 Q/K/V 取第 tt 行、第 hh 个头);z(t)Rtz^{(t)}\in\mathbb{R}^{t} 为第 tt 行的 logits;分母 mt\sum_{m\le t} 逐行不同。

能不能缓存,三条判据:

  1. causal mask 让历史行不受新 token 影响——第 tt 行只依赖 qtq_tktk_{\le t},前 t1t-1 行的结果永远不变
  2. K,VK,V 会被未来所有行反复读 → 缓存;QQ 用完即弃 → 缓存它是纯浪费
  3. 中间量 softmax()V\text{softmax}(\cdot)V 不能缓存,因为归一化分母每行都要重算;要省的是乘数,不是概率矩阵

前篇llama2.c笔记推过一遍,那里cache是一整块扁平数组,按层偏移:

1
2
3
4
// llama2.c:254-259
int loff = l * p->seq_len * kv_dim; // kv cache layer offset for convenience
s->k = s->key_cache + loff + pos * kv_dim;
s->v = s->value_cache + loff + pos * kv_dim;

pos * kv_dim就是”第pos条KV”,写入即前进指针。整个模型只有一个cache、形状在启动时定死;v5里它是每层一个对象。

然后是这本账:

KV/token=2LHkvdhb\text{KV}_{/\text{token}} = 2\cdot L\cdot H_{kv}\cdot d_h\cdot b

其中 LL 为注意力层数,HkvH_{kv} 为 KV 头数,dhd_h 为头维,bb 为每元素字节数(bf16=2),系数 2 是 K 与 V 两份——四个因子在上面那张图里全能指出来:2 是x·Wkx·Wv两条支路,LL 是”每层各自一份”,HkvdhH_{kv}d_h 就是K/V投影的宽度,bb 由dtype决定。

于是不对称就摆在这儿了:decode每走一步,按这个式子往每层追加一条KV,写是 O(1)O(1);而attention要把已经追加过的全部读一遍,读是 O(S)O(S)

LL HkvH_{kv} dhd_h 每token 上下文 满上下文cache
llama2-7B (MHA) 32 32 128 512KB 4096 2.0GB
llama2-70B (GQA) 80 8 128 320KB 4096 1.25GB
Qwen3.5-MoE (仅10层全注意力) 10 2 256 20KB 256K 5.4GB
DeepSeek-V4-Flash (K==V+压缩) 43 1 512 7.5KB 1M 7.9GB

(后两行沿用《架构变迁笔记》的维度自算,V4只算KV、indexer与gate另计;7.5KB的来处在下面「CSA/HCA」一节末尾。70B若把上下文从4096扩到128K,单个请求就是40GB——一条对话吃掉三张卡)

访存账接着算:每生成1 token要把全部权重读一遍、再把全部cache读一遍。7B的bf16权重13.4GB、cache 512KB/token,两者在26K token处交叉,此后每步读cache的字节数比读权重还多。并发更直白:一张24GB的卡装完7B剩约10GB,4K上下文每条2.0GB,只能塞5条并发

权重是定值,cache随序列和batch双向涨——长上下文时代,显存的主人是cache不是权重。

cache在一次forward里的三种身份

主角就是下面这些(models/llama/modeling_llama.py:252-277,其中两处长参数列表被合并成行):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

cos, sin = position_embeddings
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

if past_key_values is not None:
key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) # :262

attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
self.config._attn_implementation, eager_attention_forward # :264
)
attn_output, attn_weights = attention_interface(
self, query_states, key_states, value_states, attention_mask,
dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs)

三种身份:

  • 写入update返回的是合并后的全体K/V,不是增量。attention那一侧根本看不到”cache”这个概念,它只看到key_states变长了
  • 读出:kernel读返回值的dim=-2。shape约定[B, H_kv, S, d_h]——序列在第3维,头在第2维
  • 定位get_seq_length()供mask与cache_position反推位置(:387),位置信息的真相从”外部计数器”迁到了cache本身

支点只有一句话:接口只有update(k, v, layer_idx) -> (k, v)一个方法,所以换cache实现不需要动模型代码。「结构」一节的每一行、「容器」一节的每一个类,全挂在这个签名上。

注意RoPE发生在update之前——这是硬性契约,坑清单第1条会回到这里。

常见attention结构:cache的形状由谁决定

transformers里有一张官方映射表,一个config字段决定整条显存曲线(cache_utils.py:1223-1244):

1
2
3
4
5
6
7
8
9
10
11
12
13
DYNAMIC_LAYER_TYPE_MAPPING = {
"full_attention": DynamicLayer,
"sliding_attention": DynamicSlidingWindowLayer,
"chunked_attention": DynamicSlidingWindowLayer, # 对cache而言滑窗=分块,只差mask
"conv": LinearAttentionLayer, # 不存per-token KV,只存conv state
"linear_attention": LinearAttentionLayer,
"hybrid": LinearAttentionAndFullAttentionLayer,
"hybrid_sliding": LinearAttentionAndSlidingWindowAttentionLayer,
"deepseek_sparse_attention": DynamicIndexedLayer, # CSA + indexer
"qwen_sparse_attention": DynamicIndexedLayer,
"moe": LinearAttentionLayer,
"mlp": LinearAttentionLayer, # 不缓存,占空槽
}

DynamicCache(config=...)一行调用它:get_layer_types_and_kwargs(:1701)从config.layer_types逐层dispatch,layers = [DYNAMIC_LAYER_TYPE_MAPPING[t](**kw) for t in layer_types](:1793)。子类可以自带注册——_layer_type配合__init_subclass__(:38-42),在建模文件里写个类就进表,DeepSeek-V4就是这么插队的(见 CSA/HCA 一节)。

MHA → MQA → GQA:头数是cache的除数

KV/token=2LHkvdhb,Hkv={HqMHA1MQA介于两者GQA\text{KV}_{/\text{token}}=2LH_{kv}d_hb,\qquad H_{kv}=\begin{cases}H_q & \text{MHA}\\ 1 & \text{MQA}\\ \text{介于两者} & \text{GQA}\end{cases}

repeat_kv(modeling_llama.py:179-189):

1
2
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

它发生在update之后、kernel之内(eager_attention_forward:202-203)。所以cache里永远只存HkvH_{kv}份,HqH_q是读时展开出来的幻觉。sdpa更进一步,enable_gqa=True(sdpa_attention.py:102)连展开都交给kernel。llama2-70B的64/8就是1/81/8,这一刀砍的是分母。

跨层共享KV:层数也是除数

1
2
3
4
# cache_utils.py:1717-1720
num_kv_shared_layers = getattr(config, "num_kv_shared_layers", None)
if num_kv_shared_layers is not None and num_kv_shared_layers > 0:
layer_types = layer_types[: -config.num_kv_shared_layers]

Gemma3n这类模型尾部若干层不建cache、复用前层KV(CLA/YOCO一路)。对cache而言”少几层”和”少几个头”是同一个动作——都是给分母做除法。

MLA:缓存的不是K和V,是潜变量

KV/tokenMLA=L(dc+dr)b,KV/tokenMHA=LHq(dkn+dr+dv)b\text{KV}^{\text{MLA}}_{/\text{token}}=L\,(d_c+d_r)\,b,\qquad \text{KV}^{\text{MHA}}_{/\text{token}}=L\,H_q\,(d_{kn}+d_r+d_v)\,b

其中 dc=512d_c=512 为共享压缩潜变量(kv_lora_rank)、dr=64d_r=64 为带旋转的位置键(qk_rope_head_dim)、dkn=128d_{kn}=128dv=128d_v=128。注意MLA前面没有系数2——K与V同源,一份存储两用。V3取(512+64)×61×2B=70(512{+}64)\times61\times2\text{B}=70KB/token,等价MHA是128×320×61×2B=4.9128\times320\times61\times2\text{B}=4.9MB/token,差70倍。

实现(models/deepseek_v3/modeling_deepseek_v3.py:456-471,注释是原文):

1
2
3
4
5
6
7
8
9
10
11
compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
kv_nope, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
kv_nope = self.kv_a_layernorm(kv_nope)
# Both latents are viewed as single-head, 4D tensors so all cache layers handle them correctly
kv_nope = kv_nope.view(batch_size, 1, seq_length, self.kv_lora_rank)
k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)
...
# Cache read / write is performed while latent KV is still compressed
if past_key_values is not None:
kv_nope, k_rot = past_key_values.update(kv_nope, k_rot, self.layer_idx)
key_states, value_states = self.expand_kv(kv_nope, k_rot)

两个槽位被复用成”潜变量 + 带位置的头”,DynamicLayer一行没改就接住了MLA:它只要求dim=-2增长、第2维是头数(这里=1)。展开发生在cache之后(expand_kv:419-437),而且k_rot.expand(-1, k_nope.shape[1], -1, -1)注释里明说”does not affect the underlying storage”。

DeepSeek-V4-Flash更狠,K和V在cache层就是同一块内存(下面 CSA/HCA 一节里那行 self.values = self.keys)。另外V4那个”进分母即丢”的可学习sink是s_aux(flash_attention.py:37),走kernel参数,不占cache一格。

滑窗与分块:窗口即上限

1
2
3
4
5
6
7
8
# cache_utils.py:244-258 (DynamicSlidingWindowLayer.update)
self.cumulative_length += key_states.shape[-2]
full_key_states = torch.cat([self.keys, key_states], dim=-2)
full_value_states = torch.cat([self.values, value_states], dim=-2)
if not self.record_past:
self.keys = full_key_states[:, :, -self.sliding_window + 1 :, :]
self.values = full_value_states[:, :, -self.sliding_window + 1 :, :]
return full_key_states, full_value_states

两个要点:W1W-1而不是WW(新token自己占第WW个位置);返回full、缓存窗口——kernel看到的仍然是”当前行可见的全部”,驱逐对它透明。窗口左边界靠get_mask_sizes(:263)报给mask:

KV整条SWA=2Lmin(S,W1)Hkvdhb\text{KV}^{\text{SWA}}_{\text{整条}}=2L\min(S,\,W{-}1)H_{kv}d_hb

其中 WW 为滑窗大小(config 里的sliding_window:Gemma2 4096、Gemma3n 512、GPT-OSS 128,mistral 的默认值已经涨到 4096)。注意这次是整条序列的量而不是每token:一到 W1W-1 就不再随 SS 增长,摊到每token是2LHkvdhbmin(1,W1S)2LH_{kv}d_hb\cdot\min(1,\frac{W-1}{S}),序列越长越薄。代价是远处全盲,所以实践里滑窗总要配少数全注意力层——也就是 CSA/HCA 与线性注意力两节。

CSA / HCA:把历史压成条目

V4的两个新层类型靠_layer_type自动进表(models/deepseek_v4/modeling_deepseek_v4.py:185,265),HCA直接继承滑窗层:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class DeepseekV4HCACache(DynamicSlidingWindowLayer):     # :162
_layer_type = "heavily_compressed_attention" # :185
# __init__ 里:self.compress_rate = config.compress_rates["heavily_compressed_attention"] # :189

def update(self, key_states, value_states, *args, **kwargs): # :195(此处略去类型注解)
"""Shared sliding-window K=V update body. V4 uses shared-KV MQA, so `keys`
and `values` point to the same storage on every layer."""
if not self.is_initialized:
self.lazy_initialization(key_states, value_states)
self.values = self.keys
self.cumulative_length += key_states.shape[-2]
full = torch.cat([self.keys, key_states], dim=-2)
self.keys = full[:, :, -self.sliding_window + 1 :, :]
self.values = self.keys # ← 系数2在这里真的消失了
return full, full

窗口之外还挂三个dict(:190-193):buffer_kv/buffer_gate攒不够一个窗的零头、compressed_kv是已发射的条目、entry_countcompress_rate就是绝对位置——压缩条目的RoPE打在这个窗位上(rope_layer_type="compress":412)。store_compression_weights(:209)只剥最长窗对齐前缀,剩下的留在buffer里等下一次forward;CSA比HCA多一路"indexer"(:265-272)和交叠窗的overlap_kv/overlap_gate(:275-290,只存上一窗的Ca片,Cb折完即弃)。

KV/token=bd1r\text{KV}_{/\text{token}}=b\,d\sum_{\ell\in\text{层}}\frac{1}{r_\ell}

其中 d=512d=512 为该层KV维、b=2b=2 为bf16字节数、r{1,4,128}r_\ell\in\{1,4,128\} 为逐层压缩率;前面没有系数2,因为V4是K==V——上面那行self.values = self.keys把两份合成了一份。

V4-Flash的43层是2滑窗+21CSA+20HCA,1/r2+21/4+20/128=7.36\sum_\ell 1/r_\ell\approx 2+21/4+20/128=7.36,乘bd=1024bd=10247.5KB/token;不压缩的话是43×1024=4443\times1024=44KB,压缩这一刀接近6倍,滑窗层还要再受WW截断。「基线」那张表里的数字就是这么来的。

一条路线:把cache从”逐token”变成”逐条目”,检索交给indexer(msa_attention.py:230)按block打分选top-k。

线性注意力:把cache换成状态

LinearAttentionCacheLayerMixin(:891)与LinearAttentionLayer(:1003)提供另一对写入方法:update_conv_state(:1037)、update_recurrent_state(:1077),多状态时由config.number_of_conv_statesnumber_of_states(:1731)。

State=Hvdkdvb+(k1)dqkvb\text{State}=H_v\,d_k\,d_v\,b+(k{-}1)\,d_{qkv}\,b

其中第一项是记忆矩阵 StS_t、第二项是kernel=kk的因果卷积缓冲;与序列长 SS 无关,get_max_length()对这类返回-1(含义是”不适用”,不是”无限”)。

hybrid = LinearAttentionAndFullAttentionLayer(:1094):一个对象里同时挂状态和KV。前篇《架构变迁笔记》里Qwen3.5-MoE”40层只有10层产生KV”,在cache侧的真身就是这张表里hybridfull_attention按3:1交替。

彩蛋:"moe""mlp"层也被塞成LinearAttentionLayer(:1242-1243),只为get_seq_length()之类能从正确的层取到值——它们永远是空的。

谁在用:结构↔模型对照

前面所有维度都要回到config里对,那就顺手普查一次:本机models/*/configuration_*.py496个,只统计类属性里写死的默认值(不下权重)。注意llama、qwen2/3主干这类默认是None、由checkpoint的config.json决定,不在计数里。显式同时给了num_attention_headsnum_key_value_heads的101个里:GQA 72、MHA 28、MQA 1

结构 代表模型(以本机config默认为准) 认场字段
MHA gemma一代(16/16)、stablelm(32/32)、chameleon、helium、dia、moonshine_streaming H_kv == H
MQA deepseek_v4(64/1);paligemma、pi0 写在子config里的也是num_key_value_heads=1 H_kv == 1
GQA mistral/mixtral(32/8)、gemma2/3(8/4)、gemma3n(8/2)、qwen3_moe(32/4)、qwen3_next与qwen3_5_moe(16/2)、glm4(32/2)、glm4_moe(96/8)、hy_v3与gpt_oss(64/8)、jamba(32/8)、minimax_m2(48/8)、phimoe(32/8)、solar_open(64/8)、starcoder2(24/2)、inkling(64/8) H_kv < H
MLA(K==V低秩) deepseek_v2/v3/v32、minicpm3、longcat_flash、glm4_moe_lite、glm_moe_dsa、glm5_next、mistral4、youtu、axk1/axk2(共12家) kv_lora_rank
滑窗 mistral、gemma2/3、gemma3n、gpt_oss(128)、cohere2、modernbert、granite_swa、olmo3、ministral、phi3、mimo_v2_flash(sliding_window非None的44家) sliding_window
分块 llama4(attention_chunk_size=8192)、gemma4 attention_chunk_size
KV跳层共享 gemma3n(num_kv_shared_layers=15) num_kv_shared_layers
线性/hybrid mamba/mamba2/falcon_mamba、recurrent_gemma、bamba/zamba/zamba2、nemotron_h、qwen3_next、qwen3_5与qwen3_5_moe、olmo_hybrid、granitemoehybrid、lfm2、minimax、inkling、zaya(linear_attention 17家、hybrid 15家) layer_types
indexer/稀疏 deepseek_v32、deepseek_v4、qwen4_exp(qwen_sparse_attention)、minimax_m3_vl、glm_moe_dsa、glm5_next、axk2、cohere2_moe、exaone_moe layer_types + indexer
CSA/HCA压缩 deepseek_v4(default_compress_rates = {CSA:4, HCA:128} compress_rates
attention sink gpt_oss、granite_swa、mimo_v2_flash、openai_privacy_filter、deepseek_v4(8家) s_aux

三条观察:

  1. GQA是压倒性的默认。显式写了KV头数的config里GQA占72/101;真正”老实MHA”的只剩gemma一代、stablelm、chameleon这批2023前的骨干,而MHA那28个里一大半其实是MLA家族(deepseek_v3/v32、glm4_moe_lite、glm_moe_dsa、glm5_next、axk1/2、youtu)——它们的num_key_value_heads已经失去意义。类属性里写死MQA的只剩一个:deepseek_v4的64/1,而且它同时叠了K==V与压缩,跟当年MQA论文”共享一头换速度”已经不是一个路线
  2. 别只盯num_key_value_heads。MLA那12家在config里仍写heads == kv_heads(deepseek_v32是128/128),看上去像MHA,实际cache里存的是kv_lora_rank那份潜变量。判结构要kv_lora_rankcompress_rateslayer_types一起看
  3. 纯结构几乎绝迹,混合才是常态sliding_window非None 44家、layer_types里有linear_attention 17家、hybrid 15家、带indexer的稀疏8家、带sink的8家;nemotron_h的默认层型直接就是["linear_attention", "moe", "full_attention", "mlp"]的四元循环。同一个模型里每层cache长得不一样——这就是v5把cache做成逐层对象的直接理由(回到「契约」那节:Cache = list[layer]

想自己复核:grep -l "kv_lora_rank" .../transformers/models/*/configuration_*.py,再拿各家的layer_types默认值一对就行。

结构×cache主表

结构 每token KV config标签 Cache类
MHA 2LHqdhb2LH_qd_hb full_attention DynamicLayer
MQA 2Ldhb2Ld_hb 同上,Hkv=1H_{kv}=1 DynamicLayer
GQA 2LHkvdhb2LH_{kv}d_hb 同上 DynamicLayer
跨层共享 2(LLs)Hkvdhb2(L-L_s)H_{kv}d_hb num_kv_shared_layers 少建若干层
MLA(K==V) L(dc+dr)bL(d_c{+}d_r)b full_attention DynamicLayer(存潜变量)
滑窗/分块 2LHkvdhbmin(1,W1S)2LH_{kv}d_hb\min(1,\frac{W{-}1}{S}) sliding_/chunked_attention DynamicSlidingWindowLayer
CSA/HCA bd1/rbd\sum_\ell 1/r_\ell(K==V) deepseek_sparse_attention DynamicIndexedLayer/DeepseekV4CSACache
线性注意力 (Hvdkdvb+(k1)dqkvb)/S\big(H_vd_kd_vb+(k{-}1)d_{qkv}b\big)/S(整块常数) linear_attention/hybrid LinearAttentionLayer/混血
MoE/MLP层 0 moe/mlp LinearAttentionLayer(空槽)

代回基线那本账:512KB → 320KB(GQA) → 70KB(MLA) → 窗口截断 → 除以rr(压缩) → 0(线性)。layer_types一个列表就画出了一个模型的显存曲线。

常见attention kernel:eager / sdpa / flash / flex

上一节谈的是存什么,这一节谈怎么读。四个实现共享同一句结论:要dense、要连续、要全量。

eager:老老实实materialize S×S

23行,全引(modeling_llama.py:191-214):

1
2
3
4
5
6
7
8
9
10
11
12
13
key_states = repeat_kv(key, module.num_key_value_groups)
value_states = repeat_kv(value, module.num_key_value_groups)

attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
if attention_mask is not None:
attn_weights = attn_weights + attention_mask

attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
attn_output = torch.matmul(attn_weights, value_states)
attn_output = attn_output.transpose(1, 2).contiguous()

return attn_output, attn_weights

唯一返回attn_weights的实现,output_attentions=True只认它,所以它是debug用。掩膜就是基线说过的那种加性浮点数(create_causal_mask:864/eager_mask:538),不再是bool。

代价:8K序列、32头的中间量是81922×32×4B8.68192^2\times32\times4\text{B}\approx8.6GB/层——序列8K时attention矩阵比整个7B权重还大。

sdpa:fused kernel与两个flag

sdpa_attention_forward(sdpa_attention.py:79-170)92行里真正影响cache的只有两处:

1
2
3
4
5
6
7
8
if hasattr(module, "num_key_value_groups") and module.num_key_value_groups > 1:
if not use_gqa_in_sdpa(attention_mask, key, value): # :98-100
key = repeat_kv(key, module.num_key_value_groups)
value = repeat_kv(value, module.num_key_value_groups)
else:
sdpa_kwargs = {"enable_gqa": True} # :102
...
is_causal = q_length > 1 and attention_mask is None and is_causal # :124

decode阶段qlen=1q_{len}=1is_causal恒False → “不看未来”完全靠cache天然只到当前长度,mask一行都不给。顺序不能反,注释(:117-119)说先判shape是为了不把is_causal变成SymBool

第三个坑在:139-148,注释原话:sdpa的causal是upper-left对齐,而带cache偏移的滑窗需要bottom-right对齐;prefill撞上空StaticCache时(kv_length > q_length)它宁可手动key[:, :, :q_length, :]切一刀,因为flash kernel只在q_length == kv_length时才启用。

后端dispatch(flash / mem-efficient / cudnn)不看你的眼色,看shape和mask。masking_utils.py:496allow_is_causal_skip_ignore_causal_mask_sdpa(:235)——全True的mask直接返回None,省掉一个S×SS\times S

flash:不传mask,传页边界

flash_attention_forward(flash_attention.py:26-103):

1
2
3
4
5
6
7
8
seq_len = query.shape[2]
query = query.transpose(1, 2) # FA 要 [B,S,H,d]
key = key.transpose(1, 2)
value = value.transpose(1, 2)

head_dim, v_head_dim = query.shape[-1], value.shape[-1]
if v_head_dim != head_dim:
value = torch.nn.functional.pad(value, [0, head_dim - v_head_dim]) # :62-64

补零那三行就是给MLA准备的——V3的dv=128d_v=128dqk=192d_{qk}=192,FA要求K/V同头维。参数表里还有sliding_windowsoftcaps_aux(注释alias: learnable attention sink:34-37)。

padding不走mask:FA收cu_seqlens(varlen,modeling_flash_attention_utils.py里unpad→pack),padding被物理删掉而不是用-inf遮住:40-44明确拒绝output_attentions:49-54拒绝零长度维度并建议”use SDPA instead”。

和cache的关系:kernel内部tile化流水,仍然要顺序读全量K/V。基线那笔”写 O(1)O(1)、读 O(S)O(S)“的账对FA同样成立——它省的是S2S^2中间量的显存,不是cache的字节。

flex与稀疏族:mask从张量降维成函数

flex_attention_forward(flex_attention.py:262)的attention_mask位置可以放两样东西:仍是张量,或者一个BlockMask——后者由create_block_mask(mask_mod=...)(:238)造,mask_mod的签名就是(b,h,q,kv)->bool(:185注释原话)。张量那条路也不白走:它被折进score_mod(:290),与softcapposition_bias一起当逐分数变换。masking_utils.py则负责把组合子吐成mask_modcausal_mask_function(:76)、sliding_window_overlay(:92)、chunked_overlay(:104)、packed_sequence_mask_function(:182)、and_masks/or_masks(:48:62)。

两个硬约束顺带记下:flex_attention不容 dropoutdropout>0直接raise(:274-278,提示只用于 inference);而 attention sink 在score_mod里做不对——注释说得很明白(:297-299),sink 要归一化前的整行,所以得排在 flex 调用之后。

msa_attention_forward(msa_attention.py:230)是稀疏族另一半:indexer按block打分选topk_blocks(:155-159),没有indexer的层(视觉塔、无indexer的全注意力层)直接回落SDPA(:247)。前篇里DeepSeek那个”给压缩条目打分挑top-512”的Lightning Indexer,在HF侧就是这套DynamicIndexedLayer+msa kernel的组合。

一句话:eager把mask存进内存,sdpa尽量不给mask,flex把mask变成闭包

对比表与注册契约

实现 入口 要4D mask materialize S2S^2 能看权重 cache布局要求
eager eager_attention_forward dense
sdpa sdpa_attention_forward 尽量免 dense、连续
fa2/3/4 flash_attention_forward 免(cu_seqlens) dense / varlen
flex flex_attention_forward 免(mask函数) dense
paged 上述三者的paged变体 block页表(见 paged 一节)

契约(modeling_utils.py:5139-5174):签名(module, q, k, v, mask, *, dropout, scaling, **kwargs) -> (out, weights),挂上去就生效:

1
2
3
4
5
6
7
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS   # 顶层未导出,模型文件都从这拿

def my_attn(module, query, key, value, attention_mask, *, dropout, scaling, **kwargs):
...
return out.transpose(1, 2), None

ALL_ATTENTION_FUNCTIONS.register("my_attn", my_attn) # config._attn_implementation = "my_attn"

get_interface(:5161)对未注册的键名直接KeyError,"eager"是永远的下限(它同时是default实参)。

结论句:这四者对cache的要求完全一致——dense、连续、全量、dim=-2增长。唯一的例外不把cache当张量,它把它拆成了页,见 paged 那一节。

Cache容器家族

cache_utils.py2131行的类图:

1
2
3
4
5
6
7
8
9
10
11
12
13
CacheLayerMixin (:27)      update / get_mask_sizes / get_seq_length / get_max_length
├ DynamicLayer (:113) cat(dim=-2)
│ ├ DynamicSlidingWindowLayer (:203) 只留 W-1
│ ├ DynamicIndexedLayer (:319) 多挂一路 indexer key
│ └ QuantizedLayer (:703) ─ Quanto(:779) / HQQ(:834)
├ StaticLayer (:398) 预分配 + index_copy_ + mark_static_address
│ ├ StaticSlidingWindowLayer (:504) / StaticIndexedLayer (:631)
└ LinearAttentionCacheLayerMixin (:891) ─ LinearAttentionLayer (:1003)
└ LinearAttentionAnd* (:1094 / :1134 / :1166 / :1198)

Cache (:1269) = list[layer] + offload/prefetch
├ DynamicCache(:1737) StaticCache(:1829) QuantizedCache(:1884)
└ EncoderDecoderCache(:1947) MtpCache(:2102) DFlashCache(:2114)

契约:五个抽象加四个默认

抽象的5个:lazy_initialization(:53)、update(:56)、get_mask_sizes(:61)、get_seq_length(:64)、get_max_length(:67);默认可重写的4个:offload(:75)、prefetch(:81)、reset(:87)、reorder_cache(:100)。

reset()zero_()而不是置None——多轮对话复用同一个对象;reorder_cache就是index_select(0, beam_idx),beam search的全部逻辑。

Dynamic:每步一次全量拷贝

1
2
3
# cache_utils.py:144-145
self.keys = torch.cat([self.keys, key_states], dim=-2)
self.values = torch.cat([self.values, value_states], dim=-2)

一行cat的代价:长度变了→新分配+整块复制,SS步累计O(S2)O(S^2)写带宽;shape是动态量→torch.compile必然graph break。它是对照组,不是生产选项。

Static:为compile而生

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# :424-451 节选  一次开满,并把地址钉住
self.keys = torch.zeros(
(self.batch_size, self.num_heads, self.max_cache_len, self.k_head_dim),
dtype=self.dtype,
device=self.device,
)
...
if not is_torchdynamo_compiling():
torch._dynamo.mark_static_address(self.keys)
torch._dynamo.mark_static_address(self.values)
torch._dynamo.mark_static_address(self.cumulative_length)

# :478-481 就地写
cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length
self.cumulative_length.add_(kv_length) # 必须 in-place
self.keys.index_copy_(2, cache_position, key_states)

mark_static_address保住指针,就地改写才不会让cudagraph作废——注释原话”Without this, we cannot use cudagraphs!!”。cumulative_length是tensor而非int,也是同一个理由(:414注释)。

代价:空slot也占显存、max_cache_len要提前给(generation/configuration_utils.py:180,设成预期最长值以免反复重分配+重编译)、prefill不能进编译区(chunked prefill另说)。

Quantized:拿读带宽换显存

QuantizedCache(:1884,docstring直引KIVI论文)参数nbits=4q_group_size=64residual_length=128QuantizedLayer.update(:731-765)的循环:

1
2
3
4
5
6
dequant_keys   = self._dequantize(self._quantized_keys)
keys_to_return = torch.cat([dequant_keys, self.keys, key_states], dim=-2)
...
if self.keys.dim() == 4 and self.keys.shape[-2] + 1 >= self.residual_length: # :757 攒满就压
self._quantized_keys = self._quantize(keys_to_return.contiguous(), axis=self.axis_key)
self.keys = torch.tensor([], dtype=key_states.dtype, device=key_states.device)

双轨存储:最近128条保持原精度(残差窗),溢出后整批下推到量化区,K与V各自按q_group_size分组(与KIVI论文的差别在docstring里写明了)。

坏消息必须写清楚:每层每步都要dequant一遍,decode本就是访存bound,4bit省下的字节数被dequant的临时张量吃掉一截——省的是显存(=更多并发),通常不是时延。硬约束:QuantizedCache只接受纯full_attention模型(:1934-1938直接raise),滑窗/线性/压缩层一律不支持。

Offloading:搬运工与两条stream

Cache(offloading=True),核心在Cache.update(:1379-1386):

1
2
3
4
5
6
torch.cuda.default_stream(key_states.device).wait_stream(self.prefetch_stream)
self.prefetch(layer_idx + 1, self.only_non_sliding) # :1381 先预取下一层
keys, values = self.layers[layer_idx].update( # :1383
key_states, value_states, *args, **kwargs)
if self.offloading:
self.offload(layer_idx, self.only_non_sliding) # 用完这层立刻卸载

独立stream预取、找不到可搬层就绕回开头(:1335-1339的ValueError分支),形成一条循环流水线。only_non_sliding默认只搬非滑窗层——注释说得很直白:滑窗层本来就小,搬它反而亏。

用户侧入口是generation_config.cache_implementation(generation/configuration_utils.py:169-173,实体在:47-48):dynamic/static/offloaded/offloaded_static/quantized。注意offloaded_hybridoffloaded_hybrid_chunked这两个旧名字已经进了DEPRECATED_STATIC_CACHE_IMPLEMENTATIONS(:50-56),只为老配置兼容;而paged根本不在那个列表里,是校验时单独加上的(:685ALL_CACHE_IMPLEMENTATIONS + ("paged",))。

生命周期:谁在造、谁在改、谁在删

造:谁实例化Cache

1
2
3
4
5
# modeling_llama.py:383-387
if use_cache and past_key_values is None:
past_key_values = DynamicCache(config=self.config)
...
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0

生成侧_prepare_cache_for_generation(generation/utils.py:1929):用户自己传的cache与generation_config.cache_implementation互斥(:1948-1953);传tuple已经明确不支持(:1954-1957,”Please use a Cache instance”)——v4时代past_key_values: tuple[tuple[Tensor]]的老接口在这里正式埋掉。cache的名字还会按模型改(:1943-1944,mamba系叫cache_params)。

paged:cache的第三种世界观

注册表就是目录(modeling_utils.py:5147-5160):

1
2
3
4
5
"paged|flash_attention_4": paged_attention_forward,
"paged|flash_attention_3": paged_attention_forward,
"paged|flash_attention_2": paged_attention_forward,
"paged|sdpa": sdpa_attention_paged_forward,
"paged|eager": eager_paged_attention_forward,

键名用|拼——paged不是一种新算法,是三种算法共用的新布局。读侧(integrations/sdpa_paged.py:28-40):

1
2
3
4
5
6
cache: PagedAttentionCache | None = kwargs.pop("cache", None)
if cache is not None:
key, value = cache.update(key_states=key, value_states=value,
layer_idx=module.layer_idx,
read_index=kwargs["read_index"], write_index=kwargs["write_index"])
key = key.transpose(0, 1).unsqueeze(0)

read_index/write_index就是页表:物理块不必连续,读时按索引gather成逻辑连续。eager版(eager_paged.py:28-34)注释写得更清楚,shape从[1, H_kv, S_kv, d]拍平成[-1, H_kv, d](页拼成长条)。注意Cache.update(k,v,layer_idx)那个优雅协议在这里是不成立的——所以HF另起了一套:PagedAttentionCache(continuous_batching/cache.py:89,474行)不继承Cache。两套cache体系并存的原因就是这一句。

真正的价值不在”少分配显存”,在共享(continuous_batching/cache_manager.py):

  • BlockManager(:58)持_free_block_ids/_id_to_blockget_free_blocks(n, last_block_id, shareable, group_id)(:112)分配时顺手把块串成父子链
  • compute_hash(parent_hash, tokens, group_id)(:279):块指纹=父块指纹+自己的token+层组号。前缀相同的请求哈希必然相同 → search_prefix_match(cache.py:470)直接命中,第二个请求的prefill大半白嫖
  • fork_blocks(:131)的docstring画了两种情形,4块fork给2个子请求:不共享时新分配8块全复制,共享时只新分配2块(最后一个不完整块),前面完整块ref_count += num_forks
  • TP场景不能用内建hash()(每进程加盐不同,:282-288),换hashlib.blake2b(digest_size=8)
  • PagedAttentionMemoryHandler(cache.py:565):activation_peak(:621)估激活峰值→infer_max_batch_tokens_and_num_blocks(:652)反推块数——先量显存再定并发,而不是先开并发再OOM

beam search、投机解码的多分支、”两个用户的prompt前2K token一模一样”,在这套里是同一个动作:引用计数+不完整块复制。

改与删:crop、reset、reorder

crop(:165-186,负数=从尾部删N个,正数是deprecated的”截到绝对长度”)、reset(:87)、reorder_cache(:100)、batch_repeat_interleave(:190)/batch_select_indices(:196)。

滑窗层的crop有额外脾气:越过窗口之后只接受负数,且要先调activate_past_recording()(:218),否则直接raise(:286-296)。投机解码回退、多轮对话裁剪都走这条路。

手写40行:把上面压回一个能跑的最小实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import torch

class KV: # DynamicLayer 的全部语义
def __init__(self, L): self.k, self.v = [None]*L, [None]*L
def update(self, k, v, i):
self.k[i] = k if self.k[i] is None else torch.cat([self.k[i], k], dim=-2)
self.v[i] = v if self.v[i] is None else torch.cat([self.v[i], v], dim=-2)
return self.k[i], self.v[i]

def eager(q, k, v, mask): # 即 eager 那节的原样
w = q @ k.transpose(-1, -2) * q.size(-1) ** -0.5
return (w + mask).softmax(-1) @ v

def forward(x, W, cache, l, rope): # x: [Sq,d],prefill 时 Sq=S,decode 时 Sq=1
q, k, v = ((x @ w).view(-1, H, dh).transpose(0, 1) for w in W[:3]) # [H,Sq,dh]
q, k = rope(q), rope(k) # 契约:写 cache 之前旋转
k, v = cache.update(k[None], v[None], l) # ← 全篇唯一与缓存有关的一行,[1,H,Skv,dh]
Sq, Skv = q.size(-2), k.size(-2) # Skv = 旧的 + 新的
m = (torch.arange(Skv) > torch.arange(Sq)[:, None] + Skv - Sq).to(x.dtype) * -torch.inf
return eager(q[None], k, v, m) @ W[3] # bottom-right 对齐:decode 时 m 全空

# 数值对齐:一次性 prefill 与逐 token decode 必须逐位相等
c1, c2 = KV(1), KV(1)
full = forward(X, W, c1, 0, rope)[:, -1] # Sq = S,m 是下三角
for t in range(S):
step = forward(X[t:t+1], W, c2, 0, rope)[:, -1] # Sq = 1,m 全空——is_causal 空转
assert torch.allclose(full, step, atol=1e-4)
print("bytes/token =", 2 * L * Hkv * dh * 2)

对上那一下才算真懂:update那一行换成Static、滑窗、量化、paged中的任意一个,其余代码一个字都不用改——结构、kernel、容器三节的所有内容都在这一行的两端。

坑清单

  1. RoPE必须在update之前(modeling_llama.py:259-262的顺序不是偶然)。cache里的K已经带位置,decode时再按position_ids给旧K转一次=位置叠加。V4更绕:压缩条目的RoPE打在窗位entry_count*compress_rate上(modeling_deepseek_v4.py:412)。
  2. repeat_kv必须在cache之后。放前面就是往cache里存HqH_q份复制品,基线那本账直接乘Hq/HkvH_q/H_{kv}。sdpa用enable_gqa连这步都省。
  3. Dynamic cache与torch.compile互斥。cat改shape→graph break;要compile就Static+max_cache_len,那三个mark_static_address(:449-451)就是为这件事存在的。
  4. decode阶段is_causal是空转(sdpa_attention.py:124),qlen=1q_{len}=1恒False,”未来”根本不在cache里。此时还传4D mask是白给一个S×SS\times S
  5. mask对齐:sdpa的causal是upper-left,带cache偏移的滑窗需要bottom-right:139-148注释专门写了这段,两者不等价。让masking_utils去算kv_offset,别自己手搓mask。
  6. cache量化不一定更快。每层dequant一次,decode本就访存bound;而且QuantizedCache只支持纯full_attention(:1934)。先想清楚省的是显存还是时延。
  7. output_attentions=True会让sdpa/flash告警或降级(sdpa_attention.py:92-95flash_attention.py:40-44)。debug完记得关掉。

总结

cache 增长方式 上限 谁在用 代价
DynamicLayer cat,每步变长 默认 compile不友好、O(S2)O(S^2)写带宽
StaticLayer 就地index_copy_ max_cache_len 生产/compile 空slot占显存
DynamicSlidingWindowLayer W1W-1 WW Mistral/Gemma2/V4滑窗层 远处全盲
QuantizedLayer 残差窗+分组量化 位宽 长上下文低并发 dequant时延、只支持full
PagedAttentionCache 按块分配+引用计数 物理块数 serving Cache协议不通、要read/write_index
LinearAttentionLayer 不增长 固定状态 Mamba/GDN/hybrid 精确检索能力为零

三条趋势:

  1. cache从技巧变成对象:v4时代它是past_key_values的tuple,v5里它是每层一个可offload、可量化、可分页、可reset的CacheLayer,还带layer_types自动dispatch
  2. kernel与cache互相驯化:flash要varlen、sdpa要连续、compile要static、serving要paged——选哪种attention实现就是给cache选一种存储;反过来self.values = self.keys这一行也让FA的补零逻辑有了存在理由
  3. 主线是”少缓存”:MHA→GQA→MLA→滑窗驱逐→CSA/HCA压缩→线性状态,六级台阶全在把O(S)O(S)O(1)O(1)压——这正是《架构变迁笔记》那张三代表格里”KV cache”一行的来处

凯尔希。凯尔希。
多漂亮的签名啊。
她在哥伦比亚和萨尔贡之间斡旋。
她把王酋和雇佣兵都玩弄在股掌之间。
而我们,甚至都看不清她行过的道路有多长。
您说真当冬天来临的时候,
维多利亚的哪一棵栎树能免于风雪?
仍愿乌萨斯遗忘你,大公阁下。

—— 遗尘漫步 · 明日方舟

Reference

  • 前篇:llama2.c笔记(其中kvcache一节,扁平数组版)
  • 同系列:从llama2到Qwen3.5-MoE,再到DeepSeek-V4-Flash——架构变迁笔记
  • transformers 5.16.1(本机/usr/lib/python3.14/site-packages/transformers):cache_utils.pymodels/llama/modeling_llama.pymodels/deepseek_v3/modeling_deepseek_v3.pymodels/deepseek_v4/modeling_deepseek_v4.pyintegrations/{sdpa,flash,flex}_attention.pyintegrations/{eager,sdpa,flash}_paged.pymasking_utils.pygeneration/continuous_batching/{cache,cache_manager}.py
  • KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache — cache_utils.py:1887 docstring 引的
  • llama2.c: github.com/karpathy/llama2.ckey_cache/value_cache单块扁平数组)
  • FlashAttention-2 / vLLM PagedAttention / Gated DeltaNet