理解 Attention Residuals（AttnRes），需要先回答几个问题：

1. 残差连接在深度方向上累加了什么？
2. Attention 的选择对象是 Token，还是历史层输出？
3. Query 不依赖当前 Hidden State，为什么权重仍然与输入有关？
4. Full AttnRes 与 Block AttnRes 分别保留哪些状态？
5. 增加跨层访问后，显存、通信和推理执行需要怎样调整？

下面按照残差连接、深度注意力、分块实现和工程分析的顺序讲解。本文介绍 Moonshot / Kimi Team 的 Attention Residuals；概念与公式依据 原论文第 2–4 节，实现接口可对照 官方仓库。

一、残差连接：既是梯度通路，也是信息聚合

一个 PreNorm 子层可以写成：

\[ h_{l+1}=h_l+f_l(\operatorname{Norm}(h_l)) \]

这里的子层可以是 Self-Attention，也可以是 MLP / MoE。为了简化符号，下面把子层内部的 Norm 包含在 \(f_l\) 中，记子层输出为 \(v_l=f_l(h_l)\)，Embedding 为 \(v_0\)。

展开递推，可以得到：

\[ h_l=v_0+\sum_{i=1}^{l-1}v_i \]

这说明每个后续子层读到的是历史输出的累加结果。

例如，当前输入是 \(v_0+v_1+v_2\)，即使当前子层更需要 \(v_1\)，残差合并本身也没有一个单独调整 \(v_1\) 权重的接口。子层仍然可以通过自身变换处理混合结果，但无法直接读取合并前的各个来源。

随着深度增加，累加状态的幅度和新输出的相对占比会变化。这是理解 PreNorm Dilution 的入口；幅度的具体增长速度依赖输出相关性等条件，不能把它理解成所有模型都遵循同一个增长公式。

二、Attention 沿哪个维度发生？

AttnRes 的聚合算子本身不会混合不同 Token。不同 Token 可以得到不同的深度权重，但各自只聚合自己位置上的历史表示。

因此，AttnRes 与 MQA、GQA、MLA 等机制处在不同位置：后者改变 Token Attention 的实现，AttnRes 改变子层之间的信息聚合。两类设计可以组合。

还要区分三个容易混淆的 Block：

• Transformer Block：通常包含 Attention 和 MLP 两个子层。
• AttnRes Block：将一组子层输出压缩成一个摘要的深度分组。
• GPU Thread Block：GPU 执行与线程协作单位。

下面的 \(L\) 统计 Attention / MLP 子层数量，不是 Transformer Block 数量。

三、Full AttnRes：保留来源，再学习如何聚合

对于第 \(l\) 个子层，候选来源是：

\[ V_l=\{v_0,v_1,\ldots,v_{l-1}\} \]

每个子层有自己的可学习向量 \(w_l\in\mathbb R^d\)。聚合过程为：

\[ s_{i\to l}=w_l^\top\operatorname{RMSNorm}_l(v_i),\qquad \alpha_{i\to l}=\frac{\exp(s_{i\to l})}{\sum_{j=0}^{l-1}\exp(s_{j\to l})} \]

\[ h_l=\sum_{i=0}^{l-1}\alpha_{i\to l}v_i,\qquad v_l=f_l(h_l) \]

这里没有必要机械套用普通 Self-Attention 的 Q/K/V 投影矩阵或 \(1/\sqrt d\) 缩放；应以 AttnRes 自身的参数化为准。

1. 固定 Query 为什么能产生动态权重？

训练结束后，\(w_l\) 是固定参数，但是 Key 来自当前输入产生的历史表示。输入变了，Key 和分数会变化，所以权重仍然是 Input-dependent 的。

可以把 \(w_l\) 理解为该子层学习到的“选择标准”，而历史表示是每次推理中被选择的内容。

2. 加权平均与固定累加的区别

假设某个 Token 的三个来源是 \(v_0,v_1,v_2\)，权重为 \((0.1,0.7,0.2)\)，则：

\[ h_l=0.1v_0+0.7v_1+0.2v_2 \]

权重对来源不同，但同一来源的所有 Hidden 通道共享这个标量权重。这不是逐通道门控，也不是 Top-k 选择；有限分数下，Softmax 通常仍给每个候选一个非零权重。

由于权重非负、和为 1，给定这些 Value 时，输出是它们的凸组合：

\[ \|h_l\|\leq\sum_i\alpha_{i\to l}\|v_i\|\leq\max_i\|v_i\| \]

这个性质约束的是聚合本身。子层 \(f_l\) 仍然可能放大输出，因此不能由此推导整个网络一定有界或梯度一定稳定。

四、Block AttnRes：用深度分块控制状态数量

Full AttnRes 需要保留各个历史来源。Block AttnRes 将子层分为 \(N\) 组，每组内部累加子层输出，得到一个块摘要。

令 \(B_n\) 为第 \(n\) 个深度块，完成后的摘要为：

\[ b_n=\sum_{j\in B_n}v_j,\qquad b_0=v_0 \]

当前块执行到一半时，还维护已完成子层输出的部分和 \(p\)。每次聚合的候选为：

聚合后的 \(h_l\) 用作子层输入；得到 \(v_l=f_l(h_l)\) 后，更新：

\[ p\leftarrow p+v_l \]

到达块边界时，将 \(p\) 保存为一个完成块摘要，并开始新的块。

需要特别注意：块摘要是该块自身子层输出的和，不能把整个历史混合输入 \(h_l\) 每次加进摘要，否则会重复累加历史信息。

一个具体例子

假设 6 个子层被分成 3 块，每块 2 个子层：

