线性注意力并行机制解析：KDA的设计与实现

本文是关于线性注意力架构中并行计算模块的分析，因篇幅较长分为多期，本期聚焦并行相关的核心设计。

并行并非没有代价：循环神经网络（RNN）的状态递推对输入顺序敏感，交换Token位置会改变后续状态；而注意力的聚合本身不天然包含顺序信息，固定Query时，Key-Value对的重新排列不会改变单次聚合结果。

所讨论的注意力架构（KDA）同时获得两大特性：相对上下文长度为O(1)的状态存储与解码计算能力，以及chunk级并行能力。其核心是将旧状态的递归更新转化为线性或仿射变换，带来表达能力与计算效率的权衡。

局部混合机制：ShortConv、SiLU与L2归一化

线性注意力擅长将长历史压缩到固定大小的状态中，但压缩可能丢失细节，尤其是相邻Token间的短距离关系无法仅靠全局状态充分保留。因此，相关架构通常会在投影得到的q、k、v上增加一维因果短卷积，处理流程分为四步：投影得到原始q/k/v，用因果卷积融合局部信息，将卷积后的结果送入线性注意力的状态更新，由固定大小的递归状态聚合全局历史。

这类架构将局部与全局机制结合：短卷积强化相邻Token的细节与短程模式，线性注意力压缩历史到固定状态，可概括为“卷积补充局部关系，线性注意力聚合全局历史”。

原理

直觉

该架构未显式加入位置编码，如何感知相邻Token关系？ShortConv是宽度为4的因果深度卷积，让每个通道在递推前“查看”最近3个Token，注入局部位置信息。

每个通道类似一个3格的小记忆条（kernel=4对应当前+前3格），新Token到来时，将当前值与前3个值按权重混合作为输出，仅记得最近4个Token，擅长捕捉局部模式，不负责长程依赖，长程由后续递推处理。

因此，这个小记忆条成为模型感知顺序/位移的主要局部通道。

准专家视角

短卷积可理解为对q、k、v的每个通道使用固定长度滑动窗口，按相对位置加权求和，即对每个通道独立做因果1D卷积，kernel=4，无跨通道混合，输出在递推前被局部平滑。它仅提供移位信息，配合delta rule的位置感知衰减，替代RoPE的位置编码。

输入输出方面，输入为x（隐维/通道数），卷积核为W（k=4，depthwise即逐通道），输出为卷积后的特征。假设卷积核长度为4，当前输出会使用前第3、2、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])，核心有两点：窗口长度固定，局部计算成本不随上下文增长；卷积权重是训练得到的固定值，不会像注意力那样动态计算匹配分数。

因果互相关公式为y[t,c] = sum_{j=0}^{k-1} W[c,0,j] * x[t-(k-1-j),c]，仅依赖当前与过去，左侧补零保证y[t]只依赖x≤t。

专家视角

无位置编码、局部卷积、递推衰减是一套完整的位置信息方案：局部用卷积，全局用衰减序。该架构的混合模式为KDA保留顺序、近因和last-write-wins，NoPE MLA提供顺序中立的集合聚合和全局内容交互，二者互补，任务需要顺序时用KDA的偏置，需要多证据聚合时用MLA减少近因偏置。

NoPE并非解决全部顺序问题：NoPE仅意味着不加入显式位置旋转，顺序仍可能来自因果mask、ShortConv、上游KDA等，不能等同于模型完全不理解顺序。此外，NoPE避免两套位置系统竞争，KDA产生的内容相关相对距离与RoPE的固定旋转距离可能方向不一致，因此NoPE的含义是避免重复或冲突的位置归纳偏置。

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

短卷积和滑窗注意力都聚合附近Token，但权重来源不同：短卷积使用预先学习的固定权重，相同相对位置用同一组参数，计算为固定窗口内的逐通道加权求和；滑窗注意力根据当前内容动态计算权重，Query与窗口内Key匹配后决定Value的权重，特点是内容自适应但需要额外匹配计算。

二者对比：机制方面，短卷积是固定长度窗口，滑窗注意力也是固定长度窗口；权重来源，短卷积是训练得到的固定卷积权重，滑窗注意力是根据Query-Key动态计算；特点，短卷积计算简单、开销稳定、逐通道聚合，滑窗注意力内容自适应但需要额外匹配计算。

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

