轻量化复刻KIMI K3：AttnRes核心深度解析

<p style="margin: 10px 0;">本文是Attention Residuals（以下简称AttnRes）系列内容的第二部分。相关探索项目尝试基于公开资源实现轻量化大模型复刻，涉及KIMI K3、DeepSeek V4等模型，且全部采用CPU运行，不依赖GPU或大规模训练环境。</p>

<div style="margin: 15px 0; padding: 10px; background: #f5f5f5; border-radius: 5px;">
    <p style="margin: 8px 0;">关于项目的补充说明：</p>
    <ul style="margin: 8px 0 0 25px; padding: 0;">
        <li style="margin: 5px 0;">项目及系列博客是探索与学习的过程，会持续改进，也欢迎反馈意见。</li>
        <li style="margin: 5px 0;">模型组件尽量参考原始大模型设计，无开源版本则用近似组件替代。</li>
        <li style="margin: 5px 0;">每个组件会讲解解决的问题、难点、解法原理及组件间的协作逻辑，核心是理解主线而非记忆名词。</li>
        <li style="margin: 5px 0;">主要组件包含可直接运行的演示工具，还有可视化脚本可生成分析图。</li>
        <li style="margin: 5px 0;">代码行号基于本地快照，后续变化不影响函数定位。</li>
        <li style="margin: 5px 0;">为保证文章独立性，部分内容会在系列中重复，还请理解。</li>
    </ul>
</div>

<h3 style="font-size: 18px; font-weight: bold; margin: 25px 0 15px;">子知识点拆解</h3>
<p style="margin: 10px 0;">AttnRes可拆解为多个核心子知识点，各部分存在明确依赖关系：</p>
<ul style="margin: 10px 0 20px 25px; padding: 0;">
    <li style="margin: 5px 0;">S1（打分核）：核心计算单元，为后续逻辑提供基础。</li>
    <li style="margin: 5px 0;">S2（Full AttnRes）：无分块的完整形式，保留所有历史层。</li>
    <li style="margin: 5px 0;">S3（Block AttnRes）：分块优化形式，降低内存开销。</li>
    <li style="margin: 5px 0;">S4（状态机）：跨层状态管理，解决内存约束下的选择性读取问题。</li>
    <li style="margin: 5px 0;">S5（在线softmax两阶段算法）：降低推理计算开销的优化逻辑。</li>
    <li style="margin: 5px 0;">S6（生产级融合内核）：面向实际部署的高效实现。</li>
    <li style="margin: 5px 0;">S7（对比对象mHC）：横向对比的参考方案。</li>
</ul>
<p style="margin: 10px 0;">各子知识点的实现可通过参数切换，比如是否使用Full或Block形式、块大小、计算粒度以及加速算法的选择，这些配置可在项目的相关配置文件中调整。</p>

<h3 style="font-size: 18px; font-weight: bold; margin: 25px 0 15px;">深度分析</h3>

<h4 style="font-size: 16px; font-weight: bold; margin: 20px 0 10px;">共享抽象与设计动机</h4>
<p style="margin: 10px 0;">所有变体共享统一的核心抽象逻辑，即基于一组“源”集合，通过一次“查询-键打分+归一化混合”完成计算。这里的“源”是历史层输出的集合，会被显式保留而非压缩为单一状态，这是AttnRes与标准残差连接的关键区别——标准残差会将历史状态压缩为单一向量，而AttnRes保留所有历史项，从而支持选择性读取。设计这个逻辑的核心动机是，选择性读取需要保留每一项的信息，而保留全部历史的内存开销，正是后续Block形式要解决的问题。</p>

