[diffusion][feature] Add component-selective offload policies (#5929)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -30,6 +30,36 @@ def test_pipeline_declares_layerwise_offload_components() -> None:
     assert Cosmos3OmniDiffusersPipeline._resident_modules == []
     assert hasattr(Cosmos3OmniDiffusersPipeline, "enable_omni_model_cpu_offload")
 
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        Cosmos3LanguageModel,
+        Cosmos3VFMTransformer,
+    )
+
+    assert Cosmos3LanguageModel._layerwise_offload_blocks_attrs == ["layers"]
+    assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"]
+
+
+def test_component_selective_model_offload_fails_before_component_loading(monkeypatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 as pipeline_module
+
+    monkeypatch.setattr(
+        pipeline_module,
+        "get_local_device",
+        lambda: pytest.fail("component loading must not start for an unsupported selector"),
+    )
+    config = SimpleNamespace(
+        diffusion_offload_config={"mode": "module", "components": ["dit"]},
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+    )
+
+    with pytest.raises(ValueError, match="does not support the dit/text_encoder component selector"):
+        pipeline_module.Cosmos3OmniDiffusersPipeline(od_config=config)
+
 
 class StubScheduler:
     def __init__(

tests/diffusion/models/magi2/test_native_dlo_allgather.py

@@ -220,7 +220,9 @@ def _load_mmap_transform_and_reconstruct(
         tensor_transforms=backend._mmap_transforms_by_tensor_id,
     )
     assert all(not shard.requires_grad for shard in cpu_shards.values())
-    assert all(parameter.numel() == 0 for parameter in target_block.parameters())
+    # Shard preparation is non-destructive; initialize_hook clears source
+    # storage only after every shard and metadata entry has been prepared.
+    assert all(parameter.numel() > 0 for parameter in target_block.parameters())
 
     for dtype, local_shard in cpu_shards.items():
         gathered = torch.empty(

tests/diffusion/models/magi2/test_pipeline_magi2.py

@@ -578,6 +578,8 @@ class _FakeLayerwiseHook:
     def __init__(self) -> None:
         self._prev_hook = None
         self.current_slot = 0
+        self.dp_size = 1
+        self.rank_local_mmap = False
         self.prefetch_layer = Mock()
         self.get_weights = Mock()
 
@@ -625,7 +627,7 @@ def test_dlo_no_allgather_enable_preserves_staged_aux_and_streams_blocks(
     assert len(hooks) == len(pipeline.transformer.block)
     assert all(stager.offload.call_count == 1 for stager in stagers)
     assert all(stager.load.call_count == 0 for stager in stagers)
-    hooks[-1].prefetch_layer.assert_called_once_with(
-        slot=hooks[0].current_slot,
-        non_blocking=False,
-    )
+    # First prefetch stays lazy so ranks with different local groups cannot
+    # enter collectives in a different order during startup.
+    for hook in hooks:
+        hook.prefetch_layer.assert_not_called()

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -487,6 +487,10 @@ def test_dlo_offload_plan_includes_token_refiner():
 
     assert MiniMaxH3Pipeline._offload_plan.offload_submodules == {"token_refiner": "blocks"}
     assert MiniMaxH3Pipeline._offload_plan.resident_dit_paths == frozenset({"transformer"})
+    assert MiniMaxH3Pipeline._offload_plan.encoder_component_types == {"text_encoder": "text_encoder"}
+    assert MiniMaxH3Pipeline._offload_plan.encoder_block_attrs == {
+        "text_encoder": ("vision.blocks", "text_model.layers")
+    }
 
 
 def test_joint_postprocess_is_multiprocessing_picklable():
@@ -1185,6 +1189,12 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
     from vllm_omni.diffusion.models.minimax_h3.encoder import (
         MiniMaxH3Qwen3VLEncoder,
     )
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        MiniMaxH3Pipeline,
+    )
+    from vllm_omni.diffusion.offloader.component_utils import (
+        get_encoder_block_groups,
+    )
 
     encoder = MiniMaxH3Qwen3VLEncoder(
         "/nonexistent/text_encoder",
@@ -1197,6 +1207,25 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
     # The stub has no parameters, so it never contributes to the runner's
     # strict missing-parameter check on non-encoder ranks.
     assert list(encoder.named_parameters()) == []
+    assert (
+        get_encoder_block_groups(
+            encoder,
+            "text_encoder",
+            MiniMaxH3Pipeline._offload_plan,
+            strict=True,
+        )
+        == []
+    )
+
+    # An arbitrary empty module is not silently treated as a distributed
+    # stub; explicit plans still validate their declared paths.
+    with pytest.raises(ValueError, match=r"text_encoder\.vision\.blocks was not found"):
+        get_encoder_block_groups(
+            nn.Module(),
+            "text_encoder",
+            MiniMaxH3Pipeline._offload_plan,
+            strict=True,
+        )
 
 
 def test_global_quant_config_is_shared_by_dit_and_encoder():
@@ -1303,10 +1332,38 @@ def test_model_offload_uses_hooked_text_encoder_call():
 
 
 @pytest.mark.parametrize(
-    "offload_flag",
-    ["enable_layerwise_offload", "enable_distributed_layerwise_offload"],
+    ("enable_layerwise", "enable_distributed"),
+    [(True, False), (False, True)],
 )
-def test_layerwise_offload_releases_text_encoder(offload_flag):
+def test_legacy_layer_offload_preserves_minimax_stage_lifecycle(enable_layerwise, enable_distributed):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        diffusion_offload_config=None,
+        enable_cpu_offload=False,
+        enable_layerwise_offload=enable_layerwise,
+        enable_distributed_layerwise_offload=enable_distributed,
+    )
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+
+    vae = Mock()
+    with pipeline._component_on_device(vae):
+        pass
+    vae.load_to_device.assert_called_once_with()
+    vae.offload_to_cpu.assert_called_once_with()
+
+
+def test_layerwise_encoder_selection_releases_text_encoder():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
@@ -1315,8 +1372,11 @@ def test_layerwise_offload_releases_text_encoder(offload_flag):
         enable_cpu_offload=False,
         enable_layerwise_offload=False,
         enable_distributed_layerwise_offload=False,
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": ["text_encoder"],
+        },
     )
-    setattr(pipeline.od_config, offload_flag, True)
     pipeline.text_encoder = Mock()
     expected = torch.ones(2, 3)
     pipeline.text_encoder.encode_ids.return_value = expected
@@ -1328,7 +1388,7 @@ def test_layerwise_offload_releases_text_encoder(offload_flag):
     pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
 
 
-def test_distributed_layerwise_offload_releases_text_encoder():
+def test_layerwise_dit_only_keeps_text_encoder_resident():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
@@ -1336,7 +1396,8 @@ def test_distributed_layerwise_offload_releases_text_encoder():
     pipeline.od_config = SimpleNamespace(
         enable_cpu_offload=False,
         enable_layerwise_offload=False,
-        enable_distributed_layerwise_offload=True,
+        enable_distributed_layerwise_offload=False,
+        diffusion_offload_config={"mode": "layer", "components": ["dit"]},
     )
     pipeline.text_encoder = Mock()
     expected = torch.ones(2, 3)
@@ -1346,25 +1407,30 @@ def test_distributed_layerwise_offload_releases_text_encoder():
 
     assert actual is expected
     pipeline.text_encoder.load_to_device.assert_called_once_with()
-    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_not_called()
 
 
-def test_distributed_layerwise_offload_stages_vae_component():
+def test_dit_encoder_selection_keeps_vae_resident():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
         enable_layerwise_offload=False,
-        enable_distributed_layerwise_offload=True,
+        enable_distributed_layerwise_offload=False,
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": ["dit", "text_encoder"],
+        },
     )
     component = Mock()
 
     with pipeline._component_on_device(component):
-        component.load_to_device.assert_called_once_with()
-        component.offload_to_cpu.assert_not_called()
+        pass
 
-    component.offload_to_cpu.assert_called_once_with()
+    component.load_to_device.assert_not_called()
+    component.offload_to_cpu.assert_not_called()
 
 
 def test_distributed_layerwise_resident_blocks_are_stage_scoped():
@@ -1413,57 +1479,6 @@ def test_distributed_layerwise_resident_blocks_can_be_skipped():
     controller.offload_resident_layers.assert_not_called()
 
 
-def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch):
-    from vllm_omni.diffusion.models.minimax_h3.encoder import (
-        MiniMaxH3Qwen3VLEncoder,
-    )
-
-    class Stack(torch.nn.Module):
-        def __init__(self, count):
-            super().__init__()
-            self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
-
-    class TextStack(torch.nn.Module):
-        def __init__(self, count):
-            super().__init__()
-            self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
-
-    hooks = []
-
-    def fake_apply(block, next_block, device, stream, pin_memory):
-        hook = SimpleNamespace(
-            block=block,
-            next_block=next_block,
-            device=device,
-            stream=stream,
-            pin_memory=pin_memory,
-            _prev_hook=None,
-            offload_layer=Mock(),
-        )
-        hooks.append(hook)
-        return hook
-
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.offloader.layerwise_backend.apply_block_hook",
-        fake_apply,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.platforms.current_omni_platform.Stream",
-        Mock(return_value="copy-stream"),
-    )
-    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
-    torch.nn.Module.__init__(encoder)
-    encoder.device_target = torch.device("cpu")
-    encoder.vision = Stack(2)
-    encoder.text_model = TextStack(3)
-
-    encoder.enable_omni_layerwise_offload(pin_memory=False)
-
-    assert len(hooks) == 5
-    assert all(hook._prev_hook is not None for hook in hooks)
-    assert all(hook.device == torch.device("cpu") for hook in hooks)
-
-
 def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
 

tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -11,6 +13,44 @@ class _OffloadAbort(BaseException):
     pass
 
 
+def test_no_allgather_example_selects_distributed_backend():
+    from examples.offline_inference.minimax_h3.dlo_lifecycle import engine_kwargs
+    from vllm_omni.diffusion.offloader.base import OffloadConfig
+    from vllm_omni.diffusion.offloader.config import OffloadStrategy
+
+    args = SimpleNamespace(
+        model="/fake/MiniMax-H3/FL2VA",
+        mode="dlo-no-allgather",
+        dp_size=2,
+        tp_size=2,
+        batch_wait_ms=500.0,
+        init_timeout=1800.0,
+    )
+    kwargs = engine_kwargs(args)
+    od_config = SimpleNamespace(
+        diffusion_offload_config=None,
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=kwargs["enable_distributed_layerwise_offload"],
+        dlo_use_allgather=kwargs["dlo_use_allgather"],
+        dlo_resident_layers=kwargs["dlo_resident_layers"],
+        dlo_host_registration_limit_gib=0.0,
+        host_weight_runtime_mode="disabled",
+        pin_cpu_memory=True,
+        parallel_config=SimpleNamespace(
+            data_parallel_size=kwargs["data_parallel_size"],
+            sequence_parallel_size=1,
+            use_hsdp=False,
+        ),
+    )
+
+    config = OffloadConfig.from_od_config(od_config)
+
+    assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+    assert config.dlo_resident_layers == 0
+    assert not config.uses_allgather("dit")
+
+
 def test_encoder_non_block_children_use_one_shared_snapshot_stager(monkeypatch, mocker):
     from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
 
@@ -61,6 +101,7 @@ def test_manual_component_failure_forces_retained_cache_release(mocker):
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.od_config = mocker.Mock()
+    pipeline.od_config.diffusion_offload_config = None
     pipeline.od_config.enable_layerwise_offload = False
     pipeline.od_config.enable_distributed_layerwise_offload = True
     pipeline._model_cpu_offload_modules = []
