[Feat] Cosmos3 session-memory port for UND text K/V (RFC #4480 Phase 0) (#4657)

Signed-off-by: linzhenpl07 <linzhenpl07@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

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.

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)

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)

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,
             )
 

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

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()