[Model] Add Cache-DiT and CPU offload support for SANA-Video 2B (#5882)

Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>

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

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

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()

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,

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,
     ),

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(

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:

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:

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)

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",