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

vllm_omni/diffusion/models/qwen_image_21/decode_graph.py

@@ -18,9 +18,10 @@ Memory note: graph entries own an additional copy of the prefix K/V. Decode
 attention concatenates that prefix with target K/V inside the captured region;
 transient tensors use the graph memory pool.
 
-Fallbacks (all logged, all eager): CUDA unavailable, sequence/ring/tensor
-parallelism, HSDP/offload hooks, torch.compile'd blocks, KV-cache quantization,
-dynamic LoRA wrappers, padded text masks (the masked attention path branches on
+Fallbacks (all logged, all eager): CUDA unavailable, model-level CPU offload
+without persistent weight staging, sequence/ring/tensor parallelism,
+HSDP/offload hooks, torch.compile'd blocks, KV-cache quantization, dynamic
+LoRA wrappers, padded text masks (the masked attention path branches on
 mask contents, which cannot be captured), unknown graph key at decode, and
 capture failure.
 """
@@ -154,22 +155,50 @@ class QwenImage21DecodeGraphEntry:
 class QwenImage21DecodeGraphManager:
     """Owns the static KV buffers and captured graphs of one transformer."""
 
-    def __init__(self, model: QwenImage21Transformer2DModel, max_entries: int = 8):
+    def __init__(self, model: QwenImage21Transformer2DModel, max_entries: int = 8, model_level_offload: bool = False):
         self.model = model
         self.max_entries = max_entries
+        self.model_level_offload = model_level_offload
         self.entries: OrderedDict[tuple, QwenImage21DecodeGraphEntry | None] = OrderedDict()
         self._static_eligible: bool | None = None
 
     # ── eligibility ──
 
+    def _offload_reason(self) -> str | None:
+        """Reason decode graphs are incompatible with model-level offload, if any.
+
+        Model-level (sequential) offload registers its hook on the top-level
+        transformer module, not on individual blocks. With persistent staging
+        the hook keeps the DiT on fixed device storage and only rebinds
+        ``p.data`` around each generation, and the swap completes in
+        ``pre_forward`` before capture or replay begin, so captured weight
+        pointers stay valid. Without staging every swap allocates fresh
+        storage and replay would read stale pointers, so capture stays
+        disabled and decode falls back to eager.
+        """
+        registry = getattr(self.model, "_hook_registry", None)
+        hooks = registry._hooks if registry is not None else {}
+        if hooks:
+            from vllm_omni.diffusion.offloader.sequential_backend import (
+                SequentialOffloadHook,
+                sequential_offload_staging_active,
+            )
+
+            if set(hooks) <= {SequentialOffloadHook._HOOK_NAME} and sequential_offload_staging_active(self.model):
+                return None
+            return "model module carries offload/cache hooks"
+        if self.model_level_offload:
+            return "model-level CPU offload without persistent staging swaps weight storage after capture"
+        return None
+
     def _check_static_eligibility(self) -> bool:
         """One-time checks that cannot change after weights are loaded."""
         model = self.model
-        reason = None
+        reason = self._offload_reason()
         param = next(model.parameters(), None)
-        if param is None or param.device.type != "cuda":
+        if reason is None and (param is None or param.device.type != "cuda"):
             reason = f"model is not on CUDA (device={None if param is None else param.device})"
-        else:
+        elif reason is None:
             parallel_config = getattr(model, "parallel_config", None)
             if parallel_config is not None:
                 sp = getattr(parallel_config, "sequence_parallel_size", None) or 1

vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py

@@ -899,9 +899,20 @@ class QwenImage21Transformer2DModel(CachedTransformer):
         self.sequence_prepare = QwenImage21SequencePrepare(self.img_in, self.pos_embed)
 
         # Use decode graphs unless the caller requests eager execution.
+        # Model-level (sequential) offload registers its swap hook on this
+        # top-level module rather than on individual blocks; graphs stay
+        # eligible only when that hook keeps weights on persistent staging
+        # storage (see QwenImage21DecodeGraphManager._offload_reason).
         self.enable_cuda_graph_decode = not od_config.enforce_eager
         self._decode_graph_manager = (
-            QwenImage21DecodeGraphManager(self, max_entries=cuda_graph_max_decode_graphs)
+            QwenImage21DecodeGraphManager(
+                self,
+                max_entries=cuda_graph_max_decode_graphs,
+                model_level_offload=(
+                    od_config.enable_cpu_offload
+                    or getattr(od_config, "enable_distributed_layerwise_offload", False)
+                ),
+            )
             if self.enable_cuda_graph_decode
             else None
         )

vllm_omni/diffusion/offloader/module_residency.py

@@ -98,6 +98,12 @@ class PinnedModuleStager:
     A module iterable is treated as one staging group. It uses one copy stream
     and one reusable completion event. Tensors sharing storage keep their
     shapes, strides, offsets, dtypes, and aliases across every transition.
+
+    With ``retain_device_storage`` the device storages survive ``offload``:
+    ``load`` then only rebinds to the same fixed storages (weights are
+    immutable, so no H2D copy repeats). The stable storage addresses keep
+    CUDA-graph-captured weight pointers valid across offload/load swaps, at
+    the cost of keeping the device storage resident while offloaded.
     """
 
     def __init__(
@@ -108,6 +114,7 @@ class PinnedModuleStager:
         pin_memory: bool = True,
         copy_stream: Any | None = None,
         cache_retention: BoundedAllocatorCache | None = None,
+        retain_device_storage: bool = False,
     ) -> None:
         modules = (module,) if isinstance(module, nn.Module) else tuple(module)
         if not modules or not all(isinstance(item, nn.Module) for item in modules):
@@ -117,6 +124,7 @@ class PinnedModuleStager:
         self.copy_stream = copy_stream if copy_stream is not None else current_omni_platform.Stream()
         self._ready_event = current_omni_platform.Event()
         self.cache_retention = cache_retention
+        self.retain_device_storage = retain_device_storage
         self.loaded = False
         self._groups = self._snapshot_groups(modules, pin_memory=pin_memory)
         self._device_storages: list[torch.Tensor] = []
@@ -234,6 +242,12 @@ class PinnedModuleStager:
     def load(self) -> None:
         if self.loaded:
             return
+        if self._device_storages:
+            # Storage retained from a previous load still holds the immutable
+            # weights; rebinding suffices, no H2D copy repeats.
+            self._bind(self._device_storages)
+            self.loaded = True
+            return
         try:
             self._load_once()
         except torch.OutOfMemoryError:
@@ -263,7 +277,8 @@ class PinnedModuleStager:
             self.loaded = False
             self._release_cache(force=True)
             raise
-        self._device_storages.clear()
+        if not self.retain_device_storage:
+            self._device_storages.clear()
         self.loaded = False
         self._release_cache()
 

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -14,6 +14,7 @@ from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload
 from .config import DIT_COMPONENT
+from .module_residency import PinnedModuleStager
 from .plan_resolver import resolve_offload_plan
 
 logger = init_logger(__name__)
@@ -61,6 +62,7 @@ class SequentialOffloadHook(ModelHook):
         pin_memory: bool = True,
         use_hsdp: bool = False,
         offload_after_context: bool = True,
+        persistent_staging: bool = False,
     ):
         # Modules to offload to CPU before this module runs
         self.offload_targets = offload_targets
