From e3ec0812b518dc602132adb9803ee9402a35821b Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 05:32:14 +0400
Subject: [PATCH 01/13] chore: packaging fixes

---
 pyproject.toml | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index b642233..dc294a4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,7 +7,7 @@ name = "seedvr"
 version = "1.0.6"
 description = "SeedVR: Seeding Infinity in Diffusion Transformer Towards Generic Video Restoration"
 readme = "readme.md"
-requires-python = ">=3.9,<3.11"
+requires-python = ">=3.9,<3.12"
 license = {text = "Apache-2.0"}
 authors = [
     {name = "Jianyi Wang", email = "iceclear@bytedance.com"},
@@ -38,6 +38,7 @@ classifiers = [
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "Topic :: Scientific/Engineering :: Image Processing",
     "Topic :: Multimedia :: Video",
@@ -97,7 +98,10 @@ Documentation = "https://iceclear.github.io/projects/seedvr/"
 "ComfyUI Integration" = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler"
 
 [tool.setuptools]
-packages = ["seedvr"]
+include-package-data = true
+
+[tool.setuptools.packages.find]
+include = ["seedvr*"]
 
 [tool.setuptools.package-data]
 seedvr = [
@@ -147,7 +151,4 @@ addopts = [
     "--cov=seedvr",
     "--cov-report=term-missing",
     "--cov-report=html",
-]
-
-
-
+]
\ No newline at end of file

From aa0198c3c1c9075d249235c6fd565e8e08d19030 Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 05:40:56 +0400
Subject: [PATCH 02/13] perf: remove stream sync from add_noise

---
 seedvr/pipeline.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/seedvr/pipeline.py b/seedvr/pipeline.py
index 59ee187..29eca5a 100644
--- a/seedvr/pipeline.py
+++ b/seedvr/pipeline.py
@@ -375,7 +375,11 @@ def add_noise(
         Add noise to the input.
         """
         t = torch.tensor([self.sampler.timesteps.T], device=x.device) * cond_noise_scale
-        shape = torch.tensor(x.shape[1:], device=x.device).unsqueeze(0)
+        shape = (
+            torch.tensor(x.shape[1:], device="cpu", pin_memory=True)
+            .to(device=x.device, non_blocking=True)
+            .unsqueeze(0)
+        )
         t = self.timestep_transform(t, shape)
         x = self.sampler.schedule.forward(x, aug_noise, t)
         return x

From 8b46154984f481fa718ea6d53ceded001d59cdc7 Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 05:42:32 +0400
Subject: [PATCH 03/13] perf: remove stream sync from flatten()

---
 seedvr/models/dit/na.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/seedvr/models/dit/na.py b/seedvr/models/dit/na.py
index 5541dd5..12c64de 100644
--- a/seedvr/models/dit/na.py
+++ b/seedvr/models/dit/na.py
@@ -26,7 +26,15 @@ def flatten(
     torch.LongTensor,  # (b n)
 ]:
     assert len(hid) > 0
-    shape = torch.stack([torch.tensor(x.shape[:-1], device=hid[0].device) for x in hid])
+    shape = torch.stack(
+        [
+            torch.tensor(x.shape[:-1], device="cpu", pin_memory=True).to(
+                device=hid[0].device,
+                non_blocking=True,
+            )
+            for x in hid
+        ]
+    )
     hid = torch.cat([x.flatten(0, -2) for x in hid])
     return hid, shape
 

From 6af1dee2cb33fad2124afb34d7f0737589b82b9b Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 05:57:57 +0400
Subject: [PATCH 04/13] perf: remove more stream syncs

---
 seedvr/models/dit/na.py                  | 39 +++++++++++++-----------
 seedvr/models/dit/nablocks/mmsr_block.py | 20 +++++++++---
 seedvr/models/dit_v2/na.py               | 39 +++++++++++++++---------
 3 files changed, 62 insertions(+), 36 deletions(-)

diff --git a/seedvr/models/dit/na.py b/seedvr/models/dit/na.py
index 12c64de..db7feb6 100644
--- a/seedvr/models/dit/na.py
+++ b/seedvr/models/dit/na.py
@@ -28,13 +28,10 @@ def flatten(
     assert len(hid) > 0
     shape = torch.stack(
         [
-            torch.tensor(x.shape[:-1], device="cpu", pin_memory=True).to(
-                device=hid[0].device,
-                non_blocking=True,
-            )
+            torch.tensor(x.shape[:-1], device="cpu", pin_memory=True)
             for x in hid
         ]
-    )
+    ).pin_memory()
     hid = torch.cat([x.flatten(0, -2) for x in hid])
     return hid, shape
 
@@ -63,13 +60,17 @@ def concat(
 def concat_idx(
     vid_len: torch.LongTensor,  # (b)
     txt_len: torch.LongTensor,  # (b)
+    device: torch.device,
 ) -> tuple[
     Callable,
     Callable,
 ]:
-    device = vid_len.device
-    vid_idx = torch.arange(vid_len.sum(), device=device)
-    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    vid_idx = torch.arange(int(vid_len.sum()), device=device)
+    txt_idx = torch.arange(
+        len(vid_idx),
+        len(vid_idx) + int(txt_len.sum()),
+        device=device,
+    )
     tgt_idx = concat(vid_idx, txt_idx, vid_len, txt_len)
     src_idx = torch.argsort(tgt_idx)
     return (
@@ -113,13 +114,17 @@ def repeat_concat_idx(
     vid_len: torch.LongTensor,  # (n*b)
     txt_len: torch.LongTensor,  # (b)
     txt_repeat: torch.LongTensor,  # (n)
+    device: torch.device,
 ) -> tuple[
     Callable,
     Callable,
 ]:
-    device = vid_len.device
-    vid_idx = torch.arange(vid_len.sum(), device=device)
-    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    vid_idx = torch.arange(int(vid_len.sum()), device=device)
+    txt_idx = torch.arange(
+        len(vid_idx),
+        len(vid_idx) + int(txt_len.sum()),
+        device=device,
+    )
     txt_repeat_list = txt_repeat.tolist()
     tgt_idx = repeat_concat(vid_idx, txt_idx, vid_len, txt_len, txt_repeat)
     src_idx = torch.argsort(tgt_idx)
@@ -168,11 +173,10 @@ def rearrange(
 def rearrange_idx(
     hid_shape: torch.LongTensor,  # (b n)
     pattern: str,
+    device: torch.device,
     **kwargs: dict[str, int],
 ) -> tuple[Callable, Callable, torch.LongTensor]:
-    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(
-        -1
-    )
+    hid_idx = torch.arange(int(hid_shape.prod(-1).sum()), device=device).unsqueeze(-1)
     tgt_idx, tgt_shape = rearrange(hid_idx, hid_shape, pattern, **kwargs)
     tgt_idx = tgt_idx.squeeze(-1)
     src_idx = torch.argsort(tgt_idx)
@@ -235,7 +239,7 @@ def window(
 ):
     hid = unflatten(hid, hid_shape)
     hid = list(map(window_fn, hid))
-    hid_windows = torch.tensor(list(map(len, hid)), device=hid_shape.device)
+    hid_windows = torch.tensor(list(map(len, hid)), device="cpu", pin_memory=True)
     hid, hid_shape = flatten(list(chain(*hid)))
     return hid, hid_shape, hid_windows
 
@@ -243,10 +247,9 @@ def window(
 def window_idx(
     hid_shape: torch.LongTensor,  # (b n)
     window_fn: Callable[[torch.Tensor], list[torch.Tensor]],
+    device: torch.device,
 ):
-    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(
-        -1
-    )
+    hid_idx = torch.arange(int(hid_shape.prod(-1).sum()), device=device).unsqueeze(-1)
     tgt_idx, tgt_shape, tgt_windows = window(hid_idx, hid_shape, window_fn)
     tgt_idx = tgt_idx.squeeze(-1)
     src_idx = torch.argsort(tgt_idx)
diff --git a/seedvr/models/dit/nablocks/mmsr_block.py b/seedvr/models/dit/nablocks/mmsr_block.py
index 1bd6906..89f9c98 100644
--- a/seedvr/models/dit/nablocks/mmsr_block.py
+++ b/seedvr/models/dit/nablocks/mmsr_block.py
@@ -105,7 +105,7 @@ def make_window(x: torch.Tensor):
 
         window_partition, window_reverse, window_shape, window_count = cache_win(
             "win_transform",
-            lambda: na.window_idx(vid_shape, make_window),
+            lambda: na.window_idx(vid_shape, make_window, device=vid_qkv.device),
         )
         vid_qkv_win = window_partition(vid_qkv)
 
@@ -128,7 +128,13 @@ def make_window(x: torch.Tensor):
         )
         all_len_win = cache_win("all_len", lambda: vid_len_win + txt_len_win)
         concat_win, unconcat_win = cache_win(
-            "mm_pnp", lambda: na.repeat_concat_idx(vid_len_win, txt_len, window_count)
+            "mm_pnp",
+            lambda: na.repeat_concat_idx(
+                vid_len_win,
+                txt_len,
+                window_count,
+                device=vid_q.device,
+            ),
         )
 
         # window rope
@@ -141,11 +147,17 @@ def make_window(x: torch.Tensor):
             v=concat_win(vid_v, txt_v).bfloat16(),
             cu_seqlens_q=cache_win(
                 "vid_seqlens_q",
-                lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int(),
+                lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0))
+                .int()
+                .pin_memory()
+                .to(device=vid_q.device, non_blocking=True),
             ),
             cu_seqlens_k=cache_win(
                 "vid_seqlens_k",
-                lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int(),
+                lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0))
+                .int()
+                .pin_memory()
+                .to(device=vid_q.device, non_blocking=True),
             ),
             max_seqlen_q=cache_win(
                 "vid_max_seqlen_q", lambda: all_len_win.max().item()
diff --git a/seedvr/models/dit_v2/na.py b/seedvr/models/dit_v2/na.py
index 5541dd5..db7feb6 100644
--- a/seedvr/models/dit_v2/na.py
+++ b/seedvr/models/dit_v2/na.py
@@ -26,7 +26,12 @@ def flatten(
     torch.LongTensor,  # (b n)
 ]:
     assert len(hid) > 0
-    shape = torch.stack([torch.tensor(x.shape[:-1], device=hid[0].device) for x in hid])
+    shape = torch.stack(
+        [
+            torch.tensor(x.shape[:-1], device="cpu", pin_memory=True)
+            for x in hid
+        ]
+    ).pin_memory()
     hid = torch.cat([x.flatten(0, -2) for x in hid])
     return hid, shape
 
@@ -55,13 +60,17 @@ def concat(
 def concat_idx(
     vid_len: torch.LongTensor,  # (b)
     txt_len: torch.LongTensor,  # (b)
+    device: torch.device,
 ) -> tuple[
     Callable,
     Callable,
 ]:
-    device = vid_len.device
-    vid_idx = torch.arange(vid_len.sum(), device=device)
-    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    vid_idx = torch.arange(int(vid_len.sum()), device=device)
+    txt_idx = torch.arange(
+        len(vid_idx),
+        len(vid_idx) + int(txt_len.sum()),
+        device=device,
+    )
     tgt_idx = concat(vid_idx, txt_idx, vid_len, txt_len)
     src_idx = torch.argsort(tgt_idx)
     return (
@@ -105,13 +114,17 @@ def repeat_concat_idx(
     vid_len: torch.LongTensor,  # (n*b)
     txt_len: torch.LongTensor,  # (b)
     txt_repeat: torch.LongTensor,  # (n)
+    device: torch.device,
 ) -> tuple[
     Callable,
     Callable,
 ]:
-    device = vid_len.device
-    vid_idx = torch.arange(vid_len.sum(), device=device)
-    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    vid_idx = torch.arange(int(vid_len.sum()), device=device)
+    txt_idx = torch.arange(
+        len(vid_idx),
+        len(vid_idx) + int(txt_len.sum()),
+        device=device,
+    )
     txt_repeat_list = txt_repeat.tolist()
     tgt_idx = repeat_concat(vid_idx, txt_idx, vid_len, txt_len, txt_repeat)
     src_idx = torch.argsort(tgt_idx)
@@ -160,11 +173,10 @@ def rearrange(
 def rearrange_idx(
     hid_shape: torch.LongTensor,  # (b n)
     pattern: str,
+    device: torch.device,
     **kwargs: dict[str, int],
 ) -> tuple[Callable, Callable, torch.LongTensor]:
-    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(
-        -1
-    )
+    hid_idx = torch.arange(int(hid_shape.prod(-1).sum()), device=device).unsqueeze(-1)
     tgt_idx, tgt_shape = rearrange(hid_idx, hid_shape, pattern, **kwargs)
     tgt_idx = tgt_idx.squeeze(-1)
     src_idx = torch.argsort(tgt_idx)
@@ -227,7 +239,7 @@ def window(
 ):
     hid = unflatten(hid, hid_shape)
     hid = list(map(window_fn, hid))
-    hid_windows = torch.tensor(list(map(len, hid)), device=hid_shape.device)
+    hid_windows = torch.tensor(list(map(len, hid)), device="cpu", pin_memory=True)
     hid, hid_shape = flatten(list(chain(*hid)))
     return hid, hid_shape, hid_windows
 
@@ -235,10 +247,9 @@ def window(
 def window_idx(
     hid_shape: torch.LongTensor,  # (b n)
     window_fn: Callable[[torch.Tensor], list[torch.Tensor]],
+    device: torch.device,
 ):
-    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(
-        -1
-    )
+    hid_idx = torch.arange(int(hid_shape.prod(-1).sum()), device=device).unsqueeze(-1)
     tgt_idx, tgt_shape, tgt_windows = window(hid_idx, hid_shape, window_fn)
     tgt_idx = tgt_idx.squeeze(-1)
     src_idx = torch.argsort(tgt_idx)

From dbda53fda46073eb3ca645ce56dd9f19c86078a0 Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 05:58:19 +0400
Subject: [PATCH 05/13] perf: remove more stream syncs

---
 .../dit_v2/nablocks/attention/mmattn.py       | 48 +++++++++++++++----
 1 file changed, 39 insertions(+), 9 deletions(-)

diff --git a/seedvr/models/dit_v2/nablocks/attention/mmattn.py b/seedvr/models/dit_v2/nablocks/attention/mmattn.py
index 5709441..87733be 100644
--- a/seedvr/models/dit_v2/nablocks/attention/mmattn.py
+++ b/seedvr/models/dit_v2/nablocks/attention/mmattn.py
@@ -118,14 +118,29 @@ def forward(
         txt_len = cache("txt_len", lambda: txt_shape.prod(-1))
         all_len = cache("all_len", lambda: vid_len + txt_len)
 
-        concat, unconcat = cache("mm_pnp", lambda: na.concat_idx(vid_len, txt_len))
+        concat, unconcat = cache(
+            "mm_pnp",
+            lambda: na.concat_idx(vid_len, txt_len, device=vid_q.device),
+        )
 
         attn = self.attn(
             q=concat(vid_q, txt_q).bfloat16(),
             k=concat(vid_k, txt_k).bfloat16(),
             v=concat(vid_v, txt_v).bfloat16(),
-            cu_seqlens_q=cache("mm_seqlens", lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)).int()),
-            cu_seqlens_k=cache("mm_seqlens", lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)).int()),
+            cu_seqlens_q=cache(
+                "mm_seqlens",
+                lambda: safe_pad_operation(all_len.cumsum(0), (1, 0))
+                .int()
+                .pin_memory()
+                .to(device=vid_q.device, non_blocking=True),
+            ),
+            cu_seqlens_k=cache(
+                "mm_seqlens",
+                lambda: safe_pad_operation(all_len.cumsum(0), (1, 0))
+                .int()
+                .pin_memory()
+                .to(device=vid_q.device, non_blocking=True),
+            ),
             max_seqlen_q=cache("mm_maxlen", lambda: all_len.max().item()),
             max_seqlen_k=cache("mm_maxlen", lambda: all_len.max().item()),
         ).type_as(vid_q)
@@ -190,7 +205,7 @@ def make_window(x: torch.Tensor):
 
         window_partition, window_reverse, window_shape, window_count = cache_win(
             "win_transform",
-            lambda: na.window_idx(vid_shape, make_window),
+            lambda: na.window_idx(vid_shape, make_window, device=vid_qkv.device),
         )
         vid_qkv_win = window_partition(vid_qkv)
 
@@ -209,7 +224,13 @@ def make_window(x: torch.Tensor):
         txt_len_win = cache_win("txt_len", lambda: txt_len.repeat_interleave(window_count))
         all_len_win = cache_win("all_len", lambda: vid_len_win + txt_len_win)
         concat_win, unconcat_win = cache_win(
-            "mm_pnp", lambda: na.repeat_concat_idx(vid_len_win, txt_len, window_count)
+            "mm_pnp",
+            lambda: na.repeat_concat_idx(
+                vid_len_win,
+                txt_len,
+                window_count,
+                device=vid_q.device,
+            ),
         )
 
         # window rope
@@ -217,16 +238,17 @@ def make_window(x: torch.Tensor):
             if self.rope.mm:
                 # repeat text q and k for window mmrope
                 _, num_h, _ = txt_q.shape
+                window_count_list = window_count.tolist()
                 txt_q_repeat = rearrange(txt_q, "l h d -> l (h d)")
                 txt_q_repeat = na.unflatten(txt_q_repeat, txt_shape)
-                txt_q_repeat = [[x] * n for x, n in zip(txt_q_repeat, window_count)]
+                txt_q_repeat = [[x] * n for x, n in zip(txt_q_repeat, window_count_list)]
                 txt_q_repeat = list(chain(*txt_q_repeat))
                 txt_q_repeat, txt_shape_repeat = na.flatten(txt_q_repeat)
                 txt_q_repeat = rearrange(txt_q_repeat, "l (h d) -> l h d", h=num_h)
 
                 txt_k_repeat = rearrange(txt_k, "l h d -> l (h d)")
                 txt_k_repeat = na.unflatten(txt_k_repeat, txt_shape)
-                txt_k_repeat = [[x] * n for x, n in zip(txt_k_repeat, window_count)]
+                txt_k_repeat = [[x] * n for x, n in zip(txt_k_repeat, window_count_list)]
                 txt_k_repeat = list(chain(*txt_k_repeat))
                 txt_k_repeat, _ = na.flatten(txt_k_repeat)
                 txt_k_repeat = rearrange(txt_k_repeat, "l (h d) -> l h d", h=num_h)
@@ -242,10 +264,18 @@ def make_window(x: torch.Tensor):
             k=concat_win(vid_k, txt_k).bfloat16(),
             v=concat_win(vid_v, txt_v).bfloat16(),
             cu_seqlens_q=cache_win(
-                "vid_seqlens_q", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int()
+                "vid_seqlens_q",
+                lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0))
+                .int()
+                .pin_memory()
+                .to(device=vid_q.device, non_blocking=True),
             ),
             cu_seqlens_k=cache_win(
-                "vid_seqlens_k", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int()
+                "vid_seqlens_k",
+                lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0))
+                .int()
+                .pin_memory()
+                .to(device=vid_q.device, non_blocking=True),
             ),
             max_seqlen_q=cache_win("vid_max_seqlen_q", lambda: all_len_win.max().item()),
             max_seqlen_k=cache_win("vid_max_seqlen_k", lambda: all_len_win.max().item()),

From d8bec8a2c058717cf689c3cef8b4991ad607451c Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 06:39:10 +0400
Subject: [PATCH 06/13] perf: replace flash-attn with new torch varlen

---
 seedvr/models/dit/attention.py                | 44 +++++++++++--------
 seedvr/models/dit/nablocks/mmsr_block.py      |  4 +-
 seedvr/models/dit_v2/attention.py             | 29 ++++++++++--
 .../dit_v2/nablocks/attention/mmattn.py       |  4 +-
 4 files changed, 54 insertions(+), 27 deletions(-)

diff --git a/seedvr/models/dit/attention.py b/seedvr/models/dit/attention.py
index 3699453..94b82ce 100644
--- a/seedvr/models/dit/attention.py
+++ b/seedvr/models/dit/attention.py
@@ -14,21 +14,8 @@
 
 import torch
 import torch.nn.functional as F
-
-try:
-    from flash_attn_interface import flash_attn_varlen_func
-
-    print("Using FlashAttention3")
-except ImportError:
-    try:
-        from flash_attn import flash_attn_varlen_func
-
-        print("Using FlashAttention2")
-    except ImportError:
-        print("Using PyTorch Attention")
-        flash_attn_varlen_func = None
-
 from torch import nn
+from torch.nn.attention.varlen import varlen_attn
 
 
 class TorchAttention(nn.Module):
@@ -56,7 +43,7 @@ def forward(self, *args, **kwargs):
         )
 
 
-class FlashAttentionVarlen(nn.Module):
+class VarlenAttention(nn.Module):
     def tflops(self, args, kwargs, output) -> float:
         cu_seqlens_q = kwargs["cu_seqlens_q"]
         cu_seqlens_k = kwargs["cu_seqlens_k"]
@@ -66,7 +53,26 @@ def tflops(self, args, kwargs, output) -> float:
         return h * (4 * d * (seqlens_q * seqlens_k).sum())
 
     def forward(self, *args, **kwargs):
-        kwargs["deterministic"] = torch.are_deterministic_algorithms_enabled()
-        if flash_attn_varlen_func is None:
-            return TorchAttention()(*args, **kwargs)
-        return flash_attn_varlen_func(*args, **kwargs)
+        query = kwargs.pop("query", kwargs.pop("q", None))
+        key = kwargs.pop("key", kwargs.pop("k", None))
+        value = kwargs.pop("value", kwargs.pop("v", None))
+        if query is None and len(args) > 0:
+            query = args[0]
+        if key is None and len(args) > 1:
+            key = args[1]
+        if value is None and len(args) > 2:
+            value = args[2]
+        if query is None or key is None or value is None:
+            raise ValueError("query, key, value must be provided")
+
+        return varlen_attn(
+            query=query,
+            key=key,
+            value=value,
+            cu_seq_q=kwargs.pop("cu_seqlens_q"),
+            cu_seq_k=kwargs.pop("cu_seqlens_k"),
+            max_q=kwargs.pop("max_seqlen_q"),
+            max_k=kwargs.pop("max_seqlen_k"),
+            is_causal=kwargs.pop("is_causal", False),
+            return_aux=kwargs.pop("return_aux", None),
+        )
diff --git a/seedvr/models/dit/nablocks/mmsr_block.py b/seedvr/models/dit/nablocks/mmsr_block.py
index 89f9c98..06dcf14 100644
--- a/seedvr/models/dit/nablocks/mmsr_block.py
+++ b/seedvr/models/dit/nablocks/mmsr_block.py
@@ -25,7 +25,7 @@
 from seedvr.common.utils import get_torch_dtype, safe_pad_operation
 
 from .. import na
-from ..attention import FlashAttentionVarlen
+from ..attention import VarlenAttention
 from ..blocks.mmdit_window_block import MMWindowAttention, MMWindowTransformerBlock
 from ..mm import MMArg
 from ..modulation import ada_layer_type
@@ -65,7 +65,7 @@ def __init__(
             shared_qkv=shared_qkv,
         )
         self.rope = NaRotaryEmbedding3d(dim=head_dim // 2) if qk_rope else None
-        self.attn = FlashAttentionVarlen()
+        self.attn = VarlenAttention()
         self.window_op = get_window_op(window_method)
         self.rope.to(get_torch_dtype(dtype))
 
diff --git a/seedvr/models/dit_v2/attention.py b/seedvr/models/dit_v2/attention.py
index abe3c1a..78c4434 100644
--- a/seedvr/models/dit_v2/attention.py
+++ b/seedvr/models/dit_v2/attention.py
@@ -14,8 +14,8 @@
 
 import torch
 import torch.nn.functional as F
-from flash_attn import flash_attn_varlen_func
 from torch import nn
+from torch.nn.attention.varlen import varlen_attn
 
 
 class TorchAttention(nn.Module):
@@ -33,7 +33,7 @@ def forward(self, *args, **kwargs):
         return F.scaled_dot_product_attention(*args, **kwargs)
 
 
-class FlashAttentionVarlen(nn.Module):
+class VarlenAttention(nn.Module):
     def tflops(self, args, kwargs, output) -> float:
         cu_seqlens_q = kwargs["cu_seqlens_q"]
         cu_seqlens_k = kwargs["cu_seqlens_k"]
@@ -43,5 +43,26 @@ def tflops(self, args, kwargs, output) -> float:
         return h * (4 * d * (seqlens_q * seqlens_k).sum())
 
     def forward(self, *args, **kwargs):
-        kwargs["deterministic"] = torch.are_deterministic_algorithms_enabled()
-        return flash_attn_varlen_func(*args, **kwargs)
+        query = kwargs.pop("query", kwargs.pop("q", None))
+        key = kwargs.pop("key", kwargs.pop("k", None))
+        value = kwargs.pop("value", kwargs.pop("v", None))
+        if query is None and len(args) > 0:
+            query = args[0]
+        if key is None and len(args) > 1:
+            key = args[1]
+        if value is None and len(args) > 2:
+            value = args[2]
+        if query is None or key is None or value is None:
+            raise ValueError("query, key, value must be provided")
+
+        return varlen_attn(
+            query=query,
+            key=key,
+            value=value,
+            cu_seq_q=kwargs.pop("cu_seqlens_q"),
+            cu_seq_k=kwargs.pop("cu_seqlens_k"),
+            max_q=kwargs.pop("max_seqlen_q"),
+            max_k=kwargs.pop("max_seqlen_k"),
+            is_causal=kwargs.pop("is_causal", False),
+            return_aux=kwargs.pop("return_aux", None),
+        )
diff --git a/seedvr/models/dit_v2/nablocks/attention/mmattn.py b/seedvr/models/dit_v2/nablocks/attention/mmattn.py
index 87733be..d6bd9c3 100644
--- a/seedvr/models/dit_v2/nablocks/attention/mmattn.py
+++ b/seedvr/models/dit_v2/nablocks/attention/mmattn.py
@@ -23,7 +23,7 @@
 from seedvr.common.distributed.ops import gather_heads_scatter_seq, gather_seq_scatter_heads_qkv
 from seedvr.common.utils import safe_pad_operation
 from ... import na
-from ...attention import FlashAttentionVarlen
+from ...attention import VarlenAttention
 from ...mm import MMArg, MMModule
 from ...normalization import norm_layer_type
 from ...rope import get_na_rope
@@ -71,7 +71,7 @@ def __init__(
         )
 
         self.rope = get_na_rope(rope_type=rope_type, dim=rope_dim)
-        self.attn = FlashAttentionVarlen()
+        self.attn = VarlenAttention()
 
     def forward(
         self,

From 815123fa6e172a8d454dec7b5109980b3480eca6 Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 16:35:07 +0400
Subject: [PATCH 07/13] perf: remove more stream syncs

---
 seedvr/pipeline.py | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/seedvr/pipeline.py b/seedvr/pipeline.py
index 29eca5a..be5b046 100644
--- a/seedvr/pipeline.py
+++ b/seedvr/pipeline.py
@@ -144,9 +144,15 @@ def vae_encode(
             shift = getattr(self.vae.config, "shifting_factor", 0.0)
 
             if isinstance(scale, ListConfig):
-                scale = torch.tensor(scale, device=device, dtype=dtype)
+                scale = torch.tensor(scale, device="cpu", dtype=dtype, pin_memory=True).to(
+                    device=device,
+                    non_blocking=True,
+                )
             if isinstance(shift, ListConfig):
-                shift = torch.tensor(shift, device=device, dtype=dtype)
+                shift = torch.tensor(shift, device="cpu", dtype=dtype, pin_memory=True).to(
+                    device=device,
+                    non_blocking=True,
+                )
 
             # Group samples of the same shape to batches if enabled.
             if self.vae.grouping:
@@ -210,9 +216,15 @@ def vae_decode(
             shift = getattr(self.vae.config, "shifting_factor", 0.0)
 
             if isinstance(scale, ListConfig):
-                scale = torch.tensor(scale, device=device, dtype=dtype)
+                scale = torch.tensor(scale, device="cpu", dtype=dtype, pin_memory=True).to(
+                    device=device,
+                    non_blocking=True,
+                )
             if isinstance(shift, ListConfig):
-                shift = torch.tensor(shift, device=device, dtype=dtype)
+                shift = torch.tensor(shift, device="cpu", dtype=dtype, pin_memory=True).to(
+                    device=device,
+                    non_blocking=True,
+                )
 
             # Group latents of the same shape to batches if enabled.
             if self.vae.grouping:
@@ -374,7 +386,11 @@ def add_noise(
         """
         Add noise to the input.
         """
-        t = torch.tensor([self.sampler.timesteps.T], device=x.device) * cond_noise_scale
+        t = (
+            torch.tensor([self.sampler.timesteps.T], device="cpu", pin_memory=True)
+            .to(device=x.device, non_blocking=True)
+            * cond_noise_scale
+        )
         shape = (
             torch.tensor(x.shape[1:], device="cpu", pin_memory=True)
             .to(device=x.device, non_blocking=True)

From d9bf82038b88539f19e4b913cd9bc24fc8e08bdb Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 16:58:27 +0400
Subject: [PATCH 08/13] perf: remove more stream syncs

---
 seedvr/pipeline.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/seedvr/pipeline.py b/seedvr/pipeline.py
index be5b046..bc26d9b 100644
--- a/seedvr/pipeline.py
+++ b/seedvr/pipeline.py
@@ -523,7 +523,9 @@ def __call__(
             samples = [sample.unsqueeze(1) if sample.ndim == 3 else sample for sample in samples]
 
             batch_media = rearrange(batch_media, "c t h w -> t c h w").to(
-                self.wavelet_kernel.device, dtype=self.wavelet_kernel.dtype
+                self.wavelet_kernel.device,
+                dtype=self.wavelet_kernel.dtype,
+                non_blocking=True,
             )
             samples = [
                 rearrange(sample, "c t h w -> t c h w").to(

From e14bbfe6bbde64384fdcee0eca30e031deb195ab Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 17:00:29 +0400
Subject: [PATCH 09/13] perf: remove unnecessary eval call

---
 seedvr/pipeline.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/seedvr/pipeline.py b/seedvr/pipeline.py
index bc26d9b..1a592c1 100644
--- a/seedvr/pipeline.py
+++ b/seedvr/pipeline.py
@@ -283,10 +283,6 @@ def inference(
         latents, latents_shapes = flatten(noises)
         latents_cond, _ = flatten(conditions)
 
-        # Enter eval mode.
-        was_training = self.dit.training
-        self.dit.eval()
-
         # Sampling.
         latents = self.sampler.sample(
             x=latents,
@@ -310,9 +306,6 @@ def inference(
             ),
         )
 
-        # Exit eval mode.
-        self.dit.train(was_training)
-
         # Unflatten.
         latents = unflatten(latents, latents_shapes)
         latents = [latent.to(self.vae.dtype) for latent in latents]

From 746323693a72f5d42caf28083956d8ac8094a6f5 Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 18:05:17 +0400
Subject: [PATCH 10/13] perf: remove another stream sync

---
 seedvr/pipeline.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/seedvr/pipeline.py b/seedvr/pipeline.py
index 1a592c1..9a29159 100644
--- a/seedvr/pipeline.py
+++ b/seedvr/pipeline.py
@@ -522,7 +522,9 @@ def __call__(
             )
             samples = [
                 rearrange(sample, "c t h w -> t c h w").to(
-                    self.wavelet_kernel.device, dtype=self.wavelet_kernel.dtype
+                    self.wavelet_kernel.device,
+                    dtype=self.wavelet_kernel.dtype,
+                    non_blocking=True,
                 )
                 for sample in samples
             ]

From 2a3be693314400ddc07f3e961476a5cb95317a98 Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 18:18:39 +0400
Subject: [PATCH 11/13] perf: attempting to remove stream syncs from decode

---
 .../video_vae_v3/modules/attn_video_vae.py    | 60 +++++++++++++------
 1 file changed, 42 insertions(+), 18 deletions(-)

diff --git a/seedvr/models/video_vae_v3/modules/attn_video_vae.py b/seedvr/models/video_vae_v3/modules/attn_video_vae.py
index 8e4a5d8..5a18467 100644
--- a/seedvr/models/video_vae_v3/modules/attn_video_vae.py
+++ b/seedvr/models/video_vae_v3/modules/attn_video_vae.py
@@ -1446,13 +1446,9 @@ def tiled_decode(
         else:
             local_tasks = tasks
 
-        # Use GPU for distributed processing, CPU for non-distributed
-        if is_distributed:
-            data_device = device
-            computation_device = device
-        else:
-            data_device = torch.device("cpu")
-            computation_device = device
+        # Keep decode accumulation on GPU to avoid per-tile device-to-host copies.
+        data_device = device
+        computation_device = device
 
         out_T = T * 4 - 3
         weight = torch.zeros(
@@ -1483,8 +1479,8 @@ def tiled_decode(
             local_tasks, desc="Decoding", use_tqdm=use_tqdm and rank == 0
         ):
             hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_].to(computation_device)
-            hidden_states_batch = self._decode(hidden_states_batch, memory_state=memory_state).to(
-                data_device
+            hidden_states_batch = self._decode(
+                hidden_states_batch, memory_state=memory_state
             )
             hidden_states_batch.clamp_(-1, 1)
             hidden_states_batch = hidden_states_batch[:, :, :out_T]
@@ -1496,7 +1492,9 @@ def tiled_decode(
                     (size_h - stride_h) * self.spatial_downsample_factor,
                     (size_w - stride_w) * self.spatial_downsample_factor,
                 ),
-            ).to(dtype=hidden_states.dtype, device=data_device)
+                device=data_device,
+                dtype=hidden_states.dtype,
+            )
 
             target_h = h * self.spatial_downsample_factor
             target_w = w * self.spatial_downsample_factor
@@ -1531,13 +1529,13 @@ def tiled_decode(
                 # Distributed case, only return values on rank 0
                 values = values / weight
                 values = values.float().clamp_(-1, 1)
-                return values.to("cpu")
+                return values
             return None
         else:
             # Non-distributed case
             values = values / weight
             values = values.float().clamp_(-1, 1)
-            return values.to("cpu")
+            return values
 
     @torch.no_grad()
     def tiled_encode(
@@ -1699,7 +1697,13 @@ def tiled_encode(
             return values.to(device)
 
     def build_1d_mask(
-        self, length: int, left_bound: bool, right_bound: bool, border_width: int
+        self,
+        length: int,
+        left_bound: bool,
+        right_bound: bool,
+        border_width: int,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
     ) -> torch.Tensor:
         """
         Builds a 1D mask.
@@ -1710,13 +1714,17 @@ def build_1d_mask(
         :param border_width: border width
         :return: mask
         """
-        x = torch.ones((length,))
+        x = torch.ones((length,), device=device, dtype=dtype)
 
         if not left_bound:
-            x[:border_width] = (torch.arange(border_width) + 1) / border_width
+            x[:border_width] = (
+                torch.arange(border_width, device=device, dtype=dtype) + 1
+            ) / border_width
         if not right_bound:
             x[-border_width:] = torch.flip(
-                (torch.arange(border_width) + 1) / border_width, dims=(0,)
+                (torch.arange(border_width, device=device, dtype=dtype) + 1)
+                / border_width,
+                dims=(0,),
             )
 
         return x
@@ -1726,6 +1734,8 @@ def build_mask(
         data: torch.Tensor,
         is_bound: tuple[bool, bool, bool, bool],
         border_width: tuple[int, int],
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
     ) -> torch.Tensor:
         """
         :param data: data tensor [B, C, T, H, W]
@@ -1734,8 +1744,22 @@ def build_mask(
         :return: mask tensor [1, 1, 1, H, W]
         """
         _, _, _, H, W = data.shape
-        h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0])
-        w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1])
+        h = self.build_1d_mask(
+            H,
+            is_bound[0],
+            is_bound[1],
+            border_width[0],
+            device=device,
+            dtype=dtype,
+        )
+        w = self.build_1d_mask(
+            W,
+            is_bound[2],
+            is_bound[3],
+            border_width[1],
+            device=device,
+            dtype=dtype,
+        )
 
         h = repeat(h, "H -> H W", H=H, W=W)
         w = repeat(w, "W -> H W", H=H, W=W)

From bd7e00dd9dc76ad9e2a4bf4aca7163ebd6f0c396 Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 18:22:42 +0400
Subject: [PATCH 12/13] perf: attempting to remove stream syncs from encode

---
 .../video_vae_v3/modules/attn_video_vae.py    | 20 +++++++++----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/seedvr/models/video_vae_v3/modules/attn_video_vae.py b/seedvr/models/video_vae_v3/modules/attn_video_vae.py
index 5a18467..9de5056 100644
--- a/seedvr/models/video_vae_v3/modules/attn_video_vae.py
+++ b/seedvr/models/video_vae_v3/modules/attn_video_vae.py
@@ -1606,13 +1606,9 @@ def tiled_encode(
         else:
             local_tasks = tasks
 
-        # Use GPU for distributed processing, CPU for non-distributed
-        if is_distributed:
-            data_device = device
-            computation_device = device
-        else:
-            data_device = torch.device("cpu")
-            computation_device = device
+        # Keep encode accumulation on GPU to avoid per-tile device-to-host copies.
+        data_device = device
+        computation_device = device
 
         out_T = (T + 3) // 4
         weight = torch.zeros(
@@ -1644,8 +1640,8 @@ def tiled_encode(
             local_tasks, desc="Encoding", use_tqdm=use_tqdm and rank == 0
         ):
             hidden_states_batch = video[:, :, :, h:h_, w:w_].to(computation_device)
-            hidden_states_batch = self._encode(hidden_states_batch, memory_state=memory_state).to(
-                data_device
+            hidden_states_batch = self._encode(
+                hidden_states_batch, memory_state=memory_state
             )
 
             mask = self.build_mask(
@@ -1655,7 +1651,9 @@ def tiled_encode(
                     (size_h - stride_h) // self.spatial_downsample_factor,
                     (size_w - stride_w) // self.spatial_downsample_factor,
                 ),
-            ).to(dtype=video.dtype, device=data_device)
+                device=data_device,
+                dtype=video.dtype,
+            )
 
             target_h = h // self.spatial_downsample_factor
             target_w = w // self.spatial_downsample_factor
@@ -1694,7 +1692,7 @@ def tiled_encode(
         else:
             # Non-distributed case
             values = values / weight
-            return values.to(device)
+            return values
 
     def build_1d_mask(
         self,

From 28790eebf4e7ae084c5e1e936c33b110c09b70a4 Mon Sep 17 00:00:00 2001
From: Ismayil Ismayilov <39117301+readleyj@users.noreply.github.com>
Date: Sat, 7 Mar 2026 20:32:43 +0400
Subject: [PATCH 13/13] perf: rollback encoder / decoder changes

---
 .../video_vae_v3/modules/attn_video_vae.py    | 80 +++++++------------
 1 file changed, 29 insertions(+), 51 deletions(-)

diff --git a/seedvr/models/video_vae_v3/modules/attn_video_vae.py b/seedvr/models/video_vae_v3/modules/attn_video_vae.py
index 9de5056..8e4a5d8 100644
--- a/seedvr/models/video_vae_v3/modules/attn_video_vae.py
+++ b/seedvr/models/video_vae_v3/modules/attn_video_vae.py
@@ -1446,9 +1446,13 @@ def tiled_decode(
         else:
             local_tasks = tasks
 
-        # Keep decode accumulation on GPU to avoid per-tile device-to-host copies.
-        data_device = device
-        computation_device = device
+        # Use GPU for distributed processing, CPU for non-distributed
+        if is_distributed:
+            data_device = device
+            computation_device = device
+        else:
+            data_device = torch.device("cpu")
+            computation_device = device
 
         out_T = T * 4 - 3
         weight = torch.zeros(
@@ -1479,8 +1483,8 @@ def tiled_decode(
             local_tasks, desc="Decoding", use_tqdm=use_tqdm and rank == 0
         ):
             hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_].to(computation_device)
-            hidden_states_batch = self._decode(
-                hidden_states_batch, memory_state=memory_state
+            hidden_states_batch = self._decode(hidden_states_batch, memory_state=memory_state).to(
+                data_device
             )
             hidden_states_batch.clamp_(-1, 1)
             hidden_states_batch = hidden_states_batch[:, :, :out_T]
@@ -1492,9 +1496,7 @@ def tiled_decode(
                     (size_h - stride_h) * self.spatial_downsample_factor,
                     (size_w - stride_w) * self.spatial_downsample_factor,
                 ),
-                device=data_device,
-                dtype=hidden_states.dtype,
-            )
+            ).to(dtype=hidden_states.dtype, device=data_device)
 
             target_h = h * self.spatial_downsample_factor
             target_w = w * self.spatial_downsample_factor
@@ -1529,13 +1531,13 @@ def tiled_decode(
                 # Distributed case, only return values on rank 0
                 values = values / weight
                 values = values.float().clamp_(-1, 1)
-                return values
+                return values.to("cpu")
             return None
         else:
             # Non-distributed case
             values = values / weight
             values = values.float().clamp_(-1, 1)
-            return values
+            return values.to("cpu")
 
     @torch.no_grad()
     def tiled_encode(
@@ -1606,9 +1608,13 @@ def tiled_encode(
         else:
             local_tasks = tasks
 
-        # Keep encode accumulation on GPU to avoid per-tile device-to-host copies.
-        data_device = device
-        computation_device = device
+        # Use GPU for distributed processing, CPU for non-distributed
+        if is_distributed:
+            data_device = device
+            computation_device = device
+        else:
+            data_device = torch.device("cpu")
+            computation_device = device
 
         out_T = (T + 3) // 4
         weight = torch.zeros(
@@ -1640,8 +1646,8 @@ def tiled_encode(
             local_tasks, desc="Encoding", use_tqdm=use_tqdm and rank == 0
         ):
             hidden_states_batch = video[:, :, :, h:h_, w:w_].to(computation_device)
-            hidden_states_batch = self._encode(
-                hidden_states_batch, memory_state=memory_state
+            hidden_states_batch = self._encode(hidden_states_batch, memory_state=memory_state).to(
+                data_device
             )
 
             mask = self.build_mask(
@@ -1651,9 +1657,7 @@ def tiled_encode(
                     (size_h - stride_h) // self.spatial_downsample_factor,
                     (size_w - stride_w) // self.spatial_downsample_factor,
                 ),
-                device=data_device,
-                dtype=video.dtype,
-            )
+            ).to(dtype=video.dtype, device=data_device)
 
             target_h = h // self.spatial_downsample_factor
             target_w = w // self.spatial_downsample_factor
@@ -1692,16 +1696,10 @@ def tiled_encode(
         else:
             # Non-distributed case
             values = values / weight
-            return values
+            return values.to(device)
 
     def build_1d_mask(
-        self,
-        length: int,
-        left_bound: bool,
-        right_bound: bool,
-        border_width: int,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
+        self, length: int, left_bound: bool, right_bound: bool, border_width: int
     ) -> torch.Tensor:
         """
         Builds a 1D mask.
@@ -1712,17 +1710,13 @@ def build_1d_mask(
         :param border_width: border width
         :return: mask
         """
-        x = torch.ones((length,), device=device, dtype=dtype)
+        x = torch.ones((length,))
 
         if not left_bound:
-            x[:border_width] = (
-                torch.arange(border_width, device=device, dtype=dtype) + 1
-            ) / border_width
+            x[:border_width] = (torch.arange(border_width) + 1) / border_width
         if not right_bound:
             x[-border_width:] = torch.flip(
-                (torch.arange(border_width, device=device, dtype=dtype) + 1)
-                / border_width,
-                dims=(0,),
+                (torch.arange(border_width) + 1) / border_width, dims=(0,)
             )
 
         return x
@@ -1732,8 +1726,6 @@ def build_mask(
         data: torch.Tensor,
         is_bound: tuple[bool, bool, bool, bool],
         border_width: tuple[int, int],
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
     ) -> torch.Tensor:
         """
         :param data: data tensor [B, C, T, H, W]
@@ -1742,22 +1734,8 @@ def build_mask(
         :return: mask tensor [1, 1, 1, H, W]
         """
         _, _, _, H, W = data.shape
-        h = self.build_1d_mask(
-            H,
-            is_bound[0],
-            is_bound[1],
-            border_width[0],
-            device=device,
-            dtype=dtype,
-        )
-        w = self.build_1d_mask(
-            W,
-            is_bound[2],
-            is_bound[3],
-            border_width[1],
-            device=device,
-            dtype=dtype,
-        )
+        h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0])
+        w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1])
 
         h = repeat(h, "H -> H W", H=H, W=W)
         w = repeat(w, "W -> H W", H=H, W=W)