Fix LingBot packed CFG attention fallback

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

tests/diffusion/models/lingbot_video/test_lingbot_video_transformer.py

@@ -58,6 +58,43 @@ def test_tiny_transformer_depth_zero_forward_shape():
     assert torch.isfinite(out).all()
 
 
+def test_packed_attention_uses_sdpa_fallback_without_flash_varlen(monkeypatch):
+    from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
+
+    monkeypatch.setattr(module, "flash_attn_varlen_func_v3", None)
+    attn = module.LingBotVideoAttention(
+        hidden_size=8,
+        num_heads=2,
+        norm_eps=1e-6,
+        qkv_bias=False,
+        out_bias=False,
+    )
+    captured = {}
+
+    def fake_sdpa_forward(query, key, value, attn_metadata):
+        captured["mask"] = attn_metadata.attn_mask
+        return torch.zeros_like(query)
+
+    monkeypatch.setattr(attn.attn.sdpa_fallback, "forward", fake_sdpa_forward)
+    x = torch.randn(1, 5, 8)
+    rotary = torch.ones(1, 5, 2, dtype=torch.complex64)
+    packed_indices = {
+        "cu_seqlens_kv": torch.tensor([0, 2, 5], dtype=torch.int32),
+        "max_seqlen_in_batch_kv": 3,
+        "attention_mask": module._packed_block_attention_mask([2, 3], x.device),
+    }
+
+    out = attn(x, rotary, packed_indices=packed_indices)
+
+    assert out.shape == x.shape
+    mask = captured["mask"]
+    assert mask.shape == (1, 1, 5, 5)
+    assert mask[0, 0, :2, :2].all()
+    assert mask[0, 0, 2:, 2:].all()
+    assert not mask[0, 0, :2, 2:].any()
+    assert not mask[0, 0, 2:, :2].any()
+
+
 def test_tiny_transformer_rejects_invalid_rope_dims():
     from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel
 

vllm_omni/diffusion/models/lingbot_video/lingbot_video_transformer.py

@@ -15,13 +15,11 @@ from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.modeling_utils import ModelMixin
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.utils.fa import (
+    flash_attn_varlen_func as flash_attn_varlen_func_v3,
+)
 from vllm_omni.diffusion.attention.layer import Attention
 
-try:
-    from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_v3
-except Exception:  # pragma: no cover - optional CUDA kernel.
-    flash_attn_varlen_func_v3 = None
-
 LINGBOT_VIDEO_FP32_MODULES = (
     "time_embedder",
     "time_modulation",
@@ -148,6 +146,24 @@ def _cat_interleave(
     return torch.cat(blocks, dim=1)
 
 
+def _packed_block_attention_mask(sample_seq_lens: list[int], device: torch.device) -> torch.Tensor:
+    total_seq_len = sum(sample_seq_lens)
+    mask = torch.zeros(
+        1,
+        1,
+        total_seq_len,
+        total_seq_len,
+        dtype=torch.bool,
+        device=device,
+    )
+    start = 0
+    for seq_len in sample_seq_lens:
+        end = start + seq_len
+        mask[:, :, start:end, start:end] = True
+        start = end
+    return mask
+
+
 class LingBotVideoTextEmbedder(nn.Module):
     """Matches CondProjection: RMSNorm(text_dim, eps=1e-6 fixed) -> Linear-SiLU-Linear."""
 
@@ -212,8 +228,19 @@ class LingBotVideoAttention(nn.Module):
         if packed_indices is None:
             out = self.attn(q, k, v, attn_metadata=AttentionMetadata(attn_mask=attention_mask))
         else:
+            packed_attention_mask = packed_indices.get("attention_mask")
+            if packed_attention_mask is not None and parallel_config is None:
+                out = self.attn.sdpa_fallback.forward(
+                    q,
+                    k,
+                    v,
+                    AttentionMetadata(attn_mask=packed_attention_mask),
+                )
+                return self.to_out(out.flatten(2, 3).type_as(x))
             if flash_attn_varlen_func_v3 is None:
-                raise RuntimeError("flash_attn_interface.flash_attn_varlen_func is required.")
+                raise RuntimeError(
+                    "A flash attention varlen function is required for packed context parallel attention."
+                )
             if parallel_config is None:
                 result = flash_attn_varlen_func_v3(
                     q=q.reshape(-1, self.num_heads, self.head_dim),
@@ -550,6 +577,8 @@ class LingBotVideoTransformer3DModel(ModelMixin, ConfigMixin):
                 "cu_seqlens_kv": cu_seqlens,
                 "max_seqlen_in_batch_kv": max(sample_seq_lens),
             }
+            if packed_batch and not use_packed_attention:
+                packed_indices["attention_mask"] = _packed_block_attention_mask(sample_seq_lens, device)
             has_padding = False
         if has_padding:
             key_mask = torch.cat(