[optimization] Added Wan 2.2 VAE decode optimizations (#7056)

Signed-off-by: Maciej Bala <mbala@nvidia.com>
Signed-off-by: MaciejBalaNV <mbala@nvidia.com>

tests/diffusion/distributed/test_wan_vae_fastpath_install.py

@@ -0,0 +1,619 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU tests for the Wan VAE decoder fast path installer and its exact PyTorch fallbacks."""
+
+from __future__ import annotations
+
+import importlib
+from functools import wraps
+from types import SimpleNamespace
+
+import pytest
+import torch
+from diffusers.models.autoencoders import AutoencoderKLWan
+from torch import nn
+
+from vllm_omni.diffusion import registry as registry_module
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
+    REPORT_ATTR,
+    decode_frames,
+    install_wan_vae_fastpath,
+    is_installed,
+    uninstall_wan_vae_fastpath,
+)
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fastpath_forwards
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+TINY_RESIDUAL = dict(
+    base_dim=8,
+    decoder_base_dim=8,
+    z_dim=4,
+    dim_mult=[1, 1],
+    num_res_blocks=1,
+    temperal_downsample=[False, True],
+    is_residual=True,
+)
+TINY_RESIDUAL_PATCH2 = dict(TINY_RESIDUAL, patch_size=2, in_channels=12, out_channels=12)
+TINY_WAN21 = dict(
+    base_dim=8,
+    z_dim=4,
+    dim_mult=[1, 2],
+    num_res_blocks=1,
+    temperal_downsample=[False, True],
+    is_residual=False,
+)
+CONFIGS = {
+    "residual": TINY_RESIDUAL,
+    "residual_patch2": TINY_RESIDUAL_PATCH2,
+    "wan21": TINY_WAN21,
+}
+
+
+def _build_pair(config: dict, dtype: torch.dtype) -> tuple[AutoencoderKLWan, AutoencoderKLWan]:
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**config).eval().to(dtype)
+    candidate = AutoencoderKLWan(**config).eval().to(dtype)
+    candidate.load_state_dict(reference.state_dict())
+    return reference, candidate
+
+
+@pytest.mark.parametrize("config_name", sorted(CONFIGS))
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("frames", [1, 3])
+def test_lossless_fallback_paths_are_bitwise_exact(config_name: str, dtype: torch.dtype, frames: int) -> None:
+    """Without CUDA every kernel declines, so this exercises the restructured PyTorch paths."""
+    config = CONFIGS[config_name]
+    reference, candidate = _build_pair(config, dtype)
+    report = install_wan_vae_fastpath(candidate, level="lossless")
+    assert report.installed, report
+    assert report.fused_silu_dtypes == ()
+
+    torch.manual_seed(1)
+    latents = torch.randn(1, config["z_dim"], frames, 6, 8).to(dtype)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    assert actual.stride() == expected.stride()
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize("config_name", sorted(CONFIGS))
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_decode_frames_matches_reference_decode(config_name: str, dtype: torch.dtype) -> None:
+    config = CONFIGS[config_name]
+    reference, candidate = _build_pair(config, dtype)
+    install_wan_vae_fastpath(candidate)
+    for frames in (1, 2, 4):
+        torch.manual_seed(2)
+        latents = torch.randn(1, config["z_dim"], frames, 6, 8).to(dtype)
+        with torch.no_grad():
+            expected = reference._decode(latents, return_dict=False)[0]
+            actual = decode_frames(candidate, latents)
+        assert actual.stride() == expected.stride()
+        assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_omni_wan_vae_decode_override_is_bitwise_exact(dtype: torch.dtype) -> None:
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**TINY_RESIDUAL_PATCH2).eval().to(dtype)
+    candidate = OmniAutoencoderKLWan(**TINY_RESIDUAL_PATCH2).eval().to(dtype)
+    candidate.load_state_dict(reference.state_dict())
+    latents = torch.randn(1, 4, 3, 6, 8).to(dtype)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        before_install = candidate.decode(latents, return_dict=False)[0]
+        install_wan_vae_fastpath(candidate)
+        actual = candidate.decode(latents, return_dict=False)[0]
+        as_dict = candidate.decode(latents).sample
+    # The Omni wrapper runs bf16 decode under autocast (also on CPU), so compare
+    # the override against the same wrapper before installation, and against the
+    # plain diffusers decode only where no autocast is involved.
+    assert actual.stride() == before_install.stride()
+    assert torch.equal(actual, before_install)
+    assert torch.equal(as_dict, before_install)
+    if dtype is torch.float32:
+        assert torch.equal(before_install, expected)
+
+    # Tiling still dispatches to the diffusers implementation.
+    candidate.use_tiling = True
+    candidate.tile_sample_min_height = candidate.tile_sample_min_width = 16
+    candidate.tile_sample_stride_height = candidate.tile_sample_stride_width = 16
+    reference.use_tiling = True
+    reference.tile_sample_min_height = reference.tile_sample_min_width = 16
+    reference.tile_sample_stride_height = reference.tile_sample_stride_width = 16
+    with torch.no_grad():
+        expected_tiled = reference.decode(latents, return_dict=False)[0]
+        actual_tiled = candidate.decode(latents, return_dict=False)[0]
+    if dtype is torch.float32:
+        assert torch.equal(actual_tiled, expected_tiled)
+    assert actual_tiled.shape == expected_tiled.shape
+
+
+def test_omni_wan_vae_decode_preserves_autograd() -> None:
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**TINY_RESIDUAL).eval()
+    candidate = OmniAutoencoderKLWan(**TINY_RESIDUAL).eval()
+    candidate.load_state_dict(reference.state_dict())
+    install_wan_vae_fastpath(candidate)
+
+    reference_latents = torch.randn(1, 4, 2, 6, 8, requires_grad=True)
+    candidate_latents = reference_latents.detach().clone().requires_grad_()
+    expected = reference.decode(reference_latents, return_dict=False)[0]
+    actual = candidate.decode(candidate_latents, return_dict=False)[0]
+
+    torch.testing.assert_close(actual, expected)
+    expected.square().mean().backward()
+    actual.square().mean().backward()
+    torch.testing.assert_close(candidate_latents.grad, reference_latents.grad)
+
+
+def test_install_is_idempotent_and_reversible() -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    keys_before = set(vae.state_dict())
+    report = install_wan_vae_fastpath(vae)
+    assert report.installed and is_installed(vae)
+    assert report.patched["WanDecoder3d"] == 1
+    assert report.patched["WanCausalConv3d"] > 1
+    assert report.patched["WanRMS_norm"] > 1
+    assert set(vae.state_dict()) == keys_before
+    assert "forward" in vae.decoder.__dict__
+
+    assert install_wan_vae_fastpath(vae) is report
+    assert install_wan_vae_fastpath(vae, level="channels_last") is report  # warns, keeps the first level
+
+    uninstall_wan_vae_fastpath(vae)
+    assert not is_installed(vae)
+    assert "forward" not in vae.decoder.__dict__
+    assert all("forward" not in module.__dict__ for module in vae.decoder.modules())
+    assert all(fastpath_forwards.CFG_ATTR not in module.__dict__ for module in vae.decoder.modules())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("failure_stage", ["binding", "conversion", "report"])
+def test_failed_install_restores_original_state(failure_stage: str, monkeypatch) -> None:
+    installer = importlib.import_module(install_wan_vae_fastpath.__module__)
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    # Preserve mixed original layouts, an existing forward, gradients and a
+    # buffer: Module.to can replace tensors as well as change their storage.
+    vae.decoder.conv_in.to(memory_format=torch.channels_last_3d)
+    vae.decoder.forward = vae.decoder.forward
+    setattr(vae.decoder, fastpath_forwards.CFG_ATTR, object())
+    vae.decoder.conv_out.weight.grad = torch.randn_like(vae.decoder.conv_out.weight)
+    vae.post_quant_conv.register_buffer("rollback_probe", torch.randn(2, 3, 4, 5, 6))
+    attributes = {
+        module: {
+            name: module.__dict__[name] for name in ("forward", fastpath_forwards.CFG_ATTR) if name in module.__dict__
+        }
+        for module in vae.modules()
+    }
+    parameters = dict(vae.named_parameters())
+    buffers = dict(vae.named_buffers())
+    gradients = {name: parameter.grad for name, parameter in parameters.items()}
+    tensors = [*parameters.values(), *buffers.values(), *(grad for grad in gradients.values() if grad is not None)]
+    originals = [(tensor, tensor.data_ptr(), tensor.stride(), tensor.clone()) for tensor in tensors]
+    latents = torch.randn(1, 4, 2, 6, 8)
+    expected = vae.decode(latents, return_dict=False)[0]
+
+    with monkeypatch.context() as patch:
+        if failure_stage == "binding":
+            conv_type = type(vae.post_quant_conv)
+            original_setattr = conv_type.__setattr__
+
+            def fail_binding(module, name, value):
+                original_setattr(module, name, value)
+                if module is vae.post_quant_conv and name == fastpath_forwards.CFG_ATTR:
+                    raise RuntimeError("injected installation failure")
+
+            patch.setattr(conv_type, "__setattr__", fail_binding)
+        elif failure_stage == "conversion":
+            original_to = vae.post_quant_conv.to
+
+            def fail_conversion(*args, **kwargs):
+                original_to(*args, **kwargs)
+                assert vae.decoder.conv_out.weight.is_contiguous(memory_format=torch.channels_last_3d)
+                raise torch.OutOfMemoryError("injected installation failure")
+
+            patch.setattr(vae.post_quant_conv, "to", fail_conversion)
+        else:
+            original_info = installer.logger.info
+
+            def fail_report(message, *args, **kwargs):
+                if message.startswith("Wan VAE fast path (%s) installed"):
+                    assert is_installed(vae)
+                    raise RuntimeError("injected installation failure")
+                original_info(message, *args, **kwargs)
+
+            patch.setattr(installer.logger, "info", fail_report)
+
+        with pytest.raises(RuntimeError, match="injected installation failure"):
+            install_wan_vae_fastpath(vae, level="channels_last")
+
+    assert not is_installed(vae)
+    assert not hasattr(vae, REPORT_ATTR)
+    assert not hasattr(vae, installer._UNDO_ATTR)
+    for module, original in attributes.items():
+        actual = {
+            name: module.__dict__[name] for name in ("forward", fastpath_forwards.CFG_ATTR) if name in module.__dict__
+        }
+        assert actual == original
+    assert all(dict(vae.named_parameters())[name] is parameter for name, parameter in parameters.items())
+    assert all(dict(vae.named_buffers())[name] is buffer for name, buffer in buffers.items())
+    assert all(parameters[name].grad is grad for name, grad in gradients.items())
+    for tensor, pointer, strides, values in originals:
+        assert tensor.data_ptr() == pointer
+        assert tensor.stride() == strides
+        assert torch.equal(tensor, values)
+    uninstall_wan_vae_fastpath(vae)  # Safe and unnecessary after automatic rollback.
+    assert torch.equal(vae.decode(latents, return_dict=False)[0], expected)
+    assert install_wan_vae_fastpath(vae, level="channels_last").installed
+    uninstall_wan_vae_fastpath(vae)
+
+
+def test_load_state_dict_after_install_updates_patched_modules() -> None:
+    reference, candidate = _build_pair(TINY_RESIDUAL, torch.float32)
+    install_wan_vae_fastpath(candidate)
+    with torch.no_grad():
+        for parameter in reference.parameters():
+            parameter.mul_(0.5)
+    candidate.load_state_dict(reference.state_dict())
+    latents = torch.randn(1, 4, 2, 6, 8)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    assert torch.equal(actual, expected)
+
+
+def test_channels_last_level_converts_conv_weights_and_restores_them() -> None:
+    reference, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    report = install_wan_vae_fastpath(vae, level="channels_last")
+    assert report.installed and report.channels_last
+    conv3d = [m for m in vae.decoder.modules() if isinstance(m, nn.Conv3d)]
+    conv2d = [m for m in vae.decoder.modules() if isinstance(m, nn.Conv2d)]
+    assert conv3d and conv2d
+    assert all(m.weight.is_contiguous(memory_format=torch.channels_last_3d) for m in conv3d)
+    assert all(m.weight.is_contiguous(memory_format=torch.channels_last) for m in conv2d)
+    assert vae.post_quant_conv.weight.is_contiguous(memory_format=torch.channels_last_3d)
+    latents = torch.randn(1, 4, 2, 6, 8)
+    with torch.no_grad():
+        expected = reference.decode(latents, return_dict=False)[0]
+        out = vae.decode(latents, return_dict=False)[0]
+    # One temporal upsampler: 1 + 2 output frames for 2 latent frames.
+    assert expected.shape == (1, 3, 3, 12, 16)
+    assert out.shape == expected.shape
+    torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-4)
+
+    uninstall_wan_vae_fastpath(vae)
+    assert all(m.weight.is_contiguous() for m in conv3d + conv2d)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("assign", [False, True])
+def test_uninstall_restores_mixed_layouts_and_keeps_weight_updates(assign: bool) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    convs = [module for module in vae.decoder.modules() if isinstance(module, (nn.Conv2d, nn.Conv3d))]
+    convs.append(vae.post_quant_conv)
+    for index, conv in enumerate(convs):
+        if index % 3 == 1:
+            conv.to(memory_format=torch.channels_last_3d if isinstance(conv, nn.Conv3d) else torch.channels_last)
+        elif index % 3 == 2:
+            conv.weight.data = conv.weight.data.transpose(-1, -2)
+    vae.post_quant_conv.register_buffer("layout_probe", torch.randn(2, 3, 4, 5, 6).transpose(-1, -2))
+    vae.decoder.conv_out.weight.grad = torch.randn_like(vae.decoder.conv_out.weight)
+    grad_strides = vae.decoder.conv_out.weight.grad.stride()
+    strides = {name: tensor.stride() for name, tensor in vae.state_dict().items()}
+    # An explicit alias is safe to replace, but must remain an instance
+    # attribute after uninstall, with its exact original identity.
+    original_forward = vae.decoder.forward
+    vae.decoder.forward = original_forward
+    original_cfg = object()
+    setattr(vae.decoder, fastpath_forwards.CFG_ATTR, original_cfg)
+
+    assert install_wan_vae_fastpath(vae, level="channels_last").installed
+    updated = {name: tensor.contiguous() + 0.25 for name, tensor in vae.state_dict().items()}
+    vae.load_state_dict(updated, assign=assign)
+    uninstall_wan_vae_fastpath(vae)
+
+    assert vae.decoder.__dict__["forward"] is original_forward
+    assert getattr(vae.decoder, fastpath_forwards.CFG_ATTR) is original_cfg
+    for name, tensor in vae.state_dict().items():
+        assert tensor.stride() == strides[name], name
+        assert torch.equal(tensor, updated[name]), name
+    if not assign:
+        assert vae.decoder.conv_out.weight.grad.stride() == grad_strides
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("level", ["lossless", "channels_last"])
+@pytest.mark.parametrize(
+    "target",
+    [
+        "decoder",
+        "decoder.conv_in",
+        "decoder.up_blocks.0.upsampler",
+        "decoder.up_blocks.0.upsampler.resample",
+        "decoder.up_blocks.0.upsampler.resample.1",
+        "decoder.up_blocks.0.avg_shortcut",
+        "decoder.norm_out",
+        "decoder.nonlinearity",
+        "post_quant_conv",
+    ],
+)
+def test_installer_declines_forward_wrappers_it_would_bypass(target: str, level: str) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    module = vae.get_submodule(target)
+    original_forward = module.forward
+    calls = []
+
+    @wraps(original_forward)
+    def wrapped(*args, **kwargs):
+        calls.append(True)
+        return original_forward(*args, **kwargs)
+
+    module.forward = wrapped
+    originals = {name: (tensor.data_ptr(), tensor.stride()) for name, tensor in vae.state_dict().items()}
+    report = install_wan_vae_fastpath(vae, level=level)
+    assert not report.installed and target in report.reason and "custom forward" in report.reason
+    assert not is_installed(vae) and not hasattr(vae, REPORT_ATTR)
+    assert all(fastpath_forwards.CFG_ATTR not in child.__dict__ for child in vae.modules())
+    assert {name: (tensor.data_ptr(), tensor.stride()) for name, tensor in vae.state_dict().items()} == originals
+    uninstall_wan_vae_fastpath(vae)
+    assert module.forward is wrapped
+    vae.decode(torch.randn(1, 4, 2, 6, 8))
+    assert calls
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("pre_hook", [False, True])
+def test_installer_declines_convolution_hooks_it_would_bypass(pre_hook: bool) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    conv = vae.decoder.conv_in
+    calls = []
+
+    def hook(*args):
+        calls.append(True)
+
+    handle = conv.register_forward_pre_hook(hook) if pre_hook else conv.register_forward_hook(hook)
+    try:
+        report = install_wan_vae_fastpath(vae, level="channels_last")
+        assert not report.installed and "decoder.conv_in" in report.reason and "forward hooks" in report.reason
+        vae.decode(torch.randn(1, 4, 2, 6, 8))
+        assert calls
+    finally:
+        handle.remove()
+
+
+@torch.no_grad()
+def test_installer_preserves_wrappers_and_hooks_on_modules_it_still_calls() -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    original_forward = vae.decoder.mid_block.forward
+    wrapper_calls = []
+    hook_calls = []
+
+    def wrapped(*args, **kwargs):
+        wrapper_calls.append(True)
+        return original_forward(*args, **kwargs)
+
+    def hook(*args):
+        hook_calls.append(True)
+
+    vae.decoder.mid_block.forward = wrapped
+    handle = vae.decoder.register_forward_hook(hook)
+    try:
+        latents = torch.randn(1, 4, 2, 6, 8)
+        expected = vae.decode(latents).sample
+        wrapper_calls.clear()
+        hook_calls.clear()
+        assert install_wan_vae_fastpath(vae).installed
+        actual = vae.decode(latents).sample
+        assert torch.equal(actual, expected)
+        assert len(wrapper_calls) == len(hook_calls) == 2
+        uninstall_wan_vae_fastpath(vae)
+        assert vae.decoder.mid_block.forward is wrapped
+    finally:
+        handle.remove()
+
+
+def test_installer_refuses_unsupported_targets() -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+
+    assert install_wan_vae_fastpath(vae, level="off").installed is False
+    assert not is_installed(vae)
+
+    vae._vllm_omni_wan_spatial_shard_installed = True
+    report = install_wan_vae_fastpath(vae)
+    assert not report.installed and "spatial-shard" in report.reason
+    del vae._vllm_omni_wan_spatial_shard_installed
+
+    vae.distributed_executor = SimpleNamespace(parallel_mode="spatial_shard_height")
+    report = install_wan_vae_fastpath(vae)
+    assert not report.installed and "spatial_shard_height" in report.reason
+    del vae.distributed_executor
+
+    report = install_wan_vae_fastpath(nn.Linear(2, 2))
+    assert not report.installed and "AutoencoderKLWan" in report.reason
+
+    with pytest.raises(ValueError, match="vae_fast_path"):
+        install_wan_vae_fastpath(vae, level="bogus")
+
+
+@pytest.mark.parametrize(("batch", "frames"), [(1, 1), (1, 2), (2, 1)])
+def test_resample_views_keep_channels_last_recognizable(batch: int, frames: int) -> None:
+    """Regression: ``reshape`` gives a size-1 batch dim a stride the layout heuristic rejects."""
+    import torch.nn.functional as F
+
+    x = torch.randn(batch, 8, frames, 6, 10).contiguous(memory_format=torch.channels_last_3d)
+    merged = fastpath_forwards._merge_batch_and_frames(x)
+    reference = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, 8, 6, 10)
+    assert merged.data_ptr() == x.data_ptr()
+    assert torch.equal(merged, reference)
+    upsampled = F.interpolate(merged, scale_factor=(2.0, 2.0), mode="nearest-exact")
+    assert upsampled.stride(1) == 1, "nearest upsample must keep channels_last for the following Conv2d"
+    assert torch.equal(upsampled, F.interpolate(reference.contiguous(), scale_factor=(2.0, 2.0), mode="nearest-exact"))
+
+    split = fastpath_forwards._split_batch_and_frames(upsampled, batch, frames)
+    assert split.shape == (batch, 8, frames, 12, 20)
+    assert split.data_ptr() == upsampled.data_ptr()
+    assert split.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(split, upsampled.view(batch, frames, 8, 12, 20).permute(0, 2, 1, 3, 4))
+
+    plain = torch.randn(batch, 8, frames, 6, 10)
+    assert torch.equal(
+        fastpath_forwards._merge_batch_and_frames(plain), plain.permute(0, 2, 1, 3, 4).reshape(-1, 8, 6, 10)
+    )
+
+
+@torch.no_grad()
+def test_pending_conv_bias_is_added_exactly_when_no_kernel_consumes_it() -> None:
+    """Without the channels-last kernel the un-added conv bias is applied with ATen's rounding."""
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanRMS_norm
+
+    torch.manual_seed(0)
+    for dtype in (torch.float32, torch.bfloat16):
+        norm = WanRMS_norm(8, images=False).to(dtype)
+        act = nn.SiLU()
+        x = torch.randn(1, 8, 2, 4, 4).to(dtype)
+        bias = torch.randn(8).to(dtype)
+        biased = x.clone()
+        biased.add_(bias.view(1, -1, 1, 1, 1))
+        expected = act(norm(biased))
+        assert torch.equal(fastpath_forwards._norm_act(norm, act, x, pending_bias=bias), expected)
+        setattr(norm, fastpath_forwards.CFG_ATTR, fastpath_forwards.FastPathConfig(channels_last=True))
+        assert torch.equal(fastpath_forwards._norm_act(norm, act, x, pending_bias=bias), expected)
+        out = fastpath_forwards.rms_norm_fastpath(norm, x, bias=bias)
+        assert out is not None and torch.equal(out, norm(biased))
+
+
+@torch.no_grad()
+def test_resample_return_bias_is_none_without_kernels() -> None:
+    """On CPU the Conv2d keeps its bias, so the up block adds nothing twice."""
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanResample, WanResidualUpBlock
+
+    torch.manual_seed(0)
+    resample = WanResample(8, mode="upsample2d").eval()
+    assert fastpath_forwards._is_upsample_conv_pair(resample.resample)
+    assert not fastpath_forwards._is_upsample_conv_pair(WanResample(8, mode="downsample2d").resample)
+    x = torch.randn(1, 8, 2, 6, 10)
+    out, bias = fastpath_forwards.resample_forward(resample, x, [None], [0], return_bias=True)
+    assert bias is None
+    assert torch.equal(out, WanResample.forward(resample, x, feat_cache=[None], feat_idx=[0]))
+
+    # First chunk of a temporal up block: one input frame, the time conv is skipped ("Rep").
+    block = WanResidualUpBlock(8, 8, num_res_blocks=1, temperal_upsample=True, up_flag=True).eval()
+    for module in block.modules():
+        setattr(module, fastpath_forwards.CFG_ATTR, fastpath_forwards.FastPathConfig())
+    x = torch.randn(1, 8, 1, 6, 10)
+    cache_len = 6
+    expected = WanResidualUpBlock.forward(block, x, feat_cache=[None] * cache_len, feat_idx=[0], first_chunk=True)
+    actual = fastpath_forwards.residual_up_block_forward(
+        block, x, feat_cache=[None] * cache_len, feat_idx=[0], first_chunk=True
+    )
+    assert torch.equal(actual, expected)
+
+
+def test_upsample_forward_only_fuses_nearest_2x() -> None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanUpsample
+
+    assert fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"))
+    assert fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=2, mode="nearest"))
+    assert not fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=(2.0, 3.0), mode="nearest-exact"))
+    assert not fastpath_forwards._is_nearest_2x(WanUpsample(scale_factor=(2.0, 2.0), mode="bilinear"))
+    assert not fastpath_forwards._is_nearest_2x(WanUpsample(size=(12, 20), mode="nearest-exact"))
+
+    module = WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact")
+    for dtype in (torch.float32, torch.bfloat16):
+        x = torch.randn(3, 8, 6, 10).to(dtype)
+        for layout in (torch.contiguous_format, torch.channels_last):
+            x = x.contiguous(memory_format=layout)
+            expected = WanUpsample.forward(module, x)
+            actual = fastpath_forwards.upsample_forward(module, x)
+            assert actual.stride() == expected.stride()
+            assert torch.equal(actual, expected)
+
+
+def test_rms_norm_vae_substitute_is_not_matched() -> None:
+    from vllm_omni.diffusion.layers.norm import RMSNormVAE
+
+    assert not fastpath_forwards.is_diffusers_rms_norm(RMSNormVAE(8, images=False))
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    norm = vae.decoder.norm_out
+    assert fastpath_forwards.is_diffusers_rms_norm(norm)
+
+
+@torch.no_grad()
+def test_rms_norm_fastpath_declines_tensor_bias_and_dtype_mismatch() -> None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanRMS_norm
+
+    norm = WanRMS_norm(8, images=False, bias=True)
+    assert fastpath_forwards.rms_norm_fastpath(norm, torch.randn(1, 8, 1, 4, 4)) is None
+
+    norm = WanRMS_norm(8, images=False)
+    assert fastpath_forwards.rms_norm_fastpath(norm, torch.randn(1, 8, 1, 4, 4).to(torch.bfloat16)) is None
+
+    x = torch.randn(1, 8, 2, 4, 4)
+    out = fastpath_forwards.rms_norm_fastpath(norm, x)
+    assert out is not None and torch.equal(out, norm(x))
+
+
+def test_rms_norm_forward_preserves_autograd() -> None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanRMS_norm
+
+    torch.manual_seed(0)
+    reference = WanRMS_norm(8, images=False)
+    candidate = WanRMS_norm(8, images=False)
+    candidate.load_state_dict(reference.state_dict())
+    reference_input = torch.randn(1, 8, 1, 4, 4, requires_grad=True)
+    candidate_input = reference_input.detach().clone().requires_grad_()
+
+    expected = reference(reference_input)
+    assert fastpath_forwards.rms_norm_fastpath(candidate, candidate_input) is None
+    actual = fastpath_forwards.rms_norm_forward(candidate, candidate_input)
+    torch.testing.assert_close(actual, expected)
+
+    expected.square().mean().backward()
+    actual.square().mean().backward()
+    torch.testing.assert_close(candidate_input.grad, reference_input.grad)
+    torch.testing.assert_close(candidate.gamma.grad, reference.gamma.grad)
+
+
+def test_omni_diffusion_config_validates_vae_fast_path() -> None:
+    assert OmniDiffusionConfig(model="x").vae_fast_path == "lossless"
+    assert OmniDiffusionConfig(model="x", vae_fast_path="channels_last").vae_fast_path == "channels_last"
+    with pytest.raises(ValueError, match="vae_fast_path"):
+        OmniDiffusionConfig(model="x", vae_fast_path="fast")
+
+
+class _StubPipeline(nn.Module):
+    def __init__(self, vae: nn.Module) -> None:
+        super().__init__()
+        self.vae = vae
+
+
+@pytest.mark.parametrize(("level", "expected"), [("lossless", True), ("channels_last", True), ("off", False)])
+def test_registry_hook_installs_on_cuda_platform(mocker, level: str, expected: bool) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    platform = mocker.Mock()
+    platform.is_cuda.return_value = True
+    mocker.patch.object(registry_module, "current_omni_platform", platform)
+
+    registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(vae), SimpleNamespace(vae_fast_path=level))
+    assert is_installed(vae) is expected
+    if expected:
+        assert getattr(vae, REPORT_ATTR).level == level
+
+
+def test_registry_hook_skips_non_cuda_platform_and_non_wan_vaes(mocker) -> None:
+    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
+    platform = mocker.Mock()
+    platform.is_cuda.return_value = False
+    mocker.patch.object(registry_module, "current_omni_platform", platform)
+    registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(vae), SimpleNamespace(vae_fast_path="lossless"))
+    assert not is_installed(vae)
+
+    platform.is_cuda.return_value = True
+    other = nn.Linear(2, 2)
+    registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(other), SimpleNamespace(vae_fast_path="lossless"))
+    assert not hasattr(other, REPORT_ATTR)

tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py

@@ -0,0 +1,903 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""GPU bit-exactness tests for the Wan VAE decoder fast path kernels and forwards."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+from diffusers.models.autoencoders import AutoencoderKLWan
+from diffusers.models.autoencoders.autoencoder_kl_wan import (
+    CACHE_T,
+    WanCausalConv3d,
+    WanResample,
+    WanResidualBlock,
+    WanRMS_norm,
+    WanUpsample,
+)
+from torch import nn
+
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fp
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
+    install_wan_vae_fastpath,
+)
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_data_movement as dm
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_rms_norm as rn
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_rms_norm_cl as cl
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_upsample as up
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.diffusion,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+LOW_PRECISION = [torch.bfloat16, torch.float16]
+ALL_DTYPES = [torch.bfloat16, torch.float16, torch.float32]
+# Cosmos3 720p decoder stages, spatially scaled down so the suite stays fast.
+NORM_SHAPES = [
+    (1, 1024, 1, 45, 80),
+    (1, 1024, 2, 45, 80),
+    (1, 512, 4, 45, 80),
+    (1, 256, 4, 90, 160),
+]
+
+
+def _bits_equal(a: torch.Tensor, b: torch.Tensor) -> bool:
+    int_dtype = torch.int16 if a.element_size() == 2 else torch.int32
+    return a.shape == b.shape and bool((a.view(int_dtype) == b.view(int_dtype)).all())
+
+
+def _same_dense_strides(a: torch.Tensor, b: torch.Tensor) -> bool:
+    """Stride equality ignoring size-1 dims, whose stride carries no information."""
+    return a.shape == b.shape and all(
+        sa == sb for sa, sb, n in zip(a.stride(), b.stride(), a.shape, strict=True) if n > 1
+    )
+
+
+def _with_negative_zeros(x: torch.Tensor) -> torch.Tensor:
+    flat = x.view(-1)
+    flat[:: max(1, flat.numel() // 97)] = -0.0
+    return x
+
+
+def _reference_rms_norm(norm: WanRMS_norm, x: torch.Tensor) -> torch.Tensor:
+    return WanRMS_norm.forward(norm, x)
+
+
+def _make_norm(channels: int, dtype: torch.dtype, images: bool = False) -> WanRMS_norm:
+    norm = WanRMS_norm(channels, images=images).to(device="cuda", dtype=dtype)
+    with torch.no_grad():
+        norm.gamma.copy_(torch.randn_like(norm.gamma) * 0.5 + 1.0)
+    return norm
+
+
+# --------------------------------------------------------------------------- #
+# RMSNorm
+# --------------------------------------------------------------------------- #
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("shape", NORM_SHAPES)
+def test_vector_norm_from_low_precision_matches_fp32_normalize_denominator(dtype: torch.dtype, shape) -> None:
+    """The fast path relies on ATen running the same reduction for both inputs."""
+    torch.manual_seed(0)
+    x = torch.randn(shape, device="cuda").to(dtype)
+    expected = x.float().norm(2, dim=1, keepdim=True)
+    actual = torch.linalg.vector_norm(x, dim=1, keepdim=True, dtype=torch.float32)
+    assert torch.equal(actual, expected)
+    frame_major = x.permute(0, 2, 1, 3, 4).contiguous().permute(0, 2, 1, 3, 4)
+    expected = frame_major.float().norm(2, dim=1, keepdim=True)
+    actual = torch.linalg.vector_norm(frame_major, dim=1, keepdim=True, dtype=torch.float32)
+    assert torch.equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("shape", NORM_SHAPES)
+@pytest.mark.parametrize("layout", ["contiguous", "frame_major"])
+@pytest.mark.parametrize("autocast", [False, True])
+def test_rms_norm_fastpath_is_bitwise_exact(dtype: torch.dtype, shape, layout: str, autocast: bool) -> None:
+    torch.manual_seed(0)
+    norm = _make_norm(shape[1], dtype)
+    x = _with_negative_zeros(torch.randn(shape, device="cuda").to(dtype))
+    if layout == "frame_major":
+        x = x.permute(0, 2, 1, 3, 4).contiguous().permute(0, 2, 1, 3, 4)
+    with torch.autocast("cuda", dtype=dtype if dtype is not torch.float32 else torch.bfloat16, enabled=autocast):
+        expected = _reference_rms_norm(norm, x)
+        actual = fp.rms_norm_fastpath(norm, x)
+    assert actual is not None
+    assert actual.dtype == expected.dtype
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected), "RMSNorm fast path is not bit-exact (signed zeros included)"
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_rms_norm_fastpath_handles_attention_block_4d_input(dtype: torch.dtype) -> None:
+    torch.manual_seed(0)
+    norm = _make_norm(1024, dtype, images=True)
+    x = torch.randn(2, 1024, 45, 80, device="cuda").to(dtype)
+    expected = _reference_rms_norm(norm, x)
+    actual = fp.rms_norm_fastpath(norm, x)
+    assert actual is not None and _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+def test_fused_silu_epilogue_matches_silu_of_reference(dtype: torch.dtype) -> None:
+    device = torch.device("cuda")
+    if not rn.silu_epilogue_is_exact(device, dtype):
+        pytest.skip(f"fused SiLU epilogue is not bit-exact for {dtype} on this toolkit; fusion stays disabled")
+    torch.manual_seed(0)
+    norm = _make_norm(256, dtype)
+    x = _with_negative_zeros(torch.randn(1, 256, 4, 90, 160, device="cuda").to(dtype))
+    expected = F.silu(_reference_rms_norm(norm, x))
+    actual = fp.rms_norm_fastpath(norm, x, silu=True)
+    assert actual is not None and _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("shape", [(1, 1024, 2, 45, 80), (1, 256, 4, 90, 160), (2, 512, 45, 80)])
+@pytest.mark.parametrize("silu", [False, True])
+def test_channels_last_rms_norm_is_close_and_keeps_layout(dtype: torch.dtype, shape, silu: bool) -> None:
+    torch.manual_seed(0)
+    norm = _make_norm(shape[1], dtype, images=len(shape) == 4)
+    memory_format = torch.channels_last_3d if len(shape) == 5 else torch.channels_last
+    x = torch.randn(shape, device="cuda").to(dtype).contiguous(memory_format=memory_format)
+    expected = _reference_rms_norm(norm, x)
+    if silu:
+        expected = F.silu(expected)
+    actual = cl.rms_norm_channels_last(x, norm.gamma, norm.scale, silu=silu)
+    assert actual is not None
+    assert actual.is_contiguous(memory_format=memory_format)
+    assert actual.dtype == expected.dtype
+    # Fast-math kernel: approximate divide/exp and a single final rounding.
+    if dtype is torch.float32:
+        torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)
+    else:
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+    # Declines channels-first input and too many channels.
+    assert cl.rms_norm_channels_last(x.contiguous(), norm.gamma, norm.scale) is None
+    wide = torch.randn(1, 2048, 1, 4, 4, device="cuda", dtype=dtype).contiguous(memory_format=torch.channels_last_3d)
+    assert cl.rms_norm_channels_last(wide, torch.ones(2048, device="cuda", dtype=dtype), 1.0) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("silu", [False, True])
+def test_channels_last_rms_norm_absorbs_conv_bias(dtype: torch.dtype, silu: bool) -> None:
+    torch.manual_seed(0)
+    shape = (1, 256, 4, 90, 160)
+    norm = _make_norm(shape[1], dtype)
+    x = torch.randn(shape, device="cuda").to(dtype).contiguous(memory_format=torch.channels_last_3d)
+    bias = (torch.randn(shape[1], device="cuda") * 0.1).to(dtype)
+    # Reference: ATen adds the conv bias in place (one rounding) before the norm.
+    expected = _reference_rms_norm(norm, x + bias.view(1, -1, 1, 1, 1))
+    if silu:
+        expected = F.silu(expected)
+    actual = cl.rms_norm_channels_last(x, norm.gamma, norm.scale, silu=silu, bias=bias)
+    assert actual is not None
+    assert actual.is_contiguous(memory_format=torch.channels_last_3d)
+    if dtype is torch.float32:
+        torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)
+    else:
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+    assert cl.rms_norm_channels_last(x, norm.gamma, norm.scale, bias=bias[:-1]) is None
+    assert cl.rms_norm_channels_last(x, norm.gamma, norm.scale, bias=bias.view(1, -1)) is None
+    assert cl.rms_norm_channels_last(x, norm.gamma, norm.scale, bias=bias.cpu()) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_channels_last_residual_block_folds_conv1_bias(dtype: torch.dtype, monkeypatch) -> None:
+    """The conv1 bias is left un-added and consumed by norm2's kernel (rounded like ATen); output stays close."""
+    torch.manual_seed(0)
+    block = WanResidualBlock(64, 64, dropout=0.0).eval().to(device="cuda", dtype=dtype)
+    with torch.no_grad():
+        for conv in (block.conv1, block.conv2):
+            conv.bias.copy_(torch.randn_like(conv.bias) * 0.1)
+    x = torch.randn(1, 64, 4, 24, 40, device="cuda", dtype=dtype)
+    expected = WanResidualBlock.forward(block, x, feat_cache=[None, None], feat_idx=[0])
+
+    for conv in (block.conv1, block.conv2):  # as the installer does: conv weights only
+        conv.to(memory_format=torch.channels_last_3d)
+    cfg = fp.FastPathConfig(channels_last=True)
+    for module in (block, block.norm1, block.norm2):
+        setattr(module, fp.CFG_ATTR, cfg)
+    seen: list[torch.Tensor | None] = []
+    original = fp.rms_norm_fastpath
+
+    def spy(norm, x, *, silu=False, bias=None):
+        seen.append(bias)
+        return original(norm, x, silu=silu, bias=bias)
+
+    monkeypatch.setattr(fp, "rms_norm_fastpath", spy)
+    x_cl = x.contiguous(memory_format=torch.channels_last_3d)
+    actual = fp.residual_block_forward(block, x_cl, feat_cache=[None, None], feat_idx=[0])
+    assert seen[0] is None and seen[1] is block.conv1.bias, "norm2 should receive conv1's un-added bias"
+    assert actual.is_contiguous(memory_format=torch.channels_last_3d)
+    if dtype is torch.float32:
+        torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2)  # channels-last convs may run TF32
+    else:
+        # Fast-math norm plus channels-last cuDNN accumulation order: a few outputs differ by one bf16 ulp.
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+
+
+def test_rms_norm_scale_declines_bad_operands() -> None:
+    x = torch.randn(1, 8, 16, device="cuda", dtype=torch.bfloat16)
+    denom = torch.ones(1, 1, 16, device="cuda")
+    gamma = torch.ones(8, device="cuda", dtype=torch.bfloat16)
+    assert rn.rms_norm_scale(x, denom, gamma, 1.0) is not None
+    assert rn.rms_norm_scale(x.cpu(), denom.cpu(), gamma.cpu(), 1.0) is None
+    assert rn.rms_norm_scale(x, denom.to(torch.bfloat16), gamma, 1.0) is None
+    assert rn.rms_norm_scale(x, denom, gamma.float(), 1.0) is None
+    assert rn.rms_norm_scale(x.transpose(1, 2), denom, gamma, 1.0) is None
+    assert rn.rms_norm_scale(x, torch.ones(1, 1, 8, device="cuda"), gamma, 1.0) is None
+
+
+# --------------------------------------------------------------------------- #
+# Data movement
+# --------------------------------------------------------------------------- #
+
+
+def _dense_5d(shape, dtype: torch.dtype, channels_last: bool) -> torch.Tensor:
+    tensor = torch.randn(shape, device="cuda", dtype=dtype)
+    if channels_last:
+        return tensor.contiguous(memory_format=torch.channels_last_3d)
+    return tensor.contiguous()
+
+
+def _reference_cat_pad(x: torch.Tensor, cache: torch.Tensor | None, padding) -> torch.Tensor:
+    reference_padding = list(padding)
+    if cache is not None:
+        x = torch.cat([cache, x], dim=2)
+        reference_padding[4] -= cache.shape[2]
+    return F.pad(x, reference_padding)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize(
+    "channels,frames,height,width,cache_frames,padding",
+    [
+        (96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)),
+        (96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)),
+        (96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
+        (64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)),
+        (48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
+        # Non-power-of-two channels and a ragged number of rows per program.
+        (12, 3, 7, 9, 1, (1, 1, 1, 1, 2, 0)),
+        # Width wider than one block so the in-kernel width loop runs.
+        (8, 2, 3, 1300, 2, (1, 1, 1, 1, 2, 0)),
+        # Cosmos3 stage-3 shape, scaled down in height.
+        (256, 4, 8, 640, 2, (1, 1, 1, 1, 2, 0)),
+    ],
+)
+def test_cat_pad_5d_is_bitwise_and_layout_exact(
+    dtype, layout, channels, frames, height, width, cache_frames, padding
+) -> None:
+    torch.manual_seed(0)
+    channels_last = layout == "channels_last"
+    if layout == "frame_major":
+        x = torch.randn((1, frames, channels, height, width), device="cuda", dtype=dtype).permute(0, 2, 1, 3, 4)
+    else:
+        x = _dense_5d((1, channels, frames, height, width), dtype, channels_last)
+    cache = None
+    if cache_frames:
+        pad_height, pad_width = padding[2], padding[0]
+        buffer = _dense_5d(
+            (1, channels, cache_frames, height + 2 * pad_height, width + 2 * pad_width), dtype, channels_last
+        )
+        cache = buffer[:, :, :, pad_height : pad_height + height, pad_width : pad_width + width]
+
+    reference = _reference_cat_pad(x, cache, padding)
+    output = dm.cat_pad_5d(x, cache, padding)
+    assert output is not None
+    assert output.stride() == reference.stride()
+    assert torch.equal(output, reference)
+
+    pair = dm.cat_pad_5d(x, cache, padding, keep_cache_frames=CACHE_T)
+    assert pair is not None
+    output_with_cache, next_cache = pair
+    assert torch.equal(output_with_cache, reference)
+    pad_height, pad_width = padding[2], padding[0]
+    keep_frames = min(CACHE_T, reference.shape[2])
+    expected_cache = reference[:, :, -keep_frames:, pad_height : pad_height + height, pad_width : pad_width + width]
+    assert torch.equal(next_cache, expected_cache)
+    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d) is channels_last
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "frame_major"])
+@pytest.mark.parametrize("cache_layout", ["contiguous", "channels_last"])
+@pytest.mark.parametrize(
+    "channels,frames,height,width,cache_frames",
+    [
+        (96, 1, 10, 14, 0),
+        (96, 1, 10, 14, 2),
+        (12, 3, 7, 9, 1),
+        (8, 2, 3, 1300, 2),
+        (256, 4, 8, 640, 2),
+        (64, 4, 70, 70, 2),
+    ],
+)
+def test_cat_pad_5d_assembles_channels_last_output_from_channels_first_input(
+    dtype, layout, cache_layout, channels, frames, height, width, cache_frames
+) -> None:
+    """The conv_out path: channels-first activations, channels-last padded input, cache in the input layout."""
+    torch.manual_seed(0)
+    padding = (1, 1, 1, 1, 2, 0)
+    if layout == "frame_major":
+        x = torch.randn((1, frames, channels, height, width), device="cuda", dtype=dtype).permute(0, 2, 1, 3, 4)
+    else:
+        x = torch.randn((1, channels, frames, height, width), device="cuda", dtype=dtype)
+    cache = None
+    if cache_frames:
+        cache = _dense_5d((1, channels, cache_frames, height, width), dtype, cache_layout == "channels_last")
+    reference = _reference_cat_pad(x, cache, padding)
+
+    output = dm.cat_pad_5d(x, cache, padding, channels_last_output=True)
+    assert output is not None
+    assert output.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(output, reference)
+
+    pair = dm.cat_pad_5d(x, cache, padding, keep_cache_frames=CACHE_T, channels_last_output=True)
+    assert pair is not None
+    output_with_cache, next_cache = pair
+    assert output_with_cache.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(output_with_cache, reference)
+    keep_frames = min(CACHE_T, reference.shape[2])
+    assert next_cache.is_contiguous()  # follows x's layout, not the output's
+    assert torch.equal(next_cache, reference[:, :, -keep_frames:, 1 : 1 + height, 1 : 1 + width])
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("layout", ["contiguous", "frame_major"])
+def test_conv_out_channels_last_verdict_is_recorded_and_exact(layout: str) -> None:
+    torch.manual_seed(0)
+    conv = WanCausalConv3d(64, 12, 3, padding=1).to(device="cuda", dtype=torch.bfloat16)
+    if layout == "frame_major":
+        chunk = torch.randn(1, 4, 64, 24, 40, device="cuda", dtype=torch.bfloat16).permute(0, 2, 1, 3, 4)
+    else:
+        chunk = torch.randn(1, 64, 4, 24, 40, device="cuda", dtype=torch.bfloat16)
+    cache = torch.randn(1, 64, 2, 24, 40, device="cuda", dtype=torch.bfloat16)
+    reference = WanCausalConv3d.forward(conv, chunk, cache)
+    fast_cache = [cache.clone()]
+    out = fp._run_conv_out_channels_last(conv, chunk, fast_cache, 0)
+    assert out is not None  # the verdict decides which formulation's result is returned, both are exact
+    assert out.is_contiguous()
+    assert _bits_equal(out, reference)
+    verdicts = fp._CONV_OUT_LAYOUT_VERDICTS[conv]
+    assert len(verdicts) == 1
+    assert torch.equal(fast_cache[0], torch.cat([cache, chunk], dim=2)[:, :, -CACHE_T:])
+    assert fast_cache[0].is_contiguous()
+    again = fp._run_conv_out_channels_last(conv, chunk, [cache.clone()], 0)
+    if next(iter(verdicts.values())):
+        assert again is not None and _bits_equal(again, reference)
+    else:
+        assert again is None  # routed back to the standard path
+    # Channels-last activations are not this path's business.
+    assert (
+        fp._run_conv_out_channels_last(conv, chunk.contiguous(memory_format=torch.channels_last_3d), [None], 0) is None
+    )
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize(
+    "channels,frames,height,width,cache_frames,pad_front",
+    [(96, 1, 10, 14, 0, 2), (96, 1, 10, 14, 1, 2), (96, 4, 10, 14, 2, 2), (12, 3, 7, 9, 2, 2), (256, 4, 8, 640, 2, 2)],
+)
+def test_cat_time_5d_is_bitwise_and_layout_exact(
+    dtype, layout, channels, frames, height, width, cache_frames, pad_front
+):
+    torch.manual_seed(0)
+    channels_last = layout == "channels_last"
+    if layout == "frame_major":
+        x = torch.randn((1, frames, channels, height, width), device="cuda", dtype=dtype).permute(0, 2, 1, 3, 4)
+    else:
+        x = _dense_5d((1, channels, frames, height, width), dtype, channels_last)
+    cache = _dense_5d((1, channels, cache_frames, height, width), dtype, channels_last) if cache_frames else None
+
+    parts = [torch.zeros_like(x[:, :, :1]).expand(-1, -1, pad_front - cache_frames, -1, -1)]
+    if cache is not None:
+        parts.append(cache)
+    parts.append(x)
+    reference = torch.cat(parts, dim=2)
+    if channels_last:
+        reference = reference.contiguous(memory_format=torch.channels_last_3d)
+    else:
+        reference = reference.contiguous()
+
+    pair = dm.cat_time_5d(x, cache, pad_front, keep_cache_frames=CACHE_T)
+    assert pair is not None
+    assembled, next_cache = pair
+    assert assembled.stride() == reference.stride()
+    assert torch.equal(assembled, reference)
+    assert torch.equal(next_cache, reference[:, :, -CACHE_T:])
+    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d) is channels_last
+    assert torch.equal(dm.cat_time_5d(x, cache, pad_front), reference)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("cache_frames", [0, 1, 2])
+def test_cat_time_5d_exceeds_grid_y_limit(monkeypatch: pytest.MonkeyPatch, cache_frames: int) -> None:
+    # Smaller blocks reproduce the large-plane launch failure without allocating
+    # 4K feature maps. The final block is partial, exercising the plane boundary.
+    monkeypatch.setattr(dm, "_PLANE_BLOCK", 32)
+    batch, channels, frames, height, width = 2, 8, 1, 257, 1021
+    plane_size = channels * height * width
+    assert dm.triton.cdiv(plane_size, dm._PLANE_BLOCK) > 65_535
+    assert plane_size % dm._PLANE_BLOCK != 0
+
+    x = _dense_5d((batch, channels, frames, height, width), torch.bfloat16, True)
+    cache = _dense_5d((batch, channels, cache_frames, height, width), x.dtype, True) if cache_frames else None
+    parts = [torch.zeros_like(x).expand(-1, -1, 2 - cache_frames, -1, -1)]
+    if cache is not None:
+        parts.append(cache)
+    parts.append(x)
+    reference = torch.cat(parts, dim=2).contiguous(memory_format=torch.channels_last_3d)
+
+    pair = dm.cat_time_5d(x, cache, pad_front=2, keep_cache_frames=CACHE_T)
+    assert pair is not None
+    assembled, next_cache = pair
+    assert _bits_equal(assembled, reference)
+    assert assembled.stride() == reference.stride()
+    assert _bits_equal(next_cache, reference[:, :, -CACHE_T:])
+    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d)
+    assert _bits_equal(dm.cat_time_5d(x, cache, pad_front=2), reference)
+
+
+@torch.no_grad()
+def test_cached_conv_spatial_padding_verdict_is_recorded_and_exact() -> None:
+    torch.manual_seed(0)
+    conv = WanCausalConv3d(64, 64, 3, padding=1).to(device="cuda", dtype=torch.bfloat16)
+    for layout in (torch.contiguous_format, torch.channels_last_3d):
+        if layout is torch.channels_last_3d:
+            conv.to(memory_format=torch.channels_last_3d)
+        chunk = torch.randn(1, 64, 4, 24, 40, device="cuda", dtype=torch.bfloat16).contiguous(memory_format=layout)
+        cache = torch.randn(1, 64, 2, 24, 40, device="cuda", dtype=torch.bfloat16).contiguous(memory_format=layout)
+        reference_cache = [cache.clone()]
+        reference = WanCausalConv3d.forward(conv, chunk, cache)  # upstream pre-padded path
+        fast_cache = [cache.clone()]
+        out = fp._run_cached_causal_conv(conv, chunk, fast_cache, 0)
+        verdicts = fp._SPATIAL_PAD_VERDICTS[conv]
+        assert len(verdicts) >= 1
+        assert _bits_equal(out, reference)  # exact whichever path the verdict selected
+        again = fp._run_cached_causal_conv(conv, chunk, [cache.clone()], 0)
+        assert _bits_equal(again, reference)
+        assert torch.equal(fast_cache[0], torch.cat([cache, chunk], dim=2)[:, :, -CACHE_T:])
+        del reference_cache
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("main_layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize("with_bias", [False, True])
+@pytest.mark.parametrize(
+    "in_channels,out_channels,frames,height,width,factor_temporal,factor_spatial,drop_first",
+    [
+        # Cosmos3 / Wan2.2 block shapes: repeats 8 (same channels), 4 (halved), 2 (spatial only, halved).
+        (128, 64, 1, 10, 14, 2, 2, False),
+        (128, 64, 1, 10, 14, 2, 2, True),
+        (64, 32, 2, 10, 14, 1, 2, False),
+        (64, 64, 4, 10, 14, 1, 2, False),
+        # More output than input channels (repeats 8 with a 2x channel increase) and a 4x spatial factor.
+        (32, 64, 2, 10, 14, 1, 2, False),
+        (16, 16, 1, 5, 6, 1, 4, False),
+        # Odd sizes and a row wider than one column block.
+        (12, 12, 3, 7, 9, 2, 2, True),
+        (8, 8, 1, 3, 700, 1, 2, False),
+    ],
+)
+def test_dup_up3d_add_is_bitwise_and_layout_exact(
+    dtype,
+    main_layout,
+    with_bias,
+    in_channels,
+    out_channels,
+    frames,
+    height,
+    width,
+    factor_temporal,
+    factor_spatial,
+    drop_first,
+) -> None:
+    torch.manual_seed(0)
+    repeats = out_channels * factor_temporal * factor_spatial * factor_spatial // in_channels
+    source = _dense_5d((1, in_channels, frames, height, width), dtype, main_layout == "channels_last")
+    if main_layout == "frame_major":
+        # The previous block's output is frame-major too in the lossless level.
+        source = torch.randn((1, frames, in_channels, height, width), device="cuda", dtype=dtype)
+        source = source.permute(0, 2, 1, 3, 4)
+    out_frames = frames * factor_temporal - (factor_temporal - 1 if drop_first else 0)
+    shape = (1, out_channels, out_frames, height * factor_spatial, width * factor_spatial)
+    if main_layout == "channels_last":
+        main = _dense_5d(shape, dtype, True)
+    elif main_layout == "frame_major":
+        main = torch.randn((shape[0], shape[2], shape[1], shape[3], shape[4]), device="cuda", dtype=dtype)
+        main = main.permute(0, 2, 1, 3, 4)
+    else:
+        main = _dense_5d(shape, dtype, False)
+
+    shortcut = source.repeat_interleave(repeats, dim=1)
+    shortcut = shortcut.view(1, out_channels, factor_temporal, factor_spatial, factor_spatial, frames, height, width)
+    shortcut = shortcut.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+    shortcut = shortcut.view(shape[0], shape[1], frames * factor_temporal, shape[3], shape[4])
+    if drop_first:
+        shortcut = shortcut[:, :, factor_temporal - 1 :]
+    bias = None
+    biased_main = main
+    if with_bias:
+        bias = (torch.randn(out_channels, device="cuda") * 0.1).to(dtype)
+        biased_main = main.clone()
+        biased_main.add_(bias.view(1, -1, 1, 1, 1))  # ATen's conv bias add
+    reference = biased_main + shortcut
+
+    output = dm.dup_up3d_add(main, source, factor_temporal, factor_spatial, repeats, drop_first, main_bias=bias)
+    assert output is not None
+    assert _same_dense_strides(output, reference)
+    assert torch.equal(output, reference)
+
+
+def test_dup_up3d_add_declines_unsupported_cases() -> None:
+    main = torch.randn(1, 8, 2, 10, 14, device="cuda", dtype=torch.bfloat16)
+    source = torch.randn(1, 32, 2, 5, 7, device="cuda", dtype=torch.bfloat16)
+    # repeats (1) below the spatial factor: the source channel would change along a row.
+    assert dm.dup_up3d_add(main, source, 1, 2, 1, False) is None
+    source = torch.randn(1, 8, 2, 5, 7, device="cuda", dtype=torch.bfloat16)
+    assert dm.dup_up3d_add(main, source, 1, 2, 4, False) is not None
+    # Mixed layouts, wrong bias shape, wrong main shape.
+    assert dm.dup_up3d_add(main.contiguous(memory_format=torch.channels_last_3d), source, 1, 2, 4, False) is None
+    assert dm.dup_up3d_add(main, source, 1, 2, 4, False, main_bias=torch.ones(7, device="cuda")) is None
+    assert dm.dup_up3d_add(main[:, :, :1], source, 1, 2, 4, False) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("layout", ["contiguous", "channels_last", "frame_major"])
+@pytest.mark.parametrize(
+    "frames,channels,height,width",
+    [
+        (1, 96, 10, 14),
+        (4, 64, 10, 14),
+        # Non-power-of-two channels and ragged rows/columns for every block size.
+        (3, 12, 7, 9),
+        # Width wider than one column block so the in-kernel loop runs.
+        (2, 8, 3, 1300),
+        # Cosmos3 720p stage-3 input, scaled down in height.
+        (4, 256, 8, 640),
+        # Channels equal to and above the channels-last block: in-kernel channel loop.
+        (2, 1024, 5, 160),
+        (1, 1100, 3, 20),
+    ],
+)
+def test_upsample_nearest_2x_is_bitwise_and_layout_exact(dtype, layout, frames, channels, height, width) -> None:
+    torch.manual_seed(0)
+    if layout == "frame_major":
+        # The strided ``(t, c, h, w)`` view ``WanResample`` hands the upsampler in the lossless level.
+        x = fp._merge_batch_and_frames(torch.randn((1, channels, frames, height, width), device="cuda", dtype=dtype))
+        assert frames == 1 or not x.is_contiguous()
+    elif layout == "channels_last":
+        x = torch.randn((frames, channels, height, width), device="cuda", dtype=dtype)
+        x = x.contiguous(memory_format=torch.channels_last)
+    else:
+        x = torch.randn((frames, channels, height, width), device="cuda", dtype=dtype)
+    for mode in ("nearest-exact", "nearest"):
+        reference = F.interpolate(x, scale_factor=(2.0, 2.0), mode=mode)
+        output = up.upsample_nearest_2x(x)
+        assert output is not None
+        assert output.stride() == reference.stride()
+        assert torch.equal(output, reference)
+
+
+def test_upsample_nearest_2x_declines_unsupported_inputs() -> None:
+    x = torch.randn(2, 8, 6, 10, device="cuda", dtype=torch.bfloat16)
+    assert up.upsample_nearest_2x(x) is not None
+    assert up.upsample_nearest_2x(x.cpu()) is None
+    assert up.upsample_nearest_2x(x.unsqueeze(2)) is None
+    assert up.upsample_nearest_2x(x.transpose(2, 3)) is None
+    assert up.upsample_nearest_2x(x[..., ::2]) is None
+    assert up.upsample_nearest_2x(x.double()) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_upsample_forward_matches_wan_upsample(dtype: torch.dtype) -> None:
+    torch.manual_seed(0)
+    x = torch.randn(3, 64, 10, 14, device="cuda", dtype=dtype)
+    for layout in (torch.contiguous_format, torch.channels_last):
+        x = x.contiguous(memory_format=layout)
+        module = WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact")
+        expected = WanUpsample.forward(module, x)
+        actual = fp.upsample_forward(module, x)
+        assert actual.stride() == expected.stride()
+        assert torch.equal(actual, expected)
+        # Anything but a nearest 2x upsample takes the reference path.
+        other = WanUpsample(scale_factor=(3.0, 3.0), mode="nearest-exact")
+        assert torch.equal(fp.upsample_forward(other, x), WanUpsample.forward(other, x))
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("mode", ["upsample2d", "upsample3d"])
+def test_resample_forward_can_leave_conv_bias_pending(dtype: torch.dtype, mode: str) -> None:
+    torch.manual_seed(0)
+    resample = WanResample(32, mode=mode).eval().to(device="cuda", dtype=dtype)
+    setattr(resample, fp.CFG_ATTR, fp.FastPathConfig())
+    x = torch.randn(1, 32, 2, 10, 14, device="cuda", dtype=dtype)
+    cache = [torch.randn(1, 32, 2, 10, 14, device="cuda", dtype=dtype)] if mode == "upsample3d" else [None]
+    expected = WanResample.forward(
+        resample, x, feat_cache=[c.clone() if c is not None else None for c in cache], feat_idx=[0]
+    )
+    out, bias = fp.resample_forward(
+        resample, x, [c.clone() if c is not None else None for c in cache], [0], return_bias=True
+    )
+    assert bias is resample.resample[1].bias
+    assert out.stride() == expected.stride()
+    assert _bits_equal(out + bias.view(1, -1, 1, 1, 1), expected)
+    plain = fp.resample_forward(resample, x, [c.clone() if c is not None else None for c in cache], [0])
+    assert _bits_equal(plain, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+@pytest.mark.parametrize("force_prepad", [False, True])
+def test_cached_conv_deferred_bias_matches_autocast(dtype: torch.dtype, force_prepad: bool, monkeypatch) -> None:
+    """FP32 parameters and inputs must defer the bias in the convolution's output dtype."""
+    torch.manual_seed(8)
+    conv = WanCausalConv3d(32, 32, 3, padding=1).cuda().eval()
+    conv.bias.copy_(torch.randn_like(conv.bias))
+    if force_prepad:
+        monkeypatch.setattr(dm, "cat_time_5d", lambda *args, **kwargs: None)
+    cache = [None]
+    with torch.autocast("cuda", dtype=dtype):
+        for _ in range(2):
+            x = torch.randn(1, 32, 2, 6, 10, device="cuda")
+            expected = conv(x, cache[0])
+            out, bias = fp._run_cached_causal_conv(conv, x, cache, 0, return_bias=True)
+            assert out.dtype == dtype and bias is not None and bias.dtype == dtype
+            actual = dm.add_bias_residual(out, bias, torch.zeros_like(out), None)
+            assert actual is not None
+            assert _bits_equal(actual.contiguous(), expected.contiguous())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+@pytest.mark.parametrize("channels_last", [False, True])
+def test_residual_block_deferred_bias_matches_autocast(dtype: torch.dtype, channels_last: bool) -> None:
+    """Exercise conv1/conv2 and a non-identity shortcut with FP32 weights under autocast."""
+    torch.manual_seed(8)
+    block = WanResidualBlock(32, 64, dropout=0.0).cuda().eval()
+    for conv in (block.conv1, block.conv2, block.conv_shortcut):
+        conv.bias.copy_(torch.randn_like(conv.bias))
+        if channels_last:
+            conv.to(memory_format=torch.channels_last_3d)
+    x = torch.randn(1, 32, 2, 6, 10, device="cuda")
+    with torch.autocast("cuda", dtype=dtype):
+        expected = block(x, feat_cache=[None, None], feat_idx=[0])
+        setattr(block, fp.CFG_ATTR, fp.FastPathConfig(channels_last=channels_last))
+        actual = fp.residual_block_forward(block, x, feat_cache=[None, None], feat_idx=[0])
+    assert actual.dtype == expected.dtype == dtype
+    assert _bits_equal(actual.contiguous(), expected.contiguous())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", LOW_PRECISION)
+@pytest.mark.parametrize("mode", ["upsample2d", "upsample3d"])
+def test_resample_deferred_bias_matches_autocast(dtype: torch.dtype, mode: str) -> None:
+    """The upsampler's FP32 bias must be rounded before its deferred addition."""
+    torch.manual_seed(8)
+    resample = WanResample(32, mode=mode).cuda().eval()
+    resample.resample[1].bias.copy_(torch.randn_like(resample.resample[1].bias))
+    reference_cache = [None]
+    candidate_cache = [None]
+    with torch.autocast("cuda", dtype=dtype):
+        for _ in range(2):
+            x = torch.randn(1, 32, 1, 6, 10, device="cuda")
+            expected = resample(x, feat_cache=reference_cache, feat_idx=[0])
+            out, bias = fp.resample_forward(resample, x, candidate_cache, [0], return_bias=True)
+            assert out.dtype == dtype and bias is not None and bias.dtype == dtype
+            actual = out + bias.view(1, -1, 1, 1, 1)
+            assert _bits_equal(actual.contiguous(), expected.contiguous())
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("with_bias_x", [True, False])
+@pytest.mark.parametrize("with_bias_h", [True, False])
+def test_add_bias_residual_matches_aten_rounding(dtype: torch.dtype, with_bias_x: bool, with_bias_h: bool) -> None:
+    torch.manual_seed(0)
+    shape = (1, 256, 4, 45, 80)
+    x = torch.randn(shape, device="cuda").to(dtype)
+    h = torch.randn(shape, device="cuda").to(dtype)
+    bias_x = torch.randn(256, device="cuda").to(dtype) if with_bias_x else None
+    bias_h = torch.randn(256, device="cuda").to(dtype) if with_bias_h else None
+    expected_x = x.clone()
+    expected_h = h.clone()
+    if bias_x is not None:
+        expected_x.add_(bias_x.view(1, -1, 1, 1, 1))
+    if bias_h is not None:
+        expected_h.add_(bias_h.view(1, -1, 1, 1, 1))
+    expected = expected_x + expected_h
+    actual = dm.add_bias_residual(x, bias_x, h, bias_h)
+    assert actual is not None
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected)
+    x_cl = x.contiguous(memory_format=torch.channels_last_3d)
+    h_cl = h.contiguous(memory_format=torch.channels_last_3d)
+    actual_cl = dm.add_bias_residual(x_cl, bias_x, h_cl, bias_h)
+    assert actual_cl is not None and actual_cl.is_contiguous(memory_format=torch.channels_last_3d)
+    assert _bits_equal(actual_cl.contiguous(), expected)
+    assert dm.add_bias_residual(x_cl, bias_x, h, bias_h) is None
+    assert dm.add_bias_residual(x.permute(0, 2, 1, 3, 4).contiguous().permute(0, 2, 1, 3, 4), bias_x, h, bias_h) is None
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("channels_last", [False, True])
+@pytest.mark.parametrize("temporal_only", [False, True])
+@pytest.mark.parametrize("start", [None, "Rep"])
+def test_cached_conv_chunk_loop_matches_reference_including_mixed_paths(
+    monkeypatch, channels_last: bool, temporal_only: bool, start
+) -> None:
+    torch.manual_seed(0)
+    channels = 64
+    if temporal_only:
+        conv = WanCausalConv3d(channels, 2 * channels, (3, 1, 1), padding=(1, 0, 0))
+    else:
+        conv = WanCausalConv3d(channels, channels, 3, padding=1)
+    conv = conv.to(device="cuda", dtype=torch.float32)
+    chunks = [_dense_5d((1, channels, 1, 10, 14), torch.float32, channels_last) for _ in range(4)]
+
+    decline = lambda *a, **k: None  # noqa: E731
+    original_time, original_pad = fp.dm.cat_time_5d, fp.dm.cat_pad_5d
+
+    def run(mode: list[str]):
+        cache = [start]
+        outputs = []
+        for chunk, m in zip(chunks, mode, strict=True):
+            monkeypatch.setattr(fp.dm, "cat_time_5d", decline if m in ("reference", "padded") else original_time)
+            monkeypatch.setattr(fp.dm, "cat_pad_5d", decline if m == "reference" else original_pad)
+            outputs.append(fp._run_cached_causal_conv(conv, chunk, cache, 0))
+        monkeypatch.setattr(fp.dm, "cat_time_5d", original_time)
+        monkeypatch.setattr(fp.dm, "cat_pad_5d", original_pad)
+        return outputs
+
+    reference = run(["reference"] * 4)
+    for pattern in (
+        ["time"] * 4,
+        ["padded"] * 4,
+        ["time", "reference", "padded", "time"],
+        ["reference", "padded", "time", "reference"],
+    ):
+        fused = run(pattern)
+        for out, ref in zip(fused, reference, strict=True):
+            assert torch.equal(out, ref), pattern
+
+
+# --------------------------------------------------------------------------- #
+# Whole decoder
+# --------------------------------------------------------------------------- #
+
+TINY_RESIDUAL = dict(
+    base_dim=8,
+    decoder_base_dim=16,
+    z_dim=8,
+    dim_mult=[1, 2, 4, 4],
+    num_res_blocks=2,
+    temperal_downsample=[False, True, True],
+    is_residual=True,
+    patch_size=2,
+    in_channels=12,
+    out_channels=12,
+)
+TINY_WAN21 = dict(base_dim=8, z_dim=4, dim_mult=[1, 2], num_res_blocks=1, temperal_downsample=[False, True])
+
+
+def _build_pair(config: dict, dtype: torch.dtype):
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**config).eval().to(device="cuda", dtype=dtype)
+    candidate = AutoencoderKLWan(**config).eval().to(device="cuda", dtype=dtype)
+    candidate.load_state_dict(reference.state_dict())
+    return reference, candidate
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("config", [TINY_RESIDUAL, TINY_WAN21], ids=["residual_patch2", "wan21"])
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+@pytest.mark.parametrize("frames", [1, 2, 5])
+def test_lossless_decoder_is_bitwise_exact_on_cuda(config: dict, dtype: torch.dtype, frames: int) -> None:
+    reference, candidate = _build_pair(config, dtype)
+    report = install_wan_vae_fastpath(candidate, level="lossless")
+    assert report.installed, report
+    torch.manual_seed(1)
+    latents = torch.randn(1, config["z_dim"], frames, 6, 10, device="cuda").to(dtype)
+    autocast = dtype is not torch.float32
+    with torch.autocast("cuda", dtype=dtype, enabled=autocast):
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", ALL_DTYPES)
+def test_omni_wan_vae_decode_is_bitwise_exact_on_cuda(dtype: torch.dtype) -> None:
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**TINY_RESIDUAL).eval().to(device="cuda", dtype=dtype)
+    candidate = OmniAutoencoderKLWan(**TINY_RESIDUAL).eval().to(device="cuda", dtype=dtype)
+    candidate.load_state_dict(reference.state_dict())
+    install_wan_vae_fastpath(candidate)
+    latents = torch.randn(1, TINY_RESIDUAL["z_dim"], 5, 6, 10, device="cuda").to(dtype)
+    with torch.autocast("cuda", dtype=dtype, enabled=dtype is not torch.float32):
+        expected = reference.decode(latents, return_dict=False)[0]
+    actual = candidate.decode(latents, return_dict=False)[0]  # applies its own autocast context
+    assert actual.stride() == expected.stride()
+    assert _bits_equal(actual, expected)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize(
+    ("dtype", "allow_tf32", "tolerance"),
+    [
+        pytest.param(torch.bfloat16, True, 3e-2, id="bf16"),
+        # With TF32 disabled the only difference is cuDNN's channels-last algorithm
+        # choice and the norm kernel's reduction order.
+        pytest.param(torch.float32, False, 2e-4, id="fp32"),
+        # PyTorch's default allows cuDNN to run fp32 convolutions in TF32; the
+        # channels-last algorithms use it while the NCDHW reference may not, so
+        # the deviation is dominated by TF32's 10-bit mantissa.
+        pytest.param(torch.float32, True, 1e-2, id="fp32-tf32"),
+    ],
+)
+def test_channels_last_decoder_keeps_layout_and_stays_close(
+    dtype: torch.dtype, allow_tf32: bool, tolerance: float
+) -> None:
+    reference, candidate = _build_pair(TINY_RESIDUAL, dtype)
+    report = install_wan_vae_fastpath(candidate, level="channels_last")
+    assert report.installed and report.channels_last
+
+    violations: list[str] = []
+
+    def audit(name: str):
+        def hook(module: nn.Module, inputs) -> None:
+            x = inputs[0]
+            if x.dim() == 5 and x.shape[1] > 1 and not x.is_contiguous(memory_format=torch.channels_last_3d):
+                violations.append(f"{name}: {tuple(x.shape)} strides={x.stride()}")
+            if x.dim() == 4 and x.shape[1] > 1 and not x.is_contiguous(memory_format=torch.channels_last):
+                violations.append(f"{name}: {tuple(x.shape)} strides={x.stride()}")
+
+        return hook
+
+    handles = [
+        module.register_forward_pre_hook(audit(name))
+        for name, module in candidate.decoder.named_modules()
+        if isinstance(module, (nn.Conv2d, nn.Conv3d))
+    ]
+    torch.manual_seed(1)
+    latents = torch.randn(1, TINY_RESIDUAL["z_dim"], 3, 6, 10, device="cuda").to(dtype)
+    with (
+        torch.backends.cudnn.flags(enabled=True, allow_tf32=allow_tf32),
+        torch.autocast("cuda", dtype=dtype, enabled=dtype is not torch.float32),
+    ):
+        expected = reference.decode(latents, return_dict=False)[0]
+        actual = candidate.decode(latents, return_dict=False)[0]
+    for handle in handles:
+        handle.remove()
+    assert not violations, "\n".join(violations)
+    assert actual.shape == expected.shape
+    difference = (actual.float() - expected.float()).abs()
+    error = difference.max().item()
+    mse = (difference**2).mean().item()
+    psnr = float("inf") if mse == 0 else 10 * torch.log10(torch.tensor(4.0 / mse)).item()
+    assert error <= tolerance, f"max_abs_diff={error:.3e} psnr={psnr:.2f} dB tolerance={tolerance}"