feat(minimax-h3): enable 2x RTX 4090/5090 with DLO (#5764)

Signed-off-by: lishunyang12 <lishunyang03@gmail.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: lishunyang12 <lishunyang03@gmail.com>

vllm_omni/diffusion/attention/backends/abstract.py

@@ -15,6 +15,10 @@ class AttentionBackend(ABC):
 
     accept_output_buffer: bool = False
     supports_piecewise_spans: bool = False
+    # The backend can represent a contiguous valid K/V prefix by slicing the
+    # tensors instead of materializing a padding mask. Models may use this to
+    # avoid a slower masked-attention plan when tail padding is not semantic.
+    supports_prefix_kv_slicing: bool = False
 
     @classmethod
     def supports_attention_mask(cls) -> bool:
@@ -83,6 +87,8 @@ class AttentionMetadata:
     #     variable-length query/key sequences for FlashAttention.
     #   "max_seqlen_q" / "max_seqlen_k": maximum sequence lengths paired with
     #     the packed cu_seqlens tensors.
+    #   "valid_kv_length": int — contiguous valid K/V prefix length for a
+    #     backend that advertises supports_prefix_kv_slicing.
 
     # Piecewise attention metadata (mixed causal/full masks).
     # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.

vllm_omni/diffusion/attention/backends/cudnn_attn.py

@@ -17,6 +17,7 @@ logger = init_logger(__name__)
 
 class CuDNNAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
+    supports_prefix_kv_slicing: bool = True
 
     @classmethod
     def supports_attention_mask(cls) -> bool:
@@ -61,7 +62,26 @@ class CuDNNAttentionImpl(AttentionImpl):
     ) -> torch.Tensor:
         attention_mask = None
         if attn_metadata:
-            attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask, mask_mode="broadcast_k")
+            valid_kv_length = attn_metadata.extra.get("valid_kv_length")
+            if attn_metadata.attn_mask is None and isinstance(valid_kv_length, int):
+                if not 0 < valid_kv_length <= key.shape[1]:
+                    raise ValueError(
+                        "valid_kv_length must be within the K/V sequence, "
+                        f"got {valid_kv_length} for length {key.shape[1]}"
+                    )
+                # A contiguous valid prefix is mathematically equivalent to a
+                # broadcast key-padding mask. Slicing K/V keeps Q/output in the
+                # checkpoint's aligned layout while letting cuDNN select its
+                # much faster mask-free FMHA plan.
+                key = key[:, :valid_kv_length]
+                value = value[:, :valid_kv_length]
+            else:
+                attention_mask = _maybe_reshape_attn_mask(
+                    query,
+                    key,
+                    attn_metadata.attn_mask,
+                    mask_mode="broadcast_k",
+                )
 
         enable_gqa = query.shape[2] != key.shape[2]
         query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))

vllm_omni/diffusion/data.py

@@ -687,9 +687,13 @@ class OmniDiffusionConfig:
     # Distributed layer-wise offloading with H2D + AllGather overlap (RFC-1)
     enable_distributed_layerwise_offload: bool = False
     # If True: shard weights 1/dp_size + AllGather (saves CPU memory, requires
-    # concurrent requests in DP mode). If False: each rank loads full weights
-    # via H2D only (N× CPU memory, but no AllGather synchronization needed).
+    # concurrent requests in DP mode). If False: each rank streams the standard
+    # loader's rank-local tensors (including TP-local shards) via H2D only.
+    # This avoids AllGather synchronization, while host memory follows the
+    # loader's existing rank-local layout instead of adding a second DP shard.
     dlo_use_allgather: bool = True
+    # Leading main-DiT blocks kept resident by distributed layerwise offload.
+    dlo_resident_layers: int = 0
 
     pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading
 

vllm_omni/diffusion/diffusion_engine.py

@@ -237,7 +237,7 @@ class DiffusionEngine:
         if getattr(self.od_config, "parallel_config", None) is not None:
             dp_size = getattr(self.od_config.parallel_config, "data_parallel_size", 1)
         dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
-        if dp_size > 1 and dist_offload:
+        if dp_size > 1 and dist_offload and getattr(self.od_config, "dlo_use_allgather", True):
             self.scheduler.max_num_running_reqs = dp_size
             self.dp_concurrent = True
             # Set batch admission wait so concurrent requests accumulate

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -416,11 +416,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                 )
             extra_args_signatures: set = set()
             for nr in new_reqs:
