[Model] Add SANA-WM support(Stage-1 only) (#4061)

Signed-off-by: BruceLoveDecimal <liuqihao970610@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Signed-off-by: liuqihao <liuqihao970610@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for SANA-WM's camera-branch reduce-scatter.
+
+``out_proj_cam`` is row-parallel, so every rank produces a partial sum over the
+full hidden width, while the GDN stream it is added to is TP-local. Instead of
+all-reducing and then discarding all but one slice, the branch reduce-scatters
+straight to the local shard. These tests pin the shard each rank ends up with,
+including the bias slice, on both code paths:
+
+* the process-group path, which drives ``torch.distributed.reduce_scatter_tensor``
+  directly because vLLM-Omni's diffusion ``GroupCoordinator`` does not implement
+  ``reduce_scatter``;
+* the all-reduce-and-slice fallback taken when the coordinator exposes no
+  ``device_group``.
+
+CPU-only: the collectives and the rank/world-size lookups are mocked, so a
+single process can simulate an arbitrary TP degree.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmSelfAttention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
+
+HIDDEN = 24
+
+
+def _host(bias: torch.Tensor | None) -> SimpleNamespace:
+    """A stand-in ``self`` exposing only what the method reads."""
+    return SimpleNamespace(out_proj_cam=SimpleNamespace(bias=bias))
+
+
+def _fake_reduce_scatter_tensor(rank: int, tp_size: int):
+    """Model ``reduce_scatter_tensor``: split dim 0 and keep chunk ``rank``.
+
+    The test drives a single process, so the other ranks' partials are already
+    folded into the input; scattering is all this has to reproduce.
+    """
+
+    def _impl(output: torch.Tensor, input_: torch.Tensor, group=None) -> None:
+        output.copy_(input_.chunk(tp_size, dim=0)[rank])
+
+    return _impl
+
+
+def _tp_group(with_device_group: bool):
+    return SimpleNamespace(device_group=object() if with_device_group else None)
+
+
+@pytest.mark.parametrize("tp_size", [2, 4])
+@pytest.mark.parametrize("rank", [0, 1])
+def test_process_group_path_matches_all_reduce_then_slice(tp_size: int, rank: int) -> None:
+    """The scattered shard equals the slice the all-reduce path would take."""
+    local_width = HIDDEN // tp_size
+    torch.manual_seed(0)
+    contrib = torch.randn(2, 3, HIDDEN)
+    bias = torch.randn(HIDDEN)
+    reference = torch.zeros(2, 3, local_width)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
+        patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
+        patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(True)),
+        patch("torch.distributed.reduce_scatter_tensor", side_effect=_fake_reduce_scatter_tensor(rank, tp_size)),
+    ):
+        out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
+
+    start = rank * local_width
+    expected = (contrib + bias)[..., start : start + local_width]
+    assert out.shape == (2, 3, local_width)
+    torch.testing.assert_close(out, expected)
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_fallback_path_matches_the_process_group_path(rank: int) -> None:
+    """Without a device_group the method all-reduces and slices, same result."""
+    tp_size = 2
+    local_width = HIDDEN // tp_size
+    torch.manual_seed(1)
+    contrib = torch.randn(2, 3, HIDDEN)
+    bias = torch.randn(HIDDEN)
+    reference = torch.zeros(2, 3, local_width)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
+        patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
+        patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(False)),
+        patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=lambda t: t),
+    ):
+        out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
+
+    start = rank * local_width
+    torch.testing.assert_close(out, (contrib + bias)[..., start : start + local_width])
+
+
+def test_tp1_adds_bias_without_a_collective() -> None:
+    contrib = torch.randn(2, 3, HIDDEN)
+    bias = torch.randn(HIDDEN)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1),
+        patch("torch.distributed.reduce_scatter_tensor") as collective,
+    ):
+        out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, contrib)
+
+    assert collective.call_count == 0
+    torch.testing.assert_close(out, contrib + bias)
+
+
+def test_rejects_a_width_that_does_not_shard() -> None:
+    """A local width that is not 1/tp of the contribution is a wiring bug."""
+    contrib = torch.randn(1, 2, HIDDEN)
+    reference = torch.zeros(1, 2, HIDDEN // 3)  # tp=2 would need HIDDEN // 2
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+        patch("torch.distributed.reduce_scatter_tensor"),
+        pytest.raises(ValueError, match="camera contribution width mismatch"),
+    ):
+        SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(None), contrib, reference)

tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for SANA-WM's TensorParallelRMSNorm q/k fusion.
+
+`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor
+and issues a single tensor-parallel all-reduce instead of two. These tests lock
+in that it is numerically identical to applying `norm_q`/`norm_k` separately,
+and that it really issues exactly one (fused) all-reduce.
+
+CPU-only: the collective is mocked, so a single process simulates one rank
+(all-reduce is identity). This exercises the cat/split logic a real multi-rank
+run depends on.
+"""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import (
+    TensorParallelRMSNorm,
+    fused_qk_rms_norm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
+
+
+def _make_norms(dim: int, tp_size: int) -> tuple[TensorParallelRMSNorm, TensorParallelRMSNorm]:
+    """Two norms with distinct random weights so a swapped q/k slice is caught."""
+    torch.manual_seed(0)
+    norm_q = TensorParallelRMSNorm(dim, tp_size=tp_size)
+    norm_k = TensorParallelRMSNorm(dim, tp_size=tp_size)
+    norm_q.weight.data.normal_()
+    norm_k.weight.data.normal_()
+    return norm_q, norm_k
+
+
+def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
+    """Single-rank all-reduce: the reduction over one rank is the identity."""
+    return tensor
+
+
+def test_fused_matches_separate() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim, tp_size=2)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce):
+        ref_q, ref_k = norm_q(q), norm_k(k)
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_issues_one_all_reduce() -> None:
+    dim = 32
+    norm_q, norm_k = _make_norms(dim, tp_size=4)
+    q = torch.randn(1, 3, dim)
+    k = torch.randn(1, 3, dim)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+        fused_qk_rms_norm(norm_q, norm_k, q, k)
+        assert collective.call_count == 1
+        # Both sums travel in one tensor: last dim is 2, not 1.
+        assert collective.call_args.args[0].shape[-1] == 2
+
+        collective.reset_mock()
+        norm_q(q), norm_k(k)
+        assert collective.call_count == 2
+
+
+@pytest.mark.parametrize(
+    "norms",
+    [
+        (RMSNorm(16), RMSNorm(16)),
+        (nn.Identity(), nn.Identity()),
+    ],
+    ids=["plain_rms_norm", "identity"],
+)
+def test_falls_back_for_non_tp_norms(norms: tuple[nn.Module, nn.Module]) -> None:
+    """qk_norm=False and TP=1 both hand in norms the fused path cannot use."""
+    norm_q, norm_k = norms
+    q = torch.randn(2, 4, 16)
+    k = torch.randn(2, 4, 16)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+        assert collective.call_count == 0
+
+    torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)
+
+
+def test_tp1_skips_the_collective() -> None:
+    norm_q, norm_k = _make_norms(16, tp_size=1)
+    q = torch.randn(2, 4, 16)
+    k = torch.randn(2, 4, 16)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+        assert collective.call_count == 0
+
+    torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)

tests/diffusion/models/sana_wm/test_request.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for SANA-WM request normalization."""
+
+import pytest
+
+from vllm_omni.diffusion.models.sana_wm.request import (
+    SANA_WM_DEFAULT_HEIGHT,
+    SANA_WM_DEFAULT_NUM_FRAMES,
+    SANA_WM_DEFAULT_WIDTH,
+    normalize_sana_wm_payload,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _prompt(**sana_wm):
+    """A valid request, with the action rollout kept in step with num_frames."""
+    block = {"num_frames": SANA_WM_DEFAULT_NUM_FRAMES, **sana_wm}
+    block.setdefault("action", f"w-{block['num_frames'] - 1}")
+    return {
+        "prompt": "a quiet city street",
+        "multi_modal_data": {"image": object()},
+        "sana_wm": block,
+    }
+
+
+def _identity_poses(count):
+    return [[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]] for _ in range(count)]
+
+
+def _camera_prompt(**camera):
+    prompt = _prompt(num_frames=9)
+    prompt["sana_wm"].pop("action")
+    prompt["sana_wm"]["camera"] = {"poses": _identity_poses(9), **camera}
+    return prompt
+
+
+def test_defaults_are_latent_aligned():
+    payload = normalize_sana_wm_payload(_prompt())["additional_information"]["sana_wm"]
+    assert payload["num_frames"] == SANA_WM_DEFAULT_NUM_FRAMES
+    assert payload["height"] == SANA_WM_DEFAULT_HEIGHT
+    assert payload["width"] == SANA_WM_DEFAULT_WIDTH
+
+
+@pytest.mark.parametrize("num_frames", [9, 33, 161])
+def test_accepts_aligned_geometry(num_frames):
+    payload = normalize_sana_wm_payload(_prompt(num_frames=num_frames))["additional_information"]["sana_wm"]
+    assert payload["num_frames"] == num_frames
+
+
+@pytest.mark.parametrize("num_frames", [8, 10, 160])
+def test_rejects_unaligned_num_frames(num_frames):
+    # The VAE compresses 8x temporally; an unaligned count would be floored.
+    with pytest.raises(ValueError, match="num_frames"):
+        normalize_sana_wm_payload(_prompt(num_frames=num_frames))
+
+
+@pytest.mark.parametrize(("height", "width"), [(700, 1280), (704, 1270), (703, 1279)])
+def test_rejects_unaligned_resolution(height, width):
+    with pytest.raises(ValueError, match="divisible by 32"):
+        normalize_sana_wm_payload(_prompt(height=height, width=width))
+
+
+def test_normalization_is_idempotent():
+    once = normalize_sana_wm_payload(_prompt(num_frames=33))
+    twice = normalize_sana_wm_payload(once)
+    assert once["additional_information"]["sana_wm"] == twice["additional_information"]["sana_wm"]
+
+
+def test_rejects_action_shorter_than_num_frames():
+    # Used to shrink the camera tensor silently, then fail deep in the
+    # transformer on a camera/latent frame mismatch.
+    with pytest.raises(ValueError, match="rolls out 33 frames but num_frames is 161"):
+        normalize_sana_wm_payload(_prompt(action="w-32"))
+
+
+def test_rejects_action_longer_than_num_frames():
+    with pytest.raises(ValueError, match="rolls out 161 frames but num_frames is 9"):
+        normalize_sana_wm_payload(_prompt(num_frames=9, action="w-160"))
+
+
+def test_rejects_oversized_action_without_expanding_it():
+    # A 12-byte request body: expanding this rollout before checking the frame
+    # count allocated ~16 GB. Rejection has to be arithmetic, so this test also
+    # fails (by hanging/OOM) if the duration is ever expanded again.
+    with pytest.raises(ValueError, match="rolls out 1000000001 frames"):
+        normalize_sana_wm_payload(_prompt(action="w-1000000000"))
+
+
+@pytest.mark.parametrize("action", ["", "   ", "w-0", "w-abc", "w", "x-4"])
+def test_rejects_malformed_action(action):
+    with pytest.raises(ValueError):
+        normalize_sana_wm_payload(_prompt(action=action))
+
+
+@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf")])
+def test_rejects_unusable_focal_length(value):
+    # compute_raymap divides by fx/fy; a zero or non-finite focal length turns
+    # the whole Plucker map into NaN and the response into a NaN video.
+    intrinsics = {"fx": value, "fy": 640, "cx": 640, "cy": 352}
+    with pytest.raises(ValueError, match="fx"):
+        normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf")])
+def test_rejects_non_finite_principal_point(value):
+    intrinsics = {"fx": 640, "fy": 640, "cx": value, "cy": 352}
+    with pytest.raises(ValueError, match="must be finite"):
+        normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf")])
+def test_rejects_non_finite_speeds(value):
+    # `parsed <= 0` is False for NaN, so the positivity check alone let it pass.
+    with pytest.raises(ValueError, match="must be finite"):
+        normalize_sana_wm_payload(_prompt(translation_speed=value))
+
+
+def test_rejects_non_finite_pose_entry():
+    prompt = _camera_prompt()
+    prompt["sana_wm"]["camera"]["poses"][3][2][1] = float("nan")
+    with pytest.raises(ValueError, match="must be finite"):
+        normalize_sana_wm_payload(prompt)
+
+
+def test_accepts_default_camera_metadata():
+    payload = normalize_sana_wm_payload(_camera_prompt())["additional_information"]["sana_wm"]
+    assert payload["camera"]["poses"]
+    # The unread fields are not carried into the canonical payload.
+    assert set(payload["camera"]) == {"poses"}
+
+
+@pytest.mark.parametrize(
+    ("field", "value"),
+    [("format", "w2c_4x4"), ("coordinate_system", "opengl")],
+)
+def test_rejects_unsupported_camera_convention(field, value):
+    # Nothing downstream reads these, so accepting one would mean interpreting
+    # the poses under the wrong convention.
+    with pytest.raises(ValueError, match=field.split("_")[0]):
+        normalize_sana_wm_payload(_camera_prompt(**{field: value}))