KV 才是长上下文的真正瓶颈

一句话：所有 query head 共享同一份 K 与 V。 KV cache 立刻除以 $H_q$ —— Llama-3-70B 上从 64× 缩到 1×。 代价：质量略掉，长上下文 retrieval 退化。

MHA 与 MQA 的中间地带：把 $H_q$ 个 query head 分到 $H_{kv} \ll H_q$ 个组 （Llama-3 选 8）。质量与 MHA 持平、KV cache 缩 $H_q / H_{kv} = 8$×。 几乎所有现代开源模型（Llama 2/3、Qwen2、Mistral、Yi、DeepSeek 早期）默认 GQA。

# GQA: 一组 query head 共享一份 K,V
# 假设 H_q = 64, H_kv = 8 → group size = 8
def gqa(q, k, v):
    # q: [B, H_q,  S, D]   k,v: [B, H_kv, S, D]
    # 把 K,V 在 head 维复制 group_size 次（实际 kernel 是 broadcast，不真复制）
    k = k.repeat_interleave(H_q // H_kv, dim=1)
    v = v.repeat_interleave(H_q // H_kv, dim=1)
    return flash_attn(q, k, v)

关键想法：与其共享 head，不如把每个 token 的 K 和 V 共同压缩到 一个低秩潜变量 $c_t \in \mathbb{R}^{r}$（DeepSeek-V3: $r = 512$）。 推理时只缓存 $c_t$；attention 时再用一个 up-proj 把 $c_t$ 投回完整 K, V。 再配合 rotary 的处理（保留一段非压缩部分），训得到与 MHA 接近的 quality。 KV 体积近似 7%–14% 的 MHA——这是 DeepSeek 能开 128k 上下文的根本。

MLA 解耦了"建模容量"与"缓存体积"——历史上头一回。 缺点：实现复杂，必须配套 RoPE 的特殊处理； serving 时需要专门的 kernel（vLLM、SGLang 都补了）。

关键观察：K 的分布在 channel 维上方差极大（有大 outlier）， V 在 token 维上方差大。 所以——K 按 channel 量化（per-channel），V 按 token 量化（per-token）。 2-bit KV、量化 + 反量化在 attention 中即时做， 在 Llama / Mistral 上几乎不掉点。

# KIVI: K per-channel, V per-token
def kivi_quant_K(K):
    # K: [B, H, S, D]  → quantize along D-axis groups
    scale = K.abs().max(dim=2, keepdim=True).values   # per (B, H, channel)
    qK = (K / scale * 127).round().clip(-127, 127).to(int8)
    return qK, scale

def kivi_quant_V(V):
    # V: [B, H, S, D]  → quantize along D
    scale = V.abs().max(dim=-1, keepdim=True).values   # per (B, H, S)
    qV = (V / scale * 127).round().clip(-127, 127).to(int8)
    return qV, scale

几招组合：non-uniform 量化 grid（用 sensitivity 学）、 pre-RoPE quant、dense-and-sparse decomposition（异常值保留 fp16）。 打到 ~1 bit + 异常值时仍能跑 1M+ context。

经验观察：注意力的累积权重在长上下文里集中在少数 "重击者" token 上。 H2O 保留这些"高得分历史"以及一段"近期 window"，其余的 KV 直接驱逐。 在 OPT / Llama 上把 KV 减 5× 而 longbench 不掉。

一个奇怪的现象：drop 掉最早几个 token 会让模型"崩"—— 因为 softmax 必须把概率质量倒在某处，最早的 token 充当"sink"。 StreamingLLM 永远保留开头若干个 token + 最近 window， 其余驱逐，从而开启无限流式生成而不退化。 这一观察后来在 SoftMax-attention sink 的工程里反复被复用（FA-3 也用了）。

# StreamingLLM: keep first-k sinks + last-w window
def streaming_kv_select(K, V, k_sink=4, w_window=2048):
    S = K.shape[2]
    if S <= k_sink + w_window:
        return K, V
    K_sink, V_sink = K[..., :k_sink, :], V[..., :k_sink, :]
    K_win,  V_win  = K[..., -w_window:, :], V[..., -w_window:, :]
    return torch.cat([K_sink, K_win], dim=2), torch.cat([V_sink, V_win], dim=2)

在 prefill 结束时，用 prompt 末尾几行的 attention 模式来推测"哪些早期 token 重要"，对每层选 top-K KV 保留。后续 decode 都用这份压缩版。 在长 prompt（RAG / 多文档）场景把 KV 缩到 5%–15%，几乎不损质量。

Agent / RAG / few-shot 工作流里， 多个请求经常共享一段 system prompt 或文档前缀。 RadixAttention 把所有当前活跃的 prefix 维护成一棵 基数树（每条边是一段 token 序列）， 一份 KV cache 物理上挂一次， 新请求只需"在最长共享前缀上分叉"。 命中 prefix → 跳过整段 prefill。

多轮对话场景 5–10× 吞吐。SGLang 把 RadixAttention 与 structured decoding（regex / JSON schema 强约束）结合， 成为 agent 系统的事实标准之一。