-                ea = getattr(nr.req, "extra_args", None)
-                if ea and isinstance(ea, dict):
-                    extra_args_signatures.add(json.dumps(ea, sort_keys=True))
-                else:
-                    extra_args_signatures.add(None)
+                ea = getattr(nr.req.sampling_params, "extra_args", None)
+                extra_args_signatures.add(json.dumps(ea, sort_keys=True, default=repr) if ea is not None else None)
             if len(extra_args_signatures) > 1:
                 raise ValueError(
                     "DP multi-concurrency requires all concurrent requests to "
@@ -521,6 +518,18 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         if len(scheduler_output.scheduled_new_reqs) <= 1:
             return self.execute_request(scheduler_output)
 
+        parallel_config = getattr(self.od_config, "parallel_config", None)
+        dp_size = getattr(parallel_config, "data_parallel_size", 1)
+        if (
+            dp_size > 1
+            and getattr(self.od_config, "enable_distributed_layerwise_offload", False)
+            and getattr(self.od_config, "dlo_use_allgather", True)
+        ):
+            # DLO DP uses one independent request per DP replica.  It is not a
+            # fused pipeline request batch, so models such as MiniMax-H3 do not
+            # need to advertise supports_request_batch=True.
+            return self.execute_request(scheduler_output)
+
         result = self.collective_rpc(
             "execute_model_batch",
             args=(scheduler_output, self.od_config),
@@ -599,7 +608,12 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         # so the D2H copy runs on a side stream in the worker background
         # thread while the default stream is free for the next forward.
         # pump routes the result back to a Future via rpc_id.
-        if not self.od_config.step_execution and method in ("execute_model", "execute_model_batch"):
+        multi_rank_reply = unique_reply_rank is None and exec_all_ranks
+        if (
+            not self.od_config.step_execution
+            and method in ("execute_model", "execute_model_batch")
+            and not multi_rank_reply
+        ):
             rpc_id = self._next_rpc_id()
             rpc_request["rpc_id"] = rpc_id
             fut: concurrent.futures.Future = concurrent.futures.Future()

vllm_omni/diffusion/models/minimax_h3/encoder.py

@@ -1029,16 +1029,84 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module):
     def load_to_device(self) -> None:
         if not self.is_loaded:
             return
+        if getattr(self, "_omni_layerwise_enabled", False):
+            for name, child in self.vision.named_children():
+                if name != "blocks":
+                    child.to(self.device_target)
+            for name, child in self.text_model.named_children():
+                if name != "layers":
+                    child.to(self.device_target)
+            return
         self.vision.to(self.device_target)
         self.text_model.to(self.device_target)
 
     def offload_to_cpu(self) -> None:
         if not self.is_loaded:
             return
+        if getattr(self, "_omni_layerwise_enabled", False):
+            for hook in self._omni_layerwise_hooks:
+                hook.offload_layer()
+            for name, child in self.vision.named_children():
+                if name != "blocks":
+                    child.to("cpu")
+            for name, child in self.text_model.named_children():
+                if name != "layers":
+                    child.to("cpu")
+            torch.accelerator.empty_cache()
+            return
         self.vision.to("cpu")
         self.text_model.to("cpu")
         torch.accelerator.empty_cache()
 
+    def enable_omni_layerwise_offload(self, *, pin_memory: bool = True) -> None:
+        """Stream the TP-local Qwen vision/text blocks for low-HBM serving.
+
+        The encoder has its own TP process group, so these blocks must remain
+        rank-local.  Reusing the DiT AllGather group would concatenate
+        different TP shards and corrupt the encoder weights.
+        """
+        if not self.is_loaded or getattr(self, "_omni_layerwise_enabled", False):
+            return
+
+        from vllm_omni.diffusion.offloader.layerwise_backend import apply_block_hook
+        from vllm_omni.platforms import current_omni_platform
+
+        self._omni_layerwise_hooks = []
+        self._omni_layerwise_block_groups = []
+        copy_stream = current_omni_platform.Stream()
+        for blocks in (self.vision.blocks, self.text_model.layers):
+            if len(blocks) <= 1:
+                continue
+            last_hook = apply_block_hook(
+                blocks[-1],
+                blocks[0],
+                self.device_target,
+                copy_stream,
+                pin_memory,
+            )
+            hooks = [last_hook]
+            for index, block in enumerate(blocks[:-1]):
+                hooks.append(
+                    apply_block_hook(
+                        block,
+                        blocks[index + 1],
+                        self.device_target,
+                        copy_stream,
+                        pin_memory,
+                    )
+                )
+            for index, hook in enumerate(hooks):
+                hook._prev_hook = hooks[index - 1]
+            self._omni_layerwise_hooks.extend(hooks)
+            self._omni_layerwise_block_groups.append(blocks)
+
+        self._omni_layerwise_enabled = bool(self._omni_layerwise_hooks)
+        logger.info(
+            "MiniMax H3 encoder layerwise offload enabled on %d blocks across %d stacks",
+            sum(len(blocks) for blocks in self._omni_layerwise_block_groups),
+            len(self._omni_layerwise_block_groups),
+        )
+
     def _encode(
         self,
         input_ids: torch.Tensor,

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -376,18 +376,30 @@ class MiniMaxH3Attention(nn.Module):
         *,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
+        packed_total: int,
     ) -> torch.Tensor:
         """Run packed attention as a small eager island.
 
-        The scalar packed-layout metadata and the CuTe FlashAttention-4 DSL
-        are intentionally opaque to Dynamo. Keeping this boundary narrow lets
-        regional compile fuse projections, norms, RoPE, and the surrounding
-        DiT block without repeatedly graph-breaking inside the FA4 compiler.
+        The scalar packed-layout metadata and backend-specific attention
+        kernels are intentionally opaque to Dynamo. Keeping this boundary
+        narrow lets regional compile fuse projections, norms, RoPE, and the
+        surrounding DiT block without repeated graph breaks.
         """
-        used = int(cu_seqlens[1].item())
-        packed_total = int(cu_seqlens[-1].item())
+        # max_seqlen is already the first (real) packed document length. Do
+        # not read the CUDA cu_seqlens scalars here: this function runs once
+        # per layer and .item() would serialize every attention launch.
+        if not 0 < max_seqlen <= packed_total:
+            raise ValueError(
+                f"max_seqlen must be within the packed sequence, got {max_seqlen} for length {packed_total}"
+            )
+        used = min(max_seqlen, packed_total)
         attn_mask = None
-        if used < packed_total:
+        # Ring attention can dispatch to a different implementation from the
+        # configured backend, so this no-mask fast path is local-only.
+        prefix_slice = (
+            not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing
+        )
+        if used < packed_total and not prefix_slice:
             attn_mask = torch.arange(packed_total, device=q.device)[None] < used
         metadata = AttentionMetadata(
             attn_mask=attn_mask,
@@ -396,6 +408,7 @@ class MiniMaxH3Attention(nn.Module):
                 "cu_seqlens_k": cu_seqlens,
                 "max_seqlen_q": max_seqlen,
                 "max_seqlen_k": max_seqlen,
+                "valid_kv_length": used,
             },
         )
         return self.attention(
@@ -412,6 +425,7 @@ class MiniMaxH3Attention(nn.Module):
         rope_freqs: torch.Tensor | None,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
+        packed_total: int | None = None,
         sp_seq_lens: list[int] | None = None,
     ) -> torch.Tensor:
         """x: [T, hidden] packed thd rows -> [T, hidden].
@@ -448,6 +462,10 @@ class MiniMaxH3Attention(nn.Module):
             v,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
+            # Before Ulysses, q contains only this rank's row shard. The
+            # backend receives the global sequence after all-to-all, so carry
+            # its Python length explicitly instead of inferring it from q.
+            packed_total=packed_total if packed_total is not None else q.shape[0],
         )
         out = out.reshape(total, self.num_heads * self.head_dim)
         out, _ = self.out_proj(out)
@@ -655,6 +673,7 @@ class MiniMaxH3DiTBlock(nn.Module):
         rope_freqs: torch.Tensor,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
+        packed_total: int,
         sp_seq_lens: list[int] | None = None,
     ) -> torch.Tensor:
         """x: [T, H]; t_emb: [M, t_dim]; combined_indices: [T]
@@ -681,6 +700,7 @@ class MiniMaxH3DiTBlock(nn.Module):
             rope_freqs=rope_freqs,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
+            packed_total=packed_total,
             sp_seq_lens=sp_seq_lens,
         )
         x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE)
@@ -1127,6 +1147,7 @@ class MiniMaxH3DiTModel(nn.Module):
                 rope_freqs=block_rope,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
+                packed_total=seq_len,
             )
         hidden = self.sp_gather(hidden)
 

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -8,6 +8,7 @@ import math
 import os
 import tempfile
 from collections.abc import Iterable, Mapping, Sequence
+from contextlib import contextmanager
 from pathlib import Path
 from typing import Any, ClassVar
 
@@ -36,6 +37,7 @@ from vllm_omni.diffusion.models.interface import (
     SupportsComponentDiscovery,
 )
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.offloader import OffloadPlan
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
     DiffusionPipelineProfilerMixin,
 )
@@ -454,8 +456,25 @@ class MiniMaxH3Pipeline(
     """CFG-distilled joint video/audio generation for MiniMax H3."""
 
     _dit_modules: ClassVar[list[str]] = ["transformer"]
+    # TODO(offload): Reconcile this static plan with the modular H3 pipeline
+    # (#5720). Its combined mode has ``transformer`` and ``transformers_ref``;
+    # build the resident/staging plan from the selected task partition so a
+    # consumer-GPU single-task server never stages both DiTs.
     _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
     _vae_modules: ClassVar[list[str]] = ["video_vae", "audio_vae"]
+    _offload_plan: ClassVar[OffloadPlan] = OffloadPlan(
+        offload_submodules={"token_refiner": "blocks"},
+        resident_dit_paths=frozenset({"transformer"}),
+        encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")},
+        on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}),
+    )
+    # H3's regular loader performs checkpoint-layout conversions (grouped QKV
+    # and fused MLP). Keep it on that path until mmap can run the same loader
+    # callbacks for every affected parameter.
+    _supports_mmap_loading: ClassVar[bool] = False
+    # TODO(offload): Re-enable after the generic rank-local mmap path can run
+    # this model's grouped-QKV reorder and fused-MLP packing before TP sharding.
+    # Do not bypass the regular loader until that equivalence is tested.
     _PROFILER_TARGETS: ClassVar[list[str]] = [
         "_prepare_reference_videos",
         "encode_prompt",
@@ -542,13 +561,19 @@ class MiniMaxH3Pipeline(
             load_model=rank < text_encoder_tp_size,
             encoder_group=self.text_encoder_group,
         )
+        stage_components = bool(
+            od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False)
+        )
+        component_load_device = torch.device("cpu") if stage_components else self.device
         self.video_vae = MiniMaxH3VideoVAE(
             os.path.join(model_path, "video_vae"),
             device=self.device,
+            load_device=component_load_device,
         )
         self.audio_vae = MiniMaxH3AudioVAE(
             os.path.join(model_path, "audio_vae"),
             device=self.device,
+            load_device=component_load_device,
         )
         # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``.
         self.vae = self.video_vae
