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

2.3.2 对比基线（block vs full）

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

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 的关键伪代码：

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管理两个数据结构：

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

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

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

┌─────────────────────────────────────────────────────────────┐
│  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），每层都要保留之前所有层的完整输出：

# 不用 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：退化为标准残差（失去选择性）

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

# 标准残差（无 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 注意力"的状态机，它解决了：

2.4.8 AttnResState 如何解决问题

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

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

代码：attn_res.py:178-195

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）：

问题 2：选择性 → mix() 的深度维 softmax

代码：attn_res.py:157-176

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

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

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 的窗口：

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 个问题

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 对比基线

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 对比基线

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 对比基线表

三条路线并排示意图（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。

# 单机视角（你想象的）
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：

# 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 个完整块）

关键区别

2.8.3 Cross-stage Caching 原理

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

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

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

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 的状态机：

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"性质完全一样：

所以可以缓存：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）：

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)P2Npd

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

回到 AttnResState：缓存的工程映射

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

@dataclass
class AttnResState:
    blocks: list[Tensor]      # ★ 这就是 cache：已归档的块表示
    partial: Tensor | None    # ★ 正在累计的（还没定型，不能缓存）
    units_in_block: int       # 进度计数

• blocks：跨 stage 缓存的对象（只读、可复用）
• partial：当前 stage 正在产生的增量（未定型，必须传输）

pipeline 边界的传输就是：

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

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

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

一句话总结

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 就是避免重复发送已经共享的前缀

一句话总结

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：

注：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.