NPU Wan2.2-I2V模型推理优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本文档主要介绍Wan2.2-I2V模型基于NPU的推理优化策略和实现。NPU npu_fused_infer_attention_score算子适配首先需要import torch_npu以及相关package在generate.py(L30)后加上import torch_npu torch_npu.npu.set_compile_mode(jit_compileFalse) torch.npu.config.allow_internal_formatFalse from torch_npu.contrib import transfer_to_npu本样例使用torch_npu内置的npu_fused_infer_attention_score融合算子替代FlashAttention算子该算子详细可见Ascend社区文档。在wan/modules/attention.py(L69)的npu_fused_attention函数中使能fused_infer_attention_score算子attention_out, _ torch_npu.npu_fused_infer_attention_score( q, k, v, actual_seq_lengthsactual_seq_lengths, actual_seq_lengths_kvactual_seq_lengths_kv, num_headsN, scalefloat(softmax_scale), input_layoutBNSD, num_key_value_headsnum_key_value_heads, pre_tokens65535, next_tokens65535 if not causal else 0, sparse_mode0, inner_precise0, ) attention_out attention_out.transpose(1, 2).contiguous() return attention_out.to(out_dtype)NPU rotary_mul算子适配本样例使用torch_npu内置的npu_rotary_mul融合算子替换源代码中的小算子实现npu_rotary_mul详细可见Ascend社区文档。在/wan/modules/model.py(L58)的rope_apply函数中使能npu_rotary_mul融合算子torch.amp.autocast(cuda, enabledFalse) def rope_apply(x, grid_sizes, freqs_list): s, n, c x.size(1), x.size(2), x.size(3) output [] for i, (f, h, w) in enumerate(grid_sizes.tolist()): x_i x[i, :s].reshape(1, s, n, c) if not x_i.is_contiguous(): x_ix_i.contiguous() cos, sin freqs_list[i] if cos.dim() 3: cos cos.unsqueeze(0) sin sin.unsqueeze(0) cos cos.to(dtypex_i.dtype, devicex_i.device) sin sin.to(dtypex_i.dtype, devicex_i.device) x_i torch_npu.npu_rotary_mul( inputx_i, r1cos, r2sin, rotary_modeinterleave ) output.append(x_i)NPU rms_norm算子适配本样例使用torch_npu内置的npu_rms_norm融合算子替换源代码中的小算子实现。npu_rms_norm详细可见Ascend社区文档。在/wan/modules/model.py(L87)的WanRMSNorm.forward中使能了npu_rms_norm融合算子class WanRMSNorm(nn.Module): def __init__(self, dim, eps1e-5): super().__init__() self.dim dim self.eps eps self.weight nn.Parameter(torch.ones(dim)) def forward(self, x): r Args: x(Tensor): Shape [B, L, C] return torch_npu.npu_rms_norm(x, self.weight, epsilonself.eps)[0]NPU layer_norm_eval算子适配本样例使用torch_npu内置的npu_layer_norm_eval融合算子替换源代码中的小算子实现。npu_layer_norm_eval详细可见Ascend社区文档。在/wan/modules/model.py(L103)的WanLayerNorm.forward中使能了npu_layer_norm_eval融合算子class WanLayerNorm(nn.LayerNorm): def __init__(self, dim, eps1e-6, elementwise_affineFalse): super().__init__(dim, elementwise_affineelementwise_affine, epseps) self.dim dim def forward(self, x): r Args: x(Tensor): Shape [B, L, C] return torch_npu.npu_layer_norm_eval( x, normalized_shape[self.dim], weightself.weight, biasself.bias, epsself.eps )另外本样例对部分layer norm(LN)和modulate操作进行了融合同样使用npu_layer_norm_eval融合算子在/wan/modules/model.py(L119)的FusedLayerNormModulate.forward中使能了相关融合操作class FusedLayerNormModulate(nn.Module): def __init__(self, dim, eps1e-6): super().__init__() self.dim dim self.eps eps def forward(self, x, weight, shift): r Args: x(Tensor) weight 1.0 scale bias shift return torch_npu.npu_layer_norm_eval( x, normalized_shape[self.dim], weightweight, biasbias, epsself.eps )模型每个block内有4个LN操作其中3个LN操作的gemma和beta都是固定的初始1和0那么通过(x-mean)/var10)(scale1)shift调用NPU融合算子npu_layer_norm_eval将3个LN和对应的modulate步骤融合。8卡VAE并行本样例对模型的VAE并行进行了使能。通过空间并行的方式实现将大尺寸图像在高度和宽度维度上切分成多个块分配给不同的NPU进程并行处理可以加速原模型VAE推理。在wan/vae_patch_parallel.py脚本实现此优化。以下VAE并行流程图展示了空间并行处理的完整执行路径首先将输入张量按空间维度切分到多个进程每个进程处理自己的局部块然后在计算过程中根据不同操作的特点采用相应的通信策略——卷积操作通过与邻居交换边界数据来获取上下文注意力操作通过全局收集所有K,V张量来保证计算完整性插值操作通过扩展边界、计算后再裁剪来处理上采样最后通过两阶段的收集过程将各个进程的局部结果按照原始的空间位置重新拼接成完整的输出张量。CFG并行本样例对模型的CFG并行进行了使能。在/wan/iamge2video.py(L432)实现for step_idx, t in enumerate(tqdm(timesteps)): latent_model_input [latent.to(self.device)] timestep [t] timestep torch.stack(timestep).to(self.device) model self._prepare_model_for_timestep(t, boundary, offload_model) sample_guide_scale guide_scale[1] if t.item() boundary else guide_scale[0] extra_kwargs {t_idx: step_idx} if hasattr(self, use_sp) and self.use_sp else {} if get_classifier_free_guidance_world_size() 2: noise_pred model( latent_model_input, ttimestep, **arg_all, **extra_kwargs)[0].to( torch.device(cpu) if offload_model else self.device) noise_pred_cond, noise_pred_uncond get_cfg_group().all_gather( noise_pred, separate_tensorsTrue ) if offload_model: torch.cuda.empty_cache() else: noise_pred_cond model( latent_model_input, ttimestep, **arg_c, **extra_kwargs)[0] if offload_model: torch.cuda.empty_cache() noise_pred_uncond model( latent_model_input, ttimestep, **arg_null, **extra_kwargs)[0] if offload_model: torch.cuda.empty_cache() noise_pred noise_pred_uncond sample_guide_scale * ( noise_pred_cond - noise_pred_uncond)当检测到CFG并行环境时它会将条件生成和无条件生成这两个任务分配到两个不同的进程上同时执行每个进程只需要运行一次模型推理然后通过all_gather通信操作让两个进程互相获取对方的计算结果最后使用CFG公式将条件预测和无条件预测按照引导系数混合得到最终的噪声预测结果如果没有启用并行环境则会退化到传统模式顺序执行两次模型推理。Ring Attention Overlap在多卡Ring Attention中通过异步启动AllGather收集其他卡的KV数据同时立即使用本地KV计算第一个FA让NPU计算与网络通信并行执行以掩盖通信延迟待AllGather完成后将其他多个远程chunks合并为一个长序列一次性计算第二个FA最后用LSELog-Sum-Exp算法正确合并两次attention输出达成将通信时间隐藏在本地计算中。Dit-CacheDIT-Cache作为扩散模型推理加速的缓存框架通过复用/预测已有的结果减少冗余前向计算。其加速逻辑可清晰的分为Step-level和Block-level范式Step-level通过判断不同采样步数step间的特定特征差异通过阈值比较决定是否跳过完整的step计算直接复用或者预测缓存结果Block-level以block为粒度通常是attention模块和mlp模块判断是否直接复用或者预测缓存结果。本样例集成了Step-level的Dit-Cache方案支持FBCache 以及 TeaCache。Step-level典型方法在Step-level加速范畴内FBCache的原理是基于First Block L1误差比较第一个Block输出残差与上一步的第一个Block输出残差之间的差异如果首块输出误差与上一轮首块输出误差差异小于指定阈值就跳过当前步计算复用残差对当前步的输出进行估计。TeaCache。启动方式Dit-Cache 配置已收敛到启动 YAML 的dit_cache段。仓库在models/wan2.2-i2v/config/下预置了14b_single_fbcache.yaml/14b_single_teacache.yaml两份配置切换infer.sh中的YAML_FILE_NAME即可启用对应策略。参数示例dit_cache: method: TeaCache # NoCache / FBCache / TeaCache enable_separate_cfg: true # CFG 分离开关true 时 cond/uncond 分开管理 params: # 仅需覆盖的键其余沿用 DEFAULT_CACHE_CONFIG rel_l1_thresh: 0.1 # TeaCache 阈值越大跳越多精度损失越大 coefficients: [733.226126, -401.131952, 67.5869174, -3.149879, 0.0961237896] warmup: 2 # 前 N 步强制完整计算mm_function.sh在拉起时会将该 YAML 作为--cache_config路径传给generate.py由module/dit_cache/cache_method.py的load_cache_config()解析dit_cache段并合并内置默认值。各方法完整默认值见module/dit_cache/cache_method.py:DEFAULT_CACHE_CONFIG也可参考 多模态推理统一拉起设计 §5。框架位置使用module/dit_cache/作为自定义库在模型 forward 处导入具体如下cann-recipes-infer --- models # 模型目录 | --- wan2.2-i2v | --- infer.sh # 统一入口 | --- config # 启动 YAML含 dit_cache 段 | --- wan | --- cache # cache 适配模型接口 | --- cache_block.py # Dit-Cache 适配双流模块 --- module --- dit_cache # step-level 缓存实现 --- cache_method.py # CacheManager / FBCache / TeaCache / TaylorSeer / NoCache【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考