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>

docs/contributing/ci/test_examples/l4_performance_tests.inc.md

@@ -200,4 +200,4 @@ You can add any benchmark running parameters you need here. For all optional par
 | num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
 | request_rate    | float / array | No  | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
 | max_concurrency | int / array | No       | 1, [1, 2, 3]    | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode).                                                                                                                                                                                                             |
-| baseline        | object      | No       | see below       | Optional thresholds. Prefer hardware-nested form: `{"H100": {"throughput_qps": …}, "A3": {…}}` (selected via runtime hardware label). Legacy flat `{metric: threshold}` still works. Per metric: scalar, list (sweep step), or object keyed by concurrency / QPS.  
+| baseline        | object      | No       | see below       | Optional thresholds. Prefer hardware-nested form: `{"H100": {"throughput_qps": …}, "A3": {…}}` (selected via runtime hardware label). Legacy flat `{metric: threshold}` still works. Per metric: scalar, list (sweep step), or object keyed by concurrency / QPS.

recipes/MiniMaxAI/MiniMax-H3-5090.md

@@ -0,0 +1,61 @@
+# MiniMax-H3 on RTX 5090
+
+This recipe uses BF16 weights, tiled VAE decode, tensor parallelism where a
+second GPU is available, and distributed layerwise offload (DLO). It is a
+memory-first serving configuration; lower resident counts reduce HBM use and
+increase CPU-to-GPU transfer time.
+
+## Capacity requirements
+
+| Resource | One RTX 5090 | Two RTX 5090s |
+| --- | ---: | ---: |
+| GPU HBM | 32 GiB | 32 GiB per GPU |
+| Checkpoint storage | 135 GiB per partition | 135 GiB per partition |
+| Available system RAM | 200 GiB minimum | 200 GiB minimum |
+| Recommended system RAM | 384 GiB | 384 GiB |
+
+`FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one
+server at a time. DLO keeps rank-local weights in pinned host memory; increasing
+`--dlo-resident-layers` improves latency but does **not** reduce host RAM in the
+current implementation because resident layers retain pinned CPU master copies.
+
+> **Modular H3:** after #5720 lands, preserve this recipe's one-partition
+> behavior with `--task-type fl2va` or `--task-type ref2va`.
+
+## One RTX 5090: 1344x768, 5 seconds
+
+Use 12 resident DiT layers. A 50-step B300 allocation test with this exact
+single-rank topology peaked at 26.50 GiB; re-measure peak HBM on the target
+card before increasing the resident count.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 vllm serve /path/to/MiniMax-H3/FL2VA \
+  --omni --trust-remote-code --host 0.0.0.0 --port 8000 \
+  --num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
+  --usp 1 --ring 1 --vae-patch-parallel-size 1 \
+  --vae-parallel-mode tile --vae-use-tiling \
+  --enable-distributed-layerwise-offload --dlo-no-use-allgather \
+  --dlo-resident-layers 12 --enforce-eager \
+  --diffusion-attention-backend CUDNN_ATTN
+```
+
+## Two RTX 5090s: 1344x768, 5 seconds
+
+Use TP2 and 20 resident DiT layers. The two-rank B300 capacity run peaked at
+27,726 MiB per rank for this shape and 50 steps. This is a memory/correctness
+proxy, not a consumer-GPU latency claim.
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/MiniMax-H3/FL2VA \
+  --omni --trust-remote-code --host 0.0.0.0 --port 8000 \
+  --num-gpus 2 --tensor-parallel-size 2 --text-encoder-tp-size 2 \
+  --usp 1 --ring 1 --vae-patch-parallel-size 2 \
+  --vae-parallel-mode tile --vae-use-tiling \
+  --enable-distributed-layerwise-offload --dlo-no-use-allgather \
+  --dlo-resident-layers 20 --enforce-eager \
+  --diffusion-attention-backend CUDNN_ATTN
+```
+
+For Ref2VA, stop the FL2VA server and restart the same command with
+`/path/to/MiniMax-H3/Ref2VA`. Ref2VA reference video count and prompt length
+can increase activation memory; begin with one request at a time.

recipes/MiniMaxAI/MiniMax-H3.md

@@ -33,8 +33,10 @@ export MODEL_ROOT=/path/to/MiniMax-H3
 hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
 ```
 
