KDA：Kimi K3 如何按 Channel 控制遗忘——轻量化复刻大模型解析

【轻量化复刻大模型 / KIMI K3】KDA：按 Channel 控制遗忘

• 【轻量化复刻大模型 / KIMI K3】KDA：按 Channel 控制遗忘
  • 0x00 概要
  • 0x01 演进脉络与定位
    • 1.1 前身
    • 1.2 同期竞品
    • 1.3 KDA 论文
    • 1.4 定位
    • 1.5 KimiDeltaAttention
  • 0x02 原理讲解
    • 2.1 直觉类比视角
    • 2.2 专家视角
    • 2.3 KDA 原理图
    • 2.4 深入讨论
    • 2.5 图例
    • 2.6 各种机制对比
  • 0x03 子知识点拆解
  • 0x04 子知识点深度分析
    • 4.1 状态递推与 delta rule（Eq.1）
      • 4.1.1 原理
      • 4.1.2 对比基线
      • 4.1.3 计算复杂度分析
      • 4.1.4 参数敏感性
      • 4.1.5 反方辩论
      • 4.1.6 工程实现
      • 4.1.7 常见疑问 & 讨论
    • 4.2 有界衰减 log-decay（Eq.5）
      • 4.2.1 原理
      • 4.2.2 对比基线
      • 4.2.4 计算复杂度分析
      • 4.2.5 参数敏感性
      • 4.2.6 工程实现
      • 4.2.7 数值稳定性验证实验
      • 4.2.8 讨论
      • 4.2.9 反方辩手
      • 4.2.10 图例
    • 4.3 channel-wise decay 与低秩参数化（Eq.2）
      • 4.3.1 原理
        • 逻辑图
        • alpha_rank 设置
        • 如何压缩到 alpha_rank？（低秩分解）
        • 如何映射到每个 head？（广播 + 逐 head 调制）
      • 4.3.2 对比基线
      • 4.3.3 反方辩手
      • 4.3.4 计算复杂度
      • 4.3.5 参数敏感性
      • 4.3.6 工程实现
      • 4.3.7 讨论
    • 4.4 Head-wise RMSNorm（按 head 归一化的 RMSNorm）
      • 4.4.1 前身
      • 4.4.2 同期竞品
      • 4.4.3 原理讲解
        • 直觉类比视角
        • 准专业水平视角
        • 专家视角
      • 4.4.4 实现 & 问题
    • 4.5 输出门（Eq.6）
      • 4.5.1 原理
      • 4.5.2 对比基线
      • 4.5.3 反方辩手
      • 4.5.4 工程实现
    • 4.6 Gated MLA 与 KDA/MLA 互补
      • 4.6.1 MLA 负责什么
      • 4.6.2 Gate 负责什么
      • 4.6.3 组合收益
      • 4.6.4 KDA 的职责
      • 4.6.5 MLA 的职责
      • 4.6.6 为什么不能只用 KDA
  • 0xFF 参考

0x00 概要

本文档是复刻 KIMI K3 KDA（Kimi Delta Attention）的第一部分，因为文章实在太长了，所以分为两期发出。

不同于 GDN 的 head-wise 粗粒度遗忘和 vanilla 线性注意力的无选择性累积，KDA 通过 Diag(α_t) channel-wise 有界衰减 + delta-rule 写入，在固定大小状态矩阵 S 上实现了可控记忆、可稠密 chunkwise 并行、且无位置编码也自洽的线性注意力。

注：

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

0x01 演进脉络与定位

1.1 前身

• Vanilla 线性注意力（Linear Transformer）：S = Σ k v^T，矩阵加法可交换。缺陷：固定状态容量受限、表达力弱、无选择性遗忘。
• Gated Linear Attention（GLA）：channel-wise gating 的代表，但沿用的是线性注意力（外积累积），没有 delta-rule 的"写入时先减去旧贡献"机制。
• Gated DeltaNet（GDN）：把 delta rule（DeltaNet，Mamba-2 同源）加上 head-wise forget gate。缺陷：gate 是每个 head 一个标量，遗忘粒度太粗，无法对不同 feature 维度做差异化记忆控制。

KDA 的目标是：在 GDN 的 delta rule 基础上，把遗忘粒度细化到 channel 级——这需要在状态更新的转移矩阵中加入对角 gate Diag(α_t)，同时保持 chunkwise 可分块并行（DPLR / Diagonal-Plus-Low-Rank，对角加低秩结构）。

1.2 同期竞品

1.3 KDA 论文

KDA 有两篇直接相关的论文，本文档以技术报告（K3 报告）为主线、Kimi Linear 报告为算法细节补充：

• KDA 首先在 Kimi Linear 技术报告中被完整定义：它扩展 Gated DeltaNet（GDN），把粗粒度 head-wise forget gate 换成 channel-wise 对角 gate Diag(α_t)。
• K3 报告 §2.1.1 进一步给出 KDA 在 K3 中的实现细节（有界衰减 Eq.5、全秩输出门 Eq.6），并在 §5.1 系统阐述了 KDA 的算法-系统协同设计（FlashKDA、KCP、KDA-aware 前缀缓存）。

1.4 定位

Attention 的抽象不是"必须出现 Softmax"，而是：

根据 Key 决定如何读取 Value

KDA 通过状态矩阵完成：y_t = q_tᵀ S_t。所以它仍然是"Query 根据 Key-Value 关系读取内容"，只是：

• Full Attention: 保存 token 级 KV，读取时计算 Softmax；
• Linear/KDA: 保存压缩状态，读取时用核 / 矩阵乘法；
• KDA 还加入了带遗忘的在线更新。

因此可以说：KDA 是注意力的一种递推、核化、带门控的实现。 注意力的本质是内容寻址，不是 Softmax 这个具体函数。

1.5 KimiDeltaAttention

KimiDeltaAttention 代码如下：

