diff --git a/configs/distill/wan22/wan_moe_i2v_distill_int8_4step_ulysses_npu.json b/configs/distill/wan22/wan_moe_i2v_distill_int8_4step_ulysses_npu.json new file mode 100644 index 000000000..f32405ca1 --- /dev/null +++ b/configs/distill/wan22/wan_moe_i2v_distill_int8_4step_ulysses_npu.json @@ -0,0 +1,51 @@ +{ + "infer_steps": 4, + "target_video_length": 81, + "text_len": 512, + "target_height": 720, + "target_width": 1280, + "use_rainfusion": true, + "rainfusion": { + "sparsity": 0.8, + "skip_timesteps": -1 + }, + "self_attn_1_type": "npu_flash_attn", + "cross_attn_1_type": "npu_flash_attn", + "cross_attn_2_type": "npu_flash_attn", + "sample_guide_scale": [ + 3.5, + 3.5 + ], + "sample_shift": 5.0, + "enable_cfg": false, + "cpu_offload": false, + "offload_granularity": "block", + "t5_cpu_offload": false, + "vae_cpu_offload": false, + "use_image_encoder": false, + "boundary": 0.900, + "boundary_step_index": 2, + "denoising_step_list": [ + 1000, + 750, + 500, + 250 + ], + "dit_quantized": true, + "dit_quant_scheme": "int8-npu", + "modulate_type": "torch", + "rope_type": "npu_rope", + "layer_norm_type": "npu_layer_norm", + "rms_norm_type": "npu_rms_norm", + "high_noise_quantized_ckpt": "models/Wan2.2-Distill-Models/wan2.2_i2v_A14b_high_noise_int8_lightx2v_4step.safetensors", + "low_noise_quantized_ckpt": "models/Wan2.2-Distill-Models/wan2.2_i2v_A14b_low_noise_int8_lightx2v_4step.safetensors", + "video_frame_interpolation": { + "algo": "rife", + "target_fps": 30, + "model_path": "models/rife/train_log/flownet.pkl" + }, + "parallel": { + "seq_p_size": 2, + "seq_p_attn_type": "ulysses" + } +} diff --git a/lightx2v/models/networks/wan/infer/offload/transformer_infer.py b/lightx2v/models/networks/wan/infer/offload/transformer_infer.py index e7305a932..233614c5d 100755 --- a/lightx2v/models/networks/wan/infer/offload/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/offload/transformer_infer.py @@ -130,6 +130,7 @@ def infer_phase(self, cur_phase_idx, cur_phase, x, pre_infer_out): x, self.phase_params["shift_msa"], self.phase_params["scale_msa"], + grid_sizes=pre_infer_out.grid_sizes.tuple if getattr(pre_infer_out, "grid_sizes", None) is not None else None, ) elif cur_phase_idx == 1: x, self.phase_params["attn_out"] = self.infer_cross_attn( diff --git a/lightx2v/models/networks/wan/infer/transformer_infer.py b/lightx2v/models/networks/wan/infer/transformer_infer.py index 92030d16f..6e7064a23 100755 --- a/lightx2v/models/networks/wan/infer/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/transformer_infer.py @@ -79,6 +79,12 @@ def rope_wrapper(xq, xk, cos_sin_cache): self._mxfp8_fuse_available = self._probe_mxfp8_fuse_availability() if self.mxfp8_fuse_enable else False + self._rf_mgr = None + if self.config.get("use_rainfusion", False) and "npu" in AI_DEVICE: + from lightx2v_platform.ops.attn.ascend_npu.npu_rainfusion_attn import RainfusionManager + + self._rf_mgr = RainfusionManager(self.config) + @torch.no_grad() def reset_post_adapter_states(self): pass @@ -88,12 +94,16 @@ def reset_infer_states(self): self.cross_attn_cu_seqlens_q = None self.cross_attn_cu_seqlens_kv = None self.cross_attn_cu_seqlens_kv_img = None + if self._rf_mgr is not None: + self._rf_mgr.reset() if self.has_post_adapter: self.reset_post_adapter_states() @torch.no_grad() def infer(self, weights, pre_infer_out): self.cos_sin = pre_infer_out.cos_sin + if self._rf_mgr is not None and getattr(pre_infer_out, "grid_sizes", None) is not None: + self._rf_mgr.update_grid_sizes(pre_infer_out.grid_sizes.tuple) self.reset_infer_states() x = self.infer_main_blocks(weights.blocks, pre_infer_out) return self.infer_non_blocks(weights, x, pre_infer_out.embed) @@ -145,6 +155,7 @@ def infer_block(self, block, x, pre_infer_out): x, shift_msa, scale_msa, + grid_sizes=pre_infer_out.grid_sizes.tuple if getattr(pre_infer_out, "grid_sizes", None) is not None else None, ) x, attn_out = self.infer_cross_attn( block.compute_phases[1], @@ -177,7 +188,7 @@ def pre_process(self, modulation, embed0): return shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa - def infer_self_attn(self, phase, x, shift_msa, scale_msa): + def infer_self_attn(self, phase, x, shift_msa, scale_msa, grid_sizes=None): cos_sin = self.cos_sin norm1_quant = None norm1_scale = None @@ -233,7 +244,31 @@ def infer_self_attn(self, phase, x, shift_msa, scale_msa): "scheduler": self.scheduler, } - if self.config["seq_parallel"]: + if self._rf_mgr is not None and not self.config["seq_parallel"]: + self._rf_mgr.update_step_index(self.scheduler.step_index) + q_4d = q.unsqueeze(0) + k_4d = k.unsqueeze(0) + v_4d = v.unsqueeze(0) + attn_out = self._rf_mgr.apply_non_sp(q_4d, k_4d, v_4d, grid_sizes) + attn_out = attn_out.reshape(-1, self.num_heads * self.head_dim) + elif self._rf_mgr is not None and self.config["seq_parallel"]: + self._rf_mgr.update_step_index(self.scheduler.step_index) + attn_module = self._rf_mgr.get_sp_adapter() + attn_out = phase.self_attn_1_parallel.apply( + q=q, + k=k, + v=v, + slice_qkv_len=img_qkv_len, + cu_seqlens_qkv=self.self_attn_cu_seqlens_qkv, + attention_module=attn_module, + seq_p_group=self.seq_p_group, + use_fp8_comm=self.seq_p_fp8_comm, + use_fp4_comm=self.seq_p_fp4_comm, + use_tensor_fusion=self.seq_p_tensor_fusion, + enable_head_parallel=self.enable_head_parallel, + **attn_running_args, + ) + elif self.config["seq_parallel"]: attn_out = phase.self_attn_1_parallel.apply( q=q, k=k, diff --git a/lightx2v/models/networks/worldmirror/models/layers/attention.py b/lightx2v/models/networks/worldmirror/models/layers/attention.py index 869f08c5d..ddfd991bf 100755 --- a/lightx2v/models/networks/worldmirror/models/layers/attention.py +++ b/lightx2v/models/networks/worldmirror/models/layers/attention.py @@ -6,16 +6,30 @@ import torch.nn.functional as F from torch import Tensor, nn +from ...comm.communication import _All2All +from ...comm.padding import depad_by_length, pad_by_length + try: from flash_attn_interface import flash_attn_func as flash_attn_func_v3 _USE_FLASH_ATTN_V3 = True except ImportError: + _USE_FLASH_ATTN_V3 = False + +try: from flash_attn.flash_attn_interface import flash_attn_func as flash_attn_func_v2 +except ImportError: + flash_attn_func_v2 = None - _USE_FLASH_ATTN_V3 = False -from ...comm.communication import _All2All -from ...comm.padding import depad_by_length, pad_by_length +_USE_FLASH_ATTN_V2 = flash_attn_func_v2 is not None + +try: + import torch_npu + + _USE_NPU_FLASH_ATTN = torch.npu.is_available() +except ImportError: + torch_npu = None + _USE_NPU_FLASH_ATTN = False class Attention(nn.Module): @@ -54,6 +68,21 @@ def _compute_qkv(self, x: Tensor): q, k = self.q_norm(q).to(v.dtype), self.k_norm(k).to(v.dtype) return q, k, v, B, N, C + def _npu_flash_attn(self, q: Tensor, k: Tensor, v: Tensor, dropout_p: float = 0.0) -> Tensor: + keep_prob = 1.0 - dropout_p + x = torch_npu.npu_fusion_attention( + q, + k, + v, + self.num_heads, + pse=None, + atten_mask=None, + scale=self.scale, + keep_prob=keep_prob, + input_layout="BSND", + ) + return x[0] + def _apply_attention(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: if q.dtype == torch.bfloat16 or q.dtype == torch.float16: if q.is_contiguous(): @@ -70,8 +99,12 @@ def _apply_attention(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: v = v.transpose(1, 2).contiguous() if _USE_FLASH_ATTN_V3: x = flash_attn_func_v3(q, k, v) - else: + elif _USE_FLASH_ATTN_V2: x = flash_attn_func_v2(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0) + elif _USE_NPU_FLASH_ATTN: + x = self._npu_flash_attn(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0) + else: + x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0) if x.is_contiguous(): x = x.transpose(1, 2) else: diff --git a/lightx2v/models/vfi/rife/model/warplayer.py b/lightx2v/models/vfi/rife/model/warplayer.py index 096d48fed..8ea6fb33a 100644 --- a/lightx2v/models/vfi/rife/model/warplayer.py +++ b/lightx2v/models/vfi/rife/model/warplayer.py @@ -1,6 +1,9 @@ +import importlib.util + import torch -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +_HAS_NPU = importlib.util.find_spec("torch_npu") is not None and torch.npu.is_available() +device = torch.device("cuda") if torch.cuda.is_available() else (torch.device("npu") if _HAS_NPU else torch.device("cpu")) backwarp_tenGrid = {} diff --git a/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py b/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py index 0911c7259..d5afa7fe1 100644 --- a/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py +++ b/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py @@ -1,3 +1,4 @@ +import importlib.util import os from typing import List, Optional, Tuple @@ -13,7 +14,11 @@ class RIFEWrapper: BASE_DIR = os.path.dirname(os.path.abspath(__file__)) def __init__(self, model_path, device: Optional[torch.device] = None): - self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") + if device is not None: + self.device = device + else: + _HAS_NPU = importlib.util.find_spec("torch_npu") is not None and torch.npu.is_available() + self.device = torch.device("cuda") if torch.cuda.is_available() else (torch.device("npu") if _HAS_NPU else torch.device("cpu")) # Setup torch for optimal performance torch.set_grad_enabled(False) @@ -28,7 +33,7 @@ def __init__(self, model_path, device: Optional[torch.device] = None): with ProfilingContext4DebugL2("Load RIFE model"): self.model.load_model(model_path, -1) self.model.eval() - self.model.device() + self.model.flownet.to(self.device) @ProfilingContext4DebugL2("Interpolate frames") def interpolate_frames( @@ -74,7 +79,7 @@ def interpolate_frames( for source_idx1, source_idx2, interp_factor in frame_positions: if interp_factor == 0.0 or source_idx1 == source_idx2: # No interpolation needed, use the source frame directly - output_frames.append(images[source_idx1]) + output_frames.append(images[source_idx1].float().cpu()) else: # Get frames to interpolate frame1 = images[source_idx1] @@ -82,8 +87,8 @@ def interpolate_frames( # Convert ComfyUI format [H, W, C] to RIFE format [1, C, H, W] # Also convert from [0, 1] to [0, 1] (already in correct range) - I0 = frame1.permute(2, 0, 1).unsqueeze(0).to(self.device) - I1 = frame2.permute(2, 0, 1).unsqueeze(0).to(self.device) + I0 = frame1.permute(2, 0, 1).unsqueeze(0).to(self.device).float() + I1 = frame2.permute(2, 0, 1).unsqueeze(0).to(self.device).float() # Pad images I0 = F.pad(I0, padding) @@ -98,7 +103,6 @@ def interpolate_frames( interpolated_frame = interpolated[0, :, :height, :width].permute(1, 2, 0).cpu() output_frames.append(interpolated_frame) - # Stack all frames return torch.stack(output_frames, dim=0) def _calculate_target_frame_positions(self, source_fps: float, target_fps: float, total_source_frames: int) -> List[Tuple[int, int, float]]: diff --git a/lightx2v/models/vfi/rife/train_log/RIFE_HDv3.py b/lightx2v/models/vfi/rife/train_log/RIFE_HDv3.py index 29d5caa80..bcf6dc021 100644 --- a/lightx2v/models/vfi/rife/train_log/RIFE_HDv3.py +++ b/lightx2v/models/vfi/rife/train_log/RIFE_HDv3.py @@ -1,3 +1,5 @@ +import importlib.util + import torch from torch.nn.parallel import DistributedDataParallel as DDP from torch.optim import AdamW @@ -5,7 +7,8 @@ from ..model.loss import * from .IFNet_HDv3 import * -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +_HAS_NPU = importlib.util.find_spec("torch_npu") is not None and torch.npu.is_available() +device = torch.device("cuda") if torch.cuda.is_available() else (torch.device("npu") if _HAS_NPU else torch.device("cpu")) class Model: diff --git a/lightx2v_platform/ops/__init__.py b/lightx2v_platform/ops/__init__.py index b56ad794a..163b91baa 100755 --- a/lightx2v_platform/ops/__init__.py +++ b/lightx2v_platform/ops/__init__.py @@ -15,6 +15,8 @@ elif PLATFORM == "ascend_npu": from .attn.ascend_npu import * from .mm.ascend_npu import * + from .norm.ascend_npu import * + from .rope.ascend_npu import * elif PLATFORM == "metax_cuda": from .attn.metax_cuda import * elif PLATFORM == "enflame_gcu": diff --git a/lightx2v_platform/ops/attn/ascend_npu/__init__.py b/lightx2v_platform/ops/attn/ascend_npu/__init__.py index 46951c01f..b04781439 100644 --- a/lightx2v_platform/ops/attn/ascend_npu/__init__.py +++ b/lightx2v_platform/ops/attn/ascend_npu/__init__.py @@ -1 +1,2 @@ from .npu_flash_attn import * +from .npu_rainfusion_attn import * diff --git a/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py new file mode 100644 index 000000000..70850f081 --- /dev/null +++ b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py @@ -0,0 +1,346 @@ +import math + +import torch +from einops import rearrange + +try: + import mindiesd + from mindiesd.layers.flash_attn.attention_forward import attention_forward + + _HAS_MINDIESD = True +except ImportError: + mindiesd = None + attention_forward = None + _HAS_MINDIESD = False + +from lightx2v_platform.ops.attn.template import AttnWeightTemplate +from lightx2v_platform.registry_factory import PLATFORM_ATTN_WEIGHT_REGISTER + +try: + import torch_npu +except ImportError: + torch_npu = None + +""" +RainFusion is a sparse attention acceleration algorithm. +RainFusion requires MindIE-SD to be installed. Installation guide: https://gitcode.com/Ascend/MindIE-SD/blob/master/docs/zh/installation.md +""" + + +@PLATFORM_ATTN_WEIGHT_REGISTER("npu_rainfusion_attn") +class NpuRainfusionAttnWeight(AttnWeightTemplate): + def __init__(self, grid_size=None, pool_size=128, sparsity=0.8, skip_timesteps=-1, txt_len=0, txt_first=False): + """ + 参数: + grid_size (list): latents的THW网格大小。 + sparsity (float, optional): 稀疏度, 取值范围[0, 1]，默认为 0.50。 + """ + self.config = {} + assert torch_npu is not None, "torch_npu is not installed." + assert _HAS_MINDIESD, "mindiesd is not installed. RainFusion requires MindIE-SD." + + self.pool_size = pool_size + self.sparsity = sparsity + self.skip_timesteps = skip_timesteps + self.text_len = txt_len + self.txt_first = txt_first + + if grid_size is not None: + self.update_grid_size(grid_size) + + def update_grid_size(self, grid_size): + self.grid_size = grid_size + self.frame_num = self.grid_size[0] + self.num_tokens_per_frame = self.grid_size[1] * self.grid_size[2] + self.first_frame_len = self.num_tokens_per_frame + + @staticmethod + def get_grid_size(latent_size, patch_size): + t, h, w = latent_size[-3:] + return [t // patch_size[0], h // patch_size[1], w // patch_size[2]] + + def avgpool(self, input_tensor, pool_size=128): + batch, seqlen, headnum, dim = input_tensor.shape + + num_full_blocks = seqlen // pool_size + tail_size = seqlen % pool_size + + pooled_tensors = [] + if num_full_blocks > 0: + full_blocks = input_tensor[:, : num_full_blocks * pool_size, :, :] + full_blocks_reshaped = full_blocks.reshape(batch, num_full_blocks, pool_size, headnum, dim) + pooled_tensors.append(full_blocks_reshaped.mean(dim=2)) + if tail_size > 0: + tail_block = input_tensor[:, num_full_blocks * pool_size :, :, :] + tail_reshaped = tail_block.reshape(batch, 1, tail_size, headnum, dim) + pooled_tensors.append(tail_reshaped.mean(dim=2)) + + return torch.cat(pooled_tensors, dim=1) + + def get_mask_index(self, mask): + B, N, S, _ = mask.shape + device = mask.device + + mask_reshaped = mask.reshape(-1, S, S) + batch_size = mask_reshaped.shape[0] + + row_indices = torch.arange(S, device=device).view(1, 1, S).expand(batch_size, S, S) + sorted_vals = torch.where(mask_reshaped, row_indices, S) + sorted_vals, _ = torch.sort(sorted_vals, dim=-1) + valid_count = mask_reshaped.sum(dim=-1, keepdim=True) + keep_mask = row_indices < valid_count + result = torch.where(keep_mask, sorted_vals, -1) + + pos_matrix = result.reshape(B, N, S, S) + return pos_matrix + + def get_blockwise_mask(self, score_matrix, sparsity): + batch_size, num_heads, rows, cols = score_matrix.shape + + keep_len = math.ceil(cols * (1 - sparsity)) + keep_len = max(1, min(keep_len, cols)) + topk_values, _ = torch.topk(score_matrix, k=keep_len, dim=-1) + thresholds = topk_values[..., -1:] + mask = score_matrix >= thresholds + + total_len = self.frame_num * self.num_tokens_per_frame + self.text_len + protected_start_idx = total_len - self.first_frame_len - self.text_len + protected_start_block = protected_start_idx // self.pool_size + protect_len = cols - protected_start_block + + if protect_len > 0: + mask[:, :, -protect_len:, :] = True + mask[:, :, :, -protect_len:] = True + + selectIdx = self.get_mask_index(mask) + selectIdx = selectIdx[0].transpose(0, 1) + selectNumIdx = mask[0].transpose(0, 1).sum(dim=-1) + return selectIdx, selectNumIdx + + def rearrange_with_remaining(self, tensor): + b, s, n, d = tensor.shape + h = self.grid_size[1] + w = self.grid_size[2] + h_res_len, w_res_len = 0, 0 + first_frame_num = self.first_frame_len // h // w + + tensor_hwt = rearrange(tensor, "b (f h w) n d -> (b n) f h w d", f=self.frame_num - first_frame_num, h=h, w=w) + if h % 8 != 0: + tensor_hwt, tensor_h_r = torch.split(tensor_hwt, [h - (h % 8), h % 8], dim=2) + tensor_h_r = tensor_h_r.reshape(b * n, -1, d) + h_res_len = tensor_h_r.shape[1] + if w % 8 != 0: + tensor_hwt, tensor_w_r = torch.split(tensor_hwt, [w - (w % 8), w % 8], dim=3) + tensor_w_r = tensor_w_r.reshape(b * n, -1, d) + w_res_len = tensor_w_r.shape[1] + remaining_frames = self.frame_num - first_frame_num + if remaining_frames % 2 != 0: + raise ValueError(f"The number of remaining frames ({remaining_frames}) must be even to be rearranged with block size 2.") + tensor_hwt = rearrange(tensor_hwt, "bn (fn fb) (hn hb) (wn wb) d -> bn (fn hn wn fb hb wb) d", fn=remaining_frames // 2, fb=2, hb=8, wb=8, hn=h // 8, wn=w // 8) + if h % 8 != 0: + tensor_hwt = torch.cat((tensor_hwt, tensor_h_r), dim=1) + if w % 8 != 0: + tensor_hwt = torch.cat((tensor_hwt, tensor_w_r), dim=1) + tensor_hwt = rearrange(tensor_hwt, "(b n) s d -> b s n d", b=b, n=n) + return tensor_hwt, h_res_len, w_res_len + + def inv_rearrange_with_remaining(self, tensor, h_res_len, w_res_len): + b, s, n, d = tensor.shape + h = self.grid_size[1] + w = self.grid_size[2] + h_sr, w_sr = h % 8, w % 8 + first_frame_num = self.first_frame_len // (h * w) + remaining_frames = self.frame_num - first_frame_num + + tensor = rearrange(tensor, "b s n d->(b n) s d", b=b, n=n) + tensor_hwt, tensor_h, tensor_w = torch.split(tensor, [s - h_res_len - w_res_len, h_res_len, w_res_len], dim=1) + tensor_hwt = rearrange(tensor_hwt, "bn (fn hn wn fb hb wb) d -> bn (fn fb) (hn hb) (wn wb) d", fn=remaining_frames // 2, fb=2, hb=8, wb=8, hn=h // 8, wn=w // 8) + if w_res_len != 0: + tensor_w = tensor_w.reshape(b * n, remaining_frames, h - h_sr, w_sr, d) + tensor_hwt = torch.cat((tensor_hwt, tensor_w), dim=3) + if h_sr != 0: + tensor_h = tensor_h.reshape(b * n, remaining_frames, h_sr, w, d) + tensor_hwt = torch.cat((tensor_hwt, tensor_h), dim=2) + tensor_hwt = tensor_hwt.reshape(b * n, -1, d) + tensor_hwt = rearrange(tensor_hwt, "(b n) s d -> b s n d", b=b, n=n) + return tensor_hwt + + def do_tensor_rearrange_pooling(self, tensor): + b, s, n, d = tensor.shape + if s < self.text_len + self.first_frame_len: + raise ValueError(f"Sequence length {s} is too small for text_len {self.text_len} and first_frame_len {self.first_frame_len}") + if self.txt_first: + tensor_t, tensor_f, tensor_i = torch.split(tensor, [self.text_len, self.first_frame_len, s - self.text_len - self.first_frame_len], dim=1) + else: + tensor_f, tensor_i, tensor_t = torch.split(tensor, [self.first_frame_len, s - self.text_len - self.first_frame_len, self.text_len], dim=1) + tensor_i_2, h_res_len, w_res_len = self.rearrange_with_remaining(tensor_i) + tensor = torch.cat((tensor_i_2, tensor_f, tensor_t), dim=1) + tensor_pool = self.avgpool(tensor, pool_size=128) + return tensor, tensor_pool, h_res_len, w_res_len + + def do_tensor_inv_rearrange(self, tensor, h_res_len, w_res_len): + b, s, n, d = tensor.shape + tensor_i, tensor_f, tensor_t = torch.split(tensor, [s - self.text_len - self.first_frame_len, self.first_frame_len, self.text_len], dim=1) + tensor_i = self.inv_rearrange_with_remaining(tensor_i, h_res_len, w_res_len) + + if self.txt_first: + tensor = torch.cat((tensor_t, tensor_f, tensor_i), dim=1) + else: + tensor = torch.cat((tensor_f, tensor_i, tensor_t), dim=1) + return tensor + + def apply(self, q, k, v, grid_size=None, t_b_idx=None, base_blockmask=None, **kwds): + if grid_size is not None: + self.update_grid_size(grid_size) + t_idx = t_b_idx[0] + + if t_idx < self.skip_timesteps: + base_blockmask = None + x = attention_forward(q, k, v, opt_mode="manual", op_type="ascend_laser_attention", layout="BNSD") + else: + batch, qSeqlen, numHeads, headDim = q.shape + if batch != 1: + raise ValueError("NpuRainfusionAttnWeight currently only supports batch size 1.") + _, kvSeqlen, _, _ = k.shape + blockShapeX, blockShapeY = self.pool_size, self.pool_size + scale = headDim**-0.5 + + blockShape = [blockShapeX, blockShapeY] + actualSeqLengthsHost = [qSeqlen for _ in range(batch)] + actualSeqLengthsKvHost = [kvSeqlen for _ in range(batch)] + + sparsity = self.sparsity + + h_res_len, w_res_len = 0, 0 + qkv = torch.cat((q, k, v), dim=0) + qkv, qkv_pool, h_res_len, w_res_len = self.do_tensor_rearrange_pooling(qkv) + q, k, v = torch.chunk(qkv, 3, dim=0) + + if base_blockmask is None: + query_pool, key_pool, value_pool = torch.chunk(qkv_pool, 3, dim=0) + + attn_scores_head = torch.einsum("blnd,bsnd->bnls", query_pool, key_pool) * scale + attn_scores_fake = torch.nn.functional.softmax(attn_scores_head, dim=-1) + + selectIdx, selectNumIdx = self.get_blockwise_mask(attn_scores_fake, sparsity) + base_blockmask = [selectIdx, selectNumIdx] + else: + selectIdx = base_blockmask[0] + selectNumIdx = base_blockmask[1] + + q_bnsd = q.transpose(1, 2) + k_bnsd = k.transpose(1, 2) + v_bnsd = v.transpose(1, 2) + x = mindiesd.layers.flash_attn.sparse_flash_attn_rf_v2.rain_fusion_attention( + q_bnsd, + k_bnsd, + v_bnsd, + scale=scale, + head_num=numHeads, + input_layout="BNSD", + select_idx=selectIdx, + select_num_idx=selectNumIdx, + blockshape=blockShape, + actual_seq_lengths=actualSeqLengthsHost, + actual_seq_lengths_kv=actualSeqLengthsKvHost, + ) + + x = x.transpose(1, 2).reshape(batch, qSeqlen, numHeads, headDim) + + x = self.do_tensor_inv_rearrange(x, h_res_len, w_res_len) + + return x, base_blockmask + + +class RainfusionAdapter: + def __init__(self, rf_attn, manager): + self.rf = rf_attn + self.manager = manager + + def _pad_to_expected(self, tensor, actual_seqlen, expected_seqlen): + if actual_seqlen > expected_seqlen: + return tensor[:expected_seqlen], True + elif actual_seqlen < expected_seqlen: + diff = expected_seqlen - actual_seqlen + pad = tensor.new_zeros((diff,) + tensor.shape[1:]) + return torch.cat([tensor, pad], dim=0), False + else: + return tensor, None + + def apply(self, q, k, v, **kwargs): + grid_sizes = self.manager.grid_sizes + expected_seqlen = grid_sizes[0] * grid_sizes[1] * grid_sizes[2] + actual_seqlen = q.shape[0] + + q_valid, trunc_input = self._pad_to_expected(q, actual_seqlen, expected_seqlen) + k_valid, _ = self._pad_to_expected(k, actual_seqlen, expected_seqlen) + v_valid, _ = self._pad_to_expected(v, actual_seqlen, expected_seqlen) + + attn_out, self.manager._base_blockmask = self.rf.apply( + q=q_valid.unsqueeze(0), + k=k_valid.unsqueeze(0), + v=v_valid.unsqueeze(0), + grid_size=list(grid_sizes), + t_b_idx=(self.manager.step_index, 0), + base_blockmask=self.manager._base_blockmask, + ) + attn_out = attn_out.squeeze(0).reshape(attn_out.shape[1], -1) + + if trunc_input is True: + pad_out = attn_out.new_zeros((actual_seqlen - expected_seqlen, attn_out.shape[1])) + attn_out = torch.cat([attn_out, pad_out], dim=0) + elif trunc_input is False: + attn_out = attn_out[:actual_seqlen] + + return attn_out + + +class RainfusionManager: + def __init__(self, config): + self.config = config + self._rf_cfg = config.get("rainfusion", {}) + self._rf = None + self._base_blockmask = None + self.grid_sizes = None + self.step_index = 0 + + def update_grid_sizes(self, grid_sizes): + self.grid_sizes = grid_sizes + self._init_rf() + + def update_step_index(self, step_index): + self.step_index = step_index + + def reset(self): + self._base_blockmask = None + + def _init_rf(self): + if self._rf is not None and list(self._rf.grid_size) == list(self.grid_sizes): + return + self._base_blockmask = None + self._rf = NpuRainfusionAttnWeight( + grid_size=list(self.grid_sizes), + pool_size=128, + sparsity=self._rf_cfg.get("sparsity", 0.8), + skip_timesteps=self._rf_cfg.get("skip_timesteps", -1), + txt_len=0, + txt_first=False, + ) + + def apply_non_sp(self, q, k, v, grid_sizes): + attn_out, self._base_blockmask = self._rf.apply( + q=q, + k=k, + v=v, + grid_size=list(grid_sizes), + t_b_idx=(self.step_index, 0), + base_blockmask=self._base_blockmask, + ) + return attn_out + + def get_sp_adapter(self): + return RainfusionAdapter(self._rf, self) + + @property + def is_active(self): + return self._rf is not None diff --git a/lightx2v_platform/ops/norm/ascend_npu/__init__.py b/lightx2v_platform/ops/norm/ascend_npu/__init__.py new file mode 100644 index 000000000..91fa05c54 --- /dev/null +++ b/lightx2v_platform/ops/norm/ascend_npu/__init__.py @@ -0,0 +1,4 @@ +from .npu_layer_norm import NpuLayerNormWeight +from .npu_rms_norm import NpuRmsNormWeight + +__all__ = ["NpuRmsNormWeight", "NpuLayerNormWeight"] diff --git a/lightx2v_platform/ops/norm/ascend_npu/npu_layer_norm.py b/lightx2v_platform/ops/norm/ascend_npu/npu_layer_norm.py new file mode 100644 index 000000000..136c9a19e --- /dev/null +++ b/lightx2v_platform/ops/norm/ascend_npu/npu_layer_norm.py @@ -0,0 +1,47 @@ +import torch + +from lightx2v_platform.ops.norm.norm_template import LayerNormWeightTemplate +from lightx2v_platform.registry_factory import PLATFORM_LAYERNORM_WEIGHT_REGISTER + +try: + import torch_npu +except ImportError: + torch_npu = None + + +@PLATFORM_LAYERNORM_WEIGHT_REGISTER("npu_layer_norm") +class NpuLayerNormWeight(LayerNormWeightTemplate): + def __init__(self, weight_name=None, bias_name=None, create_cuda_buffer=False, create_cpu_buffer=False, lazy_load=False, lazy_load_file=None, is_post_adapter=False, eps=1e-6, **kwargs): + super().__init__(weight_name, bias_name, create_cuda_buffer, create_cpu_buffer, lazy_load, lazy_load_file, is_post_adapter, eps) + + def apply(self, input_tensor): + if torch_npu is not None and hasattr(torch_npu, "npu_layer_norm") and self.weight is not None and self.bias is not None: + if self.sensitive_layer_dtype != self.infer_dtype: + out = torch_npu.npu_layer_norm( + input_tensor.to(self.sensitive_layer_dtype), + (input_tensor.shape[-1],), + self.weight.to(self.sensitive_layer_dtype), + self.bias.to(self.sensitive_layer_dtype), + self.eps, + ) + if isinstance(out, tuple): + out = out[0] + return out.to(self.infer_dtype) + out = torch_npu.npu_layer_norm( + input_tensor, + (input_tensor.shape[-1],), + self.weight, + self.bias, + self.eps, + ) + if isinstance(out, tuple): + out = out[0] + return out + output = torch.nn.functional.layer_norm( + input_tensor, + (input_tensor.shape[-1],), + self.weight, + self.bias, + self.eps, + ) + return output diff --git a/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py b/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py new file mode 100644 index 000000000..d2e78a2fd --- /dev/null +++ b/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py @@ -0,0 +1,32 @@ +import torch + +from lightx2v_platform.ops.norm.norm_template import RMSWeightTemplate +from lightx2v_platform.registry_factory import PLATFORM_RMS_WEIGHT_REGISTER + +try: + import torch_npu +except ImportError: + torch_npu = None + + +@PLATFORM_RMS_WEIGHT_REGISTER("npu_rms_norm") +class NpuRmsNormWeight(RMSWeightTemplate): + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + def apply(self, input_tensor): + weight = self.weight + if torch_npu is not None and hasattr(torch_npu, "npu_rms_norm") and weight is not None: + if self.sensitive_layer_dtype != self.infer_dtype: + output_tensor, _ = torch_npu.npu_rms_norm( + input_tensor.to(self.sensitive_layer_dtype), + weight.to(self.sensitive_layer_dtype), + self.eps, + ) + return output_tensor.to(self.infer_dtype) + output_tensor, _ = torch_npu.npu_rms_norm(input_tensor, weight, self.eps) + return output_tensor + x = self._norm(input_tensor.float()) + if weight is not None: + return (x.float() * weight.float()).to(self.infer_dtype) + return x.to(self.infer_dtype) diff --git a/lightx2v_platform/ops/rope/ascend_npu/__init__.py b/lightx2v_platform/ops/rope/ascend_npu/__init__.py new file mode 100644 index 000000000..ba664c3de --- /dev/null +++ b/lightx2v_platform/ops/rope/ascend_npu/__init__.py @@ -0,0 +1,3 @@ +from .npu_rope import NpuRope + +__all__ = ["NpuRope"] diff --git a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py new file mode 100644 index 000000000..346209283 --- /dev/null +++ b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py @@ -0,0 +1,85 @@ +from functools import lru_cache + +import torch + +from lightx2v_platform.ops.rope.rope_template import GET_DTYPE, RopeTemplate +from lightx2v_platform.registry_factory import PLATFORM_ROPE_REGISTER + +try: + import torch_npu +except ImportError: + torch_npu = None + + +@lru_cache(maxsize=None) +def GET_SENSITIVE_DTYPE(): + import os + + DTYPE_MAP = { + "BF16": torch.bfloat16, + "FP16": torch.float16, + "FP32": torch.float32, + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, + "torch.bfloat16": torch.bfloat16, + "torch.float16": torch.float16, + "torch.float32": torch.float32, + } + flag = os.getenv("SENSITIVE_LAYER_DTYPE", "None") + if flag == "None": + return GET_DTYPE() + if flag not in DTYPE_MAP: + raise ValueError(f"Unsupported SENSITIVE_LAYER_DTYPE: {flag}. Expected one of {list(DTYPE_MAP.keys())}") + return DTYPE_MAP[flag] + + +@PLATFORM_ROPE_REGISTER("npu_rope") +class NpuRope(RopeTemplate): + def __init__(self): + super().__init__() + self.sensitive_layer_dtype = GET_SENSITIVE_DTYPE() + + def _apply_rope_fp32(self, xq, xk, cos_sin_cache): + n = xq.size(1) + seq_len = cos_sin_cache.size(0) + xq_fp32 = torch.view_as_complex(xq[:seq_len].to(torch.float32).reshape(seq_len, n, -1, 2).contiguous()) + xk_fp32 = torch.view_as_complex(xk[:seq_len].to(torch.float32).reshape(seq_len, n, -1, 2).contiguous()) + xq_rot = torch.view_as_real(xq_fp32 * cos_sin_cache).flatten(2) + xk_rot = torch.view_as_real(xk_fp32 * cos_sin_cache).flatten(2) + if xq.size(0) > seq_len: + xq_rot = torch.cat([xq_rot, xq[seq_len:]], dim=0) + xk_rot = torch.cat([xk_rot, xk[seq_len:]], dim=0) + return xq_rot.to(self.infer_dtype), xk_rot.to(self.infer_dtype) + + def apply(self, xq: torch.Tensor, xk: torch.Tensor, cos_sin_cache: torch.Tensor): + s, n, d = xq.shape + seq_len = cos_sin_cache.size(0) + cos = cos_sin_cache.real + sin = cos_sin_cache.imag + cos = cos.repeat_interleave(2, dim=-1) + sin = sin.repeat_interleave(2, dim=-1) + xq_part = xq[:seq_len] + xk_part = xk[:seq_len] + + if torch_npu is not None and hasattr(torch_npu, "npu_rotary_mul"): + if self.sensitive_layer_dtype != self.infer_dtype: + xq_part = xq_part.to(self.sensitive_layer_dtype) + xk_part = xk_part.to(self.sensitive_layer_dtype) + cos = cos.to(self.sensitive_layer_dtype) + sin = sin.to(self.sensitive_layer_dtype) + if not xq_part.is_contiguous(): + xq_part = xq_part.contiguous() + if not xk_part.is_contiguous(): + xk_part = xk_part.contiguous() + xq_rotated = torch_npu.npu_rotary_mul(xq_part, cos, sin, "interleave") + xk_rotated = torch_npu.npu_rotary_mul(xk_part, cos, sin, "interleave") + if s > seq_len: + xq = torch.cat([xq_rotated.to(self.infer_dtype), xq[seq_len:]], dim=0) + xk = torch.cat([xk_rotated.to(self.infer_dtype), xk[seq_len:]], dim=0) + else: + xq = xq_rotated.to(self.infer_dtype) + xk = xk_rotated.to(self.infer_dtype) + return xq, xk + + return self._apply_rope_fp32(xq, xk, cos_sin_cache) diff --git a/scripts/wan22/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh b/scripts/wan22/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh new file mode 100644 index 000000000..7d165d668 --- /dev/null +++ b/scripts/wan22/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +# set path firstly +lightx2v_path= +model_path= + +if [ -z "$lightx2v_path" ] || [ -z "$model_path" ]; then + echo "Error: Please set lightx2v_path and model_path in this script before running." + exit 1 +fi + +export PLATFORM=ascend_npu +export ASCEND_RT_VISIBLE_DEVICES=0,1 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +torchrun --nproc_per_node=2 -m lightx2v.infer \ +--model_cls wan2.2_moe_distill \ +--task i2v \ +--model_path $model_path \ +--config_json ${lightx2v_path}/configs/distill/wan22/wan_moe_i2v_distill_int8_4step_ulysses_npu.json \ +--prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ +--negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ +--image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ +--save_result_path ${lightx2v_path}/save_results/wan_moe_i2v_distill_int8_4step_ulysses_npu.mp4