【轻量化复刻大模型 / KIMI K3】KDA（Kimi Delta Attention）：并行 – 罗西的思考

与标准做法的公式级对比

专家视角

"无位置编码 + 局部卷积 + 递推衰减"是一套完整的位置信息方案：局部用卷积、全局用衰减序。

K3 的混合可以理解为：KDA 保留顺序、近因和 last-write-wins；NoPE MLA 提供顺序中立的集合聚合和全局内容交互。二者互补：任务需要"先后顺序"时，让 KDA 提供偏置（KDA 的顺序关系来自状态演化。即使计算被拆成 Chunk 并行，Chunk 之间仍然要按照状态转移关系组合，因此顺序不会因为并行而消失）；任务需要"多证据集合聚合"时，让 MLA 减少近因偏置。

边界是：NoPE MLA 本身无法单独回答"谁先发生"。顺序信息仍需来自因果 mask、ShortConv、上游 KDA 或其他表示。

NoPE

NoPE 不是解决全部顺序问题：NoPE 只意味着 MLA 不额外加入显式位置旋转。顺序仍可能来自 causal mask、ShortConv、上游 KDA、数据本身的序列结构。所以不能把 NoPE 等同于"模型完全不理解顺序"。

另外，NoPE 不只是"因为 KDA 已经提供位置，所以省掉 RoPE"。它还避免了两套位置系统互相竞争：

如果二者同时强行存在，模型可能同时面对：KDA 产生的内容相关相对距离；RoPE 产生的固定旋转距离。这两种位置归纳偏置可能方向不一致。所以 NoPE 的含义是避免重复或冲突的位置归纳偏置。

短卷积和滑窗注意力的区别

短卷积和滑窗注意力都会聚合附近 Token，看起来都像是在处理局部上下文，但它们的权重来源不同。

• 短卷积使用预先学习好的固定权重。相同的相对位置通常使用同一组卷积参数，计算过程就是固定窗口内的逐通道加权求和。
• 滑窗注意力则会根据当前内容动态计算权重。当前 Query 会分别和窗口内的 Key 做匹配，再根据匹配结果决定每个局部 Value 应该取多少。

因此，两者的差别可以这样概括：

短卷积并不取代线性注意力，而是为线性注意力提供一个便宜的局部信息通道。线性注意力的固定状态负责更远的历史，短卷积则弥补状态压缩对邻近细节建模不足的问题。

4.7.2 对比基线

4.7.3 计算复杂度

• ShortConv：O(T·d·k)（k=4），几乎可忽略。

4.7.4 参数敏感性

• conv_kernel（默认 4）：调大→局部感受野↑但"位置模糊"；调小→移位信息不足。

4.7.5 工程实现

项目代码 short_conv.py 中 ShortConv 类。初始化置恒等（weight[:,0,-1]=1.0）——卷积起点是 no-op，不扰动精心初始化的衰减动力学。decode 时窗口状态由 KDACache 的 conv_* 保存。

class ShortConv(nn.Module):
    """因果 depthwise 短卷积，自带显式状态以支持逐 token 精确 decode。

    本类同时提供两种形态：
    - forward：整段 / chunk 处理（prefill 或 chunkwise），一次算一整段输出，
      并返回"尾部的 kernel_size-1 个输入"作为下一段的状态；
    - step：单 token 递推（自回归 decode），用上一轮状态 + 当前 token 算输出。

    权重存储约定：weight[c, 0, lag] 乘的是 x[t - lag, c]（互相关习惯），
    即 weight[:, 0, -1] 对应当前 token x
    其余置 0，于是 out
    """

    def __init__(self, dim: int, kernel_size: int = 4, *, bias: bool = False):
        super().__init__()
        self.dim = dim
        self.kernel_size = kernel_size
        self.weight = nn.Parameter(torch.zeros(dim, 1, kernel_size))
        self.bias = nn.Parameter(torch.zeros(dim)) if bias else None
        self.reset_parameters()

    def reset_parameters(self) -> None:
        """恒等初始化：卷积初始为 no-op。
        做法是权重清零，再把最后一个 tap（对应当前 token 的系数）置 1。
        这样模型一开始 q/k/v 的局部混合是恒等，不会与 KDA 初始化的
        "近似恒等衰减"互相干扰。
        """
        with torch.no_grad():
            self.weight.zero_()
            self.weight[:, 0, -1] = 1.0
            if self.bias is not None:
                self.bias.zero_()

    def forward(self, x: Tensor, state: Tensor | None = None) -> tuple[Tensor, Tensor]:
        """x 是 (batch, seq, dim)。返回 (卷积后的序列, 尾部窗口状态)。"""
        b, _, d = x.shape
        pad = self.kernel_size - 1
        xt = x.transpose(1, 2)
        left = xt.new_zeros(b, d, pad) if state is None else state
        padded = torch.cat([left, xt], dim=-1) if pad else xt
        out = F.conv1d(padded, self.weight, self.bias, groups=self.dim)
        new_state = padded[..., -pad:] if pad else padded[..., :0]
        return out.transpose(1, 2), new_state

    def step(self, x: Tensor, state: Tensor) -> tuple[Tensor, Tensor]:
        """单 token 更新（自回归 decode）。x 是 (batch, dim)。"""
        window = torch.cat([state, x.unsqueeze(-1)], dim=-1)
        out = (window * self.weight.squeeze(1)).sum(-1)
        if self.bias is not None:
            out = out + self.bias
        return out, window[..., 1:]

与开源参考核对：SGLang 用 mamba 的 causal_conv1d_fn/causal_conv1d_update；语义一致，初始化策略不同（本项目恒等初始化，kimi-k3-pytorch 用 normal_(0, 0.02) 随机初始化）。

4.7.6 讨论

为什么需要卷积？为什么 q、k、v 都需要卷积？

最简单的 Linear Attention 往往把每个 token 独立映射为 q,k,v，然后直接写入状态。这会有两个问题：相邻 token 之间缺少局部交互；单个 token 的 q,k,v 对微小噪声过于敏感。短卷积的作用是先做局部混合：

当前 token + 邻近 token → 更平滑、更有上下文的 q/k/v → 再写入长期状态

它补的是 Linear Attention 的局部感受野，不是把全部历史重新做一遍 Full Attention。

• Q 需要卷积：Query 决定"现在要从状态中读什么"。如果只看当前 token，读取方向容易抖动。Q 卷积让查询带上局部上下文。
• K 需要卷积：Key 决定"这条信息写入状态的哪个方向"。K 过于尖锐或噪声过大，会让相邻 token 写入完全不同的方向，造成状态碎片化。K 卷积让写入地址更稳定。
• V 需要卷积：Value 是被写入的内容。如果 V 完全 token 独立，局部语义容易被拆散。V 卷积让写入内容带有邻域信息，类似先把一句短语整理成一个局部片段再存档。

所以 q/k/v 的卷积不是重复操作：Q 平滑读取；K 平滑地址；V 平滑内容。

为什么还需要 SiLU

卷积本身主要是线性混合。如果后面没有非线性，多个局部模式只能做线性叠加，表达能力有限。

SiLU(x) = x · σ(x)可以理解为一个平滑的"软开关"：x 很小时抑制；x 较大且有用时放行；与硬 ReLU 相比，梯度更平滑。它让 q/k/v 在写入状态前能够进行输入依赖的门控和特征重组。

代价：卷积和 SiLU 会重新引入额外算子，说明 Linear Attention 并不是"完全没有非线性和局部计算"。它省掉的是长历史的全量交互，不是所有计算。

Q、K 为什么加 L2 Norm，V 为什么不需要？