class KimiDeltaAttention(nn.Module):
    """KDA 完整层：投影（q/k/v/beta/decay）+ 核心递推 + 输出门。

    参数结构（报告 Eq. 2 / Eq. 5 / Eq. 6）：
    - ``q_proj/k_proj/v_proj``：线性投影到每 head 的 key/value 维度。
    - ``q_conv/k_conv/v_conv``：因果 depthwise ShortConv（宽度 4），给无位置编码的
      模型注入 token 级"位移"信息（报告 §2.1.1；Kimi Linear 论文称其为位置编码替代）。
    - ``beta_proj``：delta-rule 写入强度 ``beta = sigmoid(x)``，控制"本次写入覆盖
      多少旧记忆"（``beta`` 接近 1 表示完全重写该方向的记忆）。
    - ``alpha_down/alpha_up/alpha_bias``：**低秩** decay logits ``z = W_up W_down x + b``
      （Eq. 2），逐 channel 的 log 衰减来源。
    - ``a_log``：Eq. 5 的逐 head 对数尺度 ``A_h``，论文初始化为 0。
    - ``o_norm``：head-wise RMSNorm（Eq. 6 的第一步）。
    - ``g_proj``：全秩、依赖输入的输出门 ``sigmoid(W_g x)``（Eq. 6 的第二步）。
    """

    def __init__(self, dim: int, cfg: KDAConfig, *, layer_idx: int = 0, eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.cfg = cfg
        self.layer_idx = layer_idx
        h, dk, dv = cfg.num_heads, cfg.head_dim_k, cfg.head_dim_v
        self.num_heads, self.head_dim_k, self.head_dim_v = h, dk, dv
        k_dim, v_dim = cfg.k_dim, cfg.v_dim   # k_dim = H*Dk, v_dim = H*Dv

        # --- 线性投影到未分 head 的宽向量，再交给 ShortConv 做局部混合 ----------
        self.q_proj = nn.Linear(dim, k_dim, bias=False)
        self.k_proj = nn.Linear(dim, k_dim, bias=False)
        self.v_proj = nn.Linear(dim, v_dim, bias=False)
        self.q_conv = ShortConv(k_dim, cfg.conv_kernel)   # 宽度 4 的因果深度卷积
        self.k_conv = ShortConv(k_dim, cfg.conv_kernel)
        self.v_conv = ShortConv(v_dim, cfg.conv_kernel)

        # --- delta-rule 写入强度 beta 与低秩衰减 logits ----------------------
        self.beta_proj = nn.Linear(dim, h, bias=False)      # -> (B, T, H) 每 head 一个标量
        # 低秩 decay logits（Eq. 2）：z = W_up W_down x + b_alpha。
        self.alpha_down = nn.Linear(dim, cfg.alpha_rank, bias=False)   # 压到 rank
        self.alpha_up = nn.Linear(cfg.alpha_rank, k_dim, bias=False)   # 展开回 k_dim
        self.alpha_bias = nn.Parameter(torch.empty(k_dim))             # 逐 channel 偏置
        # Eq. 5 的逐 head log-scale A_h；论文初始化为 0。
        self.a_log = nn.Parameter(torch.zeros(h))

        # --- 输出门（Eq. 6）：head-wise RMSNorm + 全秩依赖输入的门 ------------
        self.o_norm = HeadRMSNorm(dv, eps)
        self.g_proj = nn.Linear(dim, v_dim, bias=False)   # 全秩输出门
        self.o_proj = nn.Linear(v_dim, dim, bias=False)

        self.reset_parameters()
        
    def forward(
        self,
        x: Tensor,
        *,
        reset_mask: Tensor | None = None,
        cache: KDACache | None = None,
        initial_state: Tensor | None = None,
    ) -> tuple[Tensor, Tensor]:
        """``x`` 是 ``(batch, seq, dim)``。返回 ``(output, final_state)``。

        形态选择：训练/prefill 走 chunkwise（块内并行、块间串行传递状态）；
        自回归 decode（``t == 1`` 且已有状态）走单步 ``kda_step``。两种形态数值上
        完全等价（``kda_ref`` 里有 ``kda_recurrent`` 作为 ground truth 校验）。
        """
        b, t, _ = x.shape
        q, k, v, beta, g, conv_states = self._project(x, cache)

        # (B, T, H, D) -> (B, H, T, D)：递推在 T 维进行，把 head 提到中间便于广播。
        qh, kh, vh = (z.transpose(1, 2) for z in (q, k, v))
        gh = g.transpose(1, 2)
        bh = beta.transpose(1, 2)

        state_in = initial_state if initial_state is not None else (cache.state if cache else None)
        # 中间计算用 fp32（activation 是 bf16/fp16 时），递推状态累积不丢精度。
        work_dtype = torch.float32 if x.dtype != torch.float64 else torch.float64
        chunk = self.effective_chunk_size(work_dtype)

        if t == 1 and state_in is not None:
            # decode 形态：单 token 单步更新。注意此分支只在"有先前状态"时启用，
            # 保证 decode 的第一 token 也能走 chunkwise（初始状态为 0）。
            out, state = kda_step(
                qh[:, :, 0].to(work_dtype),
                kh[:, :, 0].to(work_dtype),
                vh[:, :, 0].to(work_dtype),
                gh[:, :, 0].to(work_dtype),
                bh[:, :, 0].to(work_dtype),
                state_in.to(work_dtype),
            )
            out = out.unsqueeze(2)                     # (B,H,1,Dv)
        else:
            # prefill / 无状态 decode：chunkwise 并行形式（块内并行，块间串行）。
            fn = self._kernel()
            out, state = fn(
                qh,
                kh,
                vh,
                gh,
                bh,
                chunk_size=chunk,
                initial_state=state_in,
                reset_mask=reset_mask,
                dtype=work_dtype,
            )

        # Eq. 6 输出门：head-wise RMSNorm -> 全秩依赖输入的门 -> 输出投影。
        out = out.transpose(1, 2).reshape(b, t, self.num_heads, self.head_dim_v)
        out = self.o_norm(out.to(x.dtype)).reshape(b, t, -1)
        out = out * torch.sigmoid(self.g_proj(x))
        out = self.o_proj(out)

        if cache is not None:
            # decode 时把本轮状态写回 cache，供下一 token 使用。
            cache.state = state
            cache.conv_q, cache.conv_k, cache.conv_v = conv_states
        return out, state        

0x02 原理讲解

2.1 直觉类比视角

想象你是一位图书管理员，面前只有一个固定大小的"记忆白板"（状态矩阵 S），没有书架（KV cache）。

核心动作是"先按 channel 褪色，再按 delta 写入"。每来一本新书（token），你要做两件事：

1. 先遗忘：用每本书每个主题自己的"褪色速度"（channel-wise decay α）把白板上旧内容整体变淡一点。
2. 再更新：用新书内容按 delta rule 写上去——但写之前，先把你已经知道的新书内容"抹掉"原有信息（v − S^T k），只写新信息本身。这样不会重复累积。

关键区别：softmax 书架有无限大，能精确记住每本书；KDA 白板永远这么大，但每次都能精确地决定忘多少、更新什么。当上下文超过记忆容量时，KDA 靠与 MLA 全局层混合兜底——"小管家 + 定期大整理"。

2.2 专家视角

输入输出：

• 输入：x ∈ R^{B×T×d}，其中 B=batch、T=序列长、d=隐维度。
• 投影后：query/key q_t, k_t ∈ R^{d_k}（L2 归一化后），value v_t ∈ R^{d_v}，每 head 一个。
• 状态：S_t ∈ R^{d_k × d_v}，固定大小，不随 T 增长。

核心递推（Eq.1，单 head 视角）：

S_t = (I − β_t k_t k_tᵀ) Diag(α_t) S_{t-1} + β_t k_t v_tᵀ      (Eq.1)

• Diag(α_t)：channel-wise 衰减因子，α_t = exp(g_t)，g_t ≤ 0 是 log-decay。每一行（channel）独立遗忘。
• Diag(α_t) S_{t-1}：按 channel 衰减旧状态。
• β_t：delta-rule 写入强度 (0,1)，控制"覆盖旧记忆的方向"。
• (I − β_t k_t k_tᵀ)：秩-1 投影，把状态中与 k_t 对齐的成分按 β_t 削弱。即，沿 k 方向把旧状态中"和新 key 重叠的成分"减去——这就是 delta rule 的"写新前擦旧"。
• + β_t k_t v_tᵀ：写入新记忆。

输出：o_t = S_tᵀ q_t（对每个 head）。

等价改写（数值更稳，实际实现）：

S_t = S_{t-1} + β_t k_t (v_t − S_{t-1}ᵀ k_t)ᵀ

即先计算"本次要写的新信息" = v_t 减去"旧状态沿 k_t 方向已记住的内容"，只写增量。

挖掘："选择性遗忘 + 在线学习率"的 delta rule 与线性回归的递归最小二乘同构，把 k_t 看作输入、v_t 看作监督、S 看作线性映射权重，S_tᵀ k_t 是预测，v_t − S_tᵀ k_t 是残差，β_t 是学习率。因此 KDA 把注意力建模成"可变学习率的在线最小二乘"，理论上表达力强于纯累积的线性注意力，这让它天然有"压缩即智能"的理论支撑。

2.3 KDA 原理图

先看公式（Eq.1）

S_t = (I - beta_t * k_t k_t^T) Diag(alpha_t) S_{t-1} + beta_t * k_t v_t^T
      └─────────── 逐 token 过渡矩阵 ───────────────┘    └── 写入新信息 ──┘
alpha_t = exp(g_t)     逐 channel 保留因子（0 < alpha <= 1）

K3 架构总览，左下为 KDA 模块图

2.4 深入讨论

• 理论上为什么 work：KDA 把状态转移矩阵参数化为 DPLR 特例——Diag(α_t)（对角）+ β_t k_t k_tᵀ（秩-1）——这正是"可 chunkwise 并行"的结构前提。相比一般 DPLR，KDA 的转移矩阵与经典 delta rule 一致（S 的列沿 key 方向衰减），物理语义清晰（每列是"key 记忆的遗忘曲线"）。
• 下界衰减的必要性：Kimi Linear 用无界 softplus g = −exp(A)·softplus(z)，累积衰减乘积下溢、其倒数（K/Γ 的重缩放）上溢，迫使对角 tile 走显式 position-pair 路径；K3 改用 g = g_min·sigmoid(exp(A)z) 把 log-decay 锁在 (g_min, 0)，g_min=−5 时 exp(80)=5.5e34 仍在 BF16 范围内——每个因果 tile（含对角）都能用稠密 Tensor-Core matmul 计算。
• NoPE 自洽性：KDA 的衰减和 delta 更新天然位置感知（recency 由衰减累积编码），因此 MLA 层无需 RoPE 也能全局混合；上下文扩展到 1M 不需要任何位置参数手术。
• 论文未解决的理论问题：固定状态容量对超长程精确检索的根本上限（论文承认线性注意力检索弱，用 MLA 混合弥补）；状态矩阵的容量-精度权衡（state expansion 方向仍在探索）。

2.5 图例

图例如下：

图中 2×2 四格，全部由真实递推（kda_recurrent，全项目 KDA ground truth）跑出：

左上 —— 状态矩阵 S（记忆）：KDA 的"记忆"就是一个固定大小的 Dk × Dv 矩阵。图中热图是真实跑完 32 个 token 后的最终状态；T=32 个 token 只产生一个 16×16 矩阵——这就是它对比 KV cache 的核心卖点：内存 O(1)，不随序列长度增长。

右上 —— 单步三件事拆解：每来一个 token，状态只做三件事：

1. 衰减 Diag(alpha) S：旧记忆先按通道乘一个 alpha=exp(g)<=1 的保留因子（旧信息"忘掉一点"）；
2. 检索 S^T k：用当前 key 在记忆里"查一下"已存了什么；
3. 写入 beta * k (v - S^T k)^T：只把误差 v - 检索到的记忆 写进状态（delta rule 的核心）。

左下 —— delta rule 的"纠错"本质：人为构造"同一个 key 反复出现、v 逐步变化"的场景，展示 key 检索到的值如何收敛到最新的 v。普通线性注意力（S += k v^T）会"混在一起"，而 delta rule 通过 v - S^T k 把冲突压下去——这就是"记忆能持续更新、不互相污染"的原因。图中三条线是真实 kda_step 跑出的检索值。

右下 —— 有界 log-decay（Eq.5）：g = g_min * sigmoid(exp(A) z) 把 log-decay 严格锁在 (g_min, 0)，即 g_min=-5 时 g ∈ (-5, 0)。这个"下界"是 Q4/Q5/Q6 的数值稳定性前提。

一句话结论：KDA = 固定大小的记忆矩阵 S（O(1) 内存）+ 每 token 三件事（衰减→检索→只写误差）+ 有界衰减（保证 chunkwise 不溢出）。与标准 MHA 的 KV cache 随 T 线性增长 形成本质对比。

2.6 各种机制对比

这张表的主线是：不同方法主要在两个地方产生差异：

• 旧状态如何衰减；
• 当前 Value 如何写入以及是否使用残差纠错。

线性注意力没有遗忘和纠错；DeltaNet 增加残差写入；GDN 增加按 Head 的统一遗忘；KDA 则把遗忘细化到每个 Key 通道。

0x03 子知识点拆解

KDA（Kimi Delta Attention）
    ├── 状态递推与 delta rule（核心算子）
    ├── 有界衰减 log-decay（Eq.5）
    ├── channel-wise decay 与低秩参数化（Eq.2）
    ├── ShortConv（无位置编码的局部混合）
    ├── chunkwise 并行（Eq.3-4，DPLR）
    ├── 单步递推 kda_step（decode 形态）
    ├── 输出门（Eq.6）
    ├── 解码缓存 KDACache 与 clone() 分支
    └── KDA 上下文并行（KCP，报告 §5.1.2）

代码对应如下：

use_triton=False 时 _kernel() 返回纯 PyTorch kda_chunkwise。

统一符号约定（本文档所有小节遵循）：

• S_t ∈ R^{d_k × d_v}：递推状态矩阵（每 head 一个）。
• q, k ∈ R^{d_k}（已 L2 归一化）、v ∈ R^{d_v}。
• g_t ∈ R^{d_k}：逐 channel log-decay（≤0）；α_t = exp(g_t) 为 retention 因子。
• β_t ∈ (0,1)：delta-rule 写入强度。
• chunk 尺寸 C（配置 64，tile 16），tile 是数值安全的基本衰减块。

0x04 子知识点深度分析

4.1 状态递推与 delta rule（Eq.1）

KDA 保留了 Delta Rule 的残差覆写机制，但把 GDN 的标量遗忘门扩展成了一个 d_k 维向量。

4.1.1 原理

KDA 的基本假设是：在每一层、每一个 Head 中，每个请求都维护一个固定大小的状态矩阵 S_t。这个矩阵代表当前请求到目前位置为止的全局记忆。它不是为每个历史 Token 单独保留一条记录，而是把所有历史 Token 的上下文信息压缩到同一个矩阵里。

当 KDA 处理当前 Token 时，模型会根据当前 Token 的 Hidden State 和对应的投影参数，计算出：

• q_t：Query，用于从记忆中读取信息；
• k_t：Key，用于定义寻址模式；
• v_t：Value，用于提供要写入的内容；
• α_t：遗忘门，决定各个 Key 通道保留多少旧信息；
• β_t：写门，决定当前残差写入多少；
• γ_t：输出门，控制输出信息的强度。

KDA 的状态更新可以写成：

\[S_t =S_{t-1}\operatorname{Diag}(\alpha_t) +\beta_tv_tk_t^\top -\beta_tS_{t-1}(\alpha_t\odot k_t)k_t^\top \]

(1) 先使用 \(\bar S_t=S_{t-1}\operatorname{Diag}(\alpha_t)\)

\(S_{t-1}\operatorname{Diag}(\alpha_t)\)中，遗忘门 \(\operatorname{Diag}(\alpha_t)\) 决定旧状态中的不同 Key 通道保留多少。因为 α_t 是向量，所以并不是所有列都以相同比例衰减。第 j 个 Key 通道拥有独立的遗忘速度 α_{t,j} ，重要的通道可以保留更多，应该淡出的通道可以衰减得更快。

如果状态矩阵的形状是：\(S_t\in\mathbb R^{d_v\times d_k}\)。那么右乘 Diag(α_t) 会对状态矩阵的每一列分别缩放。每一列对应一个 Key 通道，因此 KDA 实现的是“按 Head、按 Key Channel”的遗忘。

门的规模也因此从 GDN 的：

num_heads

扩展为：

num_heads × head_dim

如果没有遗忘门，纯线性注意力可能只是：\(S_t=S_{t-1}+v_tk_t^\top\)。由于矩阵加法满足交换律，输入顺序 [A,B] 和 [B,A] 会产生相同的最终状态。这意味着简单的累加状态天然缺少位置信息：它知道有哪些关联被写入，却不一定知道这些关联的先后顺序。KDA 中的遗忘可以部分恢复这种顺序信息。距离当前越远的 Token，经历的衰减次数越多，它的信号就可能越弱。因此，遗忘过程可以看成一种隐式的相对位置编码。

(2) 再执行 Delta Rule：\(S_t=\bar S_t+\beta_t\big(v_t-\bar S_tk_t\big)k_t^\top\)。

\(\beta_tv_tk_t^\top\) 是当前 Value 的写入，即当前 Token 带来的原始新知识。它把 Value v_t 和寻址模式 k_t 绑定起来，准备写入状态。

\(-\beta_tS_{t-1}(\alpha_t\odot k_t)k_t^\top\) 这一项会减掉衰减后的旧状态中、与当前 Key 模式匹配的那部分内容。它的作用不是简单地把新内容加进去，而是避免同一个特征在长上下文中被不断重复累加。先删除或修正旧的匹配内容，再写入新的目标 Value，状态因此具备更精确的覆写能力。

与标准做法的公式级对比：

4.1.2 对比基线

改动带来的收益：O(1) 显存、O(T) 时间；代价：检索变为压缩近似。

4.1.3 计算复杂度分析

KDA 是典型"递推序列模型"：训练/prefill 用全序列 chunkwise 前向（每个位置都要输出，用于算 loss）；推理 decode 用单步递推（自回归逐 token 生成，状态传递）。桥梁是 chunkwise（块内并行、块间串行）。

• FLOPs（单 token）：KDA = O(d_k·d_v)（状态更新 + 输出读出）vs softmax = O(T·d_k)（需与全部历史比较）。T 大时 KDA 完胜。
• 内存：KDA 状态 = d_k×d_v/head（固定）；softmax KV = T×(d_k+d_v)/head（线性增长）。
• 参数：KDA 无新增参数在递推本身（投影共享）。

4.1.4 参数敏感性

• head_dim_k（默认 128）：决定状态方阵大小与"key 记忆方向"的容量。调大→状态容量↑但 FLOPs↑；过小→容量不足。
• head_dim_v（默认 128）：value 信息容量。
• 论文经验：d_k=d_v=128 是 Kimi Linear 经过调参后的默认，配合 96 head（d=7168）。

4.1.5 反方辩论

• 最强反对理由：固定矩阵的压缩必然导致信息丢失，超长上下文的精确检索（如"在第 40 万 token 提到过的某个具体数字"）会失败。
• 若失败，最可能错在假设：假设"可压缩的记忆足够表达任务所需信息"。对需要精确逐 token 回溯的任务（复杂 agentic 轨迹、代码库检索），此假设脆弱。
• 胜出条件：记忆可近似、语义密集、长上下文且 token 语义冗余高（对话、长文总结）。失败条件：需要精确位置检索、稀疏的关键信息。
• 更稳健做法：状态扩容（state expansion）、或 KDA+MLA 混合比例调高（3:1 是论文最优，但任务敏感）。
• 受伤方：长程精确检索任务方；受益方：1M 上下文推理吞吐方。

4.1.6 工程实现

伪代码（逐行）：

for t in 1..T:
    state *= exp(g[:, :, t])                          # Diag(alpha) S      channel 衰减
    sk = einsum('bhd,bhdv->bhv', k[:, :, t], state)   # S^T k  预测
    delta = beta[:, :, t] * (v[:, :, t] - sk)         # 残差 × 学习率
    state += k[:, :, t] * delta                       # rank-1 外积更新
    out = einsum('bhd,bhdv->bhv', q[:, :, t], state)  # 读出

项目真实代码（kda_ref.py 中 kda_recurrent）：

def kda_recurrent(q, k, v, g, beta, *, initial_state=None,
                  reset_mask=None, dtype=torch.float64):
    """Eq. 1 的字面递推。是其他所有 KDA 路径的 ground truth。"""
    qf, kf, vf, gf, bf = _as_float(q, k, v, g, beta, dtype=dtype)
    b, h, t, dk = qf.shape
    dv = vf.shape[-1]
    state = (torch.zeros(b, h, dk, dv, dtype=dtype, device=q.device)
             if initial_state is None else initial_state.to(dtype).clone())
    if t == 0:
        return vf.new_zeros(b, h, 0, dv), state

    outputs = []
    for i in range(t):
        if reset_mask is not None:
            r = reset_mask[:, i].view(b, 1, 1, 1)
            state = torch.where(r, torch.zeros_like(state), state)
        alpha = gf[:, :, i].exp()                    # 逐 channel 保留因子
        state = alpha.unsqueeze(-1) * state          # Diag(alpha) S
        ki = kf[:, :, i]                             # 当前 key
        bi = bf[:, :, i].unsqueeze(-1)               # 当前写入强度
        sk = torch.einsum("bhd,bhdv->bhv", ki, state)   # S^T k
        delta = bi * (vf[:, :, i] - sk)              # 写入方向 = 强度 * 残差
        state = state + ki.unsqueeze(-1) * delta.unsqueeze(-2)  # 外积更新
        outputs.append(torch.einsum("bhd,bhdv->bhv", qf[:, :, i], state))
    return torch.stack(outputs, dim=2), state

与开源参考实现逐位核对：

• 本项目 kda_recurrent 与 Kimi Linear 论文 Eq.1 完全一致；与 kimi3-master 的 naive_recurrent_kda 对拍 max diff = 0.0（同一公式、同一 dtype 路径）。
• SGLang fused_recurrent_gated_delta_rule_fwd_kernel 的 IS_KDA 分支：与论文的 KDA 递推等价；本项目 CPU 版与其逻辑逐位对齐（差异仅在 dtype/融合）。

4.1.7 常见疑问 & 讨论

Q：in-context learning 即隐式优化

KDA 每处理一个 token，就用它对状态做一次带遗忘的在线更新。状态相当于一个小模型的参数，当前上下文就是它的训练数据。因此上下文学习可以看成：读入 token → 对内部状态做一次小步优化 → 用更新后的状态处理下一个 token。衰减则类似在线学习中的正则化，防止状态参数无限增长。

Q：为什么按 head 独立

KDA 每个 head 维护独立状态矩阵 S，体现在三处：

• 初始化：状态张量形状是 (b, h, dk, dv)——第 1 维是 head 数。每个 head 一张 dk×dv 矩阵。
• 递推更新：衰减 alpha 是 (b, h, dk)（逐 head 逐 channel），state 的更新沿 head 维独立进行。head 之间没有任何交换——这是"独立"的核心。
• 缓存：KDACache.state 形状 (batch, num_heads, head_dim_k, head_dim_v)，即每 head 一个状态矩阵。

不同 head 可以学到不同的记忆模式（有的关注近程、有的关注长程、有的关注特定内容）。如果跨 head 共享状态，这些分化就会被抹平。这也正是 HeadRMSNorm 存在的理由——既然每个 head 的 S 是独立演化的，它们输出的尺度就可能差异很大，需要按 head 归一化。

Q：为什么 beta 是逐 head 标量？

在线学习视角里，β_t 就是学习率：S_t = S_{t-1} + β_t · gradient。如果 beta 做成逐通道向量，表达力会更强，但代价是：状态更新需要更多门、内存和带宽增加、chunkwise 分块公式更复杂、kernel 更难规整。逐 head 标量是一个折中：每个 head 决定本次更新整体有多大；通道之间的差异交给 Diag(α_t) 负责；beta 仍然可以看成 head 级学习率。

Q：为什么 beta 必须只依赖当前 token？

如果 beta 依赖完整历史，就需要先完成整段递推才能决定更新幅度，chunk 内并行会被破坏。让 beta 由当前 token 即时预测，牺牲一部分全局自适应能力，换取：单步可计算、chunkwise 可展开、状态可并行归约、kernel 更规则。这是并行化前提，不是随意简化。

Q：衰减已有了，为什么还要 Delta 擦除？

它们的分工不同：衰减是慢慢忘掉旧状态；Delta 是沿当前 Key 快速覆盖局部关联。类比：衰减像房间整体逐渐降噪，Delta 像拿橡皮擦掉某一行。一个负责"时间"，一个负责"方向"。但衰减有下界，它不能一步清空；缺少的快速清除能力由 Delta 项补充。但 Delta 一次通常只能擦一个方向，所以"又快、又有选择性、又能清除任意子空间"无法在单步中同时实现。

Q：递推公式会不会数值不稳定？

不会：中间计算强制 fp32/fp64，alpha 恒 < 1，无上溢。chunkwise vs recurrent 对拍 chunk=1..64 全部一致。

Q：beta 接近 1 时会发生什么？

完全覆盖旧记忆该方向，类似"写盘覆盖"。diagnostics 中有 beta 饱和度监控。

Q：固定 state 为什么是 KDA 的根本限制？

一个固定 state 必须容纳任意长历史。历史中的独立事实越多，越容易出现：两段不同历史映射到相同状态、相近 Key 共享状态方向、远古信息被衰减、Query 无法找回被压缩的 token 身份。所以 KDA 的固定状态既是高效来源，也是信息容量上限。

Q：query 只能在已经画好的天花板下挑方向

KDA 在写入历史时就决定了每条信息的存活程度。未来 Query 只能在剩下的状态里挑方向。如果某条信息写入时 beta 很小、后续被多次衰减、与其他信息发生覆盖，那么即使未来 Query 非常相关，也没有原始 token 可供重新读取。

4.2 有界衰减 log-decay（Eq.5）

有界门控把 1/gamma 变成数据无关硬上界 exp(-tile*g_min)=exp(80)=5.54e34 < BF16 上限，于是每个因果 tile（含对角）都能无条件编译成稠密 Tensor-Core matmul——数值安全换全稠密调度。

4.2.1 原理

直觉：Kimi Linear 的遗忘速度是"无底洞"（softplus 无下界），记 16 个 token 就衰减到 float 下溢；K3 给遗忘速度设了"下限" g_min，保证衰减不会太快、倒数不会溢出。

专家视角：

g_t = g_min · sigmoid(exp(A_h) · z_t)            (Eq.5)

• z_t = W_up W_down x_t + b_α：低秩 decay logits。
• A_h：逐 head 对数尺度（a_log，初始 0，所以初始 exp(A)=1）。
• sigmoid 把 exp(A)·z 压到 (0,1)，再乘 g_min<0 → g ∈ (g_min, 0)。

下界设计是 FlashKDA 稠密化的关键前提。因为 g ≥ g_min，一个 tile（16 token）内的累积 log-decay ∈ (16·g_min, 0)，倒数 1/γ ≤ exp(−16·g_min)=exp(80)≈5.5e34，BF16 最大 3.4e38 可容纳。无下界则 1/γ 溢出，无法稠密化。

有界衰减 = 把"数值可稠密化"作为模型设计的硬约束，是算法-硬件协同设计的典型：不是先定数学再找硬件适配，而是数学形态直接由硬件精度决定。

4.2.2 对比基线

4.2.4 计算复杂度分析

• 无额外 FLOPs（sigmoid 是廉价逐元素）。
• 参数：a_log（每 head 1 个）→ 新增 h 个标量，可忽略。

4.2.5 参数敏感性

• g_min（默认 −5.0）：调大（如 −3）→ 更保守衰减，数值更安全但遗忘慢；调小（如 −8）→ 衰减更快，tile 内 exp(−16·8)=exp(128) 溢出 BF16，safe_chunk_size 掉到 10（数值实测见 §4.2.7）。
• alpha_init（默认 0.99）：初始 retention 因子，alpha_bias = log(log(α_init)/g_min / (1−…)) 让模型从"近似恒等衰减"起步，避免初始训练时状态快速崩溃。

4.2.6 工程实现

伪代码：

z = alpha_up(alpha_down(x)) + alpha_bias     # (B,T,H·Dk)
scale = exp(a_log)                            # (H,)
g = g_min * sigmoid(scale * z.view(B,T,H,Dk))

项目真实代码（kda.py 中 log_decay）：

def log_decay(self, x: Tensor) -> Tensor:
    """Eq. 5：g = g_min * sigmoid(exp(A_h) z)，范围严格在 (g_min, 0)。

    这个下界是本实现能与 SGLang/FLA 的 lower_bound 机制对齐的关键：
    正因为 log-decay 有下界，chunkwise 需要的 1/gamma 才不溢出，稠密化才成立。
    """
    b, t, _ = x.shape
    z = self.alpha_up(self.alpha_down(x)) + self.alpha_bias   # (B, T, H*Dk)
    z = z.view(b, t, self.num_heads, self.head_dim_k)         # 拆回逐 head
    scale = self.a_log.exp().view(1, 1, self.num_heads, 1)    # exp(A_h) 广播
    return self.cfg.g_min * torch.sigmoid(scale * z)

safe_chunk_size 函数（kda_ref.py）依据 g_min 算出当前 dtype 的最大安全 chunk（默认 g_min=−5 时 bf16/fp32 返回 16 = 论文 tile）。

def safe_chunk_size(g_min: float, dtype: torch.dtype, *, headroom: float = 1e3) -> int:
    """求"其累积倒数衰减仍能被 dtype 装下"的最大 tile 长度。"""
    limit = _DTYPE_MAX[dtype] / headroom
    return max(1, int(math.log(limit) / -g_min))

实测数值：

• max_reciprocal_decay(-5, 16) = 5.54e+34（= exp(80)）
• safe_chunk_size(-5, bf16) = 16（正好是论文 tile！）
• safe_chunk_size(-5, fp32) = 16
• safe_chunk_size(-5, fp16) = 1（fp16 上限 65504 太小，tile=1 才能装下）
• safe_chunk_size(-8, bf16) = 10（g_min 越负，安全 chunk 越小）

4.2.7 数值稳定性验证实验

失败 case：无界衰减下 chunkwise 溢出

用 softplus 无界衰减跑 chunk=16，观察 1/γ 上溢 inf：

fp32 溢出临界约在 z≈6（exp(16*6) 已超过 3.4e38）。这是数据相关的——同样 tile=16，z 稍大就崩。

chunk=64 —— 连有界版也溢出（印证 safe_chunk_size 边界）

chunk=64 时 exp(64*5)=exp(320) 远超 fp32 上限，有界版同样 inf/NaN。这正好说明 safe_chunk_size(-5, fp32)=16 的边界是真的——超出安全 tile 必崩。

4.2.8 讨论

Q1：为什么需要下界有界衰减？

如果衰减可以无限接近 0，一步就可能把整个状态几乎清空：S_t ≈ 0。后果是：梯度可能突然消失；长程信息可能被意外抹掉；低精度下容易出现数值不稳定；不同 token 之间的状态尺度变化过大。因此需要限制衰减不能小于某个安全下界。它牺牲"一步彻底忘光"的能力，换来状态连续性和数值安全。

Q2：为什么要在 log 空间计算衰减

多个 token 的衰减会连乘：α_{s→t} = ∏_{i=s+1}^{t} α_i。如果直接乘很多小于 1 的数，容易下溢成 0。于是改在 log 空间：log α_{s→t} = Σ_{i=s+1}^{t} log α_i。乘法变加法，数值范围更容易控制。

Q3：为什么 Kimi Linear 用无界 softplus？

先拆公式

Kimi Linear（KDA v1）的门控：

g = -exp(A) * softplus(z)      （无下界：z -> +inf 时 softplus(z) -> z，g 无底洞）
alpha = exp(g)                  逐 channel 保留因子，0 < alpha <= 1

chunkwise 并行公式（Eq.3-4）里，关键量是"从 chunk 起点到位置 r 的累积衰减连乘"：

gamma^{1->r} = prod_{s=1..r} alpha_s = exp(sum g_s)     <= 1
khat_r = k_r / gamma^{1->r}      （Eq.4 的 K/Gamma 重缩放）

两个方向同时出问题

① 累积衰减乘积下溢：gamma^{1->r} = exp(sum g_s)。因为 g = -softplus(z) 无下界，只要 z 偏大（模型学到"快速遗忘"），sum g_s 就会非常负，exp 直接落到 0（下溢）。

② 其倒数上溢：chunkwise 要算 khat_r = k_r / gamma，等价于 k_r * (1/gamma)。1/gamma = exp(-sum g_s)，当 sum g_s 很负时 1/gamma 无界增长——超过 float 上限就是 inf。

结果就是：

• 显式 position-pair 路径：即 chunk 内位置两两显式配对（position-pair 是指"row/col 都是 chunk 内具体位置"的对角块），逐对计算注意力/写入关系，是 FlashAttention 式稀疏/显式路径。
• 因果三角分解时，对角 tile 需要处理"token 与它自己及同 chunk 内其它位置"的相互作用。如果 1/gamma 会溢出，就不能安全地把对角 tile 塞进一个普通的稠密矩阵乘（Tensor-Core 里中间值爆表）。
• 所以 Kimi Linear 只能绕开：对角 tile 不用稠密 matmul，而是走显式的 position-pair 计算路径（对每一对位置单独算，避免一次性放大 1/gamma）。代价是失去张量核加速，对角 tile 变慢。

即，无界 softplus 让 g 无底洞 → 1/gamma=exp(-sum g) 无上界（数据相关，z 稍大就 inf）→ 对角 tile 不能安全用稠密 matmul，只能退回显式 position-pair 路径（慢）。

Q4：K3 如何消除慢速路径？

早期 chunkwise 计算为了防止衰减数值极端，需要对角 tile 特殊处理。

K3 的思路是从源头限制 log‑decay：衰减 logit → 设置下界 g_min → 倒数重标定因子始终可表示 → 对角与非对角 tile 使用统一矩阵形式。

如果 g_min=-5，意味着最极端的衰减也不会进入难以表示的数值区域。这样可以避免单独的慢速修补路径；让对角 tile 也更接近规则 GEMM；统一 Tensor Core 计算；用少量精度约束换取更高硬件利用率。这不是"数学上消灭了因果关系"，而是让因果关系的数值表达足够稳定，从而可以用同一套矩阵 kernel 处理。

Q5：为什么"每个因果 tile（含对角）都能用稠密 Tensor-Core matmul？

因果分解后，A = Tril[(Q*Gamma)(K/Gamma)^T] 的每个对角 tile 都要算一次块内 matmul，中间量含 K/Gamma（即 1/gamma 放大）。既然 1/gamma <= 5.54e34 是数据无关、BF16 装得下的硬上界，那么：

• 编译器/内核可以无条件（对任何输入、任何 token）把这个对角 tile 编译成一个稠密 Tensor-Core matmul，不需要运行时判断"这里会不会溢出"；
• 无需退回显式 position-pair 路径，对角 tile 和普通 tile 一样走张量核加速。

这就是 K3 能比 Kimi Linear 更"稠密"、更快的原因——数值安全性买来全稠密调度。

Q6：为什么有界衰减把数据相关上界变成数据无关硬上界

K3 的 log_decay：

g = g_min * sigmoid(exp(A_h) * z)       (Eq.5, g_min = -5)

sigmoid 把输出压到 (0,1)，再乘 g_min<0 → g ∈ (g_min, 0) 严格成立，与输入 z 无关。

于是跨过一个 tile（16 token）的累积 log-decay：

tile 内累积 log-decay = sum_{16 个位置} g_s ∈ (16*g_min, 0) = (-80, 0)
1/gamma = exp(-sum g_s) <= exp(80) ≈ 5.5406e34     （数据无关的硬上界！）

BF16 最大有限值 ≈ 3.3895e38，FP32 ≈ 3.4028e38。5.54e34 < 3.39e38 → 装得下。这就是 exp(80)=5.5e34 仍在 BF16 范围内的来历。

Q7：a_log 和 dt_bias 的作用

在状态空间模型和 KDA 类门控中，常见做法是把"容易训练但无约束的参数"映射成"满足稳定性要求的衰减量"。

• a_log：通常用于表示基础时间尺度的对数参数。用对数存储的好处是：参数可以覆盖多个数量级；优化时更容易表达快衰减和慢衰减；通过指数映射可以保证基础速率为正。直观上，A_log 决定"这个通道原本记忆多久"。
• dt_bias：是对每个输入步长或时间间隔的偏移。它控制当前 token 到来时，基础时间尺度如何被调节。直观上：A_log 是默认记忆寿命；dt_bias 是当前步的时间推进偏移；输入相关门控制当前 token 是否需要更快或更慢地更新。

三者为什么能构成下界有界衰减：通常会经过类似 A_log → 正的基础衰减率；dt_bias + 当前输入 → 正的有效步长；有效步长 × 衰减率 → 映射到 log‑decay → 施加最小 log‑decay 下界 → 得到 alpha 的流程。如果使用 softplus、exp 或等价正值映射，可以保证时间步长和衰减率不出现非法符号；如果再对 log‑decay 做下界截断，就能保证衰减不会一步清空、衰减仍然可以随输入变化、长时间递推更容易保持稳定。

4.2.9 反方辩手

• 下界 g_min=−5 意味着每个 token 至少保留 e^{-5}≈0.0067 的衰减比例——对超长上下文，信息"忘不掉"可能累积噪声。
• 若失败，错在假设"衰减下限不会显著伤害 recency 建模"。
• 若把 g_min 调小（如 -8）去加速遗忘，safe_chunk_size 掉到 10，tile=16 就溢出 → 稠密化失效（见 §4.2.10 图例左下）。
• 胜出：上下文需稳定渐进衰减；失败：需要强烈 recency 偏差的任务。
• 稳健替代：把 g_min 设为可学习。

4.2.10 图例

四格读法：

• 左上：1/gamma 上界 vs tile 大小（对数）。有界版是一条数据无关的直线 exp(-g_min*tile)，tile=16 时在 BF16 上限下方；无界版没有这条线（可以无限往上长）。
• 右上：实测 tile=16——无界 z+6/z+9 爆成 inf（红柱），有界版安全。
• 左下：safe_chunk_size vs g_min——g_min 越负，允许 chunk 越小（-5→16，-8→10）。
• 右下：失败 case——chunk=64 时连有界版也溢出（超出安全边界必崩）。

4.3 channel-wise decay 与低秩参数化（Eq.2）

设计哲学：KDA 是"12288 个 channel 旋钮 + 96 个 head 温度"——channel-wise 提供细粒度控制，head-wise 提供粗粒度调节，低秩分解让参数量可控。

低秩 decay 是"LoRA 思维在模型架构内部的重用"——不只在微调时压缩，而是在生成参数时就压缩。这是参数高效设计泛化到架构级的一个例证。

4.3.1 原理

直觉：GDN 一个 head 只用一个遗忘旋钮；KDA 给每个channel 一个旋钮（共 d_k 个）。但 d_k 个旋钮如果每个都从 x 直接生成，参数太多，所以用低秩瓶颈压缩。

专家视角：

z = W_up W_down x + b_α,   W_down ∈ R^{rank×d}, W_up ∈ R^{d_k×rank}   (Eq.2)

rank = alpha_rank（默认 512）远小于 d_k = 96×128=12288。q = q_proj(x) 后做 SiLU + L2 归一化。低秩让 decay 参数共享一个低维子空间，同时保留 channel 级差异化——这是"用参数效率换细粒度控制"。

"衰减模式的秩被限制为 128（或配置的 alpha_rank）"可以直观理解为：模型没有为每个 head 单独学习一套完全独立的遗忘规则，而是先学习一个最多 alpha_rank 维的"遗忘意图空间"，各个 head 从这个空间读取自己的衰减方式。像一组乐队共用 128 个旋钮：每个 head 可以有自己的音量；但所有 head 的变化都来自同一组基础旋钮；参数更少、结构更稳定；不同 head 的遗忘行为不是完全独立。

优点：降低参数和计算；让不同 head 共享统计规律；有利于硬件矩阵计算。

缺点：如果真实任务需要超过 rank 种互相独立的遗忘模式，低秩空间会形成瓶颈；不同 head 之间可能出现相关性过强；这是表达力换效率的典型折中。

逻辑图

alpha_rank 设置

当前配置（config.py:31）：

alpha_rank: int = 512

相关维度：

设置逻辑（config.py:290）：

_check(0 < kda.alpha_rank <= kda.k_dim, "alpha_rank must be a low rank of the key dim")

• alpha_rank 必须满足 0 < alpha_rank <= k_dim（即 ≤ 12288）
• 当前 512 是 k_dim/24，压缩率约 4.2%

参数量对比：

如何压缩到 alpha_rank？（低秩分解）

关键设计：alpha_up 初始权重置零（kda.py:144）

self.alpha_up.weight.zero_()  # 初始 log-decay 恰好等于偏置

这意味着：

• 训练初期：z = 0 + alpha_bias = alpha_bias，模型从"近似恒等衰减"开始
• 训练过程中：W_up 逐渐学习，通过 W_down 的压缩表示调制 alpha_bias
• LoRA 风格：先让模型跑起来，再逐步引入低秩修正

如何映射到每个 head？（广播 + 逐 head 调制）

这是两级结构：

第一级：channel-wise（共享基底）
─────────────────────────────────

z = W_up @ W_down @ x + alpha_bias
         │
         ▼
    (B, T, H*Dk) = (B, T, 12288)
         │
         ▼
    view(B, T, H, Dk) = (B, T, 96, 128)
         │
         └──► 每个 channel 有自己的衰减偏置
              （12288 个"旋钮"）


第二级：head-wise（逐头缩放）
─────────────────────────────────

a_log: (H,) = (96,)  ←── 可学习参数，初始化为 0
    │
    ▼
exp(A_h): (H,)  ←── 初始为 1，可正可负
    │
    ▼
view(1, 1, H, 1)  ←── 广播到 (B, T, H, Dk)
    │
    ▼
scale * z: (B, T, H, Dk)
    │
    ▼
g = g_min * sigmoid(scale * z)

数学表达：

对于 head h 的第 d 个 channel：

g[b,t,h,d] = g_min * sigmoid( exp(A_h) * z[b,t,h,d] )
           = g_min * sigmoid( exp(A_h) * (W_up @ W_down @ x + b)[b,t,h*Dk+d] )

直觉理解：

• alpha_bias：12288 个 channel 各自的"基准衰减偏好"（共享参数）
• W_down/W_up：根据输入 x 动态调制这些偏好（低秩压缩）
• exp(A_h)：每个 head 的"全局温度"，控制该 head 整体衰减的敏感度

4.3.2 对比基线

4.3.3 反方辩手

• 低秩瓶颈限制了 decay 的多样性（模型可能学到冗余的 rank 子空间）。
• 若失败，错在假设 rank 足够捕获 channel 差异。
• 胜出：channel 间 decay 高度相关；失败：需要正交 channel 记忆。
• 稳健替代：直接投影（参数多但更灵活）。
• 受伤方：极端长上下文记忆任务。

4.3.4 计算复杂度

• 参数：d×rank + rank×d_k。默认 rank=512：4096×512 + 512×12288 ≈ 8.4M，占比小。

4.3.5 参数敏感性

• alpha_rank（默认 512）：调大→表达力↑参数量↑；调小→欠表达。

4.3.6 工程实现

项目代码 kda.py 中 alpha_down/alpha_up/alpha_bias，初始化 alpha_up 权重置零（LoRA 惯例）。初始化 reset_parameters 中求解方程让 retention 从 alpha_init 起步：

# 低秩 decay logits（Eq. 2）：z = W_up W_down x + b_alpha。
self.alpha_down = nn.Linear(dim, cfg.alpha_rank, bias=False)   # 压到 rank
self.alpha_up = nn.Linear(cfg.alpha_rank, k_dim, bias=False)   # 展开回 k_dim
self.alpha_bias = nn.Parameter(torch.empty(k_dim))             # 逐 channel 偏置

def reset_parameters(self) -> None:
    """定制初始化，让模型一开始处于"近似恒等衰减"的稳定状态。
    - ``alpha_up`` 权重置零：初始 log-decay 恰好等于偏置（LoRA 惯例：先把
      up 投影关掉，梯度仍能到达它，一旦非零 down 投影就开始起作用）。
    - 求解 ``g_min * sigmoid(z) = log(alpha_init)`` 得到 z，使 retention 因子
      从 ``cfg.alpha_init`` 起步，同时尊重 ``exp(g_min)`` 下界：
        ``alpha_bias = log( target/(1-target) ),  target = log(alpha_init)/g_min``
    """
    with torch.no_grad():
        self.alpha_up.weight.zero_()
        target = math.log(self.cfg.alpha_init) / self.cfg.g_min
        self.alpha_bias.fill_(math.log(target / (1.0 - target)))

与开源参考核对：kimi3-master 同样用低秩 decay projection（decay_down/up）。

4.3.7 讨论

Q1：为什么 g_proj 要全秩，而衰减门保留低秩

K3 将新增的输出门 g_proj 从低秩扩大为全秩，核心原因是：Output gate 需要对最终输出的每个通道独立决定"写入多少"。AttnRes 的 value 路不做归一化，分块摘要主要靠求和。于是 output gate 的幅度会直接影响该层在深层记忆中的份额。如果输出门太低秩，很多通道只能共享同一个调节方向，无法细粒度控制。

衰减门不同。衰减描述的是"这个 token 想让哪些状态方向长期保留、哪些方向快速遗忘"。这种"遗忘意图"通常比最终输出通道低维得多。因此可以先用低秩向量表达遗忘意图，再由不同 head / 通道做线性读出。

总结：

Q2：为什么 diag(alpha) 放在梯度步之前意味着 KDA 更强？

如果所有坐标共享同一个 alpha，状态空间各方向用相同速度更新，近似各向同性。diag(alpha) 让每个通道拥有自己的有效学习率：某些方向快速接受新信息；某些方向慢速保留旧信息；状态流形可以各向异性地更新。因此 KDA 比 GDN 有更细的时间尺度和记忆分工。

Q3：alpha_bias 和 dt_bias 的区别

alpha_bias 和 dt_bias不是同一个参数，但有相似的"加性偏置"形式。

在 Model-Viz 中，它们属于不同实现和不同命名体系：

alpha_bias 的计算路径

KDA 模块中：

z = self.alpha_up(self.alpha_down(x)) + self.alpha_bias
g = g_min * sigmoid(exp(A_log) * z)
alpha = exp(g)

其中：

• alpha_bias 的形状是 (H * Dk,), 即每个 key channel 一个偏置；
• 它直接参与 decay-logit z；
• alpha_bias 在初始化时根据 alpha_init 和 g_min 求出，使初始 retention factor alpha 接近目标值；
• 它是 KimiDeltaAttention 的持久可训练参数。

dt_bias 的计算路径

线性注意力变体中：

z = gate_logit
if dt_bias is not None:
    z = z + dt_bias
g = g_min * sigmoid(exp(A_log) * z)  # v2
alpha = exp(g)

其中：

• dt_bias 可为 (H, Dk) 或 (Dk,);
• 它是 kimi_linear_step() 的可选输入；
• 代码本身不在该函数内创建或初始化它；
• 它用于对齐 SGLang/KDA 参考实现中的 gate bias。

为什么容易混淆？

在采用相同 v2 门控公式时，两者都可能出现在：

逐 channel 的 logits
+ 一个可学习或外部传入的 bias
→ sigmoid
→ g
→ exp(g) 得到衰减因子 alpha

因此，它们在数学作用上都属于“改变衰减门控输入的加性偏置”。但这只能说明它们功能相似或可对应，不能说明它们是同一个变量。

结论

更准确地说：

• alpha_bias是Model-VizKDA主模块中明确声明并初始化的decay-logit参数；
• dt_bias是线性注意力step函数接收的可选gatebias，主要用于与外部参考实现对齐；
• 两者可以在某些公式对照中承担相近角色，但参数归属、调用接口、初始化方式和代码位置都不同；
• 不能把alpha_bias直接改名为dt_bias，也不能据此断言两个实现完全等价。

Q4：为什么不直接让每个状态元素都有独立遗忘门

既然可以从“每个 Head 一个标量”进化到“每个 Head 的每个 Key 通道一个值”，自然还可以继续设想：

\[\alpha_t\in\mathbb R^{d_v\times d_k} \]

让状态矩阵中的每一个元素都有独立的遗忘系数，并执行逐元素衰减：

\[\bar S_t=S_{t-1}\odot\alpha_t \]

这种设计的表达能力更强，因为每个 Value 行、每个 Key 列的组合都可以独立决定保留多少。

但表达能力不是唯一目标。逐元素遗忘会带来明显的计算和硬件代价：

1. 每个 Head、每个 Token 的门值数量，会从 head_dim 增加到 head_dim_v × head_dim_k；
2. 门的生成成本更高；
3. 状态读写的数据量更大；
4. 更重要的是，它破坏了当前更新中可复用的代数结构。

按 Key 通道遗忘时，所有 Value 行共享同一个右乘对角矩阵。这样，状态更新仍然具有“对角矩阵加秩一修正”的结构，能够设计高效的分块并行算法。

如果每个元素都独立遗忘，这种共享结构通常就不存在了，原有的 Chunkwise 算法不能直接照搬。

所以，这里不是“不想把表达能力做得更强”，而是在表达能力、并行效率、门生成成本和显存读写之间做取舍。代数结构决定了并行上限，硬件成本则限制了可以使用的表达能力。

4.4 Head-wise RMSNorm（按 head 归一化的 RMSNorm）

HeadwiseRMSNorm 是 RMSNorm 的一种变体：只按每个 head 的最后一维归一化，不跨 head。K3 用它在 KDA 输出门第一步把各 head 的输出重缩放到可比尺度，再接全秩 sigmoid 门。‘

为什么按 head 归一化而非整体归一化：KDA 每个 head 维护自己的递推状态矩阵 S，head 之间的输出尺度可能差异很大（不同 head 关注不同频率/内容）。若整体归一化，大尺度的 head 会主导归一化因子，把小 head 的信号压没；按 head 归一化则保证每个 head 的输出都被重缩放到可比尺度，再统一被后面的全秩门调制。

4.4.1 前身

• LayerNorm：先减均值再缩方差。缺陷：多一次 reduction、有 bias。
• RMSNorm：只按 RMS 缩放，去掉 re-centering。缺陷：跨整个向量归一化，若向量是多 head 拼接的，大 head 会主导归一化因子。
• HeadwiseRMSNorm（K3）：把 RMSNorm 限制到每个 head 内部，解决"跨 head 归一化时大 head 主导"的问题。

4.4.2 同期竞品

4.4.3 原理讲解

直觉类比视角

想象一群同学（多个 head）回答问题，每个人的音量（尺度）差别很大 —— 有人低声细语（小 head），有人震耳欲聋（大 head）。如果大家一起统计平均音量再调，大嗓门会盖过小嗓门。

HeadwiseRMSNorm 的做法是每人单独把音量调到标准值（各自 RMS=1），让所有人的"音量"可比。这样，之后一个统一的门（sigmoid）就能公平地决定每人读多少，而不会被大嗓门天然压过。

准专业水平视角

输入输出

• 输入：x ∈ R^{B×T×H×D_h}（H=head 数，D_h=head 维）
• 输出：y ∈ R^{B×T×H×D_h}

公式

rms(x[b,t,h]) = sqrt( (1/D_h) * sum_{j} x[b,t,h,j]^2 )     # 每 head 独立
y[b,t,h,j] = x[b,t,h,j] / sqrt( rms(x[b,t,h])^2 + eps ) * weight[j]

• 归一化只作用于最后一维（D_h），不跨 head（mean(-1) 而非整体 mean）。

与标准 RMSNorm 对比

专家视角

• 为什么 head-wise：KDA 每个 head 维护独立状态矩阵 S，head 间输出尺度差异可能很大（不同 head 关注不同频率/内容）。跨 head 归一化会让大 head 主导归一化因子，把小 head 压没；按 head 归一化则每 head 独立可比。
• 为什么用于输出门：K3 的 KDA 输出门 = sigmoid(W_g x) ⊙ RMSNorm_head(o)（Eq.6）。先归一化保证门读到的输入尺度稳定，门才可能公平调制。
• 为什么带 weight 仿射：归一化后乘可学习 weight，恢复每 head 需要的尺度自由度。
• 论文未解决问题：head 间的尺度是否应完全解耦？是否应每 head 独立 weight？开放。

4.4.4 实现 & 问题

Q1：HeadRMSNorm 和 RMSNorm 代码结构不是一模一样吗？只是 dim 维度不同？

是的，代码结构几乎一样。对比：

两者都调同一个 rms_norm，都只对最后一维归一化。唯一的实质区别就是 weight 的形状（dim vs head_dim），而 weight 的形状又由输入张量的最后一维是什么决定。

def rms_norm(x: Tensor, weight: Tensor | None = None, eps: float = 1e-6) -> Tensor:
    """RMS 归一化的核心计算（fp32 累积）。

    公式：``y = x / sqrt(mean(x^2) + eps) * weight``，其中 ``mean`` 取最后一维。

    三个实现细节：
    1. **fp32 升精度**：``x.float()`` 先把输入转成 fp32 再做平方求均值。因为
       bf16/fp16 下 ``x^2`` 会严重丢失精度甚至下溢，均值统计必须用更高精度，
       否则归一化因子本身就有偏。
    2. **先降回原 dtype、再乘 weight**：权重乘在降精度之后，这样训练与部署的
       数值行为一致（量化误差特征可复现）。
    3. ``weight`` 是逐通道的仿射参数，可为 None（此时等价于纯归一化）。
    """
    dtype = x.dtype
    x32 = x.float()                                          # 升到 fp32 算统计量
    out = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps)  # x / rms
    out = out.to(dtype)                                      # 降回原 dtype
    if weight is not None:
        out = out * weight                                   # 逐通道仿射
    return out

class RMSNorm(nn.Module):
    """标准 RMSNorm 层：对最后一维做单位 RMS 归一化 + 可选逐通道仿射。

    典型用法（pre-norm）：
        attn_norm = RMSNorm(dim, eps)      # model.py 的 Kimi3Block
        h = attn_norm(x)                   # 归一化后再进注意力/FFN
    也用于低秩潜在空间：MLA 的 ``RMSNorm(q_lora_rank)``、MoE 的 ``RMSNorm(latent_dim)``。
    """

    def __init__(self, dim: int, eps: float = 1e-6, *, elementwise_affine: bool = True):
        super().__init__()
        self.dim = dim
        self.eps = eps
        # 逐通道可学习缩放，初始全 1（归一化开始时是纯单位缩放，不扰动初始化）。
        self.weight = nn.Parameter(torch.ones(dim)) if elementwise_affine else None

    def forward(self, x: Tensor) -> Tensor:
        return rms_norm(x, self.weight, self.eps)

    def extra_repr(self) -> str:
        return f"dim={self.dim}, eps={self.eps}"

class HeadRMSNorm(nn.Module):
    """head-wise RMSNorm：按每个 head 独立归一化。

    K3 用它作为 KDA 输出门的第一步（kda.py:128 ``self.o_norm``）。输入形状为
    ``(batch, seq, num_heads, head_dim)``，归一化只作用于最后一维 ``head_dim``，
    **不跨 head** —— 这正是 "head-wise" 的含义。

    为什么按 head 归一化而非整体归一化：KDA 每个 head 维护自己的递推状态矩阵 S，
    head 之间的输出尺度可能差异很大（不同 head 关注不同频率/内容）。若整体
    归一化，大尺度的 head 会主导归一化因子，把小 head 的信号压没；按 head 归一化
    则保证每个 head 的输出都被重缩放到可比尺度，再统一被后面的全秩门调制。

    注意：本实现的 ``weight`` 形状是 ``(head_dim,)`` 且在所有 head 间共享
    （``(num_heads, head_dim)`` 权重变体见 ``models/llm/core/norms.py`` 的
    ``HeadRMSNorm``，接口不同，未在本模块重复导出）。
    """

    def __init__(self, head_dim: int, eps: float = 1e-6):
        super().__init__()
        self.head_dim = head_dim
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(head_dim))

    def forward(self, x: Tensor) -> Tensor:
        # x 形状 (..., head_dim)；rms_norm 只对最后一维归一化，weight 广播到 head 轴。
        return rms_norm(x, self.weight, self.eps)

    def extra_repr(self) -> str:
        return f"head_dim={self.head_dim}, eps={self.eps}" 

