From 1f01d05603e6565de5b5fb1057d8c2f0bafee52f Mon Sep 17 00:00:00 2001 From: wjsuijlenh <127192506+wjsuijlenh@users.noreply.github.com> Date: Mon, 27 Jul 2026 07:17:38 +0200 Subject: [PATCH] [Core] Session state manager for AR-diffusion world models (RFC #4480 Phase 0) (#4487) Signed-off-by: Wijnand Suijlen --- docs/features/session_state_manager.md | 160 +++++++++ .../test_session_state_accounting.py | 201 +++++++++++ .../test_session_state_equivalence.py | 257 ++++++++++++++ .../dreamzero/test_session_state_lifecycle.py | 115 +++++++ vllm_omni/diffusion/data.py | 7 + .../models/dreamzero/pipeline_dreamzero.py | 95 +++++- .../models/dreamzero/state_dreamzero.py | 12 + .../experimental/world_models/__init__.py | 11 + .../world_models/adapters/__init__.py | 5 + .../adapters/state_dreamzero_adapter.py | 315 ++++++++++++++++++ .../world_models/session_state/__init__.py | 35 ++ .../world_models/session_state/accounting.py | 92 +++++ .../world_models/session_state/attrs.py | 70 ++++ .../world_models/session_state/base.py | 160 +++++++++ .../world_models/session_state/manager.py | 265 +++++++++++++++ .../world_models/session_state/objects.py | 76 +++++ 16 files changed, 1859 insertions(+), 17 deletions(-) create mode 100644 docs/features/session_state_manager.md create mode 100644 tests/dreamzero/test_session_state_accounting.py create mode 100644 tests/dreamzero/test_session_state_equivalence.py create mode 100644 tests/dreamzero/test_session_state_lifecycle.py create mode 100644 vllm_omni/experimental/world_models/__init__.py create mode 100644 vllm_omni/experimental/world_models/adapters/__init__.py create mode 100644 vllm_omni/experimental/world_models/adapters/state_dreamzero_adapter.py create mode 100644 vllm_omni/experimental/world_models/session_state/__init__.py create mode 100644 vllm_omni/experimental/world_models/session_state/accounting.py create mode 100644 vllm_omni/experimental/world_models/session_state/attrs.py create mode 100644 vllm_omni/experimental/world_models/session_state/base.py create mode 100644 vllm_omni/experimental/world_models/session_state/manager.py create mode 100644 vllm_omni/experimental/world_models/session_state/objects.py diff --git a/docs/features/session_state_manager.md b/docs/features/session_state_manager.md new file mode 100644 index 000000000..c0c055fcb --- /dev/null +++ b/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. diff --git a/tests/dreamzero/test_session_state_accounting.py b/tests/dreamzero/test_session_state_accounting.py new file mode 100644 index 000000000..443a02404 --- /dev/null +++ b/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) diff --git a/tests/dreamzero/test_session_state_equivalence.py b/tests/dreamzero/test_session_state_equivalence.py new file mode 100644 index 000000000..a346a4039 --- /dev/null +++ b/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) diff --git a/tests/dreamzero/test_session_state_lifecycle.py b/tests/dreamzero/test_session_state_lifecycle.py new file mode 100644 index 000000000..a67578cc7 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 078354d96..a001c665c 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index 711b345d2..b2668b0a5 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/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: diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py index 4c9209808..bfd4f7aa9 100644 --- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py +++ b/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) diff --git a/vllm_omni/experimental/world_models/__init__.py b/vllm_omni/experimental/world_models/__init__.py new file mode 100644 index 000000000..9b01de3d0 --- /dev/null +++ b/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 diff --git a/vllm_omni/experimental/world_models/adapters/__init__.py b/vllm_omni/experimental/world_models/adapters/__init__.py new file mode 100644 index 000000000..7f218e2c1 --- /dev/null +++ b/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 diff --git a/vllm_omni/experimental/world_models/adapters/state_dreamzero_adapter.py b/vllm_omni/experimental/world_models/adapters/state_dreamzero_adapter.py new file mode 100644 index 000000000..ea355b471 --- /dev/null +++ b/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 diff --git a/vllm_omni/experimental/world_models/session_state/__init__.py b/vllm_omni/experimental/world_models/session_state/__init__.py new file mode 100644 index 000000000..953e2615a --- /dev/null +++ b/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", +] diff --git a/vllm_omni/experimental/world_models/session_state/accounting.py b/vllm_omni/experimental/world_models/session_state/accounting.py new file mode 100644 index 000000000..d8304efdd --- /dev/null +++ b/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 diff --git a/vllm_omni/experimental/world_models/session_state/attrs.py b/vllm_omni/experimental/world_models/session_state/attrs.py new file mode 100644 index 000000000..708b5857d --- /dev/null +++ b/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(, 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 diff --git a/vllm_omni/experimental/world_models/session_state/base.py b/vllm_omni/experimental/world_models/session_state/base.py new file mode 100644 index 000000000..be1d04016 --- /dev/null +++ b/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 diff --git a/vllm_omni/experimental/world_models/session_state/manager.py b/vllm_omni/experimental/world_models/session_state/manager.py new file mode 100644 index 000000000..bb08c064c --- /dev/null +++ b/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) diff --git a/vllm_omni/experimental/world_models/session_state/objects.py b/vllm_omni/experimental/world_models/session_state/objects.py new file mode 100644 index 000000000..e9e1f4152 --- /dev/null +++ b/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