[Bugfix][BAGEL] Keep replicated prefill and cache-update phases off the SP strategy (#8177)

Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>

tests/diffusion/models/bagel/test_sp_replicated_phases.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""BAGEL must keep its replicated phases off the sequence-parallel strategy.
+
+The text prefill, the ViT/VAE cache updates and the CFG cache-update forward
+run on sequences every SP rank holds in full; only the sharded denoise path
+may go through the strategy (Ulysses would otherwise all-to-all the identical
+per-rank copies as if they were shards).
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.bagel.bagel_transformer as bagel_mod
+from vllm_omni.diffusion.models.bagel.bagel_transformer import NaiveCache, PackedAttentionMoT
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+NUM_HEADS, NUM_KV_HEADS, HEAD_DIM = 4, 2, 8
+HIDDEN = NUM_HEADS * HEAD_DIM
+SEQ = 6  # [start_of_image, 4 vae tokens, end_of_image]
+
+
+class _RecordingAttention:
+    """Stands in for a DiffusionAttention layer: records calls, returns the query."""
+
+    def __init__(self, joint: bool = False) -> None:
+        self.joint = joint
+        self.calls: list = []
+
+    def __call__(self, q, k, v, attn_metadata=None):
+        self.calls.append(attn_metadata)
+        if self.joint and attn_metadata is not None and attn_metadata.joint_query is not None:
+            return torch.cat([attn_metadata.joint_query, q], dim=1)
+        return q
+
+
+def test_constructor_builds_local_layers_for_replicated_phases(monkeypatch: pytest.MonkeyPatch) -> None:
+    class _Attn(torch.nn.Module):
+        def __init__(self, **kwargs) -> None:
+            super().__init__()
+            self.kwargs = kwargs
+
+    monkeypatch.setattr(bagel_mod, "DiffusionAttention", _Attn)
+    for name in ("MoTQKVParallelLinear", "MoTRowParallelLinear", "MoTRMSNorm", "RotaryEmbedding"):
+        monkeypatch.setattr(bagel_mod, name, lambda *a, **k: torch.nn.Identity())
+    monkeypatch.setattr(bagel_mod, "get_tensor_model_parallel_world_size", lambda: 1)
+    cfg = SimpleNamespace(
+        hidden_size=HIDDEN, num_attention_heads=NUM_HEADS, num_key_value_heads=NUM_KV_HEADS, rms_norm_eps=1e-6
+    )
+
+    attn = PackedAttentionMoT(cfg, layer_idx=0)
+
+    assert attn.attn_causal.kwargs["causal"] is True
+    assert attn.attn_causal.kwargs["skip_sequence_parallel"] is True
+    assert attn.attn_noncausal.kwargs["causal"] is False
+    assert attn.attn_noncausal.kwargs.get("skip_sequence_parallel", False) is False
+    assert attn.attn_noncausal_local.kwargs["causal"] is False
+    assert attn.attn_noncausal_local.kwargs["skip_sequence_parallel"] is True
+
+
+def _stub(sp_active: bool) -> PackedAttentionMoT:
+    attn = PackedAttentionMoT.__new__(PackedAttentionMoT)
+    torch.nn.Module.__init__(attn)
+    attn.layer_idx = 0
+    attn.hidden_size = HIDDEN
+    attn.num_heads, attn.num_kv_heads, attn.head_dim = NUM_HEADS, NUM_KV_HEADS, HEAD_DIM
+    attn.q_size, attn.kv_size = NUM_HEADS * HEAD_DIM, NUM_KV_HEADS * HEAD_DIM
+    width = attn.q_size + 2 * attn.kv_size
+    attn.qkv_proj = lambda x, *idx: (torch.zeros(x.shape[0], width), None)
+    attn.q_norm = lambda x, *idx: x
+    attn.k_norm = lambda x, *idx: x
+    attn.rotary_op = lambda x, cos, sin: x
+    attn.o_proj = lambda x, *idx: (x[:, :HIDDEN], None)
+    attn.attn_causal = _RecordingAttention()
+    attn.attn_noncausal = _RecordingAttention(joint=True)
+    attn.attn_noncausal_local = _RecordingAttention()
+    attn._is_sp_active = lambda: sp_active
+    return attn
+
+
+def _gen_inputs():
+    packed = torch.zeros(SEQ, HIDDEN)
+    pos = (torch.zeros(SEQ, HEAD_DIM), torch.zeros(SEQ, HEAD_DIM))
+    return dict(
+        packed_query_sequence=packed,
+        query_lens=torch.tensor([SEQ]),
+        packed_query_position_embeddings=pos,
+        past_key_values=NaiveCache(1),
+        packed_vae_token_indexes=torch.arange(1, SEQ - 1),
+        packed_text_indexes=torch.tensor([0, SEQ - 1]),
+    )
+
+
+def test_gen_cache_update_under_sp_runs_locally() -> None:
+    attn = _stub(sp_active=True)
+    inputs = _gen_inputs()
+
+    out, cache = attn._forward_gen(**inputs, update_past_key_values=True)
+
+    assert attn.attn_noncausal.calls == []
+    assert len(attn.attn_noncausal_local.calls) == 1
+    assert out.shape == (SEQ, HIDDEN)
+    assert cache.key_cache[0].shape == (SEQ, NUM_KV_HEADS, HEAD_DIM)
+
+
+def test_gen_denoise_under_sp_still_uses_the_strategy() -> None:
+    attn = _stub(sp_active=True)
+
+    out, _ = attn._forward_gen(**_gen_inputs(), update_past_key_values=False)
+
+    assert attn.attn_noncausal_local.calls == []
+    (metadata,) = attn.attn_noncausal.calls
+    assert metadata.joint_strategy == "front"
+    assert metadata.joint_query.shape == (1, 2, NUM_HEADS, HEAD_DIM)
+    assert metadata.joint_key.shape == (1, 2, NUM_KV_HEADS, HEAD_DIM)
+    assert out.shape == (SEQ, HIDDEN)
+
+
+def test_gen_without_sp_runs_locally() -> None:
+    attn = _stub(sp_active=False)
+
+    attn._forward_gen(**_gen_inputs(), update_past_key_values=False)
+
+    assert attn.attn_noncausal.calls == []
+    assert len(attn.attn_noncausal_local.calls) == 1
+
+
+@pytest.mark.parametrize("is_causal", [True, False])
+def test_und_prefill_never_reaches_the_strategy(is_causal: bool) -> None:
+    attn = _stub(sp_active=True)
+    pos = (torch.zeros(SEQ, HEAD_DIM), torch.zeros(SEQ, HEAD_DIM))
+
+    out, cache = attn._forward_und(torch.zeros(SEQ, HIDDEN), pos, NaiveCache(1), is_causal=is_causal)
+
+    assert attn.attn_noncausal.calls == []
+    used = attn.attn_causal if is_causal else attn.attn_noncausal_local
+    assert len(used.calls) == 1
+    assert out.shape == (SEQ, HIDDEN)
+    assert cache.key_cache[0].shape == (SEQ, NUM_KV_HEADS, HEAD_DIM)

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -490,6 +490,7 @@ class PackedAttentionMoT(nn.Module):
             softmax_scale=1.0 / (self.head_dim**0.5),
             causal=True,
             num_kv_heads=self.total_num_kv_heads,
+            skip_sequence_parallel=True,
         )
         self.attn_noncausal = DiffusionAttention(
             num_heads=self.total_num_heads,
@@ -498,6 +499,14 @@ class PackedAttentionMoT(nn.Module):
             causal=False,
             num_kv_heads=self.total_num_kv_heads,
         )
