[Core] Session state manager for AR-diffusion world models (RFC #4480 Phase 0) (#4487)

Signed-off-by: Wijnand Suijlen <wijnand.suijlen@huawei.com>

docs/features/session_state_manager.md

@@ -0,0 +1,160 @@
+# Session State Manager (experimental)
+
+!!! warning
+    This feature is **experimental** (RFC [#4480](https://github.com/vllm-project/vllm-omni/issues/4480)).
+    The code lives under `vllm_omni/experimental/world_models/` and its APIs may
+    change without notice. It is **off by default**; when disabled, model
+    behavior is byte-for-byte unchanged.
+
+Autoregressive diffusion world models (DreamZero, Cosmos3, and similar) keep
+per-session state across `forward()` calls: accumulated video latents, frame
+stitching buffers, encode-once conditioning, counters, and reset heuristics.
+Historically each model hand-rolled this state. The session state manager
+extracts it into a shared, typed contract:
+
+- `StateObject` — one typed unit of session state (e.g. `LatentBuffer` for
+  accumulated latents and bounded rings), with a uniform lifecycle
+  (`allocate` / `commit` / `view` / `reset` / `evict` / `nbytes_by_device`).
+- `SessionStateManager` — maps `session_id -> {name: StateObject}`, caps the
+  number of retained sessions by evicting the least recently used one first
+  (matching the bespoke caches it replaces), and reports per-device byte usage
+  via `stats()` for observability.
+
+Attention KV is *not* managed here: for DreamZero it is owned by the
+AR-Diffusion engine's paged KV pool (PR
+[#4534](https://github.com/vllm-project/vllm-omni/pull/4534)). The manager
+covers the model's non-KV session state; integrating the engine's per-session
+KV handle into the same byte accounting is RFC #4480 Phase 1.
+
+## Enabling
+
+Via config:
+
+```yaml
+# deploy config
+enable_session_state_manager: true
+```
+
+or programmatically with `OmniDiffusionConfig(enable_session_state_manager=True)`.
+
+Via environment variables (no config change needed):
+
+```bash
+export OMNI_DIFFUSION_SESSION_STATE_MANAGER=1            # 1/0/true/false/yes/no/on/off
+export OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS=64   # optional, positive int
+```
+
+**Precedence:** a *set* environment variable overrides the config value, in
+both directions — `OMNI_DIFFUSION_SESSION_STATE_MANAGER=0` force-disables the
+manager even if the config enables it. Unset or unparsable values fall back
+to the config (an unparsable `MAX_SESSIONS` logs a warning and is ignored).
+
+**Why both a config field and an environment variable?** The config field is
+the API. The environment override exists for A/B equivalence validation: the
+manager-backed path must be bit-identical to the bespoke path, and the way to
+verify that is to run the *same* deploy config twice with only the flag
+flipped per process — no config-file edits, so the two runs cannot drift
+apart. (Same pattern as the `DIFFUSION_CACHE_BACKEND` fallback.)
+
+**Why is `MAX_SESSIONS` not a config field?** The retained-session cap is not
+a public tuning knob: it mirrors the bespoke `MAX_DREAMZERO_SESSIONS = 64`
+constant (which is not configurable either), so the manager-backed path evicts
+exactly like the path it replaces. The environment variable exists only to
+stress eviction in tests and experiments. It is deliberately not promoted to
+the config, because RFC #4480 Phase 1 replaces count-based capping with a byte
+budget — that budget, not this cap, is the knob that deserves a config field.
+
+## Choosing between an attribute and a `StateObject`
+
+A session holds its values in two places. Named `StateObject`s live in the
+session's object map; everything else lives in `SessionState.attrs`, declared
+one line per value with the `SessionAttr` descriptor. The rule for choosing is:
+
+> A value belongs in a `StateObject` when the **manager** must be able to act on
+> it — release it under budget pressure, accumulate into it, stage and commit
+> it, or rebuild it from a recorded source. It belongs in `attrs` when the
+> manager only needs to **carry** it: the model assigns it wholesale, reads it
+> back, and clears it, and no one but the model can decide its fate.
+
+Byte accounting is independent of that choice. A session reports every byte it
+holds in either bucket, because the number describes what the process is
+holding, not what the manager may do about it.
+
+The consequence is worth stating plainly: **a value left in `attrs` is a value
+the byte-budget planner cannot free.** In Phase 0, which accounts but does not
+enforce, that is legitimate — but it should be a recorded decision rather than
+an accident, which is what the table below records.
+
+Note what the rule does *not* say. "Is it big" is not the criterion, and
+neither is "does it grow". The largest allocation a DreamZero session holds —
+`vae_enc_feat_map`, 603 MiB — is an *attribute*, because nothing the manager
+can do to it is useful: it is replaced wholesale, it has no recompute source,
+and freeing it mid-session breaks the encoder stream. Ending the session is the
+only release available, and that reclaims attributes anyway. Size alone argues
+for counting a value, which the accounting does regardless of bucket; it does
+not argue for a lifecycle.
+
+What would change that answer is a release the manager can perform without
+losing the value — offloading it to host memory and copying it back is
+transparent, needs no recomputation, and suits a large contiguous cache. The
+contract has no such operation: `evict()` frees storage, it does not park it.
+So the bucket follows the operations that exist, and moves if they do.
+
+### Byte accounting
+
+`SessionStateManager.stats()` reports byte totals keyed by device
+(`nbytes:cuda:0`, `nbytes:cpu`) alongside `total_nbytes`. The split is not
+cosmetic: device memory is the contended resource and host memory generally is
+not, so a single figure covering both answers no question. Which pool a budget
+applies to is a question for whoever enforces it, and is not decided here.
+
+Two details make the numbers honest. Bytes are counted per **storage**, not as
+`numel() * element_size()`, so a narrow slice reports the whole allocation it
+keeps alive rather than the part it exposes. And storages are **deduplicated**
+across the whole session, so two views of one buffer — or an alias between an
+attribute and a state object — are counted once.
+
+## Scope and guarantees
+
+- **Opt-in and equivalent.** With the flag off, models use their bespoke state
+  paths unchanged. With the flag on, the manager-backed path is validated
+  bit-identical to the bespoke path (CPU equivalence tests under
+  `tests/dreamzero/test_session_state_equivalence.py`, plus GPU A/B runs).
+- **Session eviction is count-based (LRU).** Evicting a session drops it from
+  the lookup table only; an adapter still holding the session keeps its state
+  (matching bespoke behavior, where the caller holds the state object).
+- **Byte budget is recorded, not enforced.** `SessionStateManager.stats()`
+  reports per-manager, per-device byte totals; budget *enforcement* (and
+  eviction driven by it) is a later phase of RFC #4480.
+
+## Supported models
+
+### DreamZero
+
+Sizes measured on one A100 at the shipped deploy config (`180x320` per camera,
+three cameras stitched to `352x640`, bfloat16), over a 400-step session.
+
+| Value | Bucket | Size | Notes |
+|---|---|---|---|
+| `vae_enc_feat_map` | `attrs` | **603 MiB** | Wan encoder causal-conv cache, 24 entries. Constant from ~step 50. The largest thing a session holds: 64 sessions is 37.7 GiB. Counted, not releasable — see the rule above. |
+| `video_latents_across_time` | `LatentBuffer` | grows | AR video latent chunks, host memory (`.cpu()`), kept for decode. |
+| `stitched_buffer` | `LatentBuffer` | bounded | Stitched pixel frames, host memory, ring of `FRAMES_PER_CHUNK`. |
+| `vae_encoder_out` | `LatentBuffer` | 0.43 MiB | Bounded ring of `num_frame_per_block` latent frames — see below. |
+| `clip_feas`, `ys`, `language`, `prompt_embeds` | `attrs` | small | Conditioning tensors, written once per session and reread. |
+| `vae_pending_body_frames` | `attrs` | ≤2.6 MiB | Raw frames awaiting a 4-frame body chunk; bounded by construction. |
+| `call_count`, `current_start_frame`, `vae_stream_initialized` | `attrs` | — | Counters and flags. |
+
+Attention KV does not appear here: it is engine-owned (PR #4534).
+
+**The bounded VAE encoder ring.** The bespoke state concatenates every encoded
+chunk into one tensor and clears it only on a *session* reset — a window reset
+deliberately keeps it — so it grows for as long as the session lives, at
+0.215 MiB per latent frame (one per four steps). Its only reader,
+`_vae_stream_get_observation_latents`, quantises the whole history and then
+keeps the last `num_frame_per_block` latent frames. Since `quant_conv` is a
+`WanCausalConv3d` of kernel size 1, pointwise in time, those frames do not
+depend on the discarded ones. The manager-backed state therefore retains a ring
+of `num_frame_per_block` frames instead: identical output, bounded memory, and
+the per-step quantisation no longer scales with session length. Asking for more
+frames than the ring retains raises rather than silently returning a shorter
+window. The bespoke path is unchanged.

tests/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/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/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

vllm_omni/diffusion/data.py

@@ -625,6 +625,13 @@ class OmniDiffusionConfig:
     enable_prompt_embed_cache: bool = False
     prompt_embed_cache_size: int = 32
 
+    # Opt-in: route per-session world-model state through the shared
+    # SessionStateManager (RFC #4480) instead of the model's bespoke cache.
+    # Default off; the bespoke path remains the default. A *set*
+    # OMNI_DIFFUSION_SESSION_STATE_MANAGER environment variable overrides this
+    # in both directions. See docs/features/session_state_manager.md.
+    enable_session_state_manager: bool = False
+
     # Distributed executor backend
     distributed_executor_backend: str = "mp"
     nccl_port: int | None = None

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -61,11 +61,22 @@ from vllm_omni.experimental.ar_diffusion.capability import (
     ARDiffusionKVBranchSpec,
     ARDiffusionKVCacheSpec,
 )
+from vllm_omni.experimental.world_models.adapters.state_dreamzero_adapter import DreamZeroStateAdapter
+from vllm_omni.experimental.world_models.session_state import (
+    SessionStateManager,
+    resolve_session_state_config,
+)
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 logger = logging.getLogger(__name__)
 MAX_DREAMZERO_SESSIONS = 64
 
+# The pipeline's per-session state is a bespoke ``DreamZeroState`` by default, or
+# a ``DreamZeroStateAdapter`` view when the opt-in session manager is enabled.
+# The adapter mirrors ``DreamZeroState``'s surface, so every helper that reads or
+# writes session state accepts either.
+DreamZeroSessionState = DreamZeroState | DreamZeroStateAdapter
+
 
 class VideoActionScheduler:
     """Wraps video + action schedulers into single .step() interface."""
@@ -110,7 +121,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     _POSITIVE_BRANCH = "positive"
     _NEGATIVE_BRANCH = "negative"
     _ar_diffusion_kv_state = None
-    state: DreamZeroState | None
+    state: DreamZeroSessionState | None
 
     def ar_diffusion_kv_cache_spec(self) -> ARDiffusionKVCacheSpec:
         """Describe DreamZero's local KV geometry to the generic runner."""
@@ -163,7 +174,21 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     def _drop_ar_diffusion_session_state(self, session_id: str) -> None:
         """Remove model state and clear the compatibility alias when it points there."""
-        removed = self._states.pop(str(session_id or "default"), None)
+        key = str(session_id or "default")
+        # Local binding narrows the Optional and guards lightweight test fixtures
+        # that build the pipeline via __new__ without setting _memory_manager.
+        manager = getattr(self, "_memory_manager", None)
+        if manager is not None:
+            # Manager-backed path: the session lives in the manager, not in
+            # ``_states``. The runner's explicit close/reset is the session's
+            # end-of-life signal, so release it there (freeing its buffers) and
+            # drop the alias if it still views this session.
+            manager.drop_session(key)
+            current = getattr(self, "state", None)
+            if isinstance(current, DreamZeroStateAdapter) and str(current.session_id or "default") == key:
+                self.state = None
+            return
+        removed = self._states.pop(key, None)
         if removed is not None and getattr(self, "state", None) is removed:
             self.state = None
 
@@ -433,7 +458,22 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             use_dynamic_shifting=False,
         )
 
+        # Read before the first `_get_or_create_state` below: the manager-backed
+        # state bounds its VAE encoder history to this many latent frames.
+        self.num_frame_per_block: int = ah_config["num_frame_per_block"]
+
         self._states: OrderedDict[str, DreamZeroState] = OrderedDict()
+        # Opt-in: back per-session state with the shared SessionStateManager
+        # (RFC #4480). Default off -> the bespoke DreamZeroState path above.
+        self._use_memory_manager, mm_max_sessions = resolve_session_state_config(
+            enable=od_config.enable_session_state_manager,
+            max_sessions=MAX_DREAMZERO_SESSIONS,
+        )
+        self._memory_manager: SessionStateManager | None = (
+            SessionStateManager(max_sessions=mm_max_sessions) if self._use_memory_manager else None
+        )
+        if self._use_memory_manager:
+            logger.info("DreamZero: session state manager enabled (max_sessions=%d)", mm_max_sessions)
         self.state = self._get_or_create_state("default")
 
         # DiT step cache is configured by StepCacheBackend
@@ -447,7 +487,6 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         self.cfg_scale: float = model_config.get("cfg_scale", DEFAULT_CFG_SCALE)
         self.sigma_shift: float = model_config.get("sigma_shift", DEFAULT_SIGMA_SHIFT)
         self.num_frames: int = ah_config["num_frames"]
-        self.num_frame_per_block: int = ah_config["num_frame_per_block"]
         self.action_horizon: int = ah_config["action_horizon"]
 
         self.decouple_inference_noise: bool = ah_config["decouple_inference_noise"]
@@ -502,7 +541,18 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             ),
         ]
 
-    def _get_or_create_state(self, session_id: str | None) -> DreamZeroState:
+    def _get_or_create_state(self, session_id: str | None) -> DreamZeroState | DreamZeroStateAdapter:
+        # getattr guards lightweight test fixtures that build the pipeline via
+        # __new__ and seed only the bespoke fields, never setting _memory_manager.
+        if getattr(self, "_memory_manager", None) is not None:
+            # The manager owns session lifecycle and LRU; the adapter is a thin
+            # per-call view over it (no second LRU to drift).
+            return DreamZeroStateAdapter(
+                session_id,
+                self._memory_manager,
+                vae_encoder_window=self.num_frame_per_block,
+            )
+
         session_key = str(session_id or "default")
         state = self._states.get(session_key)
         if state is None:
@@ -807,7 +857,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         height: int,
         width: int,
         *,
-        state: DreamZeroState | None = None,
+        state: DreamZeroSessionState | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Encode first frame via CLIP + VAE.
         Returns: (clip_feas, ys, image_latent)
@@ -898,7 +948,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         inv_std = self.vae_latents_inv_std.to(device=mu.device, dtype=mu.dtype)
         return (mu - mean) * inv_std
 
-    def _vae_stream_seed(self, state: DreamZeroState, first_frame: torch.Tensor) -> None:
+    def _vae_stream_seed(self, state: DreamZeroSessionState, first_frame: torch.Tensor) -> None:
         """Seed incremental VAE encode with the first observation frame."""
         state.reset_vae_encoder_stream()
         feat_map = self._vae_init_enc_feat_map()
@@ -910,7 +960,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         state.vae_pending_body_frames = None
         state.vae_stream_initialized = True
 
-    def _vae_stream_append_frame(self, state: DreamZeroState, new_frame: torch.Tensor) -> None:
+    def _vae_stream_append_frame(self, state: DreamZeroSessionState, new_frame: torch.Tensor) -> None:
         """Append one pixel frame and encode a 4-frame body chunk when ready."""
         if not state.vae_stream_initialized or state.vae_enc_feat_map is None or state.vae_encoder_out is None:
             raise RuntimeError("VAE encoder stream is not initialized.")
@@ -929,19 +979,27 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                 state.vae_pending_body_frames = None
             chunk = self._vae_patchify(body)
             encoder_chunk, feat_map = self._vae_encode_encoder_chunk(chunk, state.vae_enc_feat_map)
-            state.vae_encoder_out = torch.cat([state.vae_encoder_out, encoder_chunk], dim=2)
+            state.append_vae_encoder_chunk(encoder_chunk)
             state.vae_enc_feat_map = feat_map
 
     def _vae_stream_get_observation_latents(
         self,
-        state: DreamZeroState,
+        state: DreamZeroSessionState,
         num_latent_frames: int,
         *,
         dtype: torch.dtype,
     ) -> torch.Tensor:
-        if state.vae_encoder_out is None:
+        # Read once: the manager-backed state rebuilds this from its ring.
+        encoder_out = state.vae_encoder_out
+        if encoder_out is None:
             raise RuntimeError("VAE encoder stream has no accumulated encoder output.")
-        latents = self._vae_quantize_encoder_out(state.vae_encoder_out).to(dtype=dtype)
+        # A state that bounds its encoder history keeps only the most recent
+        # frames. Asking for more than it retains would silently return a
+        # shorter (and differently padded) window, so fail instead.
+        window = state.vae_encoder_window
+        if window is not None and num_latent_frames > window:
+            raise ValueError(f"Requested {num_latent_frames} latent frames but the session retains at most {window}.")
+        latents = self._vae_quantize_encoder_out(encoder_out).to(dtype=dtype)
         if latents.shape[2] >= num_latent_frames:
             return latents[:, :, -num_latent_frames:]
         pad_count = num_latent_frames - latents.shape[2]
@@ -962,7 +1020,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     def _encode_observation_latents(
         self,
-        state: DreamZeroState,
+        state: DreamZeroSessionState,
         videos: torch.Tensor,
         *,
         latent_dtype: torch.dtype,
@@ -1029,7 +1087,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         frame_seqlen: int,
         seq_len: int,
         do_true_cfg: bool,
-        state: DreamZeroState,
+        state: DreamZeroSessionState,
     ) -> None:
         """Prefill KV cache with first frame and/or current observation.
 
@@ -1155,7 +1213,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         negative_prompt_embeds: torch.Tensor | None,
         video_action_scheduler: VideoActionScheduler,
         do_true_cfg: bool,
-        state: DreamZeroState,
+        state: DreamZeroSessionState,
         **kwargs,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Denoising loop with CFG parallel support.
@@ -1211,10 +1269,13 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             )
 
             csf = state.current_start_frame
-            if csf + self.num_frame_per_block <= state.ys.shape[2]:
-                y = state.ys[:, :, csf : csf + self.num_frame_per_block]
+            ys = state.ys
+            if ys is None:
+                raise RuntimeError("diffuse() requires state.ys, populated by the forward pass before denoising")
+            if csf + self.num_frame_per_block <= ys.shape[2]:
+                y = ys[:, :, csf : csf + self.num_frame_per_block]
             else:
-                y = state.ys[:, :, -self.num_frame_per_block :]
+                y = ys[:, :, -self.num_frame_per_block :]
 
             run_dit = _step_cache is None or _step_cache.should_run_step(_prev_predictions)
             if run_dit:

vllm_omni/diffusion/models/dreamzero/state_dreamzero.py

@@ -134,6 +134,18 @@ class DreamZeroState:
         self.vae_encoder_out: torch.Tensor | None = None
         self.vae_pending_body_frames: torch.Tensor | None = None
 
+    # How many latent frames of encoder output this state retains. ``None``
+    # means "everything", which is what this path has always done -- the
+    # accumulated output is cleared only by a session reset, so it grows for as
+    # long as the session lives. The manager-backed adapter bounds it instead.
+    vae_encoder_window: int | None = None
+
+    def append_vae_encoder_chunk(self, encoder_chunk: torch.Tensor) -> None:
+        """Append one encoded chunk to the accumulated encoder output."""
+        if self.vae_encoder_out is None:
+            raise RuntimeError("VAE encoder stream is not seeded; call reset_vae_encoder_stream first.")
+        self.vae_encoder_out = torch.cat([self.vae_encoder_out, encoder_chunk], dim=2)
+
     def reset_inference_state(self) -> None:
         """Reset KV/frame state after local attention rolls without dropping video latents."""
         self.reset(clear_video_latents=False)

vllm_omni/experimental/world_models/__init__.py

@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""World-model session state (RFC #4480).
+
+Model-agnostic session state for AR-diffusion world models: a typed
+``StateObject`` contract and a ``SessionStateManager`` that owns the objects
+per session, plus per-model adapters that present a model's own per-session
+cache over that contract. APIs may change without notice.
+"""
+
+from __future__ import annotations

vllm_omni/experimental/world_models/adapters/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-model adapters onto the world-model session state contract (RFC #4480)."""
+
+from __future__ import annotations

vllm_omni/experimental/world_models/adapters/state_dreamzero_adapter.py

@@ -0,0 +1,315 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Adapter presenting ``DreamZeroState``'s surface over the session manager.
+
+``DreamZeroStateAdapter`` exposes the exact public methods and attributes that
+``pipeline_dreamzero.py`` touches on its state object, so the pipeline can use
+it interchangeably with the bespoke ``DreamZeroState`` (behind the opt-in flag).
+
+DreamZero's attention KV lives in the AR-Diffusion engine's paged pool
+(``self._ar_diffusion_kv_state``) and is *not* managed here; this adapter
+covers the model's non-KV session state. Storage is delegated to typed
+``StateObject`` instances owned by the ``SessionStateManager``:
+
+    * the stitched frame buffer            -> ``LatentBuffer`` (ring)
+    * accumulated AR video latents         -> ``LatentBuffer`` (append, CPU)
+    * recent VAE encoder output            -> ``LatentBuffer`` (ring, bounded)
+
+The remaining metadata (counters, conditioning tensors, the Wan encoder
+causal-conv cache) lives in the session's ``attrs``, so a freshly constructed
+adapter for an existing session sees the same data (the manager is the single
+source of truth and the single LRU authority). Which values earn a
+``StateObject`` and which stay attributes follows the rule in
+``docs/features/session_state_manager.md``: a ``StateObject`` when the manager
+must be able to act on the value, an attribute when it only carries it.
+
+The public methods mirror ``DreamZeroState`` statement for statement, so the two
+can be diffed to check equivalence. The deviations are forced by the buffer
+members being ``StateObject``s rather than a plain deque/list/tensor: a read
+uses ``.view()`` where the original reads ``list(...)`` or indexes/iterates the
+container, and a wholesale replacement uses ``self._session.put(...)`` where the
+original reassigns the attribute. The one behavioural difference is deliberate
+and documented on ``_ensure_encoder_buffer``: the encoder output is bounded
+here and unbounded in the original, which the pointwise ``quant_conv`` makes
+numerically indistinguishable to the only reader.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.state_dreamzero import FRAMES_PER_CHUNK
+from vllm_omni.experimental.world_models.session_state.attrs import SessionAttr
+from vllm_omni.experimental.world_models.session_state.manager import SessionStateManager
+from vllm_omni.experimental.world_models.session_state.objects import LatentBuffer
+
+logger = logging.getLogger(__name__)
+
+_FRAMES = "frames"
+_VIDEO_LATENTS = "video_latents"
+_VAE_ENCODER_OUT = "vae_encoder_out"
+
+# Latent frames of VAE encoder output retained by default. DreamZero passes its
+# own ``num_frame_per_block``; this is only the fallback for callers that build
+# an adapter directly (tests, other entry points).
+DEFAULT_VAE_ENCODER_WINDOW = 2
+
+
+class DreamZeroStateAdapter:
+    """Drop-in replacement for ``DreamZeroState`` backed by the manager.
+
+    Session-scoped metadata is declared once per attribute via ``SessionAttr``
+    (reads fall back to the default; writes go through to the pinned session).
+    """
+
+    call_count = SessionAttr[int](default=0, coerce=int)
+    current_start_frame = SessionAttr[int](default=0, coerce=int)
+    clip_feas = SessionAttr[torch.Tensor | None](default=None)
+    ys = SessionAttr[torch.Tensor | None](default=None)
+    language = SessionAttr[torch.Tensor | None](default=None)
+    prompt_embeds = SessionAttr[torch.Tensor | None](default=None)
+    # -- incremental VAE encoder stream --
+    vae_stream_initialized = SessionAttr[bool](default=False, coerce=bool)
+    # The Wan encoder's per-layer causal convolution cache: 24 entries, 603 MiB
+    # per session, constant from the first few steps onward and by far the
+    # largest thing a DreamZero session holds. It is an attribute all the same:
+    # it is replaced wholesale, it has no recompute source (rebuilding it means
+    # re-encoding the whole observation history, which is not retained), and
+    # freeing it mid-session breaks the stream, so ending the session is the
+    # only release -- which reclaims attributes anyway.
+    vae_enc_feat_map = SessionAttr[list[torch.Tensor | None] | None](default=None)
+    vae_pending_body_frames = SessionAttr[torch.Tensor | None](default=None)
+
+    def __init__(
+        self,
+        session_id: str | None,
+        manager: SessionStateManager,
+        *,
+        vae_encoder_window: int = DEFAULT_VAE_ENCODER_WINDOW,
+    ) -> None:
+        self._session_id = session_id
+        if vae_encoder_window <= 0:
+            raise ValueError(f"vae_encoder_window must be positive, got {vae_encoder_window}")
+        self.vae_encoder_window = vae_encoder_window
+        # Pin the session for this adapter's lifetime. The manager may evict the
+        # session from its lookup table to bound memory, but an adapter that is
+        # mid-generation keeps its own reference, so the in-progress state is not
+        # lost -- matching the bespoke DreamZeroState, which the caller holds even
+        # after it leaves the table. A fresh adapter is built per forward, so the
+        # session is still marked recently-used on each step.
+        self._session = manager.get_or_create_session(session_id)
+        self._ensure_frame_buffer()
+
+    @property
+    def session_id(self) -> str | None:
+        """The session this adapter views, as passed at construction.
+
+        Lets the pipeline match the alias against a runner-supplied session id
+        when the engine signals that session's close/reset.
+        """
+        return self._session_id
+
+    # -- session / metadata plumbing ------------------------------------
+
+    @staticmethod
+    def _new_frame_buffer() -> LatentBuffer[np.ndarray]:
+        buffer: LatentBuffer[np.ndarray] = LatentBuffer()
+        buffer.allocate(maxlen=FRAMES_PER_CHUNK)
+        return buffer
+
+    @staticmethod
+    def _new_video_buffer() -> LatentBuffer[torch.Tensor]:
+        buffer: LatentBuffer[torch.Tensor] = LatentBuffer()
+        buffer.allocate(maxlen=None)
+        return buffer
+
+    def _ensure_frame_buffer(self) -> LatentBuffer[np.ndarray]:
+        """The stitched-frame ring for this session (created on first use)."""
+        return self._session.get_or_create(_FRAMES, self._new_frame_buffer, LatentBuffer)
+
+    def _ensure_video_buffer(self) -> LatentBuffer[torch.Tensor]:
+        """The accumulated AR video-latent chunks (unbounded, CPU)."""
+        return self._session.get_or_create(_VIDEO_LATENTS, self._new_video_buffer, LatentBuffer)
+
+    def _new_encoder_buffer(self) -> LatentBuffer[torch.Tensor]:
+        buffer: LatentBuffer[torch.Tensor] = LatentBuffer()
+        buffer.allocate(maxlen=self.vae_encoder_window)
+        return buffer
+
+    def _ensure_encoder_buffer(self) -> LatentBuffer[torch.Tensor]:
+        """The recent VAE encoder output, as a ring of per-chunk tensors.
+
+        The bespoke state concatenates every encoded chunk into one tensor that
+        is cleared only by a *session* reset -- a window reset deliberately
+        keeps it -- so it grows for as long as the session lives. Its only
+        reader, ``_vae_stream_get_observation_latents``, quantises the whole
+        history and then keeps the last ``num_frame_per_block`` latent frames.
+        ``quant_conv`` is a ``WanCausalConv3d`` of kernel size 1, pointwise in
+        time, so those last frames do not depend on the discarded ones: keeping
+        a ring of ``vae_encoder_window`` frames is numerically identical and
+        bounded.
+        """
+        return self._session.get_or_create(_VAE_ENCODER_OUT, self._new_encoder_buffer, LatentBuffer)
+
+    @property
+    def vae_encoder_out(self) -> torch.Tensor | None:
+        """The retained encoder output as one tensor, or ``None`` when empty."""
+        chunks = self._ensure_encoder_buffer().view()
+        if not chunks:
+            return None
+        if len(chunks) == 1:
+            return chunks[0]
+        return torch.cat(chunks, dim=2)
+
+    @vae_encoder_out.setter
+    def vae_encoder_out(self, value: torch.Tensor | None) -> None:
+        # Seeding (``_vae_stream_seed``) and clearing (``reset``) both assign
+        # wholesale; growth goes through ``append_vae_encoder_chunk``.
+        buffer = self._new_encoder_buffer()
+        if value is not None:
+            buffer.append(value)
+        self._session.put(_VAE_ENCODER_OUT, buffer)
+
+    def append_vae_encoder_chunk(self, encoder_chunk: torch.Tensor) -> None:
+        """Append one encoded chunk, evicting the oldest beyond the window."""
+        self._ensure_encoder_buffer().append(encoder_chunk)
+
+    @property
+    def stitched_buffer(self) -> LatentBuffer[np.ndarray]:
+        return self._ensure_frame_buffer()
+
+    @property
+    def video_latents_across_time(self) -> LatentBuffer[torch.Tensor]:
+        # Read access only; a wholesale replace goes through ``_session.put``
+        # (see reset / clear_video_latents).
+        return self._ensure_video_buffer()
+
+    def reset_vae_encoder_stream(self) -> None:
+        """Clear incremental Wan VAE encoder state used across AR steps."""
+        self.vae_stream_initialized = False
+        self.vae_enc_feat_map = None
+        self.vae_encoder_out = None
+        self.vae_pending_body_frames = None
+
+    # -- frame accumulation ----------------------------------------------
+
+    def accumulate_frames(self, stitched: np.ndarray) -> np.ndarray:
+        """Accumulate stitched frames and return multi-frame video."""
+        if stitched.ndim == 3:
+            self.stitched_buffer.append(stitched)
+        elif stitched.ndim == 4:
+            self.stitched_buffer.extend(list(stitched))
+        else:
+            raise ValueError(f"Expected 3D or 4D stitched, got {stitched.ndim}D")
+
+        num_frames = 1 if self.call_count == 0 else FRAMES_PER_CHUNK
+
+        buffer_frames = self.stitched_buffer.view()
+        if len(buffer_frames) >= num_frames:
+            frames = buffer_frames[-num_frames:]
+        else:
+            frames = buffer_frames
+            while len(frames) < num_frames:
+                frames.insert(0, buffer_frames[0])
+
+        self.call_count += 1
+        return np.stack(frames, axis=0)
+
+    # -- video-latent accumulation ---------------------------------------
+
+    def append_video_latents(self, video_out: torch.Tensor) -> None:
+        """Append one AR chunk of normalized video latents for later decode."""
+        if video_out.dim() != 5:
+            raise ValueError(f"Expected 5D video_out, got shape {tuple(video_out.shape)}")
+        # Upstream ``torch.cat(..., dim=2)`` uses (B, C, T, H, W).
+        chunk = video_out.transpose(1, 2).detach().cpu()
+        self.video_latents_across_time.append(chunk)
+        logger.info(
+            "append_video_latents: chunk_shape=%s total_chunks=%d total_latent_t=%d",
+            tuple(chunk.shape),
+            len(self.video_latents_across_time),
+            int(sum(c.shape[2] for c in self.video_latents_across_time.view())),
+        )
+
+    def get_concatenated_video_latents(self) -> torch.Tensor | None:
+        """Return all accumulated chunks concatenated along the time dimension."""
+        chunks = self.video_latents_across_time.view()
+        if not chunks:
+            return None
+        if len(chunks) == 1:
+            return chunks[0]
+        return torch.cat(chunks, dim=2)
+
+    def clear_video_latents(self) -> None:
+        """Drop accumulated video latents without resetting frame/VAE state."""
+        self._session.put(_VIDEO_LATENTS, self._new_video_buffer())
+
+    # -- reset / should_reset --------------------------------------------
+
+    def reset(self, *, clear_video_latents: bool = True) -> None:
+        """Clear session state.
+
+        A session reset (``clear_video_latents=True``) drops everything,
+        including the prompt-embed cache and the incremental VAE encoder stream.
+        A window ("inference") reset (``False``) keeps the accumulated video
+        latents, ``language``, ``prompt_embeds``, and the VAE encoder stream --
+        the prompt is unchanged and the Wan feat_cache history is independent of
+        the DiT attention window. The survivors are kept the way the original
+        keeps them: by leaving them untouched inside the skipped ``if`` block.
+        """
+        saved_video_latents = [] if clear_video_latents else self._ensure_video_buffer().view()
+        self._session.put(_FRAMES, self._new_frame_buffer())
+        self.call_count = 0
+        restored_video = self._new_video_buffer()
+        restored_video.extend(saved_video_latents)
+        self._session.put(_VIDEO_LATENTS, restored_video)
+
+        self.current_start_frame = 0
+
+        self.clip_feas = None
+        self.ys = None
+        if clear_video_latents:
+            # Session reset: drop the prompt-embed cache and VAE encoder stream.
+            self.language = None
+            self.prompt_embeds = None
+            self.reset_vae_encoder_stream()
+
+    def reset_inference_state(self) -> None:
+        """Reset window/frame state after local attention rolls without dropping video latents."""
+        self.reset(clear_video_latents=False)
+
+    def reset_reason(
+        self,
+        text_tokens: torch.Tensor | None,
+        num_video_frames: int,
+        local_attn_size: int,
+    ) -> str | None:
+        """Return why state should reset before the next forward(), if any."""
+        if self.language is None:
+            logger.info("language is None, resetting")
+            return "session"
+
+        if text_tokens is not None and not torch.equal(self.language, text_tokens):
+            logger.info("language changed, resetting")
+            return "session"
+
+        if num_video_frames == 1 and self.call_count > 1:
+            logger.info("single frame input after first call, resetting")
+            return "session"
+
+        if local_attn_size != -1 and self.current_start_frame >= local_attn_size:
+            logger.info(
+                "current_start_frame %d >= local_attn_size %d, resetting inference state",
+                self.current_start_frame,
+                local_attn_size,
+            )
+            return "inference"
+
+        return None
+
+    def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
+        """Determine if state should be reset before this forward()."""
+        return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None

vllm_omni/experimental/world_models/session_state/__init__.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Session memory for AR-diffusion world models (RFC #4480).
+
+Typed ``StateObject`` instances owned per session by ``SessionStateManager``.
+This package is model-agnostic; a model opts in by adapting its own per-session
+cache onto this contract (e.g. ``state_dreamzero_adapter.DreamZeroStateAdapter``).
+See RFC #4480 for the full design and roadmap.
+"""
+
+from vllm_omni.experimental.world_models.session_state.accounting import (
+    SeenStorages,
+    device_bytes,
+)
+from vllm_omni.experimental.world_models.session_state.attrs import SessionAttr
+from vllm_omni.experimental.world_models.session_state.base import StateObject
+from vllm_omni.experimental.world_models.session_state.manager import (
+    DEFAULT_MAX_SESSIONS,
+    SessionState,
+    SessionStateManager,
+    resolve_session_state_config,
+)
+from vllm_omni.experimental.world_models.session_state.objects import LatentBuffer
+
+__all__ = [
+    "DEFAULT_MAX_SESSIONS",
+    "LatentBuffer",
+    "SeenStorages",
+    "StateObject",
+    "SessionAttr",
+    "SessionState",
+    "SessionStateManager",
+    "device_bytes",
+    "resolve_session_state_config",
+]

vllm_omni/experimental/world_models/session_state/accounting.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-device byte accounting for session state (RFC #4480).
+
+A session holds bytes in more than one place -- device tensors, host tensors,
+NumPy frame buffers -- and those are different resources: a GPU pool that is
+scarce and contended, and host memory that is not. Summing them into a single
+number answers no question anyone has, so the accounting here is keyed by the
+device the bytes live on and is never reduced to a scalar by this module.
+
+Two details make the numbers honest:
+
+* **Storages, not elements.** ``numel() * element_size()`` measures the logical
+  view, but a narrow slice of a large buffer keeps the whole allocation alive.
+  Counting ``untyped_storage().nbytes()`` reports what is actually pinned.
+* **Views deduplicated.** Two tensors sharing one storage are counted once, via
+  the caller-supplied ``seen`` set. Pass the *same* set across every value in a
+  session so an alias between two state objects (or between an object and an
+  attribute) does not double-count.
+
+Anything that is not a tensor, an array, or a container of those contributes
+nothing. That is not a heuristic: in this codebase the only device memory a
+session can hold is a ``torch.Tensor``. CUDA graph pools, collective buffers
+and the paged KV blocks of PR #4534 are owned by the engine, never by a session.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+import numpy as np
+import torch
+
+# ``(device string, storage address)`` of a storage already counted.
+SeenStorages = set[tuple[str, int]]
+
+_HOST = "cpu"
+
+
+def device_bytes(value: object, seen: SeenStorages) -> dict[str, int]:
+    """Bytes pinned by ``value``, keyed by the device they live on.
+
+    ``seen`` accumulates the storages already counted and is mutated in place;
+    share one set across a whole session to deduplicate aliases.
+    """
+    totals: dict[str, int] = {}
+    _accumulate(value, seen, totals)
+    return totals
+
+
+def merge_bytes(into: dict[str, int], other: Mapping[str, int]) -> dict[str, int]:
+    """Add ``other`` into ``into`` per device, in place."""
+    for device, count in other.items():
+        into[device] = into.get(device, 0) + count
+    return into
+
+
+def _accumulate(value: object, seen: SeenStorages, totals: dict[str, int]) -> None:
+    if isinstance(value, torch.Tensor):
+        storage = value.untyped_storage()
+        device = str(value.device)
+        key = (device, storage.data_ptr())
+        if key in seen:
+            return
+        seen.add(key)
+        totals[device] = totals.get(device, 0) + storage.nbytes()
+        return
+
+    if isinstance(value, np.ndarray):
+        # ``base is not None`` marks a view; NumPy views are cheap to alias and
+        # carry no address we can key on as reliably as a torch storage, so the
+        # owning array is what gets counted.
+        owner = value if value.base is None else value.base
+        if isinstance(owner, np.ndarray):
+            key = (_HOST, owner.__array_interface__["data"][0])
+            if key in seen:
+                return
+            seen.add(key)
+            totals[_HOST] = totals.get(_HOST, 0) + int(owner.nbytes)
+        return
+
+    if isinstance(value, Mapping):
+        for item in value.values():
+            _accumulate(item, seen, totals)
+        return
+
+    # ``str``/``bytes`` are Sequences but hold no device memory; excluding them
+    # also stops the recursion from walking a string one character at a time.
+    if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
+        for item in value:
+            _accumulate(item, seen, totals)
+        return

vllm_omni/experimental/world_models/session_state/attrs.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Typed session-attribute descriptor (RFC #4480).
+
+``SessionAttr`` lets an adapter declare each piece of session-scoped scalar or
+tensor metadata as one class-level line instead of a hand-written property
+pair. The descriptor reads and writes ``SessionState.attrs`` on the host
+object's pinned ``_session``, using the attribute's own name as the key, so
+each key string exists exactly once. Because the read falls back to the
+declared default when the key is absent, an adapter clears a value simply by
+assigning the default (or by clearing ``attrs``), keeping its method bodies a
+line-for-line mirror of the model's original state class.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Generic, Protocol, TypeVar, cast, overload
+
+from vllm_omni.experimental.world_models.session_state.manager import SessionState
+
+T = TypeVar("T")
+
+
+class HasSessionState(Protocol):
+    """Host contract: any object holding a pinned ``SessionState``."""
+
+    _session: SessionState
+
+
+class SessionAttr(Generic[T]):
+    """Data descriptor for one session-scoped metadata value.
+
+    Declared at class level on an adapter::
+
+        class SomeAdapter:
+            call_count = SessionAttr[int](default=0, coerce=int)
+            language = SessionAttr[torch.Tensor | None](default=None)
+
+    Reads return ``attrs.get(<name>, default)``; writes go through ``coerce``
+    (when given) and straight into the session's ``attrs``, so a freshly
+    constructed adapter for an existing session sees the same data.
+    """
+
+    def __init__(
+        self,
+        *,
+        default: T,
+        coerce: Callable[[T], T] | None = None,
+    ) -> None:
+        self.default = default
+        self.coerce = coerce
+        self.name = ""  # assigned by ``__set_name__``
+
+    def __set_name__(self, owner: type, name: str) -> None:
+        self.name = name
+
+    @overload
+    def __get__(self, obj: None, objtype: type | None = None) -> SessionAttr[T]: ...
+
+    @overload
+    def __get__(self, obj: HasSessionState, objtype: type | None = None) -> T: ...
+
+    def __get__(self, obj: HasSessionState | None, objtype: type | None = None) -> SessionAttr[T] | T:
+        if obj is None:
+            return self
+        return cast("T", obj._session.attrs.get(self.name, self.default))
+
+    def __set__(self, obj: HasSessionState, value: T) -> None:
+        obj._session.attrs[self.name] = self.coerce(value) if self.coerce is not None else value

vllm_omni/experimental/world_models/session_state/base.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The ``StateObject`` contract for session state (RFC #4480).
+
+A session is a named collection of ``StateObject`` instances. The interface
+deliberately speaks only of bytes, growth, read views, commitment, and
+evictability -- never tokens, layers, or attention. Writes are two-phase:
+``stage()`` holds data for an in-progress window, ``commit()`` promotes it to
+persistent context, and ``discard()`` drops it.
+
+A concrete class may back its storage with plain (non-paged) buffers; the
+contract itself does not require that. Several methods are
+intentionally inert and documented as such; they exist so that paged-KV
+backing and speculative forking can be added without changing any call site.
+See RFC #4480 for the full design and roadmap.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Generic, TypeVar
+
+from vllm_omni.experimental.world_models.session_state.accounting import SeenStorages
+
+# What writers pass into ``stage()`` / ``commit()`` / ``append()``.
+PayloadT = TypeVar("PayloadT")
+# What ``view()`` hands to the consumer (not necessarily the payload type:
+# an append buffer views as a list of payloads).
+ViewT = TypeVar("ViewT")
+
+
+class StateObject(ABC, Generic[PayloadT, ViewT]):
+    """One typed unit of session state with a uniform lifecycle.
+
+    A session's state is a named collection of these objects; the kinds of
+    memory they represent include attention K/V (growing self-attention
+    caches, or conditioning K/V written once and reread), constant-size
+    recurrent state (e.g. a linear-attention carry), frame or latent buffers,
+    and retrieval pools. A concrete class binds ``PayloadT`` and ``ViewT`` to
+    its real payload and view types and implements the abstract methods; the
+    non-abstract methods provide the default behaviour described in their
+    docstrings.
+
+    The contract is storage-agnostic. A concrete class may hold its bytes in
+    a buffer it owns, or it may be a handle to storage owned by another
+    component, such as a paged KV pool. In the handle case it delegates its
+    lifecycle operations to that owner instead of copying the data.
+    """
+
+    # The dependency edge: what this object rebuilds from. ``None`` means it is
+    # never recomputable (it must be preempted or rejected, not silently
+    # dropped).
+    recompute_source: StateObject | None = None
+    # Estimated work to rebuild, given ``recompute_source`` is resident.
+    recompute_cost: int = 0
+    # Payload held by ``stage()`` until ``discard()`` clears it.
+    _staged: PayloadT | None = None
+
+    @abstractmethod
+    def allocate(self) -> None:
+        """Size and initialise the backing storage to its empty state.
+
+        Concrete classes declare their sizing parameters as typed keyword-only
+        arguments (e.g. ``LatentBuffer.allocate(maxlen=...)``); a mistyped
+        keyword raises ``TypeError`` instead of being silently ignored.
+        """
+
+    def stage(self, payload: PayloadT) -> None:
+        """Hold data for an in-progress window.
+
+        Single-window writers go straight through ``commit()``, so the default
+        implementation only records the payload and is otherwise inert.
+        Speculative forking is the intended consumer of staged state.
+        """
+        self._staged = payload
+
+    @abstractmethod
+    def commit(self, payload: PayloadT | None = None) -> None:
+        """Promote data to persistent context.
+
+        For a plain-buffer object this is the synchronous write into its
+        committed storage (e.g. ``cache[i] = updated_kv.clone()``).
+        """
+
+    def discard(self) -> None:
+        """Drop staged data (speculative rejection).
+
+        Clears any recorded staged payload; committed bytes are never touched.
+        """
+        self._staged = None
+
+    def append(self, payload: PayloadT) -> None:
+        """``stage()`` + ``commit()``; single-shot writers use only this."""
+        self.commit(payload)
+
+    @abstractmethod
+    def view(self, *, include_staged: bool = True) -> ViewT:
+        """Return what the consumer (attention metadata, pipeline) reads."""
+
+    # ``policy`` deliberately stays ``Any``: the policy protocol belongs to
+    # the byte-budget eviction planner (RFC #4480), and defining it here
+    # would invent an interface that design owns.
+    def evict(self, policy: Any = None) -> int:
+        """Release this object's backing storage so it can be reused; return the
+        bytes freed.
+
+        This is the backend-agnostic release hook. For a plain buffer it drops
+        the buffer and lets the garbage collector reclaim it; for a paged backend
+        it returns the blocks to the pool so they can be recycled. Never reclaims
+        staged bytes.
+
+        Precondition: only call this when the object is no longer in use.
+        Releasing storage that is still referenced is a correctness bug (another
+        consumer could reuse it). Count-based session eviction therefore does NOT
+        call this — it only drops the lookup-table entry and lets an unreferenced
+        session be garbage-collected. The intended driver of this hook is a
+        byte-budget eviction planner with the recompute-DAG safety described in
+        RFC #4480.
+        """
+        freed = self.nbytes
+        self.reset()
+        return freed
+
+    @abstractmethod
+    def reset(self) -> None:
+        """Clear back to the unallocated state."""
+
+    @abstractmethod
+    def nbytes_by_device(self, seen: SeenStorages) -> dict[str, int]:
+        """Bytes held in this object's committed storage, keyed by device.
+
+        Staged payloads are not counted. ``seen`` deduplicates storages across
+        the whole session -- pass it through to ``device_bytes`` rather than
+        starting a fresh set, or two views on one buffer are counted twice.
+        """
+
+    @property
+    def nbytes(self) -> int:
+        """Total bytes held, summed over every device.
+
+        Only meaningful when the caller genuinely does not care where the bytes
+        live (``evict()`` reporting what it freed, for instance). Anything
+        deciding against a budget wants ``nbytes_by_device`` instead: host and
+        device memory are different resources and their sum is not a quantity.
+        """
+        return sum(self.nbytes_by_device(set()).values())
+
+    @property
+    @abstractmethod
+    def resident(self) -> bool:
+        """Whether data is currently in memory (not yet evicted)."""
+
+    @property
+    def recomputable(self) -> bool:
+        """Droppable-and-rebuildable *right now*.
+
+        Derived, not declared: true only while ``recompute_source`` is still
+        resident, so the answer changes as other objects are evicted.
+        """
+        return self.recompute_source is not None and self.recompute_source.resident

vllm_omni/experimental/world_models/session_state/manager.py

@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Session memory manager (RFC #4480).
+
+``SessionStateManager`` maps ``session_id -> {name: StateObject}`` and is the
+single authority for session lifecycle and eviction. It evicts by session count
+(LRU), matching the per-model session caches it adapts (e.g. DreamZero's), so
+an adapted model keeps the same eviction behaviour. A byte budget is recorded
+for observability; enforcing it is left to a byte-budget eviction planner (see
+RFC #4480).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import threading
+from collections import OrderedDict
+from collections.abc import Callable
+from contextlib import AbstractContextManager
+from typing import TypeVar
+
+from vllm_omni.experimental.world_models.session_state.accounting import (
+    SeenStorages,
+    device_bytes,
+    merge_bytes,
+)
+from vllm_omni.experimental.world_models.session_state.base import StateObject
+
+logger = logging.getLogger(__name__)
+
+# Matches DreamZero's MAX_DREAMZERO_SESSIONS cap so an adapted DreamZero
+# evicts identically.
+DEFAULT_MAX_SESSIONS = 64
+
+M = TypeVar("M", bound=StateObject)
+
+
+class SessionState:
+    """The named ``StateObject`` collection for one session."""
+
+    def __init__(self) -> None:
+        self._objects: dict[str, StateObject] = {}
+        # Session-scoped scalar/tensor metadata that is not itself a typed
+        # StateObject (e.g. counters, cached conditioning tensors). Persists
+        # across the per-call adapters that read/write it.
+        # ``object`` (not ``Any``) so every read site must narrow explicitly.
+        self.attrs: dict[str, object] = {}
+
+    def get(self, name: str) -> StateObject | None:
+        return self._objects.get(name)
+
+    def put(self, name: str, obj: StateObject) -> StateObject:
+        self._objects[name] = obj
+        return obj
+
+    def get_or_create(self, name: str, factory: Callable[[], M], cls: type[M]) -> M:
+        """Return the resident object of type ``cls`` under ``name``, building
+        it via ``factory`` when absent, wrong-typed, or no longer resident."""
+        obj = self._objects.get(name)
+        if not isinstance(obj, cls) or not obj.resident:
+            obj = factory()
+            self._objects[name] = obj
+        return obj
+
+    def names(self) -> list[str]:
+        return list(self._objects)
+
+    def reset(self) -> None:
+        for obj in self._objects.values():
+            obj.reset()
+        # Clear session-scoped metadata too, so a reset leaves nothing stale
+        # (otherwise language/current_start_frame/clip_feas/ys would survive).
+        self.attrs.clear()
+
+    def nbytes_by_device(self, seen: SeenStorages | None = None) -> dict[str, int]:
+        """Bytes this session holds, keyed by device.
+
+        Both buckets are walked. A ``StateObject`` and a session attribute are
+        different in what the *manager* may do with them -- release, accumulate,
+        stage and commit versus merely carry -- but not in whether their bytes
+        occupy memory, so the accounting covers both. One ``seen`` set spans the
+        whole session, so a tensor aliased between the two buckets is counted
+        once.
+        """
+        shared: SeenStorages = set() if seen is None else seen
+        totals: dict[str, int] = {}
+        for obj in self._objects.values():
+            merge_bytes(totals, obj.nbytes_by_device(shared))
+        merge_bytes(totals, device_bytes(self.attrs, shared))
+        return totals
+
+    @property
+    def nbytes(self) -> int:
+        """Total bytes held, summed over every device -- see
+        ``StateObject.nbytes`` on why a budget wants the per-device split."""
+        return sum(self.nbytes_by_device().values())
+
+
+class SessionStateManager:
+    """Owns per-session state and arbitrates the (count-based) LRU.
+
+    The manager lives beside ``DiffusionRequestState`` (owned by the pipeline),
+    not inside its ``extra`` dict: the manager is cross-request and long-lived,
+    whereas ``DiffusionRequestState`` is per-request and transient.
+    """
+
+    def __init__(
+        self,
+        max_sessions: int = DEFAULT_MAX_SESSIONS,
+        byte_budget: int | None = None,
+        lock_factory: Callable[[], AbstractContextManager[object]] = threading.Lock,
+    ) -> None:
+        if max_sessions <= 0:
+            raise ValueError(f"max_sessions must be positive, got {max_sessions}")
+        self.max_sessions = max_sessions
+        # Recorded for observability; enforcement is left to an eviction
+        # planner (see RFC #4480). Not scoped to a device: which pool a budget
+        # applies to is a question for whoever enforces it, so read
+        # ``nbytes_by_device()`` and pick the pool at the point of enforcement.
+        self.byte_budget = byte_budget
+        self._sessions: OrderedDict[str, SessionState] = OrderedDict()
+        # The locking strategy is injected rather than hard-coded: callers
+        # that run the manager from worker threads keep the default
+        # ``threading.Lock``, while an async or single-threaded owner can pass
+        # e.g. ``contextlib.nullcontext`` (or a custom strategy) without any
+        # call-site change -- the manager only ever does ``with self._lock:``.
+        # Heavier operations (paged allocation, eviction planning) keep this
+        # seam.
+        self._lock = lock_factory()
+        self.hits = 0
+        self.misses = 0
+        self.evictions = 0
+
+    @staticmethod
+    def _key(session_id: str | None) -> str:
+        return str(session_id or "default")
+
+    def get_or_create_session(self, session_id: str | None) -> SessionState:
+        key = self._key(session_id)
+        with self._lock:
+            session = self._sessions.get(key)
+            if session is None:
+                self.misses += 1
+                session = SessionState()
+                self._sessions[key] = session
+                while len(self._sessions) > self.max_sessions:
+                    # Drop the oldest from the table only; do not free its
+                    # buffers. An adapter still using the session keeps its own
+                    # reference, so its state survives (matching a model's own
+                    # per-session state object, which the caller holds even
+                    # after it leaves the table). Unreferenced sessions are
+                    # GC-reclaimed.
+                    self._sessions.popitem(last=False)
+                    self.evictions += 1
+            else:
+                self.hits += 1
+                self._sessions.move_to_end(key)
+            return session
+
+    def drop_session(self, session_id: str | None) -> bool:
+        """Release one session explicitly and free its buffers.
+
+        This is the end-of-session signal a model's engine raises on close,
+        eviction, or reset -- distinct from the count-based LRU in
+        ``get_or_create_session``, which only drops the *table* entry and lets an
+        adapter still holding the session keep it alive. Here the session is
+        finished, so its ``StateObject`` buffers and carried attributes are reset
+        to release their memory. Returns whether a session was present.
+        """
+        key = self._key(session_id)
+        with self._lock:
+            session = self._sessions.pop(key, None)
+            if session is None:
+                return False
+            session.reset()
+            return True
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._sessions)
+
+    def __contains__(self, session_id: str | None) -> bool:
+        with self._lock:
+            return self._key(session_id) in self._sessions
+
+    def nbytes_by_device(self) -> dict[str, int]:
+        """Bytes held across every live session, keyed by device."""
+        with self._lock:
+            return self._nbytes_by_device_locked()
+
+    def _nbytes_by_device_locked(self) -> dict[str, int]:
+        # One ``seen`` set across all sessions: sessions can legitimately share
+        # a storage (a conditioning tensor broadcast to several of them), and
+        # the total is what the process holds, not the sum of what each session
+        # would hold alone.
+        shared: SeenStorages = set()
+        totals: dict[str, int] = {}
+        for session in self._sessions.values():
+            merge_bytes(totals, session.nbytes_by_device(shared))
+        return totals
+
+    def stats(self) -> dict[str, int]:
+        """Counters plus per-device byte totals.
+
+        Byte keys are prefixed ``nbytes:`` and named by device -- ``nbytes:cpu``,
+        ``nbytes:cuda:0`` -- because host and device memory are different
+        resources and their sum answers no question. ``total_nbytes`` is kept
+        for callers that only want a single number.
+        """
+        with self._lock:
+            by_device = self._nbytes_by_device_locked()
+            stats: dict[str, int] = {
+                "sessions": len(self._sessions),
+                "max_sessions": self.max_sessions,
+                "total_nbytes": sum(by_device.values()),
+                "hits": self.hits,
+                "misses": self.misses,
+                "evictions": self.evictions,
+            }
+            for device, count in sorted(by_device.items()):
+                stats[f"nbytes:{device}"] = count
+            return stats
+
+
+def resolve_session_state_config(
+    enable: bool | None = None,
+    max_sessions: int | None = None,
+) -> tuple[bool, int]:
+    """Combine explicit args with env-var overrides.
+
+    Environment variables (for quick enablement without touching config files):
+
+        ``OMNI_DIFFUSION_SESSION_STATE_MANAGER`` (``1``/``0``/``true``/``false``)
+        ``OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS`` (positive int)
+
+    Precedence: a *set* environment variable wins over the value passed in
+    (i.e. over ``OmniDiffusionConfig.enable_session_state_manager``), in both
+    directions -- ``...SESSION_MEMORY_MANAGER=0`` force-disables the manager
+    even when the config enables it. Unset or unparsable values leave the
+    passed-in value untouched (an unparsable ``MAX_SESSIONS`` logs a warning
+    and is ignored). This mirrors how other omni feature toggles allow quick
+    A/B flips on a deployed config.
+    """
+    env_enable = os.environ.get("OMNI_DIFFUSION_SESSION_STATE_MANAGER")
+    if env_enable is not None:
+        parsed = env_enable.strip().lower()
+        if parsed in ("1", "true", "yes", "on"):
+            enable = True
+        elif parsed in ("0", "false", "no", "off"):
+            enable = False
+
+    env_size = os.environ.get("OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS")
+    if env_size is not None:
+        try:
+            env_size_int = int(env_size)
+            if env_size_int > 0:
+                max_sessions = env_size_int
+        except ValueError:
+            logger.warning(
+                "Ignoring non-integer OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS=%r.",
+                env_size,
+            )
+
+    return bool(enable), int(max_sessions or DEFAULT_MAX_SESSIONS)

vllm_omni/experimental/world_models/session_state/objects.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Concrete memory objects (RFC #4480).
+
+``LatentBuffer`` backs its storage with a plain (monolithic) buffer. Attention
+KV is deliberately absent: for DreamZero it is owned by the AR-Diffusion engine's
+paged pool (PR #4534), and the RFC's ``PagedKV`` arrives in Phase 1 as a
+handle wrapping that engine state rather than a dense buffer here. The RFC's
+``FixedState`` and ``RetrievalStore`` land with their first consumers, but the
+names are reserved.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from collections.abc import Iterable
+from typing import TypeVar
+
+from vllm_omni.experimental.world_models.session_state.accounting import (
+    SeenStorages,
+    device_bytes,
+)
+from vllm_omni.experimental.world_models.session_state.base import StateObject
+
+# The frame/chunk type a ``LatentBuffer`` holds (e.g. ``np.ndarray`` pixel
+# frames, ``torch.Tensor`` latent chunks).
+ItemT = TypeVar("ItemT")
+
+
+class LatentBuffer(StateObject[ItemT, list[ItemT]]):
+    """Append / ring buffer of latent or pixel frames.
+
+    A bounded ``deque`` (``maxlen`` set at ``allocate()`` time). The object
+    only stores and views the frames; model-specific stacking stays in the
+    caller.
+    """
+
+    def __init__(self) -> None:
+        self._buf: deque[ItemT] | None = None
+
+    def allocate(self, *, maxlen: int | None = None) -> None:
+        self._buf = deque(maxlen=maxlen)
+
+    def append(self, payload: ItemT) -> None:
+        if self._buf is None:
+            raise RuntimeError("LatentBuffer is not allocated; call allocate() first.")
+        self._buf.append(payload)
+
+    def extend(self, payloads: Iterable[ItemT]) -> None:
+        if self._buf is None:
+            raise RuntimeError("LatentBuffer is not allocated; call allocate() first.")
+        self._buf.extend(payloads)
+
+    def commit(self, payload: ItemT | None = None) -> None:
+        if payload is not None:
+            self.append(payload)
+
+    def view(self, *, include_staged: bool = True) -> list[ItemT]:
+        if self._buf is None:
+            raise RuntimeError("LatentBuffer is not allocated; call allocate() first.")
+        return list(self._buf)
+
+    def __len__(self) -> int:
+        return 0 if self._buf is None else len(self._buf)
+
+    def reset(self) -> None:
+        self._buf = None
+
+    def nbytes_by_device(self, seen: SeenStorages) -> dict[str, int]:
+        if self._buf is None:
+            return {}
+        return device_bytes(list(self._buf), seen)
+
+    @property
+    def resident(self) -> bool:
+        return self._buf is not None