[Bugfix] Raise ValueError when joint_strategy='rear' and causal=True in Ring Attention (#767)

Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Signed-off-by: mxuax <mxuax@connect.ust.hk>

vllm_omni/diffusion/attention/backends/ring_flash_attn.py

@@ -28,6 +28,20 @@ def ring_flash_attn_forward(
     joint_tensor_value=None,
     joint_strategy="front",
 ):
+    # Validate causal + joint_strategy combination
+    # When causal=True and joint_strategy="rear", the causal mask would incorrectly
+    # prevent local query tokens from attending to joint key tokens (which are
+    # concatenated at the end). This breaks the semantics where joint tokens
+    # (e.g., text conditioning) should be visible to all local tokens.
+    if causal and joint_tensor_key is not None and joint_strategy == "rear":
+        raise ValueError(
+            "joint_strategy='rear' is not compatible with causal=True in Ring Attention. "
+            "When using causal attention with joint tokens, use joint_strategy='front' "
+            "to ensure joint tokens act as a visible prefix for all local tokens. "
+            "With 'rear' strategy, the causal mask would incorrectly block local tokens "
+            "from seeing the joint tokens."
+        )
+
     comm = RingComm(process_group)
 
     out = None

vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py

@@ -62,6 +62,20 @@ class RingAttentionFunc(torch.autograd.Function):
         joint_tensor_value=None,
         joint_strategy="front",
     ):
+        # Validate causal + joint_strategy combination
+        # When causal=True and joint_strategy="rear", the causal mask would incorrectly
+        # prevent local query tokens from attending to joint key tokens (which are
+        # concatenated at the end). This breaks the semantics where joint tokens
+        # (e.g., text conditioning) should be visible to all local tokens.
+        if is_causal and joint_tensor_key is not None and joint_strategy == "rear":
+            raise ValueError(
+                "joint_strategy='rear' is not compatible with causal=True in Ring Attention. "
+                "When using causal attention with joint tokens, use joint_strategy='front' "
+                "to ensure joint tokens act as a visible prefix for all local tokens. "
+                "With 'rear' strategy, the causal mask would incorrectly block local tokens "
+                "from seeing the joint tokens."
+            )
+
         comm = RingComm(group)
         # Ensure tensors are contiguous for P2P communication
         q = q.contiguous()