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

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

tests/diffusion/cache/test_seacache.py

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

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

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

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

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