Kimi Linear: An Expressive, Efficient Attention Architecture

整体架构

Kimi Linear 将 KDA 和 MLA 放在 Transformer 的 token-mixing 位置, 并采用 3 个 KDA layers + 1 个 MLA layer 的排列; KDA 用固定大小的 recurrent state 处理历史信息, MLA 则保留 full attention 与 latent KV cache; 每个 token-mixing layer 后仍然连接一个 MoE layer, 整体数据流保持 Norm -> KDA 或 MLA -> residual -> Norm -> MoE -> residual;

KDA layer 的输入, 状态更新与输出

生成第 $t+1$ 个 token 时, KDA layer $l$ 接收两个输入: 当前 token 从上一层传来的 hidden vector $h_t^{(l)}$ , 以及本层处理完上一个 token 后留下的 recurrent state $S_{t-1}^{(l)}$ ; 前者沿模型深度方向在 layers 之间传递, 后者沿 token 时间方向保留在同一个 KDA layer 内;

本层固定的 projection weights $W_Q,W_K,W_V$ 先将 $h_t^{(l)}$ 投影为当前 token 的 $q_t,k_t,v_t$ ; 其中 $k_t,v_t$ 用来将当前 token 的关联写入旧 state, 同时结合 KDA 的遗忘与修正规则得到新 state $S_t^{(l)}$ ; $q_t$ 再从更新后的 state 中读取历史关联, 得到当前 token 的 attention output; 该 output 经过 output projection 和 residual 后继续进入 MoE 与下一层, 而 $S_t^{(l)}$ 留给下一个 token 使用;

当前 hidden vector h_t
        -> 固定 projection weights 生成 q_t, k_t, v_t

旧 state S_(t-1) + 当前 k_t, v_t
        -> 更新为新 state S_t

当前 q_t 查询 S_t
        -> attention output
        -> residual, MoE 和下一层

Recurrent state 在数学形式上像一组可以被 query 读取的临时 fast weights, 但它不是训练得到的 model weights; projection weights 决定如何读写记忆, 在推理期间保持固定, recurrent state 则从零状态或 checkpoint 开始, 随当前请求的每个 token 动态更新; 每个 KDA layer 只保存最新的一个 state, 但这个 state 汇总了该 layer 对整个历史 prefix 的表示;

MLA layer 不读取 KDA layer 的 state, 而是维护自己这一层的 latent KV cache; 不同 layers 的历史状态彼此独立, 当前 token 的 hidden vector 则会依次经过所有 layers, 将前面 KDA 和 MoE 的计算结果带到后面的 MLA layer;

上下文修改与 checkpoint

KDA 的 recurrent state 是按照 token 顺序递推得到的, 因而它只对应生成该 state 时的那段确定 prefix; 如果 agent 删除或修改了中间文本, 修改位置之后的 state 就与新文本不再对应, 不能继续沿用; 系统需要回到与新旧文本仍然一致的最晚 checkpoint, 再从该位置重新处理后续 token;

在通常按 token 边界恢复的 checkpoint 中, 需要保存每个 KDA layer 的 recurrent state, 以及每个 MLA layer 对应同一 prefix 的 latent KV cache 和位置等必要元数据; 不需要额外保存所有历史 token 的 layer activation $X$ , 因为它们只是一次 forward 中沿模型深度传递的临时结果; 只有在某个 token 的 forward 尚未完成时暂停, 才可能需要保存当前的中间 activation; 因此, Kimi Linear 虽然混合了 KDA 与 MLA, 回滚思路仍然直观: 同时丢弃断点之后的 recurrent states 和 KV cache, 恢复到最晚未被修改的 prefix, 然后对新的 suffix 执行 reprefill;