145 lines
5.5 KiB
Python
145 lines
5.5 KiB
Python
# 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)
|