[Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951)

Signed-off-by: wangyu <410167048@qq.com>

tests/diffusion/models/dreamzero/test_crossattn_cache.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests that DreamZero's cross-attention KV cache is actually init'd and reused.
+
+DreamZero's I2V cross-attention caches the text k/v and image k_img/v_img in a
+per-layer ``crossattn_cache`` dict: they are computed on the first forward of a
+session (``is_init`` False -> True) and reused on every later step, because the
+text/image context is session-invariant. Only the query (which depends on the
+per-step hidden state) is recomputed.
+
+These tests lock in that ``WanI2VCrossAttention.forward`` initialises the cache
+on the first call and reuses it (not recompute) on later calls, and recomputes
+every call when no cache is provided.
+
+CPU-only: ``WanI2VCrossAttention.__init__`` needs a tensor-parallel group (it
+builds ColumnParallel/RowParallel linears), so we bypass ``__init__`` and stub
+the projections/attention to count calls and exercise only the caching logic.
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.causal_wan_model import WanI2VCrossAttention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_CLIP_LEN = 257  # image tokens the I2V cross-attn splits off the front of context
+
+
+def _make_cross_attn(n_heads: int = 2, head_dim: int = 4) -> WanI2VCrossAttention:
+    """A WanI2VCrossAttention whose submodules are call-counting stubs.
+
+    ``__init__`` is bypassed (it needs a TP group); we only wire what
+    ``forward`` touches. Each projection returns a [B, S, n*head_dim] tensor so
+    the subsequent ``unflatten(2, (n, head_dim))`` works, and is a MagicMock so
+    we can assert how many times it was invoked.
+    """
+    inner = n_heads * head_dim
+    module = WanI2VCrossAttention.__new__(WanI2VCrossAttention)
+    module.tp_num_heads = n_heads
+    module.head_dim = head_dim
+
+    def _proj() -> MagicMock:
+        return MagicMock(side_effect=lambda t: torch.zeros(t.shape[0], t.shape[1], inner))
+
+    module.q = _proj()
+    module.k = _proj()
+    module.v = _proj()
+    module.k_img = _proj()
+    module.v_img = _proj()
+    module.norm_q = lambda t: t
+    module.norm_k = lambda t: t
+    module.norm_k_img = lambda t: t
+    # attn returns the query unchanged ([B, S, n, d]); flatten(2) downstream is happy.
+    module.attn = lambda q, k, v: q
+    module.o = lambda t: t
+    return module
+
+
+def _fresh_cache() -> dict:
+    return {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None}
+
+
+def _inputs(inner: int = 8):
+    x = torch.randn(1, 3, inner)  # B=1, query_len=3
+    context = torch.randn(1, _CLIP_LEN + 5, inner)  # 257 image tokens + 5 text tokens
+    return x, context
+
+
+def test_cache_inits_on_first_call_and_reuses_on_second() -> None:
+    module = _make_cross_attn()
+    cache = _fresh_cache()
+    x, context = _inputs()
+
+    # First call: cache is initialised and every k/v projection runs once.
+    module.forward(x, context, crossattn_cache=cache)
+    assert cache["is_init"] is True
+    for key in ("k", "v", "k_img", "v_img"):
+        assert cache[key] is not None, f"{key} should be cached after first call"
+    assert module.k.call_count == 1
+    assert module.v.call_count == 1
+    assert module.k_img.call_count == 1
+    assert module.v_img.call_count == 1
+
+    cached = {key: cache[key] for key in ("k", "v", "k_img", "v_img")}
+
+    # Second call: k/v/k_img/v_img are read from cache, NOT recomputed.
+    module.forward(x, context, crossattn_cache=cache)
+    assert module.k.call_count == 1
+    assert module.v.call_count == 1
+    assert module.k_img.call_count == 1
+    assert module.v_img.call_count == 1
+    for key, tensor in cached.items():
+        assert cache[key] is tensor, f"{key} should be the same cached object"
+
+    # The query always depends on the per-step hidden state, so it is recomputed.
+    assert module.q.call_count == 2
+
+
+def test_no_cache_recomputes_every_call() -> None:
+    module = _make_cross_attn()
+    x, context = _inputs()
+
+    module.forward(x, context, crossattn_cache=None)
+    module.forward(x, context, crossattn_cache=None)
+
+    # Without a cache, the image/text k/v are recomputed on every call.
+    assert module.k_img.call_count == 2
+    assert module.v_img.call_count == 2
+    assert module.k.call_count == 2
+    assert module.v.call_count == 2

