[Perf][Cosmos3] Add SeaCache support for Cosmos3 (#6922)

Signed-off-by: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>

recipes/cosmos3/Cosmos3-Nano.md

@@ -32,13 +32,13 @@ mode is selected per request:
   T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
   the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
 - **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks:
-  - `forward_dynamics` — given a first frame or video **and** an action trajectory,
+    - `forward_dynamics` — given a first frame or video **and** an action trajectory,
     roll out the resulting video. Synchronous: `POST /v1/videos/sync`.
-  - `policy` — given a first frame or video and a language instruction,
+    - `policy` — given a first frame or video and a language instruction,
     **predict** the action trajectory (and a rollout video). Use the async
     `POST /v1/videos` endpoint and read the predicted action from the top-level
     `action` field.
-  - `inverse_dynamics` — given a video, **recover** the action trajectory. Use
+    - `inverse_dynamics` — given a video, **recover** the action trajectory. Use
     the async `POST /v1/videos` endpoint and read the recovered action from
     the top-level `action` field
     (`{data, shape, dtype, raw_action_dim, domain_id}`).
@@ -106,6 +106,25 @@ vllm serve nvidia/Cosmos3-Nano \
   --init-timeout 1800
 ```
 
+For Cosmos3, SeaCache is the recommended default choice when opting into
+diffusion caching. Caching remains opt-in; add `--cache-backend sea_cache` to
+the command above. Its default threshold and maximum cached-step streak are
+tuned for Cosmos3, so no `--cache-config` is required.
+
+Override individual defaults with a JSON cache configuration; for example:
+
+```bash
+vllm serve nvidia/Cosmos3-Nano \
+  --omni \
+  --cache-backend sea_cache \
+  --cache-config '{"sea_threshold":0.2,"sea_max_consecutive_cached":3}'
+```
+
+Lower `sea_threshold` values and smaller `sea_max_consecutive_cached` caps are
+more conservative. Higher values allow more cached steps and may improve
+speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to
+`0` removes the streak cap.
+
 To run **without** guardrails (you are responsible for license compliance),
 add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use
 `--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`;
@@ -382,12 +401,12 @@ OpenPI route uses the normal API authentication middleware. RoboLab's
 
 #### Notes
 
-- **Measured latency (1x B300, bf16, guardrails off):**
-  - T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s**
-  - T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s**
-  - I2V 1280×720, 189 frames @ 35 steps → ~**99 s**
-  - Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**.
-  - Guardrails-on overhead: ~8% on T2I, negligible on video.
+- **Measured latency (1x B300, bf16, guardrails off, diffusion cache off):**
+    - T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s**
+    - T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s**
+    - I2V 1280×720, 189 frames @ 35 steps → ~**99 s**
+    - Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**.
+    - Guardrails-on overhead: ~8% on T2I, negligible on video.
 - **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 GPU;
   full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
   ~33 GB on disk.
@@ -443,12 +462,12 @@ OpenPI route uses the normal API authentication middleware. RoboLab's
   `conditioning_fps`, `action_chunk_size`, `raw_action_dim`, `deterministic_seed`,
   and `session_id`.
 - **Known limitations:**
-  - Guardrails-on requires `cosmos-guardrail` **and** access to the gated
+    - Guardrails-on requires `cosmos-guardrail` **and** access to the gated
     `nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise
     the server fails at pipeline build with a gated-repo / safety-checker error.
-  - A guardrail-blocked prompt currently returns HTTP 500
+    - A guardrail-blocked prompt currently returns HTTP 500
     (`"Guardrail blocked prompt"`).
-  - Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
+    - Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
     online. Use async `POST /v1/videos` when you need the predicted/recovered
     action payload under the top-level `action` field; sync `/v1/videos/sync`
     returns raw MP4 bytes and does not expose action metadata in the response body.
@@ -579,14 +598,14 @@ so an md5 comparison across two runs also works as a smoke check.
   a ~390 s build of the VAE decode path *for that shape*, so warming 1024² images
   does nothing for 189-frame video. At 720p / 189 frames / 35 steps this is
   **540 s cold vs 161 s warm, with byte-identical output**.
-- **Measured on 1x MI350X (bf16, guardrails off, warm):** T2I 1024² @ 50 steps
+- **Measured on 1x MI350X (bf16, guardrails off, diffusion cache off, warm):** T2I 1024² @ 50 steps
   **~2.7 s**; T2V 1280×720 / 189 frames @ 35 steps **~161 s**, of which ~92% is
   the DiT denoise loop and ~4% VAE decode, so optimization effort belongs in the
   denoise loop. The optional flags act on different terms of the memory bill and
   are therefore complementary rather than redundant:
 
   | Flag | Latency | Peak reserved | Peak allocated | Acts on |
-  |---|---|---|---|---|
+  | --- | --- | --- | --- | --- |
   | *(none)* | 161 s | 120 GiB | 95 GiB | — |
   | `--vae-use-tiling` | 183 s (+13.5%) | **38 GiB** (−68%) | **36 GiB** (−62%) | decode activations |
   | `--enable-layerwise-offload` | 162 s (+0.8%) | 84 GiB (−30%) | 69 GiB (−27%) | weights |
@@ -768,11 +787,11 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
 
 #### Notes
 
-- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off):**
-  - T2I 1024² — 10 steps → ~8 s
-  - T2V 1280×720 @ 20 steps — 49 frames → ~55 s
-  - I2V 1280×720 @ 10 steps — 25 frames → ~25 s
-  - V2V 480×320 @ 10 steps — 17 frames → ~12 s
+- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off, diffusion cache off):**
+    - T2I 1024² — 10 steps → ~8 s
+    - T2V 1280×720 @ 20 steps — 49 frames → ~55 s
+    - I2V 1280×720 @ 10 steps — 25 frames → ~25 s
+    - V2V 480×320 @ 10 steps — 17 frames → ~12 s
 - **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 NPU;
   full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
   ~33 GB on disk.
@@ -785,9 +804,9 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
   (for model loading), `--tensor-parallel-size 8` for multi-NPU, and
   `--model-class-name Cosmos3OmniDiffusersPipeline` to force the pipeline class.
 - **Known limitations:**
-  - Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a
+    - Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a
     resolution-parsing bug; basic V2V without transfer hints works.
-  - FP8 online quantization and layerwise offload are not supported on NPU.
+    - FP8 online quantization and layerwise offload are not supported on NPU.
 
 ### 1x Ascend 910B / 910C (Atlas A2 / A3) — Offline generation
 

recipes/cosmos3/Cosmos3-Super.md

@@ -62,6 +62,21 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs;
 generation from ~83 GB to ~55 GB per GPU (2-GPU) with BF16-level quality (T2V
 composition can shift at the same seed).
 
+SeaCache is the recommended default choice when opting into diffusion caching,
+which remains opt-in. Add `--cache-backend sea_cache` to either command above;
+its defaults are tuned for Cosmos3.
+
+Override individual defaults with `--cache-config`; for example:
+
+```bash
+--cache-config '{"sea_threshold":0.2,"sea_max_consecutive_cached":3}'
+```
+
+Lower `sea_threshold` values and smaller `sea_max_consecutive_cached` caps are
+more conservative. Higher values allow more cached steps and may improve
+speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to
+`0` removes the streak cap.
+
 #### Verification
 
 Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full
@@ -109,12 +124,12 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
 
 #### Notes
 
