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

罗西的思考2026-11-11 20:56

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

目录

  • [0x00 概要](#0x00 概要)
  • [0x01 子知识点拆解](#0x01 子知识点拆解)
  • [0x02 子知识点深度分析](#0x02 子知识点深度分析)
    • [2.0 共享抽象与设计动机](#2.0 共享抽象与设计动机)
    • [2.1 打分核（scoring kernel）与伪查询](#2.1 打分核（scoring kernel）与伪查询)
      • [2.1.1 原理（三视角）](#2.1.1 原理（三视角）)
      • [2.1.2 对比基线](#2.1.2 对比基线)
      • [2.1.3 反方辩手](#2.1.3 反方辩手)
      • [2.1.4 计算复杂度](#2.1.4 计算复杂度)
      • [2.1.5 参数敏感性](#2.1.5 参数敏感性)
      • [2.1.6 工程实现](#2.1.6 工程实现)
      • [2.1.7 探讨代码](#2.1.7 探讨代码)
      • [2.1.8 验证 demo](#2.1.8 验证 demo)
    • [2.2 Full AttnRes](#2.2 Full AttnRes)
      • [2.2.1 原理（三视角）](#2.2.1 原理（三视角）)
      • [2.2.2 对比基线](#2.2.2 对比基线)
      • [2.2.3 反方辩手](#2.2.3 反方辩手)
      • [2.2.4 复杂度](#2.2.4 复杂度)
      • [2.2.5 参数敏感性](#2.2.5 参数敏感性)
      • [2.2.6 工程实现](#2.2.6 工程实现)
      • [2.2.7-2.2.8 探讨/验证](#2.2.7-2.2.8 探讨/验证)
    • [2.3 Block AttnRes](#2.3 Block AttnRes)
      • [2.3.1 原理](#2.3.1 原理)
      • [2.3.2 对比基线（block vs full）](#2.3.2 对比基线（block vs full）)
      • [2.3.3 反方辩手](#2.3.3 反方辩手)
      • [2.3.4 复杂度](#2.3.4 复杂度)
      • [2.3.5 参数敏感性](#2.3.5 参数敏感性)
      • [2.3.6 工程实现（论文 Fig. 2 伪代码 vs 本项目）](#2.3.6 工程实现（论文 Fig. 2 伪代码 vs 本项目）)
      • [2.3.7 探讨/验证](#2.3.7 探讨/验证)
    • [2.4 状态机（AttnResState ------ 深度历史如何在层间流动）](#2.4 状态机（AttnResState —— 深度历史如何在层间流动）)
      • [2.4.1 原理](#2.4.1 原理)
      • [2.4.2 对比基线](#2.4.2 对比基线)
      • [2.4.3 反方辩手](#2.4.3 反方辩手)
      • [2.4.4 复杂度](#2.4.4 复杂度)
      • [2.4.5 参数敏感性](#2.4.5 参数敏感性)
      • [2.4.6 工程实现](#2.4.6 工程实现)
      • [2.4.7 如果不用 AttnResState，会造成什么困难？](#2.4.7 如果不用 AttnResState，会造成什么困难？)
        • [困难 1：内存爆炸（Full AttnRes 的 O(Ld)）](#困难 1：内存爆炸（Full AttnRes 的 O(Ld)）)
        • [困难 2：退化为标准残差（失去选择性）](#困难 2：退化为标准残差（失去选择性）)
        • [困难 3：手写状态机容易出错](#困难 3：手写状态机容易出错)
        • [困难 4：失去可解释性窗口](#困难 4：失去可解释性窗口)
        • 一句话总结
      • [2.4.8 AttnResState 如何解决问题](#2.4.8 AttnResState 如何解决问题)
        • [问题 1：内存爆炸 → add_output() 的块塌缩](#问题 1：内存爆炸 → add_output() 的块塌缩)
        • [问题 2：选择性 → mix() 的深度维 softmax](#问题 2：选择性 → mix() 的深度维 softmax)
        • [问题 3：正确性 → finalize() 的兜底逻辑](#问题 3：正确性 → finalize() 的兜底逻辑)
        • [问题 4：可解释性 → attention_weights() 的暴露](#问题 4：可解释性 → attention_weights() 的暴露)
        • [总结：状态机的 4 个方法对应 4 个问题](#总结：状态机的 4 个方法对应 4 个问题)
    • [2.5 online-softmax 两阶段算法（论文 Algorithm 1）](#2.5 online-softmax 两阶段算法（论文 Algorithm 1）)
      • [2.5.1 原理](#2.5.1 原理)
      • [2.5.2 对比基线](#2.5.2 对比基线)
      • [2.5.3 反方辩手](#2.5.3 反方辩手)
      • [2.5.4 复杂度](#2.5.4 复杂度)
      • [2.5.5 参数敏感性](#2.5.5 参数敏感性)
      • [2.5.6 工程实现](#2.5.6 工程实现)
    • [2.6 生产级 fused kernel 原理（SGLang；TMA + online softmax + 通信融合）](#2.6 生产级 fused kernel 原理（SGLang；TMA + online softmax + 通信融合）)
      • [2.6.1 原理](#2.6.1 原理)
      • [2.6.2 对比基线](#2.6.2 对比基线)
      • [2.6.3 反方辩手](#2.6.3 反方辩手)
      • [2.6.4 复杂度](#2.6.4 复杂度)
      • [2.6.5 参数敏感性](#2.6.5 参数敏感性)
      • [2.6.6 工程实现](#2.6.6 工程实现)
    • [2.7 对比对象 ------ mHC（Manifold-Constrained Hyper-Connections）](#2.7 对比对象 —— mHC（Manifold-Constrained Hyper-Connections）)
      • [2.7.1 原理](#2.7.1 原理)
      • [2.7.2 对比基线表](#2.7.2 对比基线表)
      • [2.7.3 反方辩手](#2.7.3 反方辩手)
      • [2.7.4 复杂度](#2.7.4 复杂度)
      • [2.7.5 参数敏感性](#2.7.5 参数敏感性)
      • [2.7.6 工程实现](#2.7.6 工程实现)
      • [2.7.7 对比图（标准残差 vs mHC vs AttnRes）](#2.7.7 对比图（标准残差 vs mHC vs AttnRes）)
    • [2.8 基础设施：跨阶段缓存 + 高效 Prefill（论文 §4）](#2.8 基础设施：跨阶段缓存 + 高效 Prefill（论文 §4）)
      • [2.8.1 训练：跨阶段缓存（Cross-stage Caching）](#2.8.1 训练：跨阶段缓存（Cross-stage Caching）)
      • [2.8.2 并行状态下的 AttnResState](#2.8.2 并行状态下的 AttnResState)
      • [2.8.3 Cross-stage Caching 原理](#2.8.3 Cross-stage Caching 原理)
        • [① Cross-stage Caching 到底缓存什么？](#① Cross-stage Caching 到底缓存什么？)
        • [② 为什么这些可以缓存？](#② 为什么这些可以缓存？)
        • [③ 新阶段会产生什么？](#③ 新阶段会产生什么？)
        • [回到 AttnResState：缓存的工程映射](#回到 AttnResState：缓存的工程映射)
        • 一句话总结
        • [Cross-stage Caching 的精确含义](#Cross-stage Caching 的精确含义)
        • 一句话总结
      • [2.8.4 讨论](#2.8.4 讨论)
      • [2.8.5 推理：两阶段计算（Algorithm 1）](#2.8.5 推理：两阶段计算（Algorithm 1）)
      • [2.8.6 推理：内存高效 Prefill（sequence sharding）](#2.8.6 推理：内存高效 Prefill（sequence sharding）)
      • [2.8.7 反方辩手](#2.8.7 反方辩手)
    • [2.9 学到的深度注意模式（论文 §5.4.2，Fig. 8）](#2.9 学到的深度注意模式（论文 §5.4.2，Fig. 8）)
    • [2.10 统一视角：残差连接作为结构化矩阵（论文 §6.2，Table 5）](#2.10 统一视角：残差连接作为结构化矩阵（论文 §6.2，Table 5）)
      • [2.10.1 与序列-深度对偶（论文 §6.1）](#2.10.1 与序列-深度对偶（论文 §6.1）)
    • [2.11 最优架构分析（论文 §5.4.1，Fig. 7）](#2.11 最优架构分析（论文 §5.4.1，Fig. 7）)
  • [0xFF 参考](#0xFF 参考)

0x00 概要

本文为 AttnRes 的第二部分。

注：

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

0x01 子知识点拆解

AttnRes 拆解为以下子知识点（树状结构，标注依赖与项目代码位置）：

c 复制代码 
AttnRes
├── S1  打分核（scoring kernel）：φ(q,k) = exp(q^T·RMSNorm(k))
│   ├── S1.1 伪查询向量 pseudo-query w_l ∈ R^d（每层一个可学习参数，零初始化）
│   ├── S1.2  key = value = 历史层输出（含 embedding 源 h_1）
│   └── S1.3  RMSNorm 在 key 上（防大范数垄断）
│       代码：attn_res.py 的 mix() 函数
├── S2  Full AttnRes(O(L^2 d) 算术 + O(Ld) 内存，保留所有历史层）
│       代码：attn_res.py 的 FullAttentionResiduals（block_units=1 等价）
├── S3  Block AttnRes(O(Nd) 内存；分块 + partial sum 运行态）
│   ├── S3.1 块表示 b_n = Σ_{j∈B_n} f_j(h_j)
│   ├── S3.2 运行 partial sum（当前块内累计）
│   ├── S3.3 attention 范围 = 已完成的块表示 + 当前 partial（Eq. 6）
│   └── S3.4 embedding 永远是源 b_0 = h_1
│       代码：attn_res.py 的 AttentionResiduals 类（init_state/mix/add_output/finalize）
│       集成：model.py 的 Kimi3Model 集成
├── S4  状态机（服务端/推理维度：sources 的演进）
│       代码：attn_res.py 的 AttnResState 数据类
├── S5  在线 softmax 两阶段算法（分块间并行 + 块内顺序）
│       参考代码：kimi-k3-pytorch-main 的 online_softmax.py / two_phase.py
├── S6  生产级 fused kernel（SGLang；TMA + online softmax + 通信融合）
│       参考代码：sglang-main 的 attn_res.py + fused_tma.cuh
└── S7  对比对象：mHC（Manifold-Constrained Hyper-Connections）

依赖关系：S1 → S2 / S3（所有 mix 用打分核）；S3 → S4（block 演进状态）；S2/S3 → S5（两阶段加速）；S5 → S6（fused kernel）。S7 为横向对比，不依赖其它子知识点。

各子知识点的实现/变体切换入口：

变体 切换参数 位置
full vs block AttnResConfig.full（默认 False→block） config.py 的 AttnResConfig
block size AttnResConfig.block_size（默认 4） 同上
granularity sublayer vs layer AttnResConfig.granularity 同上
eager vs two_phase kimi-k3-pytorch 的 backend 字段 kimi-k3-pytorch-main 的 config.py

关于论文 48B 配置 vs K3 配置的澄清：

  • AttnRes 论文主实验（§5.2） ：基于 Kimi Linear 48B（27 Transformer blocks = 54 层），Block AttnRes 每块 6 层 → 9 个块 + embedding = 10 个深度源。
  • K3） ：K3 技术报告 §2.2 明确 "we partition its layers into 8 blocks with 12-layer size , giving a partial final block and 9 total blocks when counting the embedding layer"。K3 总层数为 93 层。
  • 需注意分清：9 源对应 K3（12 层/块） ，论文主实验是 10 源（6 层/块） 。本项目测试 test_paper_partition_for_the_2p8t_config 验证的正是 K3 的 93 层 → 8 块 → 9 源切分。

0x02 子知识点深度分析

2.0 共享抽象与设计动机

所有变体共享同一抽象：一个"源（source）"集合 + 一次"查询-键打分 + 归一化混合"。

  • 源（sources） ：V = [v_0, v_1, ..., v_{n-1}]，是历史层输出的集合 （显式保留，不压缩成一个状态）。这是与标准残差的分水岭：标准残差把历史"压成"h_{l-1}，AttnRes 把历史"留作"列表。
  • 打分 ：每层一个 w_l，对所有源求 exp(w_l^T·RMSNorm(v_i)) 后 softmax，得到 α_i。
  • 混合 ：h_l = Σ_i α_i v_i。
  • 设计动机 ：为什么把历史"留作列表"而非"压成状态"？因为"选择性读取"必须保得住"每一项"才能选------压成一个状态就丢了选择性。代价是保留全部历史的内存，这正是 Block 形式要解决的问题。

统一公式（Full / Block 都是它的实例）：

python 复制代码 
V   = [v_0, ..., v_{n-1}]        # 源列表
K   = RMSNorm(V)                 # key = value，RMSNorm 保证可比
w_l = learned query of layer l
α   = softmax(V·w_l)             # 每源权重
h_l = Σ_i α_i · v_i              # 混合结果

Block 形式唯一的差异是：V 被压缩为 [b_0, ..., b_{N-1}, partial]（块表示 + 运行 partial），即 n = N+1 个源而非 L 个。

2.1 打分核（scoring kernel）与伪查询

2.1.1 原理（三视角）

  • 直觉 ：每个源是一个"词汇"，每层是"搜索的 query"。伪查询向量 w_l 表示"这一层现在想要什么样的表示"，RMSNorm 把每个 key 缩放到同一长度，保证比的是"方向"而不是"长度"。
  • 准专业 ：α_i = exp(w_l^T·RMSNorm(v_i)) / Σ_j exp(w_l^T·RMSNorm(v_j))。w_l∈R^d 是层级参数，训练时学出「这一层该主要读哪些 early 层的输出」。
  • 专家 ：φ(q,k)=exp(q^T RMSNorm(k)) 是 inner-product 点积核，softmax 保证 α 非负且和=1，即每一层是一个「往历史层上的概率分布」。RMSNorm(k) 的作用论文有消融验证：去掉它性能下降（Full 1.743 vs 1.737，Block 1.750 vs 1.746），因为它防止了「范数大的层输出靠量级垄断权重」------也就是防"权重只看 magnitude 不看 semantic dimension"。对 Block 更关键（块表示跨层累积，量级差异更大，论文 §2.3 明确说明）。

2.1.2 对比基线

维度 标准做法（残差） 打分核（伪查询）
权重来源 固定常数 1 每层的 w_l 可学习向量
是否 input-dependent 否 是（经 v_i 内容影响）
额外参数 0 d per layer (`
归一化 无 softmax + RMSNorm(key)

2.1.3 反方辩手

  • a) 最强反对：伪查询是每层一个静态向量，无法根据输入 token 动态切换 。论文消融确实提出把 w_l 变成 input-dependent 投影（1.731，更低），但有 d×d 投影/层的参数代价和顺序访存代价。
  • b) 失败假设：假设每层最重要的历史是输入无关的 （或近似静态）。若每 layer 需要 "task / sequence" 级别的动态读取，固定 w_l 就力不从心。
  • c) 胜出条件：当"深度方向的读取模式本身比较稳定"（例如固定关注 embedding / 邻居层）时静态伪查询更省参且稳定；失败条件：需要动态（per-token）改变读取。
  • d) 更稳做法：用一个小的 input projection 替换 w_l（hybrid），保留两者优点。
  • e) 受损方：无------参数极小（每层 d 个标量），唯独增加 PP/通信（块表示跨 stage）。

2.1.4 计算复杂度

  • 增加：每层一次 w_l·RMSNorm(v_i) 是 d 点积 × n 源 = O(nd)；总 L 层 O(nLd)。
  • n 个源：Full n=L（K3 以 sublayer 计 2L）；Block n=N+1。
  • 内存开销：每层存 w_l∈R^d，各 key&value 来自前存输出，零新增。

2.1.5 参数敏感性

  • w_l 的初始值：论文 §5 明确"所有伪查询向量必须初始化为零 "------这保证初始 w=0 → α=uniform（softmax(0)=均匀） → 初始状态 = 标准残差（等权平均）！论文原文："This reduces AttnRes to an equal-weight average at the start of training and prevents training volatility, as we validated empirically." 本项目 queries 用 torch.zeros(...) 初始化正是此设计。
  • 若 w_l 初始为随机大值，会跳过均匀漂移直接学两极性，易造成层间读数剧烈变化。因此 zeros 是必要工程要点。

2.1.6 工程实现

attn_res.py 中的 mix() 方法逐行对应。rms_norm 来自同目录 norm.py。权重两个可训练参数：self.queries（形状 (n_units+1, d)，多一个 final 聚合），self.key_norm_weight（形状 (d,)）。

python 复制代码 
    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)

2.1.7 探讨代码

测试已覆盖：test_attention_weights_form_a_distribution（验证 α 非负、和为 1）与 test_key_rmsnorm_prevents_magnitude_domination（把某个源放大 1000x 后权重不变）。结论：RMSNorm 确实让权重对源 magnitude 不敏感------验证了论文论点。

2.1.8 验证 demo

tests/math/test_attn_res.py 全部通过（可与 test_key_rmsnorm_prevents_magnitude_domination 配合在 notebook 上可视化）。完成度：已实现并运行通过。

2.2 Full AttnRes

2.2.1 原理（三视角）

  • 直觉：无分块，全历史可读------"之前每一页都是独立卡片，随便读哪一张"。
  • 准专业 ：n_units = L（或 sub-layer），每个源都是单一层输出；h_l = softmax 混合所有早期层 + embedding（Eq. 4）。
  • 专家 ：M 稠密 rank-L；内容相关。O(L²d) 算术在 L<100 时可负担（论文：L<100 是 Full 形式可用的前提）；瓶颈是 O(Ld) 内存（每层都要持久保存所有层的输出，PP 下还要跨 stage 传输）。

2.2.2 对比基线

Full vs 标准残差：所有源都可见+可选择性 vs 只可见前一状态。

2.2.3 反方辩手

  • a) 反对最强理由：O(Ld) 内存让 PP 无法承受，L 大时退化。
  • b) 失败假设：假设 L 足够小（L<100），所以 O(L²) 可负担。若 L 数百，算术也可爆。
  • c) 胜出条件：L ≤ ~40 单卡（教育/小模型）；失败：PP/大数据。
  • d) 更稳做法：Block 形式（见下）。
  • e) 受损方：训练方（内存/通信）；演进为 Block。

2.2.4 复杂度

  • 算术：训练 O(L²d)（层间全对全 attention）；推理在"两阶段调度"下可降到 O(L·(S+N)·d)（附录 C）。
  • 内存：存储所有源 O(Ld) vs 标准 O(d)。
  • 参数：新增 0（只有 w_l 固定，d/L）。

2.2.5 参数敏感性

  • n_units 必须 ≤ 层数× granularity。若设大，多出的 queries 永不使用。
  • 论文：伪查询必须零初始化。

2.2.6 工程实现

FullAttentionResiduals(AttentionResiduals) 仅设 block_units=1、full=True，等价性由测试 test_block_size_one_recovers_the_full_form 断言（逐对 allclose atol=1e-6）。

可以看到，在 Kimi3Model 的前向传播中，会用 self.attn_res.add_output(self._attn_res_state, h) 往self._attn_res_state中增加模型每层生成的中间状态。

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

2.2.7-2.2.8 探讨/验证

测试 ：test_block_form_matches_the_eq10_oracle（naive 独立转写比对）；test_block_size_one_recovers_the_full_form；test_gradients_reach_every_depth；test_paper_partition_for_the_2p8t_config（93 层/12 → 8 块/9 源）。

2.3 Block AttnRes

2.3.1 原理

  • 直觉 ：把 L 层分成 N 个块，每个块"收官"时压缩成一个总表示 b_n。读历史时只读 N 个块 + 当前块的运行 partial。
  • 准专业（论文 Eq. 5-6） ：设块 n，层集合 B_n，块表示 b_n = Σ_{j∈B_n} f_j(h_j)，当前块前 i-1 层部分和 b_n^{i-1}。则
    • V = [b_0, ..., b_{n-1}]（第一层，还没有 partial）
    • V = [b_0, ..., b_{n-1}, b_n^{i-1}]（i≥2 时再附加 partial）
    • embedding b_0 = h_1。
  • 专家 ：内存从 O(Ld) → O(Nd)；1<N<L 时介于"等权相加"（N=1）和 Full（N=L）之间；论文的实验：N≈8 复原几乎全部（论文图 4 中 Full 与 Block 差距随规模缩小，最大规模仅 0.001）；K3 用 8 blocks × 12 层 → 9 源（加 embedding）。

2.3.2 对比基线（block vs full）

维度 Block Full
内存 O(Nd) O(Ld)
算术 O(L·N·d) O(L²d)
精度 最接近 最优
针对性 面向训练内存/通信 参照（理论最优）

Block size 扫描（论文 Fig. 6，16 层模型）：S 从 1（即 Full）粗化到 32，loss 优雅退化：

Block size S 1 (Full) 2 4 8 16 32
Loss 1.737 1.746 1.746 1.746 1.757 1.753

S=2/4/8 都落在 1.746 附近，S=16/32 靠近 baseline（1.766）。论文因此固定 ≈ 8 个块。

2.3.3 反方辩手

  • a) 反对：块内压缩丢失层间差异（块内所有层共享同一个 key/value），block 是"近似 full"的折扣------本质上是 N 的选择。
  • b) 失败假设：假设块内所有层输出可以被同等压缩（相加）；极端情况可能丢信息。
  • c) 胜出条件：随着模型变大 L 越来越大、块数 N 固定（≈8）→ 每个块更"稠密"，信息损失越来越小；失败：块过小（内存仍大）或块太大（早期块太稠密，读到的已是压缩）。
  • d) 更稳做法：soft-block 或轻量跨块（论文未做）。
  • e) 受损方：训练（PP 通信必须在块边界）。

2.3.4 复杂度

  • 内存：O(Nd)（K3: 8 块）。
  • 算术：每层 O(N d)（attention over blocks），总 O(L N d)。
  • 对比 Full：算术 O(L²) → O(LN)，N=8 时 ~L/8 倍省（L=93 时 ~11.6×）。

2.3.5 参数敏感性

  • block_size（每块层数）：本项目 block_size 直接当 block_units（默认 4），kimi3-master 用 block_units = block_size * per_layer（per_layer=2，默认 24）。论文建议 block_size≈12（层）/ 24（单元）对应 K3。
  • 论文经验：N≈8「recover most of the benefit」；增加 block count 一般会提升但受限于通信/内存。

2.3.6 工程实现（论文 Fig. 2 伪代码 vs 本项目）

论文 Fig. 2 的关键伪代码：

python 复制代码 
def block_attn_res(blocks, partial_block, proj, norm) -> Tensor:
    # blocks: N tensors [B,T,D]; partial_block: [B,T,D] (intra-block partial sum)
    V = torch.stack(blocks + [partial_block])          # [N+1, B, T, D]
    K = norm(V)
    logits = torch.einsum('d, n b t d -> n b t', proj.weight.squeeze(), K)
    h = torch.einsum('n b t, n b t d -> b t d', logits.softmax(0), V)
    return h

def forward(self, blocks, hidden_states):
    partial_block = hidden_states
    # apply block attnres before attn
    h = block_attn_res(blocks, partial_block, self.attn_res_proj, self.attn_res_norm)
    # if reaches block boundary, start new block (block_size counts ATTN+MLP)
    if self.layer_number % (self.block_size // 2) == 0:
        blocks.append(partial_block); partial_block = None
    attn_out = self.attn(self.attn_norm(h))
    partial_block = partial_block + attn_out if partial_block is not None else attn_out
    # apply block attnres before MLP
    h = block_attn_res(blocks, partial_block, self.mlp_res_proj, self.mlp_res_norm)
    mlp_out = self.mlp(self.mlp_norm(h))
    partial_block = partial_block + mlp_out
    return blocks, partial_block

与本项目 attn_res.py 的逐行对应：

  • block_attn_res ↔ mix()：V=stack(sources)、K=rms_norm、logits=einsum('d,nbtd->nbt')、h=einsum('nbt,nbtd->btd') 完全同构。
  • 块边界判断 layer_number % (block_size//2) == 0 ↔ add_output 的 units_in_block >= block_units。
  • 双 block_attn_res 调用（attn 前 + MLP 前）↔ model.py 的每 block 两次 mix。

集成 ：model.py 中的 Kimi3Model ------ 先把 h 进 init_state，对每个 block 先 mix(state, unit_idx)（把历史混合作为该层输入），执行 block 后 add_output。granularity="sublayer"（默认）时，每个 transformer 层有 attention 读 + MLP 读两个单位，故 n_units = n_layers*2；

不同库的块单位换算不同 ：Model-Viz 把 block_size 直接当 block_units（默认 4），而 kimi3-master 用 block_units = block_size * per_layer（per_layer=2 → 默认 24）。

2.3.7 探讨/验证

test_block_form_matches_the_eq10_oracle（naive 独立转写比对）；test_memory_is_bounded_by_block_count；test_gradients_reach_every_depth；test_paper_partition_for_the_2p8t_config（93 层/12 → 8 块/9 源）。全部通过。完成度：已实现并运行通过。

Q1：activation recomputation 为什么被影响

标准 activation checkpointing 可以丢掉中间激活，反向时从 checkpoint 重算。但 Full AttnRes 要让后续很多层继续访问旧层输出；如果丢掉它们，后续层无法直接算深度注意力，或者必须频繁重算历史输出。于是它把"可释放的中间态"变成"跨多层被读取的 KV‑like 源"。

Block AttnRes 的妥协是只保存每个块的聚合表示和当前块 partial prefix。这样 activation checkpointing 仍可丢掉 softmax 中间量和大量块内细节，保留少量块级 bank。

Q2：pipeline parallelism 为什么变贵

普通 PP stage 边界只传 hidden_states 和常规 residual。AttnRes 需要下游 stage 也能看到历史块，所以边界还要传 block bank。

若 Full AttnRes 传每层历史，成本随层数涨；Block AttnRes 传块表示，成本随块数涨。

论文提出 cross‑stage caching: 同一物理 stage 在交错流水线的不同 virtual stage 会重复看到部分历史块，因此可缓存已见块，只传新增块，减少冗余通信。

2.4 状态机（AttnResState ------ 深度历史如何在层间流动）

AttnResState 是 Block Attention Residuals 的跨层状态机 ，解决的核心问题是：如何在 O(Nd) 内存约束下（而不是 O(Ld)），让每一层都能"选择性读取"之前所有层的输出。

2.4.1 原理

问题：attn 需要在遍历 93 层时"记住"哪些源可用、块进行到哪一步、partial 是什么。

难点：这是一个"顺序推进的、跨层的可变状态"，且需要支持 full/block 两模式。

思路 ：用一个可变 state 对象，每过一层由 add_output 更新。blocks = 已完成的块表示列表；partial = 当前块内运行和；units_in_block = 当前块已累积的单位数。sources() 根据 partial 是否为 None 组装可读列表。

原理：把 block 折叠的时机编码进状态机，而非 per-layer 的全局逻辑。

下图给出 AttnResState 在 AttentionResiduals 中的状态演进（Block 形式）：init → 逐子层主循环（mix/block/add_output）→ 满块折叠 → finalize。

循环：每个 unit 走一遍 mix → 输出 → add_output

  • partial 累计本块输出，units_in_block 计数
  • 满块：partial 归档进 blocks\[\]，partial 置 None、计数清零（O(Ld)→O(Nd)）
  • finalize 把尾部不全块并入后做最终 mix

2.4.2 对比基线

标准残差没有"状态"------只用局部 h_{l-1}。

AttnRes 需要显式的跨层状态。这是「深度上像 RNN 一样需要 carry 状态」的体现，与序列建模中 RNN 隐状态是同构的。

2.4.3 反方辩手

  • a) 反对：状态机的存在让并行（PP/SP）变麻烦------状态必须跨 stage 传递。论文用 cache-based pipeline communication 缓解（见 2.8）。
  • b) 失败假设：假设状态很小（N 个块 + 1 partial），可被高效序列化；当 N 变大时仍要承担。
  • c) 胜出条件：N 固定小（8）；失败：大 N。
  • d) 更稳做法：把状态换成"固定长度 buffer"（预分配）。
  • e) 受损方：推理/训练基础设施（需同步状态）。

2.4.4 复杂度

内存：state = O(N·B·T·d)（N 个块表示）。partial 为 1×B×T×d。总 O(N·B·T·d) 与 Full O(L·B·T·d) 相比 N≪L。

2.4.5 参数敏感性

block_units 直接决定 units_in_block 的翻转阈值。若设太小，partial 频繁折叠成 block（→ 退化为 Full 般大 N）；设太大，块表示太稠密（选择性降低）。论文经验：K3 取 12 层/块（24 单元）。

2.4.6 工程实现

AttnResState管理两个数据结构：

字段 含义 内存复杂度
blocks: list[Tensor] 已完成 的块塌缩表示（blocks[0] 恒为 token embedding） O(Nd)，N=块数
partial: Tensor 当前块正在累计的运行和 O(d)
units_in_block: int 当前块已累计的子层数 O(1)

每过一个子层调用一次 add_output()：把该子层输出累加进 partial，块写满后把 partial 归档进 blocks 并清空------这就是"块塌缩"发生的时刻，把 L 层输出压缩到 N 个块表示。

init_state(embedding) 设 blocks=[embedding]、partial=None；add_output 折叠逻辑；finalize 把残余 partial 也 append 后 mix。对应 kimi3-master 完全一致。

一屏示意图（对应全过程的状态演进）：

python 复制代码 
┌─────────────────────────────────────────────────────────────┐
│  AttentionResiduals 数据流（AttentionState --- Block 形式）     │
└─────────────────────────────────────────────────────────────┘

  ① init_state(h_embed)                        states：blocks = [h_embed]
        inputs: partial = None, units_in_block = 0
        （embedding 恒为第 0 个源）
            │
            ▼
  ② for each unit_idx in [attention 子层, FFN 子层, ...]      # n_units
            │
            ▼
  ③ mix(unit) ：本 unit 的 深度 softmax 注意力
        sources   = state.sources()          # blocks + [partial]（partial None 时仅 blocks）
        keys      = RNorm(sources)           # 防"嗓门大"的源垄断权重
        alpha     = softmax( w_l · keys )    # 对"源"维归一化（Eq.9）
        h_l       = sum_i alpha_i · sources_i
            │
            ▼
  ④ 执行当前 sublayer（按真实模型跑正常 sublayer）
        out = sublayer( h_l )                # attention 层 / FFN 层
            │
            ▼
  ⑤ add_output( out )
        partial        += out
        units_in_block += 1
            │
            ▼
  ⑥ units_in_block == block_units ?   ── 满块？
        ├─ 否 ──────────────────────────────► 回到 ③（本块下一个 unit）
        └─ 是 ──── 折叠归档：
              blocks.append( partial )
              partial = None ; units_in_block = 0
              ▼
          （下一块第一个 unit 从 ③ 开始）
            │
            ▼ （循环结束后）
  ⑦ finalize()
        if partial is not None: blocks.append( partial)
        out = mix( final 伪查询 )            # 对全部 blocks 再混合一次
            │
            ▼
  ⑧ 返回末层输出 h_out

直观解读：③ mix 是"读"（深度维度 attention），④ 是"处理"（真实模型），⑤ 是"写"（折入状态） ；⑥ 决定何时把「块与块之间的运行和 partial」归档成「不可变的块表示」。这就是 AttnResState 这台"深度流水线状态机"的完整轨迹，与 online-softmax 两阶段加速完全对应。

2.4.7 如果不用 AttnResState，会造成什么困难？

困难 1：内存爆炸（Full AttnRes 的 O(Ld)）

如果不分块（Block），每层都要保留之前所有层的完整输出：

python 复制代码 
# 不用 AttnResState（朴素 Full AttnRes）
all_outputs = [embedding]  # 每加一层 append 一次
for layer in layers:
    h = attention_mix(all_outputs)  # 对所有历史做 softmax
    all_outputs.append(layer(h))
# 内存 = L * d，L=96 层时 96 份完整激活

K3 用 96 层，96 份 hidden state 常驻内存 → pipeline stage 间通信量爆炸，GPU 内存吃不消。

困难 2：退化为标准残差（失去选择性）

如果根本不做"跨层注意"，就退化成普通残差连接：

python 复制代码 
# 标准残差（无 AttnRes）
h = embedding
for layer in layers:
    h = h + layer(h)  # 无差别累加

这等价于 alpha_i = 1/L（均匀权重）------每一层对历史的影响完全相同，模型无法学习"第 30 层应该多读 embedding、少读第 5 层"这种选择性。

困难 3：手写状态机容易出错

Block 塌缩的边界条件很微妙：

  • 何时归档 partial？（满 block_units 时）
  • 最后一个不满的 partial 怎么办？（finalize() 兜底）
  • 如何对 gradient checkpoint 友好？（detach() 深拷贝）

手写这些逻辑每个模型都要重新实现一遍，容易写出 off-by-one 错误（比如少归档最后一个 partial、或 partial 没清空）。

困难 4：失去可解释性窗口

AttnResState.attention_weights() 暴露 alpha 分布用于分析/可视化------能看到"第 L 层到底读了哪些源"。没有这个抽象，调试"为什么模型学不出选择性"就只能盲猜。

一句话总结

AttnResState 是用 O(Nd) 内存实现"深度维 softmax 注意力"的状态机，它解决了：

问题 不用它 用它
内存 O(Ld) = 96 份激活 O(Nd) = 8 份激活
选择性 退化为均匀残差 每层学习专属 alpha
正确性 手写状态机易错 统一的 init/mix/add/finalize
可解释性 黑盒 可读 alpha 分布

2.4.8 AttnResState 如何解决问题

FullAttentionResiduals 就是 block_units=1 的特例------证明 Block 和 Full 在数学上等价，只是内存策略不同。这是理解整个设计的钥匙。

问题 1：内存爆炸 → add_output() 的块塌缩

代码 ：attn_res.py:178-195

python 复制代码 
def add_output(self, state: AttnResState, output: Tensor) -> AttnResState:
    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

机制 ：每个子层的输出先累加进 partial（running sum），写满 block_units 个后才"归档"为 blocks 中的一个条目。K3 用 block_units=12，意味着 12 个子层输出被压缩成 1 个块表示。

数学等价性 ：blocks[n] = Σ_{i∈block_n} output_i（求和塌缩）。虽然求和有损，但实验表明它保留了足够信息让后续的 softmax 注意仍能区分块。

内存对比（K3, d=7168）：

形式 内存 96 层 × 7168
Full L × d 688 KB/token（爆炸）
Block (N=8) N × d 57 KB/token（可接受）
问题 2：选择性 → mix() 的深度维 softmax

代码 ：attn_res.py:157-176

python 复制代码 
def mix(self, state: AttnResState, unit_idx: int) -> Tensor:
    sources = state.sources()  # [emb, b_1, ..., b_n, partial?]
    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)   # 防垄断
    query = self.queries[unit_idx]                   # (d,) 每层专属
    scores = torch.einsum("nbtd,d->nbt", keys, query)
    weights = torch.softmax(scores.float(), dim=0)   # 沿源维归一化
    return torch.einsum("nbt,nbtd->btd", weights, stacked)

机制 ：每个 unit（子层）有独立的可学习伪查询 self.queries[unit_idx]，对所有可见源（已完成块 + 当前 partial）打分，softmax 归一化后加权混合。

这就是"选择性"的物理实现：

  • queries[unit_idx] 高维向量是可学习的"地址"
  • 训练后模型可能学到：第 30 层应该多读 blocks[0]（embedding），少读 blocks[3]
  • 标准残差只是 weights = 1/n 的特例 ------mix() 把它推广为可学习的分布

RMSNorm 的作用 （attn_res.py:172）：防止某层输出幅值过大导致 scores 被它垄断，保证 softmax 不塌缩成 one-hot。

问题 3：正确性 → finalize() 的兜底逻辑

代码 ：attn_res.py:197-208

python 复制代码 
def finalize(self, state: AttnResState) -> Tensor:
    if state.partial is not None:           # ★ 兜底：尾巴的 partial 也要归档
        state.blocks.append(state.partial)
        state.partial = None
        state.units_in_block = 0
    return self.mix(state, self.n_units)    # 用第 n_units 个 query 做最终聚合

机制 ：K3 有 96 层 ÷ 12/块 = 8 个完整块，但如果改成 100 层 ÷ 12/块 = 8 块 + 4 层尾巴，第 9 个块是不满的 。finalize() 检查 partial is not None，把这个不满的尾巴也归档进 blocks，确保所有子层输出都被纳入最终聚合，一个不漏。

这就是手写状态机容易出错的地方：

  • 忘记归档尾巴 → 最后 4 层输出被丢弃
  • 忘记清空 partial → 下次 forward 状态污染
  • 忘记用 n_units 作为 query index → 用了某个中间层的 query 做最终聚合

AttnResState 把这 3 个坑全部封装在 finalize() 里。

问题 4：可解释性 → attention_weights() 的暴露

代码 ：attn_res.py:210-221

python 复制代码 
def attention_weights(self, state: AttnResState, unit_idx: int) -> Tensor:
    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)  # (n, b, t)

机制 ：与 mix() 共享前 4 行 （打分核完全相同），但不做加权求和 ，直接返回 alpha 分布。形状 (n, b, t) 告诉你"unit_idx 这个子层，对每个源、每个 batch、每个 token 位置分配了多少注意力权重"。

这就是 debug 的窗口：

python 复制代码 
state = attn_res.init_state(embedding)
for i, layer in enumerate(layers):
    h = attn_res.mix(state, i)
    out = layer(h)
    state = attn_res.add_output(state, out)
    if i % 12 == 0:  # 每块结束看一次
        alpha = attn_res.attention_weights(state, i)
        print(f"layer {i}: alpha over {len(state.sources())} sources = {alpha}")

没有这个方法，你只能看 mix() 的最终输出，完全不知道模型在深度维上"读"了哪些层。

总结：状态机的 4 个方法对应 4 个问题
方法 行号 解决什么问题 关键操作
init_state() 151-155 初始化（embedding 永远是源） blocks=[embedding]
mix() 157-176 选择性（softmax 深度注意） queries[unit_idx] × RMSNorm × softmax
add_output() 178-195 内存（块塌缩 O(Ld)→O(Nd)） 满 block_units 时归档 partial
finalize() 197-208 正确性（兜底尾巴） 不满块也归档 + 最终聚合
attention_weights() 210-221 可解释性（暴露 alpha） 只打分不混合

Block 和 Full 的等价性 （attn_res.py:230-244）：FullAttentionResiduals 就是 block_units=1------每个输出立即归档成"块"，数学上和 Block 完全等价，只是内存策略不同。这证明了块塌缩不损失表达能力，只是工程上的内存优化。

2.5 online-softmax 两阶段算法（论文 Algorithm 1）

2.5.1 原理

问题 ：Full/Block AttnRes 在推理时是逐 token 的（decode）。要算 h_l = softmax(w_l·K)·V，理论上需要完整 softmax（读所有源）。但 Block 形式下，「已完成块的表示」是固定不变的（前块已经写完，不会变），只有「当前 partial」是逐 token 在变。

难点：不想每 token 都对所有块重新做完整 softmax。

思路（两阶段）：

  • Phase 1（跨块，可并行） ：对当前块的所有 S 个 site（该块内的各子层），一次性算出它们对所有已完成块的 logits → 得到每 site 的「块间统计量」（online-softmax 的 (max, sum-exp, weighted-sum)）。
  • Phase 2（块内，顺序）：每个 site 只需要额外把「当前 partial」的 logit 与 Phase1 统计量 merge（在线合并），再 normalize → 得到最终输出。因为 partial 每 token 变，而块间统计量复用，decode 成本大幅降低。

专家 ：这是把 softmax 拆成「可复用的块间 part + 变化的块内 part」的 online softmax 技巧，与 flash attention 的同款思想（把 attention 分解成可 merge 的 (m, l, Σ) 三元组）。

三个库对比：

  • kimi3-master：朴素 torch（无两阶段）。
  • kimi-k3-pytorch：实现了完整 online-softmax 两阶段（online_softmax.py + two_phase.py），且 BlockAttentionResidualController 提供 backend="eager"|"two_phase" 开关。
  • SGLang：在 CUDA kernel 里把「对 bank 的 online softmax」直接 fused 进单 kernel。

2.5.2 对比基线

维度 eager two_phase（online）
对已完成块的重复计算 每次 site 都重算 只算一次（Phase1 向量化）
需要的数据结构 sources 列表 (max, Σexp, weighted_sum) 统计量
decode 单 token 成本 O(N·d)（每 token 重扫全块） O(d)（只 merge partial）

2.5.3 反方辩手

  • a) 反对：多了一套统计量状态，复杂；且 phase1 需要在块开头预先跑所有 site 的块间 logits（内存 = 块内 site 数 × 块数 × B×T）。
  • b) 失败假设：假设「已完成块」在 decode 期间不参与 partial 的在线更新；若 partial 与已完成块之间存在耦合（例如分层 RMSNorm），online 合并的数学成立性可能受考验。
  • c) 胜出条件：decode 场景、长序列、块数 N 固定。失败：训练/prefill 全量序列（full softmax 直接算即可，省复杂度）。
  • d) 更稳做法：放弃两阶段，每 token 重算（简单可靠，但 O(N) 每 token）。
  • e) 受损方：实现者/推理引擎（复杂度）；换来性能。

2.5.4 复杂度

  • eager decode：每 token 每 site 重扫 N 个块 → O(N·B·T·d)（T 为 batch tokens）。
  • two_phase decode：Phase1 每块一次 O(N·B·T·d)，Phase2 每 token 每 site O(d)。节省量随「一个块内的 site 数 S」增大而增加。
  • 论文报告：inference overhead < 2%（典型 workload，Table 1）。

2.5.5 参数敏感性

  • backend 开关（eager/two_phase）直接控制此算法。S 越大（块内 site 越多），two_phase 收益越大。
  • logits_in_fp32 / weighted_sum_in_fp32（kimi-k3-pytorch config）控制精度策略------默认 fp32 累加，保证两阶段 merge 的数值稳定性。

2.5.6 工程实现

代码（参考库） ：kimi-k3-pytorch-main 的 online_softmax.py（depth_softmax_stats / single_source_stats / merge_depth_softmax_stats / normalize_depth_softmax_stats / weights_from_stats），以及 two_phase.py（precompute_inter_block_stats / score_single_partial）。

本项目对应：本项目采用 kimi3 风格（朴素 torch 实现），未实现两阶段。若要参考生产做法，可对照上述文件。

完成度：参考实现已提供（kimi-k3-pytorch），本项目未实现（采用朴素 torch，与 kimi3 一致）。

2.6 生产级 fused kernel 原理（SGLang；TMA + online softmax + 通信融合）

2.6.1 原理

问题 ：K3 真实部署时，AttnRes 每 token 要对 bank（最多 8 个块表示快照 + 1 个 prefix）做 online softmax 聚合 + 输出 RMSNorm。在 GB200/SM100 上这是高频、低算量、高访存操作，纯 PyTorch 逐层调用会卡在 kernel launch / 内存带宽。

难点：把「读 bank 行 → score → online softmax → weighted combine → 输出 RMSNorm」整条链 fused 成一个 kernel，还要处理 TP（张量并行）下的通信。

思路（SGLang）：

  • 单 kernel：warp-specialized TMA kernel。一个 persistent CTA/SM：producer warp 用 bulk TMA 一次拉一行 bank 到 smem ring（2 个 chunk slot 双缓冲），8 个 consumer warp 做 score + online softmax fold + 写 TMEM。
  • 参数 ：bank[T, NB, H]（NB=最多块表示数=8+1，H=7168），cw（score_norm_weight×proj_weight 预乘），ow（输出 RMSNorm weight），eps（两个 norm 共用）。nvb = 有效 bank 行数（1..8）。
  • 通信融合 ：attn_res_fused_pull_rs（NVLS pull 规约 + local residual + fused）与 attn_res_fused_direct_ag（直接 multicast all-gather 融合）------把 AttnRes 聚合与 TP all-reduce 合二为一，省一次整张往返。

专家视角 ：这是「online softmax 状态 (m,l,Σ) 天然可 merge」特性在 kernel 层的极致利用：bank 行是"已完成块"，其统计可被并行 fold；write_prefix 把当前 prefix 行也写回 bank，支持后续复用。_TMA_BEST_CONFIG 是每 nvb 的调参表（GB200/GB300 实测），occupancy/consumer_regs 是 CTA 调度权衡。

2.6.2 对比基线

维度 朴素（kimi3 torch） SGLang fused TMA
kernel 数 多次（norm/score/softmax/mix/norm） 1 个 fused kernel
通信 独立 all-reduce fused 进 kernel（pull_rs/direct_ag）
内存 多次 HBM 往返 TMA 批量 + TMEM staging，最少往返
平台 任意 SM100a+（B100/B300/GB200 等）
参数限制 任意 d H==7168 模板参数

2.6.3 反方辩手

  • a) 反对：这套 kernel 复杂度极高，且 H=7168 硬编码、SM100 专用------不可移植、难维护。
  • b) 失败假设：假设"每 token 聚合 bank"值得如此大的 kernel 工程投入；若该算子只占 <1% 时间，就是过度工程。
  • c) 胜出条件：大 batch、多 token、SM100 集群（GB200 为主力时）；失败：小规模 / 非 SM100。
  • d) 更稳做法：用 Triton 通用 kernel 折中（本项目可接受）。
  • e) 受损方：运维方（版本约束、架构绑定）。

2.6.4 复杂度

  • HBM 流量：bank 读 = nvb×H 每 token；输出 = H。TMA 批量把读合并。通信：pull_rs/direct_ag 把 all-reduce 从独立 O(T·H) 变 fused（省一次 HBM 写读）。
  • 具体数字需跑 benchmark；论文报告 Block AttnRes inference overhead <2%。

2.6.5 参数敏感性

  • nvb（有效块数）、chunk_rows（每 chunk 行数）、occupancy（CTA/SM）、consumer_regs。_TMA_BEST_CONFIG 表（sglang attn_res.py）给出每 nvb 的最优 (chunk_rows, occupancy, consumer_regs)。
  • write_prefix：是否把 prefix 行回写 bank（后续可作 prefix 复用）。

2.6.6 工程实现

代码（参考） ：sglang-main 的 python/sglang/srt/layers/attn_residual.py（Python 层调度，300+ 行）+ python/sglang/srt/kernels/kimi_k3/attn_res.py（TMA kernel wrapper）+ jit/csrc/kimi_k3/attn_res/fused_tma.cuh（CUDA 内核，946 行）。

本项目：本项目采用朴素 torch（与 kimi3 一致），未实现 fused kernel；本文档将其列为「生产级参考」。

完成度：参考实现已提供（SGLang），本项目未采用。

2.7 对比对象 ------ mHC（Manifold-Constrained Hyper-Connections）

2.7.1 原理

问题 ：HC（Hyper-Connections）把残差变成可学习的多条平行流，但丧失了恒等映射性质（identity mapping 被破坏），训练不稳定、难扩展。

mHC 的解决 ：把残差映射矩阵投影到 Birkhoff 多胞体（双随机矩阵） ------行列和都等于 1。这样 P x 成为输入特征的凸组合（权重非负、和=1），特征均值守恒、信号范数被严格限制，组合闭包保恒等稳定性。

  • 直觉：HC 是"多车道"，mHC 给每车道加了"所有车道权重加和为 1"的护栏（每列、每行都归一），既保住了多车道的信息流通，又不让信号越飙越大。
  • 准专业 ：H_l^{res} x_l 其中 H_l^{res} 是双随机矩阵 （非负、行和=1、列和=1）。投影用 Sinkhorn-Knopp 算法（迭代行/列归一化）把任意矩阵投影到 Birkhoff 多胞体。
  • 专家 ：m 条平行流 H_l ∈ R^{d×m}，更新 H_l = H_{l-1} A_l + f_{l-1}(H_{l-1} α_{l-1}) β_{l-1}^T。展开后有效权重 M_{i→l} = β_i^T A_{i+1→l} α_l，是 m-semiseparable 的深度维 linear attention （state expansion 到 d×m，提高 semiseparable rank）。双随机约束使 A_l 乘积 A_{i→l} 保持性质（范数守恒、composition closure），恢复 identity 稳定性。

AttnRes 与 mHC 的本质差异（按论文 Table 5 与 §6.2）：

  • 两者都是 input-dependent（动态） 。真正的区别是注意力范式 ：
    • mHC：深度维 linear attention （m-semiseparable，低秩）。权重是「学习 + 双随机约束投影」得到的线性变换 （M_{i→l}=β^T_i A_{i+1→l} α_l），softmax 的竞争性归一化缺失------各源概率质量不竞争。
    • AttnRes：深度维 softmax attention （稠密 rank-L）。权重 α_{i→l} 是 softmax 归一化的 input-dependent （直接对 v_i 内容做 exp 打分），概率质量在源间竞争，形成更锐利的选择。
  • 论文 Table 4 中 AttnRes（1.737/1.746）优于 mHC（1.747）的原因就是：softmax 的竞争性归一化 + 跨层直接检索。

mHC 会维护多条并行的 Residual Stream。每一层先动态混合这些信息流，形成这一层的输入，再把这一层的输出分发回不同的信息流。它还会约束混合过程，避免模型变深后信息和梯度被无限放大。

Attention Residuals 则更像为每一层提供一份历史目录，让它可以直接找到前面某个阶段产生的信息，并选择性地读取。

只看理论表达能力，Attention Residuals 的上限可能更高。mHC 仍然要把历史信息递归压缩进固定数量的 Residual Stream；Full Attention Residuals 则可以让每个 token 直接关注所有较浅层的输出。打个比方，mHC 更像 RNN（循环神经网络）或线性注意力，信息通过有限状态逐层传递；Full Attention Residuals 更像 Transformer，可以直接跨层读取信息。

但实际效果仍取决于具体实现。K3 使用的并不是所有层之间直连的 Full Attention Residuals，而是 Block Attention Residuals：先把模型层划分成多个 block，在每个 block 内汇总信息，再在不同 block 之间做 Attention。

2.7.2 对比基线表

维度 Standard Residual mHC AttnRes (Block)
更新公式 h_l = h_{l-1} + f_{l-1}(h_{l-1}) H_l 多流 + 连接矩阵 h_l = Σ α_{i→l} v_i
状态数量 1 m（如 4 条流） 1 伪查询 + N 块表示
权重来源 固定单位 1 学习 + 双随机投影（Sinkhorn） 学习 + 输入相关 softmax
是否 input-dependent 否 是（Dynamic，Table 5） 是
注意力范式 深度维等权加法 深度维 linear attention（m-semiseparable） 深度维 softmax attention（dense softmax）
恒等稳定性 好 双随机保证（identity 保真） zeros 初始化平滑出发（等效均匀）
参数开销 0 中等（连接矩阵 A + 转换 α/β） 低（每层一个 w_l）
内存 I/O ~3d/层 ~34d/层（m=4） ~5.5d/层（Block）
性能（16 层，论文 Table 4） 1.766 1.747 1.737 / 1.746
并行 friendly 是 一般 Block 形式友好（块边界通信）
可解释性 无 一般 好（α 可直接可视化）
亲密度（本项目） 基线 未实现（仅对照） 已实现

三条路线并排示意图（ASCII，重点：传给"深度历史"的信息流差异）：

2.7.3 反方辩手

  • a) 反对 AttnRes：固定伪查询无法做 input-dependent query；mHC 至少在结构上是动态的。且 AttnRes 的 softmax 权重在推理时主要由静态向量 w_l 决定（query 不含 token 信息），所谓 input-dependent 主要体现在 key 上。
  • b) 失败假设：假设"每层的读取模式静态即可"。论文承认 input-dependent query 能再涨（1.731），说明固定伪查询仍有提升空间。
  • c) 胜出条件：深度读模式相对稳定（关注 embedding / 邻居层 / 特定块）的任务；失败：需要 per-token 动态深度读。
  • d) 更稳做法：AttnRes + input-dependent query（论文指出可做但需权衡）。
  • e) 受损方：无显著（两者都小开销）。

2.7.4 复杂度

  • mHC：状态 d×m，每层 I/O ~34d（m=4）远高于 AttnRes ~5.5d（Block，Table 1）。
  • AttnRes：算术 O(N·d)（Block）/ O(L·d)（Full），内存 O(N·B·T·d)（Block）。
  • 结论：AttnRes 在「表达力相仿或更好」的同时，I/O 低约 6 倍（34d vs 5.5d）。

2.7.5 参数敏感性

  • mHC：流数 m（2/4/8）、Sinkhorn 迭代数、use_dynamic_mapping。
  • AttnRes：块大小 N、伪查询初始 w=0、RMSNorm eps。
  • 论文：N≈8 最优，mHC 的 m 视任务（通常 4）。

2.7.6 工程实现

  • 两者都是「残差可学习化」家族，但一个走约束矩阵（线性注意）、一个走软注意力。本项目选择 AttnRes（K3 采用），因其为 softmax attention（与序列建模同构）、实现简单、I/O 低、可解释性好。

2.7.7 对比图（标准残差 vs mHC vs AttnRes）

  • 上排 3 图：架构示意（标准残差蓝 / mHC 琥珀 / AttnRes 红主角）。
  • 下排 3 图：深度-时间对偶性、每层 I/O（Table 1：3d vs 34d vs 5.5d）、Table 4 困惑度（1.766 vs 1.747 vs 1.737/1.746）。
  • 图表关键读数（实测值回填）：
    • Full 1.737 / Block 1.746 / mHC 1.747 / PreNorm 1.766（Table 4）
    • 每层 I/O：标准残差 ≈3d、mHC(m=4) ≈34d、AttnRes Block ≈5.5d（Table 1）
  • 关键推论：图上 AttnRes 同时满足"更低 loss + 更低 I/O"，直接支撑"softmax attention 兼具表达力与效率"观点。

2.8 基础设施：跨阶段缓存 + 高效 Prefill（论文 §4）

这是 "Block AttnRes 成为标准残差即插即用替代品" 的工程前提。

2.8.1 训练：跨阶段缓存（Cross-stage Caching）

问题 ：标准残差下 pipeline parallelism 只在相邻 stage 间传一个固定大小的 hidden state（与 pipeline 深度无关）。Block AttnRes 需要每个 stage 的所有已累积块表示来做跨块注意，naïvely 每次 stage 过渡都传完整历史 → 冗余通信。

推导（论文 Eq. 7-8） ：设 interleaved pipeline 有 P 个物理 stage、每物理 stage V 个虚拟 stage，C = P·V 个 chunk，第 j 个 chunk 累积 j·N_p 个块（N_p = 每物理 stage 每 token 产生的块数）。

  • naïve 通信（每次过渡传全部累积块）：

    复制代码 
    Comm_naive = Σ_{j=1}^{C-1} j·N_p·d = C(C-1)/2 · N_p·d      (Eq. 7)

  • 缓存后（块在本地位存，跨 stage 只传增量；第一虚拟 stage 无缓存正常累积，v≥2 只传 ~P·N_p 增量块）：

    复制代码 
    Comm_cached = P(P-1)/2 · N_p·d   (第一虚拟 stage)
            + (V-1)·P²·N_p·d     (后续虚拟 stage)           (Eq. 8)

效果 ：每过渡的峰值通信从 O(C) 降到 O(P)，是 V× 提升，且能完全与稳态 1F1B 的计算重叠。论文 Fig. 3 用 P=4、V=2 举例：第二虚拟 stage 时缓存消除了 6 次冗余块传输（每 stage 只在块边界才传增量 +b1,b2 而非全部 b0,b1,b2,...）。注：块边界无需对齐物理 stage 边界（Fig. 3 中每块跨越两个物理 stage）。

训练开销：无 PP 时 Block AttnRes 训练开销可忽略；有 PP 时端到端 overhead < 4%（论文 §4.1）。

2.8.2 并行状态下的 AttnResState

AttnResState 不是分布式的，每个 stage 有自己的副本。

每个 pipeline stage 都持有自己的 AttnResState 实例 ，blocks 字段的内容随 stage 递增 ，即 blocks 列表内容是从上游复制/继承来的。

单机视角

单机上，整个模型只有一个 AttnResState。

python 复制代码 
# 单机视角（你想象的）
state = attn_res.init_state(emb)
for layer in all_96_layers:
    h = attn_res.mix(state, i)
    out = layer(h)
    state = attn_res.add_output(state, out)  # blocks 累积到 9 个
final = attn_res.finalize(state)
分布式视角
复制代码 
物理 Rank 0                物理 Rank 1                物理 Rank 2
┌─────────────────┐       ┌─────────────────┐       ┌─────────────────┐
│ AttnResState    │       │ AttnResState    │       │ AttnResState    │
│  blocks=[b0]    │──→    │  blocks=[b0,b1] │──→    │  blocks=[b0,b1,b2]
│  partial=None   │  传输 │  partial=b2_acc │  增量 │  partial=None   │
│  units_in_block=0       │  units_in_block=1       │  units_in_block=0
└─────────────────┘       └─────────────────┘       └─────────────────┘
   处理 layer 0-1            处理 layer 2-3            处理 layer 4-5

pipeline parallelism 下，每个物理 rank 只跑一部分层，所以每个 rank 需要自己的 state：

python 复制代码 
# Rank 0（layer 0-23）
state_r0 = attn_res.init_state(emb)
for i in range(24):
    h = attn_res.mix(state_r0, i)
    out = layer_r0[i](h)
    state_r0 = attn_res.add_output(state_r0, out)
# state_r0.blocks = [emb, b1, b2]（2 个完整块 + partial）

# Rank 1（layer 24-47）------ 不重建，继承 r0 的 blocks
state_r1 = AttnResState(blocks=state_r0.blocks.copy(), partial=state_r0.partial, ...)
for i in range(24, 48):
    h = attn_res.mix(state_r1, i)
    out = layer_r1[i-24](h)
    state_r1 = attn_res.add_output(state_r1, out)
# state_r1.blocks = [emb, b1, b2, b3, b4]（4 个完整块）
关键区别
概念 是不是分布式的 说明
AttentionResiduals（nn.Module） 是（共享参数） 所有 stage 共享同一份 queries 和 key_norm_weight 权重（数据并行同步）
AttnResState（dataclass） 否（每 stage 独立） 每个 stage 一个实例，blocks 内容递增
blocks 列表 跨 stage 复制 每个 stage 的 blocks 包含上游所有块 + 自己新产生的块

2.8.3 Cross-stage Caching 原理

① Cross-stage Caching 到底缓存什么？

缓存的是 Block 表示 bn（不是 layer 输出、不是 hidden state、不是梯度）。

回看 attn_res.py:178-195 的 add_output()：

python 复制代码 
if state.units_in_block >= self.block_units:
    state.blocks.append(state.partial)   # ★ 这就是要被缓存的东西

每个 blocks[i] 是一个已完成块的塌缩表示 b_n = Σ_{j∈B_n} f_j(h_j)，shape (B, T, d)。

在 pipeline parallelism 下，物理 rank 处理多个虚拟 stage，每个虚拟 stage 会产生若干块表示。Cross-stage caching 就是把这些 b_n 留在本地内存，下次遇到同一 micro-batch 时不再重传。

② 为什么这些可以缓存？

关键洞察：块表示 bn 一旦归档，就不再变化。

看 attn_res.py 的状态机：

python 复制代码 
state.partial = output if state.partial is None else state.partial + output
# 满 block_units 后才归档
state.blocks.append(state.partial)
state.partial = None  # 归档后 partial 清空，开始新块

b_n 是求和塌缩 ------一旦块写满（units_in_block >= block_units），这个块的表示就定型了，后续 layer 不会再修改它。这跟 KV cache 的"append-only"性质完全一样：

特性 Block 表示 bnb**n KV cache 的 K/V
何时定型 块写满时 token 写入时
后续是否修改 否（只读） 否（只读）
内存生命周期 整个 forward/backward 整个 decode session

所以可以缓存 ：rank 1 在第 1 个虚拟 stage 把 [b_0, b_1] 发给 rank 2，rank 2 处理完存下来；第 2 个虚拟 stage 时 rank 1 又产生新的 [b_2]，只需发 +[b_2]，rank 2 把它和本地缓存的 [b_0, b_1] 拼起来就是完整历史 [b_0, b_1, b_2]。

③ 新阶段会产生什么？

每个新虚拟 stage 会产生 ∼P/N_p 个新块表示（P=物理 stage 数，N_p=平均每 stage 的块数）。

用论文 Figure 3 的例子（P=4, V=2）：

python 复制代码 
Virtual Stage 0（无缓存）：
  Rank 0 产生 [b_0]
  Rank 1 收到 [b_0]，产生 [b_1]      → 发给 Rank 2: [b_0, b_1]
  Rank 2 收到 [b_0, b_1]，产生 [b_2]  → 发给 Rank 3: [b_0, b_1, b_2]
  Rank 3 收到 [b_0, b_1, b_2]，产生 [b_3]

Virtual Stage 1（有缓存）：
  Rank 0 已有 [b_0]，产生新的 [b_0']  → 发给 Rank 1: +[b_0']
  Rank 1 已有 [b_0, b_1]，收到 +[b_0']，产生 [b_1'] 
       → 发给 Rank 2: +[b_0', b_1']   （只发增量！）
  Rank 2 已有 [b_0, b_1, b_2]，收到 +[b_0', b_1']，产生 [b_2']
       → 发给 Rank 3: +[b_0', b_1', b_2']

关键差别：

  • Naïve（无缓存） ：每次发完整历史 [b_0, b_1, b_2, ...]，通信量随 stage 数平方增长
  • Cached ：只发增量 +[b_new]，通信量随 stage 数线性增长

论文 Eq. 7→8 给出量化：

  • Naïve：Commnaı¨ve=C(C−1)2NpdCommnaı¨ve=2C (C −1)Npd
  • Cached：Commcached=P(P−1)2Npd+(V−1)P2NpdCommcached=2P (P −1)Npd +(V −1)P 2Npd

当 V=2、P=4 时，消除 6 次冗余传输（论文原话）。

回到 AttnResState：缓存的工程映射

AttnResState.blocks 就是 Cross-stage Caching 的本地缓存：

python 复制代码 
@dataclass
class AttnResState:
    blocks: list[Tensor]      # ★ 这就是 cache：已归档的块表示
    partial: Tensor | None    # ★ 正在累计的（还没定型，不能缓存）
    units_in_block: int       # 进度计数
  • blocks：跨 stage 缓存的对象（只读、可复用）
  • partial ：当前 stage 正在产生的增量（未定型，必须传输）

pipeline 边界的传输就是：

python 复制代码 
# 发送方：把当前 stage 新产生的 blocks 发给下一个 rank
new_blocks = state.blocks[len(receiver_cached_blocks):]
send(new_blocks)

# 接收方：把新块追加到本地缓存
state.blocks.extend(new_blocks)

partial 永远不跨 stage 传输 ------它必须在当前 stage 内完成塌缩成 blocks 才能进入缓存。这就是"为什么 partial 不能缓存"：它还在变，不是只读的。

一句话总结
问题 答案
缓存什么 已归档的 Block 表示 b_n（AttnResState.blocks）
为什么能缓存 块一旦写满就只读不变（求和塌缩的 append-only 性质）
新阶段产生什么 ∼P/N_p 个增量块 +[b_new]，只发增量不重发历史
Cross-stage Caching 的精确含义

缓存的不是"AttnResState 的一部分"，而是"blocks 列表的前缀"。

Rank 1 在第 1 个虚拟 stage 结束后，把自己的 blocks=[b0, b1] 发给 Rank 2。Rank 2 存进自己的 state.blocks。第 2 个虚拟 stage 时，Rank 1 又产生新的 b2，只发 +[b2] ，Rank 2 把它 append 到自己的 state.blocks。

所以：

  • blocks 物理上分布在多个 rank 的内存里（每个 rank 一份副本）
  • 但逻辑上它们是同一个列表的不同前缀
  • Cross-stage caching 就是避免重复发送已经共享的前缀
一句话总结
问题 答案
缓存的是 AttnResState 的一部分？ 不准确。缓存的是**blocks 列表的前缀**（跨 rank 共享的已归档块）
AttnResState 是分布式的？ 不是。每个 rank 有自己的独立实例
系统中应该只有一个 AttnResState？ 单机是 ，pipeline 下每个 rank 一个
blocks 遍布在每个 stage？ 是的 ，但内容递增------每个 rank 的 blocks 包含上游所有块 + 自己新产生的

2.8.4 讨论

为什么 activation recomputation 被影响

标准 activation checkpointing 可以丢掉中间激活，反向时从 checkpoint 重算。但 Full AttnRes 要让后续很多层继续访问旧层输出；如果丢掉它们，后续层无法直接算深度注意力，或者必须频繁重算历史输出。于是它把"可释放的中间态"变成"跨多层被读取的 KV-like 源"。

Block AttnRes 的妥协是只保存每个块的聚合表示和当前块 partial prefix。这样 activation checkpointing 仍可丢掉 softmax 中间量和大量块内细节，保留少量块级 bank。

为什么 pipeline parallelism 变贵

普通 PP stage 边界只传 hidden_states 和常规 residual。AttnRes 需要下游 stage 也能看到历史块，所以边界还要传 block bank。若 Full AttnRes 传每层历史，成本随层数涨；Block AttnRes 传块表示，成本随块数涨。

cross-stage caching: 同一物理 stage 在交错流水线的不同 virtual stage 会重复看到部分历史块，因此可缓存已见块，只传新增块，减少冗余通信。

2.8.5 推理：两阶段计算（Algorithm 1）

问题 ：Block AttnRes 的逐层注意计算类似自回归解码------块表示是被所有层共享复用的"KV cache"。naïve 实现每层都要全扫一遍所有前置块 → O(L·N) 内存访问。关键洞察 ：伪查询 w_l 与前向计算解耦（是学习参数），所以一个块内的所有 S 个 query 可以批量成一个矩阵乘，把 S 次读摊销成 1 次。

Algorithm 1 两阶段：

  • Phase 1（跨块，并行） ：对块 n 的所有 S 层 query [w_l]_{l∈B_n} 一次性对块表示 [b_0..b_{n-1}] 算注意，返回输出 + softmax 统计（max 和 log-sum-exp）。

    块 n 的 S 个 query {w_l} 在执行前全部已知，所以前面 (n-1)S 个 key-value 对只需从 HBM 读一次，被 S 个 query 复用：

    Read_inter(n) = 2(n-1)S·d # 因子 2 = key+value Read_inter = Σ_n 2(n-1)S·d = 2Sd·N(N-1)/2 = dL(N-1) Write_inter = L·d

  • Phase 2（块内，顺序） ：逐层用不断演化的 partial sum 算块内注意，然后通过 online softmax 与 Phase 1 输出合并（m_l = max(m^(1), m^(2))，h_l = (e^{m1-m}·o1 + e^{m2-m}·o2) / (e^{m1-m}·l1 + e^{m2-m}·l2)）。因为 online-softmax 合并是逐元素操作，Phase 2 可自然与周边算子 kernel 融合，进一步降 I/O。

    块内第 t 层（1≤t≤S）依赖前面 t-1 个块内 key-value：

    Read_intra(n) = Σ_{t=1}^S 2(t-1)d = S(S-1)d Write_intra = S·d

  • 总摊销 I/O（除以 L）：

    Read per layer = dL(N-1)/L + N·S(S-1)d/L = (N-1)d + (S-1)d = (S+N-2)d Write per layer = 2d Total I/O = (S+N)d

    结论：批量跨块读把 Full AttnRes 的每层 I/O 从 O(L) 降到 O(S+N)。典型值 L=128, N=8, S=16 → Full = 24d（对比 Block 的 (N/S+5)d=5.5d，Table 1）。Block 形式进一步把存储的表示数从 L 降到 N（每块压成一个向量），故 I/O 更低。

推理两阶段的直觉读法：

  1. 批量计算当前 block 中多层 query 对历史 blocks 的注意力统计。
  2. 层内顺序处理当前 partial prefix，并用 online softmax 合并历史块统计。

核心收益：历史块只读一次或少数几次，避免每层重复扫完整 bank；online softmax 只维护 max、denominator、weighted sum，不必物化大矩阵。

结果（Table 1）：典型参数 L=128, N=8, S=16, m=4：

方案 Read Write Total I/O
标准残差 2d d 3d
mHC（m 流） --- --- 34d
AttnRes Full（两阶段） --- --- (S+N)d = 24d
AttnRes Block（两阶段） --- --- (N/S+5)d = 5.5d

注：Full 形式用两阶段调度也可从 O(Ld) 摊薄到 (S+N)d（见附录 C），只是块数 N 远大于 Block。

2.8.6 推理：内存高效 Prefill（sequence sharding）

问题 ：prefill 阶段存块表示需 N·T·d 元素------128K token 序列、8 块时高达 15 GB。

方案 ：把块表示沿序列维 shard 到 P 个 tensor-parallel 设备（每设备 N·(T/P)·d）。Phase 1 在本地序列分片上独立执行；Phase 2 的 online-softmax 合并融入标准 TP all-reduce 通信路径（output reduce-scatter → 本地合并 → all-gather 重建），并可自然与 RMSNorm 等 kernel 融合。

效果 ：每设备内存从 15 GB 降到 ~1.9 GB（128K/8 块/P=8）；配合 chunked prefill（如 16K chunk）进一步降到 <0.3 GB/设备。

2.8.7 反方辩手

  • a) 反对：这套基础设施（缓存 + 两阶段 + sharding）本身复杂度高，对小模型/单卡无意义。
  • b) 失败假设：假设 PP 是大规模训练标配、块表示必须跨 stage 传。若未来互联带宽暴涨，Full AttnRes 的 O(Ld) 通信也可行（论文明确说这是 Full 形式的潜在未来）。
  • c) 胜出条件：大规模 PP 训练 + 长上下文推理；失败：单卡教育/小模型。
  • d) 更稳做法：不做优化，Block AttnRes 直接朴素实现（本项目采用）。
  • e) 受损方：基础设施工程师（复杂度）。

完成度：本项目的朴素实现未包含跨 stage 缓存/两阶段/序列分片（与 kimi3-master 一致）；本文档记录论文方案作参考。

2.9 学到的深度注意模式（论文 §5.4.2，Fig. 8）

论文可视化 16 层模型（Table 2 436M 行）在 Full 与 Block (N=8) 形式下学到的权重 α_{i→l}（Pre-Attn 与 Pre-MLP 分开展示）。三个关键观察：

  1. 保留局部性（Preserved locality） ：每层主要注意直接前驱 （对角占优），但出现选择性的非对角集中------如 layer 4 注意早期源、layers 15-16 在 Block 设置下回溯远处 → 超出标准残差路径的学到的 skip connections。
  2. 层级专门化（Layer specialization） ：embedding h_1 全程保留非平凡权重，尤其 Pre-Attn 层（深度维 attention sink，对应 §6.2 与序列维 attention sink 同源）；Pre-MLP 输入的 diag 依赖更锐利（读最近表示），Pre-Attn 输入的感受野更广（attention 跨层路由信息，MLP 局部工作）。
  3. Block 保留结构 ：对角占优、embedding 持续、层级专门化都从 Full 转移到 Block → 块级压缩充当隐式正则，同时保留关键信息通路。

教学价值 ：本项目的 attention_weights()（attn_res.py）正是为此设计的分析窗口------可复现论文 Fig. 8 的"哪一层读了哪个源"。实验 demo 已演示学到权重收敛到 embedding 源（见 0x07）。

完成度：论文实验（Fig. 8），本项目可复现（attention_weights() + demo），但未做 16 层模型级可视化。

2.10 统一视角：残差连接作为结构化矩阵（论文 §6.2，Table 5）

深度混合矩阵 M ：把残差变体统一表示为 h_l = Σ_{i=0}^{l-1} M_{i→l} v_i（v_0=h_1，v_i=f_i(h_i)）。M 是 R^{L×L} 下三角矩阵，M_{i→l} 是层 l 给层 i 输出的权重。各变体的区别在：权重如何产生（fixed / learned / input-dependent）+ M 是否受低秩约束。semiseparable rank 是统一比较的标尺。

论文 Table 5（更新机制分类）：

结构化矩阵视角（Fig. 9，L=4 示例，Block S=2）：

  • 标准残差：M = 全 1 下三角矩阵。
  • Highway ：M_{0→l}=γ_{1→l}、M_{i→l}=g_{i+1}·γ_{i+1→l}（γ_{i→l}=Π_{j=i+1}^l (1-g_j) 是 carry product）。因累积积因子化通过标量门，M 是 1-semiseparable （与标准残差同 rank，但 input-dependent）。权重自动和为 1 → Highway 是 softmax-free 的深度维 stick-breaking attention。
  • (m)HC ：M_{i→l}=β^T_i A_{i+1→l} α_l（A_{i→j}=Π_{k=i+1}^j A_k）。m×m 转移使 M 是 m-semiseparable 。mHC 进一步要求每个 A_l 双随机，稳定跨深度累积积。
  • Full AttnRes ：M_{i→l}=α_{i→l} 经 φ(w_l,k_i)=exp(w_l^T RMSNorm(k_i)) 归一化 → 稠密 rank-L。
  • Block AttnRes ：已完成块 B_n 内所有源 i 共享块级 key/value b_n → M_{i→l}=α_{n→l} ∀i∈B_n；当前块内每层额外注意演化 partial sum b_n^{i-1} → 每块内多一个额外源。有效 rank 介于 N 和 N+S 之间，在标准残差（N=1）与 Full（N=L）之间插值。

两个分析洞察（论文 §6.2）：

  1. 深度维 attention sink：input-dependent 的 M 揭示某些层无论输入如何都持续吸引高权重（2.10 的 embedding 持续），与序列维 attention sink 同源。
  2. φ 可分解性 → 递推塌缩 ：当 φ(q,k)=φ(q)^T φ(k) 可分解时，深度注意塌缩成递推------正是 MRLA↔GLA、DDL↔DeltaNet 对应的结构根源；不可分解的 softmax 核让 AttnRes 保持稠密可检索。

2.10.1 与序列-深度对偶（论文 §6.1）

论文把深度递推与序列递推 形式化对齐：TTT 把每个递推步视为对自监督 loss 的梯度下降 W_t = W_{t-1} - η∇ℓ(W_{t-1}; x_t)，f 线性时退化为线性注意 S_t = S_{t-1} + k_t v_t^T；标准残差沿深度方向有同样加性形式（h_l 是状态，每层 f_l 是一个"梯度步"）。

跨尺度的对应关系（论文指出）：

  • 序列维 data-dependent gates（RetNet, GLA）↔ 深度维 Highway
  • delta rule ↔ DDL
  • MRLA ↔ GLA（gated linear attention）
  • 线性注意 S_t = S_{t-1} + k_t v_t^T ↔ 标准残差 h_l = h_{l-1} + f_{l-1}(h_{l-1})

这些方法都停留在递推范式内。AttnRes 更进一步：用直接跨层注意替换深度递推------正如 Transformer 用 self-attention 替换时间递推。论文明确："Since the number of layers in current architectures remains well within the practical regime of softmax attention, we adopt vanilla depth-wise attention"，并把"更高效（线性复杂度）替代"列为 future work。

2.11 最优架构分析（论文 §5.4.1，Fig. 7）

问题 ：AttnRes 是否改变最优的 depth-width-attention 权衡？尤其------既然它在深度维度上有优势，是否偏好更深的模型？

方法 ：固定计算预算（FLOPs ≈ 6.5×10^19）和激活参数量（≈ 2.3×10^8），枚举 5×5 网格：d_model/L_b ∈ {15,30,45,60,75} × H/L_b ∈ {0.3,0.4,0.5,0.6,0.7}（L_b=L/2 为 Transformer 块数，H 为注意力头数）。固定每专家 MLP 扩展比 d_ff/d_model ≈ 0.45，隔离纯架构影响。

结果（Fig. 7）：

  • 两方法共享大趋势：loss 随 d_model/L_b 增大、H/L_b 减小而下降，都在 H/L_b≈0.3 达优。
  • AttnRes 在全部 25 个配置上低于 baseline，优势 0.019~0.063。
  • 最优位置移动 ：baseline 在 d_model/L_b≈60（loss 1.847）达优；AttnRes 移到 d_model/L_b≈45（loss 1.802）。

推论 ：固定参数预算下，更小的 d_model/L_b = 更深更窄的网络 → AttnRes 能更有效利用额外深度 。但论文明确警示：这不是部署推荐（更深模型推理延迟更高），而是揭示"AttnRes 在哪里收益最大"的诊断，可作为选架构时与推理成本一起权衡的因子。

完成度：论文实验（Fig. 7），本项目未复现（CPU 无法枚举 25 配置）。

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.
热门推荐
01国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）02历年考研数学一、数学二、数学三真题试卷及答案PDF03如何新建文件夹？ 电脑新建文件夹的4种方法04几个好用的ip纯净度检测网站05微信历史版本含下载地址（ Windows PC | 安卓 | MAC ）及设置微信不更新06GitHub 镜像站点07昇腾910B系列芯片完全指南：四款型号对比、服务器选型与大模型部署实战08苹果iPA游戏软件资源下载网站09蜘蛛磁力 搜索引擎大全，如何使用蜘蛛磁力查找磁力链接10Dify：一站式AI应用开发平台全解析与Windows部署指南