diff --git a/recipes/NVIDIA/SANA-Video-2B.md b/recipes/NVIDIA/SANA-Video-2B.md index ca4b44ee8..f26692867 100644 --- a/recipes/NVIDIA/SANA-Video-2B.md +++ b/recipes/NVIDIA/SANA-Video-2B.md @@ -202,6 +202,75 @@ Use the native `SanaVideoPipeline` and `SanaImageToVideoPipeline` for the primary SANA execution paths. The Diffusers adapter is retained as a validated compatibility/reference backend. +#### Native Cache-DiT and CPU offload + +The native T2V and I2V pipelines support Cache-DiT, model-level CPU offload, +and layerwise CPU offload through the common diffusion flags. Layerwise mode +keeps non-block transformer modules, the text encoder, and the VAE on the +runtime device while prefetching DiT blocks in order. + +Add one of the following flag sets to either native offline command above: + +```bash +# Cache-DiT +--cache-backend cache_dit + +# Model-level CPU offload +--enable-cpu-offload + +# Layerwise CPU offload +--enable-layerwise-offload + +# Cache-DiT plus model-level offload +--cache-backend cache_dit --enable-cpu-offload + +# Cache-DiT plus layerwise offload +--cache-backend cache_dit --enable-layerwise-offload + +# Distributed layerwise offload (weights sharded across the DP group) +--enable-distributed-layerwise-offload +``` + +If both CPU offload flags are supplied, the common offloader keeps its existing +layerwise precedence. Cache-DiT refreshes each request from the explicit +`num_inference_steps`; when that field is omitted, both native SANA pipelines +default to 50 inference steps. + +Cache-DiT and CPU offload require TP1, CFG1, and SP1. Combining them with +tensor, CFG, or sequence parallelism raises before checkpoint components are +loaded, as do other cache backends such as TeaCache. Cache-DiT cannot be +combined with distributed layerwise offload: per-rank cache skips would +desynchronize the weight AllGather. + +Measured on one A800-SXM4-80GB (driver 580.126.09, CUDA 13.0, PyTorch +2.11.0+cu130, Diffusers 0.38.0) with `DIFFUSION_ATTENTION_BACKEND=CUDNN_ATTN`, +81 frames, 50 steps, seed 42, each configuration in its own process: + +| Configuration | 480p T2V latency | 480p T2V generation peak | +|---|---:|---:| +| Baseline | 120.19 s | 24.07 GiB | +| Cache-DiT | 77.11 s (1.56x) | 24.07 GiB | +| Model CPU offload | 121.15 s | 16.81 GiB (-7.26 GiB) | +| Layerwise offload | 118.63 s | 20.75 GiB (-3.32 GiB) | + +Cache-DiT gives 1.54x on 720p I2V (37.0 s to 24.1 s). Latency is the median of +three measured runs after one warmup. Offload trades startup peak too: 14.28 +GiB baseline against 6.14 GiB for model-level and 6.80 GiB for layerwise. + +Cache-DiT is approximate: it skips block computations, so its output differs +from an uncached run. Similarity against a reference run is not used as a gate +here because this pipeline does not reproduce its own trajectory on this +hardware -- two runs of identical code differ by SSIM 0.9712 / rel_l2 8.7e-2 at +50 steps. Cache-DiT output was checked frame by frame against the uncached run +on 480p T2V and 720p I2V instead, and matches it in subject, composition, +sharpness, and lighting. + +The numbers above use the offline example's default cache configuration. +Lowering `residual_diff_threshold` in `cache_config` caches fewer steps, +trading speed for a trajectory closer to the uncached run; see the +[Cache-DiT guide](../../docs/user_guide/diffusion/cache_acceleration/cache_dit.md) +for the tuning knobs. + #### Notes - Key flags: select I2V explicitly with `--model-class-name @@ -219,8 +288,13 @@ validated compatibility/reference backend. Diffusers `DPMSolverMultistepScheduler`. - Known limitations: - Tensor parallelism (up to 2 GPUs) and CFG parallelism are supported for - the native pipeline. Sequence parallelism, Cache-DiT, TeaCache, and step - execution are not validated. + the native pipeline. Sequence parallelism, TeaCache, and step execution + are not validated. + - Cache-DiT and CPU offload are limited to TP1/CFG1/SP1. Cache-DiT with + distributed layerwise offload is not supported by the native pipeline. + - Cache-DiT speedup and offload memory numbers above are single-GPU A800 + measurements; other hardware will differ. Distributed layerwise offload is + only exercised single-rank here. - The Diffusers backend is a compatibility path and does not provide native vLLM-Omni parallelism or continuous batching. - Native describes pipeline and Transformer ownership, not a zero-Diffusers diff --git a/tests/diffusion/models/sana_video/test_cache_offload.py b/tests/diffusion/models/sana_video/test_cache_offload.py new file mode 100644 index 000000000..9ec779f11 --- /dev/null +++ b/tests/diffusion/models/sana_video/test_cache_offload.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass +from types import SimpleNamespace + +import cache_dit +import pytest +import torch +from cache_dit import ForwardPattern +from torch import nn + +import vllm_omni.diffusion.cache.cachedit.backend as cachedit_backend_module +import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module +from tests.diffusion.models.sana_video.test_transformer_sana_video import ( # noqa: F401 + _TINY_CONFIG, + _init_distributed, +) +from tests.diffusion.offloader.test_layerwise_backend import DummyEvent, DummyStream, dummy_stream +from vllm_omni.diffusion.attention import selector as attention_selector +from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend +from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.models.sana_video import pipeline_sana_video as pipeline_module +from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import ( + SanaVideoPipeline, + _validate_cache_offload_parallelism, +) +from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaVideoTransformer3DModel +from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@dataclass +class _PipelineWithTransformer: + transformer: nn.Module + + +def _tiny_transformer(monkeypatch) -> SanaVideoTransformer3DModel: + monkeypatch.setattr( + attention_selector, + "_cached_get_backend_cls", + lambda *_args, **_kwargs: SDPABackend, + ) + return SanaVideoTransformer3DModel(**(_TINY_CONFIG | {"num_layers": 2})) + + +def _record_cache_adapters(monkeypatch): + enabled_adapters = [] + original_enable_cache = cachedit_backend_module.cache_dit.enable_cache + + def record_enable_cache(block_adapter, **kwargs): + enabled_adapters.append(block_adapter) + return original_enable_cache(block_adapter, **kwargs) + + monkeypatch.setattr(cachedit_backend_module.cache_dit, "enable_cache", record_enable_cache) + return enabled_adapters + + +def test_sana_video_declares_cache_and_layerwise_metadata(): + adapter_config = SanaVideoTransformer3DModel._cache_dit_adapter_config + + assert adapter_config.block_forward_patterns == { + "transformer_blocks": ForwardPattern.Pattern_3, + } + assert adapter_config.has_separate_cfg is False + assert adapter_config.cached_adapter_cls is None + assert adapter_config.check_forward_pattern is True + assert SanaVideoTransformer3DModel._layerwise_offload_blocks_attrs == ["transformer_blocks"] + assert SanaVideoPipeline.default_num_inference_steps == 50 + + +def test_two_layer_tiny_transformer_enables_cache_dit(monkeypatch): + transformer = _tiny_transformer(monkeypatch) + pipeline = _PipelineWithTransformer(transformer=transformer) + enabled_adapters = _record_cache_adapters(monkeypatch) + backend = CacheDiTBackend() + + try: + backend.enable(pipeline) + + assert backend.is_enabled() + assert transformer._is_cached is True + assert len(enabled_adapters) == 1 + selected_blocks = cache_dit.BlockAdapter.flatten(enabled_adapters[0].blocks) + assert selected_blocks == [transformer.transformer_blocks] + assert len(selected_blocks[0]) == 2 + finally: + if enabled_adapters: + cache_dit.disable_cache(enabled_adapters[0]) + + +def test_second_engine_enables_cache_dit_after_undisabled_first(monkeypatch): + """An engine never disables Cache-DiT on shutdown, and cache_dit marks the + shared wrapper-pipe class as cached; a second engine in the same process + must still get a full enable with a working cache context.""" + enabled_adapters = _record_cache_adapters(monkeypatch) + first = _PipelineWithTransformer(transformer=_tiny_transformer(monkeypatch)) + second = _PipelineWithTransformer(transformer=_tiny_transformer(monkeypatch)) + backend = CacheDiTBackend() + + try: + CacheDiTBackend().enable(first) + backend.enable(second) + + assert backend.is_enabled() + assert second.transformer._is_cached is True + finally: + for adapter in enabled_adapters: + cache_dit.disable_cache(adapter) + + +def test_cache_dit_pattern_mismatch_fails_closed(monkeypatch): + transformer = _tiny_transformer(monkeypatch) + transformer.transformer_blocks[1].forward = lambda unexpected_input: unexpected_input + pipeline = _PipelineWithTransformer(transformer=transformer) + backend = CacheDiTBackend() + + with pytest.raises(AssertionError, match="No block forward pattern matched"): + backend.enable(pipeline) + + assert backend.is_enabled() is False + assert not getattr(transformer, "_is_cached", False) + + +@pytest.mark.parametrize( + "feature_flags", + [ + {"cache_backend": "cache_dit"}, + {"enable_cpu_offload": True}, + {"enable_layerwise_offload": True}, + {"enable_distributed_layerwise_offload": True}, + ], +) +@pytest.mark.parametrize( + "parallel_config", + [ + DiffusionParallelConfig(tensor_parallel_size=2), + DiffusionParallelConfig(cfg_parallel_size=2), + DiffusionParallelConfig(ulysses_degree=2), + ], +) +def test_cache_offload_distributed_combinations_fail_closed(feature_flags, parallel_config): + config = OmniDiffusionConfig(parallel_config=parallel_config, **feature_flags) + + with pytest.raises(NotImplementedError, match="supported only with TP1, CFG1, and SP1"): + _validate_cache_offload_parallelism(config) + + +def test_sana_video_rejects_unvalidated_cache_backends(): + with pytest.raises(NotImplementedError, match="Cache backend 'tea_cache' is not supported"): + _validate_cache_offload_parallelism(OmniDiffusionConfig(cache_backend="tea_cache")) + + +def test_sana_video_rejects_cache_dit_with_distributed_layerwise_offload(): + config = OmniDiffusionConfig(cache_backend="cache_dit", enable_distributed_layerwise_offload=True) + + with pytest.raises(NotImplementedError, match="Cache-DiT with distributed layerwise offload"): + _validate_cache_offload_parallelism(config) + + +def test_sana_video_accepts_distributed_layerwise_offload(): + _validate_cache_offload_parallelism(OmniDiffusionConfig(enable_distributed_layerwise_offload=True)) + + +@pytest.mark.parametrize( + ("od_config", "match"), + [ + ( + OmniDiffusionConfig( + cache_backend="cache_dit", + parallel_config=DiffusionParallelConfig(tensor_parallel_size=2), + ), + "tensor_parallel_size", + ), + ( + OmniDiffusionConfig(cache_backend="cache_dit", enable_distributed_layerwise_offload=True), + "Cache-DiT with distributed layerwise offload", + ), + ], +) +def test_validation_fails_before_component_loading(monkeypatch, od_config, match): + load_calls = [] + monkeypatch.setattr(pipeline_module, "get_local_device", lambda: torch.device("cpu")) + monkeypatch.setattr( + SanaVideoPipeline, + "_load_components", + lambda *args, **kwargs: load_calls.append((args, kwargs)), + ) + + with pytest.raises(NotImplementedError, match=match): + SanaVideoPipeline(od_config=od_config) + + assert load_calls == [] + + +class _TrackingModule: + def __init__(self): + self.to_calls = [] + + def to(self, *args, **kwargs): + self.to_calls.append((args, kwargs)) + return self + + +def test_component_loading_uses_loader_device_not_runtime_device(monkeypatch): + component_load_device = torch.device("cpu") + text_encoder = _TrackingModule() + vae = _TrackingModule() + transformer = _TrackingModule() + tokenizer = object() + scheduler = object() + loaded_components = iter([text_encoder, vae]) + + monkeypatch.setattr(torch, "get_default_device", lambda: component_load_device) + monkeypatch.setattr(pipeline_module, "prefetch_subfolders", lambda *args, **kwargs: None) + monkeypatch.setattr(pipeline_module, "_load_sana_tokenizer", lambda *args, **kwargs: tokenizer) + monkeypatch.setattr( + pipeline_module, + "from_pretrained_with_prefetch", + lambda *args, **kwargs: next(loaded_components), + ) + monkeypatch.setattr( + pipeline_module, + "_load_json", + lambda _model, filename, _local: ({"vae": [None, "FakeVAE"]} if filename == "model_index.json" else {}), + ) + monkeypatch.setattr( + pipeline_module, + "_resolve_vae_class_and_dtype", + lambda *_args: (SimpleNamespace(from_pretrained=None), torch.float32), + ) + monkeypatch.setattr( + pipeline_module.SanaVideoTransformer3DModel, + "from_config", + classmethod(lambda _cls, _config: transformer), + ) + monkeypatch.setattr( + pipeline_module.DPMSolverMultistepScheduler, + "from_pretrained", + lambda *args, **kwargs: scheduler, + ) + + pipeline = object.__new__(SanaVideoPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cuda:7") + pipeline.weights_sources = [] + config = OmniDiffusionConfig(model="local-sana-video", dtype=torch.bfloat16) + + loaded = pipeline._load_components(config, prefix="") + + assert loaded == (tokenizer, text_encoder, vae, transformer, scheduler) + assert pipeline.device == torch.device("cuda:7") + assert text_encoder.to_calls == [((component_load_device,), {})] + assert vae.to_calls == [((component_load_device,), {})] + assert transformer.to_calls == [ + ((), {"dtype": torch.bfloat16, "device": component_load_device}), + ] + + +def _patch_layerwise_platform(monkeypatch) -> None: + monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream) + monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent) + monkeypatch.setattr( + layerwise_backend_module.current_omni_platform, + "current_stream", + lambda: DummyStream(), + ) + monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream) + + +def _tiny_pipeline(monkeypatch) -> SanaVideoPipeline: + pipeline = object.__new__(SanaVideoPipeline) + nn.Module.__init__(pipeline) + pipeline.transformer = _tiny_transformer(monkeypatch) + pipeline.text_encoder = nn.Linear(2, 2) + pipeline.vae = nn.Linear(2, 2) + return pipeline + + +def _enable_layerwise_offload(pipeline: SanaVideoPipeline) -> LayerWiseOffloadBackend: + backend = LayerWiseOffloadBackend( + OffloadConfig(strategy=OffloadStrategy.LAYER_WISE, pin_cpu_memory=False), + device=torch.device("cpu"), + ) + backend.enable(pipeline) + return backend + + +def test_two_layer_offload_recovers_from_cached_block_skip_and_cleans_up(monkeypatch): + _patch_layerwise_platform(monkeypatch) + pipeline = _tiny_pipeline(monkeypatch) + discovered = ModuleDiscovery.discover(pipeline) + assert discovered.dits == [pipeline.transformer] + assert discovered.encoders == [pipeline.text_encoder] + assert discovered.vaes == [pipeline.vae] + + backend = _enable_layerwise_offload(pipeline) + blocks = list(pipeline.transformer.transformer_blocks) + assert backend.is_enabled() + assert len(blocks) == 2 + + block_1_hook = blocks[1]._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME) + assert block_1_hook is not None + assert block_1_hook.is_materialized is False + + # Cache-DiT may skip block 0, leaving block 1 without the normal prefetch. + block_1_hook.pre_forward(blocks[1]) + assert block_1_hook.is_materialized is True + block_1_hook.post_forward(blocks[1], None) + assert block_1_hook.is_materialized is False + + backend.disable() + assert backend.is_enabled() is False + for block in blocks: + assert block._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME) is None + + +def test_layerwise_offload_then_cache_dit_install_and_cleanup(monkeypatch): + _patch_layerwise_platform(monkeypatch) + pipeline = _tiny_pipeline(monkeypatch) + offload_backend = _enable_layerwise_offload(pipeline) + enabled_adapters = _record_cache_adapters(monkeypatch) + cache_backend = CacheDiTBackend() + blocks = list(pipeline.transformer.transformer_blocks) + + try: + cache_backend.enable(pipeline) + + assert offload_backend.is_enabled() + assert cache_backend.is_enabled() + assert pipeline.transformer._is_cached is True + for block in blocks: + assert block._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME) is not None + finally: + if enabled_adapters: + cache_dit.disable_cache(enabled_adapters[0]) + offload_backend.disable() + + assert not getattr(pipeline.transformer, "_is_cached", False) + assert offload_backend.is_enabled() is False + for block in blocks: + assert block._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME) is None diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 2626b1258..fcb823c92 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -14,6 +14,7 @@ from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request +from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.diffusion] @@ -133,12 +134,7 @@ class _CompileTrackingModel: def _make_request(): - sampling_params = SimpleNamespace( - generator=None, - seed=None, - generator_device=None, - num_inference_steps=4, - ) + sampling_params = OmniDiffusionSamplingParams(num_inference_steps=4) return SimpleNamespace( request_id="req-test", prompt="a prompt", @@ -249,6 +245,58 @@ def _make_compile_runner( return runner +class _EnabledCacheBackend: + def __init__(self): + self.refresh_calls = [] + + def is_enabled(self): + return True + + def refresh(self, pipeline, num_inference_steps, verbose=True): + self.refresh_calls.append((pipeline, num_inference_steps, verbose)) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_refresh_cache_prefers_request_steps_then_schedule_then_pipeline_default(): + cache_backend = _EnabledCacheBackend() + runner = _make_runner(cache_backend=cache_backend, cache_backend_name="cache_dit") + runner.pipeline.default_num_inference_steps = 50 + req = _make_request() + + custom_timesteps = torch.linspace(999.0, 0.0, 8) + custom_sigmas = [0.9, 0.5, 0.1] + cases = [ + (20, None, None), + (30, custom_timesteps, custom_sigmas), + (None, custom_timesteps, None), + (None, custom_timesteps, custom_sigmas), + (None, None, custom_sigmas), + (None, None, None), + ] + for num_inference_steps, timesteps, sigmas in cases: + req.sampling_params.num_inference_steps = num_inference_steps + req.sampling_params.timesteps = timesteps + req.sampling_params.sigmas = sigmas + DiffusionModelRunner._refresh_cache_for_requests(runner, [req], od_config=runner.od_config) + + assert [steps for _, steps, _ in cache_backend.refresh_calls] == [20, 30, 8, 8, 3, 50] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_refresh_cache_without_request_or_pipeline_default_warns(caplog): + cache_backend = _EnabledCacheBackend() + runner = _make_runner(cache_backend=cache_backend, cache_backend_name="cache_dit") + req = _make_request() + req.sampling_params.num_inference_steps = None + + DiffusionModelRunner._refresh_cache_for_requests(runner, [req], od_config=runner.od_config) + + assert cache_backend.refresh_calls == [] + assert "requires num_inference_steps to be passed explicitly" in caplog.text + + @pytest.mark.core_model @pytest.mark.cpu def test_update_states_carries_prepared_layout() -> None: @@ -451,16 +499,6 @@ def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypa @pytest.mark.core_model @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeypatch): - class _EnabledCacheBackend: - def __init__(self): - self.refresh_calls = [] - - def is_enabled(self): - return True - - def refresh(self, pipeline, num_inference_steps, verbose=True): - self.refresh_calls.append((pipeline, num_inference_steps, verbose)) - cache_backend = _EnabledCacheBackend() runner = _make_runner(cache_backend=cache_backend, cache_backend_name="cache_dit") req = _make_request() diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index 094ef7d84..4ef3f6c33 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/tests/model_tests/diffusion/diff_model_builders.py @@ -88,7 +88,7 @@ def _build_tiny_sana_video(pipeline_name: str) -> str: out_channels=16, num_attention_heads=2, attention_head_dim=16, - num_layers=1, + num_layers=2, num_cross_attention_heads=2, cross_attention_head_dim=16, cross_attention_dim=32, diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index 5cba1e519..122138dc2 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/tests/model_tests/diffusion/model_settings.py @@ -57,6 +57,11 @@ DIFFUSION_TEST_SETTINGS = { [DiffusionAccs.TENSOR_PARALLEL], [DiffusionAccs.CFG_PARALLEL], [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL], + [DiffusionAccs.CACHE_DIT], + [DiffusionAccs.CPU_OFFLOAD], + [DiffusionAccs.LAYERWISE_OFFLOAD], + [DiffusionAccs.CACHE_DIT, DiffusionAccs.CPU_OFFLOAD], + [DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD], ], check_multi_output=False, ), diff --git a/vllm_omni/diffusion/cache/cachedit/backend.py b/vllm_omni/diffusion/cache/cachedit/backend.py index bfc496bdc..9442ac46e 100644 --- a/vllm_omni/diffusion/cache/cachedit/backend.py +++ b/vllm_omni/diffusion/cache/cachedit/backend.py @@ -212,13 +212,20 @@ def _maybe_build_block_adapter( if missing_attributes: raise AttributeError(f"Missing Cache-DiT block attributes: {missing_attributes}") - return BlockAdapter( + block_adapter = BlockAdapter( transformer=transformer, blocks=[getattr(transformer, block_attribute) for block_attribute in block_attributes], forward_pattern=list(forward_patterns), has_separate_cfg=adapter_config.has_separate_cfg, check_forward_pattern=adapter_config.check_forward_pattern, ) + # cache_dit marks the wrapper pipe's class as cached, and for bare-transformer + # adapters that class is shared process-wide; a marker left by a previous + # engine would make this enable skip cache-context installation. + pipe_cls = type(block_adapter.pipe) + if "_is_cached" in vars(pipe_cls): + del pipe_cls._is_cached + return block_adapter def _maybe_get_cached_adapter_cls( diff --git a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py index 32a28ea60..48db6254f 100644 --- a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py +++ b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py @@ -272,6 +272,39 @@ def retrieve_timesteps( return timesteps, num_inference_steps +def _validate_cache_offload_parallelism(od_config: OmniDiffusionConfig) -> None: + cache_backend = od_config.cache_backend + if cache_backend not in ("none", "cache_dit"): + raise NotImplementedError( + f"Cache backend {cache_backend!r} is not supported by the native SANA-Video pipeline; " + "use 'cache_dit' or 'none'." + ) + if cache_backend == "cache_dit" and od_config.enable_distributed_layerwise_offload: + raise NotImplementedError( + "SANA-Video does not support Cache-DiT with distributed layerwise offload; " + "per-rank cache skips desynchronize the weight AllGather." + ) + + offload_enabled = ( + od_config.enable_cpu_offload + or od_config.enable_layerwise_offload + or od_config.enable_distributed_layerwise_offload + ) + if cache_backend == "none" and not offload_enabled: + return + + parallel_config = od_config.parallel_config + for name, size in ( + ("tensor_parallel_size", parallel_config.tensor_parallel_size), + ("cfg_parallel_size", parallel_config.cfg_parallel_size), + ("sequence_parallel_size", parallel_config.sequence_parallel_size), + ): + if size > 1: + raise NotImplementedError( + f"SANA-Video cache/offload is currently supported only with TP1, CFG1, and SP1; got {name}={size}." + ) + + class SanaVideoPipeline( nn.Module, CFGParallelMixin, @@ -300,6 +333,7 @@ class SanaVideoPipeline( _encoder_modules = ["text_encoder"] _vae_modules = ["vae"] supports_step_execution = False + default_num_inference_steps = 50 def __init__( self, @@ -321,6 +355,7 @@ class SanaVideoPipeline( self.device = get_local_device() self.weights_sources: list[DiffusersPipelineLoader.ComponentSource] = [] if od_config is not None: + _validate_cache_offload_parallelism(od_config) tokenizer, text_encoder, vae, transformer, scheduler = self._load_components(od_config, prefix) if tokenizer is None or text_encoder is None or vae is None or transformer is None or scheduler is None: @@ -368,7 +403,9 @@ class SanaVideoPipeline( model = od_config.model local_files_only = os.path.exists(model) dtype = getattr(od_config, "dtype", torch.bfloat16) - device = self.device + # The loader loads components under a default-device context that is CPU + # when offload is enabled; runtime placement stays with the offloader. + component_load_device = torch.get_default_device() # Transformer weights are streamed by DiffusersPipelineLoader below. # Prefetch only components loaded through from_pretrained here. component_subfolders = ["tokenizer", "text_encoder", "vae", "scheduler"] @@ -386,7 +423,7 @@ class SanaVideoPipeline( prefetch_list=component_subfolders, local_files_only=local_files_only, torch_dtype=dtype, - ).to(device) + ).to(component_load_device) model_index = _load_json(model, "model_index.json", local_files_only) vae_class_name = model_index.get("vae", [None, "AutoencoderKLWan"])[1] @@ -401,10 +438,13 @@ class SanaVideoPipeline( prefetch_list=component_subfolders, local_files_only=local_files_only, torch_dtype=vae_dtype, - ).to(device) + ).to(component_load_device) transformer_config = _load_json(model, "transformer/config.json", local_files_only) - transformer = SanaVideoTransformer3DModel.from_config(transformer_config).to(dtype=dtype, device=device) + transformer = SanaVideoTransformer3DModel.from_config(transformer_config).to( + dtype=dtype, + device=component_load_device, + ) scheduler = DPMSolverMultistepScheduler.from_pretrained( model, subfolder="scheduler", @@ -859,7 +899,11 @@ class SanaVideoPipeline( default_num_frames=81, is_dummy_run=req.is_dummy_run(), ) - num_steps = sampling.num_inference_steps if sampling.num_inference_steps is not None else 50 + num_steps = ( + sampling.num_inference_steps + if sampling.num_inference_steps is not None + else self.default_num_inference_steps + ) guidance_scale = sampling.guidance_scale if sampling.guidance_scale_provided else 6.0 generator = sampling.generator if generator is None and sampling.seed is not None: diff --git a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py index 9475f306b..6835a7320 100644 --- a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py +++ b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py @@ -196,7 +196,11 @@ class SanaImageToVideoPipeline(SanaVideoPipeline, SupportImageInput): default_num_frames=81, is_dummy_run=req.is_dummy_run(), ) - num_steps = sampling.num_inference_steps if sampling.num_inference_steps is not None else 50 + num_steps = ( + sampling.num_inference_steps + if sampling.num_inference_steps is not None + else self.default_num_inference_steps + ) guidance_scale = sampling.guidance_scale if sampling.guidance_scale_provided else 6.0 generator = sampling.generator if generator is None and sampling.seed is not None: diff --git a/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py b/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py index db38cc099..1b64790db 100644 --- a/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py +++ b/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py @@ -18,6 +18,7 @@ from collections.abc import Iterable from dataclasses import dataclass, fields import torch +from cache_dit import ForwardPattern from torch import nn from vllm.distributed import ( get_tensor_model_parallel_rank, @@ -30,6 +31,7 @@ from vllm.model_executor.utils import set_weight_attrs from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention as OmniAttention +from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig def validate_sana_video_parallel_config(parallel_config) -> None: @@ -962,6 +964,12 @@ class SanaVideoTransformer3DModel(nn.Module): """ _no_split_modules = ["SanaVideoTransformerBlock", "SanaModulatedNorm"] + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "transformer_blocks": ForwardPattern.Pattern_3, + }, + ) + _layerwise_offload_blocks_attrs = ["transformer_blocks"] def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return AutoWeightsLoader(self).load_weights(weights) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 61f45c76e..de415297e 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -631,6 +631,12 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # RequestBatchSamplingParamsKey, so the first request's num_inference_steps applies # to the whole runner batch. num_inference_steps = first_req.sampling_params.num_inference_steps + if num_inference_steps is None and first_req.sampling_params.timesteps is not None: + num_inference_steps = len(first_req.sampling_params.timesteps) + if num_inference_steps is None and first_req.sampling_params.sigmas is not None: + num_inference_steps = len(first_req.sampling_params.sigmas) + if num_inference_steps is None: + 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", "step_cache",