用 PyTorch 从零复现 Kimi K3：开发者拆解万亿模型的工程骨架

一位开发者近日在 Reddit 的 r/MachineLearning 社区发布了《Implementing Kimi K3 from scratch in PyTorch》，尝试用 PyTorch 从零实现 Kimi K3 的核心架构与训练流程。这不是把 2.8 万亿参数模型在个人显卡上重新训练一遍，而是将 K3 的关键设计拆成可以阅读、调试和做小规模实验的模块。

这个项目的价值不在于“复刻出一个同等能力的 K3”，而在于把一篇技术报告里的架构描述，翻译成开发者能真正运行的代码。对想研究长上下文、稀疏 MoE 和新型残差结构的人来说，这比又多一个模型权重下载链接更有用。

截至 2026 年 8 月 30 日，Kimi K3 已在 7 月 27 日开放模型权重、技术报告以及 MoonEP、FlashKDA、AgentEnv 等关键基础设施。社区复现项目则把关注点进一步下沉：如果不依赖完整的工业级训练系统，K3 的核心机制究竟应该怎样在 PyTorch 中组织？

先说结论：这是架构复现，不是能力复刻

架构复现是用等价或近似代码重建模型的计算结构；能力复刻则要求训练数据、算力、优化细节和后训练流程都足够接近原模型。 Reddit 项目更接近前者，不能据此宣称个人开发者已经“训练出了 Kimi K3”。

K3 的完整规模包括 2.8 万亿总参数、896 个专家、每个 token 激活 16 个专家、100 万 token 上下文窗口，以及原生视觉理解能力。即使 MoE 只激活少数专家，训练这样一个模型仍然需要大规模数据、数千张或更多高端加速卡、稳定的专家并行通信，以及复杂的故障恢复和检查点系统。

社区实现真正适合做的是缩小实验：把专家数量从 896 降到 8 或 16，把激活专家数从 16 降到 1 或 2，把隐藏维度和层数压缩到单机可承受范围，再验证损失曲线、路由分布、KV Cache 占用和长上下文行为是否符合预期。

这一区分很重要。很多“从零实现大模型”的内容，最后只是写出了一个 Transformer 类；而 K3 的难点不在于把几个 Linear 层串起来，而在于注意力、残差、专家路由、量化感知训练和分布式执行之间的耦合。

Kimi K3 的核心规格

Kimi K3 是一个采用混合注意力与稀疏混合专家架构的 2.8 万亿参数开源模型。 官方资料显示，K3 基于 Kimi Delta Attention（KDA）和 Attention Residuals（AttnRes）构建，支持原生视觉理解和 100 万 token 上下文窗口。

| 项目 | Kimi K3 公开信息 | 对复现的影响 | |---|---:|---| | 总参数量 | 2.8 万亿 | 个人环境只能做缩小版实验 | | 架构 | 稀疏 MoE | 需要实现路由、容量控制和专家并行 | | 专家数量 | 896 个 | 小规模复现可先使用 8～32 个 | | 每 token 激活专家 | 16 个 | 决定单 token 的计算量和通信量 | | 上下文窗口 | 100 万 token | 需要 KDA、缓存和长序列训练支持 | | 注意力结构 | 3:1 混合 KDA 与 Gated MLA | 不能简单替换成标准全注意力 | | 残差结构 | Attention Residuals | 需要保存并选择性读取跨层表示 | | 量化训练 | MXFP4 权重、MXFP8 激活 | 从 SFT 阶段开始进行量化感知训练 | | 公开时间 | 2026 年 7 月 27 日开放权重 | 目前已有官方资料和社区复现可参考 |

K3 的总参数量很大，但单 token 的实际计算量并不等于把 2.8 万亿参数全部跑一遍。MoE 的基本思路是把模型拆成多个专家，每个 token 只进入其中一小部分专家。K3 在 896 个专家中激活 16 个，按专家数量计算的激活比例约为 1.79%。不过，这个数字不等于端到端成本的降幅，因为路由、跨设备通信、共享模块和全局注意力仍然需要付出代价。