所以真正的问题是：调用时传的 x 最后一维是什么。

• RMSNorm(dim) 用在 (b, t, dim) 的隐藏状态上 → 最后一维是 dim（整个隐藏维）。
• HeadRMSNorm(head_dim) 用在 (b, t, num_heads, head_dim_v) 的输出上 → 最后一维是 head_dim_v。

"head-wise"的语义不是靠类本身实现的，而是靠调用方把张量 reshape 成 (..., num_heads, head_dim) 形状，然后 mean(-1) 自然只对每个 head 内部求平均。所以 HeadRMSNorm 本质上就是"RMSNorm，但专门作用于 per-head 张量" —— 命名传达的是使用场景，而非不同的数学。

Q2：与参考代码的异同

与 kimi-k3-pytorch-main 差异：kimi_primitives/headwise_rmsnorm.py 的 HeadwiseRMSNorm 权重形状 (num_heads, head_dim)（per-head affine），本项目为 (head_dim,)（head 间共享）；二者都只按最后一维归一化。

与 SGLang 核对：SGLang 用 FusedRMSNormGated（fla/fused_norm_gate.py）把归一化 + sigmoid 门融合成单 kernel（kimi_linear.py:318，o_norm = FusedRMSNormGated(head_dim, activation="sigmoid")）。本项目 CPU 版是 rms_norm + 独立的 sigmoid(g_proj(x))，语义一致（见 KDA 文档）。

