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/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,