-- **Measured (2x B300, bf16, guardrails off, official 2-GPU config above):**
-  - T2I 1024², 50 steps → **~6 s**
-  - T2V 1280×720, 189 frames, 35 steps → **~197 s**
-  - I2V 1280×720, 189 frames, 35 steps → **~200 s**
-  - T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
-  - (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
+- **Measured (2x B300, bf16, guardrails off, diffusion cache off, official 2-GPU config above):**
+    - T2I 1024², 50 steps → **~6 s**
+    - T2V 1280×720, 189 frames, 35 steps → **~197 s**
+    - I2V 1280×720, 189 frames, 35 steps → **~200 s**
+    - T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
+    - (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
 - **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
 - Same generation defaults, supported sizes, V2V reference-video controls
   (`condition_frame_indexes_vision`, `condition_video_keep`), and

tests/diffusion/cache/test_seacache.py

@@ -0,0 +1,635 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import datetime
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.cache.seacache import (
+    SeaCacheBackend,
+    SeaCacheConfig,
+    SeaCacheRootHook,
+    apply_sea_cache_hook,
+)
+from vllm_omni.diffusion.cache.seacache.sea_filter import (
+    apply_sea_filter,
+    extrapolate_residual,
+    indicator_distance,
+)
+from vllm_omni.diffusion.cache.selector import get_cache_backend
+from vllm_omni.diffusion.cache.teacache.extractors import EXTRACTOR_REGISTRY, CacheContext, get_extractor
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TinyCosmos3Transformer(torch.nn.Module):
+    """Small model that exposes a cache-neutral execution boundary."""
+
+    def _run_gen_layers(self, hidden_gen: torch.Tensor) -> torch.Tensor:
+        return hidden_gen * 2
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        text_ids: torch.Tensor | None = None,
+        text_mask: torch.Tensor | None = None,
+        video_shape: tuple[int, int, int] | None = None,
+        noisy_frame_mask: torch.Tensor | None = None,
+        control_latents: list[torch.Tensor] | torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        del timestep, text_ids, text_mask, video_shape, noisy_frame_mask
+        if kwargs:
+            raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}")
+        controls = (
+            []
+            if control_latents is None
+            else [control_latents]
+            if isinstance(control_latents, torch.Tensor)
+            else list(control_latents)
+        )
+        inputs = [*controls, hidden_states]
+        gen_input = torch.cat(
+            [value.movedim(1, -1).flatten(1, 3) for value in inputs],
+            dim=1,
+        )
+        return self._run_gen_layers(gen_input)
+
+
+def _extract_tiny_cosmos3_context(
+    module: TinyCosmos3Transformer,
+    hidden_states: torch.Tensor,
+    timestep: torch.Tensor,
+    text_ids: torch.Tensor | None = None,
+    text_mask: torch.Tensor | None = None,
+    video_shape: tuple[int, int, int] | None = None,
+    noisy_frame_mask: torch.Tensor | None = None,
+    control_latents: list[torch.Tensor] | torch.Tensor | None = None,
+    **kwargs,
+) -> CacheContext:
+    del timestep, text_ids, text_mask, video_shape
+    if kwargs:
+        raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}")
+    controls = (
+        []
+        if control_latents is None
+        else [control_latents]
+        if isinstance(control_latents, torch.Tensor)
+        else list(control_latents)
+    )
+    vision_items = [*controls, hidden_states]
+    gen_input = torch.cat(
+        [value.movedim(1, -1).flatten(1, 3) for value in vision_items],
+        dim=1,
+    )
+    return CacheContext(
+        modulated_input=gen_input,
+        hidden_states=gen_input,
+        encoder_hidden_states=None,
+        temb=torch.zeros_like(gen_input[:, 0]),
+        run_transformer_blocks=lambda: (module._run_gen_layers(gen_input),),
+        postprocess=lambda output: output,
+        extra_states={
+            "sea_cache_latents": vision_items,
+            "sea_cache_noisy_frame_mask": noisy_frame_mask,
+        },
+    )
+
+
+class Cosmos3OmniDiffusersPipeline:
+    def __init__(self) -> None:
+        self.transformer = TinyCosmos3Transformer()
+        self._current_step_index: int | None = None
+        self._current_sigma: float | None = None
+        self._num_timesteps: int | None = None
+
+    @property
+    def current_step_index(self) -> int | None:
+        return self._current_step_index
+
+    @property
+    def current_sigma(self) -> float | None:
+        return self._current_sigma
+
+    @property
+    def num_timesteps(self) -> int | None:
+        return self._num_timesteps
+
+
+def _latent(value: float) -> torch.Tensor:
+    return torch.full((1, 2, 2, 2, 2), value)
+
+
+def _run_step(
+    transformer: TinyCosmos3Transformer,
+    timestep: int,
+    value: float,
+    *,
+    hook: SeaCacheRootHook | None = None,
+    context: str = "cond",
+    control: float | None = None,
+    noisy_frame_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    controls = None if control is None else [_latent(control)]
+    cache_context = hook.cache_context(context) if hook is not None else nullcontext()
+    with torch.inference_mode(), cache_context:
+        return transformer(
+            hidden_states=_latent(value),
+            timestep=torch.tensor([timestep]),
+            noisy_frame_mask=noisy_frame_mask,
+            control_latents=controls,
+        )
+
+
+def _apply_test_hook(
+    transformer: TinyCosmos3Transformer,
+    metadata: SimpleNamespace,
+    config: SeaCacheConfig | None = None,
+) -> SeaCacheRootHook:
+    return apply_sea_cache_hook(
+        transformer,
+        config or SeaCacheConfig(threshold=100.0),
+        current_step_callback=lambda: metadata.step,
+        current_sigma_callback=lambda: metadata.sigma,
+        num_inference_steps_callback=lambda: metadata.num_steps,
+        extractor_fn=_extract_tiny_cosmos3_context,
+    )
+
+
+def test_config_validation() -> None:
+    assert SeaCacheConfig().threshold == 0.25
+    with pytest.raises(ValueError, match="residual_order"):
+        SeaCacheConfig(residual_order=-1)
+    with pytest.raises(ValueError, match="max_consecutive_cached"):
+        SeaCacheConfig(max_consecutive_cached=-1)
+
+
+def test_sea_filter_matches_reference_equation() -> None:
+    hidden = torch.randn(3, 4, 5, 2, dtype=torch.float32)
+    sigma = 0.4
+    power_exp = 3.0
+
+    spectrum = torch.fft.fftn(hidden, dim=(0, 1, 2))
+    gain = None
+    for axis in (0, 1, 2):
+        frequencies = torch.fft.fftfreq(hidden.shape[axis], dtype=torch.float32)
+        clean_power = 1.0 / (frequencies.abs().pow(power_exp) + 1e-16)
+        axis_gain = (1.0 - sigma) * clean_power / ((1.0 - sigma) ** 2 * clean_power + sigma**2 + 1e-16)
+        shape = [1] * hidden.ndim
+        shape[axis] = hidden.shape[axis]
+        gain = axis_gain.reshape(shape) if gain is None else gain * axis_gain.reshape(shape)
+    assert gain is not None
+    gain = gain / gain.mean()
+    expected = torch.fft.ifftn(spectrum * gain, dim=(0, 1, 2)).real
+
+    torch.testing.assert_close(
+        apply_sea_filter(hidden, sigma=sigma, power_exp=power_exp),
+        expected,
+    )
+
+
+def test_indicator_distance_and_linear_extrapolation() -> None:
+    previous = [torch.ones(2, 2)]
+    current = [torch.full((2, 2), 1.5)]
+    assert indicator_distance(current, previous) == pytest.approx(0.5)
+    assert indicator_distance([torch.ones(3)], previous) == float("inf")
+
+    history = [
+        (0, torch.full((2, 2), 2.0)),
+        (2, torch.full((2, 2), 6.0)),
+    ]
+    torch.testing.assert_close(
+        extrapolate_residual(history, step=3, order=1),
+        torch.full((2, 2), 8.0),
+    )
+    torch.testing.assert_close(
+        extrapolate_residual(history, step=3, order=0),
+        torch.full((2, 2), 6.0),
+    )
+    quadratic_history = [
+        (0, torch.zeros(2, 2)),
+        (1, torch.ones(2, 2)),
+        (2, torch.full((2, 2), 4.0)),
+    ]
+    torch.testing.assert_close(
+        extrapolate_residual(quadratic_history, step=3, order=2),
+        torch.full((2, 2), 9.0),
+    )
+
+
+def test_hook_skips_middle_steps_and_forces_endpoints() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=4)
+    hook = _apply_test_hook(
+        transformer,
+        metadata,
+        SeaCacheConfig(threshold=100.0, max_consecutive_cached=2),
+    )
+    hook.refresh(transformer)
+
+    for step, timestep in enumerate((1000, 750, 500, 250)):
+        metadata.step = step
+        metadata.sigma = timestep / 1000
+        _run_step(transformer, timestep, 1.0 - step * 0.01, hook=hook)
+
+    assert hook.full_count == 2
+    assert hook.skip_count == 2
+    assert [step for step, _ in hook.state_manager._states["cond"].history] == [0, 3]
+
+
+def test_parameter_sharded_hook_skips_when_all_shard_ranks_agree(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook._parameter_sharded = True
+    fs_group = object()
+    reduced_groups: list[object] = []
+
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(
+        parallel_state,
+        "get_fs_group",
+        lambda: SimpleNamespace(world_size=2, device_group=fs_group),
+    )
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1)
+
+    def all_reduce(decision, *, op, group):
+        assert op == torch.distributed.ReduceOp.MAX
+        reduced_groups.append(group)
+
+    monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce)
+
+    _run_step(transformer, 1000, 1.0, hook=hook)
+    metadata.step = 1
+    metadata.sigma = 0.5
+    _run_step(transformer, 500, 0.99, hook=hook)
+
+    assert hook.full_count == 1
+    assert hook.skip_count == 1
+    assert reduced_groups == [fs_group, fs_group]
+
+
+def test_parameter_sharded_peer_forces_full_compute(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook._parameter_sharded = True
+
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(
+        parallel_state,
+        "get_fs_group",
+        lambda: SimpleNamespace(world_size=2, device_group=object()),
+    )
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1)
+    reduce_count = 0
+
+    def all_reduce(decision, *, op, group):
+        nonlocal reduce_count
+        del op, group
+        reduce_count += 1
+        if reduce_count == 2:
+            decision.fill_(1)
+
+    monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce)
+
+    _run_step(transformer, 1000, 1.0, hook=hook)
+    metadata.step = 1
+    metadata.sigma = 0.5
+    _run_step(transformer, 500, 0.99, hook=hook)
+
+    assert hook.full_count == 2
+    assert hook.skip_count == 0
+    assert hook.state_manager._states["cond"].accumulated_distance == 0.0
+
+
+def _uneven_cfg_sharded_worker(
+    rank: int,
+    init_method: str,
+    result_queue: torch.multiprocessing.Queue,
+) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    torch.distributed.init_process_group(
+        "gloo",
+        init_method=init_method,
+        rank=rank,
+        world_size=4,
+        timeout=datetime.timedelta(seconds=10),
+    )
+    try:
+        fs_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        sp_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        cfg_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 2], [1, 3])
+        ]
+        fs_group = next(group for ranks, group in fs_groups if rank in ranks)
+        sp_group = next(group for ranks, group in sp_groups if rank in ranks)
+        cfg_group = next(group for ranks, group in cfg_groups if rank in ranks)
+        parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group)
+        parallel_state._SP = SimpleNamespace(world_size=2, device_group=sp_group)
+
+        hook = SeaCacheRootHook(SeaCacheConfig())
+        hook._parameter_sharded = True
+        hook._synchronize_compute(False, torch.device("cpu"))
+        if rank in (0, 1):
+            hook._synchronize_compute(False, torch.device("cpu"))
+
+        gathered = [torch.zeros(1) for _ in range(2)]
+        torch.distributed.all_gather(
+            gathered,
+            torch.tensor([rank], dtype=torch.float32),
+            group=cfg_group,
+        )
+        result_queue.put((rank, [int(value.item()) for value in gathered]))
+    finally:
+        parallel_state._FS = None
+        parallel_state._SP = None
+        torch.distributed.destroy_process_group()
+
+
+def test_parameter_sharded_hook_handles_uneven_cfg_branch_dispatch() -> None:
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    result_queue = manager.Queue()
+    torch.multiprocessing.spawn(
+        _uneven_cfg_sharded_worker,
+        args=(get_distributed_init_method("seacache_uneven_cfg_"), result_queue),
+        nprocs=4,
+    )
+
+    results = sorted(result_queue.get(timeout=2) for _ in range(4))
+    assert results == [
+        (0, [0, 2]),
+        (1, [1, 3]),
+        (2, [0, 2]),
+        (3, [1, 3]),
+    ]
+
+
+def _hybrid_sp_worker(
+    rank: int,
+    init_method: str,
+    result_queue: torch.multiprocessing.Queue,
+) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    torch.distributed.init_process_group(
+        "gloo",
+        init_method=init_method,
+        rank=rank,
+        world_size=4,
+        timeout=datetime.timedelta(seconds=10),
+    )
+    try:
+        fs_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        ulysses_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        ring_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 2], [1, 3])
+        ]
+        sp_group = torch.distributed.new_group(
+            [0, 1, 2, 3],
+            timeout=datetime.timedelta(seconds=10),
+        )
+        fs_group = next(group for ranks, group in fs_groups if rank in ranks)
+        ulysses_group = next(group for ranks, group in ulysses_groups if rank in ranks)
+        ring_group = next(group for ranks, group in ring_groups if rank in ranks)
+        parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group)
+        parallel_state._SP = SimpleNamespace(
+            world_size=4,
+            device_group=sp_group,
+            ulysses_world_size=2,
+            ulysses_group=ulysses_group,
+            ring_world_size=2,
+            ring_group=ring_group,
+        )
+
+        compute = rank == 3
+        non_sharded_hook = SeaCacheRootHook(SeaCacheConfig())
+        non_sharded_result = non_sharded_hook._synchronize_compute(compute, torch.device("cpu"))
+
+        sharded_hook = SeaCacheRootHook(SeaCacheConfig())
+        sharded_hook._parameter_sharded = True
+        sharded_result = sharded_hook._synchronize_compute(compute, torch.device("cpu"))
+        result_queue.put((rank, non_sharded_result, sharded_result))
+    finally:
+        parallel_state._FS = None
+        parallel_state._SP = None
+        torch.distributed.destroy_process_group()
+
+
+def test_hook_synchronizes_full_hybrid_sequence_parallel_group() -> None:
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    result_queue = manager.Queue()
+    torch.multiprocessing.spawn(
+        _hybrid_sp_worker,
+        args=(get_distributed_init_method("seacache_hybrid_sp_"), result_queue),
+        nprocs=4,
+    )
+
+    results = sorted(result_queue.get(timeout=2) for _ in range(4))
+    assert results == [
+        (0, True, True),
+        (1, True, True),
+        (2, True, True),
+        (3, True, True),
+    ]
+
+
+def test_parameter_sharded_hook_fails_open_without_distributed_state(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    hook = SeaCacheRootHook(SeaCacheConfig())
+    hook._parameter_sharded = True
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
+
+    assert hook._synchronize_compute(False, torch.device("cpu")) is True
+
+
+def test_hook_keeps_three_transfer_branches_separate() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+
+    contexts = ("cond", "cond_no_control", "uncond")
+    for step, (timestep, value) in enumerate(((1000, 1.0), (500, 0.9))):
+        metadata.step = step
+        metadata.sigma = timestep / 1000
+        _run_step(transformer, timestep, value, hook=hook, context=contexts[0], control=0.5)
+        _run_step(transformer, timestep, value, hook=hook, context=contexts[1])
+        _run_step(transformer, timestep, value, hook=hook, context=contexts[2], control=0.5)
+
+    assert set(hook.state_manager._states) == set(contexts)
+    assert hook.full_count == 3
+    assert hook.skip_count == 3
+    assert len(hook.state_manager._states["cond"].previous_indicator) == 2
+    assert len(hook.state_manager._states["cond_no_control"].previous_indicator) == 1
+
+
+def test_hook_fails_open_without_noisy_vision() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+    all_clean = torch.zeros(1, 1, 2, 1, 1)
+
+    _run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean)
+    metadata.step = 1
+    metadata.sigma = 0.5
+    _run_step(transformer, 500, 0.9, hook=hook, noisy_frame_mask=all_clean)
+
+    assert hook.full_count == 0
+    assert hook.skip_count == 0
+
+
+def test_conditioning_only_model_failure_runs_forward_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+    forward_calls = 0
+
+    def fail_forward(hidden_gen: torch.Tensor) -> torch.Tensor:
+        nonlocal forward_calls
+        forward_calls += 1
+        raise RuntimeError("model forward failed")
+
+    monkeypatch.setattr(transformer, "_run_gen_layers", fail_forward)
+    all_clean = torch.zeros(1, 1, 2, 1, 1)
+
+    with pytest.raises(RuntimeError, match="model forward failed"):
+        _run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean)
+
+    assert forward_calls == 1
+
+
+def test_hook_fails_open_without_explicit_context() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=2)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+
+    _run_step(transformer, 1000, 1.0)
+
+    assert hook.full_count == 0
+    assert hook.skip_count == 0
+    assert hook.state_manager._states == {}
+
+
+def test_hook_uses_exact_sigma_callback(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.cache.seacache import hook as hook_module
+
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=0.37, num_steps=2)
+    observed_sigmas: list[float] = []
+    original_filter = hook_module.apply_sea_filter
+
+    def recording_filter(hidden_states: torch.Tensor, sigma: float, power_exp: float) -> torch.Tensor:
+        observed_sigmas.append(sigma)
+        return original_filter(hidden_states, sigma, power_exp)
+
+    monkeypatch.setattr(hook_module, "apply_sea_filter", recording_filter)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+    _run_step(transformer, timestep=999, value=1.0, hook=hook)
+
+    assert observed_sigmas
+    assert all(sigma == pytest.approx(0.37) for sigma in observed_sigmas)
+
+
+def test_backend_selector_and_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context)
+
+    class FSDPTinyCosmos3Transformer(TinyCosmos3Transformer):
+        pass
+
+    assert get_extractor(FSDPTinyCosmos3Transformer) is _extract_tiny_cosmos3_context
+    backend = get_cache_backend(
+        "sea_cache",
+        {
+            "sea_threshold": 0.4,
+            "sea_residual_order": 0,
+        },
+    )
+    assert isinstance(backend, SeaCacheBackend)
+    assert backend.config.sea_threshold == 0.4
+
+    pipeline = Cosmos3OmniDiffusersPipeline()
+    backend.enable(pipeline)
+    hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME)
+    assert isinstance(hook, SeaCacheRootHook)
+    assert callable(getattr(pipeline, "_cache_context_factory", None))
+    for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"):
+        assert not hasattr(pipeline.transformer, name)
+    pipeline._current_step_index = 0
+    pipeline._current_sigma = 1.0
+    pipeline._num_timesteps = 7
+    _run_step(pipeline.transformer, 1000, 1.0, hook=hook)
+    assert hook.full_count == 1
+
+    backend.refresh(pipeline, num_inference_steps=7)
+    assert hook.full_count == 0
+    assert hook.skip_count == 0
+    assert hook.state_manager._states == {}
+
+
+def test_backend_uses_resolved_pipeline_metadata_not_refresh_argument(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context)
+    pipeline = Cosmos3OmniDiffusersPipeline()
+    backend = SeaCacheBackend(DiffusionCacheConfig())
+    backend.enable(pipeline)
+    hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME)
+    assert isinstance(hook, SeaCacheRootHook)
+    backend.refresh(pipeline, num_inference_steps=35)
+    pipeline._num_timesteps = 50
+    pipeline._current_step_index = 0
+    pipeline._current_sigma = 1.0
+    _run_step(pipeline.transformer, 1000, 1.0, hook=hook)
+    assert hook.full_count == 1
+    assert hook.num_inference_steps_callback is not None
+    assert hook.num_inference_steps_callback() == 50
+
+    pipeline._current_step_index = 1
+    pipeline._current_sigma = 0.98
+    _run_step(pipeline.transformer, 980, 0.99, hook=hook)
+    assert hook.full_count == 1
+    assert hook.skip_count == 1
+
+
+def test_shared_config_defaults() -> None:
+    config = DiffusionCacheConfig()
+    assert config.sea_threshold == 0.25
+    assert config.sea_residual_order == 1
+    assert config.sea_max_consecutive_cached == 2
+    assert config.sea_power_exp == 3.0

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -1,11 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
 import json
 import sys
 import types