@@ -83,6 +124,7 @@ def test_manual_component_offload_failure_forces_retained_cache_release(mocker):
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.od_config = mocker.Mock()
+    pipeline.od_config.diffusion_offload_config = None
     pipeline.od_config.enable_layerwise_offload = False
     pipeline.od_config.enable_distributed_layerwise_offload = True
     pipeline._model_cpu_offload_modules = []

tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py

@@ -247,6 +247,55 @@ def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
     assert captured["unsupported_offload_mode"] is None
 
 
+@pytest.mark.parametrize(
+    ("mode", "components", "expected_mode"),
+    [
+        ("module", ["text_encoder"], None),
+        ("layer", ["text_encoder"], None),
+        ("module", ["dit"], "model-level CPU offload"),
+        ("layer", ["dit"], "layerwise offload"),
+    ],
+)
+def test_h3_turbo_offload_compatibility_follows_dit_selection(
+    monkeypatch,
+    mode,
+    components,
+    expected_mode,
+):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "fl2va"
+    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._native_lora_adapter_ids = set()
+    pipeline._lora_sigma_schedules = {}
+    pipeline.od_config = SimpleNamespace(
+        diffusion_offload_config={"mode": mode, "components": components},
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+    )
+    captured = {}
+
+    def load_turbo(**kwargs):
+        captured.update(kwargs)
+        return object(), object()
+
+    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
+
+    assert pipeline._load_diffusion_lora_adapter(
+        lora_request=_request("turbo"),
+        lora_path="turbo",
+        dtype=torch.bfloat16,
+    )
+    assert captured["unsupported_offload_mode"] == expected_mode
+
+
 def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
     unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
     _write_tiny_turbo(unsupported)

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from contextlib import contextmanager
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import Mock
 
 import pytest
@@ -111,6 +112,66 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
     remove_offload.assert_called_once_with([*dits, *stages])
 
 
+def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.transformer = torch.nn.Linear(2, 2)
+    pipeline.transformers_ref = torch.nn.Linear(2, 2)
+    pipeline.text_encoder = torch.nn.Linear(2, 2)
+    pipeline.video_vae = torch.nn.Linear(2, 2)
+    pipeline.audio_vae = torch.nn.Linear(2, 2)
+    apply_offload = Mock()
+    monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
+
+    pipeline.enable_omni_model_cpu_offload(
+        device=torch.device("cpu"),
+        pin_memory=False,
+        use_hsdp=False,
+        offload_components=frozenset({"dit", "text_encoder"}),
+    )
+
+    dits = [pipeline.transformer, pipeline.transformers_ref]
+    stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae]
+    apply_offload.assert_called_once_with(
+        dit_modules=dits,
+        encoder_modules=stages,
+        device=torch.device("cpu"),
+        pin_memory=False,
+        use_hsdp=False,
+        offload_initial_dits=True,
+        offload_dit_modules=dits,
+        offload_encoder_modules=[pipeline.text_encoder],
+    )
+
+
+def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.transformer = torch.nn.Linear(2, 2)
+    pipeline.transformers_ref = torch.nn.Linear(2, 2)
+    pipeline.text_encoder = None
+    pipeline.video_vae = torch.nn.Linear(2, 2)
+    pipeline.audio_vae = torch.nn.Linear(2, 2)
+    apply_offload = Mock()
+    monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
+
+    with pytest.raises(ValueError, match="no loaded text encoder"):
+        pipeline.enable_omni_model_cpu_offload(
+            device=torch.device("cpu"),
+            pin_memory=False,
+            use_hsdp=False,
+            offload_components=frozenset({"text_encoder"}),
+        )
+
+    apply_offload.assert_not_called()
+
+
 @pytest.mark.parametrize("decode_fails", [False, True])
 def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -119,7 +180,7 @@ def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     component = torch.nn.Linear(2, 2)
