Merge main into dev/vllm-align for vLLM 0.31 release

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/diffusion/models/auk/test_auk_cudagraph.py

@@ -61,7 +61,10 @@ def test_single_request_graph_wrapper_cpu_falls_back_to_eager(cfg_strength: floa
 
     # The conditioning is prepared once per request and reused by every step.
     assert prepare_spy.call_count == 1
-    assert prepare_spy.call_args.args[-1] is (cfg_strength >= 1e-5)
+    # _prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
+    assert prepare_spy.call_args.args[6] is (cfg_strength >= 1e-5)
+    # The time grid reaches prepare, so the adaLN modulations are computed once.
+    assert prepare_spy.call_args.args[7].numel() == 2
     assert step_spy.call_count == 2
     capture_spy.assert_not_called()
     assert not wrapper._cache
@@ -84,7 +87,8 @@ def test_graph_inputs_use_bounded_length_buckets() -> None:
     assert padded[4].shape == (1, 100, 4)
     assert padded[5].shape == (1, 100)
     assert [mask.sum().item() for mask in (padded[1], padded[3], padded[5])] == [65, 65, 51]
-    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False)
+    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False, 0)
+    assert wrapper._key(padded[0], padded[2], padded[4], False, 32) == (96, 96, 100, False, 32)
     assert wrapper.max_graphs == 32
 
 
@@ -153,18 +157,21 @@ def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_streng
             inputs["ref"],
             inputs["ref_mask"],
         )
-        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5)
+        # sample_latents passes the grid, so the key carries its step count (two steps here).
+        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5, 2)
         assert key in wrapper._cache
 
         # The static context was refreshed for this request's conditioning.
         entry = wrapper._cache[key]
-        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5)
+        grid = torch.tensor([0.0, 0.4], device="cuda")
+        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5, grid)
         for static, want in zip(entry.static_ctx.tensors(), fresh.tensors(), strict=True):
             if want is None:
                 assert static is None
             else:
                 torch.testing.assert_close(static, want)
         torch.testing.assert_close(entry.static_timestep, torch.tensor(0.4, device="cuda"))
+        assert int(entry.static_step) == 1
 
     assert len(wrapper._cache) == 1
 

