Fix decode cuda graph for Qwen-Image-2.1 model-level cpu offload (#7801)

tests/diffusion/models/qwen_image_21/test_decode_graph.py

@@ -6,7 +6,9 @@ import weakref
 import pytest
 import torch
 
+from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
 from vllm_omni.diffusion.models.qwen_image_21 import decode_graph
+from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
@@ -71,6 +73,41 @@ def decode(manager, cache, layout=(1, 2, 2)):
     )
 
 
+@pytest.mark.cpu
+def test_model_level_offload_without_staging_disables_decode_graphs():
+    manager = decode_graph.QwenImage21DecodeGraphManager(DecodeModel(), model_level_offload=True)
+    assert manager._offload_reason() is not None
+    assert manager.eligible() is False
+
+
+@pytest.mark.cpu
+def test_top_level_offload_hook_without_staging_disables_decode_graphs():
+    model = DecodeModel()
+    manager = decode_graph.QwenImage21DecodeGraphManager(model)
+    registry = HookRegistry.get_or_create(model)
+    # A present-but-empty registry (e.g. after hooks were removed) must not disqualify.
+    assert manager._offload_reason() is None
+    registry.register_hook("sequential_offload", ModelHook())
+    assert manager._offload_reason() is not None
+    assert manager.eligible() is False
+
+
+@pytest.mark.cpu
+def test_persistent_staging_makes_model_level_offload_graph_eligible():
+    model = DecodeModel()
+    manager = decode_graph.QwenImage21DecodeGraphManager(model, model_level_offload=True)
+    hook = SequentialOffloadHook(offload_targets=[], device=torch.device("cpu"))
+    HookRegistry.get_or_create(model).register_hook(SequentialOffloadHook._HOOK_NAME, hook)
+    # No staging storage yet: a swap would allocate fresh storage per generation.
+    assert manager._offload_reason() is not None
+    # Staging established: fixed device storage keeps captured pointers valid.
+    hook._stager = object()
+    assert manager._offload_reason() is None
+    # A second, non-offload hook on the top-level module still disqualifies.
+    model._hook_registry.register_hook("teacache", ModelHook())
+    assert manager._offload_reason() is not None
+
+
 @pytest.mark.cpu
 def test_allocation_oom_releases_partial_buffers_and_remembers_failure(graph_case, monkeypatch):
     manager, device, cleanup = graph_case

tests/diffusion/offloader/test_sequential_backend.py

@@ -284,6 +284,49 @@ def test_component_selective_model_offload_requires_swap_counterpart(component,
         backend.enable(pipeline)
 
 
+def test_persistent_dit_staging_reuses_fixed_device_storage(accelerator_device) -> None:
+    dit = _create_simple_module().to(accelerator_device)
+    encoder = _create_simple_module().to(accelerator_device)
+    original_weight = dit.linear.weight.detach().clone()
+
+    apply_sequential_offload(
+        dit_modules=[dit],
+        encoder_modules=[encoder],
+        device=accelerator_device,
+        pin_memory=False,
+        persistent_dit_staging=True,
+    )
+
+    dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    encoder_hook = encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert dit_hook._stager is not None
+    assert encoder_hook._stager is None  # encoders keep plain move semantics
+
+    # Registration snapshots the CPU master and rebinds to it.
+    assert dit.linear.weight.device.type == "cpu"
+
+    dit_hook._to_gpu(dit)
+    first_ptr = dit.linear.weight.data_ptr()
+    assert dit.linear.weight.device == accelerator_device
+    torch.testing.assert_close(dit.linear.weight.data, original_weight)
+
+    # Offload via a *different* module's hook (the cross-module swap path),
+    # then reload: the same fixed device storage must be reused so that
+    # CUDA-graph-captured weight pointers stay valid.
+    encoder_hook._to_cpu(dit)
+    assert dit.linear.weight.device.type == "cpu"
+    dit_hook._to_gpu(dit)
+    assert dit.linear.weight.data_ptr() == first_ptr
+    torch.testing.assert_close(dit.linear.weight.data, original_weight)
+
+    # A second offload/load cycle through the DiT hook itself stays put too.
+    dit_hook._to_cpu(dit)
+    dit_hook._to_gpu(dit)
+    assert dit.linear.weight.data_ptr() == first_ptr
+
+    remove_sequential_offload([dit, encoder])
+
+
 def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
     dit = _create_simple_module()
     encoder = _create_simple_module()