@@ -935,12 +960,16 @@ class MiniMaxH3Pipeline(
         input_ids: torch.Tensor,
         vision_kwargs: dict[str, torch.Tensor],
     ) -> torch.Tensor:
-        if self.od_config.enable_cpu_offload:
+        if self.od_config.enable_cpu_offload and not getattr(
+            self.od_config, "enable_distributed_layerwise_offload", False
+        ):
             # Invoke nn.Module.__call__ so the generic model-level offloader
             # swaps the resident DiT and encoder.
             return self.text_encoder(input_ids, **vision_kwargs)
 
-        if self.od_config.enable_layerwise_offload:
+        if self.od_config.enable_layerwise_offload or getattr(
+            self.od_config, "enable_distributed_layerwise_offload", False
+        ):
             # Layerwise DiT offload already provides the low-residency encoder
             # phase used by the checkpoint reference.
             self.text_encoder.load_to_device()
@@ -955,12 +984,33 @@ class MiniMaxH3Pipeline(
         self.text_encoder.load_to_device()
         return self.text_encoder.encode_ids(input_ids, **vision_kwargs)
 
+    def _uses_manual_component_offload(self) -> bool:
+        od_config = getattr(self, "od_config", None)
+        return bool(
+            getattr(od_config, "enable_layerwise_offload", False)
+            or getattr(od_config, "enable_distributed_layerwise_offload", False)
+        )
+
+    @contextmanager
+    def _component_on_device(self, component: nn.Module):
+        staged = self._uses_manual_component_offload()
+        if staged:
+            component.load_to_device()
+        try:
+            yield
+        finally:
+            if staged:
+                component.offload_to_cpu()
+
     def _encode_visual_condition(
         self,
         image: Image.Image,
     ) -> torch.Tensor:
         _, rank, _ = _dit_rank_world()
-        rows = self.video_vae.encode_image(image) if rank == 0 else None
+        rows = None
+        if rank == 0:
+            with self._component_on_device(self.video_vae):
+                rows = self.video_vae.encode_image(image)
         return _broadcast_tensor(
             rows,
             dtype=torch.float32,
@@ -996,7 +1046,8 @@ class MiniMaxH3Pipeline(
         rows = None
         audio_t = 0
         if rank == 0:
-            rows, audio_t = self.audio_vae.encode_waveform(*audio)
+            with self._component_on_device(self.audio_vae):
+                rows, audio_t = self.audio_vae.encode_waveform(*audio)
         audio_t_tensor = torch.tensor(
             [audio_t],
             dtype=torch.long,
@@ -1042,6 +1093,15 @@ class MiniMaxH3Pipeline(
         prepared_videos: list[dict[str, Any]] | None,
         *,
         count: int,
+    ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]:
+        with self._component_on_device(self.video_vae):
+            return self._encode_video_conditions_resident(prepared_videos, count=count)
+
+    def _encode_video_conditions_resident(
+        self,
+        prepared_videos: list[dict[str, Any]] | None,
+        *,
+        count: int,
     ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]:
         group, rank, world_size = _dit_rank_world()
         distributed_encode = self.video_vae.is_distributed_enabled()
@@ -1089,6 +1149,15 @@ class MiniMaxH3Pipeline(
         prepared_videos: list[dict[str, Any]] | None,
         *,
         has_audio: list[bool],
+    ) -> tuple[torch.Tensor | None, list[int]]:
+        with self._component_on_device(self.audio_vae):
+            return self._encode_video_audio_conditions_resident(prepared_videos, has_audio=has_audio)
+
+    def _encode_video_audio_conditions_resident(
+        self,
+        prepared_videos: list[dict[str, Any]] | None,
+        *,
+        has_audio: list[bool],
     ) -> tuple[torch.Tensor | None, list[int]]:
         _, rank, _ = _dit_rank_world()
         count = sum(has_audio)
@@ -1156,6 +1225,17 @@ class MiniMaxH3Pipeline(
         )
         return video_rows, audio_rows
 
+    @contextmanager
+    def _resident_dit_layers_on_device(self):
+        controller = getattr(self, "_dlo_residency_controller", None)
+        if controller is not None:
+            controller.load_resident_layers()
+        try:
+            yield
+        finally:
+            if controller is not None:
+                controller.offload_resident_layers()
+
     def diffuse(
         self,
         *,
@@ -1277,21 +1357,22 @@ class MiniMaxH3Pipeline(
             num_steps=num_steps,
             shift_scale=audio_shift,
         )
-        with self.progress_bar(total=len(video_sigmas) - 1) as progress:
-            video_rows, audio_rows = minimax_h3_denoise_loop(
-                model=self.transformer,
-                positive=branch,
-                initial_video_rows=initial_video,
-                initial_audio_rows=initial_audio,
-                keyframe_cond_rows=visual_anchor,
-                audio_ref_rows=audio_anchor,
-                sigmas_video=video_sigmas,
-                sigmas_audio=audio_sigmas,
-                device=self.device,
-                imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP),
-                audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP),
-                on_step=lambda step, video, audio: progress.update(),
-            )
+        with self._resident_dit_layers_on_device():
+            with self.progress_bar(total=len(video_sigmas) - 1) as progress:
+                video_rows, audio_rows = minimax_h3_denoise_loop(
+                    model=self.transformer,
+                    positive=branch,
+                    initial_video_rows=initial_video,
+                    initial_audio_rows=initial_audio,
+                    keyframe_cond_rows=visual_anchor,
+                    audio_ref_rows=audio_anchor,
+                    sigmas_video=video_sigmas,
+                    sigmas_audio=audio_sigmas,
+                    device=self.device,
+                    imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP),
+                    audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP),
+                    on_step=lambda step, video, audio: progress.update(),
+                )
 
         target_video = video_rows[branch.update_mask_dev]
         video_latent = minimax_h3_unpatchify_video_tokens(
@@ -1320,14 +1401,16 @@ class MiniMaxH3Pipeline(
         height: int,
         width: int,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        with current_omni_platform.create_autocast_context(
-            device_type=self.device.type,
-            dtype=torch.float16,
-            enabled=True,
-        ):
-            video = self.video_vae.decode_latent(video_latent)
+        with self._component_on_device(self.video_vae):
+            with current_omni_platform.create_autocast_context(
+                device_type=self.device.type,
+                dtype=torch.float16,
+                enabled=True,
+            ):
+                video = self.video_vae.decode_latent(video_latent)
         video = video[..., :height, :width].contiguous()
-        audio = self.audio_vae.decode_latent(audio_latent)
+        with self._component_on_device(self.audio_vae):
+            audio = self.audio_vae.decode_latent(audio_latent)
         return video, audio
 
     @torch.no_grad()

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -19,6 +19,7 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     DistributedVaeMixin,
 )
 from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
+from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
 
 from .packed_tokens import minimax_h3_patchify_video_latent
 
@@ -102,8 +103,10 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         component_path: str,
         *,
         device: torch.device,
+        load_device: torch.device | None = None,
     ) -> None:
         super().__init__()
+        self._device_target = device
         self.config_dict = _load_component_config(component_path)
         self.remote = _load_remote_component(
             component_path,
@@ -112,13 +115,34 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         # Match the reference loader contract: video VAE weights stay FP32.
         # Keyframe encoding is numerically sensitive to first casting the
         # checkpoint through FP16; decode still runs under FP16 autocast.
-        self.remote.eval().to(device=device, dtype=torch.float32)
+        initial_device = load_device or device
+        self.remote.eval().to(device=initial_device, dtype=torch.float32)
+        self._stager = None
+        if initial_device.type == "cpu" and device.type not in ("cpu", "meta"):
+            self._stager = PinnedModuleStager(
+                self.remote,
+                device,
+                pin_memory=True,
+            )
         self.model = self.remote.model
         self.use_tiling = True
         self.use_slicing = False
         self.parallel_size = 1
         self.device_module = torch.get_device_module()
 
+    def load_to_device(self) -> None:
+        if self._stager is not None:
+            self._stager.load()
+        else:
+            self.remote.to(self._device_target)
+
+    def offload_to_cpu(self) -> None:
+        if self._stager is not None:
+            self._stager.offload()
+        else:
+            self.remote.to("cpu")
+            torch.accelerator.empty_cache()
+
     def set_parallel_size(
         self,
         parallel_size: int,
@@ -277,8 +301,10 @@ class MiniMaxH3AudioVAE(nn.Module):
         component_path: str,
         *,
         device: torch.device,
+        load_device: torch.device | None = None,
     ) -> None:
         super().__init__()
+        self._device_target = device
         self.config_dict = _load_component_config(component_path)
         self.remote = _load_remote_component(
             component_path,
@@ -286,10 +312,31 @@ class MiniMaxH3AudioVAE(nn.Module):
         )
         # The checkpoint's audio VAE contract is FP32 for both reference
         # encoding and waveform decoding.
-        self.remote.eval().to(device=device, dtype=torch.float32)
+        initial_device = load_device or device
+        self.remote.eval().to(device=initial_device, dtype=torch.float32)
+        self._stager = None
+        if initial_device.type == "cpu" and device.type not in ("cpu", "meta"):
+            self._stager = PinnedModuleStager(
+                self.remote,
+                device,
+                pin_memory=True,
+            )
         self.model = self.remote.model
         self.sample_rate = int(self.config_dict["sample_rate"])
 
+    def load_to_device(self) -> None:
+        if self._stager is not None:
+            self._stager.load()
+        else:
+            self.remote.to(self._device_target)
+
+    def offload_to_cpu(self) -> None:
+        if self._stager is not None:
+            self._stager.offload()
+        else:
+            self.remote.to("cpu")
+            torch.accelerator.empty_cache()
+
     @torch.inference_mode()
     def encode_waveform(
         self,

vllm_omni/diffusion/offloader/__init__.py

@@ -16,6 +16,7 @@ from .distributed_layerwise_backend import (
     remove_distributed_block_hook,
 )
 from .layerwise_backend import LayerWiseOffloadBackend
+from .module_residency import PinnedModuleStager
 from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading
 from .sequential_backend import (
     ModelLevelOffloadBackend,
@@ -41,6 +42,7 @@ __all__ = [
     "DistributedLayerwiseOffloadBackend",
     "DistributedLayerwiseOffloadHook",
     "ModelLevelOffloadBackend",
+    "PinnedModuleStager",
     "apply_sequential_offload",
     "remove_sequential_offload",
     "apply_distributed_block_hook",

vllm_omni/diffusion/offloader/base.py

@@ -27,7 +27,10 @@ class OffloadConfig:
     pin_cpu_memory: bool = True
     use_hsdp: bool = False
     dp_size: int = 1  # derived from parallel_config, not user-configurable
-    dlo_use_allgather: bool = True  # if True: shard + AllGather; if False: full weights H2D only
+    # True: add DP sharding + AllGather. False: stream the standard loader's
+    # rank-local tensors (including TP-local shards) with H2D only.
+    dlo_use_allgather: bool = True
+    dlo_resident_layers: int = 0  # leading DiT layers kept on device
     model_path: str | None = None  # checkpoint path for mmap weight loading
 
     @classmethod
@@ -94,17 +97,26 @@ class OffloadConfig:
         else:
             strategy = OffloadStrategy.NONE
 
-        # When dlo_use_allgather=False, each rank loads full weights (no sharding,
-        # no AllGather).  This avoids AllGather synchronization requirements
-        # (concurrent requests, dummy run skip) at the cost of N× CPU memory.
+        # With dlo_use_allgather=False, do not add another DP shard. Each rank
+        # streams the tensors produced by the standard loader, which may
+        # already be TP-local shards. This avoids AllGather synchronization
+        # requirements (concurrent requests, dummy run skip).
         dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
+        dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
+        if dlo_resident_layers < 0:
+            raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
+        if dlo_resident_layers and dlo_use_allgather:
+            raise ValueError(
+                "dlo_resident_layers currently requires --dlo-no-use-allgather so "
+                "resident blocks use weights prepared by the standard TP-aware loader"
+            )
 
         # If dlo_use_allgather=False, force dp_size=1 (each rank independent)
         if enable_distributed_layerwise_offload and not dlo_use_allgather:
             dp_size = 1
             logger.info(
                 "Distributed layerwise offload: dlo_use_allgather=False, "
-                "each rank loads full weights (no shard, no AllGather)"
+                "streaming standard-loader rank-local weights (no DP shard or AllGather)"
             )
 
         # HSDP already shards parameters into DTensors.  Running distributed
@@ -114,8 +126,8 @@ class OffloadConfig:
             raise ValueError(
                 "Distributed layerwise offload with AllGather is incompatible with "
                 "HSDP: HSDP parameters are already sharded DTensors, and the offloader "
-                "would double-shard them. Use --dlo-no-use-allgather (full weights per "
-                "rank) or disable HSDP."
+                "would double-shard them. Use --dlo-no-use-allgather (standard-loader "
+                "rank-local weights) or disable HSDP."
             )
 
         return cls(
@@ -124,6 +136,7 @@ class OffloadConfig:
             use_hsdp=use_hsdp,
             dp_size=dp_size,
             dlo_use_allgather=dlo_use_allgather,
+            dlo_resident_layers=dlo_resident_layers,
             model_path=getattr(od_config, "model", None),
         )
 

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -219,6 +219,15 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             weights_with_local = []
             for name, t in name2weights.items():
                 local_t = t.to_local() if hasattr(t, "to_local") else t
+                mmap_transform = getattr(t, "mmap_weight_transform", None)
+                if callable(mmap_transform) and getattr(t, "mmap_weight_transform_pending", False):
+                    # Some checkpoints use a layout that is converted by the
+                    # regular weight loader (for example MiniMax-H3 grouped
+                    # QKV).  Apply that conversion one block at a time while
+                    # copying the rank-local CPU shard.  Keeping the raw
+                    # parameter as an mmap view avoids a private full-model
+                    # copy in every worker.
+                    local_t = mmap_transform(local_t)
                 weights_with_local.append((name, t, local_t))
 
             total_numel = sum(local.numel() for _, _, local in weights_with_local)
@@ -331,9 +340,9 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             with current_omni_platform.stream(self.comm_stream):
                 for dtype, local_shard in gpu_shards.items():
                     # Slice the shared (max-sized) output buffer down to this
-                    # block's actual AllGather output size.  The buffers are
-                    # sized to the *largest* block across all groups, so for
-                    # any smaller block the full buffer would violate the
+                    # block's actual AllGather output size. The buffers are
+                    # sized to the *largest* block across all groups, so for any
+                    # smaller block the full buffer would violate the
                     # all_gather_into_tensor contract
                     # (output.numel() == world_size * input.numel()).
                     # Repoint offsets are relative to the block's flattened
@@ -485,6 +494,111 @@ def remove_distributed_block_hook(module: nn.Module) -> None:
         logger.debug("Removed distributed offload hook from %s", module.__class__.__name__)
 
 
+class PinnedResidentLayerGroup:
+    """Keep selected layers in pinned host memory between requests.
+
+
+    TODO(offload): Extract this alongside PinnedModuleStager after the
+    distributed shard-and-pin operation becomes a shared storage primitive.
+    It currently remains here because it depends on DLO's local-shard layout.
+    Unlike ``module.to(device)``/``module.to("cpu")``, this group retains a
+    pinned CPU master copy and never copies generated device weights back to
+    host.  Entering the denoise stage performs one asynchronous H2D pass;
+    leaving it only restores zero-sized placeholders and releases the device
+    buffers.  This lets the following VAE stage reuse the same HBM.
+
+    The resident path is intentionally local-shard only.  With tensor
+    parallelism, the regular model loader has already produced each rank's TP
+    shard, so no DP AllGather is required or desirable here.
+    """
+
+    def __init__(
+        self,
+        blocks: list[nn.Module],
+        device: torch.device,
+        copy_stream: Any,
+        pin_memory: bool,
+    ) -> None:
+        self.device = device
+        self.copy_stream = copy_stream
+        self.loaded = False
+        self._states: list[dict[str, Any]] = []
+        self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
+
+        for block in blocks:
+            params = dict(block.named_parameters())
+            bufs = dict(block.named_buffers())
+            targets: dict[str, torch.Tensor] = {**params, **bufs}
+            cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin(
+                params,
+                bufs,
+                dp_size=1,
+                rank=0,
+                pin_memory=pin_memory,
+            )
+            self._states.append(
+                {
+                    "targets": targets,
+                    "cpu_shards": cpu_shards,
+                    "metadata": metadata,
+                }
+            )
+
+    def load(self) -> None:
+        if self.loaded:
+            return
+
+        # Allocate on the compute stream so the caching allocator can reuse
+        # blocks released by the encoder stage.  Allocating on the copy stream
+        # would create a separate stream-local pool and inflate peak HBM.
+        gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
+        for state in self._states:
+            block_buffers: dict[torch.dtype, torch.Tensor] = {}
+            for dtype, cpu_shard in state["cpu_shards"].items():
+                block_buffers[dtype] = torch.empty(
+                    cpu_shard.shape,
+                    dtype=dtype,
+                    device=self.device,
+                )
+            gpu_buffers.append(block_buffers)
+
+        self.copy_stream.wait_stream(current_omni_platform.current_stream())
+        ready = current_omni_platform.Event()
+        with current_omni_platform.stream(self.copy_stream):
+            for state, block_buffers in zip(self._states, gpu_buffers):
+                for dtype, cpu_shard in state["cpu_shards"].items():
+                    block_buffers[dtype].copy_(cpu_shard, non_blocking=True)
+            ready.record(self.copy_stream)
+
+        for state, block_buffers in zip(self._states, gpu_buffers):
+            targets = state["targets"]
+            for dtype, metas in state["metadata"].items():
+                gpu_buffer = block_buffers[dtype]
+                for meta in metas:
+                    set_tensor_storage(
+                        targets[meta["name"]],
+                        gpu_buffer[meta["offset"] : meta["offset"] + meta["numel"]].view(meta["shape"]),
+                    )
+
+        current_omni_platform.current_stream().wait_event(ready)
+        self._gpu_buffers = gpu_buffers
+        self.loaded = True
+
+    def offload(self) -> None:
+        if not self.loaded:
+            return
+
+        # Denoise kernels consume these buffers on the current compute stream.
+        # Synchronize once at the stage boundary; no D2H copy is necessary
+        # because the pinned CPU master weights were never overwritten.
+        current_omni_platform.synchronize()
+        for state in self._states:
+            for target in state["targets"].values():
+                set_tensor_storage(target, make_offload_placeholder(target))
+        self._gpu_buffers.clear()
+        self.loaded = False
+
+
 # ---------------------------------------------------------------------- #
 #  Backend                                                                #
 # ---------------------------------------------------------------------- #
@@ -502,6 +616,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
     overlapped with computation.
     """
 
+    _MMAP_PARAM_ATTRS = ("weight_loader", "mmap_weight_transform")
+
     def __init__(self, config: OffloadConfig, device: torch.device):
         super().__init__(config, device)
 
@@ -512,6 +628,41 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self.rank = 0
         self._blocks: list[list[nn.Module]] = []
         self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = []
+        self._resident_blocks: list[nn.Module] = []
+        self._resident_layer_group: PinnedResidentLayerGroup | None = None
+
+    def load_resident_layers(self) -> None:
+        """Load the model-declared leading blocks for the denoise stage."""
+        if self._resident_layer_group is not None:
+            self._resident_layer_group.load()
+
+    def offload_resident_layers(self) -> None:
+        """Release leading blocks before VAE decode to bound peak HBM."""
+        if self._resident_layer_group is not None:
+            self._resident_layer_group.offload()
+
+    def _remember_mmap_param_attrs(self, pipeline: nn.Module) -> None:
+        """Save loader metadata before ``to_empty`` replaces Parameters."""
+        self._mmap_param_attrs = {
+            name: {attr: value for attr in self._MMAP_PARAM_ATTRS if (value := getattr(param, attr, None)) is not None}
+            for name, param in pipeline.named_parameters()
+        }
+
+    def _attach_mmap_param_attrs(
+        self,
+        name: str,
+        replacement: nn.Parameter,
+        source: nn.Parameter | None = None,
+    ) -> None:
+        attrs = getattr(self, "_mmap_param_attrs", {}).get(name, {})
+        for attr in self._MMAP_PARAM_ATTRS:
+            value = attrs.get(attr)
+            if value is None and source is not None:
+                value = getattr(source, attr, None)
+            if value is not None:
+                setattr(replacement, attr, value)
+                if attr == "mmap_weight_transform":
+                    replacement.mmap_weight_transform_pending = True
 
     def _load_weights_via_mmap(self, pipeline: nn.Module, modules) -> None:
         """Load DiT weights from safetensors via mmap views (no RSS).
@@ -614,6 +765,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 "instead of TP, or disable distributed layerwise offload."
             )
 
+        # ``Module.to_empty`` creates new Parameter objects and drops custom
+        # attributes installed by vLLM weight loaders.  Preserve them by full
+        # pipeline name so mmap replacements can reproduce checkpoint layout
+        # transforms such as MiniMax-H3 grouped QKV.
+        self._remember_mmap_param_attrs(pipeline)
+
         #
         # Important: when DiT modules are nested (e.g. transformer contains
         # transformer.language_model), to_empty("meta") on the parent
@@ -738,7 +895,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     parent = parent[int(part)]
                 else:
                     parent = getattr(parent, part)
-            parent._parameters[parts[-1]] = torch.nn.Parameter(tensor, requires_grad=param.requires_grad)
+            replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad)
+            self._attach_mmap_param_attrs(name, replacement, param)
+            parent._parameters[parts[-1]] = replacement
             loaded += 1
             loaded_names.add(name)
 
@@ -793,7 +952,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                             parent = parent[int(part)]
                         else:
                             parent = getattr(parent, part)
-                    parent._parameters[bparts[-1]] = torch.nn.Parameter(tensor, requires_grad=bparam.requires_grad)
+                    replacement = torch.nn.Parameter(tensor, requires_grad=bparam.requires_grad)
+                    self._attach_mmap_param_attrs(full_param_name, replacement, bparam)
+                    parent._parameters[bparts[-1]] = replacement
                     block_loaded += 1
                     loaded_names.add(full_param_name)
 
@@ -904,7 +1065,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     parent = parent[int(part)]
                 else:
                     parent = getattr(parent, part)
-            parent._parameters[parts[-1]] = torch.nn.Parameter(tensor, requires_grad=param.requires_grad)
+            replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad)
+            self._attach_mmap_param_attrs(full_name, replacement, param)
+            parent._parameters[parts[-1]] = replacement
             loaded += 1
 
         if loaded > 0:
@@ -962,24 +1125,74 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             coord.ranks,
         )
 
-    def _register_on_demand_hook(self, module: nn.Module, label: str) -> None:
-        """Move *module* to GPU permanently.
+    def _register_on_demand_hook(self, module: nn.Module, label: str, *, stage_on_demand: bool = False) -> None:
+        """Prepare a pipeline-managed stage component or keep it resident.
 
-        Previously, on-demand hooks moved the module to GPU before forward
-        and back to CPU after forward (with synchronize + empty_cache).
-        However, _post_forward's global synchronize() and empty_cache()
-        disrupt the async prefetch pipeline (copy_stream / comm_stream)
-        used by distributed layerwise offload, causing garbled output.
-
-        Moving the module to GPU permanently avoids this issue at the
-        cost of slightly higher HBM usage (VAE/encoders are small).
+        Components that expose an explicit stage lifecycle are initially
+        offloaded and loaded by their pipeline only around encode/decode.
+        Other models retain the conservative resident behavior because a
+        generic post-forward hook can disrupt the DiT prefetch streams.
         """
+        offload_to_cpu = getattr(module, "offload_to_cpu", None)
+        if stage_on_demand and callable(offload_to_cpu):
+            offload_to_cpu()
+            logger.info("Prepared %s (%s) for pipeline-managed staged offload", label, module.__class__.__name__)
+            return
         module.to(self.device)
         logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__)
 
+    def _try_layerwise_offload_encoder(self, module: nn.Module, name: str, plan: OffloadPlan | None) -> bool:
+        """Stream plan-declared encoder blocks on each rank without AllGather."""
+        if plan is None or name not in plan.encoder_block_attrs:
+            return False
+        if getattr(module, "_omni_layerwise_enabled", False):
+            return True
+
+        from operator import attrgetter
+
+        from vllm_omni.diffusion.offloader.layerwise_backend import apply_block_hook
+
+        hooks = []
+        block_groups = []
+        copy_stream = current_omni_platform.Stream()
+        for block_path in plan.encoder_block_attrs[name]:
+            try:
+                blocks = attrgetter(block_path)(module)
+            except AttributeError:
+                logger.warning("Encoder offload path %s.%s was not found", name, block_path)
+                continue
+            if not isinstance(blocks, nn.ModuleList) or len(blocks) <= 1:
+                logger.warning("Encoder offload path %s.%s is not a streamable block list", name, block_path)
+                continue
+            group_hooks = [
+                apply_block_hook(blocks[-1], blocks[0], self.device, copy_stream, self.config.pin_cpu_memory)
+            ]
+            group_hooks.extend(
+                apply_block_hook(block, blocks[index + 1], self.device, copy_stream, self.config.pin_cpu_memory)
+                for index, block in enumerate(blocks[:-1])
+            )
+            for index, hook in enumerate(group_hooks):
+                hook._prev_hook = group_hooks[index - 1]
+            hooks.extend(group_hooks)
+            block_groups.append(blocks)
+
+        if not hooks:
+            return False
+        # The component lifecycle uses these generic attributes to keep only
+        # non-block encoder state resident during the encode phase.
+        module._omni_layerwise_hooks = hooks
+        module._omni_layerwise_block_groups = block_groups
+        module._omni_layerwise_enabled = True
+        logger.info(
+            "Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)",
+            name,
+            sum(len(blocks) for blocks in block_groups),
+            len(block_groups),
+        )
+        return True
+
     def _try_layerwise_offload_submodule(self, module: nn.Module, name: str, plan: OffloadPlan | None = None) -> bool:
         """Try to apply layerwise offload to a large submodule's blocks.
-
         Resolution order:
         1. OffloadPlan.offload_submodules (declarative, if plan is provided)
         2. Heuristic search for common block-list attributes
@@ -1089,6 +1302,78 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._blocks.append(blocks)
         return True
 
+    def _prepare_dit_non_block_modules(
+        self,
+        dit_module: nn.Module,
+        dit_name: str,
+        blocks_attr_names: list[str],
+        all_dit_modules: set[int],
+        plan: OffloadPlan | None,
+    ) -> None:
+        """Place or hook the DiT parts that are outside its repeated blocks.
+
+        This must run even when every repeated block is resident.  Otherwise
+        an all-resident stage skips placement for modules such as H3's token
+        refiner and enters the forward pass with CPU or meta tensors.
+        """
+        _ON_DEMAND_THRESHOLD = _ON_DEMAND_THRESHOLD_MB
+        for name, module in dit_module.named_children():
+            if name in blocks_attr_names:
+                logger.debug("Skipped blocks module %s", name)
+                continue
+
+            has_meta = any(getattr(param, "is_meta", False) for param in module.parameters())
+            if has_meta:
+                self._load_module_weights_from_mmap(module, dit_name, name)
+            module_mb = (
+                sum(
+                    param.nelement() * param.element_size() if not getattr(param, "is_meta", False) else 0
+                    for param in module.parameters()
+                )
+                / 1048576
+            )
+            explicitly_planned = plan is not None and name in plan.offload_submodules
+            if explicitly_planned or module_mb > _ON_DEMAND_THRESHOLD:
+                if id(module) in all_dit_modules:
+                    logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name)
+                elif self._try_layerwise_offload_submodule(module, name, plan):
+                    pass
+                else:
+                    self._register_on_demand_hook(module, name)
+                continue
+
+            try:
+                module.to(self.device)
+            except (NotImplementedError, RuntimeError):
+                self._load_module_weights_from_mmap(module, dit_name, name)
+                # Non-persistent buffers such as RoPE frequencies do not
+                # exist in the checkpoint and must be reconstructed.
+                has_meta_buffer = any(getattr(buffer, "is_meta", False) for buffer in module.buffers(recurse=True))
+                if has_meta_buffer:
+                    saved_params = {
+                        param_name: param.data.clone()
+                        for param_name, param in module.named_parameters()
+                        if not getattr(param, "is_meta", False)
+                    }
+                    module.to_empty(device=self.device)
+                    for submodule in module.modules():
+                        if hasattr(submodule, "reset_parameters"):
+                            submodule.reset_parameters()
+                    for param_name, param in module.named_parameters():
+                        if param_name in saved_params:
+                            param.data.copy_(saved_params[param_name])
+                try:
+                    module.to(self.device)
+                except Exception:
+                    logger.warning("Module %s still has meta params after mmap load", name)
+
+        for param in dit_module._parameters.values():
+            if param is not None and not getattr(param, "is_meta", False):
+                param.data = param.data.to(self.device, non_blocking=True)
+        for buffer in dit_module._buffers.values():
+            if buffer is not None:
+                buffer.data = buffer.data.to(self.device, non_blocking=True)
+
     def enable(self, pipeline: nn.Module) -> None:
         if self.enabled:
             logger.warning("DistributedLayerwiseOffloadBackend already enabled")