第一层拆解：KDA 为什么适合百万 Token

Kimi Delta Attention（KDA）是一种通过递归状态压缩历史信息的线性注意力机制。 与标准全注意力需要让当前 token 访问整个历史序列不同，线性注意力会把历史压缩成相对固定大小的状态，因此在超长序列下不必为每个新 token 保存完整的键值矩阵。

标准 Transformer 的注意力计算可以简化理解为：当前 token 拿着 Query，逐个查看历史 token 的 Key 和 Value。序列越长，KV Cache 越大，解码阶段的显存和访存压力也越高。KDA 则更像一个带门控的高速记事本：模型不会永久保存所有原文，而是持续更新一组状态，并决定哪些信息应该写入、保留或覆盖。

KDA 的关键改动是通道级门控。此前一些 DeltaNet 类方法更多以注意力头为单位控制更新，KDA 将门控细化到特征维度，允许同一个注意力头中的不同通道采用不同的记忆策略。对于代码、表格和长文档，这种粒度更细的状态更新理论上更适合处理不同类型的信息。

但 K3 并没有把所有层都换成线性注意力。其公开设计采用 3:1 混合布局：每 4 层中，3 层使用 KDA，1 层保留带门控的 MLA 全局注意力。这样做的取舍很直接：KDA 负责低成本地持续积累信息，全局注意力负责定期重新查看远距离 token，避免递归状态在长序列中丢失关键细节。

公开资料引用的 Kimi Linear 研究结果显示，类似配置最多可减少 75% 的 KV Cache；K3 方面则给出了在 100 万 token 上下文下最高 6.3 倍解码速度提升的说法。具体收益会随硬件、批大小、序列长度和缓存实现变化，不能简单理解为所有场景固定加速 6.3 倍。

在 PyTorch 复现时，最容易犯的错误是把 KDA 写成一个“少算几次矩阵乘法”的普通注意力层。真正需要验证的是递归状态更新是否保持因果性、门控是否逐通道生效、状态在 chunk 切分后是否与整段计算一致，以及训练和推理模式下的数值结果是否稳定。

第二层拆解：AttnRes 重做了残差连接

Attention Residuals（AttnRes）是一种用注意力选择前序层表示的残差机制。 传统残差连接通常把输入和当前层输出直接相加，即“上一层状态加上这一层增量”；AttnRes 则允许当前层根据输入内容，从更早的多个层表示中选择性读取信息。

标准残差的优点是简单、稳定、易于并行，但它对所有历史层采用了固定的加法关系。随着网络变深，隐藏状态不断累积，某些层的有效贡献可能被稀释。AttnRes 的思路是把“均匀求和”换成“带权读取”：不同层的表示先经过投影或压缩，再由一个 Softmax 权重决定当前层更应该依赖哪些深度的状态。

从实现角度看，AttnRes 不只是把 x + residual 改成另一个加法。开发者需要处理三个问题：第一，前序层输出如何保存，显存是否会随层数线性增长；第二，注意力权重如何计算，是否需要低秩投影或块级聚合；第三，训练时如何避免所有层都偏向最近的一个表示。

K3 的公开介绍将 AttnRes 描述为块级注意力残差，并提到相关 48B 研究模型在不到 2% 额外计算成本下获得约 25% 的训练效率提升，也就是约 1.25 倍效率。这个结果来自特定模型、数据和训练设置，不应直接外推到所有网络，但它说明残差连接仍然可能是提升大模型扩展效率的杠杆。

对于个人实验，建议先实现一个简化版本：把若干 Transformer block 分成一个残差块，只保留每个 block 的压缩表示，再用轻量的内容相关权重进行融合。先观察深度增加后训练损失是否更平滑，再考虑完整的跨层检索和显存优化。这样比一开始照搬万亿模型的所有工程细节更容易定位问题。