Q/K 的点积决定相似度和路由。若范数变化过大，模型可能用"放大向量"代替真正的方向匹配，造成 logits、衰减和索引不稳定。L2 Norm 让 Q/K 更接近"方向相似度"：

q_hat = q / ||q||_2,   k_hat = k / ||k||_2

V 不负责决定权重，它是被读取的内容。V 的幅度本身可能携带信息，因此不应强行归一化。

Decode 时如何保存卷积状态

在普通 Decode 中，每次通常只处理一个新 Token。为了计算这个 Token 的短卷积输出，需要知道它前面的 conv_kernel-1 个卷积输入。

因此，每个请求只需要保存最近的：

conv_kernel - 1

个“卷积前的 qkv”。

这里特别需要区分两种数据：

• 缓存的是卷积之前的原始 qkv；
• 卷积之后的输出不会再次写入 conv_state。

原因是下一步卷积需要把新的原始 qkv 与历史窗口中的原始 qkv 重新按照卷积核计算，而不是把已经卷积过的结果再次作为输入叠加。

整个缓存池的形状可以写成：

[num_cache_slots, 3 × num_heads × head_dim, conv_kernel - 1]

其中：

• num_cache_slots 是缓存池能够容纳的物理槽位数量；
• 3 对应 q、k、v 三类投影；
• num_heads × head_dim 是每一类向量的总通道数；
• conv_kernel - 1 是每个请求需要保留的历史窗口长度。

每个请求不会直接拥有一个独立的动态数组，而是通过缓存索引找到自己对应的物理槽位。新 Token 到来时，系统使用该请求的槽位读取最近的卷积输入，计算当前卷积输出，再把新的卷积前 qkv 放入窗口，供后续 Decode 使用。

因此，短卷积状态和递归状态 S 的职责不同：

• conv_state 只保存最近几个 Token 的局部原始输入；
• recurrent_state 把更长历史压缩到固定大小的全局记忆矩阵中。

这也解释了为什么推测解码时，conv_state 比递归状态 S 更容易回滚：卷积状态保存的是彼此独立的局部向量窗口，而 S 已经把多个 Token 混合进同一个矩阵。

4.7.7 反方辩手

• 局部卷积提供的"位置"极其有限，长程位置信息完全靠衰减隐式编码。
• 若失败，错在假设衰减累积足以编码任意位置关系。
• 胜出：语义稠密、位置精度要求不高的任务。
• 稳健替代：MLA 层加 RoPE 的混合变体（论文也评估过 Kimi Linear-RoPE 基线）。
• 受伤方：精确坐标/时间戳引用任务。

4.8 单步递推

我们首先看看代码中的两个单步递推实现。

4.8.1 KDA 版本

KDA 的门控在推理/串行递推与训练/chunk 并行两种形态下采用两种实现，分别对应不同论文阶段：

• v1（Kimi Linear 原版）：门控用无界 softplus：

  g = -exp(A_log) * softplus(gate_logit + dt_bias)
  

  g 无下界（可以非常负），alpha = exp(g) 可以任意接近 0。这是 Kimi Linear 技术报告中的原始形式，也是 SGLang 的默认路径。

• v2（K3 有界版）：门控用有界 sigmoid：

  g = g_min * sigmoid(exp(A_log) * (gate_logit + dt_bias))
  

  g 严格在 (g_min, 0) 之间，因此 alpha 在 (exp(g_min), 1) 范围内。这是 K3 技术报告中的最终形态，也是本项目 kimi_k3 模块实际使用的版本。

两者都已在项目中实现，由 kimi_linear_step(..., version="v1"/"v2") 切换。kimi_linear_step 默认（也是 SGLang 默认）是 KDA v1；v2 是 K3 最终采用的"安全下界"版本。

4.8.2 kimi_linear_step vs KimiDeltaAttention

我们来解释下 kimi_linear_step 和 KimiDeltaAttention 有何差异和联系。

# ============================================================
# 6. Kimi Linear (Moonshot AI, 2025)
#    Core idea: KDA v1 with softplus gating
#    S_t = Diag(exp(g)) * S_{t-1} + beta*k*(v - S^T*k)^T
#    g = -exp(A_log) * softplus(z) — per-head-per-channel, unbounded below
#
#    【SGLang 对齐说明】
#    SGLang 递推核 fused_sigmoid_gating_recurrent.py
#    fused_sigmoid_gating_delta_rule_update (L274-288) 与 fused_recurrent.py
#    IS_KDA=True 分支 (L94-107) 用完全相同的时序：
#        a) b_h *= exp(g_k)              # 先衰减（per-K 通道门控）
#        b) b_v -= sum(b_h * b_k)        # 用衰减后状态算误差
#        c) b_v *= b_beta
#        d) b_h += b_v * b_k
#        e) b_o = sum(b_h * q_scaled)    # q 已乘 scale
#    门控激活形式（SGLang kda.py L1001-1006）：
#        v1: g = -exp(A_log) * softplus(raw + dt_bias)          （默认）
#        v2: g = lower_bound * sigmoid(exp(A_log) * (raw + dt_bias)) （安全下界）
#    本函数因此有两处修改：
#       [1] 门控加 dt_bias 项（原版漏掉该偏置，SGLang 两版都会加）。
#       [2] 读出缩放 scale = Dk**-0.5（原版漏乘，数值验证显示输出差 1/scale ≈ 2.83 倍）。
#       [3] 衰减时序先于误差检索（与 gated_delta_net_step 同理）。
# ============================================================
def kimi_linear_step(S: torch.Tensor, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                     gate_logit: torch.Tensor, beta_logit: torch.Tensor,
                     A_log: torch.Tensor, g_min: float = -5.0,
                     dt_bias: Optional[torch.Tensor] = None,
                     version: str = "v1") -> tuple[torch.Tensor, torch.Tensor]:
    """Single step of Kimi Linear / KDA.

    gate_logit (B, H, Dk) — raw gate input
    A_log (H,) — learnable log-scale per head
    dt_bias (H, Dk) or (Dk,) — per-key bias added to gate input (SGLang 有)
    g_min (float) — v2 门控的下界（衰减因子上界 alpha_max = exp(g_min)）
    version — "v1"：softplus 门控，g 无下界（无界衰减，SGLang 默认）
              "v2"：sigmoid 门控，g ∈ (g_min, 0)，保证 chunkwise 数值安全（K3 实际使用）

    KDA v1: g = -exp(A_log) * softplus(gate_logit + dt_bias)   (unbounded below)
    KDA v2: g = g_min * sigmoid(exp(A_log) * (gate_logit + dt_bias))  (bounded)
    """
    q = l2_normalize(q)
    k = l2_normalize(k)

    # 【对齐 SGLang】支持 dt_bias（kda.py L1003 / fused_recurrent.py L479, L486），
    # 缺失时退化为 0 偏置，不影响原调用方（测试未传 dt_bias 也通过）。
    z = gate_logit
    if dt_bias is not None:
        z = z + dt_bias

    if version == "v1":
        # v1: softplus 门控，无下界。softplus 无上界 → g 可非常负 → alpha = exp(g) 可任意接近 0
        g = -torch.exp(A_log).view(1, -1, 1) * F.softplus(z)
    elif version == "v2":
        # v2: sigmoid 门控，有下界。sigmoid ∈ (0,1) → g ∈ (g_min, 0)，
        # 等价 alpha ∈ (exp(g_min), 1)，给衰减因子一个硬上界，防 chunkwise 累积溢出。
        g = g_min * torch.sigmoid(torch.exp(A_log).view(1, -1, 1) * z)
    else:
        raise ValueError(f"unknown KDA version: {version!r} (expected 'v1' or 'v2')")

    alpha = torch.exp(g)  # decay factor in (0, 1)
    beta = torch.sigmoid(beta_logit).unsqueeze(-1)

    # 【对齐 SGLang】读出缩放（fused_recurrent.py L88），标准 softmax attention 习惯
    scale = q.shape[-1] ** -0.5
    q = q * scale

    # 【对齐 SGLang】先衰减，再用衰减后状态算误差（fused_recurrent.py L95-99）
    S = alpha.unsqueeze(-1) * S
    retrieved = torch.einsum('bhd,bhdv->bhv', k, S)
    error = v - retrieved
    S_new = S + beta.unsqueeze(-1) * torch.einsum('bhd,bhv->bhdv', k, error)
    out = torch.einsum('bhd,bhdv->bhv', q, S_new)
    return out, S_new, retrieved, error