tests/diffusion/models/auk/test_auk_dit_fastpaths.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The AuK DiT fast paths keep the reference numerics.
+
+Two paths are checked: the grouped conv position embedding as one batched
+GEMM, and the adaLN modulations precomputed for a whole time grid, eager and
+under the per-step CUDA graph.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.auk import auk_transformer as T
+from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, build_time_grid
+from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_dit(device: str = "cpu") -> AuKTransformer:
+    torch.manual_seed(5)
+    return (
+        AuKTransformer(
+            dim=32,
+            heads=2,
+            dim_head=16,
+            ff_mult=2,
+            latent_dim=4,
+            text_hidden_dim=8,
+            num_layers=2,
+            num_single_layers=2,
+        )
+        .eval()
+        .to(device)
+    )
+
+
+def _reference_conv_pos_embed(module: T.ConvPosEmbedding, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
+    """The channels-first conv stack as the reference implementation writes it."""
+    keep = None if mask is None else mask.unsqueeze(1)
+    x = x.transpose(1, 2)
+    if keep is not None:
+        x = x.masked_fill(~keep, 0.0)
+    for layer in module.conv1d:
+        x = layer(x)
+        if keep is not None and isinstance(layer, torch.nn.Conv1d):
+            x = x.masked_fill(~keep, 0.0)
+    return x.transpose(1, 2)
+
+
+@torch.inference_mode()
+def test_grouped_conv_as_batched_gemm_matches_conv1d() -> None:
+    torch.manual_seed(0)
+    conv = torch.nn.Conv1d(64, 64, 31, groups=16, padding=15)
+    x = torch.randn(2, 37, 64)
+    expected = conv(x.transpose(1, 2)).transpose(1, 2)
+    torch.testing.assert_close(T._grouped_conv1d_btc(x, conv), expected, atol=1e-5, rtol=1e-5)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("masked", [False, True])
+def test_conv_pos_embedding_matches_the_channels_first_stack(masked: bool) -> None:
+    torch.manual_seed(1)
+    module = T.ConvPosEmbedding(64).eval()
+    x = torch.randn(2, 37, 64)
+    mask = None
+    if masked:
+        mask = torch.ones(2, 37, dtype=torch.bool)
+        mask[1, 30:] = False
+    torch.testing.assert_close(module(x, mask), _reference_conv_pos_embed(module, x, mask), atol=1e-5, rtol=1e-5)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_infer", [False, True])
+def test_modulation_table_matches_per_step_modulation(cfg_infer: bool) -> None:
+    dit = _make_dit()
+    text = torch.randn(1, 5, 8)
+    ref = torch.randn(1, 3, 4)
+    x = torch.randn(1, 6, 4)
+    grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cpu")
+    plain = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer)
+    tabled = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer, timesteps=grid[:-1])
+    assert tabled.modulation is not None and tabled.modulation.shape[0] == 4
+    for i in range(4):
+        expected = dit.step(x, grid[i], plain)
+        torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=i), expected, atol=1e-5, rtol=1e-5)
+        index = torch.tensor(i)
+        torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=index), expected, atol=1e-5, rtol=1e-5)
+
+
+def _padded_inputs(device: str) -> dict[str, torch.Tensor]:
+    generator = torch.Generator(device="cpu").manual_seed(3)
+    text = torch.randn(1, 7, 8, generator=generator)
+    c_mask = torch.ones(1, 7, dtype=torch.bool)
+    c_mask[:, 5:] = False
+    text[:, 5:] = 0.0
+    ref = torch.randn(1, 6, 4, generator=generator)
+    ref_mask = torch.ones(1, 6, dtype=torch.bool)
+    ref_mask[:, 4:] = False
+    mask = torch.ones(1, 9, dtype=torch.bool)
+    mask[:, 7:] = False
+    inputs = {"text": text, "c_mask": c_mask, "ref": ref, "ref_mask": ref_mask, "mask": mask}
+    return {name: value.to(device) for name, value in inputs.items()}
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
+def test_graph_with_time_grid_matches_eager_with_time_grid(cfg_strength: float) -> None:
+    dit = _make_dit("cuda")
+    inputs = _padded_inputs("cuda")
+    grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cuda")
+    x = torch.randn(1, 9, 4, device="cuda")
+    wrapper = AuKCUDAGraphWrapper(dit)
+    common = dict(
+        text=inputs["text"],
+        c_mask=inputs["c_mask"],
+        ref=inputs["ref"],
+        ref_mask=inputs["ref_mask"],
+        cfg_strength=cfg_strength,
+        timesteps=grid[:-1],
+    )
+    for i in range(4):
+        graph = wrapper(x=x, timestep=grid[i], new_request=i == 0, step_index=i, **common)
+        ctx = dit.prepare(
+            inputs["text"],
+            target_len=9,
+            c_mask=inputs["c_mask"],
+            ref=inputs["ref"],
+            ref_mask=inputs["ref_mask"],
+            cfg_infer=cfg_strength >= 1e-5,
+            timesteps=grid[:-1],
+        )
+        eager = wrapper._step(x, grid[i], ctx, cfg_strength, i)
+        torch.testing.assert_close(graph, eager, atol=1e-4, rtol=1e-4)
+    # One graph serves every step of the grid.
+    assert len(wrapper._cache) == 1
+    assert next(iter(wrapper._cache))[-1] == 4

tests/diffusion/models/auk/test_auk_step_context.py

@@ -127,7 +127,9 @@ def test_hoisted_euler_loop_matches_per_step_forward(cfg_strength: float) -> Non
             v = dit(x, time=timesteps[i], **inputs)
         x = x + (timesteps[i + 1] - timesteps[i]) * v
 
-    torch.testing.assert_close(got, x, rtol=0, atol=0)
+    # sample_latents precomputes the adaLN modulations of the whole grid in one
+    # GEMM per layer, whose rows round differently from one-row GEMMs in fp32.
+    torch.testing.assert_close(got, x, rtol=0, atol=1e-6)
     assert dit.text_cond is None and dit.text_uncond is None
 
 

tests/diffusion/models/auk/test_auk_vae_fastpath.py

@@ -11,10 +11,29 @@ eager off CUDA) and the tiled decode of long clips.
 import pytest
 import torch
 
-from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.models.auk import auk_vae
+from vllm_omni.diffusion.models.auk.auk_vae import AliasFreeActivation, AuKVAE, LowPass, SnakeBeta, Upsample
 from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles
+from vllm_omni.platforms import current_omni_platform
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model]
+
+
+@pytest.fixture(autouse=True)
+def _bound_cpu_decode_threads(request):
+    if request.node.get_closest_marker("cpu") is None:
+        yield
+        return
+    # These narrow convolutions run in shared CI pods. A host-sized Torch
+    # pool spends more time coordinating workers than decoding small clips.
+    previous_threads = torch.get_num_threads()
+    torch.set_num_threads(1)
+    print(f"AuK CPU decode threads: {previous_threads} -> {torch.get_num_threads()}")
+    try:
+        yield
+    finally:
+        torch.set_num_threads(previous_threads)
 
 
 def _small_vae() -> AuKVAE:
@@ -45,6 +64,23 @@ def _reference_snake(module: SnakeBeta, x: torch.Tensor) -> torch.Tensor:
     return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2)
 
 
+def _assert_plain_graph_matches_eager(replay: torch.Tensor, eager: torch.Tensor) -> None:
+    assert replay.shape == eager.shape and replay.dtype == eager.dtype
+    assert torch.isfinite(replay).all() and torch.isfinite(eager).all()
+    error = (replay - eager).abs()
+    rms = error.square().mean().sqrt().item()
+    print(f"AuK plain graph: max_abs={error.max().item():.9g}, rms={rms:.9g}")
+    if current_omni_platform.is_rocm():
+        # MI300 build 13040 measured max_abs <= 1.64e-7 and RMS <= 4.13e-8.
+        # Bound both peak and aggregate FP32 error; near-zero samples must
+        # not hide behind a relative tolerance. NVIDIA retains bit parity.
+        torch.testing.assert_close(replay, eager, atol=1e-6, rtol=0.0)
+        assert rms <= 1e-7, rms
+    else:
+        assert torch.equal(replay, eager)
+
+
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_window_padding_decodes_to_raw_zero() -> None:
     """Bucket padding is the normalized latent whose denormalized value is zero."""
@@ -57,6 +93,7 @@ def test_window_padding_decodes_to_raw_zero() -> None:
     torch.testing.assert_close(raw, torch.zeros_like(raw), atol=1e-6, rtol=0.0)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_eager_snake_matches_the_reference_formula() -> None:
     module = SnakeBeta(6, alpha_logscale=True)
@@ -69,6 +106,7 @@ def test_eager_snake_matches_the_reference_formula() -> None:
     assert torch.equal(module(x), _reference_snake(module, x))
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_filter_cache_leaves_the_filters_and_output_unchanged() -> None:
     lowpass = LowPass(cutoff=0.25, half_width=0.3, stride=2, kernel_size=12, causal=True)
@@ -88,6 +126,7 @@ def test_filter_cache_leaves_the_filters_and_output_unchanged() -> None:
         assert module._expanded.shape == (3, 1, 12)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_decode_fast_paths_reproduce_the_plain_decode() -> None:
     vae = _small_vae()
@@ -103,6 +142,7 @@ def test_decode_fast_paths_reproduce_the_plain_decode() -> None:
     assert all(module._cached for module in vae.modules() if isinstance(module, SnakeBeta))
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_graph_wrapper_falls_back_to_eager_off_cuda(mocker) -> None:
     vae = _small_vae()
@@ -119,6 +159,7 @@ def test_graph_wrapper_falls_back_to_eager_off_cuda(mocker) -> None:
     assert not wrapper._cache and not wrapper._compiled
 
 
