161 lines
8.8 KiB
Markdown
161 lines
8.8 KiB
Markdown
# 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.
|