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

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

vllm_omni/diffusion/cache/seacache/__init__.py

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

vllm_omni/diffusion/cache/seacache/backend.py

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

vllm_omni/diffusion/cache/seacache/config.py

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

vllm_omni/diffusion/cache/seacache/hook.py

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

vllm_omni/diffusion/cache/seacache/sea_filter.py

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

vllm_omni/diffusion/cache/seacache/state.py

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

vllm_omni/diffusion/cache/selector.py

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

vllm_omni/diffusion/cache/teacache/extractors.py

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