二者是不同抽象层次的关系：kimi_linear_step 是最小的递推算子函数，KimiDeltaAttention 是完整的神经网络层，内部调用更底层的 kda_step/kda_chunkwise。

联系（数学核心一致）

二者的递推核心都是 Eq.1 的 delta rule + channel-wise decay：先衰减、再检索、再修正、最后读出。状态更新的数学形式完全相同。

对拍验证：将 kimi_linear_step(version="v2") 与 kda_ref.py 中的 kda_step 在相同输入下对比——状态矩阵 S 的 max diff 为 1.11e-16（机器精度），说明二者的状态更新完全等价。

差异

一句话总结：kimi_linear_step 是 KDA 的"数学内核"教学演示，KimiDeltaAttention 是 KDA 的"工程完整版"——把投影、卷积、归一化、输出门、缓存管理全部包装好的实际使用层。数学核心相同，外层包裹了完整网络层所需的一切基础设施。KimiDeltaAttention 是完整的神经网络层，内部调用更底层的 kda_step/kda_chunkwise。

4.8.3 kda_step（decode 形态）

kda_step 是 v2 的"单步 decode 形态"，即只做一次递推（不是整段序列）。

v2 和 kda_step 的区别只是工程封装：

1. kimi_linear_step 输入 gate_logit，自己算 g；kda_step 输入预计算好的 g（上层 log_decay 算）
2. kimi_linear_step 对 q 乘了 1/√Dk 的 scale（对齐 SGLang fused_recurrent）；kda_step 不加（尺度由 HeadRMSNorm + 输出门吸收）
3. 状态更新的数学本质完全一样

原理

decode 时每步只有 1 个 token，直接走 Eq.1 的单步更新（无 chunk 结构开销）。且因为首个 token 无先前状态，走 chunkwise（初始状态 0）保证与 prefill 完全一致。

计算复杂度为：单 tokenO(d_k·d_v)（一次 matmul 更新 + 一次读出）。

对比基线

工程实现

项目代码（kda_ref.py 中 kda_step）：

def kda_step(q, k, v, g, beta, state):
    """单 token 更新，decode 阶段（一次只看一个 token）专用。

    q, k, g 形状 (B, H, Dk)，v 形状 (B, H, Dv)，beta 形状 (B, H)，
    state 形状 (B, H, Dk, Dv)。逻辑与 kda_recurrent 的单步完全一致，
    只是把时间维去掉、直接在 batch 上张量化地执行。
    """
    state = g.exp().unsqueeze(-1) * state              # 逐 channel 衰减旧状态
    sk = torch.einsum("bhd,bhdv->bhv", k, state)       # S^T k：检索记忆
    delta = beta.unsqueeze(-1) * (v - sk)              # 写入方向 beta (v - S^T k)
    state = state + k.unsqueeze(-1) * delta.unsqueeze(-2)  # 外积更新 S
    out = torch.einsum("bhd,bhdv->bhv", q, state)      # 读出 o = S^T q
    return out, state

调用处在 kda.py 的 forward 中（t==1 and state_in is not None 才走单步，保证首个 decode token 也能走 chunkwise 初始化）。

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)

与开源参考核对：SGLang fused_sigmoid_gating_delta_rule_update 即此单步 kernel；本项目 CPU 版等价。

4.9 chunkwise 并行（Eq.3-4，DPLR）

chunkwise 揭示了递推模型的普遍范式：能分块 = 能并行。只要转移矩阵是"对角+低秩"（DPLR 类），就存在 UT 变换把块内串行展开为矩阵运算。Mamba-2 的 SSD、GLA 都属此族。这是"RNN 复兴"的工程基石。

4.9.1 原理

直觉：递推天然串行（每步依赖上一步），但训练要并行。chunkwise 把序列切成块：块内用矩阵运算并行算出所有贡献，块间才串行传递状态。

专家视角（块长 C，进入块的状态 S_0）：

O = (Γ ⊙ Q) S_0 + A Ṽ (Eq.4a 输出)
(I + M) Ṽ = Diag(β)(V − K̄ S_0) (Eq.4b UT 变换)

其中 Γ 是块内累积衰减，Q̄=ΓQ、K̄=ΓK、K̂=K/Γ，A = tril[(Q̄)(K̂)ᵀ] 是块内"注意力"矩阵，M = tril(Diag(β)K̄K̂ᵀ, −1)。(I+M) 是单位下三角 → 可精确三角求解（UT 变换）。

这是 DPLR 特例的分解：对角部分（Γ）逐元素累积，低秩部分（K̄K̂ᵀ 的 rank-1 累积）用稠密 matmul。

关键：因为衰减有下界，所有 tile（含对角）都可用 Tensor-Core 稠密乘。

4.9.2 chunkwise 形式的推导

单 Token 状态更新

设状态采用项目代码和技术报告的布局：

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

其中 k_t\in\mathbb R^{d_k}、v_t\in\mathbb R^{d_v} 是列向量，\alpha_t 和
\beta_t 是当前 Token 的门控量。论文 Eq. 1 的状态更新为：

\[S_t
=\left(I-\beta_t k_t k_t^\top\right)
\operatorname{Diag}(\alpha_t)S_{t-1}
+\beta_t k_t v_t^\top
\]

令：

\[M_t=
\left(I-\beta_t k_t k_t^\top\right)
\operatorname{Diag}(\alpha_t)
\]

则：

\[S_t=M_tS_{t-1}+\beta_t k_t v_t^\top
\]

这里：

• M_t 是当前 Token 对旧状态施加的衰减与 delta-rule 修正；
• \beta_t k_t v_t^\top 是当前 Token 的写入项；
• M_t 和写入项由当前 Token 的 k_t、v_t、\alpha_t、\beta_t 决定；
• 它们不需要等待 S_{t-1} 计算完成后才准备。

因此，虽然 KDA 的门控和投影可以包含非线性，但在当前 Token 的系数确定后，从 S_{t-1} 到 S_t 的状态更新是一个仿射变换。

Chunk 级状态转移

设第 i+1 个 Chunk 包含一段连续 Token。将这一段中的单 Token 仿射更新按顺序组合后，整个 Chunk 仍然可以表示为一个仿射变换：

\[\boxed{S_{[i+1]}=M_{[i+1]}S_{[i]}+\widetilde{S}_{[i+1]}}
\]

这条公式是 Chunkwise 实现的核心，也是论文中采用的状态传递形式。其中：

• S_{[i]}：第 i 个 Chunk 的输入状态；
• S_{[i+1]}：第 i+1 个 Chunk 的输出状态；
• M_{[i+1]}：第 i+1 个 Chunk 的累积状态转移矩阵；
• \widetilde{S}_{[i+1]}：第 i+1 个 Chunk 在输入状态取零时，
  由该 Chunk 自己生成的状态；
• M_{[i+1]}S_{[i]}：输入状态经过该 Chunk 后对输出产生的影响；
• \widetilde{S}_{[i+1]}：该 Chunk 内部 Token 独立产生的增量。

