Reiew pr #1159 and refactor codes by Watebear · Pull Request #1231 · ModelTC/LightX2V

Watebear · 2026-07-07T10:54:11Z

No description provided.

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

…s on Ascend NPU.

Resolve attention.py import conflict by keeping pr-1159's separate flash_attn v2 import and NPU fallback chain. Co-authored-by: Cursor <cursoragent@cursor.com>

gemini-code-assist

Code Review

This pull request introduces support for the RainFusion sparse attention acceleration algorithm on Ascend NPU, along with NPU-optimized implementations of LayerNorm, RMSNorm, RoPE, and Flash Attention. It also adds NPU support for RIFE video frame interpolation and provides configuration and execution scripts for the Wan2.2 MoE I2V Distill pipeline. The review feedback highlights critical issues where torch.npu.is_available() is called without importing torch_npu first, which will lead to runtime errors. Additionally, a layout mismatch was identified in the NPU RainFusion dense attention operator, and a cleaner tensor flattening approach was suggested in the RainFusion attention implementation.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-07T10:56:51Z

+import importlib.util
+
 import torch

-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+_HAS_NPU = importlib.util.find_spec("torch_npu") is not None and torch.npu.is_available()
+device = torch.device("cuda") if torch.cuda.is_available() else (torch.device("npu") if _HAS_NPU else torch.device("cpu"))


Using importlib.util.find_spec("torch_npu") to check for NPU availability without actually importing torch_npu will result in an AttributeError: module 'torch' has no attribute 'npu' when torch.npu.is_available() is called. In PyTorch, the npu attribute is dynamically registered on the torch module only after torch_npu is imported. Please import torch_npu inside a try-except block to safely check NPU availability.

Suggested change

import importlib.util

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

_HAS_NPU = importlib.util.find_spec("torch_npu") is not None and torch.npu.is_available()

device = torch.device("cuda") if torch.cuda.is_available() else (torch.device("npu") if _HAS_NPU else torch.device("cpu"))

import torch

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

gemini-code-assist · 2026-07-07T10:56:52Z

+            _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"))


Using importlib.util.find_spec("torch_npu") to check for NPU availability without actually importing torch_npu will result in an AttributeError: module 'torch' has no attribute 'npu' when torch.npu.is_available() is called. In PyTorch, the npu attribute is dynamically registered on the torch module only after torch_npu is imported. Please import torch_npu inside a try-except block to safely check NPU availability.

Suggested change

_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"))

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

gemini-code-assist · 2026-07-07T10:56:52Z

+_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"))


Using importlib.util.find_spec("torch_npu") to check for NPU availability without actually importing torch_npu will result in an AttributeError: module 'torch' has no attribute 'npu' when torch.npu.is_available() is called. In PyTorch, the npu attribute is dynamically registered on the torch module only after torch_npu is imported. Please import torch_npu inside a try-except block to safely check NPU availability.

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

gemini-code-assist · 2026-07-07T10:56:52Z

+    def dense_attention(self, q, k, v):
+        return attention_forward(q, k, v, opt_mode="manual", op_type="ascend_laser_attention", layout="BNSD")


In dense_attention, layout="BNSD" is passed to attention_forward, but the input tensors q, k, v are in BSND layout (shape (batch, seqlen, num_heads, head_dim)). This layout mismatch will cause incorrect results or crashes. Since attention_forward supports BSND layout, please change layout="BNSD" to layout="BSND".

Suggested change

def dense_attention(self, q, k, v):

return attention_forward(q, k, v, opt_mode="manual", op_type="ascend_laser_attention", layout="BNSD")

def dense_attention(self, q, k, v):

return attention_forward(q, k, v, opt_mode="manual", op_type="ascend_laser_attention", layout="BSND")

gemini-code-assist · 2026-07-07T10:56:52Z

+            )
+            out = self.do_tensor_inv_rearrange(out, h_res_len, w_res_len)
+
+        out = out.squeeze(0).reshape(out.shape[1], -1)


Using out.shape[1] of the pre-squeezed tensor in the same expression where out is reassigned is slightly confusing and prone to bugs if the tensor shape changes. It is much cleaner and more idiomatic to use out = out.squeeze(0).flatten(1).

Suggested change

out = out.squeeze(0).reshape(out.shape[1], -1)

out = out.squeeze(0).flatten(1)

liaopingbo and others added 13 commits June 17, 2026 11:22

feat(npu): Add 4 Ascend NPU custom operators for sparse video inference

74146ff

feat(npu): Add 4 Ascend NPU custom operators for sparse video inference

d49ef69

feat(npu): Add 4 Ascend NPU custom operators for sparse video inference

4fab6f6

feat(npu): Add 4 Ascend NPU custom operators for sparse video inference

c07b320

feat(npu): Add 4 Ascend NPU custom operators for sparse video inference

8f0d06a

Update lightx2v_platform/ops/attn/ascend_npu/rainfusion_blockwise.py

0a65ead

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

feat(npu): Add 4 Ascend NPU custom operators for sparse video inference

349c350

feat(npu): Perform inference with sparse FA and NPU affinity operator…

b3555b8

…s on Ascend NPU.

feat(npu): Perform inference with sparse FA and NPU affinity operator…

0846547

…s on Ascend NPU.

feat(npu): Perform inference with sparse FA and NPU affinity operator…

5ed56f4

…s on Ascend NPU.

feat(npu): pre-commit fix

df8577a

Merge branch 'pr-1159' into main

9d51e1d

Resolve attention.py import conflict by keeping pr-1159's separate flash_attn v2 import and NPU fallback chain. Co-authored-by: Cursor <cursoragent@cursor.com>

update codes

160d7c3

gemini-code-assist Bot reviewed Jul 7, 2026

View reviewed changes

Watebear added 3 commits July 8, 2026 08:27

temp

ea74ce6

support nv backend rainfusion

611eae3

rm personal path

8a5573b

helloyongyang approved these changes Jul 9, 2026

View reviewed changes

helloyongyang merged commit 6f436da into main Jul 9, 2026
2 checks passed

helloyongyang deleted the dev_pr1159 branch July 9, 2026 03:21

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Reiew pr #1159 and refactor codes#1231

Reiew pr #1159 and refactor codes#1231
helloyongyang merged 16 commits into
mainfrom
dev_pr1159

Watebear commented Jul 7, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		_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"))

-            _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"))
+            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"))

		_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"))

		def dense_attention(self, q, k, v):
		return attention_forward(q, k, v, opt_mode="manual", op_type="ascend_laser_attention", layout="BNSD")

	out = out.squeeze(0).reshape(out.shape[1], -1)
	out = out.squeeze(0).flatten(1)

Uh oh!

Conversation

Watebear commented Jul 7, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants