Attention Sink 与 Streaming KV Cache

它是什么，不是什么

Attention Sink 现象指自回归模型常把显著 attention mass 分配给少量早期 tokens，即使这些 tokens 未必有强语义内容。StreamingLLM 的工程做法是：当固定窗口向前滚动时，保留少量初始 sink tokens，同时保留最近 tokens，而不是纯粹丢弃最旧位置。

它位于 attention 可见性和 KV Cache生命周期层，不是：

• RoPE scaling：后者改变 position→angle schedule，不选择哪些历史可见；
• PagedAttention：后者映射逻辑/物理 KV blocks，不定义 sink mask；
• Prefix Cache：sink 是同一流式请求内长期保留的早期状态，不是请求间相同 prefix 的共享命中；
• 完整长上下文：被滚出 recent 区且不在 sink 中的内容已经不可直接访问。

可见集合与容量定义

令缓存总容量为 $C$，sink token 数为 $S<C$，recent 容量 $W=C-S$。对 Query 位置 $i$，本页定义：

有些 API 把 window_length 定义为总容量 $C$，另一些文字会把 window 指 recent 部分 $W$。部署必须核对字段语义，避免容量多算/少算 $S$。

$C=6,S=2,T=10$ 的 worked example

query i    visible logical positions
0          [0]
1          [0,1]
2          [0,1,2]
3          [0,1,2,3]
4          [0,1,2,3,4]
5          [0,1,2,3,4,5]
6          [0,1,3,4,5,6]
7          [0,1,4,5,6,7]
8          [0,1,5,6,7,8]
9          [0,1,6,7,8,9]

位置 2～5 在 $i=9$ 时已不可见。所以 incremental cache 与“sink+recent sparse mask”可以等价，却不会与保留 [0..9] 的 full causal attention 等价。

在 tensor parallel 下，Hkv/head_dim 的本地 shard 规则依模型布局而定；每个 rank 必须保留相同 logical position set。Paged allocator 可让 sink blocks pinned、recent blocks 循环驱逐，但 sharing/refcount 和 block rounding 会让实际容量高于纯 payload。

RoPE：absolute positions 与 local rerotation 是两种实现

若 cache 永远保存按真实 absolute position 旋转的 Keys，新 Query 也用真实 position，那么 recent token 即使物理 slot 回绕，也不应修改 K；page table/ring metadata 保留 logical positions 即可。

另一类有限窗口实现把每次送入 attention 的 local positions 重映射到 [0,C-1]，使窗口滑动后 recent token 从位置 $p$ 变为 $p-1$。若缓存的是已应用 RoPE 的 $R_p k$，必须重旋转：

对二维旋转块，$R_aR_b=R_{a+b}$，所以不必恢复 raw K 也能 shift。Values 没有 RoPE，通常只搬移/索引；sink Keys 留在原 local positions，不随 recent 区 shift。

可运行的 mask、增量缓存与 rerotation 实验

实验问题：sink ∪ recent 的增量 K/V selection 是否与同一 sparse mask 对齐？它是否真的不同于 full causal？已旋转 K 从 local position [3,4,5] shift 到 [2,3,4] 时，相对旋转是否等价于从 raw K 重算？

import math,platform,torch
torch.manual_seed(229)
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE=torch.float64

def rotate_half(x):
    return torch.stack((-x[...,1::2],x[...,0::2]),-1).flatten(-2)

def rope(x,positions,base=10_000.):
    d=x.shape[-1]
    inv=base**(-torch.arange(0,d,2,device=x.device,dtype=x.dtype)/d)
    angles=positions.to(x.dtype)[...,None]*inv
    cos=torch.repeat_interleave(angles.cos(),2,-1)
    sin=torch.repeat_interleave(angles.sin(),2,-1)
    return x*cos+rotate_half(x)*sin

def visible_indices(position,capacity,sink_tokens):
    if position<capacity: return list(range(position+1))
    recent=capacity-sink_tokens
    start=max(sink_tokens,position-recent+1)
    return list(range(sink_tokens))+list(range(start,position+1))

def dense_sink(q,k,v,capacity,sink_tokens):
    t=q.shape[-2]; scores=q@k.transpose(-2,-1)/math.sqrt(q.shape[-1])
    visible=torch.zeros(t,t,device=q.device,dtype=torch.bool)
    for position in range(t): visible[position,visible_indices(position,capacity,sink_tokens)]=True
    return torch.softmax(scores.masked_fill(~visible,-torch.inf),-1)@v,visible

def incremental_sink(q,k,v,capacity,sink_tokens):
    outputs=[]; cache_positions=[]
    for position in range(q.shape[-2]):
        indices=visible_indices(position,capacity,sink_tokens)
        index=torch.tensor(indices,device=q.device)
        kc=k.index_select(-2,index); vc=v.index_select(-2,index)
        scores=q[...,position:position+1,:]@kc.transpose(-2,-1)/math.sqrt(q.shape[-1])
        outputs.append(torch.softmax(scores,-1)@vc); cache_positions.append(indices)
    return torch.cat(outputs,-2),cache_positions

def full_causal(q,k,v):
    t=q.shape[-2]; scores=q@k.transpose(-2,-1)/math.sqrt(q.shape[-1])
    mask=torch.tril(torch.ones(t,t,device=q.device,dtype=torch.bool))
    return torch.softmax(scores.masked_fill(~mask,-torch.inf),-1)@v

