From 5b540f55d3b8f0e029350e2243d507a51ca5e6b2 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sat, 19 Sep 2026 00:10:53 +0800 Subject: [PATCH] Fix decode cuda graph for Qwen-Image-2.1 model-level cpu offload (#7801) --- .../models/qwen_image_21/test_decode_graph.py | 37 +++++++++++++ .../offloader/test_sequential_backend.py | 43 +++++++++++++++ .../models/qwen_image_21/decode_graph.py | 43 ++++++++++++--- .../qwen_image_21_transformer.py | 13 ++++- .../diffusion/offloader/module_residency.py | 17 +++++- .../diffusion/offloader/sequential_backend.py | 55 +++++++++++++++++++ 6 files changed, 199 insertions(+), 9 deletions(-) diff --git a/tests/diffusion/models/qwen_image_21/test_decode_graph.py b/tests/diffusion/models/qwen_image_21/test_decode_graph.py index bc377dbdd..ff4ea1a0a 100644 --- a/tests/diffusion/models/qwen_image_21/test_decode_graph.py +++ b/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 diff --git a/tests/diffusion/offloader/test_sequential_backend.py b/tests/diffusion/offloader/test_sequential_backend.py index 2d3e0db3b..043ca2236 100644 --- a/tests/diffusion/offloader/test_sequential_backend.py +++ b/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() diff --git a/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py b/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py index 1f568780a..018e8f35e 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py +++ b/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 diff --git a/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py b/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py index 07cf31645..b84c2912d 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py +++ b/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 ) diff --git a/vllm_omni/diffusion/offloader/module_residency.py b/vllm_omni/diffusion/offloader/module_residency.py index ab584c790..0e579e931 100644 --- a/vllm_omni/diffusion/offloader/module_residency.py +++ b/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() diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 1187c8cee..3dc32320b 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/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: