From 74146ff3c00bd1fa3943a8daa2aa86ea51adec0c Mon Sep 17 00:00:00 2001
From: liaopingbo <liaopingbo@noreply.gitcode.com>
Date: Tue, 16 Jun 2026 14:38:07 +0800
Subject: [PATCH 01/15] feat(npu): Add 4 Ascend NPU custom operators for sparse
 video inference

---
 lightx2v_platform/ops/__init__.py             |   2 +
 .../attn/ascend_npu/rainfusion_blockwise.py   | 270 ++++++++++++++++++
 .../ops/norm/ascend_npu/__init__.py           |   4 +
 .../ops/norm/ascend_npu/npu_layer_norm.py     |  36 +++
 .../ops/norm/ascend_npu/npu_rms_norm.py       |  28 ++
 .../ops/rope/ascend_npu/__init__.py           |   3 +
 .../ops/rope/ascend_npu/npu_rope.py           |  80 ++++++
 7 files changed, 423 insertions(+)
 create mode 100644 lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
 create mode 100644 lightx2v_platform/ops/norm/ascend_npu/__init__.py
 create mode 100644 lightx2v_platform/ops/norm/ascend_npu/npu_layer_norm.py
 create mode 100644 lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py
 create mode 100644 lightx2v_platform/ops/rope/ascend_npu/__init__.py
 create mode 100644 lightx2v_platform/ops/rope/ascend_npu/npu_rope.py

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/rainfusion_blockwise.py b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
new file mode 100644
index 000000000..382404300
--- /dev/null
+++ b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
@@ -0,0 +1,270 @@
+import os
+import math
+import torch
+import torch_npu
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+import time
+import mindiesd
+from mindiesd.layers.flash_attn.attention_forward import attention_forward
+
+
+class Rainfusion_blockwise(nn.Module):
+    def __init__(
+            self,
+            grid_size: list,
+            pool_size: int = 128,
+            sparsity: float = 0.9,
+            skip_timesteps: int = 0,
+            txt_len: int = 0,
+            txt_first: bool = False,
+    ) -> None:
+        """
+        参数:
+            grid_size (list): latents的THW网格大小。
+            sparsity (float, optional): 稀疏度, 取值范围[0, 1]，默认为 0.50。
+        """
+        super().__init__()
+
+        # Rainfusion_param
+        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
+
+        self.pool_size = pool_size
+        self.sparsity = sparsity
+        self.skip_timesteps = skip_timesteps
+        self.text_len = txt_len
+        self.txt_first = txt_first
+
+    @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):  # BSND in,  BSND out
+        batch, seqlen, headnum, dim = input_tensor.shape
+
+        num_full_blocks = seqlen // pool_size
+        tail_size = seqlen % pool_size
+
+        if num_full_blocks > 0:
+            full_blocks = input_tensor[:, :num_full_blocks * pool_size, :, :]
+            full_blocks_reshaped = full_blocks.view(batch, num_full_blocks, pool_size, headnum, dim)
+            full_pooled = full_blocks_reshaped.mean(dim=2)
+        else:
+            full_pooled = torch.empty(0, device=input_tensor.device)
+        if tail_size > 0:
+            tail_block = input_tensor[:, num_full_blocks * pool_size:, :, :]
+            tail_reshaped = tail_block.view(batch, 1, tail_size, headnum, dim)
+            tail_pooled = tail_reshaped.mean(dim=2)
+        else:
+            tail_pooled = torch.empty(0, device=input_tensor.device)
+
+        if num_full_blocks > 0 and tail_size > 0:
+            output_tensor = torch.cat([full_pooled, tail_pooled], dim=1)
+        elif num_full_blocks > 0:
+            output_tensor = full_pooled
+        else:
+            output_tensor = tail_pooled
+
+        return output_tensor
+
+    def get_mask_index(self, mask):
+        B, N, S, _ = mask.shape
+        device = mask.device
+
+        # 1. 重塑维度 → (B*N)×S×S
+        mask_reshaped = mask.reshape(-1, S, S)
+        batch_size = mask_reshaped.shape[0]
+
+        # 2. 生成行索引  标记False位置为S（大于所有有效索引）
+        row_indices = torch.arange(S, device=device).expand(batch_size, S, -1)  # (B*N, S, S)
+        sorted_vals = torch.where(mask_reshaped, row_indices, 1e9).to(torch.float32)
+        sorted_vals, _ = torch.sort(sorted_vals, dim=-1)
+        valid_count = mask_reshaped.sum(dim=-1, keepdim=True)  # 每行True的个数
+        keep_mask = row_indices < valid_count  # 前valid_count个位置保留索引，其余填-1
+        result = torch.where(keep_mask, sorted_vals, -1)
+
+        pos_matrix = result.reshape(B, N, S, S).to(torch.int64)
+        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))
+        topk_values, _ = torch.topk(score_matrix, k=keep_len, dim=-1)
+        thresholds = topk_values[..., -1:]
+        mask = score_matrix >= thresholds
+
+        protect_len = (self.first_frame_len + self.text_len + self.pool_size - 1) // self.pool_size
+
+        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):  # BSND in ,  BSND out
+        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), 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), dim=3)
+            tensor_w_r = tensor_w_r.reshape(b * n, -1, d)
+            w_res_len = tensor_w_r.shape[1]
+        tensor_hwt = rearrange(tensor_hwt, 'b (fn fb) (hn hb) (wn wb) d -> b (fn hn wn fb hb wb) d',
+                               fn=self.frame_num // 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):  # BSND in ,  BSND out
+        b, s, n, d = tensor.shape
+        h = self.grid_size[1]
+        w = self.grid_size[2]
+        h_sr, w_sr = h % 8, w % 8
+
+        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, 'b (fn hn wn fb hb wb) d -> b (fn fb) (hn hb) (wn wb) d',
+                               fn=self.frame_num // 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, self.frame_num - 1, 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, self.frame_num - 1, 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):  # BSND in ,  BSND out
+        b, s, n, d = tensor.shape
+        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.concat((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)
+        tensor = torch.concat((tensor_t, tensor_i), dim=1)
+
+        if self.txt_first:
+            tensor = torch.concat((tensor_t, tensor_f, tensor_i), dim=1)
+        else:
+            tensor = torch.concat((tensor_f, tensor_i, tensor_t), dim=1)
+        return tensor
+
+    def do_tensor_pooling(self, tensor):
+        tensor_t = tensor[:, :self.text_len, :, :]
+        tensor_i = tensor[:, self.text_len:, :, :]
+
+        tensor_i_pool = self.avgpool(tensor_i, pool_size=128)
+        tensor_t_pool = self.avgpool(tensor_t, pool_size=128)
+
+        tensor_pool = torch.concat((tensor_t_pool, tensor_i_pool), dim=1)
+        return tensor_pool
+
+    def forward(
+            self,
+            q,  # BSND
+            k,
+            v,
+            t_b_idx,
+            base_blockmask,
+    ):
+        t_idx = t_b_idx[0]
+        b_idx = t_b_idx[1]
+        device = q.device
+
+        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
+            _, kvSeqlen, _, _ = k.shape
+            blockShapeX, blockShapeY = self.pool_size, self.pool_size
+            scale = headDim ** -0.5
+
+            totalQTokens = batch * qSeqlen
+            totalKvTokens = batch * kvSeqlen
+            qBlockNum = math.ceil(qSeqlen / blockShapeX)
+            kvBlockNum = math.ceil(kvSeqlen / blockShapeY)
+            totalQBlocks = qBlockNum
+            maxKvBlockNum = kvBlockNum
+            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
+            if base_blockmask is None:
+                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)
+                # qkv_pool = self.do_tensor_pooling(qkv)
+                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]
+                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)
+
+            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).view(batch, qSeqlen, numHeads, headDim)
+
+            x = self.do_tensor_inv_rearrange(x, h_res_len, w_res_len)
+            base_blockmask = None
+
+        return x, base_blockmask
\ No newline at end of file
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..e522b460b
--- /dev/null
+++ b/lightx2v_platform/ops/norm/ascend_npu/__init__.py
@@ -0,0 +1,4 @@
+from .npu_rms_norm import NpuRmsNormWeight
+from .npu_layer_norm import NpuLayerNormWeight
+
+__all__ = ["NpuRmsNormWeight", "NpuLayerNormWeight"]
\ No newline at end of file
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..871c5d28f
--- /dev/null
+++ b/lightx2v_platform/ops/norm/ascend_npu/npu_layer_norm.py
@@ -0,0 +1,36 @@
+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'):
+            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
\ No newline at end of file
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..bf01fefc7
--- /dev/null
+++ b/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py
@@ -0,0 +1,28 @@
+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"):
+            if self.sensitive_layer_dtype != self.infer_dtype:
+                output_tensor, _ = torch_npu.npu_rms_norm(
+                    input_tensor.float(), weight.float(), 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())
+        return (x.float() * weight.float()).to(self.infer_dtype)
\ No newline at end of file
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..2d205eb83
--- /dev/null
+++ b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
@@ -0,0 +1,80 @@
+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
+
+
+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,
+    }
+    flag = os.getenv("SENSITIVE_LAYER_DTYPE", "None")
+    if flag == "None":
+        return GET_DTYPE()
+    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)
+        )
+        xk_fp32 = torch.view_as_complex(
+            xk[:seq_len].to(torch.float32).reshape(seq_len, n, -1, 2)
+        )
+        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.float()
+                xk_part = xk_part.float()
+                cos = cos.float() if cos.dtype != torch.float32 else cos
+                sin = sin.float() if sin.dtype != torch.float32 else sin
+            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, xq[seq_len:]], dim=0)
+                xk = torch.cat([xk_rotated, xk[seq_len:]], dim=0)
+            else:
+                xq = xq_rotated
+                xk = xk_rotated
+            return xq.to(self.infer_dtype), xk.to(self.infer_dtype)
+
+        return self._apply_rope_fp32(xq, xk, cos_sin_cache)

From d49ef69d7dff96399ca3a1f9746f210faf3c195c Mon Sep 17 00:00:00 2001
From: liaopingbo <liaopingbo@foxmail.com>
Date: Tue, 16 Jun 2026 15:58:18 +0800
Subject: [PATCH 02/15] feat(npu): Add 4 Ascend NPU custom operators for sparse
 video inference

---
 .../attn/ascend_npu/rainfusion_blockwise.py   | 22 ++++++-------------
 .../ops/norm/ascend_npu/npu_layer_norm.py     |  2 +-
 .../ops/norm/ascend_npu/npu_rms_norm.py       |  2 +-
 .../ops/rope/ascend_npu/npu_rope.py           |  3 +++
 4 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
index 382404300..c8dd5d027 100644
--- a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
+++ b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
@@ -82,13 +82,13 @@ def get_mask_index(self, mask):
 
         # 2. 生成行索引  标记False位置为S（大于所有有效索引）
         row_indices = torch.arange(S, device=device).expand(batch_size, S, -1)  # (B*N, S, S)
-        sorted_vals = torch.where(mask_reshaped, row_indices, 1e9).to(torch.float32)
+        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)  # 每行True的个数
-        keep_mask = row_indices < valid_count  # 前valid_count个位置保留索引，其余填-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).to(torch.int64)
+        pos_matrix = result.reshape(B, N, S, S)
         return pos_matrix
 
     def get_blockwise_mask(self, score_matrix, sparsity):