Q3：out = self.o_norm(out.to(x.dtype)).reshape(b, t, -1) 里 out 是 per-head 的吗？从哪开始 per-head？

是 per-head 的，关键在第 252 行：

out = out.transpose(1, 2).reshape(b, t, self.num_heads, self.head_dim_v)  # ← 拆成 per-head
out = self.o_norm(out.to(x.dtype)).reshape(b, t, -1)              # ← per-head 归一化后拼回

per-head 是从第 252 行的 reshape(b, t, num_heads, head_dim_v) 开始的。但在那之前，out 在整个递推过程中一直是 per-head 的，只是排列不同：

完整流转：

递推输出 out: (b, h, t, dv)          ← 从 kda_chunkwise/kda_step 返回（per-head，head 在第 1 维）
   ↓ transpose(1,2)                  → (b, t, h, dv)
   ↓ reshape(b, t, h, dv)            → (b, t, h, dv)  ← 仍是 per-head，只是改成 (t, h) 排列
out.to(x.dtype) → o_norm             → 对最后一维 dv 归一化（每 head 独立）
   ↓ reshape(b, t, -1)               → (b, t, h*dv)  ← 把 head 拼回，不再是 per-head
out * sigmoid(g_proj(x))             → 全秩门
o_proj                               → (b, t, dim)