+from contextlib import contextmanager
 from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Any
@@ -85,6 +86,7 @@ class StubScheduler:
         flow_shift: float = 1.0,
     ) -> None:
         self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
+        self.sigmas = self.timesteps.float() / 10
         self.config = SimpleNamespace(
             num_train_timesteps=1000,
             flow_shift=flow_shift,
@@ -113,6 +115,7 @@ class StubScheduler:
         else:
             assert num_inference_steps is not None
             self.timesteps = torch.arange(num_inference_steps, 0, -1, dtype=torch.int64, device=device)
+        self.sigmas = self.timesteps.float()
 
     def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs):
         del kwargs
@@ -250,7 +253,7 @@ def passthrough_progress_bar(iterable):
 
 @pytest.fixture(autouse=True)
 def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch):
-    module = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
+    module: Any = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
     module.is_guardrails_enabled = lambda od_config, sampling_params=None: False
     module.ensure_initialized = lambda od_config: None
     module.check_text_safety = lambda text: None
@@ -283,6 +286,8 @@ def make_cosmos3_pipeline():
         pipeline.is_edge_model = False
         pipeline._guidance_scale = None
         pipeline._num_timesteps = None
+        pipeline._current_step_index = None
+        pipeline._current_sigma = None
         pipeline._cosmos3_branch_caches = None
         pipeline._cache_dit_requires_paired_cfg = False
         pipeline._sound_tokenizer = None
@@ -300,7 +305,7 @@ def sequential_cfg_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
 
 
 def make_sampling_params(**overrides: Any) -> SimpleNamespace:
