基于Atlas A3集群的DeepSeek-R1模型decode阶段推理性能优化【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer概述本文主要介绍基于Atlas A3 系列产品的DeepSeek-R1模型decode阶段的并行部署和性能优化策略。本文将介绍如何在Atlas A3 系列产品上进行大EPExpert Parallel推理及其他并行技术和优化方案。性能优化通用优化DeepSeek-R1结构中的非MoE部分与Llama类似通用优化点可参考Llama的改动如固定KV Cache大小、cos/sin优化、AddRMSNorm融合、全量优化LM Head计算量等。MLA (Multi-Head Latent Attention)低秩压缩优化使能融合算子参考Deepseek论文中提及的低秩压缩方法可以减少KV cache占用的内存提升推理效率。在实现MLA低秩压缩后可以通过使能融合kernel实现性能优化相关实现可以参考DeepseekV3Attention类中的forward_page_attention_mla_prolog函数。MLA前置计算性能优化使能npu_mla_prolog_v3融合kernel替换attention计算前的计算其中包含Q、K、V的线性层计算、旋转位置编码 (ROPE)、RmsNorm计算及KV Cache更新等计算处理Attention性能优化使能npu_fused_infer_attention_score融合kernel实现对MLA计算的加速。MoE模块实现Expert Parallel (EP)及使能融合算子在MoE模块中如采用原始的MoE实现将通过for循环处理expert_num个FFN专家效率较低。针对MoE层中的token路由和专家计算等操作CANN提供了一系列融合算子提升计算效率。Router计算优化使用torch_npu.npu_moe_gating_top_k算子对router计算的结果排序并选取前top-k个专家高性能专家计算使用torch_npu.npu_grouped_matmul算子可以同时处理多个专家的计算提高计算和搬运效率多卡间高性能通信路由使能torch_npu.npu_moe_distribute_dispatch_v2 和torch_npu.npu_moe_distribute_combine_v2算子实现EP并行下多卡间的通信。在使用前可参考上述的算子文档检查HCCL_BUFFSIZE等环境变量的配置是否合理了解该算子的使用场景和约束。MLP线性层计算合并原始DeepseekV3MLP实现中存在gate_proj、up_proj与down_proj三个matmul运算可通过将gate_proj与up_proj进行合并计算得到gate_up_proj提升整体计算效率。torchair图模式Decode场景一般时延要求较高在pytorch的eager_mode下算子很可能面临比较严重的host下发bound此时可以通过torch.compile使能图模式进行优化。TorchAirTorch Ascend Intermediate Representation是Ascend Extension for PyTorchtorch_npu的图模式能力扩展库提供了昇腾设备亲和的torch.compile图模式后端实现了PyTorch网络在昇腾NPU上的图模式推理加速以及性能优化。TorchAir提供了max-autotune和reduce-overhead两种实现模式具体可参考文档。本样例在decode场景下对这两种模式都支持可以通过配置yaml文件中exe_mode使能其中ge_graph对应max-autotune模式acl_graph对应reduce-overhead模式。支持Multi-Token Prediction (MTP)根据Deepseek论文中介绍的MTP方法实现了MTP投机推理在未达到计算bound的场景下MTP计算可以实现较好的推理加速效果。可通过next_n参数使能MTP。集合通信使能AIV展开利用Device的Vector Core计算单元来加速通信操作的执行可参考HCCL_OP_EXPANSION_MODE环境变量export HCCL_OP_EXPANSION_MODEAIVSuperKernel优化在decode使能ge_graph图模式的场景下支持对模型的计算图按照用户定义的范围进行SuperKernel优化。SuperKernel技术的详细介绍请参考官方文档。在本示例中可通过enable_superkernel开关使能支持将所有的decode layers优化在一个SuperKernel scope内实现对任务调度的等待时间和调度开销的优化提升整体性能。如希望修改SuperKernel scope的范围可以在modeling_deepseek.py中调整示例如下from executor.utils import superkernel_scope switch enable_superkernel scope fscope_name options fscope_options with superkernel_scope(switch, scope, options): your modules其中scope表示当前融合范围的SuperKernel名称options表示SuperKernel编译的自定义选项。关于scope和options的具体描述和使用范围请参考SuperKernel文档。Benchmark基于Atlas A3本实践对DeepSeek-R1 W8A8C8 量化版本进行了性能Benchmark测试。 |Quant Mode| Global Batch Size | Seq Length | Chips | TPOT (ms) | Throughput (tokens/p/s) | |-------| ----------------- | ---------- | ----- | --------- | ----------------------- | |W8A8C8 | 6144 | 4096 | 64 | 44.9 |2138 |注性能数据基于 MTP1 与 perfect eplb 配置采集平均 1 个 draft token 中 accept token 为 0.7 个。附录环境部署以及样例执行【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考