From 8e20b33b389f0cfb1121974a079b1e25de29412a Mon Sep 17 00:00:00 2001 From: XU Mingshi <91017482+mxuax@users.noreply.github.com> Date: Tue, 13 Jan 2026 17:13:31 +0800 Subject: [PATCH] [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 --- .../attention/backends/ring_flash_attn.py | 14 ++++++++++++++ .../attention/backends/ring_pytorch_attn.py | 14 ++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/vllm_omni/diffusion/attention/backends/ring_flash_attn.py b/vllm_omni/diffusion/attention/backends/ring_flash_attn.py index e163ef1ac..dd27f88a8 100644 --- a/vllm_omni/diffusion/attention/backends/ring_flash_attn.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py b/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py index 43ee35f70..9ed2c2076 100644 --- a/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py +++ b/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()