diff --git a/tests/diffusion/models/bagel/test_sp_replicated_phases.py b/tests/diffusion/models/bagel/test_sp_replicated_phases.py new file mode 100644 index 000000000..b42ea93c4 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 47ffda310..09449f4fa 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/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),