@@ -127,7 +127,7 @@ def rearrange_with_remaining(self, tensor):  # BSND in ,  BSND out
             tensor_w_r = tensor_w_r.reshape(b * n, -1, d)
             w_res_len = tensor_w_r.shape[1]
         tensor_hwt = rearrange(tensor_hwt, 'b (fn fb) (hn hb) (wn wb) d -> b (fn hn wn fb hb wb) d',
-                               fn=self.frame_num // 2, fb=2, hb=8, wb=8, hn=h // 8, wn=w // 8)
+                               fn=(self.frame_num - first_frame_num) // 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:
@@ -144,7 +144,7 @@ def inv_rearrange_with_remaining(self, tensor, h_res_len, w_res_len):  # BSND in
         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, 'b (fn hn wn fb hb wb) d -> b (fn fb) (hn hb) (wn wb) d',
-                               fn=self.frame_num // 2, fb=2, hb=8, wb=8, hn=h // 8, wn=w // 8)
+                               fn=(self.frame_num - 1) // 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, self.frame_num - 1, h - h_sr, w_sr, d)
             tensor_hwt = torch.cat((tensor_hwt, tensor_w), dim=3)
@@ -175,7 +175,6 @@ def do_tensor_inv_rearrange(self, tensor, h_res_len, w_res_len):
                                                    [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)
-        tensor = torch.concat((tensor_t, tensor_i), dim=1)
 
         if self.txt_first:
             tensor = torch.concat((tensor_t, tensor_f, tensor_i), dim=1)
@@ -202,8 +201,6 @@ def forward(
             base_blockmask,
     ):
         t_idx = t_b_idx[0]
-        b_idx = t_b_idx[1]
-        device = q.device
 
         if t_idx < self.skip_timesteps:
             base_blockmask = None
@@ -211,16 +208,11 @@ def forward(
                                   opt_mode="manual", op_type="ascend_laser_attention", layout="BNSD")
         else:
             batch, qSeqlen, numHeads, headDim = q.shape
+            assert batch == 1, "Rainfusion_blockwise currently only supports batch size 1."
             _, kvSeqlen, _, _ = k.shape
             blockShapeX, blockShapeY = self.pool_size, self.pool_size
             scale = headDim ** -0.5
 
-            totalQTokens = batch * qSeqlen
-            totalKvTokens = batch * kvSeqlen
-            qBlockNum = math.ceil(qSeqlen / blockShapeX)
-            kvBlockNum = math.ceil(kvSeqlen / blockShapeY)
-            totalQBlocks = qBlockNum
-            maxKvBlockNum = kvBlockNum
             blockShape = [blockShapeX, blockShapeY]
             actualSeqLengthsHost = [qSeqlen for _ in range(batch)]
             actualSeqLengthsKvHost = [kvSeqlen for _ in range(batch)]
diff --git a/lightx2v_platform/ops/norm/ascend_npu/npu_layer_norm.py b/lightx2v_platform/ops/norm/ascend_npu/npu_layer_norm.py
index 871c5d28f..465424d61 100644
--- a/lightx2v_platform/ops/norm/ascend_npu/npu_layer_norm.py
+++ b/lightx2v_platform/ops/norm/ascend_npu/npu_layer_norm.py
@@ -15,7 +15,7 @@ def __init__(self, weight_name=None, bias_name=None, create_cuda_buffer=False, c
         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'):
+        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:
             out = torch_npu.npu_layer_norm(
                 input_tensor,
                 (input_tensor.shape[-1],),
diff --git a/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py b/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py
index bf01fefc7..2f94674f9 100644
--- a/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py
+++ b/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py
@@ -16,7 +16,7 @@ def _norm(self, x):
 
     def apply(self, input_tensor):
         weight = self.weight
-        if torch_npu is not None and hasattr(torch_npu, "npu_rms_norm"):
+        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.float(), weight.float(), self.eps,
diff --git a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
index 2d205eb83..a4188d4b3 100644
--- a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
+++ b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
@@ -18,6 +18,9 @@ def GET_SENSITIVE_DTYPE():
         "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":

From 4fab6f61e47823d4c5ffd44f28928c1c88299dd4 Mon Sep 17 00:00:00 2001
From: liaopingbo <liaopingbo@foxmail.com>
Date: Tue, 16 Jun 2026 16:38:07 +0800
Subject: [PATCH 03/15] feat(npu): Add 4 Ascend NPU custom operators for sparse
 video inference

---
 .../ops/attn/ascend_npu/rainfusion_blockwise.py  | 16 ++++++++--------
 .../ops/norm/ascend_npu/npu_rms_norm.py          |  4 +++-
 .../ops/rope/ascend_npu/npu_rope.py              | 10 +++++-----
 3 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
index c8dd5d027..7f10b156e 100644
--- a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
+++ b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
@@ -126,8 +126,11 @@ def rearrange_with_remaining(self, tensor):  # BSND in ,  BSND out
             tensor_hwt, tensor_w_r = torch.split(tensor_hwt, w - (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, 'b (fn fb) (hn hb) (wn wb) d -> b (fn hn wn fb hb wb) d',
-                               fn=(self.frame_num - first_frame_num) // 2, fb=2, hb=8, wb=8, hn=h // 8, wn=w // 8)
+                               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:
@@ -220,11 +223,11 @@ def forward(
             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:
-                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)
-                # qkv_pool = self.do_tensor_pooling(qkv)
                 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
@@ -235,9 +238,6 @@ def forward(
             else:
                 selectIdx = base_blockmask[0]
                 selectNumIdx = base_blockmask[1]
-                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)
 
             q_bnsd = q.transpose(1, 2)
             k_bnsd = k.transpose(1, 2)
diff --git a/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py b/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py
index 2f94674f9..fcac7669e 100644
--- a/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py
+++ b/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py
@@ -25,4 +25,6 @@ def apply(self, input_tensor):
             output_tensor, _ = torch_npu.npu_rms_norm(input_tensor, weight, self.eps)
             return output_tensor
         x = self._norm(input_tensor.float())
-        return (x.float() * weight.float()).to(self.infer_dtype)
\ No newline at end of file
+        if weight is not None:
+            return (x.float() * weight.float()).to(self.infer_dtype)
+        return x.to(self.infer_dtype)
\ No newline at end of file
diff --git a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
index a4188d4b3..ec01dc50e 100644
--- a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
+++ b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
@@ -73,11 +73,11 @@ def apply(self, xq: torch.Tensor, xk: torch.Tensor, cos_sin_cache: torch.Tensor)
             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, xq[seq_len:]], dim=0)
-                xk = torch.cat([xk_rotated, xk[seq_len:]], dim=0)
+                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
-                xk = xk_rotated
-            return xq.to(self.infer_dtype), xk.to(self.infer_dtype)
+                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)

From c07b3203ac0e91ea874bdd7393651cd0035050de Mon Sep 17 00:00:00 2001
From: liaopingbo <liaopingbo@foxmail.com>
Date: Tue, 16 Jun 2026 17:15:16 +0800
Subject: [PATCH 04/15] feat(npu): Add 4 Ascend NPU custom operators for sparse
 video inference

---
 .../attn/ascend_npu/rainfusion_blockwise.py   | 22 ++++++++-----------
 .../ops/rope/ascend_npu/npu_rope.py           |  4 ++++
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
index 7f10b156e..d0267f829 100644
--- a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
+++ b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
@@ -55,22 +55,15 @@ def avgpool(self, input_tensor, pool_size=128):  # BSND in,  BSND out
             full_blocks_reshaped = full_blocks.view(batch, num_full_blocks, pool_size, headnum, dim)
             full_pooled = full_blocks_reshaped.mean(dim=2)
         else:
-            full_pooled = torch.empty(0, device=input_tensor.device)
+            full_pooled = torch.empty(batch, 0, headnum, dim, device=input_tensor.device)
         if tail_size > 0:
             tail_block = input_tensor[:, num_full_blocks * pool_size:, :, :]
             tail_reshaped = tail_block.view(batch, 1, tail_size, headnum, dim)
             tail_pooled = tail_reshaped.mean(dim=2)
         else:
-            tail_pooled = torch.empty(0, device=input_tensor.device)
+            tail_pooled = torch.empty(batch, 0, headnum, dim, device=input_tensor.device)
 
-        if num_full_blocks > 0 and tail_size > 0:
-            output_tensor = torch.cat([full_pooled, tail_pooled], dim=1)
-        elif num_full_blocks > 0:
-            output_tensor = full_pooled
-        else:
-            output_tensor = tail_pooled
-
-        return output_tensor
+        return torch.cat([full_pooled, tail_pooled], dim=1)
 
     def get_mask_index(self, mask):
         B, N, S, _ = mask.shape
@@ -95,6 +88,7 @@ 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
@@ -143,16 +137,18 @@ def inv_rearrange_with_remaining(self, tensor, h_res_len, w_res_len):  # BSND in
         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, 'b (fn hn wn fb hb wb) d -> b (fn fb) (hn hb) (wn wb) d',
-                               fn=(self.frame_num - 1) // 2, fb=2, hb=8, wb=8, hn=h // 8, wn=w // 8)
+                               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, self.frame_num - 1, h - h_sr, w_sr, d)
+            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, self.frame_num - 1, h_sr, w, d)
+            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)
diff --git a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
index ec01dc50e..56daaf4fc 100644
--- a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
+++ b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
@@ -1,4 +1,5 @@
 import torch
+from functools import lru_cache
 
 from lightx2v_platform.ops.rope.rope_template import GET_DTYPE, RopeTemplate
 from lightx2v_platform.registry_factory import PLATFORM_ROPE_REGISTER
@@ -9,6 +10,7 @@
     torch_npu = None
 
 
+@lru_cache(maxsize=None)
 def GET_SENSITIVE_DTYPE():
     import os
     DTYPE_MAP = {
@@ -25,6 +27,8 @@ def GET_SENSITIVE_DTYPE():
     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]
 
 

From 8f0d06a67693cf9195434fa54471982c164b27ad Mon Sep 17 00:00:00 2001
From: liaopingbo <liaopingbo@foxmail.com>
Date: Tue, 16 Jun 2026 17:53:23 +0800
Subject: [PATCH 05/15] feat(npu): Add 4 Ascend NPU custom operators for sparse
 video inference

---
 .../attn/ascend_npu/rainfusion_blockwise.py   | 95 ++++++++-----------
 .../ops/norm/ascend_npu/__init__.py           |  4 +-
 .../ops/norm/ascend_npu/npu_layer_norm.py     | 15 ++-
 .../ops/norm/ascend_npu/npu_rms_norm.py       |  6 +-
 .../ops/rope/ascend_npu/npu_rope.py           | 20 ++--
 5 files changed, 70 insertions(+), 70 deletions(-)

diff --git a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
index d0267f829..277132ffe 100644
--- a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
+++ b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
@@ -1,24 +1,21 @@
-import os
 import math
+
+import mindiesd
 import torch
-import torch_npu
 import torch.nn as nn
-import torch.nn.functional as F
 from einops import rearrange
-import time
-import mindiesd
 from mindiesd.layers.flash_attn.attention_forward import attention_forward
 
 
 class Rainfusion_blockwise(nn.Module):
     def __init__(
-            self,
-            grid_size: list,
-            pool_size: int = 128,
-            sparsity: float = 0.9,
-            skip_timesteps: int = 0,
-            txt_len: int = 0,
-            txt_first: bool = False,
+        self,
+        grid_size: list,
+        pool_size: int = 128,
+        sparsity: float = 0.9,
+        skip_timesteps: int = 0,
+        txt_len: int = 0,
+        txt_first: bool = False,
     ) -> None:
         """
         参数:
@@ -50,20 +47,17 @@ def avgpool(self, input_tensor, pool_size=128):  # BSND in,  BSND out
         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 = input_tensor[:, : num_full_blocks * pool_size, :, :]
             full_blocks_reshaped = full_blocks.view(batch, num_full_blocks, pool_size, headnum, dim)
-            full_pooled = full_blocks_reshaped.mean(dim=2)
-        else:
-            full_pooled = torch.empty(batch, 0, headnum, dim, device=input_tensor.device)
+            pooled_tensors.append(full_blocks_reshaped.mean(dim=2))
         if tail_size > 0:
-            tail_block = input_tensor[:, num_full_blocks * pool_size:, :, :]
+            tail_block = input_tensor[:, num_full_blocks * pool_size :, :, :]
             tail_reshaped = tail_block.view(batch, 1, tail_size, headnum, dim)
-            tail_pooled = tail_reshaped.mean(dim=2)
-        else:
-            tail_pooled = torch.empty(batch, 0, headnum, dim, device=input_tensor.device)
+            pooled_tensors.append(tail_reshaped.mean(dim=2))
 
-        return torch.cat([full_pooled, tail_pooled], dim=1)
+        return torch.cat(pooled_tensors, dim=1)
 
     def get_mask_index(self, mask):
         B, N, S, _ = mask.shape
@@ -74,7 +68,7 @@ def get_mask_index(self, mask):
         batch_size = mask_reshaped.shape[0]
 
         # 2. 生成行索引  标记False位置为S（大于所有有效索引）
-        row_indices = torch.arange(S, device=device).expand(batch_size, S, -1)  # (B*N, S, S)
+        row_indices = torch.arange(S, device=device).view(1, 1, S).expand(batch_size, S, S)  # (B*N, 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)
@@ -111,7 +105,7 @@ def rearrange_with_remaining(self, tensor):  # BSND in ,  BSND out
         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)
+        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), dim=2)
             tensor_h_r = tensor_h_r.reshape(b * n, -1, d)
@@ -123,13 +117,12 @@ def rearrange_with_remaining(self, tensor):  # BSND in ,  BSND out
         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, 'b (fn fb) (hn hb) (wn wb) d -> b (fn hn wn fb hb wb) d',
-                               fn=remaining_frames // 2, fb=2, hb=8, wb=8, hn=h // 8, wn=w // 8)
+        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)
+        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):  # BSND in ,  BSND out
@@ -140,10 +133,9 @@ def inv_rearrange_with_remaining(self, tensor, h_res_len, w_res_len):  # BSND in
         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 = 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, 'b (fn hn wn fb hb wb) d -> b (fn fb) (hn hb) (wn wb) d',
-                               fn=remaining_frames // 2, fb=2, hb=8, wb=8, hn=h // 8, wn=w // 8)
+        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)
@@ -151,18 +143,17 @@ def inv_rearrange_with_remaining(self, tensor, h_res_len, w_res_len):  # BSND in
             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)
+        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):  # BSND in ,  BSND out
         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)
+            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_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.concat((tensor_i_2, tensor_f, tensor_t), dim=1)
         tensor_pool = self.avgpool(tensor, pool_size=128)
@@ -170,9 +161,7 @@ def do_tensor_rearrange_pooling(self, tensor):  # BSND in ,  BSND out
 
     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, 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:
@@ -182,8 +171,8 @@ def do_tensor_inv_rearrange(self, tensor, h_res_len, w_res_len):
         return tensor
 
     def do_tensor_pooling(self, tensor):
-        tensor_t = tensor[:, :self.text_len, :, :]
-        tensor_i = tensor[:, self.text_len:, :, :]
+        tensor_t = tensor[:, : self.text_len, :, :]
+        tensor_i = tensor[:, self.text_len :, :, :]
 
         tensor_i_pool = self.avgpool(tensor_i, pool_size=128)
         tensor_t_pool = self.avgpool(tensor_t, pool_size=128)
@@ -192,25 +181,24 @@ def do_tensor_pooling(self, tensor):
         return tensor_pool
 
     def forward(
-            self,
-            q,  # BSND
-            k,
-            v,
-            t_b_idx,
-            base_blockmask,
+        self,
+        q,  # BSND
+        k,
+        v,
+        t_b_idx,
+        base_blockmask,
     ):
         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")
+            x = attention_forward(q, k, v, opt_mode="manual", op_type="ascend_laser_attention", layout="BNSD")
         else:
             batch, qSeqlen, numHeads, headDim = q.shape
             assert batch == 1, "Rainfusion_blockwise currently only supports batch size 1."
             _, kvSeqlen, _, _ = k.shape
             blockShapeX, blockShapeY = self.pool_size, self.pool_size
-            scale = headDim ** -0.5
+            scale = headDim**-0.5
 
             blockShape = [blockShapeX, blockShapeY]
             actualSeqLengthsHost = [qSeqlen for _ in range(batch)]
@@ -239,7 +227,9 @@ def forward(
             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,
+                q_bnsd,
+                k_bnsd,
+                v_bnsd,
                 scale=scale,
                 head_num=numHeads,
                 input_layout="BNSD",
@@ -247,12 +237,11 @@ def forward(
                 select_num_idx=selectNumIdx,
                 blockshape=blockShape,
                 actual_seq_lengths=actualSeqLengthsHost,
-                actual_seq_lengths_kv=actualSeqLengthsKvHost
+                actual_seq_lengths_kv=actualSeqLengthsKvHost,
             )
 
             x = x.transpose(1, 2).view(batch, qSeqlen, numHeads, headDim)
 
             x = self.do_tensor_inv_rearrange(x, h_res_len, w_res_len)
-            base_blockmask = None
 
-        return x, base_blockmask
\ No newline at end of file
+        return x, base_blockmask
diff --git a/lightx2v_platform/ops/norm/ascend_npu/__init__.py b/lightx2v_platform/ops/norm/ascend_npu/__init__.py
index e522b460b..91fa05c54 100644
--- a/lightx2v_platform/ops/norm/ascend_npu/__init__.py
+++ b/lightx2v_platform/ops/norm/ascend_npu/__init__.py
@@ -1,4 +1,4 @@
-from .npu_rms_norm import NpuRmsNormWeight
 from .npu_layer_norm import NpuLayerNormWeight
+from .npu_rms_norm import NpuRmsNormWeight
 
-__all__ = ["NpuRmsNormWeight", "NpuLayerNormWeight"]
\ No newline at end of file
+__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
index 465424d61..136c9a19e 100644
--- a/lightx2v_platform/ops/norm/ascend_npu/npu_layer_norm.py
+++ b/lightx2v_platform/ops/norm/ascend_npu/npu_layer_norm.py
@@ -15,7 +15,18 @@ def __init__(self, weight_name=None, bias_name=None, create_cuda_buffer=False, c
         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 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],),
@@ -33,4 +44,4 @@ def apply(self, input_tensor):
             self.bias,
             self.eps,
         )
-        return output
\ No newline at end of file
+        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
index fcac7669e..d2e78a2fd 100644
--- a/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py
+++ b/lightx2v_platform/ops/norm/ascend_npu/npu_rms_norm.py
@@ -19,7 +19,9 @@ def apply(self, input_tensor):
         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.float(), weight.float(), self.eps,
+                    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)
@@ -27,4 +29,4 @@ def apply(self, input_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)
\ No newline at end of file
+        return x.to(self.infer_dtype)
diff --git a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
index 56daaf4fc..b08116811 100644
--- a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
+++ b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
@@ -1,6 +1,7 @@
-import torch
 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
 
@@ -13,6 +14,7 @@
 @lru_cache(maxsize=None)
 def GET_SENSITIVE_DTYPE():
     import os
+
     DTYPE_MAP = {
         "BF16": torch.bfloat16,
         "FP16": torch.float16,
@@ -41,12 +43,8 @@ def __init__(self):
     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)
-        )
-        xk_fp32 = torch.view_as_complex(
-            xk[:seq_len].to(torch.float32).reshape(seq_len, n, -1, 2)
-        )
+        xq_fp32 = torch.view_as_complex(xq[:seq_len].to(torch.float32).reshape(seq_len, n, -1, 2))
+        xk_fp32 = torch.view_as_complex(xk[:seq_len].to(torch.float32).reshape(seq_len, n, -1, 2))
         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:
@@ -66,10 +64,10 @@ def apply(self, xq: torch.Tensor, xk: torch.Tensor, cos_sin_cache: torch.Tensor)
 
         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.float()
-                xk_part = xk_part.float()
-                cos = cos.float() if cos.dtype != torch.float32 else cos
-                sin = sin.float() if sin.dtype != torch.float32 else sin
+                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():

From 0a65ead8626a8906288fc06f5497baaa13e753c2 Mon Sep 17 00:00:00 2001
From: liaopingbo <liaopingbo@foxmail.com>
Date: Tue, 16 Jun 2026 18:04:37 +0800
Subject: [PATCH 06/15] Update
 lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
index 277132ffe..b9dfd619f 100644
--- a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
+++ b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
@@ -240,7 +240,7 @@ def forward(
                 actual_seq_lengths_kv=actualSeqLengthsKvHost,
             )
 
-            x = x.transpose(1, 2).view(batch, qSeqlen, numHeads, headDim)
+            x = x.transpose(1, 2).reshape(batch, qSeqlen, numHeads, headDim)
 
             x = self.do_tensor_inv_rearrange(x, h_res_len, w_res_len)
 

From 349c350d0e5333fcf2023fe4c0d8e33d37b1d89e Mon Sep 17 00:00:00 2001
From: liaopingbo <liaopingbo@foxmail.com>
Date: Tue, 16 Jun 2026 18:05:29 +0800
Subject: [PATCH 07/15] feat(npu): Add 4 Ascend NPU custom operators for sparse
 video inference

---
 .../attn/ascend_npu/rainfusion_blockwise.py   | 21 +++++++------------
 1 file changed, 7 insertions(+), 14 deletions(-)

diff --git a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
index b9dfd619f..ec9077942 100644
--- a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
+++ b/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
@@ -87,7 +87,10 @@ def get_blockwise_mask(self, score_matrix, sparsity):
         thresholds = topk_values[..., -1:]
         mask = score_matrix >= thresholds
 
-        protect_len = (self.first_frame_len + self.text_len + self.pool_size - 1) // self.pool_size
+        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
@@ -155,7 +158,7 @@ def do_tensor_rearrange_pooling(self, tensor):  # BSND in ,  BSND out
         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.concat((tensor_i_2, tensor_f, tensor_t), dim=1)
+        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
 
@@ -165,21 +168,11 @@ def do_tensor_inv_rearrange(self, tensor, h_res_len, w_res_len):
         tensor_i = self.inv_rearrange_with_remaining(tensor_i, h_res_len, w_res_len)
 
         if self.txt_first:
-            tensor = torch.concat((tensor_t, tensor_f, tensor_i), dim=1)
+            tensor = torch.cat((tensor_t, tensor_f, tensor_i), dim=1)
         else:
-            tensor = torch.concat((tensor_f, tensor_i, tensor_t), dim=1)
+            tensor = torch.cat((tensor_f, tensor_i, tensor_t), dim=1)
         return tensor
 
-    def do_tensor_pooling(self, tensor):
-        tensor_t = tensor[:, : self.text_len, :, :]
-        tensor_i = tensor[:, self.text_len :, :, :]
-
-        tensor_i_pool = self.avgpool(tensor_i, pool_size=128)
-        tensor_t_pool = self.avgpool(tensor_t, pool_size=128)
-
-        tensor_pool = torch.concat((tensor_t_pool, tensor_i_pool), dim=1)
-        return tensor_pool
-
     def forward(
         self,
         q,  # BSND

From b3555b8342b4af88ea0410c31e99238c56d947f1 Mon Sep 17 00:00:00 2001
From: liaopingbo <liaopingbo@foxmail.com>
Date: Wed, 17 Jun 2026 11:19:48 +0800
Subject: [PATCH 08/15] feat(npu): Perform inference with sparse FA and NPU
 affinity operators on Ascend NPU.

---
 ...oe_i2v_distill_int8_4step_ulysses_npu.json |  51 ++++++
 .../wan/infer/offload/transformer_infer.py    |   1 +
 .../networks/wan/infer/transformer_infer.py   |  38 +++-
 .../worldmirror/models/layers/attention.py    |  40 +++-
 lightx2v/models/vfi/rife/model/warplayer.py   |   7 +-
 .../models/vfi/rife/rife_comfyui_wrapper.py   |  23 ++-
 .../models/vfi/rife/train_log/RIFE_HDv3.py    |   7 +-
 .../ops/attn/ascend_npu/__init__.py           |   1 +
 ...on_blockwise.py => npu_rainfusion_attn.py} | 171 ++++++++++++++----
 .../ops/rope/ascend_npu/npu_rope.py           |   4 +-
 ..._moe_i2v_distill_int8_4step_ulysses_npu.sh |  21 +++
 11 files changed, 309 insertions(+), 55 deletions(-)
 create mode 100644 configs/distill/wan22/wan_moe_i2v_distill_int8_4step_ulysses_npu.json
 rename lightx2v_platform/ops/attn/ascend_npu/{rainfusion_blockwise.py => npu_rainfusion_attn.py} (65%)
 create mode 100644 scripts/wan22/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh

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..581dcebe6
--- /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"
+    }
+}
\ No newline at end of file
diff --git a/lightx2v/models/networks/wan/infer/offload/transformer_infer.py b/lightx2v/models/networks/wan/infer/offload/transformer_infer.py
index e7305a932..449714433 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,
             )
         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..e3303d0f9 100755
--- a/lightx2v/models/networks/wan/infer/transformer_infer.py
+++ b/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:
+            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,
diff --git a/lightx2v/models/networks/worldmirror/models/layers/attention.py b/lightx2v/models/networks/worldmirror/models/layers/attention.py
index 869f08c5d..7c987f6a9 100755
--- a/lightx2v/models/networks/worldmirror/models/layers/attention.py
+++ b/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 = True
+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:
diff --git a/lightx2v/models/vfi/rife/model/warplayer.py b/lightx2v/models/vfi/rife/model/warplayer.py
index 096d48fed..70565b0a7 100644
--- a/lightx2v/models/vfi/rife/model/warplayer.py
+++ b/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 = {}
 
 
diff --git a/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py b/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py
index 0911c7259..cc82a9f29 100644
--- a/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py
+++ b/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,7 +82,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].to(self.device).float())
             else:
                 # Get frames to interpolate
                 frame1 = images[source_idx1]
@@ -82,8 +90,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)
@@ -95,11 +103,10 @@ def interpolate_frames(
 
                 # Convert back to ComfyUI format [H, W, C]
                 # Crop to original size and permute dimensions
-                interpolated_frame = interpolated[0, :, :height, :width].permute(1, 2, 0).cpu()
+                interpolated_frame = interpolated[0, :, :height, :width].permute(1, 2, 0)
                 output_frames.append(interpolated_frame)
 
-        # Stack all frames
-        return torch.stack(output_frames, dim=0)
+        return torch.stack(output_frames, dim=0).cpu()
 
     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..47dda9656 100644
--- a/lightx2v/models/vfi/rife/train_log/RIFE_HDv3.py
+++ b/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:
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/rainfusion_blockwise.py b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
similarity index 65%
rename from lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
rename to lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
index ec9077942..2fb91efc3 100644
--- a/lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py
+++ b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
@@ -2,33 +2,31 @@
 
 import mindiesd
 import torch
-import torch.nn as nn
 from einops import rearrange
 from mindiesd.layers.flash_attn.attention_forward import attention_forward
 
-
-class Rainfusion_blockwise(nn.Module):
-    def __init__(
-        self,
-        grid_size: list,
-        pool_size: int = 128,
-        sparsity: float = 0.9,
-        skip_timesteps: int = 0,
-        txt_len: int = 0,
-        txt_first: bool = False,
-    ) -> None:
+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。
         """
