企业官网建设流程全解析

基于Atlas A3集群的DeepSeek-R1模型decode阶段推理性能优化

【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法，提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer

概述

本文主要介绍基于Atlas A3 系列产品的DeepSeek-R1模型decode阶段的并行部署和性能优化策略。本文将介绍如何在Atlas A3 系列产品上进行大EP（Expert 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_v2融合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使能图模式进行优化。TorchAir（Torch Ascend Intermediate Representation）是Ascend Extension for PyTorch（torch_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_MODE=AIV

SuperKernel优化

在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 = f"scope_name" options = f"scope_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