【轻量化复刻大模型 / KIMI K3】Attention Residuals 完整梳理（1）

[0x00 概要](#0x00 概要)
- [0.1 脉络](#0.1 脉络)
- [0.2 论文信息](#0.2 论文信息)
- [0.3 关键图/表索引](#0.3 关键图/表索引)
- [0.4 注意说明](#0.4 注意说明)
[0x01 先懂 AttnRes：残差为什么是瓶颈](#0x01 先懂 AttnRes：残差为什么是瓶颈)
- [1.1 标准残差的展开](#1.1 标准残差的展开)
- [1.2 PreNorm 稀释与信息瓶颈](#1.2 PreNorm 稀释与信息瓶颈)
- [1.3 一句话痛点](#1.3 一句话痛点)
- [1.4 思路：把治疗时间维度的药，用在深度维度上](#1.4 思路：把治疗时间维度的药，用在深度维度上)
- [1.5 一句话直觉](#1.5 一句话直觉)
[0x02 演进脉络与定位](#0x02 演进脉络与定位)
- [2.1 前身：标准残差连接及其痛点](#2.1 前身：标准残差连接及其痛点)
- [2.2 同期竞品：都是"残差路由"的玩家](#2.2 同期竞品：都是"残差路由"的玩家)
- [2.3 深度信息流的五派方案](#2.3 深度信息流的五派方案)
- [2.4 论文的关键消融结论](#2.4 论文的关键消融结论)
- [2.5 一个常见的理解误区](#2.5 一个常见的理解误区)
- [2.6 定位一句话](#2.6 定位一句话)
[0x03 原理](#0x03 原理)
- [3.1 直觉类比视角](#3.1 直觉类比视角)
- [3.2 准专业水平视角](#3.2 准专业水平视角)
- [3.3 专家视角（信息论 / 优化 / 计算复杂性）](#3.3 专家视角（信息论 / 优化 / 计算复杂性）)
- [3.4 AttnRes 的数学机制](#3.4 AttnRes 的数学机制)
  - 从"求和"改成"期望"
  - [伪查询、RMSNorm 与 score](#伪查询、RMSNorm 与 score)
  - [Full AttnRes](#Full AttnRes)
  - [Block AttnRes](#Block AttnRes)
- [3.5 伪代码](#3.5 伪代码)
- [3.6 真实代码](#3.6 真实代码)
- [3.7 伪代码 vs 本项目实现对照](#3.7 伪代码 vs 本项目实现对照)
- [3.8 如何使用](#3.8 如何使用)
[0x04 生产级实现与源码证据](#0x04 生产级实现与源码证据)
- [4.1 SGLang 实现](#4.1 SGLang 实现)
  - [4.1.1 Python 层：sglang-main/python/sglang/srt/layers/attn_residual.py](#4.1.1 Python 层：sglang-main/python/sglang/srt/layers/attn_residual.py)
  - [4.1.2 模型集成：sglang-main/python/sglang/srt/models/kimi_k3.py](#4.1.2 模型集成：sglang-main/python/sglang/srt/models/kimi_k3.py)
  - [4.1.3 SP 集合通信融合：sglang-main/python/sglang/srt/layers/k3_sp_collective.py](#4.1.3 SP 集合通信融合：sglang-main/python/sglang/srt/layers/k3_sp_collective.py)
  - [4.1.4 SGLang Kernel 层（TMA + HIP）](#4.1.4 SGLang Kernel 层（TMA + HIP）)
- [4.2 vLLM 实现](#4.2 vLLM 实现)
  - [4.2.1 NVIDIA 模型路径：vllm-main/vllm/models/kimi_k3/nvidia/model.py](#4.2.1 NVIDIA 模型路径：vllm-main/vllm/models/kimi_k3/nvidia/model.py)
  - [4.2.2 NVIDIA operator：vllm-main/vllm/models/kimi_k3/nvidia/ops/attn_res.py](#4.2.2 NVIDIA operator：vllm-main/vllm/models/kimi_k3/nvidia/ops/attn_res.py)
  - [4.2.3 AMD operator: vllm-main/vllm/models/kimi_k3/amd/linear.py 与 amd/ops/attn_res.py](#4.2.3 AMD operator: vllm-main/vllm/models/kimi_k3/amd/linear.py 与 amd/ops/attn_res.py)
  - [4.2.4 vLLM warmup、env 与配置](#4.2.4 vLLM warmup、env 与配置)
- [4.3 源码伪代码总览（基于 SGLang/vLLM 共同语义）](#4.3 源码伪代码总览（基于 SGLang/vLLM 共同语义）)
- [4.4 与 Kimi K3 架构的协同](#4.4 与 Kimi K3 架构的协同)
[0xFF 参考](#0xFF 参考)

0x00 概要

一句话定位 ：标准残差是"深度维度的等权加法"；AttnRes 把深度维的聚合换成"深度维的 softmax 注意力"------每层用一个可学习的伪查询向量 （pseudo-query）去"注意"之前的各层输出，权重由内容和深度共同决定。这是对 RNN→Transformer 那场范式跨越在深度轴上的完整复刻。

0.1 脉络

标准 PreNorm Transformer 的残差流像"所有旧信息都以 1 倍权重倒进同一个桶"：层数越深，桶越大，新层信号越像一滴水。

Attention Residual (AttnRes)把这件事改成"当前层先问：我现在更需要哪些历史层/历史块？"再用 softmax 概率加权取一个混合表示，因此深度方向的信息流从固定累加变成了可学习的选择。

Block AttnRes 是工程化版本：不保存每一层输出，而是每若干层形成一个块表示，再在"历史块 + 当前块部分和"上做深度方向注意力。它牺牲块内分辨率，换来训练、推理和流水线并行可落地。

0.2 论文信息

项	值
标题	Attention Residuals（TECHNICAL REPORT OF ATTENTION RESIDUALS）
作者	Kimi Team
机构	Moonshot AI（月之暗面 / Kimi）
类型	技术报告（Technical Report，非会议论文）
代码仓库	https://github.com/MoonshotAI/Attention-Residuals

0.3 关键图/表索引

以下是论文的关键图/表索引

编号	内容
Figure 1	三种残差结构总览（(a)标准 / (b)Full / (c)Block）
Figure 2	Block AttnRes 的 PyTorch 伪代码（`block_attn_res` + `forward`）
Figure 3	Cache-based pipeline 通信示例（P=4, V=2）
Figure 4	Scaling law 曲线（Baseline/Full/Block 三线拟合）
Figure 5	训练动态：(a)验证 loss (b)每块输出幅值 (c)每块梯度幅值
Figure 6	Block size S 对验证 loss 的影响（16 层模型）
Figure 7	固定算力架构扫描热图（d_model/L_b × H/L_b）
Figure 8	学到的深度注意力权重 α 分布（Full 上 / Block 下）
Figure 9	四种残差变体的深度混合矩阵 M
Table 1	每 token 每层的记忆访问成本（残差机制 I/O）
Table 2	Scaling law 模型配置与验证 loss（5 规模）
Table 3	48B 模型下游任务对比（Baseline vs AttnRes）
Table 4	16 层模型消融实验
Table 5	残差更新机制统一视角（Weight/Source 分类）
Algorithm 1	两阶段计算（Phase 1 跨块并行 + Phase 2 块内顺序 + online softmax）
Appendix B	Full AttnRes 优化推理 I/O 的完整推导

0.4 注意说明

项目https://github.com/RossiLH/Model-Viz 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
涉及 KIMI K3，DeepSeek V4，PI0.5，PI0.7 等模型。后续可能会跟踪新模型。
全 CPU 复刻，不涉及 GPU，也不涉及大规模训练（这是另外一个庞杂深邃的领域）和硬件。
模型的组件尽量是原始大模型的缩小版，如果组件没有开源版本，则用近似组件替代。
会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么，以及组件之间如何咬合成一个整体。重点不是记名词，而是理解一条主线。
每个主要组件都有可直接运行的 Gradio demo------无需 GPU。另有若干可视化脚本（plot_*.py），可生成静态图。
所有文件行号基于本地源码快照；源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。
为了保证每个文章可以尽量独立成文，所以有些内容会在不同文章中都存在，可能在系列层面上存在冗余，还请大家理解。

0x01 先懂 AttnRes：残差为什么是瓶颈

现代 Transformer 几乎每一层都在用同一句话结束：x = x + F(x)。这句话叫 残差连接（residual connection），它在 2015 年（ResNet）就被证明能训很深的网络，2017 年被 Transformer 拿走。我们经常忘了它其实有一个很深的缺陷。

1.1 标准残差的展开

标准残差可写为：

$ h_{l+1}=h_l+F_l(\mathrm{Norm}(h_l)) $

展开后近似是：

$ h_L=h_0+\sum_{i=0}^{L-1}F_i(\mathrm{Norm}(h_i)) $

大白话：每一层看到的都是全部先前信息的"无差别大杂烩" ------所有早期层的输出被 固定的单位权重 1 加在一起。最后一层拿到的不是"最近一层的干净结论"，而是从 embedding 到所有层输出的总和。

1.2 PreNorm 稀释与信息瓶颈

论文 §2.1 明确指出，标准残差带来三个问题：

无选择性访问（no selective access）：不同类型的层（attention vs MLP）拿到同一个聚合状态，尽管它们可能从不同的加权方式中获益。
不可逆损失（irreversible loss）：聚合中丢失的信息，无法在更深层被选择性地找回。
输出增长（output growth）：为了让自己的输出在累积残差里更有影响力，后面的层被迫学出越来越大的输出，这会破坏训练稳定。

PreNorm 会把每层模块的输入归一化，所以每层新产生的残差增量常常有相近尺度。若这些增量方向大体独立，平方范数会随层数近似线性增长，L2 范数约随$\sqrt{L}$ 增长；若方向更同向，增长可更接近$L$。

这不等于模型一定发散，但意味着一个固定尺度的新特征进入一个越来越大的残差流后，相对占比会下降------这就是深层网络的信息稀释：残差连接保护了梯度不消失，却没有给前向信息流提供选择机制。

1.3 一句话痛点

残差连接把「历史上所有计算结果」压成一个"和"（深度方向的瓶颈），并且没有给网络"决定读哪一层"的能力。

1.4 思路：把治疗时间维度的药，用在深度维度上

我们回想序列建模领域有一个著名的"换药事件"：RNN 把整段历史压缩进一个隐状态（h_t），这是"时间维度的瓶颈"；Transformer 把它替换成 attention over time ------每个位置可以选择性地读取所有过去位置，权重由内容决定（data-dependent weights）。

AttnRes 说：模型深度其实和序列时间是一模一样的困境。 残差连接把历史的计算压成一个状态（h_l），是无法选择的、等量的。那为什么不用同样的药？------把 attention 用在深度上：让第 l 层也把前面所有层当作一个"序列"，用 softmax 注意力去挑选它真正想读的早期层。

1.5 一句话直觉

标准 PreNorm Transformer 的残差流像"所有旧信息都以 1 倍权重倒进同一个桶"：层数越深，桶越大，新层信号越像一滴水。Attention Residual (AttnRes) 把这件事改成"当前层先问：我现在更需要哪些历史层/历史块？"再用 softmax 概率加权取一个混合表示，因此深度方向的信息流从固定累加变成了可学习的选择。

0x02 演进脉络与定位

2.1 前身：标准残差连接及其痛点

前代方案 ：Pre-Residual（x = x + F(x)），2015 年 ResNet 提出、Transformer 沿用至今。
缺陷：固定等权相加 → 深度稀释、hidden 范数随深度 O(L) 无界增长（PreNorm dilution）、无法按需选择性读取早期层。

标准残差可写为：$h_{l+1}=h_l+F_l(\mathrm{Norm}(h_l))$展，开后近似是：$h_L=h_0+\sum_{i=0}^{L-1}F_i(\mathrm{Norm}(h_i))$

大白话：最后一层拿到的不是"最近一层的干净结论"，而是从 embedding 到所有层输出的总和；每个历史项默认权重都是 1。

论文证据：
- 图 5(b) ：Baseline 的每块输出幅值随深度单调爬升（PreNorm 稀释的直接可视化）；Block AttnRes 则呈现有界周期性模式（在块边界被重置）。
- 图 5(c)：Baseline 的最浅层梯度幅值异常大（残差权重全为 1，无法调控跨深度梯度流）；Block AttnRes 的梯度分布明显更均匀。
- Scaling law（图 4）：Baseline 的 loss 曲线始终高于 AttnRes。
- 论文引用「深层被裁剪掉一部分损失很小」作为"早期信息被埋没、无法选择找回"的旁证。

2.2 同期竞品：都是"残差路由"的玩家

论文 Table 5 把所有残差变体分成三大类（单状态递推 / 多状态递推 / 跨层访问）

以下按"权重来源 × 源访问范围"组织：

方案	核心机制	权重来源	可访问源	优势	劣势
标准残差	`x + F(x)`	固定单位权重	仅 `h_{l-1}`	简单、稳定、已被验证	无选择性；逐层稀释
ReZero	每层可学习标量 λ_l	静态标量	仅 `h_{l-1}`	极简、深层稳定	无选择性
LayerScale	`diag(λ)` 逐通道缩放	静态对角向量	仅 `h_{l-1}`	稳定性好	无选择性
Highway	`(1-g)·x + g·F(x)`	输入相关标量门	仅 `h_{l-1}`	输入相关	单状态；深度维 1-semiseparable
SiameseNorm	双流 PreNorm/PostNorm 共享参数	固定	2 条流	保恒等梯度 + 有界表示	缓解而非选择
mHC	m 条并行残差流 + 学习混合矩阵	输入相关（动态）	m 条流	状态扩张 rank↑、稳定	I/O 高（~34d/层）；仍是线性注意
DenseFormer	各层访问所有先前输出，加权求和	静态标量系数（不随输入）	`[h_1..h_{l-1}]`	跨层访问	权重与输入无关（消融证实无收益）
MRLA	逐层 sigmoid 门控（可分离 QK 点积）	输入相关	`[h_1..h_{l-1}]`	跨层访问	更接近 linear attention 而非 softmax
AttnRes（Full/Block）	深度维 softmax 注意力	输入相关（softmax）	`[h_1..f_{l-1}]` / `[b_0..b_{n-1}]`	选择性最强；显式深度检索	Full 内存 O(Ld)；Block 块内压缩

2.3 深度信息流的五派方案

从更宏观的视角，可以把深度维度上的信息处理思路分成几大学派：

(1) 稳定化/缩放派

代表思路：PostNorm、PreNorm、残差缩放、ReZero、DeepNorm、LayerScale 等。它们主要调节信号幅度、初始化或梯度稳定性。优点是改动轻、系统友好；缺点是信息路径仍近似一条递推链，历史选择能力有限。

(2) 门控递推派

Highway、LSTM/GRU 类思想在深度方向的对应物，是为残差通道加门。可以把 Highway 视作一种不显式 softmax 的 stick‑breaking 式深度注意力：权重可以依赖输入，但仍沿递推传播。

(3) 多状态流派

mHC、MRLA、DDL 等方法维护多个并行状态或线性注意力式状态。它们把"单管道"拓宽成多管道，降低压缩瓶颈。局限是仍然更像递推系统：信息要经过状态转移，不是每层都直接访问任意历史层。

(4) 跨层访问派

DenseNet、ELMo、DenseFormer、MUDDFormer 等允许跨层组合。静态权重方法缺少输入依赖；部分动态方法更像线性注意力。AttnRes 的位置是：在深度轴上使用 softmax attention，使跨层访问既密集又竞争。

(5) 结构化矩阵视角

也可把第 l 层输入看成历史输出的加权和：

$ x_l=\sum_{j<l}M_{\{l,j\}}y_j $

标准残差：M 是全 1 下三角矩阵。
Highway：M 由门控连乘产生，低秩但输入依赖。
mHC：M 有多条状态流，半可分秩约为 m。
Full AttnRes：M 可接近稠密，按内容动态产生。
Block AttnRes：M 介于标准残差和 Full AttnRes 之间，块内共享粗粒度源。

2.4 论文的关键消融结论

DenseFormer（静态权重）1.767 ≈ baseline 1.766 → "可学习但静态"不够，必须输入相关。
换成 input-independent 可学习标量混合 1.749 > Full 1.737 → 再次证明输入相关性是增益来源。
mHC（动态）1.747 介于 DenseFormer 与 AttnRes 之间 → 动态但仍是线性注意；AttnRes 用软注意跨层选择性检索再进一步。

2.5 一个常见的理解误区

过去一些解读把 mHC 的权重描述为"可学习矩阵 + 正交约束"，甚至直接说它"与输入无关"。这与论文 Table 5 不符------mHC 是 Dynamic（输入相关） 的：其更新规则 H_l = H_{l-1}A_l + f_{l-1}(H_{l-1}α_{l-1})β^T_{l-1} 中，有效权重 M_{i→l}=β^T_i A_{i+1→l} α_l 含输入相关项（α、β 都依赖 H_{l-1}，而 H_{l-1} 由之前的输入决定）。

AttnRes 与 mHC 的真正区别不在"是否输入相关"，而在：

注意力范式 ：mHC 是深度维 linear attention （m-semiseparable，低秩）；AttnRes 是深度维 softmax attention （稠密 rank-L）------softmax 的竞争性归一化（probability mass 在源间竞争）带来更锐利的选择。
状态形态：mHC 把历史压进固定 m 条流（RNN 式有限状态）；AttnRes 保留块表示列表（可跨层检索）。

论文 Table 4 中 AttnRes（1.737/1.746）优于 mHC（1.747）的原因就是：softmax 的竞争性归一化 + 跨层直接检索。

2.6 定位一句话

不同于 mHC 的"多流矩阵状态 + m-semiseparable 线性注意"、也不同于 DenseFormer 的固定系数加权，本方案（AttnRes）通过逐层可学习伪查询向量 + 深度维 softmax 注意力 ，让每一层内容相关 地选择性读取所有先前层输出------把"深度维度的 linear attention"升级为"深度维度的 softmax attention"，完成与序列维度"linear → softmax"相同的范式跨越。（论文核心观点 §6.2：把 HC/mHC 等统一看成"深度维的 linear attention"，AttnRes 是深度维 softmax attention 的第一个实例，即"深度-时间对偶性"。）

0x03 原理

3.1 直觉类比视角

想象一个大学课堂上，讲台上正在讲一个很长的证明。标准残差连接就像老师每讲完一步，就把"之前所有内容的摘要"工工整整地写在同一张白纸上------越往后，白纸越乱，每一行旧字都被新字盖得几乎看不见（稀释），而且老师从不回头重读某一页（无选择）。

AttnRes 的做法是：把之前每一步的推导分别存成独立的小卡片 （不是糊成一团的总和），第 l 步讲话之前，老师快速"扫一眼"所有卡片，用注意力加权挑出跟当前最相关的几张（例如上一步、以及最初的嵌入），按权重拼成这一讲的讲义。既不会稀释，又能精准引用早先的结果。

Block AttnRes 则是"每讲完一个知识块就把块里的卡片合并成一张总卡"，只留 N 张总卡在桌上------白板没那么乱了，代价是块内的细致区分被牺牲。

一句话：标准残差是"一张黑板的累积字迹"，AttnRes 是"卡片目录上的检索注意力"。

3.2 准专业水平视角

输入输出的 shape 与含义

设 h_l ∈ R^{B×T×d} 为第 l 层（attention 或 MLP，论文把每个 self-attention 或 MLP 视为一个独立 layer）的输入（B：batch, T：序列长度, d：hidden dim）。本文全程使用 h_1 表示 token embedding（下标从 1 开始，与论文一致）。f_l 是第 l 层的变换函数。

标准做法的公式

python 复制代码

h_l = h_{l-1} + f_{l-1}(h_{l-1})          （PreNorm: x + F(x)）

展开后它是等权累加：

python 复制代码

h_l = h_1 + Σ_{i=1}^{l-1} f_i(h_i)        # 每个累加项的权重都是固定 1

AttnRes 的做法（论文 Eq. 2-4，即 Full 形式）

第 l 层拥有一个逐层可学习伪查询向量 q_l = w_l ∈ R^d。keys 与 values 取同样的东西------所有之前层的输出（以及 token embedding）：

python 复制代码

k_i = v_i = { h_1                  if i = 0            ← 嵌入永远是一个 source
              f_i(h_i)             if 1 ≤ i ≤ l-1 }    ← 第 i 层的输出

注意力权重：
α_{i→l} = φ(q_l, k_i) / Σ_j φ(q_l, k_j)
φ(q, k) = exp(q^T · RMSNorm(k))          # 核函数

h_l = Σ_i α_{i→l} · v_i                   # 加权和，而非单位加

α 是非负、且随 l 变化的 softmax 分布 （对每一 l，Σ_i α_{i→l} = 1）。
RMSNorm 的作用：防止某个出口量级巨大的层靠"嗓门大"（大范数）垄断权重 ------归一化后每层方向等比同步，φ 只表达方向相似度、不表达量级。论文消融证实「去掉 RMSNorm 会掉点」（Full 1.743 vs 1.737，Block 1.750 vs 1.746）。
h_l 与标准残差的关键差别 ：标准残差的 h_l 是逐层累加的既有量 + 一位新贡献；AttnRes 的 h_l 本身就是一次关于历史层的混合 ------它是"从历史中选出来的加权组合"，而不是"上一个状态再加点什么"。

公式级对比表

维度	标准残差	AttnRes（Eq. 2-4）
输入	只用 `h_{l-1}`	所有历史：`h_1` 到 `f_{l-1}(h_{l-1})`
权重	固定单位 1	内容相关的 softmax `α_{i→l}`
语义	等权累加 `Σ v_i`	加权和 `Σ α·v`
选择性	无	有（不同 task 学不同读取）
参数开销	0	每层一个 `w_l ∈ R^d`
记忆	O(d)	Full：O(Ld)（保留所有历史层输出）

3.3 专家视角（信息论 / 优化 / 计算复杂性）

为什么它理论上有用（深度-时间对偶） 。论文提出一个统一视角：把「深度维的聚合」类比「时间维的递推」，把残差本身看成"深度维度的 linear attention"：

标准残差 = 权重矩阵为全 1 的下三角（M_{i→l}=1），即"fixed uniform linear attention over depth"。
Highway = 标量门 g，权重可因子化成 M_{i→l} = γ_{i+1→l}·g_{i+1}（γ_{i→l} = Π_{j=i+1}^l (1-g_j)），1-semiseparable（秩 1），是 softmax-free 的 stick-breaking attention。
(m)HC / mHC = m 个平行流 + 学习的段间转移矩阵 A ，展开即 M_{i→l} = β^T_i A_{i+1→l} α_l，m-semiseparable ，相当于"深度维的、具有 m 维状态（state expansion，RNN 状态从 d 扩到 d×m）的 linear attention"。mHC 进一步要求每个 A_l 双随机（doubly stochastic），使跨深度的累积积稳定（Sinkhorn-Knopp 投影到 Birkhoff 多胞体）。
AttnRes = 深度维度的 softmax（dense rank-L）attention ：M_{i→l} = α_{i→l} 由 input-dependent softmax 定义，矩阵稠密。

也就是说：序列建模走了 "linear attention（RNN）→ softmax attention（Transformer）"，而深度建模在 AttnRes 里完成了同一个过渡------把"深度维 linear attention"（residual/Highway/mHC）升级为"深度维 softmax attention"。

信息论视角 ：等权残差是"信息最差情况下的融合"（平均）；softmax 让每层有选择地提取与任务最相关的深度表示，降低"单层输出的熵"。

下图是三个残差实现的对比图。

论文尺度证据（ablation，Table 4，16 层模型）

变体	Loss	相对 Full 说明
Baseline (PreNorm)	1.766	参照
DenseFormer	1.767	静态权重跨层，无增益
mHC	1.747	动态但线性注意
Full AttnRes	1.737	最优
Block (S=4)	1.746	接近 Full
SWA (W=1+8)	1.764	只保留最近 8 层，远不如 Block
w/ input-dependent query	1.731	理论更优但参数/顺序访存代价
w/ input-independent mixing	1.749	去掉输入相关性掉点
w/ sigmoid	1.741	无竞争性归一化，掉点
w/o RMSNorm	1.743	RMSNorm 必要
w/ multihead (H=16)	1.752	多头反而有害

关键解读：

输入依赖是杠杆点：DenseFormer（固定权重）看不到任何收益，证明 gain 来自 input-dependent 而非"多读一层"。
softmax vs sigmoid ：softmax 的竞争性归一化（probability mass 在源间竞争）比 sigmoid 更锐利------用 sigmoid 掉点（1.741）。
多头有害（1.752 > 1.746）：深度混合在各通道上基本均匀，多头引入冗余。论文解释："当一层输出相关时，它是整体相关的"。
RMSNorm 必要：Full 去掉掉到 1.743、Block 掉到 1.750。对 Block 更关键，因为块表示跨更多层累积，量级差异更大。
SWA 反直觉发现 ：只读最近 W=8 层（+embedding）远不如 Block（1.764 vs 1.746）------说明"选择性访问远层"比"访问很多近层"更重要。
缩放曲线（图 4） ：Baseline L=1.891·C^-0.057，Block AttnRes 1.870·C^-0.058，Full 1.865·C^-0.057------斜率一致但截距更低，在整个算力区间持续拉开差距。按 5.6 PFLOP-s-days，Block AttnRes 达 1.692 vs baseline 1.714，等价 1.25× compute advantage。

未解决 / 待探索的理论问题（论文 §6.1-6.2 明确列出的 future work）：

既然序列建模又被 linear attention（Mamba、GLA、DeltaNet 等）挑战，深度维度是否也有对应的"更高效 linear-complexity 深度注意"？论文明确说"incorporating more expressive yet memory-efficient (e.g. linear-complexity) alternatives is a natural direction"。
φ 可分解性揭示的规律 （§6.2）：当 φ(q,k)=φ(q)^T φ(k) 可分解时，深度注意会塌缩成递推------这正是 MRLA↔GLA、DDL↔DeltaNet 对应的结构根源。反过来，不可分解的 softmax 核是 AttnRes 保持"稠密可检索"的原因。
Block size / N 与学到权重的交互（跨规模的实证尚未系统化）。
论文明确："Attention sinks in depth" 与序列维 attention sink 同源（§6.2），但尚无理论刻画。

3.4 AttnRes 的数学机制

从"求和"改成"期望"

AttnRes 的目标不是把所有历史信号相加，而是在历史源集合 $S_{1}$ 上形成概率分布：

\ $\\alpha_{1,j}=\\mathrm{softmax}_j(s_{\\{1,j\\}}),\\quad\\tilde h_1=\\sum_{j\\in S_1}\\alpha_{1,j}v_j \\$

其中$v_j$ 是历史层或历史块表示，$s_{\{1,j\}}$ 是当前层对第$j$ 个源的打分。

伪查询、RMSNorm 与 score

文章强调一种关键工程选择：查询向量可以是每层的可学习参数，而不必由当前 hidden state 动态投影得到。这样做的收益是：查询与当前前向值解耦，允许把一个 block 中多层对历史块的查询批处理，降低反复读历史 bank 的 I/O。常见 score 形式可简化为：

$ s_{\{1,j\}}=q_1^\top\mathrm{RMSNorm}(v_j) $

RMSNorm 的作用不是改变 value，而是让 score 不被某个天然范数更大的历史块垄断。

Full AttnRes

Full AttnRes 让每层直接看所有前面层输出：$S_{l}=\{h_0,\,y_0,\,y_1,\,\dots,\,y_{l-1}\}$

优势是跨深度检索最细；缺点是每层都要保留、访问、跨 stage 传递更多历史层输出。文章指出，深度 $L$ 相比序列长度小，所以算术量本身不吓人，真正痛点是激活保存、重计算失效和流水线通信。

Block AttnRes

Block AttnRes 把层分成块。每个块内部仍维护局部前缀和；跨块只保存一个块表示。对某块内层而言，候选源约为：$S=\{h_0,\,b_0,\,b_1,\dots,b_{k-1},\,p_{\{k,t\}}\}$。其中 $b_i$ 是已完成块的表示，$p_{\{k,t\}}$ 是当前块内前若干层的部分和。

直观讲：跨楼层查档案时，不看每一页原始记录，而看每个文件夹的摘要；当前文件夹还没封档时，看当前草稿。

3.5 伪代码

3.6 真实代码

python 复制代码

r"""Block Attention Residuals (paper section 2.2, Eqs. 8-10).

A standard residual stream compresses everything a network has computed into one
vector per position and hands it to the next layer -- a depth-wise bottleneck
directly analogous to the hidden state of an RNN over time.  Attention Residuals
apply the transformer's own fix to depth: layer ``l`` *attends* over the outputs
of all preceding layers instead of accumulating them uniformly.

【中文释义】本模块实现了 Attention Residuals 的核心逻辑（论文 §2.2，Eqs. 8-10）。
它把「深度方向」看得与「时间方向」同构：标准残差流把网络算过的一切压缩成一个
逐位置的向量交给下一层------即"深度维的瓶颈"，无差异、不可选择。AttnRes 的办法
是把 attention 迁移到深度维：第 ``l`` 层用逐层可学习的伪查询向量去**注意**之前
所有层的输出，而不是等权相加。从深度维的 linear attention（mHC）升级成
「深度维的 softmax attention」，是本实现的核心立意。

    q_l  = w_l                                  learned pseudo-query, per layer
    k_i  = v_i = f_i(h_i),   k_0 = v_0 = h_1    the embedding is always a source
    phi(q, k) = exp(q^T RMSNorm(k))             RMSNorm stops a layer with large
    alpha_{i->l} = phi(q_l, k_i) / sum_j phi     outputs from dominating
    h_l  = sum_i alpha_{i->l} v_i               Eq. 9

Note what this replaces: ``h_l`` is the mixture itself, not "previous state plus
something".  Uniform accumulation is one particular setting of ``alpha``.

The full form costs ``O(L^2 d)`` arithmetic -- affordable at ``L < 100`` -- but
``O(Ld)`` memory, which is not, once every layer output must stay resident across
pipeline stages.  The block form fixes that: partition ``L`` layers into ``N``
blocks, reduce each block to a single summed representation ``b_n``, and attend
over block representations plus the running partial sum of the current block.
Memory drops from ``O(Ld)`` to ``O(Nd)``.  For K3, ``N = 8`` blocks of 12 layers
(the last partial), giving 9 sources once the embedding is counted.

【中文释义】数学结构：
- 伪查询：``q_l = w_l``（每层一个可学习向量）；key = value = 各层输出，
  且 embedding ``h_1`` 永远是第一个源（source）。
- 打分核：``phi(q,k) = exp(q^T RMSNorm(k))``，RMSNorm 防止某层输出幅值过大
  垄断权重。
- 混合：``h_l = sum_i alpha_{i->l} v_i``（Eq. 9），注意它不是"上一状态 + 新增"，
  而是**每次都用一次加权混合重新从历史中"选出"当前输入**------均匀累加只是
  ``alpha`` 的一个退化设置。

复杂度权衡：完整（Full）形式算术 ``O(L^2 d)``（``L < 100`` 时可负担）但内存
``O(Ld)``------一旦每层输出都要在 pipeline stage 间常驻内存，就无法接受。
块（Block）形式解决问题：把 ``L`` 层切成 ``N`` 块，每块塌缩成单个求和表示
 ``b_n``，注意力只发生在"块表示 + 当前块的 running partial"上，内存降到
 ``O(Nd)``。对 K3：``N = 8`` 块、每块 12 层（末块保留 partial），算上 embedding
共 9 个 source。
"""

@dataclass
class AttnResState:
    """Sources visible to the next sublayer.

    ``blocks[0]`` is the token embedding (``b_0 = h_1``), followed by one entry per
    completed block.  ``partial`` is the running sum inside the current block.

    【中文释义】跨层传播的状态机。当模型逐子层前向时，这个对象才是最真实的
    "深度历史"：
    - ``blocks``：已归档（完成）的块表示列表，``blocks[0]`` 恒为 token embedding
      （``b_0 = h_1``），其后每项是一个完成了的块塌缩表示；
    - ``partial``：当前块内"正在累计"的运行和（running sum），一旦块写满就
      并入 ``blocks`` 并清空；
    - ``units_in_block``：当前块已经累计了多少个子层单位。
    ``sources()`` 负责把"可读的历史"组装成深度注意力所见的 key/value 列表。
    """

    blocks: list[Tensor] = field(default_factory=list)
    partial: Tensor | None = None
    units_in_block: int = 0

    def sources(self) -> list[Tensor]:
        """Return every source the next sublayer can attend over.

        已完成的块 + （若存在）当前块的运行 partial。只有一个块（即只含
        embedding，尚未产出任何层输出）时就直接返回 ``blocks``。
        """
        return self.blocks if self.partial is None else [*self.blocks, self.partial]

    def detach(self) -> "AttnResState":
        """Detach a deep copy; used where gradients must not flow through history."""
        return AttnResState(
            blocks=[b.detach() for b in self.blocks],
            partial=None if self.partial is None else self.partial.detach(),
            units_in_block=self.units_in_block,
        )


class AttentionResiduals(nn.Module):
    """Owns the per-unit pseudo-queries and the block bookkeeping.

    ``n_units`` counts the addressing units: sublayers when
    ``granularity="sublayer"`` (Figure 2 draws a residual arrow per module), or
    decoder layers otherwise.  ``block_units`` is the block size measured in the
    same unit.

    【中文释义】这是 AttnRes 的完整实现。它只管理两件事：
    1. 每子层（unit）一个可学习伪查询 ``w_l``，以及 key 归一化用的 RMSNorm
       weight ------ 即 S1 打分核；
    2. Block 的台账逻辑（init_state/mix/add_output/finalize），即源列表如何
       演进（S3/S4）。
    ``n_units`` 是"可寻址单位"数：granularity="sublayer" 时按子层计（论文
    Figure 2 每模块画一条残差箭头）；否则按 decoder 层计。``block_units`` 是
    以相同单位计的块大小。
    """

    def __init__(
        self,
        dim: int,
        n_units: int,
        block_units: int,
        *,
        eps: float = 1e-6,
        full: bool = False,
    ):
        """
        Args:
            dim: hidden size（隐藏维度，``d``）。
            n_units: 可寻址的单位数（sublayer 数量）。
            block_units: 块大小，以 unit（sublayer）为单位计。
            eps: RMSNorm 的 epsilon。
            full: 若为 True，则等价于 Full 形式（测试用 oracle）。

        【中文释义】参数初始化：``queries`` 形状为 ``(n_units + 1, dim)``------每个
        unit 一个伪查询，**多出的一个**供 ``finalize`` 时在所有块表示上做最终聚合；
        ``key_norm_weight`` 是 RMSNorm 的逐维缩放权重。
        """
        super().__init__()
        self.dim = dim
        self.n_units = n_units
        self.block_units = block_units
        self.eps = eps
        self.full = full
        # One learned pseudo-query per unit, plus one extra for the final
        # aggregation over all block representations.
        # 每个 unit 分配一个可学习的伪查询，外加一个给 finalize 的最终聚合项。
        self.queries = nn.Parameter(torch.zeros(n_units + 1, dim))
        self.key_norm_weight = nn.Parameter(torch.ones(dim))

    def init_state(self, embedding: Tensor) -> AttnResState:
        """``b_0 = h_1``: the token embedding is always an available source."""
        # 【中文释义】状态初始化：block 0 固定为 token embedding，保证 embedding
        # 永远是第一个可读源；partial 为空、计数为 0。
        return AttnResState(blocks=[embedding], partial=None, units_in_block=0)

    def mix(self, state: AttnResState, unit_idx: int) -> Tensor:
        """Eq. 9 -- the representation handed to unit ``unit_idx``.

        【中文释义】深度维 softmax 注意的核心算子（对应论文 Eq. 9）：把当前可读
        源列表（``state.sources()``）叠成一个张量，其 key 统一过 RMSNorm，
        用本 unit 的伪查询 ``self.queries[unit_idx]`` 打分得到 ``alpha``（softmax
        归一化），最后加权求和解组成交给该 unit 的表示 ``h_l``。
        当只有 embedding 一个源时直接返回（避免无谓的 einsum/softmax）。
        shape 说明：source 叠成 ``(n, b, t, d)``，scores 为 ``(n, b, t)``，
        softmax 沿源维 ``n`` 归一化，输出回 ``(b, t, d)``。
        """
        sources = state.sources()
        if len(sources) == 1:
            return sources[0]
        stacked = torch.stack(sources, dim=0)                     # (n, b, t, d)
        keys = rms_norm(stacked, self.key_norm_weight, self.eps)
        query = self.queries[unit_idx]                            # (d,)
        scores = torch.einsum("nbtd,d->nbt", keys, query)
        weights = torch.softmax(scores.float(), dim=0).to(stacked.dtype)
        return torch.einsum("nbt,nbtd->btd", weights, stacked)

    def add_output(self, state: AttnResState, output: Tensor) -> AttnResState:
        """Fold a unit's output into the current block's partial sum, rolling the
        block over when it is full.

        【中文释义】接收一个子层的输出 ``output``，把它累加进当前块的运行和：
        1. ``partial`` = 已累计的输出之和（从 ``None`` 起）；
        2. ``units_in_block`` 计数 +1；
        3. 一旦达到 ``block_units``，就把这个 partial 打成包归档进 ``blocks``，
           并清空 partial/count，开始下一个块。
        这就是"块塌缩"发生的时刻------也是内存从 O(Ld) 降到 O(Nd) 的关键步骤。
        """
        state.partial = output if state.partial is None else state.partial + output
        state.units_in_block += 1
        if state.units_in_block >= self.block_units:
            state.blocks.append(state.partial)
            state.partial = None
            state.units_in_block = 0
        return state

    def finalize(self, state: AttnResState) -> Tensor:
        """Aggregate every block representation for the output layer.

        【中文释义】整个模型走完所有子层后调用：若尾巴还有个未满的 partial，先
        把它也并入 ``blocks``，然后用第 ``n_units`` 个伪查询（final 聚合项）在
        所有块表示上再调一次 ``mix``，产出最终输出层看到的表示。
        """
        if state.partial is not None:
            state.blocks.append(state.partial)
            state.partial = None
            state.units_in_block = 0
        return self.mix(state, self.n_units)

    def attention_weights(self, state: AttnResState, unit_idx: int) -> Tensor:
        """Expose ``alpha`` for analysis: which depths a unit actually reads.

        【中文释义】只算打分核而**不**做加权混合，返回形状 ``(n, b, t)`` 的
        ``alpha`` 分布------用于分析/可视化：这个 unit 在深度维上到底读了哪些源、
        权重多大。这是教学与调试（观察"选择性读取"）的窗口。
        """
        sources = state.sources()
        stacked = torch.stack(sources, dim=0)
        keys = rms_norm(stacked, self.key_norm_weight, self.eps)
        scores = torch.einsum("nbtd,d->nbt", keys, self.queries[unit_idx])
        return torch.softmax(scores.float(), dim=0)

    def extra_repr(self) -> str:
        return (
            f"dim={self.dim}, units={self.n_units}, block_units={self.block_units}, "
            f"full={self.full}"
        )


class FullAttentionResiduals(AttentionResiduals):
    """The unblocked ``O(L^2 d)`` form of Eqs. 8-9, used as a test oracle.

    Every unit output is kept as its own source, so ``block_units = 1`` recovers
    this behaviour from the block implementation -- which is exactly the
    equivalence the test suite asserts.

    【中文释义】Full 形式（Eqs. 8-9，``O(L^2 d)`` 算术）------每个子层输出都保留为
    独立 source，不做块塌缩。实现上仅靠 ``block_units=1`` 即可从 Block 实现
    复原（每收到一个输出立即归档成"块"），因此本类主要作为**测试 oracle**：
    测试套件断言两者逐位等价。
    """

    def __init__(self, dim: int, n_units: int, *, eps: float = 1e-6):
        super().__init__(dim, n_units, block_units=1, eps=eps, full=True)

3.7 伪代码 vs 本项目实现对照

论文 Fig.2 元素	本项目实现	对应位置
`block_attn_res(blocks, partial_block, proj, norm)`	`AttentionResiduals.mix(state, unit_idx)`	attn_res.py 的 mix() 方法
`V = torch.stack(blocks + [partial_block])`	`stacked = torch.stack(state.sources(), dim=0)`	attn_res.py 的 torch.stack 行
`K = norm(V)`	`keys = rms_norm(stacked, ...)`	attn_res.py 的 rms_norm 调用
`logits = einsum('d,n b t d->n b t', proj.weight.squeeze(), K)`	`scores = einsum('nbtd,d->nbt', keys, query)`	attn_res.py 的 einsum 行
`h = einsum('n b t,n b t d->b t d', softmax, V)`	`return einsum('nbt,nbtd->btd', weights, stacked)`	attn_res.py 的加权求和 einsum
`partial_block = hidden_states`（块初态）	`init_state(embedding)` → `blocks=[embedding]`	attn_res.py 的 init_state()
`if layer_number % (block_size//2) == 0: blocks.append(partial_block)`	`add_output`：`units_in_block >= block_units` 时归档	attn_res.py 的 add_output()
双 `block_attn_res`（attn 前 + MLP 前）	每 block 两次 `mix`（model.py 集成）	model.py 的 Kimi3Model.forward
`proj.weight` / `norm`（每 unit 独立投影+norm）	`self.queries[unit_idx]` / `self.key_norm_weight`	attn_res.py 的 queries 与 key_norm_weight

实现差异说明：论文伪代码把投影矩阵 proj（query 生成）与 norm 作为每个 unit 的属性传入；本项目等价地用 self.queries（(n_units+1, d) 参数矩阵，按 unit_idx 取行）+ 共享 key_norm_weight。数学上 proj.weight.squeeze()（一行）= self.queries[unit_idx]，完全等价。块边界判断的差异源于论文按 decoder 层计（block_size//2 因为每层 ATTN+MLP 两个 block_attn_res 调用），本项目按 sublayer/unit 直接计。

3.8 如何使用

在 Kimi3Model 中，给出了使用示例。

python 复制代码

class Kimi3Model(nn.Module):
    """K3 剔除 lm_head 的完整主杆：embedding + 全部 block + 最后 norm。

    额外管理一个可选的 AttentionResiduals 模块（跨层的深度方向注意力）。
    AttnRes 启用时，每个 block 被**调用两次**（block 内前半个单元 da 次），
    AttnRes 在其间插入"层间混合"（brand-residual mix）。
    """

    def __init__(self, config: Kimi3Config):
        super().__init__()
        self.config = config
        self.tok_embed = nn.Embedding(config.model.vocab_size, config.model.dim)
        # 每层交替：偶数层 KDA（use_kda=True），奇数层 MLA（use_kda=False）。
        # 这决定了 init_caches 返回的列表 K/ML 交替出现。
        self.blocks = nn.ModuleList([
            Kimi3Block(config, i, use_kda=(i % 2 == 0))
            for i in range(config.model.n_layers)
        ])
        if config.model.attn_res.enabled:
            # 启用 AttentionResiduals：n_units = 2*n_layers（每 block 两个
            # unit：attn+ffn），block_size 决定多少个 unit 共享一个残差状态。
            n_units = config.model.n_layers * 2
            block_units = config.model.attn_res.block_size
            self.attn_res = AttentionResiduals(
                config.model.dim, n_units, block_units,
                full=config.model.attn_res.full,
            )
            self._attn_res_state = None  # 前向过程中暂存的 AttnRes 状态缓存
        else:
            self.attn_res = None
        self.norm = RMSNorm(config.model.dim, config.model.norm_eps)

    def forward(
        self,
        tokens: Tensor,
        doc_ids: Optional[Tensor] = None,
        caches: Optional[list] = None,
    ) -> Tensor:
        """完整模型前向。

        ``caches`` 应与 ``self.blocks`` 等长（每 block 一个 cache）。
        解码时与服务端相同：传一个 token，返回一层 logits，内部自动调用
        cache 完成增量前向。训练时一次性给整段 tokens，cache 每层重新
        复算或者用 cache 减少重复计算。
        """
        h = self.tok_embed(tokens)
        reset_mask, attn_mask = doc_masks(
            doc_ids, tokens.shape[1], tokens.device
        )

        if self.attn_res is not None:
            # ---- 带 AttnRes 的路 =========================================
            self._attn_res_state = self.attn_res.init_state(h)
            unit_idx = 0
            for i, block in enumerate(self.blocks):
                cache = caches[i] if caches is not None else None  # 每 block 一个 cache
                # 每个 block 被调度两次，中间混入一次 AttnRes 深度混合：
                #   mix -> attn -> add_output -> mix -> ffn -> add_output
                h = self.attn_res.mix(self._attn_res_state, unit_idx)
                h = block(h, cache=cache, reset_mask=reset_mask, attn_mask=attn_mask)
                self._attn_res_state = self.attn_res.add_output(self._attn_res_state, h)
                unit_idx += 1
                h = self.attn_res.mix(self._attn_res_state, unit_idx)
                h = block(h, cache=cache, reset_mask=reset_mask, attn_mask=attn_mask)
                self._attn_res_state = self.attn_res.add_output(self._attn_res_state, h)
                unit_idx += 1
            h = self.attn_res.finalize(self._attn_res_state)
        else:
            # 不带 AttnRes：每个 block 只调用一次。cache 也是"每 block 一个"。
            for i, block in enumerate(self.blocks):
                cache = caches[i] if caches is not None else None
                h = block(h, cache=cache, reset_mask=reset_mask, attn_mask=attn_mask)
        return self.norm(h)   # 最后一层 RMSNorm，喂给 lm_head

子知识点的详解将在下一篇完成。

0x04 生产级实现与源码证据

以下为直接源码证据；「Infra 含义」为基于源码数据流的解读。本节把 SGLang 与 vLLM 两个推理框架的 AttnRes 实现合并到一张综合索引表中，避免同一份信息重复展示。

4.1 SGLang 实现

4.1.1 Python 层：`sglang-main/python/sglang/srt/layers/attn_residual.py`

这是 SGLang 中 AttnRes 的核心运行时层。文件注释明确声明：此文件实现 Kimi‑K3 Attention Residual ------ snapshot bank + aggregation；流程为 score rows → softmax → weighted sum → RMSNorm，并按硬件分发 fast/hip/fused 多条路径。dispatch(score→softmax→mix→norm)。SGLang 把 AttnRes 做成独立 runtime primitive。

证据项	位置	直接源码证据	简化逻辑	Infra 含义
TMA fast gate	行 36‑50	`_supports_attn_res_tma()` 要求 major >= 10 且不是 12; `_use_fast()` 还要求 hidden_size == 7168。	`if SM100-like and H=7168: use fast`	NVIDIA 快路径高度绑定 K3 形状与新硬件。
score 权重缓存	行 66‑83	`get_cw()` 缓存 `norm.weight * proj.weight.squeeze()`, 分别支持 bf16/fp32。	`cw = RMSNormWeight ⊙ queryWeight`	把 RMSNorm scale 与投影预乘，减少每次聚合开销。
TMA fast 调用	行 86‑117	`_aggregate_fast()` 创建 `out`, 调用 `attn_res_fused_tma(prefix_sum, bank, cw, out_norm.weight, out, nvb, eps, write_prefix=...)`。	`out = fused_tma(prefix, bank, cw, out_norm, nvb)`	fast path 把 score、online softmax、加权和、输出 RMS 融进 bank write 融合。
Triton score kernel	行 120‑155	`_score_kernel()` 对每个 token 和 source row 扫 hidden 维，算 sumsq 与 dot, 输出 `dot * rsqrt(mean_square + eps)`。	`score[j]=dot(v_j,cw)/rms(v_j)`	score 实际就是归一化后的源向量与伪查询的内积。
Triton combine kernel	行 157‑207	`_combine_kernel()` 对 scores 做 softmax, 再对 bank rows 与 prefix row 加权求和。	`p=softmax(scores); out=sum_j p_j*v_j`	深度方向 softmax 在 source 维，而非序列维。
NPU fallback	行 209‑230	`_mix_fused()` 在 `is_npu()` 时导入 `sgl_kernel_npu.kimi_k3.attn_residual.mix_fused`。	`if NPU: external mix_fused(...)`	SGLang 预留 NPU kernel 入口，但实现来自外部包。
两段 Triton fallback	行 231‑280	`_mix_fused()` / `_aggregate_fused()` 先 `_score_kernel` 再 `_combine_kernel`; `_aggregate_fused()` 再接 `out_norm(...)`。	`scores = score(); mixed = combine(); norm(mixed)`	fallback 牺牲一次中间 scores 读写，换跨硬件可用。
ROCm single kernel	行 283‑315	`_aggregate_hip()` 调用 `attn_res_hip()`, 可折叠 pending add; bank snapshot、output RMSNorm。	`hip(prefix, addend, bank, cw, out, write)`	ROCm 避免两 kernel 和重复 HBM 读取。
调度入口	行 400‑440	`_aggregate()` 空 token 直接返回；fast 优先，其次 HIP，最后 Triton fused。	`fast → hip → triton`	同一 API 下硬件选择最优路径。
bank 生命周期	行 443‑465	`AttnResidual.__init__()` 分配 `[T, block_num, H]`, 若来自上个 PP rank，则拷贝已有 `block_residual`。	`bank=new_empty; copy(old)`	bank 是跨层/跨 PP 的状态。
block 写入	行 467‑475	`write()` 把 `prefix_sum` 写入 `bank[num_valid_blocks, :]`, 然后计数加一；SP 时只写本地 rows。	`bank[nvb]=prefix; nvb++`	block boundary snapshot。
聚合 API	行 477‑539	`forward()` 处理 `nvb==0`、pending add、fast/hip fused write、fallback write。	`prefix = hidden or prefix+delta; aggregate; maybe write`	同一个函数服务 pre‑attn、pre‑MLP、输出侧聚合。
SP all‑gather 融合	行 541‑576	`forward_sp_all_gather()` 本地聚合后调用 `k3_sp_collective.attn_res_all_gather()`，可写 bank。	`local aggregate + multicast AG`	SP 场景下先 gather 的冗余。
SP reduce‑scatter 融合	行 578‑608	`forward_sp_reduce_scatter()` 调用 `k3_sp_collective.reduce_scatter_attn_res()`。	`o_proj RS + prefix add + attn_res`	把 attention 输出通信与 AttnRes 聚合合并。

伪代码（基于该文件抽象）：

python 复制代码

class AttnResidual:
    bank = empty(tokens, num_blocks, hidden)
    nvb = inherited_rows
    def forward(hidden, prefix, proj, score_norm, out_norm, write=False):
        if nvb == 0:
            if write: bank[nvb] = hidden; nvb += 1
            return out_norm(hidden), hidden
        prefix = hidden if prefix is None else prefix + hidden
        scores = dot(score_norm([bank[:nvb], prefix]), proj_weight)
        probs = softmax(scores, dim='source')
        mixed = sum(probs * [bank[:nvb], prefix])
        if write: bank[nvb] = prefix; nvb += 1
        return out_norm(mixed), prefix

4.1.2 模型集成：`sglang-main/python/sglang/srt/models/kimi_k3.py`

证据项	位置	直接源码证据	简化逻辑	Infra 含义
KDA/MLA 选择	行 2129‑2151	`config.is_kda_layer(layer_idx)` 选 `KimiK3DeltaAttention`, 否则选 `KimiK3MLAAttention`。	`layer = KDA if listed else MLA`	AttnRes 包在 KDA/MLA 层外围，和混合注意力共存。
AttnRes 参数	行 2181‑2204	若 `config.attn_res_block_size` 非空，创建 `self_attention_res_norm/proj` 与 `mlp_res_norm/proj`。	`each layer has two q/norm pairs`	pre‑attn 与 pre‑MLP 各有独立深度聚合。
forward 分流	行 2345‑2367	`forward()` 检查 `attn_res is not None` 后转入 `_forward_attn_residual()`。	`if enabled: attnres path else standard path`	标准残差路径仍保留。
聚合 1: attention 侧	行 2388‑2444	`_forward_attn_residual()` 先用 `self_attention_res_*` 调 `attn_res.forward()`; block write layer 同调用写入 bank; 写后 `prefix_sum=None`。	`x = AttnRes(prefix, bank); if boundary: bank.write(prefix)`	attention 的深度聚合后的输入。
attention 后通信	行 2450‑2503	SP‑MoE 优先 `forward_sp_reduce_scatter()`; 否则可走 `k3_ar_fusion.all_reduce(hidden_states, prefix_sum)`。	`finish o_proj reduce, maybe fold residual`	AttnRes 与 TP/SP 通信路径交织，不是纯算子替换。
聚合 2: MLP 侧	行 2504‑2519	未被 SP 融合时，使用 `mlp_res_*` 再次 `attn_res.forward()`, 然后 MLP 消费 `prefix_sum`。	`x = AttnRes(after attn, bank); y = MLP(x, prefix_sum)`	每个 Transformer block 内 attention 前和 MLP 前都可重路由残差信息。
模型级 bank	行 2640‑2648	若开启 AttnRes，计算 `attn_res_block_num`, 构造 `AttnResidual(hidden_states, block_num, block_residual=residual)` 并清空标准 residual。	`residual becomes block bank`	PP 入口的 residual 语义变为块 bank。
SP AttnRes 开关	行 2650‑2659	`sp_attn_res` 要求 AttnRes、SGLANG_K3_SP_ATTN_RES、无 PP、无 dspark capture 且 collective enabled。	`enable local shared carry only in safe mode`	源码限制了融合路径适用场景。
PP stage 输出	行 2685‑2695	非最后 PP rank 若有 AttnRes，会物化 delayed add, 并把 `attn_res.block_residual` 放进 `'PPProxyTensors'` 的 `residual`。	`send hidden_states + bank to next stage`	pipeline 边界确实要携带 bank。
最终输出聚合	行 2697‑2729	最后一段使用 `output_attn_res_proj/norm` 调 `attn_res.forward()`。	`final = AttnRes(prefix, bank, output q/norm)`	输出侧也按 AttnRes 混合，而非简单 `hidden + residual`。
aux capture	行 2740‑2767	`_dspark_capture_stream()` 对下游 consumer 的权重调用 `aggregate_stream()`, 输出 pre‑norm mixture。	`capture re what next consumer reads`	辅助隐藏态要匹配 AttnRes 语义，不能直接拿 prefix。
post‑load warm cache	行 3078‑3092	`post_load_weights()` 预热 `get_cw()` 的 bf16/fp32 缓存。	`precompute cw before cuda graph`	避免 lazy multiply 被捕获进 CUDA graph replay。

4.1.3 SP 集合通信融合：`sglang-main/python/sglang/srt/layers/k3_sp_collective.py`

证据项	位置	直接源码证据	简化逻辑	Infra 含义
初始化条件	行 54‑119	`_init_state()` 要求 SM103、TP4/TP8、MegaMoE/DeepEP、CustomAllReduceV2 multicast；可由 `SGLANG_K3_SP_COLLECTIVE` 控制。	`if hardware/topology ok: register comm`	SP 集合通信融合只在验证过的拓扑上启用。
对称 buffer	行 126‑132	`_symm_buffer()` 使用 `k3_ar_fusion.symm_buffer` 创建命名持久对称缓冲。	`named symmetric storage`	为 NVLS/direct multicast 保持跨 rank 地址一致。
reduce‑scatter + AttnRes	行 300‑360	`reduce_scatter_attn_res()` 检查 `SGLANG_K3_SP_ATTN_RES`、形状、bank/cw/ow 后调用 `attn_res.attn_res_fused_pull_rs()`。	`pull RS → local add → aggregate → norm`	把通信和深度聚合成一个局部 shard 输出。
all‑gather + AttnRes	行 413‑469	`attn_res_all_gather()` 本地聚合并将规范化向量 multicast 到全输出。	`local aggregate → direct AG`	避免完整 token 张量先聚合再通信。

相关配置：sglang-main/python/sglang/srt/environ.py 行 1476‑1492 定义 SGLANG_K3_AR_FUSION、SGLANG_K3_SP_COLLECTIVE、SGLANG_K3_SP_ATTN_RES、SGLANG_K3_GEMM_AR。其中注释明确说 SP AttnRes 在 consecutive SP‑MoE 层间保持 token‑shared residual stream，下一次 AttnRes 聚合和 bank write 在本地 shard 上运行，然后只 all‑gather 归一化后的 attention 输入。

配置入口：sglang-main/python/sglang/srt/configs/kimi_linear.py 行 45‑55 接收 attn_res_block_size, 行 109 保存到配置；行 117‑120 校验 linear_attn_config 的 KDA/full attention 层表，行 156‑168 提供 is_kda_layer()、linear_layer_ids、full_attention_layer_ids。

4.1.4 SGLang Kernel 层（TMA + HIP）

attn_res.py（kernels 目录下）：

_jit_fused_tma_module() 行 30‑61：JIT 编译 SM100 TMA kernel wrappers。NVIDIA fast path 来自 CUDA JIT。
_TMA_BEST_CONFIG / _tuning() 行 64‑97：按 nvb 选择 chunk/occupancy/registers。针对 bank rows 调参。
attn_res_fused_tma() 行 153‑190：TMA kernel 做 score、online softmax、combine、RMSNorm、可选 write。GB200/GB300 类 fast aggregation。
attn_res_fused_direct_ag() 行 226‑260：本地聚合后 direct multicast all‑gather。聚合与 AG 合并。
attn_res_fused_pull_rs() 行 263‑300：pull reduce‑scatter、local residual add、TMA aggregation。聚合与 RS 合并。

attn_res_hip.py（ROCm 路径）：

supports_attn_res_hip() 行 26‑34：用 register tile budget 判断 ROCm 是否可走单 kernel。ROCm 快路径受寄存器容量限制。
_agg_kernel() 行 37‑133：一个 CTA/token 内完成 score、softmax、mix、norm、write。ROCm 用 register tile 复用 bank。
attn_res_hip() 行 136‑212：Python wrapper 检查 nvb/shape 并 launch kernel。ROCm single kernel 入口。

4.2 vLLM 实现

4.2.1 NVIDIA 模型路径：`vllm-main/vllm/models/kimi_k3/nvidia/model.py`

证据项	位置	直接源码证据	简化逻辑	Infra 含义
operator 导入	行 100‑109	从 `vllm.models.kimi_k3.nvidia.ops` 导入 `attn_res`。	`attn_res = nvidia op wrapper`	NVIDIA 路径使用专门 op，而非普通 PyTorch 拼装。
层级配置	行 971‑996	`attn_res_block_size` 非空则设置 `is_block_write_layer`、`block_write_idx`、`prev_valid_blocks`, 并创建两套 norm/proj。	`block_idx = layer_idx // S`	与论文 Block AttnRes 的块边界一致。
pre‑attn 聚合	行 1016‑1045	`_pre_attn_norm()` 在 AttnRes 开启时调用 `attn_res(prefix_sum, hidden_states, residual, self_attention_res_norm.weight, self_attention_res_proj.weight, input_layernorm.weight, ...)`。	`x = op(prefix, delta, bank, q/norm, out‑norm)`	把 pending delta、bank 写入和输入 RMSNorm 合成一个 op。
post‑attn 聚合	行 1047‑1078	`_post_attn_norm()` 在 block write layer attention 输出重启 prefix，否则作为 delta；随后用 `mlp_res_*` 调 `attn_res()`。	`if boundary: prefix=attn_out else delta=attn_out`	块内 partial prefix 由 attention/MLP 输出共同演化。
模型级 bank	行 1350‑1361	forward 开启 AttnRes 时分配 `block_residual = new_empty(T, num_blocks, H)`, 继承 PP residual 后把 `prefix_sum=hidden_states`, `hidden_states=None`。	`residual slot becomes block bank`	vLLM 和 SGLang 一样改写 residual 语义。
PP 边界	行 1388‑1396	非最后 PP rank 返回 `IntermediateTensors(hidden_states=hidden_states, 'residual': residual)`, 若有 prefix 则先加回。	`send current hidden + block bank`	PP 必须传 AttnRes bank。
最终聚合	行 1398‑1411	最后 rank 调 `attn_res()` 使用 `output_attn_res_norm/proj`, `block_write_idx=-1`。	`final = output‑side depth mix`	输出不是标准 residual 加法。
sequence parallel	行 1332‑1408	SP 时先 `sp‑shard`, 层内 attention 需要 all‑gather, 末尾并 all‑gather。	`shard tokens, gather when needed`	vLLM NVIDIA 支持 AttnRes 与 SP 的显式编排，但代码断言 SP 不支持 PP 同时开启。
aux hidden state	行 1230‑1247、1249‑1307	`_set_aux_hidden_state_layers()` 记录 mode; `_capture_aux_hidden_stream()` 可在 `VLLM_KIMI_K3_AUX_ATTN_RES_STREAM` 开启时输出 pre‑norm AttnRes mixture，否则 `prefix‑only`。	`aux = prefix or attnres‑stream`	speculative/drafter 需要知道拿到的是哪种语义的 hidden。
final norm 位置	行 1430‑1434	注释称 final norm 在 `compute_logits` 中做，以便 MTP draft model 收到 pre‑norm hidden。	`return pre‑norm hidden`	AttnRes 输出还要配合 vLLM 的 logits/MTP 生命周期。

4.2.2 NVIDIA operator：`vllm-main/vllm/models/kimi_k3/nvidia/ops/attn_res.py`

证据项	位置	直接源码证据	简化逻辑	Infra 含义
warmup profiles	行 18‑31	`get_attn_res_triton_warmup_profiles()` 生成不同 `num_blocks/has_delta/block_write/apply_norm` 组合。	`profiles = common decode/prefill paths`	warmup 覆盖 fallback Triton 编译形态。
Triton kernel 参数	行 34‑59	`attn_res_triton()` 接收 `prefix、delta、blocks、norm_weight、qk_weight、output_norm_weight、num_blocks、block_write_idx`。	`one kernel for add/write/mix/norm`	op 直接表达 Block AttnRes 数据面。
prefix + delta	行 67‑85	若 `HAS_DELTA`, 加载 delta 加到 prefix，并写回 prefix。	`prefix += delta; store prefix`	delayed residual add 在 kernel 内物化。
block write	行 86‑94	`WRITE_BLOCK` 时把 `updated‑prefix` 写入 `blocks[:, block_write_idx, :]`。	`blocks[write_idx] = prefix`	块边界 snapshot 同内核。
source softmax	行 95‑148	`num_blocks==0` 时 softmax 退化为 prefix；否则遍历 `num_blocks+1` 源，RMSNorm score, online softmax 累加 mixed。	`mixed = online_softmax(bank rows + prefix)`	使用 online softmax 避免物化大 attention 矩阵。
output norm	行 154‑167	若 `APPLY_OUTPUT_NORM`, 对 mixed 做 RMSNorm。	`out = norm(mixed)`	与 pre‑attn/pre‑MLP 的 LayerNorm/RMSNorm 是 epilogue 融合。
native vs Triton	行 170‑248	`attn_res()` 在 H=7168、有 delta、有输出 norm、num_blocks>0、非 block write、SM100 且 `_C.kimi_k3_attn_res` 存在时走 native；否则走 Triton。	`if native‑eligible: C++/CUDA, else triton`	常见 decode 路径走 C++/CUDA，边界 / 最终输出走 Triton。
C++/CUDA 路径	`vllm-custom-ops.py` 行 2850‑2874 包装 `torch.ops._C.kimi_k3_attn_res`; `csrc/libtorch_stable/ops.h` 行 409‑418 声明在 `VLLM_ENABLE_KIMI_K3_ATTN_RES` 下编译；`csrc/libtorch_stable/kimi_k3_attn_res_kernel.cu` 行 916‑1006 检查 SM100、dense row stride，并按 `num_blocks: 1..8` 专门化 launch。		native kernel 限制清楚。

4.2.3 AMD operator: `vllm-main/vllm/models/kimi_k3/amd/linear.py` 与 `amd/ops/attn_res.py`

证据项	位置	直接源码证据	简化逻辑	Infra 含义
AMD op 导入	`amd/linear.py` 行 60‑66	导入 `vllm.models.kimi_k3.amd.ops.attn_res.attn_res`。	`use ROCm Triton op`	AMD 路径不走 NVIDIA native op。
wrapper	`amd/linear.py` 行 140‑162	`_apply_attn_res()` 把 `RMSNorm.weight`、`proj.weight.squeeze()`、可选 `output_norm` 传给 `attn_res()`。	`op(prefix, delta, bank, norm_w, q_w, out‑norm_w)`	Python 层保持和 NVIDIA 的调用语义。
层级配置	`amd/linear.py` 行 555‑579	设置 `use_attn_residuals`、block write、valid blocks，并创建 pre‑attn/pre‑MLP norm/proj。	`same block schedule`	AMD 模型层支持同一 Block AttnRes 结构。
forward_attn_residual	`amd/linear.py` 行 619‑662	先 pre‑attn `_apply_attn_res`, block 边界清 prefix; attention 后 pre‑MLP `_apply_attn_res`, MLP 后更新 `prefix_sum` 并返回。	`pre‑attn mix → attn → pre‑MLP mix → mlp → prefix`	AMD 推理链路，便于验证输出。
模型 bank/PP	`amd/linear.py` 行 810‑850	创建 `block_residual`, 继承 PP residual, 循环层；非最后 PP rank 返回 `hidden/bank`；最后 `_apply_attn_res` 输出聚合。	`bank crosses PP as residual`	PP 通信语义与 NVIDIA/SGLang 一致。
AMD Triton kernel	`amd/ops/attn_res.py` 行 16‑137	kernel 处理 prefix+delta、block write、source softmax、output norm。	`single Triton on kernel`	ROCm 路径用 Triton 直接覆盖全部核心逻辑。
AMD op wrapper	`amd/ops/attn_res.py` 行 139‑193	Python wrapper 检查 contiguous, 选择 `BLOCK‑L/warps`, launch kernel。	`num_tokens>=256 or blocks<=1 → source tile 1 else 4`	decode/prefill 使用不同 source tiling。

4.2.4 vLLM warmup、env 与配置

证据项	位置	直接源码证据	简化逻辑	Infra 含义
KimiLinearConfig	`vllm/transformers_utils/configs/kimi_linear.py` 行 45‑52、110	`attn_res_block_size` 是文本配置字段，默认 `None`。	`None disables AttnRes`	是否启用由模型 config 控制。
KimiK3Config 包装	`vllm/transformers_utils/configs/kimi_k3.py` 行 86‑131	多模态配置内嵌 `KimiLinearConfig`。	`text_config carries AttnRes`	K3 文本侧继承该配置。
env 默认值	`vllm/env.py` 行 200‑214	`VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=False`、`GEMM_AR=True`、`GEMM_RS=False`。	`aux off, GEMM‑AR on by default`	性能路径可环境开关控制。
env 解释	`vllm/env.py` 行 1597‑1612	注释说明 aux stream 改变 drafter 看到的数值，GEMM‑AR/RS 要求同一 NVLink domain。	`opt‑in numeric mode + topology limits`	AttnRes 与 speculative / 通信优化有显式风险开关。
warmup 调度	`vllm/model_executor/warmup/kernel_warmup.py` 行 34‑36、153‑155	导入并在 `enable_jit_warmup` 时调用 `kimi_k3_triton_warmup(worker)`。	`global warmup → K3 warmup`	请求前预编译 Triton kernels。
AttnRes warmup	`vllm/model_executor/warmup/kimi_k3_triton_warmup.py` 行 42‑89	`warm_attn_res()` 读取 `attn_res_block_size`, 构造零张量，对所有 warmup‑profile 调 `attn_res()`。	`compile common AttnRes profiles`	避免线上首次遇到某种 block/delta/norm 组合时 JIT 抖动。
MTP 禁用 AttnRes	`vllm/models/kimi_k3/nvidia/mtp.py` 行 88‑90; `amd/mtp.py` 行 75‑84	MTP block 浅拷贝 config 后设 `attn_res_block_size=None`。	`draft block starts without bank`	推测模型不继承主模型 AttnRes 状态，避免状态不匹配。

4.3 源码伪代码总览（基于 SGLang/vLLM 共同语义）

python 复制代码

# 模型入口
hidden = embed(input) or recv_pp_hidden
bank = empty(T, ceil(end_layer / block_size), H)
if recv_pp_bank:
    bank[:, :old_rows] = recv_pp_bank
prefix = hidden

for layer in layers:
    # 聚合 1: attention 输入
    write_idx = layer.block_write_idx if layer.is_block_write_layer else -1
    x = attn_res(prefix, delta=None_or_pending, bank,
                 layer.self_attention_res_norm.weight,
                 layer.self_attention_res_proj.weight,
                 layer.input_layernorm.weight,
                 num_blocks=layer.prev_valid_blocks,
                 block_write_idx=write_idx)

    if layer.is_block_write_layer:
        prefix = None  # 新块开始
    attn_out = attention(x)

    # 聚合 2: MLP 输入
    if prefix is None:
        prefix, delta = attn_out, None
    else:
        delta = attn_out
        x = attn_res(prefix, delta, bank,
                     layer.mlp_res_norm.weight,
                     layer.mlp_res_proj.weight,
                     layer.post_attention_layernorm.weight,
                     num_blocks=layer.prev_valid_blocks + layer.is_block_write_layer,
                     block_write_idx=-1)

    mlp_out = mlp(x)
    prefix = prefix + mlp_out

# 输出侧
hidden = attn_res(prefix, pending_mlp, bank,
                  output_attn_res_norm.weight,
                  output_attn_res_proj.weight,
                  final_norm_or_none,
                  num_blocks=all_blocks,
                  block_write_idx=-1)

4.4 与 Kimi K3 架构的协同

AttnRes 不是替代 KDA 或 MLA 的 attention 算法；它重写的是 residual stream 的输入选择。KDA/MLA 负责序列维度信息混合，AttnRes 负责深度维度信息混合，MoE/MLP 负责非线性变换和容量扩展。

源码层面可验证：SGLang 配置通过 linear_attn_config['kda_layers'] 决定哪些层是 KDA，其他为 MLA；模型构造时在 KimiK3DecoderLayer 内二选一。vLLM NVIDIA 路径也导入 KimiK3DeltaAttention 与 MultiHeadLatentAttention，并在同一 decoder layer 的外层挂 AttnRes 参数。

下一篇我们进行子知识点深度分析。

0xFF 参考

$1$ Kimi Team. Attention Residuals（技术报告）, arXiv:2602.15031, 2026. 仓库：https://github.com/MoonshotAI/Attention-Residuals
$2$ Kimi Team. Kimi K3 Technical Report §2.2（Eqs. 8-10 与 8 块×12 层=9 源配置）, arXiv:2607.24653.
$3$ Zhenda Xie et al. mHC: Manifold-Constrained Hyper-Connections, arXiv:2512.24880.
$4$ Yongyi Yang & Jianyang Gao. mHC-lite: You Don't Need 20 Sinkhorn-Knopp Iterations, arXiv:2601.05732.
$5$ Yu Zhang et al. Kimi Linear: An Expressive, Efficient Attention Architecture, arXiv:2510.26692.
参考实现（两阶段）：kimi-k3-pytorch-main/src/attention_residuals/
参考实现（fused kernel / 生产推理）：sglang-main、vllm-main 中的 K3 AttnRes 相关文件

【轻量化复刻大模型 / KIMI K3】Attention Residuals 完整梳理（1）

【轻量化复刻大模型 / KIMI K3】Attention Residuals 完整梳理（1）

0x00 概要

0.1 脉络

0.2 论文信息

0.3 关键图/表索引

0.4 注意说明

0x01 先懂 AttnRes：残差为什么是瓶颈

1.1 标准残差的展开

1.2 PreNorm 稀释与信息瓶颈

1.3 一句话痛点

1.4 思路：把治疗时间维度的药，用在深度维度上

1.5 一句话直觉

0x02 演进脉络与定位

2.1 前身：标准残差连接及其痛点

2.2 同期竞品：都是"残差路由"的玩家

2.3 深度信息流的五派方案

2.4 论文的关键消融结论

2.5 一个常见的理解误区

2.6 定位一句话

0x03 原理

3.1 直觉类比视角

3.2 准专业水平视角

3.3 专家视角（信息论 / 优化 / 计算复杂性）

3.4 AttnRes 的数学机制

从"求和"改成"期望"

伪查询、RMSNorm 与 score

Full AttnRes

Block AttnRes

3.5 伪代码

3.6 真实代码

3.7 伪代码 vs 本项目实现对照

3.8 如何使用

0x04 生产级实现与源码证据

4.1 SGLang 实现

4.1.1 Python 层：sglang-main/python/sglang/srt/layers/attn_residual.py

4.1.2 模型集成：sglang-main/python/sglang/srt/models/kimi_k3.py

4.1.3 SP 集合通信融合：sglang-main/python/sglang/srt/layers/k3_sp_collective.py

4.1.4 SGLang Kernel 层（TMA + HIP）

4.2 vLLM 实现

4.2.1 NVIDIA 模型路径：vllm-main/vllm/models/kimi_k3/nvidia/model.py

4.2.2 NVIDIA operator：vllm-main/vllm/models/kimi_k3/nvidia/ops/attn_res.py

4.2.3 AMD operator: vllm-main/vllm/models/kimi_k3/amd/linear.py 与 amd/ops/attn_res.py

4.2.4 vLLM warmup、env 与配置

4.3 源码伪代码总览（基于 SGLang/vLLM 共同语义）

4.4 与 Kimi K3 架构的协同

0xFF 参考

4.1.1 Python 层：`sglang-main/python/sglang/srt/layers/attn_residual.py`

4.1.2 模型集成：`sglang-main/python/sglang/srt/models/kimi_k3.py`

4.1.3 SP 集合通信融合：`sglang-main/python/sglang/srt/layers/k3_sp_collective.py`

4.2.1 NVIDIA 模型路径：`vllm-main/vllm/models/kimi_k3/nvidia/model.py`

4.2.2 NVIDIA operator：`vllm-main/vllm/models/kimi_k3/nvidia/ops/attn_res.py`

4.2.3 AMD operator: `vllm-main/vllm/models/kimi_k3/amd/linear.py` 与 `amd/ops/attn_res.py`