def main():
    b,h,t,d=1,2,10,8; capacity,sinks=6,2
    positions=torch.arange(t,device=DEVICE)
    qraw=torch.randn(b,h,t,d,device=DEVICE,dtype=DTYPE)
    kraw=torch.randn_like(qraw); v=torch.randn_like(qraw)
    q=rope(qraw,positions); k=rope(kraw,positions)
    dense,visible=dense_sink(q,k,v,capacity,sinks)
    incremental,cache_positions=incremental_sink(q,k,v,capacity,sinks)
    torch.testing.assert_close(incremental,dense,rtol=1e-12,atol=1e-12)
    assert visible.sum(-1).tolist()==[1,2,3,4,5,6,6,6,6,6]
    assert cache_positions[-1]==[0,1,6,7,8,9]
    full=full_causal(q,k,v)
    sparse_vs_full=(dense[...,capacity:,:]-full[...,capacity:,:]).abs().max()
    assert sparse_vs_full>1e-3

    raw=torch.randn(3,d,device=DEVICE,dtype=DTYPE)
    old_pos=torch.tensor([3,4,5],device=DEVICE); new_pos=old_pos-1
    old_rotated=rope(raw,old_pos)
    rerotated=rope(old_rotated,new_pos-old_pos)
    recomputed=rope(raw,new_pos)
    torch.testing.assert_close(rerotated,recomputed,rtol=1e-12,atol=1e-12)
    query=rope(torch.randn(d,device=DEVICE,dtype=DTYPE),torch.tensor(5,device=DEVICE))
    wrong=query@old_rotated.T/math.sqrt(d); correct=query@rerotated.T/math.sqrt(d)
    shift_diff=(correct-wrong).abs().max(); assert shift_diff>1e-3
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={DTYPE} Q=K=V={(b,h,t,d)}")
    print(f"capacity={capacity} sink_tokens={sinks} recent_capacity={capacity-sinks}")
    print(f"visible_keys_per_query={visible.sum(-1).tolist()}")
    print(f"final_cache_logical_positions={cache_positions[-1]}")
    print(f"incremental_vs_dense_sink_max_abs_diff={(incremental-dense).abs().max():.3e}")
    print(f"sink_window_vs_full_causal_tail_max_abs_diff={sparse_vs_full:.6f}")
    print(f"rerotation_vs_recompute_max_abs_diff={(rerotated-recomputed).abs().max():.3e}")
    print(f"missing_rerotation_score_max_abs_diff={shift_diff:.6f}")
    print(f"full_KV_elements={2*b*h*t*d} capped_sink_KV_elements={2*b*h*capacity*d}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/attention_sink.py。实际 incremental 与 dense sink mask 最大差 4.441e-16；sink-window 与 full causal tail 最大差 1.237495。rerotation 与从 raw K 重算最大差 2.220e-16；若 local positions 已 shift 却继续使用旧 K rotations，score 最大差为 0.119146。

实验使用 MHA、FP64 和 logical gather，只证明可见集合与旋转群不变量；它没有训练模型、attention sink mass 统计或生产 kernel 性能。真实 GQA/MLA、partial rotary、低精度 cos/sin、paged blocks 与 fused decode kernel 都需按实际实现验证。

容量、读写与质量成本

若 $C$ 是每层总上限，标准 MHA/GQA cache payload 封顶：

相对 full cache $2LBTH_{kv}Ds$，当 $T\gg C$ 时 payload/read source 从 $T$ 封顶到 $C$；sink 不额外超过 $C$，只是占用其中 $S$ 个 slots，使 recent history 只有 $C-S$。若某实现把 recent window $W$ 和 sinks 分开配置，总上限才是 $S+W$。

收益成立还需要 runtime 真的驱逐旧 blocks、decode kernel 只读取 sink+recent、position contract 正确，并且任务容忍丢失中段内容。成本/风险包括：

• sink blocks 长期 pinned，可能降低 block allocator 的自由度；
• recent 驱逐、page-table/ring updates 与可选 rerotation 增加地址/写入工作；
• 非连续 sink + recent source 需要 gather 或能理解两段/block table 的 kernel；fallback 物化 concat 会增加 IO；
• $C$ 小会减少 KV/attention 读取，却缩短可恢复历史；长对话事实、代码上下文和 document QA 可能失效；
• sink 数过大挤压 recent 区，过小又可能无法维持稳定性；没有跨模型通用常数；
• prefix sharing、beam/tree tentative branches、offload 或 P/D transfer 仍需独立 refcount/ownership。

失败场景与常见误解

• 把 full-attention 模型部署时直接换成 sink cache，并声称输出语义不变；
• 只保留 BOS token，不验证该模型/layer/head 是否存在 sink 行为；
• window_length 总容量与 recent window 两种 API 口径混淆；
• logical position 随物理 ring slot 回绕，却既没保持 absolute RoPE，也没 rerotate；
• 对 YaRN/dynamic NTK 使用按标准 RoPE 推导的 shift table，却未核对 frequency schedule 是否固定；
• KV 元素数下降，但 backend 仍 gather/concat 或读 full cache，延迟没有兑现；
• 生成“仍流畅”就宣称记住了被驱逐内容，没有做 retrieval/copy/task accuracy 与不同丢失位置评测。

与上下游词条的连接

• Sliding Window Attention只保留 recent band；Attention Sink 在同一容量内固定一组全局早期 tokens，visible mask 不同。
• KV Cache定义跨 decode step 的 K/V 状态；本页给出一种有损但有上界的保留/驱逐 policy。
• RoPE、PI、NTK-aware 与 YaRN定义 position→rotation；是否保持 absolute positions 或做 local rerotation，必须匹配该 schedule。
• PagedAttention可把 sink blocks pinned，并滚动 recent block table；分页解决物理管理，不证明 sink 质量。
• FlashAttention可让 kernel 跳过不可见 tiles 或读取两段 KV；online softmax 仍需覆盖全部 sink+recent scores。

参考资料

• Xiao et al., Efficient Streaming Language Models with Attention Sinks
• Hugging Face Transformers — Cache strategies / Sink cache
• Hugging Face Transformers — SinkCache API