对比基线

维度对比：RoPE提供全局相对位置，需额外缩放适配长上下文；ShortConv提供局部邻域混合，无需任何处理。

计算复杂度

ShortConv的计算复杂度为O(T·d·k)（k=4），几乎可忽略。

参数敏感性

卷积核默认4，调大则局部感受野增加但位置模糊，调小则移位信息不足。

工程实现

项目代码中ShortConv类的实现：初始化置恒等，卷积起点为无操作，不扰动精心初始化的衰减动力学；decode时窗口状态由conv_*保存。代码示例为Python类，实现forward（整段/块处理，返回尾部窗口状态）和step（单token递推，用上一轮状态+当前token计算输出）。

与开源参考核对：类似mamba的causal_conv1d_fn，语义一致，初始化策略不同（本项目恒等初始化，其他项目用正态随机初始化）。

讨论

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

线性注意力的局部细节缺失问题：独立映射q/k/v写入状态，相邻Token缺少局部交互，单个Token对噪声敏感。短卷积的作用是局部混合，让当前Token+邻近Token生成更平滑、有上下文的q/k/v，再写入长期状态，弥补线性注意力的局部感受野，而非重复全量注意力。

Q需要卷积：Query决定读取方向，只看当前Token易抖动，卷积让查询带上局部上下文；K需要卷积：Key决定写入方向，尖锐或噪声大的Key会让相邻Token写入方向碎片化，卷积让写入地址更稳定；V需要卷积：Value是被写入的内容，完全独立的V会拆散局部语义，卷积让写入内容带有邻域信息，类似整理成局部片段再存档。

因此，q/k/v的卷积是互补操作：Q平滑读取，K平滑地址，V平滑内容。

为什么还需要SiLU？

卷积是线性混合，无非线性则表达能力有限，SiLU作为平滑软开关，x小时抑制，x大且有用时放行，梯度比硬ReLU更平滑，让q/k/v在写入状态前进行输入依赖的门控和特征重组。

代价是引入额外算子，说明线性注意力并非完全没有非线性和局部计算，仅省去长历史的全量交互，而非所有计算。

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

Q/K的点积决定相似度和路由，范数变化大时，模型可能用放大向量代替方向匹配，造成logits、衰减和索引不稳定，因此用L2 Norm让Q/K更接近方向相似度，公式为q_hat = q / ||q||₂，k_hat = k / ||k||₂。

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

Decode时如何保存卷积状态

普通Decode每次处理一个新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使用。

卷积状态和递归状态职责不同：conv_state保存最近几个Token的局部原始输入，recurrent_state把更长历史压缩到固定大小的全局记忆矩阵。推测解码时，conv_state比递归状态更容易回滚，因为卷积状态是独立的局部向量窗口，而递归状态已混合多个Token到同一矩阵。

反方观点

局部卷积提供的位置信息有限，长程位置信息靠衰减隐式编码，若失败则假设衰减累积足以编码任意位置关系不成立，胜出的是语义稠密、位置精度要求不高的任务，稳健替代是MLA层加RoPE的混合变体，精确坐标/时间戳引用任务会受影响。

单步递推

首先看两个单步递推实现，KDA的门控在推理/串行递推和训练/chunk并行两种形态下采用两种实现：v1用无界softplus门控，v2用有界sigmoid门控。

v1版本：门控为g = -exp(A_log) * softplus(gate_logit + dt_bias)，g无下界（可非常负，alpha=exp(g)可接近0），是技术报告原始形式，也是默认路径。

v2版本：门控为g = g_min * sigmoid(exp(A_log) * (gate_logit + dt_bias))，g严格在(g_min,0)之间，因此alpha在(exp(g_min),1)范围内，给衰减因子硬上界，防止chunkwise累积溢出，是最终形态，也是实际使用的版本。

二者由kimi_linear_step(version="v1"/"v2")切换，默认是v1，v2是安全下界版本。

kimi_linear_step与KimiDeltaAttention的差异和联系

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

联系：数学核心一致，递推核心是delta rule + channel-wise衰减，状态更新数学形式相同，对拍验证显示状态矩阵的max diff为机器精度，说明完全等价。

