【轻量化复刻大模型 / KIMI K3】KDA（Kimi Delta Attention）：并行
0x00 概要
本文档是复刻 KIMI K3 KDA（Kimi Delta Attention）的第二部分，因为文章实在太长了，所以分为两期发出。
本篇主要分析并行相关的部分。
Chunkwise 并行并不是没有代价。RNN 的状态递推通常对输入顺序敏感：交换两个 Token 的位置，会改变它们经过的状态，也会改变后续计算。而 Attention 的一次内容聚合并不天然表达顺序关系。固定一个 Query，如果只把可见范围里的 Key-Value 对重新排列，单次聚合结果本身可能不发生变化。
KDA 能够同时获得：
- 相对上下文长度为 O(1) 的状态存储和 Decode 计算；
- Chunkwise 并行能力；
根本原因是 KDA 把对旧状态的递归更新写成线性或仿射变换：。这使得多个 Token 的状态转移可以组合成一个 Chunk 级变换。这里体现出的是表达能力和计算效率之间的取舍。
注：
- 项目github.com/RossiLH/Mod… 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
- 涉及 KIMI K3，DeepSeek V4，PI0.5，PI0.7 等模型。后续可能会跟踪新模型。
- 全 CPU 复刻，不涉及 GPU，也不涉及大规模训练（这是另外一个庞杂深邃的领域）和硬件。
- 模型的组件尽量是原始大模型的缩小版，如果组件没有开源版本，则用近似组件替代。
- 会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么，以及组件之间如何咬合成一个整体。重点不是记名词，而是理解一条主线。
- 每个主要组件都有可直接运行的 Gradio demo——无需 GPU。另有若干可视化脚本（plot_*.py），可生成静态图。
- 所有文件行号基于本地源码快照；源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。
- 为了保证每个文章可以尽量独立成文，所以有些内容会在不同文章中都存在，可能在系列层面上存在冗余，还请大家理解。
0x04 子知识点深度分析
4.7 ShortConv、SiLU 与 L2 归一化（无位置编码的局部混合）
线性注意力擅长把很长的历史压缩到固定大小的状态中，但这种压缩也意味着细节可能丢失。尤其是相邻 Token 之间的短距离关系，不一定能够仅靠全局状态被充分保留下来。因此，DeltaNet、GDN、KDA 等架构通常会在投影得到的 q、k、v 上增加一维因果短卷积（Causal Conv1d）。处理流程可以理解为：
- 先从当前 Hidden State 投影得到原始的 q/k/v；
- 用因果短卷积融合当前 Token 和附近若干历史 Token 的局部信息；
- 把卷积后的 q/k/v 送入线性注意力的状态更新和读取过程；
- 由固定大小的递归状态负责聚合更长距离的全局历史。
因此，这类架构实际上把两种互补机制组合在了一起：
- 短卷积负责局部信息：强化相邻几个 Token 之间的细节和短程模式；
- 线性注意力负责全局信息：把不断增长的历史压缩到固定大小的状态中。
可以把它概括为：卷积补充局部关系，线性注意力聚合全局历史。
4.7.1 原理
直觉
KDA 没有位置编码，模型怎么知道"相邻 token"关系？ShortConv 是一个宽度 4 的因果深度卷积，让每个 channel 在进入递推前先"看"最近 3 个 token——局部位置信息由此注入。
想象每个通道有一个3 格的小记忆条（kernel 4 = 当前 + 前 3 格）。每来一个新 token，它把"当前值 + 前 3 个值"按一组小权重混合，作为新的输出。这个小记忆条只记得非常近的过去（4 个 token），所以它擅长捕捉局部模式（如字符级/词级的位置关系），但不负责长程 —— 长程由后面的递推（KDA）负责。
因为没有位置编码（NoPE），这个小记忆条成了模型感知"顺序/位移"的主要局部通道。
准专家视角
短卷积可以直接理解为：对 q、k、v 的每个通道分别使用一个固定长度的滑动窗口，并对窗口内的值做加权求和。即，对每个 channel 独立做因果 1D 卷积，kernel=4，无跨 channel 混合。输出 q/k/v 在进入递推前被局部平滑。它是“因果”的，意味着当前位置的输出只能使用当前 Token 和过去的 Token，不能访问未来 Token。
这是混合位置编码思路（卷积为局部位置敏感、递推为全局压缩）的关键一环。它只提供"移位"（shift）信息，配合 delta rule 的位置感知衰减共同替代 RoPE。
输入输出
- 输入：x ∈ R^{B×T×d}（d=隐维/通道数）
- 卷积核：W ∈ R^{d×1×k}（k=kernel_size=4，depthwise 即逐通道）
- 输出：y ∈ R^{B×T×d}
假设卷积核长度为 4，那么当前输出会使用：
- 前面第 3 个 Token；
- 前面第 2 个 Token；
- 前面第 1 个 Token；
- 当前 Token。
每个相对位置有一个训练得到的卷积权重。以某一个通道为例，卷积和激活可以写成：
卷积输出 = SiLU(
前第3个Token的qkv * conv_weight[:, 0] + 前第2个Token的qkv * conv_weight[:, 1] +
前第1个Token的qkv * conv_weight[:, 2] + 当前Token的qkv * conv_weight[:, 3]
)
这里的重点有两个：
- 窗口长度是固定的，所以局部计算成本不会随着总上下文长度增长；
- 卷积权重是在训练中学习得到的，但在具体计算时是固定权重，不会像注意力那样根据当前 Query 和局部 Key 的内容动态计算一组匹配分数。
公式（因果互相关）
y[t,c] = sum_{j=0}^{k-1} W[c,0,j] * x[t - (k-1-j), c] # 只看当前与过去
- 因果性：左侧补 k-1 个零/状态，保证 y[t] 只依赖 x[≤t]。
与标准做法的公式级对比
专家视角
"无位置编码 + 局部卷积 + 递推衰减"是一套完整的位置信息方案：局部用卷积、全局用衰减序。
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[t]。初始化把最后一个 tap 置 1、
其余置 0，于是 out[t] = 1 * x[t]——恒等输出。
"""
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-K3-Eq1
对拍验证：将 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 的区别只是工程封装：
- kimi_linear_step 输入 gate_logit，自己算 g；kda_step 输入预计算好的 g（上层 log_decay 算）
- kimi_linear_step 对 q 乘了 1/√Dk 的 scale（对齐 SGLang fused_recurrent）；kda_step 不加（尺度由 HeadRMSNorm + 输出门吸收）
- 状态更新的数学本质完全一样
原理
decode 时每步只有 1 个 token，直接走 Eq.1 的单步更新（无 chunk 结构开销）。且因为首个 token 无先前状态，走 chunkwise（初始状态 0）保证与 pre
本文由 AI图鉴自动抓取并翻译，版权归原作者所有；如需查看原文请点下方链接。
KDA (Kimi Delta Attention):
0x00 Summary
This document is the second part of the recapitulation of KIMI K3 KDA (Kimi Delta Act), as the article is too long to be issued in two stages.
The main analysis in this chapter is the parallel and related part.
Chunkwise's parallel is not free. RNN's status is usually sensitive to input sequences: the exchange of two Token positions changes the state they pass through and changes subsequent calculations. An Attention content aggregation is not a natural sequence. Sets a Query, and if only the Key-Value in the visible range is rearranged, the single aggregation result itself may not change.
KDA can access both:
- State storage and Decode calculation relative to context length O(1);
- Chunkwise parallel capacity;
The underlying reason is that KDA writes an update of the back-to-back state as linear or imitating transformation: This allows multiple token status shifts to be combined into a Chink class transformation. This reflects the trade-off between expressional competence and computational efficiency.
Notes:
- Project Github.com/RossiLH/Mod... and series of blogs are an attempt and a learning process. I will myself reflect on improvements and hope that you will make valuable comments, in the absence of certainty as to whether the needs of the reader are met.
- Models such as KIMI K3, DeepSeek V4, PI0.5, PI0.7 etc. New models may be followed up.
- All CPUs, not GPUs, not large-scale training (an additional area of depth) and hardware.
- The components of the model are as narrow as possible from the original LLM and, if there is no open-source version of the components, are replaced by approximate components.
- Explains how each component solves a problem, where it is difficult, why it can only be so solved, what the rationale for the solution is, and how the components combine into a whole. The focus is not on naming, but on understanding a thread.
- Each major component has a direct running Gradio demo - no GPU. There are also several visual scripts (plot *.py) that generate static maps.
- All file line numbers are based on local source snapshots; if the source code changes, the line numbers may be moved, but the function name and call relationship can still serve as the location entry.
- In order to ensure that each article can be written as independently as possible, some elements will be present in different articles and may be redundant at the level of the series.
0x04 sub-knowledge point depth analysis
4.7 ShortConv, SiLU and L2 homogenization (local mix without location code)
Linear attention is good at compressing a long history into a fixed-size state, but this compression also means that details may be lost. In particular, the short distance between neighbouring Token is not necessarily sufficiently preserved by the overall state alone. As a result, structures such as DeltaNet, GDN, and KDA usually add one-dimensional complication (Causal Conv1d) to their projection. Processes can be understood as:
- First get the original q/k/v from the current Hidden State projection;
- Combining local information from the current Token and a number of nearby history Tokens with cause- and effect-short volumes;
- Updates and readings of the state of q/k/v, which brings the volume back to linear attention;
- The global history of longer distances is brought together by a fixed-size regression.
Thus, such structures actually combine two complementary mechanisms:
- Short volumes for local information: reinforce details and short-range models between several adjacent Tokens;
- Linear attention is responsible for global information: condensing the growing history to fixed sizes.
It can be summed up as a volume supplement to local relationships and a linear focus that brings together global history.
4.7.1 Rationale
Intuition.
KDA doesn't have a location code. How does the model know about the next token relationship? ShortConv is a cause-and-effect-deep volume of 4 widths that allows each channel to "see" three recent tokens before entering the thrust -- local location information is injected.
Imagine each channel having a 3-gauge small memory bar (kernel 4 = current + first 3-gauge). For every new token, it mixs "current value + first three values" into a set of small weights as a new output. This little memory note only remembers the very near past (4 tokens), so it's good at capturing local patterns (e.g. character/word level position relationships), but it's not responsible for long distances - long distances are the responsibility of KDA.
Because there is no location code (NOPE), this small memory bar is the main local channel for model perception "sequence/place shift".
Para-expert perspective
Short volume can be understood directly as using a slide window of a fixed length for each q, k, and v and weighting the values inside the window. That is, each channel is independent of cause and effect 1D volume, kernel = 4, no cross channel mix. Output q/k/v is partially smoothed before entry. It's "causal" and it means that the output of the current position can only use the current and past Token, without access to the future Token.
This is a key link in the idea of coding a mixed position (and volume is locally sensitive and cumulatively global compression). It only provides "shift" information, with the co-replacement of RoPE with delta rule's position perception decay.
Input Output
- Input: x ∈R{BXTxd}(d=Invisible/passage Number)
- Volume core: W ∈ R{dx1xk}(k=kernel size=4, channel by channel)
- Output: y ∈R{BXTxd}
Assuming volume core length is four, the current output will be used:
- 3rd in front of Token;
- 2nd in front of Token;
- Number 1 above Token;
- Current Token.
Each relative location has a volume weight that has been trained. For example, volume and activation can be written in one channel:
Volume Output = SiLU()
Qkv * conv right[:, 0] * conv right[, 1] * conv right [:, 1] +
Qkv * conv right[:, 2] * conv right[:, 3]
)
The focus here is twofold:
- The length of the window is fixed so that the partial costing does not increase with the total length of context;
- Volume weights are learned in training, but are fixed in specific calculations, and do not compute a set of points based on current Query and local Key content dynamics, as is the case with attention.
Formula (causes and effects are interrelated)
♪ Just look at the present and the past ♪
- Causal: k-1 zero/state to the left, guaranteed y[t] only x[t].
Formula level comparison with standard practice
Expert perspective
"No Location Encoding + Local Volume + Gradually Decline" is a complete set of location information programs: Local Volume, Global Decline Order.
The K3 mix can be understood as: KDA retention order, proximate cause and last-write-wins; NoPE MLA provides a sequence-neutral aggregation and cross-section. The two complement each other: when tasks need to be "sequenced," let KDA provide bias (KDA 's sequencing relationship comes from state evolution). Even if the calculations were to be broken down to the Chunk parallel, the chunk would still have to be combined in a state shift relationship, so the sequence would not disappear in parallel; and MLA would be reduced by proximate bias when the task required "multi-evidence aggregation".
The boundary is: Nope MLA itself cannot answer "who happens first". Sequencing information still needs to be derived from causal mask, ShortConv, upstream KDA or other indications.
NoPE
NoPE does not solve the whole order problem: NoPE just means that MLA does not add an extra visible rotation. Orders may still come from causal mask, ShortConv, upstream KDA, and the sequence structure of the data itself. So the NOPE can't be equated with "the model doesn't understand the order."
Besides, Nope isn't just "because KDA has provided a location, so save the RoPE". It also avoided competition between two location systems:
If both are imposed, the model may be faced with: relative distance associated with the content generated by KDA; fixed rotation distance generated by RoPE. These two positions may differ in direction. So what the NOPE means is to avoid duplication or conflicting positions.
The difference between short volumes and slider attention.
Short rolls and slide-window attention all converge around Token, which looks like dealing with local contexts, but they have different sources of weight.
- Short volumes use a fixed weight of pre-learning. The same relative position usually uses the same set of volume parameters, and the calculation process is a channel-by-line weighted sum within a fixed window.
- Sliding attention calculates weights based on current content dynamics. The current Query will match the key in the window separately, then determine how much each local Value should take.
The difference can therefore be summarized as follows:
Short volumes do not replace linear attention but provide a cheap local information channel for linear attention. The fixed state of linear attention is responsible for a more distant history, while the short volume compensates for the lack of modelling of adjacent details.
4.7.2 Comparison of baselines
4.7.3 Computation complexity degrees
- ShortConv: O (T.d.k) (k=4), almost negligible.
4.7.4 Parameter sensitivity
- conv kernel (default 4): Turning up local feeling but "silent"; Tuning migration information is insufficient.
4.7.5 Project realization
ShortConv class in project code Initial constants (right[, 0, 1] = 1.0) - Volume starting with no-op, no disturbance of carefully initialized decay dynamics. The status of the window at decode is saved by the conv * of KDACache.
class ShortConv(nn.Module):
""cause and effect depthwise short volume with self-expression to support token exact decode.
This category offers two forms:
-Forward: whole / chunk processing (prefill or chunkwise), counting an entire output at a time,
and returns "Kernel size-1 input at the end" as the next paragraph;
-step: single token intensifier (regression decode) with the last round + current token output.
Weight storage agreement: weight[c, 0, lag] multiplied by x[t - lag, c] (relative custom),
That's right. Initialize the last tap in 1,
The rest is 0, so out [t] = 1 * x [t] - constant output.
"""
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:
Initiation of the constant equation: no-op at the beginning of the volume.
The approach is zero weighting, and the last tap (the current token coefficient) is set to 1.
The local mixing of the model at the start of the q/k/v is constant and will not be initiated with KDA
"Almost constant decay" interferes with each other.
"""
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 is." Returns (sequence after volume, end window state). "'"
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]:
""List token update (regression decode). "X is (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:]
Open source reference check: SGLang uses mamba ' s causal conv1d fn/causal conv1d update; semantic consistency, with different initialization strategies (consistent initialization of the project, Kimi-k3-pytorch with random initial (0, 0.02).
4.7.6 Discussion
Why do you need volume? Why do you need q, k, v?
The simplest Linear Attention often maps each token independently toq,k,v and then writes directly to the state. There are two problems: the lack of local interaction between neighbouring tokens; and the sensitivity of individual token 's q,k,v to small noises. The effect of a short volume is to start with local mixing:
Current token + adjacent to token → smoother and context q/k/v write long-term status
It complements Linear Attention's local feeling, not the whole history.
- Q needs volume: Query decides "What to read now from the state." If you look only at the current token, the reading direction is easily shaking. Q Volume allows queries to be accompanied by partial context.
- K needs volume: Key decides which direction this information is written in. K is too sharp or too loud to allow neighbouring token to write in completely different directions, resulting in state fragmentation. K volume makes the writing address more stable.
- V needs volume: Value is written. If V is completely token independent, local semantics can be broken down. Volume V allows the writing to contain information on a neighbourhood, similar to organizing a phrase into a local segment before archiving.
So the volume of q/k/v is not a duplicate operation: Q smooth read; K smooth address; V smooth content.
Why do you need SiLU?
The volume itself is largely linear. If there is no non-linearness in the back, multiple local models can only be added linearly, with limited expression.
SiLU(x) = x & x can be understood as a smooth "soft switch": x is inhibited at very hours; x is released when larger and useful; the gradient is smoother than hard ReLU. It allows q/k/v to reorganise the doors and features that depend on for input before writing.
Cost: Volume and SiLU will re-introduce extra algorithms to show that Linear Attention is not "totally non-linear and local." It saves a long history of full-scale interaction, not all calculations.
Q, K, why add L2 Norm, V, why not?
The size of Q/K determines the similarity and route. If the paradigm changes too much, the model may substitute "magnify vector" for a true orientation match, resulting in logits, decay and index instability. L2 Norm brings Q/K closer to "directional resemblance":
q_hat = q / ||q||_2, k_hat = k / ||k||_2
V is not responsible for determining the weight, it is read. The range of V itself may carry information and should not be subject to forced unification.
How to save volume status while Decode
In regular Decode, only a new Token is usually handled every time. To calculate the short volume output of this Token, you need to know the conv kernel-1 volume input before it.
Therefore, each request requires only the most recent:
conv_kernel - 1
A "qkv before volume".
In particular, two types of data need to be distinguished:
- cache is the original qkv before the volume;
- The output after volume will not be written again in conv state.
The reason for this is that next step volume requires recalculating the new original qkv with the original qkv in the historical window to the volume core, rather than superimplating the already accumulated result as an input.
The shape of the entire cache pool can be written as:
[num_cache_slots, 3 × num_heads × head_dim, conv_kernel - 1]
Of which:
- Num cache slots is the number of physical slots that can be accommodated by the cache pool;
- 3 for q, k, v;
- Num heads x head dim is the total number of corridors for each vector;
- conv kernel -1 is the length of the historical window that each request needs to keep.
Each request will not have a separate dynamic array directly, but will find its own physical slot via the cache index. When the new Token arrives, the system will use the requested slot to read the latest volume input, calculate the current volume output and place the new pre-volume qkv in the window for subsequent Decode.
Therefore, the functions of S are different from those of Short Volume and Recursive Status:
- conv state only saves some recent local raw inputs for Token;
- Recurrent state compresses longer history into a fixed-sized global memory matrix.
This explains why it is easier to roll back when decoded: the volume state saves a separate local vector window, and S has mixed multiple Tokens into the same matrix.
4.7.7 Counterpartists
- Local volumes provide very limited "location" and long-range position information is based entirely on decaying hidden codes.
- If failure occurs, error is assumed to accumulate enough to encode any position relationship.
- Winner: a task that is highly semantic and less precise in its location.
- Steady substitution: MLA layer plus a mixed variant of RoPE (a paper also evaluated Kimi Linear-RoPE baseline).
- Wounded party: precise coordinates/time stamp reference mission.
4.8 Step over time
Let's start by looking at two single steps in the code.
4.8.1 KDA Version
The KDA door control uses two types of realization in parallel with the reasoning/strength and training/chunk, corresponding to the different paper stages:
-v1 (Kimi Linear original): door unbound softplus: g = -exp (A log) * softplus (gate logit + dt bias)
g No lower bounds (can be very negative), alpha = exp(g) can approach 0 at will. This is the original form of the Kimi Linear technical report and the default path for SGLang.
-v2 (K3 bounded): gate control bounds sigmoid: g = g min * sigmoid (exp (A log)* (gate logit + dt bias))
g Strictly between (g min, 0), so alpha is within (exp(g min), 1). This is the final form of the K3 technical report and the actual version of the Kimi k3 module of this project.
Both have been achieved in the project, switched by Kimi linear step (..., version= "v1"/ "v2 " ). Kimi linear step default (also SGLang default) is KDA v1; v2 is the final version of K3 'undersecurity'.
4.8.2 kimi_linear_step vs KimiDeltaAttention
Let's explain the difference and connection between Kimi linear step and 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 alignment note]
# SGLang
#used sigmoid gating delta rule update (L274-288) and used recurrent.py
# IS KDA=True Branch (L94-107) uses exactly the same time series:
#a) b h* = exp(g k) # predecided (per-K channel door control)
# b) b v - = sum(b h * b k) # calculate error using decay status
# c) b_v *= b_beta
# d) b_h += b_v * b_k
# e) b o = sum(b h*q scaled) q times scale
# Door-controlled mode of activation (SGLang kda.py L1001-1006):
# v1: g = -exp (A log) * softplus (raw + dt bias) (default)
# v2: g = lower bound * sigmoid (exp(A log)* (raw + dt bias))
# There are two modifications to this function:
# [1] Door control plus dt bias entry (the original version missing the bias, both SGLang).
# [2] Read zoom scale = Dk**-0.5 (absorption of original version, value validation indicates output difference 1/scale ≈ 2.83 times).
# [3] Decay time sequence precedes error search (same as gaed 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 died to Gate input (SGLang have)
g min (float) — lower bounds of v2 door control (degrading factor alpha max = ext (g min))
version - "v1": softplus door control, g no lower bounds (unlimited decay, SGLang default)
"v2": sigmoid door control, g ∈ (g min, 0), chunkwise numerical security (K3 actually used)
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)
# Support dt bias (kda.pyL1003/ / used recurrent.pyL479,L486)
# Degrades to zero bias when missing, without prejudice to the original caller (the test is not passed without dt bias).
z = gate_logit
if dt_bias is not None:
z = z + dt_bias
if version == "v1":
# v1: Softplus door control, no bottom line. Softplus no upper bounds g can be very negative → alpha = ext(g) can approach 0 at any time
g = -torch.exp(A_log).view(1, -1, 1) * F.softplus(z)
elif version == "v2":
# v2: Sigmoid Door Control, There's the bottom line. Sigmoid (0,1) g∈ (g min, 0),
# Equivalent alpha ∈ (exp(g min), 1) with a hard upper bound for the decay factor to protect the chunkwise from cumulative spills.
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)
# [Coordinate SGLang] Read zoom (used recurent.py L88), standard softmax attachment
scale = q.shape[-1] ** -0.5
q = q * scale
# [Coordinate SGLang] Decline first, then calculate the error (used 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
The relationship is different in abstraction: Kimi linear step is the smallest incisor sub-function, KimiDeltaAttention is the complete neural network layer, with an internal call to the lower kda step/kda chunkwise.
Contacts (mathematic core consistent)
The thrust core of the two is the delta rule + channel-wise decay of Eq.1: decay, re-research, re-correction, final reading. The mathematical form of the status update is identical.
KIMI-K3-Eq1
Verification: Compare Kimi linear step (version = "v2") with kda step in kda ref.py under the same input - max diff for Status Matrix S is 1.11e-16 (machine accuracy), indicating that the state of both is up to full parity.
Variance
One sentence: Kimi linear step is a teaching demonstration of KDA's "Mathematic Core" and KimiDeltaAttention is a complete version of KDA's "Project" -- projection, collage, integration, output door, cache managed all the fully packaged operational layers. The mathematical core is the same and the outer layer covers all the infrastructure required for the complete network layer. KimiDeltaAttention is a complete neural layer, with a lower-level internal call.
4.8.3 kda step (decode form)
kda step is the "one-step decode form" of v2, i.e. only one push (not the whole sequence).
The difference between v2 and kda step is only the project envelope:
-Kimi linear step Enter Gate logit, calculate g; kda step Enter g (mode decay)
- Kimi linear step pair q with 1/ 1/ √Dk scale (SSGLang used recurrent); kda step not (scale absorbed by HeadRMSNorm+ output)
- The mathematical nature of the status update is identical.
Rationale
Decode has only 1 token per step, taking the Eq.1 one step update (no chunk structure cost). chunkwise (initial state 0) guarantees and pre-state
Auto-fetched and translated by AI Atlas. All rights belong to the original author.