<h4 style="font-size: 16px; font-weight: bold; margin: 20px 0 10px;">打分核（scoring kernel）与伪查询</h4>
<p style="margin: 10px 0;">打分核是AttnRes的核心计算单元，每层对应一个伪查询向量，对所有源计算打分后通过Softmax归一化得到权重。</p>
<p style="margin: 10px 0;">从原理来看，有三个视角：</p>
<ul style="margin: 10px 0 20px 25px; padding: 0;">
    <li style="margin: 5px 0;">直觉视角：每个源对应“词汇”，每层对应“搜索的查询”，伪查询向量表示当前层的读取偏好，RMSNorm会将每个源的向量缩放到统一长度，保证比较的是方向而非量级。</li>
    <li style="margin: 5px 0;">准专业视角：权重计算是将伪查询与经过RMSNorm的源向量做点积后取指数，再归一化得到分布，伪查询向量是可学习的层级参数，训练时会学习出该层的读取偏好。</li>
    <li style="margin: 5px 0;">专家视角：这是一种内积点积核，保证权重非负且和为1，即每层对历史层的概率分布。RMSNorm的作用经过实验验证，去掉它会导致性能下降，因为它防止了范数大的源垄断权重，这对Block形式尤为重要。</li>
</ul>
<p style="margin: 10px 0;">与标准残差连接对比，打分核的权重来源是每层的可学习向量，而非固定常数1；它的权重是输入相关的，会受源内容的影响；额外参数是每层的伪查询向量；同时会结合Softmax和RMSNorm做归一化。</p>

<h4 style="font-size: 16px; font-weight: bold; margin: 20px 0 10px;">Full AttnRes</h4>
<p style="margin: 10px 0;">Full AttnRes是无分块的形式，所有历史层都可读，相当于“每一页都是独立卡片，可自由选择读取”。从准专业角度，它的计算单元数量等于层数或子层数，对所有早期层和embedding做混合；从专家视角，它的计算复杂度是O(L²d)，在层数小于100时可负担，瓶颈是内存开销，需要保存所有层的输出，这在并行训练中会带来较大压力。</p>
<p style="margin: 10px 0;">与Block形式相比，Full AttnRes的内存开销是O(Ld)，而Block形式是O(Nd)；算术复杂度上Full AttnRes是O(L²d)，Block形式是O(LNd)；精度上Block形式接近Full，针对性面向内存和通信优化。</p>

<h4 style="font-size: 16px; font-weight: bold; margin: 20px 0 10px;">Block AttnRes</h4>
<p style="margin: 10px 0;">Block AttnRes是为解决内存问题提出的分块形式，将所有层分成N个块，每个块收官时压缩为一个总表示，读取历史时只读已完成的块和当前块的运行部分。从准专业角度，每个块的表示是块内各层输出的求和，当前块会保留运行中的部分和；从专家视角，它的内存开销是O(Nd)，介于等权相加和Full形式之间，实验显示当块数约为8时，能保留大部分性能，K3模型就是采用8个块，每个块12层，对应9个源。</p>
<p style="margin: 10px 0;">Block形式的算术复杂度是O(LNd)，对比Full AttnRes的O(L²d)，当层数较大时，Block形式能大幅降低计算开销。</p>

<h4 style="font-size: 16px; font-weight: bold; margin: 20px 0 10px;">状态机（AttnResState）</h4>
<p style="margin: 10px 0;">AttnResState是跨层的状态机，解决内存约束下的选择性读取问题，核心是用可变状态对象，记录已完成的块、当前块的运行和、当前块的子层数，状态机的演进过程是初始化、逐子层处理、满块折叠、最终聚合，这样的结构避免了手写状态机的错误，同时降低了内存开销。</p>
<p style="margin: 10px 0;">如果不用状态机，会面临几个问题：内存爆炸（Full AttnRes的O(Ld)内存）、退化为均匀残差（失去选择性）、手写状态机容易出错、可解释性差等。</p>

<h4 style="font-size: 16px; font-weight: bold; margin: 20px 0 10px;">在线softmax两阶段算法</h4>
<p style="margin: 10px 0;">在线softmax两阶段算法是为了降低推理时的计算开销，把softmax拆成跨块和块内两部分，跨块部分可并行计算，块内部分逐token处理，这样在解码时的计算成本大幅降低。对比eager模式，两阶段算法能复用跨块的统计量，减少重复计算，提升效率。</p>