-        super().__init__()
-
-        # Rainfusion_param
-        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
+        self.config = {}
+        assert torch_npu is not None, "torch_npu is not installed."
 
         self.pool_size = pool_size
         self.sparsity = sparsity
@@ -36,12 +34,21 @@ def __init__(
         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):  # BSND in,  BSND out
+    def avgpool(self, input_tensor, pool_size=128):
         batch, seqlen, headnum, dim = input_tensor.shape
 
         num_full_blocks = seqlen // pool_size
@@ -49,12 +56,12 @@ def avgpool(self, input_tensor, pool_size=128):  # BSND in,  BSND out
 
         pooled_tensors = []
         if num_full_blocks > 0:
-            full_blocks = input_tensor[:, : num_full_blocks * pool_size, :, :]
-            full_blocks_reshaped = full_blocks.view(batch, num_full_blocks, pool_size, headnum, dim)
+            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.view(batch, 1, tail_size, headnum, dim)
+            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)
@@ -63,12 +70,10 @@ def get_mask_index(self, mask):
         B, N, S, _ = mask.shape
         device = mask.device
 
-        # 1. 重塑维度 → (B*N)×S×S
         mask_reshaped = mask.reshape(-1, S, S)
         batch_size = mask_reshaped.shape[0]
 
-        # 2. 生成行索引  标记False位置为S（大于所有有效索引）
-        row_indices = torch.arange(S, device=device).view(1, 1, S).expand(batch_size, S, S)  # (B*N, S, S)
+        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)
@@ -101,7 +106,7 @@ def get_blockwise_mask(self, score_matrix, sparsity):
         selectNumIdx = mask[0].transpose(0, 1).sum(dim=-1)
         return selectIdx, selectNumIdx
 
-    def rearrange_with_remaining(self, tensor):  # BSND in ,  BSND out
+    def rearrange_with_remaining(self, tensor):
         b, s, n, d = tensor.shape
         h = self.grid_size[1]
         w = self.grid_size[2]
@@ -128,7 +133,7 @@ def rearrange_with_remaining(self, tensor):  # BSND in ,  BSND out
         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):  # BSND in ,  BSND out
+    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]
@@ -149,7 +154,7 @@ def inv_rearrange_with_remaining(self, tensor, h_res_len, w_res_len):  # BSND in
         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):  # BSND in ,  BSND out
+    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}")
@@ -173,14 +178,9 @@ def do_tensor_inv_rearrange(self, tensor, h_res_len, w_res_len):
             tensor = torch.cat((tensor_f, tensor_i, tensor_t), dim=1)
         return tensor
 
-    def forward(
-        self,
-        q,  # BSND
-        k,
-        v,
-        t_b_idx,
-        base_blockmask,
-    ):
+    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:
@@ -188,7 +188,8 @@ def forward(
             x = attention_forward(q, k, v, opt_mode="manual", op_type="ascend_laser_attention", layout="BNSD")
         else:
             batch, qSeqlen, numHeads, headDim = q.shape
-            assert batch == 1, "Rainfusion_blockwise currently only supports batch size 1."
+            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
@@ -238,3 +239,95 @@ def forward(
             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 = torch.zeros((diff, tensor.shape[1], tensor.shape[2]), dtype=tensor.dtype, device=tensor.device)
+            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 = torch.zeros((actual_seqlen - expected_seqlen, attn_out.shape[1]), dtype=attn_out.dtype, device=attn_out.device)
+            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/rope/ascend_npu/npu_rope.py b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
index b08116811..346209283 100644
--- a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
+++ b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py
@@ -43,8 +43,8 @@ def __init__(self):
     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))
-        xk_fp32 = torch.view_as_complex(xk[:seq_len].to(torch.float32).reshape(seq_len, n, -1, 2))
+        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:
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..38dfbd5ba
--- /dev/null
+++ b/scripts/wan22/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh
@@ -0,0 +1,21 @@
+#!/bin/bash
+
+# set path firstly
+lightx2v_path=
+model_path=
+
+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
\ No newline at end of file

From 08465474b63a72807ef21b295edea49c5879633b Mon Sep 17 00:00:00 2001
From: liaopingbo <liaopingbo@foxmail.com>
Date: Wed, 17 Jun 2026 11:43:03 +0800
Subject: [PATCH 09/15] feat(npu): Perform inference with sparse FA and NPU
 affinity operators on Ascend NPU.

---
 .../models/networks/wan/infer/offload/transformer_infer.py  | 2 +-
 lightx2v/models/networks/wan/infer/transformer_infer.py     | 2 +-
 .../models/networks/worldmirror/models/layers/attention.py  | 2 +-
 lightx2v/models/vfi/rife/rife_comfyui_wrapper.py            | 6 +++---
 .../run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh     | 5 +++++
 5 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/lightx2v/models/networks/wan/infer/offload/transformer_infer.py b/lightx2v/models/networks/wan/infer/offload/transformer_infer.py
index 449714433..233614c5d 100755
--- a/lightx2v/models/networks/wan/infer/offload/transformer_infer.py
+++ b/lightx2v/models/networks/wan/infer/offload/transformer_infer.py
@@ -130,7 +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,
+                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 e3303d0f9..da4a2e843 100755
--- a/lightx2v/models/networks/wan/infer/transformer_infer.py
+++ b/lightx2v/models/networks/wan/infer/transformer_infer.py
@@ -101,7 +101,7 @@ def reset_infer_states(self):
     @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:
+        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)
diff --git a/lightx2v/models/networks/worldmirror/models/layers/attention.py b/lightx2v/models/networks/worldmirror/models/layers/attention.py
index 7c987f6a9..93a438673 100755
--- a/lightx2v/models/networks/worldmirror/models/layers/attention.py
+++ b/lightx2v/models/networks/worldmirror/models/layers/attention.py
@@ -22,7 +22,7 @@
 
 try:
     import torch_npu
-    _USE_NPU_FLASH_ATTN = True
+    _USE_NPU_FLASH_ATTN = torch.npu.is_available()
 except ImportError:
     torch_npu = None
     _USE_NPU_FLASH_ATTN = False
diff --git a/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py b/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py
index cc82a9f29..e3bc6251a 100644
--- a/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py
+++ b/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py
@@ -82,7 +82,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].to(self.device).float())
+                output_frames.append(images[source_idx1].float())
             else:
                 # Get frames to interpolate
                 frame1 = images[source_idx1]
@@ -103,10 +103,10 @@ def interpolate_frames(
 
                 # Convert back to ComfyUI format [H, W, C]
                 # Crop to original size and permute dimensions
-                interpolated_frame = interpolated[0, :, :height, :width].permute(1, 2, 0)
+                interpolated_frame = interpolated[0, :, :height, :width].permute(1, 2, 0).cpu()
                 output_frames.append(interpolated_frame)
 
-        return torch.stack(output_frames, dim=0).cpu()
+        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/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
index 38dfbd5ba..3be4fbf21 100644
--- 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
@@ -4,6 +4,11 @@
 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
 

From 5ed56f4fce94760f5261e371677703fe6c181d90 Mon Sep 17 00:00:00 2001
From: liaopingbo <liaopingbo@foxmail.com>
Date: Wed, 17 Jun 2026 15:04:36 +0800
Subject: [PATCH 10/15] feat(npu): Perform inference with sparse FA and NPU
 affinity operators on Ascend NPU.

---
 .../networks/wan/infer/transformer_infer.py   |  2 +-
 .../worldmirror/models/layers/attention.py    | 14 +++++--------
 .../models/vfi/rife/rife_comfyui_wrapper.py   |  2 +-
 .../attn/ascend_npu/npu_rainfusion_attn.py    | 20 +++++++++++++------
 4 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/lightx2v/models/networks/wan/infer/transformer_infer.py b/lightx2v/models/networks/wan/infer/transformer_infer.py
index da4a2e843..5fe35c418 100755
--- a/lightx2v/models/networks/wan/infer/transformer_infer.py
+++ b/lightx2v/models/networks/wan/infer/transformer_infer.py
@@ -154,7 +154,7 @@ def infer_block(self, block, x, pre_infer_out):
             x,
             shift_msa,
             scale_msa,
-            grid_sizes=pre_infer_out.grid_sizes.tuple,
+            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],
diff --git a/lightx2v/models/networks/worldmirror/models/layers/attention.py b/lightx2v/models/networks/worldmirror/models/layers/attention.py
index 93a438673..7c44e1690 100755
--- a/lightx2v/models/networks/worldmirror/models/layers/attention.py
+++ b/lightx2v/models/networks/worldmirror/models/layers/attention.py
@@ -68,23 +68,19 @@ def _compute_qkv(self, x: Tensor):
         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,
+            q,
+            k,
+            v,
             self.num_heads,
             pse=None,
             atten_mask=None,
             scale=self.scale,
             keep_prob=keep_prob,
-            input_layout="BNSD",
+            input_layout="BSND",
         )
-        x = x[0].transpose(1, 2)
-        return x
+        return x[0]
 
     def _apply_attention(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
         if q.dtype == torch.bfloat16 or q.dtype == torch.float16:
diff --git a/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py b/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py
index e3bc6251a..fb844c077 100644
--- a/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py
+++ b/lightx2v/models/vfi/rife/rife_comfyui_wrapper.py
@@ -82,7 +82,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].float())
+                output_frames.append(images[source_idx1].float().cpu())
             else:
                 # Get frames to interpolate
                 frame1 = images[source_idx1]
diff --git a/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
index 2fb91efc3..6aebd54b5 100644
--- a/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
+++ b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
@@ -1,9 +1,16 @@
 import math
 
-import mindiesd
 import torch
 from einops import rearrange
