From 981fe043f6e2db1d0390fe1aff717391f7df5255 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Mon, 31 Aug 2026 03:17:48 -0700
Subject: [PATCH 01/14] Add optional FP32 sampling state for Cosmos3

---
 .../cosmos/before_denoise.py                  | 35 ++++++++--
 .../modular_pipelines/cosmos/denoise.py       | 16 ++++-
 .../modular_blocks_cosmos3_distilled.py       | 14 +++-
 .../pipelines/cosmos/pipeline_cosmos3_omni.py | 19 +++++-
 ...test_modular_pipeline_cosmos3_distilled.py | 67 ++++++++++++++++++-
 tests/pipelines/cosmos/test_cosmos3.py        | 31 +++++++++
 6 files changed, 170 insertions(+), 12 deletions(-)

diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py
index 7bf431aa855b..298312e450c9 100644
--- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py
+++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py
@@ -68,6 +68,10 @@ def description(self) -> str:
     def expected_components(self) -> list[ComponentSpec]:
         return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
 
+    @property
+    def expected_configs(self) -> list[ConfigSpec]:
+        return [ConfigSpec(name="default_use_fp32_sampling_state", default=False)]
+
     @property
     def inputs(self) -> list[InputParam]:
         return [
@@ -98,6 +102,15 @@ def inputs(self) -> list[InputParam]:
                 description="Pre-generated noisy vision latents.",
             ),
             InputParam.template("generator"),
+            InputParam(
+                name="use_fp32_sampling_state",
+                type_hint=bool | None,
+                default=None,
+                description=(
+                    "Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the "
+                    "pipeline's `default_use_fp32_sampling_state` config."
+                ),
+            ),
         ]
 
     @property
@@ -121,13 +134,20 @@ def intermediate_outputs(self) -> list[OutputParam]:
                 type_hint=torch.Tensor,
                 description="Clean encoded vision latents used to re-anchor image conditioning each step.",
             ),
+            OutputParam(
+                "use_fp32_sampling_state",
+                type_hint=bool,
+                description="Whether vision sampling state is kept in float32.",
+            ),
         ]
 
     @torch.no_grad()
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         device = components._execution_device
-        dtype = components.transformer.dtype
+        if block_state.use_fp32_sampling_state is None:
+            block_state.use_fp32_sampling_state = components.config.default_use_fp32_sampling_state
+        sampling_dtype = torch.float32 if block_state.use_fp32_sampling_state else components.transformer.dtype
 
         x0_tokens_vision = block_state.x0_tokens_vision
         if x0_tokens_vision is None:
@@ -151,21 +171,26 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
 
         block_state.fps_vision = float(block_state.fps)
         condition_frames = block_state.vision_condition_frames or []
-        block_state.vision_condition_mask = torch.zeros((x0_tokens_vision.shape[2], 1, 1), device=device, dtype=dtype)
+        block_state.vision_condition_mask = torch.zeros(
+            (x0_tokens_vision.shape[2], 1, 1), device=device, dtype=sampling_dtype
+        )
         for frame_idx in condition_frames:
             if 0 <= frame_idx < block_state.vision_condition_mask.shape[0]:
                 block_state.vision_condition_mask[frame_idx, 0, 0] = 1.0
 
         if block_state.latents is None:
             pure_noise = randn_tensor(
-                tuple(x0_tokens_vision.shape), generator=block_state.generator, device=device, dtype=dtype
+                tuple(x0_tokens_vision.shape),
+                generator=block_state.generator,
+                device=device,
+                dtype=sampling_dtype,
             )
             block_state.latents = (
-                block_state.vision_condition_mask * x0_tokens_vision.to(device=device, dtype=dtype)
+                block_state.vision_condition_mask * x0_tokens_vision.to(device=device, dtype=sampling_dtype)
                 + (1.0 - block_state.vision_condition_mask) * pure_noise
             )
         else:
-            block_state.latents = block_state.latents.to(device=device, dtype=dtype)
+            block_state.latents = block_state.latents.to(device=device, dtype=sampling_dtype)
 
         vision_condition_indexes = torch.nonzero(
             block_state.vision_condition_mask[:, 0, 0] > 0, as_tuple=False
diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py
index eda37c8e99cf..ee28ea935883 100644
--- a/src/diffusers/modular_pipelines/cosmos/denoise.py
+++ b/src/diffusers/modular_pipelines/cosmos/denoise.py
@@ -319,6 +319,12 @@ def inputs(self) -> list[InputParam]:
                 description="Indexes of conditioned vision latent frames; non-empty for image-to-video.",
             ),
             InputParam.template("generator"),
+            InputParam(
+                name="use_fp32_sampling_state",
+                type_hint=bool,
+                default=False,
+                description="Whether to keep the distilled vision scheduler state in float32.",
+            ),
         ]
 
     @property
@@ -327,11 +333,17 @@ def intermediate_outputs(self) -> list[OutputParam]:
 
     @torch.no_grad()
     def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+        velocity_vision = block_state.velocity_vision
+        latents = block_state.latents
+        if block_state.use_fp32_sampling_state:
+            velocity_vision = velocity_vision.float()
+            latents = latents.float()
+
         # Pass the generator so the scheduler's stochastic (SDE) re-noising is seedable/reproducible.
         block_state.latents = components.scheduler.step(
-            block_state.velocity_vision.unsqueeze(0),
+            velocity_vision.unsqueeze(0),
             t,
-            block_state.latents.unsqueeze(0),
+            latents.unsqueeze(0),
             generator=block_state.generator,
             return_dict=False,
         )[0].squeeze(0)
diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
index e168cc24d8cd..a0de9459e5ec 100644
--- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
+++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
@@ -86,7 +86,8 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks):
           transformer (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`)
 
       Configs:
-          is_distilled (default: True) distilled_sigmas (default: None)
+          default_use_fp32_sampling_state (default: False) is_distilled (default: True) distilled_sigmas (default:
+          None)
 
       Inputs:
           cond_input_ids (`None`):
@@ -109,6 +110,9 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
+          use_fp32_sampling_state (`bool | NoneType`, *optional*):
+              Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the pipeline's
+              `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`, *optional*):
               The number of denoising steps.
           guidance_scale (`float`, *optional*):
@@ -165,8 +169,9 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks):
           (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`)
 
       Configs:
-          default_use_system_prompt (default: True) enable_safety_checker (default: True) is_distilled (default: True)
-          distilled_sigmas (default: None)
+          default_use_system_prompt (default: True) enable_safety_checker (default: True)
+          default_use_fp32_sampling_state (default: False) is_distilled (default: True) distilled_sigmas (default:
+          None)
 
       Inputs:
           prompt (`str`):
@@ -201,6 +206,9 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
+          use_fp32_sampling_state (`bool | NoneType`, *optional*):
+              Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the pipeline's
+              `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`, *optional*):
               The number of denoising steps.
           guidance_scale (`float`, *optional*):
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
index 589e0ed3d6b0..6c025531fa22 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
@@ -1340,6 +1340,7 @@ def __call__(
         add_resolution_template: bool = True,
         add_duration_template: bool = True,
         enable_safety_check: bool = True,
+        use_fp32_sampling_state: bool = False,
     ) -> Cosmos3OmniPipelineOutput:
         r"""
         Run the Cosmos 3 omni pipeline end-to-end: encode the (optional) conditioning image/video, denoise vision and
@@ -1435,6 +1436,11 @@ def __call__(
                 When `True` and a `CosmosSafetyChecker` is attached, runs the text guardrail on the prompt before
                 generation and the video guardrail on the decoded frames. Set to `False` to skip both for this call;
                 the checker remains loaded for subsequent calls.
+            use_fp32_sampling_state (`bool`, *optional*, defaults to `False`):
+                When `True`, keeps vision, sound, and action denoising latents plus classifier-free-guidance arithmetic
+                in `torch.float32`. Transformer inputs are still cast to the transformer's dtype before each forward.
+                This improves sampling-state precision at the cost of additional memory and preserves the existing
+                model-dtype behavior when disabled.
 
         Returns:
             [`Cosmos3OmniPipelineOutput`] or `tuple`:
@@ -1497,6 +1503,7 @@ def __call__(
 
         device = self._get_execution_device()
         dtype = self.transformer.dtype
+        sampling_dtype = torch.float32 if use_fp32_sampling_state else dtype
 
         if enable_safety_check and isinstance(self.safety_checker, CosmosSafetyChecker):
             self.safety_checker.to(device)
@@ -1557,7 +1564,7 @@ def __call__(
             action_latents=action_latents,
             generator=generator,
             device=device,
-            dtype=dtype,
+            dtype=sampling_dtype,
             enable_sound=enable_sound,
             action=action,
         )
@@ -1786,6 +1793,14 @@ def __call__(
                         raw_action_dim=raw_action_dim_resolved,
                     )
 
+                if use_fp32_sampling_state:
+                    cond_v_vision = cond_v_vision.float()
+                    cond_v_sound = cond_v_sound.float() if cond_v_sound is not None else None
+                    cond_v_action = cond_v_action.float() if cond_v_action is not None else None
+                    uncond_v_vision = uncond_v_vision.float() if uncond_v_vision is not None else None
+                    uncond_v_sound = uncond_v_sound.float() if uncond_v_sound is not None else None
+                    uncond_v_action = uncond_v_action.float() if uncond_v_action is not None else None
+
                 # --- CFG combine + per-modality scheduler step ---
                 # UniPC's multistep_uni_p_bh_update einsum ("k,bkc...->bc...") requires sample
                 # to carry a batch dim; per-modality latents have no batch axis, so wrap for the step.
@@ -1830,6 +1845,8 @@ def __call__(
                         callback_kwargs[key] = locals()[key]
                     callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
                     latents = callback_outputs.pop("latents", latents)
+                    if use_fp32_sampling_state:
+                        latents = latents.float()
 
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                     progress_bar.update()
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
index 97bb2b478989..11e6900d12a1 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
@@ -17,8 +17,12 @@
 import torch
 from PIL import Image
 
-from diffusers import ModularPipeline
+from diffusers import FlowMatchEulerDiscreteScheduler, ModularPipeline
 from diffusers.modular_pipelines import Cosmos3DistilledBlocks, Cosmos3DistilledModularPipeline
+from diffusers.modular_pipelines.cosmos.before_denoise import Cosmos3VisionPrepareLatentsStep
+from diffusers.modular_pipelines.cosmos.denoise import Cosmos3DistilledVisionLoopSchedulerStep
+from diffusers.modular_pipelines.cosmos.encoders import Cosmos3DistilledTextEncoderStep
+from diffusers.modular_pipelines.modular_pipeline import BlockState, PipelineState
 
 from ...testing_utils import torch_device
 from ..testing_utils import (
@@ -113,6 +117,67 @@ def test_declares_distilled_configs(self):
         pipe = self.pipeline_class()
         assert pipe.config.is_distilled is True
         assert pipe.config.distilled_sigmas is None
+        assert pipe.config.default_use_system_prompt is True
+        assert pipe.config.default_use_fp32_sampling_state is False
+    @pytest.mark.parametrize(
+        ("config_enabled", "input_enabled", "expected_dtype"),
+        [
+            (False, None, torch.bfloat16),
+            (True, None, torch.float32),
+            (True, False, torch.bfloat16),
+        ],
+    )
+    def test_prepare_vision_latents_fp32_sampling_state(
+        self, config_enabled, input_enabled, expected_dtype
+    ):
+        components = BlockState(
+            _execution_device=torch.device("cpu"),
+            transformer=BlockState(dtype=torch.bfloat16),
+            config=BlockState(default_use_fp32_sampling_state=config_enabled),
+            vae_scale_factor_spatial=16,
+            vae_scale_factor_temporal=4,
+            num_channels_latents=4,
+        )
+        state = PipelineState()
+        state.set("x0_tokens_vision", None)
+        state.set("vision_condition_frames", None)
+        state.set("num_frames", 5)
+        state.set("height", 32)
+        state.set("width", 32)
+        state.set("fps", 24.0)
+        state.set("latents", None)
+        state.set("generator", torch.Generator("cpu").manual_seed(0))
+        state.set("use_fp32_sampling_state", input_enabled)
+
+        Cosmos3VisionPrepareLatentsStep()(components, state)
+
+        assert state.get("latents").dtype == expected_dtype
+        assert state.get("vision_condition_mask").dtype == expected_dtype
+        assert state.get("use_fp32_sampling_state") is (config_enabled if input_enabled is None else input_enabled)
+
+    @pytest.mark.parametrize(
+        ("use_fp32_sampling_state", "expected_dtype"),
+        [(False, torch.bfloat16), (True, torch.float32)],
+    )
+    def test_distilled_scheduler_fp32_state(self, use_fp32_sampling_state, expected_dtype):
+        scheduler = FlowMatchEulerDiscreteScheduler(stochastic_sampling=True)
+        scheduler.set_timesteps(sigmas=[1.0, 0.5])
+        latents = torch.zeros((1, 2, 1, 1, 1), dtype=torch.bfloat16)
+        block_state = BlockState(
+            latents=latents,
+            velocity_vision=torch.zeros_like(latents),
+            vision_condition_mask=torch.zeros((1, 1, 1), dtype=expected_dtype),
+            vision_conditioning_latents=None,
+            vision_condition_indexes_for_pack=[],
+            generator=torch.Generator("cpu").manual_seed(0),
+            use_fp32_sampling_state=use_fp32_sampling_state,
+        )
+
+        Cosmos3DistilledVisionLoopSchedulerStep()(
+            BlockState(scheduler=scheduler), block_state, i=0, t=scheduler.timesteps[0]
+        )
+
+        assert block_state.latents.dtype == expected_dtype
 
     def test_vae_encoder_rejects_image_and_video_together(self):
         vae_encoder = Cosmos3DistilledBlocks().sub_blocks["vae_encoder"]
diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py
index dfc1dd0a1ea1..cbe30dd41e5d 100644
--- a/tests/pipelines/cosmos/test_cosmos3.py
+++ b/tests/pipelines/cosmos/test_cosmos3.py
@@ -118,6 +118,37 @@ def test_inference(self):
 
         self.assertEqual(video.shape, (1, 16, 16, 3))
 
+    def test_fp32_sampling_state_keeps_transformer_inputs_in_model_dtype(self):
+        pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
+        pipeline.transformer.to(dtype=torch.float16)
+        pipeline.set_progress_bar_config(disable=None)
+        transformer_input_dtypes = []
+        callback_latent_dtypes = []
+
+        def transformer_forward(**kwargs):
+            vision_tokens = kwargs["vision_tokens"][0]
+            transformer_input_dtypes.append(vision_tokens.dtype)
+            return ([torch.zeros_like(vision_tokens)], None, None)
+
+        def callback_on_step_end(_pipeline, _step_index, _timestep, callback_kwargs):
+            callback_latent_dtypes.append(callback_kwargs["latents"].dtype)
+            return {"latents": callback_kwargs["latents"].to(torch.float16)}
+
+        inputs = self.get_dummy_inputs(torch_device)
+        inputs.update(
+            output_type="latent",
+            enable_safety_check=False,
+            use_fp32_sampling_state=True,
+            callback_on_step_end=callback_on_step_end,
+        )
+        with mock.patch.object(pipeline.transformer, "forward", side_effect=transformer_forward):
+            latents = pipeline(**inputs).video
+
+        assert transformer_input_dtypes
+        assert all(dtype == torch.float16 for dtype in transformer_input_dtypes)
+        assert callback_latent_dtypes
+        assert all(dtype == torch.float32 for dtype in callback_latent_dtypes)
+        assert latents.dtype == torch.float32
     def test_cosmos3_tokenize_prompt_uses_checkpoint_system_prompt_default(self):
         components = self.get_dummy_components()
         components["default_use_system_prompt"] = False

From 1dac4d42fde0a12621cf15fe05945e32e5deb89f Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Mon, 31 Aug 2026 03:43:50 -0700
Subject: [PATCH 02/14] Add cache execution contexts to Cosmos3 pipelines

---
 src/diffusers/models/cache_utils.py           |  10 +-
 .../transformers/transformer_cosmos3.py       |   3 +-
 .../modular_pipelines/cosmos/denoise.py       |  60 +++++---
 .../cosmos/modular_pipeline.py                |  24 ++++
 .../pipelines/cosmos/pipeline_cosmos3_omni.py | 133 +++++++++++-------
 .../cosmos/test_modular_pipeline_cosmos3.py   |  23 +++
 tests/pipelines/cosmos/test_cosmos3.py        |  25 ++++
 7 files changed, 203 insertions(+), 75 deletions(-)

diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py
index 5aa189987ba2..5bf161b2cc08 100644
--- a/src/diffusers/models/cache_utils.py
+++ b/src/diffusers/models/cache_utils.py
@@ -69,6 +69,7 @@ def enable_cache(self, config) -> None:
         from ..hooks import (
             FasterCacheConfig,
             FirstBlockCacheConfig,
+            HookRegistry,
             MagCacheConfig,
             PyramidAttentionBroadcastConfig,
             TaylorSeerCacheConfig,
@@ -102,6 +103,7 @@ def enable_cache(self, config) -> None:
             raise ValueError(f"Cache config {type(config)} is not supported.")
 
         self._cache_config = config
+        HookRegistry.check_if_exists_or_initialize(self)._child_registries_cache = None
 
     def disable_cache(self) -> None:
         from ..hooks import (
@@ -145,6 +147,7 @@ def disable_cache(self) -> None:
             raise ValueError(f"Cache config {type(self._cache_config)} is not supported.")
 
         self._cache_config = None
+        registry._child_registries_cache = None
 
     def _reset_stateful_cache(self, recurse: bool = True) -> None:
         from ..hooks import HookRegistry
@@ -159,6 +162,7 @@ def cache_context(self, name: str):
         registry = HookRegistry.check_if_exists_or_initialize(self)
         registry._set_context(name)
 
-        yield
-
-        registry._set_context(None)
+        try:
+            yield
+        finally:
+            registry._set_context(None)
diff --git a/src/diffusers/models/transformers/transformer_cosmos3.py b/src/diffusers/models/transformers/transformer_cosmos3.py
index f7cfc317bc79..07c6bd12d0dc 100644
--- a/src/diffusers/models/transformers/transformer_cosmos3.py
+++ b/src/diffusers/models/transformers/transformer_cosmos3.py
@@ -23,6 +23,7 @@
 from ...utils import BaseOutput
 from ..attention import AttentionMixin, AttentionModuleMixin
 from ..attention_dispatch import dispatch_attention_fn
+from ..cache_utils import CacheMixin
 from ..embeddings import TimestepEmbedding, Timesteps
 from ..modeling_utils import ModelMixin
 from ..normalization import RMSNorm
@@ -370,7 +371,7 @@ def forward(
         return residual_und + mlp_out_und, residual_gen + mlp_out_gen
 
 
-class Cosmos3OmniTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin):
+class Cosmos3OmniTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin, CacheMixin):
     _supports_gradient_checkpointing = True
     _no_split_modules = ["Cosmos3VLTextMoTDecoderLayer"]
     _repeated_blocks = ["Cosmos3VLTextMoTDecoderLayer"]
diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py
index ee28ea935883..ce6ba0a4bee5 100644
--- a/src/diffusers/modular_pipelines/cosmos/denoise.py
+++ b/src/diffusers/modular_pipelines/cosmos/denoise.py
@@ -1,4 +1,5 @@
 import inspect
+from contextlib import nullcontext
 
 import torch
 
@@ -14,6 +15,11 @@
 from .modular_pipeline import Cosmos3OmniModularPipeline
 
 
+def _cache_context(transformer: torch.nn.Module, name: str):
+    cache_context = getattr(transformer, "cache_context", None)
+    return cache_context(name) if callable(cache_context) else nullcontext()
+
+
 class Cosmos3VisionLoopPrepareStep(ModularPipelineBlocks):
     model_name = "cosmos3-omni"
 
@@ -215,7 +221,10 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
             transformer_kwargs = {
                 name: value for name, value in transformer_kwargs.items() if name in transformer_args
             }
-            preds_vision, preds_sound, preds_action = components.transformer(**transformer_kwargs, return_dict=False)
+            with _cache_context(components.transformer, pass_name):
+                preds_vision, preds_sound, preds_action = components.transformer(
+                    **transformer_kwargs, return_dict=False
+                )
             velocities[pass_name] = components._mask_velocity_predictions(
                 preds_vision,
                 preds_sound,
@@ -480,13 +489,23 @@ def loop_inputs(self) -> list[InputParam]:
     @torch.no_grad()
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
+        components._current_step_index = None
+        components._current_sigma = None
+        components._num_timesteps = len(block_state.timesteps)
         with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
             for i, t in enumerate(block_state.timesteps):
+                components._current_step_index = i
+                scheduler_sigmas = getattr(components.scheduler, "sigmas", None)
+                components._current_sigma = (
+                    scheduler_sigmas[i] if scheduler_sigmas is not None and i < len(scheduler_sigmas) else None
+                )
                 components, block_state = self.loop_step(components, block_state, i=i, t=t)
                 if i == len(block_state.timesteps) - 1 or (
                     (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
                 ):
                     progress_bar.update()
+        components._current_step_index = None
+        components._current_sigma = None
         self.set_block_state(state, block_state)
         return components, state
 
@@ -718,21 +737,22 @@ def intermediate_outputs(self) -> list[OutputParam]:
         return [OutputParam("velocity", type_hint=torch.Tensor, description="Predicted (masked) transfer velocity.")]
 
     @staticmethod
-    def _forward(components, static, vision_tokens, vision_timesteps):
-        preds_vision, _, _ = components.transformer(
-            input_ids=static["input_ids"],
-            text_indexes=static["text_indexes"],
-            position_ids=static["position_ids"],
-            und_len=static["und_len"],
-            sequence_length=static["sequence_length"],
-            vision_tokens=vision_tokens,
-            vision_token_shapes=static["vision_token_shapes"],
-            vision_sequence_indexes=static["vision_sequence_indexes"],
-            vision_mse_loss_indexes=static["vision_mse_loss_indexes"],
-            vision_timesteps=vision_timesteps,
-            vision_noisy_frame_indexes=static["vision_noisy_frame_indexes"],
-            return_dict=False,
-        )
+    def _forward(components, static, vision_tokens, vision_timesteps, context_name):
+        with _cache_context(components.transformer, context_name):
+            preds_vision, _, _ = components.transformer(
+                input_ids=static["input_ids"],
+                text_indexes=static["text_indexes"],
+                position_ids=static["position_ids"],
+                und_len=static["und_len"],
+                sequence_length=static["sequence_length"],
+                vision_tokens=vision_tokens,
+                vision_token_shapes=static["vision_token_shapes"],
+                vision_sequence_indexes=static["vision_sequence_indexes"],
+                vision_mse_loss_indexes=static["vision_mse_loss_indexes"],
+                vision_timesteps=vision_timesteps,
+                vision_noisy_frame_indexes=static["vision_noisy_frame_indexes"],
+                return_dict=False,
+            )
         return preds_vision[-1]
 
     @torch.no_grad()
@@ -757,7 +777,11 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
         uncond_full_static = denoiser_input_fields["uncond_full_static"]
 
         cond_full = self._forward(
-            components, cond_full_static, block_state.vision_tokens_full, block_state.vision_timesteps
+            components,
+            cond_full_static,
+            block_state.vision_tokens_full,
+            block_state.vision_timesteps,
+            "cond",
         )
 
         cond_no_control = None
@@ -767,6 +791,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
                 cond_no_control_static,
                 block_state.vision_tokens_target,
                 block_state.vision_timesteps,
+                "cond_no_control",
             )
 
         uncond_full = None
@@ -776,6 +801,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
                 uncond_full_static,
                 block_state.vision_tokens_full,
                 block_state.vision_timesteps,
+                "uncond",
             )
 
         if needs_control_cfg and needs_text_cfg:
diff --git a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py
index d6c09703c12e..5efba85de64f 100644
--- a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py
+++ b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py
@@ -18,6 +18,30 @@ class Cosmos3OmniModularPipeline(ModularPipeline):
     inverse_image_resolution_template = "This image is not of {height}x{width} resolution."
     inverse_video_resolution_template = "This video is not of {height}x{width} resolution."
 
+    @property
+    def current_step_index(self):
+        return getattr(self, "_current_step_index", None)
+
+    @property
+    def current_sigma(self):
+        return getattr(self, "_current_sigma", None)
+
+    @property
+    def num_timesteps(self):
+        return getattr(self, "_num_timesteps", None)
+
+    def __call__(self, *args, **kwargs):
+        transformer = getattr(self, "transformer", None)
+        if hasattr(transformer, "_reset_stateful_cache"):
+            transformer._reset_stateful_cache()
+        try:
+            return super().__call__(*args, **kwargs)
+        finally:
+            self._current_step_index = None
+            self._current_sigma = None
+            if hasattr(transformer, "_reset_stateful_cache"):
+                transformer._reset_stateful_cache()
+
     @property
     def vae_scale_factor_spatial(self):
         if getattr(self, "vae", None) is not None:
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
index 6c025531fa22..d7c1bffd850e 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
@@ -16,6 +16,7 @@
 import json
 import math
 from collections.abc import Iterable
+from contextlib import nullcontext
 from dataclasses import dataclass
 from typing import Any, Callable, Literal
 
@@ -41,6 +42,11 @@
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
+def _cache_context(transformer: torch.nn.Module, name: str):
+    cache_context = getattr(transformer, "cache_context", None)
+    return cache_context(name) if callable(cache_context) else nullcontext()
+
+
 if is_cosmos_guardrail_available():
     from cosmos_guardrail import CosmosSafetyChecker
 else:
@@ -1292,6 +1298,14 @@ def _apply_video_safety_check(self, video: Any, output_type: str, device: torch.
     def current_timestep(self):
         return self._current_timestep
 
+    @property
+    def current_step_index(self):
+        return self._current_step_index
+
+    @property
+    def current_sigma(self):
+        return self._current_sigma
+
     @property
     def guidance_scale(self):
         return self._guidance_scale
@@ -1448,6 +1462,9 @@ def __call__(
                 `sound` (`torch.Tensor` of shape `[C, N]`, or `None` when `enable_sound=False`). Otherwise a tuple
                 `(video, sound)` with the same fields.
         """
+        if hasattr(self.transformer, "_reset_stateful_cache"):
+            self.transformer._reset_stateful_cache()
+
         if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
             callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
 
@@ -1492,6 +1509,8 @@ def __call__(
             )
 
         self._current_timestep = None
+        self._current_step_index = None
+        self._current_sigma = None
         self._interrupt = False
         self._guidance_scale = guidance_scale
 
@@ -1698,6 +1717,8 @@ def __call__(
                     continue
 
                 self._current_timestep = t
+                self._current_step_index = i
+                self._current_sigma = self.scheduler.sigmas[i]
                 timestep = t.item()
 
                 # The transformer projections (proj_in / audio_proj_in) are bf16; cast the per-step
@@ -1716,33 +1737,34 @@ def __call__(
                 )
 
                 # --- Conditional pass ---
-                preds_vision, preds_sound, preds_action = self.transformer(
-                    input_ids=cond_packed_static["input_ids"],
-                    text_indexes=cond_packed_static["text_indexes"],
-                    position_ids=cond_packed_static["position_ids"],
-                    und_len=cond_packed_static["und_len"],
-                    sequence_length=cond_packed_static["sequence_length"],
-                    vision_tokens=[vision_tokens],
-                    vision_token_shapes=cond_packed_static["vision_token_shapes"],
-                    vision_sequence_indexes=cond_packed_static["vision_sequence_indexes"],
-                    vision_mse_loss_indexes=cond_packed_static["vision_mse_loss_indexes"],
-                    vision_timesteps=vision_timesteps,
-                    vision_noisy_frame_indexes=cond_packed_static["vision_noisy_frame_indexes"],
-                    sound_tokens=[sound_tokens] if sound_tokens is not None else None,
-                    sound_token_shapes=cond_packed_static.get("sound_token_shapes"),
-                    sound_sequence_indexes=cond_packed_static.get("sound_sequence_indexes"),
-                    sound_mse_loss_indexes=cond_packed_static.get("sound_mse_loss_indexes"),
-                    sound_timesteps=sound_timesteps,
-                    sound_noisy_frame_indexes=cond_packed_static.get("sound_noisy_frame_indexes"),
-                    action_tokens=[action_tokens] if action_tokens is not None else None,
-                    action_token_shapes=cond_packed_static.get("action_token_shapes"),
-                    action_sequence_indexes=cond_packed_static.get("action_sequence_indexes"),
-                    action_mse_loss_indexes=cond_packed_static.get("action_mse_loss_indexes"),
-                    action_timesteps=action_timesteps,
-                    action_noisy_frame_indexes=cond_packed_static.get("action_noisy_frame_indexes"),
-                    action_domain_ids=[action_domain_id] if action_domain_id is not None else None,
-                    return_dict=False,
-                )
+                with _cache_context(self.transformer, "cond"):
+                    preds_vision, preds_sound, preds_action = self.transformer(
+                        input_ids=cond_packed_static["input_ids"],
+                        text_indexes=cond_packed_static["text_indexes"],
+                        position_ids=cond_packed_static["position_ids"],
+                        und_len=cond_packed_static["und_len"],
+                        sequence_length=cond_packed_static["sequence_length"],
+                        vision_tokens=[vision_tokens],
+                        vision_token_shapes=cond_packed_static["vision_token_shapes"],
+                        vision_sequence_indexes=cond_packed_static["vision_sequence_indexes"],
+                        vision_mse_loss_indexes=cond_packed_static["vision_mse_loss_indexes"],
+                        vision_timesteps=vision_timesteps,
+                        vision_noisy_frame_indexes=cond_packed_static["vision_noisy_frame_indexes"],
+                        sound_tokens=[sound_tokens] if sound_tokens is not None else None,
+                        sound_token_shapes=cond_packed_static.get("sound_token_shapes"),
+                        sound_sequence_indexes=cond_packed_static.get("sound_sequence_indexes"),
+                        sound_mse_loss_indexes=cond_packed_static.get("sound_mse_loss_indexes"),
+                        sound_timesteps=sound_timesteps,
+                        sound_noisy_frame_indexes=cond_packed_static.get("sound_noisy_frame_indexes"),
+                        action_tokens=[action_tokens] if action_tokens is not None else None,
+                        action_token_shapes=cond_packed_static.get("action_token_shapes"),
+                        action_sequence_indexes=cond_packed_static.get("action_sequence_indexes"),
+                        action_mse_loss_indexes=cond_packed_static.get("action_mse_loss_indexes"),
+                        action_timesteps=action_timesteps,
+                        action_noisy_frame_indexes=cond_packed_static.get("action_noisy_frame_indexes"),
+                        action_domain_ids=[action_domain_id] if action_domain_id is not None else None,
+                        return_dict=False,
+                    )
                 cond_v_vision, cond_v_sound, cond_v_action = self._mask_velocity_predictions(
                     preds_vision,
                     preds_sound,
@@ -1756,33 +1778,34 @@ def __call__(
                 # --- Unconditional pass (Skip if not using CFG) ---
                 uncond_v_vision = uncond_v_sound = uncond_v_action = None
                 if self.do_classifier_free_guidance:
-                    preds_vision, preds_sound, preds_action = self.transformer(
-                        input_ids=uncond_packed_static["input_ids"],
-                        text_indexes=uncond_packed_static["text_indexes"],
-                        position_ids=uncond_packed_static["position_ids"],
-                        und_len=uncond_packed_static["und_len"],
-                        sequence_length=uncond_packed_static["sequence_length"],
-                        vision_tokens=[vision_tokens],
-                        vision_token_shapes=uncond_packed_static["vision_token_shapes"],
-                        vision_sequence_indexes=uncond_packed_static["vision_sequence_indexes"],
-                        vision_mse_loss_indexes=uncond_packed_static["vision_mse_loss_indexes"],
-                        vision_timesteps=vision_timesteps,
-                        vision_noisy_frame_indexes=uncond_packed_static["vision_noisy_frame_indexes"],
-                        sound_tokens=[sound_tokens] if sound_tokens is not None else None,
-                        sound_token_shapes=uncond_packed_static.get("sound_token_shapes"),
-                        sound_sequence_indexes=uncond_packed_static.get("sound_sequence_indexes"),
-                        sound_mse_loss_indexes=uncond_packed_static.get("sound_mse_loss_indexes"),
-                        sound_timesteps=sound_timesteps,
-                        sound_noisy_frame_indexes=uncond_packed_static.get("sound_noisy_frame_indexes"),
-                        action_tokens=[action_tokens] if action_tokens is not None else None,
-                        action_token_shapes=uncond_packed_static.get("action_token_shapes"),
-                        action_sequence_indexes=uncond_packed_static.get("action_sequence_indexes"),
-                        action_mse_loss_indexes=uncond_packed_static.get("action_mse_loss_indexes"),
-                        action_timesteps=action_timesteps,
-                        action_noisy_frame_indexes=uncond_packed_static.get("action_noisy_frame_indexes"),
-                        action_domain_ids=[action_domain_id] if action_domain_id is not None else None,
-                        return_dict=False,
-                    )
+                    with _cache_context(self.transformer, "uncond"):
+                        preds_vision, preds_sound, preds_action = self.transformer(
+                            input_ids=uncond_packed_static["input_ids"],
+                            text_indexes=uncond_packed_static["text_indexes"],
+                            position_ids=uncond_packed_static["position_ids"],
+                            und_len=uncond_packed_static["und_len"],
+                            sequence_length=uncond_packed_static["sequence_length"],
+                            vision_tokens=[vision_tokens],
+                            vision_token_shapes=uncond_packed_static["vision_token_shapes"],
+                            vision_sequence_indexes=uncond_packed_static["vision_sequence_indexes"],
+                            vision_mse_loss_indexes=uncond_packed_static["vision_mse_loss_indexes"],
+                            vision_timesteps=vision_timesteps,
+                            vision_noisy_frame_indexes=uncond_packed_static["vision_noisy_frame_indexes"],
+                            sound_tokens=[sound_tokens] if sound_tokens is not None else None,
+                            sound_token_shapes=uncond_packed_static.get("sound_token_shapes"),
+                            sound_sequence_indexes=uncond_packed_static.get("sound_sequence_indexes"),
+                            sound_mse_loss_indexes=uncond_packed_static.get("sound_mse_loss_indexes"),
+                            sound_timesteps=sound_timesteps,
+                            sound_noisy_frame_indexes=uncond_packed_static.get("sound_noisy_frame_indexes"),
+                            action_tokens=[action_tokens] if action_tokens is not None else None,
+                            action_token_shapes=uncond_packed_static.get("action_token_shapes"),
+                            action_sequence_indexes=uncond_packed_static.get("action_sequence_indexes"),
+                            action_mse_loss_indexes=uncond_packed_static.get("action_mse_loss_indexes"),
+                            action_timesteps=action_timesteps,
+                            action_noisy_frame_indexes=uncond_packed_static.get("action_noisy_frame_indexes"),
+                            action_domain_ids=[action_domain_id] if action_domain_id is not None else None,
+                            return_dict=False,
+                        )
                     uncond_v_vision, uncond_v_sound, uncond_v_action = self._mask_velocity_predictions(
                         preds_vision,
                         preds_sound,
@@ -1852,6 +1875,8 @@ def __call__(
                     progress_bar.update()
 
         self._current_timestep = None
+        self._current_step_index = None
+        self._current_sigma = None
 
         # 8. Postprocess + decode
         sound = self.decode_sound(sound_latents) if sound_latents is not None else None
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
index b5bf5fb04393..2750838897ab 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
@@ -13,6 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from contextlib import contextmanager
+from unittest import mock
+
 import numpy as np
 import pytest
 import torch
@@ -167,6 +170,26 @@ def test_num_images_per_prompt(self):
     def test_float16_inference(self):
         pass
 
+    def test_transformer_cache_contexts_receive_exact_scheduler_metadata(self):
+        pipe = self.get_pipeline().to(torch_device)
+        observed = []
+
+        @contextmanager
+        def record_context(name):
+            observed.append((name, pipe.current_step_index, pipe.current_sigma))
+            yield
+
+        with mock.patch.object(pipe.transformer, "cache_context", side_effect=record_context):
+            pipe(**self.get_dummy_inputs(), output=self.output_name)
+
+        assert [name for name, _, _ in observed] == ["cond", "uncond", "cond", "uncond"]
+        for call_index, (_, step_index, sigma) in enumerate(observed):
+            expected_step = call_index // 2
+            assert step_index == expected_step
+            torch.testing.assert_close(sigma, pipe.scheduler.sigmas[expected_step])
+        assert pipe.current_step_index is None
+        assert pipe.current_sigma is None
+
     def test_vae_encoder_is_standalone_and_validates_conditioning_inputs(self):
         pipe = self.get_pipeline()
         vae_encoder = pipe.blocks.sub_blocks["vae_encoder"]
diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py
index cbe30dd41e5d..eb5c28e17efe 100644
--- a/tests/pipelines/cosmos/test_cosmos3.py
+++ b/tests/pipelines/cosmos/test_cosmos3.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import unittest
+from contextlib import contextmanager
 from unittest import mock
 
 import numpy as np
@@ -149,6 +150,30 @@ def callback_on_step_end(_pipeline, _step_index, _timestep, callback_kwargs):
         assert callback_latent_dtypes
         assert all(dtype == torch.float32 for dtype in callback_latent_dtypes)
         assert latents.dtype == torch.float32
+
+    def test_transformer_cache_contexts_receive_exact_scheduler_metadata(self):
+        pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
+        pipeline.set_progress_bar_config(disable=None)
+        observed = []
+
+        @contextmanager
+        def record_context(name):
+            observed.append((name, pipeline.current_step_index, pipeline.current_sigma))
+            yield
+
+        inputs = self.get_dummy_inputs(torch_device)
+        inputs["guidance_scale"] = 2.0
+        with mock.patch.object(pipeline.transformer, "cache_context", side_effect=record_context):
+            pipeline(**inputs)
+
+        assert [name for name, _, _ in observed] == ["cond", "uncond", "cond", "uncond"]
+        for call_index, (_, step_index, sigma) in enumerate(observed):
+            expected_step = call_index // 2
+            assert step_index == expected_step
+            torch.testing.assert_close(sigma, pipeline.scheduler.sigmas[expected_step])
+        assert pipeline.current_step_index is None
+        assert pipeline.current_sigma is None
+
     def test_cosmos3_tokenize_prompt_uses_checkpoint_system_prompt_default(self):
         components = self.get_dummy_components()
         components["default_use_system_prompt"] = False

From 613e44161175088afe24220d968374cce30ba395 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Mon, 31 Aug 2026 03:51:52 -0700
Subject: [PATCH 03/14] Add SeaCache support for Cosmos3

---
 src/diffusers/__init__.py                     |    6 +
 src/diffusers/hooks/__init__.py               |    1 +
 src/diffusers/hooks/_helpers.py               |   15 +
 src/diffusers/hooks/sea_cache.py              | 1111 +++++++++++++++++
 src/diffusers/models/cache_utils.py           |   31 +-
 src/diffusers/utils/dummy_pt_objects.py       |   23 +
 tests/hooks/test_sea_cache.py                 |  443 +++++++
 .../test_models_transformer_cosmos3.py        |  142 ++-
 8 files changed, 1770 insertions(+), 2 deletions(-)
 create mode 100644 src/diffusers/hooks/sea_cache.py
 create mode 100644 tests/hooks/test_sea_cache.py

diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py
index 1fc34e6bdbf6..37ea9c1db39f 100644
--- a/src/diffusers/__init__.py
+++ b/src/diffusers/__init__.py
@@ -204,6 +204,7 @@
             "LayerSkipConfig",
             "MagCacheConfig",
             "PyramidAttentionBroadcastConfig",
+            "SeaCacheConfig",
             "SmoothedEnergyGuidanceConfig",
             "TaylorSeerCacheConfig",
             "TextKVCacheConfig",
@@ -212,8 +213,10 @@
             "apply_layer_skip",
             "apply_mag_cache",
             "apply_pyramid_attention_broadcast",
+            "apply_sea_cache",
             "apply_taylorseer_cache",
             "apply_text_kv_cache",
+            "get_sea_cache_stats",
         ]
     )
     _import_structure["image_processor"] = [
@@ -1081,6 +1084,7 @@
             LayerSkipConfig,
             MagCacheConfig,
             PyramidAttentionBroadcastConfig,
+            SeaCacheConfig,
             SmoothedEnergyGuidanceConfig,
             TaylorSeerCacheConfig,
             TextKVCacheConfig,
@@ -1089,8 +1093,10 @@
             apply_layer_skip,
             apply_mag_cache,
             apply_pyramid_attention_broadcast,
+            apply_sea_cache,
             apply_taylorseer_cache,
             apply_text_kv_cache,
+            get_sea_cache_stats,
         )
         from .image_processor import (
             InpaintProcessor,
diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py
index d999ab32d6d7..2e9c7395c2b1 100644
--- a/src/diffusers/hooks/__init__.py
+++ b/src/diffusers/hooks/__init__.py
@@ -25,6 +25,7 @@
     from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook
     from .mag_cache import MagCacheConfig, apply_mag_cache
     from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
+    from .sea_cache import SeaCacheConfig, apply_sea_cache, get_sea_cache_stats
     from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig
     from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache
     from .tensor_parallel import apply_tensor_parallel
diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py
index 9cbe5bc8108f..bfc6c59477c8 100644
--- a/src/diffusers/hooks/_helpers.py
+++ b/src/diffusers/hooks/_helpers.py
@@ -27,6 +27,8 @@ class TransformerBlockMetadata:
     return_hidden_states_index: int = None
     return_encoder_hidden_states_index: int = None
     hidden_states_argument_name: str = "hidden_states"
+    encoder_hidden_states_argument_name: str = "encoder_hidden_states"
+    hidden_states_norm_module_name: str = None
 
     _cls: Type = None
     _cached_parameter_indices: dict[str, int] = None
@@ -174,6 +176,7 @@ def _register_transformer_blocks_metadata():
     from ..models.transformers.cogvideox_transformer_3d import CogVideoXBlock
     from ..models.transformers.transformer_bria import BriaTransformerBlock
     from ..models.transformers.transformer_cogview4 import CogView4TransformerBlock
+    from ..models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer
     from ..models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock
     from ..models.transformers.transformer_hunyuan_video import (
         HunyuanVideoSingleTransformerBlock,
@@ -230,6 +233,18 @@ def _register_transformer_blocks_metadata():
         ),
     )
 
+    # Cosmos 3
+    TransformerBlockRegistry.register(
+        model_class=Cosmos3VLTextMoTDecoderLayer,
+        metadata=TransformerBlockMetadata(
+            return_hidden_states_index=1,
+            return_encoder_hidden_states_index=0,
+            hidden_states_argument_name="gen_seq",
+            encoder_hidden_states_argument_name="und_seq",
+            hidden_states_norm_module_name="input_layernorm_moe_gen",
+        ),
+    )
+
     # Flux
     TransformerBlockRegistry.register(
         model_class=FluxTransformerBlock,
diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py
new file mode 100644
index 000000000000..c5d243a32721
--- /dev/null
+++ b/src/diffusers/hooks/sea_cache.py
@@ -0,0 +1,1111 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import inspect
+import math
+import time
+from dataclasses import dataclass
+from typing import Any, Callable, Literal, Sequence
+
+import torch
+
+from ..utils import logging
+from ..utils.torch_utils import unwrap_module
+from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS
+from ._helpers import TransformerBlockMetadata, TransformerBlockRegistry
+from .hooks import BaseState, HookRegistry, ModelHook, StateManager
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+_SEA_CACHE_ROOT_HOOK = "sea_cache_root"
+_SEA_CACHE_LEADER_BLOCK_HOOK = "sea_cache_leader_block"
+_SEA_CACHE_BLOCK_HOOK = "sea_cache_block"
+_SEA_CACHE_POST_NORM_HOOK = "sea_cache_post_norm"
+
+
+@dataclass
+class SeaCacheConfig:
+    r"""
+    Configuration for [SeaCache](https://huggingface.co/papers/2602.18993).
+
+    SeaCache compares Spectral-Evolution-Aware (SEA) indicators between scheduler steps. If their accumulated relative
+    change stays below `threshold`, the expensive language-model hidden transform is replaced with a cached residual.
+    For Cosmos 3, the residual spans the decoder stack and final pathway normalization; input packing and modality
+    prediction heads still execute.
+
+    Args:
+        threshold (`float`, defaults to `0.35`):
+            Accumulated relative-L1 budget. Larger values reuse the cache more often.
+        residual_order (`int`, defaults to `0`):
+            Order used to predict the generation-stream language-model residual. `0` directly reuses the most recent
+            residual and `1` linearly extrapolates from the two most recent full executions.
+        retention_steps (`int`, defaults to `1`):
+            Number of initial scheduler steps that always execute in full.
+        cache_end_steps (`int`, defaults to `1`):
+            Number of final scheduler steps that always execute in full.
+        power_exp (`float`, defaults to `3.0`):
+            Exponent of the SEA clean-signal power prior. SeaCache uses `3.0` for video features.
+        indicator_source (`str`, defaults to `"first_block"`):
+            Feature source used to construct the SEA indicator. `"first_block"` follows the published method and
+            filters the timestep-modulated pre-attention input of the first transformer block.
+            `"raw_vision_latents"` is an opt-in lower-cost approximation that filters the complete raw vision latent,
+            including clean conditioning frames in I2V. Thresholds are not generally transferable between the two
+            sources.
+        current_step_callback (`Callable[[], int]`):
+            Callback returning the current scheduler step index.
+        current_sigma_callback (`Callable[[], float]`):
+            Callback returning the exact current scheduler sigma in `[0, 1]`.
+        num_inference_steps_callback (`Callable[[], int]`):
+            Callback returning the number of scheduler steps in the current pipeline call.
+        metadata_callback (`Callable`, *optional*):
+            Advanced model adapter returning a list of `(indices, (T, H, W))` entries that locate projected noisy
+            vision tokens in the generation stream for the `"first_block"` indicator. Cosmos 3 uses its native
+            adapter when this is omitted.
+        raw_vision_callback (`Callable`, *optional*):
+            Advanced model adapter returning raw vision latents with shape `(C, T, H, W)` for the
+            `"raw_vision_latents"` indicator. Cosmos 3 uses its native adapter when this is omitted.
+        gate_schedule (`Sequence[bool]`, *optional*):
+            Per-step full-compute decisions to replay for controlled residual-prediction ablations. SEA still evaluates
+            its natural decision and reports mismatches; invalid or unsafe calls always fail open to full compute.
+
+    Example:
+        ```python
+        >>> from diffusers import Cosmos3OmniPipeline, SeaCacheConfig
+
+        >>> pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano")
+        >>> pipe.transformer.enable_cache(
+        ...     SeaCacheConfig(
+        ...         current_step_callback=lambda: pipe.current_step_index,
+        ...         current_sigma_callback=lambda: pipe.current_sigma,
+        ...         num_inference_steps_callback=lambda: pipe.num_timesteps,
+        ...     )
+        ... )
+        ```
+    """
+
+    threshold: float = 0.35
+    residual_order: int = 0
+    retention_steps: int = 1
+    cache_end_steps: int = 1
+    power_exp: float = 3.0
+    indicator_source: Literal["first_block", "raw_vision_latents"] = "first_block"
+    current_step_callback: Callable[[], int] = None
+    current_sigma_callback: Callable[[], float] = None
+    num_inference_steps_callback: Callable[[], int] = None
+    metadata_callback: Callable[
+        [torch.nn.Module, tuple[Any, ...], dict[str, Any]],
+        list[tuple[torch.Tensor, tuple[int, int, int]]] | None,
+    ] = None
+    raw_vision_callback: Callable[
+        [torch.nn.Module, tuple[Any, ...], dict[str, Any]],
+        list[torch.Tensor] | None,
+    ] = None
+    gate_schedule: Sequence[bool] = None
+
+    def __post_init__(self):
+        if not math.isfinite(self.threshold) or self.threshold < 0:
+            raise ValueError(f"`threshold` must be non-negative, got {self.threshold}.")
+        if self.residual_order not in (0, 1):
+            raise ValueError(f"`residual_order` must be 0 or 1, got {self.residual_order}.")
+        if self.retention_steps < 0:
+            raise ValueError(f"`retention_steps` must be non-negative, got {self.retention_steps}.")
+        if self.cache_end_steps < 0:
+            raise ValueError(f"`cache_end_steps` must be non-negative, got {self.cache_end_steps}.")
+        if not math.isfinite(self.power_exp) or self.power_exp <= 0:
+            raise ValueError(f"`power_exp` must be positive, got {self.power_exp}.")
+        if self.indicator_source not in ("first_block", "raw_vision_latents"):
+            raise ValueError(
+                f"`indicator_source` must be 'first_block' or 'raw_vision_latents', got {self.indicator_source!r}."
+            )
+        for name in (
+            "current_step_callback",
+            "current_sigma_callback",
+            "num_inference_steps_callback",
+            "metadata_callback",
+            "raw_vision_callback",
+        ):
+            callback = getattr(self, name)
+            if callback is not None and not callable(callback):
+                raise TypeError(f"`{name}` must be callable or `None`.")
+        if self.gate_schedule is not None:
+            self.gate_schedule = tuple(self.gate_schedule)
+            if any(not isinstance(value, bool) for value in self.gate_schedule):
+                raise TypeError("`gate_schedule` must contain only boolean values.")
+
+
+@dataclass
+class _SeaCacheForwardMetadata:
+    step_index: int
+    sigma: float
+    num_inference_steps: int
+    vision_layout: list[tuple[torch.Tensor, tuple[int, int, int]]] | None = None
+    raw_vision: list[torch.Tensor] | None = None
+
+
+class SeaCacheContextState(BaseState):
+    def __init__(self):
+        self.history: list[tuple[int, torch.Tensor, torch.Tensor]] = []
+        self.gate_key: tuple[int, float] | None = None
+        self.gate_should_compute = True
+        self.previous_indicator: list[torch.Tensor] | None = None
+        self.accumulated_distance = 0.0
+        self.skip_remaining = False
+        self.full_execution_pending = False
+        self.cacheable_execution = False
+        self.step_index: int | None = None
+        self.und_input: torch.Tensor | None = None
+        self.gen_input: torch.Tensor | None = None
+        self.und_output: torch.Tensor | None = None
+        self.cached_und_output: torch.Tensor | None = None
+        self.cached_gen_residual: torch.Tensor | None = None
+        self.full_started_at: float | None = None
+
+    def reset_forward(self):
+        self.skip_remaining = False
+        self.full_execution_pending = False
+        self.cacheable_execution = False
+        self.step_index = None
+        self.und_input = None
+        self.gen_input = None
+        self.und_output = None
+        self.cached_und_output = None
+        self.cached_gen_residual = None
+        self.full_started_at = None
+
+    def reset(self):
+        self.history = []
+        self.gate_key = None
+        self.gate_should_compute = True
+        self.previous_indicator = None
+        self.accumulated_distance = 0.0
+        self.reset_forward()
+
+
+class SeaCacheSharedState:
+    def __init__(self):
+        self._warned_messages: set[str] = set()
+        self.reset()
+
+    def reset(self):
+        self.forward_metadata: _SeaCacheForwardMetadata | None = None
+        self.transformer_calls = 0
+        self.gate_evaluations = 0
+        self.gate_full_decisions = 0
+        self.gate_skip_decisions = 0
+        self.gate_trace: list[bool] = []
+        self.gate_schedule_mismatches = 0
+        self.num_full_steps = 0
+        self.num_cached_steps = 0
+        self.fail_open_calls = 0
+        self.indicator_seconds = 0.0
+        self.decision_seconds = 0.0
+        self.full_seconds = 0.0
+        self.branch_full_executions: dict[str, int] = {}
+        self.branch_reuses: dict[str, int] = {}
+
+    def warn_once(self, message: str):
+        if message not in self._warned_messages:
+            logger.warning(message)
+            self._warned_messages.add(message)
+
+    def mark_fail_open(self, message: str):
+        self.fail_open_calls += 1
+        self.warn_once(message)
+
+    def resolve_gate(
+        self,
+        state: SeaCacheContextState,
+        metadata: _SeaCacheForwardMetadata,
+        indicator: list[torch.Tensor] | None,
+        config: SeaCacheConfig,
+    ) -> bool:
+        gate_key = (metadata.step_index, metadata.sigma)
+        if state.gate_key == gate_key:
+            return state.gate_should_compute
+
+        self.gate_evaluations += 1
+        is_non_adjacent = state.gate_key is not None and metadata.step_index != state.gate_key[0] + 1
+        if is_non_adjacent:
+            self.mark_fail_open("SeaCache received non-adjacent scheduler steps; running full.")
+        is_retained = metadata.step_index < config.retention_steps
+        is_in_cache_end = metadata.step_index >= metadata.num_inference_steps - config.cache_end_steps
+        is_first_observation = state.previous_indicator is None
+        invalid_gate = is_non_adjacent or indicator is None
+        forced_compute = invalid_gate or is_retained or is_in_cache_end or is_first_observation
+        candidate_accumulated_distance = 0.0
+
+        if forced_compute:
+            natural_should_compute = True
+        else:
+            if len(indicator) != len(state.previous_indicator) or not indicator:
+                distance = float("inf")
+                invalid_gate = True
+            else:
+                distance = 0.0
+                for current, previous in zip(indicator, state.previous_indicator):
+                    if (
+                        current.shape != previous.shape
+                        or current.device != previous.device
+                        or current.dtype != previous.dtype
+                    ):
+                        distance = float("inf")
+                        invalid_gate = True
+                        break
+                    numerator = (current.float() - previous.float()).abs().mean()
+                    denominator = previous.float().abs().mean() + 1e-16
+                    distance += float((numerator / denominator).detach().cpu())
+                distance /= len(indicator)
+
+            if not math.isfinite(distance):
+                invalid_gate = True
+                self.mark_fail_open("SeaCache indicator history changed shape, device, or dtype; running full.")
+            candidate_accumulated_distance = state.accumulated_distance + distance
+            natural_should_compute = invalid_gate or candidate_accumulated_distance >= config.threshold
+
+        should_compute = natural_should_compute
+
+        if config.gate_schedule is not None:
+            schedule_is_valid = len(config.gate_schedule) == metadata.num_inference_steps
+            scheduled_compute = bool(config.gate_schedule[metadata.step_index]) if schedule_is_valid else True
+            can_replay_skip = not (invalid_gate or is_retained or is_in_cache_end or is_first_observation)
+            if not schedule_is_valid:
+                self.mark_fail_open(
+                    "SeaCache gate schedule length does not match the number of inference steps; running full."
+                )
+                should_compute = True
+            elif scheduled_compute or can_replay_skip:
+                if natural_should_compute != scheduled_compute:
+                    self.gate_schedule_mismatches += 1
+                should_compute = scheduled_compute
+            else:
+                self.mark_fail_open("SeaCache gate schedule requested an unsafe cache hit; running full.")
+                should_compute = True
+
+        state.accumulated_distance = 0.0 if should_compute else candidate_accumulated_distance
+
+        state.gate_key = gate_key
+        state.gate_should_compute = should_compute
+        state.previous_indicator = None if indicator is None else [value.detach() for value in indicator]
+        self.gate_trace.append(should_compute)
+        if should_compute:
+            self.gate_full_decisions += 1
+        else:
+            self.gate_skip_decisions += 1
+        return should_compute
+
+
+def _get_block_inputs(
+    metadata: TransformerBlockMetadata, args: tuple[Any, ...], kwargs: dict[str, Any]
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    hidden_states = metadata._get_parameter_from_args_kwargs(metadata.hidden_states_argument_name, args, kwargs)
+    encoder_hidden_states = None
+    if metadata.return_encoder_hidden_states_index is not None:
+        encoder_hidden_states = metadata._get_parameter_from_args_kwargs(
+            metadata.encoder_hidden_states_argument_name, args, kwargs
+        )
+    return hidden_states, encoder_hidden_states
+
+
+def _build_block_output(
+    metadata: TransformerBlockMetadata,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor | None,
+) -> torch.Tensor | tuple[torch.Tensor, ...]:
+    if metadata.return_encoder_hidden_states_index is None:
+        return hidden_states
+
+    output = [None] * (max(metadata.return_hidden_states_index, metadata.return_encoder_hidden_states_index) + 1)
+    output[metadata.return_hidden_states_index] = hidden_states
+    output[metadata.return_encoder_hidden_states_index] = encoder_hidden_states
+    return tuple(output)
+
+
+def _get_block_outputs(
+    metadata: TransformerBlockMetadata, output: torch.Tensor | tuple[torch.Tensor, ...]
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    if isinstance(output, tuple):
+        hidden_states = output[metadata.return_hidden_states_index]
+        encoder_hidden_states = (
+            output[metadata.return_encoder_hidden_states_index]
+            if metadata.return_encoder_hidden_states_index is not None
+            else None
+        )
+        return hidden_states, encoder_hidden_states
+    return output, None
+
+
+def _record_full_execution(
+    config: SeaCacheConfig,
+    state_manager: StateManager,
+    shared_state: SeaCacheSharedState,
+    state: SeaCacheContextState,
+    gen_output: torch.Tensor,
+    und_output: torch.Tensor | None,
+) -> None:
+    shared_state.num_full_steps += 1
+    branch = state_manager._current_context
+    shared_state.branch_full_executions[branch] = shared_state.branch_full_executions.get(branch, 0) + 1
+    if state.full_started_at is not None:
+        shared_state.full_seconds += time.perf_counter() - state.full_started_at
+    if (
+        state.cacheable_execution
+        and state.step_index is not None
+        and state.gen_input is not None
+        and und_output is not None
+        and gen_output.shape == state.gen_input.shape
+    ):
+        state.history.append(
+            (
+                state.step_index,
+                und_output.detach().clone(),
+                (gen_output - state.gen_input).detach().clone(),
+            )
+        )
+        state.history = state.history[-(config.residual_order + 1) :]
+    state.reset_forward()
+
+
+def _prepare_cosmos3_vision_metadata(
+    module: torch.nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]
+) -> list[tuple[torch.Tensor, tuple[int, int, int]]] | None:
+    module = unwrap_module(module)
+    bound_arguments = inspect.signature(module.__class__.forward).bind_partial(module, *args, **kwargs).arguments
+    vision_tokens = bound_arguments.get("vision_tokens")
+    vision_token_shapes = bound_arguments.get("vision_token_shapes")
+    vision_sequence_indexes = bound_arguments.get("vision_sequence_indexes")
+    vision_timesteps = bound_arguments.get("vision_timesteps")
+    vision_noisy_frame_indexes = bound_arguments.get("vision_noisy_frame_indexes")
+    und_len = bound_arguments.get("und_len")
+
+    if (
+        not isinstance(vision_tokens, (list, tuple))
+        or not isinstance(vision_token_shapes, (list, tuple))
+        or not isinstance(vision_sequence_indexes, torch.Tensor)
+        or not isinstance(vision_timesteps, torch.Tensor)
+        or vision_timesteps.numel() == 0
+        or not isinstance(vision_noisy_frame_indexes, (list, tuple))
+        or und_len is None
+        or len(vision_tokens) != len(vision_token_shapes)
+        or len(vision_tokens) != len(vision_noisy_frame_indexes)
+    ):
+        return None
+
+    vision_sequence_indexes = vision_sequence_indexes.flatten()
+    layout = []
+    offset = 0
+    for token, token_shape, noisy_frame_indexes in zip(vision_tokens, vision_token_shapes, vision_noisy_frame_indexes):
+        if (
+            not isinstance(token, torch.Tensor)
+            or not isinstance(noisy_frame_indexes, torch.Tensor)
+            or len(token_shape) != 3
+        ):
+            return None
+
+        temporal, height, width = (int(value) for value in token_shape)
+        item_numel = temporal * height * width
+        item_indexes = vision_sequence_indexes[offset : offset + item_numel]
+        if item_indexes.numel() != item_numel:
+            return None
+        offset += item_numel
+
+        noisy_frame_indexes = noisy_frame_indexes.flatten().to(device=item_indexes.device, dtype=torch.long)
+        if noisy_frame_indexes.numel() == 0:
+            continue
+        if torch.any(noisy_frame_indexes < 0) or torch.any(noisy_frame_indexes >= temporal):
+            return None
+
+        item_indexes = item_indexes.reshape(temporal, height, width)
+        generation_indexes = item_indexes[noisy_frame_indexes].flatten() - int(und_len)
+        if torch.any(generation_indexes < 0):
+            return None
+        layout.append(
+            (
+                generation_indexes,
+                (int(noisy_frame_indexes.numel()), height, width),
+            )
+        )
+
+    if offset != vision_sequence_indexes.numel() or not layout:
+        return None
+    return layout
+
+
+def _prepare_cosmos3_raw_vision_metadata(
+    module: torch.nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]
+) -> list[torch.Tensor] | None:
+    module = unwrap_module(module)
+    bound_arguments = inspect.signature(module.__class__.forward).bind_partial(module, *args, **kwargs).arguments
+    vision_tokens = bound_arguments.get("vision_tokens")
+    vision_noisy_frame_indexes = bound_arguments.get("vision_noisy_frame_indexes")
+
+    if (
+        not isinstance(vision_tokens, (list, tuple))
+        or not isinstance(vision_noisy_frame_indexes, (list, tuple))
+        or len(vision_tokens) != len(vision_noisy_frame_indexes)
+    ):
+        return None
+
+    raw_vision = []
+    has_noisy_vision = False
+    for latent, noisy_frame_indexes in zip(vision_tokens, vision_noisy_frame_indexes):
+        if not isinstance(latent, torch.Tensor) or not isinstance(noisy_frame_indexes, torch.Tensor):
+            return None
+        if latent.ndim == 5:
+            if latent.shape[0] != 1:
+                return None
+            latent = latent.squeeze(0)
+        if latent.ndim != 4:
+            return None
+
+        noisy_frame_indexes = noisy_frame_indexes.flatten().to(device=latent.device, dtype=torch.long)
+        if torch.any(noisy_frame_indexes < 0) or torch.any(noisy_frame_indexes >= latent.shape[1]):
+            return None
+        has_noisy_vision = has_noisy_vision or noisy_frame_indexes.numel() > 0
+        raw_vision.append(latent)
+
+    return raw_vision if raw_vision and has_noisy_vision else None
+
+
+def _apply_sea_filter(
+    hidden_states: torch.Tensor,
+    sigma: float,
+    power_exp: float,
+) -> torch.Tensor:
+    hidden_states_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, sigma))
+    signal_scale = 1.0 - sigma
+    noise_scale = sigma
+    gain = 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]
+        gain = axis_gain.reshape(axis_shape) if gain is None else gain * axis_gain.reshape(axis_shape)
+
+    # SeaCache Eq. (7): density-normalize the combined spatiotemporal response
+    # to unit mean so cache distances are comparable across scheduler steps.
+    mean_gain = gain.mean()
+    if torch.isfinite(mean_gain) and mean_gain > 0:
+        gain = gain / mean_gain
+    return torch.fft.ifftn(spectrum * gain, dim=dimensions).real.to(hidden_states_dtype)
+
+
+def _is_parameter_sharded(module: torch.nn.Module) -> bool:
+    """Whether a block is managed by a parameter-sharding runtime that SeaCache cannot safely bypass."""
+
+    for submodule in unwrap_module(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):
+    _is_stateful = True
+
+    def __init__(
+        self,
+        config: SeaCacheConfig,
+        state_manager: StateManager,
+        shared_state: SeaCacheSharedState,
+        metadata_callback: Callable,
+        raw_vision_callback: Callable,
+        residual_boundary: str,
+    ):
+        super().__init__()
+        self.config = config
+        self.state_manager = state_manager
+        self.shared_state = shared_state
+        self.metadata_callback = metadata_callback
+        self.raw_vision_callback = raw_vision_callback
+        self.residual_boundary = residual_boundary
+        self._last_stats: dict[str, Any] | None = None
+
+    @property
+    def num_full_steps(self) -> int:
+        return self.shared_state.num_full_steps
+
+    @property
+    def num_cached_steps(self) -> int:
+        return self.shared_state.num_cached_steps
+
+    def stats(self) -> dict[str, Any]:
+        if self.shared_state.transformer_calls == 0 and self._last_stats is not None:
+            return copy.deepcopy(self._last_stats)
+
+        persistent_cache_bytes = 0
+        per_branch = {}
+        for branch, state in sorted(self.state_manager._state_cache.items()):
+            if state.previous_indicator is not None:
+                persistent_cache_bytes += sum(
+                    value.numel() * value.element_size() for value in state.previous_indicator
+                )
+            for _, und_output, gen_residual in state.history:
+                persistent_cache_bytes += und_output.numel() * und_output.element_size()
+                persistent_cache_bytes += gen_residual.numel() * gen_residual.element_size()
+            per_branch[branch] = {
+                "full_calls": self.shared_state.branch_full_executions.get(branch, 0),
+                "reuse_calls": self.shared_state.branch_reuses.get(branch, 0),
+            }
+
+        opportunities = self.shared_state.num_full_steps + self.shared_state.num_cached_steps
+        return {
+            "indicator_source": self.config.indicator_source,
+            "residual_order": self.config.residual_order,
+            "residual_boundary": self.residual_boundary,
+            "transformer_calls": self.shared_state.transformer_calls,
+            "gate_evaluations": self.shared_state.gate_evaluations,
+            "gate_full_decisions": self.shared_state.gate_full_decisions,
+            "gate_skip_decisions": self.shared_state.gate_skip_decisions,
+            "gate_trace": list(self.shared_state.gate_trace),
+            "gate_schedule_replayed": self.config.gate_schedule is not None,
+            "gate_schedule_mismatches": self.shared_state.gate_schedule_mismatches,
+            "actual_full_executions": self.shared_state.num_full_steps,
+            "actual_reuses": self.shared_state.num_cached_steps,
+            "actual_reuse_rate": (self.shared_state.num_cached_steps / opportunities if opportunities else 0.0),
+            "fail_open_calls": self.shared_state.fail_open_calls,
+            "sea_indicator_seconds": self.shared_state.indicator_seconds,
+            "sea_decision_seconds": self.shared_state.decision_seconds,
+            "sea_seconds": (self.shared_state.indicator_seconds + self.shared_state.decision_seconds),
+            "full_seconds": self.shared_state.full_seconds,
+            "timing_note": "host wall time; CUDA work is asynchronous",
+            "persistent_cache_bytes": persistent_cache_bytes,
+            "branch_full_executions": dict(sorted(self.shared_state.branch_full_executions.items())),
+            "branch_reuses": dict(sorted(self.shared_state.branch_reuses.items())),
+            "per_branch": per_branch,
+        }
+
+    def pre_forward(self, module: torch.nn.Module, *args, **kwargs):
+        self.shared_state.forward_metadata = None
+        self.shared_state.transformer_calls += 1
+        if torch.is_grad_enabled():
+            self.shared_state.mark_fail_open(
+                "SeaCache is inference-only; calls with autograd enabled run in fail-open mode."
+            )
+            return args, kwargs
+        if self.state_manager._current_context is None:
+            self.shared_state.mark_fail_open(
+                "SeaCache requires a cache context for each transformer call; running in fail-open mode."
+            )
+            return args, kwargs
+        callbacks = (
+            self.config.current_step_callback,
+            self.config.current_sigma_callback,
+            self.config.num_inference_steps_callback,
+        )
+        if any(callback is None for callback in callbacks):
+            self.shared_state.mark_fail_open(
+                "SeaCache is running in fail-open mode because scheduler step, sigma, and step-count callbacks are "
+                "required."
+            )
+            return args, kwargs
+
+        try:
+            step_index = self.config.current_step_callback()
+            sigma = self.config.current_sigma_callback()
+            num_inference_steps = self.config.num_inference_steps_callback()
+            if isinstance(step_index, torch.Tensor):
+                step_index = step_index.item()
+            if isinstance(sigma, torch.Tensor):
+                sigma = sigma.item()
+            if isinstance(num_inference_steps, torch.Tensor):
+                num_inference_steps = num_inference_steps.item()
+            step_index = int(step_index)
+            sigma = float(sigma)
+            num_inference_steps = int(num_inference_steps)
+        except (IndexError, TypeError, ValueError, RuntimeError) as error:
+            self.shared_state.mark_fail_open(
+                f"SeaCache scheduler metadata is unavailable; running in fail-open mode: {error}"
+            )
+            return args, kwargs
+
+        if (
+            step_index < 0
+            or num_inference_steps <= 0
+            or step_index >= num_inference_steps
+            or not math.isfinite(sigma)
+            or not 0.0 <= sigma <= 1.0
+        ):
+            self.shared_state.mark_fail_open(
+                "SeaCache scheduler metadata is invalid; expected a valid step index and exact sigma in [0, 1]."
+            )
+            return args, kwargs
+
+        vision_layout = None
+        raw_vision = None
+        try:
+            if self.config.indicator_source == "first_block":
+                vision_layout = (
+                    self.metadata_callback(module, args, kwargs) if self.metadata_callback is not None else None
+                )
+            else:
+                raw_vision = (
+                    self.raw_vision_callback(module, args, kwargs) if self.raw_vision_callback is not None else None
+                )
+        except (TypeError, ValueError, RuntimeError) as error:
+            self.shared_state.mark_fail_open(
+                f"SeaCache model metadata is unavailable; running in fail-open mode: {error}"
+            )
+            return args, kwargs
+        if self.config.indicator_source == "first_block" and not vision_layout:
+            self.shared_state.mark_fail_open(
+                "SeaCache requires noisy vision tokens; action-only, sound-only, and conditioning-only calls run in "
+                "fail-open mode."
+            )
+            return args, kwargs
+        if self.config.indicator_source == "raw_vision_latents" and not raw_vision:
+            self.shared_state.mark_fail_open(
+                "SeaCache requires raw noisy vision latents for the selected indicator source; action-only, sound-only, "
+                "conditioning-only, and unsupported model calls run in fail-open mode."
+            )
+            return args, kwargs
+
+        self.shared_state.forward_metadata = _SeaCacheForwardMetadata(
+            step_index=step_index,
+            sigma=sigma,
+            num_inference_steps=num_inference_steps,
+            vision_layout=vision_layout,
+            raw_vision=raw_vision,
+        )
+        return args, kwargs
+
+    def post_forward(self, module: torch.nn.Module, output: Any) -> Any:
+        self.shared_state.forward_metadata = None
+        if self.state_manager._current_context is not None:
+            self.state_manager.get_state().reset_forward()
+        return output
+
+    def reset_state(self, module: torch.nn.Module):
+        if self.shared_state.transformer_calls > 0:
+            self._last_stats = self.stats()
+        self.state_manager.reset()
+        self.shared_state.reset()
+        return module
+
+
+class SeaCacheLeaderBlockHook(ModelHook):
+    def __init__(
+        self,
+        config: SeaCacheConfig,
+        state_manager: StateManager,
+        shared_state: SeaCacheSharedState,
+        post_norm_boundary: bool = False,
+    ):
+        super().__init__()
+        self.config = config
+        self.state_manager = state_manager
+        self.shared_state = shared_state
+        self.post_norm_boundary = post_norm_boundary
+        self._metadata = None
+        self._normalization = None
+
+    def initialize_hook(self, module: torch.nn.Module):
+        module = unwrap_module(module)
+        self._metadata = TransformerBlockRegistry.get(module.__class__)
+        if self._metadata.hidden_states_norm_module_name is not None:
+            self._normalization = getattr(module, self._metadata.hidden_states_norm_module_name)
+        return module
+
+    def _build_indicator(
+        self,
+        hidden_states: torch.Tensor,
+        forward_metadata: _SeaCacheForwardMetadata,
+    ) -> list[torch.Tensor] | None:
+        if self.config.indicator_source == "raw_vision_latents":
+            if not forward_metadata.raw_vision:
+                return None
+            indicator = []
+            for latent in forward_metadata.raw_vision:
+                raw_vision = latent.movedim(0, -1)
+                indicator.append(
+                    _apply_sea_filter(
+                        raw_vision,
+                        sigma=forward_metadata.sigma,
+                        power_exp=self.config.power_exp,
+                    ).detach()
+                )
+            return indicator
+
+        if self._normalization is None:
+            return None
+        if not forward_metadata.vision_layout:
+            return None
+        normalized_hidden_states = self._normalization(hidden_states)
+        indicator = []
+        for indexes, shape in forward_metadata.vision_layout:
+            indexes = indexes.to(device=normalized_hidden_states.device, dtype=torch.long)
+            if (
+                indexes.numel() != math.prod(shape)
+                or torch.any(indexes < 0)
+                or torch.any(indexes >= normalized_hidden_states.shape[0])
+            ):
+                return None
+            noisy_vision = normalized_hidden_states.index_select(0, indexes).reshape(*shape, -1)
+            indicator.append(
+                _apply_sea_filter(
+                    noisy_vision,
+                    sigma=forward_metadata.sigma,
+                    power_exp=self.config.power_exp,
+                ).detach()
+            )
+        return indicator
+
+    @torch.compiler.disable
+    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
+        hidden_states, encoder_hidden_states = _get_block_inputs(self._metadata, args, kwargs)
+        context_is_set = self.state_manager._current_context is not None
+        state = self.state_manager.get_state() if context_is_set else None
+        if state is not None:
+            state.reset_forward()
+            state.full_execution_pending = True
+            state.gen_input = hidden_states
+            state.und_input = encoder_hidden_states
+
+        forward_metadata = self.shared_state.forward_metadata
+        if state is None or forward_metadata is None:
+            if state is not None:
+                state.full_started_at = time.perf_counter()
+            return self.fn_ref.original_forward(*args, **kwargs)
+
+        state.step_index = forward_metadata.step_index
+        state.cacheable_execution = True
+        indicator_started = time.perf_counter()
+        indicator_error_reported = False
+        if _is_parameter_sharded(module):
+            self.shared_state.mark_fail_open(
+                "SeaCache cannot safely bypass parameter-sharded transformer blocks; running in fail-open mode."
+            )
+            indicator = None
+            indicator_error_reported = True
+        else:
+            try:
+                indicator = self._build_indicator(hidden_states, forward_metadata)
+            except (TypeError, ValueError, RuntimeError) as error:
+                self.shared_state.mark_fail_open(
+                    f"SeaCache could not construct its vision indicator; running in fail-open mode: {error}"
+                )
+                indicator = None
+                indicator_error_reported = True
+        self.shared_state.indicator_seconds += time.perf_counter() - indicator_started
+        if indicator is None and not indicator_error_reported:
+            self.shared_state.mark_fail_open(
+                "SeaCache could not construct its vision indicator; running in fail-open mode."
+            )
+        decision_started = time.perf_counter()
+        should_compute = self.shared_state.resolve_gate(state, forward_metadata, indicator, self.config)
+        self.shared_state.decision_seconds += time.perf_counter() - decision_started
+
+        if should_compute or not state.history:
+            if not should_compute:
+                state.accumulated_distance = 0.0
+                self.shared_state.mark_fail_open(
+                    "SeaCache selected a cache hit without residual history; running in fail-open mode."
+                )
+            state.full_started_at = time.perf_counter()
+            return self.fn_ref.original_forward(*args, **kwargs)
+
+        residual_history = state.history[-(self.config.residual_order + 1) :]
+        _, cached_und, cached_residual = residual_history[-1]
+        if (
+            any(
+                residual.shape != hidden_states.shape
+                or residual.device != hidden_states.device
+                or residual.dtype != hidden_states.dtype
+                for _, _, residual in residual_history
+            )
+            or encoder_hidden_states is None
+            or cached_und.shape != encoder_hidden_states.shape
+            or cached_und.device != encoder_hidden_states.device
+            or cached_und.dtype != encoder_hidden_states.dtype
+        ):
+            state.history = []
+            state.accumulated_distance = 0.0
+            self.shared_state.mark_fail_open(
+                "SeaCache residual history changed shape, device, or dtype; running in fail-open mode."
+            )
+            state.full_started_at = time.perf_counter()
+            return self.fn_ref.original_forward(*args, **kwargs)
+
+        if self.config.residual_order == 1 and len(residual_history) >= 2:
+            previous_step, _, previous_residual = residual_history[-2]
+            latest_step, _, latest_residual = residual_history[-1]
+            if latest_step != previous_step:
+                step_scale = (forward_metadata.step_index - latest_step) / (latest_step - previous_step)
+                cached_residual = latest_residual + (latest_residual - previous_residual) * step_scale
+
+        state.skip_remaining = True
+        state.full_execution_pending = False
+        state.cached_und_output = cached_und
+        state.cached_gen_residual = cached_residual
+        self.shared_state.num_cached_steps += 1
+        branch = self.state_manager._current_context
+        self.shared_state.branch_reuses[branch] = self.shared_state.branch_reuses.get(branch, 0) + 1
+        if self.post_norm_boundary:
+            return _build_block_output(self._metadata, hidden_states, encoder_hidden_states)
+        return _build_block_output(self._metadata, hidden_states + cached_residual, cached_und)
+
+
+class SeaCacheBlockHook(ModelHook):
+    def __init__(
+        self,
+        config: SeaCacheConfig,
+        state_manager: StateManager,
+        shared_state: SeaCacheSharedState,
+        is_tail: bool = False,
+        post_norm_boundary: bool = False,
+    ):
+        super().__init__()
+        self.config = config
+        self.state_manager = state_manager
+        self.shared_state = shared_state
+        self.is_tail = is_tail
+        self.post_norm_boundary = post_norm_boundary
+        self._metadata = None
+
+    def initialize_hook(self, module: torch.nn.Module):
+        self._metadata = TransformerBlockRegistry.get(unwrap_module(module).__class__)
+        return module
+
+    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
+        if self.state_manager._current_context is None:
+            return self.fn_ref.original_forward(*args, **kwargs)
+
+        state: SeaCacheContextState = self.state_manager.get_state()
+        if state.skip_remaining:
+            hidden_states, encoder_hidden_states = _get_block_inputs(self._metadata, args, kwargs)
+            return _build_block_output(self._metadata, hidden_states, encoder_hidden_states)
+
+        output = self.fn_ref.original_forward(*args, **kwargs)
+        if not self.is_tail or state.skip_remaining or not state.full_execution_pending:
+            return output
+        if self.post_norm_boundary:
+            return output
+
+        hidden_states, encoder_hidden_states = _get_block_outputs(self._metadata, output)
+        _record_full_execution(
+            self.config,
+            self.state_manager,
+            self.shared_state,
+            state,
+            gen_output=hidden_states,
+            und_output=encoder_hidden_states,
+        )
+        return output
+
+
+class SeaCachePostNormHook(ModelHook):
+    def __init__(
+        self,
+        config: SeaCacheConfig,
+        state_manager: StateManager,
+        shared_state: SeaCacheSharedState,
+        pathway: Literal["und", "gen"],
+    ):
+        super().__init__()
+        self.config = config
+        self.state_manager = state_manager
+        self.shared_state = shared_state
+        self.pathway = pathway
+
+    @torch.compiler.disable
+    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
+        if self.state_manager._current_context is None:
+            return self.fn_ref.original_forward(*args, **kwargs)
+
+        state: SeaCacheContextState = self.state_manager.get_state()
+        if state.skip_remaining:
+            if self.pathway == "und":
+                if state.cached_und_output is not None:
+                    return state.cached_und_output
+            elif state.gen_input is not None and state.cached_gen_residual is not None:
+                output = state.gen_input + state.cached_gen_residual
+                state.reset_forward()
+                return output
+
+            self.shared_state.mark_fail_open(
+                "SeaCache post-normalization state is incomplete after the decoder stack was skipped."
+            )
+            output = self.fn_ref.original_forward(*args, **kwargs)
+            state.reset_forward()
+            return output
+
+        output = self.fn_ref.original_forward(*args, **kwargs)
+        if not state.full_execution_pending:
+            return output
+        if self.pathway == "und":
+            state.und_output = output
+            return output
+
+        _record_full_execution(
+            self.config,
+            self.state_manager,
+            self.shared_state,
+            state,
+            gen_output=output,
+            und_output=state.und_output,
+        )
+        return output
+
+
+def apply_sea_cache(module: torch.nn.Module, config: SeaCacheConfig) -> None:
+    r"""
+    Apply SeaCache to a supported transformer.
+
+    The hook caches the transformer's expensive language-model hidden transform. For Cosmos 3, the cache stores
+    post-normalization understanding output and a generation residual from the pre-block input to the
+    post-normalization output. Modality prediction heads continue to run normally. Other model adapters fall back to
+    caching the complete repeated-block stack.
+
+    Args:
+        module (`torch.nn.Module`):
+            Transformer module to cache.
+        config (`SeaCacheConfig`):
+            SeaCache configuration.
+    """
+    from ..models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
+
+    unwrapped_module = unwrap_module(module)
+    is_cosmos3 = isinstance(unwrapped_module, Cosmos3OmniTransformer)
+    metadata_callback = config.metadata_callback
+    raw_vision_callback = config.raw_vision_callback
+    if metadata_callback is None or raw_vision_callback is None:
+        if is_cosmos3:
+            if metadata_callback is None:
+                metadata_callback = _prepare_cosmos3_vision_metadata
+            if raw_vision_callback is None:
+                raw_vision_callback = _prepare_cosmos3_raw_vision_metadata
+
+    post_norm_modules = None
+    if is_cosmos3:
+        und_norm = getattr(unwrapped_module, "norm", None)
+        gen_norm = getattr(unwrapped_module, "norm_moe_gen", None)
+        if isinstance(und_norm, torch.nn.Module) and isinstance(gen_norm, torch.nn.Module):
+            post_norm_modules = (("und", "norm", und_norm), ("gen", "norm_moe_gen", gen_norm))
+        else:
+            logger.warning(
+                "SeaCache could not locate the Cosmos 3 final pathway normalizations; falling back to the repeated-block "
+                "residual boundary."
+            )
+    post_norm_boundary = post_norm_modules is not None
+    residual_boundary = "post_language_model_norm" if post_norm_boundary else "repeated_block_stack"
+
+    blocks = []
+    for name, submodule in unwrapped_module.named_children():
+        if name not in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS or not isinstance(submodule, torch.nn.ModuleList):
+            continue
+        blocks.extend((f"{name}.{index}", block) for index, block in enumerate(submodule))
+
+    if not blocks:
+        raise ValueError("SeaCache found no repeated transformer blocks on the model.")
+
+    state_manager = StateManager(SeaCacheContextState)
+    shared_state = SeaCacheSharedState()
+    root_registry = HookRegistry.check_if_exists_or_initialize(module)
+    registrations: list[tuple[HookRegistry, str]] = []
+
+    def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None:
+        registry = HookRegistry.check_if_exists_or_initialize(target)
+        registry.register_hook(hook, name)
+        registrations.append((registry, name))
+
+    try:
+        register_hook(
+            module,
+            SeaCacheRootHook(
+                config,
+                state_manager,
+                shared_state,
+                metadata_callback,
+                raw_vision_callback,
+                residual_boundary,
+            ),
+            _SEA_CACHE_ROOT_HOOK,
+        )
+
+        leader_name, leader = blocks[0]
+        logger.debug(f"Applying SeaCache leader hook to '{leader_name}'.")
+        register_hook(
+            leader,
+            SeaCacheLeaderBlockHook(config, state_manager, shared_state, post_norm_boundary=post_norm_boundary),
+            _SEA_CACHE_LEADER_BLOCK_HOOK,
+        )
+
+        for name, block in blocks[1:-1]:
+            logger.debug(f"Applying SeaCache identity hook to '{name}'.")
+            register_hook(
+                block,
+                SeaCacheBlockHook(config, state_manager, shared_state, post_norm_boundary=post_norm_boundary),
+                _SEA_CACHE_BLOCK_HOOK,
+            )
+
+        tail_name, tail = blocks[-1]
+        logger.debug(f"Applying SeaCache tail hook to '{tail_name}'.")
+        register_hook(
+            tail,
+            SeaCacheBlockHook(
+                config,
+                state_manager,
+                shared_state,
+                is_tail=True,
+                post_norm_boundary=post_norm_boundary,
+            ),
+            _SEA_CACHE_BLOCK_HOOK,
+        )
+        if post_norm_modules is not None:
+            for pathway, name, norm_module in post_norm_modules:
+                logger.debug(f"Applying SeaCache post-normalization hook to '{name}'.")
+                register_hook(
+                    norm_module,
+                    SeaCachePostNormHook(config, state_manager, shared_state, pathway=pathway),
+                    _SEA_CACHE_POST_NORM_HOOK,
+                )
+    except Exception:
+        for registry, name in reversed(registrations):
+            registry.remove_hook(name, recurse=False)
+        root_registry._child_registries_cache = None
+        raise
+
+    root_registry._child_registries_cache = None
+
+
+def get_sea_cache_stats(module: torch.nn.Module) -> dict[str, Any]:
+    """Return statistics for the SeaCache instance currently attached to ``module``."""
+
+    registry = getattr(module, "_diffusers_hook", None)
+    root_hook = registry.get_hook(_SEA_CACHE_ROOT_HOOK) if registry is not None else None
+    if not isinstance(root_hook, SeaCacheRootHook):
+        raise ValueError("SeaCache is not enabled on this module.")
+    return root_hook.stats()
diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py
index 5bf161b2cc08..6f6542df11c2 100644
--- a/src/diffusers/models/cache_utils.py
+++ b/src/diffusers/models/cache_utils.py
@@ -28,6 +28,7 @@ class CacheMixin:
         - [Pyramid Attention Broadcast](https://huggingface.co/papers/2408.12588)
         - [FasterCache](https://huggingface.co/papers/2410.19355)
         - [FirstBlockCache](https://github.com/chengzeyi/ParaAttention/blob/7a266123671b55e7e5a2fe9af3121f07a36afc78/README.md#first-block-cache-our-dynamic-caching)
+        - [SeaCache](https://huggingface.co/papers/2602.18993)
     """
 
     _cache_config = None
@@ -41,11 +42,12 @@ def enable_cache(self, config) -> None:
         Enable caching techniques on the model.
 
         Args:
-            config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | TextKVCacheConfig`):
+            config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | SeaCacheConfig | TextKVCacheConfig`):
                 The configuration for applying the caching technique. Currently supported caching techniques are:
                     - [`~hooks.PyramidAttentionBroadcastConfig`]
                     - [`~hooks.FasterCacheConfig`]
                     - [`~hooks.FirstBlockCacheConfig`]
+                    - [`~hooks.SeaCacheConfig`]
                     - [`~hooks.TextKVCacheConfig`]
 
         Example:
@@ -72,12 +74,14 @@ def enable_cache(self, config) -> None:
             HookRegistry,
             MagCacheConfig,
             PyramidAttentionBroadcastConfig,
+            SeaCacheConfig,
             TaylorSeerCacheConfig,
             TextKVCacheConfig,
             apply_faster_cache,
             apply_first_block_cache,
             apply_mag_cache,
             apply_pyramid_attention_broadcast,
+            apply_sea_cache,
             apply_taylorseer_cache,
             apply_text_kv_cache,
         )
@@ -97,6 +101,8 @@ def enable_cache(self, config) -> None:
             apply_text_kv_cache(self, config)
         elif isinstance(config, PyramidAttentionBroadcastConfig):
             apply_pyramid_attention_broadcast(self, config)
+        elif isinstance(config, SeaCacheConfig):
+            apply_sea_cache(self, config)
         elif isinstance(config, TaylorSeerCacheConfig):
             apply_taylorseer_cache(self, config)
         else:
@@ -112,6 +118,7 @@ def disable_cache(self) -> None:
             HookRegistry,
             MagCacheConfig,
             PyramidAttentionBroadcastConfig,
+            SeaCacheConfig,
             TaylorSeerCacheConfig,
             TextKVCacheConfig,
         )
@@ -119,6 +126,12 @@ def disable_cache(self) -> None:
         from ..hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK
         from ..hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK
         from ..hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK
+        from ..hooks.sea_cache import (
+            _SEA_CACHE_BLOCK_HOOK,
+            _SEA_CACHE_LEADER_BLOCK_HOOK,
+            _SEA_CACHE_POST_NORM_HOOK,
+            _SEA_CACHE_ROOT_HOOK,
+        )
         from ..hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK
         from ..hooks.text_kv_cache import _TEXT_KV_CACHE_BLOCK_HOOK, _TEXT_KV_CACHE_TRANSFORMER_HOOK
 
@@ -141,6 +154,11 @@ def disable_cache(self) -> None:
         elif isinstance(self._cache_config, TextKVCacheConfig):
             registry.remove_hook(_TEXT_KV_CACHE_TRANSFORMER_HOOK, recurse=True)
             registry.remove_hook(_TEXT_KV_CACHE_BLOCK_HOOK, recurse=True)
+        elif isinstance(self._cache_config, SeaCacheConfig):
+            registry.remove_hook(_SEA_CACHE_POST_NORM_HOOK, recurse=True)
+            registry.remove_hook(_SEA_CACHE_BLOCK_HOOK, recurse=True)
+            registry.remove_hook(_SEA_CACHE_LEADER_BLOCK_HOOK, recurse=True)
+            registry.remove_hook(_SEA_CACHE_ROOT_HOOK, recurse=True)
         elif isinstance(self._cache_config, TaylorSeerCacheConfig):
             registry.remove_hook(_TAYLORSEER_CACHE_HOOK, recurse=True)
         else:
@@ -149,6 +167,17 @@ def disable_cache(self) -> None:
         self._cache_config = None
         registry._child_registries_cache = None
 
+    def get_cache_stats(self) -> dict:
+        """Return statistics for the currently enabled cache implementation."""
+
+        from ..hooks import SeaCacheConfig, get_sea_cache_stats
+
+        if self._cache_config is None:
+            raise ValueError("Caching is not enabled.")
+        if isinstance(self._cache_config, SeaCacheConfig):
+            return get_sea_cache_stats(self)
+        raise NotImplementedError(f"Cache statistics are not available for {type(self._cache_config).__name__}.")
+
     def _reset_stateful_cache(self, recurse: bool = True) -> None:
         from ..hooks import HookRegistry
 
diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py
index 1598814f835a..324c92f3d507 100644
--- a/src/diffusers/utils/dummy_pt_objects.py
+++ b/src/diffusers/utils/dummy_pt_objects.py
@@ -272,6 +272,21 @@ def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["torch"])
 
 
+class SeaCacheConfig(metaclass=DummyObject):
+    _backends = ["torch"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["torch"])
+
+    @classmethod
+    def from_config(cls, *args, **kwargs):
+        requires_backends(cls, ["torch"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["torch"])
+
+
 class SmoothedEnergyGuidanceConfig(metaclass=DummyObject):
     _backends = ["torch"]
 
@@ -337,6 +352,10 @@ def apply_pyramid_attention_broadcast(*args, **kwargs):
     requires_backends(apply_pyramid_attention_broadcast, ["torch"])
 
 
+def apply_sea_cache(*args, **kwargs):
+    requires_backends(apply_sea_cache, ["torch"])
+
+
 def apply_taylorseer_cache(*args, **kwargs):
     requires_backends(apply_taylorseer_cache, ["torch"])
 
@@ -345,6 +364,10 @@ def apply_text_kv_cache(*args, **kwargs):
     requires_backends(apply_text_kv_cache, ["torch"])
 
 
+def get_sea_cache_stats(*args, **kwargs):
+    requires_backends(get_sea_cache_stats, ["torch"])
+
+
 class InpaintProcessor(metaclass=DummyObject):
     _backends = ["torch"]
 
diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py
new file mode 100644
index 000000000000..85db1fbc22c1
--- /dev/null
+++ b/tests/hooks/test_sea_cache.py
@@ -0,0 +1,443 @@
+# Copyright 2026 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from diffusers import SeaCacheConfig
+from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry
+from diffusers.hooks.hooks import HookRegistry, ModelHook
+from diffusers.hooks.sea_cache import (
+    _SEA_CACHE_BLOCK_HOOK,
+    _SEA_CACHE_LEADER_BLOCK_HOOK,
+    _SEA_CACHE_ROOT_HOOK,
+    _apply_sea_filter,
+)
+from diffusers.models.cache_utils import CacheMixin
+
+
+class CountingIdentity(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.calls = 0
+
+    def forward(self, hidden_states):
+        self.calls += 1
+        return hidden_states
+
+
+class DummySeaBlock(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.indicator_norm = CountingIdentity()
+        self.calls = 0
+
+    def forward(self, und_seq, gen_seq, rotary_emb=None):
+        self.calls += 1
+        return und_seq + 1, gen_seq * 2
+
+
+class DummySeaTransformer(torch.nn.Module, CacheMixin):
+    def __init__(self, num_layers=3):
+        super().__init__()
+        self.layers = torch.nn.ModuleList([DummySeaBlock() for _ in range(num_layers)])
+
+    def forward(self, hidden_states):
+        shape = hidden_states.shape
+        gen_seq = hidden_states.reshape(-1, shape[-1])
+        und_seq = torch.zeros(1, shape[-1], device=hidden_states.device, dtype=hidden_states.dtype)
+        for layer in self.layers:
+            und_seq, gen_seq = layer(und_seq, gen_seq)
+        return und_seq, gen_seq.reshape(shape)
+
+
+@pytest.fixture(autouse=True)
+def register_dummy_sea_block():
+    TransformerBlockRegistry.register(
+        DummySeaBlock,
+        TransformerBlockMetadata(
+            return_hidden_states_index=1,
+            return_encoder_hidden_states_index=0,
+            hidden_states_argument_name="gen_seq",
+            encoder_hidden_states_argument_name="und_seq",
+            hidden_states_norm_module_name="indicator_norm",
+        ),
+    )
+
+
+def _metadata_callback(module, args, kwargs):
+    hidden_states = kwargs.get("hidden_states", args[0] if args else None)
+    if hidden_states is None:
+        return None
+    temporal, height, width = hidden_states.shape[:3]
+    indexes = torch.arange(temporal * height * width, device=hidden_states.device)
+    return [(indexes, (temporal, height, width))]
+
+
+def _make_config(runtime, **kwargs):
+    config_kwargs = {
+        "threshold": 100.0,
+        "retention_steps": 1,
+        "cache_end_steps": 0,
+        "current_step_callback": lambda: runtime["step"],
+        "current_sigma_callback": lambda: runtime["sigma"],
+        "num_inference_steps_callback": lambda: runtime["num_steps"],
+        "metadata_callback": _metadata_callback,
+    }
+    config_kwargs.update(kwargs)
+    return SeaCacheConfig(**config_kwargs)
+
+
+def _get_root_hook(model):
+    return model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK)
+
+
+@torch.no_grad()
+def test_sea_cache_uses_independent_gates_and_histories_per_context():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 4}
+    model = DummySeaTransformer()
+    model.enable_cache(_make_config(runtime, threshold=0.5))
+    root_hook = _get_root_hook(model)
+
+    first_input = torch.ones(2, 2, 2, 1)
+    with model.cache_context("cond"):
+        _, first_output = model(first_input)
+    torch.testing.assert_close(first_output, first_input * 8)
+    assert [block.calls for block in model.layers] == [1, 1, 1]
+
+    with model.cache_context("uncond"):
+        _, first_uncond_output = model(first_input)
+    torch.testing.assert_close(first_uncond_output, first_input * 8)
+    assert [block.calls for block in model.layers] == [2, 2, 2]
+
+    runtime.update(step=1)
+    with model.cache_context("cond"):
+        _, cached_output = model(first_input)
+    torch.testing.assert_close(cached_output, first_input * 8)
+    assert [block.calls for block in model.layers] == [2, 2, 2]
+
+    changed_uncond_input = first_input * 100
+    with model.cache_context("uncond"):
+        _, uncond_output = model(changed_uncond_input)
+    torch.testing.assert_close(uncond_output, changed_uncond_input * 8)
+    assert [block.calls for block in model.layers] == [3, 3, 3]
+    assert model.layers[0].indicator_norm.calls == 4
+    assert root_hook.num_full_steps == 3
+    assert root_hook.num_cached_steps == 1
+    stats = model.get_cache_stats()
+    assert stats["indicator_source"] == "first_block"
+    assert stats["residual_order"] == 0
+    assert stats["residual_boundary"] == "repeated_block_stack"
+    assert stats["transformer_calls"] == 4
+    assert stats["gate_evaluations"] == 4
+    assert stats["gate_trace"] == [True, True, False, True]
+    assert stats["branch_full_executions"] == {"cond": 1, "uncond": 2}
+    assert stats["branch_reuses"] == {"cond": 1}
+    assert stats["persistent_cache_bytes"] > 0
+
+    model._reset_stateful_cache()
+    assert root_hook.num_full_steps == 0
+    assert root_hook.num_cached_steps == 0
+    archived_stats = model.get_cache_stats()
+    assert archived_stats["transformer_calls"] == 4
+    assert archived_stats["gate_trace"] == [True, True, False, True]
+
+    model.disable_cache()
+    assert not model.is_cache_enabled
+
+
+@torch.no_grad()
+def test_sea_cache_raw_vision_indicator_includes_conditioning_frames():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
+    model = DummySeaTransformer()
+
+    def raw_vision(module, args, kwargs):
+        hidden_states = kwargs.get("hidden_states", args[0] if args else None)
+        latent = hidden_states.permute(3, 0, 1, 2)
+        return [latent]
+
+    model.enable_cache(
+        _make_config(
+            runtime,
+            threshold=1e-6,
+            indicator_source="raw_vision_latents",
+            raw_vision_callback=raw_vision,
+        )
+    )
+
+    first_input = torch.tensor([1.0, 2.0]).reshape(2, 1, 1, 1)
+    with model.cache_context("cond"):
+        model(first_input)
+
+    # The noisy frame stays fixed, but changing the conditioning frame changes the complete raw-latent indicator.
+    runtime.update(step=1, sigma=0.9)
+    second_input = torch.tensor([100.0, 2.0]).reshape(2, 1, 1, 1)
+    with model.cache_context("cond"):
+        model(second_input)
+
+    stats = model.get_cache_stats()
+    assert stats["indicator_source"] == "raw_vision_latents"
+    assert stats["gate_trace"] == [True, True]
+    assert stats["actual_full_executions"] == 2
+    assert stats["actual_reuses"] == 0
+    assert model.layers[0].indicator_norm.calls == 0
+
+
+@torch.no_grad()
+def test_sea_cache_residual_order_one_uses_actual_full_step_history():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 4}
+    config = _make_config(runtime, residual_order=1, threshold=0.0)
+    model = DummySeaTransformer()
+    model.enable_cache(config)
+
+    with model.cache_context("cond"):
+        model(torch.ones(1, 1, 1, 1))
+
+    runtime.update(step=1, sigma=0.7)
+    with model.cache_context("cond"):
+        model(torch.full((1, 1, 1, 1), 2.0))
+
+    config.threshold = 100.0
+    runtime.update(step=2, sigma=0.5)
+    with model.cache_context("cond"):
+        _, output = model(torch.full((1, 1, 1, 1), 3.0))
+
+    # Full residuals are 7 and 14 at steps 0 and 1, so linear extrapolation predicts 21 at step 2.
+    torch.testing.assert_close(output, torch.full_like(output, 24.0))
+    root_hook = _get_root_hook(model)
+    assert root_hook.num_full_steps == 2
+    assert root_hook.num_cached_steps == 1
+
+
+@torch.no_grad()
+def test_sea_cache_replays_gate_schedule_and_reports_natural_mismatches():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 3}
+    model = DummySeaTransformer()
+    model.enable_cache(
+        _make_config(
+            runtime,
+            threshold=0.0,
+            gate_schedule=(True, False, True),
+        )
+    )
+
+    for step, sigma in enumerate((0.9, 0.6, 0.3)):
+        runtime.update(step=step, sigma=sigma)
+        with model.cache_context("cond"):
+            model(torch.full((1, 1, 1, 1), float(step + 1)))
+
+    stats = model.get_cache_stats()
+    assert stats["gate_trace"] == [True, False, True]
+    assert stats["gate_schedule_replayed"]
+    assert stats["gate_schedule_mismatches"] == 1
+    assert stats["actual_full_executions"] == 2
+    assert stats["actual_reuses"] == 1
+
+
+@torch.no_grad()
+def test_cache_context_registry_is_refreshed_when_cache_is_enabled_after_an_uncached_call():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
+    model = DummySeaTransformer()
+
+    # Pipeline cache contexts may be entered before a cache is enabled (for example, during the baseline).
+    with model.cache_context("baseline"):
+        model(torch.ones(1, 1, 1, 1))
+
+    model.enable_cache(_make_config(runtime))
+    with model.cache_context("cond"):
+        model(torch.ones(1, 1, 1, 1))
+
+    assert model.get_cache_stats()["branch_full_executions"] == {"cond": 1}
+
+
+@torch.no_grad()
+def test_sea_cache_fails_open_without_vision_metadata():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
+    config = _make_config(runtime)
+    config.metadata_callback = lambda module, args, kwargs: None
+    model = DummySeaTransformer()
+    model.enable_cache(config)
+
+    with model.cache_context("cond"):
+        model(torch.ones(1, 1, 1, 1))
+    runtime.update(step=1, sigma=0.5)
+    with model.cache_context("cond"):
+        model(torch.ones(1, 1, 1, 1))
+
+    assert [block.calls for block in model.layers] == [2, 2, 2]
+    root_hook = _get_root_hook(model)
+    assert root_hook.num_full_steps == 2
+    assert root_hook.num_cached_steps == 0
+
+
+@torch.no_grad()
+def test_sea_cache_fails_open_for_non_adjacent_steps_and_shape_changes():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 4}
+    model = DummySeaTransformer()
+    model.enable_cache(_make_config(runtime, residual_order=1))
+
+    with model.cache_context("cond"):
+        model(torch.ones(1, 1, 1, 1))
+
+    runtime.update(step=2, sigma=0.5)
+    with model.cache_context("cond"):
+        model(torch.ones(1, 2, 1, 1))
+
+    runtime.update(step=3, sigma=0.2)
+    with model.cache_context("cond"):
+        model(torch.ones(1, 2, 1, 1))
+
+    root_hook = _get_root_hook(model)
+    assert root_hook.num_full_steps == 3
+    assert root_hook.num_cached_steps == 0
+    assert [block.calls for block in model.layers] == [3, 3, 3]
+
+
+def test_sea_cache_is_inference_only_and_fails_open_with_autograd():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
+    model = DummySeaTransformer()
+    model.enable_cache(_make_config(runtime))
+
+    with model.cache_context("cond"):
+        model(torch.ones(1, 1, 1, 1, requires_grad=True))
+    runtime.update(step=1, sigma=0.5)
+    with model.cache_context("cond"):
+        model(torch.ones(1, 1, 1, 1, requires_grad=True))
+
+    root_hook = _get_root_hook(model)
+    assert root_hook.num_full_steps == 2
+    assert root_hook.num_cached_steps == 0
+
+
+@torch.no_grad()
+def test_sea_filter_density_normalizes_gain_to_unit_mean():
+    impulse = torch.zeros(2, 3, 4, 1)
+    impulse[0, 0, 0, 0] = 1
+
+    filtered = _apply_sea_filter(impulse, sigma=0.5, power_exp=3.0)
+    recovered_gain = torch.fft.fftn(filtered.float(), dim=(0, 1, 2))
+
+    torch.testing.assert_close(
+        recovered_gain.real.mean(),
+        torch.tensor(1.0),
+        atol=1e-5,
+        rtol=1e-5,
+    )
+    torch.testing.assert_close(
+        recovered_gain.imag,
+        torch.zeros_like(recovered_gain.imag),
+        atol=1e-5,
+        rtol=0,
+    )
+
+
+@pytest.mark.parametrize("sigma", [0.0, 1.0])
+def test_sea_filter_is_finite_at_scheduler_endpoints(sigma):
+    filtered = _apply_sea_filter(torch.randn(2, 2, 2, 4), sigma=sigma, power_exp=2.0)
+
+    assert torch.isfinite(filtered).all()
+    assert filtered.abs().sum() > 0
+
+
+@torch.no_grad()
+def test_sea_cache_single_block_supports_full_and_cached_execution_then_disables_cleanly():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
+    model = DummySeaTransformer(num_layers=1)
+    model.enable_cache(_make_config(runtime))
+
+    with model.cache_context("cond"):
+        _, first_output = model(torch.ones(1, 1, 1, 1))
+    torch.testing.assert_close(first_output, torch.full_like(first_output, 2.0))
+
+    runtime.update(step=1, sigma=0.5)
+    with model.cache_context("cond"):
+        _, cached_output = model(torch.full((1, 1, 1, 1), 2.0))
+    torch.testing.assert_close(cached_output, torch.full_like(cached_output, 3.0))
+    assert model.layers[0].calls == 1
+    assert model.get_cache_stats()["actual_full_executions"] == 1
+    assert model.get_cache_stats()["actual_reuses"] == 1
+
+    model.disable_cache()
+
+    _, output = model(torch.ones(1, 1, 1, 1))
+
+    torch.testing.assert_close(output, torch.full_like(output, 2.0))
+    assert model.layers[0].calls == 2
+    assert model.layers[0]._diffusers_hook.hooks == {}
+
+
+@torch.no_grad()
+def test_sea_cache_fails_open_for_parameter_sharded_blocks():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
+    model = DummySeaTransformer()
+    model.layers[0]._get_fsdp_state = lambda: object()
+    model.enable_cache(_make_config(runtime))
+
+    with model.cache_context("cond"):
+        model(torch.ones(1, 1, 1, 1))
+    runtime.update(step=1, sigma=0.5)
+    with model.cache_context("cond"):
+        model(torch.ones(1, 1, 1, 1))
+
+    stats = model.get_cache_stats()
+    assert stats["actual_full_executions"] == 2
+    assert stats["actual_reuses"] == 0
+    assert stats["fail_open_calls"] == 2
+    assert [block.calls for block in model.layers] == [2, 2, 2]
+
+
+def test_sea_cache_failed_enable_rolls_back_only_new_hooks():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
+    model = DummySeaTransformer()
+    existing_hook = ModelHook()
+    middle_registry = HookRegistry.check_if_exists_or_initialize(model.layers[1])
+    middle_registry.register_hook(existing_hook, _SEA_CACHE_BLOCK_HOOK)
+
+    with pytest.raises(ValueError, match="already exists"):
+        model.enable_cache(_make_config(runtime))
+
+    assert not model.is_cache_enabled
+    assert model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) is None
+    assert model.layers[0]._diffusers_hook.get_hook(_SEA_CACHE_LEADER_BLOCK_HOOK) is None
+    assert middle_registry.get_hook(_SEA_CACHE_BLOCK_HOOK) is existing_hook
+    assert not hasattr(model.layers[2], "_diffusers_hook")
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "message"),
+    [
+        ({"threshold": -1.0}, "threshold"),
+        ({"residual_order": 2}, "residual_order"),
+        ({"retention_steps": -1}, "retention_steps"),
+        ({"cache_end_steps": -1}, "cache_end_steps"),
+        ({"power_exp": 0.0}, "power_exp"),
+        ({"indicator_source": "raw"}, "indicator_source"),
+        ({"threshold": float("nan")}, "threshold"),
+        ({"power_exp": float("inf")}, "power_exp"),
+    ],
+)
+def test_sea_cache_config_validation(kwargs, message):
+    with pytest.raises(ValueError, match=message):
+        SeaCacheConfig(**kwargs)
+
+
+def test_sea_cache_gate_schedule_validation():
+    with pytest.raises(TypeError, match="gate_schedule"):
+        SeaCacheConfig(gate_schedule=(True, 1))
+
+
+@pytest.mark.parametrize("callback_name", ["metadata_callback", "raw_vision_callback"])
+def test_sea_cache_callback_validation(callback_name):
+    with pytest.raises(TypeError, match=callback_name):
+        SeaCacheConfig(**{callback_name: 1})
diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py
index 0ead8cb1a6bb..977ea6b52a1a 100644
--- a/tests/models/transformers/test_models_transformer_cosmos3.py
+++ b/tests/models/transformers/test_models_transformer_cosmos3.py
@@ -16,11 +16,15 @@
 import pytest
 import torch
 
-from diffusers import Cosmos3OmniTransformer
+from diffusers import Cosmos3OmniTransformer, SeaCacheConfig
+from diffusers.hooks._helpers import TransformerBlockRegistry
+from diffusers.hooks.sea_cache import _SEA_CACHE_ROOT_HOOK
+from diffusers.models.cache_utils import CacheMixin
 from diffusers.models.transformers.transformer_cosmos3 import (
     Cosmos3NemotronRMSNorm,
     Cosmos3OmniTransformerOutput,
     Cosmos3PackedMoTAttention,
+    Cosmos3VLTextMoTDecoderLayer,
 )
 from diffusers.utils.torch_utils import randn_tensor
 
@@ -103,6 +107,142 @@ def output_shape(self) -> tuple[int, ...]:
 
 
 class TestCosmos3OmniTransformerModel(Cosmos3OmniTransformerTesterConfig, ModelTesterMixin):
+    @pytest.mark.parametrize("indicator_source", ["first_block", "raw_vision_latents"])
+    def test_cosmos3_supports_sea_cache_without_changing_state_dict_keys(self, indicator_source):
+        model = self.model_class(**self.get_init_dict()).to(torch_device).eval()
+        state_dict_keys = set(model.state_dict())
+        norm_calls = {"und": 0, "gen": 0}
+        original_und_norm_forward = model.norm.forward
+        original_gen_norm_forward = model.norm_moe_gen.forward
+
+        def counted_und_norm_forward(*args, **kwargs):
+            norm_calls["und"] += 1
+            return original_und_norm_forward(*args, **kwargs)
+
+        def counted_gen_norm_forward(*args, **kwargs):
+            norm_calls["gen"] += 1
+            return original_gen_norm_forward(*args, **kwargs)
+
+        model.norm.forward = counted_und_norm_forward
+        model.norm_moe_gen.forward = counted_gen_norm_forward
+        runtime = {"step": 0, "sigma": 0.9, "num_steps": 3}
+        config = SeaCacheConfig(
+            threshold=100.0,
+            cache_end_steps=0,
+            indicator_source=indicator_source,
+            current_step_callback=lambda: runtime["step"],
+            current_sigma_callback=lambda: runtime["sigma"],
+            num_inference_steps_callback=lambda: runtime["num_steps"],
+        )
+
+        assert isinstance(model, CacheMixin)
+        model.enable_cache(config)
+        assert set(model.state_dict()) == state_dict_keys
+
+        inputs = self.get_dummy_inputs()
+        with torch.no_grad(), model.cache_context("cond"):
+            model(**inputs)
+        assert norm_calls == {"und": 1, "gen": 1}
+
+        runtime.update(step=1, sigma=0.6)
+        cached_inputs = self.get_dummy_inputs()
+        cached_inputs["vision_tokens"] = [cached_inputs["vision_tokens"][0] + 0.1]
+        with torch.no_grad(), model.cache_context("cond"):
+            output = model(**cached_inputs)
+
+        root_hook = model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK)
+        # A hit bypasses both final pathway normalizations but still runs the prediction heads.
+        assert norm_calls == {"und": 1, "gen": 1}
+        assert output.sample[0].shape == self.output_shape
+        assert root_hook.num_full_steps == 1
+        assert root_hook.num_cached_steps == 1
+        stats = model.get_cache_stats()
+        assert stats["indicator_source"] == indicator_source
+        assert stats["residual_boundary"] == "post_language_model_norm"
+
+        model.disable_cache()
+        assert set(model.state_dict()) == state_dict_keys
+        with torch.no_grad():
+            model(**self.get_dummy_inputs())
+        assert norm_calls == {"und": 2, "gen": 2}
+
+    def test_cosmos3_sea_cache_post_norm_boundary_numeric_semantics(self):
+        model = self.model_class(**self.get_init_dict()).to(torch_device).eval()
+        runtime = {"step": 0, "sigma": 0.9, "num_steps": 3}
+        config = SeaCacheConfig(
+            threshold=100.0,
+            cache_end_steps=0,
+            current_step_callback=lambda: runtime["step"],
+            current_sigma_callback=lambda: runtime["sigma"],
+            num_inference_steps_callback=lambda: runtime["num_steps"],
+        )
+        first_block_gen_inputs = []
+        returned_und_outputs = []
+        returned_gen_outputs = []
+        projection_head_inputs = []
+        projection_head_calls = 0
+
+        def capture_first_block_gen_input(module, args):
+            first_block_gen_inputs.append(args[0].detach().clone())
+
+        def capture_und_output(module, args, output):
+            returned_und_outputs.append(output.detach().clone())
+
+        def capture_gen_output(module, args, output):
+            returned_gen_outputs.append(output.detach().clone())
+
+        original_projection_forward = model.proj_out.forward
+
+        def counted_projection_forward(hidden_states):
+            nonlocal projection_head_calls
+            projection_head_calls += 1
+            projection_head_inputs.append(hidden_states.detach().clone())
+            return original_projection_forward(hidden_states)
+
+        model.layers[0].input_layernorm_moe_gen.register_forward_pre_hook(capture_first_block_gen_input)
+        model.norm.register_forward_hook(capture_und_output)
+        model.norm_moe_gen.register_forward_hook(capture_gen_output)
+        model.proj_out.forward = counted_projection_forward
+        model.enable_cache(config)
+
+        with torch.no_grad(), model.cache_context("cond"):
+            model(**self.get_dummy_inputs())
+
+        root_hook = model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK)
+        state = root_hook.state_manager._state_cache["cond"]
+        cached_step, cached_und_output, cached_gen_residual = state.history[-1]
+        assert cached_step == 0
+        torch.testing.assert_close(cached_und_output, returned_und_outputs[0])
+        torch.testing.assert_close(
+            cached_gen_residual,
+            returned_gen_outputs[0] - first_block_gen_inputs[0],
+        )
+
+        runtime.update(step=1, sigma=0.6)
+        cached_inputs = self.get_dummy_inputs()
+        cached_inputs["vision_tokens"] = [cached_inputs["vision_tokens"][0] + 0.1]
+        with torch.no_grad(), model.cache_context("cond"):
+            model(**cached_inputs)
+
+        torch.testing.assert_close(returned_und_outputs[1], cached_und_output)
+        torch.testing.assert_close(
+            returned_gen_outputs[1],
+            first_block_gen_inputs[-1] + cached_gen_residual,
+        )
+        torch.testing.assert_close(projection_head_inputs[1], returned_gen_outputs[1])
+        assert projection_head_calls == 2
+        assert root_hook.num_full_steps == 1
+        assert root_hook.num_cached_steps == 1
+
+    def test_cosmos3_decoder_layer_cache_metadata_tracks_generation_stream(self):
+        metadata = TransformerBlockRegistry.get(Cosmos3VLTextMoTDecoderLayer)
+
+        assert metadata.return_hidden_states_index == 1
+        assert metadata.return_encoder_hidden_states_index == 0
+        assert metadata.hidden_states_argument_name == "gen_seq"
+        assert metadata.encoder_hidden_states_argument_name == "und_seq"
+        assert metadata.hidden_states_norm_module_name == "input_layernorm_moe_gen"
+
     def test_output_format(self):
         model = self.model_class(**self.get_init_dict()).to(torch_device).eval()
 

From 12f66f6d8814cc0e2bf4577cd860a30191512b20 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Mon, 31 Aug 2026 04:16:09 -0700
Subject: [PATCH 04/14] Cap consecutive SeaCache reuses

---
 src/diffusers/hooks/sea_cache.py | 41 +++++++++++++++++++++++++++++---
 tests/hooks/test_sea_cache.py    | 41 ++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py
index c5d243a32721..264ba4545059 100644
--- a/src/diffusers/hooks/sea_cache.py
+++ b/src/diffusers/hooks/sea_cache.py
@@ -47,7 +47,7 @@ class SeaCacheConfig:
     prediction heads still execute.
 
     Args:
-        threshold (`float`, defaults to `0.35`):
+        threshold (`float`, defaults to `0.25`):
             Accumulated relative-L1 budget. Larger values reuse the cache more often.
         residual_order (`int`, defaults to `0`):
             Order used to predict the generation-stream language-model residual. `0` directly reuses the most recent
@@ -56,6 +56,9 @@ class SeaCacheConfig:
             Number of initial scheduler steps that always execute in full.
         cache_end_steps (`int`, defaults to `1`):
             Number of final scheduler steps that always execute in full.
+        max_consecutive_cached (`int`, defaults to `2`):
+            Maximum consecutive residual reuses per cache context before forcing a full execution. `0` disables the
+            limit.
         power_exp (`float`, defaults to `3.0`):
             Exponent of the SEA clean-signal power prior. SeaCache uses `3.0` for video features.
         indicator_source (`str`, defaults to `"first_block"`):
@@ -96,10 +99,11 @@ class SeaCacheConfig:
         ```
     """
 
-    threshold: float = 0.35
+    threshold: float = 0.25
     residual_order: int = 0
     retention_steps: int = 1
     cache_end_steps: int = 1
+    max_consecutive_cached: int = 2
     power_exp: float = 3.0
     indicator_source: Literal["first_block", "raw_vision_latents"] = "first_block"
     current_step_callback: Callable[[], int] = None
@@ -124,6 +128,14 @@ def __post_init__(self):
             raise ValueError(f"`retention_steps` must be non-negative, got {self.retention_steps}.")
         if self.cache_end_steps < 0:
             raise ValueError(f"`cache_end_steps` must be non-negative, got {self.cache_end_steps}.")
+        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 positive, got {self.power_exp}.")
         if self.indicator_source not in ("first_block", "raw_vision_latents"):
@@ -162,6 +174,9 @@ def __init__(self):
         self.gate_should_compute = True
         self.previous_indicator: list[torch.Tensor] | None = None
         self.accumulated_distance = 0.0
+        self.consecutive_cached = 0
+        self.max_consecutive_cached_observed = 0
+        self.max_consecutive_forced_full = 0
         self.skip_remaining = False
         self.full_execution_pending = False
         self.cacheable_execution = False
@@ -191,6 +206,9 @@ def reset(self):
         self.gate_should_compute = True
         self.previous_indicator = None
         self.accumulated_distance = 0.0
+        self.consecutive_cached = 0
+        self.max_consecutive_cached_observed = 0
+        self.max_consecutive_forced_full = 0
         self.reset_forward()
 
 
@@ -243,8 +261,11 @@ def resolve_gate(
         is_retained = metadata.step_index < config.retention_steps
         is_in_cache_end = metadata.step_index >= metadata.num_inference_steps - config.cache_end_steps
         is_first_observation = state.previous_indicator is None
+        is_max_consecutive = bool(
+            config.max_consecutive_cached and state.consecutive_cached >= config.max_consecutive_cached
+        )
         invalid_gate = is_non_adjacent or indicator is None
-        forced_compute = invalid_gate or is_retained or is_in_cache_end or is_first_observation
+        forced_compute = invalid_gate or is_retained or is_in_cache_end or is_first_observation or is_max_consecutive
         candidate_accumulated_distance = 0.0
 
         if forced_compute:
@@ -294,6 +315,10 @@ def resolve_gate(
                 self.mark_fail_open("SeaCache gate schedule requested an unsafe cache hit; running full.")
                 should_compute = True
 
+        if is_max_consecutive:
+            should_compute = True
+            state.max_consecutive_forced_full += 1
+
         state.accumulated_distance = 0.0 if should_compute else candidate_accumulated_distance
 
         state.gate_key = gate_key
@@ -356,6 +381,7 @@ def _record_full_execution(
     und_output: torch.Tensor | None,
 ) -> None:
     shared_state.num_full_steps += 1
+    state.consecutive_cached = 0
     branch = state_manager._current_context
     shared_state.branch_full_executions[branch] = shared_state.branch_full_executions.get(branch, 0) + 1
     if state.full_started_at is not None:
@@ -579,6 +605,8 @@ def stats(self) -> dict[str, Any]:
             per_branch[branch] = {
                 "full_calls": self.shared_state.branch_full_executions.get(branch, 0),
                 "reuse_calls": self.shared_state.branch_reuses.get(branch, 0),
+                "max_consecutive_cached_observed": state.max_consecutive_cached_observed,
+                "max_consecutive_forced_full": state.max_consecutive_forced_full,
             }
 
         opportunities = self.shared_state.num_full_steps + self.shared_state.num_cached_steps
@@ -586,6 +614,11 @@ def stats(self) -> dict[str, Any]:
             "indicator_source": self.config.indicator_source,
             "residual_order": self.config.residual_order,
             "residual_boundary": self.residual_boundary,
+            "max_consecutive_cached": self.config.max_consecutive_cached,
+            "max_consecutive_cached_observed": max(
+                (stats["max_consecutive_cached_observed"] for stats in per_branch.values()), default=0
+            ),
+            "max_consecutive_forced_full": sum(stats["max_consecutive_forced_full"] for stats in per_branch.values()),
             "transformer_calls": self.shared_state.transformer_calls,
             "gate_evaluations": self.shared_state.gate_evaluations,
             "gate_full_decisions": self.shared_state.gate_full_decisions,
@@ -870,6 +903,8 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
         state.full_execution_pending = False
         state.cached_und_output = cached_und
         state.cached_gen_residual = cached_residual
+        state.consecutive_cached += 1
+        state.max_consecutive_cached_observed = max(state.max_consecutive_cached_observed, state.consecutive_cached)
         self.shared_state.num_cached_steps += 1
         branch = self.state_manager._current_context
         self.shared_state.branch_reuses[branch] = self.shared_state.branch_reuses.get(branch, 0) + 1
diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py
index 85db1fbc22c1..d6b3798d548a 100644
--- a/tests/hooks/test_sea_cache.py
+++ b/tests/hooks/test_sea_cache.py
@@ -157,6 +157,37 @@ def test_sea_cache_uses_independent_gates_and_histories_per_context():
     assert not model.is_cache_enabled
 
 
+@torch.no_grad()
+def test_sea_cache_max_consecutive_cached_forces_full_per_context():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 8}
+    model = DummySeaTransformer()
+    model.enable_cache(
+        _make_config(
+            runtime,
+            threshold=100.0,
+            retention_steps=0,
+            cache_end_steps=0,
+            max_consecutive_cached=2,
+        )
+    )
+
+    for step in range(runtime["num_steps"]):
+        runtime.update(step=step, sigma=0.9 - step * 0.1)
+        with model.cache_context("cond"):
+            model(torch.ones(1, 1, 1, 1))
+
+    stats = model.get_cache_stats()
+    assert stats["gate_trace"] == [True, False, False, True, False, False, True, False]
+    assert stats["actual_full_executions"] == 3
+    assert stats["actual_reuses"] == 5
+    assert stats["max_consecutive_cached"] == 2
+    assert stats["max_consecutive_cached_observed"] == 2
+    assert stats["max_consecutive_forced_full"] == 2
+    assert stats["per_branch"]["cond"]["max_consecutive_cached_observed"] == 2
+    assert stats["per_branch"]["cond"]["max_consecutive_forced_full"] == 2
+    assert [block.calls for block in model.layers] == [3, 3, 3]
+
+
 @torch.no_grad()
 def test_sea_cache_raw_vision_indicator_includes_conditioning_frames():
     runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
@@ -414,6 +445,13 @@ def test_sea_cache_failed_enable_rolls_back_only_new_hooks():
     assert not hasattr(model.layers[2], "_diffusers_hook")
 
 
+def test_sea_cache_config_defaults():
+    config = SeaCacheConfig()
+
+    assert config.threshold == 0.25
+    assert config.max_consecutive_cached == 2
+
+
 @pytest.mark.parametrize(
     ("kwargs", "message"),
     [
@@ -421,6 +459,9 @@ def test_sea_cache_failed_enable_rolls_back_only_new_hooks():
         ({"residual_order": 2}, "residual_order"),
         ({"retention_steps": -1}, "retention_steps"),
         ({"cache_end_steps": -1}, "cache_end_steps"),
+        ({"max_consecutive_cached": -1}, "max_consecutive_cached"),
+        ({"max_consecutive_cached": 1.5}, "max_consecutive_cached"),
+        ({"max_consecutive_cached": True}, "max_consecutive_cached"),
         ({"power_exp": 0.0}, "power_exp"),
         ({"indicator_source": "raw"}, "indicator_source"),
         ({"threshold": float("nan")}, "threshold"),

From 54c4cc5d46f4b4b09e90696cd7c7c05af8c21768 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Mon, 31 Aug 2026 04:47:29 -0700
Subject: [PATCH 05/14] Align SeaCache defaults with Cosmos3 inference

---
 src/diffusers/__init__.py                     |   2 -
 src/diffusers/hooks/__init__.py               |   2 +-
 src/diffusers/hooks/sea_cache.py              | 186 ++----------------
 src/diffusers/models/cache_utils.py           |  11 --
 src/diffusers/utils/dummy_pt_objects.py       |   4 -
 tests/hooks/test_sea_cache.py                 |  95 +--------
 .../test_models_transformer_cosmos3.py        |  12 +-
 7 files changed, 24 insertions(+), 288 deletions(-)

diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py
index 37ea9c1db39f..68789af008f5 100644
--- a/src/diffusers/__init__.py
+++ b/src/diffusers/__init__.py
@@ -216,7 +216,6 @@
             "apply_sea_cache",
             "apply_taylorseer_cache",
             "apply_text_kv_cache",
-            "get_sea_cache_stats",
         ]
     )
     _import_structure["image_processor"] = [
@@ -1096,7 +1095,6 @@
             apply_sea_cache,
             apply_taylorseer_cache,
             apply_text_kv_cache,
-            get_sea_cache_stats,
         )
         from .image_processor import (
             InpaintProcessor,
diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py
index 2e9c7395c2b1..70399f9ce805 100644
--- a/src/diffusers/hooks/__init__.py
+++ b/src/diffusers/hooks/__init__.py
@@ -25,7 +25,7 @@
     from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook
     from .mag_cache import MagCacheConfig, apply_mag_cache
     from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
-    from .sea_cache import SeaCacheConfig, apply_sea_cache, get_sea_cache_stats
+    from .sea_cache import SeaCacheConfig, apply_sea_cache
     from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig
     from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache
     from .tensor_parallel import apply_tensor_parallel
diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py
index 264ba4545059..827648709518 100644
--- a/src/diffusers/hooks/sea_cache.py
+++ b/src/diffusers/hooks/sea_cache.py
@@ -12,12 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import copy
 import inspect
 import math
-import time
 from dataclasses import dataclass
-from typing import Any, Callable, Literal, Sequence
+from typing import Any, Callable, Literal
 
 import torch
 
@@ -49,7 +47,7 @@ class SeaCacheConfig:
     Args:
         threshold (`float`, defaults to `0.25`):
             Accumulated relative-L1 budget. Larger values reuse the cache more often.
-        residual_order (`int`, defaults to `0`):
+        residual_order (`int`, defaults to `1`):
             Order used to predict the generation-stream language-model residual. `0` directly reuses the most recent
             residual and `1` linearly extrapolates from the two most recent full executions.
         retention_steps (`int`, defaults to `1`):
@@ -61,12 +59,11 @@ class SeaCacheConfig:
             limit.
         power_exp (`float`, defaults to `3.0`):
             Exponent of the SEA clean-signal power prior. SeaCache uses `3.0` for video features.
-        indicator_source (`str`, defaults to `"first_block"`):
-            Feature source used to construct the SEA indicator. `"first_block"` follows the published method and
-            filters the timestep-modulated pre-attention input of the first transformer block.
-            `"raw_vision_latents"` is an opt-in lower-cost approximation that filters the complete raw vision latent,
-            including clean conditioning frames in I2V. Thresholds are not generally transferable between the two
-            sources.
+        indicator_source (`str`, defaults to `"raw_vision_latents"`):
+            Feature source used to construct the SEA indicator. `"raw_vision_latents"` filters the complete raw vision
+            latent, including clean conditioning frames in I2V. `"first_block"` filters the timestep-modulated
+            pre-attention input of the first transformer block. Thresholds are not generally transferable between the
+            two sources.
         current_step_callback (`Callable[[], int]`):
             Callback returning the current scheduler step index.
         current_sigma_callback (`Callable[[], float]`):
@@ -80,10 +77,6 @@ class SeaCacheConfig:
         raw_vision_callback (`Callable`, *optional*):
             Advanced model adapter returning raw vision latents with shape `(C, T, H, W)` for the
             `"raw_vision_latents"` indicator. Cosmos 3 uses its native adapter when this is omitted.
-        gate_schedule (`Sequence[bool]`, *optional*):
-            Per-step full-compute decisions to replay for controlled residual-prediction ablations. SEA still evaluates
-            its natural decision and reports mismatches; invalid or unsafe calls always fail open to full compute.
-
     Example:
         ```python
         >>> from diffusers import Cosmos3OmniPipeline, SeaCacheConfig
@@ -100,12 +93,12 @@ class SeaCacheConfig:
     """
 
     threshold: float = 0.25
-    residual_order: int = 0
+    residual_order: int = 1
     retention_steps: int = 1
     cache_end_steps: int = 1
     max_consecutive_cached: int = 2
     power_exp: float = 3.0
-    indicator_source: Literal["first_block", "raw_vision_latents"] = "first_block"
+    indicator_source: Literal["first_block", "raw_vision_latents"] = "raw_vision_latents"
     current_step_callback: Callable[[], int] = None
     current_sigma_callback: Callable[[], float] = None
     num_inference_steps_callback: Callable[[], int] = None
@@ -117,7 +110,6 @@ class SeaCacheConfig:
         [torch.nn.Module, tuple[Any, ...], dict[str, Any]],
         list[torch.Tensor] | None,
     ] = None
-    gate_schedule: Sequence[bool] = None
 
     def __post_init__(self):
         if not math.isfinite(self.threshold) or self.threshold < 0:
@@ -152,10 +144,6 @@ def __post_init__(self):
             callback = getattr(self, name)
             if callback is not None and not callable(callback):
                 raise TypeError(f"`{name}` must be callable or `None`.")
-        if self.gate_schedule is not None:
-            self.gate_schedule = tuple(self.gate_schedule)
-            if any(not isinstance(value, bool) for value in self.gate_schedule):
-                raise TypeError("`gate_schedule` must contain only boolean values.")
 
 
 @dataclass
@@ -175,8 +163,6 @@ def __init__(self):
         self.previous_indicator: list[torch.Tensor] | None = None
         self.accumulated_distance = 0.0
         self.consecutive_cached = 0
-        self.max_consecutive_cached_observed = 0
-        self.max_consecutive_forced_full = 0
         self.skip_remaining = False
         self.full_execution_pending = False
         self.cacheable_execution = False
@@ -186,7 +172,6 @@ def __init__(self):
         self.und_output: torch.Tensor | None = None
         self.cached_und_output: torch.Tensor | None = None
         self.cached_gen_residual: torch.Tensor | None = None
-        self.full_started_at: float | None = None
 
     def reset_forward(self):
         self.skip_remaining = False
@@ -198,7 +183,6 @@ def reset_forward(self):
         self.und_output = None
         self.cached_und_output = None
         self.cached_gen_residual = None
-        self.full_started_at = None
 
     def reset(self):
         self.history = []
@@ -207,8 +191,6 @@ def reset(self):
         self.previous_indicator = None
         self.accumulated_distance = 0.0
         self.consecutive_cached = 0
-        self.max_consecutive_cached_observed = 0
-        self.max_consecutive_forced_full = 0
         self.reset_forward()
 
 
@@ -219,20 +201,6 @@ def __init__(self):
 
     def reset(self):
         self.forward_metadata: _SeaCacheForwardMetadata | None = None
-        self.transformer_calls = 0
-        self.gate_evaluations = 0
-        self.gate_full_decisions = 0
-        self.gate_skip_decisions = 0
-        self.gate_trace: list[bool] = []
-        self.gate_schedule_mismatches = 0
-        self.num_full_steps = 0
-        self.num_cached_steps = 0
-        self.fail_open_calls = 0
-        self.indicator_seconds = 0.0
-        self.decision_seconds = 0.0
-        self.full_seconds = 0.0
-        self.branch_full_executions: dict[str, int] = {}
-        self.branch_reuses: dict[str, int] = {}
 
     def warn_once(self, message: str):
         if message not in self._warned_messages:
@@ -240,7 +208,6 @@ def warn_once(self, message: str):
             self._warned_messages.add(message)
 
     def mark_fail_open(self, message: str):
-        self.fail_open_calls += 1
         self.warn_once(message)
 
     def resolve_gate(
@@ -254,7 +221,6 @@ def resolve_gate(
         if state.gate_key == gate_key:
             return state.gate_should_compute
 
-        self.gate_evaluations += 1
         is_non_adjacent = state.gate_key is not None and metadata.step_index != state.gate_key[0] + 1
         if is_non_adjacent:
             self.mark_fail_open("SeaCache received non-adjacent scheduler steps; running full.")
@@ -269,7 +235,7 @@ def resolve_gate(
         candidate_accumulated_distance = 0.0
 
         if forced_compute:
-            natural_should_compute = True
+            should_compute = True
         else:
             if len(indicator) != len(state.previous_indicator) or not indicator:
                 distance = float("inf")
@@ -294,41 +260,13 @@ def resolve_gate(
                 invalid_gate = True
                 self.mark_fail_open("SeaCache indicator history changed shape, device, or dtype; running full.")
             candidate_accumulated_distance = state.accumulated_distance + distance
-            natural_should_compute = invalid_gate or candidate_accumulated_distance >= config.threshold
-
-        should_compute = natural_should_compute
-
-        if config.gate_schedule is not None:
-            schedule_is_valid = len(config.gate_schedule) == metadata.num_inference_steps
-            scheduled_compute = bool(config.gate_schedule[metadata.step_index]) if schedule_is_valid else True
-            can_replay_skip = not (invalid_gate or is_retained or is_in_cache_end or is_first_observation)
-            if not schedule_is_valid:
-                self.mark_fail_open(
-                    "SeaCache gate schedule length does not match the number of inference steps; running full."
-                )
-                should_compute = True
-            elif scheduled_compute or can_replay_skip:
-                if natural_should_compute != scheduled_compute:
-                    self.gate_schedule_mismatches += 1
-                should_compute = scheduled_compute
-            else:
-                self.mark_fail_open("SeaCache gate schedule requested an unsafe cache hit; running full.")
-                should_compute = True
-
-        if is_max_consecutive:
-            should_compute = True
-            state.max_consecutive_forced_full += 1
+            should_compute = invalid_gate or candidate_accumulated_distance >= config.threshold
 
         state.accumulated_distance = 0.0 if should_compute else candidate_accumulated_distance
 
         state.gate_key = gate_key
         state.gate_should_compute = should_compute
         state.previous_indicator = None if indicator is None else [value.detach() for value in indicator]
-        self.gate_trace.append(should_compute)
-        if should_compute:
-            self.gate_full_decisions += 1
-        else:
-            self.gate_skip_decisions += 1
         return should_compute
 
 
@@ -374,18 +312,11 @@ def _get_block_outputs(
 
 def _record_full_execution(
     config: SeaCacheConfig,
-    state_manager: StateManager,
-    shared_state: SeaCacheSharedState,
     state: SeaCacheContextState,
     gen_output: torch.Tensor,
     und_output: torch.Tensor | None,
 ) -> None:
-    shared_state.num_full_steps += 1
     state.consecutive_cached = 0
-    branch = state_manager._current_context
-    shared_state.branch_full_executions[branch] = shared_state.branch_full_executions.get(branch, 0) + 1
-    if state.full_started_at is not None:
-        shared_state.full_seconds += time.perf_counter() - state.full_started_at
     if (
         state.cacheable_execution
         and state.step_index is not None
@@ -569,7 +500,6 @@ def __init__(
         shared_state: SeaCacheSharedState,
         metadata_callback: Callable,
         raw_vision_callback: Callable,
-        residual_boundary: str,
     ):
         super().__init__()
         self.config = config
@@ -577,73 +507,9 @@ def __init__(
         self.shared_state = shared_state
         self.metadata_callback = metadata_callback
         self.raw_vision_callback = raw_vision_callback
-        self.residual_boundary = residual_boundary
-        self._last_stats: dict[str, Any] | None = None
-
-    @property
-    def num_full_steps(self) -> int:
-        return self.shared_state.num_full_steps
-
-    @property
-    def num_cached_steps(self) -> int:
-        return self.shared_state.num_cached_steps
-
-    def stats(self) -> dict[str, Any]:
-        if self.shared_state.transformer_calls == 0 and self._last_stats is not None:
-            return copy.deepcopy(self._last_stats)
-
-        persistent_cache_bytes = 0
-        per_branch = {}
-        for branch, state in sorted(self.state_manager._state_cache.items()):
-            if state.previous_indicator is not None:
-                persistent_cache_bytes += sum(
-                    value.numel() * value.element_size() for value in state.previous_indicator
-                )
-            for _, und_output, gen_residual in state.history:
-                persistent_cache_bytes += und_output.numel() * und_output.element_size()
-                persistent_cache_bytes += gen_residual.numel() * gen_residual.element_size()
-            per_branch[branch] = {
-                "full_calls": self.shared_state.branch_full_executions.get(branch, 0),
-                "reuse_calls": self.shared_state.branch_reuses.get(branch, 0),
-                "max_consecutive_cached_observed": state.max_consecutive_cached_observed,
-                "max_consecutive_forced_full": state.max_consecutive_forced_full,
-            }
-
-        opportunities = self.shared_state.num_full_steps + self.shared_state.num_cached_steps
-        return {
-            "indicator_source": self.config.indicator_source,
-            "residual_order": self.config.residual_order,
-            "residual_boundary": self.residual_boundary,
-            "max_consecutive_cached": self.config.max_consecutive_cached,
-            "max_consecutive_cached_observed": max(
-                (stats["max_consecutive_cached_observed"] for stats in per_branch.values()), default=0
-            ),
-            "max_consecutive_forced_full": sum(stats["max_consecutive_forced_full"] for stats in per_branch.values()),
-            "transformer_calls": self.shared_state.transformer_calls,
-            "gate_evaluations": self.shared_state.gate_evaluations,
-            "gate_full_decisions": self.shared_state.gate_full_decisions,
-            "gate_skip_decisions": self.shared_state.gate_skip_decisions,
-            "gate_trace": list(self.shared_state.gate_trace),
-            "gate_schedule_replayed": self.config.gate_schedule is not None,
-            "gate_schedule_mismatches": self.shared_state.gate_schedule_mismatches,
-            "actual_full_executions": self.shared_state.num_full_steps,
-            "actual_reuses": self.shared_state.num_cached_steps,
-            "actual_reuse_rate": (self.shared_state.num_cached_steps / opportunities if opportunities else 0.0),
-            "fail_open_calls": self.shared_state.fail_open_calls,
-            "sea_indicator_seconds": self.shared_state.indicator_seconds,
-            "sea_decision_seconds": self.shared_state.decision_seconds,
-            "sea_seconds": (self.shared_state.indicator_seconds + self.shared_state.decision_seconds),
-            "full_seconds": self.shared_state.full_seconds,
-            "timing_note": "host wall time; CUDA work is asynchronous",
-            "persistent_cache_bytes": persistent_cache_bytes,
-            "branch_full_executions": dict(sorted(self.shared_state.branch_full_executions.items())),
-            "branch_reuses": dict(sorted(self.shared_state.branch_reuses.items())),
-            "per_branch": per_branch,
-        }
 
     def pre_forward(self, module: torch.nn.Module, *args, **kwargs):
         self.shared_state.forward_metadata = None
-        self.shared_state.transformer_calls += 1
         if torch.is_grad_enabled():
             self.shared_state.mark_fail_open(
                 "SeaCache is inference-only; calls with autograd enabled run in fail-open mode."
@@ -742,8 +608,6 @@ def post_forward(self, module: torch.nn.Module, output: Any) -> Any:
         return output
 
     def reset_state(self, module: torch.nn.Module):
-        if self.shared_state.transformer_calls > 0:
-            self._last_stats = self.stats()
         self.state_manager.reset()
         self.shared_state.reset()
         return module
@@ -829,13 +693,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
 
         forward_metadata = self.shared_state.forward_metadata
         if state is None or forward_metadata is None:
-            if state is not None:
-                state.full_started_at = time.perf_counter()
             return self.fn_ref.original_forward(*args, **kwargs)
 
         state.step_index = forward_metadata.step_index
         state.cacheable_execution = True
-        indicator_started = time.perf_counter()
         indicator_error_reported = False
         if _is_parameter_sharded(module):
             self.shared_state.mark_fail_open(
@@ -852,14 +713,11 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
                 )
                 indicator = None
                 indicator_error_reported = True
-        self.shared_state.indicator_seconds += time.perf_counter() - indicator_started
         if indicator is None and not indicator_error_reported:
             self.shared_state.mark_fail_open(
                 "SeaCache could not construct its vision indicator; running in fail-open mode."
             )
-        decision_started = time.perf_counter()
         should_compute = self.shared_state.resolve_gate(state, forward_metadata, indicator, self.config)
-        self.shared_state.decision_seconds += time.perf_counter() - decision_started
 
         if should_compute or not state.history:
             if not should_compute:
@@ -867,7 +725,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
                 self.shared_state.mark_fail_open(
                     "SeaCache selected a cache hit without residual history; running in fail-open mode."
                 )
-            state.full_started_at = time.perf_counter()
             return self.fn_ref.original_forward(*args, **kwargs)
 
         residual_history = state.history[-(self.config.residual_order + 1) :]
@@ -889,7 +746,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
             self.shared_state.mark_fail_open(
                 "SeaCache residual history changed shape, device, or dtype; running in fail-open mode."
             )
-            state.full_started_at = time.perf_counter()
             return self.fn_ref.original_forward(*args, **kwargs)
 
         if self.config.residual_order == 1 and len(residual_history) >= 2:
@@ -904,10 +760,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
         state.cached_und_output = cached_und
         state.cached_gen_residual = cached_residual
         state.consecutive_cached += 1
-        state.max_consecutive_cached_observed = max(state.max_consecutive_cached_observed, state.consecutive_cached)
-        self.shared_state.num_cached_steps += 1
-        branch = self.state_manager._current_context
-        self.shared_state.branch_reuses[branch] = self.shared_state.branch_reuses.get(branch, 0) + 1
         if self.post_norm_boundary:
             return _build_block_output(self._metadata, hidden_states, encoder_hidden_states)
         return _build_block_output(self._metadata, hidden_states + cached_residual, cached_und)
@@ -952,8 +804,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
         hidden_states, encoder_hidden_states = _get_block_outputs(self._metadata, output)
         _record_full_execution(
             self.config,
-            self.state_manager,
-            self.shared_state,
             state,
             gen_output=hidden_states,
             und_output=encoder_hidden_states,
@@ -1006,8 +856,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
 
         _record_full_execution(
             self.config,
-            self.state_manager,
-            self.shared_state,
             state,
             gen_output=output,
             und_output=state.und_output,
@@ -1055,7 +903,6 @@ def apply_sea_cache(module: torch.nn.Module, config: SeaCacheConfig) -> None:
                 "residual boundary."
             )
     post_norm_boundary = post_norm_modules is not None
-    residual_boundary = "post_language_model_norm" if post_norm_boundary else "repeated_block_stack"
 
     blocks = []
     for name, submodule in unwrapped_module.named_children():
@@ -1085,7 +932,6 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None:
                 shared_state,
                 metadata_callback,
                 raw_vision_callback,
-                residual_boundary,
             ),
             _SEA_CACHE_ROOT_HOOK,
         )
@@ -1134,13 +980,3 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None:
         raise
 
     root_registry._child_registries_cache = None
-
-
-def get_sea_cache_stats(module: torch.nn.Module) -> dict[str, Any]:
-    """Return statistics for the SeaCache instance currently attached to ``module``."""
-
-    registry = getattr(module, "_diffusers_hook", None)
-    root_hook = registry.get_hook(_SEA_CACHE_ROOT_HOOK) if registry is not None else None
-    if not isinstance(root_hook, SeaCacheRootHook):
-        raise ValueError("SeaCache is not enabled on this module.")
-    return root_hook.stats()
diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py
index 6f6542df11c2..2d7e0309db3b 100644
--- a/src/diffusers/models/cache_utils.py
+++ b/src/diffusers/models/cache_utils.py
@@ -167,17 +167,6 @@ def disable_cache(self) -> None:
         self._cache_config = None
         registry._child_registries_cache = None
 
-    def get_cache_stats(self) -> dict:
-        """Return statistics for the currently enabled cache implementation."""
-
-        from ..hooks import SeaCacheConfig, get_sea_cache_stats
-
-        if self._cache_config is None:
-            raise ValueError("Caching is not enabled.")
-        if isinstance(self._cache_config, SeaCacheConfig):
-            return get_sea_cache_stats(self)
-        raise NotImplementedError(f"Cache statistics are not available for {type(self._cache_config).__name__}.")
-
     def _reset_stateful_cache(self, recurse: bool = True) -> None:
         from ..hooks import HookRegistry
 
diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py
index 324c92f3d507..94700aa2689c 100644
--- a/src/diffusers/utils/dummy_pt_objects.py
+++ b/src/diffusers/utils/dummy_pt_objects.py
@@ -364,10 +364,6 @@ def apply_text_kv_cache(*args, **kwargs):
     requires_backends(apply_text_kv_cache, ["torch"])
 
 
-def get_sea_cache_stats(*args, **kwargs):
-    requires_backends(get_sea_cache_stats, ["torch"])
-
-
 class InpaintProcessor(metaclass=DummyObject):
     _backends = ["torch"]
 
diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py
index d6b3798d548a..d68885420137 100644
--- a/tests/hooks/test_sea_cache.py
+++ b/tests/hooks/test_sea_cache.py
@@ -90,6 +90,7 @@ def _make_config(runtime, **kwargs):
         "threshold": 100.0,
         "retention_steps": 1,
         "cache_end_steps": 0,
+        "indicator_source": "first_block",
         "current_step_callback": lambda: runtime["step"],
         "current_sigma_callback": lambda: runtime["sigma"],
         "num_inference_steps_callback": lambda: runtime["num_steps"],
@@ -99,16 +100,11 @@ def _make_config(runtime, **kwargs):
     return SeaCacheConfig(**config_kwargs)
 
 
-def _get_root_hook(model):
-    return model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK)
-
-
 @torch.no_grad()
 def test_sea_cache_uses_independent_gates_and_histories_per_context():
     runtime = {"step": 0, "sigma": 0.9, "num_steps": 4}
     model = DummySeaTransformer()
     model.enable_cache(_make_config(runtime, threshold=0.5))
-    root_hook = _get_root_hook(model)
 
     first_input = torch.ones(2, 2, 2, 1)
     with model.cache_context("cond"):
@@ -133,25 +129,11 @@ def test_sea_cache_uses_independent_gates_and_histories_per_context():
     torch.testing.assert_close(uncond_output, changed_uncond_input * 8)
     assert [block.calls for block in model.layers] == [3, 3, 3]
     assert model.layers[0].indicator_norm.calls == 4
-    assert root_hook.num_full_steps == 3
-    assert root_hook.num_cached_steps == 1
-    stats = model.get_cache_stats()
-    assert stats["indicator_source"] == "first_block"
-    assert stats["residual_order"] == 0
-    assert stats["residual_boundary"] == "repeated_block_stack"
-    assert stats["transformer_calls"] == 4
-    assert stats["gate_evaluations"] == 4
-    assert stats["gate_trace"] == [True, True, False, True]
-    assert stats["branch_full_executions"] == {"cond": 1, "uncond": 2}
-    assert stats["branch_reuses"] == {"cond": 1}
-    assert stats["persistent_cache_bytes"] > 0
 
     model._reset_stateful_cache()
-    assert root_hook.num_full_steps == 0
-    assert root_hook.num_cached_steps == 0
-    archived_stats = model.get_cache_stats()
-    assert archived_stats["transformer_calls"] == 4
-    assert archived_stats["gate_trace"] == [True, True, False, True]
+    with model.cache_context("cond"):
+        model(first_input)
+    assert [block.calls for block in model.layers] == [4, 4, 4]
 
     model.disable_cache()
     assert not model.is_cache_enabled
@@ -176,15 +158,6 @@ def test_sea_cache_max_consecutive_cached_forces_full_per_context():
         with model.cache_context("cond"):
             model(torch.ones(1, 1, 1, 1))
 
-    stats = model.get_cache_stats()
-    assert stats["gate_trace"] == [True, False, False, True, False, False, True, False]
-    assert stats["actual_full_executions"] == 3
-    assert stats["actual_reuses"] == 5
-    assert stats["max_consecutive_cached"] == 2
-    assert stats["max_consecutive_cached_observed"] == 2
-    assert stats["max_consecutive_forced_full"] == 2
-    assert stats["per_branch"]["cond"]["max_consecutive_cached_observed"] == 2
-    assert stats["per_branch"]["cond"]["max_consecutive_forced_full"] == 2
     assert [block.calls for block in model.layers] == [3, 3, 3]
 
 
@@ -217,11 +190,7 @@ def raw_vision(module, args, kwargs):
     with model.cache_context("cond"):
         model(second_input)
 
-    stats = model.get_cache_stats()
-    assert stats["indicator_source"] == "raw_vision_latents"
-    assert stats["gate_trace"] == [True, True]
-    assert stats["actual_full_executions"] == 2
-    assert stats["actual_reuses"] == 0
+    assert [block.calls for block in model.layers] == [2, 2, 2]
     assert model.layers[0].indicator_norm.calls == 0
 
 
@@ -246,34 +215,7 @@ def test_sea_cache_residual_order_one_uses_actual_full_step_history():
 
     # Full residuals are 7 and 14 at steps 0 and 1, so linear extrapolation predicts 21 at step 2.
     torch.testing.assert_close(output, torch.full_like(output, 24.0))
-    root_hook = _get_root_hook(model)
-    assert root_hook.num_full_steps == 2
-    assert root_hook.num_cached_steps == 1
-
-
-@torch.no_grad()
-def test_sea_cache_replays_gate_schedule_and_reports_natural_mismatches():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 3}
-    model = DummySeaTransformer()
-    model.enable_cache(
-        _make_config(
-            runtime,
-            threshold=0.0,
-            gate_schedule=(True, False, True),
-        )
-    )
-
-    for step, sigma in enumerate((0.9, 0.6, 0.3)):
-        runtime.update(step=step, sigma=sigma)
-        with model.cache_context("cond"):
-            model(torch.full((1, 1, 1, 1), float(step + 1)))
-
-    stats = model.get_cache_stats()
-    assert stats["gate_trace"] == [True, False, True]
-    assert stats["gate_schedule_replayed"]
-    assert stats["gate_schedule_mismatches"] == 1
-    assert stats["actual_full_executions"] == 2
-    assert stats["actual_reuses"] == 1
+    assert [block.calls for block in model.layers] == [2, 2, 2]
 
 
 @torch.no_grad()
@@ -289,7 +231,7 @@ def test_cache_context_registry_is_refreshed_when_cache_is_enabled_after_an_unca
     with model.cache_context("cond"):
         model(torch.ones(1, 1, 1, 1))
 
-    assert model.get_cache_stats()["branch_full_executions"] == {"cond": 1}
+    assert model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) is not None
 
 
 @torch.no_grad()
@@ -307,9 +249,6 @@ def test_sea_cache_fails_open_without_vision_metadata():
         model(torch.ones(1, 1, 1, 1))
 
     assert [block.calls for block in model.layers] == [2, 2, 2]
-    root_hook = _get_root_hook(model)
-    assert root_hook.num_full_steps == 2
-    assert root_hook.num_cached_steps == 0
 
 
 @torch.no_grad()
@@ -329,9 +268,6 @@ def test_sea_cache_fails_open_for_non_adjacent_steps_and_shape_changes():
     with model.cache_context("cond"):
         model(torch.ones(1, 2, 1, 1))
 
-    root_hook = _get_root_hook(model)
-    assert root_hook.num_full_steps == 3
-    assert root_hook.num_cached_steps == 0
     assert [block.calls for block in model.layers] == [3, 3, 3]
 
 
@@ -346,9 +282,7 @@ def test_sea_cache_is_inference_only_and_fails_open_with_autograd():
     with model.cache_context("cond"):
         model(torch.ones(1, 1, 1, 1, requires_grad=True))
 
-    root_hook = _get_root_hook(model)
-    assert root_hook.num_full_steps == 2
-    assert root_hook.num_cached_steps == 0
+    assert [block.calls for block in model.layers] == [2, 2, 2]
 
 
 @torch.no_grad()
@@ -396,8 +330,6 @@ def test_sea_cache_single_block_supports_full_and_cached_execution_then_disables
         _, cached_output = model(torch.full((1, 1, 1, 1), 2.0))
     torch.testing.assert_close(cached_output, torch.full_like(cached_output, 3.0))
     assert model.layers[0].calls == 1
-    assert model.get_cache_stats()["actual_full_executions"] == 1
-    assert model.get_cache_stats()["actual_reuses"] == 1
 
     model.disable_cache()
 
@@ -421,10 +353,6 @@ def test_sea_cache_fails_open_for_parameter_sharded_blocks():
     with model.cache_context("cond"):
         model(torch.ones(1, 1, 1, 1))
 
-    stats = model.get_cache_stats()
-    assert stats["actual_full_executions"] == 2
-    assert stats["actual_reuses"] == 0
-    assert stats["fail_open_calls"] == 2
     assert [block.calls for block in model.layers] == [2, 2, 2]
 
 
@@ -449,7 +377,9 @@ def test_sea_cache_config_defaults():
     config = SeaCacheConfig()
 
     assert config.threshold == 0.25
+    assert config.residual_order == 1
     assert config.max_consecutive_cached == 2
+    assert config.indicator_source == "raw_vision_latents"
 
 
 @pytest.mark.parametrize(
@@ -473,11 +403,6 @@ def test_sea_cache_config_validation(kwargs, message):
         SeaCacheConfig(**kwargs)
 
 
-def test_sea_cache_gate_schedule_validation():
-    with pytest.raises(TypeError, match="gate_schedule"):
-        SeaCacheConfig(gate_schedule=(True, 1))
-
-
 @pytest.mark.parametrize("callback_name", ["metadata_callback", "raw_vision_callback"])
 def test_sea_cache_callback_validation(callback_name):
     with pytest.raises(TypeError, match=callback_name):
diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py
index 977ea6b52a1a..c03d39f2b4c1 100644
--- a/tests/models/transformers/test_models_transformer_cosmos3.py
+++ b/tests/models/transformers/test_models_transformer_cosmos3.py
@@ -150,15 +150,9 @@ def counted_gen_norm_forward(*args, **kwargs):
         with torch.no_grad(), model.cache_context("cond"):
             output = model(**cached_inputs)
 
-        root_hook = model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK)
         # A hit bypasses both final pathway normalizations but still runs the prediction heads.
         assert norm_calls == {"und": 1, "gen": 1}
         assert output.sample[0].shape == self.output_shape
-        assert root_hook.num_full_steps == 1
-        assert root_hook.num_cached_steps == 1
-        stats = model.get_cache_stats()
-        assert stats["indicator_source"] == indicator_source
-        assert stats["residual_boundary"] == "post_language_model_norm"
 
         model.disable_cache()
         assert set(model.state_dict()) == state_dict_keys
@@ -183,7 +177,7 @@ def test_cosmos3_sea_cache_post_norm_boundary_numeric_semantics(self):
         projection_head_calls = 0
 
         def capture_first_block_gen_input(module, args):
-            first_block_gen_inputs.append(args[0].detach().clone())
+            first_block_gen_inputs.append(args[1].detach().clone())
 
         def capture_und_output(module, args, output):
             returned_und_outputs.append(output.detach().clone())
@@ -199,7 +193,7 @@ def counted_projection_forward(hidden_states):
             projection_head_inputs.append(hidden_states.detach().clone())
             return original_projection_forward(hidden_states)
 
-        model.layers[0].input_layernorm_moe_gen.register_forward_pre_hook(capture_first_block_gen_input)
+        model.layers[0].register_forward_pre_hook(capture_first_block_gen_input)
         model.norm.register_forward_hook(capture_und_output)
         model.norm_moe_gen.register_forward_hook(capture_gen_output)
         model.proj_out.forward = counted_projection_forward
@@ -231,8 +225,6 @@ def counted_projection_forward(hidden_states):
         )
         torch.testing.assert_close(projection_head_inputs[1], returned_gen_outputs[1])
         assert projection_head_calls == 2
-        assert root_hook.num_full_steps == 1
-        assert root_hook.num_cached_steps == 1
 
     def test_cosmos3_decoder_layer_cache_metadata_tracks_generation_stream(self):
         metadata = TransformerBlockRegistry.get(Cosmos3VLTextMoTDecoderLayer)

From 058c0e98c9ba52ee35a26b72ed10f5bdd6cb1c3c Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Mon, 31 Aug 2026 05:27:43 -0700
Subject: [PATCH 06/14] Enable SeaCache by default for Cosmos3

---
 .../cosmos/modular_pipeline.py                | 13 +++++
 .../pipelines/cosmos/pipeline_cosmos3_omni.py | 53 +++++++++++++++++++
 .../cosmos/test_modular_pipeline_cosmos3.py   | 17 +++++-
 tests/pipelines/cosmos/test_cosmos3.py        | 29 +++++++++-
 4 files changed, 110 insertions(+), 2 deletions(-)

diff --git a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py
index 5efba85de64f..5121e331306e 100644
--- a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py
+++ b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py
@@ -30,7 +30,20 @@ def current_sigma(self):
     def num_timesteps(self):
         return getattr(self, "_num_timesteps", None)
 
+    def _prepare_sea_cache_config(self, config=None):
+        return Cosmos3OmniPipeline._prepare_sea_cache_config(self, config)
+
+    def enable_sea_cache(self, config=None):
+        return Cosmos3OmniPipeline.enable_sea_cache(self, config)
+
+    def disable_sea_cache(self):
+        return Cosmos3OmniPipeline.disable_sea_cache(self)
+
+    def _maybe_enable_sea_cache(self):
+        return Cosmos3OmniPipeline._maybe_enable_sea_cache(self)
+
     def __call__(self, *args, **kwargs):
+        self._maybe_enable_sea_cache()
         transformer = getattr(self, "transformer", None)
         if hasattr(transformer, "_reset_stateful_cache"):
             transformer._reset_stateful_cache()
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
index d7c1bffd850e..ff6aeb3b439e 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
@@ -27,6 +27,7 @@
 from transformers import AutoTokenizer, BatchEncoding
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
+from ...hooks import SeaCacheConfig
 from ...models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer
 from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan
 from ...models.transformers.transformer_cosmos3 import (
@@ -474,9 +475,55 @@ def __init__(
 
         # Recommended quality-control negative prompts are documented in the Cosmos3 docs
         # page (text2video / image2video). When the caller passes None we fall back to "".
+        self._is_sea_cache_enabled = True
+
+    def _prepare_sea_cache_config(self, config: SeaCacheConfig | None = None) -> SeaCacheConfig:
+        config = SeaCacheConfig() if config is None else config
+        config.current_step_callback = lambda: self.current_step_index
+        config.current_sigma_callback = lambda: self.current_sigma
+        config.num_inference_steps_callback = lambda: self.num_timesteps
+        return config
+
+    def enable_sea_cache(self, config: SeaCacheConfig | None = None) -> None:
+        """Enable SeaCache for subsequent pipeline calls."""
+        transformer = getattr(self, "transformer", None)
+        if transformer is None:
+            raise ValueError("SeaCache requires a loaded transformer.")
+        current_config = getattr(transformer, "_cache_config", None)
+        if current_config is not None:
+            if isinstance(current_config, SeaCacheConfig):
+                self._prepare_sea_cache_config(current_config)
+                self._is_sea_cache_enabled = True
+                return
+            raise ValueError(
+                f"Caching is already enabled with {type(current_config).__name__}. Disable it before enabling SeaCache."
+            )
+
+        transformer.enable_cache(self._prepare_sea_cache_config(config))
+        self._is_sea_cache_enabled = True
+
+    def disable_sea_cache(self) -> None:
+        """Disable the default SeaCache optimization for subsequent pipeline calls."""
+        self._is_sea_cache_enabled = False
+        transformer = getattr(self, "transformer", None)
+        if isinstance(getattr(transformer, "_cache_config", None), SeaCacheConfig):
+            transformer.disable_cache()
+
+    def _maybe_enable_sea_cache(self) -> None:
+        transformer = getattr(self, "transformer", None)
+        if transformer is None or not getattr(self, "_is_sea_cache_enabled", True):
+            return
+        if isinstance(getattr(transformer, "_cache_config", None), SeaCacheConfig):
+            self._prepare_sea_cache_config(transformer._cache_config)
+            return
+        if transformer.is_cache_enabled:
+            return
+        self.enable_sea_cache()
 
     # TODO YiYi & Daniel: fix for this use case in the base class
     def _get_execution_device(self) -> torch.device:
+        from ...hooks.group_offloading import _get_group_onload_device
+
         # `self._execution_device` walks `self.components` and ultimately falls back to
         # `self.device`, which iterates modules in sorted order and ignores
         # `_exclude_from_cpu_offload`. With `safety_checker` registered, that path picks
@@ -488,6 +535,11 @@ def _get_execution_device(self) -> torch.device:
             if not isinstance(component, torch.nn.Module):
                 continue
 
+            try:
+                return _get_group_onload_device(component)
+            except ValueError:
+                pass
+
             for module in component.modules():
                 hook = getattr(module, "_hf_hook", None)
                 execution_device = getattr(hook, "execution_device", None)
@@ -1462,6 +1514,7 @@ def __call__(
                 `sound` (`torch.Tensor` of shape `[C, N]`, or `None` when `enable_sound=False`). Otherwise a tuple
                 `(video, sound)` with the same fields.
         """
+        self._maybe_enable_sea_cache()
         if hasattr(self.transformer, "_reset_stateful_cache"):
             self.transformer._reset_stateful_cache()
 
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
index 2750838897ab..395f32d9cc86 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
@@ -21,7 +21,7 @@
 import torch
 from PIL import Image
 
-from diffusers import ModularPipeline, UniPCMultistepScheduler
+from diffusers import ModularPipeline, SeaCacheConfig, UniPCMultistepScheduler
 from diffusers.modular_pipelines import (
     Cosmos3OmniBlocks,
     Cosmos3OmniModularPipeline,
@@ -189,6 +189,21 @@ def record_context(name):
             torch.testing.assert_close(sigma, pipe.scheduler.sigmas[expected_step])
         assert pipe.current_step_index is None
         assert pipe.current_sigma is None
+        assert isinstance(pipe.transformer._cache_config, SeaCacheConfig)
+
+    def test_default_sea_cache_can_be_disabled_and_reenabled(self):
+        pipe = self.get_pipeline()
+        pipe._maybe_enable_sea_cache()
+
+        assert isinstance(pipe.transformer._cache_config, SeaCacheConfig)
+        pipe.disable_sea_cache()
+        assert not pipe.transformer.is_cache_enabled
+
+        pipe._maybe_enable_sea_cache()
+        assert not pipe.transformer.is_cache_enabled
+
+        pipe.enable_sea_cache()
+        assert isinstance(pipe.transformer._cache_config, SeaCacheConfig)
 
     def test_vae_encoder_is_standalone_and_validates_conditioning_inputs(self):
         pipe = self.get_pipeline()
diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py
index eb5c28e17efe..2ce74bcf94a4 100644
--- a/tests/pipelines/cosmos/test_cosmos3.py
+++ b/tests/pipelines/cosmos/test_cosmos3.py
@@ -21,7 +21,13 @@
 from PIL import Image
 from transformers import AutoTokenizer
 
-from diffusers import AutoencoderKLWan, Cosmos3OmniPipeline, Cosmos3OmniTransformer, UniPCMultistepScheduler
+from diffusers import (
+    AutoencoderKLWan,
+    Cosmos3OmniPipeline,
+    Cosmos3OmniTransformer,
+    SeaCacheConfig,
+    UniPCMultistepScheduler,
+)
 from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import _preprocess_conditioning_image
 
 from ...testing_utils import enable_full_determinism, torch_device
@@ -118,6 +124,27 @@ def test_inference(self):
         video = pipeline(**self.get_dummy_inputs(torch_device)).video
 
         self.assertEqual(video.shape, (1, 16, 16, 3))
+        self.assertIsInstance(pipeline.transformer._cache_config, SeaCacheConfig)
+
+    def test_default_sea_cache_can_be_disabled_and_reenabled(self):
+        components = self.get_dummy_components()
+        pipeline = self.pipeline_class(**components)
+        pipeline._maybe_enable_sea_cache()
+
+        self.assertIsInstance(pipeline.transformer._cache_config, SeaCacheConfig)
+        pipeline.disable_sea_cache()
+        self.assertFalse(pipeline.transformer.is_cache_enabled)
+
+        pipeline._maybe_enable_sea_cache()
+        self.assertFalse(pipeline.transformer.is_cache_enabled)
+
+        pipeline.enable_sea_cache()
+        self.assertIsInstance(pipeline.transformer._cache_config, SeaCacheConfig)
+
+        replacement_pipeline = self.pipeline_class(**components)
+        replacement_pipeline._current_step_index = 3
+        replacement_pipeline._maybe_enable_sea_cache()
+        self.assertEqual(replacement_pipeline.transformer._cache_config.current_step_callback(), 3)
 
     def test_fp32_sampling_state_keeps_transformer_inputs_in_model_dtype(self):
         pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device)

From 9985bf7c422fd3fc374548f5c01d932a4f759da9 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Mon, 31 Aug 2026 05:45:42 -0700
Subject: [PATCH 07/14] Document SeaCache for Cosmos3

---
 docs/source/en/api/cache.md             |  6 ++++
 docs/source/en/api/pipelines/cosmos3.md | 38 +++++++++++++++++++++++++
 docs/source/en/optimization/cache.md    | 35 +++++++++++++++++++++++
 3 files changed, 79 insertions(+)

diff --git a/docs/source/en/api/cache.md b/docs/source/en/api/cache.md
index a5ed8751118d..5d0d16585013 100644
--- a/docs/source/en/api/cache.md
+++ b/docs/source/en/api/cache.md
@@ -46,3 +46,9 @@ Cache methods speedup diffusion transformers by storing and reusing intermediate
 [[autodoc]] MagCacheConfig
 
 [[autodoc]] apply_mag_cache
+
+## SeaCacheConfig
+
+[[autodoc]] SeaCacheConfig
+
+[[autodoc]] apply_sea_cache
diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md
index ae8de9ca66c8..1430dbbaa4bb 100644
--- a/docs/source/en/api/pipelines/cosmos3.md
+++ b/docs/source/en/api/pipelines/cosmos3.md
@@ -660,6 +660,44 @@ if result.action is not None:
 </hfoption>
 </hfoptions>
 
+## SeaCache
+
+Cosmos 3 inference uses [`SeaCacheConfig`] by default. SeaCache reuses transformer residuals when the
+Spectral-Evolution-Aware indicator changes slowly, reducing the number of full transformer executions. The default
+configuration filters raw vision latents, linearly extrapolates cached residuals, uses a `0.25` threshold, and forces a
+full execution after at most two consecutive cached steps.
+
+SeaCache is approximate and can change generated outputs. Disable it before inference when you need every denoising
+step to execute the full transformer, for example when producing an uncached baseline:
+
+```python
+pipe.disable_sea_cache()
+result = pipe(
+    prompt=prompt,
+    num_frames=189,
+    height=720,
+    width=1280,
+)
+```
+
+The same method works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and
+[`Cosmos3DistilledModularPipeline`]. Reenable the default configuration with `pipe.enable_sea_cache()`, or pass a
+custom configuration:
+
+```python
+from diffusers import SeaCacheConfig
+
+pipe.enable_sea_cache(
+    SeaCacheConfig(
+        threshold=0.2,
+        max_consecutive_cached=2,
+    )
+)
+```
+
+The pipeline supplies SeaCache with the active scheduler step, sigma, and number of inference steps. Cache state is
+reset for each pipeline call, and conditional and unconditional guidance branches keep independent histories.
+
 ## Context parallelism
 
 For long videos or high resolutions, a single forward pass can exceed the memory and latency budget of one GPU. Cosmos 3 supports **context parallelism (CP)** to shard the sequence dimension across multiple GPUs, splitting the attention computation so each device holds only a slice of the tokens.
diff --git a/docs/source/en/optimization/cache.md b/docs/source/en/optimization/cache.md
index 079f073b73f0..4db41753507f 100644
--- a/docs/source/en/optimization/cache.md
+++ b/docs/source/en/optimization/cache.md
@@ -68,6 +68,41 @@ config = FasterCacheConfig(
 pipeline.transformer.enable_cache(config)
 ```
 
+## SeaCache
+
+[SeaCache](https://huggingface.co/papers/2602.18993) compares Spectral-Evolution-Aware (SEA) indicators between
+successive denoising steps. When the accumulated indicator change remains below a threshold, it skips the expensive
+transformer block stack and predicts its output from cached residuals.
+
+[`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and [`Cosmos3DistilledModularPipeline`] enable SeaCache
+automatically when inference starts. The default [`SeaCacheConfig`] filters raw vision latents, linearly extrapolates
+the cached residual, uses `threshold=0.25`, and allows at most two consecutive cached steps before forcing a full
+transformer execution.
+
+SeaCache is an approximate optimization and may change generated outputs. Disable it for full transformer execution,
+such as when measuring an uncached baseline:
+
+```python
+from diffusers import Cosmos3OmniPipeline
+
+pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano")
+pipe.disable_sea_cache()
+```
+
+Call [`~Cosmos3OmniPipeline.enable_sea_cache`] to reenable it or provide a custom configuration. The pipeline fills in
+the scheduler callbacks required by SeaCache.
+
+```python
+from diffusers import SeaCacheConfig
+
+pipe.enable_sea_cache(
+    SeaCacheConfig(
+        threshold=0.2,
+        max_consecutive_cached=2,
+    )
+)
+```
+
 ## FirstBlockCache
 
 [FirstBlock Cache](https://huggingface.co/docs/diffusers/main/en/api/cache#diffusers.FirstBlockCacheConfig) checks how much the early layers of the denoiser changes from one timestep to the next. If the change is small, the model skips the expensive later layers and reuses the previous output.

From e1852251e74f4b0c16ba540d64f6db8a934fdb44 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Mon, 31 Aug 2026 07:14:16 -0700
Subject: [PATCH 08/14] Use FP32 sampling state by default for Cosmos3

---
 docs/source/en/api/pipelines/cosmos3.md                   | 6 ++++++
 src/diffusers/modular_pipelines/cosmos/before_denoise.py  | 2 +-
 .../cosmos/modular_blocks_cosmos3_distilled.py            | 4 ++--
 src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py   | 8 ++++----
 .../cosmos/test_modular_pipeline_cosmos3_distilled.py     | 3 ++-
 tests/pipelines/cosmos/test_cosmos3.py                    | 1 -
 6 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md
index 1430dbbaa4bb..84752ace65b8 100644
--- a/docs/source/en/api/pipelines/cosmos3.md
+++ b/docs/source/en/api/pipelines/cosmos3.md
@@ -698,6 +698,12 @@ pipe.enable_sea_cache(
 The pipeline supplies SeaCache with the active scheduler step, sigma, and number of inference steps. Cache state is
 reset for each pipeline call, and conditional and unconditional guidance branches keep independent histories.
 
+## Sampling precision
+
+Cosmos 3 keeps denoising latents and classifier-free-guidance arithmetic in `torch.float32` by default while casting
+transformer inputs to the model dtype. Pass `use_fp32_sampling_state=False` to an inference call to keep the sampling
+state in the model dtype instead.
+
 ## Context parallelism
 
 For long videos or high resolutions, a single forward pass can exceed the memory and latency budget of one GPU. Cosmos 3 supports **context parallelism (CP)** to shard the sequence dimension across multiple GPUs, splitting the attention computation so each device holds only a slice of the tokens.
diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py
index 298312e450c9..a717e97518f9 100644
--- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py
+++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py
@@ -70,7 +70,7 @@ def expected_components(self) -> list[ComponentSpec]:
 
     @property
     def expected_configs(self) -> list[ConfigSpec]:
-        return [ConfigSpec(name="default_use_fp32_sampling_state", default=False)]
+        return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)]
 
     @property
     def inputs(self) -> list[InputParam]:
diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
index a0de9459e5ec..53b140b2960b 100644
--- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
+++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
@@ -86,7 +86,7 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks):
           transformer (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`)
 
       Configs:
-          default_use_fp32_sampling_state (default: False) is_distilled (default: True) distilled_sigmas (default:
+          default_use_fp32_sampling_state (default: True) is_distilled (default: True) distilled_sigmas (default:
           None)
 
       Inputs:
@@ -170,7 +170,7 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks):
 
       Configs:
           default_use_system_prompt (default: True) enable_safety_checker (default: True)
-          default_use_fp32_sampling_state (default: False) is_distilled (default: True) distilled_sigmas (default:
+          default_use_fp32_sampling_state (default: True) is_distilled (default: True) distilled_sigmas (default:
           None)
 
       Inputs:
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
index ff6aeb3b439e..1322434d17c7 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
@@ -1406,7 +1406,7 @@ def __call__(
         add_resolution_template: bool = True,
         add_duration_template: bool = True,
         enable_safety_check: bool = True,
-        use_fp32_sampling_state: bool = False,
+        use_fp32_sampling_state: bool = True,
     ) -> Cosmos3OmniPipelineOutput:
         r"""
         Run the Cosmos 3 omni pipeline end-to-end: encode the (optional) conditioning image/video, denoise vision and
@@ -1502,11 +1502,11 @@ def __call__(
                 When `True` and a `CosmosSafetyChecker` is attached, runs the text guardrail on the prompt before
                 generation and the video guardrail on the decoded frames. Set to `False` to skip both for this call;
                 the checker remains loaded for subsequent calls.
-            use_fp32_sampling_state (`bool`, *optional*, defaults to `False`):
+            use_fp32_sampling_state (`bool`, *optional*, defaults to `True`):
                 When `True`, keeps vision, sound, and action denoising latents plus classifier-free-guidance arithmetic
                 in `torch.float32`. Transformer inputs are still cast to the transformer's dtype before each forward.
-                This improves sampling-state precision at the cost of additional memory and preserves the existing
-                model-dtype behavior when disabled.
+                This improves sampling-state precision at the cost of additional memory. Set it to `False` to keep
+                sampling state in the model dtype.
 
         Returns:
             [`Cosmos3OmniPipelineOutput`] or `tuple`:
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
index 11e6900d12a1..566343115ca8 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
@@ -118,7 +118,8 @@ def test_declares_distilled_configs(self):
         assert pipe.config.is_distilled is True
         assert pipe.config.distilled_sigmas is None
         assert pipe.config.default_use_system_prompt is True
-        assert pipe.config.default_use_fp32_sampling_state is False
+        assert pipe.config.default_use_fp32_sampling_state is True
+
     @pytest.mark.parametrize(
         ("config_enabled", "input_enabled", "expected_dtype"),
         [
diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py
index 2ce74bcf94a4..da7c2725018d 100644
--- a/tests/pipelines/cosmos/test_cosmos3.py
+++ b/tests/pipelines/cosmos/test_cosmos3.py
@@ -166,7 +166,6 @@ def callback_on_step_end(_pipeline, _step_index, _timestep, callback_kwargs):
         inputs.update(
             output_type="latent",
             enable_safety_check=False,
-            use_fp32_sampling_state=True,
             callback_on_step_end=callback_on_step_end,
         )
         with mock.patch.object(pipeline.transformer, "forward", side_effect=transformer_forward):

From 2cf229ca87092fbde0c68200c4d6be9a9e9e4040 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Mon, 31 Aug 2026 10:56:32 -0700
Subject: [PATCH 09/14] Fix SeaCache reconfiguration, transfer state isolation,
 and FP32 sampling

---
 src/diffusers/hooks/sea_cache.py              | 182 +++++++++++++++++-
 .../cosmos/before_denoise.py                  | 102 ++++++++--
 .../modular_pipelines/cosmos/denoise.py       |  26 ++-
 .../cosmos/modular_blocks_cosmos3.py          | 105 +++++++---
 .../modular_blocks_cosmos3_distilled.py       |  44 +++--
 .../pipelines/cosmos/pipeline_cosmos3_omni.py |  15 +-
 tests/hooks/test_sea_cache.py                 | 133 +++++++++++--
 .../cosmos/test_modular_pipeline_cosmos3.py   | 119 +++++++++++-
 ...test_modular_pipeline_cosmos3_distilled.py | 103 +++++-----
 tests/pipelines/cosmos/test_cosmos3.py        |  21 ++
 10 files changed, 711 insertions(+), 139 deletions(-)

diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py
index 827648709518..93e2164dfccc 100644
--- a/src/diffusers/hooks/sea_cache.py
+++ b/src/diffusers/hooks/sea_cache.py
@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import inspect
 import math
+import time
 from dataclasses import dataclass
-from typing import Any, Callable, Literal
+from typing import Any, Callable, Literal, Sequence
 
 import torch
 
@@ -72,11 +74,15 @@ class SeaCacheConfig:
             Callback returning the number of scheduler steps in the current pipeline call.
         metadata_callback (`Callable`, *optional*):
             Advanced model adapter returning a list of `(indices, (T, H, W))` entries that locate projected noisy
-            vision tokens in the generation stream for the `"first_block"` indicator. Cosmos 3 uses its native
-            adapter when this is omitted.
+            vision tokens in the generation stream for the `"first_block"` indicator. Cosmos 3 uses its native adapter
+            when this is omitted.
         raw_vision_callback (`Callable`, *optional*):
             Advanced model adapter returning raw vision latents with shape `(C, T, H, W)` for the
             `"raw_vision_latents"` indicator. Cosmos 3 uses its native adapter when this is omitted.
+        gate_schedule (`Sequence[bool]`, *optional*):
+            Per-step full-compute decisions to replay for controlled residual-prediction ablations. SEA still evaluates
+            its natural decision and reports mismatches; invalid or unsafe calls always fail open to full compute.
+
     Example:
         ```python
         >>> from diffusers import Cosmos3OmniPipeline, SeaCacheConfig
@@ -110,6 +116,7 @@ class SeaCacheConfig:
         [torch.nn.Module, tuple[Any, ...], dict[str, Any]],
         list[torch.Tensor] | None,
     ] = None
+    gate_schedule: Sequence[bool] = None
 
     def __post_init__(self):
         if not math.isfinite(self.threshold) or self.threshold < 0:
@@ -144,6 +151,10 @@ def __post_init__(self):
             callback = getattr(self, name)
             if callback is not None and not callable(callback):
                 raise TypeError(f"`{name}` must be callable or `None`.")
+        if self.gate_schedule is not None:
+            self.gate_schedule = tuple(self.gate_schedule)
+            if any(not isinstance(value, bool) for value in self.gate_schedule):
+                raise TypeError("`gate_schedule` must contain only boolean values.")
 
 
 @dataclass
@@ -163,34 +174,41 @@ def __init__(self):
         self.previous_indicator: list[torch.Tensor] | None = None
         self.accumulated_distance = 0.0
         self.consecutive_cached = 0
+        self.max_consecutive_cached_observed = 0
+        self.max_consecutive_forced_full = 0
         self.skip_remaining = False
         self.full_execution_pending = False
         self.cacheable_execution = False
         self.step_index: int | None = None
-        self.und_input: torch.Tensor | None = None
         self.gen_input: torch.Tensor | None = None
         self.und_output: torch.Tensor | None = None
         self.cached_und_output: torch.Tensor | None = None
         self.cached_gen_residual: torch.Tensor | None = None
+        self.full_started_at: float | None = None
 
     def reset_forward(self):
         self.skip_remaining = False
         self.full_execution_pending = False
         self.cacheable_execution = False
         self.step_index = None
-        self.und_input = None
         self.gen_input = None
         self.und_output = None
         self.cached_und_output = None
         self.cached_gen_residual = None
+        self.full_started_at = None
 
-    def reset(self):
+    def reset_trajectory(self):
         self.history = []
         self.gate_key = None
         self.gate_should_compute = True
         self.previous_indicator = None
         self.accumulated_distance = 0.0
         self.consecutive_cached = 0
+
+    def reset(self):
+        self.reset_trajectory()
+        self.max_consecutive_cached_observed = 0
+        self.max_consecutive_forced_full = 0
         self.reset_forward()
 
 
@@ -201,6 +219,20 @@ def __init__(self):
 
     def reset(self):
         self.forward_metadata: _SeaCacheForwardMetadata | None = None
+        self.transformer_calls = 0
+        self.gate_evaluations = 0
+        self.gate_full_decisions = 0
+        self.gate_skip_decisions = 0
+        self.gate_trace: list[bool] = []
+        self.gate_schedule_mismatches = 0
+        self.num_full_steps = 0
+        self.num_cached_steps = 0
+        self.fail_open_calls = 0
+        self.indicator_seconds = 0.0
+        self.decision_seconds = 0.0
+        self.full_seconds = 0.0
+        self.branch_full_executions: dict[str, int] = {}
+        self.branch_reuses: dict[str, int] = {}
 
     def warn_once(self, message: str):
         if message not in self._warned_messages:
@@ -208,6 +240,7 @@ def warn_once(self, message: str):
             self._warned_messages.add(message)
 
     def mark_fail_open(self, message: str):
+        self.fail_open_calls += 1
         self.warn_once(message)
 
     def resolve_gate(
@@ -221,9 +254,11 @@ def resolve_gate(
         if state.gate_key == gate_key:
             return state.gate_should_compute
 
+        self.gate_evaluations += 1
         is_non_adjacent = state.gate_key is not None and metadata.step_index != state.gate_key[0] + 1
         if is_non_adjacent:
             self.mark_fail_open("SeaCache received non-adjacent scheduler steps; running full.")
+            state.reset_trajectory()
         is_retained = metadata.step_index < config.retention_steps
         is_in_cache_end = metadata.step_index >= metadata.num_inference_steps - config.cache_end_steps
         is_first_observation = state.previous_indicator is None
@@ -235,7 +270,7 @@ def resolve_gate(
         candidate_accumulated_distance = 0.0
 
         if forced_compute:
-            should_compute = True
+            natural_should_compute = True
         else:
             if len(indicator) != len(state.previous_indicator) or not indicator:
                 distance = float("inf")
@@ -260,13 +295,41 @@ def resolve_gate(
                 invalid_gate = True
                 self.mark_fail_open("SeaCache indicator history changed shape, device, or dtype; running full.")
             candidate_accumulated_distance = state.accumulated_distance + distance
-            should_compute = invalid_gate or candidate_accumulated_distance >= config.threshold
+            natural_should_compute = invalid_gate or candidate_accumulated_distance >= config.threshold
+
+        should_compute = natural_should_compute
+
+        if config.gate_schedule is not None:
+            schedule_is_valid = len(config.gate_schedule) == metadata.num_inference_steps
+            scheduled_compute = bool(config.gate_schedule[metadata.step_index]) if schedule_is_valid else True
+            can_replay_skip = not (invalid_gate or is_retained or is_in_cache_end or is_first_observation)
+            if not schedule_is_valid:
+                self.mark_fail_open(
+                    "SeaCache gate schedule length does not match the number of inference steps; running full."
+                )
+                should_compute = True
+            elif scheduled_compute or can_replay_skip:
+                if natural_should_compute != scheduled_compute:
+                    self.gate_schedule_mismatches += 1
+                should_compute = scheduled_compute
+            else:
+                self.mark_fail_open("SeaCache gate schedule requested an unsafe cache hit; running full.")
+                should_compute = True
+
+        if is_max_consecutive:
+            should_compute = True
+            state.max_consecutive_forced_full += 1
 
         state.accumulated_distance = 0.0 if should_compute else candidate_accumulated_distance
 
         state.gate_key = gate_key
         state.gate_should_compute = should_compute
         state.previous_indicator = None if indicator is None else [value.detach() for value in indicator]
+        self.gate_trace.append(should_compute)
+        if should_compute:
+            self.gate_full_decisions += 1
+        else:
+            self.gate_skip_decisions += 1
         return should_compute
 
 
@@ -312,11 +375,18 @@ def _get_block_outputs(
 
 def _record_full_execution(
     config: SeaCacheConfig,
+    state_manager: StateManager,
+    shared_state: SeaCacheSharedState,
     state: SeaCacheContextState,
     gen_output: torch.Tensor,
     und_output: torch.Tensor | None,
 ) -> None:
+    shared_state.num_full_steps += 1
     state.consecutive_cached = 0
+    branch = state_manager._current_context
+    shared_state.branch_full_executions[branch] = shared_state.branch_full_executions.get(branch, 0) + 1
+    if state.full_started_at is not None:
+        shared_state.full_seconds += time.perf_counter() - state.full_started_at
     if (
         state.cacheable_execution
         and state.step_index is not None
@@ -500,6 +570,7 @@ def __init__(
         shared_state: SeaCacheSharedState,
         metadata_callback: Callable,
         raw_vision_callback: Callable,
+        residual_boundary: str,
     ):
         super().__init__()
         self.config = config
@@ -507,9 +578,73 @@ def __init__(
         self.shared_state = shared_state
         self.metadata_callback = metadata_callback
         self.raw_vision_callback = raw_vision_callback
+        self.residual_boundary = residual_boundary
+        self._last_stats: dict[str, Any] | None = None
+
+    @property
+    def num_full_steps(self) -> int:
+        return self.shared_state.num_full_steps
+
+    @property
+    def num_cached_steps(self) -> int:
+        return self.shared_state.num_cached_steps
+
+    def stats(self) -> dict[str, Any]:
+        if self.shared_state.transformer_calls == 0 and self._last_stats is not None:
+            return copy.deepcopy(self._last_stats)
+
+        persistent_cache_bytes = 0
+        per_branch = {}
+        for branch, state in sorted(self.state_manager._state_cache.items()):
+            if state.previous_indicator is not None:
+                persistent_cache_bytes += sum(
+                    value.numel() * value.element_size() for value in state.previous_indicator
+                )
+            for _, und_output, gen_residual in state.history:
+                persistent_cache_bytes += und_output.numel() * und_output.element_size()
+                persistent_cache_bytes += gen_residual.numel() * gen_residual.element_size()
+            per_branch[branch] = {
+                "full_calls": self.shared_state.branch_full_executions.get(branch, 0),
+                "reuse_calls": self.shared_state.branch_reuses.get(branch, 0),
+                "max_consecutive_cached_observed": state.max_consecutive_cached_observed,
+                "max_consecutive_forced_full": state.max_consecutive_forced_full,
+            }
+
+        opportunities = self.shared_state.num_full_steps + self.shared_state.num_cached_steps
+        return {
+            "indicator_source": self.config.indicator_source,
+            "residual_order": self.config.residual_order,
+            "residual_boundary": self.residual_boundary,
+            "max_consecutive_cached": self.config.max_consecutive_cached,
+            "max_consecutive_cached_observed": max(
+                (stats["max_consecutive_cached_observed"] for stats in per_branch.values()), default=0
+            ),
+            "max_consecutive_forced_full": sum(stats["max_consecutive_forced_full"] for stats in per_branch.values()),
+            "transformer_calls": self.shared_state.transformer_calls,
+            "gate_evaluations": self.shared_state.gate_evaluations,
+            "gate_full_decisions": self.shared_state.gate_full_decisions,
+            "gate_skip_decisions": self.shared_state.gate_skip_decisions,
+            "gate_trace": list(self.shared_state.gate_trace),
+            "gate_schedule_replayed": self.config.gate_schedule is not None,
+            "gate_schedule_mismatches": self.shared_state.gate_schedule_mismatches,
+            "actual_full_executions": self.shared_state.num_full_steps,
+            "actual_reuses": self.shared_state.num_cached_steps,
+            "actual_reuse_rate": (self.shared_state.num_cached_steps / opportunities if opportunities else 0.0),
+            "fail_open_calls": self.shared_state.fail_open_calls,
+            "sea_indicator_seconds": self.shared_state.indicator_seconds,
+            "sea_decision_seconds": self.shared_state.decision_seconds,
+            "sea_seconds": (self.shared_state.indicator_seconds + self.shared_state.decision_seconds),
+            "full_seconds": self.shared_state.full_seconds,
+            "timing_note": "host wall time; CUDA work is asynchronous",
+            "persistent_cache_bytes": persistent_cache_bytes,
+            "branch_full_executions": dict(sorted(self.shared_state.branch_full_executions.items())),
+            "branch_reuses": dict(sorted(self.shared_state.branch_reuses.items())),
+            "per_branch": per_branch,
+        }
 
     def pre_forward(self, module: torch.nn.Module, *args, **kwargs):
         self.shared_state.forward_metadata = None
+        self.shared_state.transformer_calls += 1
         if torch.is_grad_enabled():
             self.shared_state.mark_fail_open(
                 "SeaCache is inference-only; calls with autograd enabled run in fail-open mode."
@@ -608,6 +743,8 @@ def post_forward(self, module: torch.nn.Module, output: Any) -> Any:
         return output
 
     def reset_state(self, module: torch.nn.Module):
+        if self.shared_state.transformer_calls > 0:
+            self._last_stats = self.stats()
         self.state_manager.reset()
         self.shared_state.reset()
         return module
@@ -689,14 +826,16 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
             state.reset_forward()
             state.full_execution_pending = True
             state.gen_input = hidden_states
-            state.und_input = encoder_hidden_states
 
         forward_metadata = self.shared_state.forward_metadata
         if state is None or forward_metadata is None:
+            if state is not None:
+                state.full_started_at = time.perf_counter()
             return self.fn_ref.original_forward(*args, **kwargs)
 
         state.step_index = forward_metadata.step_index
         state.cacheable_execution = True
+        indicator_started = time.perf_counter()
         indicator_error_reported = False
         if _is_parameter_sharded(module):
             self.shared_state.mark_fail_open(
@@ -713,11 +852,14 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
                 )
                 indicator = None
                 indicator_error_reported = True
+        self.shared_state.indicator_seconds += time.perf_counter() - indicator_started
         if indicator is None and not indicator_error_reported:
             self.shared_state.mark_fail_open(
                 "SeaCache could not construct its vision indicator; running in fail-open mode."
             )
+        decision_started = time.perf_counter()
         should_compute = self.shared_state.resolve_gate(state, forward_metadata, indicator, self.config)
+        self.shared_state.decision_seconds += time.perf_counter() - decision_started
 
         if should_compute or not state.history:
             if not should_compute:
@@ -725,6 +867,7 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
                 self.shared_state.mark_fail_open(
                     "SeaCache selected a cache hit without residual history; running in fail-open mode."
                 )
+            state.full_started_at = time.perf_counter()
             return self.fn_ref.original_forward(*args, **kwargs)
 
         residual_history = state.history[-(self.config.residual_order + 1) :]
@@ -746,6 +889,7 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
             self.shared_state.mark_fail_open(
                 "SeaCache residual history changed shape, device, or dtype; running in fail-open mode."
             )
+            state.full_started_at = time.perf_counter()
             return self.fn_ref.original_forward(*args, **kwargs)
 
         if self.config.residual_order == 1 and len(residual_history) >= 2:
@@ -760,6 +904,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
         state.cached_und_output = cached_und
         state.cached_gen_residual = cached_residual
         state.consecutive_cached += 1
+        state.max_consecutive_cached_observed = max(state.max_consecutive_cached_observed, state.consecutive_cached)
+        self.shared_state.num_cached_steps += 1
+        branch = self.state_manager._current_context
+        self.shared_state.branch_reuses[branch] = self.shared_state.branch_reuses.get(branch, 0) + 1
         if self.post_norm_boundary:
             return _build_block_output(self._metadata, hidden_states, encoder_hidden_states)
         return _build_block_output(self._metadata, hidden_states + cached_residual, cached_und)
@@ -804,6 +952,8 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
         hidden_states, encoder_hidden_states = _get_block_outputs(self._metadata, output)
         _record_full_execution(
             self.config,
+            self.state_manager,
+            self.shared_state,
             state,
             gen_output=hidden_states,
             und_output=encoder_hidden_states,
@@ -856,6 +1006,8 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
 
         _record_full_execution(
             self.config,
+            self.state_manager,
+            self.shared_state,
             state,
             gen_output=output,
             und_output=state.und_output,
@@ -903,6 +1055,7 @@ def apply_sea_cache(module: torch.nn.Module, config: SeaCacheConfig) -> None:
                 "residual boundary."
             )
     post_norm_boundary = post_norm_modules is not None
+    residual_boundary = "post_language_model_norm" if post_norm_boundary else "repeated_block_stack"
 
     blocks = []
     for name, submodule in unwrapped_module.named_children():
@@ -932,6 +1085,7 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None:
                 shared_state,
                 metadata_callback,
                 raw_vision_callback,
+                residual_boundary,
             ),
             _SEA_CACHE_ROOT_HOOK,
         )
@@ -980,3 +1134,13 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None:
         raise
 
     root_registry._child_registries_cache = None
+
+
+def get_sea_cache_stats(module: torch.nn.Module) -> dict[str, Any]:
+    """Return statistics for the SeaCache instance currently attached to ``module``."""
+
+    registry = getattr(module, "_diffusers_hook", None)
+    root_hook = registry.get_hook(_SEA_CACHE_ROOT_HOOK) if registry is not None else None
+    if not isinstance(root_hook, SeaCacheRootHook):
+        raise ValueError("SeaCache is not enabled on this module.")
+    return root_hook.stats()
diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py
index a717e97518f9..decbe325782b 100644
--- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py
+++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py
@@ -12,6 +12,15 @@
 from .modular_pipeline import Cosmos3OmniModularPipeline
 
 
+def _resolve_sampling_dtype(
+    components: Cosmos3OmniModularPipeline, use_fp32_sampling_state: bool | None
+) -> tuple[bool, torch.dtype]:
+    if use_fp32_sampling_state is None:
+        use_fp32_sampling_state = components.config.default_use_fp32_sampling_state
+    sampling_dtype = torch.float32 if use_fp32_sampling_state else components.transformer.dtype
+    return use_fp32_sampling_state, sampling_dtype
+
+
 class Cosmos3PrepareTextSegmentsStep(ModularPipelineBlocks):
     model_name = "cosmos3-omni"
 
@@ -107,7 +116,8 @@ def inputs(self) -> list[InputParam]:
                 type_hint=bool | None,
                 default=None,
                 description=(
-                    "Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the "
+                    "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. "
+                    "If unset, uses the "
                     "pipeline's `default_use_fp32_sampling_state` config."
                 ),
             ),
@@ -145,9 +155,9 @@ def intermediate_outputs(self) -> list[OutputParam]:
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         device = components._execution_device
-        if block_state.use_fp32_sampling_state is None:
-            block_state.use_fp32_sampling_state = components.config.default_use_fp32_sampling_state
-        sampling_dtype = torch.float32 if block_state.use_fp32_sampling_state else components.transformer.dtype
+        block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype(
+            components, block_state.use_fp32_sampling_state
+        )
 
         x0_tokens_vision = block_state.x0_tokens_vision
         if x0_tokens_vision is None:
@@ -216,6 +226,10 @@ def expected_components(self) -> list[ComponentSpec]:
             ComponentSpec("scheduler", UniPCMultistepScheduler),
         ]
 
+    @property
+    def expected_configs(self) -> list[ConfigSpec]:
+        return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)]
+
     @property
     def inputs(self) -> list[InputParam]:
         return [
@@ -228,6 +242,16 @@ def inputs(self) -> list[InputParam]:
                 description="Pre-generated noisy sound latents.",
             ),
             InputParam.template("generator"),
+            InputParam(
+                name="use_fp32_sampling_state",
+                type_hint=bool | None,
+                default=None,
+                description=(
+                    "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. "
+                    "If unset, uses the "
+                    "pipeline's `default_use_fp32_sampling_state` config."
+                ),
+            ),
         ]
 
     @property
@@ -248,7 +272,9 @@ def intermediate_outputs(self) -> list[OutputParam]:
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         device = components._execution_device
-        dtype = components.transformer.dtype
+        block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype(
+            components, block_state.use_fp32_sampling_state
+        )
 
         if not components.transformer.config.sound_gen:
             raise ValueError("Sound generation requires a transformer trained with sound_gen=True.")
@@ -258,19 +284,21 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
         n_audio_samples = int(block_state.num_frames / block_state.fps * components.sound_sampling_rate)
         hop_size = components.sound_hop_size
         t_sound = (n_audio_samples + hop_size - 1) // hop_size
-        x0_tokens_sound = torch.zeros(sound_dim, t_sound, device=device, dtype=dtype)
-        block_state.sound_condition_mask = torch.zeros((x0_tokens_sound.shape[1], 1), device=device, dtype=dtype)
+        x0_tokens_sound = torch.zeros(sound_dim, t_sound, device=device, dtype=sampling_dtype)
+        block_state.sound_condition_mask = torch.zeros(
+            (x0_tokens_sound.shape[1], 1), device=device, dtype=sampling_dtype
+        )
 
         if block_state.sound_latents is None:
             pure_noise = randn_tensor(
-                tuple(x0_tokens_sound.shape), generator=block_state.generator, device=device, dtype=dtype
+                tuple(x0_tokens_sound.shape), generator=block_state.generator, device=device, dtype=sampling_dtype
             )
             block_state.sound_latents = (
                 block_state.sound_condition_mask.T * x0_tokens_sound
                 + (1.0 - block_state.sound_condition_mask.T) * pure_noise
             )
         else:
-            block_state.sound_latents = block_state.sound_latents.to(device=device, dtype=dtype)
+            block_state.sound_latents = block_state.sound_latents.to(device=device, dtype=sampling_dtype)
 
         block_state.sound_scheduler = copy.deepcopy(components.scheduler)
 
@@ -292,6 +320,10 @@ def expected_components(self) -> list[ComponentSpec]:
             ComponentSpec("scheduler", UniPCMultistepScheduler),
         ]
 
+    @property
+    def expected_configs(self) -> list[ConfigSpec]:
+        return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)]
+
     @property
     def inputs(self) -> list[InputParam]:
         return [
@@ -314,6 +346,16 @@ def inputs(self) -> list[InputParam]:
                 description="Pre-generated noisy action latents.",
             ),
             InputParam.template("generator"),
+            InputParam(
+                name="use_fp32_sampling_state",
+                type_hint=bool | None,
+                default=None,
+                description=(
+                    "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. "
+                    "If unset, uses the "
+                    "pipeline's `default_use_fp32_sampling_state` config."
+                ),
+            ),
         ]
 
     @property
@@ -345,7 +387,9 @@ def intermediate_outputs(self) -> list[OutputParam]:
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         device = components._execution_device
-        dtype = components.transformer.dtype
+        block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype(
+            components, block_state.use_fp32_sampling_state
+        )
         action = block_state.action
 
         if not components.transformer.config.action_gen:
@@ -367,7 +411,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
             raw_actions = action.raw_actions
             if raw_actions is None:
                 raise ValueError("action_mode='forward_dynamics' requires an action tensor.")
-            raw_actions = raw_actions.to(device=device, dtype=dtype)
+            raw_actions = raw_actions.to(device=device, dtype=sampling_dtype)
             if raw_actions.shape[-1] > action_dim:
                 raise ValueError(
                     f"Cosmos3 action dimension {raw_actions.shape[-1]} exceeds model action_dim={action_dim}."
@@ -388,7 +432,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
                 raw_actions = torch.cat([raw_actions, action_padding], dim=-1)
             x0_tokens_action = raw_actions
         else:
-            x0_tokens_action = torch.zeros(action_chunk_size, action_dim, device=device, dtype=dtype)
+            x0_tokens_action = torch.zeros(action_chunk_size, action_dim, device=device, dtype=sampling_dtype)
 
         if action.domain_name not in _EMBODIMENT_TO_DOMAIN_ID:
             raise ValueError(
@@ -398,14 +442,16 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
             torch.tensor([_EMBODIMENT_TO_DOMAIN_ID[action.domain_name]], dtype=torch.long, device=device)
         ]
         condition_frames = block_state.action_condition_frame_indexes or []
-        block_state.action_condition_mask = torch.zeros((x0_tokens_action.shape[0], 1), device=device, dtype=dtype)
+        block_state.action_condition_mask = torch.zeros(
+            (x0_tokens_action.shape[0], 1), device=device, dtype=sampling_dtype
+        )
         for frame_idx in condition_frames:
             if 0 <= frame_idx < block_state.action_condition_mask.shape[0]:
                 block_state.action_condition_mask[frame_idx, 0] = 1.0
 
         if block_state.action_latents is None:
             pure_noise = randn_tensor(
-                tuple(x0_tokens_action.shape), generator=block_state.generator, device=device, dtype=dtype
+                tuple(x0_tokens_action.shape), generator=block_state.generator, device=device, dtype=sampling_dtype
             )
             block_state.action_latents = (
                 block_state.action_condition_mask * x0_tokens_action
@@ -414,7 +460,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
             if block_state.raw_action_dim_resolved is not None:
                 block_state.action_latents[:, block_state.raw_action_dim_resolved :] = 0
         else:
-            block_state.action_latents = block_state.action_latents.to(device=device, dtype=dtype)
+            block_state.action_latents = block_state.action_latents.to(device=device, dtype=sampling_dtype)
 
         block_state.action_scheduler = copy.deepcopy(components.scheduler)
 
@@ -1021,6 +1067,10 @@ def description(self) -> str:
     def expected_components(self) -> list[ComponentSpec]:
         return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
 
+    @property
+    def expected_configs(self) -> list[ConfigSpec]:
+        return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)]
+
     @property
     def inputs(self) -> list[InputParam]:
         return [
@@ -1037,6 +1087,16 @@ def inputs(self) -> list[InputParam]:
                 description="Number of pixel frames used to seed this chunk's target.",
             ),
             InputParam.template("generator"),
+            InputParam(
+                name="use_fp32_sampling_state",
+                type_hint=bool | None,
+                default=None,
+                description=(
+                    "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. "
+                    "If unset, uses the "
+                    "pipeline's `default_use_fp32_sampling_state` config."
+                ),
+            ),
         ]
 
     @property
@@ -1064,20 +1124,24 @@ def intermediate_outputs(self) -> list[OutputParam]:
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         device = components._execution_device
-        dtype = components.transformer.dtype
+        block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype(
+            components, block_state.use_fp32_sampling_state
+        )
         tcf = components.vae_scale_factor_temporal
 
-        target_x0 = block_state.x0_tokens_vision.to(device=device)
+        target_x0 = block_state.x0_tokens_vision.to(device=device, dtype=sampling_dtype)
         current_conditional_frames = block_state.current_conditional_frames
 
         # Build the noisy target latents + conditioning mask from the clean target latents.
         latent_t = target_x0.shape[2]
-        condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=dtype)
+        condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=sampling_dtype)
         latent_condition_frames = 0
         if current_conditional_frames > 0:
             latent_condition_frames = (current_conditional_frames - 1) // tcf + 1
             condition_mask[:latent_condition_frames] = 1.0
-        noise = randn_tensor(tuple(target_x0.shape), generator=block_state.generator, device=device, dtype=dtype)
+        noise = randn_tensor(
+            tuple(target_x0.shape), generator=block_state.generator, device=device, dtype=sampling_dtype
+        )
         block_state.latents = condition_mask * target_x0 + (1.0 - condition_mask) * noise
         block_state.velocity_mask = 1.0 - condition_mask
         block_state.condition_latents = condition_mask * target_x0
diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py
index ce6ba0a4bee5..5b06b4f910a2 100644
--- a/src/diffusers/modular_pipelines/cosmos/denoise.py
+++ b/src/diffusers/modular_pipelines/cosmos/denoise.py
@@ -172,6 +172,12 @@ def inputs(self) -> list[InputParam]:
                 default=6.0,
                 description="Scale for classifier-free guidance.",
             ),
+            InputParam(
+                name="use_fp32_sampling_state",
+                type_hint=bool,
+                default=False,
+                description="Whether to keep velocity and guidance arithmetic in float32.",
+            ),
         ]
 
     @property
@@ -236,6 +242,10 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
             )
 
         cond_velocity_vision, cond_velocity_sound, cond_velocity_action = velocities["cond"]
+        if block_state.use_fp32_sampling_state:
+            cond_velocity_vision = cond_velocity_vision.float()
+            cond_velocity_sound = cond_velocity_sound.float() if cond_velocity_sound is not None else None
+            cond_velocity_action = cond_velocity_action.float() if cond_velocity_action is not None else None
         if do_cfg:
             uncond_velocity_vision, uncond_velocity_sound, uncond_velocity_action = velocities["uncond"]
             block_state.velocity_vision = uncond_velocity_vision + block_state.guidance_scale * (
@@ -730,6 +740,12 @@ def inputs(self) -> list[InputParam]:
                 default=None,
                 description="Timestep interval [lo, hi] over which control guidance is active (None = always).",
             ),
+            InputParam(
+                name="use_fp32_sampling_state",
+                type_hint=bool,
+                default=False,
+                description="Whether to keep velocity and guidance arithmetic in float32.",
+            ),
         ]
 
     @property
@@ -804,6 +820,11 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
                 "uncond",
             )
 
+        if block_state.use_fp32_sampling_state:
+            cond_full = cond_full.float()
+            cond_no_control = cond_no_control.float() if cond_no_control is not None else None
+            uncond_full = uncond_full.float() if uncond_full is not None else None
+
         if needs_control_cfg and needs_text_cfg:
             control_cond = cond_no_control + step_control * (cond_full - cond_no_control)
             velocity = uncond_full + step_guidance * (control_cond - uncond_full)
@@ -878,7 +899,8 @@ class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper):
     Runs the per-chunk transfer denoising loop over scheduler timesteps.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`)
+          scheduler (`UniPCMultistepScheduler`)
 
       Inputs:
           timesteps (`Tensor`):
@@ -905,6 +927,8 @@ class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper):
               Timestep interval [lo, hi] over which text guidance is active (None = always).
           control_guidance_interval (`tuple`, *optional*):
               Timestep interval [lo, hi] over which control guidance is active (None = always).
+          use_fp32_sampling_state (`bool`, *optional*, defaults to False):
+              Whether to keep velocity and guidance arithmetic in float32.
           latents (`Tensor`):
               Noisy target latents to update.
           condition_latents (`Tensor`):
diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py
index d7c6c846fc8e..ee02e192da21 100644
--- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py
+++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py
@@ -53,11 +53,11 @@
 # auto_docstring
 class Cosmos3TransferTextBlocks(SequentialPipelineBlocks):
     """
-    Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) prompt in
-    transfer mode using the per-chunk frame count.
+    Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) prompt in transfer mode using the per-chunk frame count.
 
       Components:
-          video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`)
+          video_processor (`VideoProcessor`)
+          text_tokenizer (`AutoTokenizer`)
 
       Inputs:
           control_videos (`dict`):
@@ -121,10 +121,12 @@ class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks):
        - Cosmos3TextEncoderStep runs otherwise.
 
       Components:
-          video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`)
+          video_processor (`VideoProcessor`)
+          text_tokenizer (`AutoTokenizer`)
 
       Configs:
-          default_use_system_prompt (default: True) enable_safety_checker (default: True)
+          default_use_system_prompt (default: True)
+          enable_safety_checker (default: True)
 
       Inputs:
           control_videos (`dict`, *optional*):
@@ -204,7 +206,8 @@ class Cosmos3AutoVaeEncoderStep(ConditionalPipelineBlocks):
        - when no action, image, or video conditioning is provided, this block is skipped.
 
       Components:
-          vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+          vae (`AutoencoderKLWan`)
+          video_processor (`VideoProcessor`)
 
       Inputs:
           action (`CosmosActionCondition`, *optional*):
@@ -313,7 +316,9 @@ class Cosmos3DecodeStep(SequentialPipelineBlocks):
     Decodes denoised latents into modality outputs.
 
       Components:
-          vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) sound_tokenizer (`Cosmos3AVAEAudioTokenizer`)
+          vae (`AutoencoderKLWan`)
+          video_processor (`VideoProcessor`)
+          sound_tokenizer (`Cosmos3AVAEAudioTokenizer`)
 
       Inputs:
           latents (`Tensor`):
@@ -368,9 +373,11 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks):
     Runs the text-and-vision Cosmos3 denoising workflow.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`)
+          scheduler (`UniPCMultistepScheduler`)
 
       Configs:
+          default_use_fp32_sampling_state (default: True)
           use_native_flow_schedule (default: False)
 
       Inputs:
@@ -394,6 +401,9 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
+          use_fp32_sampling_state (`bool | NoneType`, *optional*):
+              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
+              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`):
               The number of denoising steps.
           **denoiser_input_fields (`None`, *optional*):
@@ -439,9 +449,11 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks):
     Runs the text, vision, and sound Cosmos3 denoising workflow.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`)
+          scheduler (`UniPCMultistepScheduler`)
 
       Configs:
+          default_use_fp32_sampling_state (default: True)
           use_native_flow_schedule (default: False)
 
       Inputs:
@@ -465,6 +477,9 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
+          use_fp32_sampling_state (`bool | NoneType`, *optional*):
+              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
+              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`):
               The number of denoising steps.
           sound_latents (`Tensor`, *optional*):
@@ -523,9 +538,11 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks):
     Runs the text, vision, and action Cosmos3 denoising workflow.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`)
+          scheduler (`UniPCMultistepScheduler`)
 
       Configs:
+          default_use_fp32_sampling_state (default: True)
           use_native_flow_schedule (default: False)
 
       Inputs:
@@ -549,6 +566,9 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
+          use_fp32_sampling_state (`bool | NoneType`, *optional*):
+              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
+              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`):
               The number of denoising steps.
           action (`CosmosActionCondition`):
@@ -611,9 +631,11 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks):
     Runs the text, vision, sound, and action Cosmos3 denoising workflow.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`)
+          scheduler (`UniPCMultistepScheduler`)
 
       Configs:
+          default_use_fp32_sampling_state (default: True)
           use_native_flow_schedule (default: False)
 
       Inputs:
@@ -637,6 +659,9 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
+          use_fp32_sampling_state (`bool | NoneType`, *optional*):
+              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
+              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`):
               The number of denoising steps.
           sound_latents (`Tensor`, *optional*):
@@ -707,13 +732,16 @@ def outputs(self):
 # auto_docstring
 class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks):
     """
-    Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a non-leaf
-    LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). Per-chunk
-    cross-carry (previous_output, output_chunks) lives on PipelineState.
+    Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a non-leaf LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). Per-chunk cross-carry (previous_output, output_chunks) lives on PipelineState.
 
       Components:
-          vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`Cosmos3OmniTransformer`) scheduler
-          (`UniPCMultistepScheduler`)
+          vae (`AutoencoderKLWan`)
+          video_processor (`VideoProcessor`)
+          transformer (`Cosmos3OmniTransformer`)
+          scheduler (`UniPCMultistepScheduler`)
+
+      Configs:
+          default_use_fp32_sampling_state (default: True)
 
       Inputs:
           chunk_id (`int`, *optional*, defaults to 0):
@@ -740,6 +768,9 @@ class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks):
               Number of frames each later chunk reuses from the previous chunk's tail.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
+          use_fp32_sampling_state (`bool | NoneType`, *optional*):
+              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
+              pipeline's `default_use_fp32_sampling_state` config.
           cond_text_segment (`dict`):
               Conditional text segment.
           uncond_text_segment (`dict`):
@@ -842,6 +873,8 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
         state.set("output_chunks", [])
         state.set("previous_output", None)
         for chunk_id in range(num_chunks):
+            if chunk_id > 0:
+                components.transformer._reset_stateful_cache()
             state.set("chunk_id", chunk_id)
             for _, block in self.sub_blocks.items():
                 components, state = block(components, state)
@@ -854,8 +887,13 @@ class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks):
     Transfer denoise stage: prepare shared text segments once, then run the autoregressive chunk loop.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler
-          (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`)
+          vae (`AutoencoderKLWan`)
+          video_processor (`VideoProcessor`)
+          scheduler (`UniPCMultistepScheduler`)
+
+      Configs:
+          default_use_fp32_sampling_state (default: True)
 
       Inputs:
           cond_input_ids (`None`):
@@ -886,6 +924,9 @@ class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks):
               Number of frames each later chunk reuses from the previous chunk's tail.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
+          use_fp32_sampling_state (`bool | NoneType`, *optional*):
+              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
+              pipeline's `default_use_fp32_sampling_state` config.
           fps (`float`, *optional*, defaults to 24.0):
               Frame rate of the generated video.
           num_inference_steps (`int`):
@@ -973,10 +1014,13 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks):
        - vision runs otherwise.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler
-          (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`)
+          vae (`AutoencoderKLWan`)
+          video_processor (`VideoProcessor`)
+          scheduler (`UniPCMultistepScheduler`)
 
       Configs:
+          default_use_fp32_sampling_state (default: True)
           use_native_flow_schedule (default: False)
 
       Inputs:
@@ -1008,6 +1052,9 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks):
               Number of frames each later chunk reuses from the previous chunk's tail.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
+          use_fp32_sampling_state (`bool | NoneType`, *optional*):
+              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
+              pipeline's `default_use_fp32_sampling_state` config.
           fps (`float`, *optional*, defaults to 24.0):
               Frame rate of the generated video.
           num_inference_steps (`int`):
@@ -1162,13 +1209,18 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks):
         - `action_inverse_dynamics`: requires `prompt`, `action`
 
       Components:
-          video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) transformer
-          (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) sound_tokenizer
-          (`Cosmos3AVAEAudioTokenizer`)
+          video_processor (`VideoProcessor`)
+          text_tokenizer (`AutoTokenizer`)
+          vae (`AutoencoderKLWan`)
+          transformer (`Cosmos3OmniTransformer`)
+          scheduler (`UniPCMultistepScheduler`)
+          sound_tokenizer (`Cosmos3AVAEAudioTokenizer`)
 
       Configs:
-          default_use_system_prompt (default: True) enable_safety_checker (default: True) use_native_flow_schedule
-          (default: False)
+          default_use_system_prompt (default: True)
+          enable_safety_checker (default: True)
+          default_use_fp32_sampling_state (default: True)
+          use_native_flow_schedule (default: False)
 
       Inputs:
           control_videos (`dict`, *optional*):
@@ -1213,6 +1265,9 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks):
               Number of frames the first chunk reuses from the input video.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
+          use_fp32_sampling_state (`bool | NoneType`, *optional*):
+              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
+              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`):
               The number of denoising steps.
           **denoiser_input_fields (`None`, *optional*):
diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
index 53b140b2960b..6841d2f3ff4f 100644
--- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
+++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
@@ -25,7 +25,8 @@ class Cosmos3DistilledAutoVaeEncoderStep(ConditionalPipelineBlocks):
        - when no image or video conditioning is provided, this block is skipped.
 
       Components:
-          vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+          vae (`AutoencoderKLWan`)
+          video_processor (`VideoProcessor`)
 
       Inputs:
           video (`None`, *optional*):
@@ -83,11 +84,13 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks):
     Runs the text-and-vision distilled Cosmos3 denoising workflow.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+          transformer (`Cosmos3OmniTransformer`)
+          scheduler (`FlowMatchEulerDiscreteScheduler`)
 
       Configs:
-          default_use_fp32_sampling_state (default: True) is_distilled (default: True) distilled_sigmas (default:
-          None)
+          default_use_fp32_sampling_state (default: True)
+          is_distilled (default: True)
+          distilled_sigmas (default: None)
 
       Inputs:
           cond_input_ids (`None`):
@@ -111,13 +114,13 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks):
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
           use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the pipeline's
-              `default_use_fp32_sampling_state` config.
+              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
+              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`, *optional*):
               The number of denoising steps.
           guidance_scale (`float`, *optional*):
-              Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is
-              forced to 1.0. Passing a value other than 1.0 raises an error.
+              Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is forced to
+              1.0. Passing a value other than 1.0 raises an error.
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
 
@@ -165,13 +168,18 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks):
         - `video2video`: requires `prompt`, `video`
 
       Components:
-          text_tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer
-          (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+          text_tokenizer (`AutoTokenizer`)
+          vae (`AutoencoderKLWan`)
+          video_processor (`VideoProcessor`)
+          transformer (`Cosmos3OmniTransformer`)
+          scheduler (`FlowMatchEulerDiscreteScheduler`)
 
       Configs:
-          default_use_system_prompt (default: True) enable_safety_checker (default: True)
-          default_use_fp32_sampling_state (default: True) is_distilled (default: True) distilled_sigmas (default:
-          None)
+          default_use_system_prompt (default: True)
+          enable_safety_checker (default: True)
+          default_use_fp32_sampling_state (default: True)
+          is_distilled (default: True)
+          distilled_sigmas (default: None)
 
       Inputs:
           prompt (`str`):
@@ -184,7 +192,7 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks):
               Width of the generated video or image in pixels.
           fps (`float`, *optional*, defaults to 24.0):
               Frame rate of the generated video.
-          use_system_prompt (`bool`, *optional*, defaults to True):
+          use_system_prompt (`bool | NoneType`, *optional*):
               Whether to prepend the Cosmos3 system prompt.
           add_resolution_template (`bool`, *optional*, defaults to True):
               Whether to add resolution metadata to the prompt.
@@ -207,13 +215,13 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks):
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
           use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the pipeline's
-              `default_use_fp32_sampling_state` config.
+              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
+              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`, *optional*):
               The number of denoising steps.
           guidance_scale (`float`, *optional*):
-              Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is
-              forced to 1.0. Passing a value other than 1.0 raises an error.
+              Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is forced to
+              1.0. Passing a value other than 1.0 raises an error.
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
           output_type (`str`, *optional*, defaults to pil):
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
index 1322434d17c7..22f8d34f8e55 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
@@ -492,12 +492,15 @@ def enable_sea_cache(self, config: SeaCacheConfig | None = None) -> None:
         current_config = getattr(transformer, "_cache_config", None)
         if current_config is not None:
             if isinstance(current_config, SeaCacheConfig):
-                self._prepare_sea_cache_config(current_config)
-                self._is_sea_cache_enabled = True
-                return
-            raise ValueError(
-                f"Caching is already enabled with {type(current_config).__name__}. Disable it before enabling SeaCache."
-            )
+                if config is None or config is current_config:
+                    self._prepare_sea_cache_config(current_config)
+                    self._is_sea_cache_enabled = True
+                    return
+                transformer.disable_cache()
+            else:
+                raise ValueError(
+                    f"Caching is already enabled with {type(current_config).__name__}. Disable it before enabling SeaCache."
+                )
 
         transformer.enable_cache(self._prepare_sea_cache_config(config))
         self._is_sea_cache_enabled = True
diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py
index d68885420137..895bf299ccf4 100644
--- a/tests/hooks/test_sea_cache.py
+++ b/tests/hooks/test_sea_cache.py
@@ -100,11 +100,16 @@ def _make_config(runtime, **kwargs):
     return SeaCacheConfig(**config_kwargs)
 
 
+def _get_root_hook(model):
+    return model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK)
+
+
 @torch.no_grad()
 def test_sea_cache_uses_independent_gates_and_histories_per_context():
     runtime = {"step": 0, "sigma": 0.9, "num_steps": 4}
     model = DummySeaTransformer()
     model.enable_cache(_make_config(runtime, threshold=0.5))
+    root_hook = _get_root_hook(model)
 
     first_input = torch.ones(2, 2, 2, 1)
     with model.cache_context("cond"):
@@ -129,11 +134,25 @@ def test_sea_cache_uses_independent_gates_and_histories_per_context():
     torch.testing.assert_close(uncond_output, changed_uncond_input * 8)
     assert [block.calls for block in model.layers] == [3, 3, 3]
     assert model.layers[0].indicator_norm.calls == 4
+    assert root_hook.num_full_steps == 3
+    assert root_hook.num_cached_steps == 1
+    stats = model.get_cache_stats()
+    assert stats["indicator_source"] == "first_block"
+    assert stats["residual_order"] == 1
+    assert stats["residual_boundary"] == "repeated_block_stack"
+    assert stats["transformer_calls"] == 4
+    assert stats["gate_evaluations"] == 4
+    assert stats["gate_trace"] == [True, True, False, True]
+    assert stats["branch_full_executions"] == {"cond": 1, "uncond": 2}
+    assert stats["branch_reuses"] == {"cond": 1}
+    assert stats["persistent_cache_bytes"] > 0
 
     model._reset_stateful_cache()
-    with model.cache_context("cond"):
-        model(first_input)
-    assert [block.calls for block in model.layers] == [4, 4, 4]
+    assert root_hook.num_full_steps == 0
+    assert root_hook.num_cached_steps == 0
+    archived_stats = model.get_cache_stats()
+    assert archived_stats["transformer_calls"] == 4
+    assert archived_stats["gate_trace"] == [True, True, False, True]
 
     model.disable_cache()
     assert not model.is_cache_enabled
@@ -158,6 +177,15 @@ def test_sea_cache_max_consecutive_cached_forces_full_per_context():
         with model.cache_context("cond"):
             model(torch.ones(1, 1, 1, 1))
 
+    stats = model.get_cache_stats()
+    assert stats["gate_trace"] == [True, False, False, True, False, False, True, False]
+    assert stats["actual_full_executions"] == 3
+    assert stats["actual_reuses"] == 5
+    assert stats["max_consecutive_cached"] == 2
+    assert stats["max_consecutive_cached_observed"] == 2
+    assert stats["max_consecutive_forced_full"] == 2
+    assert stats["per_branch"]["cond"]["max_consecutive_cached_observed"] == 2
+    assert stats["per_branch"]["cond"]["max_consecutive_forced_full"] == 2
     assert [block.calls for block in model.layers] == [3, 3, 3]
 
 
@@ -190,7 +218,11 @@ def raw_vision(module, args, kwargs):
     with model.cache_context("cond"):
         model(second_input)
 
-    assert [block.calls for block in model.layers] == [2, 2, 2]
+    stats = model.get_cache_stats()
+    assert stats["indicator_source"] == "raw_vision_latents"
+    assert stats["gate_trace"] == [True, True]
+    assert stats["actual_full_executions"] == 2
+    assert stats["actual_reuses"] == 0
     assert model.layers[0].indicator_norm.calls == 0
 
 
@@ -215,7 +247,34 @@ def test_sea_cache_residual_order_one_uses_actual_full_step_history():
 
     # Full residuals are 7 and 14 at steps 0 and 1, so linear extrapolation predicts 21 at step 2.
     torch.testing.assert_close(output, torch.full_like(output, 24.0))
-    assert [block.calls for block in model.layers] == [2, 2, 2]
+    root_hook = _get_root_hook(model)
+    assert root_hook.num_full_steps == 2
+    assert root_hook.num_cached_steps == 1
+
+
+@torch.no_grad()
+def test_sea_cache_replays_gate_schedule_and_reports_natural_mismatches():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 3}
+    model = DummySeaTransformer()
+    model.enable_cache(
+        _make_config(
+            runtime,
+            threshold=0.0,
+            gate_schedule=(True, False, True),
+        )
+    )
+
+    for step, sigma in enumerate((0.9, 0.6, 0.3)):
+        runtime.update(step=step, sigma=sigma)
+        with model.cache_context("cond"):
+            model(torch.full((1, 1, 1, 1), float(step + 1)))
+
+    stats = model.get_cache_stats()
+    assert stats["gate_trace"] == [True, False, True]
+    assert stats["gate_schedule_replayed"]
+    assert stats["gate_schedule_mismatches"] == 1
+    assert stats["actual_full_executions"] == 2
+    assert stats["actual_reuses"] == 1
 
 
 @torch.no_grad()
@@ -231,7 +290,7 @@ def test_cache_context_registry_is_refreshed_when_cache_is_enabled_after_an_unca
     with model.cache_context("cond"):
         model(torch.ones(1, 1, 1, 1))
 
-    assert model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) is not None
+    assert model.get_cache_stats()["branch_full_executions"] == {"cond": 1}
 
 
 @torch.no_grad()
@@ -249,26 +308,54 @@ def test_sea_cache_fails_open_without_vision_metadata():
         model(torch.ones(1, 1, 1, 1))
 
     assert [block.calls for block in model.layers] == [2, 2, 2]
+    root_hook = _get_root_hook(model)
+    assert root_hook.num_full_steps == 2
+    assert root_hook.num_cached_steps == 0
 
 
 @torch.no_grad()
-def test_sea_cache_fails_open_for_non_adjacent_steps_and_shape_changes():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 4}
+def test_sea_cache_non_adjacent_steps_start_a_new_residual_trajectory():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
+    config = _make_config(runtime, threshold=0.0, retention_steps=0, residual_order=1)
     model = DummySeaTransformer()
-    model.enable_cache(_make_config(runtime, residual_order=1))
+    model.enable_cache(config)
 
     with model.cache_context("cond"):
         model(torch.ones(1, 1, 1, 1))
+    runtime.update(step=1, sigma=0.5)
+    with model.cache_context("cond"):
+        model(torch.full((1, 1, 1, 1), 100.0))
 
-    runtime.update(step=2, sigma=0.5)
+    config.threshold = 100.0
+    runtime.update(step=0, sigma=0.9)
     with model.cache_context("cond"):
-        model(torch.ones(1, 2, 1, 1))
+        model(torch.full((1, 1, 1, 1), 2.0))
+    runtime.update(step=1, sigma=0.5)
+    with model.cache_context("cond"):
+        _, output = model(torch.full((1, 1, 1, 1), 2.0))
+
+    torch.testing.assert_close(output, torch.full_like(output, 16.0))
+    assert [block.calls for block in model.layers] == [3, 3, 3]
+    assert model.get_cache_stats()["gate_trace"] == [True, True, True, False]
+
+
+@torch.no_grad()
+def test_sea_cache_fails_open_for_shape_changes():
+    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
+    model = DummySeaTransformer()
+    model.enable_cache(_make_config(runtime, residual_order=1))
+
+    with model.cache_context("cond"):
+        model(torch.ones(1, 1, 1, 1))
 
-    runtime.update(step=3, sigma=0.2)
+    runtime.update(step=1, sigma=0.5)
     with model.cache_context("cond"):
         model(torch.ones(1, 2, 1, 1))
 
-    assert [block.calls for block in model.layers] == [3, 3, 3]
+    root_hook = _get_root_hook(model)
+    assert root_hook.num_full_steps == 2
+    assert root_hook.num_cached_steps == 0
+    assert [block.calls for block in model.layers] == [2, 2, 2]
 
 
 def test_sea_cache_is_inference_only_and_fails_open_with_autograd():
@@ -276,12 +363,17 @@ def test_sea_cache_is_inference_only_and_fails_open_with_autograd():
     model = DummySeaTransformer()
     model.enable_cache(_make_config(runtime))
 
-    with model.cache_context("cond"):
+    with torch.enable_grad(), model.cache_context("cond"):
         model(torch.ones(1, 1, 1, 1, requires_grad=True))
     runtime.update(step=1, sigma=0.5)
-    with model.cache_context("cond"):
+    with torch.enable_grad(), model.cache_context("cond"):
         model(torch.ones(1, 1, 1, 1, requires_grad=True))
 
+    root_hook = _get_root_hook(model)
+    assert root_hook.shared_state.transformer_calls == 2
+    assert root_hook.shared_state.fail_open_calls == 2
+    assert root_hook.num_full_steps == 2
+    assert root_hook.num_cached_steps == 0
     assert [block.calls for block in model.layers] == [2, 2, 2]
 
 
@@ -330,6 +422,8 @@ def test_sea_cache_single_block_supports_full_and_cached_execution_then_disables
         _, cached_output = model(torch.full((1, 1, 1, 1), 2.0))
     torch.testing.assert_close(cached_output, torch.full_like(cached_output, 3.0))
     assert model.layers[0].calls == 1
+    assert model.get_cache_stats()["actual_full_executions"] == 1
+    assert model.get_cache_stats()["actual_reuses"] == 1
 
     model.disable_cache()
 
@@ -353,6 +447,10 @@ def test_sea_cache_fails_open_for_parameter_sharded_blocks():
     with model.cache_context("cond"):
         model(torch.ones(1, 1, 1, 1))
 
+    stats = model.get_cache_stats()
+    assert stats["actual_full_executions"] == 2
+    assert stats["actual_reuses"] == 0
+    assert stats["fail_open_calls"] == 2
     assert [block.calls for block in model.layers] == [2, 2, 2]
 
 
@@ -403,6 +501,11 @@ def test_sea_cache_config_validation(kwargs, message):
         SeaCacheConfig(**kwargs)
 
 
+def test_sea_cache_gate_schedule_validation():
+    with pytest.raises(TypeError, match="gate_schedule"):
+        SeaCacheConfig(gate_schedule=(True, 1))
+
+
 @pytest.mark.parametrize("callback_name", ["metadata_callback", "raw_vision_callback"])
 def test_sea_cache_callback_validation(callback_name):
     with pytest.raises(TypeError, match=callback_name):
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
index 395f32d9cc86..d3afd2e33a97 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
@@ -21,7 +21,7 @@
 import torch
 from PIL import Image
 
-from diffusers import ModularPipeline, SeaCacheConfig, UniPCMultistepScheduler
+from diffusers import CosmosActionCondition, ModularPipeline, SeaCacheConfig, UniPCMultistepScheduler
 from diffusers.modular_pipelines import (
     Cosmos3OmniBlocks,
     Cosmos3OmniModularPipeline,
@@ -30,12 +30,16 @@
 from diffusers.modular_pipelines.cosmos.before_denoise import (
     Cosmos3ActionDenoiseInputStep,
     Cosmos3ActionPackSequenceStep,
+    Cosmos3ActionPrepareLatentsStep,
     Cosmos3SetTimestepsStep,
     Cosmos3SoundDenoiseInputStep,
+    Cosmos3SoundPrepareLatentsStep,
+    Cosmos3TransferPrepareLatentsStep,
     Cosmos3VisionDenoiseInputStep,
     Cosmos3VisionPackSequenceStep,
 )
 from diffusers.modular_pipelines.cosmos.encoders import Cosmos3TextEncoderStep
+from diffusers.modular_pipelines.cosmos.modular_blocks_cosmos3 import Cosmos3TransferChunkDenoiseStep
 
 from ...testing_utils import torch_device
 from ..testing_utils import (
@@ -205,6 +209,119 @@ def test_default_sea_cache_can_be_disabled_and_reenabled(self):
         pipe.enable_sea_cache()
         assert isinstance(pipe.transformer._cache_config, SeaCacheConfig)
 
+    def test_default_sea_cache_reuses_transformer_execution(self):
+        pipe = self.get_pipeline().to(torch_device)
+        scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, use_flow_sigmas=True)
+        pipe.update_components(scheduler=scheduler, use_native_flow_schedule=True)
+        pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6))
+        inputs = self.get_dummy_inputs()
+        inputs["num_inference_steps"] = 4
+
+        pipe(**inputs, output=self.output_name)
+
+        stats = pipe.transformer.get_cache_stats()
+        assert stats["actual_reuses"] > 0
+        assert stats["actual_full_executions"] > 0
+
+    def _get_sampling_state_block_pipe(self, block):
+        pipe = block.init_pipeline(self.pretrained_model_name_or_path)
+        pipe.load_components(torch_dtype=torch.bfloat16)
+        pipe.to(torch_device)
+        return pipe
+
+    @pytest.mark.parametrize(
+        ("use_fp32_sampling_state", "expected_dtype"),
+        [(False, torch.bfloat16), (True, torch.float32)],
+    )
+    def test_sound_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype):
+        pipe = self._get_sampling_state_block_pipe(Cosmos3SoundPrepareLatentsStep())
+
+        outputs = pipe(
+            num_frames=5,
+            fps=24.0,
+            generator=self.get_generator(0),
+            use_fp32_sampling_state=use_fp32_sampling_state,
+            output=["sound_latents", "sound_condition_mask"],
+        )
+
+        assert outputs["sound_latents"].dtype == expected_dtype
+        assert outputs["sound_condition_mask"].dtype == expected_dtype
+
+    @pytest.mark.parametrize(
+        ("use_fp32_sampling_state", "expected_dtype"),
+        [(False, torch.bfloat16), (True, torch.float32)],
+    )
+    def test_action_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype):
+        pipe = self._get_sampling_state_block_pipe(Cosmos3ActionPrepareLatentsStep())
+        action = CosmosActionCondition(
+            mode="policy",
+            chunk_size=2,
+            domain_name="av",
+            image=torch.zeros(3, 16, 16),
+        )
+
+        outputs = pipe(
+            action=action,
+            action_condition_frame_indexes=[],
+            generator=self.get_generator(0),
+            use_fp32_sampling_state=use_fp32_sampling_state,
+            output=["action_latents", "action_condition_mask"],
+        )
+
+        assert outputs["action_latents"].dtype == expected_dtype
+        assert outputs["action_condition_mask"].dtype == expected_dtype
+
+    @pytest.mark.parametrize(
+        ("use_fp32_sampling_state", "expected_dtype"),
+        [(False, torch.bfloat16), (True, torch.float32)],
+    )
+    def test_transfer_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype):
+        pipe = self._get_sampling_state_block_pipe(Cosmos3TransferPrepareLatentsStep())
+
+        outputs = pipe(
+            x0_tokens_vision=torch.zeros(1, 4, 2, 2, 2),
+            current_conditional_frames=1,
+            generator=self.get_generator(0),
+            use_fp32_sampling_state=use_fp32_sampling_state,
+            output=["latents", "velocity_mask", "condition_latents"],
+        )
+
+        assert outputs["latents"].dtype == expected_dtype
+        assert outputs["velocity_mask"].dtype == expected_dtype
+        assert outputs["condition_latents"].dtype == expected_dtype
+
+    def test_transfer_chunks_reset_stateful_cache_at_boundaries(self):
+        block = Cosmos3TransferChunkDenoiseStep()
+        child_block = mock.Mock(side_effect=lambda components, state: (components, state))
+        block.sub_blocks = {"child": child_block}
+        components = mock.Mock()
+        state = mock.Mock()
+        state.get.return_value = 3
+
+        block(components, state)
+
+        assert child_block.call_count == 3
+        assert components.transformer._reset_stateful_cache.call_count == 2
+        assert [call.args for call in state.set.call_args_list if call.args[0] == "chunk_id"] == [
+            ("chunk_id", 0),
+            ("chunk_id", 1),
+            ("chunk_id", 2),
+        ]
+
+    @pytest.mark.parametrize(
+        ("use_fp32_sampling_state", "expected_dtype"),
+        [(False, torch.bfloat16), (True, torch.float32)],
+    )
+    def test_sampling_state_controls_modular_cfg_and_scheduler_dtype(self, use_fp32_sampling_state, expected_dtype):
+        pipe = self.get_pipeline(dtype=torch.bfloat16).to(torch_device)
+        inputs = self.get_dummy_inputs()
+        inputs["use_fp32_sampling_state"] = use_fp32_sampling_state
+
+        outputs = pipe(**inputs, output=["velocity_vision", "latents"])
+
+        assert outputs["velocity_vision"].dtype == expected_dtype
+        assert outputs["latents"].dtype == expected_dtype
+
     def test_vae_encoder_is_standalone_and_validates_conditioning_inputs(self):
         pipe = self.get_pipeline()
         vae_encoder = pipe.blocks.sub_blocks["vae_encoder"]
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
index 566343115ca8..eb94dc80b010 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
@@ -17,12 +17,10 @@
 import torch
 from PIL import Image
 
-from diffusers import FlowMatchEulerDiscreteScheduler, ModularPipeline
+from diffusers import ModularPipeline, SeaCacheConfig
 from diffusers.modular_pipelines import Cosmos3DistilledBlocks, Cosmos3DistilledModularPipeline
 from diffusers.modular_pipelines.cosmos.before_denoise import Cosmos3VisionPrepareLatentsStep
-from diffusers.modular_pipelines.cosmos.denoise import Cosmos3DistilledVisionLoopSchedulerStep
 from diffusers.modular_pipelines.cosmos.encoders import Cosmos3DistilledTextEncoderStep
-from diffusers.modular_pipelines.modular_pipeline import BlockState, PipelineState
 
 from ...testing_utils import torch_device
 from ..testing_utils import (
@@ -120,6 +118,30 @@ def test_declares_distilled_configs(self):
         assert pipe.config.default_use_system_prompt is True
         assert pipe.config.default_use_fp32_sampling_state is True
 
+    def test_distilled_text_step_uses_system_prompt_config_fallback(self):
+        text_pipe = Cosmos3DistilledTextEncoderStep().init_pipeline(self.pretrained_model_name_or_path)
+        text_pipe.load_components()
+        text_pipe.disable_safety_checker()
+
+        inputs = {
+            "prompt": "A small robot moves across a table.",
+            "num_frames": 5,
+            "height": 32,
+            "width": 32,
+        }
+        default_with_system_prompt = text_pipe(**inputs, output="cond_input_ids")
+        explicit_with_system_prompt = text_pipe(**inputs, use_system_prompt=True, output="cond_input_ids")
+        explicit_without_system_prompt = text_pipe(**inputs, use_system_prompt=False, output="cond_input_ids")
+
+        text_pipe.update_components(default_use_system_prompt=False)
+        default_without_system_prompt = text_pipe(**inputs, output="cond_input_ids")
+        updated_with_system_prompt = text_pipe(**inputs, use_system_prompt=True, output="cond_input_ids")
+        updated_without_system_prompt = text_pipe(**inputs, use_system_prompt=False, output="cond_input_ids")
+
+        assert default_with_system_prompt == explicit_with_system_prompt == updated_with_system_prompt
+        assert explicit_without_system_prompt == default_without_system_prompt == updated_without_system_prompt
+        assert len(default_with_system_prompt) > len(default_without_system_prompt)
+
     @pytest.mark.parametrize(
         ("config_enabled", "input_enabled", "expected_dtype"),
         [
@@ -128,57 +150,48 @@ def test_declares_distilled_configs(self):
             (True, False, torch.bfloat16),
         ],
     )
-    def test_prepare_vision_latents_fp32_sampling_state(
-        self, config_enabled, input_enabled, expected_dtype
-    ):
-        components = BlockState(
-            _execution_device=torch.device("cpu"),
-            transformer=BlockState(dtype=torch.bfloat16),
-            config=BlockState(default_use_fp32_sampling_state=config_enabled),
-            vae_scale_factor_spatial=16,
-            vae_scale_factor_temporal=4,
-            num_channels_latents=4,
+    def test_prepare_vision_latents_fp32_sampling_state(self, config_enabled, input_enabled, expected_dtype):
+        prepare_pipe = Cosmos3VisionPrepareLatentsStep().init_pipeline(self.pretrained_model_name_or_path)
+        prepare_pipe.load_components(torch_dtype=torch.bfloat16)
+        prepare_pipe.update_components(default_use_fp32_sampling_state=config_enabled)
+        prepare_pipe.to(torch_device)
+
+        outputs = prepare_pipe(
+            num_frames=5,
+            height=32,
+            width=32,
+            fps=24.0,
+            generator=self.get_generator(0),
+            use_fp32_sampling_state=input_enabled,
+            output=["latents", "vision_condition_mask", "use_fp32_sampling_state"],
         )
-        state = PipelineState()
-        state.set("x0_tokens_vision", None)
-        state.set("vision_condition_frames", None)
-        state.set("num_frames", 5)
-        state.set("height", 32)
-        state.set("width", 32)
-        state.set("fps", 24.0)
-        state.set("latents", None)
-        state.set("generator", torch.Generator("cpu").manual_seed(0))
-        state.set("use_fp32_sampling_state", input_enabled)
-
-        Cosmos3VisionPrepareLatentsStep()(components, state)
-
-        assert state.get("latents").dtype == expected_dtype
-        assert state.get("vision_condition_mask").dtype == expected_dtype
-        assert state.get("use_fp32_sampling_state") is (config_enabled if input_enabled is None else input_enabled)
+
+        assert outputs["latents"].dtype == expected_dtype
+        assert outputs["vision_condition_mask"].dtype == expected_dtype
+        assert outputs["use_fp32_sampling_state"] is (config_enabled if input_enabled is None else input_enabled)
 
     @pytest.mark.parametrize(
         ("use_fp32_sampling_state", "expected_dtype"),
         [(False, torch.bfloat16), (True, torch.float32)],
     )
     def test_distilled_scheduler_fp32_state(self, use_fp32_sampling_state, expected_dtype):
-        scheduler = FlowMatchEulerDiscreteScheduler(stochastic_sampling=True)
-        scheduler.set_timesteps(sigmas=[1.0, 0.5])
-        latents = torch.zeros((1, 2, 1, 1, 1), dtype=torch.bfloat16)
-        block_state = BlockState(
-            latents=latents,
-            velocity_vision=torch.zeros_like(latents),
-            vision_condition_mask=torch.zeros((1, 1, 1), dtype=expected_dtype),
-            vision_conditioning_latents=None,
-            vision_condition_indexes_for_pack=[],
-            generator=torch.Generator("cpu").manual_seed(0),
-            use_fp32_sampling_state=use_fp32_sampling_state,
-        )
+        pipe = self.get_pipeline(torch_dtype=torch.bfloat16).to(torch_device)
+        inputs = self.get_dummy_inputs()
+        inputs["use_fp32_sampling_state"] = use_fp32_sampling_state
 
-        Cosmos3DistilledVisionLoopSchedulerStep()(
-            BlockState(scheduler=scheduler), block_state, i=0, t=scheduler.timesteps[0]
-        )
+        latents = pipe(**inputs, output=self.output_name)
+
+        assert latents.dtype == expected_dtype
+
+    def test_default_sea_cache_reuses_transformer_execution(self):
+        pipe = self.get_pipeline().to(torch_device)
+        pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6))
+
+        pipe(**self.get_dummy_inputs(), output=self.output_name)
 
-        assert block_state.latents.dtype == expected_dtype
+        stats = pipe.transformer.get_cache_stats()
+        assert stats["actual_reuses"] > 0
+        assert stats["actual_full_executions"] > 0
 
     def test_vae_encoder_rejects_image_and_video_together(self):
         vae_encoder = Cosmos3DistilledBlocks().sub_blocks["vae_encoder"]
diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py
index da7c2725018d..a23e774e728c 100644
--- a/tests/pipelines/cosmos/test_cosmos3.py
+++ b/tests/pipelines/cosmos/test_cosmos3.py
@@ -146,6 +146,27 @@ def test_default_sea_cache_can_be_disabled_and_reenabled(self):
         replacement_pipeline._maybe_enable_sea_cache()
         self.assertEqual(replacement_pipeline.transformer._cache_config.current_step_callback(), 3)
 
+        replacement_config = SeaCacheConfig(threshold=0.1)
+        replacement_pipeline.enable_sea_cache(replacement_config)
+        self.assertIs(replacement_pipeline.transformer._cache_config, replacement_config)
+        self.assertEqual(replacement_pipeline.transformer._cache_config.threshold, 0.1)
+        self.assertEqual(replacement_pipeline.transformer._cache_config.current_step_callback(), 3)
+
+    def test_default_sea_cache_reuses_transformer_execution(self):
+        pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
+        pipeline.set_progress_bar_config(disable=None)
+        pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True)
+        pipeline.register_to_config(use_native_flow_schedule=True)
+        pipeline.enable_sea_cache(SeaCacheConfig(threshold=1e6))
+        inputs = self.get_dummy_inputs(torch_device)
+        inputs.update(num_inference_steps=4, output_type="latent", enable_safety_check=False)
+
+        pipeline(**inputs)
+
+        stats = pipeline.transformer.get_cache_stats()
+        self.assertGreater(stats["actual_reuses"], 0)
+        self.assertGreater(stats["actual_full_executions"], 0)
+
     def test_fp32_sampling_state_keeps_transformer_inputs_in_model_dtype(self):
         pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
         pipeline.transformer.to(dtype=torch.float16)

From 3e68be91d627a1d084d33cd2cf97cd1d06c62720 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Mon, 31 Aug 2026 13:22:33 -0700
Subject: [PATCH 10/14] Remove SeaCache stats and ablation instrumentation

---
 src/diffusers/hooks/sea_cache.py              | 159 +-----------------
 tests/hooks/test_sea_cache.py                 |  92 +---------
 .../cosmos/test_modular_pipeline_cosmos3.py   |  13 +-
 ...test_modular_pipeline_cosmos3_distilled.py |  13 +-
 tests/pipelines/cosmos/test_cosmos3.py        |  13 +-
 tests/testing_utils.py                        |  32 ++++
 6 files changed, 58 insertions(+), 264 deletions(-)

diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py
index 93e2164dfccc..19d9ce0817f6 100644
--- a/src/diffusers/hooks/sea_cache.py
+++ b/src/diffusers/hooks/sea_cache.py
@@ -12,12 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import copy
 import inspect
 import math
-import time
 from dataclasses import dataclass
-from typing import Any, Callable, Literal, Sequence
+from typing import Any, Callable, Literal
 
 import torch
 
@@ -79,9 +77,6 @@ class SeaCacheConfig:
         raw_vision_callback (`Callable`, *optional*):
             Advanced model adapter returning raw vision latents with shape `(C, T, H, W)` for the
             `"raw_vision_latents"` indicator. Cosmos 3 uses its native adapter when this is omitted.
-        gate_schedule (`Sequence[bool]`, *optional*):
-            Per-step full-compute decisions to replay for controlled residual-prediction ablations. SEA still evaluates
-            its natural decision and reports mismatches; invalid or unsafe calls always fail open to full compute.
 
     Example:
         ```python
@@ -116,7 +111,6 @@ class SeaCacheConfig:
         [torch.nn.Module, tuple[Any, ...], dict[str, Any]],
         list[torch.Tensor] | None,
     ] = None
-    gate_schedule: Sequence[bool] = None
 
     def __post_init__(self):
         if not math.isfinite(self.threshold) or self.threshold < 0:
@@ -151,10 +145,6 @@ def __post_init__(self):
             callback = getattr(self, name)
             if callback is not None and not callable(callback):
                 raise TypeError(f"`{name}` must be callable or `None`.")
-        if self.gate_schedule is not None:
-            self.gate_schedule = tuple(self.gate_schedule)
-            if any(not isinstance(value, bool) for value in self.gate_schedule):
-                raise TypeError("`gate_schedule` must contain only boolean values.")
 
 
 @dataclass
@@ -174,8 +164,6 @@ def __init__(self):
         self.previous_indicator: list[torch.Tensor] | None = None
         self.accumulated_distance = 0.0
         self.consecutive_cached = 0
-        self.max_consecutive_cached_observed = 0
-        self.max_consecutive_forced_full = 0
         self.skip_remaining = False
         self.full_execution_pending = False
         self.cacheable_execution = False
@@ -184,7 +172,6 @@ def __init__(self):
         self.und_output: torch.Tensor | None = None
         self.cached_und_output: torch.Tensor | None = None
         self.cached_gen_residual: torch.Tensor | None = None
-        self.full_started_at: float | None = None
 
     def reset_forward(self):
         self.skip_remaining = False
@@ -195,7 +182,6 @@ def reset_forward(self):
         self.und_output = None
         self.cached_und_output = None
         self.cached_gen_residual = None
-        self.full_started_at = None
 
     def reset_trajectory(self):
         self.history = []
@@ -207,8 +193,6 @@ def reset_trajectory(self):
 
     def reset(self):
         self.reset_trajectory()
-        self.max_consecutive_cached_observed = 0
-        self.max_consecutive_forced_full = 0
         self.reset_forward()
 
 
@@ -219,20 +203,6 @@ def __init__(self):
 
     def reset(self):
         self.forward_metadata: _SeaCacheForwardMetadata | None = None
-        self.transformer_calls = 0
-        self.gate_evaluations = 0
-        self.gate_full_decisions = 0
-        self.gate_skip_decisions = 0
-        self.gate_trace: list[bool] = []
-        self.gate_schedule_mismatches = 0
-        self.num_full_steps = 0
-        self.num_cached_steps = 0
-        self.fail_open_calls = 0
-        self.indicator_seconds = 0.0
-        self.decision_seconds = 0.0
-        self.full_seconds = 0.0
-        self.branch_full_executions: dict[str, int] = {}
-        self.branch_reuses: dict[str, int] = {}
 
     def warn_once(self, message: str):
         if message not in self._warned_messages:
@@ -240,7 +210,6 @@ def warn_once(self, message: str):
             self._warned_messages.add(message)
 
     def mark_fail_open(self, message: str):
-        self.fail_open_calls += 1
         self.warn_once(message)
 
     def resolve_gate(
@@ -254,7 +223,6 @@ def resolve_gate(
         if state.gate_key == gate_key:
             return state.gate_should_compute
 
-        self.gate_evaluations += 1
         is_non_adjacent = state.gate_key is not None and metadata.step_index != state.gate_key[0] + 1
         if is_non_adjacent:
             self.mark_fail_open("SeaCache received non-adjacent scheduler steps; running full.")
@@ -299,37 +267,14 @@ def resolve_gate(
 
         should_compute = natural_should_compute
 
-        if config.gate_schedule is not None:
-            schedule_is_valid = len(config.gate_schedule) == metadata.num_inference_steps
-            scheduled_compute = bool(config.gate_schedule[metadata.step_index]) if schedule_is_valid else True
-            can_replay_skip = not (invalid_gate or is_retained or is_in_cache_end or is_first_observation)
-            if not schedule_is_valid:
-                self.mark_fail_open(
-                    "SeaCache gate schedule length does not match the number of inference steps; running full."
-                )
-                should_compute = True
-            elif scheduled_compute or can_replay_skip:
-                if natural_should_compute != scheduled_compute:
-                    self.gate_schedule_mismatches += 1
-                should_compute = scheduled_compute
-            else:
-                self.mark_fail_open("SeaCache gate schedule requested an unsafe cache hit; running full.")
-                should_compute = True
-
         if is_max_consecutive:
             should_compute = True
-            state.max_consecutive_forced_full += 1
 
         state.accumulated_distance = 0.0 if should_compute else candidate_accumulated_distance
 
         state.gate_key = gate_key
         state.gate_should_compute = should_compute
         state.previous_indicator = None if indicator is None else [value.detach() for value in indicator]
-        self.gate_trace.append(should_compute)
-        if should_compute:
-            self.gate_full_decisions += 1
-        else:
-            self.gate_skip_decisions += 1
         return should_compute
 
 
@@ -375,18 +320,11 @@ def _get_block_outputs(
 
 def _record_full_execution(
     config: SeaCacheConfig,
-    state_manager: StateManager,
-    shared_state: SeaCacheSharedState,
     state: SeaCacheContextState,
     gen_output: torch.Tensor,
     und_output: torch.Tensor | None,
 ) -> None:
-    shared_state.num_full_steps += 1
     state.consecutive_cached = 0
-    branch = state_manager._current_context
-    shared_state.branch_full_executions[branch] = shared_state.branch_full_executions.get(branch, 0) + 1
-    if state.full_started_at is not None:
-        shared_state.full_seconds += time.perf_counter() - state.full_started_at
     if (
         state.cacheable_execution
         and state.step_index is not None
@@ -570,7 +508,6 @@ def __init__(
         shared_state: SeaCacheSharedState,
         metadata_callback: Callable,
         raw_vision_callback: Callable,
-        residual_boundary: str,
     ):
         super().__init__()
         self.config = config
@@ -578,73 +515,9 @@ def __init__(
         self.shared_state = shared_state
         self.metadata_callback = metadata_callback
         self.raw_vision_callback = raw_vision_callback
-        self.residual_boundary = residual_boundary
-        self._last_stats: dict[str, Any] | None = None
-
-    @property
-    def num_full_steps(self) -> int:
-        return self.shared_state.num_full_steps
-
-    @property
-    def num_cached_steps(self) -> int:
-        return self.shared_state.num_cached_steps
-
-    def stats(self) -> dict[str, Any]:
-        if self.shared_state.transformer_calls == 0 and self._last_stats is not None:
-            return copy.deepcopy(self._last_stats)
-
-        persistent_cache_bytes = 0
-        per_branch = {}
-        for branch, state in sorted(self.state_manager._state_cache.items()):
-            if state.previous_indicator is not None:
-                persistent_cache_bytes += sum(
-                    value.numel() * value.element_size() for value in state.previous_indicator
-                )
-            for _, und_output, gen_residual in state.history:
-                persistent_cache_bytes += und_output.numel() * und_output.element_size()
-                persistent_cache_bytes += gen_residual.numel() * gen_residual.element_size()
-            per_branch[branch] = {
-                "full_calls": self.shared_state.branch_full_executions.get(branch, 0),
-                "reuse_calls": self.shared_state.branch_reuses.get(branch, 0),
-                "max_consecutive_cached_observed": state.max_consecutive_cached_observed,
-                "max_consecutive_forced_full": state.max_consecutive_forced_full,
-            }
-
-        opportunities = self.shared_state.num_full_steps + self.shared_state.num_cached_steps
-        return {
-            "indicator_source": self.config.indicator_source,
-            "residual_order": self.config.residual_order,
-            "residual_boundary": self.residual_boundary,
-            "max_consecutive_cached": self.config.max_consecutive_cached,
-            "max_consecutive_cached_observed": max(
-                (stats["max_consecutive_cached_observed"] for stats in per_branch.values()), default=0
-            ),
-            "max_consecutive_forced_full": sum(stats["max_consecutive_forced_full"] for stats in per_branch.values()),
-            "transformer_calls": self.shared_state.transformer_calls,
-            "gate_evaluations": self.shared_state.gate_evaluations,
-            "gate_full_decisions": self.shared_state.gate_full_decisions,
-            "gate_skip_decisions": self.shared_state.gate_skip_decisions,
-            "gate_trace": list(self.shared_state.gate_trace),
-            "gate_schedule_replayed": self.config.gate_schedule is not None,
-            "gate_schedule_mismatches": self.shared_state.gate_schedule_mismatches,
-            "actual_full_executions": self.shared_state.num_full_steps,
-            "actual_reuses": self.shared_state.num_cached_steps,
-            "actual_reuse_rate": (self.shared_state.num_cached_steps / opportunities if opportunities else 0.0),
-            "fail_open_calls": self.shared_state.fail_open_calls,
-            "sea_indicator_seconds": self.shared_state.indicator_seconds,
-            "sea_decision_seconds": self.shared_state.decision_seconds,
-            "sea_seconds": (self.shared_state.indicator_seconds + self.shared_state.decision_seconds),
-            "full_seconds": self.shared_state.full_seconds,
-            "timing_note": "host wall time; CUDA work is asynchronous",
-            "persistent_cache_bytes": persistent_cache_bytes,
-            "branch_full_executions": dict(sorted(self.shared_state.branch_full_executions.items())),
-            "branch_reuses": dict(sorted(self.shared_state.branch_reuses.items())),
-            "per_branch": per_branch,
-        }
 
     def pre_forward(self, module: torch.nn.Module, *args, **kwargs):
         self.shared_state.forward_metadata = None
-        self.shared_state.transformer_calls += 1
         if torch.is_grad_enabled():
             self.shared_state.mark_fail_open(
                 "SeaCache is inference-only; calls with autograd enabled run in fail-open mode."
@@ -743,8 +616,6 @@ def post_forward(self, module: torch.nn.Module, output: Any) -> Any:
         return output
 
     def reset_state(self, module: torch.nn.Module):
-        if self.shared_state.transformer_calls > 0:
-            self._last_stats = self.stats()
         self.state_manager.reset()
         self.shared_state.reset()
         return module
@@ -829,13 +700,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
 
         forward_metadata = self.shared_state.forward_metadata
         if state is None or forward_metadata is None:
-            if state is not None:
-                state.full_started_at = time.perf_counter()
             return self.fn_ref.original_forward(*args, **kwargs)
 
         state.step_index = forward_metadata.step_index
         state.cacheable_execution = True
-        indicator_started = time.perf_counter()
         indicator_error_reported = False
         if _is_parameter_sharded(module):
             self.shared_state.mark_fail_open(
@@ -852,14 +720,11 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
                 )
                 indicator = None
                 indicator_error_reported = True
-        self.shared_state.indicator_seconds += time.perf_counter() - indicator_started
         if indicator is None and not indicator_error_reported:
             self.shared_state.mark_fail_open(
                 "SeaCache could not construct its vision indicator; running in fail-open mode."
             )
-        decision_started = time.perf_counter()
         should_compute = self.shared_state.resolve_gate(state, forward_metadata, indicator, self.config)
-        self.shared_state.decision_seconds += time.perf_counter() - decision_started
 
         if should_compute or not state.history:
             if not should_compute:
@@ -867,7 +732,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
                 self.shared_state.mark_fail_open(
                     "SeaCache selected a cache hit without residual history; running in fail-open mode."
                 )
-            state.full_started_at = time.perf_counter()
             return self.fn_ref.original_forward(*args, **kwargs)
 
         residual_history = state.history[-(self.config.residual_order + 1) :]
@@ -889,7 +753,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
             self.shared_state.mark_fail_open(
                 "SeaCache residual history changed shape, device, or dtype; running in fail-open mode."
             )
-            state.full_started_at = time.perf_counter()
             return self.fn_ref.original_forward(*args, **kwargs)
 
         if self.config.residual_order == 1 and len(residual_history) >= 2:
@@ -904,10 +767,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
         state.cached_und_output = cached_und
         state.cached_gen_residual = cached_residual
         state.consecutive_cached += 1
-        state.max_consecutive_cached_observed = max(state.max_consecutive_cached_observed, state.consecutive_cached)
-        self.shared_state.num_cached_steps += 1
-        branch = self.state_manager._current_context
-        self.shared_state.branch_reuses[branch] = self.shared_state.branch_reuses.get(branch, 0) + 1
         if self.post_norm_boundary:
             return _build_block_output(self._metadata, hidden_states, encoder_hidden_states)
         return _build_block_output(self._metadata, hidden_states + cached_residual, cached_und)
@@ -952,8 +811,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
         hidden_states, encoder_hidden_states = _get_block_outputs(self._metadata, output)
         _record_full_execution(
             self.config,
-            self.state_manager,
-            self.shared_state,
             state,
             gen_output=hidden_states,
             und_output=encoder_hidden_states,
@@ -1006,8 +863,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
 
         _record_full_execution(
             self.config,
-            self.state_manager,
-            self.shared_state,
             state,
             gen_output=output,
             und_output=state.und_output,
@@ -1055,7 +910,6 @@ def apply_sea_cache(module: torch.nn.Module, config: SeaCacheConfig) -> None:
                 "residual boundary."
             )
     post_norm_boundary = post_norm_modules is not None
-    residual_boundary = "post_language_model_norm" if post_norm_boundary else "repeated_block_stack"
 
     blocks = []
     for name, submodule in unwrapped_module.named_children():
@@ -1085,7 +939,6 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None:
                 shared_state,
                 metadata_callback,
                 raw_vision_callback,
-                residual_boundary,
             ),
             _SEA_CACHE_ROOT_HOOK,
         )
@@ -1134,13 +987,3 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None:
         raise
 
     root_registry._child_registries_cache = None
-
-
-def get_sea_cache_stats(module: torch.nn.Module) -> dict[str, Any]:
-    """Return statistics for the SeaCache instance currently attached to ``module``."""
-
-    registry = getattr(module, "_diffusers_hook", None)
-    root_hook = registry.get_hook(_SEA_CACHE_ROOT_HOOK) if registry is not None else None
-    if not isinstance(root_hook, SeaCacheRootHook):
-        raise ValueError("SeaCache is not enabled on this module.")
-    return root_hook.stats()
diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py
index 895bf299ccf4..5a31edcd564f 100644
--- a/tests/hooks/test_sea_cache.py
+++ b/tests/hooks/test_sea_cache.py
@@ -100,16 +100,11 @@ def _make_config(runtime, **kwargs):
     return SeaCacheConfig(**config_kwargs)
 
 
-def _get_root_hook(model):
-    return model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK)
-
-
 @torch.no_grad()
 def test_sea_cache_uses_independent_gates_and_histories_per_context():
     runtime = {"step": 0, "sigma": 0.9, "num_steps": 4}
     model = DummySeaTransformer()
     model.enable_cache(_make_config(runtime, threshold=0.5))
-    root_hook = _get_root_hook(model)
 
     first_input = torch.ones(2, 2, 2, 1)
     with model.cache_context("cond"):
@@ -134,25 +129,8 @@ def test_sea_cache_uses_independent_gates_and_histories_per_context():
     torch.testing.assert_close(uncond_output, changed_uncond_input * 8)
     assert [block.calls for block in model.layers] == [3, 3, 3]
     assert model.layers[0].indicator_norm.calls == 4
-    assert root_hook.num_full_steps == 3
-    assert root_hook.num_cached_steps == 1
-    stats = model.get_cache_stats()
-    assert stats["indicator_source"] == "first_block"
-    assert stats["residual_order"] == 1
-    assert stats["residual_boundary"] == "repeated_block_stack"
-    assert stats["transformer_calls"] == 4
-    assert stats["gate_evaluations"] == 4
-    assert stats["gate_trace"] == [True, True, False, True]
-    assert stats["branch_full_executions"] == {"cond": 1, "uncond": 2}
-    assert stats["branch_reuses"] == {"cond": 1}
-    assert stats["persistent_cache_bytes"] > 0
 
     model._reset_stateful_cache()
-    assert root_hook.num_full_steps == 0
-    assert root_hook.num_cached_steps == 0
-    archived_stats = model.get_cache_stats()
-    assert archived_stats["transformer_calls"] == 4
-    assert archived_stats["gate_trace"] == [True, True, False, True]
 
     model.disable_cache()
     assert not model.is_cache_enabled
@@ -177,15 +155,7 @@ def test_sea_cache_max_consecutive_cached_forces_full_per_context():
         with model.cache_context("cond"):
             model(torch.ones(1, 1, 1, 1))
 
-    stats = model.get_cache_stats()
-    assert stats["gate_trace"] == [True, False, False, True, False, False, True, False]
-    assert stats["actual_full_executions"] == 3
-    assert stats["actual_reuses"] == 5
-    assert stats["max_consecutive_cached"] == 2
-    assert stats["max_consecutive_cached_observed"] == 2
-    assert stats["max_consecutive_forced_full"] == 2
-    assert stats["per_branch"]["cond"]["max_consecutive_cached_observed"] == 2
-    assert stats["per_branch"]["cond"]["max_consecutive_forced_full"] == 2
+    # With max_consecutive_cached=2, every third step forces a full execution: 3 full, 5 cached.
     assert [block.calls for block in model.layers] == [3, 3, 3]
 
 
@@ -218,11 +188,7 @@ def raw_vision(module, args, kwargs):
     with model.cache_context("cond"):
         model(second_input)
 
-    stats = model.get_cache_stats()
-    assert stats["indicator_source"] == "raw_vision_latents"
-    assert stats["gate_trace"] == [True, True]
-    assert stats["actual_full_executions"] == 2
-    assert stats["actual_reuses"] == 0
+    assert [block.calls for block in model.layers] == [2, 2, 2]
     assert model.layers[0].indicator_norm.calls == 0
 
 
@@ -247,34 +213,7 @@ def test_sea_cache_residual_order_one_uses_actual_full_step_history():
 
     # Full residuals are 7 and 14 at steps 0 and 1, so linear extrapolation predicts 21 at step 2.
     torch.testing.assert_close(output, torch.full_like(output, 24.0))
-    root_hook = _get_root_hook(model)
-    assert root_hook.num_full_steps == 2
-    assert root_hook.num_cached_steps == 1
-
-
-@torch.no_grad()
-def test_sea_cache_replays_gate_schedule_and_reports_natural_mismatches():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 3}
-    model = DummySeaTransformer()
-    model.enable_cache(
-        _make_config(
-            runtime,
-            threshold=0.0,
-            gate_schedule=(True, False, True),
-        )
-    )
-
-    for step, sigma in enumerate((0.9, 0.6, 0.3)):
-        runtime.update(step=step, sigma=sigma)
-        with model.cache_context("cond"):
-            model(torch.full((1, 1, 1, 1), float(step + 1)))
-
-    stats = model.get_cache_stats()
-    assert stats["gate_trace"] == [True, False, True]
-    assert stats["gate_schedule_replayed"]
-    assert stats["gate_schedule_mismatches"] == 1
-    assert stats["actual_full_executions"] == 2
-    assert stats["actual_reuses"] == 1
+    assert [block.calls for block in model.layers] == [2, 2, 2]
 
 
 @torch.no_grad()
@@ -290,7 +229,7 @@ def test_cache_context_registry_is_refreshed_when_cache_is_enabled_after_an_unca
     with model.cache_context("cond"):
         model(torch.ones(1, 1, 1, 1))
 
-    assert model.get_cache_stats()["branch_full_executions"] == {"cond": 1}
+    assert [block.calls for block in model.layers] == [2, 2, 2]
 
 
 @torch.no_grad()
@@ -308,9 +247,6 @@ def test_sea_cache_fails_open_without_vision_metadata():
         model(torch.ones(1, 1, 1, 1))
 
     assert [block.calls for block in model.layers] == [2, 2, 2]
-    root_hook = _get_root_hook(model)
-    assert root_hook.num_full_steps == 2
-    assert root_hook.num_cached_steps == 0
 
 
 @torch.no_grad()
@@ -336,7 +272,6 @@ def test_sea_cache_non_adjacent_steps_start_a_new_residual_trajectory():
 
     torch.testing.assert_close(output, torch.full_like(output, 16.0))
     assert [block.calls for block in model.layers] == [3, 3, 3]
-    assert model.get_cache_stats()["gate_trace"] == [True, True, True, False]
 
 
 @torch.no_grad()
@@ -352,9 +287,6 @@ def test_sea_cache_fails_open_for_shape_changes():
     with model.cache_context("cond"):
         model(torch.ones(1, 2, 1, 1))
 
-    root_hook = _get_root_hook(model)
-    assert root_hook.num_full_steps == 2
-    assert root_hook.num_cached_steps == 0
     assert [block.calls for block in model.layers] == [2, 2, 2]
 
 
@@ -369,11 +301,6 @@ def test_sea_cache_is_inference_only_and_fails_open_with_autograd():
     with torch.enable_grad(), model.cache_context("cond"):
         model(torch.ones(1, 1, 1, 1, requires_grad=True))
 
-    root_hook = _get_root_hook(model)
-    assert root_hook.shared_state.transformer_calls == 2
-    assert root_hook.shared_state.fail_open_calls == 2
-    assert root_hook.num_full_steps == 2
-    assert root_hook.num_cached_steps == 0
     assert [block.calls for block in model.layers] == [2, 2, 2]
 
 
@@ -422,8 +349,6 @@ def test_sea_cache_single_block_supports_full_and_cached_execution_then_disables
         _, cached_output = model(torch.full((1, 1, 1, 1), 2.0))
     torch.testing.assert_close(cached_output, torch.full_like(cached_output, 3.0))
     assert model.layers[0].calls == 1
-    assert model.get_cache_stats()["actual_full_executions"] == 1
-    assert model.get_cache_stats()["actual_reuses"] == 1
 
     model.disable_cache()
 
@@ -447,10 +372,6 @@ def test_sea_cache_fails_open_for_parameter_sharded_blocks():
     with model.cache_context("cond"):
         model(torch.ones(1, 1, 1, 1))
 
-    stats = model.get_cache_stats()
-    assert stats["actual_full_executions"] == 2
-    assert stats["actual_reuses"] == 0
-    assert stats["fail_open_calls"] == 2
     assert [block.calls for block in model.layers] == [2, 2, 2]
 
 
@@ -501,11 +422,6 @@ def test_sea_cache_config_validation(kwargs, message):
         SeaCacheConfig(**kwargs)
 
 
-def test_sea_cache_gate_schedule_validation():
-    with pytest.raises(TypeError, match="gate_schedule"):
-        SeaCacheConfig(gate_schedule=(True, 1))
-
-
 @pytest.mark.parametrize("callback_name", ["metadata_callback", "raw_vision_callback"])
 def test_sea_cache_callback_validation(callback_name):
     with pytest.raises(TypeError, match=callback_name):
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
index d3afd2e33a97..b4f9c78d6622 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
@@ -22,6 +22,7 @@
 from PIL import Image
 
 from diffusers import CosmosActionCondition, ModularPipeline, SeaCacheConfig, UniPCMultistepScheduler
+from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer
 from diffusers.modular_pipelines import (
     Cosmos3OmniBlocks,
     Cosmos3OmniModularPipeline,
@@ -41,7 +42,7 @@
 from diffusers.modular_pipelines.cosmos.encoders import Cosmos3TextEncoderStep
 from diffusers.modular_pipelines.cosmos.modular_blocks_cosmos3 import Cosmos3TransferChunkDenoiseStep
 
-from ...testing_utils import torch_device
+from ...testing_utils import count_transformer_cache_reuse, torch_device
 from ..testing_utils import (
     BaseModularPipelineTesterConfig,
     ModularLoadingTesterMixin,
@@ -213,15 +214,15 @@ def test_default_sea_cache_reuses_transformer_execution(self):
         pipe = self.get_pipeline().to(torch_device)
         scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, use_flow_sigmas=True)
         pipe.update_components(scheduler=scheduler, use_native_flow_schedule=True)
-        pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6))
         inputs = self.get_dummy_inputs()
         inputs["num_inference_steps"] = 4
 
-        pipe(**inputs, output=self.output_name)
+        with count_transformer_cache_reuse(pipe.transformer, Cosmos3VLTextMoTDecoderLayer) as counts:
+            pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6))
+            pipe(**inputs, output=self.output_name)
 
-        stats = pipe.transformer.get_cache_stats()
-        assert stats["actual_reuses"] > 0
-        assert stats["actual_full_executions"] > 0
+        assert counts["cached_steps"] > 0
+        assert counts["full_steps"] > 0
 
     def _get_sampling_state_block_pipe(self, block):
         pipe = block.init_pipeline(self.pretrained_model_name_or_path)
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
index eb94dc80b010..705ca74fd4d3 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
@@ -18,11 +18,12 @@
 from PIL import Image
 
 from diffusers import ModularPipeline, SeaCacheConfig
+from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer
 from diffusers.modular_pipelines import Cosmos3DistilledBlocks, Cosmos3DistilledModularPipeline
 from diffusers.modular_pipelines.cosmos.before_denoise import Cosmos3VisionPrepareLatentsStep
 from diffusers.modular_pipelines.cosmos.encoders import Cosmos3DistilledTextEncoderStep
 
-from ...testing_utils import torch_device
+from ...testing_utils import count_transformer_cache_reuse, torch_device
 from ..testing_utils import (
     BaseModularPipelineTesterConfig,
     ModularLoadingTesterMixin,
@@ -185,13 +186,13 @@ def test_distilled_scheduler_fp32_state(self, use_fp32_sampling_state, expected_
 
     def test_default_sea_cache_reuses_transformer_execution(self):
         pipe = self.get_pipeline().to(torch_device)
-        pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6))
 
-        pipe(**self.get_dummy_inputs(), output=self.output_name)
+        with count_transformer_cache_reuse(pipe.transformer, Cosmos3VLTextMoTDecoderLayer) as counts:
+            pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6))
+            pipe(**self.get_dummy_inputs(), output=self.output_name)
 
-        stats = pipe.transformer.get_cache_stats()
-        assert stats["actual_reuses"] > 0
-        assert stats["actual_full_executions"] > 0
+        assert counts["cached_steps"] > 0
+        assert counts["full_steps"] > 0
 
     def test_vae_encoder_rejects_image_and_video_together(self):
         vae_encoder = Cosmos3DistilledBlocks().sub_blocks["vae_encoder"]
diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py
index a23e774e728c..8c0b1e304625 100644
--- a/tests/pipelines/cosmos/test_cosmos3.py
+++ b/tests/pipelines/cosmos/test_cosmos3.py
@@ -28,9 +28,10 @@
     SeaCacheConfig,
     UniPCMultistepScheduler,
 )
+from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer
 from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import _preprocess_conditioning_image
 
-from ...testing_utils import enable_full_determinism, torch_device
+from ...testing_utils import count_transformer_cache_reuse, enable_full_determinism, torch_device
 from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
 from ..test_pipelines_common import PipelineTesterMixin
 
@@ -157,15 +158,15 @@ def test_default_sea_cache_reuses_transformer_execution(self):
         pipeline.set_progress_bar_config(disable=None)
         pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True)
         pipeline.register_to_config(use_native_flow_schedule=True)
-        pipeline.enable_sea_cache(SeaCacheConfig(threshold=1e6))
         inputs = self.get_dummy_inputs(torch_device)
         inputs.update(num_inference_steps=4, output_type="latent", enable_safety_check=False)
 
-        pipeline(**inputs)
+        with count_transformer_cache_reuse(pipeline.transformer, Cosmos3VLTextMoTDecoderLayer) as counts:
+            pipeline.enable_sea_cache(SeaCacheConfig(threshold=1e6))
+            pipeline(**inputs)
 
-        stats = pipeline.transformer.get_cache_stats()
-        self.assertGreater(stats["actual_reuses"], 0)
-        self.assertGreater(stats["actual_full_executions"], 0)
+        self.assertGreater(counts["cached_steps"], 0)
+        self.assertGreater(counts["full_steps"], 0)
 
     def test_fp32_sampling_state_keeps_transformer_inputs_in_model_dtype(self):
         pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
diff --git a/tests/testing_utils.py b/tests/testing_utils.py
index c35f975285c4..0c9766518f9f 100644
--- a/tests/testing_utils.py
+++ b/tests/testing_utils.py
@@ -198,6 +198,38 @@ def numpy_cosine_similarity_distance(a, b):
     return distance
 
 
+@contextmanager
+def count_transformer_cache_reuse(transformer, block_cls):
+    """Count real block executions and total transformer calls while running under a cache.
+
+    ``block_cls.forward`` is patched to count real executions (skipped blocks reuse cached outputs and never run),
+    and a forward pre-hook counts every transformer call. On exit the yielded dict is populated with ``full_steps``
+    (transformer calls that recomputed the blocks) and ``cached_steps`` (calls served from the cache).
+
+    The patch must be active before the cache is enabled, so enable the cache inside the ``with`` block.
+    """
+    counts = {"block_calls": 0, "transformer_calls": 0, "full_steps": 0, "cached_steps": 0}
+    original_forward = block_cls.forward
+
+    def counting_forward(self, *args, **kwargs):
+        counts["block_calls"] += 1
+        return original_forward(self, *args, **kwargs)
+
+    def counting_pre_hook(module, args):
+        counts["transformer_calls"] += 1
+
+    block_cls.forward = counting_forward
+    handle = transformer.register_forward_pre_hook(counting_pre_hook)
+    try:
+        yield counts
+    finally:
+        handle.remove()
+        block_cls.forward = original_forward
+        num_blocks = sum(1 for module in transformer.modules() if isinstance(module, block_cls))
+        counts["full_steps"] = counts["block_calls"] // num_blocks if num_blocks else 0
+        counts["cached_steps"] = counts["transformer_calls"] - counts["full_steps"]
+
+
 def check_if_dicts_are_equal(dict1, dict2):
     dict1, dict2 = dict1.copy(), dict2.copy()
 

From 21c75ca3d6b84f78dc02ba0032d36b217f7badb1 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Wed, 2 Sep 2026 07:30:33 -0700
Subject: [PATCH 11/14] Move SeaCache to transformer-level and disable by
 default

---
 docs/source/en/api/pipelines/cosmos3.md       | 35 ++++--------
 docs/source/en/optimization/cache.md          | 29 ++++------
 src/diffusers/hooks/sea_cache.py              |  3 +
 .../cosmos/modular_pipeline.py                | 15 -----
 .../pipelines/cosmos/pipeline_cosmos3_omni.py | 55 -------------------
 .../cosmos/test_modular_pipeline_cosmos3.py   | 34 +-----------
 ...test_modular_pipeline_cosmos3_distilled.py | 15 +----
 tests/pipelines/cosmos/test_cosmos3.py        | 46 +---------------
 tests/testing_utils.py                        | 32 -----------
 9 files changed, 30 insertions(+), 234 deletions(-)

diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md
index 84752ace65b8..dcfb7167db90 100644
--- a/docs/source/en/api/pipelines/cosmos3.md
+++ b/docs/source/en/api/pipelines/cosmos3.md
@@ -662,41 +662,28 @@ if result.action is not None:
 
 ## SeaCache
 
-Cosmos 3 inference uses [`SeaCacheConfig`] by default. SeaCache reuses transformer residuals when the
-Spectral-Evolution-Aware indicator changes slowly, reducing the number of full transformer executions. The default
-configuration filters raw vision latents, linearly extrapolates cached residuals, uses a `0.25` threshold, and forces a
-full execution after at most two consecutive cached steps.
-
-SeaCache is approximate and can change generated outputs. Disable it before inference when you need every denoising
-step to execute the full transformer, for example when producing an uncached baseline:
-
-```python
-pipe.disable_sea_cache()
-result = pipe(
-    prompt=prompt,
-    num_frames=189,
-    height=720,
-    width=1280,
-)
-```
-
-The same method works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and
-[`Cosmos3DistilledModularPipeline`]. Reenable the default configuration with `pipe.enable_sea_cache()`, or pass a
-custom configuration:
+SeaCache is disabled by default. Cosmos 3 supports enabling it explicitly with [`SeaCacheConfig`]. SeaCache reuses
+transformer residuals when the Spectral-Evolution-Aware indicator changes slowly, reducing the number of full
+transformer executions. Enable it on the transformer with scheduler metadata callbacks from the pipeline:
 
 ```python
 from diffusers import SeaCacheConfig
 
-pipe.enable_sea_cache(
+pipe.transformer.enable_cache(
     SeaCacheConfig(
         threshold=0.2,
         max_consecutive_cached=2,
+        current_step_callback=lambda: pipe.current_step_index,
+        current_sigma_callback=lambda: pipe.current_sigma,
+        num_inference_steps_callback=lambda: pipe.num_timesteps,
     )
 )
 ```
 
-The pipeline supplies SeaCache with the active scheduler step, sigma, and number of inference steps. Cache state is
-reset for each pipeline call, and conditional and unconditional guidance branches keep independent histories.
+The same model-level API works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and
+[`Cosmos3DistilledModularPipeline`]. SeaCache is approximate and can change generated outputs. Disable it with
+`pipe.transformer.disable_cache()` when you need every denoising step to execute the full transformer. Cache state is
+reset after each pipeline call, and conditional and unconditional guidance branches keep independent histories.
 
 ## Sampling precision
 
diff --git a/docs/source/en/optimization/cache.md b/docs/source/en/optimization/cache.md
index 4db41753507f..ad65fceeaada 100644
--- a/docs/source/en/optimization/cache.md
+++ b/docs/source/en/optimization/cache.md
@@ -74,35 +74,28 @@ pipeline.transformer.enable_cache(config)
 successive denoising steps. When the accumulated indicator change remains below a threshold, it skips the expensive
 transformer block stack and predicts its output from cached residuals.
 
-[`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and [`Cosmos3DistilledModularPipeline`] enable SeaCache
-automatically when inference starts. The default [`SeaCacheConfig`] filters raw vision latents, linearly extrapolates
-the cached residual, uses `threshold=0.25`, and allows at most two consecutive cached steps before forcing a full
-transformer execution.
-
-SeaCache is an approximate optimization and may change generated outputs. Disable it for full transformer execution,
-such as when measuring an uncached baseline:
+SeaCache is disabled by default. Enable it on the transformer and provide callbacks for the active scheduler step,
+sigma, and number of inference steps:
 
 ```python
-from diffusers import Cosmos3OmniPipeline
+from diffusers import Cosmos3OmniPipeline, SeaCacheConfig
 
 pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano")
-pipe.disable_sea_cache()
-```
-
-Call [`~Cosmos3OmniPipeline.enable_sea_cache`] to reenable it or provide a custom configuration. The pipeline fills in
-the scheduler callbacks required by SeaCache.
-
-```python
-from diffusers import SeaCacheConfig
-
-pipe.enable_sea_cache(
+pipe.transformer.enable_cache(
     SeaCacheConfig(
         threshold=0.2,
         max_consecutive_cached=2,
+        current_step_callback=lambda: pipe.current_step_index,
+        current_sigma_callback=lambda: pipe.current_sigma,
+        num_inference_steps_callback=lambda: pipe.num_timesteps,
     )
 )
 ```
 
+This model-level API works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and
+[`Cosmos3DistilledModularPipeline`]. SeaCache is an approximate optimization and may change generated outputs. Call
+`pipe.transformer.disable_cache()` when you need every denoising step to execute the full transformer.
+
 ## FirstBlockCache
 
 [FirstBlock Cache](https://huggingface.co/docs/diffusers/main/en/api/cache#diffusers.FirstBlockCacheConfig) checks how much the early layers of the denoiser changes from one timestep to the next. If the change is small, the model skips the expensive later layers and reuses the previous output.
diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py
index 19d9ce0817f6..f277f90f0ce7 100644
--- a/src/diffusers/hooks/sea_cache.py
+++ b/src/diffusers/hooks/sea_cache.py
@@ -39,6 +39,9 @@ class SeaCacheConfig:
     r"""
     Configuration for [SeaCache](https://huggingface.co/papers/2602.18993).
 
+    SeaCache is disabled by default and only activates after this configuration is passed to
+    `transformer.enable_cache(config)`.
+
     SeaCache compares Spectral-Evolution-Aware (SEA) indicators between scheduler steps. If their accumulated relative
     change stays below `threshold`, the expensive language-model hidden transform is replaced with a cached residual.
     For Cosmos 3, the residual spans the decoder stack and final pathway normalization; input packing and modality
diff --git a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py
index 5121e331306e..fe6f3cbd721d 100644
--- a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py
+++ b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py
@@ -30,23 +30,8 @@ def current_sigma(self):
     def num_timesteps(self):
         return getattr(self, "_num_timesteps", None)
 
-    def _prepare_sea_cache_config(self, config=None):
-        return Cosmos3OmniPipeline._prepare_sea_cache_config(self, config)
-
-    def enable_sea_cache(self, config=None):
-        return Cosmos3OmniPipeline.enable_sea_cache(self, config)
-
-    def disable_sea_cache(self):
-        return Cosmos3OmniPipeline.disable_sea_cache(self)
-
-    def _maybe_enable_sea_cache(self):
-        return Cosmos3OmniPipeline._maybe_enable_sea_cache(self)
-
     def __call__(self, *args, **kwargs):
-        self._maybe_enable_sea_cache()
         transformer = getattr(self, "transformer", None)
-        if hasattr(transformer, "_reset_stateful_cache"):
-            transformer._reset_stateful_cache()
         try:
             return super().__call__(*args, **kwargs)
         finally:
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
index 22f8d34f8e55..48d754b4889b 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
@@ -27,7 +27,6 @@
 from transformers import AutoTokenizer, BatchEncoding
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
-from ...hooks import SeaCacheConfig
 from ...models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer
 from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan
 from ...models.transformers.transformer_cosmos3 import (
@@ -473,56 +472,6 @@ def __init__(
         self.inverse_image_resolution_template = "This image is not of {height}x{width} resolution."
         self.inverse_video_resolution_template = "This video is not of {height}x{width} resolution."
 
-        # Recommended quality-control negative prompts are documented in the Cosmos3 docs
-        # page (text2video / image2video). When the caller passes None we fall back to "".
-        self._is_sea_cache_enabled = True
-
-    def _prepare_sea_cache_config(self, config: SeaCacheConfig | None = None) -> SeaCacheConfig:
-        config = SeaCacheConfig() if config is None else config
-        config.current_step_callback = lambda: self.current_step_index
-        config.current_sigma_callback = lambda: self.current_sigma
-        config.num_inference_steps_callback = lambda: self.num_timesteps
-        return config
-
-    def enable_sea_cache(self, config: SeaCacheConfig | None = None) -> None:
-        """Enable SeaCache for subsequent pipeline calls."""
-        transformer = getattr(self, "transformer", None)
-        if transformer is None:
-            raise ValueError("SeaCache requires a loaded transformer.")
-        current_config = getattr(transformer, "_cache_config", None)
-        if current_config is not None:
-            if isinstance(current_config, SeaCacheConfig):
-                if config is None or config is current_config:
-                    self._prepare_sea_cache_config(current_config)
-                    self._is_sea_cache_enabled = True
-                    return
-                transformer.disable_cache()
-            else:
-                raise ValueError(
-                    f"Caching is already enabled with {type(current_config).__name__}. Disable it before enabling SeaCache."
-                )
-
-        transformer.enable_cache(self._prepare_sea_cache_config(config))
-        self._is_sea_cache_enabled = True
-
-    def disable_sea_cache(self) -> None:
-        """Disable the default SeaCache optimization for subsequent pipeline calls."""
-        self._is_sea_cache_enabled = False
-        transformer = getattr(self, "transformer", None)
-        if isinstance(getattr(transformer, "_cache_config", None), SeaCacheConfig):
-            transformer.disable_cache()
-
-    def _maybe_enable_sea_cache(self) -> None:
-        transformer = getattr(self, "transformer", None)
-        if transformer is None or not getattr(self, "_is_sea_cache_enabled", True):
-            return
-        if isinstance(getattr(transformer, "_cache_config", None), SeaCacheConfig):
-            self._prepare_sea_cache_config(transformer._cache_config)
-            return
-        if transformer.is_cache_enabled:
-            return
-        self.enable_sea_cache()
-
     # TODO YiYi & Daniel: fix for this use case in the base class
     def _get_execution_device(self) -> torch.device:
         from ...hooks.group_offloading import _get_group_onload_device
@@ -1517,10 +1466,6 @@ def __call__(
                 `sound` (`torch.Tensor` of shape `[C, N]`, or `None` when `enable_sound=False`). Otherwise a tuple
                 `(video, sound)` with the same fields.
         """
-        self._maybe_enable_sea_cache()
-        if hasattr(self.transformer, "_reset_stateful_cache"):
-            self.transformer._reset_stateful_cache()
-
         if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
             callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
 
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
index b4f9c78d6622..013940fd6970 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
@@ -21,8 +21,7 @@
 import torch
 from PIL import Image
 
-from diffusers import CosmosActionCondition, ModularPipeline, SeaCacheConfig, UniPCMultistepScheduler
-from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer
+from diffusers import CosmosActionCondition, ModularPipeline, UniPCMultistepScheduler
 from diffusers.modular_pipelines import (
     Cosmos3OmniBlocks,
     Cosmos3OmniModularPipeline,
@@ -42,7 +41,7 @@
 from diffusers.modular_pipelines.cosmos.encoders import Cosmos3TextEncoderStep
 from diffusers.modular_pipelines.cosmos.modular_blocks_cosmos3 import Cosmos3TransferChunkDenoiseStep
 
-from ...testing_utils import count_transformer_cache_reuse, torch_device
+from ...testing_utils import torch_device
 from ..testing_utils import (
     BaseModularPipelineTesterConfig,
     ModularLoadingTesterMixin,
@@ -194,35 +193,6 @@ def record_context(name):
             torch.testing.assert_close(sigma, pipe.scheduler.sigmas[expected_step])
         assert pipe.current_step_index is None
         assert pipe.current_sigma is None
-        assert isinstance(pipe.transformer._cache_config, SeaCacheConfig)
-
-    def test_default_sea_cache_can_be_disabled_and_reenabled(self):
-        pipe = self.get_pipeline()
-        pipe._maybe_enable_sea_cache()
-
-        assert isinstance(pipe.transformer._cache_config, SeaCacheConfig)
-        pipe.disable_sea_cache()
-        assert not pipe.transformer.is_cache_enabled
-
-        pipe._maybe_enable_sea_cache()
-        assert not pipe.transformer.is_cache_enabled
-
-        pipe.enable_sea_cache()
-        assert isinstance(pipe.transformer._cache_config, SeaCacheConfig)
-
-    def test_default_sea_cache_reuses_transformer_execution(self):
-        pipe = self.get_pipeline().to(torch_device)
-        scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, use_flow_sigmas=True)
-        pipe.update_components(scheduler=scheduler, use_native_flow_schedule=True)
-        inputs = self.get_dummy_inputs()
-        inputs["num_inference_steps"] = 4
-
-        with count_transformer_cache_reuse(pipe.transformer, Cosmos3VLTextMoTDecoderLayer) as counts:
-            pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6))
-            pipe(**inputs, output=self.output_name)
-
-        assert counts["cached_steps"] > 0
-        assert counts["full_steps"] > 0
 
     def _get_sampling_state_block_pipe(self, block):
         pipe = block.init_pipeline(self.pretrained_model_name_or_path)
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
index 705ca74fd4d3..354b9b08288f 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
@@ -17,13 +17,12 @@
 import torch
 from PIL import Image
 
-from diffusers import ModularPipeline, SeaCacheConfig
-from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer
+from diffusers import ModularPipeline
 from diffusers.modular_pipelines import Cosmos3DistilledBlocks, Cosmos3DistilledModularPipeline
 from diffusers.modular_pipelines.cosmos.before_denoise import Cosmos3VisionPrepareLatentsStep
 from diffusers.modular_pipelines.cosmos.encoders import Cosmos3DistilledTextEncoderStep
 
-from ...testing_utils import count_transformer_cache_reuse, torch_device
+from ...testing_utils import torch_device
 from ..testing_utils import (
     BaseModularPipelineTesterConfig,
     ModularLoadingTesterMixin,
@@ -184,16 +183,6 @@ def test_distilled_scheduler_fp32_state(self, use_fp32_sampling_state, expected_
 
         assert latents.dtype == expected_dtype
 
-    def test_default_sea_cache_reuses_transformer_execution(self):
-        pipe = self.get_pipeline().to(torch_device)
-
-        with count_transformer_cache_reuse(pipe.transformer, Cosmos3VLTextMoTDecoderLayer) as counts:
-            pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6))
-            pipe(**self.get_dummy_inputs(), output=self.output_name)
-
-        assert counts["cached_steps"] > 0
-        assert counts["full_steps"] > 0
-
     def test_vae_encoder_rejects_image_and_video_together(self):
         vae_encoder = Cosmos3DistilledBlocks().sub_blocks["vae_encoder"]
         vae_pipe = vae_encoder.init_pipeline(self.pretrained_model_name_or_path)
diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py
index 8c0b1e304625..fdcb8cc53931 100644
--- a/tests/pipelines/cosmos/test_cosmos3.py
+++ b/tests/pipelines/cosmos/test_cosmos3.py
@@ -25,13 +25,11 @@
     AutoencoderKLWan,
     Cosmos3OmniPipeline,
     Cosmos3OmniTransformer,
-    SeaCacheConfig,
     UniPCMultistepScheduler,
 )
-from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer
 from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import _preprocess_conditioning_image
 
-from ...testing_utils import count_transformer_cache_reuse, enable_full_determinism, torch_device
+from ...testing_utils import enable_full_determinism, torch_device
 from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
 from ..test_pipelines_common import PipelineTesterMixin
 
@@ -125,48 +123,6 @@ def test_inference(self):
         video = pipeline(**self.get_dummy_inputs(torch_device)).video
 
         self.assertEqual(video.shape, (1, 16, 16, 3))
-        self.assertIsInstance(pipeline.transformer._cache_config, SeaCacheConfig)
-
-    def test_default_sea_cache_can_be_disabled_and_reenabled(self):
-        components = self.get_dummy_components()
-        pipeline = self.pipeline_class(**components)
-        pipeline._maybe_enable_sea_cache()
-
-        self.assertIsInstance(pipeline.transformer._cache_config, SeaCacheConfig)
-        pipeline.disable_sea_cache()
-        self.assertFalse(pipeline.transformer.is_cache_enabled)
-
-        pipeline._maybe_enable_sea_cache()
-        self.assertFalse(pipeline.transformer.is_cache_enabled)
-
-        pipeline.enable_sea_cache()
-        self.assertIsInstance(pipeline.transformer._cache_config, SeaCacheConfig)
-
-        replacement_pipeline = self.pipeline_class(**components)
-        replacement_pipeline._current_step_index = 3
-        replacement_pipeline._maybe_enable_sea_cache()
-        self.assertEqual(replacement_pipeline.transformer._cache_config.current_step_callback(), 3)
-
-        replacement_config = SeaCacheConfig(threshold=0.1)
-        replacement_pipeline.enable_sea_cache(replacement_config)
-        self.assertIs(replacement_pipeline.transformer._cache_config, replacement_config)
-        self.assertEqual(replacement_pipeline.transformer._cache_config.threshold, 0.1)
-        self.assertEqual(replacement_pipeline.transformer._cache_config.current_step_callback(), 3)
-
-    def test_default_sea_cache_reuses_transformer_execution(self):
-        pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
-        pipeline.set_progress_bar_config(disable=None)
-        pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True)
-        pipeline.register_to_config(use_native_flow_schedule=True)
-        inputs = self.get_dummy_inputs(torch_device)
-        inputs.update(num_inference_steps=4, output_type="latent", enable_safety_check=False)
-
-        with count_transformer_cache_reuse(pipeline.transformer, Cosmos3VLTextMoTDecoderLayer) as counts:
-            pipeline.enable_sea_cache(SeaCacheConfig(threshold=1e6))
-            pipeline(**inputs)
-
-        self.assertGreater(counts["cached_steps"], 0)
-        self.assertGreater(counts["full_steps"], 0)
 
     def test_fp32_sampling_state_keeps_transformer_inputs_in_model_dtype(self):
         pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
diff --git a/tests/testing_utils.py b/tests/testing_utils.py
index 0c9766518f9f..c35f975285c4 100644
--- a/tests/testing_utils.py
+++ b/tests/testing_utils.py
@@ -198,38 +198,6 @@ def numpy_cosine_similarity_distance(a, b):
     return distance
 
 
-@contextmanager
-def count_transformer_cache_reuse(transformer, block_cls):
-    """Count real block executions and total transformer calls while running under a cache.
-
-    ``block_cls.forward`` is patched to count real executions (skipped blocks reuse cached outputs and never run),
-    and a forward pre-hook counts every transformer call. On exit the yielded dict is populated with ``full_steps``
-    (transformer calls that recomputed the blocks) and ``cached_steps`` (calls served from the cache).
-
-    The patch must be active before the cache is enabled, so enable the cache inside the ``with`` block.
-    """
-    counts = {"block_calls": 0, "transformer_calls": 0, "full_steps": 0, "cached_steps": 0}
-    original_forward = block_cls.forward
-
-    def counting_forward(self, *args, **kwargs):
-        counts["block_calls"] += 1
-        return original_forward(self, *args, **kwargs)
-
-    def counting_pre_hook(module, args):
-        counts["transformer_calls"] += 1
-
-    block_cls.forward = counting_forward
-    handle = transformer.register_forward_pre_hook(counting_pre_hook)
-    try:
-        yield counts
-    finally:
-        handle.remove()
-        block_cls.forward = original_forward
-        num_blocks = sum(1 for module in transformer.modules() if isinstance(module, block_cls))
-        counts["full_steps"] = counts["block_calls"] // num_blocks if num_blocks else 0
-        counts["cached_steps"] = counts["transformer_calls"] - counts["full_steps"]
-
-
 def check_if_dicts_are_equal(dict1, dict2):
     dict1, dict2 = dict1.copy(), dict2.copy()
 

From 0d3c8d6259d57a6fda7e4c99521758407689c439 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Wed, 2 Sep 2026 08:00:36 -0700
Subject: [PATCH 12/14] Align SeaCache and tests with model-level cache
 patterns

---
 .../modular_pipelines/cosmos/denoise.py       |  10 +-
 .../pipelines/cosmos/pipeline_cosmos3_omni.py |  10 +-
 tests/hooks/test_sea_cache.py                 | 428 ------------------
 tests/models/testing_utils/__init__.py        |   4 +
 tests/models/testing_utils/cache.py           | 144 ++++++
 .../test_models_transformer_cosmos3.py        |   5 +
 6 files changed, 157 insertions(+), 444 deletions(-)
 delete mode 100644 tests/hooks/test_sea_cache.py

diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py
index 5b06b4f910a2..8cf4b377198f 100644
--- a/src/diffusers/modular_pipelines/cosmos/denoise.py
+++ b/src/diffusers/modular_pipelines/cosmos/denoise.py
@@ -1,5 +1,4 @@
 import inspect
-from contextlib import nullcontext
 
 import torch
 
@@ -15,11 +14,6 @@
 from .modular_pipeline import Cosmos3OmniModularPipeline
 
 
-def _cache_context(transformer: torch.nn.Module, name: str):
-    cache_context = getattr(transformer, "cache_context", None)
-    return cache_context(name) if callable(cache_context) else nullcontext()
-
-
 class Cosmos3VisionLoopPrepareStep(ModularPipelineBlocks):
     model_name = "cosmos3-omni"
 
@@ -227,7 +221,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
             transformer_kwargs = {
                 name: value for name, value in transformer_kwargs.items() if name in transformer_args
             }
-            with _cache_context(components.transformer, pass_name):
+            with components.transformer.cache_context(pass_name):
                 preds_vision, preds_sound, preds_action = components.transformer(
                     **transformer_kwargs, return_dict=False
                 )
@@ -754,7 +748,7 @@ def intermediate_outputs(self) -> list[OutputParam]:
 
     @staticmethod
     def _forward(components, static, vision_tokens, vision_timesteps, context_name):
-        with _cache_context(components.transformer, context_name):
+        with components.transformer.cache_context(context_name):
             preds_vision, _, _ = components.transformer(
                 input_ids=static["input_ids"],
                 text_indexes=static["text_indexes"],
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
index 48d754b4889b..b6982789a834 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
@@ -16,7 +16,6 @@
 import json
 import math
 from collections.abc import Iterable
-from contextlib import nullcontext
 from dataclasses import dataclass
 from typing import Any, Callable, Literal
 
@@ -42,11 +41,6 @@
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
-def _cache_context(transformer: torch.nn.Module, name: str):
-    cache_context = getattr(transformer, "cache_context", None)
-    return cache_context(name) if callable(cache_context) else nullcontext()
-
-
 if is_cosmos_guardrail_available():
     from cosmos_guardrail import CosmosSafetyChecker
 else:
@@ -1738,7 +1732,7 @@ def __call__(
                 )
 
                 # --- Conditional pass ---
-                with _cache_context(self.transformer, "cond"):
+                with self.transformer.cache_context("cond"):
                     preds_vision, preds_sound, preds_action = self.transformer(
                         input_ids=cond_packed_static["input_ids"],
                         text_indexes=cond_packed_static["text_indexes"],
@@ -1779,7 +1773,7 @@ def __call__(
                 # --- Unconditional pass (Skip if not using CFG) ---
                 uncond_v_vision = uncond_v_sound = uncond_v_action = None
                 if self.do_classifier_free_guidance:
-                    with _cache_context(self.transformer, "uncond"):
+                    with self.transformer.cache_context("uncond"):
                         preds_vision, preds_sound, preds_action = self.transformer(
                             input_ids=uncond_packed_static["input_ids"],
                             text_indexes=uncond_packed_static["text_indexes"],
diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py
deleted file mode 100644
index 5a31edcd564f..000000000000
--- a/tests/hooks/test_sea_cache.py
+++ /dev/null
@@ -1,428 +0,0 @@
-# Copyright 2026 HuggingFace Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import pytest
-import torch
-
-from diffusers import SeaCacheConfig
-from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry
-from diffusers.hooks.hooks import HookRegistry, ModelHook
-from diffusers.hooks.sea_cache import (
-    _SEA_CACHE_BLOCK_HOOK,
-    _SEA_CACHE_LEADER_BLOCK_HOOK,
-    _SEA_CACHE_ROOT_HOOK,
-    _apply_sea_filter,
-)
-from diffusers.models.cache_utils import CacheMixin
-
-
-class CountingIdentity(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.calls = 0
-
-    def forward(self, hidden_states):
-        self.calls += 1
-        return hidden_states
-
-
-class DummySeaBlock(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.indicator_norm = CountingIdentity()
-        self.calls = 0
-
-    def forward(self, und_seq, gen_seq, rotary_emb=None):
-        self.calls += 1
-        return und_seq + 1, gen_seq * 2
-
-
-class DummySeaTransformer(torch.nn.Module, CacheMixin):
-    def __init__(self, num_layers=3):
-        super().__init__()
-        self.layers = torch.nn.ModuleList([DummySeaBlock() for _ in range(num_layers)])
-
-    def forward(self, hidden_states):
-        shape = hidden_states.shape
-        gen_seq = hidden_states.reshape(-1, shape[-1])
-        und_seq = torch.zeros(1, shape[-1], device=hidden_states.device, dtype=hidden_states.dtype)
-        for layer in self.layers:
-            und_seq, gen_seq = layer(und_seq, gen_seq)
-        return und_seq, gen_seq.reshape(shape)
-
-
-@pytest.fixture(autouse=True)
-def register_dummy_sea_block():
-    TransformerBlockRegistry.register(
-        DummySeaBlock,
-        TransformerBlockMetadata(
-            return_hidden_states_index=1,
-            return_encoder_hidden_states_index=0,
-            hidden_states_argument_name="gen_seq",
-            encoder_hidden_states_argument_name="und_seq",
-            hidden_states_norm_module_name="indicator_norm",
-        ),
-    )
-
-
-def _metadata_callback(module, args, kwargs):
-    hidden_states = kwargs.get("hidden_states", args[0] if args else None)
-    if hidden_states is None:
-        return None
-    temporal, height, width = hidden_states.shape[:3]
-    indexes = torch.arange(temporal * height * width, device=hidden_states.device)
-    return [(indexes, (temporal, height, width))]
-
-
-def _make_config(runtime, **kwargs):
-    config_kwargs = {
-        "threshold": 100.0,
-        "retention_steps": 1,
-        "cache_end_steps": 0,
-        "indicator_source": "first_block",
-        "current_step_callback": lambda: runtime["step"],
-        "current_sigma_callback": lambda: runtime["sigma"],
-        "num_inference_steps_callback": lambda: runtime["num_steps"],
-        "metadata_callback": _metadata_callback,
-    }
-    config_kwargs.update(kwargs)
-    return SeaCacheConfig(**config_kwargs)
-
-
-@torch.no_grad()
-def test_sea_cache_uses_independent_gates_and_histories_per_context():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 4}
-    model = DummySeaTransformer()
-    model.enable_cache(_make_config(runtime, threshold=0.5))
-
-    first_input = torch.ones(2, 2, 2, 1)
-    with model.cache_context("cond"):
-        _, first_output = model(first_input)
-    torch.testing.assert_close(first_output, first_input * 8)
-    assert [block.calls for block in model.layers] == [1, 1, 1]
-
-    with model.cache_context("uncond"):
-        _, first_uncond_output = model(first_input)
-    torch.testing.assert_close(first_uncond_output, first_input * 8)
-    assert [block.calls for block in model.layers] == [2, 2, 2]
-
-    runtime.update(step=1)
-    with model.cache_context("cond"):
-        _, cached_output = model(first_input)
-    torch.testing.assert_close(cached_output, first_input * 8)
-    assert [block.calls for block in model.layers] == [2, 2, 2]
-
-    changed_uncond_input = first_input * 100
-    with model.cache_context("uncond"):
-        _, uncond_output = model(changed_uncond_input)
-    torch.testing.assert_close(uncond_output, changed_uncond_input * 8)
-    assert [block.calls for block in model.layers] == [3, 3, 3]
-    assert model.layers[0].indicator_norm.calls == 4
-
-    model._reset_stateful_cache()
-
-    model.disable_cache()
-    assert not model.is_cache_enabled
-
-
-@torch.no_grad()
-def test_sea_cache_max_consecutive_cached_forces_full_per_context():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 8}
-    model = DummySeaTransformer()
-    model.enable_cache(
-        _make_config(
-            runtime,
-            threshold=100.0,
-            retention_steps=0,
-            cache_end_steps=0,
-            max_consecutive_cached=2,
-        )
-    )
-
-    for step in range(runtime["num_steps"]):
-        runtime.update(step=step, sigma=0.9 - step * 0.1)
-        with model.cache_context("cond"):
-            model(torch.ones(1, 1, 1, 1))
-
-    # With max_consecutive_cached=2, every third step forces a full execution: 3 full, 5 cached.
-    assert [block.calls for block in model.layers] == [3, 3, 3]
-
-
-@torch.no_grad()
-def test_sea_cache_raw_vision_indicator_includes_conditioning_frames():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
-    model = DummySeaTransformer()
-
-    def raw_vision(module, args, kwargs):
-        hidden_states = kwargs.get("hidden_states", args[0] if args else None)
-        latent = hidden_states.permute(3, 0, 1, 2)
-        return [latent]
-
-    model.enable_cache(
-        _make_config(
-            runtime,
-            threshold=1e-6,
-            indicator_source="raw_vision_latents",
-            raw_vision_callback=raw_vision,
-        )
-    )
-
-    first_input = torch.tensor([1.0, 2.0]).reshape(2, 1, 1, 1)
-    with model.cache_context("cond"):
-        model(first_input)
-
-    # The noisy frame stays fixed, but changing the conditioning frame changes the complete raw-latent indicator.
-    runtime.update(step=1, sigma=0.9)
-    second_input = torch.tensor([100.0, 2.0]).reshape(2, 1, 1, 1)
-    with model.cache_context("cond"):
-        model(second_input)
-
-    assert [block.calls for block in model.layers] == [2, 2, 2]
-    assert model.layers[0].indicator_norm.calls == 0
-
-
-@torch.no_grad()
-def test_sea_cache_residual_order_one_uses_actual_full_step_history():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 4}
-    config = _make_config(runtime, residual_order=1, threshold=0.0)
-    model = DummySeaTransformer()
-    model.enable_cache(config)
-
-    with model.cache_context("cond"):
-        model(torch.ones(1, 1, 1, 1))
-
-    runtime.update(step=1, sigma=0.7)
-    with model.cache_context("cond"):
-        model(torch.full((1, 1, 1, 1), 2.0))
-
-    config.threshold = 100.0
-    runtime.update(step=2, sigma=0.5)
-    with model.cache_context("cond"):
-        _, output = model(torch.full((1, 1, 1, 1), 3.0))
-
-    # Full residuals are 7 and 14 at steps 0 and 1, so linear extrapolation predicts 21 at step 2.
-    torch.testing.assert_close(output, torch.full_like(output, 24.0))
-    assert [block.calls for block in model.layers] == [2, 2, 2]
-
-
-@torch.no_grad()
-def test_cache_context_registry_is_refreshed_when_cache_is_enabled_after_an_uncached_call():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
-    model = DummySeaTransformer()
-
-    # Pipeline cache contexts may be entered before a cache is enabled (for example, during the baseline).
-    with model.cache_context("baseline"):
-        model(torch.ones(1, 1, 1, 1))
-
-    model.enable_cache(_make_config(runtime))
-    with model.cache_context("cond"):
-        model(torch.ones(1, 1, 1, 1))
-
-    assert [block.calls for block in model.layers] == [2, 2, 2]
-
-
-@torch.no_grad()
-def test_sea_cache_fails_open_without_vision_metadata():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
-    config = _make_config(runtime)
-    config.metadata_callback = lambda module, args, kwargs: None
-    model = DummySeaTransformer()
-    model.enable_cache(config)
-
-    with model.cache_context("cond"):
-        model(torch.ones(1, 1, 1, 1))
-    runtime.update(step=1, sigma=0.5)
-    with model.cache_context("cond"):
-        model(torch.ones(1, 1, 1, 1))
-
-    assert [block.calls for block in model.layers] == [2, 2, 2]
-
-
-@torch.no_grad()
-def test_sea_cache_non_adjacent_steps_start_a_new_residual_trajectory():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
-    config = _make_config(runtime, threshold=0.0, retention_steps=0, residual_order=1)
-    model = DummySeaTransformer()
-    model.enable_cache(config)
-
-    with model.cache_context("cond"):
-        model(torch.ones(1, 1, 1, 1))
-    runtime.update(step=1, sigma=0.5)
-    with model.cache_context("cond"):
-        model(torch.full((1, 1, 1, 1), 100.0))
-
-    config.threshold = 100.0
-    runtime.update(step=0, sigma=0.9)
-    with model.cache_context("cond"):
-        model(torch.full((1, 1, 1, 1), 2.0))
-    runtime.update(step=1, sigma=0.5)
-    with model.cache_context("cond"):
-        _, output = model(torch.full((1, 1, 1, 1), 2.0))
-
-    torch.testing.assert_close(output, torch.full_like(output, 16.0))
-    assert [block.calls for block in model.layers] == [3, 3, 3]
-
-
-@torch.no_grad()
-def test_sea_cache_fails_open_for_shape_changes():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
-    model = DummySeaTransformer()
-    model.enable_cache(_make_config(runtime, residual_order=1))
-
-    with model.cache_context("cond"):
-        model(torch.ones(1, 1, 1, 1))
-
-    runtime.update(step=1, sigma=0.5)
-    with model.cache_context("cond"):
-        model(torch.ones(1, 2, 1, 1))
-
-    assert [block.calls for block in model.layers] == [2, 2, 2]
-
-
-def test_sea_cache_is_inference_only_and_fails_open_with_autograd():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
-    model = DummySeaTransformer()
-    model.enable_cache(_make_config(runtime))
-
-    with torch.enable_grad(), model.cache_context("cond"):
-        model(torch.ones(1, 1, 1, 1, requires_grad=True))
-    runtime.update(step=1, sigma=0.5)
-    with torch.enable_grad(), model.cache_context("cond"):
-        model(torch.ones(1, 1, 1, 1, requires_grad=True))
-
-    assert [block.calls for block in model.layers] == [2, 2, 2]
-
-
-@torch.no_grad()
-def test_sea_filter_density_normalizes_gain_to_unit_mean():
-    impulse = torch.zeros(2, 3, 4, 1)
-    impulse[0, 0, 0, 0] = 1
-
-    filtered = _apply_sea_filter(impulse, sigma=0.5, power_exp=3.0)
-    recovered_gain = torch.fft.fftn(filtered.float(), dim=(0, 1, 2))
-
-    torch.testing.assert_close(
-        recovered_gain.real.mean(),
-        torch.tensor(1.0),
-        atol=1e-5,
-        rtol=1e-5,
-    )
-    torch.testing.assert_close(
-        recovered_gain.imag,
-        torch.zeros_like(recovered_gain.imag),
-        atol=1e-5,
-        rtol=0,
-    )
-
-
-@pytest.mark.parametrize("sigma", [0.0, 1.0])
-def test_sea_filter_is_finite_at_scheduler_endpoints(sigma):
-    filtered = _apply_sea_filter(torch.randn(2, 2, 2, 4), sigma=sigma, power_exp=2.0)
-
-    assert torch.isfinite(filtered).all()
-    assert filtered.abs().sum() > 0
-
-
-@torch.no_grad()
-def test_sea_cache_single_block_supports_full_and_cached_execution_then_disables_cleanly():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
-    model = DummySeaTransformer(num_layers=1)
-    model.enable_cache(_make_config(runtime))
-
-    with model.cache_context("cond"):
-        _, first_output = model(torch.ones(1, 1, 1, 1))
-    torch.testing.assert_close(first_output, torch.full_like(first_output, 2.0))
-
-    runtime.update(step=1, sigma=0.5)
-    with model.cache_context("cond"):
-        _, cached_output = model(torch.full((1, 1, 1, 1), 2.0))
-    torch.testing.assert_close(cached_output, torch.full_like(cached_output, 3.0))
-    assert model.layers[0].calls == 1
-
-    model.disable_cache()
-
-    _, output = model(torch.ones(1, 1, 1, 1))
-
-    torch.testing.assert_close(output, torch.full_like(output, 2.0))
-    assert model.layers[0].calls == 2
-    assert model.layers[0]._diffusers_hook.hooks == {}
-
-
-@torch.no_grad()
-def test_sea_cache_fails_open_for_parameter_sharded_blocks():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
-    model = DummySeaTransformer()
-    model.layers[0]._get_fsdp_state = lambda: object()
-    model.enable_cache(_make_config(runtime))
-
-    with model.cache_context("cond"):
-        model(torch.ones(1, 1, 1, 1))
-    runtime.update(step=1, sigma=0.5)
-    with model.cache_context("cond"):
-        model(torch.ones(1, 1, 1, 1))
-
-    assert [block.calls for block in model.layers] == [2, 2, 2]
-
-
-def test_sea_cache_failed_enable_rolls_back_only_new_hooks():
-    runtime = {"step": 0, "sigma": 0.9, "num_steps": 2}
-    model = DummySeaTransformer()
-    existing_hook = ModelHook()
-    middle_registry = HookRegistry.check_if_exists_or_initialize(model.layers[1])
-    middle_registry.register_hook(existing_hook, _SEA_CACHE_BLOCK_HOOK)
-
-    with pytest.raises(ValueError, match="already exists"):
-        model.enable_cache(_make_config(runtime))
-
-    assert not model.is_cache_enabled
-    assert model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) is None
-    assert model.layers[0]._diffusers_hook.get_hook(_SEA_CACHE_LEADER_BLOCK_HOOK) is None
-    assert middle_registry.get_hook(_SEA_CACHE_BLOCK_HOOK) is existing_hook
-    assert not hasattr(model.layers[2], "_diffusers_hook")
-
-
-def test_sea_cache_config_defaults():
-    config = SeaCacheConfig()
-
-    assert config.threshold == 0.25
-    assert config.residual_order == 1
-    assert config.max_consecutive_cached == 2
-    assert config.indicator_source == "raw_vision_latents"
-
-
-@pytest.mark.parametrize(
-    ("kwargs", "message"),
-    [
-        ({"threshold": -1.0}, "threshold"),
-        ({"residual_order": 2}, "residual_order"),
-        ({"retention_steps": -1}, "retention_steps"),
-        ({"cache_end_steps": -1}, "cache_end_steps"),
-        ({"max_consecutive_cached": -1}, "max_consecutive_cached"),
-        ({"max_consecutive_cached": 1.5}, "max_consecutive_cached"),
-        ({"max_consecutive_cached": True}, "max_consecutive_cached"),
-        ({"power_exp": 0.0}, "power_exp"),
-        ({"indicator_source": "raw"}, "indicator_source"),
-        ({"threshold": float("nan")}, "threshold"),
-        ({"power_exp": float("inf")}, "power_exp"),
-    ],
-)
-def test_sea_cache_config_validation(kwargs, message):
-    with pytest.raises(ValueError, match=message):
-        SeaCacheConfig(**kwargs)
-
-
-@pytest.mark.parametrize("callback_name", ["metadata_callback", "raw_vision_callback"])
-def test_sea_cache_callback_validation(callback_name):
-    with pytest.raises(TypeError, match=callback_name):
-        SeaCacheConfig(**{callback_name: 1})
diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py
index 760a3fac04e0..2d7d5ae23257 100644
--- a/tests/models/testing_utils/__init__.py
+++ b/tests/models/testing_utils/__init__.py
@@ -9,6 +9,8 @@
     MagCacheTesterMixin,
     PyramidAttentionBroadcastConfigMixin,
     PyramidAttentionBroadcastTesterMixin,
+    SeaCacheConfigMixin,
+    SeaCacheTesterMixin,
     TaylorSeerCacheConfigMixin,
     TaylorSeerCacheTesterMixin,
 )
@@ -90,6 +92,8 @@
     "NunchakuLiteTesterMixin",
     "PyramidAttentionBroadcastConfigMixin",
     "PyramidAttentionBroadcastTesterMixin",
+    "SeaCacheConfigMixin",
+    "SeaCacheTesterMixin",
     "TaylorSeerCacheConfigMixin",
     "TaylorSeerCacheTesterMixin",
     "QuantizationCompileTesterMixin",
diff --git a/tests/models/testing_utils/cache.py b/tests/models/testing_utils/cache.py
index 8357d34f3077..c8d757140348 100644
--- a/tests/models/testing_utils/cache.py
+++ b/tests/models/testing_utils/cache.py
@@ -23,12 +23,19 @@
     FirstBlockCacheConfig,
     MagCacheConfig,
     PyramidAttentionBroadcastConfig,
+    SeaCacheConfig,
     TaylorSeerCacheConfig,
 )
 from diffusers.hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK
 from diffusers.hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK
 from diffusers.hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK
 from diffusers.hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK
+from diffusers.hooks.sea_cache import (
+    _SEA_CACHE_BLOCK_HOOK,
+    _SEA_CACHE_LEADER_BLOCK_HOOK,
+    _SEA_CACHE_POST_NORM_HOOK,
+    _SEA_CACHE_ROOT_HOOK,
+)
 from diffusers.hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK
 from diffusers.models.cache_utils import CacheMixin
 
@@ -433,6 +440,143 @@ def test_fbc_reset_stateful_cache(self):
         self._test_reset_stateful_cache()
 
 
+@is_cache
+class SeaCacheConfigMixin:
+    """
+    Base mixin providing SeaCache config.
+
+    Expected class attributes:
+        - model_class: The model class to test (must use CacheMixin)
+    """
+
+    SEA_CACHE_CONFIG = {
+        "threshold": 100.0,
+        "retention_steps": 0,
+        "cache_end_steps": 0,
+    }
+
+    def _get_cache_config(self):
+        runtime = {"step": 0, "sigma": 0.9, "num_steps": 3}
+        self._sea_cache_runtime = runtime
+        return SeaCacheConfig(
+            **self.SEA_CACHE_CONFIG,
+            current_step_callback=lambda: runtime["step"],
+            current_sigma_callback=lambda: runtime["sigma"],
+            num_inference_steps_callback=lambda: runtime["num_steps"],
+        )
+
+    def _get_hook_names(self):
+        return [
+            _SEA_CACHE_ROOT_HOOK,
+            _SEA_CACHE_LEADER_BLOCK_HOOK,
+            _SEA_CACHE_BLOCK_HOOK,
+            _SEA_CACHE_POST_NORM_HOOK,
+        ]
+
+
+@is_cache
+class SeaCacheTesterMixin(SeaCacheConfigMixin, CacheTesterMixin):
+    """
+    Mixin class for testing SeaCache on models.
+
+    Expected methods to be implemented by subclasses:
+        - get_init_dict(): Returns dict of arguments to initialize the model
+        - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass
+
+    Pytest mark: cache
+        Use `pytest -m "not cache"` to skip these tests
+    """
+
+    @staticmethod
+    def _unwrap_cache_output(output):
+        while isinstance(output, (list, tuple)):
+            output = output[0]
+        return output
+
+    def _get_modified_cache_inputs(self):
+        inputs = self.get_dummy_inputs()
+        value = inputs[self.cache_input_key]
+        if isinstance(value, torch.Tensor):
+            inputs[self.cache_input_key] = value + 0.1
+        else:
+            inputs[self.cache_input_key] = [tensor + 0.1 for tensor in value]
+        return inputs
+
+    @torch.no_grad()
+    def _test_cache_inference(self):
+        model = self.model_class(**self.get_init_dict()).to(torch_device).eval()
+        model.enable_cache(self._get_cache_config())
+
+        with model.cache_context("sea_cache_test"):
+            model(**self.get_dummy_inputs(), return_dict=False)
+
+        self._sea_cache_runtime.update(step=1, sigma=0.6)
+        modified_inputs = self._get_modified_cache_inputs()
+        with model.cache_context("sea_cache_test"):
+            output_with_cache = self._unwrap_cache_output(model(**modified_inputs, return_dict=False))
+
+        assert output_with_cache is not None
+        assert not torch.isnan(output_with_cache).any()
+
+        model.disable_cache()
+        output_without_cache = self._unwrap_cache_output(model(**modified_inputs, return_dict=False))
+        assert not torch.allclose(output_without_cache, output_with_cache, atol=1e-5)
+
+    @torch.no_grad()
+    def _test_cache_context_manager(self, atol=1e-5, rtol=0):
+        model = self.model_class(**self.get_init_dict()).to(torch_device).eval()
+        model.enable_cache(self._get_cache_config())
+        inputs = self.get_dummy_inputs()
+
+        with model.cache_context("context_1"):
+            output_ctx1 = self._unwrap_cache_output(model(**inputs, return_dict=False))
+        with model.cache_context("context_2"):
+            output_ctx2 = self._unwrap_cache_output(model(**inputs, return_dict=False))
+
+        assert_tensors_close(
+            output_ctx1,
+            output_ctx2,
+            atol=atol,
+            rtol=rtol,
+            msg="First pass in different cache contexts should produce the same output.",
+        )
+        model.disable_cache()
+
+    @torch.no_grad()
+    def _test_reset_stateful_cache(self):
+        model = self.model_class(**self.get_init_dict()).to(torch_device).eval()
+        model.enable_cache(self._get_cache_config())
+
+        with model.cache_context("sea_cache_test"):
+            model(**self.get_dummy_inputs(), return_dict=False)
+        model._reset_stateful_cache()
+        model.disable_cache()
+
+    @require_cache_mixin
+    def test_sea_cache_enable_disable_state(self):
+        self._test_cache_enable_disable_state()
+
+    @require_cache_mixin
+    def test_sea_cache_double_enable_raises_error(self):
+        self._test_cache_double_enable_raises_error()
+
+    @require_cache_mixin
+    def test_sea_cache_hooks_registered(self):
+        self._test_cache_hooks_registered()
+
+    @require_cache_mixin
+    def test_sea_cache_inference(self):
+        self._test_cache_inference()
+
+    @require_cache_mixin
+    def test_sea_cache_context_manager(self):
+        self._test_cache_context_manager()
+
+    @require_cache_mixin
+    def test_sea_cache_reset_stateful_cache(self):
+        self._test_reset_stateful_cache()
+
+
 @is_cache
 class FasterCacheConfigMixin:
     """
diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py
index c03d39f2b4c1..8a078b64ba76 100644
--- a/tests/models/transformers/test_models_transformer_cosmos3.py
+++ b/tests/models/transformers/test_models_transformer_cosmos3.py
@@ -34,6 +34,7 @@
     BaseModelTesterConfig,
     MemoryTesterMixin,
     ModelTesterMixin,
+    SeaCacheTesterMixin,
     TorchCompileTesterMixin,
     TrainingTesterMixin,
 )
@@ -375,6 +376,10 @@ def test_cosmos3_nemotron_rms_norm_multiplies_in_float32(self):
         torch.testing.assert_close(norm(hidden_states), expected, rtol=0, atol=0)
 
 
+class TestCosmos3OmniTransformerSeaCache(Cosmos3OmniTransformerTesterConfig, SeaCacheTesterMixin):
+    cache_input_key = "vision_tokens"
+
+
 class TestCosmos3OmniTransformerMemory(Cosmos3OmniTransformerTesterConfig, MemoryTesterMixin):
     @pytest.mark.skip("The transformer returns one tensor list per generated modality.")
     def test_layerwise_casting_training(self):

From b6b564f14e6829085054433b72cee3d312e05f33 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Thu, 3 Sep 2026 04:13:02 -0700
Subject: [PATCH 13/14] Make Cosmos3 sampling state always FP32

---
 docs/source/en/api/pipelines/cosmos3.md       |  6 --
 .../cosmos/before_denoise.py                  | 86 +------------------
 .../modular_pipelines/cosmos/denoise.py       | 44 +++-------
 .../cosmos/modular_blocks_cosmos3.py          | 36 --------
 .../modular_blocks_cosmos3_distilled.py       |  8 --
 .../pipelines/cosmos/pipeline_cosmos3_omni.py | 25 ++----
 .../cosmos/test_modular_pipeline_cosmos3.py   | 46 +++-------
 ...test_modular_pipeline_cosmos3_distilled.py | 29 ++-----
 8 files changed, 42 insertions(+), 238 deletions(-)

diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md
index dcfb7167db90..5a382c4743a8 100644
--- a/docs/source/en/api/pipelines/cosmos3.md
+++ b/docs/source/en/api/pipelines/cosmos3.md
@@ -685,12 +685,6 @@ The same model-level API works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModula
 `pipe.transformer.disable_cache()` when you need every denoising step to execute the full transformer. Cache state is
 reset after each pipeline call, and conditional and unconditional guidance branches keep independent histories.
 
-## Sampling precision
-
-Cosmos 3 keeps denoising latents and classifier-free-guidance arithmetic in `torch.float32` by default while casting
-transformer inputs to the model dtype. Pass `use_fp32_sampling_state=False` to an inference call to keep the sampling
-state in the model dtype instead.
-
 ## Context parallelism
 
 For long videos or high resolutions, a single forward pass can exceed the memory and latency budget of one GPU. Cosmos 3 supports **context parallelism (CP)** to shard the sequence dimension across multiple GPUs, splitting the attention computation so each device holds only a slice of the tokens.
diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py
index decbe325782b..7e9d83fa6316 100644
--- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py
+++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py
@@ -12,15 +12,6 @@
 from .modular_pipeline import Cosmos3OmniModularPipeline
 
 
-def _resolve_sampling_dtype(
-    components: Cosmos3OmniModularPipeline, use_fp32_sampling_state: bool | None
-) -> tuple[bool, torch.dtype]:
-    if use_fp32_sampling_state is None:
-        use_fp32_sampling_state = components.config.default_use_fp32_sampling_state
-    sampling_dtype = torch.float32 if use_fp32_sampling_state else components.transformer.dtype
-    return use_fp32_sampling_state, sampling_dtype
-
-
 class Cosmos3PrepareTextSegmentsStep(ModularPipelineBlocks):
     model_name = "cosmos3-omni"
 
@@ -77,10 +68,6 @@ def description(self) -> str:
     def expected_components(self) -> list[ComponentSpec]:
         return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
 
-    @property
-    def expected_configs(self) -> list[ConfigSpec]:
-        return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)]
-
     @property
     def inputs(self) -> list[InputParam]:
         return [
@@ -111,16 +98,6 @@ def inputs(self) -> list[InputParam]:
                 description="Pre-generated noisy vision latents.",
             ),
             InputParam.template("generator"),
-            InputParam(
-                name="use_fp32_sampling_state",
-                type_hint=bool | None,
-                default=None,
-                description=(
-                    "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. "
-                    "If unset, uses the "
-                    "pipeline's `default_use_fp32_sampling_state` config."
-                ),
-            ),
         ]
 
     @property
@@ -144,20 +121,13 @@ def intermediate_outputs(self) -> list[OutputParam]:
                 type_hint=torch.Tensor,
                 description="Clean encoded vision latents used to re-anchor image conditioning each step.",
             ),
-            OutputParam(
-                "use_fp32_sampling_state",
-                type_hint=bool,
-                description="Whether vision sampling state is kept in float32.",
-            ),
         ]
 
     @torch.no_grad()
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         device = components._execution_device
-        block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype(
-            components, block_state.use_fp32_sampling_state
-        )
+        sampling_dtype = torch.float32
 
         x0_tokens_vision = block_state.x0_tokens_vision
         if x0_tokens_vision is None:
@@ -226,10 +196,6 @@ def expected_components(self) -> list[ComponentSpec]:
             ComponentSpec("scheduler", UniPCMultistepScheduler),
         ]
 
-    @property
-    def expected_configs(self) -> list[ConfigSpec]:
-        return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)]
-
     @property
     def inputs(self) -> list[InputParam]:
         return [
@@ -242,16 +208,6 @@ def inputs(self) -> list[InputParam]:
                 description="Pre-generated noisy sound latents.",
             ),
             InputParam.template("generator"),
-            InputParam(
-                name="use_fp32_sampling_state",
-                type_hint=bool | None,
-                default=None,
-                description=(
-                    "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. "
-                    "If unset, uses the "
-                    "pipeline's `default_use_fp32_sampling_state` config."
-                ),
-            ),
         ]
 
     @property
@@ -272,9 +228,7 @@ def intermediate_outputs(self) -> list[OutputParam]:
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         device = components._execution_device
-        block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype(
-            components, block_state.use_fp32_sampling_state
-        )
+        sampling_dtype = torch.float32
 
         if not components.transformer.config.sound_gen:
             raise ValueError("Sound generation requires a transformer trained with sound_gen=True.")
@@ -320,10 +274,6 @@ def expected_components(self) -> list[ComponentSpec]:
             ComponentSpec("scheduler", UniPCMultistepScheduler),
         ]
 
-    @property
-    def expected_configs(self) -> list[ConfigSpec]:
-        return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)]
-
     @property
     def inputs(self) -> list[InputParam]:
         return [
@@ -346,16 +296,6 @@ def inputs(self) -> list[InputParam]:
                 description="Pre-generated noisy action latents.",
             ),
             InputParam.template("generator"),
-            InputParam(
-                name="use_fp32_sampling_state",
-                type_hint=bool | None,
-                default=None,
-                description=(
-                    "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. "
-                    "If unset, uses the "
-                    "pipeline's `default_use_fp32_sampling_state` config."
-                ),
-            ),
         ]
 
     @property
@@ -387,9 +327,7 @@ def intermediate_outputs(self) -> list[OutputParam]:
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         device = components._execution_device
-        block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype(
-            components, block_state.use_fp32_sampling_state
-        )
+        sampling_dtype = torch.float32
         action = block_state.action
 
         if not components.transformer.config.action_gen:
@@ -1067,10 +1005,6 @@ def description(self) -> str:
     def expected_components(self) -> list[ComponentSpec]:
         return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
 
-    @property
-    def expected_configs(self) -> list[ConfigSpec]:
-        return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)]
-
     @property
     def inputs(self) -> list[InputParam]:
         return [
@@ -1087,16 +1021,6 @@ def inputs(self) -> list[InputParam]:
                 description="Number of pixel frames used to seed this chunk's target.",
             ),
             InputParam.template("generator"),
-            InputParam(
-                name="use_fp32_sampling_state",
-                type_hint=bool | None,
-                default=None,
-                description=(
-                    "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. "
-                    "If unset, uses the "
-                    "pipeline's `default_use_fp32_sampling_state` config."
-                ),
-            ),
         ]
 
     @property
@@ -1124,9 +1048,7 @@ def intermediate_outputs(self) -> list[OutputParam]:
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         device = components._execution_device
-        block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype(
-            components, block_state.use_fp32_sampling_state
-        )
+        sampling_dtype = torch.float32
         tcf = components.vae_scale_factor_temporal
 
         target_x0 = block_state.x0_tokens_vision.to(device=device, dtype=sampling_dtype)
diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py
index 8cf4b377198f..ba6e06a15244 100644
--- a/src/diffusers/modular_pipelines/cosmos/denoise.py
+++ b/src/diffusers/modular_pipelines/cosmos/denoise.py
@@ -166,12 +166,6 @@ def inputs(self) -> list[InputParam]:
                 default=6.0,
                 description="Scale for classifier-free guidance.",
             ),
-            InputParam(
-                name="use_fp32_sampling_state",
-                type_hint=bool,
-                default=False,
-                description="Whether to keep velocity and guidance arithmetic in float32.",
-            ),
         ]
 
     @property
@@ -236,12 +230,14 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
             )
 
         cond_velocity_vision, cond_velocity_sound, cond_velocity_action = velocities["cond"]
-        if block_state.use_fp32_sampling_state:
-            cond_velocity_vision = cond_velocity_vision.float()
-            cond_velocity_sound = cond_velocity_sound.float() if cond_velocity_sound is not None else None
-            cond_velocity_action = cond_velocity_action.float() if cond_velocity_action is not None else None
+        cond_velocity_vision = cond_velocity_vision.float()
+        cond_velocity_sound = cond_velocity_sound.float() if cond_velocity_sound is not None else None
+        cond_velocity_action = cond_velocity_action.float() if cond_velocity_action is not None else None
         if do_cfg:
             uncond_velocity_vision, uncond_velocity_sound, uncond_velocity_action = velocities["uncond"]
+            uncond_velocity_vision = uncond_velocity_vision.float()
+            uncond_velocity_sound = uncond_velocity_sound.float() if uncond_velocity_sound is not None else None
+            uncond_velocity_action = uncond_velocity_action.float() if uncond_velocity_action is not None else None
             block_state.velocity_vision = uncond_velocity_vision + block_state.guidance_scale * (
                 cond_velocity_vision - uncond_velocity_vision
             )
@@ -332,12 +328,6 @@ def inputs(self) -> list[InputParam]:
                 description="Indexes of conditioned vision latent frames; non-empty for image-to-video.",
             ),
             InputParam.template("generator"),
-            InputParam(
-                name="use_fp32_sampling_state",
-                type_hint=bool,
-                default=False,
-                description="Whether to keep the distilled vision scheduler state in float32.",
-            ),
         ]
 
     @property
@@ -346,11 +336,8 @@ def intermediate_outputs(self) -> list[OutputParam]:
 
     @torch.no_grad()
     def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
-        velocity_vision = block_state.velocity_vision
-        latents = block_state.latents
-        if block_state.use_fp32_sampling_state:
-            velocity_vision = velocity_vision.float()
-            latents = latents.float()
+        velocity_vision = block_state.velocity_vision.float()
+        latents = block_state.latents.float()
 
         # Pass the generator so the scheduler's stochastic (SDE) re-noising is seedable/reproducible.
         block_state.latents = components.scheduler.step(
@@ -734,12 +721,6 @@ def inputs(self) -> list[InputParam]:
                 default=None,
                 description="Timestep interval [lo, hi] over which control guidance is active (None = always).",
             ),
-            InputParam(
-                name="use_fp32_sampling_state",
-                type_hint=bool,
-                default=False,
-                description="Whether to keep velocity and guidance arithmetic in float32.",
-            ),
         ]
 
     @property
@@ -814,10 +795,9 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
                 "uncond",
             )
 
-        if block_state.use_fp32_sampling_state:
-            cond_full = cond_full.float()
-            cond_no_control = cond_no_control.float() if cond_no_control is not None else None
-            uncond_full = uncond_full.float() if uncond_full is not None else None
+        cond_full = cond_full.float()
+        cond_no_control = cond_no_control.float() if cond_no_control is not None else None
+        uncond_full = uncond_full.float() if uncond_full is not None else None
 
         if needs_control_cfg and needs_text_cfg:
             control_cond = cond_no_control + step_control * (cond_full - cond_no_control)
@@ -921,8 +901,6 @@ class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper):
               Timestep interval [lo, hi] over which text guidance is active (None = always).
           control_guidance_interval (`tuple`, *optional*):
               Timestep interval [lo, hi] over which control guidance is active (None = always).
-          use_fp32_sampling_state (`bool`, *optional*, defaults to False):
-              Whether to keep velocity and guidance arithmetic in float32.
           latents (`Tensor`):
               Noisy target latents to update.
           condition_latents (`Tensor`):
diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py
index ee02e192da21..d15a176a0824 100644
--- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py
+++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py
@@ -377,7 +377,6 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks):
           scheduler (`UniPCMultistepScheduler`)
 
       Configs:
-          default_use_fp32_sampling_state (default: True)
           use_native_flow_schedule (default: False)
 
       Inputs:
@@ -401,9 +400,6 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
-          use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
-              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`):
               The number of denoising steps.
           **denoiser_input_fields (`None`, *optional*):
@@ -453,7 +449,6 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks):
           scheduler (`UniPCMultistepScheduler`)
 
       Configs:
-          default_use_fp32_sampling_state (default: True)
           use_native_flow_schedule (default: False)
 
       Inputs:
@@ -477,9 +472,6 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
-          use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
-              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`):
               The number of denoising steps.
           sound_latents (`Tensor`, *optional*):
@@ -542,7 +534,6 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks):
           scheduler (`UniPCMultistepScheduler`)
 
       Configs:
-          default_use_fp32_sampling_state (default: True)
           use_native_flow_schedule (default: False)
 
       Inputs:
@@ -566,9 +557,6 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
-          use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
-              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`):
               The number of denoising steps.
           action (`CosmosActionCondition`):
@@ -635,7 +623,6 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks):
           scheduler (`UniPCMultistepScheduler`)
 
       Configs:
-          default_use_fp32_sampling_state (default: True)
           use_native_flow_schedule (default: False)
 
       Inputs:
@@ -659,9 +646,6 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
-          use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
-              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`):
               The number of denoising steps.
           sound_latents (`Tensor`, *optional*):
@@ -740,9 +724,6 @@ class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks):
           transformer (`Cosmos3OmniTransformer`)
           scheduler (`UniPCMultistepScheduler`)
 
-      Configs:
-          default_use_fp32_sampling_state (default: True)
-
       Inputs:
           chunk_id (`int`, *optional*, defaults to 0):
               Index of the current chunk.
@@ -768,9 +749,6 @@ class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks):
               Number of frames each later chunk reuses from the previous chunk's tail.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
-          use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
-              pipeline's `default_use_fp32_sampling_state` config.
           cond_text_segment (`dict`):
               Conditional text segment.
           uncond_text_segment (`dict`):
@@ -892,9 +870,6 @@ class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks):
           video_processor (`VideoProcessor`)
           scheduler (`UniPCMultistepScheduler`)
 
-      Configs:
-          default_use_fp32_sampling_state (default: True)
-
       Inputs:
           cond_input_ids (`None`):
               Token IDs for the conditional prompt.
@@ -924,9 +899,6 @@ class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks):
               Number of frames each later chunk reuses from the previous chunk's tail.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
-          use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
-              pipeline's `default_use_fp32_sampling_state` config.
           fps (`float`, *optional*, defaults to 24.0):
               Frame rate of the generated video.
           num_inference_steps (`int`):
@@ -1020,7 +992,6 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks):
           scheduler (`UniPCMultistepScheduler`)
 
       Configs:
-          default_use_fp32_sampling_state (default: True)
           use_native_flow_schedule (default: False)
 
       Inputs:
@@ -1052,9 +1023,6 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks):
               Number of frames each later chunk reuses from the previous chunk's tail.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
-          use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
-              pipeline's `default_use_fp32_sampling_state` config.
           fps (`float`, *optional*, defaults to 24.0):
               Frame rate of the generated video.
           num_inference_steps (`int`):
@@ -1219,7 +1187,6 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks):
       Configs:
           default_use_system_prompt (default: True)
           enable_safety_checker (default: True)
-          default_use_fp32_sampling_state (default: True)
           use_native_flow_schedule (default: False)
 
       Inputs:
@@ -1265,9 +1232,6 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks):
               Number of frames the first chunk reuses from the input video.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
-          use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
-              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`):
               The number of denoising steps.
           **denoiser_input_fields (`None`, *optional*):
diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
index 6841d2f3ff4f..b6ec5d47eb8d 100644
--- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
+++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
@@ -88,7 +88,6 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks):
           scheduler (`FlowMatchEulerDiscreteScheduler`)
 
       Configs:
-          default_use_fp32_sampling_state (default: True)
           is_distilled (default: True)
           distilled_sigmas (default: None)
 
@@ -113,9 +112,6 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
-          use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
-              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`, *optional*):
               The number of denoising steps.
           guidance_scale (`float`, *optional*):
@@ -177,7 +173,6 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks):
       Configs:
           default_use_system_prompt (default: True)
           enable_safety_checker (default: True)
-          default_use_fp32_sampling_state (default: True)
           is_distilled (default: True)
           distilled_sigmas (default: None)
 
@@ -214,9 +209,6 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks):
               Pre-generated noisy vision latents.
           generator (`Generator`, *optional*):
               Torch generator for deterministic generation.
-          use_fp32_sampling_state (`bool | NoneType`, *optional*):
-              Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the
-              pipeline's `default_use_fp32_sampling_state` config.
           num_inference_steps (`int`, *optional*):
               The number of denoising steps.
           guidance_scale (`float`, *optional*):
diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
index b6982789a834..9f22da1035bb 100644
--- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
+++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
@@ -1352,7 +1352,6 @@ def __call__(
         add_resolution_template: bool = True,
         add_duration_template: bool = True,
         enable_safety_check: bool = True,
-        use_fp32_sampling_state: bool = True,
     ) -> Cosmos3OmniPipelineOutput:
         r"""
         Run the Cosmos 3 omni pipeline end-to-end: encode the (optional) conditioning image/video, denoise vision and
@@ -1448,11 +1447,6 @@ def __call__(
                 When `True` and a `CosmosSafetyChecker` is attached, runs the text guardrail on the prompt before
                 generation and the video guardrail on the decoded frames. Set to `False` to skip both for this call;
                 the checker remains loaded for subsequent calls.
-            use_fp32_sampling_state (`bool`, *optional*, defaults to `True`):
-                When `True`, keeps vision, sound, and action denoising latents plus classifier-free-guidance arithmetic
-                in `torch.float32`. Transformer inputs are still cast to the transformer's dtype before each forward.
-                This improves sampling-state precision at the cost of additional memory. Set it to `False` to keep
-                sampling state in the model dtype.
 
         Returns:
             [`Cosmos3OmniPipelineOutput`] or `tuple`:
@@ -1517,7 +1511,7 @@ def __call__(
 
         device = self._get_execution_device()
         dtype = self.transformer.dtype
-        sampling_dtype = torch.float32 if use_fp32_sampling_state else dtype
+        sampling_dtype = torch.float32
 
         if enable_safety_check and isinstance(self.safety_checker, CosmosSafetyChecker):
             self.safety_checker.to(device)
@@ -1811,13 +1805,12 @@ def __call__(
                         raw_action_dim=raw_action_dim_resolved,
                     )
 
-                if use_fp32_sampling_state:
-                    cond_v_vision = cond_v_vision.float()
-                    cond_v_sound = cond_v_sound.float() if cond_v_sound is not None else None
-                    cond_v_action = cond_v_action.float() if cond_v_action is not None else None
-                    uncond_v_vision = uncond_v_vision.float() if uncond_v_vision is not None else None
-                    uncond_v_sound = uncond_v_sound.float() if uncond_v_sound is not None else None
-                    uncond_v_action = uncond_v_action.float() if uncond_v_action is not None else None
+                cond_v_vision = cond_v_vision.float()
+                cond_v_sound = cond_v_sound.float() if cond_v_sound is not None else None
+                cond_v_action = cond_v_action.float() if cond_v_action is not None else None
+                uncond_v_vision = uncond_v_vision.float() if uncond_v_vision is not None else None
+                uncond_v_sound = uncond_v_sound.float() if uncond_v_sound is not None else None
+                uncond_v_action = uncond_v_action.float() if uncond_v_action is not None else None
 
                 # --- CFG combine + per-modality scheduler step ---
                 # UniPC's multistep_uni_p_bh_update einsum ("k,bkc...->bc...") requires sample
@@ -1862,9 +1855,7 @@ def __call__(
                     for key in callback_on_step_end_tensor_inputs:
                         callback_kwargs[key] = locals()[key]
                     callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-                    latents = callback_outputs.pop("latents", latents)
-                    if use_fp32_sampling_state:
-                        latents = latents.float()
+                    latents = callback_outputs.pop("latents", latents).float()
 
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                     progress_bar.update()
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
index 013940fd6970..abe293471ac8 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
@@ -200,29 +200,20 @@ def _get_sampling_state_block_pipe(self, block):
         pipe.to(torch_device)
         return pipe
 
-    @pytest.mark.parametrize(
-        ("use_fp32_sampling_state", "expected_dtype"),
-        [(False, torch.bfloat16), (True, torch.float32)],
-    )
-    def test_sound_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype):
+    def test_sound_prepare_latents_uses_fp32(self):
         pipe = self._get_sampling_state_block_pipe(Cosmos3SoundPrepareLatentsStep())
 
         outputs = pipe(
             num_frames=5,
             fps=24.0,
             generator=self.get_generator(0),
-            use_fp32_sampling_state=use_fp32_sampling_state,
             output=["sound_latents", "sound_condition_mask"],
         )
 
-        assert outputs["sound_latents"].dtype == expected_dtype
-        assert outputs["sound_condition_mask"].dtype == expected_dtype
+        assert outputs["sound_latents"].dtype == torch.float32
+        assert outputs["sound_condition_mask"].dtype == torch.float32
 
-    @pytest.mark.parametrize(
-        ("use_fp32_sampling_state", "expected_dtype"),
-        [(False, torch.bfloat16), (True, torch.float32)],
-    )
-    def test_action_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype):
+    def test_action_prepare_latents_uses_fp32(self):
         pipe = self._get_sampling_state_block_pipe(Cosmos3ActionPrepareLatentsStep())
         action = CosmosActionCondition(
             mode="policy",
@@ -235,31 +226,25 @@ def test_action_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, ex
             action=action,
             action_condition_frame_indexes=[],
             generator=self.get_generator(0),
-            use_fp32_sampling_state=use_fp32_sampling_state,
             output=["action_latents", "action_condition_mask"],
         )
 
-        assert outputs["action_latents"].dtype == expected_dtype
-        assert outputs["action_condition_mask"].dtype == expected_dtype
+        assert outputs["action_latents"].dtype == torch.float32
+        assert outputs["action_condition_mask"].dtype == torch.float32
 
-    @pytest.mark.parametrize(
-        ("use_fp32_sampling_state", "expected_dtype"),
-        [(False, torch.bfloat16), (True, torch.float32)],
-    )
-    def test_transfer_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype):
+    def test_transfer_prepare_latents_uses_fp32(self):
         pipe = self._get_sampling_state_block_pipe(Cosmos3TransferPrepareLatentsStep())
 
         outputs = pipe(
             x0_tokens_vision=torch.zeros(1, 4, 2, 2, 2),
             current_conditional_frames=1,
             generator=self.get_generator(0),
-            use_fp32_sampling_state=use_fp32_sampling_state,
             output=["latents", "velocity_mask", "condition_latents"],
         )
 
-        assert outputs["latents"].dtype == expected_dtype
-        assert outputs["velocity_mask"].dtype == expected_dtype
-        assert outputs["condition_latents"].dtype == expected_dtype
+        assert outputs["latents"].dtype == torch.float32
+        assert outputs["velocity_mask"].dtype == torch.float32
+        assert outputs["condition_latents"].dtype == torch.float32
 
     def test_transfer_chunks_reset_stateful_cache_at_boundaries(self):
         block = Cosmos3TransferChunkDenoiseStep()
@@ -279,19 +264,14 @@ def test_transfer_chunks_reset_stateful_cache_at_boundaries(self):
             ("chunk_id", 2),
         ]
 
-    @pytest.mark.parametrize(
-        ("use_fp32_sampling_state", "expected_dtype"),
-        [(False, torch.bfloat16), (True, torch.float32)],
-    )
-    def test_sampling_state_controls_modular_cfg_and_scheduler_dtype(self, use_fp32_sampling_state, expected_dtype):
+    def test_sampling_state_uses_fp32_for_modular_cfg_and_scheduler(self):
         pipe = self.get_pipeline(dtype=torch.bfloat16).to(torch_device)
         inputs = self.get_dummy_inputs()
-        inputs["use_fp32_sampling_state"] = use_fp32_sampling_state
 
         outputs = pipe(**inputs, output=["velocity_vision", "latents"])
 
-        assert outputs["velocity_vision"].dtype == expected_dtype
-        assert outputs["latents"].dtype == expected_dtype
+        assert outputs["velocity_vision"].dtype == torch.float32
+        assert outputs["latents"].dtype == torch.float32
 
     def test_vae_encoder_is_standalone_and_validates_conditioning_inputs(self):
         pipe = self.get_pipeline()
diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
index 354b9b08288f..86bc6ad4d6f2 100644
--- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
+++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
@@ -116,7 +116,6 @@ def test_declares_distilled_configs(self):
         assert pipe.config.is_distilled is True
         assert pipe.config.distilled_sigmas is None
         assert pipe.config.default_use_system_prompt is True
-        assert pipe.config.default_use_fp32_sampling_state is True
 
     def test_distilled_text_step_uses_system_prompt_config_fallback(self):
         text_pipe = Cosmos3DistilledTextEncoderStep().init_pipeline(self.pretrained_model_name_or_path)
@@ -142,18 +141,9 @@ def test_distilled_text_step_uses_system_prompt_config_fallback(self):
         assert explicit_without_system_prompt == default_without_system_prompt == updated_without_system_prompt
         assert len(default_with_system_prompt) > len(default_without_system_prompt)
 
-    @pytest.mark.parametrize(
-        ("config_enabled", "input_enabled", "expected_dtype"),
-        [
-            (False, None, torch.bfloat16),
-            (True, None, torch.float32),
-            (True, False, torch.bfloat16),
-        ],
-    )
-    def test_prepare_vision_latents_fp32_sampling_state(self, config_enabled, input_enabled, expected_dtype):
+    def test_prepare_vision_latents_uses_fp32(self):
         prepare_pipe = Cosmos3VisionPrepareLatentsStep().init_pipeline(self.pretrained_model_name_or_path)
         prepare_pipe.load_components(torch_dtype=torch.bfloat16)
-        prepare_pipe.update_components(default_use_fp32_sampling_state=config_enabled)
         prepare_pipe.to(torch_device)
 
         outputs = prepare_pipe(
@@ -162,26 +152,19 @@ def test_prepare_vision_latents_fp32_sampling_state(self, config_enabled, input_
             width=32,
             fps=24.0,
             generator=self.get_generator(0),
-            use_fp32_sampling_state=input_enabled,
-            output=["latents", "vision_condition_mask", "use_fp32_sampling_state"],
+            output=["latents", "vision_condition_mask"],
         )
 
-        assert outputs["latents"].dtype == expected_dtype
-        assert outputs["vision_condition_mask"].dtype == expected_dtype
-        assert outputs["use_fp32_sampling_state"] is (config_enabled if input_enabled is None else input_enabled)
+        assert outputs["latents"].dtype == torch.float32
+        assert outputs["vision_condition_mask"].dtype == torch.float32
 
-    @pytest.mark.parametrize(
-        ("use_fp32_sampling_state", "expected_dtype"),
-        [(False, torch.bfloat16), (True, torch.float32)],
-    )
-    def test_distilled_scheduler_fp32_state(self, use_fp32_sampling_state, expected_dtype):
+    def test_distilled_scheduler_uses_fp32_state(self):
         pipe = self.get_pipeline(torch_dtype=torch.bfloat16).to(torch_device)
         inputs = self.get_dummy_inputs()
-        inputs["use_fp32_sampling_state"] = use_fp32_sampling_state
 
         latents = pipe(**inputs, output=self.output_name)
 
-        assert latents.dtype == expected_dtype
+        assert latents.dtype == torch.float32
 
     def test_vae_encoder_rejects_image_and_video_together(self):
         vae_encoder = Cosmos3DistilledBlocks().sub_blocks["vae_encoder"]

From 9009730d7a85dca72f67887de30a3aeaab2a15fd Mon Sep 17 00:00:00 2001
From: root <yzhautouskay@nvidia.com>
Date: Thu, 3 Sep 2026 06:35:03 -0700
Subject: [PATCH 14/14] Add pipe definition to the docs snippet

---
 docs/source/en/api/pipelines/cosmos3.md | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md
index 5a382c4743a8..36f5f089339f 100644
--- a/docs/source/en/api/pipelines/cosmos3.md
+++ b/docs/source/en/api/pipelines/cosmos3.md
@@ -667,7 +667,12 @@ transformer residuals when the Spectral-Evolution-Aware indicator changes slowly
 transformer executions. Enable it on the transformer with scheduler metadata callbacks from the pipeline:
 
 ```python
-from diffusers import SeaCacheConfig
+import torch
+from diffusers import Cosmos3OmniPipeline, SeaCacheConfig
+
+pipe = Cosmos3OmniPipeline.from_pretrained(
+    "nvidia/Cosmos3-Nano", dtype=torch.bfloat16, device_map="cuda"
+)
 
 pipe.transformer.enable_cache(
     SeaCacheConfig(