@@ -1110,7 +1395,19 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # When present, replaces heuristic block discovery.
         plan = get_offload_plan(pipeline)
 
+        if self.config.dlo_resident_layers and (plan is None or not plan.resident_dit_paths):
+            logger.warning(
+                "dlo_resident_layers=%d was requested, but this model declares no "
+                "resident_dit_paths; all blocks will be streamed.",
+                self.config.dlo_resident_layers,
+            )
+
         # Load weights via mmap for DLO+AllGather.
+        # TODO(offload): Add a rank-local mmap path for dlo_no_use_allgather.
+        # It must apply declarative checkpoint-to-runtime adapters before the
+        # standard TP loader shards each block, and must cover staged encoders
+        # and VAEs as well as DiT blocks. The current AllGather mmap path is
+        # deliberately not reused because it changes the weight layout.
         # Gate condition MUST match diffusers_loader.py:
         #   supports_mmap_loading(pipeline) and not _has_online_quant
         # When the gate is False, the loader has already loaded weights via
@@ -1137,10 +1434,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9)
         # during the DiT forward pass.  They are only needed briefly for
         # text-encoding (before DiT) and VAE-decode (after DiT).
-        for enc in modules.encoders:
-            self._register_on_demand_hook(enc, "encoder")
-        for vae in modules.vaes:
-            self._register_on_demand_hook(vae, "vae")
+        for enc, enc_name in zip(modules.encoders, modules.encoder_names):
+            self._try_layerwise_offload_encoder(enc, enc_name, plan)
+            self._register_on_demand_hook(
+                enc, "encoder", stage_on_demand=plan is not None and enc_name in plan.on_demand_component_paths
+            )
+        for vae, vae_name in zip(modules.vaes, modules.vae_names):
+            self._register_on_demand_hook(
+                vae,
+                "vae",
+                stage_on_demand=(plan is not None and vae_name in plan.on_demand_component_paths),
+            )
 
         # Move resident modules to GPU (small modules needed every forward)
         for name, module in zip(modules.resident_names, modules.resident_modules):