差异：抽象层次，kimi_linear_step是单步算子函数（教学/基准），KimiDeltaAttention是完整神经网络层（生产级）；输入，kimi_linear_step输入已投影好的q/k/v等，KimiDeltaAttention输入原始x；q/k归一化，kimi_linear_step函数内做L2 normalize，KimiDeltaAttention在_project中做；q scale，kimi_linear_step乘了1/√Dk的scale，KimiDeltaAttention不额外乘，尺度由输出门控制；输出差异，kimi_linear_step输出比kda_step小Dk^-0.5倍，KimiDeltaAttention输出尺度由输出门控制；v1/v2版本，kimi_linear_step同时支持v1和v2，KimiDeltaAttention只实现v2；ShortConv，kimi_linear_step无，KimiDeltaAttention有q/k/v各一条；SiLU激活，kimi_linear_step无，KimiDeltaAttention有投影后；beta来源，kimi_linear_step传入logit，函数内sigmoid，KimiDeltaAttention从x投影；输出门，kimi_linear_step无，KimiDeltaAttention有全秩输出门；解码缓存，kimi_linear_step无，KimiDeltaAttention有KDACache；chunkwise，kimi_linear_step无，KimiDeltaAttention有块内并行。

一句话总结：kimi_linear_step是KDA的数学内核教学演示，KimiDeltaAttention是工程完整版，包装了投影、卷积、归一化、输出门、缓存管理等基础设施，数学核心相同，外层包裹了完整网络层所需的一切。

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，kda_step不加，尺度由输出门吸收；状态更新的数学本质完全一样。

原理：decode时每步只有1个token，直接走Eq.1的单步更新，首个token无先前状态，走chunkwise初始化保证与prefill一致。计算复杂度为单token O(d_k·d_v)，一次matmul更新加一次读出。

对比基线：chunkwise适用于prefill/训练，支持块内并行；kda_step适用于decode，单步；开销，chunkwise需要三角求解，kda_step最小。

工程实现

项目代码中kda_step的实现：单token更新，decode阶段专用，输入q、k、v、g、beta、state，输出out和新状态。具体步骤：状态衰减S = alpha.unsqueeze(-1) * S，检索记忆retrieved = einsum('bhd,bhdv->bhv',k,S)，误差delta = beta.unsqueeze(-1) * (v - retrieved)，状态更新S_new = S + k.unsqueeze(-1) * delta.unsqueeze(-2)，读出out = einsum('bhd,bhdv->bhv',q,S_new)，返回out, S_new, retrieved, error。

调用处在kda.py的forward中，当t==1且state_in不为None时走单步，保证首个decode token也走chunkwise初始化，否则用chunkwise并行形式。

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

chunkwise并行

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

原理

直觉：递推天然串行，训练需并行，chunkwise把序列切成块，块内用矩阵运算并行算出所有贡献，块间串行传递状态。

专家视角：块长C，进入块的状态S_0，输出O = (Γ ⊙ Q) S_0 + A Ṽ，其中I+M的变换需三角求解。

这里Γ是块内累积衰减，Q̄=ΓQ、K̄=ΓK、K̂=K/Γ，A=tril[(Q̄)(K̂)ᵀ]是块内注意力矩阵，M=tril(Diag(β)K̄K̂ᵀ,-1)，(I+M)是单位下三角，可精确三角求解。

这是DPLR特例的分解：对角部分逐元素累积，低秩部分用稠密matmul，关键是衰减有下界，所有tile可用Tensor-Core稠密乘。

chunkwise形式的推导

单Token状态更新公式为S_t = M_t S_{t-1} + β_t k_t v_t^T，块内串行更新组合后，整个Chunk仍可表示为仿射变换S_{[i+1]} = M_{[i+1]} S_{[i]} + S̃_{[i+1]}。

反方观点

局部卷积提供的位置信息有限，长程位置信息靠衰减隐式编码，若失败则假设衰减累积足以编码任意位置关系不成立，胜出的是语义稠密、位置精度要求不高的任务，稳健替代是MLA层加RoPE的混合变体，精确坐标/时间戳引用任务会受影响。