-Install vLLM-Omni from the checkout containing MiniMax H3 support. On
-Blackwell, install the optional FlashAttention-4 dependency:
+Install vLLM-Omni from the checkout containing MiniMax H3 support. The
+two-GPU RTX 5090/4090 profiles use cuDNN attention and do not need
+FlashAttention-4. Install the optional dependency only for the four-GPU
+B300/GB200 `FLASH_ATTN` profile:
 
 ```bash
 uv venv
@@ -50,6 +52,26 @@ reference-video preparation and MP4 output.
 One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
 and FL2VA requests, or to `Ref2VA` for either Ref2VA request.
 
+### Memory and storage requirements
+
+Treat GPU HBM, host RAM, and checkpoint storage as separate requirements. Each
+H3 checkpoint partition (`FL2VA` or `Ref2VA`) contains about **134 GiB** of
+BF16 safetensors (about **135 GiB** on disk). Keeping both partitions locally
+therefore needs roughly **270 GiB** of model storage, although the examples
+serve only one partition at a time.
+
+CPU offload and distributed layerwise offload reduce GPU residency; they do
+not make the model weights disappear. With `--dlo-no-use-allgather`, each
+worker retains its standard-loader rank-local weights in host memory, including
+pinned CPU buffers used for H2D streaming. Use at least **200 GiB available
+system RAM** before starting the two-GPU recipe; a **384 GiB host is
+recommended** to leave room for the OS, CUDA/PyTorch allocations, request
+inputs, and filesystem cache. Do not run the FL2VA and Ref2VA servers at the
+same time on a host sized for this minimum.
+
+The consumer-GPU profiles below are HBM budgets only. They still require the
+host-RAM budget above.
+
 ### Single GPU: accuracy and memory first
 
 The single-GPU configuration uses model-level CPU offload.
@@ -77,6 +99,98 @@ Use a GPU with enough memory for the active H3 component and enough system RAM
 for the offloaded components. CPU offload reduces GPU memory pressure but adds
 PCIe/NVLink transfer latency.
 
+### Two 24/32 GB GPUs: TP2 distributed layerwise offload
+
+For two PCIe consumer GPUs, combine TP2 with distributed layerwise offload
+(DLO). The standard loader first creates the rank-local TP shard. DLO keeps
+that shard in pinned host memory and streams the 30 tail DiT blocks through a
+shared two-buffer window without a DP AllGather. The first 20 DiT blocks are
+copied to the GPUs once per denoise stage, reused by every sampling step, and
+released before VAE decode so the decoder can reuse their HBM.
+
+```bash
+export MODEL="${MODEL_ROOT}/FL2VA"
+export PORT=8091
+
+CUDA_VISIBLE_DEVICES=0,1 \
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
+vllm serve "${MODEL}" \
+  --omni \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --trust-remote-code \
+  --num-gpus 2 \
+  --tensor-parallel-size 2 \
+  --usp 1 \
+  --ring 1 \
+  --text-encoder-tp-size 2 \
+  --vae-patch-parallel-size 2 \
+  --vae-parallel-mode tile \
+  --vae-use-tiling \
+  --enable-distributed-layerwise-offload \
+  --dlo-no-use-allgather \
+  --dlo-resident-layers 20 \
+  --enforce-eager \
+  --diffusion-attention-backend CUDNN_ATTN
+```
+
+Use the profile that matches the per-GPU memory capacity:
+
+| Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status |
+|---|---:|---:|---:|---|---|---|
+| `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated |
+| `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point |
+
+This topology uses all available parallel capacity: TP2 shards both the DiT
+and text encoder, `--dlo-no-use-allgather` streams each rank's local TP shard
+without reconstructing full blocks, and VAE patch parallelism splits tiled
+decode across both GPUs. cuDNN attention is selected explicitly for the RTX
+consumer path; the server stays eager to avoid an unqualified compile path.
+
+The resident count changes placement and transfer frequency only; it does not
+quantize or change the BF16/FP32 denoise math. Re-measure peak memory before
+increasing it on a different request shape.
+
+### RTX 5090 target-hardware validation
+
+At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on
+2 x RTX 5090 without OOM:
+
+| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation |
+|---:|---:|---:|---:|---|
+| 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed |
+
+This is a single end-to-end validation run, not a warmed multi-run latency
+benchmark. The sampled `nvidia-smi` peak is also not a CUDA allocator
+high-water mark. The environment used vLLM 0.26.0, vLLM-Omni
+`0.26.1.dev14+gae6577ea`, and PyTorch 2.11.0+cu130. The
+[run record](https://github.com/lishunyang12/vllm-omni-rankings/blob/dcd06d7e83cb069842535918c0169ee9f3f29ba0/scripts/%E5%BE%AE%E4%BF%A1%E5%9B%BE%E7%89%87_20260805000034_86_237.png)
+captures the environment, output contract, elapsed time, and sampled peak.
+
+Before the target run, both profiles were exercised on two B300 ranks as an
+allocation and correctness proxy. At 1344x768, 124 frames, and 50 steps, the
+20-layer profile peaked at 27,726 MiB per rank. At 1024x576, the 12-layer
+profile peaked at 18,888 MiB per rank in a 5-step capacity run. The resident
+and fully streamed placements produced identical decoded video-frame and audio
+hashes for the same shape, step count, prompt, and seed. The B300 result does
+not establish RTX 4090 PCIe latency; treat the 4090 profile as a conservative
+starting point until it is measured on that GPU.
+
+To run T2VA, FL2VA, image+audio Ref2VA, and two-video Ref2VA in order, validate
+every MP4's H.264/AAC streams, and retain live server and GPU-memory logs:
+
+```bash
+RUN_ROOT=/path/to/run-root \
+MODEL_ROOT=/path/to/MiniMax-H3 \
+GPU_IDS=0,1 \
+PROFILE=rtx5090 \
+bash examples/offline_inference/minimax_h3/run_h3_2gpu_all_tasks.sh
+```
+
+The script selects 20 resident layers for `PROFILE=rtx5090` and 12 for
+`PROFILE=rtx4090`; `DLO_RESIDENT_LAYERS=N` overrides either default.
+
 ### Four GPUs: measured best-practice throughput
 
 The best validated four-GPU configuration on four NVIDIA B300 GPUs is:

tests/diffusion/attention/test_cudnn_attn.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import nullcontext
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
+    observed = {}
+
+    def fake_sdpa(query, key, value, **kwargs):
+        observed.update(query=query, key=key, value=value, kwargs=kwargs)
+        return query
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
+        lambda _backends: nullcontext(),
+    )
+    monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
+    impl = CuDNNAttentionImpl(
+        num_heads=2,
+        head_size=4,
+        softmax_scale=0.5,
+    )
+    query = torch.randn(1, 8, 2, 4)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+
+    output = impl.forward_cuda(
+        query,
+        key,
+        value,
+        AttentionMetadata(extra={"valid_kv_length": 5}),
+    )
+
+    assert output.shape == query.shape
+    assert observed["query"].shape == (1, 2, 8, 4)
+    assert observed["key"].shape == (1, 2, 5, 4)
+    assert observed["value"].shape == (1, 2, 5, 4)
+    assert observed["kwargs"]["attn_mask"] is None
+
+
+def test_cudnn_rejects_invalid_valid_kv_length():
+    impl = CuDNNAttentionImpl(
+        num_heads=2,
+        head_size=4,
+        softmax_scale=0.5,
+    )
+    query = torch.randn(1, 8, 2, 4)
+
+    with pytest.raises(ValueError, match="valid_kv_length"):
+        impl.forward_cuda(
+            query,
+            query,
+            query,
+            AttentionMetadata(extra={"valid_kv_length": 9}),
+        )

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -30,6 +30,13 @@ def test_pipeline_import_registry_and_component_discovery():
     assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
 
 
+def test_dlo_offload_plan_includes_token_refiner():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    assert MiniMaxH3Pipeline._offload_plan.offload_submodules == {"token_refiner": "blocks"}
+    assert MiniMaxH3Pipeline._offload_plan.resident_dit_paths == frozenset({"transformer"})
+
+
 def test_joint_postprocess_is_multiprocessing_picklable():
     from vllm_omni.diffusion.models.minimax_h3 import (
         get_minimax_h3_post_process_func,
@@ -87,6 +94,87 @@ def test_shifted_sigma_schedule_matches_reference_values():
     )
 
 
+def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3Attention,
+    )
+
+    class FakeBackend:
+        supports_prefix_kv_slicing = True
+
+    class FakeAttention(torch.nn.Module):
+        attn_backend = FakeBackend
+
+        def __init__(self):
+            super().__init__()
+            self.metadata = None
+
+        def forward(self, query, key, value, metadata):
+            self.metadata = metadata
+            return query
+
+    attention = object.__new__(MiniMaxH3Attention)
+    torch.nn.Module.__init__(attention)
+    attention.attention = FakeAttention()
+    # Model-side Ulysses shards rows before Attention gathers them again.
+    # The Python packed_total must therefore remain global even though q is
+    # local here.
+    q = torch.randn(4, 2, 4)
+    cu_seqlens = torch.tensor([0, 5, 8], dtype=torch.int32)
+
+    output = attention._run_packed_attention(
+        q,
+        q,
+        q,
+        cu_seqlens=cu_seqlens,
+        max_seqlen=5,
+        packed_total=8,
+    )
+
+    assert output.shape == q.shape
+    assert attention.attention.metadata.attn_mask is None
+    assert attention.attention.metadata.extra["valid_kv_length"] == 5
+
+
+def test_packed_attention_keeps_padding_mask_for_other_backends():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3Attention,
+    )
+
+    class FakeBackend:
+        supports_prefix_kv_slicing = False
+
+    class FakeAttention(torch.nn.Module):
+        attn_backend = FakeBackend
+
+        def __init__(self):
+            super().__init__()
+            self.metadata = None
+
+        def forward(self, query, key, value, metadata):
+            self.metadata = metadata
+            return query
+
+    attention = object.__new__(MiniMaxH3Attention)
+    torch.nn.Module.__init__(attention)
+    attention.attention = FakeAttention()
+    q = torch.randn(8, 2, 4)
+
+    attention._run_packed_attention(
+        q,
+        q,
+        q,
+        cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
+        max_seqlen=5,
+        packed_total=8,
+    )
+
+    assert torch.equal(
+        attention.attention.metadata.attn_mask,
+        torch.tensor([[True, True, True, True, True, False, False, False]]),
+    )
+
+
 def test_reference_image_resize_contract():
     from PIL import Image
 
@@ -293,6 +381,127 @@ def test_layerwise_offload_releases_text_encoder():
     pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
 
 
+def test_distributed_layerwise_offload_releases_text_encoder():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=True,
+    )
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+
+
+def test_distributed_layerwise_offload_stages_vae_component():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=True,
+    )
+    component = Mock()
+
+    with pipeline._component_on_device(component):
+        component.load_to_device.assert_called_once_with()
+        component.offload_to_cpu.assert_not_called()
+
+    component.offload_to_cpu.assert_called_once_with()
+
+
+def test_distributed_layerwise_resident_blocks_are_stage_scoped():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    controller = Mock()
+    pipeline._dlo_residency_controller = controller
+
+    with pipeline._resident_dit_layers_on_device():
+        controller.load_resident_layers.assert_called_once_with()
+        controller.offload_resident_layers.assert_not_called()
+
+    controller.offload_resident_layers.assert_called_once_with()
+
+
+def test_distributed_layerwise_resident_blocks_release_on_failure():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    controller = Mock()
+    pipeline._dlo_residency_controller = controller
+
+    with pytest.raises(RuntimeError, match="denoise failed"):
+        with pipeline._resident_dit_layers_on_device():
+            raise RuntimeError("denoise failed")
+
+    controller.load_resident_layers.assert_called_once_with()
+    controller.offload_resident_layers.assert_called_once_with()
+
+
+def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.encoder import (
+        MiniMaxH3Qwen3VLEncoder,
+    )
+
+    class Stack(torch.nn.Module):
+        def __init__(self, count):
+            super().__init__()
+            self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
+
+    class TextStack(torch.nn.Module):
+        def __init__(self, count):
+            super().__init__()
+            self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
+
+    hooks = []
+
+    def fake_apply(block, next_block, device, stream, pin_memory):
+        hook = SimpleNamespace(
+            block=block,
+            next_block=next_block,
+            device=device,
+            stream=stream,
+            pin_memory=pin_memory,
+            _prev_hook=None,
+            offload_layer=Mock(),
+        )
+        hooks.append(hook)
+        return hook
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.offloader.layerwise_backend.apply_block_hook",
+        fake_apply,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.platforms.current_omni_platform.Stream",
+        Mock(return_value="copy-stream"),
+    )
+    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
+    torch.nn.Module.__init__(encoder)
+    encoder.device_target = torch.device("cpu")
+    encoder.vision = Stack(2)
+    encoder.text_model = TextStack(3)
+
+    encoder.enable_omni_layerwise_offload(pin_memory=False)
+
+    assert len(hooks) == 5
+    assert all(hook._prev_hook is not None for hook in hooks)
+    assert all(hook.device == torch.device("cpu") for hook in hooks)
+
+
 def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
 
@@ -324,6 +533,82 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
     assert next(video_vae.parameters()).dtype == torch.float32
 
 
+def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
+
+    class FakeRemote(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = torch.nn.Linear(1, 1)
+
+    monkeypatch.setattr(
+        vae_module,
+        "_load_component_config",
+        lambda _path: {
+            "latent_channels": 1,
+            "latents_mean": [0.0],
+            "latents_std": [1.0],
+        },
+    )
+    monkeypatch.setattr(
+        vae_module,
+        "_load_remote_component",
+        lambda _path, _config: FakeRemote(),
+    )
+
+    video_vae = vae_module.MiniMaxH3VideoVAE(
+        "unused",
+        device=torch.device("meta"),
+        load_device=torch.device("cpu"),
+    )
+
+    assert next(video_vae.parameters()).device.type == "cpu"
+    assert video_vae._device_target.type == "meta"
+
+
+def test_video_vae_uses_snapshot_stager_for_accelerator_target(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
+
+    class FakeRemote(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = torch.nn.Linear(1, 1)
+
+    stager = Mock()
+    stager_cls = Mock(return_value=stager)
+    monkeypatch.setattr(
+        vae_module,
+        "_load_component_config",
+        lambda _path: {
+            "latent_channels": 1,
+            "latents_mean": [0.0],
+            "latents_std": [1.0],
+        },
+    )
+    monkeypatch.setattr(
+        vae_module,
+        "_load_remote_component",
+        lambda _path, _config: FakeRemote(),
+    )
+    monkeypatch.setattr(vae_module, "PinnedModuleStager", stager_cls)
+
+    video_vae = vae_module.MiniMaxH3VideoVAE(
+        "unused",
+        device=torch.device("cuda"),
+        load_device=torch.device("cpu"),
+    )
+    video_vae.load_to_device()
+    video_vae.offload_to_cpu()
+
+    stager_cls.assert_called_once_with(
+        video_vae.remote,
+        torch.device("cuda"),
+        pin_memory=True,
+    )
+    stager.load.assert_called_once_with()
+    stager.offload.assert_called_once_with()
+
+
 def test_video_vae_encode_uses_configured_parallel_tiling():
     from vllm_omni.diffusion.models.minimax_h3.vae import (
         MiniMaxH3VideoVAE,

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -27,7 +27,9 @@ from vllm_omni.diffusion.offloader.block_discovery import (
 from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
     DistributedLayerwiseOffloadBackend,
     DistributedLayerwiseOffloadHook,
+    PinnedResidentLayerGroup,
 )
+from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
 from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan
 from vllm_omni.platforms import current_omni_platform
 
@@ -97,6 +99,7 @@ def patched_offload_runtime(monkeypatch):
     monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent)
     monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
     monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream)
+    monkeypatch.setattr(dist_backend_module.current_omni_platform, "synchronize", lambda: None)
 
 
 class TinyBlock(nn.Module):
@@ -308,6 +311,164 @@ class TestDistributedLayerwiseOffloadHook:
         assert shard1.numel() == 2
         assert torch.equal(shard1, torch.tensor([102.0, 0.0]))
 
+    def test_mmap_layout_transform_is_applied_before_sharding(self, dist_group, patched_offload_runtime):
+        block = nn.Module()
+        block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
+        next_block = nn.Module()
+        next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
+        next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
+        next_block.weight.mmap_weight_transform_pending = True
+
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=2,
+            rank=0,
+            pin_memory=False,
+        )
+        hook.initialize_hook(block)
+
+        assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0]))
+
+    def test_regular_loader_weight_does_not_reapply_mmap_transform(self, dist_group, patched_offload_runtime):
+        block = nn.Module()
+        block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
+        next_block = nn.Module()
+        next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
+        next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
+
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=2,
+            rank=0,
+            pin_memory=False,
+        )
+        hook.initialize_hook(block)
+
+        assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0]))
+
+
+class TestPinnedResidentLayerGroup:
+    def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
+        blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
+        expected = []
+        for index, block in enumerate(blocks):
+            block.weight.data.fill_(index + 1)
+            block.bias.data.fill_(index + 3)
+            expected.append((block.weight.detach().clone(), block.bias.detach().clone()))
+
+        group = PinnedResidentLayerGroup(
+            blocks,
+            device=torch.device("cpu"),
+            copy_stream=DummyStream(),
+            pin_memory=False,
+        )
+
+        assert all(block.weight.numel() == 0 for block in blocks)
+
+        group.load()
+        for block, (weight, bias) in zip(blocks, expected):
+            assert torch.equal(block.weight, weight)
+            assert torch.equal(block.bias, bias)
+
+        # Device-side changes must not overwrite the persistent host master.
+        blocks[0].weight.data.zero_()
+        group.offload()
+        assert all(block.weight.numel() == 0 for block in blocks)
+
+        group.load()
+        for block, (weight, bias) in zip(blocks, expected):
+            assert torch.equal(block.weight, weight)
+            assert torch.equal(block.bias, bias)
+
+        group.offload()
+
+    def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
+        module = nn.Linear(2, 2, bias=False)
+
+        def transform(tensor):
+            return tensor.flip(0)
+
+        module.weight.mmap_weight_transform = transform
+
+        backend = object.__new__(DistributedLayerwiseOffloadBackend)
+        backend._remember_mmap_param_attrs(module)
+        module.to_empty(device="meta")
+
+        replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float())
+        backend._attach_mmap_param_attrs("weight", replacement, module.weight)
+
+        assert replacement.mmap_weight_transform is transform
+        assert replacement.mmap_weight_transform_pending is True
+
+    def test_shared_allgather_output_is_narrowed_to_current_block(
+        self, dist_group, patched_offload_runtime, monkeypatch
+    ):
+        block = nn.Module()
+        block.weight = nn.Parameter(torch.zeros(3, dtype=torch.float32))
+        next_block = nn.Module()
+        next_block.weight = nn.Parameter(torch.arange(3, dtype=torch.float32))
+
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=object(),
+            dp_size=2,
+            rank=0,
+            pin_memory=False,
+        )
+        hook.initialize_hook(block)
+        # Simulate unified buffers sized for a larger heterogeneous block.
+        hook.gpu_buffers = [
+            {torch.float32: torch.empty(12)},
+            {torch.float32: torch.empty(12)},
+        ]
+        hook.gpu_shard_buffers = [
+            {torch.float32: torch.empty(6)},
+            {torch.float32: torch.empty(6)},
+        ]
+        output_sizes = []
+
+        def fake_allgather(output, local_shard, *, group):
+            del group
+            output_sizes.append(output.numel())
+            output.zero_()
+            output[: local_shard.numel()].copy_(local_shard)
+
+        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
+
+        hook.prefetch_layer(slot=0, non_blocking=False)
+
+        assert output_sizes == [4]
+
+
+class TestPinnedModuleStager:
+    def test_offload_rebinds_immutable_cpu_master(self, patched_offload_runtime):
+        module = nn.Linear(2, 2)
+        expected_weight = module.weight.detach().clone()
+        expected_bias = module.bias.detach().clone()
+        stager = PinnedModuleStager(
+            module,
+            torch.device("cpu"),
+            pin_memory=False,
+            copy_stream=DummyStream(),
+        )
+
+        stager.load()
+        module.weight.data.zero_()
+        module.bias.data.zero_()
+        stager.offload()
+
+        assert torch.equal(module.weight, expected_weight)
+        assert torch.equal(module.bias, expected_bias)
+        stager.load()
+        assert torch.equal(module.weight, expected_weight)
+        assert torch.equal(module.bias, expected_bias)
+        stager.offload()
+
 
 class _DummyBlock(nn.Module):
     def __init__(self):
@@ -592,6 +753,7 @@ class TestOffloadPlan:
         plan = OffloadPlan(
             block_attrs={"transformer": ("gen_layers",)},
             offload_submodules={"context_encoder": "layers"},
+            resident_dit_paths=frozenset({"transformer"}),
         )
 
         class PipelineWithPlan(nn.Module):
@@ -602,12 +764,14 @@ class TestOffloadPlan:
         assert result is plan
         assert result.block_attrs == {"transformer": ("gen_layers",)}
         assert result.offload_submodules == {"context_encoder": "layers"}
+        assert result.resident_dit_paths == frozenset({"transformer"})
 
     def test_offload_plan_defaults_to_empty(self):
         """OffloadPlan with no arguments should have empty dicts."""
         plan = OffloadPlan()
         assert plan.block_attrs == {}
         assert plan.offload_submodules == {}
+        assert plan.resident_dit_paths == frozenset()
 
     def test_offload_plan_is_frozen(self):
         """OffloadPlan should be immutable (frozen=True)."""
@@ -615,6 +779,50 @@ class TestOffloadPlan:
         with pytest.raises(Exception):
             plan.block_attrs = {"x": ("y",)}  # type: ignore
 
+    def test_all_resident_dit_still_prepares_planned_submodules(self, patched_offload_runtime):
+        class TokenRefiner(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+
+        class Transformer(nn.Module):
+            _layerwise_offload_blocks_attrs = ["blocks"]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+                self.token_refiner = TokenRefiner()
+
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(
+                offload_submodules={"token_refiner": "blocks"},
+                resident_dit_paths=frozenset({"transformer"}),
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+        pipeline = Pipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+                dlo_resident_layers=2,
+            ),
+            torch.device("cpu"),
+        )
+
+        backend.enable(pipeline)
+
+        assert len(backend._resident_blocks) == 2
+        assert len(backend._all_hook_groups) == 1
+        assert len(backend._all_hook_groups[0]) == 2
+        assert backend.enabled
+
+        backend.disable()
+
 
 class TestMmapValidation:
     """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
@@ -821,7 +1029,7 @@ class TestConfigValidation:
             OffloadConfig.from_od_config(FakeODConfig())
 
     def test_hsdp_without_allgather_allowed(self):
-        """HSDP + DLO + no-AllGather should be allowed (full weights per rank)."""
+        """HSDP + DLO + no-AllGather uses standard-loader rank-local weights."""
         from vllm_omni.diffusion.offloader.base import OffloadConfig
 
         class FakePC:
@@ -843,12 +1051,33 @@ class TestConfigValidation:
             enable_layerwise_offload = False
             enable_distributed_layerwise_offload = True
             dlo_use_allgather = False  # no AllGather → should be allowed
+            dlo_resident_layers = 20
             pin_cpu_memory = True
             parallel_config = FakePC()
             model = "/fake/path"
 
         config = OffloadConfig.from_od_config(FakeODConfig())
         assert config.dp_size == 1  # forced to 1 when no AllGather
+        assert config.dlo_resident_layers == 20
+
+    def test_resident_layers_with_allgather_rejected(self):
+        class FakePC:
+            data_parallel_size = 2
+            use_hsdp = False
+            sequence_parallel_size = 1
+
+        class FakeODConfig:
+            enable_cpu_offload = False
+            enable_layerwise_offload = False
+            enable_distributed_layerwise_offload = True
+            dlo_use_allgather = True
+            dlo_resident_layers = 20
+            pin_cpu_memory = True
+            parallel_config = FakePC()
+            model = "/fake/path"
+
+        with pytest.raises(ValueError, match="requires --dlo-no-use-allgather"):
+            OffloadConfig.from_od_config(FakeODConfig())
 
     def test_num_inference_steps_none_rejected(self):
         """DP multi-concurrency should reject None num_inference_steps."""

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -247,6 +247,56 @@ class TestRequestModeDispatch:
         results = {ro.request_id: ro.result.error for ro in out.runner_outputs}
         assert results == {request_id: f"result_for_{request_id}" for request_id in request_ids}
 
+    def test_dlo_dp_routes_multiple_requests_without_pipeline_batching(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.execute_request = Mock(return_value="dlo-dp")
+        executor.collective_rpc = Mock()
+        scheduler_output = _make_sched_output("A", "B")
+
+        result = executor.execute_batch(scheduler_output)
+
+        assert result == "dlo-dp"
+        executor.execute_request.assert_called_once_with(scheduler_output)
+        executor.collective_rpc.assert_not_called()
+
+    def test_dlo_dp_multi_rank_reply_uses_synchronous_rpc_collection(self):
+        executor, req_q, res_q = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+        )
+        executor._sync_result_buffer = res_q
+
+        def worker():
+            request = req_q.get(timeout=2)
+            assert "rpc_id" not in request
+            wave_id = request["wave_id"]
+            for dp_rank in range(2):
+                res_q.put(
+                    {
+                        "dp_rank": dp_rank,
+                        "output": _tagged_output(str(dp_rank)),
+                        "wave_id": wave_id,
+                    }
+                )
+
+        thread = threading.Thread(target=worker, daemon=True)
+        thread.start()
+        results = executor.collective_rpc(
+            "execute_model",
+            unique_reply_rank=None,
+            exec_all_ranks=True,
+        )
+        thread.join(timeout=2)
+
+        assert [result.error for result in results] == ["0", "1"]
+
 
 # ───────────────── concurrent collective RPC ─────────────────
 

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -262,6 +262,36 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
     assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
 
 
+def test_serve_cli_forwards_distributed_offload_residency():
+    """Ensure the two-GPU DLO placement controls reach the diffusion stage."""
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="command")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(
+        [
+            "serve",
+            "MiniMaxAI/MiniMax-H3",
+            "--omni",
+            "--enable-distributed-layerwise-offload",
+            "--dlo-no-use-allgather",
+            "--dlo-resident-layers",
+            "20",
+        ]
+    )
+
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    engine_args = stage_cfg["engine_args"]
+
+    assert args.enable_distributed_layerwise_offload is True
+    assert args.dlo_use_allgather is False
+    assert args.dlo_resident_layers == 20
+    assert engine_args["enable_distributed_layerwise_offload"] is True
+    assert engine_args["dlo_use_allgather"] is False
+    assert engine_args["dlo_resident_layers"] == 20
+
+
 def test_serve_cli_accepts_diffusion_compile_controls():
     """Ensure both compile controls reach the diffusion stage."""
     parser = TrackingArgumentParser()

vllm_omni/config/omni_config.py

@@ -479,6 +479,7 @@ class _DiffusionConfigProjection:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_resident_layers: int = Field(default=0, ge=0)
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)

vllm_omni/config/stage_config.py

@@ -413,6 +413,9 @@ class StageDeployConfig:
     enable_cpu_offload: bool | None = None
     enable_layerwise_offload: bool | None = None
 
+    enable_distributed_layerwise_offload: bool | None = None
+    dlo_use_allgather: bool | None = None
+    dlo_resident_layers: int | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None
 

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

vllm_omni/engine/arg_utils.py

@@ -495,6 +495,7 @@ class OrchestratorArgs:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_resident_layers: int = 0
     boundary_ratio: float | None = None
     flow_shift: float | None = None
     diffusion_kv_cache_dtype: str | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1032,6 +1032,7 @@ class AsyncOmniEngine:
             "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
             "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
+            "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "diffusion_compile_granularity": (
                 "regional"

vllm_omni/entrypoints/cli/serve.py

@@ -646,15 +646,25 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             default=True,
             help="Use shard + AllGather for weight reconstruction (default: True). "
-            "When disabled (--dlo-no-use-allgather), each rank loads full weights "
-            "via H2D only — no sharding, no AllGather, no concurrent request "
-            "requirement, but N× CPU memory.",
+            "When disabled (--dlo-no-use-allgather), each rank streams the "
+            "standard loader's rank-local tensors via H2D only — no additional "
+            "DP sharding, no AllGather, and no concurrent-request requirement.",
         )
         omni_config_group.add_argument(
             "--dlo-no-use-allgather",
             dest="dlo_use_allgather",
             action="store_false",
-            help="Disable AllGather: each rank loads full weights independently.",
+            help=(
+                "Disable AllGather and stream standard-loader rank-local weights "
+                "independently (including existing TP shards)."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--dlo-resident-layers",
+            type=int,
+            default=0,
+            help="Keep this many leading main-DiT blocks resident on the device "
+            "while distributed layerwise offload streams the remaining blocks.",
         )
         # Video model parameters (e.g., Wan2.2) - engine-level
         omni_config_group.add_argument(