Kimi K3.1 Pro来了:基于CANN的Day0推理支持与性能优化
10月24日,Moonshot AI发布Kimi K3.1 Pro,并于11月4日正式开源。CANN完成Kimi K3.1 Pro的Day0支持与优化,助力开发者在昇腾平台上快速部署与体验。
01 承接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状态回放机制,从计算、通信、调度和存储多个环节提升推理效率,支撑新架构的快速适配。
02 DSA融合与多流,减少解码等待
LI将相关性打分、头权重加权和Top-K选择融合执行,直接读取分页K缓存,减少中间结果写回与算子启动;SFA根据选出的索引读取主注意力KV,完成稀疏注意力计算。Prefill与Decode均支持LI的BF16和量化路径。
Indexer Prolog进一步通过三流协同,使Q投影与K归一化、Q量化与头权重投影并行执行,减少输入准备阶段的串行等待。

03 分层量化,降低缓存与搬运开销
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缓存的存储、访存开销。

04 Prefill CP,兼顾长序列切分与负载均衡
本次支持CPxTP组合部署,沿序列维度切分长输入,并结合TP组内序列并行,降低单卡激活与中间张量占用。

针对DSA因果索引计算中"序列越靠后、历史越长"的负载差异,采用Zigzag布局,将首尾配对的两段序列分配到同一卡,平衡跨卡索引计算量,完成计算后恢复原有顺序。

05 基于CANNBot-DSL,支持FlashKDA CP
超长序列场景下,KDA的Prefill计算耗时随序列长度增长,影响首Token响应时延(TTFT)。仅采用TP按注意力头切分时,每张卡仍需处理完整序列。FlashKDA CP进一步沿序列维度切分输入,将长序列计算分摊到多张卡,旨在降低超长序列下的TTFT,同时减少单卡激活与中间张量占用。
本次基于CANNBot-DSL实现FlashKDA CP,各卡先并行预计算并通过AllGather交换状态摘要,再按序合成前序分片的状态贡献,得到各自分片的起始状态,随后并行完成本段计算。

06 性能表现: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分享部署反馈、优化建议与代码贡献。期待与社区伙伴共同推进模型适配和性能优化,让更多开源创新在昇腾平台落地。
更多推荐



所有评论(0)