-    events = []
+    events: list[tuple[Any, ...]] = []
 
     @contextmanager
     def record_component(value):
@@ -262,7 +323,7 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
     torch.nn.Module.__init__(pipeline)
     pipeline.audio_vae = Mock()
     pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
-    events = []
+    events: list[tuple[Any, ...]] = []
 
     @contextmanager
     def record_component(value):

tests/diffusion/models/sana_video/test_cache_offload.py

@@ -16,7 +16,7 @@ from tests.diffusion.models.sana_video.test_transformer_sana_video import (  # n
     _TINY_CONFIG,
     _init_distributed,
 )
-from tests.diffusion.offloader.test_layerwise_backend import DummyEvent, DummyStream, dummy_stream
+from tests.diffusion.offloader.helpers 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
@@ -226,7 +226,7 @@ def test_component_loading_uses_loader_device_not_runtime_device(monkeypatch):
     monkeypatch.setattr(
         pipeline_module,
         "_load_json",
-        lambda _model, filename, _local: ({"vae": [None, "FakeVAE"]} if filename == "model_index.json" else {}),
+        lambda _model, filename, _local: {"vae": [None, "FakeVAE"]} if filename == "model_index.json" else {},
     )
     monkeypatch.setattr(
         pipeline_module,

tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py

@@ -173,9 +173,14 @@ def test_length_lives_in_device_tensors():
 # ---------------------------------------------------------------------------
 
 cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+vllm_flash_attn_only = pytest.mark.skipif(
+    current_omni_platform.is_rocm(),
+    reason="vLLM's bundled FlashAttention extension is CUDA-only",
+)
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_paged_attention_matches_sdpa_over_the_same_buffer():
@@ -217,6 +222,7 @@ def test_paged_attention_matches_sdpa_over_the_same_buffer():
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_a_captured_graph_follows_a_later_length():
@@ -278,6 +284,7 @@ def test_a_captured_graph_follows_a_later_length():
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_a_captured_graph_writes_each_step_to_its_own_slot():
@@ -339,6 +346,7 @@ def test_a_captured_graph_writes_each_step_to_its_own_slot():
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_a_whole_decode_loop_matches_the_unpaged_path():

tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py

@@ -9,6 +9,7 @@ from PIL import Image
 from torch import nn
 
 from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
     Wan22I2VPipeline,
     get_wan22_i2v_post_process_func,
@@ -30,6 +31,10 @@ def test_wan22_i2v_postprocess_honors_request_output_type() -> None:
     assert output is video
 
 
+def test_i2v_pipeline_declares_text_encoder_offload_blocks() -> None:
+    assert Wan22I2VPipeline._offload_plan is _WAN_TEXT_ENCODER_OFFLOAD_PLAN
+
+
 def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
     pipeline = object.__new__(Wan22I2VPipeline)
     nn.Module.__init__(pipeline)

tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py

@@ -18,6 +18,15 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_t2v_pipeline_declares_text_encoder_offload_blocks() -> None:
+    plan = wan22_module.Wan22Pipeline._offload_plan
+
+    assert plan.encoder_component_types == {"text_encoder": "text_encoder"}
+    assert plan.encoder_block_attrs == {"text_encoder": ("encoder.block",)}
+    assert plan.encoder_dlo_weight_replication == frozenset({"text_encoder"})
+    assert plan.on_demand_component_paths == frozenset()
+
+
 class _LatentDist:
     def sample(self, generator):
         assert isinstance(generator, torch.Generator)

tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py

@@ -392,52 +392,31 @@ def test_encode_audio_skips_unshard_reshard_when_not_fsdp():
     assert "audio_emb" in result
 
 
-def test_s2v_pipeline_skips_cpu_offload_when_hsdp_enabled():
-    """Test that transformer.to('cpu') is NOT called when HSDP is active."""
-    from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline
-
+@pytest.mark.parametrize(
+    ("components", "use_hsdp", "expected"),
+    [
+        (None, False, True),
+        (["dit"], False, True),
+        (["text_encoder"], False, False),
+        (["dit"], True, False),
+    ],
+)
+def test_s2v_dit_release_honors_component_selection(components, use_hsdp, expected):
     pipeline = object.__new__(Wan22S2VPipeline)
     nn.Module.__init__(pipeline)
+    compact = None if components is None else {"mode": "module", "components": components}
+    pipeline.od_config = SimpleNamespace(
+        diffusion_offload_config=compact,
+        enable_cpu_offload=True,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+        parallel_config=SimpleNamespace(use_hsdp=use_hsdp),
+    )
 
-    od_config = MagicMock()
-    od_config.enable_cpu_offload = True
-    parallel_config = MagicMock()
-    parallel_config.use_hsdp = True
-    od_config.parallel_config = parallel_config
-    pipeline.od_config = od_config
-
-    mock_transformer = MagicMock()
-    pipeline.transformer = mock_transformer
-
-    # Simulate the offload decision from the forward loop
-    if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False):
-        pipeline.transformer.to("cpu")
-
-    mock_transformer.to.assert_not_called()
-
-
-def test_s2v_pipeline_allows_cpu_offload_when_hsdp_disabled():
-    """Test that transformer.to('cpu') IS called when HSDP is not active."""
-    from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline
-
-    pipeline = object.__new__(Wan22S2VPipeline)
-    nn.Module.__init__(pipeline)
-
-    od_config = MagicMock()
-    od_config.enable_cpu_offload = True
-    parallel_config = MagicMock()
-    parallel_config.use_hsdp = False
-    od_config.parallel_config = parallel_config
-    pipeline.od_config = od_config
-
-    mock_transformer = MagicMock()
-    pipeline.transformer = mock_transformer
-
-    # Simulate the offload decision from the forward loop
-    if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False):
-        pipeline.transformer.to("cpu")
-
-    mock_transformer.to.assert_called_once_with("cpu")
+    assert pipeline._should_release_dit_before_decode() is expected
 
 
 def test_s2v_pipeline_hsdp_forward_complete_process():

tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass, field
 from types import SimpleNamespace
 
 import pytest
@@ -20,6 +21,46 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeypatch) -> None:
+    @dataclass
+    class RuntimeConfig:
+        flow_shift: float | None = None
+        diffusion_offload_config: dict[str, object] = field(
+            default_factory=lambda: {
+                "mode": "layer",
+                "components": ["text_encoder"],
+            }
+        )
+        enable_layerwise_offload: bool = True
+        post_init_calls: int = 0
+
+        def __post_init__(self) -> None:
+            self.post_init_calls += 1
+
+    config = RuntimeConfig()
+    resolved_offload = object()
+    config._resolved_diffusion_offload = resolved_offload
+    config._diffusion_offload_flags_materialized = True
+    captured = {}
+
+    def capture_base_init(self, *, od_config, prefix="") -> None:
+        captured.update(od_config=od_config, prefix=prefix)
+
+    monkeypatch.setattr(Wan22VACEPipeline.__mro__[1], "__init__", capture_base_init)
+
+    Wan22VACEPipeline(od_config=config, prefix="stage")
+
+    forwarded = captured["od_config"]
+    assert forwarded is not config
+    assert config.flow_shift is None
+    assert forwarded.flow_shift == 3.0
+    assert forwarded.diffusion_offload_config is config.diffusion_offload_config
+    assert forwarded._resolved_diffusion_offload is resolved_offload
+    assert forwarded._diffusion_offload_flags_materialized is True
+    assert forwarded.post_init_calls == 1
+    assert captured["prefix"] == "stage"
+
+
 def _make_vace_pipeline() -> Wan22VACEPipeline:
     pipeline = object.__new__(Wan22VACEPipeline)
     nn.Module.__init__(pipeline)