CANN Qwen3-MoE推理优化

基于Atlas A3训练/推理集群的Qwen3-MoE模型低时延推理性能优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer概述本文主要介绍Qwen3-MoE模型基于NPU的低时延推理优化策略。基于Atlas A3 训练/推理系列产品decode采用8卡部署实现BF16场景下单batch推理时延小于20ms。低时延场景Tensor Parallel (TP)优化Attention TP优化切分策略对Attention的张量切分策略可以分为对QKV头的切分和对线性层的切分。 在对QKV头切分时attention的多头计算机制可以方便进行张量切分每个头先独立计算再将结果concat起来。假设模型的attention层需要对num_heads个query按照切分数量attn_tp_size进行切分要求num_heads必须能被attn_tp_size整除每张卡放置query头个数为num_heads_per_rank num_heads // attn_tp_sizekey和value头数相等且可能小于等于query头个数在MQA和GQA的场景下会小于。为了确保每张卡至少放置一个key和value头每张卡放置的key或value头数计算方法为num_key_value_heads_per_rank max(num_key_value_heads // attn_tp_size, 1)。QKV头在多卡上的排布情况如下图所示。在对线性层o_proj进行切分时按照行切分即可。计算分解该优化策略先将Q、K、V的线性层计算合并为一次Matmul计算图中merged_qkv_proj提升计算性能。将merged_qkv_proj的输出结果按Q、K、V拆分后对Q和K进行归一化操作并使用旋转位置编码再计算attention图中Fused_infer_attention_score最后通过o_proj层输出。MoE TP优化切分策略假设模型的MoE层的切分数量为moe_tp_size专家个数为expert_num。对MoE层进行张量切分每个专家相当于一个mlp层切分方法与mlp的张量切分方法相似。具体做法是对gate_proj与up_proj进行列切分对down_proj进行行切分。同时对gate_proj与up_proj线性层采用合并计算的优化方式得到w13_weight。计算分解每个专家层存在gate_proj、up_proj与down_proj三个matmul运算具体运算为 x down( SiLU(gate(x))*up(x) )。本优化将张量切分后的gate_proj和up_proj进行concat操作再使能torch_npu.npu_swiglu融合算子接口优化该算子能完成以下两步计算将输入的x沿最后一维切分为两块即x torch.chunk(x, 2, -1)。计算并返回 SiLU(x[0]) * x[1]。本优化通过将gate_proj与up_proj合并计算提升整体计算效率。使能融合算子GMM使能和Routing优化在MoE模块中如果通过for循环处理每个专家单独计算expert_num个前馈神经网络FFN容易导致计算效率较低。CANN提供了GroupedMatmul算子可以同时计算多个专家从而提高计算和搬运效率。具体实现可参考在Qwen3MoeSparseMoeBlock类中的moe_infer_tp和moe_infer_fusion函数。快速选择专家在计算专家和token之间的路由分数时可以使用torch_npu.npu_moe_gating_top_k_softmax融合算子代替原来先topk再softmax多算子操作可以更快速地计算出token和专家的分数。高效排序和token路由使能torch_npu.npu_moe_init_routing融合算子实现MoE routing计算获取专家的排序使能torch_npu.npu_moe_compute_expert_tokens融合算子获取每个专家需要计算的token数使能torch_npu.npu_moe_finalize_routing融合算子将专家计算完成后的token重新排布并加权求和获得最终输出。高性能专家计算使能torch_npu.npu_grouped_matmul融合算子实现多个专家的矩阵乘计算提高计算和搬运效率。RmsNorm算子优化通过使能torch_npu.npu_rms_norm算子能够提升模型的推理性能。RmsNorm是大模型常用的归一化操作相比LayerNorm其去掉了减去均值的部分。flash attention融合算子优化通过使能torch.ops.npu.npu_fused_infer_attention_score推理场景下支持图模式的FlashAttention算子既可以支持全量计算场景也可支持增量计算场景。使能图模式使用静态图可以获得更好的推理性能。Qwen3MoeRunner通过覆写executor/model_runner.py中的ModelRunner的graph_compile函数将模型编译为静态图。使能图编译缓存在模型推理场景下使能图编译缓存可以缓存编译后的静态图避免每次推理都需要编译模型从而提高推理性能。可参考Qwen3MoeRunner中graph_compile函数中的使用if self.enable_cache_compile: case_name compile_cache/ os.getenv(CASE_NAME) cache_dir os.path.join(os.path.dirname(os.path.abspath(__file__)), case_name) self.model.decode tng.inference.cache_compile(self.model.decode, cache_dircache_dir, configcompiler_config, dynamicTrue, fullgraphTrue, ge_cacheTrue)缓存默认路径为./compile_cache/CASE_NAME。集合通信使能AIV展开利用Device的Vector Core计算单元来加速通信操作的执行可参考HCCL_OP_EXPANSION_MODE环境变量export HCCL_OP_EXPANSION_MODEAIV附录环境部署以及样例执行【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考