+@pytest.mark.cpu
 def test_compiled_bucket_is_the_smallest_captured_one_that_fits() -> None:
     wrapper = AuKVAEDecodeGraph(_small_vae(), compile_shapes=(16, 8, 8))
     assert wrapper.compile_shapes == [8, 16]
@@ -131,6 +172,7 @@ def test_compiled_bucket_is_the_smallest_captured_one_that_fits() -> None:
     assert wrapper.compiled_bucket(6) == 16
 
 
+@pytest.mark.cpu
 def test_plan_tiles_covers_the_clip_once_with_full_context() -> None:
     assert plan_tiles(40, 64, 10, 3) == [(0, 40, 0, 40)]
     windows = plan_tiles(1000, 512, 53, 9)
@@ -149,10 +191,16 @@ def test_plan_tiles_covers_the_clip_once_with_full_context() -> None:
     # 512 and a 256 window instead of two 512s.
     assert plan_tiles(600, 512, 53, 9, sizes=(128, 256, 512)) == [(0, 512, 0, 503), (344, 256, 503, 600)]
     assert plan_tiles(1000, 512, 53, 9, sizes=(128, 256, 512))[-1] == (872, 128, 953, 1000)
+    narrow = plan_tiles(200, 64, 53, 9)
+    assert len(narrow) == 69
+    assert [emit for _, _, emit, _ in narrow[:3]] == [0, 55, 57]
+    assert all(a[3] == b[2] for a, b in zip(narrow, narrow[1:]))
+    assert narrow[-1][3] == 200
     with pytest.raises(ValueError, match="must exceed"):
         plan_tiles(100, 60, 53, 9)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_decode_context_bounds_the_measured_receptive_field() -> None:
     vae = _small_vae()
@@ -170,25 +218,26 @@ def test_decode_context_bounds_the_measured_receptive_field() -> None:
     assert differing and all(index < left or end - start - index <= right for index in differing)
 
 
+@pytest.mark.cpu
 @torch.inference_mode()
 def test_tiled_decode_matches_the_whole_decode() -> None:
     vae = _small_vae()
     # Off the accelerator every window decodes eagerly, so this isolates the stitching.
-    wrapper = AuKVAEDecodeGraph(vae, compile_shapes=(64,), tile_frames=64)
-    assert wrapper.tile_frames == 64 and wrapper.context_frames == (53, 9)
-    for frames in (64, 65, 200, 331):
+    wrapper = AuKVAEDecodeGraph(vae, compile_shapes=(128,), tile_frames=128)
+    assert wrapper.tile_frames == 128 and wrapper.context_frames == (53, 9)
+    for frames in (128, 129, 200, 331):
         latents = torch.randn(1, frames, vae.latent_dim)
         whole = vae.decode(latents)
         tiled = wrapper(latents)
         assert tiled.shape == whole.shape
         torch.testing.assert_close(tiled, whole, atol=1e-5, rtol=0.0)
-        assert wrapper.last_mode == ("eager" if frames == 64 else "tiled")
+        assert wrapper.last_mode == ("eager" if frames == 128 else "tiled")
     # The streaming form yields the same audio in order.
     latents = torch.randn(1, 200, vae.latent_dim)
     pieces = [(start, chunk.clone()) for start, chunk in wrapper.decode_tiles(latents)]
-    # With a 64-frame tile and a 62-frame halo each interior tile adds two frames.
-    assert [start for start, _ in pieces] == [emit for _, _, emit, _ in plan_tiles(200, 64, 53, 9)]
-    assert [start for start, _ in pieces][:3] == [0, 55, 57]
+    # A 128-frame tile avoids repeatedly decoding the same 62-frame halo.
+    assert [start for start, _ in pieces] == [emit for _, _, emit, _ in plan_tiles(200, 128, 53, 9)]
+    assert [start for start, _ in pieces][:3] == [0, 119, 185]
     torch.testing.assert_close(
         torch.cat([chunk for _, chunk in pieces], dim=1), vae.decode(latents), atol=1e-5, rtol=0.0
     )
