From c1d1337afc52e17f6f724fb7294dc8997e9ea0d1 Mon Sep 17 00:00:00 2001
From: Ratish P <114130421+Ratish1@users.noreply.github.com>
Date: Mon, 16 Feb 2026 19:27:58 +0530
Subject: [PATCH] [diffusion][Wan]: fix sparse attention backends being applied
 to cross-attention (#17596)

---
 .../runtime/models/dits/causal_wanvideo.py          | 11 ++++++++++-
 .../multimodal_gen/runtime/models/dits/wanvideo.py  | 13 ++++++++-----
 .../multimodal_gen/runtime/platforms/interface.py   | 10 ++++++++++
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py
index d2e5c7ed4..159f06e5d 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py
@@ -301,7 +301,16 @@ class CausalWanTransformerBlock(nn.Module):
 
         # 2. Cross-attention
         # Only T2V for now
-        self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps)
+        cross_attn_backends = {
+            b for b in supported_attention_backends if not b.is_sparse
+        }
+        self.attn2 = WanT2VCrossAttention(
+            dim,
+            num_heads,
+            qk_norm=qk_norm,
+            eps=eps,
+            supported_attention_backends=cross_attn_backends,
+        )
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim, eps=eps, elementwise_affine=False, dtype=torch.float32
         )
diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
index f8836737b..24441a780 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
@@ -220,7 +220,6 @@ class WanI2VCrossAttention(WanSelfAttention):
         eps=1e-6,
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
     ) -> None:
-        # VSA should not be in supported_attention_backends
         super().__init__(
             dim,
             num_heads,
@@ -371,6 +370,9 @@ class WanTransformerBlock(nn.Module):
         )
 
         # 2. Cross-attention
+        cross_attn_backends = {
+            b for b in supported_attention_backends if not b.is_sparse
+        }
         if added_kv_proj_dim is not None:
             # I2V
             self.attn2 = WanI2VCrossAttention(
@@ -564,9 +566,10 @@ class WanTransformerBlock_VSA(nn.Module):
             dtype=torch.float32,
         )
 
-        if AttentionBackendEnum.VIDEO_SPARSE_ATTN in supported_attention_backends:
-            supported_attention_backends.remove(AttentionBackendEnum.VIDEO_SPARSE_ATTN)
         # 2. Cross-attention
+        cross_attn_backends = {
+            b for b in supported_attention_backends if not b.is_sparse
+        }
         if added_kv_proj_dim is not None:
             # I2V
             self.attn2 = WanI2VCrossAttention(
@@ -574,7 +577,7 @@ class WanTransformerBlock_VSA(nn.Module):
                 num_heads,
                 qk_norm=qk_norm,
                 eps=eps,
-                supported_attention_backends=supported_attention_backends,
+                supported_attention_backends=cross_attn_backends,
             )
         else:
             # T2V
@@ -583,7 +586,7 @@ class WanTransformerBlock_VSA(nn.Module):
                 num_heads,
                 qk_norm=qk_norm,
                 eps=eps,
-                supported_attention_backends=supported_attention_backends,
+                supported_attention_backends=cross_attn_backends,
             )
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py
index 1225640a7..1dcbb7962 100644
--- a/python/sglang/multimodal_gen/runtime/platforms/interface.py
+++ b/python/sglang/multimodal_gen/runtime/platforms/interface.py
@@ -41,6 +41,16 @@ class AttentionBackendEnum(enum.Enum):
     def __str__(self):
         return self.name.lower()
 
+    @property
+    def is_sparse(self) -> bool:
+        return self in {
+            AttentionBackendEnum.SLIDING_TILE_ATTN,
+            AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+            AttentionBackendEnum.VMOBA_ATTN,
+            AttentionBackendEnum.SLA_ATTN,
+            AttentionBackendEnum.SAGE_SLA_ATTN,
+        }
+
 
 class PlatformEnum(enum.Enum):
     CUDA = enum.auto()