Qwen3.5-MoE 在 NPU 上的推理性能优化实践融合算子、TP/EP 并行与图模式全解析【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-inferQwen3.5-MoE 是同时包含 full attentionMRoPE、linear attentionGated DeltaNet与 MoE 三种结构的混合专家模型推理链路既涉及 Attention/MLP 类大矩阵计算也包含大量路由、EP 通信与归一化小算子。本文以 Qwen3.5-MoE 模型推理性能优化实践 为骨架结合 Qwen3.5-MoE 推理实现 与 模型源码 中的真实实现按融合算子 → TP/EP 并行 → 图模式 → 矩阵融合四个维度完整讲解其 NPU 优化方案。读完本文你将掌握每个融合算子的启用条件与调用方式、attention/MoE 的 TP 与 EP 切分细节、Decode 阶段整图编译的触发机制以及如何在 YAML 配置 中落地这些优化。一、优化全景四类问题的四个解法Qwen3.5-MoE 的推理链路同时包含三类计算full attention 层基于 MRoPE 位置编码的稠密注意力Prefill 需要全量计算Decode 需要增量计算linear attention 层基于 Gated DeltaNet 的线性注意力Prefill 按 chunk 递推Decode 按 token 递推MoE 层256 个专家以 35B-A3B 为例Decode 阶段每个 token 只激活 top-k 专家涉及路由、EP 通信与结果聚合。因此优化不能只盯着 Attention/MLP 的矩阵计算还必须压缩 MoE 路由、EP 通信和归一化等小算子带来的开销。本样例的优化按四个维度组织融合算子用npu_rotary_mul、npu_fused_infer_attention_score_v2、npu_chunk_gated_delta_rule、npu_rms_norm、npu_mm_all_reduce_base、npu_moe_distribute_dispatch_v2/combine_v2等 NPU 融合算子替换多算子展开序列TP/EP 并行通过parallel_config分别控制 attention、MoE、Embedding/LM Head 等模块的切分粒度图模式Decode 阶段单 token 小算子多用 TorchAir 图编译成整图执行缓解 Host-bound矩阵融合将 Gated DeltaNet 的多路输入投影合并为一次MergedColumnParallelLinear。以下按这四个维度展开。二、使能融合算子从算子级削减小算子开销2.1 RoPE 融合算子优化Qwen3.5-MoE 的 full attention 层使用MRoPE位置编码。原始实现通常通过切分、旋转、拼接等多个 Tensor 算子完成 RoPE 计算算子数量多且引入额外数据搬运。本样例在Qwen3_5MoeAttention中使用torch_npu.npu_rotary_mul完成 query 和 key 的旋转位置编码计算源码位于 modeling_qwen3_5_moe.pyq_rot torch_npu.npu_rotary_mul( q_rot.unsqueeze(0), cos.unsqueeze(2), sin.unsqueeze(2), rotary_modehalf ).squeeze(0) k_rot torch_npu.npu_rotary_mul( k_rot.unsqueeze(0), cos.unsqueeze(2), sin.unsqueeze(2), rotary_modehalf ).squeeze(0)关键点在于 Qwen3.5-MoE只对 head 维度中的一部分应用 RoPErotary_dim。代码先将 query/key 拆分为旋转部分和透传部分仅对旋转部分调用融合算子最后再与透传部分拼接。rotary_modehalf表明以半旋转方式施加位置编码与 MRoPE 的部分维度旋转语义对应。该优化为通用优化在 Prefill 和 Decode 阶段均可使用用于减少 RoPE 计算中的小算子数量并提升位置编码计算效率。2.2 FlashAttention 融合算子优化full attention 层同时覆盖 Prefill 全量计算和 Decode 增量计算。若直接使用 PyTorch 原生 attention需要展开为 QK 矩阵乘、mask、softmax、PV 矩阵乘等多个算子显存访问和小算子调度开销较高。本样例使用npu_fused_infer_attention_score_v2作为推理场景下的 FlashAttention 融合算子。full attention 层将 Q/K/V 组织为TND layout并在写入 block 化 KV cache 后通过block_table、block_size以及实际序列长度信息传入算子attn_output, _ fa_ops.npu_fused_infer_attention_score_v2( query_states, k_cache_fa, v_cache_fa, num_query_headsself.num_heads_per_rank, num_key_value_headsself.num_key_value_heads_per_rank, softmax_scaleself.scale_fa, input_layoutTND, sparse_modesparse_mode, atten_maskforward_metadata.attention_mask if is_prefill else None, actual_seq_qlenactual_seq_qlen, actual_seq_kvlenactual_seq_kvlen, block_tableblock_table, block_sizeself.block_size, )实现细节见 modeling_qwen3_5_moe.pyPrefill 阶段通过sparse_mode3和 attention mask 表达因果掩码Q/K/V 以 TND 连续排布传入Decode 阶段先将当前 token 的 K/V 写入 KV cache_update_cache再通过actual_seq_kvlen控制当前 token 可访问的 cache 范围避免计算多余 mask算子路径切换图模式 Decode 场景切换到torchair.ops.npu_fused_infer_attention_score_v2路径保证可被图编译器识别普通 eager 场景使用torch.ops.npu.npu_fused_infer_attention_score_v2路径。2.3 ChunkGDR 融合算子优化Gated DeltaNet PrefillQwen3.5-MoE 的 linear attention 层使用Gated DeltaNet结构。Prefill 阶段需要对整段序列进行 chunk 计算原始 PyTorch 实现通过多次矩阵乘、decay mask 构造、下三角求解等操作完成块内递推涉及较多中间 Tensor 搬运和小算子调度。本样例在Qwen3_5MoeGatedDeltaNet中使用torch_npu.npu_chunk_gated_delta_rule作为 Prefill 阶段 Gated DeltaNet 的融合算子。启用条件为当前为 Prefill 阶段torch_npu支持npu_chunk_gated_delta_rule接口。融合路径中算子接收 TND 格式的 query、key、value 输入以及 beta 门控参数、initial_state递推状态、actual_seq_lengths序列长度信息和 decay 参数 g一次性完成 chunk 内的 L2 归一化、decay 计算、下三角求解和递推状态更新core_attn_out, last_recurrent_state torch_npu.npu_chunk_gated_delta_rule( query.to(torch.bfloat16), key.to(torch.bfloat16), value.to(torch.bfloat16), betabeta.to(torch.bfloat16), initial_stateinitial_state, actual_seq_lengthsactual_seq_lens.to(torch.int32), scalescale, gg.to(torch.float32), )该融合算子将原始 PyTorch 实现中的以下操作合并chunk decay 的 cumsum 和指数计算块内注意力矩阵的下三角求解替代逐行循环更新块间递推状态更新和输出计算。算子输出core_attn_out和last_recurrent_state前者为当前 chunk 的 attention 输出后者用于后续 chunk 计算或传递给 Decode 阶段。与之对应Decode 阶段使用torch_npu.npu_recurrent_gated_delta_rule源码位置逐 token 更新递推状态。该优化主要用于 Linear Attention Prefill 场景减少 Gated DeltaNet 块内递推计算中的小算子数量并提升计算效率。2.4 RMSNorm 融合算子优化RMSNorm 是 Qwen3.5-MoE 中高频出现的归一化操作分布在 attention 前后、linear attention 模块以及最终输出 norm 中。若使用 PyTorch 原生算子展开通常包含平方、均值、rsqrt、乘权重等多个操作。本样例在Qwen3_5MoeRMSNorm中使用torch_npu.npu_rms_norm替换展开计算源码位置return torch_npu.npu_rms_norm(x, self.norm_weight, self.eps)[0]对于带残差输入的场景样例使用torch_npu.npu_add_rms_norm融合残差相加与 RMSNorm 计算y, _, residual torch_npu.npu_add_rms_norm(residual, x, self.norm_weight, self.eps)在 Gated DeltaNet 中Qwen3_5MoeRMSNormGated同样使用torch_npu.npu_rms_norm完成归一化再与门控分支相乘。RMSNorm 相关优化属于通用优化在 TP、EP 以及不同 batch 配置下均可使用。2.5 Matmul AllReduce 融合优化在 TP 并行场景下部分线性层计算后需要执行all_reduce汇聚各卡结果。常规实现会先执行 Matmul再单独执行集合通信。对于通信和计算开销都比较敏感的场景可以使用torch_npu.npu_mm_all_reduce_base将矩阵乘与 all reduce 融合减少中间结果搬运和通信调度开销。本样例封装了qwen3_5_prefill_mm_all_reduce源码位置在满足条件时对需要 TP 汇聚的 RowParallelLinear 输出投影启用 Matmul AllReduce 融合包括 full attention 的o_proj、linear attention 的out_proj以及 dense MLP 的down_projreturn torch_npu.npu_mm_all_reduce_base(input_, layer.weight.data, hcom, reduce_opsum)该融合当前只在满足以下条件时启用开启enable_mm_all_reduce_baseYAML 中custom_params字段仅 Atlas 950 平台可用默认False当前为 Prefill 阶段线性层 TP 切分数大于 1且输入已经按 TP 切分线性层无 bias且不使用skip_bias_add当前线性层为非量化计算路径能够获取对应通信组的 HCCL group name。该优化主要用于 A3 等支持该融合能力且通信收益明显的部署场景。若条件不满足代码会回退到常规线性层计算再按需执行dist.all_reduce。2.6 MoE Dispatch/Combine 融合优化Qwen3.5-MoE 包含大量专家Decode 阶段每个 token 只激活部分专家。纯 EP 部署时token 需要按照专家归属在 EP 组内分发专家计算完成后再按原 token 顺序聚合。如果使用普通all_to_all配合多次重排、索引和 combine 操作会产生较多通信和小算子开销。本样例在纯 EP Decode 场景下提供torch_npu.npu_moe_distribute_dispatch_v2与torch_npu.npu_moe_distribute_combine_v2融合路径源码位置。启用条件为开启enable_decode_moe_dispatch_combine_v2当前为 Decode 阶段moe_ep_size 1moe_tp_size 1即纯 EP 场景。融合路径中npu_moe_distribute_dispatch_v2根据 top-k 专家选择结果完成 token 分发并输出专家计算所需的辅助信息专家计算完成后npu_moe_distribute_combine_v2根据 dispatch 阶段生成的辅助信息完成跨 EP 聚合、权重加权和原序恢复output torch_npu.npu_moe_distribute_dispatch_v2(**dispatch_args) ... return torch_npu.npu_moe_distribute_combine_v2(**combine_args)该优化减少了纯 EP Decode 链路中的手写重排、all_to_all和 finalize routing 开销更适合 Decode 阶段小 batch、低时延场景。三、TP/EP 并行优化3.1 并行策略总览与约束Qwen3.5-MoE 样例通过parallel_config分别控制不同模块的并行策略以 qwen3_5_35b_ep8.yaml 为例parallel_config: world_size: 8 attn_tp_size: 8 moe_tp_size: 1 embed_tp_size: 8 lmhead_tp_size: 8 shared_tp_size: 1attn_tp_size控制 full attention 和 linear attention 的张量并行moe_tp_size控制 MoE 专家内部矩阵的张量并行moe_ep_size world_size // moe_tp_size由框架自动计算用于控制专家并行规模无需在 YAML 中配置。当前仅支持纯 EPmoe_tp_size1、moe_ep_sizeworld_size或纯 TPmoe_tp_sizeworld_size、moe_ep_size1。混合 EPTP 配置会在模型初始化时被拦截源码中通过_validate_qwen3_5_moe_parallel_support显式校验并抛出ValueErrormodeling_qwen3_5_moe.pyif parallel_config.moe_tp_size 1 and parallel_config.moe_ep_size 1: raise ValueError( Qwen3.5 MoE does not support mixed tensor parallelism and expert parallelism: fgot moe_tp_size{parallel_config.moe_tp_size}, fmoe_ep_size{parallel_config.moe_ep_size}. Set moe_tp_size1 for pure EP or moe_tp_sizeworld_size for pure TP. )Embedding、LM Head、dense MLP 和o_proj也分别提供独立 TP 配置embed_tp_size、lmhead_tp_size、shared_tp_size便于根据模型结构和卡数做更细粒度切分。仓库提供多组现成配置EP8 的 qwen3_5_35b_ep8.yaml、TP8 的 qwen3_5_122b_tp8.yaml、TP16 的 qwen3_5_397b_tp16.yaml 以及 MXFP8 在线量化的 mxfp8/qwen3_5_35b_mxfp8_tp4.yaml。3.2 Attention TP 优化切分策略Qwen3.5-MoE 的 full attention 按照 head 维度进行 TP 切分。假设 attention 层的 query 头数为num_heads切分数量为attn_tp_size则每张卡上的 query 头数为num_heads_per_rank num_heads // attn_tp_sizekey/value 头数为num_key_value_heads每张卡上的 key/value 头数为num_key_value_heads_per_rank max(num_key_value_heads // attn_tp_size, 1)。每张卡只计算本 rank 负责的 Q/K/V headattention 输出后再通过o_proj和 all reduce 完成跨 TP 汇聚。Qwen3.5-MoE 还包含 linear attention 层。该模块同样使用attn_tp_size切分linear_num_key_heads和linear_num_value_heads每张卡计算本 rank 负责的 Gated DeltaNet head。代码中会校验linear_num_key_heads和linear_num_value_heads必须能被attn_tp_size整除避免切分后 head 分布不均。计算分解full attention 层的计算链路见下图为先通过merged_qkv_proj一次 Matmul 完成 Q、K、V 投影其中 Q 投影同时输出 attention gate随后将投影结果拆分为 query、key、value对 query 和 key 执行 RMSNorm再对 RoPE 维度调用npu_rotary_mul完成位置编码K/V 写入 block 化 KV cache 后Prefill 和 Decode 阶段统一调用npu_fused_infer_attention_score_v2完成 attention 计算Decode 阶段通过实际 KV 长度控制当前 token 可访问的 cache 范围attention 结果与 gate 相乘后进入o_proj最后通过 all reduce 或npu_mm_all_reduce_base完成 TP 汇聚。linear attention 层的计算链路为in_proj_fused投影 → depthwise conv → Gated DeltaNet 核心计算 → RMSNormGated →out_proj。其中 Q/K/V/Z/B/A 投影通过一次MergedColumnParallelLinear完成并按attn_tp_size切分Prefill 阶段使用 chunk Gated DeltaNet 计算整段序列Decode 阶段使用npu_recurrent_gated_delta_rule更新递推状态out_proj输出后同样执行 TP 汇聚。3.3 MoE TP 优化切分策略MoE TP 用于切分单个专家内部的 FFN 矩阵计算。假设 MoE 层的切分数量为moe_tp_size专家个数为num_experts每个专家都按照 MLP 的 TP 方式切分gate_proj和up_proj按列切分down_proj按行切分。同时gate_proj和up_proj在权重加载后合并为w13_weight减少专家计算中的 Matmul 次数。计算分解每个专家层的原始计算为down_proj(SiLU(gate_proj(x)) * up_proj(x))。如下图所示本样例将gate_proj和up_proj合并为一次 GroupedMatmul 计算输出结果通过npu_swiglu完成SiLU(gate) * up再通过第二次 GroupedMatmul 完成down_proj计算。专家输出经过npu_moe_finalize_routing恢复 token 顺序并按路由权重聚合当moe_tp_size 1时最后通过moe_tp_group执行 all reduce 汇聚各 TP 分片结果。该方式适合专家权重较大、需要降低单卡显存占用的场景当moe_tp_size world_size时moe_ep_size为 1所有专家在 TP 组内切分计算不发生 EP 维度的专家分发。3.4 MoE EP 优化MoE EP 用于将路由专家均匀分布到不同 rank 上。代码中每张卡只加载本 rank 负责的专家范围experts_per_rank num_experts // moe_ep_size local_expert_start moe_ep_rank * experts_per_rank local_expert_end local_expert_start experts_per_rank普通 EP 路径采用Double-Routing方式完成 token 分发、专家计算和结果恢复流程如下图所示图中编号与下方步骤一一对应主要计算步骤如下调用npu_moe_init_routing_v2源码位置对本 rank 上的 token 按照 top-k 专家进行展开和排序得到expanded_x同时获取每个专家需要计算的 token 数量tokens_per_expert以及用于恢复原始 token 顺序的expanded_row_idx。对tokens_per_expert执行 EP 组内 AlltoAll计算当前 rank 需要发送给其他 rank 的input_splits以及当前 rank 需要从其他 rank 接收的output_splits。调用dist.all_to_all_single按照input_splits和output_splits在moe_ep_group内分发 token使每张卡获取本地专家需要处理的 token。调用npu_moe_re_routing源码位置对接收到的 token 按照本地专家重新排序得到hidden_states_ordered_by_experts、tokens_per_local_expert以及用于恢复本地接收顺序的gathered_ids_unsort。调用专家计算模块使用 GroupedMatmul 完成本 rank 本地专家的 FFN 计算。使用torch.index_select按照gathered_ids_unsort恢复本地接收 token 顺序。再次调用dist.all_to_all_single将专家计算结果从本地专家 rank 发送回原始 token 所在 rank。调用npu_moe_finalize_routing源码位置根据第 1 步保存的expanded_row_idx恢复 token 原始顺序并按照 router 权重聚合 top-k 专家输出。在纯 EP Decode 场景moe_ep_size 1且moe_tp_size 1下普通 Double-Routing 链路还可以替换为 2.6 节介绍的npu_moe_distribute_dispatch_v2和npu_moe_distribute_combine_v2融合路径进一步减少 Decode 阶段通信路由开销。四、图模式Decode 阶段整图编译Qwen3.5-MoE 在 Decode 阶段以单 token 增量推理为主算子粒度较小在 NPU 上容易出现Host-bound问题。为减少 CPU 到 NPU 的逐算子下发与同步开销样例结合torch_npu的TorchAir 图模式编译能力将 Decode 阶段尽可能编译为整图执行从而降低调度开销并提升低时延推理性能。当前图模式仅支持 TP 并行配置。针对图模式的核心处理是Prefill 阶段仍使用普通model.forward完成整段输入计算和状态初始化Decode 阶段进入单 token 生成后如果exe_mode为ge_graphnpugraph_ex模式暂未验证则调用编译后的model_compiled执行if self.exe_mode in [ge_graph, npugraph_ex] and not is_prefill: logits self.model_compiled(**model_inputs) else: logits self.model(**model_inputs)图编译在warm-up 阶段触发。执行器先运行一次 Prefill 初始化 KV cache 和 linear attention 状态再使用单 token 输入触发 Decode 图编译保证正式推理时可以直接使用已编译的 Decode 图if self.exe_mode in [ge_graph, npugraph_ex]: self.main_worker.compile_model()从源码还可以看到两个与图模式配套的细节算子兼容处理为保证ge_graph能正确 lower代码中如ge_safe_softplus一类函数刻意避开aten::softplus改用torch.relu(x) torch.log1p(torch.exp(-torch.abs(x)))等价实现modeling_qwen3_5_moe.py这类图模式安全实现遍布 Decode 链路平台差异配置_configure_qwen3_5_npugraph在非 950 平台上为npugraph_ex关闭内部 format 与 jit 编译相关选项源码位置。对于固定配置反复启动的推理任务可以通过enable_cache_compileYAML 中model_config字段开启图编译缓存减少后续启动时的编译等待时间。配置示例见 qwen3_5_122b_tp8.yaml其exe_mode: ge_graph即启用 Decode 图编译。五、矩阵融合优化Gated DeltaNet 输入投影融合Qwen3.5-MoE 的 linear attention 层包含 Gated DeltaNet 结构。原始实现中in_proj_qkv、in_proj_z、in_proj_b和in_proj_a是多次独立线性层计算mixed_qkv self.in_proj_qkv(hidden_states) z self.in_proj_z(hidden_states) b self.in_proj_b(hidden_states) a self.in_proj_a(hidden_states)当前样例将上述输入投影合并为in_proj_fused通过一次MergedColumnParallelLinear同时计算 Q/K/V 分支、Z 门控分支以及 B/A 分支再按输出维度拆分fused_proj self.in_proj_fused(hidden_states) mixed_qkv, z, b, a torch.split( fused_proj, [self.key_dim * 2 self.value_dim, self.value_dim, self.num_v_heads, self.num_v_heads], dim-1, )该优化减少了 Gated DeltaNet 输入侧的 Matmul 次数也让mixed_qkv、z、b和a走同一套 TP 切分、量化和输出 dtype 路径避免多路独立投影在量化或并行切分场景下出现 dtype 不一致。随后mixed_qkv进入 depthwise conv 和 Gated DeltaNet 核心计算z作为 RMSNormGated 的门控输入参与输出归一化b和a用于生成 Gated DeltaNet 中的递推参数。六、配置速查与适用范围小结将上述优化落到实际推理时关键开关集中在 YAML 配置中。以 qwen3_5_35b_ep8.yaml 与 qwen3_5_35b_mxfp8_tp4.yaml 为例配置项位置含义与适用范围exe_modemodel_configeager/ge_graph/npugraph_exDecode 图编译开关MXFP8 在 Atlas 950 上暂不支持ge_graphenable_mm_all_reduce_basemodel_config.custom_params启用 MatmulAllReduce 融合仅 Atlas 950 平台可用enable_online_mxfp8_quantizationmodel_config.custom_params加载 BF16/FP16 权重后在线转为 MXFP8与online_mxfp8_quant_layers、online_mxfp8_ignored_layers配合enable_cache_compilemodel_config开启图编译缓存减少固定配置反复启动时的编译等待attn_tp_size/moe_tp_sizeparallel_configattention 与 MoE 的并行切分MoE 仅支持纯 EP 或纯 TPembed_tp_size/lmhead_tp_size/shared_tp_sizeparallel_configEmbedding、LM Head、dense MLP 与o_proj的独立 TP 配置各优化点的适用范围可归纳为通用优化Prefill DecodeTP/EP 均适用RoPE 融合、FlashAttention 融合、RMSNorm 融合Prefill 专用ChunkGDR 融合、Matmul AllReduce 融合Decode 专用MoE Dispatch/Combine 融合纯 EP、图模式整图编译纯 TP平台相关npu_mm_all_reduce_base仅 950 可用FP8/MXFP8 量化样例仅支持 Atlas 950 系列产品A3 平台推理时建议设置HCCL_DETERMINISTICtrue以确保精度详见 Qwen3.5-MoE 推理实现 README。综合来看Qwen3.5-MoE 的 NPU 优化本质上是按阶段Prefill/Decode、按并行策略TP/EP、按平台A3/950三重维度精确选择融合算子与执行模式的过程。读者在实际部署时可以对照上述配置速查表结合 YAML 参数描述 与 Qwen3.5-MoE 推理实现 中给出的多组现成配置快速组装出适配自身模型规模与卡数的优化组合。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考