这里的下标 [i+1] 表示 Chunk 的编号，不是单 Token 的时间下标。

为避免与单 Token 公式混淆，本文统一使用：

S_t = M_t S_{t-1} + beta_t k_t v_t^T

表示单 Token 更新，使用：

S_[i+1] = M_[i+1] S_[i] + Stilde_[i+1]

表示 Chunk 之间的状态传递。论文公式的严格数学写法为：

\[S_{[i+1]} = M_{[i+1]} S_{[i]} + \widetilde{S}_{[i+1]}
\]

M 与 \widetilde S 的含义

考虑一个 Chunk 内部依次包含多个 Token。每个 Token 都有自己的局部仿射更新：

\[S_t=M_tS_{t-1}+\beta_t k_t v_t^\top
\]

将这些更新按 Token 顺序复合，可以得到：

\[S_{\mathrm{out}}
=M_{\mathrm{chunk}}S_{\mathrm{in}}
+\widetilde{S}_{\mathrm{chunk}}
\]

累积转移矩阵 M

M_{\mathrm{chunk}} 描述 Chunk 对输入状态的整体作用。它可以理解为：

1. 将 Chunk 内所有 Token 的状态转移矩阵按顺序复合；
2. 令输入状态为单位基底；
3. 观察该 Chunk 对输入状态的最终变换。

直观地说，M 保留了“输入状态如何穿过整个 Chunk”的信息。

独立状态 \widetilde S

\widetilde{S}_{\mathrm{chunk}} 表示输入状态为零时，该 Chunk 通过自身 Token 生成的状态：

\[\widetilde{S}_{\mathrm{chunk}}
=
S_{\mathrm{out}}\big|_{S_{\mathrm{in}}=0}
\]

它不依赖前一个 Chunk 的输入状态，因此可以在 Chunk 内独立计算。于是，一个 Chunk 不必直接传递整段 Token 的中间递归状态，只需传递固定大小的(M, \widetilde S) 表示。

Chunk 间状态传递

假设序列被切成多个 Chunk：

Chunk 0, Chunk 1, Chunk 2, ...

第 i+1 个 Chunk 的状态更新为：

\[S_{[i+1]}
=M_{[i+1]}S_{[i]}
+\widetilde{S}_{[i+1]}
\]

因此，执行流程是：

1. 每个 Chunk 独立计算自己的 M_{[i+1]}；
2. 每个 Chunk 独立计算自己的 \widetilde{S}_{[i+1]}；
3. 按 Chunk 顺序取得 S_{[i]}；
4. 使用论文公式计算 S_{[i+1]}；
5. 将 S_{[i+1]} 传给下一个 Chunk。

Chunk 之间仍然存在状态依赖，但依赖的数据不再是整段 Token 的中间结果，而是固定大小的 Chunk 状态。

如果要把前 n 个 Chunk 合并成一个更大的仿射变换，也可以递归组合：

\[\begin{aligned}
S_{[i+1]} &= M_{[i+1]}S_{[i]}+\widetilde S_{[i+1]},\\
S_{[i+2]} &= M_{[i+2]}S_{[i+1]}+\widetilde S_{[i+2]}.
\end{aligned}
\]

代入得到：

\[S_{[i+2]}
=
M_{[i+2]}M_{[i+1]}S_{[i]}
+
M_{[i+2]}\widetilde S_{[i+1]}
+
\widetilde S_{[i+2]}
\]

因此，两个连续 Chunk 仍然可以合并成：

\[S_{[i+2]}
=M_{[i:i+2]}S_{[i]}
+\widetilde S_{[i:i+2]}
\]

其中：

\[M_{[i:i+2]}=M_{[i+2]}M_{[i+1]}
\]

\[\widetilde S_{[i:i+2]}
=M_{[i+2]}\widetilde S_{[i+1]}
+\widetilde S_{[i+2]}
\]

这说明 Chunk 的状态转移具有可组合性。

4.9.3 对比基线

4.9.4 计算复杂度

chunkwise 是训练/prefill 的并行桥梁；decode 切到 kda_step。

• 块内 token 并行阶段主要包含 A/M 的 C×C 矩阵构造、三角求解和输出矩阵乘；仅计算 A/M 时约为 O(C²·d_k) per block，完整代价还依赖 d_v。因此不能只用一个 d_k 项概括全部 chunkwise 计算。
• 块间串行递推：O(n_chunks · d_k·d_v)。
• 代码通过单位下三角求解避免显式求逆；实际复杂度还取决于 C、d_k、d_v
  和所使用的 kernel，不能仅凭本节简化式断言相对所有 DPLR 实现都“显著减计算”。

4.9.5 参数敏感性

• chunk_size（默认 64）：越大块内并行越充分但内存越大；不能超过 safe_chunk_size。
• tile_size（默认 16）：数值安全基本块。

4.9.5 工程实现

伪代码（简化）

for each chunk:
    log_gamma = cumsum(g) # 块内累积衰减
    gamma, gamma_inv = exp(log_gamma), exp(-log_gamma)
    qbar = q*gamma; khat = k*gamma_inv; kbar = k*gamma
    A = tril(qbar @ khatᵀ) # 块内 attention
    M = tril(beta * (kbar @ khatᵀ), -1)
    rhs = cat([beta*v, beta*kbar], -1)
    [U,W] = solve_triangular(M, rhs) # UT 变换（精确）
    Ṽ = U - W @ S0 # 与入块状态解耦
    O = qbar @ S0 + A @ Ṽ # 读出
    S_C = gamma_last*S0 + k_tail @ Ṽ # 出块状态

真实代码

项目真实代码（kda_ref.py 中 kda_chunkwise），核心流程：

• 累积衰减与倒数
• qbar/khat/kbar 三把重缩放 key/query
• M 与 A 矩阵（严格下三角 / 因果下三角）
• UT 变换（torch.linalg.solve_triangular，unitriangular=True）
• 块间串行递推