@@ -1171,93 +1475,41 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 dit_module.to(self.device)
                 continue
 
+            resident_count = 0
+            if plan is not None and dit_name in plan.resident_dit_paths:
+                resident_count = min(self.config.dlo_resident_layers, len(blocks))
+            if resident_count:
+                resident_blocks = blocks[:resident_count]
+                self._resident_blocks.extend(resident_blocks)
+                blocks = blocks[resident_count:]
+                logger.info(
+                    "Keeping %d leading blocks resident on %s; streaming %d tail blocks",
+                    resident_count,
+                    dit_name,
+                    len(blocks),
+                )
+
+            self._prepare_dit_non_block_modules(
+                dit_module,
+                dit_name,
+                blocks_attr_names,
+                all_dit_modules,
+                plan,
+            )
+
             num_blocks = len(blocks)
+            if num_blocks == 0:
+                logger.info("All blocks for %s are resident; no streaming hooks required", dit_name)
+                continue
             if num_blocks <= 1:
                 logger.warning(
-                    "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)",
+                    "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)",
                     dit_name,
                     dit_module.__class__.__name__,
                 )
-                dit_module.to(self.device)
+                self._resident_blocks.extend(blocks)
                 continue
 
-            # Move non-block modules to GPU (they stay resident).
-            # Large modules (> 1 GB) use on-demand hooks instead — they are
-            # only needed briefly (e.g. language_model for text encoding)
-            # and keeping them resident wastes HBM during the DiT loop.
-            _ON_DEMAND_THRESHOLD = _ON_DEMAND_THRESHOLD_MB  # alias for readability
-            for name, m in dit_module.named_children():
-                if name not in blocks_attr_names:
-                    _has_meta = any((p.is_meta if hasattr(p, "is_meta") else False) for p in m.parameters())
-                    if _has_meta:
-                        # Meta params: load from safetensors before moving to GPU
-                        self._load_module_weights_from_mmap(m, dit_name, name)
-                    _mb = (
-                        sum(
-                            (p.nelement() * p.element_size() if not (hasattr(p, "is_meta") and p.is_meta) else 0)
-                            for p in m.parameters()
-                        )
-                        / 1048576
-                    )
-                    if _mb > _ON_DEMAND_THRESHOLD:
-                        # Skip if this submodule is already a separate DiT
-                        # module (e.g. transformer.language_model when both
-                        # "transformer.language_model" and "transformer" are
-                        # in _dit_modules) — hooks already applied.
-                        if id(m) in all_dit_modules:
-                            logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name)
-                        elif self._try_layerwise_offload_submodule(m, name, plan):
-                            pass  # layerwise hooks applied
-                        else:
-                            self._register_on_demand_hook(m, name)
-                    else:
-                        try:
-                            m.to(self.device)
-                        except (NotImplementedError, RuntimeError):
-                            # Meta params: load from safetensors via mmap, then move to GPU
-                            self._load_module_weights_from_mmap(m, dit_name, name)
-                            # Materialize any remaining meta buffers (e.g. RoPE
-                            # inv_freq, timestep freqs) on device and
-                            # recompute them from their original formulas.
-                            # These are non-persistent buffers not present in
-                            # the checkpoint — they must be reconstructed, not
-                            # left as torch.empty.
-                            # NOTE: Do NOT call reset_parameters() here — it
-                            # would re-initialize ALL parameters (including
-                            # mmap-loaded weights) with random values.
-                            _has_meta_buf = any(hasattr(b, "is_meta") and b.is_meta for b in m.buffers(recurse=True))
-                            if _has_meta_buf:
-                                # Save mmap-loaded param values before
-                                # reset_parameters (which re-initializes
-                                # ALL params with random values).
-                                saved_params = {
-                                    n: p.data.clone()
-                                    for n, p in m.named_parameters()
-                                    if not (hasattr(p, "is_meta") and p.is_meta)
-                                }
-                                m.to_empty(device=self.device)
-                                for submod in m.modules():
-                                    if hasattr(submod, "reset_parameters"):
-                                        submod.reset_parameters()
-                                # Restore mmap-loaded weights
-                                for n, p in m.named_parameters():
-                                    if n in saved_params:
-                                        p.data.copy_(saved_params[n])
-                            try:
-                                m.to(self.device)
-                            except Exception:
-                                logger.warning("Module %s still has meta params after mmap load", name)
-                else:
-                    logger.debug(f"Skipped blocks module {name}")
-
-            # Move top-level params/buffers to GPU
-            for param in dit_module._parameters.values():
-                if param is not None and not (hasattr(param, "is_meta") and param.is_meta):
-                    param.data = param.data.to(self.device, non_blocking=True)
-            for buffer in dit_module._buffers.values():
-                if buffer is not None:
-                    buffer.data = buffer.data.to(self.device, non_blocking=True)
-
             # Register hooks in a circular sliding window:
             # last block prefetches first block, block i prefetches block (i+1)
             # All hooks share 2 global device buffers (RFC: "exactly two layers on device")