@@ -68,6 +70,21 @@ class SequentialOffloadHook(ModelHook):
         self.pin_memory = pin_memory
         self.use_hsdp = use_hsdp
         self.offload_after_context = offload_after_context
+        self.persistent_staging = persistent_staging
+        self._stager: PinnedModuleStager | None = None
+
+    def initialize_hook(self, module: nn.Module) -> nn.Module:
+        if self.persistent_staging and not self.use_hsdp:
+            # Fixed staging storage keeps CUDA-graph-captured weight pointers
+            # valid across offload/load swaps. HSDP keeps the plain move
+            # semantics (decode graphs refuse HSDP anyway).
+            self._stager = PinnedModuleStager(
+                module,
+                self.device,
+                pin_memory=self.pin_memory,
+                retain_device_storage=True,
+            )
+        return module
 
     @staticmethod
     def _move_params(
@@ -102,7 +119,23 @@ class SequentialOffloadHook(ModelHook):
                 moved = True
         return moved
 
+    @staticmethod
+    def _module_stager(module: nn.Module) -> PinnedModuleStager | None:
+        """The persistent stager owned by ``module``'s own sequential hook, if any.
+
+        Cross-module swaps call ``_to_cpu(target)`` through the *running*
+        module's hook, so staging state must be resolved on the target's hook
+        to keep its ``loaded`` bookkeeping consistent.
+        """
+        registry: HookRegistry | None = getattr(module, "_hook_registry", None)
+        hook = registry.get_hook(SequentialOffloadHook._HOOK_NAME) if registry is not None else None
+        return hook._stager if isinstance(hook, SequentialOffloadHook) else None
+
     def _to_cpu(self, module: nn.Module) -> None:
+        stager = self._module_stager(module)
+        if stager is not None:
+            stager.offload()
+            return
         # XPU's allocator doesn't respect stream dependencies in empty_cache,
         # so non-blocking copies can race with cache eviction. Use blocking
         # copies on XPU to avoid NULL pointer errors during DMA.
@@ -117,6 +150,10 @@ class SequentialOffloadHook(ModelHook):
             current_omni_platform.empty_cache()
 
     def _to_gpu(self, module: nn.Module) -> None:
+        stager = self._module_stager(module)
+        if stager is not None:
+            stager.load()
+            return
         self._move_params(module, self.device, non_blocking=False)
 
     def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]:
