二万二千五百八十。这就是一个 KimiK3（2026）中能装下多少个 GPT-2（2019）。七年间，我们把规模扩大了 22,580 倍。但这只是……规模变大吗？
这篇工作笔记会带你回顾我们是如何走到今天，以及从那时起究竟改变了多少、又有多少并未改变。我们会沿着通向 KimiK3 的主要架构演进一路追踪。
GPT-2
GPT-2 是一种纯解码器（decoder-only）架构：
tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
x = self.transformer.drop(tok_emb + pos_emb)
for block in self.transformer.h:
x = block(x)
x = self.transformer.ln_f(x)
logits = self.lm_head(x)
return logits
输入会获得 token 嵌入与位置嵌入：
放大来看，每一个 Transformer block 是这样的：
class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
self.attn = CausalSelfAttention(config)
self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
self.mlp = MLP(config)
def forward(self, x):
x = x + self.attn(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return x
注意力的计算过程：
B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
# calculate query, key, values for all heads in batch and move head forward to be the batch dim
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
# manual implementation of attention
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
att = self.attn_dropout(att)
y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
# output projection
y = self.resid_dropout(self.c_proj(y))
return y
最终的隐藏状态矩阵生成后，语言模型头会将其映射为词表上的 logits。在自回归解码时，只需最后一个位置的 logits 来选择下一个 token。
这正是纯解码器生成的一个低效之处：模型为每一个输入位置都计算了表示，但每一步解码只会消耗最后一个位置的 logits。没有缓存时，其中大量工作会在生成下一个 token 时重复进行。
KV cache 来自一个很直接的观察：把生成出的 token 追加到输入后，模型本来会为所有先前 token 重新计算投影。保存它们的 key 与 value 向量，就能避免这部分重复工作。
这份存储就是 KV cache。它保留前 N-1 个 token 的向量，并可能大到形成内存带宽瓶颈。
总体而言，我们的基准模型拥有约 50k 个可能 token、12 个 block、12 个头与 768 的嵌入维度，约为 1.24 亿参数。
vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
n_layer: int = 12
n_head: int = 12
n_embd: int = 768
KimiK3 的 2.8 万亿参数，约等于 22,580 个 GPT-2 的参数总量。
线性注意力
Softmax 注意力会在 q·k 相乘之后施加非线性，因此每个 query 都与每个 key 耦合。线性注意力则分别对 q 和 k 施加特征映射，例如 ELU+1。这使乘法可以重新结合，因此不断增长的 K、V 向量集合可折叠进一个固定的 D×D 状态。
论文把复杂度写成 O(N²) 一度让我困惑。说“Transformer 每个时间步的成本会按当前序列长度的平方增长”并不对；这正是 Flash Attention 解决的问题……后来我才看到它是在 2020 年发布的。
当时训练通常会实体化完整的 N×N 注意力矩阵，FlashAttention 尚不存在，参考实现里的自回归模型也常常不使用 KV cache 而重新计算 token 历史。
def forward(self, x, mask=None, past_kv=None):
# x is b,t,d
b,t,d=x.shape
d_head=d//self.num_heads
h=self.num_heads
qkv=self.qkv_proj(x)
q=qkv[:, :, :d].view(b,t,h,d_head).transpose(1,2)
k=qkv[:, :, d:2*d].view(b,t,h,d_head).transpose(1,2)
v=qkv[:, :, 2*d:].view(b,t,h,d_head).transpose(1,2)
# at prefill, q,k,v have shapes b,h,t,d
# at decode, shape is b, h, 1, d
# so i cat at the t dimension, dim(2)
if past_kv is not None:
k_past=past_kv[0]
v_past=past_kv[1]
k=torch.cat((k_past, k), dim=2)
v=torch.cat((v_past, v), dim=2)
scores=(q@k.transpose(-1,-2))/math.sqrt(d_head)
if past_kv is None: #we're in prefill and need to mask
causal_mask=torch.ones(t,t,dtype=bool, device=q.device)
causal_mask=torch.triu(causal_mask, diagonal=1)
scores=scores.masked_fill(causal_mask, float('-inf'))
if mask is not None:
scores=scores.masked_fill(~mask, float('-inf'))
#get attn (bhtt x bhtd)
attn=scores.softmax(-1)#bhtt
o=attn@v #bhtd
o=o.transpose(1,2).contiguous().view(b,t,d) #b,t,d
# use x to get qkv
o_proj=self.o_proj(o)
past_kv=(k, v)
return o_proj, past_kv
用图更容易理解同一过程。每一步解码都会对 HBM 执行两次 ND 读取和两次 1D 写入；而 KV cache 会随序列长度以 O(N) 线性增长。
注意这些过量的读写；这篇论文将它们替换为：
def forward(self, x, mask=None, cache=None):
# x is b,t,d
b,t,d=x.shape
d_head=d//self.num_heads
h=self.num_heads
qkv=self.qkv_proj(x)
q=qkv[:, :, :d].view(b,t,h,d_head).transpose(1,2)
k=qkv[:, :, d:2*d].view(b,t,h,d_head).transpose(1,2)
v=qkv[:, :, 2*d:].view(b,t,h,d_head).transpose(1,2)
k=F.elu(k)+1
k=k.transpose(-1,-2)
q=F.elu(q)+1
S,z=cache if cache is not None else (0.0, 0.0)
S=S+k@v
z=z+k
o=q@S #bhtd
denom=q@z
o_scaled=o/denom
o_scaled=o_scaled.transpose(1,2).contiguous().view(b,t,d)
o_proj=self.o_proj(o_scaled)
cache=(S,z)
return o_proj, cache
这里存在取舍。
这里我们用分别施加在 q、k 上的 ELU+1，取代了 softmax 在它们交互之后使用的指数函数。两种方法都会归一化所得分数，但线性注意力的特征映射对 softmax kernel 的近似表达能力较弱。这个近似可能降低保真度，不过实际精度损失取决于架构与工作负载。
请注意，我们仍会除以 qk 的总和；为简洁起见，图中省略了这一点。高层来看，注意力由三步组成：
- 让 qk 分数变为非负。线性注意力使用 ELU+1，softmax 使用指数化。
- 除以总和。
- 计算 value 的加权平均。
这样保留了注意力的基本契约，却使用表达能力较弱的特征映射让 QK 分数非负。
DeltaNet（快速权重编程器）
有限的 cache 必须覆盖或合并已经存储的信息。来自 token i-1 的状态不会拥有自己的槽位；它被加入同一个 D×D 矩阵。因此，新的 query 无法再为每一个更早 token 检索完全隔离的表示。
这份相加也是效率提升的来源。用加法而非拼接更新 cache，使它不再按 O(N) 增长，但同一个操作也会造成信息互相干扰。DeltaNet 正是为了解决这种可恢复性的损失。
Schlag 的论文《Fast Weight Programmers》有一段很优雅的表述：“当序列长度超过存储容量时，模型可能进入超容量状态。为了在这种状态下正常运行，模型应该学会与记忆内容动态交互，有选择地决定保留哪些 key-value 关联、删除哪些关联。纯加法指令可能并不适合这一目的……在有限大小的记忆中无休止地添加新关联，如式 17 所示，终将达到极限。”
线性注意力最吸引人的状态正是 N 远大于 D；但这也暴露了它的主要限制。一旦状态超过有效容量，关联便开始相互干扰，因为更新是加法式的，且没有任何东西离开 cache。
def forward(self, x, mask=None, cache=None):
# x is b,t,d
b,t,d=x.shape
d_head=d//self.num_heads
h=self.num_heads
qkv=self.qkv_proj(x)
q=qkv[:, :, :d].view(b,t,h,d_head).transpose(1,2)
k=qkv[:, :, d:2*d].view(b,t,h,d_head).transpose(1,2)
v=qkv[:, :, 2*d:].view(b,t,h,d_head).transpose(1,2)
q = F.normalize(F.silu(q), dim=-1)
k = F.normalize(F.silu(k), dim=-1)
beta = torch.sigmoid(self.w_beta(x)).view(b, 1, t, 1)
# new: per-token write strength
S = cache if cache is not None else 0.0
v_old = k @ S # read the board at this key
u = beta * (v - v_old) # the delta: only what's actually new
S = S + k.transpose(-1, -2) @ u # same outer-product write as before
o = q @ S # read, no denominator
o = o.transpose(1, 2).contiguous().view(b, t, d)
return self.o_proj(o), S
用一个视觉示例会更容易跟上。
设有一个写入的关联 S = k.T @ v。若用同一把 key 读回，得到的是 k @ (k.T @ v)，即 (k @ k.T) v，也就是 k 的平方范数乘以 v。因此读回结果会按 key 的平方范数缩放；若将 k 归一化到单位长度，或仅用范数相除，就能精确得到 v。
Q 也是一个学习到的指针。Wq 与 Wk 从同一个残差流读取，而某个事实的 query 指向该事实被写入的 key 方向。更新会先询问当前 key 能从 cache 中取回什么信息；它从要存储的 value 中减去现有信息，将 key 乘以差值后再加回去。旧信息被移除，新信息则写入原位。
DeltaNet（用 Delta Rule 并行化线性 Transformer）
这是本文最难的一节。我花了约七小时才建立起可用的理解，因此会从实现出发构建说明。简言之，DeltaNet 用广义 Householder 转移矩阵实现一阶线性递推，从而支持按 chunk 并行的前向计算，在硬件上实现高效的线性时间训练。它将输入、输出分成多个大小为 C 的 chunk，并基于前一 chunk 的最终状态，以及当前 chunk 的 query、key、value block 计算每个 chunk 的输出。
实际问题在 prefill。对含有 T 个 token 的序列直接实现 Delta rule，会像这样：
S = torch.zeros(b, h, dh, dh) if cache is None else cache
outs = []
for i in range(t):
k_i = k[:, :, i:i+1]
v_i = v[:, :, i:i+1]
b_i = beta[:, :, i:i+1]
v_old = k_i @ S
u_i = b_i * (v_i - v_old)
S = S + k_i.transpose(-1, -2) @ u_i # write
outs.append(q[:, :, i:i+1] @ S)
o = torch.cat(outs, dim=2)
不同于标准注意力，这个表述要求在每一个 key 向量处进行一次修正，因此通往并行矩阵乘法的路径并不显然。即使没有 Delta rule，直接的线性注意力 prefill 仍然是顺序执行的：
S = torch.zeros(b, h, dh, dh) if cache is None else cache
outs = []
for i in range(t):
q = q[:, :, i:i+1]
k = k[:, :, i:i+1]
v = v[:, :, i:i+1]
S=S_old+k@v
o=q@S #bhtd
o=self.norm(o)
o=o.transpose(1, 2).contiguous().view(b, t, d)
out=self.o_proj(o)
cache=S
outs.append(out)
o = torch.cat(outs, dim=2)
按 chunk 表述能提供更高效的办法。通过一个示例更容易理解其中机制：
令 C=N 就回到了标准 O(N^2) 注意力，而 C=1 则是常规线性注意力。中间值可在两者间插值：以额外的 chunk 内工作换取更好的硬件利用率。实践中 C 通常取 64 或 128，因为 tensor-core 指令在这一粒度下效率很高；UMMA 就是一个例子。
中间 tile 会作为状态更新的一部分被折叠进 S：
S = torch.zeros(b, h, dh, dh) if cache is None else cache
outs = []
for i in range(t//C):
q_c = q[:, :, i*C:(i+1)*C]
k_c = k[:, :, i*C:(i+1)*C]
v_c = v[:, :, i*C:(i+1)*C]
o_prev=q_c@S #this is everything up to this block
attn=(q_c@k_c.transpose(-1,-2)).tril() #masked attention
o_curr=attn@v_c
o=o_prev+o_curr
S_new=k_c.transpose(-1,-2)@v_c #recurrent attention
S=S+S_new
outs.append(o)
o = torch.cat(outs, dim=2)
在一个 block 内，我们做 q(kᵀv)：先算分数，这就是带 mask 的常规注意力次序。跨 block 时，我们遵循 (kᵀv)q，因此先处理递归次序、也就是状态。注意力复杂度会按 O(N²) 增长，而这里不会。block 内我执行真正的注意力（带 mask 的 QKᵀ 乘 V）；block 之间则把一切折叠进状态，再用一次 matmul 读回。因此成本分成两部分：固定项 2Ld² 是状态工作，与 C 完全无关；增长项 2LCd 则是对角线上的分数矩阵。全注意力只是 C 等于 L 的情形，这时第二项成为二次的 2L²d。所以 C 越小，FLOPs 越少。
从纯 FLOP 角度看，C=1 最便宜，但未必拥有最短的实际耗时。当工作能高效映射到 GPU 的矩阵乘法硬件时，GPU 可更快完成更多运算。
下一步是将同样的方法扩展到 DeltaNet。
根本问题很简单：纯加法注意力使用的 chunk 方法，不能直接套用于 delta 更新：
v_old = k_i @ S
u_i = b_i * (v_i - v_old)
为了计算需要减掉的信息，我们需要按顺序得到每一个状态。若没有某种数学上的重新参数化，就无法以相同方式并行化。因此作者将 delta 更新从：
u=v_new-v_old
S_t= S_(t-1)+K.T@u
o=q@S_T
改写为如下形式。前一种形式以顺序循环在每一次迭代计算一个 delta；重新参数化后的形式是：
S_t = S_{t-1}(I − β_t k_t k_tᵀ) + β_t v_t k_tᵀ
o_t = S_t q_t
这一表述让按 chunk 的代码可以一次计算全部 C 个 delta：
def chunk_delta_rule_forward(Q, K, V, beta, C):
# L: sequence length, d: head dimension
L, d = Q.shape
# chunking
Q, K, V = map(lambda x: x.reshape(-1,C,d), [Q, K, V])
beta = beta.reshape(-1, C)
K_beta = K * beta.unsqueeze(-1)
V_beta = V * beta.unsqueeze(-1)
# compute eq. 10 with vectorized forward substitution for fast inverse
T = -(K_beta @ K.t()).tril(-1)
for i in range(1, C):
T[i, :i] = T[i, :i] + (T[i, :, None] * T[:, :i]).sum(-2)
T += torch.eye(C)
W = T @ K_beta
U = T @ V_beta
# chunkwise parallel. Eq. 8-9
S = torch.zeros(d, d)
O = torch.empty_like(V)
for i in range(L//C):
q_i, k_i, w_i = Q[i], K[i], W[i]
u_i = U[i] - w_i @ S # the corrections, all of one chunk
o_inter = q_i @ S
A_i = (q_i @ k_i.t()).tril() #qk.t
o_intra = A_i @ u_i # attention @ v (with corrections, so u)
S += k_i.t() @ u_i # update state with addition
O[i] = o_intra + o_inter #update output with flash + recurrent
return O.reshape(L, d)
到此我们抵达第一个对比点：MHA Transformer 与 DeltaNet Transformer：
门控 Delta Net
现在我们已经有办法精确修改 cache。每增加一个新事实（每一个新的 key 向量），都可以准确查看该位置原先存储的信息，再将其替换为希望注意到的新信息。
但是，这个机制只能忘记某个有特定替换物的关联。它无法在上下文切换时高效清除多个关联，也无法更普遍地让记忆衰减以释放容量。
如果我们做的是纯加法线性注意力：
加入遗忘能力会很简单。我们只需要一个控制遗忘状态的参数：
S_old=cache
S_new=k@v
# cache=S_old+S_new
cache=alpha * S_old + S_new
这正是 Mamba-2 的贡献。我们让旧 cache 衰减，然后以完整强度加入新 cache，防止状态无限增长。
在每个时间步按动态比例均匀衰减所有 key-value 关联是一种可行方法，Mamba 正是这么做的。但它没有考虑不同 key-value 关联的重要性各不相同。
也就是说，若模型需要遗忘一个特定关联，所有关联都会同等被遗忘。相比之下，Delta rule 可以更新单个事实，却没有让其余事实衰减的机制。
因此，Gated Delta rule 将 Mamba 的门控更新规则与 Delta rule 结合。它加入参数 alpha：取值为一时退化为纯 Delta rule；取值为零时清空记忆。难点在于，如何用同一套并行 chunk 方法实现它。
实现沿用上一节所述的 DeltaNet 重新参数化。数学上几乎完全相同，只增加了一项：一个介于零与一之间、依赖数据的标量，用来控制前一状态的衰减。这将高效的 key-value 关联学习与自适应记忆管理结合在一起。
对应代码的变化如下：
γʳ/γⁱ 项刻画累积衰减。在时间步 x 写入的 token，在 x+t 被读取时，已经乘以 αₓαₓ₊₁αₓ₊₂…αₓ₊ₜ。这是前缀和计算的乘法类比。
所得架构如下：
KDA / Kimi Linear
到了这里，研究者开始尝试把多种注意力形式结合在同一个架构中的混合模型，例如 Gated DeltaNet 与 Mamba。
Kimi Linear 因一个核心主张而受到关注：在受控对比下，它优于全注意力。作者将它定位为可直接替换的架构：质量更好，解码吞吐量最高可提升 6 倍。
Kimi Linear 通过引入细粒度 gating 改进了 Gated DeltaNet。它不再采用单一标量衰减，而是为每个 channel 学习独立的衰减值。
KDA 的更新规则仍然类似，不过代码现在更像这样：
这里的 alpha.reshape(nb, C, d) 捕捉了论文最重要的贡献：对记忆衰减进行细粒度控制。
与 DeltaNet Transformer 对照时，Kimi Linear 架构引入了三个主要变化：
- 它使用交错排列 Multi-head Latent Attention（MLA）层的混合系统。
- 它用 Mixture-of-Experts（MoE）层替换 MLP。
- 它通过 alpha projection 为 DeltaNet 增加容量。
后面的章节会更详细讨论 MLA 与 MoE。眼下重要的是：这不是盲目的扩展。增加的容量有明确的数学用途：逐 channel 的尺度让模型能更细致地控制记忆衰减。
Scaling law 依然有效，但容量必须加在正确的位置，并采用系统能利用的形态。在这条演进路径中，每个架构都增加容量，以应对前一系统的某个具体限制。
Kimi K3
归根结底，KimiK3 的语言主干与上面的 Kimi Linear 模型相似。它包含 23 个四层 macrocycle。每个 macrocycle 中，三层使用 Kimi Delta Attention，第四层使用 Multi-head Latent Attention。第一层使用稠密的前馈网络；其余每一层使用 latent Mixture-of-Experts。
乍看之下，Kimi Linear 到 KimiK3 的变化似乎不大：
- 规模显著增加
- 每 12 层加入一次 blockwise AttnRes
- MLA query LoRA 与输出 gating
- latent-space MoE
- SiTU 激活
- Gated MLA
KDA 提供常量状态的递归记忆，周期性的 MLA 层则保留了对上下文进行完整 softmax 检索的能力。下面的简化图为之后的变化提供了一个有用参照。
我们先从更直接的变化说起：Gated MLA、latent-space MoE 和 SiTU 激活。
Gated MLA 决定从 MLA 检索到的每项特征有多少能进入残差流。它通过一个由输入投影出的 gate 与这些特征逐元素相乘来做到这一点。
在传统 MoE 中，学习得到的 router 用点积相似度，将每个 token 发往一小部分 expert network。KimiK3 总共有 898 个 expert：其中 2 个为共享 expert，处理每一个 token；剩余 896 个中，router 为每个 token 选择 16 个。
KimiK3 也改变了 expert activation。它不再对 up projection 施加 SiLU、与 gate 逐元素相乘、再应用 down projection，而是使用 SiTU：
d = x.shape[-1] // 2
gate = x[..., :d].to(torch.float32)
up = x[..., d:].to(torch.float32)
situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate)
if self.linear_beta is not None:
up = self.linear_beta * torch.tanh(up / self.linear_beta)
return (situ_a * up).to(x.dtype)
模型还会将输入下投影到共享 expert，并把它们的最终和上投影回来：
这揭示了模型推理中的一个反复出现的挑战。若没有 fused kernel，新 activation 比原路径慢将近 3 倍。一个抵消性优化是，expert 在压缩的 latent 空间中运行，使其前向传递快得多，且 FLOPs 几乎减半。
剩余变化包括 MLA query LoRA、输出 gating，以及每 12 层的 blockwise Attention Residual。AttnRes 大约增加 2% 推理延迟，但提供两个重要收益：
- 有选择地检索更早的表示，从而缓解残差稀释与隐藏状态增长
- 1.25 倍的计算优势
AttnRes 和 MLA 从不同方向解决同一项底层限制。KDA 层以常量大小的状态运行，必然不得不丢弃信息。MLA 从 token 上下文中检索，而 AttnRes 从更早的、沿深度方向的表示中检索。
AttnRes
感谢 @chloey3k 协助本节。在每次前向传递中，输入穿过一叠层。这里每层由一个 attention block（KDA 或 MLA）和一个 MLP 或 MoE block 组成。通常，每层的输入是原始 embedding 加上此前所有层输出的和，且所有部分权重相同。
这里，h_i 是第 i 层输入，h_1 是当前 token 的 embedding（截至目前序列中的最后一个 token），f_i(h_i) 是第 i 层的输出（attention 或 MLP block）。
问题在于缺乏选择性访问。不同层类型接收到的是相同的聚合状态，即便它们可能从不同的权重分配中获益。由于递推纯粹是加法式，后续层为了影响累积残差，也必须学会产生越来越大的输出，这可能使训练不稳定。AttnRes 不再平等对待所有层，而是把该和的每一项乘以专门的权重，让模型能在当前上下文中赋予最有用层更高的重要性。
每一个权重 alpha_i 都由 query-key 点积计算。query 为每层学习得到，而 key 和 value 来自更早的残差流状态。分数被归一化至总和为一，再用来形成这些状态的加权组合。
因此模型不必只以直接前驱为条件。AttnRes 让每一层都可有选择地访问更早层的输出，使其学习得到的 query 能检索当前计算最有用的表示。
下面的伪代码把相同思路用在 block 粒度。一个 block 是 12 个 decoder layer 中累积的 attention 与 MLP 输出逐元素求和，并作为单一的深度表示保存，以供之后的 AttnRes 混合。
在每一层都应用 residual attention，会增加过多训练与推理成本。只在固定 block 边界应用，则能以更低成本取得大部分收益。KimiK3 中每个边界都位于 12 个 decoder layer 之后。跨越 23 个四层 macrocycle，这会产生八个 AttnRes block，从而提升推理速度。
block_attn_res 函数中或许最重要的部分是：
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
至此，我们完成了从 GPT-2 到 KimiK3 的演进路径。
核心变化并不只是规模。每一步架构演进都改变了模型存储什么、如何更新状态，或如何检索固定大小状态无法保留的信息。
KimiK3 将常量状态的递归记忆、周期性的 softmax 检索、稀疏的 expert 容量，以及有选择的深度向残差访问结合在一起。最终，额外容量被花在了具有明确功能角色的地方。
本质上，固定容量的联想记忆（维度固定）需要一种驱逐策略：纯加法线性操作一旦达到容量上限，终会引入干扰。为此，gating、routing 或 decay 这样的学习型选择机制是必要的，而注意力则是最有效的选择性读取机制。