@@ -1314,7 +1566,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             self._all_hook_groups.append(block_hooks)
             self._blocks.append(blocks)
 
+        if self._resident_blocks:
+            self._resident_layer_group = PinnedResidentLayerGroup(
+                self._resident_blocks,
+                self.device,
+                self.copy_stream,
+                self.config.pin_cpu_memory,
+            )
+            pipeline._dlo_residency_controller = self
+
         if not self._all_hook_groups:
+            self.enabled = bool(self._resident_blocks)
             return
 
         # Unified allocation: 2 shared output buffers + 2 shared shard buffers
@@ -1424,8 +1686,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._on_demand_handles = []
         self._on_demand_shard_infos = []
 
+        self.offload_resident_layers()
         self._blocks.clear()
         self._all_hook_groups.clear()
+        self._resident_blocks.clear()
+        self._resident_layer_group = None
         self.enabled = False
         logger.info("Distributed layer-wise offloading disabled")
 

vllm_omni/diffusion/offloader/module_collector.py

@@ -19,6 +19,7 @@ class PipelineModules:
     encoders: list[nn.Module]
     encoder_names: list[str]
     vaes: list[nn.Module]
+    vae_names: list[str] = field(default_factory=list)
     resident_modules: list[nn.Module] = field(default_factory=list)
     resident_names: list[str] = field(default_factory=list)
 
