Kimi K3推理架构深度解析：线性注意力与快照策略优化

我们可以从推理维度对Kimi K3模型展开详细分析。这款模型参数规模达2.8T，采用FP4量化时显存占用约1.4TB，适配Hopper架构的表现尚可，但该架构本质是为B卡定制的，目前尚不确定国产加速卡能否适配这一架构，若无法适配则说明国产卡在相关视野上存在不足。同时不难发现，不少市场参与者对推理链路的商业流转理解得更为透彻。

Kimi K3的核心架构由KDA和Attn Res组成，其中KDA对应线性注意力机制，模型内部还融合了普通注意力，比如MLA结构并额外添加了Gated模块，整体属于线性注意力与普通注意力混合的架构。这种设计会给KV缓存管理带来一定挑战，官方团队也表示将为vllm项目提交针对prefix cache的PR，后续可以从openinfer的视角进一步探讨相关优化方向。

模型采用了896个专家网络，每次仅激活16个，这一比例符合常见的2倍激活设计，但具体的最优EP参数需要结合实际模型shape进行压测验证，由于目前尚未公布模型权重，暂时无法展开详细测算。此外，LatentMoe模块的作用是压缩MoE架构中的all-to-all通信开销，不过参考Deepseek的MegaMoE技术方案，计算过程可以与通信过程重叠，此时通信速度的影响相对有限，但关键在于如何规划计算与通信的重叠方式——MegaMoE采用按专家Wave切分的方式实现重叠，若EP规模过大则可能无法实现这种切分。

Kimi K3支持1M长度的上下文窗口，目前暂未明确其实际KV缓存占用规模，预计整体占用不会过高。针对线性注意力架构，其prefill和decode阶段的FLOPS需求、带宽要求等参数，都需要通过profile工具进行实际测算才能明确。

线性注意力架构对KV缓存带来的挑战主要可以分为三个维度：预填充与解码分离、prefix cache优化、解码效率提升。

首先是预填充与解码分离优化，参考相关最新研究成果，线性注意力架构降低了prefill和decode阶段的数据传输量，这意味着可以使用更低质量的网络完成p/d传输。延迟的计算公式可简化为数据量除以带宽，当数据量大幅减少时，即便带宽有所降低，整体延迟仍能保持稳定。这种优化的优势在于可以采用性价比更高的硬件来完成prefill阶段，比如国产加速卡、老旧显卡甚至竞价实例。但实际落地中，大部分厂商和机房仅部署单一型号的加速卡，很难实现异构硬件部署，要推进异构部署需要综合考量多方面成本，且需要说服多方配合，仅有少数头部玩家能够实现这一方案。对于多数MaaS厂商而言，不太可能为单一模型投入如此大的成本进行异构架构改造，毕竟后续可能还有其他模型推出，因此底层算力基础设施的优化还是需要基础模型公司来主导推进。

接下来探讨prefix cache优化，此前从技术交流中了解到相关概念后，对这一机制有了新的理解。传统的普通注意力机制需要保留全部的KV日志，若要回到某个历史状态，必须持有完整的日志数据，这会导致存储开销过大。类比Raft共识算法的设计，其并不会全量存储日志，而是通过快照机制来限制日志规模：当日志积累到一定程度时，对当前系统状态生成快照并丢弃此前的日志。在线性注意力架构中，这种快照机制就对应着线性注意力的状态存储。

不过Raft算法并未规定快照生成的时机，理论上可以在每个日志条目后生成快照，也可以间隔多个条目生成一次，甚至仅保留最后一个快照，这与线性注意力的状态存储逻辑一致，这也是相关技术报告中讨论的核心问题之一。我们可以将这种机制称为快照策略，最基础的实现方式是参考传统方案，每64个token生成一次快照，这种方式虽然简单但能够正常运行。

快照策略的核心问题在于如何在有限的存储空间内，通过合理的快照规划实现最高的缓存命中率。既可以选择每1个token生成一次快照，也可以选择每1024个token生成一次，具体需要结合实际业务场景下的缓存模拟器运行结果来确定，这里存在一定的学术优化空间，基础实现可以正常运行，但仍有不少可以改进的地方。

确定快照策略后，需要解决索引管理的问题。由于模型并非完全采用线性注意力，还融合了普通注意力，两者的存储方式存在差异，如果采用分页管理方式，可以通过统一的索引来完成管理。如果两者的索引不一致，就需要分别维护线性注意力的索引和KV缓存的索引，这会大幅提升处理复杂度，因此应该以线性注意力的索引为准，保持整体索引的一致性。在存储线性注意力快照的同时，保留对应的KV缓存引用，这样仅需要查询一次线性注意力索引即可完成相关操作，后续的参数调优仍围绕快照策略展开。

我认为最优的快照策略应该是动态调整的，比如在多轮对话场景中，缓存的命中边界应该以对话轮次为单位，而非固定的token长度。这种方案需要考虑共享前缀的缓存优化，可以参考相关缓存算法的思路，这一方向甚至可以作为学术论文的研究方向。

此外，还有内存分配策略的问题。根据不同的快照策略，线性注意力状态和KV缓存占用的内存比例会动态变化，比如某一时刻线性注意力占用40%的内存、KV缓存占用60%，另一时刻则可能变为线性注意力占用20%、KV缓存占用80%。这意味着无法在服务启动时粗暴地划分固定比例的内存空间，再通过两个独立的free list分别管理，虽然这种方式也能运行，但并非最优方案。目前相关开源项目都对此展开了相关探索。

相关方案的核心思路是对两种内存粒度计算最小公倍数，比如两种粒度分别为2kB和3kB，则以6kB作为统一的内存管理粒度，让大内存页能够同时适配两种分配需求，不过这种方案仍存在部分内存碎片的问题，具体取决于最小公倍数的计算结果，如果两种粒度相同，则可以完全沿用现有的free list机制，实现零内存碎片，目前暂不清楚相关团队是否考虑了这一优化方向。此前咨询相关技术工具得到的建议也是采用类似方案，这一方案的内存碎片问题并不严重，至此关于prefix cache的相关讨论暂告一段落。

最后是解码效率优化。采用线性注意力架构能够降低单次解码请求占用的HBM显存，进而支持更大的batch size。比如在200GB HBM的硬件环境下，此前256k上下文的解码请求需要占用5GB显存，最大batch size为40；采用线性注意力架构后，仅需占用2GB显存，最大batch size可以提升至100。更大的batch size会带来两方面的影响：一是WideEP场景下的all-to-all通信量会大幅增加，这也是LatentMoe模块设计的出发点之一；二是group gemm运算更容易达到计算饱和，提升硬件利用率。不过在这些降本优化的背后，Kimi K3的输出成本达到了15美元每百万tokens。

此前尝试对相关模型进行测试，但遇到了分配问题、PP加速解码问题等，最终出现异常导致测试中断。另外需要说明的是，目前相关开源项目暂未支持混合注意力架构的prefix cache优化，后续有空会继续跟进这部分工作。

[1] Prefill-as-a-Service: KVCache of Next-Generation Models Could Go Cross-Datacenter