tests/diffusion/models/dreamzero/test_deploy_configs.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Every shipped DreamZero deploy config must select the AR-Diffusion engine.
+
+``DreamZeroPipeline.__init__`` rejects any other engine, so a config missing
+``engine_backend`` fails at worker startup rather than at request time. That
+only surfaces in slow multi-GPU e2e serving tests (see issue #5046), so assert
+it here on CPU instead.
+"""
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import yaml
+
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Discover siblings of a known config so a future dreamzero*.yaml is covered
+# without editing this test.
+DEPLOY_DIR = Path(get_deploy_config_path("dreamzero.yaml")).parent
+
+
+def _dreamzero_stages() -> list[tuple[str, int, dict]]:
+    stages = []
+    for name in sorted(path.name for path in DEPLOY_DIR.glob("dreamzero*.yaml")):
+        config = yaml.safe_load(Path(get_deploy_config_path(name)).read_text())
+        for index, stage in enumerate(config.get("stages") or []):
+            stages.append((name, index, stage))
+    return stages
+
+
+def test_dreamzero_deploy_configs_exist() -> None:
+    assert _dreamzero_stages(), f"no DreamZero deploy configs found under {DEPLOY_DIR}"
+
+
+@pytest.mark.parametrize(
+    ("config_name", "stage_index", "stage"),
+    _dreamzero_stages(),
+    ids=lambda value: value if isinstance(value, str) else None,
+)
+def test_dreamzero_stage_selects_ar_diffusion_engine(config_name: str, stage_index: int, stage: dict) -> None:
+    backend = stage.get("engine_backend")
+    assert backend, (
+        f"{config_name} stage {stage_index} does not set engine_backend; "
+        "DreamZeroPipeline requires the AR-Diffusion engine"
+    )
+    # Resolving also proves the import path is real, not just a lookalike string.
+    engine_class = DiffusionEngine.resolve_engine_class(SimpleNamespace(engine_backend=backend))
+    assert issubclass(engine_class, ARDiffusionEngine), (
+        f"{config_name} stage {stage_index} resolves to {engine_class.__name__}, not an AR-Diffusion engine"
+    )

tests/diffusion/models/dreamzero/test_dreamzero_utils.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.diffusion.models.dreamzero.utils import (
+    DEFAULT_CFG_SCALE,
+    DEFAULT_EMBODIMENT_NAME_TO_ID,
+    DEFAULT_NEGATIVE_PROMPT,
+    DEFAULT_NUM_INFERENCE_STEPS,
+    DEFAULT_SEED,
+    DEFAULT_SIGMA_SHIFT,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_dreamzero_default_constants_match_source_baseline():
+    assert DEFAULT_NUM_INFERENCE_STEPS == 16
+    assert DEFAULT_CFG_SCALE == 5.0
+    assert DEFAULT_SIGMA_SHIFT == 5.0
+    assert DEFAULT_SEED == 1140
+    assert "worst quality" in DEFAULT_NEGATIVE_PROMPT
+    assert DEFAULT_EMBODIMENT_NAME_TO_ID["oxe_droid"] == 17

tests/diffusion/models/dreamzero/test_fused_qk_rms_norm.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for DreamZero's fused q/k qk-norm TP all-reduce.
+
+`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor
+and issues a single tensor-parallel all-reduce instead of two. These tests lock
+in that it is numerically identical to applying `norm_q`/`norm_k` separately,
+and that it really issues exactly one (fused) all-reduce.
+
+CPU-only: the TP world size and the collective are mocked, so a single process
+simulates one rank (all-reduce is identity). This exercises the cat/split and
+count-scaling logic that a real multi-rank run depends on.
+"""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.causal_wan_model import (
+    DistributedRMSNorm,
+    fused_qk_rms_norm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.dreamzero.causal_wan_model"
+
+
+def _make_norms(dim: int) -> tuple[DistributedRMSNorm, DistributedRMSNorm]:
+    """Two norms with distinct random weights so a swapped q/k slice is caught."""
+    torch.manual_seed(0)
+    norm_q = DistributedRMSNorm(dim)
+    norm_k = DistributedRMSNorm(dim)
+    norm_q.weight.data.normal_()
+    norm_k.weight.data.normal_()
+    return norm_q, norm_k
+
+
+def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
+    """Single-rank all-reduce: the reduction over one rank is the identity."""
+    return tensor
+
+
+def test_fused_matches_separate_tp1() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    with patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1):
+        ref_q, ref_k = norm_q(q), norm_k(k)
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_matches_separate_tp_gt1() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    # tp_size=2 exercises the reduce branch and the `global_count *= tp_size`
+    # path in both the per-norm forward and the fused helper.
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+        patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce),
+    ):
+        ref_q, ref_k = norm_q(q), norm_k(k)
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_issues_single_all_reduce_over_packed_pair() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+        patch(
+            f"{_MODULE}.tensor_model_parallel_all_reduce",
+            side_effect=_identity_all_reduce,
+        ) as mock_ar,
+    ):
+        fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    # Exactly one collective, over a tensor whose last dim packs the q+k pair.
+    assert mock_ar.call_count == 1
+    packed = mock_ar.call_args[0][0]
+    assert packed.shape[-1] == 2
+
+
+def test_fused_identity_fallback_when_qk_norm_disabled() -> None:
+    # qk_norm=False makes the norms nn.Identity; the helper must fall back to
+    # plain passthrough without touching the collective.
+    norm_q, norm_k = torch.nn.Identity(), torch.nn.Identity()
+    q = torch.randn(2, 5, 64)
+    k = torch.randn(2, 5, 64)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce") as mock_ar:
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, k, rtol=0, atol=0)
+    mock_ar.assert_not_called()