@@ -198,20 +247,90 @@ def test_tiled_decode_matches_the_whole_decode() -> None:
         AuKVAEDecodeGraph(vae, tile_frames=60)
 
 
-@pytest.mark.cuda
+@pytest.mark.cpu
+@torch.inference_mode()
+def test_fused_activation_is_skipped_off_cuda() -> None:
+    activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True))
+    x = torch.randn(1, 4, 16)
+    assert not activation._fused_eligible(x)
+    assert activation(x).shape == (1, 4, 16)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
+@torch.inference_mode()
+def test_fused_activation_is_skipped_without_triton(monkeypatch: pytest.MonkeyPatch) -> None:
+    activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True)).to("cuda")
+    x = torch.randn(1, 4, 16, device="cuda")
+    assert activation._fused_eligible(x)
+    monkeypatch.setattr(auk_vae, "HAS_TRITON", False)
+    assert not activation._fused_eligible(x)
+    assert activation(x).shape == (1, 4, 16)
+
+
+def _activation_input(layout: str, length: int) -> torch.Tensor:
+    """A ``[2, 6, length]`` input, contiguous or as a strided view of a larger buffer."""
+    if layout == "contiguous":
+        x = torch.randn(2, 6, length, device="cuda")
+    elif layout == "transposed":
+        x = torch.randn(2, length, 6, device="cuda").transpose(1, 2)
+    elif layout == "sliced":
+        x = torch.randn(2, 8, 2 * length, device="cuda")[:, 1:7, ::2]
+    else:
+        raise ValueError(layout)
+    assert x.shape == (2, 6, length)
+    if length > 1:  # contiguity of a single-sample view depends on the layout
+        assert x.is_contiguous() == (layout == "contiguous")
+    return x * 3
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
+@pytest.mark.parametrize("layout", ["contiguous", "transposed", "sliced"])
+@pytest.mark.parametrize("causal", [False, True])
+# The kernels tile the output in blocks of 1024 samples: the 2x oversampled
+# signal crosses a block boundary at 512 input samples and the decimated one at
+# 1024, so both sides of each boundary are covered, plus short and odd lengths.
+@pytest.mark.parametrize("length", [1, 7, 511, 512, 513, 1000, 1023, 1024, 1025])
+@torch.inference_mode()
+def test_fused_activation_matches_the_eager_modules(layout: str, causal: bool, length: int) -> None:
+    torch.manual_seed(0)
+    activation = AliasFreeActivation(SnakeBeta(6, alpha_logscale=True), causal=causal).to("cuda")
+    with torch.no_grad():
+        activation.act.alpha.normal_(0.0, 0.3)
+        activation.act.beta.normal_(0.0, 0.3)
+    x = _activation_input(layout, length)
+    original = x.clone()
+
+    eager = activation.downsample(activation.act(activation.upsample(x)))
+    assert activation._fused_eligible(x)
+    fused = activation(x)
+
+    # Same fp32 arithmetic, only the order in which the taps are summed differs.
+    assert fused.shape == eager.shape
+    torch.testing.assert_close(fused, eager, atol=1e-5, rtol=1e-5)
+    # Neither path writes into its (possibly strided) input.
+    torch.testing.assert_close(x, original, atol=0, rtol=0)
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
 @torch.inference_mode()
 def test_graph_replay_matches_eager_per_length() -> None:
     vae = _small_vae().to("cuda")
     wrapper = AuKVAEDecodeGraph(vae, max_graphs=2)
     pools = []
+    saved_outputs: list[tuple[torch.Tensor, torch.Tensor]] = []
     for frames in (6, 9, 6, 12, 6):
         latents = torch.randn(1, frames, vae.latent_dim, device="cuda")
         eager = vae.decode(latents)
         replay = wrapper(latents)
-        # Exact-length graphs replay the very same kernels: bit-identical,
-        # including after the generation was retired.
-        assert torch.equal(replay, eager), frames
+        _assert_plain_graph_matches_eager(replay, eager)
+        # Returned waveforms own their storage and survive later replays
+        # and retirement of the generation that produced them.
+        for previous, snapshot in saved_outputs:
+            assert torch.equal(previous, snapshot)
+        saved_outputs.append((replay, replay.clone()))
         pools.append(wrapper._plain_pool)
     # The third distinct length found the cache full, so the whole generation
     # (6 and 9) was retired together with its pool rather than one graph at a
@@ -220,7 +339,7 @@ def test_graph_replay_matches_eager_per_length() -> None:
     assert pools[0] is pools[2] and pools[3] is pools[4] and pools[2] is not pools[3]
 
 
-@pytest.mark.cuda
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
 @torch.inference_mode()
 def test_bucketed_graph_only_disturbs_the_tail() -> None:
@@ -240,7 +359,7 @@ def test_bucketed_graph_only_disturbs_the_tail() -> None:
     torch.testing.assert_close(replay, eager, atol=0.1, rtol=0.0)
 
 
-@pytest.mark.cuda
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="torch.compile + CUDA graph capture requires CUDA")
 @torch.inference_mode()
 def test_compiled_buckets_replay_within_fusion_tolerance_and_leave_longer_clips_to_plain_graphs() -> None:
@@ -266,11 +385,12 @@ def test_compiled_buckets_replay_within_fusion_tolerance_and_leave_longer_clips_
     # a longer clip gets its own plain graph.
     assert wrapper.tile_frames == 0
     long = torch.randn(1, 12, vae.latent_dim, device="cuda")
-    assert torch.equal(wrapper(long), vae.decode(long))
+    replay, eager = wrapper(long), vae.decode(long)
+    _assert_plain_graph_matches_eager(replay, eager)
     assert wrapper.last_mode == "graph" and list(wrapper._cache) == [12]
 
 
-@pytest.mark.cuda
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="torch.compile + CUDA graph capture requires CUDA")
 @torch.inference_mode()
 def test_tiles_replay_the_compiled_bucket_for_long_clips() -> None:

tests/diffusion/models/auk/test_pipeline_auk.py

@@ -291,6 +291,18 @@ class TestRequestParsing:
         with pytest.raises(ValueError, match="AuK max_dit_graphs must be a positive integer"):
             build_pipeline(model_config={"max_dit_graphs": max_dit_graphs})
 
+    @pytest.mark.parametrize(
+        "model_config, expected", [({}, 16), ({"auk_ref_cache_size": 0}, 0), ({"auk_ref_cache_size": 4}, 4)]
+    )
+    def test_ref_cache_size_is_taken_from_model_config(self, build_pipeline, model_config, expected):
+        pipeline, _ = build_pipeline(model_config=model_config)
+        assert pipeline._ref_cache_size == expected
+
+    @pytest.mark.parametrize("size", [-1, 1.9, True, False, "4", None])
+    def test_invalid_ref_cache_size_is_rejected(self, build_pipeline, size):
+        with pytest.raises(ValueError, match="AuK auk_ref_cache_size must be a non-negative integer"):
+            build_pipeline(model_config={"auk_ref_cache_size": size})
+
     def test_pipeline_declares_audio_output(self, build_pipeline):
         pipeline, _ = build_pipeline()
 
@@ -492,6 +504,40 @@ class TestRequestParsing:
         assert call["sample"] is True
         assert isinstance(call["generator"], torch.Generator)
 
+    def test_reference_latents_are_cached_by_content(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+        clip = _silence(1.0)
+
+        for _ in range(3):
+            pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+        # A different clip is a new entry, the first one still hits.
+        other = (np.full_like(clip[0], 0.25), clip[1])
+        pipeline.forward(_batch(_prompt(audio=other, knobs={"gen_seconds": 1.0}), seed=1))
+        pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+
+        assert len(pipeline.vae.encode_calls) == 2
+        assert calls[0]["ref"] is calls[1]["ref"] is calls[4]["ref"]
+
+    def test_posterior_draws_and_disabled_cache_always_encode(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+        for _ in range(2):
+            pipeline.forward(
+                _batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1)
+            )
+        assert len(pipeline.vae.encode_calls) == 2
+
+        pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 0})
+        for _ in range(2):
+            pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0}), seed=1))
+        assert len(pipeline.vae.encode_calls) == 2
+
+    def test_reference_cache_is_bounded(self, build_pipeline):
+        pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 2})
+        for level in (0.1, 0.2, 0.3):
+            clip = (np.full(SAMPLE_RATE, level, dtype=np.float32), SAMPLE_RATE)
+            pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+        assert len(pipeline._ref_cache) == 2
+
     def test_one_request_per_forward(self, build_pipeline):
         pipeline, _ = build_pipeline()
         batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1)

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -868,6 +868,93 @@ def test_qwen_gen_mlp_remains_gated() -> None:
     assert hasattr(model.gen_layers[0].mlp, "gate_proj")
 
 
+@pytest.mark.parametrize("edge", [False, True])
+@pytest.mark.parametrize(
+    "use_hsdp,strategy", [(False, "full"), (False, "pre_sharded"), (True, "full"), (True, "pre_sharded")]
+)
+def test_cosmos3_releases_blocks_only_for_pre_sharded_hsdp(edge, use_hsdp, strategy, monkeypatch):
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge
+    from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
+
+    released = []
+
+    def release_and_check(block):
+        parameters = dict(block.named_parameters())
+        metadata = {name: dict(parameter.__dict__) for name, parameter in parameters.items()}
+        release_module_parameters_to_meta(block)
+        for name, parameter in block.named_parameters():
+            assert parameter is parameters[name]
+            assert parameter.device.type == "meta"
+            assert parameter.__dict__ == metadata[name]
+        released.append(block)
+
+    monkeypatch.setattr(transformer_cosmos3, "release_module_parameters_to_meta", release_and_check)
+    monkeypatch.setattr(transformer_cosmos3_edge, "release_module_parameters_to_meta", release_and_check)
+    model_cls = (
+        transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer if edge else transformer_cosmos3.Cosmos3VFMTransformer
+    )
+    config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
+    model = model_cls(
+        SimpleNamespace(
+            tf_model_config=config(num_hidden_layers=1),
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(use_hsdp=use_hsdp),
+            hsdp_weight_load_strategy=strategy,
+        )
+    )
+    blocks = [*model.language_model.layers, *model.gen_layers]
+    expected_meta = use_hsdp and strategy == "pre_sharded"
+    assert len(released) == (len(blocks) if expected_meta else 0)
+    assert all((p.device.type == "meta") == expected_meta for block in blocks for p in block.parameters())
+    assert model.time_embedder.linear_1.weight.device.type != "meta"
+    assert all(buffer.device.type != "meta" for buffer in model.buffers())
+
+
+@pytest.mark.parametrize("edge", [False, True])
+def test_cosmos3_pre_sharded_meta_processing_accepts_real_unquantized_blocks(edge, monkeypatch):
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+    from vllm_omni.diffusion.models.cosmos3 import Cosmos3EdgeVFMTransformer, Cosmos3VFMTransformer
+
+    model_cls = Cosmos3EdgeVFMTransformer if edge else Cosmos3VFMTransformer
+    config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
+    model = model_cls(
+        SimpleNamespace(
+            tf_model_config=config(num_hidden_layers=1),
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(use_hsdp=True),
+            hsdp_weight_load_strategy="pre_sharded",
+        )
+    )
+    groups = [SimpleNamespace(module=block) for block in (*model.language_model.layers, *model.gen_layers)]
+    plan = SimpleNamespace(roots=[SimpleNamespace(module=model)], groups=groups)
+    processed = []
+    # Exercise model selection with real vLLM linears, without invoking the
+    # CPU-specific weight packing hook on meta tensors in CPU-only CI.
+    monkeypatch.setattr(
+        UnquantizedLinearMethod, "process_weights_after_loading", lambda self, layer: processed.append(layer)
+    )
+    DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    expected = [module for group in groups for module in group.module.modules() if hasattr(module, "quant_method")]
+    assert expected
+    assert processed == expected
+
+    model._hsdp_pre_sharded_meta_post_load = False
+    with pytest.raises(ValueError, match="has not declared"):
+        DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    model._hsdp_pre_sharded_meta_post_load = True
+
+    class UnsupportedMethod(UnquantizedLinearMethod):
+        pass
+
+    expected[-1].quant_method = UnsupportedMethod()
+    processed.clear()
+    with pytest.raises(ValueError, match="UnsupportedMethod"):
+        DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    assert not processed
+
+
 def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 

tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py

@@ -148,7 +148,8 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
     assert array.shape == (2, 4, 4, 3)
     video = post({"video": frames}, SimpleNamespace(output_type="pt"))
     assert set(video) == {"video"}
-    assert video["video"] is frames
+    assert isinstance(video["video"], np.ndarray)
+    np.testing.assert_array_equal(video["video"], frames.numpy())
     image = post(
         {"image": torch.ones(4, 4, 3)},
         SimpleNamespace(output_type="pt"),
@@ -162,6 +163,40 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
     assert latent_output["image"] is latent
 
 
+@pytest.mark.parametrize("output_type", [None, "pt", "np"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_postprocess_preserves_video_range_for_serving(output_type: str | None, dtype: torch.dtype) -> None:
+    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
+    from vllm_omni.entrypoints.openai.video_api_utils import _coerce_video_to_frames
+
+    # Include black and midtones: all-white frames cannot expose double normalization.
+    frames = torch.tensor([0.0, 0.25, 0.5, 0.75, 1.0], dtype=dtype)
+    frames = frames.reshape(5, 1, 1, 1).expand(5, 2, 4, 3).clone()
+    expected = frames.float().numpy().copy()
+    sampling = None if output_type is None else OmniDiffusionSamplingParams(output_type=output_type)
+    post = get_lingbot_video_post_process_func(None)
+
+    processed = post({"video": frames}, sampling)
+    encoded_frames = np.stack(_coerce_video_to_frames(processed["video"]))
+
+    assert set(processed) == {"video"}
+    assert encoded_frames.dtype == np.float32
+    np.testing.assert_array_equal(encoded_frames, expected)
+    np.testing.assert_array_equal(frames.float().numpy(), expected)
+
+
+@pytest.mark.parametrize("output_key", ["image", "video"])
+def test_postprocess_preserves_latent_payload(output_key: str) -> None:
+    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
+
+    latents = torch.tensor([-2.0, 0.0, 2.0]).reshape(1, 3, 1, 1, 1)
+    sampling = OmniDiffusionSamplingParams(output_type="latent")
+    processed = get_lingbot_video_post_process_func(None)({output_key: latents}, sampling)
+
+    assert set(processed) == {output_key}
+    assert processed[output_key] is latents
+
+
 def test_postprocess_preserves_image_contract_for_serving():
     from vllm_omni.diffusion.data import DiffusionOutput
     from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
@@ -181,7 +216,9 @@ def test_postprocess_preserves_image_contract_for_serving():
         sampling_params=sampling,
         request_id="lingbot-image-contract-test",
     )
-    raw_output = {"image": torch.ones(8, 12, 3)}
+    image = torch.zeros(8, 12, 3)
+    image[4:] = 1.0
+    raw_output = {"image": image}
     processed = get_lingbot_video_post_process_func(SimpleNamespace())(raw_output, sampling)
     normalized = normalize_diffusion_postprocess_output(processed)
 
@@ -199,6 +236,8 @@ def test_postprocess_preserves_image_contract_for_serving():
     endpoint_images = _extract_images_from_result(result)
     assert len(endpoint_images) == 1
     assert endpoint_images[0].size == (12, 8)
+    assert endpoint_images[0].getpixel((0, 0)) == (0, 0, 0)
+    assert endpoint_images[0].getpixel((0, 7)) == (255, 255, 255)
 
 
 def test_forward_resolves_t2v_sampling_and_flow_shift_alias():