NPU HunyuanImage-3.0模型推理优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer概述本文主要介绍HunyuanImage-3.0模型基于NPU的推理优化策略。本文将介绍如何在Atlas A2/A3 系列产品上进行TPTensor Parallel推理及其他并行技术和算子融合优化方案。切分策略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和V进行归一化操作并使用旋转位置编码再计算attention图中Fused_infer_attention_score最后通过o_proj层输出。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_proj与up_proj合并计算提升整体计算效率。MoE EP优化在前述TP切分策略下MoE的routing和finilizing与切分的份数无关不论RANK数为多少routing和finilizeing的开销都一致随着RANK数的增加这部分的占比逐渐提升、无法优化。因此考虑对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通信将tokens发送至对应的RANK执行LocalRerouting在每个RANK内按照专家序对token进行重排执行GroupedMatMul进行专家计算执行LocalFinalizeRerouting排序回第二次all_to_all后的排布使用all_to_all把计算完的token发送回原本的TP域内对应RANK执行FinalizeRerouting对专家计算的结果在TP域内进行累加为了优化通信量配合此EPDoubleRouting方案在q_proj前使用AllGather对hidden_states进行聚合o_proj后使用ReduceScatter对Attentiond的结果进行分发这样可以为通算融合创造机会AllGather与qkv_proj融合、o_proj与ReduceScatter融合。并行优化MoE多流并行优化MoE模块中路由专家采用TP部署。共享专家的计算与路由专家在完成MoE计算后的通信可以通过多流并行机制使二者流水掩盖使用all_reduce的异步机制实现。排布情况如下图所示。VAE并行优化本样例对模型的VAE并行进行了使能。通过空间并行的方式实现将大尺寸图像在高度和宽度维度上切分成多个块分配给不同的NPU进程并行处理可以加速原模型VAE推理。在vae_patch_parallel.py脚本实现此优化。以下VAE并行流程图展示了空间并行处理的完整执行路径首先将输入张量按空间维度切分到多个进程每个进程处理自己的局部块然后在计算过程中根据不同操作的特点采用相应的通信策略——卷积操作通过与邻居交换边界数据来获取上下文注意力操作通过全局收集所有K,V张量来保证计算完整性插值操作通过扩展边界、计算后再裁剪来处理上采样最后通过两阶段的收集过程将各个进程的局部结果按照原始的空间位置重新拼接成完整的输出张量。CFG并行优化本样例对模型的CFG并行进行了使能。 对于一张图的生成CFG(Classifier Free Guidance)需要跑一个无条件的生成和一个有条件的生成相当于网络跑了两个样本可以把这两个样本拆分到两张卡上去跑然后通过all_gather通信操作互相获取对方的计算结果得到最终的预测结果。以TP2CFGP为例排布情况如下图所示。使能融合算子通算融合在当前的EP方案中attention部分的qkv_proj之前需要先对hidden_states进行AllGather在o_proj之后又需要对attn_output进行ReduceScatter操作计算与通信串行执行的情况下会存在大量的性能浪费。查询Ascend Extension for PyTorch自定义API可以找到合适的通算融合算子以实现掩盖、缩短端到端耗时。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平台下暂未使能该融合算子)。GMM使能Routing优化在MoE模块中如果通过for循环处理每个专家单独计算expert_num个前馈神经网络FFN容易导致计算效率较低。CANN提供了GroupedMatmul算子可以同时计算多个专家从而提高计算和搬运效率。具体实现可参考在HunyuanMoE类中的npu_grouped_matmul模式下的实现。高效排序和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其去掉了减去均值的部分。消除冗余计算消除旋转位置编码冗余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)在旋转位置编码计算时query_states, key_states是实时获取的数据类型为bfloat16为了消除cos, sin从float32自动转化成bfloat16的冗余cast算子在HunyuanImage3ForCausalMM类中将cos, sin转成bfloat16。消除gate计算冗余cast算子# 屏蔽外层的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)由于外部开启了torch.autocast且数据类型设置成了bfloat16即使设置成了float32在计算时依然会被转成bflat16出现多个冗余的cast算子为了保持gate计算的精度gate计算部分屏蔽外层的torch.autocast能力。消除FA中scale的冗余计算在每次调用FA算子时默认会计算一次scale1/math.sqrt(d)这个操作是在CPU进行会导致下发中断和device的等待造成性能浪费。由于d是一个固定的整数在config文件中指定那么完全可以在初始化时将scale的值计算出来并在后续调用FA的地方复用。在本项目中对此scale的计算被提到HunyuanImage3Model.__init__()中作为一个成员变量存在后续调用flash_attn_func_npu时均直接使用、不再重复计算。消除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操作提升总体性能。附录环境部署以及样例执行【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考