From 6040bb97e2036927632151a87b0aaed94ccbce7d Mon Sep 17 00:00:00 2001 From: yzhautouskay Date: Tue, 15 Sep 2026 23:57:15 +0200 Subject: [PATCH] [Perf][Cosmos3] Add SeaCache support for Cosmos3 (#6922) Signed-off-by: Yuliya Zhautouskaya --- recipes/cosmos3/Cosmos3-Nano.md | 61 +- recipes/cosmos3/Cosmos3-Super.md | 27 +- tests/diffusion/cache/test_seacache.py | 635 ++++++++++++++++++ .../models/cosmos3/test_cosmos3_pipeline.py | 105 ++- .../cosmos3/test_cosmos3_transformer.py | 104 ++- .../diffusion/cache/seacache/__init__.py | 18 + vllm_omni/diffusion/cache/seacache/backend.py | 100 +++ vllm_omni/diffusion/cache/seacache/config.py | 36 + vllm_omni/diffusion/cache/seacache/hook.py | 397 +++++++++++ .../diffusion/cache/seacache/sea_filter.py | 93 +++ vllm_omni/diffusion/cache/seacache/state.py | 26 + vllm_omni/diffusion/cache/selector.py | 11 +- .../diffusion/cache/teacache/extractors.py | 100 ++- vllm_omni/diffusion/data.py | 8 + .../models/cosmos3/pipeline_cosmos3.py | 77 ++- .../models/cosmos3/transformer_cosmos3.py | 274 +++++--- .../worker/diffusion_model_runner.py | 49 +- vllm_omni/entrypoints/cli/serve.py | 5 +- 18 files changed, 1966 insertions(+), 160 deletions(-) create mode 100644 tests/diffusion/cache/test_seacache.py create mode 100644 vllm_omni/diffusion/cache/seacache/__init__.py create mode 100644 vllm_omni/diffusion/cache/seacache/backend.py create mode 100644 vllm_omni/diffusion/cache/seacache/config.py create mode 100644 vllm_omni/diffusion/cache/seacache/hook.py create mode 100644 vllm_omni/diffusion/cache/seacache/sea_filter.py create mode 100644 vllm_omni/diffusion/cache/seacache/state.py diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index 6ef24bcf4..29088e94c 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -32,13 +32,13 @@ mode is selected per request: T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples. - **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks: - - `forward_dynamics` — given a first frame or video **and** an action trajectory, + - `forward_dynamics` — given a first frame or video **and** an action trajectory, roll out the resulting video. Synchronous: `POST /v1/videos/sync`. - - `policy` — given a first frame or video and a language instruction, + - `policy` — given a first frame or video and a language instruction, **predict** the action trajectory (and a rollout video). Use the async `POST /v1/videos` endpoint and read the predicted action from the top-level `action` field. - - `inverse_dynamics` — given a video, **recover** the action trajectory. Use + - `inverse_dynamics` — given a video, **recover** the action trajectory. Use the async `POST /v1/videos` endpoint and read the recovered action from the top-level `action` field (`{data, shape, dtype, raw_action_dim, domain_id}`). @@ -106,6 +106,25 @@ vllm serve nvidia/Cosmos3-Nano \ --init-timeout 1800 ``` +For Cosmos3, SeaCache is the recommended default choice when opting into +diffusion caching. Caching remains opt-in; add `--cache-backend sea_cache` to +the command above. Its default threshold and maximum cached-step streak are +tuned for Cosmos3, so no `--cache-config` is required. + +Override individual defaults with a JSON cache configuration; for example: + +```bash +vllm serve nvidia/Cosmos3-Nano \ + --omni \ + --cache-backend sea_cache \ + --cache-config '{"sea_threshold":0.2,"sea_max_consecutive_cached":3}' +``` + +Lower `sea_threshold` values and smaller `sea_max_consecutive_cached` caps are +more conservative. Higher values allow more cached steps and may improve +speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to +`0` removes the streak cap. + To run **without** guardrails (you are responsible for license compliance), add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use `--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`; @@ -382,12 +401,12 @@ OpenPI route uses the normal API authentication middleware. RoboLab's #### Notes -- **Measured latency (1x B300, bf16, guardrails off):** - - T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s** - - T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s** - - I2V 1280×720, 189 frames @ 35 steps → ~**99 s** - - Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**. - - Guardrails-on overhead: ~8% on T2I, negligible on video. +- **Measured latency (1x B300, bf16, guardrails off, diffusion cache off):** + - T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s** + - T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s** + - I2V 1280×720, 189 frames @ 35 steps → ~**99 s** + - Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**. + - Guardrails-on overhead: ~8% on T2I, negligible on video. - **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 GPU; full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer) ~33 GB on disk. @@ -443,12 +462,12 @@ OpenPI route uses the normal API authentication middleware. RoboLab's `conditioning_fps`, `action_chunk_size`, `raw_action_dim`, `deterministic_seed`, and `session_id`. - **Known limitations:** - - Guardrails-on requires `cosmos-guardrail` **and** access to the gated + - Guardrails-on requires `cosmos-guardrail` **and** access to the gated `nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise the server fails at pipeline build with a gated-repo / safety-checker error. - - A guardrail-blocked prompt currently returns HTTP 500 + - A guardrail-blocked prompt currently returns HTTP 500 (`"Guardrail blocked prompt"`). - - Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported + - Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported online. Use async `POST /v1/videos` when you need the predicted/recovered action payload under the top-level `action` field; sync `/v1/videos/sync` returns raw MP4 bytes and does not expose action metadata in the response body. @@ -579,14 +598,14 @@ so an md5 comparison across two runs also works as a smoke check. a ~390 s build of the VAE decode path *for that shape*, so warming 1024² images does nothing for 189-frame video. At 720p / 189 frames / 35 steps this is **540 s cold vs 161 s warm, with byte-identical output**. -- **Measured on 1x MI350X (bf16, guardrails off, warm):** T2I 1024² @ 50 steps +- **Measured on 1x MI350X (bf16, guardrails off, diffusion cache off, warm):** T2I 1024² @ 50 steps **~2.7 s**; T2V 1280×720 / 189 frames @ 35 steps **~161 s**, of which ~92% is the DiT denoise loop and ~4% VAE decode, so optimization effort belongs in the denoise loop. The optional flags act on different terms of the memory bill and are therefore complementary rather than redundant: | Flag | Latency | Peak reserved | Peak allocated | Acts on | - |---|---|---|---|---| + | --- | --- | --- | --- | --- | | *(none)* | 161 s | 120 GiB | 95 GiB | — | | `--vae-use-tiling` | 183 s (+13.5%) | **38 GiB** (−68%) | **36 GiB** (−62%) | decode activations | | `--enable-layerwise-offload` | 162 s (+0.8%) | 84 GiB (−30%) | 69 GiB (−27%) | weights | @@ -768,11 +787,11 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \ #### Notes -- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off):** - - T2I 1024² — 10 steps → ~8 s - - T2V 1280×720 @ 20 steps — 49 frames → ~55 s - - I2V 1280×720 @ 10 steps — 25 frames → ~25 s - - V2V 480×320 @ 10 steps — 17 frames → ~12 s +- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off, diffusion cache off):** + - T2I 1024² — 10 steps → ~8 s + - T2V 1280×720 @ 20 steps — 49 frames → ~55 s + - I2V 1280×720 @ 10 steps — 25 frames → ~25 s + - V2V 480×320 @ 10 steps — 17 frames → ~12 s - **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 NPU; full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer) ~33 GB on disk. @@ -785,9 +804,9 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \ (for model loading), `--tensor-parallel-size 8` for multi-NPU, and `--model-class-name Cosmos3OmniDiffusersPipeline` to force the pipeline class. - **Known limitations:** - - Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a + - Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a resolution-parsing bug; basic V2V without transfer hints works. - - FP8 online quantization and layerwise offload are not supported on NPU. + - FP8 online quantization and layerwise offload are not supported on NPU. ### 1x Ascend 910B / 910C (Atlas A2 / A3) — Offline generation diff --git a/recipes/cosmos3/Cosmos3-Super.md b/recipes/cosmos3/Cosmos3-Super.md index c8bfb6bbd..82ae1f20a 100644 --- a/recipes/cosmos3/Cosmos3-Super.md +++ b/recipes/cosmos3/Cosmos3-Super.md @@ -62,6 +62,21 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs; generation from ~83 GB to ~55 GB per GPU (2-GPU) with BF16-level quality (T2V composition can shift at the same seed). +SeaCache is the recommended default choice when opting into diffusion caching, +which remains opt-in. Add `--cache-backend sea_cache` to either command above; +its defaults are tuned for Cosmos3. + +Override individual defaults with `--cache-config`; for example: + +```bash +--cache-config '{"sea_threshold":0.2,"sea_max_consecutive_cached":3}' +``` + +Lower `sea_threshold` values and smaller `sea_max_consecutive_cached` caps are +more conservative. Higher values allow more cached steps and may improve +speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to +`0` removes the streak cap. + #### Verification Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full @@ -109,12 +124,12 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ #### Notes -- **Measured (2x B300, bf16, guardrails off, official 2-GPU config above):** - - T2I 1024², 50 steps → **~6 s** - - T2V 1280×720, 189 frames, 35 steps → **~197 s** - - I2V 1280×720, 189 frames, 35 steps → **~200 s** - - T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo** - - (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.) +- **Measured (2x B300, bf16, guardrails off, diffusion cache off, official 2-GPU config above):** + - T2I 1024², 50 steps → **~6 s** + - T2V 1280×720, 189 frames, 35 steps → **~197 s** + - I2V 1280×720, 189 frames, 35 steps → **~200 s** + - T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo** + - (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.) - **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk. - Same generation defaults, supported sizes, V2V reference-video controls (`condition_frame_indexes_vision`, `condition_video_keep`), and diff --git a/tests/diffusion/cache/test_seacache.py b/tests/diffusion/cache/test_seacache.py new file mode 100644 index 000000000..ea11537d1 --- /dev/null +++ b/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 diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 9892a72d7..e73dd7ff6 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/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, diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 44bd8e9a7..de14ee233 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/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 diff --git a/vllm_omni/diffusion/cache/seacache/__init__.py b/vllm_omni/diffusion/cache/seacache/__init__.py new file mode 100644 index 000000000..2821c4de0 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/diffusion/cache/seacache/backend.py b/vllm_omni/diffusion/cache/seacache/backend.py new file mode 100644 index 000000000..e5e513a73 --- /dev/null +++ b/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") diff --git a/vllm_omni/diffusion/cache/seacache/config.py b/vllm_omni/diffusion/cache/seacache/config.py new file mode 100644 index 000000000..501d4653e --- /dev/null +++ b/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}") diff --git a/vllm_omni/diffusion/cache/seacache/hook.py b/vllm_omni/diffusion/cache/seacache/hook.py new file mode 100644 index 000000000..3fac58c09 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/cache/seacache/sea_filter.py b/vllm_omni/diffusion/cache/seacache/sea_filter.py new file mode 100644 index 000000000..f68b5c288 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/cache/seacache/state.py b/vllm_omni/diffusion/cache/seacache/state.py new file mode 100644 index 000000000..bdf21c1f4 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/cache/selector.py b/vllm_omni/diffusion/cache/selector.py index d58bada6c..79bb86a4c 100644 --- a/vllm_omni/diffusion/cache/selector.py +++ b/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'" ) diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 7f0242d57..3aca638ec 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/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." ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 7a7de3bb1..c333fb11a 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -541,6 +541,8 @@ class DiffusionCacheConfig: scm_steps_mask_policy, scm_steps_policy - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio, mag_ratios, mag_calibrate + - SeaCache: sea_threshold, sea_residual_order, + sea_max_consecutive_cached, sea_power_exp - step_cache: step_cache_dit_enabled, velocity_sim_thresholds, velocity_skip_countdowns, step_cache_dit_min_history @@ -560,6 +562,12 @@ class DiffusionCacheConfig: rel_l1_thresh: float | None = None coefficients: list[float] | None = None # Uses model-specific defaults if None + # SeaCache parameters [sea_cache only] + sea_threshold: float = 0.25 + sea_residual_order: int = 1 + sea_max_consecutive_cached: int = 2 + sea_power_exp: float = 3.0 + # MagCache parameters [mag_cache only] # Default: 0.24 threshold for accumulated magnitude error mag_threshold: float = 0.24 diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 5977e0eea..04f25ac58 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Cosmos3 text/image/video/sound/action pipeline for vllm-omni. One pipeline class serves the Cosmos3 family modes. Output modality is selected @@ -34,6 +34,7 @@ import math import os import time from collections.abc import Iterable, Mapping +from contextlib import nullcontext from dataclasses import fields from typing import Any, ClassVar @@ -1075,6 +1076,8 @@ class Cosmos3OmniDiffusersPipeline( self._guidance_scale = None self._num_timesteps = None + self._current_step_index = None + self._current_sigma = None self._cosmos3_branch_caches: dict[str, tuple[Any, Any]] | None = None self._robolab_transforms: dict[bool, Any] = {} @@ -1291,17 +1294,24 @@ class Cosmos3OmniDiffusersPipeline( or a tuple in video, action, sound order for multimodal generation. """ cache_key = kwargs.pop("_cosmos3_cache_key", None) - if cache_key is None: - return self.transformer(**kwargs) + context_name = str(kwargs.pop("_cache_context", "cond")) - branch_caches = self._cosmos3_branch_caches - if branch_caches is None: - return self.transformer(**kwargs) - - cache_key = str(cache_key) - self.transformer.cached_kv, self.transformer.cached_freqs_gen = branch_caches.get(cache_key, (None, None)) - prediction = self.transformer(**kwargs) - branch_caches[cache_key] = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) + context_factory = getattr(self, "_cache_context_factory", None) + context = context_factory(context_name) if callable(context_factory) else nullcontext() + with context: + if cache_key is None: + prediction = self.transformer(**kwargs) + else: + branch_caches = self._cosmos3_branch_caches + if branch_caches is None: + prediction = self.transformer(**kwargs) + else: + cache_key = str(cache_key) + self.transformer.cached_kv, self.transformer.cached_freqs_gen = branch_caches.get( + cache_key, (None, None) + ) + prediction = self.transformer(**kwargs) + branch_caches[cache_key] = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) return prediction def combine_multi_branch_cfg_noise( @@ -1840,6 +1850,29 @@ class Cosmos3OmniDiffusersPipeline( def num_timesteps(self): return self._num_timesteps + @property + def current_step_index(self): + return self._current_step_index + + @property + def current_sigma(self): + return self._current_sigma + + def _set_denoise_step_metadata( + self, + step_index: int, + timesteps: torch.Tensor, + scheduler: Any, + ) -> None: + self._current_step_index = step_index + self._num_timesteps = len(timesteps) + sigmas = getattr(scheduler, "sigmas", None) + self._current_sigma = sigmas[step_index] if sigmas is not None and step_index < len(sigmas) else None + + def _clear_denoise_step_metadata(self) -> None: + self._current_step_index = None + self._current_sigma = None + def _set_mixed_precision_step(self, step_index: int, num_steps: int) -> None: setter = getattr(self.transformer, "set_mixed_precision_step", None) if setter is not None: @@ -2848,6 +2881,7 @@ class Cosmos3OmniDiffusersPipeline( # else uncond), so session keying loads/stores only this rank's branch. cfg_rank_is_negative = get_classifier_free_guidance_rank() != 0 for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_denoise_step_metadata(step_index, timesteps, step_scheduler) self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) # Outside the interval, scale=1 makes the combined output equal @@ -2859,6 +2893,7 @@ class Cosmos3OmniDiffusersPipeline( do_true_cfg=True, true_cfg_scale=step_scale, positive_kwargs=dict( + _cache_context="cond", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -2868,6 +2903,7 @@ class Cosmos3OmniDiffusersPipeline( **shared_kwargs, ), negative_kwargs=dict( + _cache_context="uncond", hidden_states=latents, timestep=timestep, text_ids=uncond_ids, @@ -2888,13 +2924,15 @@ class Cosmos3OmniDiffusersPipeline( keep_uncond_for_cache = self._cache_requires_paired_cfg() for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_denoise_step_metadata(step_index, timesteps, step_scheduler) self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) cfg_active = _cfg_active_at(t) if not self._kv_load_und(kv_state, is_negative=False): self.transformer.cached_kv, self.transformer.cached_freqs_gen = cond_cache - noise_cond = self.transformer( + noise_cond = self.predict_noise( + _cache_context="cond", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -2911,7 +2949,8 @@ class Cosmos3OmniDiffusersPipeline( if cfg_active or keep_uncond_for_cache: if not self._kv_load_und(kv_state, is_negative=True): self.transformer.cached_kv, self.transformer.cached_freqs_gen = uncond_cache - noise_uncond = self.transformer( + noise_uncond = self.predict_noise( + _cache_context="uncond", hidden_states=latents, timestep=timestep, text_ids=uncond_ids, @@ -2943,10 +2982,12 @@ class Cosmos3OmniDiffusersPipeline( # No CFG: a single cond branch per step. Bespoke (state None) keeps # using the transformer-instance cache exactly as before. for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_denoise_step_metadata(step_index, timesteps, step_scheduler) self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) self._kv_load_und(kv_state, is_negative=False) - noise_pred = self.transformer( + noise_pred = self.predict_noise( + _cache_context="cond", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -2959,6 +3000,7 @@ class Cosmos3OmniDiffusersPipeline( self._kv_capture_und(kv_state, is_negative=False) _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents)) finally: + self._clear_denoise_step_metadata() self._reset_mixed_precision() # Cosmos3 currently receives a unique request_id rather than a # reusable rollout session id. Retaining its state would only pin @@ -3145,6 +3187,7 @@ class Cosmos3OmniDiffusersPipeline( self._cosmos3_branch_caches = {} try: for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_denoise_step_metadata(step_index, timesteps, self.scheduler) self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) step_guidance = guidance_scale if _active_at(t, guidance_interval) else 1.0 @@ -3153,6 +3196,7 @@ class Cosmos3OmniDiffusersPipeline( needs_control_cfg = step_control != 1.0 cond_full_kwargs = dict( + _cache_context="cond", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -3166,6 +3210,7 @@ class Cosmos3OmniDiffusersPipeline( branches_kwargs = [ cond_full_kwargs, dict( + _cache_context="cond_no_control", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -3175,6 +3220,7 @@ class Cosmos3OmniDiffusersPipeline( **shared_kwargs, ), dict( + _cache_context="uncond", hidden_states=latents, timestep=timestep, text_ids=uncond_ids, @@ -3200,6 +3246,7 @@ class Cosmos3OmniDiffusersPipeline( branches_kwargs = [ cond_full_kwargs, dict( + _cache_context="cond_no_control", hidden_states=latents, timestep=timestep, text_ids=cond_ids, @@ -3223,6 +3270,7 @@ class Cosmos3OmniDiffusersPipeline( branches_kwargs = [ cond_full_kwargs, dict( + _cache_context="uncond", hidden_states=latents, timestep=timestep, text_ids=uncond_ids, @@ -3257,6 +3305,7 @@ class Cosmos3OmniDiffusersPipeline( )[0] latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents finally: + self._clear_denoise_step_metadata() self._reset_mixed_precision() self._cosmos3_branch_caches = None self.transformer.reset_cache() diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 1ded3db2f..fc5a6f7cc 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Cosmos3 VFM Transformer for vllm-omni. Implements the Mixture-of-Transformers architecture with two pathways: @@ -14,7 +14,7 @@ from __future__ import annotations import math from collections.abc import Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, NamedTuple import torch import torch.distributed as dist @@ -427,7 +427,7 @@ class TimestepEmbedder(nn.Module): max_period: int = 10000, ) -> None: super().__init__() - # Following diffusers naming pattern here for checkpoint compatibility. + # Preserve checkpoint-compatible parameter names. self.linear_1 = nn.Linear(frequency_embedding_size, hidden_size, bias=True) self.act = nn.SiLU() self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True) @@ -1105,6 +1105,28 @@ class Cosmos3GenSPPrepare(nn.Module): return hidden_gen, freqs_cos, freqs_sin +class _GenPrepared(NamedTuple): + """GEN-pathway state shared by normal and cached execution.""" + + hidden_gen: torch.Tensor + time_embed: torch.Tensor + t: int + h: int + w: int + s_video: int + s_control: int + s_action: int + s_sound: int + has_control: bool + has_action: bool + has_sound: bool + action_domain_ids: torch.Tensor | None + ulysses_size: int + use_multi_control_attention: bool + multi_control_token_sizes: tuple[int, ...] | None + multi_control_weights: tuple[float, ...] | None + + class Cosmos3VFMTransformer(nn.Module): """Cosmos3 VFM Transformer: UND language model + GEN denoising layers. @@ -1675,6 +1697,61 @@ class Cosmos3VFMTransformer(nn.Module): pad = (-(base + sound_frames)) % ulysses_size return sound_frames + pad + def _run_gen_layers( + self, + hidden_gen: torch.Tensor, + *, + s_video: int, + s_control: int, + s_action: int, + s_sound: int, + has_action: bool, + has_sound: bool, + has_control: bool, + ulysses_size: int, + use_multi_control_attention: bool, + multi_control_token_sizes: tuple[int, ...] | None, + multi_control_weights: tuple[float, ...] | None, + ) -> torch.Tensor: + """Run the complete GEN decoder between full-layout boundaries.""" + if self.cached_kv is None or self.cached_freqs_gen is None: + raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") + freqs_cos, freqs_sin = self.cached_freqs_gen + if not use_multi_control_attention: + hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) + freqs_gen = (freqs_cos, freqs_sin) + + if len(self.gen_layers) == len(self.cached_kv): + for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True): + hidden_gen = layer( + hidden_gen, + k_und=k_und, + v_und=v_und, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + control_token_sizes=multi_control_token_sizes, + control_weights=multi_control_weights, + ) + # Cache-dit's block wrapper may return a tuple; unwrap it. + if isinstance(hidden_gen, tuple): + hidden_gen = hidden_gen[0] + else: + # Cache-dit patches gen_layers to a grouped wrapper. + for layer in self.gen_layers: + hidden_gen = layer( + hidden_gen, + cached_kv=self.cached_kv, + freqs_gen=freqs_gen, + control_token_sizes=multi_control_token_sizes, + control_weights=multi_control_weights, + ) + if isinstance(hidden_gen, tuple): + hidden_gen = hidden_gen[0] + + if not use_multi_control_attention: + hidden_gen = self.gen_sp_gather(hidden_gen) + return hidden_gen + # -- Forward ------------------------------------------------------------- def forward( @@ -1697,7 +1774,51 @@ class Cosmos3VFMTransformer(nn.Module): transfer_share_vision_temporal_positions: bool = True, **kwargs, ) -> torch.Tensor | tuple[torch.Tensor, ...]: + """Run the shared Cosmos3 GEN preprocess, stack, and postprocess path.""" + if kwargs: + raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}") + prep = self._gen_preprocess( + hidden_states, + timestep, + text_ids, + text_mask, + video_shape, + fps=fps, + action_latents=action_latents, + action_domain_ids=action_domain_ids, + action_noisy_mask=action_noisy_mask, + action_start_frame_offset=action_start_frame_offset, + action_fps=action_fps, + sound_latents=sound_latents, + noisy_frame_mask=noisy_frame_mask, + control_latents=control_latents, + control_weights=control_weights, + transfer_share_vision_temporal_positions=transfer_share_vision_temporal_positions, + ) + return self._gen_postprocess(self._run_gen_stack(prep), prep) + + def _gen_preprocess( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + text_ids: torch.Tensor, + text_mask: torch.Tensor, + video_shape: tuple[int, int, int], + fps: float | None = None, + action_latents: torch.Tensor | None = None, + action_domain_ids: torch.Tensor | None = None, + action_noisy_mask: torch.Tensor | None = None, + action_start_frame_offset: int = 1, + action_fps: float | None = None, + sound_latents: torch.Tensor | None = None, + noisy_frame_mask: torch.Tensor | None = None, + control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None, + control_weights: list[float] | tuple[float, ...] | torch.Tensor | None = None, + transfer_share_vision_temporal_positions: bool = True, + ) -> _GenPrepared: """ + Prepare the packed GEN sequence before its cacheable execution region. + Args: hidden_states: [B, C, t, h, w] noisy latents timestep: [B] diffusion timestep @@ -1720,13 +1841,9 @@ class Cosmos3VFMTransformer(nn.Module): transfer controls. Values are normalized to sum to one. Returns: - [B, C, t, h, w] velocity prediction, or - tuple outputs in video, action, sound order when action/sound streams - are provided. Transfer-control streams condition the video prediction - and are not returned. + Packed GEN inputs and metadata consumed by the shared execution and + postprocessing helpers. """ - if kwargs: - raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}") t, h, w = video_shape hp, wp, _, _ = self._pad_to_patch_size(h, w) text_lengths = text_mask.sum(dim=1) @@ -1926,78 +2043,81 @@ class Cosmos3VFMTransformer(nn.Module): hidden_parts.append(hidden_sound) hidden_gen = torch.cat(hidden_parts, dim=1) - # Run GEN layers. UND K/V (replicated) is passed to each layer; - # the Cosmos3CrossAttention forwards them as joint_key/value so the - # framework Attention handles the Ulysses head-slicing internally. - if self.cached_kv is None or self.cached_freqs_gen is None: - raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") - freqs_cos, freqs_sin = self.cached_freqs_gen - if not use_multi_control_attention: - hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) - freqs_gen = (freqs_cos, freqs_sin) + return _GenPrepared( + hidden_gen=hidden_gen, + time_embed=time_embed, + t=t, + h=h, + w=w, + s_video=s_video, + s_control=s_control, + s_action=s_action, + s_sound=s_sound, + has_control=has_control, + has_action=has_action, + has_sound=has_sound, + action_domain_ids=action_domain_ids, + ulysses_size=ulysses_size, + use_multi_control_attention=use_multi_control_attention, + multi_control_token_sizes=multi_control_token_sizes, + multi_control_weights=multi_control_weights, + ) - if len(self.gen_layers) == len(self.cached_kv): - for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True): - hidden_gen = layer( - hidden_gen, - k_und=k_und, - v_und=v_und, - freqs_cos=freqs_cos, - freqs_sin=freqs_sin, - control_token_sizes=multi_control_token_sizes, - control_weights=multi_control_weights, - ) - # Cache-dit's block wrapper may return a tuple; unwrap it. - if isinstance(hidden_gen, tuple): - hidden_gen = hidden_gen[0] - else: - # Cache-dit patches gen_layers to a grouped wrapper. - for layer in self.gen_layers: - hidden_gen = layer( - hidden_gen, - cached_kv=self.cached_kv, - freqs_gen=freqs_gen, - control_token_sizes=multi_control_token_sizes, - control_weights=multi_control_weights, - ) - if isinstance(hidden_gen, tuple): - hidden_gen = hidden_gen[0] + def _run_gen_stack(self, prep: _GenPrepared) -> torch.Tensor: + """Execute the cacheable full-layout GEN stack, including final norm.""" + hidden_gen = self._run_gen_layers( + prep.hidden_gen, + s_video=prep.s_video, + s_control=prep.s_control, + s_action=prep.s_action, + s_sound=prep.s_sound, + has_action=prep.has_action, + has_sound=prep.has_sound, + has_control=prep.has_control, + ulysses_size=prep.ulysses_size, + use_multi_control_attention=prep.use_multi_control_attention, + multi_control_token_sizes=prep.multi_control_token_sizes, + multi_control_weights=prep.multi_control_weights, + ) + return self.norm_moe_gen(hidden_gen) - if not use_multi_control_attention: - hidden_gen = self.gen_sp_gather(hidden_gen) + def _gen_postprocess( + self, + hidden_gen: torch.Tensor, + prep: _GenPrepared, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + """Project an already-normalized packed GEN state to model outputs.""" + if not prep.has_action and not prep.has_sound and not prep.has_control: + return self.unpatchify(self.proj_out(hidden_gen), prep.t, prep.h, prep.w) - # Final norm and project back to latent space - hidden_gen = self.norm_moe_gen(hidden_gen) - if not has_action and not has_sound and not has_control: - return self.unpatchify(self.proj_out(hidden_gen), t, h, w) - - split_sizes = [] - if has_control: - split_sizes.append(s_control) - split_sizes.append(s_video) - if has_action: - split_sizes.append(s_action) - if has_sound: - split_sizes.append(s_sound) - split_hidden = hidden_gen.split(split_sizes, dim=1) - split_idx = 0 - if has_control: - split_idx += 1 - hidden_video = split_hidden[split_idx] + split_sizes = [] + if prep.has_control: + split_sizes.append(prep.s_control) + split_sizes.append(prep.s_video) + if prep.has_action: + split_sizes.append(prep.s_action) + if prep.has_sound: + split_sizes.append(prep.s_sound) + split_hidden = hidden_gen.split(split_sizes, dim=1) + split_idx = 0 + if prep.has_control: split_idx += 1 - video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w) - if has_control: - return video_pred - outputs: list[torch.Tensor] = [video_pred] - if has_action: - hidden_action = split_hidden[split_idx] - split_idx += 1 - assert action_domain_ids is not None - outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids))) - if has_sound: - hidden_sound = split_hidden[split_idx] - outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound))) - return tuple(outputs) + hidden_video = split_hidden[split_idx] + split_idx += 1 + video_pred = self.unpatchify(self.proj_out(hidden_video), prep.t, prep.h, prep.w) + if prep.has_control: + return video_pred + + outputs: list[torch.Tensor] = [video_pred] + if prep.has_action: + hidden_action = split_hidden[split_idx] + split_idx += 1 + assert prep.action_domain_ids is not None + outputs.append(self.unpack_action(self.action_proj_out(hidden_action, prep.action_domain_ids))) + if prep.has_sound: + hidden_sound = split_hidden[split_idx] + outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound))) + return tuple(outputs) def post_load_weights(self) -> None: """Post-load processing: ensure correct dtypes.""" diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 9dfcd6f3f..5fb124a18 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -25,7 +25,7 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.cache.cachedit import cache_summary +from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend, cache_summary from vllm_omni.diffusion.cache.prompt_embed_cache import ( install_prompt_embed_cache, resolve_prompt_embed_cache_config, @@ -409,7 +409,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # may explicitly adopt the enabled Cache-DiT backend and then # own all later request-boundary enable/disable transitions. self.cache_backend.enable(self.pipeline) - if str(self.od_config.cache_backend).lower() == "cache_dit" and adopt_request_scoped_cache_dit( + if isinstance(self.cache_backend, CacheDiTBackend) and adopt_request_scoped_cache_dit( self.pipeline, self.cache_backend, ): @@ -663,12 +663,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): num_inference_steps = getattr(self.pipeline, "default_num_inference_steps", None) if num_inference_steps is None and od_config.cache_backend in ( "tea_cache", + "sea_cache", "step_cache", ): # When num_inference_steps is None, some pipelines defer to their - # own defaults. TeaCache refresh ignores this value; step_cache - # refresh is a no-op because per-chunk state resets in the denoise - # loop. Use the pipeline default when available to keep refresh + # own defaults. These backends use refresh to reset request state; + # runtime step metadata is either unused or resolved in the + # pipeline. Use the pipeline default when available to keep refresh # behavior aligned with single-request execution. num_inference_steps = getattr(self.pipeline, "num_inference_steps", 0) or 0 @@ -1046,6 +1047,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): new_request_ids: list[str], ) -> tuple[list[StepRequestState], InputBatch | None, list[RunnerOutput]]: # process new reqs + pipeline = self.pipeline + assert pipeline is not None, "Model not loaded. Call load_model() first." prepared_states: list[StepRequestState] = [] error_outputs: list[RunnerOutput] = [] for state in states: @@ -1058,12 +1061,12 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): per_req_exc: BaseException | None = None try: self._initialize_generator(state.sampling) - clear_pipeline_stage_durations(self.pipeline) + clear_pipeline_stage_durations(pipeline) # encode - self.pipeline.prepare_encode(state) + pipeline.prepare_encode(state) merge_stage_durations( state, - consume_pipeline_stage_durations(self.pipeline), + consume_pipeline_stage_durations(pipeline), ) except Exception as exc: per_req_exc = exc @@ -1247,6 +1250,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): use_hsdp = self.od_config.parallel_config.use_hsdp grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() with grad_context: + pipeline = self.pipeline + assert pipeline is not None, "Model not loaded. Call load_model() first." had_active_states = bool(self.state_cache) states, new_request_ids = self._update_states(scheduler_output) is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 @@ -1261,7 +1266,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): states, input_batch, runner_output_list = self._prepare_batch_inputs(states, new_request_ids) if input_batch is None: return BatchRunnerOutput.from_list(runner_output_list) - attn_metadata = {} + attn_metadata: dict[str, Any] = {} kv_backend = getattr(self, "diffusion_kv_backend", None) paged_kv_runtime = kv_backend if getattr(kv_backend, "paged_attention_adapter", None) is not None else None @@ -1272,16 +1277,16 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): paged_kv_runtime=paged_kv_runtime, in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile, ): - clear_pipeline_stage_durations(self.pipeline) - noise_pred = self.pipeline.denoise_step(input_batch, states=states) - denoise_stage_durations = consume_pipeline_stage_durations(self.pipeline) + clear_pipeline_stage_durations(pipeline) + noise_pred = pipeline.denoise_step(input_batch, states=states) + denoise_stage_durations = consume_pipeline_stage_durations(pipeline) for state in states: merge_stage_durations( state, denoise_stage_durations, ) - pipeline_interrupted = getattr(self.pipeline, "interrupt", False) + pipeline_interrupted = getattr(pipeline, "interrupt", False) if noise_pred is None and pipeline_interrupted: for state in states: runner_output_list.append( @@ -1296,9 +1301,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): else: offset = 0 for req in states: + if req.latents is None: + raise RuntimeError(f"Stepwise request {req.request_id} has no latent state.") row_num = req.latents.shape[0] try: - self.pipeline.step_scheduler( + pipeline.step_scheduler( req, noise_pred[offset : offset + row_num] if noise_pred is not None else None ) if self.od_config.streaming_output: @@ -1307,8 +1314,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): should_decode = req.denoise_completed if should_decode: - clear_pipeline_stage_durations(self.pipeline) - result = self.pipeline.post_decode(req) + clear_pipeline_stage_durations(pipeline) + result = pipeline.post_decode(req) if result is not None: result = self._prepare_output_for_transport(result, req.sampling) self._attach_stepwise_metadata( @@ -1317,8 +1324,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): ) # After consuming this chunk's interaction metadata, apply pending interactions and # prepare the next chunk (prepare_next_chunk may be a no-op---depending on pipeline) - if supports_interaction_apply(self.pipeline) and not req.request_denoise_completed: - pipe = cast(SupportsInteractionApply, self.pipeline) + if supports_interaction_apply(pipeline) and not req.request_denoise_completed: + pipe = cast(SupportsInteractionApply, pipeline) pipe.apply_interaction_at_chunk_boundary(req) pipe.prepare_next_chunk(req) else: @@ -1370,12 +1377,12 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): for runner_output in runner_output_list: if runner_output.result is None: continue - state = states_by_id.get(runner_output.request_id) - if state is None: + matched_state = states_by_id.get(runner_output.request_id) + if matched_state is None: continue runner_output.result.peak_memory_mb = max( runner_output.result.peak_memory_mb, - state.peak_memory_mb, + matched_state.peak_memory_mb, ) terminal_request_ids = [ diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 0166cb4c1..bc7f81132 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -712,7 +712,10 @@ class OmniServeCommand(CLISubcommand): "--cache-backend", type=str, default="none", - help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache', 'step_cache'", + help=( + "Cache backend for diffusion models, options: 'tea_cache', " + "'cache_dit', 'mag_cache', 'sea_cache', 'step_cache'" + ), ) omni_config_group.add_argument( "--cache-config",