def kda_chunkwise(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    g: Tensor,
    beta: Tensor,
    *,
    chunk_size: int = 16,
    initial_state: Tensor | None = None,
    reset_mask: Tensor | None = None,
    dtype: torch.dtype | None = None,
    return_chunk_states: bool = False,
) -> tuple[Tensor, Tensor] | tuple[Tensor, Tensor, list[Tensor]]:
    """Eqs. 3-4 的 chunkwise 并行形式。

    与 :func:`kda_recurrent` 数学等价；区别只在调度：块内 token 之间并行（一次矩阵乘
    解决），块与块之间串行传递状态。``chunk_size`` 不得超过工作 dtype 下的
    :func:`safe_chunk_size`（理由见模块 docstring 的"数值范围"段）。

    ``return_chunk_states`` 额外返回进入每个 chunk 时的状态快照——KDA 上下文并行
    （KCP）正是用它来校验自己的前缀扫描（prefix scan）是否正确。
    """
    # 工作精度：默认 fp64（输入本身是 fp64 时），否则一律 fp32 做中间计算。
    dtype = dtype or (torch.float64 if q.dtype == torch.float64 else torch.float32)
    qf, kf, vf, gf, bf = _as_float(q, k, v, g, beta, dtype=dtype)
    b, h, t, dk = qf.shape
    dv = vf.shape[-1]
    device = q.device

    state = (
        torch.zeros(b, h, dk, dv, dtype=dtype, device=device)
        if initial_state is None
        else initial_state.to(dtype).clone()
    )
    # 向上取整到整块数：“-(-t // c)” 等价于 ceil(t / c)。
    n_chunks = -(-t // chunk_size)
    pad = n_chunks * chunk_size - t
    if pad:
        # padding 的参数特意取 beta=0（不写状态）、g=0（无衰减），后期不携带任何信息，
        # 因此对递推状态是严格的 no-op——剥离 padding 后结果与不 padding 完全一致。
        qf = F.pad(qf, (0, 0, 0, pad)) # (b,h,t+pad,dk)：在时间维尾部补 0
        kf = F.pad(kf, (0, 0, 0, pad))
        vf = F.pad(vf, (0, 0, 0, pad))
        gf = F.pad(gf, (0, 0, 0, pad))
        bf = F.pad(bf, (0, pad)) # beta 是 (b,h,t)，只补时间维
        if reset_mask is not None:
            reset_mask = F.pad(reset_mask, (0, pad), value=False) # 补的 padding 不是新文档起点

    c = chunk_size
    # 重排成 (batch, head, chunk 数, chunk 内位置, 特征)，让"块内位置"变成一个可矩阵化的维。
    q_c = qf.view(b, h, n_chunks, c, dk)
    k_c = kf.view(b, h, n_chunks, c, dk)
    v_c = vf.view(b, h, n_chunks, c, dv)
    g_c = gf.view(b, h, n_chunks, c, dk)
    b_c = bf.view(b, h, n_chunks, c)

    # ---- token 并行阶段 ----------------------------------------------------
    # 所有"不依赖传入状态 S_0"的量都能对全部 chunk 一次性算完——这正是 fused kernel
    # 把这段空排成与递推无关的独立调度阶段的原因（论文 §5.1.1）。
    # （下面 gamma/khat/kbar 等符号全部对应模块 docstring 的推导记号。）

    # 每个 chunk 开始起的累积 log-decay（含当前位置）：log gamma^{1->r}。
    log_gamma = g_c.cumsum(dim=3)
    gamma = log_gamma.exp() # gamma = 保留因子连乘，<= 1（因 g<=0）
    gamma_inv = (-log_gamma).exp() # 1/gamma >= 1，上界 exp(-c * g_min)

    # 三把重缩放 key/query，分别为 Eq. 4 的 Gamma*Q、K/Gamma、K*Gamma。
    q_bar = q_c * gamma # Q * Gamma
    k_hat = k_c * gamma_inv # K / Gamma
    k_bar = k_c * gamma # K * Gamma

    # chunk 内的文档分段：不同文档的位置绝不能互相读；只有第一个文档段能读携带状态。
    # seg 取值 = "相对 chunk 起点的文档起始计数"（_segment_ids 做的事 + 减去块首偏置）。
    seg = _segment_ids(reset_mask, b, n_chunks * c, device).view(b, n_chunks, c)
    seg = seg - seg[:, :, :1] # 把每个 chunk 的第一个位置计成 0
    if reset_mask is not None:
        # 若 chunk 首 token 自己就是新文档起点，seg 会整体 -1 成负值；这里 +1 把它
        # 抬回 0、该 chunk 从"新文档"开始（不允许读携带状态，由 carry_ok 处理）。
        seg = seg + reset_mask.view(b, n_chunks, c)[:, :, :1].long()
    same_seg = (seg[:, :, :, None] == seg[:, :, None, :]).unsqueeze(1) # (b,1,nc,c,c) 块内同文档对
    carry_ok = (seg == 0).to(dtype)[:, None, :, :, None] # (b,1,nc,c,1) 只有文档段 0 能接续旧状态
    last_seg = seg[:, :, -1:] # (b,nc,1) 每个 chunk 的末位置文档号
    in_last = (seg == last_seg).to(dtype)[:, None, :, :, None] # (b,1,nc,c,1) 标记"属于末段"的位置
    keep = (last_seg == 0).to(dtype)[:, None, :, :, None] # (b,1,nc,1,1) 末位置仍在文档 0 才传递状态

    # 因果 / 严格下三角掩码，用于构造 A 与 M。
    causal = torch.ones(c, c, dtype=torch.bool, device=device).tril() # 下三角（含对角线）
    strict = causal & ~torch.eye(c, dtype=torch.bool, device=device) # 严格下三角（不含对角线）

    # M = tril(Diag(b) Kbar Khat^T, -1)，用 strict & same_seg 同时遮挡"对角线"与"跨文档"：
    # (k_bar @ k_hat^T) 是 (b,h,nc,c,c) 的块内点积矩阵，先乘 beta 再按掩码过滤。
    m = (b_c.unsqueeze(-1) * (k_bar @ k_hat.transpose(-1, -2))) * (strict & same_seg)
    # A = Tril[(Q*Gamma)(K/Gamma)^T]，因果 & 同文档即可（允许对角线，对应 o_i 能读到自身贡献）。
    a = (q_bar @ k_hat.transpose(-1, -2)) * (causal & same_seg)

    # UT 变换，两套右端项合并成一次三角求解：
    # U = (I+M)^-1 Diag(beta) V ， W = (I+M)^-1 Diag(beta) Kbar
    # 于是 Vt = U - W S_0，其中 S_0 由逐 chunk 的递推在串行阶段提供。
    # I + M 是单位下三角，所以 solve_triangular 的求解是精确且 O(c^2) 廉价的。
    # 注意第二列右端项乘了 carry_ok：只有文档段 0 的位置才携带 Kbar 去算 W（供接续旧状态）。
    rhs = torch.cat(
        [b_c.unsqueeze(-1) * v_c, b_c.unsqueeze(-1) * (carry_ok * k_bar)], dim=-1
    )
    uw = torch.linalg.solve_triangular(m, rhs, upper=False, unitriangular=True)
    u_mat, w_mat = uw[..., :dv], uw[..., dv:] # 拆回 U 与 W，供下面的逐 chunk 递推用

    # 块末状态更新要用的三个量：
    log_tail = log_gamma[:, :, :, -1:] - log_gamma # log gamma^{r+1->C} <= 0（从 r 到块末的衰减）
    k_tail = (k_c * log_tail.exp() * in_last).transpose(-1, -2) # (b,h,nc,dk,c) 加权后的 K^T
    gamma_last = gamma[:, :, :, -1].unsqueeze(-1) # (b,h,nc,dk,1) 整块累积衰减 gamma^{1->C}

    # ---- head 并行递推：唯一的串行部分 ---------------------------------------
    outputs = []
    chunk_states: list[Tensor] = []
    for i in range(n_chunks):
        if return_chunk_states:
            chunk_states.append(state.clone()) # 记录进入第 i 块前的状态快照
        v_tilde = u_mat[:, :, i] - w_mat[:, :, i] @ state # Vt = U - W S_0
        outputs.append(
            carry_ok[:, :, i] * (q_bar[:, :, i] @ state) + a[:, :, i] @ v_tilde
        ) # O = (Gamma*Q) S_0 + A Vt：携带项只对文档段 0 生效，其余位置只有块内注意力
        # 块末状态：S_C = Diag(gamma^{1->C}) S_0 + sum_r (gamma^{r+1->C} k_r) Vt_r^T
        # keep 只在该 chunk 末位置仍属于文档 0 时才传递状态，避免跨文档泄漏。
        state = keep[:, :, i] * (gamma_last[:, :, i] * state) + k_tail[:, :, i] @ v_tilde

    out = torch.cat(outputs, dim=2) # 沿 chunk 维拼回 (b, h, n_chunks*c, dv)
    if pad:
        out = out[:, :, :t] # 剥掉 padding 恢复原始序列长度
    if return_chunk_states:
        return out, state, chunk_states
    return out, state

与开源参考核对：与 kimi3-master 的 kda_ref.py 的 chunkwise 语义一致；与 SGLang chunk_kda（FLA kernel）逻辑逐位对齐（CPU 版用 PyTorch 实现同一数学）。

图例

下面是 kda_chunkwise（kda_ref.py:225）的 ASCII 流程图与数据流图。

流程图（控制流）

数据流图（张量形状演变）

关键点：

• 阶段 1（全并行）：所有不依赖携带状态 S 的量（A、M、U、W、掩码、衰减）一次性算完。
• 阶段 2（串行）：仅 n_chunks 步递推，每步 Vt = U − W·S、O = ΓQ·S + A·Vt、更新 S。
• carry_ok/keep/same_seg 三个掩码保证跨文档不泄漏（reset_mask 语义）。
• padding 取 β=0, g=0，对递推严格 no-op。

4.9.6 反方辩手

• chunkwise 的 1/γ（K/Γ 重缩放）在无界衰减下溢出——这就是必须改有界衰减的原因。
• 若失败，错在假设三角求解与块间递推的数值稳定性。
• 胜出：硬件充裕、需训练并行。
• 稳健替代：纯递推（慢但简单）。
• 受伤方：精度敏感的 tiny 模型。

4.9.7 讨论

Q1：为什么 chunk 还要细分成 16 token tile

一个大 chunk 内 token 越多，相对衰减跨度越大。直接对整个 chunk 展开，数值范围和中间矩阵都可能过大。细分为 16 token tile 后：每个 tile 内的衰减范围较小；tile 之间可以用规则矩阵乘法；更容易装入 Tensor Core 的 tile 形状；中间结果更容易保持在可控范围。

Q2：为什么 Tril 要保留对角线

chunkwise 计算通常会构造一个 token‑pair 关系矩阵：第 t 个 token 读取哪些更早 token 的写入。如果只保留下三角但去掉对角线，表示的是：当前 token 只能看到过去写入的内容。但 KDA 的实际定义是：当前 token 先把自己的内容写进去，然后读取"包含自己写入之后"的状态。所以 token 自己对自己的作用必须保留，对应矩阵对角线必须是 1 或保留该项。

通俗说：不保留对角线，当前员工看不到自己刚提交的修改；保留对角线，当前员工的修改会立刻进入本次结果。

Q3：为什么非对角 tile 容易 GEMM，而对角 tile 较难

非对角 tile 表示：一整块 token 对另一整块 token 的影响。它们之间没有"当前 token 不能看未来 token"的复杂边界，可以直接组成稠密矩阵乘。

对角 tile 则包含同一小块内部的因果关系：token 0 不能看 token 1；token 1 可以看 token 0；token 1 还要看到自己的写入；每个位置的衰减长度不同。因此对角 tile 必须显式计算 position‑pair，难以完全当作普通稠密 GEMM。

Q4：每个 token 都是一个仿射变换

把 KDA 的单步更新写成：S_t = M_t S_{t-1} + b_t，其中 M_t 是旧状态经过多少衰减和擦除，b_t 是当前 token 写入的新内容。两个 token 连起来：

S_2 = M_2(M_1 S_0 + b_1) + b_2
S_2 = (M_2 M_1) S_0 + (M_2 b_1 + b_2)

所以两步可以合并成一个新的仿射变换：(M_2, b_2) ∘ (M_1, b_1) = (M_2 M_1, M_2 b_1 + b_2)。

Q5：稳定性主要在 M_t，表达力主要在 b_t

S_t = M_t S_{t-1} + b_t

• M_t 决定旧状态如何传播、衰减和变形，关系到爆炸、消失和数值稳定；
• b_t 决定新信息写入哪些方向、写入多少，关系到能表达多少内容。

因此可以粗略说：M_t 管"旧记忆会不会坏掉"，b_t 管"新知识能不能写进去"。

Q6：交叉项是表达力和并行难度的根源

两段合并时：b_C = M_2 b_1 + b_2。前一段写入的 b_1 必须接受后一段 M_2 的继续作用。这让 KDA 比纯加法更能表达"后续 token 修改前面 token 的影响"，但也使段间合并必须保持顺序，不能直接无序求和。

Q7：单位元的意义

仿射变换的单位元是：(I, 0)。它表示状态保持不变。因此空 token 序列可以视为一个合法的空 chunk：不改变状态；仍可参与 chunk 合并；边界不需要额外 if 分支；分段算法更规整。

Q8：结合律为什么重要

如果一段 token 的净效果可以浓缩成一个 (M_C, b_C)，就可以：先算 chunk A；先算 chunk B；再把 A、B 的结果合并。而且 ((A∘B)∘C) = (A∘(B∘C))，也就是说，加括号的方式不影响最终结果。这就是结合律。

它允许：把长序列切成多个 chunk；chunk 内部并行计算；用树形归约或分层归约合并；用 Tensor Core 处理块矩阵；形成 FlashKDA、KCP 一类的并行策略。

Q9：为什么不满足交换律

交换两个 token：(M_2, b_2)∘(M_1, b_1) ≠ (M_1, b_1)∘(M_2, b_2)。原因就在交叉项：M_2 b_1 + b_2。前一个 token 写入的 b_1，必须先被后一个 token 的 M_2 衰减或擦除，再和 b_2 相加。如果交换顺序，交叉项就变成 M_1 b_2 + b_1，一般不会相等。

Q10：非交换性为什么正好编码时序

语言中的顺序本来就有意义。KDA 的"后写入继续作用于先写入"的机制，使得新信息可以覆盖旧信息；越近的信息通常影响更大；中间发生过什么会影响最终状态；顺序信息保留在状态转移乘积中。这就是 KDA 的近因偏置和时序表达能力。

Q11：为什么 (I+M) 容易求解？ → I+M 是单位下三角矩阵，因此可以用forward substitution 直接求解，不需要显式矩阵求逆。这里的“精确”是代数结构上的精确解法；实际浮点误差仍受 khat 的放大、dtype 和 chunk/tile 大小影响，不能称为“无条件数值稳定”。

Q12：padding 是否影响结果？ → 不会：padding 位置 β=0、g=0，对状态是 no-op。

Q13：跨文档（reset_mask）怎么处理？ → _segment_ids 用 cumsum 标记文档边界，跨文档位置强制不交互。

4.10 KDA 上下文并行（KCP）

本节是对前面一节的扩展。

4.10.1 原理

直觉：

多卡并行时，序列切成 N 段分给 N 卡。对 softmax attention，跨 rank 的 KV 交换量通常随需要访问的上下文长度增长；

KDA 状态不能简单相加，因为递推的转移矩阵不可交换。

解法：每段分解出"自己的贡献"（Stilde）与"传递矩阵"（M），再按 S = M·S + Stilde 前缀扫描重组。

专家视角：

S_{[i+1]} = M_{[i+1]} S_{[i]} + Stilde_{[i+1]}

这里是论文 Eq.17 在“本段处理完全部本地 Token”时的简化写法。论文的更一般形式是在第 i+1 个 rank 内处理到本地位置 t：

\[S^{t}_{[i+1]}
=\widetilde S^{t}_{[i+1]}
+M^{t\leftarrow1}_{[i+1]}S^{T_i}_{[i]}
\]

当 t=T_{i+1} 时，M^{T_{i+1}\leftarrow1}_{[i+1]} 和 \widetilde S^{T_{i+1}}_{[i+1]} 就是该 rank 对外交换的两个局部量，从而得到上面的段末状态公式。

• Stilde_{[i]}：段 [i] 从 S=0 起算的自生成状态（可并行算）。
• M_{[i]}：段 [i] 的累积转移矩阵（v=0、S0=I 下跑完该段）。
• 通信张量只包含 (M,\widetilde S)；在状态维度、dtype 和并行度固定时，张量形状与序列长度无关。

报告 Eq.17 的准确含义是：在 d_k、d_v、batch/head 数以及 world size 固定时，每个 rank 只需交换固定形状的 M 和 \widetilde S；通信张量大小不随本地序列长度增长，但 all-gather 的总流量仍会随 rank 数和实现细节变化。这是 KDA KCP 与需要交换增长中的 KV 块的 softmax CP 的关键区别。

4.10.2 解读

KDA 的状态更新看起来是逐 Token 的递归过程，但它可以被整理成适合Chunkwise 并行的仿射状态转移。

4.10.2.1 Chunkwise 实现的三个阶段

Chunkwise 可以被分成三个阶段。

阶段 2 保留了 Chunk 之间的顺序依赖，但其并行度仍然来自：

• 请求数量；
• Head 数量；
• Value 轴方向上的分块。

阶段 1 和阶段 3 则可以利用序列维度上的并行性。

4.10.2.2 两层并行

宏观并行：Chunk 之间

每个 Chunk 可以先独立算出：

• 对输入状态的累积转移矩阵 M_{[i]}；
• 输入状态为零时的独立状态 \widetilde S_{[i]}。

之后再按照：

\[S_{[i+1]} = M_{[i+1]}S_{[i]}+\widetilde S_{[i+1]}
\]

顺序传递 Chunk 状态。

因此，Chunk 的独立预计算可以并行，只有固定大小的状态传递保留顺序。

微观并行：Chunk 内部

Chunk 内部的残差更新本来是递归的。通过代数整理，可以将其改写成下三角系统和矩阵运算。

文中使用的辅助矩阵形式为：

\[U=(I+A)^{-1}BV
\]

\[W=(I+A)^{-1}BK_{\mathrm{in}}
\]

\[E=U-WS_0^\top
\]

这些式子用于把大量块内递推转成矩阵乘法，使硬件更容易并行执行。其中，S_0 表示当前 Chunk 的输入状态；它由 Chunk 间的论文公式提供，而不是在Chunk 内重新计算得到。需要注意，U/W/E 的具体记号来自 chunkwise 推导，而 Eq.17 的 KCP 只负责段级 (M,\widetilde S) 重组，二者不是同一层次的公式。

4.10.2.3 为什么这种形式适合上下文并行

对于第 i+1 个 Chunk，只要知道：

(M_[i+1], Stilde_[i+1])

以及前一个 Chunk 的输出状态 S_[i]，就可以计算：

\[S_{[i+1]} = M_{[i+1]}S_{[i]}+\widetilde S_{[i+1]}
\]

通信只需交换固定大小的状态表示，而不需要交换该 Chunk 内每个 Token 的全部中间递归结果。

因此，当状态维度和并行度固定时，单个 rank 的 KCP 交换张量形状与序列长度无关。序列变长时，主要增加本地计算量；如果改变 rank 数，all-gather 的总通信量当然会变化。这正是论文 Eq.17 所体现的固定状态通信性质，也是递推模型上下文并行与 softmax attention 上下文并行的重要区别。

4.10.2.4 与注意力类型的对比

Softmax attention 通常需要处理跨 Token 的归一化和全局依赖，序列分片后需要交换更多与序列长度相关的统计量或中间结果。

KDA 的递归状态则可以在 Chunk 级别表示为：

\[S_{[i+1]} = M_{[i+1]}S_{[i]}+\widetilde S_{[i+1]}
\]

每个 Chunk 只需提供一个固定形状的累积转移矩阵和一个固定形状的独立状态。这使得：

• Chunk 内 Token 可以通过矩阵运算并行处理；
• Chunk 之间只需传递固定大小的状态；
• 上下文并行可以将顺序依赖压缩到 Chunk 级别；
• 状态 checkpoint 可以自然放在 Chunk 边界。

关键不在于整个模型是否线性，而在于：

1. 递归状态如何进入下一步更新；
2. 更新系数是否可以在不依赖递归状态的情况下提前准备；
3. 状态更新是否能被写成可组合的线性或仿射形式。

对 KDA 而言，单 Token 级别的仿射形式是：

\[S_t=M_tS_{t-1}+\beta_t k_t v_t^\top
\]

Chunk 级别、也是论文采用的形式是：

\[\boxed{
S_{[i+1]} = M_{[i+1]}S_{[i]}+\widetilde S_{[i+1]}
}
\]

前者解释局部递归如何组合，后者解释 Chunk 如何进行状态传递与上下文并行。

4.10.3 对比基线

4.10.3 反方辩手

• a) KCP 的前缀扫描是串行的（段间依赖），极端多卡下通信重叠收益递减。
• b) 若失败，错在假设固定状态值得交换成本。
• c) 胜出：超长上下文 + 多卡。
• d) 稳健替代：直接全量（序列不长时）。

4.10.4 工程实现

kda_chunkwise(..., return_chunk_states=True) 专门提供段入状态快照——是 KCP 的验证接口。参考实现 kimi3-master/parallel/kcp.py 实现报告 Eq.17 的前缀扫描；本项目parallel/kcp.py 是 CPU/单进程的等价模拟，不是已经接入 torch.distributed 的生产多卡实现。真正的跨 rank 交换需要把本地 (M,\widetilde S) 接入 all-gather，再按文档边界处理前缀扫描。

核心代码：

def local_transition_and_state(k, v, g, beta, *, chunk_size=16, dtype=None):
    # Stilde：段从 S=0 开始，自身激活产生的状态
    _, s_tilde = kda_chunkwise(k, k, v, g, beta, chunk_size=chunk_size, dtype=dtype)
    # M：homogeneous 响应。v=0 后写入项消失，只剩过渡项，给定 S0=I，
    # 最终 state 就等于段内累积过渡矩阵
    eye = torch.eye(dk, dtype=dtype, device=k.device).expand(b, h, dk, dk).contiguous()
    zero_v = torch.zeros_like(k)
    _, m_cum = kda_chunkwise(k, k, zero_v, g, beta, chunk_size=chunk_size,
                            initial_state=eye, dtype=dtype)
    return m_cum, s_tilde

def prefix_scan_states(transitions, zero_states):
    # transitions: (n,B,H,Dk,Dk) 每段过渡矩阵 M；zero_states: (n,B,H,Dk,Dv) 每段 Stilde
    n = transitions.shape[0]
    incoming = torch.zeros_like(zero_states)
    state = torch.zeros_like(zero_states[0])
    for i in range(n):
        incoming[i] = state
        state = transitions[i] @ state + zero_states[i]
    return incoming

def kcp_forward(q, k, v, g, beta, *, chunk_size=16, n_segments=4, dtype=None):
    # 阶段 1：每段局部 (M, Stilde)（分布式下在各自 rank 本地并行执行）
    m_all, s_all = [], []
    for sl in segs:
        m_, st_ = local_transition_and_state(k[:,:,sl], v[:,:,sl], g[:,:,sl], beta[:,:,sl], ...)
        m_all.append(m_); s_all.append(st_)
    Ms = torch.stack(m_all, dim=0)
    Ss = torch.stack(s_all, dim=0)
    incoming = prefix_scan_states(Ms, Ss) # 前缀扫描重组出各段入状态
    # 阶段 2：逐段用 incoming 重新计算真实输出
    outs = []; state = incoming[0]
    for i in range(n_segments):
        out, state = kda_chunkwise(q[:,:,segs[i]], ..., initial_state=incoming[i], ...)
        outs.append(out)
    return torch.cat(outs, dim=2), state

通信张量：每个 rank 每个 KDA 层交换两个定长张量 (M,\widetilde S)；其元素数约为 d_k^2+d_kd_v，字节数还要乘以 batch、head 和 dtype 的 item size。它对序列长度不变，但 all-gather 的总字节数会受 world size 和通信算法影响。

4.10.5 实测数值

直接相加 vs KCP 前缀扫描（64 token, 4 段）：

• 直接相加误差：9.782e-01（几乎完全错误）——证明"跨段不可加"。
• KCP 前缀扫描误差: 1.110e-16（机器精度）——证明 M/Stilde 分解 + 前缀扫描是精确的。
• state rms：0.11986。

通信量对比：原文中的具体字节数必须绑定配置后再引用。按本项目communication_volume(batch, heads, dk, dv, world, itemsize) 的定义，单层基础张量大小为batch * heads * (dk^2 + dk*dv) * itemsize；函数返回值还乘以 world，表示集群总量。
例如 batch=heads=1、dk=dv=16 时，float32 为 2048 bytes，float64 为4096 bytes，不能脱离 dtype 统一写成“恒定 2048 bytes”。softmax 的 1MB/1GB数字也应标明 KV 维度、dtype、rank 数和统计口径，否则只能作为某组实验配置下的示例。

4.10.6 讨论

KCP 是"状态压缩 → 通信恒定"的典范：当每个 worker 只携带固定状态，分布式成本与上下文无关，这是线性注意力在系统层面的杀手锏——也是 K3 报告 §5.1 花大篇幅讲系统协同设计的理由。

Q1：为什么状态不能像 vanilla 线性注意力那样相加？ → 更新里有 (I−βkk^T) 作用于 S，各段转移算子不可交换，S_seg1 + S_seg2 无意义（实测：直接相加误差 0.978）。

4.11 解码缓存 KDACache 与 clone() 分支

4.11.1 原理

直觉：KV cache 可逐 token 截断回滚；KDA 状态是整体矩阵，回滚只能靠快照。

专家视角：KDACache = {state, conv_q, conv_k, conv_v}，clone() 深拷贝。clone() 是 beam search / 投机解码"分支状态"的接口——候选续写互不污染，失败候选丢弃副本。这是 SGLang ReplaySSM 在工程上解决的同款问题。

4.11.2 对比基线

4.11.3 反方辩手

• a) 每分支都 clone 整个状态矩阵，beam 宽度大时显存爆炸。
• b) 若失败，错在假设分支数小。

4.11.4 工程实现

项目代码（kda.py 中 KDACache）：

class KDACache:
    """每层的解码状态：递推矩阵 + 三条 ShortConv 的滑动窗口。

    这是 KDA 的"KV cache 等价物"，但与 softmax attention 的 cache 有本质区别：
    - 尺寸固定：state 是 (batch, num_heads, head_dim_k, head_dim_v)，
      不随已生成 token 数增长；只有 conv_* 的窗口随 kernel_size 固定。
    - 是整体状态而非逐 token 条目，因此"回滚到第 k 步"无法靠简单截断实现，
      只能靠保存/恢复快照（clone()）。

    支持 clone()：beam search 和投机解码需要在不同候选续写之间分支状态而不互相污染。
    """

    __slots__ = ("state", "conv_q", "conv_k", "conv_v")

    def __init__(self, state, conv_q, conv_k, conv_v):
        self.state = state          # (B, H, Dk, Dv) 递推状态矩阵 S
        self.conv_q = conv_q        # (B, k_dim, kernel-1) q 卷积窗口
        self.conv_k = conv_k        # (B, k_dim, kernel-1) k 卷积窗口
        self.conv_v = conv_v        # (B, v_dim, kernel-1) v 卷积窗口

    def clone(self):
        return KDACache(
            state=self.state.clone(),
            conv_q=self.conv_q.clone(),
            conv_k=self.conv_k.clone(),
            conv_v=self.conv_v.clone(),
        )

init_cache 方法：为 decode 初始化空状态，之后每步由 forward(cache=...) 原地更新。

模型级使用示例：init_caches 逐 block 调 block.attn.init_cache(...) → 返回 KDACache 或 MLACache 的混合列表；generate 逐 token 复用同一份 caches。

cache 列表长度 = n_layers（block 数），但 contents 是 KDA/MLA 两类交替。每层 attention 只消费 caches[i]，list 下标 i 恰好对应该 block。

4.11.5 问题 & 讨论

Q1：KDACache 中，state 是针对每个 head 一个吗？

精确说是每 batch × 每 head 一个递推矩阵：

state       (B, H, Dk, Dv)   # B=批次, H=head 数, Dk/Dv=head 维 key/value

即每个 head 持有自己独立的 Dk×Dv 矩阵 S。逐 head 独立——S_h ∈ R^{Dk×Dv}，delta-rule 递推 S ← γS + beta·k·vᵀ 在 H 维各自进行，互不混合。

conv_q/conv_k/conv_v 却不是逐 head 的：形状 (B, k_dim, kernel-1)、(B, v_dim, kernel-1)，其中 k_dim = H·Dk、v_dim = H·Dv——卷积窗口按扁平宽存，head 维被折叠进第二维，因为 ShortConv 先于 head 切分。

所以结构是：head 各自的空间给 state（矩阵），共享的空间给卷积窗口（扁平宽）。这与 KDA 的语义一致——递推状态是每头的关联记忆，卷积只做 token 级局部混合、作用在整个宽向量上、head 无关闭。

Q2：K3 是每个 block 一个 Cache？

是的——但 cache 的类型随 layer 交替。因为每个 block 只有一个注意力子层 self.attn，它要么是 KimiDeltaAttention 要么是 GatedMLA。所以：

与开源参考核对：kimi3-master 用 serve/decode.py::SpeculativeKDAState + ReplaySSM 处理回滚（不存每步快照，而是重放轨迹）；本项目 clone() 是简单但正确的 CPU 等价。

Q3：Block KDA 的 State 描述是否合理？

合理，但应理解为实现语义变化：

• 纯 KDA: implicit + long-term，状态隐式包含历史；
• block KDA: 把一段历史的状态效果显式封装为 block‑level transition；
• 配合 Full/Sparse: 提供 explicit/addressable + short‑term 访问。

它不是把 KDA 变成了普通 KV Cache，而是把部分隐式状态管理显式化，以支持缓存、分块和并行。

0x05 可视化与实验

5.1 可视化脚本

5.2 图说明

plot_kda_principle.png 四格：

• 左上：状态矩阵 S（记忆）热图。展示固定大小 Dk × Dv 矩阵，T 个 token 只产生一个矩阵。
• 右上：单步三件事拆解（衰减→检索→只写误差）。
• 左下：delta rule 的"纠错"本质（同一 key 下 v 逐渐收敛到最新值）。
• 右下：有界 log-decay（g ∈ (g_min, 0)）曲线。

plot_kcp.png 两格：

• 左图：直接相加 vs KCP 前缀扫描的误差（log 轴）。
• 右图：KDA 固定通信量 vs softmax 线性增长（log-log）。

plot_decay_stability.png 四格：

• 左上：1/gamma 上界 vs tile 大小（有界版是数据无关直线）。
• 右上：实测 tile=16——无界 z+6/z+9 爆成 inf。
• 左下：safe_chunk_size vs g_min。
• 右下：失败 case——chunk=64 时连有界版也溢出。

5.3 关键实测数值回填

KCP 实验：

• 直接相加误差：9.782e-01
• KCP 前缀扫描误差：1.110e-16
• 通信量：KDA 恒定 2048 字节；softmax 从 1k 到 1M 序列线性到 1GB

数值稳定性实验：

• max_reciprocal_decay(-5, 16) = 5.5406e+34 = exp(80)
• safe_chunk_size(-5, bf16) = 16 = 论文 tile
• safe_chunk_size(-5, fp16) = 1
• safe_chunk_size(-8, bf16) = 10
• 无界 softplus z+6 tile=16 → inf/NaN
• chunk=64 连有界版也溢出（印证安全边界）

5.4 测试验证

合计 37 项测试全部通过。

0xFF 参考

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