feat(npu): Perform inference with sparse FA and NPU fusion operators on Ascend NPU

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"
+    }
+}

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(

lightx2v/models/networks/wan/infer/transformer_infer.py

@@ -79,6 +79,11 @@ 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 +93,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 +154,7 @@ def infer_block(self, block, x, pre_infer_out):
             x,
             shift_msa,
             scale_msa,
+            grid_sizes=pre_infer_out.grid_sizes.tuple,
         )
         x, attn_out = self.infer_cross_attn(
             block.compute_phases[1],
@@ -177,7 +187,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 +243,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,

lightx2v/models/networks/worldmirror/models/layers/attention.py

@@ -11,9 +11,22 @@
 
     _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_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
 
-    _USE_FLASH_ATTN_V3 = False
 from ...comm.communication import _All2All
 from ...comm.padding import depad_by_length, pad_by_length
 
@@ -54,6 +67,25 @@ 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:
+        q_t = q.transpose(1, 2).contiguous()
+        k_t = k.transpose(1, 2).contiguous()
+        v_t = v.transpose(1, 2).contiguous()
+        keep_prob = 1.0 - dropout_p
+        x = torch_npu.npu_fusion_attention(
+            q_t,
+            k_t,
+            v_t,
+            self.num_heads,
+            pse=None,
+            atten_mask=None,
+            scale=self.scale,
+            keep_prob=keep_prob,
+            input_layout="BNSD",
+        )
+        x = x[0].transpose(1, 2)
+        return x
+
     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 +102,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:

lightx2v/models/vfi/rife/model/warplayer.py

@@ -1,6 +1,11 @@
 import torch
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+try:
+    import torch_npu
+    _HAS_NPU = torch.npu.is_available()
+except ImportError:
+    _HAS_NPU = False
+device = torch.device("cuda") if torch.cuda.is_available() else (torch.device("npu") if _HAS_NPU else torch.device("cpu"))
 backwarp_tenGrid = {}
 
 

lightx2v/models/vfi/rife/rife_comfyui_wrapper.py

@@ -13,7 +13,15 @@ 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:
+            try:
+                import torch_npu
+                _HAS_NPU = torch.npu.is_available()
+            except ImportError:
+                _HAS_NPU = False
+            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 +36,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,16 +82,16 @@ 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())
             else:
                 # Get frames to interpolate
                 frame1 = images[source_idx1]
                 frame2 = images[source_idx2]
 
                 # 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 +106,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]]:

lightx2v/models/vfi/rife/train_log/RIFE_HDv3.py

@@ -5,7 +5,12 @@
 from ..model.loss import *
 from .IFNet_HDv3 import *
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+try:
+    import torch_npu
+    _HAS_NPU = torch.npu.is_available()
+except ImportError:
+    _HAS_NPU = False
+device = torch.device("cuda") if torch.cuda.is_available() else (torch.device("npu") if _HAS_NPU else torch.device("cpu"))
 
 
 class Model:

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":

lightx2v_platform/ops/attn/ascend_npu/__init__.py

@@ -1 +1,2 @@
 from .npu_flash_attn import *
+from .npu_rainfusion_attn import *