-    values = {
+    values: dict[str, Any] = {
         "height": None,
         "width": None,
         "num_frames": None,
@@ -814,7 +819,7 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file(
     t_list = [1.0, 0.75, 0.5, 0.25]
     scheduler_dir = tmp_path / "scheduler"
     scheduler_dir.mkdir()
-    scheduler_config = {"_class_name": scheduler_class_name}
+    scheduler_config: dict[str, Any] = {"_class_name": scheduler_class_name}
     if expected_distilled:
         scheduler_config["fixed_step_sampler_config"] = {"sample_type": "sde", "t_list": t_list}
     (scheduler_dir / "scheduler_config.json").write_text(json.dumps(scheduler_config))
@@ -1531,7 +1536,7 @@ def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.Monke
         fused_add_rms_norm: list[str]
         custom_op: list[str]
 
-    fake_kernel = types.ModuleType("vllm.config.kernel")
+    fake_kernel: Any = types.ModuleType("vllm.config.kernel")
     fake_kernel.IrOpPriorityConfig = FakeIrOpPriorityConfig
     monkeypatch.setitem(sys.modules, fake_kernel.__name__, fake_kernel)
 
@@ -2154,6 +2159,47 @@ 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_publishes_exact_seacache_metadata_and_cfg_contexts(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.scheduler.sigmas = torch.tensor([0.91, 0.17])
+    observations: list[tuple[str, int | None, float | None, int | None]] = []
+
+    class RecordingHook:
+        @contextmanager
+        def cache_context(self, name: str):
+            sigma = pipeline.current_sigma
+            observations.append(
+                (
+                    name,
+                    pipeline.current_step_index,
+                    None if sigma is None else float(sigma),
+                    pipeline.num_timesteps,
+                )
+            )
+            yield
+
+    pipeline._cache_context_factory = RecordingHook().cache_context
+    pipeline.diffuse(
+        latents=torch.zeros(1, 2, 1, 1, 1),
+        timesteps=torch.tensor([900, 100]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
+    )
+
+    assert observations == [
+        ("cond", 0, pytest.approx(0.91), 2),
+        ("uncond", 0, pytest.approx(0.91), 2),
+        ("cond", 1, pytest.approx(0.17), 2),
+        ("uncond", 1, pytest.approx(0.17), 2),
+    ]
+    assert pipeline.current_step_index is None
+    assert pipeline.current_sigma is None
+
+
 def test_diffuse_drops_session_when_progress_iteration_fails(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     pipeline._use_session_state = True
@@ -2213,6 +2259,53 @@ 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_uses_named_seacache_contexts(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.scheduler.sigmas = torch.tensor([0.42])
+    contexts: list[tuple[str, int | None, float | None, int | None]] = []
+
+    class RecordingHook:
+        @contextmanager
+        def cache_context(self, name: str):
+            sigma = pipeline.current_sigma
+            contexts.append(
+                (
+                    name,
+                    pipeline.current_step_index,
+                    None if sigma is None else float(sigma),
+                    pipeline.num_timesteps,
+                )
+            )
+            yield
+
+    pipeline._cache_context_factory = RecordingHook().cache_context
+    latents = torch.zeros(1, 2, 1, 1, 1)
+    velocity_mask = torch.ones(1, 1, 1, 1, 1)
+    pipeline.diffuse_transfer(
+        latents=latents,
+        timesteps=torch.tensor([7]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        control_guidance=1.5,
+        control_guidance_interval=None,
+        control_latents=[torch.zeros_like(latents)],
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask},
+        velocity_mask=velocity_mask,
+        condition_latents=torch.zeros_like(latents),
+    )
+
+    assert contexts == [
+        ("cond", 0, pytest.approx(0.42), 1),
+        ("cond_no_control", 0, pytest.approx(0.42), 1),
+        ("uncond", 0, pytest.approx(0.42), 1),
+    ]
+    assert pipeline.current_step_index is None
+    assert pipeline.current_sigma is None
+
+
 def test_diffuse_transfer_rejects_session_state_manager(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     pipeline._use_session_state = True
@@ -2583,7 +2676,7 @@ def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -
 
 class TestForwardRouting:
     def _install_forward_stubs(self, pipeline):
-        captured: dict[str, object] = {"diffuse_calls": [], "prepare_calls": []}
+        captured: dict[str, Any] = {"diffuse_calls": [], "prepare_calls": []}
 
         def fake_format(
             prompt,

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -513,6 +513,108 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
+def test_cache_execution_residual_spans_final_gen_norm(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.cache.teacache.extractors import extract_cosmos3_context
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    class TrackingNorm(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.calls = 0
+
+        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            self.calls += 1
+            return hidden_states + 5.0
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    norm = TrackingNorm()
+    model.norm_moe_gen = norm
+    captured: dict[str, torch.Tensor] = {}
+
+    def run_gen_layers(hidden_gen: torch.Tensor, **kwargs) -> torch.Tensor:
+        del kwargs
+        captured["input"] = hidden_gen.detach().clone()
+        return hidden_gen + 2.0
+
+    monkeypatch.setattr(model, "_run_gen_layers", run_gen_layers)
+    forward_kwargs = {
+        "hidden_states": torch.zeros(1, 2, 1, 2, 2),
+        "timestep": torch.tensor([1.0]),
+        "text_ids": torch.tensor([[1, 2]], dtype=torch.long),
+        "text_mask": torch.ones(1, 2, dtype=torch.long),
+        "video_shape": (1, 2, 2),
+        "fps": 24.0,
+    }
+
+    full_output = model(**forward_kwargs)
+
+    assert norm.calls == 1
+    norm.calls = 0
+
+    ctx = extract_cosmos3_context(model, **forward_kwargs)
+    execution_input = ctx.hidden_states.detach().clone()
+    execution_output = ctx.run_transformer_blocks()[0]
+    residual = execution_output - execution_input
+    extracted_output = ctx.postprocess(execution_output)
+
+    assert norm.calls == 1
+    expected_residual = ((captured["input"] + 2.0) + 5.0) - captured["input"]
+    assert torch.equal(residual, expected_residual)
+    assert torch.equal(extracted_output, full_output)
+
+    norm.calls = 0
+
+    def fail_if_gen_layers_run(*args, **kwargs):
+        del args, kwargs
+        raise AssertionError("SeaCache hit unexpectedly executed GEN layers")
+
+    monkeypatch.setattr(model, "_run_gen_layers", fail_if_gen_layers_run)
+
+    cached_ctx = extract_cosmos3_context(model, **forward_kwargs)
+    cached_output = cached_ctx.postprocess(cached_ctx.hidden_states + residual)
+
+    assert norm.calls == 0
+    assert torch.equal(cached_output, full_output)
+    for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"):
+        assert not hasattr(model, name)
+
+
+def test_no_cache_still_runs_final_gen_norm_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    class CountingNorm(nn.Module):
+        def __init__(self, original: nn.Module) -> None:
+            super().__init__()
+            self.original = original
+            self.calls = 0
+
+        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            self.calls += 1
+            return self.original(hidden_states)
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    norm = CountingNorm(model.norm_moe_gen)
+    model.norm_moe_gen = norm
+
+    output = model(
+        hidden_states=torch.zeros(1, 2, 1, 2, 2),
+        timestep=torch.tensor([1.0]),
+        text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+        text_mask=torch.ones(1, 2, dtype=torch.long),
+        video_shape=(1, 2, 2),
+        fps=24.0,
+    )
+
+    assert norm.calls == 1
+    assert tuple(output.shape) == (1, 2, 1, 2, 2)
+
+
 def test_qwen_gen_mlp_remains_gated() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer
 

vllm_omni/diffusion/cache/seacache/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from vllm_omni.diffusion.cache.seacache.backend import SeaCacheBackend
+from vllm_omni.diffusion.cache.seacache.config import SeaCacheConfig
+from vllm_omni.diffusion.cache.seacache.hook import (
+    SeaCacheRootHook,
+    apply_sea_cache_hook,
+)
+from vllm_omni.diffusion.cache.seacache.state import SeaCacheState
+
+__all__ = [
+    "SeaCacheBackend",
+    "SeaCacheConfig",
+    "SeaCacheRootHook",
+    "SeaCacheState",
+    "apply_sea_cache_hook",
+]

vllm_omni/diffusion/cache/seacache/backend.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from typing import Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.cache.base import CacheBackend
+from vllm_omni.diffusion.cache.seacache.config import SeaCacheConfig
+from vllm_omni.diffusion.cache.seacache.hook import (
+    SeaCacheRootHook,
+    apply_sea_cache_hook,
+)
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+logger = init_logger(__name__)
+
+
+def enable_cosmos3_seacache(
+    pipeline: Any,
+    config: DiffusionCacheConfig,
+) -> SeaCacheRootHook:
+    transformer = getattr(pipeline, "transformer", None)
+    if transformer is None:
+        raise ValueError("SeaCache requires a pipeline with a transformer")
+
+    sea_config = SeaCacheConfig(
+        threshold=config.sea_threshold,
+        residual_order=config.sea_residual_order,
+        max_consecutive_cached=config.sea_max_consecutive_cached,
+        power_exp=config.sea_power_exp,
+    )
+    hook = apply_sea_cache_hook(
+        transformer,
+        sea_config,
+        current_step_callback=lambda: getattr(pipeline, "current_step_index", None),
+        current_sigma_callback=lambda: getattr(pipeline, "current_sigma", None),
+        num_inference_steps_callback=lambda: getattr(
+            pipeline,
+            "num_timesteps",
+            None,
+        ),
+    )
+    logger.info(
+        "SeaCache enabled for %s (threshold=%s, residual_order=%d, max_consecutive_cached=%d, power_exp=%s)",
+        pipeline.__class__.__name__,
+        sea_config.threshold,
+        sea_config.residual_order,
+        sea_config.max_consecutive_cached,
+        sea_config.power_exp,
+    )
+    return hook
+
+
+CUSTOM_SEACACHE_ENABLERS = {
+    "Cosmos3OmniDiffusersPipeline": enable_cosmos3_seacache,
+    "Cosmos3OmniPipeline": enable_cosmos3_seacache,
+}
+
+
+class SeaCacheBackend(CacheBackend):
+    """Backend for spectral-evolution-aware diffusion caching."""
+
+    def __init__(self, config: DiffusionCacheConfig):
+        super().__init__(config)
+        self._transformer_id: int | None = None
+
+    def enable(self, pipeline: Any) -> None:
+        pipeline_type = pipeline.__class__.__name__
+        enabler = CUSTOM_SEACACHE_ENABLERS.get(pipeline_type)
+        if enabler is None:
+            raise ValueError(f"SeaCache does not support pipeline type {pipeline_type}")
+        hook = enabler(pipeline, self.config)
+        self._transformer_id = id(pipeline.transformer)
+        self.enabled = True
+        pipeline._cache_context_factory = hook.cache_context
+
+    def refresh(
+        self,
+        pipeline: Any,
+        num_inference_steps: int,
+        verbose: bool = True,
+    ) -> None:
+        del num_inference_steps
+        transformer = getattr(pipeline, "transformer", None)
+        if transformer is None:
+            raise ValueError("SeaCache requires a pipeline with a transformer")
+        if not self.enabled or self._transformer_id != id(transformer):
+            self.enable(pipeline)
+
+        registry = getattr(transformer, "_hook_registry", None)
+        hook = registry.get_hook(SeaCacheRootHook._HOOK_NAME) if registry is not None else None
+        if not isinstance(hook, SeaCacheRootHook):
+            raise RuntimeError("SeaCache hook is not installed on the pipeline transformer")
+        hook.refresh(transformer)
+        pipeline._cache_context_factory = hook.cache_context
+        if verbose:
+            logger.debug("SeaCache state refreshed")

vllm_omni/diffusion/cache/seacache/config.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class SeaCacheConfig:
+    """Configuration for SeaCache.
+
+    Defaults are tuned for Cosmos3 and may require adjustment for other models.
+    """
+
+    threshold: float = 0.25
+    residual_order: int = 1
+    max_consecutive_cached: int = 2
+    power_exp: float = 3.0
+
+    def __post_init__(self) -> None:
+        if not math.isfinite(self.threshold) or self.threshold < 0:
+            raise ValueError(f"threshold must be finite and non-negative, got {self.threshold}")
+        if isinstance(self.residual_order, bool) or not isinstance(self.residual_order, int) or self.residual_order < 0:
+            raise ValueError(f"residual_order must be a non-negative integer, got {self.residual_order!r}")
+        if (
+            isinstance(self.max_consecutive_cached, bool)
+            or not isinstance(self.max_consecutive_cached, int)
+            or self.max_consecutive_cached < 0
+        ):
+            raise ValueError(
+                f"max_consecutive_cached must be a non-negative integer, got {self.max_consecutive_cached!r}"
+            )
+        if not math.isfinite(self.power_exp) or self.power_exp <= 0:
+            raise ValueError(f"power_exp must be finite and positive, got {self.power_exp}")

vllm_omni/diffusion/cache/seacache/hook.py

@@ -0,0 +1,397 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.cache.seacache.config import SeaCacheConfig
+from vllm_omni.diffusion.cache.seacache.sea_filter import (
+    apply_sea_filter,
+    extrapolate_residual,
+    indicator_distance,
+)
+from vllm_omni.diffusion.cache.seacache.state import SeaCacheState
+from vllm_omni.diffusion.cache.teacache.extractors import CacheContext, get_extractor
+from vllm_omni.diffusion.hooks import HookRegistry, ModelHook, StateManager
+
+logger = init_logger(__name__)
+
+
+def _is_parameter_sharded(module: torch.nn.Module) -> bool:
+    """Detect parameter-sharding runtimes whose collectives cannot be skipped."""
+    for submodule in module.modules():
+        module_type = type(submodule)
+        if callable(getattr(submodule, "_get_fsdp_state", None)):
+            return True
+        if module_type.__name__ == "FullyShardedDataParallel" and module_type.__module__.startswith(
+            "torch.distributed.fsdp"
+        ):
+            return True
+        for parameter in submodule.parameters(recurse=False):
+            parameter_type = type(parameter)
+            if (
+                parameter_type.__name__ == "FlatParameter"
+                and parameter_type.__module__.startswith("torch.distributed.fsdp")
+            ) or (
+                parameter_type.__name__ == "DTensor"
+                and parameter_type.__module__.startswith("torch.distributed.tensor")
+            ):
+                return True
+    return False
+
+
+class SeaCacheRootHook(ModelHook):
+    """Drive SeaCache gating and transformer forward control."""
+
+    _HOOK_NAME = "sea_cache"
+
+    def __init__(
+        self,
+        config: SeaCacheConfig,
+        *,
+        current_step_callback: Callable[[], int | torch.Tensor | None] | None = None,
+        current_sigma_callback: Callable[[], float | torch.Tensor | None] | None = None,
+        num_inference_steps_callback: Callable[[], int | torch.Tensor | None] | None = None,
+        extractor_fn: Callable[..., CacheContext] | None = None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.current_step_callback = current_step_callback
+        self.current_sigma_callback = current_sigma_callback
+        self.num_inference_steps_callback = num_inference_steps_callback
+        self.state_manager = StateManager(SeaCacheState)
+        self._warned_messages: set[str] = set()
+        self.full_count = 0
+        self.skip_count = 0
+        self.extractor_fn = extractor_fn
+        self._parameter_sharded = False
+        self._collective_skip_groups: list[torch.distributed.ProcessGroup] = []
+
+    def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
+        if self.extractor_fn is None:
+            self.extractor_fn = get_extractor(type(module))
+        self._parameter_sharded = _is_parameter_sharded(module)
+        seen_groups: set[int] = set()
+        for block in getattr(module, "gen_layers", ()):
+            registry = getattr(block, "_hook_registry", None)
+            dlo_hook = registry.get_hook("distributed_layerwise_offload") if registry is not None else None
+            group = getattr(dlo_hook, "dp_group", None)
+            if group is not None and int(getattr(dlo_hook, "dp_size", 1)) > 1 and id(group) not in seen_groups:
+                seen_groups.add(id(group))
+                self._collective_skip_groups.append(group)
+        return module
+
+    def _warn_once(self, message: str) -> None:
+        if message not in self._warned_messages:
+            logger.warning(message)
+            self._warned_messages.add(message)
+
+    @contextmanager
+    def cache_context(self, name: str) -> Iterator[None]:
+        previous_context = self.state_manager._context
+        self.state_manager.set_context(name)
+        try:
+            yield
+        finally:
+            self.state_manager.set_context(previous_context)
+
+    def _build_indicator(
+        self,
+        vision_items: list[torch.Tensor] | None,
+        sigma: float,
+    ) -> list[torch.Tensor] | None:
+        if not vision_items:
+            return None
+        hidden_states = vision_items[-1]
+        if not isinstance(hidden_states, torch.Tensor) or hidden_states.ndim != 5:
+            return None
+
+        # Controls precede the denoised target in packed vision-token order,
+        # and each vision item is filtered independently.
+        if any(
+            not isinstance(item, torch.Tensor)
+            or item.ndim != 5
+            or item.shape[0] != hidden_states.shape[0]
+            or item.shape[1:] != hidden_states.shape[1:]
+            for item in vision_items
+        ):
+            return None
+
+        indicator = []
+        for batch_index in range(hidden_states.shape[0]):
+            for latent in vision_items:
+                thwc = latent[batch_index].movedim(0, -1)
+                indicator.append(
+                    apply_sea_filter(
+                        thwc,
+                        sigma=sigma,
+                        power_exp=self.config.power_exp,
+                    ).detach()
+                )
+        return indicator or None
+
+    def _resolve_gate(
+        self,
+        state: SeaCacheState,
+        indicator: list[torch.Tensor] | None,
+        step: int,
+        num_inference_steps: int,
+    ) -> bool:
+        if state.last_step is not None and step != state.last_step + 1:
+            state.reset()
+        state.last_step = step
+        max_consecutive = bool(
+            self.config.max_consecutive_cached and state.consecutive_cached >= self.config.max_consecutive_cached
+        )
+        forced_compute = (
+            step < 1
+            or step >= num_inference_steps - 1
+            or max_consecutive
+            or not state.history
+            or indicator is None
+            or state.previous_indicator is None
+        )
+        if forced_compute:
+            state.accumulated_distance = 0.0
+            state.previous_indicator = None if indicator is None else [value.detach() for value in indicator]
+            return True
+
+        assert indicator is not None
+        assert state.previous_indicator is not None
+        distance = indicator_distance(indicator, state.previous_indicator)
+        state.previous_indicator = [value.detach() for value in indicator]
+        if not math.isfinite(distance):
+            state.accumulated_distance = 0.0
+            self._warn_once("SeaCache indicator history changed shape, device, or dtype; running full.")
+            return True
+
+        state.accumulated_distance += distance
+        if state.accumulated_distance < self.config.threshold:
+            return False
+        state.accumulated_distance = 0.0
+        return True
+
+    def _synchronize_compute(self, compute: bool, device: torch.device) -> bool:
+        if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+            return True if self._parameter_sharded else compute
+        decision = torch.tensor(int(compute), dtype=torch.int32, device=device)
+        if self._parameter_sharded:
+            from vllm_omni.diffusion.distributed.parallel_state import (
+                get_fs_group,
+                get_sequence_parallel_world_size,
+                get_sp_group,
+            )
+
+            fs_group = get_fs_group()
+            if fs_group.world_size > 1:
+                torch.distributed.all_reduce(
+                    decision,
+                    op=torch.distributed.ReduceOp.MAX,
+                    group=fs_group.device_group,
+                )
+            if get_sequence_parallel_world_size() > 1:
+                torch.distributed.all_reduce(
+                    decision,
+                    op=torch.distributed.ReduceOp.MAX,
+                    group=get_sp_group().device_group,
+                )
+            return bool(decision.item())
+
+        for group in self._collective_skip_groups:
+            torch.distributed.all_reduce(
+                decision,
+                op=torch.distributed.ReduceOp.MAX,
+                group=group,
+            )
+        from vllm_omni.diffusion.distributed.parallel_state import (
+            get_sequence_parallel_world_size,
+            get_sp_group,
+        )
+
+        if get_sequence_parallel_world_size() > 1:
+            torch.distributed.all_reduce(
+                decision,
+                op=torch.distributed.ReduceOp.MAX,
+                group=get_sp_group().device_group,
+            )
+        return bool(decision.item())
+
+    @torch.compiler.disable
+    def new_forward(
+        self,
+        module: torch.nn.Module,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        if self.extractor_fn is None:
+            raise RuntimeError("SeaCache extractor was not initialized")
+        ctx = self.extractor_fn(module, *args, **kwargs)
+
+        if torch.is_grad_enabled():
+            self._warn_once("SeaCache is inference-only; autograd-enabled calls run in full.")
+            return self._run_uncached(ctx)
+        if self.state_manager._current_context is None:
+            self._warn_once("SeaCache requires an explicit cache context; running full.")
+            return self._run_uncached(ctx)
+        callbacks = (
+            self.current_step_callback,
+            self.current_sigma_callback,
+            self.num_inference_steps_callback,
+        )
+        if any(callback is None for callback in callbacks):
+            self._warn_once("SeaCache requires scheduler step, sigma, and step-count callbacks; running full.")
+            return self._run_uncached(ctx)
+        assert self.current_step_callback is not None
+        assert self.current_sigma_callback is not None
+        assert self.num_inference_steps_callback is not None
+
+        try:
+            extra_states = ctx.extra_states or {}
+            vision_items = extra_states.get("sea_cache_latents")
+            if not isinstance(vision_items, list):
+                raise ValueError("extractor did not provide SeaCache vision inputs")
+            noisy_frame_mask = extra_states.get("sea_cache_noisy_frame_mask")
+            conditioning_only = isinstance(noisy_frame_mask, torch.Tensor) and not bool(
+                torch.any(noisy_frame_mask != 0).item()
+            )
+
+            step = self.current_step_callback()
+            sigma = self.current_sigma_callback()
+            num_inference_steps = self.num_inference_steps_callback()
+            if isinstance(step, torch.Tensor):
+                step = step.item()
+            if isinstance(sigma, torch.Tensor):
+                sigma = sigma.item()
+            if isinstance(num_inference_steps, torch.Tensor):
+                num_inference_steps = num_inference_steps.item()
+            if step is None or sigma is None or num_inference_steps is None:
+                raise ValueError("scheduler metadata is unavailable")
+            step = int(step)
+            sigma = float(sigma)
+            num_inference_steps = int(num_inference_steps)
+            if (
+                step < 0
+                or num_inference_steps <= 0
+                or step >= num_inference_steps
+                or not math.isfinite(sigma)
+                or not 0.0 <= sigma <= 1.0
+            ):
+                raise ValueError("expected a valid step index and exact sigma in [0, 1]")
+        except (IndexError, TypeError, ValueError, RuntimeError) as error:
+            self._warn_once(f"SeaCache metadata is invalid; running full: {error}")
+            return self._run_uncached(ctx)
+
+        if conditioning_only:
+            self._warn_once("SeaCache requires noisy vision; conditioning-only calls run in full.")
+            return self._run_uncached(ctx)
+
+        state: SeaCacheState = self.state_manager.get_state()
+        try:
+            indicator = self._build_indicator(vision_items, sigma)
+        except (TypeError, ValueError, RuntimeError) as error:
+            self._warn_once(f"SeaCache could not construct its vision indicator; running full: {error}")
+            indicator = None
+
+        local_compute = self._resolve_gate(state, indicator, step, num_inference_steps)
+        should_compute = self._synchronize_compute(local_compute, ctx.hidden_states.device)
+        if should_compute and not local_compute:
+            state.accumulated_distance = 0.0
+
+        if should_compute:
+            self.full_count += 1
+            output = self._run_full_stack(ctx)
+            result = ctx.postprocess(output)
+            self._record_execution(state, step, ctx.hidden_states, output)
+            return result
+
+        residual = extrapolate_residual(
+            state.history,
+            step,
+            self.config.residual_order,
+        )
+        if residual.device != ctx.hidden_states.device:
+            residual = residual.to(ctx.hidden_states.device)
+        state.consecutive_cached += 1
+        self.skip_count += 1
+
+        can_reuse = (
+            residual.shape == ctx.hidden_states.shape
+            and residual.device == ctx.hidden_states.device
+            and residual.dtype == ctx.hidden_states.dtype
+        )
+        if can_reuse:
+            return ctx.postprocess(ctx.hidden_states + residual)
+
+        output = self._run_full_stack(ctx)
+        result = ctx.postprocess(output)
+        self._record_execution(state, step, ctx.hidden_states, output)
+        return result
+
+    @staticmethod
+    def _run_full_stack(ctx: CacheContext) -> torch.Tensor:
+        outputs = ctx.run_transformer_blocks()
+        if not outputs:
+            raise RuntimeError("Cache extractor returned no transformer outputs")
+        return outputs[0]
+
+    @staticmethod
+    def _run_uncached(ctx: CacheContext) -> Any:
+        return ctx.postprocess(SeaCacheRootHook._run_full_stack(ctx))
+
+    def _record_execution(
+        self,
+        state: SeaCacheState,
+        step: int,
+        execution_input: torch.Tensor,
+        output: torch.Tensor,
+    ) -> None:
+        if (
+            output.shape == execution_input.shape
+            and output.device == execution_input.device
+            and output.dtype == execution_input.dtype
+        ):
+            state.history.append((step, (output - execution_input).detach().clone()))
+            state.history = state.history[-(self.config.residual_order + 1) :]
+            state.consecutive_cached = 0
+            return
+
+        state.history.clear()
+        state.accumulated_distance = 0.0
+        self._warn_once("SeaCache execution boundary returned an incompatible tensor; clearing cache history.")
+
+    def reset_state(self, module: torch.nn.Module) -> torch.nn.Module:
+        self.state_manager.reset()
+        self.full_count = 0
+        self.skip_count = 0
+        return module
+
+    def refresh(self, module: torch.nn.Module) -> None:
+        self.reset_state(module)
+
+
+def apply_sea_cache_hook(
+    module: torch.nn.Module,
+    config: SeaCacheConfig,
+    *,
+    current_step_callback: Callable[[], int | torch.Tensor | None] | None = None,
+    current_sigma_callback: Callable[[], float | torch.Tensor | None] | None = None,
+    num_inference_steps_callback: Callable[[], int | torch.Tensor | None] | None = None,
+    extractor_fn: Callable[..., CacheContext] | None = None,
+) -> SeaCacheRootHook:
+    registry = HookRegistry.get_or_create(module)
+    hook = SeaCacheRootHook(
+        config,
+        current_step_callback=current_step_callback,
+        current_sigma_callback=current_sigma_callback,
+        num_inference_steps_callback=num_inference_steps_callback,
+        extractor_fn=extractor_fn,
+    )
+    registry.register_hook(SeaCacheRootHook._HOOK_NAME, hook)
+    return hook

vllm_omni/diffusion/cache/seacache/sea_filter.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+
+# SeaCache paper: "Spectral-Evolution-Aware Cache for Accelerating Diffusion Models"
+# https://arxiv.org/abs/2602.18993
+def apply_sea_filter(hidden_states: torch.Tensor, sigma: float, power_exp: float = 3.0) -> torch.Tensor:
+    """Apply the separable SEA Wiener filter over ``(T, H, W)``."""
+    original_dtype = hidden_states.dtype
+    hidden_states = hidden_states.contiguous().float()
+    dimensions = (0, 1, 2)
+    spectrum = torch.fft.fftn(hidden_states, dim=dimensions)
+
+    sigma = max(1e-6, min(1.0 - 1e-6, float(sigma)))
+    signal_scale = 1.0 - sigma
+    noise_scale = sigma
+    gain: torch.Tensor | None = None
+    for axis in dimensions:
+        frequencies = torch.fft.fftfreq(
+            hidden_states.shape[axis],
+            device=hidden_states.device,
+            dtype=torch.float32,
+        )
+        clean_power = 1.0 / (frequencies.abs().pow(power_exp) + 1e-16)
+        axis_gain = signal_scale * clean_power / (signal_scale**2 * clean_power + noise_scale**2 + 1e-16)
+        axis_shape = [1] * hidden_states.ndim
+        axis_shape[axis] = axis_gain.shape[0]
+        reshaped_gain = axis_gain.reshape(axis_shape)
+        gain = reshaped_gain if gain is None else gain * reshaped_gain
+
+    assert gain is not None
+    gain = gain / gain.mean()
+    return torch.fft.ifftn(spectrum * gain, dim=dimensions).real.to(original_dtype)
+
+
+def relative_l1_distance(current: torch.Tensor, previous: torch.Tensor, eps: float = 1e-16) -> float:
+    numerator = (current.float() - previous.float()).abs().mean()
+    denominator = previous.float().abs().mean() + eps
+    return float((numerator / denominator).detach().cpu())
+
+
+def indicator_distance(
+    current: list[torch.Tensor] | None,
+    previous: list[torch.Tensor] | None,
+) -> float:
+    """Mean relative-L1 distance, or infinity for incompatible indicators."""
+    if current is None or previous is None or not current or len(current) != len(previous):
+        return float("inf")
+
+    total = 0.0
+    for current_value, previous_value in zip(current, previous, strict=True):
+        if (
+            current_value.shape != previous_value.shape
+            or current_value.device != previous_value.device
+            or current_value.dtype != previous_value.dtype
+        ):
+            return float("inf")
+        total += relative_l1_distance(current_value, previous_value)
+    distance = total / len(current)
+    return distance if math.isfinite(distance) else float("inf")
+
+
+def extrapolate_residual(
+    history: list[tuple[int, torch.Tensor]],
+    step: int,
+    order: int,
+) -> torch.Tensor:
+    """Evaluate a Newton residual polynomial at ``step``."""
+    if not history:
+        raise ValueError("SeaCache residual extrapolation requires non-empty history")
+
+    polynomial_order = max(0, min(order, len(history) - 1))
+    window = history[-(polynomial_order + 1) :]
+    steps = [history_step for history_step, _ in window]
+    coefficients = [residual.clone() for _, residual in window]
+
+    for level in range(1, polynomial_order + 1):
+        for index in range(polynomial_order, level - 1, -1):
+            denominator = steps[index] - steps[index - level]
+            if denominator == 0:
+                return window[-1][1].clone()
+            coefficients[index] = (coefficients[index] - coefficients[index - 1]) / float(denominator)
+
+    result = coefficients[polynomial_order]
+    for index in range(polynomial_order - 1, -1, -1):
+        result = result * float(step - steps[index]) + coefficients[index]
+    return result

vllm_omni/diffusion/cache/seacache/state.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import torch
+
+
+@dataclass(slots=True)
+class SeaCacheState:
+    """Per-cache-context trajectory state."""
+
+    last_step: int | None = None
+    accumulated_distance: float = 0.0
+    previous_indicator: list[torch.Tensor] | None = None
+    history: list[tuple[int, torch.Tensor]] = field(default_factory=list)
+    consecutive_cached: int = 0
+
+    def reset(self) -> None:
+        self.last_step = None
+        self.accumulated_distance = 0.0
+        self.previous_indicator = None
+        self.history.clear()
+        self.consecutive_cached = 0

vllm_omni/diffusion/cache/selector.py

@@ -1,8 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from typing import Any
 
 from vllm_omni.diffusion.cache.base import CacheBackend
 from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend
 from vllm_omni.diffusion.cache.magcache import MagCacheBackend
+from vllm_omni.diffusion.cache.seacache import SeaCacheBackend
 from vllm_omni.diffusion.cache.stepcache import StepCacheBackend
 from vllm_omni.diffusion.cache.teacache import TeaCacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
@@ -15,11 +19,12 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
     - cache_dit: Uses CacheDiTBackend with enable()/refresh() interface
     - tea_cache: Uses TeaCacheBackend with enable()/refresh() interface
     - mag_cache: Uses MagCacheBackend with enable()/refresh() interface
+    - sea_cache: Uses SeaCacheBackend for spectral residual caching
     - step_cache: Uses StepCacheBackend for DreamZero velocity step skip
 
     Args:
         cache_backend: Cache backend name ("cache_dit", "tea_cache",
-            "mag_cache", "step_cache", or None).
+            "mag_cache", "sea_cache", "step_cache", or None).
         cache_config: Cache configuration (dict or DiffusionCacheConfig instance).
 
     Returns:
@@ -40,10 +45,12 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
         return TeaCacheBackend(cache_config)
     elif cache_backend == "mag_cache":
         return MagCacheBackend(cache_config)
+    elif cache_backend == "sea_cache":
+        return SeaCacheBackend(cache_config)
     elif cache_backend in ("step_cache", "stepcache", "step_cache_dit"):
         return StepCacheBackend(cache_config)
     else:
         raise ValueError(
             f"Unsupported cache backend: {cache_backend}. "
-            "Supported: 'cache_dit', 'tea_cache', 'mag_cache', 'step_cache'"
+            "Supported: 'cache_dit', 'tea_cache', 'mag_cache', 'sea_cache', 'step_cache'"
         )

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -199,7 +199,7 @@ def extract_qwen_context(
     hidden_states, vid_freqs, txt_freqs = module.image_rope_prepare(hidden_states, img_shapes, txt_seq_lens)
     image_rotary_emb = (vid_freqs, txt_freqs)
 
-    timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype)
+    timestep = torch.as_tensor(timestep, device=hidden_states.device, dtype=hidden_states.dtype)
 
     # Call modulate_index_prepare instead of handling timestep directly.
     # For zero_cond_t=False: timestep unchanged, modulate_index=None.
@@ -618,6 +618,8 @@ def extract_flux2_klein_context(
     # ============================================================================
     dtype = hidden_states.dtype
 
+    if encoder_hidden_states is None:
+        raise ValueError("Flux2Klein requires encoder_hidden_states")
     num_txt_tokens = encoder_hidden_states.shape[1]
 
     timestep = timestep.to(dtype=dtype) * 1000
@@ -1466,6 +1468,77 @@ def extract_minimax_h3_context(
     )
 
 
+def extract_cosmos3_context(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+    timestep: torch.Tensor,
+    text_ids: torch.Tensor,
+    text_mask: torch.Tensor,
+    video_shape: tuple[int, int, int],
+    fps: float | None = None,
+    action_latents: torch.Tensor | None = None,
+    action_domain_ids: torch.Tensor | None = None,
+    action_noisy_mask: torch.Tensor | None = None,
+    action_start_frame_offset: int = 1,
+    action_fps: float | None = None,
+    sound_latents: torch.Tensor | None = None,
+    noisy_frame_mask: torch.Tensor | None = None,
+    control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None,
+    control_weights: list[float] | tuple[float, ...] | torch.Tensor | None = None,
+    transfer_share_vision_temporal_positions: bool = True,
+    **kwargs: Any,
+) -> CacheContext:
+    """Build the shared cache execution context for Cosmos3's GEN pathway."""
+    if kwargs:
+        raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}")
+
+    prep = module._gen_preprocess(
+        hidden_states,
+        timestep,
+        text_ids,
+        text_mask,
+        video_shape,
+        fps=fps,
+        action_latents=action_latents,
+        action_domain_ids=action_domain_ids,
+        action_noisy_mask=action_noisy_mask,
+        action_start_frame_offset=action_start_frame_offset,
+        action_fps=action_fps,
+        sound_latents=sound_latents,
+        noisy_frame_mask=noisy_frame_mask,
+        control_latents=control_latents,
+        control_weights=control_weights,
+        transfer_share_vision_temporal_positions=transfer_share_vision_temporal_positions,
+    )
+
+    def run_transformer_blocks() -> tuple[torch.Tensor, ...]:
+        return (module._run_gen_stack(prep),)
+
+    def postprocess(hidden: torch.Tensor) -> Any:
+        return module._gen_postprocess(hidden, prep)
+
+    if control_latents is None:
+        controls: list[torch.Tensor] = []
+    elif isinstance(control_latents, torch.Tensor):
+        controls = [control_latents]
+    else:
+        controls = list(control_latents)
+
+    return CacheContext(
+        # SeaCache uses the separate inputs in extra_states for its decision.
+        modulated_input=prep.hidden_gen,
+        hidden_states=prep.hidden_gen,
+        encoder_hidden_states=None,
+        temb=prep.time_embed,
+        run_transformer_blocks=run_transformer_blocks,
+        postprocess=postprocess,
+        extra_states={
+            "sea_cache_latents": [*controls, hidden_states],
+            "sea_cache_noisy_frame_mask": noisy_frame_mask,
+        },
+    )
+
+
 # Registry for model-specific extractors
 # Key: Transformer class name
 # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext
@@ -1475,6 +1548,8 @@ def extract_minimax_h3_context(
 EXTRACTOR_REGISTRY: dict[str, Callable] = {
     "QwenImageTransformer2DModel": extract_qwen_context,
     "Bagel": extract_bagel_context,
+    "Cosmos3EdgeVFMTransformer": extract_cosmos3_context,
+    "Cosmos3VFMTransformer": extract_cosmos3_context,
     "ZImageTransformer2DModel": extract_zimage_context,
     "Flux2Klein": extract_flux2_klein_context,
     "StableAudioDiTModel": extract_stable_audio_context,
@@ -1521,16 +1596,15 @@ def register_extractor(transformer_cls_name: str, extractor_fn: Callable) -> Non
     EXTRACTOR_REGISTRY[transformer_cls_name] = extractor_fn
 
 
-def get_extractor(transformer_cls_name: str) -> Callable:
+def get_extractor(transformer_type: str | type) -> Callable:
     """
     Get extractor function for given transformer class.
 
-    This function looks up the extractor based on the exact transformer_cls_name string,
-    which should match the transformer type in the pipeline (i.e., pipeline.transformer.__class__.__name__).
+    String lookups match an exact registered name. Type lookups also check base
+    classes so runtime wrappers such as FSDP retain the underlying model extractor.
 
     Args:
-        transformer_cls_name: Transformer class name (e.g., "QwenImageTransformer2DModel")
-                                Must exactly match a key in EXTRACTOR_REGISTRY.
+        transformer_type: Transformer class name or runtime type.
 
     Returns:
         Extractor function with signature (module, *args, **kwargs) -> CacheContext
@@ -1543,14 +1617,18 @@ def get_extractor(transformer_cls_name: str) -> Callable:
         >>> extractor = get_extractor("QwenImageTransformer2DModel")
         >>> ctx = extractor(transformer, hidden_states, encoder_hidden_states, timestep, ...)
     """
-    # Direct lookup - no substring matching
-    if transformer_cls_name in EXTRACTOR_REGISTRY:
-        return EXTRACTOR_REGISTRY[transformer_cls_name]
+    candidate_names: tuple[str, ...]
+    if isinstance(transformer_type, str):
+        candidate_names = (transformer_type,)
+    else:
+        candidate_names = tuple(candidate.__name__ for candidate in transformer_type.__mro__)
+    for candidate_name in candidate_names:
+        if candidate_name in EXTRACTOR_REGISTRY:
+            return EXTRACTOR_REGISTRY[candidate_name]
 
-    # No match found
     available_types = list(EXTRACTOR_REGISTRY.keys())
     raise ValueError(
-        f"Unknown model type: '{transformer_cls_name}'. "
+        f"Unknown model type: '{transformer_type}'. "
         f"Available types: {available_types}\n"
         f"To add support for a new model, use register_extractor() or add to EXTRACTOR_REGISTRY."
     )

vllm_omni/diffusion/data.py

@@ -541,6 +541,8 @@ class DiffusionCacheConfig:
                     scm_steps_mask_policy, scm_steps_policy
         - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio,
                     mag_ratios, mag_calibrate
+        - SeaCache: sea_threshold, sea_residual_order,
+                    sea_max_consecutive_cached, sea_power_exp
         - step_cache: step_cache_dit_enabled, velocity_sim_thresholds,
                           velocity_skip_countdowns, step_cache_dit_min_history
 
@@ -560,6 +562,12 @@ class DiffusionCacheConfig:
     rel_l1_thresh: float | None = None
     coefficients: list[float] | None = None  # Uses model-specific defaults if None
 
+    # SeaCache parameters [sea_cache only]
+    sea_threshold: float = 0.25
+    sea_residual_order: int = 1
+    sea_max_consecutive_cached: int = 2
+    sea_power_exp: float = 3.0
+
     # MagCache parameters [mag_cache only]
     # Default: 0.24 threshold for accumulated magnitude error
     mag_threshold: float = 0.24

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Cosmos3 text/image/video/sound/action pipeline for vllm-omni.
 
 One pipeline class serves the Cosmos3 family modes. Output modality is selected
@@ -34,6 +34,7 @@ import math
 import os
 import time
 from collections.abc import Iterable, Mapping
+from contextlib import nullcontext
 from dataclasses import fields
 from typing import Any, ClassVar
 
@@ -1075,6 +1076,8 @@ class Cosmos3OmniDiffusersPipeline(
 
         self._guidance_scale = None
         self._num_timesteps = None
+        self._current_step_index = None
+        self._current_sigma = None
         self._cosmos3_branch_caches: dict[str, tuple[Any, Any]] | None = None
         self._robolab_transforms: dict[bool, Any] = {}
 
@@ -1291,17 +1294,24 @@ class Cosmos3OmniDiffusersPipeline(
         or a tuple in video, action, sound order for multimodal generation.
         """
         cache_key = kwargs.pop("_cosmos3_cache_key", None)
-        if cache_key is None:
-            return self.transformer(**kwargs)
+        context_name = str(kwargs.pop("_cache_context", "cond"))
 
-        branch_caches = self._cosmos3_branch_caches
-        if branch_caches is None:
-            return self.transformer(**kwargs)
-
-        cache_key = str(cache_key)
-        self.transformer.cached_kv, self.transformer.cached_freqs_gen = branch_caches.get(cache_key, (None, None))
-        prediction = self.transformer(**kwargs)
-        branch_caches[cache_key] = (self.transformer.cached_kv, self.transformer.cached_freqs_gen)
+        context_factory = getattr(self, "_cache_context_factory", None)
+        context = context_factory(context_name) if callable(context_factory) else nullcontext()
+        with context:
+            if cache_key is None:
+                prediction = self.transformer(**kwargs)
+            else:
+                branch_caches = self._cosmos3_branch_caches
+                if branch_caches is None:
+                    prediction = self.transformer(**kwargs)
+                else:
+                    cache_key = str(cache_key)
+                    self.transformer.cached_kv, self.transformer.cached_freqs_gen = branch_caches.get(
+                        cache_key, (None, None)
+                    )
+                    prediction = self.transformer(**kwargs)
+                    branch_caches[cache_key] = (self.transformer.cached_kv, self.transformer.cached_freqs_gen)
         return prediction
 
     def combine_multi_branch_cfg_noise(
@@ -1840,6 +1850,29 @@ class Cosmos3OmniDiffusersPipeline(
     def num_timesteps(self):
         return self._num_timesteps
 
+    @property
+    def current_step_index(self):
+        return self._current_step_index
+
+    @property
+    def current_sigma(self):
+        return self._current_sigma
+
+    def _set_denoise_step_metadata(
+        self,
+        step_index: int,
+        timesteps: torch.Tensor,
+        scheduler: Any,
+    ) -> None:
+        self._current_step_index = step_index
+        self._num_timesteps = len(timesteps)
+        sigmas = getattr(scheduler, "sigmas", None)
+        self._current_sigma = sigmas[step_index] if sigmas is not None and step_index < len(sigmas) else None
+
+    def _clear_denoise_step_metadata(self) -> None:
+        self._current_step_index = None
+        self._current_sigma = None
+
     def _set_mixed_precision_step(self, step_index: int, num_steps: int) -> None:
         setter = getattr(self.transformer, "set_mixed_precision_step", None)
         if setter is not None:
@@ -2848,6 +2881,7 @@ class Cosmos3OmniDiffusersPipeline(
                 # else uncond), so session keying loads/stores only this rank's branch.
                 cfg_rank_is_negative = get_classifier_free_guidance_rank() != 0
                 for step_index, t in enumerate(self.progress_bar(timesteps)):
+                    self._set_denoise_step_metadata(step_index, timesteps, step_scheduler)
                     self._set_mixed_precision_step(step_index, len(timesteps))
                     timestep = t.unsqueeze(0)
                     # Outside the interval, scale=1 makes the combined output equal
@@ -2859,6 +2893,7 @@ class Cosmos3OmniDiffusersPipeline(
                         do_true_cfg=True,
                         true_cfg_scale=step_scale,
                         positive_kwargs=dict(
+                            _cache_context="cond",
                             hidden_states=latents,
                             timestep=timestep,
                             text_ids=cond_ids,
@@ -2868,6 +2903,7 @@ class Cosmos3OmniDiffusersPipeline(
                             **shared_kwargs,
                         ),
                         negative_kwargs=dict(
+                            _cache_context="uncond",
                             hidden_states=latents,
                             timestep=timestep,
                             text_ids=uncond_ids,
@@ -2888,13 +2924,15 @@ class Cosmos3OmniDiffusersPipeline(
                 keep_uncond_for_cache = self._cache_requires_paired_cfg()
 
                 for step_index, t in enumerate(self.progress_bar(timesteps)):
+                    self._set_denoise_step_metadata(step_index, timesteps, step_scheduler)
                     self._set_mixed_precision_step(step_index, len(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(
+                    noise_cond = self.predict_noise(
+                        _cache_context="cond",
                         hidden_states=latents,
                         timestep=timestep,
                         text_ids=cond_ids,
@@ -2911,7 +2949,8 @@ class Cosmos3OmniDiffusersPipeline(
                     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(
+                        noise_uncond = self.predict_noise(
+                            _cache_context="uncond",
                             hidden_states=latents,
                             timestep=timestep,
                             text_ids=uncond_ids,
@@ -2943,10 +2982,12 @@ class Cosmos3OmniDiffusersPipeline(
                 # No CFG: a single cond branch per step. Bespoke (state None) keeps
                 # using the transformer-instance cache exactly as before.
                 for step_index, t in enumerate(self.progress_bar(timesteps)):
+                    self._set_denoise_step_metadata(step_index, timesteps, step_scheduler)
                     self._set_mixed_precision_step(step_index, len(timesteps))
                     timestep = t.unsqueeze(0)
                     self._kv_load_und(kv_state, is_negative=False)
-                    noise_pred = self.transformer(
+                    noise_pred = self.predict_noise(
+                        _cache_context="cond",
                         hidden_states=latents,
                         timestep=timestep,
                         text_ids=cond_ids,
@@ -2959,6 +3000,7 @@ class Cosmos3OmniDiffusersPipeline(
                         self._kv_capture_und(kv_state, is_negative=False)
                     _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents))
         finally:
+            self._clear_denoise_step_metadata()
             self._reset_mixed_precision()
             # Cosmos3 currently receives a unique request_id rather than a
             # reusable rollout session id. Retaining its state would only pin
@@ -3145,6 +3187,7 @@ class Cosmos3OmniDiffusersPipeline(
         self._cosmos3_branch_caches = {}
         try:
             for step_index, t in enumerate(self.progress_bar(timesteps)):
+                self._set_denoise_step_metadata(step_index, timesteps, self.scheduler)
                 self._set_mixed_precision_step(step_index, len(timesteps))
                 timestep = t.unsqueeze(0)
                 step_guidance = guidance_scale if _active_at(t, guidance_interval) else 1.0
@@ -3153,6 +3196,7 @@ class Cosmos3OmniDiffusersPipeline(
                 needs_control_cfg = step_control != 1.0
 
                 cond_full_kwargs = dict(
+                    _cache_context="cond",
                     hidden_states=latents,
                     timestep=timestep,
                     text_ids=cond_ids,
@@ -3166,6 +3210,7 @@ class Cosmos3OmniDiffusersPipeline(
                     branches_kwargs = [
                         cond_full_kwargs,
                         dict(
+                            _cache_context="cond_no_control",
                             hidden_states=latents,
                             timestep=timestep,
                             text_ids=cond_ids,
@@ -3175,6 +3220,7 @@ class Cosmos3OmniDiffusersPipeline(
                             **shared_kwargs,
                         ),
                         dict(
+                            _cache_context="uncond",
                             hidden_states=latents,
                             timestep=timestep,
                             text_ids=uncond_ids,
@@ -3200,6 +3246,7 @@ class Cosmos3OmniDiffusersPipeline(
                     branches_kwargs = [
                         cond_full_kwargs,
                         dict(
+                            _cache_context="cond_no_control",
                             hidden_states=latents,
                             timestep=timestep,
                             text_ids=cond_ids,
@@ -3223,6 +3270,7 @@ class Cosmos3OmniDiffusersPipeline(
                     branches_kwargs = [
                         cond_full_kwargs,
                         dict(
+                            _cache_context="uncond",
                             hidden_states=latents,
                             timestep=timestep,
                             text_ids=uncond_ids,
@@ -3257,6 +3305,7 @@ class Cosmos3OmniDiffusersPipeline(
                 )[0]
                 latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents
         finally:
+            self._clear_denoise_step_metadata()
             self._reset_mixed_precision()
             self._cosmos3_branch_caches = None
             self.transformer.reset_cache()

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Cosmos3 VFM Transformer for vllm-omni.
 
 Implements the Mixture-of-Transformers architecture with two pathways:
@@ -14,7 +14,7 @@ from __future__ import annotations
 import math
 from collections.abc import Iterator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, NamedTuple
 
 import torch
 import torch.distributed as dist
@@ -427,7 +427,7 @@ class TimestepEmbedder(nn.Module):
         max_period: int = 10000,
     ) -> None:
         super().__init__()
-        # Following diffusers naming pattern here for checkpoint compatibility.
+        # Preserve checkpoint-compatible parameter names.
         self.linear_1 = nn.Linear(frequency_embedding_size, hidden_size, bias=True)
         self.act = nn.SiLU()
         self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True)
@@ -1105,6 +1105,28 @@ class Cosmos3GenSPPrepare(nn.Module):
         return hidden_gen, freqs_cos, freqs_sin
 
 
+class _GenPrepared(NamedTuple):
+    """GEN-pathway state shared by normal and cached execution."""
+
+    hidden_gen: torch.Tensor
+    time_embed: torch.Tensor
+    t: int
+    h: int
+    w: int
+    s_video: int
+    s_control: int
+    s_action: int
+    s_sound: int
+    has_control: bool
+    has_action: bool
+    has_sound: bool
+    action_domain_ids: torch.Tensor | None
+    ulysses_size: int
+    use_multi_control_attention: bool
+    multi_control_token_sizes: tuple[int, ...] | None
+    multi_control_weights: tuple[float, ...] | None
+
+
 class Cosmos3VFMTransformer(nn.Module):
     """Cosmos3 VFM Transformer: UND language model + GEN denoising layers.
 
@@ -1675,6 +1697,61 @@ class Cosmos3VFMTransformer(nn.Module):
         pad = (-(base + sound_frames)) % ulysses_size
         return sound_frames + pad
 
+    def _run_gen_layers(
+        self,
+        hidden_gen: torch.Tensor,
+        *,
+        s_video: int,
+        s_control: int,
+        s_action: int,
+        s_sound: int,
+        has_action: bool,
+        has_sound: bool,
+        has_control: bool,
+        ulysses_size: int,
+        use_multi_control_attention: bool,
+        multi_control_token_sizes: tuple[int, ...] | None,
+        multi_control_weights: tuple[float, ...] | None,
+    ) -> torch.Tensor:
+        """Run the complete GEN decoder between full-layout boundaries."""
+        if self.cached_kv is None or self.cached_freqs_gen is None:
+            raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.")
+        freqs_cos, freqs_sin = self.cached_freqs_gen
+        if not use_multi_control_attention:
+            hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin)
+        freqs_gen = (freqs_cos, freqs_sin)
+
+        if len(self.gen_layers) == len(self.cached_kv):
+            for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True):
+                hidden_gen = layer(
+                    hidden_gen,
+                    k_und=k_und,
+                    v_und=v_und,
+                    freqs_cos=freqs_cos,
+                    freqs_sin=freqs_sin,
+                    control_token_sizes=multi_control_token_sizes,
+                    control_weights=multi_control_weights,
+                )
+                # Cache-dit's block wrapper may return a tuple; unwrap it.
+                if isinstance(hidden_gen, tuple):
+                    hidden_gen = hidden_gen[0]
+        else:
+            # Cache-dit patches gen_layers to a grouped wrapper.
+            for layer in self.gen_layers:
+                hidden_gen = layer(
+                    hidden_gen,
+                    cached_kv=self.cached_kv,
+                    freqs_gen=freqs_gen,
+                    control_token_sizes=multi_control_token_sizes,
+                    control_weights=multi_control_weights,
+                )
+                if isinstance(hidden_gen, tuple):
+                    hidden_gen = hidden_gen[0]
+
+        if not use_multi_control_attention:
+            hidden_gen = self.gen_sp_gather(hidden_gen)
+        return hidden_gen
+
     # -- Forward -------------------------------------------------------------
 
     def forward(
@@ -1697,7 +1774,51 @@ class Cosmos3VFMTransformer(nn.Module):
         transfer_share_vision_temporal_positions: bool = True,
         **kwargs,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        """Run the shared Cosmos3 GEN preprocess, stack, and postprocess path."""
+        if kwargs:
+            raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}")
+        prep = self._gen_preprocess(
+            hidden_states,
+            timestep,
+            text_ids,
+            text_mask,
+            video_shape,
+            fps=fps,
+            action_latents=action_latents,
+            action_domain_ids=action_domain_ids,
+            action_noisy_mask=action_noisy_mask,
+            action_start_frame_offset=action_start_frame_offset,
+            action_fps=action_fps,
+            sound_latents=sound_latents,
+            noisy_frame_mask=noisy_frame_mask,
+            control_latents=control_latents,
+            control_weights=control_weights,
+            transfer_share_vision_temporal_positions=transfer_share_vision_temporal_positions,
+        )
+        return self._gen_postprocess(self._run_gen_stack(prep), prep)
+
+    def _gen_preprocess(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        text_ids: torch.Tensor,
+        text_mask: torch.Tensor,
+        video_shape: tuple[int, int, int],
+        fps: float | None = None,
+        action_latents: torch.Tensor | None = None,
+        action_domain_ids: torch.Tensor | None = None,
+        action_noisy_mask: torch.Tensor | None = None,
+        action_start_frame_offset: int = 1,
+        action_fps: float | None = None,
+        sound_latents: torch.Tensor | None = None,
+        noisy_frame_mask: torch.Tensor | None = None,
+        control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None,
+        control_weights: list[float] | tuple[float, ...] | torch.Tensor | None = None,
+        transfer_share_vision_temporal_positions: bool = True,
+    ) -> _GenPrepared:
         """
+        Prepare the packed GEN sequence before its cacheable execution region.
+
         Args:
             hidden_states: [B, C, t, h, w] noisy latents
             timestep: [B] diffusion timestep
@@ -1720,13 +1841,9 @@ class Cosmos3VFMTransformer(nn.Module):
                 transfer controls. Values are normalized to sum to one.
 
         Returns:
-            [B, C, t, h, w] velocity prediction, or
-            tuple outputs in video, action, sound order when action/sound streams
-            are provided. Transfer-control streams condition the video prediction
-            and are not returned.
+            Packed GEN inputs and metadata consumed by the shared execution and
+            postprocessing helpers.
         """
-        if kwargs:
-            raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}")
         t, h, w = video_shape
         hp, wp, _, _ = self._pad_to_patch_size(h, w)
         text_lengths = text_mask.sum(dim=1)
@@ -1926,78 +2043,81 @@ class Cosmos3VFMTransformer(nn.Module):
                 hidden_parts.append(hidden_sound)
             hidden_gen = torch.cat(hidden_parts, dim=1)
 
-            # Run GEN layers.  UND K/V (replicated) is passed to each layer;
-            # the Cosmos3CrossAttention forwards them as joint_key/value so the
-            # framework Attention handles the Ulysses head-slicing internally.
-            if self.cached_kv is None or self.cached_freqs_gen is None:
-                raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.")
-            freqs_cos, freqs_sin = self.cached_freqs_gen
-            if not use_multi_control_attention:
-                hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin)
-            freqs_gen = (freqs_cos, freqs_sin)
+            return _GenPrepared(
+                hidden_gen=hidden_gen,
+                time_embed=time_embed,
+                t=t,
+                h=h,
+                w=w,
+                s_video=s_video,
+                s_control=s_control,
+                s_action=s_action,
+                s_sound=s_sound,
+                has_control=has_control,
+                has_action=has_action,
+                has_sound=has_sound,
+                action_domain_ids=action_domain_ids,
+                ulysses_size=ulysses_size,
+                use_multi_control_attention=use_multi_control_attention,
+                multi_control_token_sizes=multi_control_token_sizes,
+                multi_control_weights=multi_control_weights,
+            )
 
-            if len(self.gen_layers) == len(self.cached_kv):
-                for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True):
-                    hidden_gen = layer(
-                        hidden_gen,
-                        k_und=k_und,
-                        v_und=v_und,
-                        freqs_cos=freqs_cos,
-                        freqs_sin=freqs_sin,
-                        control_token_sizes=multi_control_token_sizes,
-                        control_weights=multi_control_weights,
-                    )
-                    # Cache-dit's block wrapper may return a tuple; unwrap it.
-                    if isinstance(hidden_gen, tuple):
-                        hidden_gen = hidden_gen[0]
-            else:
-                # Cache-dit patches gen_layers to a grouped wrapper.
-                for layer in self.gen_layers:
-                    hidden_gen = layer(
-                        hidden_gen,
-                        cached_kv=self.cached_kv,
-                        freqs_gen=freqs_gen,
-                        control_token_sizes=multi_control_token_sizes,
-                        control_weights=multi_control_weights,
-                    )
-                    if isinstance(hidden_gen, tuple):
-                        hidden_gen = hidden_gen[0]
+    def _run_gen_stack(self, prep: _GenPrepared) -> torch.Tensor:
+        """Execute the cacheable full-layout GEN stack, including final norm."""
+        hidden_gen = self._run_gen_layers(
+            prep.hidden_gen,
+            s_video=prep.s_video,
+            s_control=prep.s_control,
+            s_action=prep.s_action,
+            s_sound=prep.s_sound,
+            has_action=prep.has_action,
+            has_sound=prep.has_sound,
+            has_control=prep.has_control,
+            ulysses_size=prep.ulysses_size,
+            use_multi_control_attention=prep.use_multi_control_attention,
+            multi_control_token_sizes=prep.multi_control_token_sizes,
+            multi_control_weights=prep.multi_control_weights,
+        )
+        return self.norm_moe_gen(hidden_gen)
 
-            if not use_multi_control_attention:
-                hidden_gen = self.gen_sp_gather(hidden_gen)
+    def _gen_postprocess(
+        self,
+        hidden_gen: torch.Tensor,
+        prep: _GenPrepared,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        """Project an already-normalized packed GEN state to model outputs."""
+        if not prep.has_action and not prep.has_sound and not prep.has_control:
+            return self.unpatchify(self.proj_out(hidden_gen), prep.t, prep.h, prep.w)
 
-            # Final norm and project back to latent space
-            hidden_gen = self.norm_moe_gen(hidden_gen)
-            if not has_action and not has_sound and not has_control:
-                return self.unpatchify(self.proj_out(hidden_gen), t, h, w)
-
-            split_sizes = []
-            if has_control:
-                split_sizes.append(s_control)
-            split_sizes.append(s_video)
-            if has_action:
-                split_sizes.append(s_action)
-            if has_sound:
-                split_sizes.append(s_sound)
-            split_hidden = hidden_gen.split(split_sizes, dim=1)
-            split_idx = 0
-            if has_control:
-                split_idx += 1
-            hidden_video = split_hidden[split_idx]
+        split_sizes = []
+        if prep.has_control:
+            split_sizes.append(prep.s_control)
+        split_sizes.append(prep.s_video)
+        if prep.has_action:
+            split_sizes.append(prep.s_action)
+        if prep.has_sound:
+            split_sizes.append(prep.s_sound)
+        split_hidden = hidden_gen.split(split_sizes, dim=1)
+        split_idx = 0
+        if prep.has_control:
             split_idx += 1
-            video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w)
-            if has_control:
-                return video_pred
-            outputs: list[torch.Tensor] = [video_pred]
-            if has_action:
-                hidden_action = split_hidden[split_idx]
-                split_idx += 1
-                assert action_domain_ids is not None
-                outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids)))
-            if has_sound:
-                hidden_sound = split_hidden[split_idx]
-                outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound)))
-            return tuple(outputs)
+        hidden_video = split_hidden[split_idx]
+        split_idx += 1
+        video_pred = self.unpatchify(self.proj_out(hidden_video), prep.t, prep.h, prep.w)
+        if prep.has_control:
+            return video_pred
+
+        outputs: list[torch.Tensor] = [video_pred]
+        if prep.has_action:
+            hidden_action = split_hidden[split_idx]
+            split_idx += 1
+            assert prep.action_domain_ids is not None
+            outputs.append(self.unpack_action(self.action_proj_out(hidden_action, prep.action_domain_ids)))
+        if prep.has_sound:
+            hidden_sound = split_hidden[split_idx]
+            outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound)))
+        return tuple(outputs)
 
     def post_load_weights(self) -> None:
         """Post-load processing: ensure correct dtypes."""

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -25,7 +25,7 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cachedit import cache_summary
+from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend, cache_summary
 from vllm_omni.diffusion.cache.prompt_embed_cache import (
     install_prompt_embed_cache,
     resolve_prompt_embed_cache_config,
@@ -409,7 +409,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 # may explicitly adopt the enabled Cache-DiT backend and then
                 # own all later request-boundary enable/disable transitions.
                 self.cache_backend.enable(self.pipeline)
-                if str(self.od_config.cache_backend).lower() == "cache_dit" and adopt_request_scoped_cache_dit(
+                if isinstance(self.cache_backend, CacheDiTBackend) and adopt_request_scoped_cache_dit(
                     self.pipeline,
                     self.cache_backend,
                 ):
@@ -663,12 +663,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             num_inference_steps = getattr(self.pipeline, "default_num_inference_steps", None)
         if num_inference_steps is None and od_config.cache_backend in (
             "tea_cache",
+            "sea_cache",
             "step_cache",
         ):
             # When num_inference_steps is None, some pipelines defer to their
-            # own defaults. TeaCache refresh ignores this value; step_cache
-            # refresh is a no-op because per-chunk state resets in the denoise
-            # loop. Use the pipeline default when available to keep refresh
+            # own defaults. These backends use refresh to reset request state;
+            # runtime step metadata is either unused or resolved in the
+            # pipeline. Use the pipeline default when available to keep refresh
             # behavior aligned with single-request execution.
             num_inference_steps = getattr(self.pipeline, "num_inference_steps", 0) or 0
 
@@ -1046,6 +1047,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         new_request_ids: list[str],
     ) -> tuple[list[StepRequestState], InputBatch | None, list[RunnerOutput]]:
         # process new reqs
+        pipeline = self.pipeline
+        assert pipeline is not None, "Model not loaded. Call load_model() first."
         prepared_states: list[StepRequestState] = []
         error_outputs: list[RunnerOutput] = []
         for state in states:
@@ -1058,12 +1061,12 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 per_req_exc: BaseException | None = None
                 try:
                     self._initialize_generator(state.sampling)
-                    clear_pipeline_stage_durations(self.pipeline)
+                    clear_pipeline_stage_durations(pipeline)
                     # encode
-                    self.pipeline.prepare_encode(state)
+                    pipeline.prepare_encode(state)
                     merge_stage_durations(
                         state,
-                        consume_pipeline_stage_durations(self.pipeline),
+                        consume_pipeline_stage_durations(pipeline),
                     )
                 except Exception as exc:
                     per_req_exc = exc
@@ -1247,6 +1250,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         use_hsdp = self.od_config.parallel_config.use_hsdp
         grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
         with grad_context:
+            pipeline = self.pipeline
+            assert pipeline is not None, "Model not loaded. Call load_model() first."
             had_active_states = bool(self.state_cache)
             states, new_request_ids = self._update_states(scheduler_output)
             is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
@@ -1261,7 +1266,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             states, input_batch, runner_output_list = self._prepare_batch_inputs(states, new_request_ids)
             if input_batch is None:
                 return BatchRunnerOutput.from_list(runner_output_list)
-            attn_metadata = {}
+            attn_metadata: dict[str, Any] = {}
 
             kv_backend = getattr(self, "diffusion_kv_backend", None)
             paged_kv_runtime = kv_backend if getattr(kv_backend, "paged_attention_adapter", None) is not None else None
@@ -1272,16 +1277,16 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 paged_kv_runtime=paged_kv_runtime,
                 in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
             ):
-                clear_pipeline_stage_durations(self.pipeline)
-                noise_pred = self.pipeline.denoise_step(input_batch, states=states)
-                denoise_stage_durations = consume_pipeline_stage_durations(self.pipeline)
+                clear_pipeline_stage_durations(pipeline)
+                noise_pred = pipeline.denoise_step(input_batch, states=states)
+                denoise_stage_durations = consume_pipeline_stage_durations(pipeline)
                 for state in states:
                     merge_stage_durations(
                         state,
                         denoise_stage_durations,
                     )
 
-                pipeline_interrupted = getattr(self.pipeline, "interrupt", False)
+                pipeline_interrupted = getattr(pipeline, "interrupt", False)
                 if noise_pred is None and pipeline_interrupted:
                     for state in states:
                         runner_output_list.append(
@@ -1296,9 +1301,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 else:
                     offset = 0
                     for req in states:
+                        if req.latents is None:
+                            raise RuntimeError(f"Stepwise request {req.request_id} has no latent state.")
                         row_num = req.latents.shape[0]
                         try:
-                            self.pipeline.step_scheduler(
+                            pipeline.step_scheduler(
                                 req, noise_pred[offset : offset + row_num] if noise_pred is not None else None
                             )
                             if self.od_config.streaming_output:
@@ -1307,8 +1314,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                                 should_decode = req.denoise_completed
 
                             if should_decode:
-                                clear_pipeline_stage_durations(self.pipeline)
-                                result = self.pipeline.post_decode(req)
+                                clear_pipeline_stage_durations(pipeline)
+                                result = pipeline.post_decode(req)
                                 if result is not None:
                                     result = self._prepare_output_for_transport(result, req.sampling)
                                     self._attach_stepwise_metadata(
@@ -1317,8 +1324,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                                     )
                                     # After consuming this chunk's interaction metadata, apply pending interactions and
                                     # prepare the next chunk (prepare_next_chunk may be a no-op---depending on pipeline)
-                                    if supports_interaction_apply(self.pipeline) and not req.request_denoise_completed:
-                                        pipe = cast(SupportsInteractionApply, self.pipeline)
+                                    if supports_interaction_apply(pipeline) and not req.request_denoise_completed:
+                                        pipe = cast(SupportsInteractionApply, pipeline)
                                         pipe.apply_interaction_at_chunk_boundary(req)
                                         pipe.prepare_next_chunk(req)
                             else:
@@ -1370,12 +1377,12 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     for runner_output in runner_output_list:
                         if runner_output.result is None:
                             continue
-                        state = states_by_id.get(runner_output.request_id)
-                        if state is None:
+                        matched_state = states_by_id.get(runner_output.request_id)
+                        if matched_state is None:
                             continue
                         runner_output.result.peak_memory_mb = max(
                             runner_output.result.peak_memory_mb,
-                            state.peak_memory_mb,
+                            matched_state.peak_memory_mb,
                         )
 
                 terminal_request_ids = [

vllm_omni/entrypoints/cli/serve.py

@@ -712,7 +712,10 @@ class OmniServeCommand(CLISubcommand):
             "--cache-backend",
             type=str,
             default="none",
-            help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache', 'step_cache'",
+            help=(
+                "Cache backend for diffusion models, options: 'tea_cache', "
+                "'cache_dit', 'mag_cache', 'sea_cache', 'step_cache'"
+            ),
         )
         omni_config_group.add_argument(
             "--cache-config",