-from mindiesd.layers.flash_attn.attention_forward import attention_forward
+
+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
@@ -27,6 +34,7 @@ def __init__(self, grid_size=None, pool_size=128, sparsity=0.8, skip_timesteps=-
         """
         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
@@ -115,11 +123,11 @@ def rearrange_with_remaining(self, tensor):
 
         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), dim=2)
+            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), dim=3)
+            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
@@ -251,7 +259,7 @@ def _pad_to_expected(self, tensor, actual_seqlen, expected_seqlen):
             return tensor[:expected_seqlen], True
         elif actual_seqlen < expected_seqlen:
             diff = expected_seqlen - actual_seqlen
-            pad = torch.zeros((diff, tensor.shape[1], tensor.shape[2]), dtype=tensor.dtype, device=tensor.device)
+            pad = tensor.new_zeros((diff,) + tensor.shape[1:])
             return torch.cat([tensor, pad], dim=0), False
         else:
             return tensor, None
@@ -276,7 +284,7 @@ def apply(self, q, k, v, **kwargs):
         attn_out = attn_out.squeeze(0).reshape(attn_out.shape[1], -1)
 
         if trunc_input is True:
-            pad_out = torch.zeros((actual_seqlen - expected_seqlen, attn_out.shape[1]), dtype=attn_out.dtype, device=attn_out.device)
+            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]

From df8577a4923649a3d129c9c6295f546f04292f90 Mon Sep 17 00:00:00 2001
From: liaopingbo <liaopingbo@foxmail.com>
Date: Wed, 17 Jun 2026 17:16:33 +0800
Subject: [PATCH 11/15] feat(npu): pre-commit fix

---
 .../wan_moe_i2v_distill_int8_4step_ulysses_npu.json   |  2 +-
 .../models/networks/wan/infer/transformer_infer.py    |  1 +
 .../networks/worldmirror/models/layers/attention.py   |  7 ++++---
 lightx2v/models/vfi/rife/model/warplayer.py           |  8 +++-----
 lightx2v/models/vfi/rife/rife_comfyui_wrapper.py      |  7 ++-----
 lightx2v/models/vfi/rife/train_log/RIFE_HDv3.py       |  8 +++-----
 .../ops/attn/ascend_npu/npu_rainfusion_attn.py        | 11 ++++++++---
 ...un_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh |  2 +-
 8 files changed, 23 insertions(+), 23 deletions(-)

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
index 581dcebe6..f32405ca1 100644
--- 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
@@ -48,4 +48,4 @@
         "seq_p_size": 2,
         "seq_p_attn_type": "ulysses"
     }
-}
\ No newline at end of file
+}
diff --git a/lightx2v/models/networks/wan/infer/transformer_infer.py b/lightx2v/models/networks/wan/infer/transformer_infer.py
index 5fe35c418..6e7064a23 100755
--- a/lightx2v/models/networks/wan/infer/transformer_infer.py
+++ b/lightx2v/models/networks/wan/infer/transformer_infer.py
@@ -82,6 +82,7 @@ def rope_wrapper(xq, xk, cos_sin_cache):
         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()
diff --git a/lightx2v/models/networks/worldmirror/models/layers/attention.py b/lightx2v/models/networks/worldmirror/models/layers/attention.py
index 7c44e1690..ddfd991bf 100755
--- a/lightx2v/models/networks/worldmirror/models/layers/attention.py
+++ b/lightx2v/models/networks/worldmirror/models/layers/attention.py
@@ -6,6 +6,9 @@
 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
 
@@ -22,14 +25,12 @@
 
 try:
     import torch_npu
+
     _USE_NPU_FLASH_ATTN = torch.npu.is_available()
 except ImportError:
     torch_npu = None
     _USE_NPU_FLASH_ATTN = False
 
-from ...comm.communication import _All2All
-from ...comm.padding import depad_by_length, pad_by_length
-
 
 class Attention(nn.Module):
     def __init__(
diff --git a/lightx2v/models/vfi/rife/model/warplayer.py b/lightx2v/models/vfi/rife/model/warplayer.py
index 70565b0a7..8ea6fb33a 100644
--- a/lightx2v/models/vfi/rife/model/warplayer.py
+++ b/lightx2v/models/vfi/rife/model/warplayer.py
@@ -1,10 +1,8 @@
+import importlib.util
+
 import torch
 
-try:
-    import torch_npu
-    _HAS_NPU = torch.npu.is_available()
-except ImportError:
-    _HAS_NPU = False
+_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 fb844c077..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
 
@@ -16,11 +17,7 @@ def __init__(self, model_path, device: Optional[torch.device] = None):
         if device is not None:
             self.device = device
         else:
-            try:
-                import torch_npu
-                _HAS_NPU = torch.npu.is_available()
-            except ImportError:
-                _HAS_NPU = False
+            _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
diff --git a/lightx2v/models/vfi/rife/train_log/RIFE_HDv3.py b/lightx2v/models/vfi/rife/train_log/RIFE_HDv3.py
index 47dda9656..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,11 +7,7 @@
 from ..model.loss import *
 from .IFNet_HDv3 import *
 
-try:
-    import torch_npu
-    _HAS_NPU = torch.npu.is_available()
-except ImportError:
-    _HAS_NPU = False
+_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"))
 
 
diff --git a/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
index 6aebd54b5..70850f081 100644
--- a/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
+++ b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
@@ -6,6 +6,7 @@
 try:
     import mindiesd
     from mindiesd.layers.flash_attn.attention_forward import attention_forward
+
     _HAS_MINDIESD = True
 except ImportError:
     mindiesd = None
@@ -24,6 +25,8 @@
 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):
@@ -64,11 +67,11 @@ def avgpool(self, input_tensor, pool_size=128):
 
         pooled_tensors = []
         if num_full_blocks > 0:
-            full_blocks = input_tensor[:, :num_full_blocks * pool_size, :, :]
+            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_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))
 
@@ -326,7 +329,9 @@ def _init_rf(self):
 
     def apply_non_sp(self, q, k, v, grid_sizes):
         attn_out, self._base_blockmask = self._rf.apply(
-            q=q, k=k, v=v,
+            q=q,
+            k=k,
+            v=v,
             grid_size=list(grid_sizes),
             t_b_idx=(self.step_index, 0),
             base_blockmask=self._base_blockmask,
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
index 3be4fbf21..7d165d668 100644
--- 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
@@ -23,4 +23,4 @@ torchrun --nproc_per_node=2 -m lightx2v.infer \
 --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
\ No newline at end of file
+--save_result_path ${lightx2v_path}/save_results/wan_moe_i2v_distill_int8_4step_ulysses_npu.mp4

From 160d7c3819922dd6e400fbd610126905f2da70e9 Mon Sep 17 00:00:00 2001
From: wushuo <540295877@qq.com>
Date: Tue, 7 Jul 2026 10:38:46 +0000
Subject: [PATCH 12/15] update codes

---
 ...oe_i2v_distill_int8_4step_ulysses_npu.json |   6 +-
 lightx2v/common/ops/attn/__init__.py          |   1 +
 lightx2v/common/ops/attn/rainfusion_attn.py   | 307 +++++++++++++++
 .../networks/wan/infer/transformer_infer.py   |  45 +--
 .../wan/weights/transformer_weights.py        |   5 +
 .../attn/ascend_npu/npu_rainfusion_attn.py    | 355 ++----------------
 6 files changed, 360 insertions(+), 359 deletions(-)
 create mode 100644 lightx2v/common/ops/attn/rainfusion_attn.py

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
index f32405ca1..0ae50f3cb 100644
--- 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
@@ -4,12 +4,12 @@
     "text_len": 512,
     "target_height": 720,
     "target_width": 1280,
-    "use_rainfusion": true,
-    "rainfusion": {
+    "rainfusion_attn_setting": {
+        "pool_size": 128,
         "sparsity": 0.8,
         "skip_timesteps": -1
     },
-    "self_attn_1_type": "npu_flash_attn",
+    "self_attn_1_type": "rainfusion_attn",
     "cross_attn_1_type": "npu_flash_attn",
     "cross_attn_2_type": "npu_flash_attn",
     "sample_guide_scale": [
diff --git a/lightx2v/common/ops/attn/__init__.py b/lightx2v/common/ops/attn/__init__.py
index aeea39507..64eff2e52 100755
--- a/lightx2v/common/ops/attn/__init__.py
+++ b/lightx2v/common/ops/attn/__init__.py
@@ -3,6 +3,7 @@
 from .general_sparse_attn import GeneralSparseAttnWeight
 from .nbhd_attn import NbhdAttnWeight, NbhdAttnWeightFlashInfer
 from .radial_attn import RadialAttnWeight
+from .rainfusion_attn import RainfusionAttnWeight
 from .ring_attn import RingAttnWeight
 from .sage_attn import SageAttn2KInt8VFP8Weight, SageAttn2Weight, SageAttn3Weight, SparseSageAttn2Weight, SparseSageAttn3Weight
 from .sla_attn import SlaAttnWeight
diff --git a/lightx2v/common/ops/attn/rainfusion_attn.py b/lightx2v/common/ops/attn/rainfusion_attn.py
new file mode 100644
index 000000000..4e532b3ba
--- /dev/null
+++ b/lightx2v/common/ops/attn/rainfusion_attn.py
@@ -0,0 +1,307 @@
+import math
+
+import torch
+from einops import rearrange
+
+from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER
+from lightx2v_platform.base.global_var import AI_DEVICE
+
+from .template import AttnWeightTemplate
+
+
+@ATTN_WEIGHT_REGISTER("rainfusion_attn")
+class RainfusionAttnWeight(AttnWeightTemplate):
+    pool_size = 128
+    sparsity = 0.8
+    skip_timesteps = -1
+    text_len = 0
+    txt_first = False
+
+    _operator = None
+    _base_blockmask = None
+    _grid_size = None
+    _step_index = None
+
+    def __init__(self):
+        self.config = {}
+
+    @classmethod
+    def configure(cls, setting):
+        cls.pool_size = setting.get("pool_size", cls.pool_size)
+        cls.sparsity = setting.get("sparsity", cls.sparsity)
+        cls.skip_timesteps = setting.get("skip_timesteps", cls.skip_timesteps)
+        cls.text_len = setting.get("text_len", setting.get("txt_len", cls.text_len))
+        cls.txt_first = setting.get("txt_first", cls.txt_first)
+
+    @classmethod
+    def reset_state(cls):
+        cls._base_blockmask = None
+        cls._step_index = None
+
+    @classmethod
+    def _get_operator(cls):
+        if cls._operator is not None:
+            return cls._operator
+        if "npu" not in AI_DEVICE:
+            raise RuntimeError("rainfusion_attn currently requires ascend_npu platform.")
+        from lightx2v_platform.ops.attn.ascend_npu.npu_rainfusion_attn import NpuRainfusionOperator
+
+        cls._operator = NpuRainfusionOperator()
+        return cls._operator
+
+    @classmethod
+    def update_grid_size(cls, grid_size):
+        grid_size = tuple(grid_size)
+        if cls._grid_size != grid_size:
+            cls._grid_size = grid_size
+            cls._base_blockmask = None
+
+    @classmethod
+    def _get_grid_shape(cls):
+        if cls._grid_size is None:
+            raise ValueError("rainfusion_attn requires grid_sizes in apply kwargs.")
+        frame_num, h, w = cls._grid_size
+        return frame_num, h, w, h * w
+
+    @classmethod
+    def _get_expected_seqlen(cls):
+        frame_num, _, _, num_tokens_per_frame = cls._get_grid_shape()
+        return frame_num * num_tokens_per_frame
+
+    @classmethod
+    def avgpool(cls, input_tensor):
+        batch, seqlen, headnum, dim = input_tensor.shape
+        num_full_blocks = seqlen // cls.pool_size
+        tail_size = seqlen % cls.pool_size
+
+        pooled_tensors = []
+        if num_full_blocks > 0:
+            full_blocks = input_tensor[:, : num_full_blocks * cls.pool_size, :, :]
+            full_blocks_reshaped = full_blocks.reshape(batch, num_full_blocks, cls.pool_size, headnum, dim)
+            pooled_tensors.append(full_blocks_reshaped.mean(dim=2))
+        if tail_size > 0:
+            tail_block = input_tensor[:, num_full_blocks * cls.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)
+
+    @classmethod
+    def get_mask_index(cls, 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
+
+    @classmethod
+    def get_blockwise_mask(cls, score_matrix):
+        batch_size, num_heads, rows, cols = score_matrix.shape
+
+        keep_len = math.ceil(cols * (1 - cls.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
+
+        frame_num, _, _, num_tokens_per_frame = cls._get_grid_shape()
+        first_frame_len = num_tokens_per_frame
+        total_len = frame_num * num_tokens_per_frame + cls.text_len
+        protected_start_idx = total_len - first_frame_len - cls.text_len
+        protected_start_block = protected_start_idx // cls.pool_size
+        protect_len = cols - protected_start_block
+
+        if protect_len > 0:
+            mask[:, :, -protect_len:, :] = True
+            mask[:, :, :, -protect_len:] = True
+
+        select_idx = cls.get_mask_index(mask)
+        select_idx = select_idx[0].transpose(0, 1)
+        select_num_idx = mask[0].transpose(0, 1).sum(dim=-1)
+        return select_idx, select_num_idx
+
+    @classmethod
+    def rearrange_with_remaining(cls, tensor):
+        b, s, n, d = tensor.shape
+        frame_num, h, w, first_frame_len = cls._get_grid_shape()
+        h_res_len, w_res_len = 0, 0
+        first_frame_num = first_frame_len // h // w
+
+        tensor_hwt = rearrange(tensor, "b (f h w) n d -> (b n) f h w d", f=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 = 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
+
+    @classmethod
+    def inv_rearrange_with_remaining(cls, tensor, h_res_len, w_res_len):
+        b, s, n, d = tensor.shape
+        frame_num, h, w, first_frame_len = cls._get_grid_shape()
+        h_sr, w_sr = h % 8, w % 8
+        first_frame_num = first_frame_len // (h * w)
+        remaining_frames = 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
+
+    @classmethod
+    def do_tensor_rearrange_pooling(cls, tensor):
+        b, s, n, d = tensor.shape
+        _, _, _, first_frame_len = cls._get_grid_shape()
+        if s < cls.text_len + first_frame_len:
+            raise ValueError(f"Sequence length {s} is too small for text_len {cls.text_len} and first_frame_len {first_frame_len}")
+        if cls.txt_first:
+            tensor_t, tensor_f, tensor_i = torch.split(tensor, [cls.text_len, first_frame_len, s - cls.text_len - first_frame_len], dim=1)
+        else:
+            tensor_f, tensor_i, tensor_t = torch.split(tensor, [first_frame_len, s - cls.text_len - first_frame_len, cls.text_len], dim=1)
+        tensor_i_2, h_res_len, w_res_len = cls.rearrange_with_remaining(tensor_i)
+        tensor = torch.cat((tensor_i_2, tensor_f, tensor_t), dim=1)
+        tensor_pool = cls.avgpool(tensor)
+        return tensor, tensor_pool, h_res_len, w_res_len
+
+    @classmethod
+    def do_tensor_inv_rearrange(cls, tensor, h_res_len, w_res_len):
+        b, s, n, d = tensor.shape
+        _, _, _, first_frame_len = cls._get_grid_shape()
+        tensor_i, tensor_f, tensor_t = torch.split(tensor, [s - cls.text_len - first_frame_len, first_frame_len, cls.text_len], dim=1)
+        tensor_i = cls.inv_rearrange_with_remaining(tensor_i, h_res_len, w_res_len)
+
+        if cls.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
+
+    @staticmethod
+    def _pad_to_expected(tensor, expected_seqlen):
+        actual_seqlen = tensor.shape[0]
+        if actual_seqlen > expected_seqlen:
+            return tensor[:expected_seqlen], True
+        if 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
+        return tensor, None
+
+    @classmethod
+    def _get_step_index(cls, scheduler):
+        if scheduler is None:
+            return 0
+        return getattr(scheduler, "step_index", 0)
+
+    def apply(
+        self,
+        q,
+        k,
+        v,
+        cu_seqlens_q=None,
+        cu_seqlens_kv=None,
+        max_seqlen_q=None,
+        max_seqlen_kv=None,
+        **kwargs,
+    ):
+        grid_sizes = kwargs.get("grid_sizes", None)
+        if grid_sizes is None:
+            raise ValueError("rainfusion_attn requires grid_sizes in apply kwargs.")
+        cls = type(self)
+        cls.update_grid_size(grid_sizes)
+
+        scheduler = kwargs.get("scheduler", None)
+        step_index = cls._get_step_index(scheduler)
+        if cls._step_index != step_index:
+            cls._base_blockmask = None
+            cls._step_index = step_index
+
+        if len(q.shape) == 4:
+            bs = q.shape[0]
+            if bs != 1:
+                raise ValueError("rainfusion_attn currently only supports batch size 1.")
+            q_4d, k_4d, v_4d = q, k, v
+            actual_seqlen = q.shape[1]
+            trunc_input = None
+        else:
+            expected_seqlen = self._get_expected_seqlen()
+            q_valid, trunc_input = self._pad_to_expected(q, expected_seqlen)
+            k_valid, _ = self._pad_to_expected(k, expected_seqlen)
+            v_valid, _ = self._pad_to_expected(v, expected_seqlen)
+            actual_seqlen = q.shape[0]
+            q_4d = q_valid.unsqueeze(0)
+            k_4d = k_valid.unsqueeze(0)
+            v_4d = v_valid.unsqueeze(0)
+
+        if step_index < self.skip_timesteps:
+            out = cls._get_operator().dense_attention(q_4d, k_4d, v_4d)
+        else:
+            batch, q_seqlen, num_heads, head_dim = q_4d.shape
+            if batch != 1:
+                raise ValueError("rainfusion_attn currently only supports batch size 1.")
+            kv_seqlen = k_4d.shape[1]
+            scale = head_dim**-0.5
+
+            h_res_len, w_res_len = 0, 0
+            qkv = torch.cat((q_4d, k_4d, v_4d), dim=0)
+            qkv, qkv_pool, h_res_len, w_res_len = self.do_tensor_rearrange_pooling(qkv)
+            q_4d, k_4d, v_4d = torch.chunk(qkv, 3, dim=0)
+
+            if cls._base_blockmask is None:
+                query_pool, key_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)
+                cls._base_blockmask = cls.get_blockwise_mask(attn_scores_fake)
+
+            select_idx, select_num_idx = cls._base_blockmask
+            out = cls._get_operator().sparse_attention(
+                q_4d,
+                k_4d,
+                v_4d,
+                scale=scale,
+                head_num=num_heads,
+                select_idx=select_idx,
+                select_num_idx=select_num_idx,
+                block_shape=[self.pool_size, self.pool_size],
+                actual_seq_lengths=[q_seqlen for _ in range(batch)],
+                actual_seq_lengths_kv=[kv_seqlen for _ in range(batch)],
+            )
+            out = self.do_tensor_inv_rearrange(out, h_res_len, w_res_len)
+
+        out = out.squeeze(0).reshape(out.shape[1], -1)
+        if trunc_input is True:
+            pad_out = out.new_zeros((actual_seqlen - out.shape[0], out.shape[1]))
+            out = torch.cat([out, pad_out], dim=0)
+        elif trunc_input is False:
+            out = out[:actual_seqlen]
+        return out
diff --git a/lightx2v/models/networks/wan/infer/transformer_infer.py b/lightx2v/models/networks/wan/infer/transformer_infer.py
index 6e7064a23..b4136c4b8 100755
--- a/lightx2v/models/networks/wan/infer/transformer_infer.py
+++ b/lightx2v/models/networks/wan/infer/transformer_infer.py
@@ -79,12 +79,6 @@ 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
@@ -94,17 +88,21 @@ 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()
 
+    def reset_attention_states(self, blocks):
+        for block in blocks:
+            self_attn = block.compute_phases[0].self_attn_1
+            reset_state = getattr(self_attn, "reset_state", None)
+            if reset_state is not None:
+                reset_state()
+
     @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()
+        self.reset_attention_states(weights.blocks)
         x = self.infer_main_blocks(weights.blocks, pre_infer_out)
         return self.infer_non_blocks(weights, x, pre_infer_out.embed)
 
@@ -242,33 +240,10 @@ def infer_self_attn(self, phase, x, shift_msa, scale_msa, grid_sizes=None):
         attn_running_args = {
             "block_idx": self.block_idx,
             "scheduler": self.scheduler,
+            "grid_sizes": grid_sizes,
         }
 
-        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"]:
+        if self.config["seq_parallel"]:
             attn_out = phase.self_attn_1_parallel.apply(
                 q=q,
                 k=k,
diff --git a/lightx2v/models/networks/wan/weights/transformer_weights.py b/lightx2v/models/networks/wan/weights/transformer_weights.py
index 8d09b4c1e..23cb4753c 100755
--- a/lightx2v/models/networks/wan/weights/transformer_weights.py
+++ b/lightx2v/models/networks/wan/weights/transformer_weights.py
@@ -359,6 +359,11 @@ def __init__(
                 if "min_width" in self.config["nbhd_attn_setting"]:
                     attention_weights_cls.min_width = self.config["nbhd_attn_setting"]["min_width"]
 
+        # rainfusion_attn setting
+        if self.config["self_attn_1_type"] == "rainfusion_attn":
+            rainfusion_config = self.config.get("rainfusion_attn_setting", self.config.get("rainfusion", {}))
+            attention_weights_cls.configure(rainfusion_config)
+
         # draft_attn setting
         if self.config["self_attn_1_type"] == "draft_attn":
             attention_weights_cls.sparsity_ratio = self.config.get("draft_attn_sparsity_ratio", 0.75)
diff --git a/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
index 70850f081..da83fdb0b 100644
--- a/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
+++ b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
@@ -1,8 +1,3 @@
-import math
-
-import torch
-from einops import rearrange
-
 try:
     import mindiesd
     from mindiesd.layers.flash_attn.attention_forward import attention_forward
@@ -13,9 +8,6 @@
     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:
@@ -27,320 +19,41 @@
 """
 
 
-@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 = {}
+class NpuRainfusionOperator:
+    def __init__(self):
         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 dense_attention(self, q, k, v):
+        return attention_forward(q, k, v, opt_mode="manual", op_type="ascend_laser_attention", layout="BNSD")
+
+    def sparse_attention(
+        self,
+        q,
+        k,
+        v,
+        scale,
+        head_num,
+        select_idx,
+        select_num_idx,
+        block_shape,
+        actual_seq_lengths,
+        actual_seq_lengths_kv,
+    ):
+        q_bnsd = q.transpose(1, 2)
+        k_bnsd = k.transpose(1, 2)
+        v_bnsd = v.transpose(1, 2)
+        out = mindiesd.layers.flash_attn.sparse_flash_attn_rf_v2.rain_fusion_attention(
+            q_bnsd,
+            k_bnsd,
+            v_bnsd,
+            scale=scale,
+            head_num=head_num,
+            input_layout="BNSD",
+            select_idx=select_idx,
+            select_num_idx=select_num_idx,
+            blockshape=block_shape,
+            actual_seq_lengths=actual_seq_lengths,
+            actual_seq_lengths_kv=actual_seq_lengths_kv,
         )
-
-    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
+        return out.transpose(1, 2).reshape(q.shape)

From ea74ce658a56d16e96252c662d062fc937f522c0 Mon Sep 17 00:00:00 2001
From: wushuo <540295877@qq.com>
Date: Wed, 8 Jul 2026 08:27:48 +0000
Subject: [PATCH 13/15] temp

---
 .../attentions/wan_i2v_rainfusion_480p.json   |  24 ++++
 lightx2v/common/ops/attn/rainfusion_attn.py   | 110 +++++++++++++++---
 lightx2v/common/ops/attn/sparse_operator.py   |  36 +++++-
 .../attn/ascend_npu/npu_rainfusion_attn.py    |  13 ++-
 .../attentions/run_wan_i2v_nbhd_attn_480p.sh  |   4 +-
 .../run_wan_i2v_rainfusion_attn_480p.sh       |  20 ++++
 scripts/wan/run_wan_t2v.sh                    |   4 +-
 7 files changed, 184 insertions(+), 27 deletions(-)
 create mode 100755 configs/attentions/wan_i2v_rainfusion_480p.json
 create mode 100755 scripts/attentions/run_wan_i2v_rainfusion_attn_480p.sh

diff --git a/configs/attentions/wan_i2v_rainfusion_480p.json b/configs/attentions/wan_i2v_rainfusion_480p.json
new file mode 100755
index 000000000..967694443
--- /dev/null
+++ b/configs/attentions/wan_i2v_rainfusion_480p.json
@@ -0,0 +1,24 @@
+{
+    "infer_steps": 40,
+    "target_video_length": 81,
+    "target_height": 480,
+    "target_width": 832,
+    "self_attn_1_type": "rainfusion_attn",
+    "rainfusion_attn_setting": {
+        "backend": "cuda_sparse",
+        "sparse_operator": "flashinfer_operator",
+        "dense_attn_type": "flash_attn3",
+        "pool_size": 128,
+        "sparsity": 0.8,
+        "skip_timesteps": -1
+    },
+    "cross_attn_1_type": "flash_attn3",
+    "cross_attn_2_type": "flash_attn3",
+    "sample_guide_scale": 5,
+    "sample_shift": 3,
+    "enable_cfg": true,
+    "cpu_offload": false,
+    "t5_cpu_offload": true,
+    "clip_cpu_offload": true,
+    "vae_cpu_offload": true
+}
diff --git a/lightx2v/common/ops/attn/rainfusion_attn.py b/lightx2v/common/ops/attn/rainfusion_attn.py
index 4e532b3ba..b215e6b6c 100644
--- a/lightx2v/common/ops/attn/rainfusion_attn.py
+++ b/lightx2v/common/ops/attn/rainfusion_attn.py
@@ -3,12 +3,62 @@
 import torch
 from einops import rearrange
 
-from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER
+from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, SPARSE_OPERATOR_REGISTER
 from lightx2v_platform.base.global_var import AI_DEVICE
 
 from .template import AttnWeightTemplate
 
 
+class CudaRainfusionSparseOperator:
+    def __init__(self, sparse_operator, dense_attn_type="torch_sdpa", pool_size=128, operator_setting=None):
+        self.pool_size = pool_size
+        self.dense_attn_type = dense_attn_type
+        self.operator_setting = operator_setting or {}
+        operator_cls = SPARSE_OPERATOR_REGISTER[sparse_operator]
+        try:
+            self.sparse_operator = operator_cls(q_block_size=pool_size, k_block_size=pool_size, operator_setting=self.operator_setting)
+        except TypeError:
+            self.sparse_operator = operator_cls(self.operator_setting)
+            if getattr(self.sparse_operator, "q_block_size", pool_size) != pool_size or getattr(self.sparse_operator, "k_block_size", pool_size) != pool_size:
+                raise ValueError(f"{sparse_operator} does not support RainFusion pool_size={pool_size}.")
+        self.dense_attn = ATTN_WEIGHT_REGISTER[dense_attn_type]()
+
+    def dense_attention(self, q, k, v):
+        b, q_seqlen, num_heads, head_dim = q.shape
+        if b != 1:
+            raise ValueError("CudaRainfusionSparseOperator currently only supports batch size 1.")
+        out = self.dense_attn.apply(q, k, v, max_seqlen_q=q_seqlen, max_seqlen_kv=k.shape[1])
+        return out.reshape(b, q_seqlen, num_heads, head_dim)
+
+    def sparse_attention(
+        self,
+        q,
+        k,
+        v,
+        scale,
+        head_num,
+        block_mask,
+        select_idx=None,
+        select_num_idx=None,
+        block_shape=None,
+        actual_seq_lengths=None,
+        actual_seq_lengths_kv=None,
+    ):
+        b, q_seqlen, num_heads, head_dim = q.shape
+        if b != 1:
+            raise ValueError("CudaRainfusionSparseOperator currently only supports batch size 1.")
+        out = self.sparse_operator(
+            q.squeeze(0),
+            k.squeeze(0),
+            v.squeeze(0),
+            block_mask,
+            max_seqlen_q=q_seqlen,
+            max_seqlen_kv=k.shape[1],
+            softmax_scale=scale,
+        )
+        return out.reshape(b, q_seqlen, num_heads, head_dim)
+
+
 @ATTN_WEIGHT_REGISTER("rainfusion_attn")
 class RainfusionAttnWeight(AttnWeightTemplate):
     pool_size = 128
@@ -16,8 +66,13 @@ class RainfusionAttnWeight(AttnWeightTemplate):
     skip_timesteps = -1
     text_len = 0
     txt_first = False
+    backend = "auto"
+    sparse_operator = "flashinfer_operator"
+    dense_attn_type = "torch_sdpa"
+    operator_setting = {}
 
     _operator = None
+    _operator_key = None
     _base_blockmask = None
     _grid_size = None
     _step_index = None
@@ -32,6 +87,13 @@ def configure(cls, setting):
         cls.skip_timesteps = setting.get("skip_timesteps", cls.skip_timesteps)
         cls.text_len = setting.get("text_len", setting.get("txt_len", cls.text_len))
         cls.txt_first = setting.get("txt_first", cls.txt_first)
+        cls.backend = setting.get("backend", cls.backend)
+        cls.sparse_operator = setting.get("sparse_operator", cls.sparse_operator)
+        cls.dense_attn_type = setting.get("dense_attn_type", cls.dense_attn_type)
+        cls.operator_setting = setting.get("operator_setting", cls.operator_setting)
+        cls._operator = None
+        cls._operator_key = None
+        cls._base_blockmask = None
 
     @classmethod
     def reset_state(cls):
@@ -40,13 +102,26 @@ def reset_state(cls):
 
     @classmethod
     def _get_operator(cls):
-        if cls._operator is not None:
+        backend = cls.backend
+        if backend == "auto":
+            backend = "npu" if "npu" in AI_DEVICE else "cuda_sparse"
+        operator_key = (backend, cls.sparse_operator, cls.dense_attn_type, cls.pool_size, repr(cls.operator_setting))
+        if cls._operator is not None and cls._operator_key == operator_key:
             return cls._operator
-        if "npu" not in AI_DEVICE:
-            raise RuntimeError("rainfusion_attn currently requires ascend_npu platform.")
-        from lightx2v_platform.ops.attn.ascend_npu.npu_rainfusion_attn import NpuRainfusionOperator
-
-        cls._operator = NpuRainfusionOperator()
+        if backend == "npu":
+            from lightx2v_platform.ops.attn.ascend_npu.npu_rainfusion_attn import NpuRainfusionOperator
+
+            cls._operator = NpuRainfusionOperator()
+        elif backend == "cuda_sparse":
+            cls._operator = CudaRainfusionSparseOperator(
+                sparse_operator=cls.sparse_operator,
+                dense_attn_type=cls.dense_attn_type,
+                pool_size=cls.pool_size,
+                operator_setting=cls.operator_setting,
+            )
+        else:
+            raise ValueError(f"Unsupported rainfusion_attn backend: {backend}")
+        cls._operator_key = operator_key
         return cls._operator
 
     @classmethod
@@ -105,7 +180,7 @@ def get_mask_index(cls, mask):
         return pos_matrix
 
     @classmethod
-    def get_blockwise_mask(cls, score_matrix):
+    def build_blockwise_mask(cls, score_matrix):
         batch_size, num_heads, rows, cols = score_matrix.shape
 
         keep_len = math.ceil(cols * (1 - cls.sparsity))
@@ -125,6 +200,15 @@ def get_blockwise_mask(cls, score_matrix):
             mask[:, :, -protect_len:, :] = True
             mask[:, :, :, -protect_len:] = True
 
+        return mask
+
+    @classmethod
+    def get_blockwise_mask(cls, score_matrix):
+        mask = cls.build_blockwise_mask(score_matrix)
+        return cls.get_select_indices(mask)
+
+    @classmethod
+    def get_select_indices(cls, mask):
         select_idx = cls.get_mask_index(mask)
         select_idx = select_idx[0].transpose(0, 1)
         select_num_idx = mask[0].transpose(0, 1).sum(dim=-1)
@@ -242,9 +326,7 @@ def apply(
 
         scheduler = kwargs.get("scheduler", None)
         step_index = cls._get_step_index(scheduler)
-        if cls._step_index != step_index:
-            cls._base_blockmask = None
-            cls._step_index = step_index
+        cls._step_index = step_index
 
         if len(q.shape) == 4:
             bs = q.shape[0]
@@ -281,15 +363,17 @@ def apply(
                 query_pool, key_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)
-                cls._base_blockmask = cls.get_blockwise_mask(attn_scores_fake)
+                cls._base_blockmask = cls.build_blockwise_mask(attn_scores_fake)
 
-            select_idx, select_num_idx = cls._base_blockmask
+            block_mask = cls._base_blockmask
+            select_idx, select_num_idx = cls.get_select_indices(block_mask)
             out = cls._get_operator().sparse_attention(
                 q_4d,
                 k_4d,
                 v_4d,
                 scale=scale,
                 head_num=num_heads,
+                block_mask=block_mask,
                 select_idx=select_idx,
                 select_num_idx=select_num_idx,
                 block_shape=[self.pool_size, self.pool_size],
diff --git a/lightx2v/common/ops/attn/sparse_operator.py b/lightx2v/common/ops/attn/sparse_operator.py
index 6dfa70db7..eda2823e4 100644
--- a/lightx2v/common/ops/attn/sparse_operator.py
+++ b/lightx2v/common/ops/attn/sparse_operator.py
@@ -234,6 +234,7 @@ def __call__(
             q_ranges=q_ranges,
             k_ranges=k_ranges,
             attn_type_map=attn_type_map,
+            softmax_scale=kwargs.get("softmax_scale", None),
             auto_range_merge=True,
         )[0]
 
@@ -287,14 +288,19 @@ def __call__(
 class FlashinferOperator:
     sparse_wrapper = None
     mask = None
+    plan_key = None
 
     def __init__(self, q_block_size=128, k_block_size=128, operator_setting={}):
         self.q_block_size = q_block_size
         self.k_block_size = k_block_size
         self.operator_setting = operator_setting
-        if FlashinferOperator.sparse_wrapper is None:
+        self._ensure_sparse_wrapper()
+
+    @classmethod
+    def _ensure_sparse_wrapper(cls):
+        if cls.sparse_wrapper is None:
             float_workspace_buffer = torch.empty(1024 * 1024 * 1024, dtype=torch.uint8, device=AI_DEVICE)
-            FlashinferOperator.sparse_wrapper = flashinfer.sparse.VariableBlockSparseAttentionWrapper(float_workspace_buffer, backend="fa2")
+            cls.sparse_wrapper = flashinfer.sparse.VariableBlockSparseAttentionWrapper(float_workspace_buffer, backend="fa2")
 
     def __call__(
         self,
@@ -308,16 +314,33 @@ def __call__(
         max_seqlen_kv=None,
         **kwargs,
     ):
+        self._ensure_sparse_wrapper()
         seqlen, head_num, head_dim = q.shape
+        kv_seqlen = k.shape[0]
+        softmax_scale = kwargs.get("softmax_scale", None)
         # (B, H, Q_block_num, K_block_num) -> (H, Q_block_num, K_block_num)
         mask = mask.squeeze(0)
-        if FlashinferOperator.mask is None or not torch.equal(mask, FlashinferOperator.mask):
+        plan_key = (
+            tuple(mask.shape),
+            seqlen,
+            kv_seqlen,
+            head_num,
+            head_dim,
+            q.dtype,
+            k.dtype,
+            str(q.device),
+            str(k.device),
+            self.q_block_size,
+            self.k_block_size,
+            softmax_scale,
+        )
+        if FlashinferOperator.mask is None or FlashinferOperator.plan_key != plan_key or not torch.equal(mask, FlashinferOperator.mask):
             _, q_block_num, k_block_num = mask.shape
             block_row_sz = torch.ones(q_block_num, dtype=torch.int32, device=q.device) * self.q_block_size
             block_row_sz[-1] = seqlen - self.q_block_size * (q_block_num - 1)
             block_row_sz = block_row_sz.unsqueeze(0).repeat(head_num, 1)
             block_col_sz = torch.ones(k_block_num, dtype=torch.int32, device=k.device) * self.k_block_size
-            block_col_sz[-1] = seqlen - self.k_block_size * (k_block_num - 1)
+            block_col_sz[-1] = kv_seqlen - self.k_block_size * (k_block_num - 1)
             block_col_sz = block_col_sz.unsqueeze(0).repeat(head_num, 1)
             FlashinferOperator.sparse_wrapper.plan(
                 block_mask_map=mask,
@@ -327,8 +350,11 @@ def __call__(
                 num_kv_heads=head_num,
                 head_dim=head_dim,
                 q_data_type=q.dtype,
+                kv_data_type=k.dtype,
+                sm_scale=softmax_scale,
             )
-            FlashinferOperator.mask = mask
+            FlashinferOperator.mask = mask.clone()
+            FlashinferOperator.plan_key = plan_key
 
         q = q.transpose(0, 1)
         k = k.transpose(0, 1)
diff --git a/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
index da83fdb0b..f7bad13f1 100644
--- a/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
+++ b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py
@@ -34,12 +34,15 @@ def sparse_attention(
         v,
         scale,
         head_num,
-        select_idx,
-        select_num_idx,
-        block_shape,
-        actual_seq_lengths,
-        actual_seq_lengths_kv,
+        block_mask=None,
+        select_idx=None,
+        select_num_idx=None,
+        block_shape=None,
+        actual_seq_lengths=None,
+        actual_seq_lengths_kv=None,
     ):
+        if select_idx is None or select_num_idx is None:
+            raise ValueError("NpuRainfusionOperator requires select_idx and select_num_idx.")
         q_bnsd = q.transpose(1, 2)
         k_bnsd = k.transpose(1, 2)
         v_bnsd = v.transpose(1, 2)
diff --git a/scripts/attentions/run_wan_i2v_nbhd_attn_480p.sh b/scripts/attentions/run_wan_i2v_nbhd_attn_480p.sh
index 5dde95a1c..c78773006 100755
--- a/scripts/attentions/run_wan_i2v_nbhd_attn_480p.sh
+++ b/scripts/attentions/run_wan_i2v_nbhd_attn_480p.sh
@@ -1,8 +1,8 @@
 #!/bin/bash
 
 # set path firstly
-lightx2v_path=
-model_path=
+lightx2v_path=/data/nvme1/wushuo/LightX2V
+model_path=/data/nvme1/wushuo/hf_models/Wan2.1-I2V-14B-480P
 
 export CUDA_VISIBLE_DEVICES=0
 
diff --git a/scripts/attentions/run_wan_i2v_rainfusion_attn_480p.sh b/scripts/attentions/run_wan_i2v_rainfusion_attn_480p.sh
new file mode 100755
index 000000000..8b08b7404
--- /dev/null
+++ b/scripts/attentions/run_wan_i2v_rainfusion_attn_480p.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+# set path firstly
+lightx2v_path=/data/nvme1/wushuo/LightX2V
+model_path=/data/nvme1/wushuo/hf_models/Wan2.1-I2V-14B-480P
+
+export CUDA_VISIBLE_DEVICES=0
+
+# set environment variables
+source ${lightx2v_path}/scripts/base/base.sh
+
+python -m lightx2v.infer \
+--model_cls wan2.1 \
+--task i2v \
+--model_path $model_path \
+--config_json ${lightx2v_path}/configs/attentions/wan_i2v_rainfusion_480p.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/output_lightx2v_wan_i2v_rainfusion_attn_480p.mp4
diff --git a/scripts/wan/run_wan_t2v.sh b/scripts/wan/run_wan_t2v.sh
index a3436d47a..d5bad848b 100755
--- a/scripts/wan/run_wan_t2v.sh
+++ b/scripts/wan/run_wan_t2v.sh
@@ -1,8 +1,8 @@
 #!/bin/bash
 
 # set path firstly
-lightx2v_path=
-model_path=
+lightx2v_path=/data/nvme1/wushuo/LightX2V
+model_path=/data/nvme1/wushuo/hf_models/models/Wan2.1-T2V-1.3B
 
 export CUDA_VISIBLE_DEVICES=0
 

From 611eae35dc13f3bb7858d40475e30a715ed6be85 Mon Sep 17 00:00:00 2001
From: wushuo <540295877@qq.com>
Date: Wed, 8 Jul 2026 12:12:43 +0000
Subject: [PATCH 14/15] support nv backend rainfusion

---
 .../attentions/wan_i2v_rainfusion_480p.json   | 24 -----------
 ...ll_model_4step_cfg_ulysses_rainfusion.json | 42 +++++++++++++++++++
 .../run_wan_i2v_rainfusion_attn_480p.sh       | 20 ---------
 ...oe_i2v_distill_4step_ulysses_rainfusion.sh | 20 +++++++++
 ..._moe_i2v_distill_int8_4step_ulysses_npu.sh |  4 +-
 5 files changed, 64 insertions(+), 46 deletions(-)
 delete mode 100755 configs/attentions/wan_i2v_rainfusion_480p.json
 create mode 100644 configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json
 delete mode 100755 scripts/attentions/run_wan_i2v_rainfusion_attn_480p.sh
 create mode 100644 scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh

diff --git a/configs/attentions/wan_i2v_rainfusion_480p.json b/configs/attentions/wan_i2v_rainfusion_480p.json
deleted file mode 100755
index 967694443..000000000
--- a/configs/attentions/wan_i2v_rainfusion_480p.json
+++ /dev/null
@@ -1,24 +0,0 @@
-{
-    "infer_steps": 40,
-    "target_video_length": 81,
-    "target_height": 480,
-    "target_width": 832,
-    "self_attn_1_type": "rainfusion_attn",
-    "rainfusion_attn_setting": {
-        "backend": "cuda_sparse",
-        "sparse_operator": "flashinfer_operator",
-        "dense_attn_type": "flash_attn3",
-        "pool_size": 128,
-        "sparsity": 0.8,
-        "skip_timesteps": -1
-    },
-    "cross_attn_1_type": "flash_attn3",
-    "cross_attn_2_type": "flash_attn3",
-    "sample_guide_scale": 5,
-    "sample_shift": 3,
-    "enable_cfg": true,
-    "cpu_offload": false,
-    "t5_cpu_offload": true,
-    "clip_cpu_offload": true,
-    "vae_cpu_offload": true
-}
diff --git a/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json b/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json
new file mode 100644
index 000000000..14acc6e63
--- /dev/null
+++ b/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json
@@ -0,0 +1,42 @@
+{
+    "infer_steps": 4,
+    "target_video_length": 81,
+    "text_len": 512,
+    "target_height": 720,
+    "target_width": 1280,
+    "rainfusion_attn_setting": {
+        "backend": "cuda_sparse",
+        "sparse_operator": "flashinfer_operator",
+        "dense_attn_type": "flash_attn3",
+        "pool_size": 128,
+        "sparsity": 0.8,
+        "skip_timesteps": -1
+    },
+    "self_attn_1_type": "rainfusion_attn",
+    "cross_attn_1_type": "flash_attn3",
+    "cross_attn_2_type": "flash_attn3",
+    "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_step_index": 2,
+    "denoising_step_list": [
+        1000,
+        750,
+        500,
+        250
+    ],
+    "high_noise_original_ckpt": "/data/nvme1/wushuo/hf_models/Wan2.2-Distill-Models/wan2.2_i2v_A14b_high_noise_lightx2v_4step.safetensors",
+    "low_noise_original_ckpt": "/data/nvme1/wushuo/hf_models/Wan2.2-Distill-Models/wan2.2_i2v_A14b_low_noise_lightx2v_4step.safetensors",
+    "parallel": {
+        "seq_p_size": 4,
+        "seq_p_attn_type": "ulysses"
+    }
+}
diff --git a/scripts/attentions/run_wan_i2v_rainfusion_attn_480p.sh b/scripts/attentions/run_wan_i2v_rainfusion_attn_480p.sh
deleted file mode 100755
index 8b08b7404..000000000
--- a/scripts/attentions/run_wan_i2v_rainfusion_attn_480p.sh
+++ /dev/null
@@ -1,20 +0,0 @@
-#!/bin/bash
-
-# set path firstly
-lightx2v_path=/data/nvme1/wushuo/LightX2V
-model_path=/data/nvme1/wushuo/hf_models/Wan2.1-I2V-14B-480P
-
-export CUDA_VISIBLE_DEVICES=0
-
-# set environment variables
-source ${lightx2v_path}/scripts/base/base.sh
-
-python -m lightx2v.infer \
---model_cls wan2.1 \
---task i2v \
---model_path $model_path \
---config_json ${lightx2v_path}/configs/attentions/wan_i2v_rainfusion_480p.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/output_lightx2v_wan_i2v_rainfusion_attn_480p.mp4
diff --git a/scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh b/scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh
new file mode 100644
index 000000000..229427fa5
--- /dev/null
+++ b/scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+# set path firstly
+lightx2v_path=/data/nvme1/wushuo/LightX2V
+model_path=/data/nvme1/wushuo/hf_models/Wan2.2-I2V-A14B
+
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+# set environment variables
+source ${lightx2v_path}/scripts/base/base.sh
+
+torchrun --nproc_per_node=4 -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_model_4step_cfg_ulysses_rainfusion.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_4step_cfg_ulysses_rainfusion.mp4
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
index 7d165d668..b749f4ca7 100644
--- 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
@@ -1,8 +1,8 @@
 #!/bin/bash
 
 # set path firstly
-lightx2v_path=
-model_path=
+lightx2v_path=/data/nvme1/wushuo/LightX2V
+model_path=/data/nvme1/wushuo/hf_models/Wan2.2-I2V-A14B
 
 if [ -z "$lightx2v_path" ] || [ -z "$model_path" ]; then
     echo "Error: Please set lightx2v_path and model_path in this script before running."

From 8a5573b7c90011de8c06db531460d6def543e0e3 Mon Sep 17 00:00:00 2001
From: wushuo <540295877@qq.com>
Date: Wed, 8 Jul 2026 12:15:36 +0000
Subject: [PATCH 15/15] rm personal path

---
 ...an_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json | 4 ++--
 .../run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh     | 4 ++--
 .../run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh       | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json b/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json
index 14acc6e63..4ec6e8c01 100644
--- a/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json
+++ b/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json
@@ -33,8 +33,8 @@
         500,
         250
     ],
-    "high_noise_original_ckpt": "/data/nvme1/wushuo/hf_models/Wan2.2-Distill-Models/wan2.2_i2v_A14b_high_noise_lightx2v_4step.safetensors",
-    "low_noise_original_ckpt": "/data/nvme1/wushuo/hf_models/Wan2.2-Distill-Models/wan2.2_i2v_A14b_low_noise_lightx2v_4step.safetensors",
+    "high_noise_original_ckpt": "/Wan2.2-Distill-Models/wan2.2_i2v_A14b_high_noise_lightx2v_4step.safetensors",
+    "low_noise_original_ckpt": "/Wan2.2-Distill-Models/wan2.2_i2v_A14b_low_noise_lightx2v_4step.safetensors",
     "parallel": {
         "seq_p_size": 4,
         "seq_p_attn_type": "ulysses"
diff --git a/scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh b/scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh
index 229427fa5..8ac5ee6a3 100644
--- a/scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh
+++ b/scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh
@@ -1,8 +1,8 @@
 #!/bin/bash
 
 # set path firstly
-lightx2v_path=/data/nvme1/wushuo/LightX2V
-model_path=/data/nvme1/wushuo/hf_models/Wan2.2-I2V-A14B
+lightx2v_path=/LightX2V
+model_path=/Wan2.2-I2V-A14B
 
 export CUDA_VISIBLE_DEVICES=0,1,2,3
 
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
index b749f4ca7..c397df405 100644
--- 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
@@ -1,8 +1,8 @@
 #!/bin/bash
 
 # set path firstly
-lightx2v_path=/data/nvme1/wushuo/LightX2V
-model_path=/data/nvme1/wushuo/hf_models/Wan2.2-I2V-A14B
+lightx2v_path=/LightX2V
+model_path=/Wan2.2-I2V-A14B
 
 if [ -z "$lightx2v_path" ] || [ -z "$model_path" ]; then
     echo "Error: Please set lightx2v_path and model_path in this script before running."