anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[BugFix][NPU] Wan2.2-S2V RoPE: replace complex64 advanced indexing with index_select (#6320)

Browse Source 
Signed-off-by: Weiming Liao <liaowm5@gmail.com>
This commit is contained in:
Weiming Liao authored and GitHub committed 2026-08-19 00:06:29 +08:00
1 parent 284e05c88b
commit f208f907bc
1 file changed
+10 -6
Download Patch File Download Diff File Expand all files Collapse all files
vllm_omni/diffusion/layers/rope.py
+10 -6
View File
@@ -423,16 +423,16 @@ class WanS2VRotaryPosEmbed(torch.nn.Module):
h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long()
w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long()
freqs_0 = freqs_split[0][f_sam] if f_o >= 0 else freqs_split[0][f_sam].conj()
freqs_0 = torch.index_select(freqs_split[0] if f_o >= 0 else freqs_split[0].conj(), 0, f_sam)
freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1)
freqs_i = torch.cat(
[
freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1),
freqs_split[1][h_sam]
torch.index_select(freqs_split[1], 0, h_sam)
.view(1, seq_h_int, 1, -1)
.expand(seq_f_int, seq_h_int, seq_w_int, -1),
freqs_split[2][w_sam]
torch.index_select(freqs_split[2], 0, w_sam)
.view(1, 1, seq_w_int, -1)
.expand(seq_f_int, seq_h_int, seq_w_int, -1),
],
@@ -523,14 +523,18 @@ class RotaryEmbeddingS2VGrid(torch.nn.Module):
h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long()
w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long()
freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj()
freqs_0 = torch.index_select(freqs[0] if f_o >= 0 else freqs[0].conj(), 0, f_sam)
freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1)
freqs_i = torch.cat(
[
freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1),
freqs[1][h_sam].view(1, seq_h_int, 1, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1),
freqs[2][w_sam].view(1, 1, seq_w_int, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1),
torch.index_select(freqs[1], 0, h_sam)
.view(1, seq_h_int, 1, -1)
.expand(seq_f_int, seq_h_int, seq_w_int, -1),
torch.index_select(freqs[2], 0, w_sam)
.view(1, 1, seq_w_int, -1)
.expand(seq_f_int, seq_h_int, seq_w_int, -1),
],
dim=-1,
).reshape(seg_len, 1, -1)