tests/diffusion/models/dreamzero/test_pipeline_state.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import OrderedDict
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _empty_pipeline() -> DreamZeroPipeline:
+    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipeline._states = OrderedDict()
+    return pipeline
+
+
+def test_dreamzero_pipeline_state_is_session_keyed() -> None:
+    pipeline = _empty_pipeline()
+
+    session_a = pipeline._get_or_create_state("session-a")
+    session_b = pipeline._get_or_create_state("session-b")
+    session_a.call_count = 7
+    session_b.call_count = 3
+
+    assert pipeline._get_or_create_state("session-a") is session_a
+    assert pipeline._get_or_create_state("session-b") is session_b
+    assert session_a.call_count == 7
+    assert session_b.call_count == 3
+
+
+def test_dreamzero_pipeline_state_follows_runner_lifecycle_notifications() -> None:
+    pipeline = _empty_pipeline()
+
+    session_a = pipeline._get_or_create_state("session-a")
+    session_b = pipeline._get_or_create_state("session-b")
+    pipeline.state = session_b
+
+    pipeline.close_ar_diffusion_session("session-a")
+    assert pipeline.state is session_b
+    pipeline.reset_ar_diffusion_session("session-b")
+
+    assert not pipeline._states
+    assert pipeline.state is None
+    assert pipeline._get_or_create_state("session-a") is not session_a
+    assert pipeline._get_or_create_state("session-b") is not session_b
+
+
+def test_close_session_clears_active_state_alias() -> None:
+    pipeline = _empty_pipeline()
+    session_a = pipeline._get_or_create_state("session-a")
+    pipeline.state = session_a
+
+    pipeline.close_ar_diffusion_session("session-a")
+
+    assert not pipeline._states
+    assert pipeline.state is None
+
+
+def test_dreamzero_warmup_provider_builds_session_scoped_requests() -> None:
+    provider = SimpleNamespace(
+        ar_diffusion_kv_cache_spec=lambda: SimpleNamespace(window_frames=5, frames_per_block=4),
+        od_config=SimpleNamespace(
+            ar_diffusion_kv_config=None,
+            model_config={"policy_server_config": {"image_resolution": [8, 16]}},
+        ),
+        _ar_warmup_robot_obs=DreamZeroPipeline._ar_warmup_robot_obs,
+    )
+
+    requests = list(DreamZeroPipeline.ar_diffusion_warmup_requests(provider, "warmup-session"))
+
+    assert [request.prompt for request in requests] == ["warmup", "warmup"]
+    assert all(request.sampling_params.extra_args["session_id"] == "warmup-session" for request in requests)
+    assert requests[0].sampling_params.extra_args["robot_obs"]["observation/exterior_image_0_left"].shape == (
+        8,
+        16,
+        3,
+    )
+    assert requests[1].sampling_params.extra_args["robot_obs"]["observation/exterior_image_0_left"].shape == (
+        4,
+        8,
+        16,
+        3,
+    )
+
+
+def test_dreamzero_state_owns_no_kv_caches() -> None:
+    """KV (self- and cross-attention) is engine-owned since the AR-Diffusion
+    paged backend: the model-local cache accessors must be gone so nothing can
+    silently bypass the engine pool."""
+    state = DreamZeroState()
+
+    for removed in ("get_kv_caches", "create_kv_caches", "update_kv_cache", "get_crossattn_caches"):
+        assert not hasattr(state, removed)

tests/diffusion/models/dreamzero/test_qkv_fusion.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence test for DreamZero self-attn QKV fusion.
+
+CausalWanSelfAttention now projects q/k/v with a single QKVParallelLinear and
+splits the result, instead of three separate ColumnParallelLinear. The DreamZero
+checkpoint still ships separate q/k/v weights, so load_weights routes each into
+the packed `qkv` param with its shard id. This test proves the chain
+"load separate q/k/v via shard id -> single GEMM -> split" reproduces the three
+independent projections, given identical weights.
+
+CPU-only: a single-rank (world_size=1) tensor-parallel group is initialized and
+the GEMM dispatcher is forced to the default (CPU-compatible) path, following
+tests/diffusion/layers/test_adalayernorm.py.
+"""
+
+import os
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _init_distributed():
+    """Minimal world_size=1 TP group required to build QKVParallelLinear."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29507")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(autouse=True)
+def _force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for CPU test tensors."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+def test_fused_qkv_load_and_split_matches_separate_projections():
+    from vllm.model_executor.layers.linear import QKVParallelLinear
+
+    torch.manual_seed(0)
+    dim, head_dim, num_heads = 16, 4, 4  # inner_dim == dim, no GQA
+
+    qkv = QKVParallelLinear(
+        hidden_size=dim,
+        head_size=head_dim,
+        total_num_heads=num_heads,
+        bias=True,
+    )
+
+    # Distinct random weight + bias per projection, loaded the exact way
+    # DreamZero.load_weights does: param.weight_loader(param, tensor, shard_id).
+    refs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+    for shard_id in ("q", "k", "v"):
+        w = torch.randn(dim, dim)
+        b = torch.randn(dim)
+        refs[shard_id] = (w, b)
+        qkv.weight.weight_loader(qkv.weight, w, shard_id)
+        qkv.bias.weight_loader(qkv.bias, b, shard_id)
+
+    x = torch.randn(2, 3, dim)
+    fused, _ = qkv(x)
+    q, k, v = fused.split([dim, dim, dim], dim=-1)
+
+    for got, shard_id in zip((q, k, v), ("q", "k", "v"), strict=True):
+        w, b = refs[shard_id]
+        ref = torch.nn.functional.linear(x, w, b)
+        torch.testing.assert_close(got, ref)
+
+
+def test_fused_qkv_shard_ids_are_not_transposed():
+    # Guard against q/k/v being routed to the wrong slice: give each projection a
+    # constant, distinct weight so a swapped shard id would surface immediately.
+    from vllm.model_executor.layers.linear import QKVParallelLinear
+
+    dim, head_dim, num_heads = 16, 4, 4
+    qkv = QKVParallelLinear(
+        hidden_size=dim,
+        head_size=head_dim,
+        total_num_heads=num_heads,
+        bias=True,
+    )
+
+    scale = {"q": 1.0, "k": 2.0, "v": 3.0}
+    for shard_id, s in scale.items():
+        qkv.weight.weight_loader(qkv.weight, torch.full((dim, dim), s), shard_id)
+        qkv.bias.weight_loader(qkv.bias, torch.zeros(dim), shard_id)
+
+    x = torch.ones(1, dim)
+    fused, _ = qkv(x)
+    q, k, v = fused.split([dim, dim, dim], dim=-1)
+
+    # F.linear(ones, full(s)) == s * dim per element.
+    torch.testing.assert_close(q, torch.full((1, dim), scale["q"] * dim))
+    torch.testing.assert_close(k, torch.full((1, dim), scale["k"] * dim))
+    torch.testing.assert_close(v, torch.full((1, dim), scale["v"] * dim))

tests/diffusion/models/dreamzero/test_session_state_accounting.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-device byte accounting and the bounded VAE encoder ring (RFC #4480).
+
+Two things are pinned here. First, that a session reports the bytes it actually
+holds: both buckets walked, storages rather than logical element counts, and
+aliases counted once. Second, that bounding the adapter's VAE encoder history
+changes no numbers -- ``quant_conv`` is pointwise in time, so the frames the
+reader keeps do not depend on the ones the ring drops.
+
+All CPU, tiny tensors, no model.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.experimental.world_models.adapters.state_dreamzero_adapter import (
+    DreamZeroStateAdapter,
+)
+from vllm_omni.experimental.world_models.session_state import (
+    LatentBuffer,
+    SessionStateManager,
+    device_bytes,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _bytes(value: object) -> dict[str, int]:
+    return device_bytes(value, set())
+
+
+# -- device_bytes ------------------------------------------------------------
+
+
+def test_counts_tensor_storage_under_its_device() -> None:
+    tensor = torch.zeros(64, dtype=torch.float32)
+    assert _bytes(tensor) == {"cpu": 256}
+
+
+def test_views_on_one_storage_counted_once() -> None:
+    tensor = torch.zeros(64, dtype=torch.float32)
+    seen: set[tuple[str, int]] = set()
+    first = device_bytes(tensor, seen)
+    second = device_bytes(tensor[:8], seen)
+    assert first == {"cpu": 256}
+    assert second == {}
+
+
+def test_slice_reports_the_storage_it_pins() -> None:
+    # The whole point of counting storages: this slice keeps all 256 bytes
+    # alive, and ``numel() * element_size()`` would claim 32.
+    sliced = torch.zeros(64, dtype=torch.float32)[:8]
+    assert _bytes(sliced) == {"cpu": 256}
+
+
+def test_containers_recurse_and_scalars_contribute_nothing() -> None:
+    payload = {
+        "tensors": [torch.zeros(16, dtype=torch.float32), None],
+        "count": 7,
+        "flag": True,
+        "name": "not-a-buffer",
+    }
+    assert _bytes(payload) == {"cpu": 64}
+
+
+def test_numpy_arrays_land_on_the_host() -> None:
+    assert _bytes(np.zeros(32, dtype=np.float32)) == {"cpu": 128}
+
+
+# -- SessionState / manager --------------------------------------------------
+
+
+def test_session_counts_both_buckets_and_dedupes_across_them() -> None:
+    manager = SessionStateManager()
+    session = manager.get_or_create_session("s")
+
+    buffer: LatentBuffer[torch.Tensor] = LatentBuffer()
+    buffer.allocate(maxlen=None)
+    shared = torch.zeros(64, dtype=torch.float32)  # 256 bytes
+    buffer.append(shared)
+    session.put("latents", buffer)
+    session.attrs["conditioning"] = torch.zeros(16, dtype=torch.float32)  # 64 bytes
+
+    assert session.nbytes_by_device() == {"cpu": 320}
+
+    # The same storage reachable from both buckets is one allocation, not two.
+    session.attrs["alias"] = shared
+    assert session.nbytes_by_device() == {"cpu": 320}
+
+
+def test_stats_reports_per_device_totals() -> None:
+    manager = SessionStateManager()
+    manager.get_or_create_session("s").attrs["x"] = torch.zeros(64, dtype=torch.float32)
+
+    stats = manager.stats()
+    assert stats["nbytes:cpu"] == 256
+    assert stats["total_nbytes"] == 256
+
+
+def test_attribute_bytes_were_previously_invisible() -> None:
+    """The defect behind tzhouam's review comment on PR #4487.
+
+    Every byte here lives in ``attrs``, and the old ``nbytes`` walked only the
+    typed objects -- so a session holding real memory reported zero.
+    """
+    manager = SessionStateManager()
+    session = manager.get_or_create_session("s")
+    session.attrs["conv_cache"] = [torch.zeros(1024, dtype=torch.float32)]
+
+    assert session.nbytes_by_device() == {"cpu": 4096}
+
+
+# -- the largest value DreamZero holds ---------------------------------------
+
+
+def test_feat_map_bytes_are_counted() -> None:
+    """``vae_enc_feat_map`` is the largest thing a DreamZero session holds --
+    603 MiB at the shipped config -- and it stays an attribute, so this is
+    exactly the memory the old accounting missed."""
+    manager = SessionStateManager()
+    adapter = DreamZeroStateAdapter("s", manager)
+    adapter.vae_enc_feat_map = [torch.zeros(1024, dtype=torch.float32), None]
+
+    assert manager.stats()["nbytes:cpu"] == 4096
+
+    # Freeing it piecemeal is not something the manager can do: the cache has no
+    # recompute source and the stream breaks without it. Ending the stream (or
+    # the session) is the only release, and that reclaims attributes too.
+    adapter.reset_vae_encoder_stream()
+    assert manager.stats().get("nbytes:cpu", 0) == 0
+
+
+# -- bounded VAE encoder ring ------------------------------------------------
+
+
+def _quantize(encoder_out: torch.Tensor) -> torch.Tensor:
+    """Stand-in for ``_vae_quantize_encoder_out``.
+
+    The real one is ``WanCausalConv3d(z_dim*2, z_dim*2, 1)`` followed by a
+    per-channel affine rescale. Every operation is pointwise in time, which is
+    the property the ring relies on; a 1x1x1 convolution reproduces it.
+    """
+    weight = torch.arange(1, encoder_out.shape[1] + 1, dtype=encoder_out.dtype).view(1, -1, 1, 1, 1)
+    return encoder_out * weight + 0.5
+
+
+def test_ring_is_bounded() -> None:
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=2)
+    adapter.vae_encoder_out = torch.randn(1, 4, 1, 2, 2)
+    for _ in range(10):
+        adapter.append_vae_encoder_chunk(torch.randn(1, 4, 1, 2, 2))
+
+    assert adapter.vae_encoder_out is not None
+    assert adapter.vae_encoder_out.shape[2] == 2
+
+
+def test_ring_matches_unbounded_history_for_the_frames_that_are_read() -> None:
+    window = 2
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=window)
+    seed = torch.randn(1, 4, 1, 2, 2)
+    chunks = [torch.randn(1, 4, 1, 2, 2) for _ in range(9)]
+
+    adapter.vae_encoder_out = seed
+    unbounded = seed
+    for chunk in chunks:
+        adapter.append_vae_encoder_chunk(chunk)
+        unbounded = torch.cat([unbounded, chunk], dim=2)
+
+    # What the reader actually returns: quantise, then keep the last `window`.
+    from_ring = _quantize(adapter.vae_encoder_out)[:, :, -window:]
+    from_history = _quantize(unbounded)[:, :, -window:]
+    torch.testing.assert_close(from_ring, from_history, rtol=0, atol=0)
+
+
+def test_ring_holds_a_single_frame_before_the_first_body_chunk() -> None:
+    # After `_vae_stream_seed` there is exactly one latent frame, which is what
+    # sends the reader down its padding branch.
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=2)
+    seed = torch.randn(1, 4, 1, 2, 2)
+    adapter.vae_encoder_out = seed
+
+    assert adapter.vae_encoder_out is not None
+    assert adapter.vae_encoder_out.shape[2] == 1
+    torch.testing.assert_close(adapter.vae_encoder_out, seed, rtol=0, atol=0)
+
+
+def test_resetting_the_stream_clears_the_ring() -> None:
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=2)
+    adapter.vae_encoder_out = torch.randn(1, 4, 1, 2, 2)
+    adapter.reset_vae_encoder_stream()
+
+    assert adapter.vae_encoder_out is None
+
+
+def test_window_must_be_positive() -> None:
+    with pytest.raises(ValueError, match="vae_encoder_window"):
+        DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=0)

tests/diffusion/models/dreamzero/test_session_state_equivalence.py

@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Bit-equivalence tests: DreamZeroStateAdapter vs bespoke DreamZeroState.
+
+These drive both state objects directly with tiny CPU tensors -- no model, no
+GPU -- and assert the adapter (backed by the SessionStateManager) stores and
+returns exactly what the bespoke DreamZeroState does. This gates routing
+DreamZero through the shared session state manager (RFC #4480).
+
+DreamZero's attention KV is engine-owned (AR-Diffusion paged pool, PR #4534)
+and out of the adapter's scope; the state surface tested here is the model's
+non-KV session state.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.state_dreamzero import (
+    FRAMES_PER_CHUNK,
+    DreamZeroState,
+)
+from vllm_omni.experimental.world_models.adapters.state_dreamzero_adapter import DreamZeroStateAdapter
+from vllm_omni.experimental.world_models.session_state import SessionStateManager
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _create_both() -> tuple[DreamZeroState, DreamZeroStateAdapter]:
+    return DreamZeroState(), DreamZeroStateAdapter("session-0", SessionStateManager())
+
+
+def test_accumulate_frames_and_call_count_match() -> None:
+    bespoke, adapter = _create_both()
+    rng = np.random.default_rng(0)
+    for _ in range(2 * FRAMES_PER_CHUNK + 1):
+        frame = rng.integers(0, 255, size=(5, 6, 3), dtype=np.uint8)
+        b_out = bespoke.accumulate_frames(frame)
+        a_out = adapter.accumulate_frames(frame)
+        np.testing.assert_array_equal(b_out, a_out)
+        assert bespoke.call_count == adapter.call_count
+        assert len(bespoke.stitched_buffer) == len(adapter.stitched_buffer)
+    # Deque rollover capped at FRAMES_PER_CHUNK.
+    assert len(adapter.stitched_buffer) == FRAMES_PER_CHUNK
+
+
+def test_accumulate_frames_4d_extend_matches() -> None:
+    bespoke, adapter = _create_both()
+    clip = np.random.default_rng(1).integers(0, 255, size=(3, 5, 6, 3), dtype=np.uint8)
+    np.testing.assert_array_equal(bespoke.accumulate_frames(clip), adapter.accumulate_frames(clip))
+    assert len(bespoke.stitched_buffer) == len(adapter.stitched_buffer)
+
+
+def test_video_latents_roundtrip_matches() -> None:
+    bespoke, adapter = _create_both()
+    assert bespoke.get_concatenated_video_latents() is None
+    assert adapter.get_concatenated_video_latents() is None
+
+    rng = torch.Generator().manual_seed(0)
+    for _ in range(3):
+        # (B, T, C, H, W) as returned by the denoise loop.
+        chunk = torch.randn(1, 2, 3, 4, 4, generator=rng)
+        bespoke.append_video_latents(chunk)
+        adapter.append_video_latents(chunk)
+        torch.testing.assert_close(
+            bespoke.get_concatenated_video_latents(),
+            adapter.get_concatenated_video_latents(),
+        )
+
+    bespoke.clear_video_latents()
+    adapter.clear_video_latents()
+    assert bespoke.get_concatenated_video_latents() is None
+    assert adapter.get_concatenated_video_latents() is None
+
+
+def test_vae_stream_fields_match_bespoke() -> None:
+    bespoke, adapter = _create_both()
+    for state in (bespoke, adapter):
+        assert state.vae_stream_initialized is False
+        assert state.vae_enc_feat_map is None
+        assert state.vae_encoder_out is None
+        assert state.vae_pending_body_frames is None
+
+    feat_map: list[torch.Tensor | None] = [torch.ones(1, 2), None]
+    encoder_out = torch.randn(1, 4, 2, 3, 3)
+    pending = torch.randn(1, 3, 2, 8, 8)
+    for state in (bespoke, adapter):
+        state.vae_enc_feat_map = feat_map
+        state.vae_encoder_out = encoder_out
+        state.vae_pending_body_frames = pending
+        state.vae_stream_initialized = True
+
+    assert adapter.vae_stream_initialized is bespoke.vae_stream_initialized is True
+    torch.testing.assert_close(adapter.vae_encoder_out, bespoke.vae_encoder_out)
+    torch.testing.assert_close(adapter.vae_pending_body_frames, bespoke.vae_pending_body_frames)
+    assert adapter.vae_enc_feat_map is feat_map
+
+    for state in (bespoke, adapter):
+        state.reset_vae_encoder_stream()
+        assert state.vae_stream_initialized is False
+        assert state.vae_enc_feat_map is None
+        assert state.vae_encoder_out is None
+        assert state.vae_pending_body_frames is None
+
+
+def test_reset_clears_state_and_should_reset_matches() -> None:
+    bespoke, adapter = _create_both()
+    tokens = torch.tensor([1, 2, 3])
+    # language unset -> both want reset.
+    assert bespoke.should_reset(tokens, 4, -1) is adapter.should_reset(tokens, 4, -1) is True
+
+    for state in (bespoke, adapter):
+        state.language = tokens
+        state.current_start_frame = 2
+    # same language -> no reset; both agree.
+    assert bespoke.should_reset(tokens, 4, -1) == adapter.should_reset(tokens, 4, -1) is False
+    # language change -> reset.
+    other = torch.tensor([9, 9])
+    assert bespoke.should_reset(other, 4, -1) == adapter.should_reset(other, 4, -1) is True
+    # local_attn_size exceeded -> reset ("inference", not "session").
+    assert bespoke.reset_reason(tokens, 4, 1) == adapter.reset_reason(tokens, 4, 1) == "inference"
+
+    # A session reset clears everything, including the prompt-embed cache and
+    # the VAE encoder stream (mirrors DreamZeroState.reset()).
+    for state in (bespoke, adapter):
+        state.prompt_embeds = torch.ones(1, 2)
+        state.vae_encoder_out = torch.ones(1, 4, 1, 2, 2)
+        state.vae_stream_initialized = True
+        state.reset()
+        assert state.language is None
+        assert state.prompt_embeds is None
+        assert state.current_start_frame == 0
+        assert state.call_count == 0
+        assert state.vae_stream_initialized is False
+        assert state.vae_encoder_out is None
+
+
+def test_window_reset_preserves_latents_prompt_and_vae_stream() -> None:
+    # reset(clear_video_latents=False) is the window ("inference") reset: the
+    # prompt is unchanged and the VAE feat_cache history is independent of the
+    # DiT attention window, so both survive alongside the video latents.
+    bespoke, adapter = _create_both()
+    tokens = torch.tensor([1, 2, 3])
+    embeds = torch.randn(1, 4)
+    encoder_out = torch.randn(1, 4, 2, 3, 3)
+    chunk = torch.randn(1, 2, 3, 4, 4)
+    for state in (bespoke, adapter):
+        state.language = tokens
+        state.prompt_embeds = embeds
+        state.vae_encoder_out = encoder_out
+        state.vae_stream_initialized = True
+        state.current_start_frame = 7
+        state.append_video_latents(chunk)
+
+        state.reset_inference_state()
+
+        assert state.current_start_frame == 0
+        assert state.call_count == 0
+        assert state.clip_feas is None
+        torch.testing.assert_close(state.language, tokens)
+        torch.testing.assert_close(state.prompt_embeds, embeds)
+        assert state.vae_stream_initialized is True
+        torch.testing.assert_close(state.vae_encoder_out, encoder_out)
+        assert state.get_concatenated_video_latents() is not None
+
+    torch.testing.assert_close(
+        bespoke.get_concatenated_video_latents(),
+        adapter.get_concatenated_video_latents(),
+    )
+
+
+def test_metadata_fields_persist_across_adapter_instances() -> None:
+    # A fresh adapter for the same session must see prior metadata (the manager
+    # is the single source of truth).
+    manager = SessionStateManager()
+    first = DreamZeroStateAdapter("s", manager)
+    first.current_start_frame = 5
+    first.clip_feas = torch.ones(2, 2)
+    first.append_video_latents(torch.ones(1, 2, 3, 4, 4))
+
+    second = DreamZeroStateAdapter("s", manager)
+    assert second.current_start_frame == 5
+    torch.testing.assert_close(second.clip_feas, torch.ones(2, 2))
+    # Video latents appended via `first` are visible via `second`.
+    assert second.get_concatenated_video_latents() is not None
+
+
+def test_session_lru_caps_retained_sessions() -> None:
+    manager = SessionStateManager(max_sessions=3)
+    for i in range(5):
+        DreamZeroStateAdapter(f"s{i}", manager)
+    assert len(manager) == 3
+    # Oldest two evicted; newest three retained.
+    assert "s0" not in manager
+    assert "s1" not in manager
+    for i in (2, 3, 4):
+        assert f"s{i}" in manager
+
+
+def test_adapter_keeps_state_when_its_session_is_evicted() -> None:
+    # Regression: an adapter mid-generation must not lose its data when the
+    # manager evicts its session to make room for newer ones. The adapter pins
+    # the session, so its stored latents survive eviction from the lookup table.
+    manager = SessionStateManager(max_sessions=2)
+    adapter = DreamZeroStateAdapter("active", manager)
+    chunk = torch.randn(1, 2, 3, 4, 4)
+    adapter.append_video_latents(chunk)
+    adapter.language = torch.tensor([1, 2, 3])
+
+    # Push "active" out of the manager's table with newer sessions.
+    for i in range(3):
+        DreamZeroStateAdapter(f"other{i}", manager)
+    assert "active" not in manager
+
+    # The in-use adapter still returns its own data.
+    assert adapter.get_concatenated_video_latents() is not None
+    torch.testing.assert_close(adapter.language, torch.tensor([1, 2, 3]))
+
+
+def test_session_reset_clears_metadata() -> None:
+    # Regression: resetting a session must not leave stale metadata behind, so a
+    # fresh adapter for the same session sees defaults, not the old values.
+    manager = SessionStateManager()
+    adapter = DreamZeroStateAdapter("s", manager)
+    adapter.language = torch.tensor([1, 2, 3])
+    adapter.current_start_frame = 5
+    adapter.clip_feas = torch.ones(2, 2)
+    adapter.prompt_embeds = torch.ones(1, 2)
+
+    adapter.reset()
+
+    fresh = DreamZeroStateAdapter("s", manager)
+    assert fresh.language is None
+    assert fresh.current_start_frame == 0
+    assert fresh.clip_feas is None
+    assert fresh.prompt_embeds is None
+
+
+def test_bespoke_path_without_memory_manager() -> None:
+    # Regression: a pipeline built via __new__ (lightweight fixtures) seeds only
+    # the bespoke fields and never sets _memory_manager; _get_or_create_state
+    # must fall through to the bespoke DreamZeroState path without AttributeError.
+    from collections import OrderedDict
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
+        MAX_DREAMZERO_SESSIONS,
+        DreamZeroPipeline,
+    )
+
+    pipe = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipe._states = OrderedDict()
+    pipe._max_session_states = MAX_DREAMZERO_SESSIONS
+
+    state = DreamZeroPipeline._get_or_create_state(pipe, "x")
+    assert isinstance(state, DreamZeroState)

tests/diffusion/models/dreamzero/test_session_state_lifecycle.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Session end-of-life routing (RFC #4480, adapting to #5271).
+
+The AR-Diffusion runner raises an explicit end-of-session signal on reset,
+close, eviction, and failed forwards (``reset_ar_diffusion_session`` /
+``close_ar_diffusion_session``). On the bespoke path those pop ``_states``; on
+the opt-in manager path the session lives in the manager instead, so the hooks
+must release it there or every closed session leaks its buffers.
+
+All CPU, tiny tensors, no model.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
+from vllm_omni.experimental.world_models.adapters.state_dreamzero_adapter import (
+    DreamZeroStateAdapter,
+)
+from vllm_omni.experimental.world_models.session_state import (
+    LatentBuffer,
+    SessionStateManager,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# -- SessionStateManager.drop_session ---------------------------------------
+
+
+def test_drop_session_absent_returns_false() -> None:
+    manager = SessionStateManager()
+    assert manager.drop_session("missing") is False
+
+
+def test_drop_session_removes_and_frees() -> None:
+    manager = SessionStateManager()
+    session = manager.get_or_create_session("s")
+    buffer: LatentBuffer[torch.Tensor] = LatentBuffer()
+    buffer.allocate(maxlen=None)
+    buffer.append(torch.zeros(256, dtype=torch.float32))
+    session.put("payload", buffer)
+    assert manager.nbytes_by_device().get("cpu", 0) >= 1024
+
+    assert manager.drop_session("s") is True
+    assert "s" not in manager
+    assert len(manager) == 0
+    # The session is gone from the table and its buffers were reset, so the
+    # manager reports no held bytes.
+    assert manager.nbytes_by_device() == {}
+
+
+def test_drop_session_recreates_fresh() -> None:
+    manager = SessionStateManager()
+    manager.get_or_create_session("s").attrs["k"] = 1
+    manager.drop_session("s")
+    assert manager.get_or_create_session("s").attrs == {}
+
+
+# -- pipeline hook routing ---------------------------------------------------
+
+
+def _manager_pipe(manager: SessionStateManager, session_id: str) -> DreamZeroPipeline:
+    """A pipeline built via __new__ whose state is a manager-backed adapter."""
+    pipe = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipe._states = OrderedDict()
+    pipe._memory_manager = manager
+    pipe.state = DreamZeroStateAdapter(session_id, manager, vae_encoder_window=2)
+    return pipe
+
+
+@pytest.mark.parametrize("hook", ["close_ar_diffusion_session", "reset_ar_diffusion_session"])
+def test_hook_releases_manager_session_and_clears_alias(hook: str) -> None:
+    manager = SessionStateManager()
+    pipe = _manager_pipe(manager, "sess")
+    assert "sess" in manager
+
+    getattr(DreamZeroPipeline, hook)(pipe, "sess")
+
+    assert "sess" not in manager
+    # The alias viewed the released session, so it is dropped and rebuilt lazily.
+    assert pipe.state is None
+
+
+def test_hook_keeps_alias_for_a_different_session() -> None:
+    manager = SessionStateManager()
+    pipe = _manager_pipe(manager, "default")
+    manager.get_or_create_session("other")
+
+    DreamZeroPipeline.close_ar_diffusion_session(pipe, "other")
+
+    assert "other" not in manager
+    # The alias views "default", not the closed session, so it survives.
+    assert isinstance(pipe.state, DreamZeroStateAdapter)
+    assert pipe.state.session_id == "default"
+
+
+def test_bespoke_hook_pops_states_and_clears_alias() -> None:
+    pipe = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipe._states = OrderedDict()
+    pipe._memory_manager = None
+    state = DreamZeroState()
+    pipe._states["sess"] = state
+    pipe.state = state
+
+    DreamZeroPipeline.close_ar_diffusion_session(pipe, "sess")
+
+    assert "sess" not in pipe._states
+    assert pipe.state is None

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -9,6 +9,8 @@ from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransforme
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def test_ltx_rms_norm_no_affine_matches_official_torch_module():
     norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)

tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py

@@ -6,6 +6,8 @@ import torch
 
 from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def _pipeline_without_init() -> SenseNovaU1Pipeline:
     return object.__new__(SenseNovaU1Pipeline)