基于Atlas A3、950训练/推理集群的Qwen3-MoE模型低时延推理性能优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer概述本文主要介绍Qwen3-MoE模型基于NPU的低时延推理优化策略。基于Atlas A3 训练/推理系列产品decode采用8卡部署实现BF16场景下单batch推理时延小于20ms。基于Atlas 950 PR 训练/推理系列产品decode采用4卡部署实现W4A8混精推理4K场景下单batch Prefill时延小于280ms,Decode时延小于24ms。低时延场景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高吞吐场景Prefill优化Attention SPTP, MoE EP概述在Prefill推理场景下模型需要一次性处理完整的输入序列计算密集度高。为在TTFT时延约束下实现最大吞吐采用Attention序列并行SP张量并行TP与MoE专家并行EP的组合策略结合MXFP8/MXFP4量化与多算子融合显著降低Prefill阶段的计算与通信开销。Attention SPTP优化并行策略选择Prefill阶段算子多为计算boundTP并行在切分后每卡的序列长度缩短但AllReduce通信量与完整序列成正比。SP并行沿序列维度切分每卡仅处理seq_len // attn_tp_size长度的子序列通信算子采用AlltoAll通信数据量更小且可配合低bit量化。综合考虑通信数据量、数据类型和额外引入算子三个因素Prefill阶段Attention采用SPTP混合并行FA前通信使用AllGather将各Rank的子序列聚合成完整序列配合MXFP8量化减小通信体积FA后通信使用AlltoAll将o_proj的输出按序列维度重新分发到各Rank配合MXFP8量化。切分策略输入序列沿token维度按attn_tp_size均分每个Rank处理seq_len // attn_tp_size个token。要求输入序列长度能被attn_tp_size整除不足时进行padding对齐。Rank 0: [token_0, ..., token_{S/P-1}] Rank 1: [token_{S/P}, ..., token_{2S/P-1}] ... Rank P-1: [token_{(P-1)S/P}, ..., token_{S-1}]其中S为序列长度P为attn_tp_size。计算分解Prefill阶段的Attention计算流程如下序列切分与Embedding输入input_ids按attn_tp_size切分每个Rank仅对自身持有的子序列做EmbeddingAllGather聚合通过dist.all_gather_into_tensor将各Rank的hidden_states聚合为完整序列若使能MXFP8量化则同时聚合hidden_states_scalemerged_qkv_proj对聚合后的完整序列执行QKV线性层计算FlashAttention使用npu_fused_infer_attention_score_v2完成注意力计算o_proj AlltoAllo_proj使用ReplicatedLinear每卡持有完整权重输出经MXFP8量化后通过AlltoAll按序列维度重新分发每个Rank仅保留自身负责的子序列部分MoE输入准备AlltoAll后的输出shape为(seq_len // attn_tp_size, hidden_size)直接作为MoE层的输入。MoE EP优化并行策略选择MoE部分支持TP和EP两种并行方式。TP并行会将GMM的K或N轴切得过小导致计算性能下降。EP并行每个Rank持有部分专家通过通信完成token与专家的路由匹配在ep_rank数较小时通信开销可控。EP并行进一步分为AllGather方案和Double RoutingAlltoAll方案并行策略通信算子通信数据量适用场景EP (AllGather)AllGather ReduceScatter(B,S,H) × 2ep_rank较小、top_k ep_rank时数据量更小EP (Double Routing)AlltoAll AlltoAll(BS×top_k/ep_rank, H) × 2大EP场景、ep top_k时数据量更小当前Qwen3-MoE模型top_k8在ep_rank4时top_k ep_rankAllGather方案通信数据量更小为默认优选方案。计算分解AllGather方案W4A8量化当使能W4A8MXFP4量化且attn_tp_size moe_ep_size时MoE EP采用AllGather ReduceScatter通信模式具体流程如下MXFP8量化对输入hidden_states执行npu_dynamic_mx_quant得到MXFP8量化数据及scaleAllGather通过dist.all_gather_into_tensor聚合各Rank的量化数据、scale、topk_ids和topk_weight得到完整序列的token信息Routing使用npu_moe_init_routing_v2对聚合后的token进行专家路由仅筛选属于本Rank的专家范围内的tokenGMM计算通过npu_grouped_matmul封装在FusedMoEGMM中完成多专家并行矩阵乘计算Finalize Routing使用npu_moe_finalize_routing将专家计算结果按路由信息重新排布并加权求和ReduceScatter通过dist.reduce_scatter_tensor将各Rank的部分结果归约并分发每个Rank获得自身负责的子序列输出。计算分解Double Routing方案当未使能W4A8量化时MoE EP采用Double RoutingAlltoAll通信模式具体流程如下Routing使用npu_moe_init_routing_v2对token进行初始路由和排序DispatchAlltoAll通过dist.all_to_all_single将token按专家归属分发到对应Rank同时交换tokens_per_expert信息用于确定接收数量Re-routing使用npu_moe_re_routing对接收到的token进行二次排序GMM计算通过npu_grouped_matmul完成专家计算CombineAlltoAll将专家计算结果通过dist.all_to_all_single回传到原始RankFinalize Routing使用npu_moe_finalize_routing完成最终的结果归约。融合算子优化Prefill场景下为减少量化引入的额外算子开销将归一化、残差加法与量化操作进行融合融合RMSNorm MX量化通过使能torch_npu.npu_add_rms_norm_dynamic_mx_quant融合算子将残差加法、RMSNorm归一化和MXFP8动态量化合并为一次计算避免中间结果的反复读写hidden_states_mx, residual, hidden_states_scale, _ \ torch_npu.npu_add_rms_norm_dynamic_mx_quant( hidden_states, past_residual, norm_weight, betaNone, epsilonvariance_epsilon, round_moderint, dst_typetorch.float8_e4m3fn )当残差为None时首层使用torch_npu.npu_rms_norm_dynamic_mx_quant完成RMSNorm MX量化融合。融合RMSNorm CastMoE前的归一化需要同时输出FP32用于gate计算和BF16用于专家计算两种精度通过使能torch_npu.npu_add_rms_norm_cast融合算子将残差加法、RMSNorm和数据类型转换合并hidden_states_fp32, hidden_states_bf16, _, residual \ torch_npu.npu_add_rms_norm_cast( hidden_states, residual, norm_weight, epsilon )配置说明在YAML配置文件中设置Prefill SPTP与MoE EP相关参数parallel_config: attn_tp_size: 4 # Attention SP/TP并行度Prefill时同时作为序列并行度 moe_tp_size: 1 # MoE TP大小EP场景下设为1 moe_ep_size: 4 # MoE EP大小需与attn_tp_size相等以启用SP优化路径参数约束参数约束条件说明attn_tp_sizeworld_size % attn_tp_size 0Attention并行度moe_ep_sizeworld_size % moe_ep_size 0MoE专家并行度attn_tp_size moe_ep_size推荐相等启用Prefill SP优化路径prefill_optseq_len需能被attn_tp_size整除不足时自动padding对齐附录环境部署以及样例执行【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考