块级 KV 复用的交叉注意力修复方案

5.1 序列与分块

\(S = B_1 \oplus B_2 \oplus \cdots \oplus B_m\)，其中第 \(i\) 块为 token 区间

\[ B_i = x_{s_i : e_i} \]

\(s_i\)：块 \(B_i\) 在整段序列中的起始全局下标（含）；
\(e_i\)：结束下标（不含，半开区间，与 Python 切片一致）。
例：\(s_1=0,\, e_1=384\) 表示 \(B_1\) 占据位置 \(0,1,\ldots,383\)。

5.2 混合缓存 \(L\)

对每个全局位置 \(t\)、层 \(\ell\)：

\[ L[t,\ell] = \bigl(\hat K_{t,\ell},\; \hat V_{t,\ell},\; m_t\bigr) \]

• \(\hat K_{t,\ell}, \hat V_{t,\ell}\)：拼接加载后的 KV；
• \(m_t = (\type_t,\, b_t,\, g_t)\)：
  • \(\type_t \in \{\mathrm{global},\mathrm{local}\}\) — 该位置 KV 是否已等价于 full-prefill oracle；
  • \(b_t\) — 所属块编号；
  • \(g_t\) — 全局位置（RoPE 与 \(F\) 使用）。

紧凑写法（说明型，非公式定义）：

\[ L \;\approx\; [\;\underbrace{K^{\mathrm{global}}V^{\mathrm{global}}}_{B_1\;\text{全局KV}},\;\underbrace{K^{\mathrm{local}}V^{\mathrm{local}}}_{B_2\;\text{加载时为局部KV}},\;\ldots\;] \]

上标 global / local 表示 KV 的质量类型， 下标表示块编号——并非每个块固定为某一类： B₁ 在真前缀条件下通常全为 global；B₂ 加载时多为 local，经 \(\mathcal{T}_\tau\) 上 \(F\) 写回后可变为 global。

5.3 重算集合 \(\mathcal{T}_\tau\)

接缝 \(B_i \mid B_{i+1}\) 的对称窗口：

\[ \mathcal{T}^{(i)}_{\mathrm{seam}} = [\,s_i,\, e_i)\cap[e_i-\tau,\,e_i) \;\cup\; [\,s_{i+1},\, e_{i+1})\cap[s_{i+1},\, s_{i+1}+\tau) \]

非对称规则（B₁ 为真前缀）：接缝 \(B_1|B_2\) 不重算 \(B_1\) 尾部 \(\tau\) token，预算补至 \(B_m\) 尾部 \(\tau\) token。

重算占比 \(|\mathcal{T}_\tau|/N\)；工程目标：\(\le 5\%\) 且精度达标。

5.4 算子 \(F(Q,L)\) 的功能性表述

\(F\) 是在标准自注意力之上的结构感知扩展，输入 query 集合 \(Q\) 与混合缓存 \(L\)，输出与 oracle 一致的隐藏表示：

\[ F(Q,\,L) \;=\; \Attn\!\bigl(Q,\, \hat K,\, \hat V;\, \mathcal{M}_{\mathrm{corr}}(L)\bigr) \;\approx\; \Attn(Q,\, K^{\star},\, V^{\star}) \]

其中 \(\mathcal{M}_{\mathrm{corr}}(L)\) 为由 \(L\) 元数据导出的纠偏项，明确为：

\[ \mathcal{M}_{\mathrm{corr}}(L) = \{M_{\mathrm{local}}(j),\; M_{\mathrm{sink}}(j)\} \]

\[ M_{\mathrm{local}}(j)=\mathbb{1}[\type_j=\mathrm{local}], \qquad M_{\mathrm{sink}}(j)=\mathbb{1}[j\in\mathcal{S}'] \]

• Prefill · 物理修复模式：对 \(t\in\mathcal{T}_\tau\)，\(F\) 执行完整层间因果前向，写回 \((\hat K_{t,\ell},\hat V_{t,\ell})\) 并将 \(\type_t\leftarrow\mathrm{global}\)；
• Prefill / Decode · 推理模式：对 query 仅做 attention 前向（不写回），用 \(\mathcal{M}_{\mathrm{corr}}(L)\) 纠偏 \(E_3\)；
• 退化：当 \(\tau \ge |B|/2\) 且非对称覆盖 \(B_{2:m}\) 时，\(L\approx L^{\star}\)，\(\mathcal{M}_{\mathrm{corr}}(L)=\varnothing\)，\(F\) 等价于标准 \(\mathrm{Softmax}(QK^{\top}/\sqrt{d})V\)。

\(L^{\star}\)：对完整 \(S\) 一次 full-prefill 的 oracle 缓存（对照基准，非在线必算）。

Logits 形式（单层）

\[ s_{ij} = \frac{\langle q_i,\,\hat k_j\rangle}{\sqrt{d}} + \underbrace{\alpha\,\mathbb{1}[\type_j=\mathrm{local}]}_{\text{局部 KV 倾斜项}} + \underbrace{\beta\,\mathbb{1}[j\in\mathcal{S}']}_{\text{假 Sink 抑制项}} \]

\[ o_i = \sum_j \mathrm{softmax}_j(s_{ij})\,\hat v_j \]

• \(\alpha < 0\)（或对全局 key 等价增权）：引导 mass 离开尚未修复的局部 KV；
• \(\beta < 0\)：抑制非首块块首 token 集合 \(\mathcal{S}'\)，保留序列最前 \(s\) 个真 sink；
• 步骤 ③ 写回后 \(\type_j=\mathrm{global}\) 增多 → \(\alpha\) 项自动减弱。

\(F\) 的退化（非独立 regime，而是算子内项为空）

定义纠偏掩码集合（与 §5.4 一致）：

\[ \mathcal{M}_{\mathrm{corr}}(L) = \{M_{\mathrm{local}}(j),\;M_{\mathrm{sink}}(j)\}, \;\; M_{\mathrm{local}}(j)=\mathbb{1}[\type_j=\mathrm{local}], \;\; M_{\mathrm{sink}}(j)=\mathbb{1}[j\in\mathcal{S}'] \]

当 \(\mathcal{T}_\tau\) 足够大（如 \(\tau\ge |B|/2\) 且非对称覆盖 \(B_{2:m}\)）时：

• 步骤 ③ 已将相关位置写回为全局KV → \(\type_j=\mathrm{local}\) 的集合为空；
• 假 Sink 经抑制与写回消除 → \(\mathcal{S}'\) 有效为空；
• 故 \(\alpha\) 项与 \(\beta\) 项均为空集上的零映射，

\[ F(Q,L) \;\equiv\; \mathrm{Softmax}\!\left(\frac{QK^{\top}}{\sqrt{d}}\right)V \]

即自动退化为标准 Transformer 自注意力。 此时计算上仍对后缀做了完整因果 prefill，复用上主要保留 \(B_1\) 缓存收益。