企业官网建设流程全解析

Qwen3.5-MoE 在 NPU 上的推理性能优化实践：融合算子、TP/EP 并行与图模式全解析

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

Qwen3.5-MoE 是同时包含 full attention（MRoPE）、linear attention（Gated 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 通信和归一化等小算子带来的开销。本样例的优化按四个维度组织：

1. 融合算子：用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 融合算子替换多算子展开序列；
2. TP/EP 并行：通过parallel_config分别控制 attention、MoE、Embedding/LM Head 等模块的切分粒度；
3. 图模式：Decode 阶段单 token 小算子多，用 TorchAir 图编译成整图执行，缓解 Host-bound；
4. 矩阵融合：将 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.py：

q_rot = torch_npu.npu_rotary_mul( q_rot.unsqueeze(0), cos.unsqueeze(2), sin.unsqueeze(2), rotary_mode='half' ).squeeze(0) k_rot = torch_npu.npu_rotary_mul( k_rot.unsqueeze(0), cos.unsqueeze(2), sin.unsqueeze(2), rotary_mode='half' ).squeeze(0)

关键点在于 Qwen3.5-MoE只对 head 维度中的一部分应用 RoPE（rotary_dim）。代码先将 query/key 拆分为旋转部分和透传部分，仅对旋转部分调用融合算子，最后再与透传部分拼接。rotary_mode='half'表明以"半旋转"方式施加位置编码，与 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_heads=self.num_heads_per_rank, num_key_value_heads=self.num_key_value_heads_per_rank, softmax_scale=self.scale_fa, input_layout="TND", sparse_mode=sparse_mode, atten_mask=forward_metadata.attention_mask if is_prefill else None, actual_seq_qlen=actual_seq_qlen, actual_seq_kvlen=actual_seq_kvlen, block_table=block_table, block_size=self.block_size, )

实现细节（见 modeling_qwen3_5_moe.py）：

• Prefill 阶段：通过sparse_mode=3和 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 Prefill）

Qwen3.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), beta=beta.to(torch.bfloat16), initial_state=initial_state, actual_seq_lengths=actual_seq_lens.to(torch.int32), scale=scale, g=g.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_proj：

return torch_npu.npu_mm_all_reduce_base(input_, layer.weight.data, hcom, reduce_op="sum")

该融合当前只在满足以下条件时启用：

• 开启enable_mm_all_reduce_base（YAML 中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 > 1；
• moe_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: 1

• attn_tp_size：控制 full attention 和 linear attention 的张量并行；
• moe_tp_size：控制 MoE 专家内部矩阵的张量并行；
• moe_ep_size = world_size // moe_tp_size：由框架自动计算，用于控制专家并行规模，无需在 YAML 中配置。

当前仅支持纯 EP（moe_tp_size=1、moe_ep_size=world_size）或纯 TP（moe_tp_size=world_size、moe_ep_size=1）。混合 EP+TP 配置会在模型初始化时被拦截，源码中通过_validate_qwen3_5_moe_parallel_support显式校验并抛出ValueError（modeling_qwen3_5_moe.py）：

if 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: " f"got moe_tp_size={parallel_config.moe_tp_size}, " f"moe_ep_size={parallel_config.moe_ep_size}. " "Set moe_tp_size=1 for pure EP or moe_tp_size=world_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_size；key/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 head，attention 输出后再通过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 分发、专家计算和结果恢复，流程如下图所示（图中编号与下方步骤一一对应）：

主要计算步骤如下：

1. 调用npu_moe_init_routing_v2（源码位置），对本 rank 上的 token 按照 top-k 专家进行展开和排序，得到expanded_x，同时获取每个专家需要计算的 token 数量tokens_per_expert以及用于恢复原始 token 顺序的expanded_row_idx。
2. 对tokens_per_expert执行 EP 组内 AlltoAll，计算当前 rank 需要发送给其他 rank 的input_splits，以及当前 rank 需要从其他 rank 接收的output_splits。
3. 调用dist.all_to_all_single，按照input_splits和output_splits在moe_ep_group内分发 token，使每张卡获取本地专家需要处理的 token。
4. 调用npu_moe_re_routing（源码位置），对接收到的 token 按照本地专家重新排序，得到hidden_states_ordered_by_experts、tokens_per_local_expert以及用于恢复本地接收顺序的gathered_ids_unsort。
5. 调用专家计算模块，使用 GroupedMatmul 完成本 rank 本地专家的 FFN 计算。
6. 使用torch.index_select按照gathered_ids_unsort恢复本地接收 token 顺序。
7. 再次调用dist.all_to_all_single，将专家计算结果从本地专家 rank 发送回原始 token 所在 rank。
8. 调用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_graph（npugraph_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_compile（YAML 中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 为例：

各优化点的适用范围可归纳为：

• 通用优化（Prefill + Decode，TP/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_DETERMINISTIC=true以确保精度（详见 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