第三层拆解：Stable LatentMoE 的难点是路由稳定

Mixture of Experts（MoE）是将模型中的前馈网络拆成多个专家，并由路由器为每个 token 选择少数专家的稀疏架构。 MoE 能在保持较低激活计算量的同时扩大总参数量，但它的瓶颈通常不是专家层本身，而是路由是否均衡、通信是否高效、训练是否稳定。

K3 使用 Stable LatentMoE，在 896 个专家中激活 16 个，并结合 SiTU-GLU 与 Quantile Balancing。SiTU-GLU 可以理解为一种带 Sigmoid-Tanh 控制的门控激活结构，用于增强专家内部的非线性与信息筛选；Quantile Balancing 则不完全依赖一个容易失控的固定负载损失，而是根据路由分数的分位数来分配 token。

普通 Top-K 路由在训练早期可能出现“赢家通吃”：少数专家被大量 token 访问，其他专家几乎没有梯度。结果是热门专家容量溢出，冷门专家学不到东西，系统还要为丢弃 token、填充容量和跨卡搬运付出额外成本。

分位数平衡的直觉类似把考生按成绩分桶，而不是规定所有桶必须使用同一个阈值。不同 batch、不同阶段的路由分数分布可能变化，动态分位数比固定阈值更能适应训练过程。复现时不能只看总 loss，还应该记录以下指标：

• 每个专家接收的 token 数量；
• 专家负载的最大值、最小值和变异系数；
• token 被丢弃或二次路由的比例；
• Top-1 与 Top-K 路由的一致性；
• 路由辅助损失占总损失的比例；
• 不同设备之间的 all-to-all 通信时间。

如果一个小模型的训练 loss 在下降，但专家负载长期集中在少数几个专家，说明它还不能算稳定的 MoE 实现。模型可能暂时学会了，但扩展到更多专家、更大 batch 或更长序列后会迅速失效。

从 PyTorch 代码看，最小复现应该包含什么

社区教程的实际意义，是把 K3 的模块拆成可以独立替换的组件。一个可读、可测的实现至少应包含：词元嵌入、混合注意力层、AttnRes、MoE 前馈层、因果掩码、KV 或递归状态缓存，以及训练和评估脚本。

下面是一个简化的结构示意，不代表 Kimi K3 的完整官方实现，也没有包含分布式并行、量化内核和视觉编码器：

class K3Block(nn.Module):
    def __init__(self, cfg, layer_id):
        super().__init__()
        self.norm1 = RMSNorm(cfg.hidden_size)
        self.attn = KDA(cfg) if layer_id % 4 != 3 else GatedMLA(cfg)
        self.norm2 = RMSNorm(cfg.hidden_size)
        self.moe = StableLatentMoE(cfg)

    def forward(self, x, state=None, previous_layers=None):
        h = self.attn(self.norm1(x), state=state)
        x = self.attn_residual(x, h, previous_layers)
        h = self.moe(self.norm2(x))
        return self.attn_residual(x, h, previous_layers)

这段代码只表达模块关系。真正可训练的版本需要补上参数初始化、位置编码、路由容量、状态缓存、混合精度、梯度裁剪和检查点恢复。尤其是 KDA 的 state 不能被当作普通的 KV Cache 直接处理；在 chunk-by-chunk 推理时，前一个 chunk 的递归状态必须准确传递到下一个 chunk。

训练流程也不应只写成一个大循环。一个更合理的实验拆分是：

1. 单模块单元测试：验证 KDA 的因果性、AttnRes 的权重归一化和 MoE 的路由结果。
2. 小数据过拟合：用几百条样本让模型快速过拟合，确认反向传播、缓存和 loss 计算没有错误。
3. 短上下文预训练：先使用 2K 或 4K token，观察语言建模损失与专家负载。
4. 长度扩展实验：逐步增加到 16K、64K，再测试状态传递、显存占用和吞吐变化。
5. 监督微调：从 SFT 阶段引入 MXFP4 权重和 MXFP8 激活的量化感知训练近似方案。
6. Agent 与编程评估：加入工具调用、终端任务、代码修复和多轮规划，而不是只测通用问答。