+        self.attn_noncausal_local = DiffusionAttention(
+            num_heads=self.total_num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            num_kv_heads=self.total_num_kv_heads,
+            skip_sequence_parallel=True,
+        )
 
     def _is_sp_active(self) -> bool:
         """Check if SP is active for this attention layer."""
@@ -583,7 +592,7 @@ class PackedAttentionMoT(nn.Module):
 
         # NOTE: we reshape to batched (1, S, H, D) for diffusion Attention
         # attn_out should be: (1, text_len + local_vae_len, H, D)
-        if self._is_sp_active():
+        if self._is_sp_active() and not update_past_key_values:
             # Joint mechanism keeps text+cache replicated across SP ranks
             attn_out = self.attn_noncausal(
                 vae_q.unsqueeze(0),
@@ -631,7 +640,7 @@ class PackedAttentionMoT(nn.Module):
                 k_4d = torch.stack([torch.cat([t, v]) for t, v in zip(text_k_parts, vae_k_parts)])
                 v_4d = torch.stack([torch.cat([t, v]) for t, v in zip(text_v_parts, vae_v_parts)])
                 metadata = None
-            attn_out = self.attn_noncausal(q_4d, k_4d, v_4d, metadata)
+            attn_out = self.attn_noncausal_local(q_4d, k_4d, v_4d, metadata)
 
         attn_out = attn_out.reshape(num_branches, -1, self.q_size)
         text_attn = attn_out[:, :text_per_branch].reshape(-1, self.q_size)
@@ -747,7 +756,7 @@ class PackedAttentionMoT(nn.Module):
             )
             attn_out = attn_out_4d.permute(0, 2, 1, 3)
         else:
-            attn = self.attn_causal if is_causal else self.attn_noncausal
+            attn = self.attn_causal if is_causal else self.attn_noncausal_local
             attn_out = attn(
                 q.unsqueeze(0),
                 full_k.unsqueeze(0),