Kimi K3.1 Pro来了：基于CANN的Day0推理支持与性能优化
CSDN · 软件开发
CSDN 2026-11-06 14:21 江苏 CANN完成Kimi K3.1 Pro的Day0支持与优化，助力开发者在昇腾平台上快速部署与体验。 10月24日，Moonshot AI发布Kimi K3.1 Pro，并于11月4日正式开源。CANN完成Kimi K3.1 Pro的Day0支持与优化，助力开发者在昇腾平台上快速部署与体验。 承接Kimi K3积累，快速适配新架构 Kimi K3.1 Pro交错使用KDA与DSA：KDA通过状态递推处理序列；DSA在MLA基础上引入索引选择，由LI（Lightning Indexer）筛选Top-K历史位置，再由SFA（Sparse Flash Attention）读取对应KV完成稀疏注意力计算。模型继续通过Attention Residuals聚合跨层表示。 依托K3阶段的优化积累，K3.1 Pro复用FlashKDA/MegaKDA、AttnRes两阶段融合及MegaMoE/MC2等算子优化，结合混合并行、通信优化与多流编排，以及SuperKernel、图模式、编译缓存和NZ权重布局，并沿用DSpark投机推理与ReplaySSM状态回放机制，从计算、通信、调度和存储多个环节提升推理效率，支撑新架构的快速适配。 DSA融合与多流，减少解码等待 LI将相关性打分、头权重加权和Top-K选择融合执行，直接读取分页K缓存，减少中间结果写回与算子启动；SFA根据选出的索引读取主注意力KV，完成稀疏注意力计算。Prefill与Decode均支持LI的BF16和量化路径。 Indexer Prolog进一步通过三流协同，使Q投影与K归一化、Q量化与头权重投影并行执行，减少输入准备阶段的串行等待。 分层量化，降低缓存与搬运开销 DSA保留BF16投影计算。LI全量化采用MXFP8，按per-group方式对索引Q/K进行动态量化，group size为32，即每32个元素共享一个缩放因子。SFA的KVCache采用FP8，量化粒度为per-token-head-tile，tile size为128，即每个Token、每个Head内按128个元素分块量化。该策略降低索引数据与 KV缓存的存储、访存开销。 Prefill CP，兼顾长序列切分与负载均衡 本次支持CPxTP组合部署，沿序列维度切分长输入，并结合TP组内序列并行，降低单卡激活与中间张量占用。 针对DSA因果索引计算中"序列越靠后、历史越长"的负载差异，采用Zigzag布局，将首尾配对的两段序列分配到同一卡，平衡跨卡索引计算量，完成计算后恢复原有顺序。 基于CANNBot-DSL，支持FlashKDA CP 超长序列场景下，KDA的Prefill计算耗时随序列长度增长，影响首Token响应时延（TTFT）。仅采用TP按注意力头切分时，每张卡仍需处理完整序列。FlashKDA CP进一步沿序列维度切分输入，将长序列计算分摊到多张卡，旨在降低超长序列下的TTFT，同时减少单卡激活与中间张量占用。 本次基于CANNBot-DSL实现FlashKDA CP，各卡先并行预计算并通过AllGather交换状态摘要，再按序合成前序分片的状态贡献，得到各自分片的起始状态，随后并行完成本段计算。 性能表现：100K长序列场景下，解码时延低于10ms 得益于DSA稀疏Attention结构，并结合DSpark投机推理等优化，在昇腾950DT平台32卡部署、100K输入长度、Global Batch 32的场景下，模型平均Decode计算时延低于 10 ms/Token（仅统计模型耗时，DSpark接受长度=4）。 相关链接 CANN推理实践 · cann-recipes-infer https://gitcode.com/cann/cann-recipes-infer CANNBot-DSL · 融合算子开发 https://gitcode.com/cann/cannbot-dsl ops-transformer · Transformer算子库 https://gitcode.com/cann/ops-transformer 欢迎开发者体验Kimi K3.1 Pro在昇腾平台上的推理实践，基于CANNBot-DSL探索融合算子开发，并通过Issue、PR分享部署反馈、优化建议与代码贡献。期待与社区伙伴共同推进模型适配和性能优化，让更多开源创新在昇腾平台落地。 跳转微信打开