训练流程：架构只是起点

大模型训练流程是从预训练、监督微调到强化学习和评估的连续工程链路。 K3 的公开资料提到，模型在通用推理、通用 Agent 和编程 Agent 三个方向进行大规模任务合成，并建设了面向百万 token 上下文的强化学习基础设施。

预训练阶段决定模型的基础知识、代码能力和长文本建模能力。SFT 阶段则把模型从“能续写”推向“能按照任务要求完成工作”。对于编程 Agent，训练样本不能只有代码补全，还需要包含终端操作、测试失败、错误修复、文件修改和任务验收等完整轨迹。

强化学习阶段的难点是环境。传统语言模型只需要对下一个 token 计算损失，Agent 则需要真的运行代码、访问沙箱、等待测试结果并根据反馈继续行动。AgentEnv 的意义就在于提供这类可扩展环境，让训练过程从静态文本变成可验证的任务执行。

个人开发者不必一开始就搭建完整的 Agent 强化学习集群。一个更实际的路径，是先用本地沙箱做可重复的代码修复任务：给模型一个有测试用例的仓库，让它修改文件、运行测试，再把测试通过率作为奖励信号。这样可以验证规划、工具调用和错误恢复是否有效。

量化与并行：为什么“能跑”不等于“能训练”

量化感知训练（QAT）是在训练过程中模拟低精度计算误差，让模型提前适应部署时的量化格式。 K3 从 SFT 阶段开始使用 MXFP4 权重和 MXFP8 激活，目标是让模型更适合不同硬件环境，而不是训练结束后再简单地把浮点权重压缩一遍。

4 bit 权重可以显著降低显存占用，但专家模型的权重访问本身就不规则。若量化格式没有匹配硬件矩阵乘法，理论上的显存节省可能换来更高的解码延迟。因此，复现项目在消费级 GPU 上可以先使用 BF16 或 FP16 验证数值正确性，再逐步引入模拟量化；不要把“加一个 fake quant 函数”误认为已经实现了 MXFP4 训练。

并行方面，K3 官方建议在由 64 个或更多加速器组成的 supernode 配置上部署。原因不只是参数太多，还包括专家并行中的 all-to-all 通信，以及长上下文下注意力状态的高带宽访问。MoonEP 面向专家并行通信，FlashKDA 面向 KDA 算子，AgentEnv 面向分布式强化学习环境，它们分别对应模型训练链路中的不同瓶颈。

这也是社区 PyTorch 实现与官方系统之间最大的差距：前者强调可读性，后者强调吞吐、通信和容错。一个单卡版本即使能在功能上跑通，也不能据此判断在 64 卡或 256 卡集群上能否线性扩展。

和标准 Transformer、其他编程模型相比，K3 到底强在哪

K3 的优势主要集中在长程编程、系统级任务和 Agent 工作流，而不是所有短文本问答都一定领先。公开资料提到，K3 在 Terminal Bench 上取得 88.3 分；这一类测试更接近命令行操作、环境配置、Dockerfile、Kubernetes 和 CI/CD 等任务。

| 模型或工具 | 主要优势 | 长上下文/Agent特点 | 更适合的场景 | |---|---|---|---| | Kimi K3 | KDA、AttnRes、稀疏 MoE | 100 万 token，上下文与 Agent 协同 | 长程编程、终端任务、复杂知识工作 | | GitHub Copilot | GitHub 代码生态和语言覆盖 | IDE 内补全与仓库级辅助 | 日常代码补全、团队开发 | | Cursor | 编辑器集成和代码库索引 | Composer、Agent 工作流 | 全栈项目、快速迭代 | | Claude Code | 多轮推理与子 Agent | 适合拆解复杂编程任务 | 代码审查、重构、自动化开发 | | 标准稠密 Transformer | 结构成熟、工具链丰富 | 长上下文成本随 KV Cache 增长 | 教学、基线和中小模型训练 |