@@ -153,7 +154,7 @@ class ModuleDiscovery:
 
         dit_modules, dit_names = ModuleDiscovery._collect_modules(pipeline, dit_attrs, warn_missing=declared)
         encoders, encoder_names = ModuleDiscovery._collect_modules(pipeline, enc_attrs, warn_missing=declared)
-        vaes, _ = ModuleDiscovery._collect_modules(pipeline, vae_attrs, warn_missing=declared)
+        vaes, vae_names = ModuleDiscovery._collect_modules(pipeline, vae_attrs, warn_missing=declared)
         residents, resident_names = ModuleDiscovery._collect_modules(pipeline, res_attrs, warn_missing=declared)
 
         return PipelineModules(
@@ -163,5 +164,6 @@ class ModuleDiscovery:
             encoder_names=encoder_names,
             vaes=vaes,
             resident_modules=residents,
+            vae_names=vae_names,
             resident_names=resident_names,
         )

vllm_omni/diffusion/offloader/module_residency.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""On-demand module staging backed by an immutable pinned CPU snapshot."""
+
+from __future__ import annotations
+
+from itertools import chain
+from typing import Any
+
+import torch
+from torch import nn
+
+from vllm_omni.platforms import current_omni_platform
+
+from .tensor_utils import set_tensor_storage
+
+
+class PinnedModuleStager:
+    """Stage an immutable module without copying device weights back to CPU.
+
+    ``nn.Module.to("cpu")`` performs a device-to-host copy for every parameter
+    after each forward. In inference the weights are immutable, so retain one
+    pinned CPU master instead. ``load`` materializes fresh device tensors from
+    that master; ``offload`` only rebinds the Parameters/Buffers to it and
+    releases the device allocations.
+    """
+
+    def __init__(
+        self,
+        module: nn.Module,
+        device: torch.device,
+        *,
+        pin_memory: bool = True,
+        copy_stream: Any | None = None,
+    ) -> None:
+        self.device = device
+        self.copy_stream = copy_stream or current_omni_platform.Stream()
+        self.loaded = False
+        self._entries: list[tuple[torch.Tensor, torch.Tensor]] = []
+        self._device_tensors: list[torch.Tensor] = []
+
+        # parameters()/buffers() remove duplicate objects by default, which
+        # preserves tied weights and buffer aliases when their storage moves.
+        for target in chain(module.parameters(), module.buffers()):
+            master = target.detach()
+            if master.device.type != "cpu":
+                master = master.to("cpu")
+            if pin_memory and not master.is_pinned():
+                master = master.pin_memory()
+            set_tensor_storage(target, master)
+            self._entries.append((target, master))
+
+    def load(self) -> None:
+        if self.loaded:
+            return
+
+        # Allocate on the compute stream so the caching allocator can reuse
+        # memory released by the preceding encoder/DiT stage.
+        device_tensors = [torch.empty_like(master, device=self.device) for _, master in self._entries]
+        compute_stream = current_omni_platform.current_stream()
+        self.copy_stream.wait_stream(compute_stream)
+        ready = current_omni_platform.Event()
+        with current_omni_platform.stream(self.copy_stream):
+            for device_tensor, (_, master) in zip(device_tensors, self._entries):
+                device_tensor.copy_(master, non_blocking=master.is_pinned())
+            ready.record(self.copy_stream)
+
+        for device_tensor, (target, _) in zip(device_tensors, self._entries):
+            set_tensor_storage(target, device_tensor)
+        compute_stream.wait_event(ready)
+        self._device_tensors = device_tensors
+        self.loaded = True
+
+    def offload(self) -> None:
+        if not self.loaded:
+            return
+
+        # The module has completed on the compute stream. Synchronize once at
+        # the stage boundary, then release HBM by rebinding to the unchanged
+        # host masters. No device-to-host transfer is necessary.
+        current_omni_platform.synchronize()
+        for target, master in self._entries:
+            set_tensor_storage(target, master)
+        self._device_tensors.clear()
+        self.loaded = False
+        current_omni_platform.empty_cache()
+
+
+__all__ = ["PinnedModuleStager"]

vllm_omni/diffusion/offloader/offload_plan.py

@@ -35,10 +35,21 @@ class OffloadPlan:
             for large non-DiT submodules within a DiT that should be
             independently offloaded with their own hooks.
             e.g. ``{"context_encoder": "layers"}``
+        resident_dit_paths: DiT paths whose leading blocks may be kept on the
+            device when ``dlo_resident_layers`` is nonzero. Keeping this
+            model-declared avoids applying a consumer-GPU tuning knob to
+            auxiliary or dual DiTs unintentionally.
+        encoder_block_attrs: Maps encoder paths to rank-local block-list paths.
+            These blocks are streamed with ordinary layerwise hooks, never
+            with the DiT AllGather group.
     """
 
+    on_demand_component_paths: frozenset[str] = field(default_factory=frozenset)
+
     block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict)
     offload_submodules: dict[str, str] = field(default_factory=dict)
+    resident_dit_paths: frozenset[str] = field(default_factory=frozenset)
+    encoder_block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict)
 
 
 def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None:
@@ -58,4 +69,6 @@ def supports_mmap_loading(pipeline: nn.Module) -> bool:
     (to decide whether to call ``_load_weights_via_mmap``), ensuring both
     paths use the same gate condition.
     """
-    return any(callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules())
+    return bool(getattr(pipeline, "_supports_mmap_loading", True)) and any(
+        callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules()
+    )