From 8e2e9b6b53e86e6a479ed2c0a53782f655f60e04 Mon Sep 17 00:00:00 2001 From: linzhenpl07 <60651557+linzhenpl07@users.noreply.github.com> Date: Sun, 9 Aug 2026 20:53:17 +0800 Subject: [PATCH] [Feat] Cosmos3 session-memory port for UND text K/V (RFC #4480 Phase 0) (#4657) Signed-off-by: linzhenpl07 Co-authored-by: Claude Opus 4.8 --- docs/features/session_state_manager.md | 70 ++++- .../test_session_memory_equivalence.py | 192 ++++++++++++ .../models/cosmos3/test_cosmos3_pipeline.py | 73 +++++ .../models/cosmos3/pipeline_cosmos3.py | 276 +++++++++++++----- .../models/cosmos3/transformer_cosmos3.py | 18 +- .../adapters/state_cosmos3_adapter.py | 167 +++++++++++ 6 files changed, 703 insertions(+), 93 deletions(-) create mode 100644 tests/cosmos3/test_session_memory_equivalence.py create mode 100644 vllm_omni/experimental/world_models/adapters/state_cosmos3_adapter.py diff --git a/docs/features/session_state_manager.md b/docs/features/session_state_manager.md index c0c055fcb..87f74613a 100644 --- a/docs/features/session_state_manager.md +++ b/docs/features/session_state_manager.md @@ -20,11 +20,12 @@ extracts it into a shared, typed contract: (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. +Autoregressive generation 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)). Cosmos3's fixed, +encode-once UND conditioning K/V is different and is covered by the dense +Cosmos3 adapter described below. Integrating the engine's per-session paged KV +handle into the same byte accounting is RFC #4480 Phase 1. ## Enabling @@ -158,3 +159,62 @@ 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. + +### Cosmos3 + +When `enable_session_state_manager` is enabled, the Cosmos3 adapter stores the +dense, encode-once UND text K/V for every transformer layer and active CFG +branch, together with that branch's GEN M-RoPE cosine/sine tensors. These are +GPU-resident tensors and can contribute to CUDA OOM. A useful description of +the retained storage on one rank is: + +```text +sum_over_branches( + sum_over_layers(storage_bytes(K) + storage_bytes(V)) + + storage_bytes(freqs_gen_cos) + storage_bytes(freqs_gen_sin) +) +``` + +This is not a peak-memory upper bound: construction of the next branch can +temporarily overlap the already retained branch with UND activations, the new +full K/V result, RoPE tensors, and generator activations. K/V slices also keep +their underlying allocation alive, so physical storage can follow the padded +text length rather than only the exposed trimmed length. The manager's byte +accounting deliberately measures storage for this reason. + +Branch placement changes the per-rank cost: + +- no CFG retains one conditional branch; +- sequential CFG can retain both conditional and unconditional branches on the + same GPU; +- CFG-parallel retains one branch on each CFG rank; +- UND K/V is replicated across sequence-parallel ranks rather than divided by + the SP degree. + +Consequently, risk increases with text batch/sequence length, media token count +(which sizes `freqs_gen`), sequential CFG, dtype width, and concurrent live +requests. Phase 0 records `nbytes:cuda:*` but does not enforce `byte_budget`, so +the manager cannot prevent an allocation from exhausting the device. The +count-based session limit is not an OOM guarantee. + +The adapter transfers ownership by tensor reference; it does not clone K/V or +RoPE tensors. For a normal serial request this should not materially increase +tensor storage relative to Cosmos3's existing bespoke cond/uncond cache, which +retains the same branch values. It does add small CPU-side session/object +metadata. This claim is based on the reference flow in the implementation; +there is not yet a current-head GPU peak-memory A/B result for this PR. + +Request-finally cleanup drops the manager-owned session, so branches do not +accumulate with completed request IDs. The transformer handoff field still +references the last loaded branch until `reset_cache()` runs at the start of +the next generation (or the pipeline is destroyed). Thus this is bounded +post-request residency, not growth per completed request, but manager +`stats()` after `drop_session()` no longer includes that last transformer-held +alias. A profiler that only reads manager statistics can therefore under-report +live process memory. + +For OOM validation, measure complete cache-off and true flag-on generations in +otherwise identical fresh processes. Synchronize CUDA and report process-level +peak memory plus PyTorch peak allocated and peak reserved memory. Include the +period between requests when checking cleanup; do not infer safety from the +manager byte total, a single denoising step, or an operator-only profile. diff --git a/tests/cosmos3/test_session_memory_equivalence.py b/tests/cosmos3/test_session_memory_equivalence.py new file mode 100644 index 000000000..6cef232a3 --- /dev/null +++ b/tests/cosmos3/test_session_memory_equivalence.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Equivalence and session-isolation tests for Cosmos3StateAdapter (RFC #4480). + +These drive the adapter directly with tiny CPU tensors -- no model, no GPU -- +and assert it stores/returns the per-branch UND K/V exactly and isolates +durable state between sessions. They do not cover the transformer's mutable +load-to-forward handoff. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.experimental.world_models.adapters.state_cosmos3_adapter import ( + Cosmos3StateAdapter, +) +from vllm_omni.experimental.world_models.session_state import SessionStateManager + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +LAYERS = 3 +SEQ, DIM = 3, 4 + + +def _fake_cached_kv(seed: int) -> list[tuple[torch.Tensor, torch.Tensor]]: + """A list of per-layer (K, V); contents seeded so branches/sessions differ.""" + g = torch.Generator().manual_seed(seed) + return [(torch.randn(1, SEQ, DIM, generator=g), torch.randn(1, SEQ, DIM, generator=g)) for _ in range(LAYERS)] + + +def _adapter(session_id: str, manager: SessionStateManager | None = None) -> Cosmos3StateAdapter: + # `manager is not None`, not `manager or ...`: an empty SessionStateManager is + # falsy (it defines __len__), so `or` would silently swap in a fresh manager. + return Cosmos3StateAdapter(session_id, manager if manager is not None else SessionStateManager()) + + +def _assert_kv_equal(got: list | None, expected: list) -> None: + assert got is not None and len(got) == len(expected) == LAYERS + for (k, v), (ke, ve) in zip(got, expected): + torch.testing.assert_close(k, ke) + torch.testing.assert_close(v, ve) + + +class _FakeTransformer: + cached_kv = None + cached_freqs_gen = None + + +# ----------------------------- equivalence ----------------------------- + + +def test_branch_starts_uninitialized() -> None: + a = _adapter("s0") + for is_neg in (False, True): + assert a.is_branch_initialized(is_neg) is False + assert a.get_branch_kv(is_neg) is None + + +def test_set_then_get_roundtrip() -> None: + a = _adapter("s0") + kv = _fake_cached_kv(1) + a.set_branch_kv(False, kv) + assert a.is_branch_initialized(False) is True + _assert_kv_equal(a.get_branch_kv(False), kv) + + +def test_cond_uncond_branches_isolated() -> None: + a = _adapter("s0") + kv_pos, kv_neg = _fake_cached_kv(1), _fake_cached_kv(2) + a.set_branch_kv(False, kv_pos) + a.set_branch_kv(True, kv_neg) + _assert_kv_equal(a.get_branch_kv(False), kv_pos) + _assert_kv_equal(a.get_branch_kv(True), kv_neg) + + +def test_encode_once_capture_is_noop_when_initialized() -> None: + a = _adapter("s0") + kv = _fake_cached_kv(1) + a.set_branch_kv(False, kv) + tr = _FakeTransformer() + tr.cached_kv = _fake_cached_kv(999) # different content + a.capture_from_transformer(tr, False) + _assert_kv_equal(a.get_branch_kv(False), kv) # still the first one + + +def test_capture_writes_when_uninitialized() -> None: + a = _adapter("s0") + tr = _FakeTransformer() + tr.cached_kv = _fake_cached_kv(5) + tr.cached_freqs_gen = (torch.ones(2), torch.zeros(2)) + a.capture_from_transformer(tr, False) + _assert_kv_equal(a.get_branch_kv(False), _fake_cached_kv(5)) + + +def test_capture_without_freqs_does_not_partially_initialize() -> None: + a = _adapter("s0") + tr = _FakeTransformer() + tr.cached_kv = _fake_cached_kv(5) + tr.cached_freqs_gen = None + with pytest.raises(RuntimeError, match="did not produce GEN RoPE frequencies"): + a.capture_from_transformer(tr, False) + assert a.is_branch_initialized(False) is False + + +def test_load_into_transformer_without_freqs_sets_kv() -> None: + a = _adapter("s0") + kv = _fake_cached_kv(1) + a.set_branch_kv(False, kv) + tr = _FakeTransformer() + a.load_into_transformer(tr, False) + _assert_kv_equal(tr.cached_kv, kv) + assert tr.cached_freqs_gen is None + + +def test_set_then_get_preserves_per_layer_values() -> None: + # The stored value must reflect the tensors passed in (per-layer). + a = _adapter("s0") + kv = _fake_cached_kv(3) + a.set_branch_kv(False, kv) + _assert_kv_equal(a.get_branch_kv(False), kv) + + +def test_load_into_transformer_restores_kv_and_freqs_gen() -> None: + a = _adapter("s0") + tr = _FakeTransformer() + tr.cached_kv = _fake_cached_kv(3) + tr.cached_freqs_gen = (torch.ones(2), torch.zeros(2)) + a.capture_from_transformer(tr, False) + tr.cached_kv = None + tr.cached_freqs_gen = None + a.load_into_transformer(tr, False) + _assert_kv_equal(tr.cached_kv, _fake_cached_kv(3)) + assert tr.cached_freqs_gen is not None + torch.testing.assert_close(tr.cached_freqs_gen[0], torch.ones(2)) + + +def test_partial_branch_raises() -> None: + a = _adapter("s0") + a.set_branch_kv(False, _fake_cached_kv(1)) + a._session.get("und_kv/1").reset() # type: ignore[union-attr] + with pytest.raises(RuntimeError, match="partially initialized"): + a.get_branch_kv(False) + + +# ----------------------------- reset ----------------------------- + + +def test_reset_clears_both_branches() -> None: + a = _adapter("s0") + a.set_branch_kv(False, _fake_cached_kv(1)) + a.set_branch_kv(True, _fake_cached_kv(2)) + a.reset() + for is_neg in (False, True): + assert a.is_branch_initialized(is_neg) is False + assert a.get_branch_kv(is_neg) is None + + +# ----------------------- concurrency isolation (core) ----------------------- + + +def test_two_sessions_do_not_clobber() -> None: + manager = SessionStateManager() + a, b = _adapter("A", manager), _adapter("B", manager) + kv_a, kv_b = _fake_cached_kv(10), _fake_cached_kv(20) + a.set_branch_kv(False, kv_a) + assert b.is_branch_initialized(False) is False # B not seeing A's state + b.set_branch_kv(False, kv_b) + _assert_kv_equal(a.get_branch_kv(False), kv_a) + _assert_kv_equal(b.get_branch_kv(False), kv_b) + + +def test_fresh_adapter_same_session_sees_prior_state() -> None: + manager = SessionStateManager() + kv = _fake_cached_kv(7) + _adapter("s", manager).set_branch_kv(False, kv) + again = Cosmos3StateAdapter("s", manager) + _assert_kv_equal(again.get_branch_kv(False), kv) + + +def test_adapter_survives_session_lru_eviction() -> None: + # An in-use adapter keeps its state even after the manager evicts its + # session id from the lookup table (LRU bounded by max_sessions). + manager = SessionStateManager(max_sessions=2) + a = _adapter("active", manager) + kv = _fake_cached_kv(11) + a.set_branch_kv(False, kv) + for i in range(3): + _adapter(f"other{i}", manager) + assert "active" not in manager + _assert_kv_equal(a.get_branch_kv(False), kv) diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 3093e6ceb..796de06f1 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -16,6 +16,7 @@ from PIL import Image from torch import nn from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.experimental.world_models.session_state import SessionStateManager pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -425,6 +426,29 @@ def test_distilled_generation_rejects_robolab_policy(make_cosmos3_pipeline) -> N pipeline._forward_robolab_policy(make_sampling_params(), None, 0.0) +def test_forward_threads_request_id_to_robolab(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + robolab_inputs = object() + captured: dict[str, Any] = {} + expected = object() + pipeline._build_robolab_policy_inputs = lambda sp, prompt, request_id: robolab_inputs + + def fake_forward_robolab(sp, inputs, pipeline_start, session_id=None): + del sp, inputs, pipeline_start + captured["session_id"] = session_id + return expected + + pipeline._forward_robolab_policy = fake_forward_robolab + request = SimpleNamespace( + prompts=["policy"], + sampling_params=make_sampling_params(), + request_id="robolab-request-7", + ) + + assert pipeline.forward(request) is expected + assert captured["session_id"] == "robolab-request-7" + + @pytest.mark.parametrize( ("prompt", "sampling_params", "message"), [ @@ -595,6 +619,7 @@ def _make_od_config( return SimpleNamespace( enable_cpu_offload=False, enable_diffusion_pipeline_profiler=False, + enable_session_state_manager=False, model="/nonexistent/model/path", dtype=torch.float32, flow_shift=None, @@ -1673,6 +1698,32 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result)) +def test_diffuse_drops_session_when_progress_iteration_fails(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + pipeline._use_session_state = True + pipeline._memory_manager = SessionStateManager() + + class FailingProgress: + def __iter__(self): + raise RuntimeError("progress failed before first step") + + pipeline.progress_bar = lambda timesteps: FailingProgress() + with pytest.raises(RuntimeError, match="progress failed"): + pipeline.diffuse( + latents=torch.zeros(1, 2, 1, 1, 1), + timesteps=torch.tensor([7]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=1.0, + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0}, + session_id="request-that-fails", + ) + + assert "request-that-fails" not in pipeline._memory_manager + + def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_cfg_parallel) -> None: pipeline = make_cosmos3_pipeline() latents = torch.zeros(1, 2, 1, 1, 1) @@ -1702,6 +1753,28 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_ torch.testing.assert_close(result, torch.full_like(latents, 254.0)) +def test_diffuse_transfer_rejects_session_state_manager(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + pipeline._use_session_state = True + + with pytest.raises(NotImplementedError, match="richer cache keys"): + pipeline.diffuse_transfer( + latents=torch.zeros(1, 2, 1, 1, 1), + timesteps=torch.tensor([7]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=1.0, + control_guidance=1.0, + control_guidance_interval=None, + control_latents=[], + shared_kwargs={}, + velocity_mask=torch.ones(1, 1, 1, 1, 1), + condition_latents=torch.zeros(1, 2, 1, 1, 1), + ) + + def test_diffuse_transfer_skips_idle_cfg_branches(make_cosmos3_pipeline, sequential_cfg_parallel) -> None: latents = torch.zeros(1, 2, 1, 1, 1) velocity_mask = torch.ones(1, 1, 1, 1, 1) diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index fcc2a8b6c..97cf92e26 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -51,6 +51,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_rank, get_classifier_free_guidance_world_size, ) from vllm_omni.diffusion.distributed.utils import get_local_device @@ -71,6 +72,13 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.entrypoints.openai.video_api_utils import positive_float +from vllm_omni.experimental.world_models.adapters.state_cosmos3_adapter import ( + Cosmos3StateAdapter, +) +from vllm_omni.experimental.world_models.session_state import ( + SessionStateManager, + resolve_session_state_config, +) from vllm_omni.inputs.data import OmniDiffusionSamplingParams from .action import ( @@ -910,6 +918,24 @@ class Cosmos3OmniDiffusersPipeline( ) self.is_edge_model = transformer_cls is Cosmos3EdgeVFMTransformer + # Opt-in: route the per-CFG-branch UND text K/V through the shared + # SessionStateManager (RFC #4480) instead of the transformer-instance + # cache. This establishes request-keyed state isolation; it does not by + # itself make the mutable pipeline safe for concurrent execution. + # Default off -> the bespoke transformer-instance path below is unchanged. + self._use_session_state, mm_max_sessions = resolve_session_state_config( + enable=od_config.enable_session_state_manager, + ) + self._memory_manager: SessionStateManager | None = ( + SessionStateManager(max_sessions=mm_max_sessions) if self._use_session_state else None + ) + # Forward declaration for the future BDE paged path (Stage C/D). Nothing + # writes this field in Phase 0; a non-None value is rejected by the guards + # below. None means the dense session adapter or the bespoke path. + self._bde_kv_state = None + if getattr(self, "_use_session_state", False): + logger.info("Cosmos3: session state manager enabled (max_sessions=%d)", mm_max_sessions) + # --- Scheduler --- # Distilled model differs from regular one only by scheduler, # distilled one uses FlowMatchEulerDiscreteScheduler, while @@ -1463,6 +1489,7 @@ class Cosmos3OmniDiffusersPipeline( sp: OmniDiffusionSamplingParams, inputs: RoboLabPolicyInputs, pipeline_start: float, + session_id: str | None = None, ) -> DiffusionOutput: if self.is_distilled_model: raise self._distilled_unsupported_error("RoboLab/action policy requests are unsupported.") @@ -1576,6 +1603,7 @@ class Cosmos3OmniDiffusersPipeline( guidance_interval=None, raw_action_dim=raw_action_dim, scheduler=scheduler, + session_id=session_id, generator=generator, ) @@ -2415,6 +2443,49 @@ class Cosmos3OmniDiffusersPipeline( action_velocity_mask = 1.0 - condition_mask return action_latents, action_velocity_mask, clean_action, raw_action_dim + # -- Session-memory bypass for UND text K/V (RFC #4480) ----------------- + # Mirrors DreamZero's `_kv_*` pattern: route through the BDE pool when set + # (Stage C/D), else the session adapter (Stage B), else the bespoke + # transformer-instance cache. All gated so the default path is unchanged. + + def _new_cosmos3_state(self, session_id: str | None) -> Cosmos3StateAdapter | None: + if getattr(self, "_memory_manager", None) is None: + return None + return Cosmos3StateAdapter(session_id, self._memory_manager) + + def _kv_load_und(self, state: Cosmos3StateAdapter | None, is_negative: bool) -> bool: + """Load this branch's UND K/V onto the transformer before forward. + + Returns True if the session/BDE path handled the load (caller skips the + bespoke ``cond_cache`` assignment), False for the bespoke path. + """ + if getattr(self, "_bde_kv_state", None) is not None: + raise NotImplementedError("BDE pool path for Cosmos3 UND K/V is Stage C/D") + if state is not None: + state.load_into_transformer(self.transformer, is_negative) + return True + return False + + def _kv_capture_und(self, state: Cosmos3StateAdapter | None, is_negative: bool) -> None: + """After forward, store the freshly computed UND K/V for this branch (encode-once).""" + if getattr(self, "_bde_kv_state", None) is not None: + raise NotImplementedError("BDE pool path for Cosmos3 UND K/V is Stage C/D") + if state is not None: + state.capture_from_transformer(self.transformer, is_negative) + + def _kv_reset_und(self, state: Cosmos3StateAdapter | None) -> None: + """Reset UND K/V at the start of a generation (replaces ``transformer.reset_cache()``). + + Always clears the transformer-instance cache so no stale UND K/V leaks into + paths that read it directly (cfg-parallel / no-CFG / bespoke fallback); also + clears the session objects when session state is active. + """ + if getattr(self, "_bde_kv_state", None) is not None: + raise NotImplementedError("BDE pool path for Cosmos3 UND K/V is Stage C/D") + self.transformer.reset_cache() + if state is not None: + state.reset() + # -- Denoising loop (shared by T2V and I2V) ----------------------------- def diffuse( @@ -2438,6 +2509,7 @@ class Cosmos3OmniDiffusersPipeline( guidance_interval: tuple[float, float] | None = None, raw_action_dim: int | None = None, scheduler: Any | None = None, + session_id: str | None = None, generator: torch.Generator | None = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Denoising loop with 3-mode CFG support (parallel, sequential, none). @@ -2470,7 +2542,9 @@ class Cosmos3OmniDiffusersPipeline( do_cfg = guidance_scale > 1.0 cfg_parallel = self._cfg_parallel_active() and do_cfg step_scheduler = scheduler if scheduler is not None else self.scheduler - self.transformer.reset_cache() + # Session-keyed UND K/V (RFC #4480); None => bespoke transformer-instance cache. + kv_state = self._new_cosmos3_state(session_id) + self._kv_reset_und(kv_state) def _cfg_active_at(t: torch.Tensor) -> bool: if guidance_interval is None: @@ -2606,18 +2680,57 @@ class Cosmos3OmniDiffusersPipeline( if sound_latents is not None: sound_latents = step_out[idx] - if cfg_parallel: - for t in self.progress_bar(timesteps): - timestep = t.unsqueeze(0) - # Out-of-interval steps run with effective scale 1.0 so the - # combined output equals the cond branch (uncond is dropped). - # All ranks still execute both branches; no CFG-Parallel - # divergence. - step_scale = guidance_scale if _cfg_active_at(t) else 1.0 - noise_pred = self.predict_noise_maybe_with_cfg( - do_true_cfg=True, - true_cfg_scale=step_scale, - positive_kwargs=dict( + try: + if cfg_parallel: + # Each CFG-parallel rank runs exactly one branch (rank 0 -> cond, + # else uncond), so session keying loads/stores only this rank's branch. + cfg_rank_is_negative = get_classifier_free_guidance_rank() != 0 + for t in self.progress_bar(timesteps): + timestep = t.unsqueeze(0) + # Outside the interval, scale=1 makes the combined output equal + # the cond branch. Every rank remains on the same iteration and + # collective schedule. + step_scale = guidance_scale if _cfg_active_at(t) else 1.0 + self._kv_load_und(kv_state, is_negative=cfg_rank_is_negative) + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=step_scale, + positive_kwargs=dict( + hidden_states=latents, + timestep=timestep, + text_ids=cond_ids, + text_mask=cond_mask, + action_latents=action_latents, + sound_latents=sound_latents, + **shared_kwargs, + ), + negative_kwargs=dict( + hidden_states=latents, + timestep=timestep, + text_ids=uncond_ids, + text_mask=uncond_mask, + action_latents=action_latents, + sound_latents=sound_latents, + **shared_kwargs, + ), + cfg_normalize=False, + ) + if kv_state is not None: + self._kv_capture_und(kv_state, is_negative=cfg_rank_is_negative) + _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents)) + + elif do_cfg: + cond_cache: tuple = (None, None) + uncond_cache: tuple = (None, None) + keep_uncond_for_cache = self._cache_requires_paired_cfg() + + for t in self.progress_bar(timesteps): + timestep = t.unsqueeze(0) + cfg_active = _cfg_active_at(t) + + if not self._kv_load_und(kv_state, is_negative=False): + self.transformer.cached_kv, self.transformer.cached_freqs_gen = cond_cache + noise_cond = self.transformer( hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -2625,79 +2738,67 @@ class Cosmos3OmniDiffusersPipeline( action_latents=action_latents, sound_latents=sound_latents, **shared_kwargs, - ), - negative_kwargs=dict( + ) + if kv_state is not None: + self._kv_capture_und(kv_state, is_negative=False) + elif cond_cache[0] is None: + cond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) + + if cfg_active or keep_uncond_for_cache: + if not self._kv_load_und(kv_state, is_negative=True): + self.transformer.cached_kv, self.transformer.cached_freqs_gen = uncond_cache + noise_uncond = self.transformer( + hidden_states=latents, + timestep=timestep, + text_ids=uncond_ids, + text_mask=uncond_mask, + action_latents=action_latents, + sound_latents=sound_latents, + **shared_kwargs, + ) + if kv_state is not None: + self._kv_capture_und(kv_state, is_negative=True) + elif uncond_cache[0] is None: + uncond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) + # Outside the interval, scale=1.0 makes the combined result + # equal to noise_cond; the uncond pass is computed only to + # preserve cache-dit's cond/uncond parity. + step_scale = guidance_scale if cfg_active else 1.0 + noise_pred = self.combine_cfg_noise( + noise_cond, + noise_uncond, + step_scale, + cfg_normalize=False, + ) + else: + noise_pred = noise_cond + + _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents)) + + else: + # No CFG: a single cond branch per step. Bespoke (state None) keeps + # using the transformer-instance cache exactly as before. + for t in self.progress_bar(timesteps): + timestep = t.unsqueeze(0) + self._kv_load_und(kv_state, is_negative=False) + noise_pred = self.transformer( hidden_states=latents, timestep=timestep, - text_ids=uncond_ids, - text_mask=uncond_mask, - action_latents=action_latents, - sound_latents=sound_latents, - **shared_kwargs, - ), - cfg_normalize=False, - ) - _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents)) - - elif do_cfg: - cond_cache: tuple = (None, None) - uncond_cache: tuple = (None, None) - - keep_uncond_for_cache = self._cache_requires_paired_cfg() - - for t in self.progress_bar(timesteps): - timestep = t.unsqueeze(0) - cfg_active = _cfg_active_at(t) - - self.transformer.cached_kv, self.transformer.cached_freqs_gen = cond_cache - noise_cond = self.transformer( - hidden_states=latents, - timestep=timestep, - text_ids=cond_ids, - text_mask=cond_mask, - action_latents=action_latents, - sound_latents=sound_latents, - **shared_kwargs, - ) - if cond_cache[0] is None: - cond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) - - if cfg_active or keep_uncond_for_cache: - self.transformer.cached_kv, self.transformer.cached_freqs_gen = uncond_cache - noise_uncond = self.transformer( - hidden_states=latents, - timestep=timestep, - text_ids=uncond_ids, - text_mask=uncond_mask, + text_ids=cond_ids, + text_mask=cond_mask, action_latents=action_latents, sound_latents=sound_latents, **shared_kwargs, ) - if uncond_cache[0] is None: - uncond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) - # Outside the interval, scale=1.0 makes the combined result - # equal to noise_cond; the uncond pass is computed only to - # preserve cache-dit's cond/uncond parity. - step_scale = guidance_scale if cfg_active else 1.0 - noise_pred = self.combine_cfg_noise(noise_cond, noise_uncond, step_scale, cfg_normalize=False) - else: - noise_pred = noise_cond - - _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents)) - - else: - for t in self.progress_bar(timesteps): - timestep = t.unsqueeze(0) - noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep, - text_ids=cond_ids, - text_mask=cond_mask, - action_latents=action_latents, - sound_latents=sound_latents, - **shared_kwargs, - ) - _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents)) + if kv_state is not None: + self._kv_capture_und(kv_state, is_negative=False) + _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents)) + finally: + # Cosmos3 currently receives a unique request_id rather than a + # reusable rollout session id. Retaining its state would only pin + # K/V buffers on device after this generation finishes. + if kv_state is not None: + self._memory_manager.drop_session(session_id) outputs = [latents] if action_latents is not None: @@ -2794,6 +2895,13 @@ class Cosmos3OmniDiffusersPipeline( guidance_interval: tuple[float, float] | None = None, generator: torch.Generator | None = None, ) -> torch.Tensor: + if getattr(self, "_use_session_state", False): + raise NotImplementedError( + "Cosmos3 transfer does not yet support SessionStateManager: " + "its control/text CFG branches require richer cache keys. " + "Disable the session state manager for transfer requests." + ) + def _active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool: if interval is None: return True @@ -3178,7 +3286,12 @@ class Cosmos3OmniDiffusersPipeline( sp = req.sampling_params robolab_inputs = self._build_robolab_policy_inputs(sp, prompt_data, getattr(req, "request_id", None)) if robolab_inputs is not None: - return self._forward_robolab_policy(sp, robolab_inputs, pipeline_start) + return self._forward_robolab_policy( + sp, + robolab_inputs, + pipeline_start, + session_id=getattr(req, "request_id", None), + ) if isinstance(prompt_data, str): prompt = prompt_data @@ -3543,6 +3656,7 @@ class Cosmos3OmniDiffusersPipeline( guidance_interval=guidance_interval, raw_action_dim=raw_action_dim, scheduler=scheduler, + session_id=getattr(req, "request_id", None), generator=generator, ) diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 475c02889..cea5a32ff 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -1674,8 +1674,11 @@ class Cosmos3VFMTransformer(nn.Module): sound_tokens = self.pack_sound(sound_latents) s_sound = sound_tokens.shape[1] - # Run UND pathway once and cache K/V (replicated across all ranks) - if self.cached_kv is None: + # Run UND pathway once and cache K/V (replicated across all ranks). + # freqs_gen (M-RoPE cos/sin) is derived purely from shape/fps and is + # cached alongside UND K/V by both the bespoke and session-state paths. + need_kv = self.cached_kv is None + if need_kv or self.cached_freqs_gen is None: freqs_und, freqs_gen = self._compute_rope_freqs( text_mask, t, @@ -1691,13 +1694,14 @@ class Cosmos3VFMTransformer(nn.Module): num_vision_items=len(control_latent_list) + 1, share_vision_temporal_positions=transfer_share_vision_temporal_positions, ) - with self._offload_context("reasoner"): - cached_kv_full = self.language_model(text_ids, freqs_und) self.cached_freqs_gen = freqs_gen - # Trim to real text length (remove padding). K/V stay replicated; - # the framework Attention layer head-slices them via joint_key/value. - self.cached_kv = [(k[:, :max_real_len], v[:, :max_real_len]) for k, v in cached_kv_full] + if need_kv: + with self._offload_context("reasoner"): + cached_kv_full = self.language_model(text_ids, freqs_und) + # Trim to real text length (remove padding). K/V stay replicated; + # the framework Attention layer head-slices them via joint_key/value. + self.cached_kv = [(k[:, :max_real_len], v[:, :max_real_len]) for k, v in cached_kv_full] with self._offload_context("generator"): # Patchify latents and project to hidden space after UND cache diff --git a/vllm_omni/experimental/world_models/adapters/state_cosmos3_adapter.py b/vllm_omni/experimental/world_models/adapters/state_cosmos3_adapter.py new file mode 100644 index 000000000..b1fee68cc --- /dev/null +++ b/vllm_omni/experimental/world_models/adapters/state_cosmos3_adapter.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Adapter routing Cosmos3's per-CFG-branch UND text K/V through session state. + +Cosmos3's cross-step cache is the UND (reasoner) text K/V, currently kept as +instance fields on ``Cosmos3VFMTransformer`` (``cached_kv``). +``Cosmos3StateAdapter`` moves its durable ownership into the shared +``SessionStateManager`` (RFC #4480), behind the +``enable_session_state_manager`` opt-in flag. The load-to-transformer handoff +still uses the transformer's fields, so this is session-state groundwork rather +than synchronization for overlapping forwards. + +Each (layer, CFG branch) pair is a model-specific ``Cosmos3UNDKV`` state object. +The shared package deliberately does not define a dense attention-KV object: +DreamZero's attention KV belongs to the paged AR-Diffusion engine, whereas this +is Cosmos3's fixed, encode-once conditioning KV. Keeping the concrete object in +the model adapter preserves that ownership boundary while retaining the shared +lifecycle and byte accounting. + +``freqs_gen`` (M-RoPE cos/sin for the GEN pathway) is also session-scoped. It +is constant for one request, so retaining it in ``SessionState.attrs`` avoids +recomputing both UND and GEN RoPE frequencies on every denoise step. +""" + +from __future__ import annotations + +from typing import Protocol + +import torch + +from vllm_omni.experimental.world_models.session_state.accounting import ( + SeenStorages, + device_bytes, +) +from vllm_omni.experimental.world_models.session_state.base import StateObject +from vllm_omni.experimental.world_models.session_state.manager import ( + SessionStateManager, +) + +KVPair = tuple[torch.Tensor, torch.Tensor] + + +class _Cosmos3KVTransformer(Protocol): + cached_kv: list[KVPair] | None + cached_freqs_gen: tuple[torch.Tensor, torch.Tensor] | None + + +class Cosmos3UNDKV(StateObject[KVPair, KVPair]): + """One layer of fixed, encode-once Cosmos3 UND conditioning K/V.""" + + def __init__(self) -> None: + self._kv: KVPair | None = None + + def allocate(self) -> None: + self._kv = None + + def commit(self, payload: KVPair | None = None) -> None: + value = self._staged if payload is None else payload + if value is None: + raise RuntimeError("Cosmos3UNDKV.commit requires K/V payload") + self._kv = value + self.discard() + + def view(self, *, include_staged: bool = True) -> KVPair: + if include_staged and self._staged is not None: + return self._staged + if self._kv is None: + raise RuntimeError("Cosmos3UNDKV is not resident") + return self._kv + + def reset(self) -> None: + self._kv = None + self.discard() + + def nbytes_by_device(self, seen: SeenStorages) -> dict[str, int]: + return {} if self._kv is None else device_bytes(self._kv, seen) + + @property + def resident(self) -> bool: + return self._kv is not None + + +def _layer_key(layer_index: int, is_negative: bool) -> str: + """Session key for one (layer, CFG branch) UND K/V object (cf. DreamZero ``_xattn_key``).""" + return f"und_kv_neg/{layer_index}" if is_negative else f"und_kv/{layer_index}" + + +def _freqs_key(is_negative: bool) -> str: + return "und_freqs_gen_neg" if is_negative else "und_freqs_gen" + + +class Cosmos3StateAdapter: + """Session-keyed view over Cosmos3's per-branch UND text K/V. + + A fresh adapter is built per forward; the durable state lives in the manager + (mirroring ``DreamZeroStateAdapter``). + """ + + def __init__(self, session_id: str | None, manager: SessionStateManager) -> None: + self._session_id = session_id + # Pin the session: the manager may evict it from its lookup table under + # LRU pressure, but an adapter holding a reference keeps its state alive. + self._session = manager.get_or_create_session(session_id) + + @property + def _num_layers(self) -> int | None: + value = self._session.attrs.get("_num_layers") + return value if isinstance(value, int) else None + + def is_branch_initialized(self, is_negative: bool) -> bool: + """Whether this CFG branch's UND K/V has been stored (encode-once).""" + if self._num_layers is None: + return False + obj = self._session.get(_layer_key(0, is_negative)) + return obj is not None and obj.resident + + def get_branch_kv(self, is_negative: bool) -> list[KVPair] | None: + """Return the branch's ``cached_kv`` as ``list[(K, V)]`` by layer, or ``None``.""" + if not self.is_branch_initialized(is_negative): + return None + num_layers = self._num_layers + assert num_layers is not None + out: list[KVPair] = [] + for i in range(num_layers): + obj = self._session.get(_layer_key(i, is_negative)) + if not isinstance(obj, Cosmos3UNDKV) or not obj.resident: + raise RuntimeError("Cosmos3 UND K/V is partially initialized.") + out.append(obj.view()) + return out + + def set_branch_kv(self, is_negative: bool, cached_kv: list[KVPair]) -> None: + """Store this branch's freshly computed ``cached_kv`` once.""" + self._session.attrs["_num_layers"] = len(cached_kv) + for i, (k, v) in enumerate(cached_kv): + obj = self._session.get(_layer_key(i, is_negative)) + if not isinstance(obj, Cosmos3UNDKV): + obj = Cosmos3UNDKV() + obj.allocate() + self._session.put(_layer_key(i, is_negative), obj) + obj.commit((k, v)) + + def get_branch_freqs_gen(self, is_negative: bool) -> tuple[torch.Tensor, torch.Tensor] | None: + value = self._session.attrs.get(_freqs_key(is_negative)) + if isinstance(value, tuple) and len(value) == 2 and all(isinstance(freq, torch.Tensor) for freq in value): + return value + return None + + def load_into_transformer(self, transformer: _Cosmos3KVTransformer, is_negative: bool) -> None: + """Assign this branch's UND K/V and GEN RoPE frequencies onto the transformer.""" + transformer.cached_kv = self.get_branch_kv(is_negative) + transformer.cached_freqs_gen = self.get_branch_freqs_gen(is_negative) + + def capture_from_transformer(self, transformer: _Cosmos3KVTransformer, is_negative: bool) -> None: + """Capture the freshly computed ``cached_kv`` into the session if not already stored.""" + if not self.is_branch_initialized(is_negative): + cached_kv = transformer.cached_kv + if cached_kv is None: + raise RuntimeError("Cosmos3 transformer did not produce UND K/V") + freqs_gen = transformer.cached_freqs_gen + if freqs_gen is None: + raise RuntimeError("Cosmos3 transformer did not produce GEN RoPE frequencies") + self.set_branch_kv(is_negative, cached_kv) + self._session.attrs[_freqs_key(is_negative)] = freqs_gen + + def reset(self) -> None: + """Clear all UND K/V objects + metadata for this session (replaces ``reset_cache()``).""" + self._session.reset()