K3 的长上下文并不意味着它能无损记住 100 万 token 中的每个细节。上下文窗口是模型可以接收的最大序列长度，实际有效利用率还取决于信息位置、任务结构、检索能力和训练分布。对开发者而言，更值得测试的是：模型能否定位跨文件依赖，能否记住早期约束，能否在多轮修改后保持接口一致。

这个复现项目适合谁，不适合谁

它适合三类人：第一类是想理解线性注意力和 MoE 内部机制的研究者；第二类是需要自己改造模型结构的训练工程师；第三类是希望搭建长上下文代码 Agent、但不想只停留在调用现成模型层面的开发者。

它不适合把教程当作低成本训练 K3 的捷径。没有高质量数据、分布式训练经验和足够算力，复现出来的模型最多用于验证机制，无法接近原始 K3 的通用能力。更现实的目标是做一个几亿到几十亿参数的缩小模型，比较 KDA 与全注意力在长文本上的速度和质量，再研究 AttnRes 或 MoE 是否值得保留。

建议的学习顺序是：先实现标准 Transformer 作为基线，再替换一部分层为 KDA；之后加入简化 AttnRes，最后引入 MoE 和路由均衡。每次只改一个变量，并固定数据、训练步数和评测集。否则当 loss 变化时，很难判断收益来自注意力结构、残差结构，还是仅仅来自参数量增加。

OpenAI Hub 判断：最有价值的是“拆解”，不是“复刻”

这次社区实现的意义，在于它把 K3 从一个“2.8 万亿参数、100 万上下文”的宣传规格，变成了一组可以被开发者逐项验证的工程问题：KDA 能否降低长上下文缓存压力，AttnRes 是否改善深层信息流，Stable LatentMoE 能否在极高稀疏度下保持均衡，以及量化和并行是否真的带来吞吐收益。

K3 最值得关注的地方也不是参数量最大，而是它同时改动了注意力、残差和训练基础设施。传统 Transformer 的很多默认组件已经成为工程习惯，但 K3 说明这些组件并非不可替代。尤其对长程编程和 Agent 来说，模型能否稳定处理跨文件、跨工具、跨小时的任务，可能比短答案基准上的几个百分点更重要。

不过，开发者需要保持边界感：社区 PyTorch 代码可以帮助理解架构，不能替代官方权重、训练数据和分布式 Infra。把小模型跑通是第一步；用严格消融实验验证每个改动，才是这类复现项目真正的技术含金量。

截至今天，最值得动手的实验不是“训练一个 K3”，而是用统一配置比较三组结果：标准全注意力、3:1 KDA 混合注意力，以及加入 AttnRes 后的混合结构。记录 4K、16K、64K 和更长上下文下的显存、吞吐、损失与代码任务准确率，你会比单纯阅读模型规格更快理解 K3 为什么这样设计。

参考来源

• Reddit：Implementing Kimi K3 from scratch in PyTorch：本文讨论的社区 PyTorch 复现项目及开发者原帖。
• IT之家：K3 开源重塑竞争逻辑：介绍 K3 在编程工具、Terminal Bench、模型适配和 Agent 工作流中的定位。
• 知乎：Kimi K3 技术解析：对 KDA、AttnRes、Stable LatentMoE、量化训练和长上下文设计的第三方拆解。
• 腾讯云开发者社区：Kimi K3 架构详解：补充 KDA 与 AttnRes 的结构背景，以及社区可验证的研究 checkpoint 信息。

注：Kimi K3 的完整模型规格、权重和 Infra 细节应以月之暗面官方技术报告及开源仓库的最新版本为准。社区复现代码可能采用缩小配置或近似实现，不代表官方训练实现。