vllm-omni/tests/diffusion/distributed/test_wan_decoder_fast_path.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""The exact Wan decoder fast path must be bit-identical to the plain streaming decode.

Two decoders are driven frame by frame through ``WanStreamingDecoder`` over two chunks of one session and
a fresh session (the persistent buffers' causal front frames go zeros -> cache -> zeros), on CPU: plain
diffusers modules in one process, and the spatially sharded wrappers across two gloo ranks with real halo
exchange, each in fp32 and under bf16 autocast (where the parameter cast applies). Every output must
``torch.equal`` the untouched decoder's on every rank.
"""

from __future__ import annotations

import copy
import os
from types import SimpleNamespace

import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d, WanDecoder3d

from vllm_omni.diffusion.distributed.autoencoders.wan_decoder_fast_path import install_wan_decoder_fast_path
from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import (
    WanDistCausalConv3d,
    WanDistConv2d,
    install_wan_spatial_shard_decode,
)
from vllm_omni.experimental.ar_diffusion.streaming_decode import WanStreamingDecoder


def _make_vae(seed: int = 0) -> SimpleNamespace:
    torch.manual_seed(seed)
    # Three levels: the default temporal-upsample list gives one upsample2d and one upsample3d resample.
    decoder = WanDecoder3d(dim=8, z_dim=4, dim_mult=[1, 2, 4], num_res_blocks=1, attn_scales=[])
    post_quant_conv = WanCausalConv3d(4, 4, 1)
    for p in list(decoder.parameters()) + list(post_quant_conv.parameters()):
        p.data.uniform_(-0.5, 0.5)
    decoder.eval()
    count = sum(1 for m in decoder.modules() if m.__class__.__name__ == "WanCausalConv3d")
    return SimpleNamespace(decoder=decoder, post_quant_conv=post_quant_conv, _cached_conv_counts={"decoder": count})


def _run(vae, latents: list[torch.Tensor], autocast: bool) -> list[torch.Tensor]:
    decoder = WanStreamingDecoder(vae)
    outs = []
    ctx = torch.autocast("cpu", dtype=torch.bfloat16) if autocast else torch.no_grad()
    with torch.no_grad(), ctx:
        state = decoder.new_decode_state("a")
        outs.append(decoder.decode_chunk(latents[0], state))
        outs.append(decoder.decode_chunk(latents[1], state))
        fresh = decoder.new_decode_state("b")
        outs.append(decoder.decode_chunk(latents[2], fresh))
    return [o.float() for o in outs]


def _latents() -> list[torch.Tensor]:
    torch.manual_seed(1)
    return [torch.randn(1, 4, 2, 6, 10), torch.randn(1, 4, 1, 6, 10), torch.randn(1, 4, 2, 6, 10)]


def _check_pair(reference, candidate, autocast: bool, sharded: bool, level: str = "exact") -> None:
    dtype = torch.bfloat16 if autocast else None
    counts = install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level=level)
    assert counts["level"] == level
    assert counts["upsamples"] == 2
    convs = sum(1 for m in candidate.decoder.modules() if isinstance(m, torch.nn.Conv3d | torch.nn.Conv2d)) + 1
    assert counts["conv_params_cast"] == (convs if autocast else 0)
    dist_convs = sum(1 for m in candidate.decoder.modules() if isinstance(m, WanDistCausalConv3d | WanDistConv2d))
    assert counts["persistent_input_buffers"] == dist_convs
    assert (dist_convs > 0) == sharded
    if level == "fused":
        assert (
            counts["fused_norm_silu"]
            == 2 * sum(1 for m in candidate.decoder.modules() if m.__class__.__name__ == "WanResidualBlock") + 1
        )
        assert counts["channels_last_convs"] == convs
        assert counts["resample_forwards"] == 2
        with pytest.raises(ValueError, match="already installed"):
            install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level="exact")
    # Idempotent: a second install is a no-op that reports the same counts.
    assert install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level=level) is counts
    latents = _latents()
    expected = _run(reference, latents, autocast)
    actual = _run(candidate, latents, autocast)
    assert len(expected) == len(actual) == 3
    for e, a in zip(expected, actual):
        assert e.shape == a.shape
        if level == "exact":
            assert torch.equal(e, a)
        else:
            # The fused level changes layouts only on CPU (the kernels need CUDA); the eager norm's reduction
            # order may differ with the layout, so the outputs are close, not identical.
            assert torch.allclose(e, a, atol=1e-2 if autocast else 1e-5, rtol=1e-2 if autocast else 1e-4), (
                (e - a).abs().max()
            )
    if sharded:
        bufs = [m._input_buf for m in candidate.decoder.modules() if isinstance(m, WanDistCausalConv3d | WanDistConv2d)]
        assert all(b is not None for b in bufs)


@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("autocast", [False, True])
def test_fast_path_is_bit_identical(autocast: bool) -> None:
    reference = _make_vae()
    candidate = copy.deepcopy(reference)
    _check_pair(reference, candidate, autocast, sharded=False)


def _swap_norms_to_rmsnorm_vae(vae) -> int:
    """Ensure the decoder uses RMSNormVAE and return the number of matching norms.

    Collecting pipeline tests can already apply patch_wan_rms_norm process-wide before this test runs.
    """
    from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE

    count = 0
    for module in list(vae.decoder.modules()):
        for name, child in list(module.named_children()):
            if isinstance(child, RMSNormVAE):
                count += 1
                continue
            if child.__class__.__name__ == "WanRMS_norm":
                images = child.gamma.dim() == 3  # (dim, 1, 1) for the attention block's 4D norm
                norm = RMSNormVAE(child.gamma.shape[0], channel_first=child.channel_first, images=images, bias=False)
                norm.gamma.data.copy_(child.gamma.data.reshape(norm.gamma.shape))
                setattr(module, name, norm)
                count += 1
    return count


@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("norm", ["default", "RMSNormVAE"])
@pytest.mark.parametrize("autocast", [False, True])
def test_fused_level_plumbing_on_cpu(autocast: bool, norm: str) -> None:
    reference = _make_vae()
    if norm == "RMSNormVAE":
        assert _swap_norms_to_rmsnorm_vae(reference) > 0
    candidate = copy.deepcopy(reference)
    _check_pair(reference, candidate, autocast, sharded=False, level="fused")


def _sharded_worker(rank: int, world_size: int, port: str, autocast: bool, return_dict, level: str = "exact") -> None:
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = port
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    try:
        torch.manual_seed(0)
        reference = _make_vae()
        candidate = copy.deepcopy(reference)
        for vae in (reference, candidate):
            install_wan_spatial_shard_decode(vae, dist.group.WORLD, split_dim="width", dst=None)
        _check_pair(reference, candidate, autocast, sharded=True, level=level)
        return_dict[rank] = "ok"
    except BaseException as exc:  # report, the parent asserts
        return_dict[rank] = f"{type(exc).__name__}: {exc}"
        raise
    finally:
        dist.destroy_process_group()


@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("level", ["exact", "fused"])
@pytest.mark.parametrize("autocast", [False, True])
def test_sharded_fast_path_on_every_rank(autocast: bool, level: str) -> None:
    manager = mp.get_context("spawn").Manager()
    return_dict = manager.dict()
    port = str(29610 + int(autocast) + 2 * (level == "fused"))
    mp.spawn(_sharded_worker, args=(2, port, autocast, return_dict, level), nprocs=2, join=True)
    for rank in range(2):
        assert return_dict.get(rank) == "ok", f"rank {rank}: {return_dict.get(rank)}"


@pytest.mark.core_model
@pytest.mark.cpu
def test_conv_dtype_requires_half() -> None:
    with pytest.raises(ValueError, match="conv_dtype"):
        install_wan_decoder_fast_path(_make_vae(), conv_dtype=torch.float32)