1. 第一块先读取 \(v_0\)，产生 \(v_1\)；第二个子层读取 \(\{v_0,v_1\}\)，产生 \(v_2\)，完成 \(b_1=v_1+v_2\)。
2. 第二块先读取 \(\{v_0,b_1\}\)，产生 \(v_3\)；下一个子层读取 \(\{v_0,b_1,v_3\}\)，产生 \(v_4\)，完成 \(b_2=v_3+v_4\)。
3. 第三块按相同方式读取历史摘要并维护自己的部分和。

第一块结束后，后续子层可以选择 \(b_1\)，但不能再分别选择 \(v_1\) 和 \(v_2\)。这是状态压缩带来的选择粒度变化。

五、一个可读的 PyTorch 聚合示例

下面仅实现候选来源的聚合算子，省略可学习的 RMSNorm Gain、完整 Transformer 与系统优化。sources 可以是 Full 的历史子层输出，也可以是 Block 的摘要与部分和。

import torch


def attention_residual(sources, query, eps=1e-6):
    # sources: 非空列表，每个 Tensor 的 Shape 为 [B, T, D]
    # query: 当前子层的可学习参数，Shape 为 [D]
    values = torch.stack(sources, dim=0)  # [R, B, T, D]
    values_fp32 = values.float()
    rms_inv = torch.rsqrt(values_fp32.square().mean(dim=-1, keepdim=True) + eps)
    keys = values_fp32 * rms_inv
    scores = (keys * query.float()).sum(dim=-1)  # [R, B, T]
    weights = torch.softmax(scores, dim=0)      # 沿来源轴归一化
    mixed = (weights.unsqueeze(-1) * values_fp32).sum(dim=0)
    return mixed.to(values.dtype), weights


sources = [torch.randn(2, 3, 8) for _ in range(4)]
query = torch.nn.Parameter(torch.zeros(8))
mixed, weights = attention_residual(sources, query)
assert mixed.shape == (2, 3, 8)
torch.testing.assert_close(weights.sum(dim=0), torch.ones(2, 3))
torch.testing.assert_close(mixed, torch.stack(sources).mean(dim=0))

\(w_l=0\) 时，分数都是 0，初始权重是均匀分布。这个例子由此验证输出为来源的平均值。均匀平均与标准残差的求和相差来源数量的比例，不能称为数值完全相同。

工程实现还要决定最终输出的聚合、Norm 参数、块边界及初始化方式。用 stack 展开来源适合解释 Shape，但高性能实现应分析它的额外分配和访存成本。

六、显存与计算成本如何估算？

设当前处理的 Tensor 大小为 \([B,T,d]\)，存储每个元素需要 \(e\) 字节。以下估算针对前向聚合使用的来源状态，不包含模型权重、Token Attention KV Cache 或完整反向激活。

Block 的计算估算来自“\(L\) 个子层分别访问最多约 \(N\) 个来源”，不应仅因为状态数降到 \(N\)，就把整个网络直接实现的计算量写成与 \(L\) 无关。

例如，假设 \(B=1,T=8192,d=4096,N=8,e=2\)，只按 8 个摘要 Tensor 估算：

\[ M=BTNd e=536{,}870{,}912\text{ bytes}=512\text{ MiB} \]

这不包含 Embedding、当前部分和及临时 Buffer。即使块数不多，长上下文 Prefill 的临时状态仍可能很大。

训练中，如果某些激活本来就为反向传播保留，来源状态与其可能重叠；采用重计算或流水线并行时，应重新分析状态生命周期和跨 Stage 传输，不能直接把上表当作额外显存的精确值。

七、为什么可以采用两阶段计算？

对于当前块，已经完成的历史摘要不再变化，而块内所有子层的 Query 都是已知参数。因此，可以先批量计算这些 Query 对历史摘要的打分和聚合，再随子层执行依次处理当前部分和。

但 Self-Attention 和 MLP 本身仍依赖前面的计算，不能因为 Query 已知，就把整个 Transformer 的子层并行执行。

对两组候选，分别保存稳定 Softmax 的三个统计量：

\[ m=\max_i s_i,\quad z=\sum_i e^{s_i-m},\quad u=\sum_i e^{s_i-m}v_i \]

合并两组统计量时，令 \(m=\max(m_a,m_b)\)：

\[ z=e^{m_a-m}z_a+e^{m_b-m}z_b,\qquad u=e^{m_a-m}u_a+e^{m_b-m}u_b,\qquad h=u/z \]

这说明两组候选可以分别处理后再合并，保持同一个 Softmax 的数学结果；浮点实现仍可能有舍入差异。不能把两组各自归一化后的输出直接相加，因为它们的归一化分母不同。

对照 GPU Arch，这类优化主要减少重复 HBM 读取并改变 Kernel 调度，而不是改变子层要执行的数学变换。

八、与其他残差机制如何比较？

这个对比描述设计接口，不代表模型效果排序。比较实际模型需要控制参数量、训练数据、训练预算和执行配置。

九、工程实践中需要检查什么？

AttnRes 来源状态服务于当前 Token 沿深度的计算；Token KV Cache 服务于未来 Token 的序列 Attention。不要把两个维度的缓存混为一谈，也不要默认所有深度摘要都需要跨生成步永久保留。

这是模型结构的改变。把已训练模型中的残差加法直接替换成加权平均，会改变各层输入分布；要得到有效模型，需要相应训练或经过验证的适配过程。

学习时可以先手工展开三个子层的残差，再实现来源轴上的 Softmax，最后加入 Block 摘要。每一步都检查：保存了哪些来源、哪些来源还可被单独选择，以及为了读取它们付出了多少存储与访存成本。