尧图网站建设 
HunyuanImage-3.0 在 CANN NPU 上的推理优化实践TP/EP 并行切分、通算融合与冗余计算消除【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本技术指南基于 cann-recipes-infer 仓库中的优化文档系统梳理 HunyuanImage-3.0 原生多模态模型在 Atlas A2/A3 系列产品上的 NPU 推理优化方案。文章覆盖 Attention/MoE 的 TP 与 EP 切分策略、VAE/CFG 多流并行、torch_npu 融合算子使能以及冗余计算消除四大部分并补充仓库内真实配置与源码作为佐证。读完本文你将掌握该模型在 NPU 上的并行切分原理、关键融合算子用法与可落地的性能调优思路。一、优化背景与总体思路HunyuanImage-3.0 是腾讯开源的、在自回归框架内统一多模态理解与生成任务的原生多模态模型其文生图能力在开源模型中具备较强竞争力。在 CANN 平台上适配该模型推理时性能优化的总体思路分为四条主线切分策略对 Attention 与 MoE 层分别设计 TPTensor Parallel张量并行或 EPExpert Parallel专家并行切分方案让大规模权重与计算分摊到多张 NPU 上并行优化在 MoE 通信、VAE 解码、CFG 采样等环节引入多流与空间并行机制让通信与计算互相掩盖使能融合算子用 torch_npu 提供的高性能融合算子替换朴素算子序列减少算子下发次数、提高搬运效率消除冗余计算清理由数据类型隐式转换、CPU 侧同步计算带来的冗余 cast 算子与 Host2Device 开销。这四条主线在原优化文档hunyuan_image_3_optimization.md中被完整阐述下文逐一展开。二、切分策略Attention TP 与 MoE TP/EP2.1 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层按行切分即可完成张量并行下的输出合并。计算分解在 TP 切分的基础上优化策略先将 Q、K、V 三个线性层合并为一次 Matmul 计算merged_qkv_proj以提升计算性能。merged_qkv_proj的输出按 Q、K、V 拆分后对 Q 和 V 进行归一化操作并使用旋转位置编码再计算 attentionFused_infer_attention_score最后经o_proj输出。该分解在models/hunyuan-image-3.0/adaptor_patches/下的适配补丁中落地对应文档中的attention_calcu.png图示见 figures 目录。2.2 MoE TP 优化切分策略假设 MoE 层的切分数量为moe_tp_size专家个数为expert_num。对 MoE 层进行张量切分时对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 与 up 合并计算减少算子下发与中间张量搬运从而提升整体计算效率。具体实现位于HunyuanMoE类的相关模式中配合npu_grouped_matmul使用详见后文 4.2 节。2.3 MoE EP 优化DoubleRouting 方案在前述 TP 切分策略下MoE 的 routing 与 finalizing 开销与切分份数无关——不论 RANK 数为多少这两部分开销都保持一致。随着 RANK 数增加这部分固定开销占比逐渐提升且无法通过 TP 优化。因此对 MoE 部分改用EP 切分方案对 MoE 的 gating 切分 S 轴之后的 routing 与 finalizing 计算量均变为原来的1/ep_size倍。此时 MoE 部分需要采用DoubleRouting方案完整流程如下在 GatingTopK 之后使用InitRouting确定当前 RANK 的 token 需要分发给哪些专家使用all_to_all通信将tokens_per_expert发送至对应 RANK使用all_to_all通信将 token 发送至对应 RANK执行LocalRerouting在每个 RANK 内按照专家序对 token 进行重排执行GroupedMatMul进行专家计算执行LocalFinalizeRerouting排序回第二次all_to_all后的排布使用all_to_all把计算完的 token 发送回原本的 TP 域内对应 RANK执行FinalizeRerouting对专家计算的结果在 TP 域内进行累加。为了优化通信量配合此 EP DoubleRouting 方案在q_proj前使用AllGather对hidden_states进行聚合o_proj后使用ReduceScatter对 Attention 的结果进行分发。这样为通算融合创造了机会——AllGather可与qkv_proj融合、o_proj可与ReduceScatter融合详细说明见 4.1 节。需要说明的是MoE 的 TP 与 EP 切分互斥权重转换脚本 convert_model.py 通过--tp-moe与--ep两个参数分别指定二者只能选一且 Attention 部分的 TP 并行度需与 MoE 部分的并行度保持一致。三、并行优化多流、空间并行与 CFG 并行3.1 MoE 多流并行优化在 MoE 模块中路由专家采用 TP 部署。共享专家的计算与路由专家完成 MoE 计算后的通信可以通过多流并行机制使二者流水掩盖使用all_reduce的异步机制实现。多流的排布方式moe_mlp_parallel.png展示了共享专家计算流与 MoE 通信流如何并行交错、互相掩盖从而压缩关键路径耗时。多流机制在仓库executor/utils/stream_utils.py等工具中有配套的流管理实现供参考。3.2 VAE 并行优化本样例对 VAE 并行进行了使能通过空间并行的方式实现将大尺寸图像在高度和宽度维度上切分成多个块分配给不同的 NPU 进程并行处理从而加速原模型 VAE 推理。该优化在 module/vae_patch_parallel.py 中实现。从源码结构看Parallel_VAE_SP类以h_split × w_split的网格将进程组织起来输入张量沿空间维度切分宽度维dim-1按w_split切、高度维dim-2按h_split切每个进程只处理自己的局部块计算过程中按操作特点选择通信策略卷积操作通过与邻居进程交换边界数据来获取上下文注意力操作通过全局收集所有 K、V 张量来保证计算完整性插值操作通过扩展边界、计算后再裁剪来处理上采样最后通过两阶段的收集过程按行/列通信组还原将各进程的局部结果按照原始的空间位置重新拼接成完整的输出张量。VAE 空间并行的完整执行路径即文档中的vae_parallel.png流程图。在 pipeline 适配代码 hunyuan_image_3_pipeline.py 的init_vae_parallel中会根据分布式 world size 自动确定切分网格2 卡为1×2、4 卡为2×2、8 卡为4×2其他卡数按h_split sqrt(world_size)估算并通过set_vae_patch_parallel对 VAE 的decoder.forward打补丁启用。该并行由环境变量USE_VAE_PARALLEL控制见 ep8_cfg.yaml。3.3 CFG 并行优化CFGClassifier Free Guidance在生成一张图时需要同时执行一个无条件生成和一个有条件生成相当于网络跑了两个样本。本样例的 CFG 并行把这些样本拆分到两张卡上分别执行然后通过all_gather通信互相获取对方的计算结果得到最终的预测结果。以 TP2 CFGP 为例排布方式如下在 hunyuan_image_3_pipeline.py 的采样循环中可以看到当self.model.cfg_parallel_size 1时cfg_factor保持为 1batch 无需再乘 2条件与无条件输出分别由不同 rank 产生随后通过torch.distributed.all_gather汇聚成pred_cond与pred_uncond再做 guidance 融合。开启 CFG 并行后world_size需为原 TP 规模的 2 倍——预置配置 ep8_cfg.yaml 中CFG_PARALLEL: 1且world_size: 16Attn TP8 × 2即基于此设计。四、使能融合算子4.1 通算融合AllGatherMM 与 MMReduceScatter在 EP 方案中attention 部分的qkv_proj之前需要先对hidden_states进行AllGathero_proj之后又需要对attn_output进行ReduceScatter。若计算与通信串行执行会存在大量性能浪费。torch_npu 提供的通算融合算子可以同时完成通信与矩阵乘以掩盖通信开销、缩短端到端耗时。AllGatherMM对于AllGather与qkv_proj可考虑使用torch_npu.npu_all_gather_base_mm替换。使用时需注意该算子仅支持 2 维 tensor 作为输入且只能沿着第 0 轴切分调用前需要先将hidden_states的 layout 由[b s h]转为[(s b) h]调用之后再转回[b s h]调用时 s 轴必须在最内侧否则会出现精度问题。MMReduceScatter对于o_proj与ReduceScatter可考虑使用torch_npu.npu_mm_reduce_scatter_base替换。使用时需注意同样仅支持 2 维 tensor 输入只能沿第 0 轴切分调用前需要先将hidden_states的 layout 由[b n s d]转为[(s b) (n d)]调用之后再转回[b s (n d)]s 轴必须在最内侧否则会出现精度问题该融合算子在 A2 平台下暂未使能使用前需确认目标产品型号。4.2 GMM 使能与 Routing 优化在 MoE 模块中如果通过 for 循环逐个处理每个专家、单独计算expert_num个 FFN计算效率较低。CANN 提供GroupedMatmul算子可以同时计算多个专家提高计算和搬运效率。具体实现可参考HunyuanMoE类中npu_grouped_matmul模式下的实现对应推理参数--moe-impl npu_grouped_matmul。配合该模式Routing 阶段使能以下 torch_npu 融合算子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实现多个专家的矩阵乘计算提高计算和搬运效率。这组算子与 2.3 节 DoubleRouting 流程中的 InitRouting / LocalRerouting / LocalFinalizeRerouting / FinalizeRerouting 一一对应构成 EP 方案中 token 路由与专家计算的高性能实现。4.3 RmsNorm 算子优化通过使能torch_npu.npu_rms_norm算子能够提升模型推理性能。RmsNorm 是大模型常用的归一化操作相比 LayerNorm其去掉了减去均值的部分仅对均方根进行归一化在 NPU 上以融合算子方式实现可以显著减少中间张量与算子下发开销。五、消除冗余计算5.1 消除旋转位置编码冗余 cast 算子在旋转位置编码计算时query_states、key_states是实时获取的数据类型为 bfloat16。原始实现中cos、sin为 float32会触发从 float32 自动转换为 bfloat16 的冗余 cast 算子# 将cos, sin转成bfloat16 cos real_batched_index_select(cos, dim1, idxposition_ids).to(torch.bfloat16) sin real_batched_index_select(sin, dim1, idxposition_ids).to(torch.bfloat16) # 旋转位置编码计算使用cos, sin if self.use_rotary_pos_emb: cos, sin custom_pos_emb query_states, key_states apply_rotary_pos_emb(query_states, key_states, cos, sin)为消除冗余 cast 算子在HunyuanImage3ForCausalMM类中将cos、sin提前转为 bfloat16使编码计算全程保持同一数据类型减少 graph 中的转换节点。5.2 消除 gate 计算冗余 cast 算子外部开启了torch.autocast且数据类型设置为 bfloat16 时即使中间数据被设置为 float32计算时依然会被转成 bfloat16从而出现多个冗余的 cast 算子。为保持 gate 计算的精度gate 计算部分屏蔽外层的torch.autocast能力# 屏蔽外层的torch.autocast能力 with torch.npu.amp.autocast(enabledFalse): if self.wg.weight.dtype torch.float32: hidden_states hidden_states.float() logits self.wg(hidden_states)这样既保证了 gate 输出精度的稳定也避免了无关的 dtype 转换节点进入计算图。5.3 消除 FA 中 scale 的冗余计算每次调用 FA 算子时默认会计算一次scale 1 / math.sqrt(d)。这个操作在 CPU 上进行会导致下发中断和 device 的等待造成性能浪费。由于d是一个在 config 文件中指定的固定整数完全可以在初始化时将scale计算出来并在后续调用 FA 的地方复用。在本项目中此scale的计算被提前到HunyuanImage3Model.__init__()中作为成员变量存在后续调用flash_attn_func_npu时均直接使用、不再重复计算。5.4 消除 timestep_index 的冗余计算在 DecoderLayer 中计算 FA 前需要获取timestep_index用于确定casual_len其获取方式是timestep_index gen_timestep_scatter_index[0, 0].item()此处的.item()操作是一个强同步操作会造成 device 计算的中断。事实上gen_timestep_scatter_index的值在HunyuanImage3Text2ImagePipeline.__call__()的开头即确定下来并不再发生改变。因此可以将timestep_index的初始化提前至HunyuanImage3Text2ImagePipeline.__call__()的开头并透传至casual_len的计算处避免反复的 Host2Device 操作。在仓库适配代码 hunyuan_image_3_pipeline.py 的__call__实现中可以看到timestep_index gen_timestep_scatter_index[0, 0].item()在采样循环开始前被一次性计算并通过model_kwargs[timestep_index] timestep_index透传给后续各层正是这一优化思路的落地形态。六、工程落地配置、权重转换与推理执行优化文档附录指向环境部署与样例执行说明models/hunyuan-image-3.0/README.md此处补充与上述优化点配套的关键工程环节6.1 环境与权重准备要点支持 Atlas A2/A3 系列产品依赖 CANN 9.0.0-beta.1 软件包、Ascend Extension for PyTorch v7.3.1PyTorch 2.7.1Python 3.11需要将 HunyuanImage-3.0 开源仓库固定 commit62da220178f4b0b7d83e91665a46a20a3ee4f7cd以“非覆盖模式”复制到models/hunyuan-image-3.0/目录再由adaptor_patches/中的补丁代码model_adaptor.py完成模型、pipeline 的 NPU 适配权重转换使用 weight_convert.sh 拉起 convert_model.py关键参数为--tp-attnAttention TP 份数、--tp-moeMoE TP 份数、--epMoE EP 份数、--max-shard-size单块权重上限单位 G。约束为MoE 的 TP 与 EP 只能二选一Attention TP 并行度需与 MoE 并行度一致。6.2 推理配置与启动推理参数集中在 config/ep8_cfg.yaml 维护由 infer.sh 调用executor/scripts/mm_function.sh拉起。预置配置为Attn TP8 MoE EP8 CFG 并行 VAE 并行共 16 卡关键字段如下model_name: hunyuan-image-3.0 world_size: 16 master_port: 10086 entry_script: run_image_gen.py env_vars: CFG_PARALLEL: 1 # 开启 CFG 并行nproc_per_node 需为原来的 2 倍 USE_VAE_PARALLEL: 1 # 开启 VAE 并行 CPU_AFFINITY_CONF: 2 # 自动绑核缓解 HOST 下发瓶颈 model_args: reproduce: true model-id: ./ckpts/weight_ep8 # 指向权重转换后的目录 prompt: A cinematic medium shot captures a single Asian woman seated on a chair within a dimly lit room, creating an intimate and theatrical atmosphere. attn-impl: npu # 使能 NPU 上的 Attention 实现 moe-impl: npu_grouped_matmul # 使能 NPU 上的 MoE 实现 moe-ep: true # MoE 使用 EP与 moe-tp 互斥 seed: 42 diff-infer-steps: 50 image-size: 1024x1024 verbose: 0其中与本文优化点直接对应的是attn-impl: npu使能 NPU 的 Attention 实现对应 2.1 节 Attention 计算分解moe-impl: npu_grouped_matmul使能 NPU 的 MoE 实现对应 4.2 节 GMM 使能moe-ep: trueMoE 采用 EP 切分对应 2.3 节 DoubleRouting与moe-tp互斥需与权重转换时选择的并行方式一致CFG_PARALLEL/USE_VAE_PARALLEL对应 3.2、3.3 节的 VAE 与 CFG 并行开关。YAML 字段会透传为run_image_gen.py入口脚本的命令行参数torchrun --nproc_per_nodeworld_size run_image_gen.py model_args。注意模型最少需 4 个 Device 正常运行若开启CFG_PARALLEL则需 8 个 Device。根据仓库 README.md 记录在 A3 环境、CFG 并行与 VAE 并行均开启的条件下Attn TP8 MoE TP8 的端到端耗时约为 10.3sAttn TP8 MoE EP8 约为 9.9s可作为性能基线参考。七、总结HunyuanImage-3.0 在 CANN NPU 上的推理优化是一套组合拳并行切分层面Attention 以 TP 为主QKV 头切分 线性层行切分 QKV 合并计算MoE 在 TP 基础上演进为 EP DoubleRouting使 routing/finalizing 开销随 EP 份数线性下降并通过 AllGather/ReduceScatter 与矩阵乘融合压缩通信成本并行机制层面MoE 共享专家与通信多流掩盖、VAE 空间并行、CFG 双样本跨卡并行共同缩短端到端时间算子层面npu_swiglu、npu_grouped_matmul、MoE routing 系列融合算子、npu_rms_norm、npu_all_gather_base_mm、npu_mm_reduce_scatter_base等 torch_npu 融合算子替代朴素算子序列计算图层面消除 cos/sin 与 gate 计算中的冗余 cast、FA scale 的 CPU 重复计算、timestep_index的强同步.item()操作等冗余节点。上述优化均可在仓库的适配补丁adaptor_patches/、VAE 并行实现module/vae_patch_parallel.py、配置文件config/ep8_cfg.yaml与入口脚本run_image_gen.py中找到对应实现是理解并复现该模型 NPU 推理加速路径的完整参考。环境部署与样例执行的详细步骤可继续阅读 models/hunyuan-image-3.0/README.md。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 HunyuanImage-3.0 在 CANN NPU 上的推理优化实践：TP/EP 并行切分、通算融合与冗余计算消除
HunyuanImage-3.0 在 CANN NPU 上的推理优化实践TP/EP 并行切分、通算融合与冗余计算消除【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本技术指南基于 cann-recipes-infer 仓库中的优化文档系统梳理 HunyuanImage-3.0 原生多模态模型在 Atlas A2/A3 系列产品上的 NPU 推理优化方案。文章覆盖 Attention/MoE 的 TP 与 EP 切分策略、VAE/CFG 多流并行、torch_npu 融合算子使能以及冗余计算消除四大部分并补充仓库内真实配置与源码作为佐证。读完本文你将掌握该模型在 NPU 上的并行切分原理、关键融合算子用法与可落地的性能调优思路。一、优化背景与总体思路HunyuanImage-3.0 是腾讯开源的、在自回归框架内统一多模态理解与生成任务的原生多模态模型其文生图能力在开源模型中具备较强竞争力。在 CANN 平台上适配该模型推理时性能优化的总体思路分为四条主线切分策略对 Attention 与 MoE 层分别设计 TPTensor Parallel张量并行或 EPExpert Parallel专家并行切分方案让大规模权重与计算分摊到多张 NPU 上并行优化在 MoE 通信、VAE 解码、CFG 采样等环节引入多流与空间并行机制让通信与计算互相掩盖使能融合算子用 torch_npu 提供的高性能融合算子替换朴素算子序列减少算子下发次数、提高搬运效率消除冗余计算清理由数据类型隐式转换、CPU 侧同步计算带来的冗余 cast 算子与 Host2Device 开销。这四条主线在原优化文档hunyuan_image_3_optimization.md中被完整阐述下文逐一展开。二、切分策略Attention TP 与 MoE TP/EP2.1 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层按行切分即可完成张量并行下的输出合并。计算分解在 TP 切分的基础上优化策略先将 Q、K、V 三个线性层合并为一次 Matmul 计算merged_qkv_proj以提升计算性能。merged_qkv_proj的输出按 Q、K、V 拆分后对 Q 和 V 进行归一化操作并使用旋转位置编码再计算 attentionFused_infer_attention_score最后经o_proj输出。该分解在models/hunyuan-image-3.0/adaptor_patches/下的适配补丁中落地对应文档中的attention_calcu.png图示见 figures 目录。2.2 MoE TP 优化切分策略假设 MoE 层的切分数量为moe_tp_size专家个数为expert_num。对 MoE 层进行张量切分时对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 与 up 合并计算减少算子下发与中间张量搬运从而提升整体计算效率。具体实现位于HunyuanMoE类的相关模式中配合npu_grouped_matmul使用详见后文 4.2 节。2.3 MoE EP 优化DoubleRouting 方案在前述 TP 切分策略下MoE 的 routing 与 finalizing 开销与切分份数无关——不论 RANK 数为多少这两部分开销都保持一致。随着 RANK 数增加这部分固定开销占比逐渐提升且无法通过 TP 优化。因此对 MoE 部分改用EP 切分方案对 MoE 的 gating 切分 S 轴之后的 routing 与 finalizing 计算量均变为原来的1/ep_size倍。此时 MoE 部分需要采用DoubleRouting方案完整流程如下在 GatingTopK 之后使用InitRouting确定当前 RANK 的 token 需要分发给哪些专家使用all_to_all通信将tokens_per_expert发送至对应 RANK使用all_to_all通信将 token 发送至对应 RANK执行LocalRerouting在每个 RANK 内按照专家序对 token 进行重排执行GroupedMatMul进行专家计算执行LocalFinalizeRerouting排序回第二次all_to_all后的排布使用all_to_all把计算完的 token 发送回原本的 TP 域内对应 RANK执行FinalizeRerouting对专家计算的结果在 TP 域内进行累加。为了优化通信量配合此 EP DoubleRouting 方案在q_proj前使用AllGather对hidden_states进行聚合o_proj后使用ReduceScatter对 Attention 的结果进行分发。这样为通算融合创造了机会——AllGather可与qkv_proj融合、o_proj可与ReduceScatter融合详细说明见 4.1 节。需要说明的是MoE 的 TP 与 EP 切分互斥权重转换脚本 convert_model.py 通过--tp-moe与--ep两个参数分别指定二者只能选一且 Attention 部分的 TP 并行度需与 MoE 部分的并行度保持一致。三、并行优化多流、空间并行与 CFG 并行3.1 MoE 多流并行优化在 MoE 模块中路由专家采用 TP 部署。共享专家的计算与路由专家完成 MoE 计算后的通信可以通过多流并行机制使二者流水掩盖使用all_reduce的异步机制实现。多流的排布方式moe_mlp_parallel.png展示了共享专家计算流与 MoE 通信流如何并行交错、互相掩盖从而压缩关键路径耗时。多流机制在仓库executor/utils/stream_utils.py等工具中有配套的流管理实现供参考。3.2 VAE 并行优化本样例对 VAE 并行进行了使能通过空间并行的方式实现将大尺寸图像在高度和宽度维度上切分成多个块分配给不同的 NPU 进程并行处理从而加速原模型 VAE 推理。该优化在 module/vae_patch_parallel.py 中实现。从源码结构看Parallel_VAE_SP类以h_split × w_split的网格将进程组织起来输入张量沿空间维度切分宽度维dim-1按w_split切、高度维dim-2按h_split切每个进程只处理自己的局部块计算过程中按操作特点选择通信策略卷积操作通过与邻居进程交换边界数据来获取上下文注意力操作通过全局收集所有 K、V 张量来保证计算完整性插值操作通过扩展边界、计算后再裁剪来处理上采样最后通过两阶段的收集过程按行/列通信组还原将各进程的局部结果按照原始的空间位置重新拼接成完整的输出张量。VAE 空间并行的完整执行路径即文档中的vae_parallel.png流程图。在 pipeline 适配代码 hunyuan_image_3_pipeline.py 的init_vae_parallel中会根据分布式 world size 自动确定切分网格2 卡为1×2、4 卡为2×2、8 卡为4×2其他卡数按h_split sqrt(world_size)估算并通过set_vae_patch_parallel对 VAE 的decoder.forward打补丁启用。该并行由环境变量USE_VAE_PARALLEL控制见 ep8_cfg.yaml。3.3 CFG 并行优化CFGClassifier Free Guidance在生成一张图时需要同时执行一个无条件生成和一个有条件生成相当于网络跑了两个样本。本样例的 CFG 并行把这些样本拆分到两张卡上分别执行然后通过all_gather通信互相获取对方的计算结果得到最终的预测结果。以 TP2 CFGP 为例排布方式如下在 hunyuan_image_3_pipeline.py 的采样循环中可以看到当self.model.cfg_parallel_size 1时cfg_factor保持为 1batch 无需再乘 2条件与无条件输出分别由不同 rank 产生随后通过torch.distributed.all_gather汇聚成pred_cond与pred_uncond再做 guidance 融合。开启 CFG 并行后world_size需为原 TP 规模的 2 倍——预置配置 ep8_cfg.yaml 中CFG_PARALLEL: 1且world_size: 16Attn TP8 × 2即基于此设计。四、使能融合算子4.1 通算融合AllGatherMM 与 MMReduceScatter在 EP 方案中attention 部分的qkv_proj之前需要先对hidden_states进行AllGathero_proj之后又需要对attn_output进行ReduceScatter。若计算与通信串行执行会存在大量性能浪费。torch_npu 提供的通算融合算子可以同时完成通信与矩阵乘以掩盖通信开销、缩短端到端耗时。AllGatherMM对于AllGather与qkv_proj可考虑使用torch_npu.npu_all_gather_base_mm替换。使用时需注意该算子仅支持 2 维 tensor 作为输入且只能沿着第 0 轴切分调用前需要先将hidden_states的 layout 由[b s h]转为[(s b) h]调用之后再转回[b s h]调用时 s 轴必须在最内侧否则会出现精度问题。MMReduceScatter对于o_proj与ReduceScatter可考虑使用torch_npu.npu_mm_reduce_scatter_base替换。使用时需注意同样仅支持 2 维 tensor 输入只能沿第 0 轴切分调用前需要先将hidden_states的 layout 由[b n s d]转为[(s b) (n d)]调用之后再转回[b s (n d)]s 轴必须在最内侧否则会出现精度问题该融合算子在 A2 平台下暂未使能使用前需确认目标产品型号。4.2 GMM 使能与 Routing 优化在 MoE 模块中如果通过 for 循环逐个处理每个专家、单独计算expert_num个 FFN计算效率较低。CANN 提供GroupedMatmul算子可以同时计算多个专家提高计算和搬运效率。具体实现可参考HunyuanMoE类中npu_grouped_matmul模式下的实现对应推理参数--moe-impl npu_grouped_matmul。配合该模式Routing 阶段使能以下 torch_npu 融合算子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实现多个专家的矩阵乘计算提高计算和搬运效率。这组算子与 2.3 节 DoubleRouting 流程中的 InitRouting / LocalRerouting / LocalFinalizeRerouting / FinalizeRerouting 一一对应构成 EP 方案中 token 路由与专家计算的高性能实现。4.3 RmsNorm 算子优化通过使能torch_npu.npu_rms_norm算子能够提升模型推理性能。RmsNorm 是大模型常用的归一化操作相比 LayerNorm其去掉了减去均值的部分仅对均方根进行归一化在 NPU 上以融合算子方式实现可以显著减少中间张量与算子下发开销。五、消除冗余计算5.1 消除旋转位置编码冗余 cast 算子在旋转位置编码计算时query_states、key_states是实时获取的数据类型为 bfloat16。原始实现中cos、sin为 float32会触发从 float32 自动转换为 bfloat16 的冗余 cast 算子# 将cos, sin转成bfloat16 cos real_batched_index_select(cos, dim1, idxposition_ids).to(torch.bfloat16) sin real_batched_index_select(sin, dim1, idxposition_ids).to(torch.bfloat16) # 旋转位置编码计算使用cos, sin if self.use_rotary_pos_emb: cos, sin custom_pos_emb query_states, key_states apply_rotary_pos_emb(query_states, key_states, cos, sin)为消除冗余 cast 算子在HunyuanImage3ForCausalMM类中将cos、sin提前转为 bfloat16使编码计算全程保持同一数据类型减少 graph 中的转换节点。5.2 消除 gate 计算冗余 cast 算子外部开启了torch.autocast且数据类型设置为 bfloat16 时即使中间数据被设置为 float32计算时依然会被转成 bfloat16从而出现多个冗余的 cast 算子。为保持 gate 计算的精度gate 计算部分屏蔽外层的torch.autocast能力# 屏蔽外层的torch.autocast能力 with torch.npu.amp.autocast(enabledFalse): if self.wg.weight.dtype torch.float32: hidden_states hidden_states.float() logits self.wg(hidden_states)这样既保证了 gate 输出精度的稳定也避免了无关的 dtype 转换节点进入计算图。5.3 消除 FA 中 scale 的冗余计算每次调用 FA 算子时默认会计算一次scale 1 / math.sqrt(d)。这个操作在 CPU 上进行会导致下发中断和 device 的等待造成性能浪费。由于d是一个在 config 文件中指定的固定整数完全可以在初始化时将scale计算出来并在后续调用 FA 的地方复用。在本项目中此scale的计算被提前到HunyuanImage3Model.__init__()中作为成员变量存在后续调用flash_attn_func_npu时均直接使用、不再重复计算。5.4 消除 timestep_index 的冗余计算在 DecoderLayer 中计算 FA 前需要获取timestep_index用于确定casual_len其获取方式是timestep_index gen_timestep_scatter_index[0, 0].item()此处的.item()操作是一个强同步操作会造成 device 计算的中断。事实上gen_timestep_scatter_index的值在HunyuanImage3Text2ImagePipeline.__call__()的开头即确定下来并不再发生改变。因此可以将timestep_index的初始化提前至HunyuanImage3Text2ImagePipeline.__call__()的开头并透传至casual_len的计算处避免反复的 Host2Device 操作。在仓库适配代码 hunyuan_image_3_pipeline.py 的__call__实现中可以看到timestep_index gen_timestep_scatter_index[0, 0].item()在采样循环开始前被一次性计算并通过model_kwargs[timestep_index] timestep_index透传给后续各层正是这一优化思路的落地形态。六、工程落地配置、权重转换与推理执行优化文档附录指向环境部署与样例执行说明models/hunyuan-image-3.0/README.md此处补充与上述优化点配套的关键工程环节6.1 环境与权重准备要点支持 Atlas A2/A3 系列产品依赖 CANN 9.0.0-beta.1 软件包、Ascend Extension for PyTorch v7.3.1PyTorch 2.7.1Python 3.11需要将 HunyuanImage-3.0 开源仓库固定 commit62da220178f4b0b7d83e91665a46a20a3ee4f7cd以“非覆盖模式”复制到models/hunyuan-image-3.0/目录再由adaptor_patches/中的补丁代码model_adaptor.py完成模型、pipeline 的 NPU 适配权重转换使用 weight_convert.sh 拉起 convert_model.py关键参数为--tp-attnAttention TP 份数、--tp-moeMoE TP 份数、--epMoE EP 份数、--max-shard-size单块权重上限单位 G。约束为MoE 的 TP 与 EP 只能二选一Attention TP 并行度需与 MoE 并行度一致。6.2 推理配置与启动推理参数集中在 config/ep8_cfg.yaml 维护由 infer.sh 调用executor/scripts/mm_function.sh拉起。预置配置为Attn TP8 MoE EP8 CFG 并行 VAE 并行共 16 卡关键字段如下model_name: hunyuan-image-3.0 world_size: 16 master_port: 10086 entry_script: run_image_gen.py env_vars: CFG_PARALLEL: 1 # 开启 CFG 并行nproc_per_node 需为原来的 2 倍 USE_VAE_PARALLEL: 1 # 开启 VAE 并行 CPU_AFFINITY_CONF: 2 # 自动绑核缓解 HOST 下发瓶颈 model_args: reproduce: true model-id: ./ckpts/weight_ep8 # 指向权重转换后的目录 prompt: A cinematic medium shot captures a single Asian woman seated on a chair within a dimly lit room, creating an intimate and theatrical atmosphere. attn-impl: npu # 使能 NPU 上的 Attention 实现 moe-impl: npu_grouped_matmul # 使能 NPU 上的 MoE 实现 moe-ep: true # MoE 使用 EP与 moe-tp 互斥 seed: 42 diff-infer-steps: 50 image-size: 1024x1024 verbose: 0其中与本文优化点直接对应的是attn-impl: npu使能 NPU 的 Attention 实现对应 2.1 节 Attention 计算分解moe-impl: npu_grouped_matmul使能 NPU 的 MoE 实现对应 4.2 节 GMM 使能moe-ep: trueMoE 采用 EP 切分对应 2.3 节 DoubleRouting与moe-tp互斥需与权重转换时选择的并行方式一致CFG_PARALLEL/USE_VAE_PARALLEL对应 3.2、3.3 节的 VAE 与 CFG 并行开关。YAML 字段会透传为run_image_gen.py入口脚本的命令行参数torchrun --nproc_per_nodeworld_size run_image_gen.py model_args。注意模型最少需 4 个 Device 正常运行若开启CFG_PARALLEL则需 8 个 Device。根据仓库 README.md 记录在 A3 环境、CFG 并行与 VAE 并行均开启的条件下Attn TP8 MoE TP8 的端到端耗时约为 10.3sAttn TP8 MoE EP8 约为 9.9s可作为性能基线参考。七、总结HunyuanImage-3.0 在 CANN NPU 上的推理优化是一套组合拳并行切分层面Attention 以 TP 为主QKV 头切分 线性层行切分 QKV 合并计算MoE 在 TP 基础上演进为 EP DoubleRouting使 routing/finalizing 开销随 EP 份数线性下降并通过 AllGather/ReduceScatter 与矩阵乘融合压缩通信成本并行机制层面MoE 共享专家与通信多流掩盖、VAE 空间并行、CFG 双样本跨卡并行共同缩短端到端时间算子层面npu_swiglu、npu_grouped_matmul、MoE routing 系列融合算子、npu_rms_norm、npu_all_gather_base_mm、npu_mm_reduce_scatter_base等 torch_npu 融合算子替代朴素算子序列计算图层面消除 cos/sin 与 gate 计算中的冗余 cast、FA scale 的 CPU 重复计算、timestep_index的强同步.item()操作等冗余节点。上述优化均可在仓库的适配补丁adaptor_patches/、VAE 并行实现module/vae_patch_parallel.py、配置文件config/ep8_cfg.yaml与入口脚本run_image_gen.py中找到对应实现是理解并复现该模型 NPU 推理加速路径的完整参考。环境部署与样例执行的详细步骤可继续阅读 models/hunyuan-image-3.0/README.md。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考