更早的源头：per-head 的形状其实从投影时就开始了。

• kda.py:180-182：q/k/v 先由宽向量 view(b, t, num_heads, head_dim_k) 拆成 per-head，然后 l2_normalize/silu。
• kda.py:215：transpose(1,2) → (b, h, t, dk)，送入递推。
• 递推 kda_ref.py 全程按 (b, h, ...) 操作。

一句话总结：out 从 kda.py:252 的 reshape 起明确是 (b, t, h, dv) 四维 per-head 张量，o_norm 对它做按 head 归一化；而它的 per-head 本质其实从 kda.py:180 拆 head 就开始了，只是排列在 (t,h) 与 (h,t) 间切换。

4.5 输出门（Eq.6）

输出门 = "读出时的注意力选择"：模型对写入（delta rule）和读出（gate）分别做选择性，这是记忆系统设计的两端自由度，与 Mixture-of-Experts 的"路由即选择"同理。

4.5.1 原理

直觉：KDA 输出是"压缩记忆的读出"，质量参差。输出门让每个 token 自己决定"从记忆里读哪些 channel"。

o_t = W_o( G(x_t) ⊙ RMSNorm_head(o_t') )        (Eq.6)

• o_t' = S_tᵀ q_t：raw 读出。
• 先按 head 做 RMSNorm（对 d_v 维），再乘 sigmoid(W_g x_t) 全秩门，最后 W_o 投影回 d。

相比 Kimi Linear 的低秩门，K3 用全秩门（W_g ∈ R^{d×d}）——成本更高但可控 channel 更多。这是"read 选择性"的表达力增强。

为什么还需要全局 Output Gate

衰减门决定"旧状态留下多少"，但它不负责"本层输出写进下一层多少"。如果没有全局 output gate：某些层可能输出过强；某些层可能输出过弱；不同层的状态贡献难以独立调节；深层残差叠加可能出现稀释或爆炸。全局 Output Gate 是一层输出总量控制器；衰减门是历史记忆的寿命控制器。二者职责不同。

分块摘要本质上是求和，求和不会自动抹掉幅度；AttnRes 的 value 路也不归一化。所以输出门的每个通道幅度，会直接决定该层对深层记忆的写入份额。全秩门能逐通道控制这个份额。

4.5.2 对比基线

4.5.3 反方辩手

• a) 全秩门增加 d² 参数与算力，收益存疑。
• b) 若失败，错在假设 token 级选择 gate 有用。
• c) 胜出：需要 token 差异化读取的任务。