@@ -148,6 +185,7 @@ def apply_sequential_offload(
     offload_initial_dits: bool = False,
     offload_dit_modules: Collection[nn.Module] | None = None,
     offload_encoder_modules: Collection[nn.Module] | None = None,
+    persistent_dit_staging: bool = False,
 ) -> None:
     """Apply sequential offloading hooks to DiT and encoder modules.
 
@@ -166,6 +204,10 @@ def apply_sequential_offload(
             every DiT for backward compatibility.
         offload_encoder_modules: Encoder/stage modules allowed to move to CPU.
             None selects every supplied module for backward compatibility.
+        persistent_dit_staging: Keep DiT weights on fixed device staging
+            storage across swaps, so CUDA-graph-captured weight pointers stay
+            valid. Encoders keep plain move semantics so their memory is
+            actually freed while the DiT runs.
 
     Example:
         >>> apply_sequential_offload(
@@ -194,6 +236,7 @@ def apply_sequential_offload(
                 pin_memory=pin_memory,
                 use_hsdp=use_hsdp,
                 offload_after_context=id(dit_mod) in selected_dit_ids,
+                persistent_staging=persistent_dit_staging,
             )
             registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook)
             logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__)
@@ -262,6 +305,15 @@ def _get_sequential_offload_hook(module: nn.Module) -> SequentialOffloadHook:
     return hook
 
 
+def sequential_offload_staging_active(module: nn.Module) -> bool:
+    """True when the module's sequential offload hook serves fixed device storage.
+
+    Fixed staging storage keeps CUDA-graph-captured weight pointers valid
+    across offload/load swaps.
+    """
+    return SequentialOffloadHook._module_stager(module) is not None
+
+
 @contextmanager
 def sequential_offload_component(module: nn.Module) -> Iterator[None]:
     """Activate and release a hooked component called outside ``forward``."""
@@ -362,6 +414,9 @@ class ModelLevelOffloadBackend(OffloadBackend):
                 use_hsdp=self.config.use_hsdp,
                 offload_dit_modules=(dits if self.config.offloads(DIT_COMPONENT) else ()),
                 offload_encoder_modules=selected_encoders,
+                # Fixed DiT staging keeps decode-graph weight pointers valid;
+                # other platforms keep plain move semantics.
+                persistent_dit_staging=self.device.type == "cuda",
             )
         except BaseException:
             try: