# 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.