4.5.4 工程实现

项目代码（kda.py 中 forward 末尾）：

out = out.transpose(1, 2).reshape(b, t, self.num_heads, self.head_dim_v)
out = self.o_norm(out.to(x.dtype)).reshape(b, t, -1)
out = out * torch.sigmoid(self.g_proj(x))
out = self.o_proj(out)

4.6 Gated MLA 与 KDA/MLA 互补

4.6.1 MLA 负责什么

MLA 通过低秩潜变量压缩 KV：降低 KV Cache；降低历史读取带宽；保留 token‑addressable 的全局交互；让 Query 仍可以选择不同历史位置。

4.6.2 Gate 负责什么

即便 MLA 找到了某些内容，也不代表这些内容应该全部写入下一层。Gate 可以控制：

• 当前注意力输出的通道幅度；
• 不同输入下 MLA 分支的可信度；
• KDA 与 MLA 输出的相对贡献；
• 深层残差的尺度稳定性。

4.6.3 组合收益

MLA: 我能读到哪些全局内容
Gate: 这些内容要写进去多少

这比单纯扩大 MLA 表达力更直接，因为它同时解决了"读取"和"写入强度"两个问题。

4.6.4 KDA 的职责

KDA 以固定状态承担大多数长序列 mixing：历史状态大小固定；decode 不需要读取所有历史 token；带宽和 KV Cache 成本较低；非交换更新天然保留时序和近因；适合持续状态跟踪。

4.6.5 MLA 的职责

MLA 周期性提供 token‑addressable global interaction：历史 token 仍保留地址；当前 Query 可以重新选择远程 token；适合 needle retrieval；适合多实体、多文档、多跳证据组合；修复纯固定状态的精确检索上限。

4.6.6 为什么不能只用 KDA

固定状态必须把无限历史压成有限矩阵。即使状态更新很聪明，也无法让每个 token 永远拥有独立地址。它像一个容量有限的工作记忆：适合保存趋势、模式、主题和近期状态；不适合保证任意远古 token 都能被精确点名。

MLA 就像偶尔打开原始档案库：成本更高；但能直接找到具体 token；让整个混合系统不至于完全受固定状态瓶颈限制。

这就是：KDA 负责低成本的连续记忆，MLA 负责高价值的精确检索。

0xFF 参考

AI Infra进阶：3万字深入理解线性注意力机制