NPU DeepSeek-V3.2-Exp推理优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-inferDeepSeek团队发布了最新的模型DeepSeek-V3.2-Exp可利用稀疏架构DeepSeek Sparse Attention(DSA)来提高长序列的计算效率降低推理成本。长上下文场景和其新颖的DSA结构共同对推理优化系统提出了新诉求。本文主要介绍基于A3集群的DeepSeek-V3.2-Exp模型的Prefill和Decode推理优化首先分析了DeepSeek-V3.2-Exp模型的结构特点实现了长序列亲和模型并行策略并设计了新的NPU融合Kernel和多流并行优化。基于这些优化点本实践0day实现了DeepSeek-V3.2-Exp BF16模型部署1day实现W8A8C16模型部署目前已支持C8量化未来将持续完成低比特量化算法。4~8K常规序列继承原有DeepSeek-V3.1优化16K~166K长序列推理场景结合DSA稀疏收益性能超越DeepSeek-V3.1模型和融合Kernel均已开源。针对该模型涉及到的DSA结构中的Lightning Indexer(LI)和Sparse Flash Attention(SFA)新算子本次开源不仅包含AscendC实现并且首次公布了自研PyPTO实现仅需几百行代码即可完成动态Shape算子编程。PyPTO是CANN 推出的大融合算子的编程体系提供面向 MegaKernel 的编程范式。其核心思想是使用 Tensor/Tile 作为数据的基本表达方式借助一系列对 Tensor 的基本运算来描述和组装完整的计算流程采用human-in-the-loop和白盒化的方式能够充分利用SRAM的空间实现对多核的高效利用。同时TileLang开源社区也已对接CANN软件栈同步支持了DSA结构中的LI和SFA新算子。Highlights整体部署策略沿用DeepSeek的大EP并行方案针对稀疏DSA结构叠加实现长序列亲和的CP并行策略兼顾时延和吞吐。模型推理代码已开源同时也适配了主流开源推理框架vLLM和SGLang基于AscendC实现NPU LI和SFA融合Kernel发挥稀疏计算潜力AscendC Kernel技术文档和代码已开源在CANN/ops-transformer仓对应算子目录包含Lightning Indexer和Sparse Flash Attention等基于自研PyPTO框架实现NPU DSA提高融合算子编程易用性。不仅实现了LI融合Kernel同时实现了更大范围的Decode Attention融合PyPTO Kernel技术文档和代码已开源开源社区TileLang同步支持DSA结构中的LI和SFA算子TileLang Kernel技术文档和代码已开源基于上述优化点CANN已0day支持DeepSeek-V3.2-Exp BF16推理部署1day支持Int8量化。Prefill和Decode的参考性能采用BF16精度无损方式64卡128K长序列TTFT小于2s无缓存命中TPOT小于30msW8A8C8量化场景64卡128K长序列TPOT小于20msOutlineDeepSeek-V3.2-Exp vs V3.1并行策略MTP融合Kernel量化策略多流并行优化KVCache OffloadBenchmarkFuture PlanDeepSeek-V3.2-Exp vs V3.1相较DeepSeek-V3.1DeepSeek-V3.2-Exp主要新增了Lightning Indexer稀疏模块旨在从长序列中稀疏选择Topk个Token用于计算注意力。权重内存分析Lightning Indexer是一个轻量的类MQA结构新增了q_b projwk proj和weight proj三个Linear层其权重大小分别为 $$ \begin{aligned} \mathrm{Param}{\mathrm{q_b_proj}}\mathrm{q_lora_rank} \times \mathrm{Indexer_head_num} \times \mathrm{Indexer_head_dim} \ \mathrm{Param}{\mathrm{wk_proj}} \mathrm{hidden_size} \times \mathrm{Indexer_head_dim}\ \mathrm{Param}_{\mathrm{weight_proj}} \mathrm{hidden_size} \times \mathrm{Indexer_head_num} \end{aligned} $$ 综合所有Layer权重参数量新增0.85B左右。KVCache内存分析为了保障模型效果DeepSeek-V3.2-Exp仍然缓存了完整MLA的Full KVCache。与此同时为了提高推理效率避免decode阶段重复计算Indexer Key需要新增缓存Indexer Key Cache其内存大小为 $$ \mathrm{batch_size}\times \mathrm{kv_length} \times \mathrm{Indexer_head_dim} \times \mathrm{storage_bytes} \times \mathrm{num_layers} $$ 以每个rank处理4batch 64K序列长度为例BF16场景新增的Indexer Key Cache约为4GB左右。MLA Naive vs Absorb在Prefill阶段以1batch 64K推理为例Lightning Indexer为每个q token选择TopK2048个kv tokenMLA的计算流程可以有三种选择方案一MLA Naive Sparse MaskPrefill MLA使用Naive模式和原始DeepSeek V3保持一致。每个q token和所有的历史kv token计算Attention仅在softmax前通过Attention Mask将不属于TopK的token过滤掉。Attention中的两个BatchMatmul Shape如下方案一BatchMKNBMM1(Q*K^T)1*12864K19264KBMM2(P*V)1*12864K64K128该方案计算量和原始的Full Attention一致但是无法拿到DSA的稀疏计算收益长序列场景下性能不佳。方案二MLA Naive Sparse AttentionPrefill MLA使用Naive模式每个q token与TopK2048个kv token计算Attention。方案二BatchMKNBMM1(Q*K^T)1*64K*12811922048BMM2(P*V)1*64K*12812048128方案二的优点在于BMM的计算量较小相对原始的Full Attention计算量减小了64K/204832倍但是存在以下问题BMM的M轴为1矩阵乘法计算效率较低BMM的HBM访存量相较原始的Full Attention激增2048倍将会面临访存瓶颈方案三MLA Absorb Sparse AttentionPrefill MLA使用Absorb模式与Decode保持一致。同样地每个q token与TopK2048个kv token计算Attention。方案三BatchMKNBMM1(Q*K^T)1*64K1285762048BMM2(P*V)1*64K1282048512方案三与方案二对比如下方案三的计算量增加了3倍左右但其BMM的M轴为128对于矩阵乘法更为友好方案三的HBM访存量相对方案二降低几十倍访存耗时更低综合考虑计算和访存耗时以及长序列应用场景本实践选择基于方案三(MLA Absorb Sparse Attention)来完成Prefill部署从而Prefill和Decode的MLA计算流可以归一。为了更优的推理时延和吞吐针对长序列DeepSeek-V3.2-Exp模型本实践设计了NPU亲和的并行策略并实现了NPU Lightning Indexer(LI)和Sparse Flash Attention(SFA)融合kernel细节参考后续章节。并行策略Atlas A3推荐部署策略如下图所示Prefill使用M个节点部署Decode使用N个节点部署每个节点包含8卡。其中BF16场景下推荐根据资源数量、SLA等约束M在4~8N在4~16内动态调整。Prefill并行策略DeepSeek-V3.2-Exp新增的DSA主要面向长序列推理场景Prefill阶段设备内存OOM风险较高而且推理服务部署时TTFT也会面临巨大的挑战因此优化内存占用和TTFT是并行策略设计的主要目的。如果沿用DeepSeek R1 Blog里的Attention DP策略每个rank都需要推理不同的超长序列总计算量较大TTFT较高用户体验较差。Tensor Parallel(TP)是优化推理TTFT的重要手段但DeepSeek-V3.2-Exp中DSA新增的Indexer模块计算量和$S^2$成正比是长序列场景下的主要瓶颈该模块针对head轴做了Reduce Sum操作TP切分会产生大量的通信开销影响整网性能。综合上述两点DeepSeek-V3.2-Exp模型Prefill Attention选用Context Parallel(CP)并行多个rank均摊长序列的计算单rank的计算量和activation内存都较小TTFT较为可控用户体验更好。MoE模块则沿用DeepSeek-V3.1的EP并行兼顾吞吐与时延。Prefill的并行策略可以设计为下图形式Embedding并行策略为了节省Embedding Table的内存占用Embedding计算使用TP并行tp_size控制在单Node内之后通过AllReduce聚合。Attention并行策略Attention整体选用CP并行以64K输入推理为例每个rank处理64K/cp_size1K个token在计算完kv之后对所有CP域的kv token进行AllGather得到完整的kv结果。每个rank拿到64K/cp_size的q token和完整的kv token进行后续的Attention计算。Attention计算需要遵循因果注意力如果CP简单按照rank顺序进行切片可能会面临计算负载均衡问题。如第一个rank关注到的历史kv token很少计算量较小最后一个rank关注到的历史kv token较多计算量较大。为了降低负载不均带来的影响需要将输入的input_ids按照cp_size*2进行切片如上图右侧所示每个rank负责计算头尾对称的两个切片每层Lightning Indexer和Sparse Flash Attention计算前通过Token重排将kv还原回因果顺序。考虑到O_proj模块内存占用较大本实践选择对其局部进行TP切分因此O_proj前后引入AlltoAll和ReduceScatter算子。MoE并行策略MoE模块沿用了和DeepSeekV3相同的EP并行部署EP复用CP通信域。LM Head并行策略为了节省LM Head的内存占用LM Head选用TP切分tp_size控制在单Node内。LM_Head模块采用TP切分同样需要在TP域内做AllGather但此处为整个CP域内的AllGather并通过Token重排将最终输入LM_Head的feature map还原到原始序列排布。Decode并行策略Decode阶段依旧沿用DeepSeek-V3.1的部署策略选用Attention DP MoE EP部署。特别地由于O_proj和LM_Head权重内存较大且在Decode阶段表现为明显的访存瓶颈本实践选用局部TP并行。同时为了降低设备内存占用Embedding层同样使用TP切分。为了尽可能地减小TP并行带来的通信开销TP域控制在高速互联HCCS域内。MTPDeepSeek-V3.2-Exp依然提供了原生的Multi-Token Prediction(MTP)机制MTP机制允许在一次主模型推理过程中同时推理多个Token在相似的数据搬运下进行更多的计算以充分利用芯片的算力提升模型等效时延和吞吐。DeepSeek-V3.2-Exp由于采用了新的DSA结构MTP加速相对于DeepSeek-V3.1会更复杂。Full Attention场景的MLA无论是否使用MTP其KV Cache的搬运量相同。因此MTP能够提高计算访存比从而提高算力利用率CANN中MLA算子在MTP1场景的算力利用比可以达到70%以上。而对于DeepSeek-V3.2-Exp采用的Sparse Flash Attention每个q token都会选择2048个激活的kv token极端情况下MTP1场景需要搬运的KV Cache是非MTP场景的两倍反而增加了离散访存的代价。64k序列长度 4batch的搬运量和计算量KV搬运量MBCube计算量GFlopsQK^T搬运量MBMLA非MTP14419.3332MTP114438.6564SFA非MTP4.50.601MTP191.212LI非MTP322.1516MTP1324.2932尽管Sparse Flash Attention很难利用MTP实现可观的加速效果但是Lightning Indexer算子在长序列场景中耗时占比更高其K Cache搬运与稠密MLA相似可以利用MTP机制提高计算访存比。另一方面由于DSA稀疏计算显著降低了长序列的计算量削弱了Attention在整个推理的占比而剩下的算子如Matmul等都能达到更好的计算访存比因此使用MTP对整个推理有比较可观的加速。融合KernelDSA的计算过程可分为MLAProlog、IndexerProlog、Lightning Indexer、Sparse Flash Attention、MLAEpilog五部分如下图所示Lightning Indexer包含Score Batchmatmul、ReLU、ReduceSum、TopK等操作长序列场景TopK操作会成为瓶颈可用TopK计算耗时流水掩盖掉其他操作的耗时从而拿到计算流水收益Sparse Flash Attention包含了从完整KVCache里选取TopK相关Token及计算稀疏Flash Attention操作。此处的耗时瓶颈点主要在于离散聚合访存可用此耗时掩盖其他操作的耗时流水并行加速MLAProlog和IndexerProlog包含Q/KV的LoRA、RoPE、Norm、KVCache更新等操作存在较多Cube/Vector并行的流水空间MLAEpilog包含O_proj及V升维操作目前NPU已实现并在CANN/ops-transformer仓开源lightning_indexer和量化版本quant_lightning_indexer以及sparse_flash_attention和量化版本kv_quant_sparse_flash_attention。量化策略相对于BF16推理Int8量化可以有效降低端到端时延提升系统吞吐。目前本实践已经支持W8A8C8/W4A8C8量化量化架构如下其中MLAProlog、MLAEpilog、IndexerProlog、LightningIndexer三个量化融合算子如下MLAProlog除Q_b_proj使用W8A8其他Linear均不量化KVCache量化到C8Sparse Flash AttentionKVCache Int8存储BF16计算IndexerProlog除Q_b_proj使用W8A8其他Linear均不量化Indexer Q使用A8量化Indexer Cache使用C8量化Lightning Indexer: BatchMatmtul使用Int8计算MoE路由专家使用W8A8/W4A8量化共享专家使用W8A8量化MLAEpilogO_proj使用W8A8量化LM_Head暂不量化。注 W8A8W8指权重使用静态Per-Channel Int8量化A8指数据使用动态Per-Token Int8量化 A8C8A8表示Lightning Indexer中的Q使用动态Per-Token-Head Int8量化Indexer Cache使用动态Per-Token-Head Int8量化 MLAEpilogO_proj使用W8A8量化 KVCache C8表示KVCache 使用动态Per-Token-Head-Tile-128 Int8量化W8A8C8对线性层量化数量较少MLA线性层只量化了q_b_proj和w_o_projIndexer线性层只量化wq_b_proj。主要原因是IndexerProlog融合算子设计成weights_proj出fp16且不做量化因此MLA输入关联的Linear统一不做量化好处是可将同一份BF16数据输入IndexerProlog和MLAProlog。其次MLAProlog KVCache的量化策略使用了动态存8算16。在超长序列情况下W8A8C8量化精度接近无损同时权重内存占用优化2倍。MLA C8算16获取内存收益可以打高吞吐量。另一方面LightningIndexer的A8C8获取计算收益降低LI计算时延TTFT和TPOT也同步优化。W4A8C8量化版本针对DeepSeek-V3.2-Exp使用基于学习的量化算法优化Clamp参数缓解W4A8离群值量化困难的问题实现了较优的量化模型精度。同时W4A8C8版本比W8A8C8节约MoE权重显存2x因此在大EP场景下利用W4A8 MoEGMM算子同一张卡可以装下更多的专家节约资源优化计算访存比实现单机部署。量化模型精度表现模型MMLUGPQADROPMGSMDeepSeek-V3.2-BF1690.875.8587.5790.9DeepSeek-V3.2-Exp-W8A8C890.6275.9087.4190.87DeepSeek-V3.2-Exp-W4A8C890.6676.587.6891.3多流并行优化CANN提供了多流并行的机制可支持计算/通信并行等。由于A3芯片为CV分离架构也可支持Cube/Vector并行。针对MoE模型可利用CANN的多流并行能力提高芯片资源利用率。MoE模块中路由专家采用EP部署共享专家采用DP部署。共享专家的计算与路由专家的Gating、Dispatch、计算没有依赖关系可以通过多流并行机制来使二者流水掩盖。由于Gating函数的权重较小且Dispatch算子为通信算子二者的HBM带宽占用率比较低可以与共享专家带宽瓶颈进行并发掩盖掉共享专家的耗时。简单代码示例如下with torchair.scope.npu_stream_switch(stream_tag, stream_priority): # shared_expert use another stream hidden_states_share self.shared_experts(hidden_states) # router_experts use main stream hidden_states_router self.router_experts(hidden_states) hidden_states hidden_states_share hidden_states_router除了MoE模块Lightning Indexer前的一系列Cube和Vector算子IndexerProlog也可使用多流并行流水示意图如下KVCache Offload使能W8A8C8量化后由于受限于设备内存在长序列场景 batch size依然无法进一步提升。为了解决设备内存紧张问题可将占比较高的MLA Full KVCache内存卸载到Host内存上完成KVCache的Offload释放设备内存压力达到提高batch size的目的。 KVCache Offload方案涉及Prefill阶段和Decode阶段。Prefill StagePrefill阶段在设备上只申请一层MLA KVCache同时每卡在Host内存上申请61层KVCache内存。MLA计算完后通过多流的方式将当前层设备上的KVCache卸载到Host侧的内存上。Decode StageDecode阶段每一层计算得到的当前Token的Cache更新到Host内存上新增的GatherSelectionKvCache算子根据Lightning Indexer模块返回的TopK个相关Token的位置信息从Host侧的Full KVCache选出对应的需要参与计算的kv token送给后续的SFA计算。其中 GatherSelectionKvCache 算子充分利用了前后Token在同一层上TopK 命中率相似度特性将已经在设备上的Cache做了复用。也即如果当前步的TopK2048个kv token和前一步的TopK2048个kv token有60%一样那只需要从Host侧读取40%的没有复用的kv token60%可以复用的kv token直接从设备上读取从而减少了H2D的耗时。BenchmarkW8A8C8基于 Atlas A3本实践对 DeepSeek-V3.2-Exp 与 DeepSeek-V3.1 W8A8C8 版本进行了性能Benchmark 测试。从吞吐对比曲线可见随着序列长度增加DeepSeek-V3.2-Exp 的性能优势逐步扩大当序列长度达到 128K 时其吞吐达到 DeepSeek-V3.1 的 450%。下表展示了未启用 Offload 时不同 batch size 和序列长度下的推理时延和吞吐值。Global Batch SizeSeq LengthChipsTPOT (ms)Throughput (tokens/p/s)256655366419.84202512655366422.783511281310726418.851065121310726425.8310注性能数据基于 MTP3 与 perfect eplb 配置采集平均 3 个 draft token 中 accept token 为 1.44 个。MoE W4A8 Attention W8A8C8本实践新增了MoE部分W4A8量化的支持针对权重较大存在搬运bound的场景W4A8能够有效减少内存占用并降低权重搬运耗时从而显著提升推理性能。基于Atlas A3环境本实践对MoE W4A8量化特性进行了Benchmark测试。相同集群规模和配置下相比W8A8量化模型推理性能得到了一定的提升且单芯片上的权重越大性能提升效果越明显。MoE Quant ModeGlobal Batch SizeChipsTPOT (ms)Throughput (tokens/p/s)W8A8643220.8196W4A8643220.1899W8A8321624.3982W4A8321622.5189注性能数据基于相同输入序列长度65536MTP及其他配置同上。KVCache Offload基于 Atlas A3 环境在 DeepSeek-V3.2-Exp 模型使能 W8A8C8 量化的基础上对 KVCache Offload 特性进行了 Benchmark 测试。同等序列长度下通过使能 KVCache Offload 技术方案相比非 Offload模型推理支持的最大 batch size 可以翻倍。下图展示了 Offload 的内存收益固定序列长度为 64K相比非 Offload在使用 Offload 时 global batch size 可从 1024 增长到 2048固定 global batch size 为 128使用 Offload 时序列长度可从 256K 增长到 384K。注内存数据基于 MTP3 与 perfect eplb 配置采集。Future Plan量化目前支持BF16/Int8版本推理。未来进一步开发低比特量化版本探索KVCache量化压缩算法软硬协同优化NPU计算效率降低系统时延KVCache Offload利用前后Token命中率关系或跨layer预取等手段进一步降低GatherSelectionKvCache耗时MegaKernelDecode阶段仍然存在较多流水并行空间可通过PyPTO实现更大范围的MegaKernel完成多核MPMD并行调度提升计算效率【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考