[Diffusion] Document DLO mmap model opt-in

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/design/feature/distributed_layerwise_offload.md

@@ -5,6 +5,8 @@ models. DLO keeps only a small number of DiT blocks on the accelerator and
 streams the remaining blocks from host memory. The distributed backend can
 either shard those host-side weights across an existing parallel group or keep
 the standard loader's rank-local weights and avoid an additional collective.
+The rank-local mmap storage design is tracked in
+[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
 
 For user-facing commands, see the
 [distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
@@ -16,8 +18,10 @@ DLO is implemented for multi-device diffusion execution. The default
 AllGather path is the primary path for DP and SP deployments. The
 `--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
 for standard-loader sharding, workstation bring-up, and systems where an
-additional DLO collective is undesirable, but it does not reduce host weight
+additional DLO collective is undesirable. For explicitly mmap-compatible
-storage across ranks.
+models in a TP1, non-HSDP topology, this path shares file-backed DiT checkpoint
+pages across processes and keeps only bounded staging buffers private. TP,
+HSDP, and other models retain standard-loader rank-local host storage.
 
 The compatibility matrix below describes the current implementation. The
 unit-level guards are covered, but not every parallelism combination has a
@@ -49,6 +53,10 @@ block's shard is copied to a device buffer and reconstructed with
 `all_gather_into_tensor` on a communication stream while the current block is
 executing.
 
+Multi-rank AllGather behavior is unchanged. When the effective DLO group size
+is one, the backend now uses rank-local mmap staging; this closes a previous
+loader/backend mismatch that could leave parameters on the `meta` device.
+
 ```text
 Compute:    [Block N]             [Block N+1]          [Block N+2]
 H2D:                      [shard N+1]           [shard N+2]
@@ -68,9 +76,18 @@ requests must take the same execution path at every denoising step.
 ### Rank-local path without DLO AllGather
 
 With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
-one. The regular model loader remains responsible for preparing each rank's
+one and streams complete runtime blocks using H2D copies only. Host storage is
-weights, including TP-local tensors or HSDP-managed parameters. DLO then
+selected independently from that transfer protocol:
-streams those rank-local tensors block by block using H2D copies only.
+
+- For an explicitly mmap-compatible model with TP1 and no HSDP, each worker
+  retains immutable safetensors views backed by the node's shared OS page
+  cache. Two pinned host buffers, sized to the largest streamed block, stage
+  mmap pages for asynchronous H2D. Checkpoint-to-runtime adapters run while a
+  block is packed, so transformed weights do not become a persistent private
+  model copy.
+- With TP, HSDP, online quantization, or a model without mmap support, the
+  regular loader remains responsible for preparing rank-local tensors. This
+  preserves TP/HSDP/quantization loader callbacks.
 
 This mode means:
 
@@ -80,19 +97,26 @@ This mode means:
   not shard weights across SP ranks.
 - TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
   DLO's additional weight AllGather is disabled.
-- Pure DP deployments keep a full host-side model copy per rank, subject to
+- Pure DP deployments with a compatible mmap checkpoint share the streamed
-  shared OS page-cache behavior.
+  DiT's file-backed pages. Other components and fallback paths may still keep
+  rank-private host weights.
 - The scheduler does not require a synchronized DP request wave for DLO.
 
-This path is intentionally not implemented by reusing the AllGather mmap
+The mmap source remains immutable in both transfer modes. AllGather copies a
-loader. It relies on the standard loader so model-specific TP/HSDP transforms
+persistent `1 / group_size` shard from it and then closes the mappings;
-remain intact.
+rank-local mode keeps the mappings open and copies a complete block through
+bounded staging without adding a collective.
+
+The current implementation packs mmap views into a staging slot synchronously
+on the host; the subsequent pinned H2D transfer overlaps with block compute.
+This keeps host memory bounded but makes CPU packing bandwidth part of the
+rank-local path's performance profile.
 
 ## Parallelism compatibility
 
 | Parallelism | DLO + AllGather | DLO without AllGather |
 |---|---|---|
-| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
+| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight collective. Compatible TP1 models share file-backed DiT pages across replicas. |
 | **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
 | **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
 | **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
@@ -122,11 +146,35 @@ AllGather DP multi-concurrency requires:
 The no-AllGather path does not impose these DLO-specific synchronized-wave
 requirements.
 
-The mmap loader is used only by the supported DLO+AllGather path when the
+The mmap loader is used when the model explicitly declares a compatible
-model has a compatible checkpoint layout. Online quantization is incompatible
+checkpoint layout. It is independent of AllGather for TP1/non-HSDP execution.
-with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
+Models can declare a per-parameter checkpoint-to-runtime adapter; the adapter
-quantization. Models that do not meet the mmap requirements use the regular
+must preserve dtype and element count so it can be applied while packing one
-loader path.
+staging block. MiniMax-H3 uses this mechanism for grouped QKV. Online
+quantization and TP/HSDP no-AllGather execution use the regular loader path.
+
+### Adding mmap support to another model
+
+A new pipeline should opt in only after its ordinary loader contract can be
+represented without materializing a private full-model copy:
+
+1. Populate `weights_sources` with safetensors component sources and prefixes
+   that place every DiT checkpoint key in the pipeline parameter namespace.
+2. If prefixed checkpoint keys already equal runtime parameter names, declare
+   `_supports_mmap_loading = True`. If names differ, provide
+   `_remap_ckpt_key` to map checkpoint names to runtime names (or return
+   `None` for intentionally ignored keys).
+3. If a parameter needs a TP1 layout conversion, attach a callable
+   `mmap_weight_transform` to that parameter during model construction. The
+   transform must preserve dtype and element count because it runs while one
+   bounded block is packed.
+4. Verify that mmap and the regular loader produce the same runtime tensors,
+   that every expected DiT tensor maps successfully, and that TP, HSDP, and
+   online-quantized configurations retain their standard-loader behavior.
+
+Setting `_supports_mmap_loading = True` alone is not sufficient: explicit
+opt-ins fail during initialization if any expected DiT tensor remains on the
+`meta` device after checkpoint mapping.
 
 ## Validation coverage
 
@@ -136,6 +184,8 @@ Current source-level validation includes:
 - HSDP + DLO without AllGather acceptance at configuration level;
 - TP rejection in the DLO+AllGather mmap path;
 - resident-layer requests requiring no-AllGather;
+- rank-local mmap source retention, layout transforms, and bounded host
+  staging;
 - DP request-wave validation for denoising-step compatibility;
 - sharding, double-buffer, AllGather-size, and heterogeneous-block regression
   tests.
@@ -150,7 +200,8 @@ on the target CUDA/NCCL or CANN/HCCL hardware.
   scaling path.
 - Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
   not the goal.
-- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
+- Use **no-AllGather** when replicas must avoid DLO collectives. TP1 models with
-  the expectation of higher host-memory use and lower validation confidence.
+  mmap support can share DiT checkpoint pages; TP/HSDP and unsupported models
+  retain the higher-memory standard-loader path.
 - Prefer **HSDP alone** for production HSDP deployments until the combined
   HSDP + DLO no-AllGather path has broader end-to-end coverage.

docs/user_guide/diffusion/cpu_offload.md

@@ -189,7 +189,10 @@ multi-device deployments. In the default mode, each rank stores only
 reconstructed at runtime via **AllGather** on a dedicated communication stream,
 overlapped with computation via a fixed double-buffer scheme. The
 `--dlo-no-use-allgather` mode instead streams standard-loader rank-local
-weights independently and does not shard weights across ranks.
+weights independently and does not shard weights across ranks. For explicitly
+mmap-compatible TP1 models, the same mode can retain node-shared file-backed
+DiT weights and use bounded per-worker host staging instead of a private full
+DiT copy.
 
 For the implementation design and compatibility matrix, see
 [Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
@@ -201,8 +204,9 @@ For the implementation design and compatibility matrix, see
   time, regardless of model size
 - **DP multi-concurrency**: N concurrent requests processed in parallel
   (AllGather only gathers weight shards, which are request-independent)
-- **mmap weight loading**: weights loaded as mmap views pointing to shared OS
+- **mmap weight storage**: compatible checkpoints use views into the shared OS
-  page cache, eliminating O(dp_size × model_size) RSS during model creation
+  page cache. AllGather copies persistent rank shards; rank-local mode keeps
+  the views and packs complete blocks through two bounded host staging slots
 - **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU
   (CANN/HCCL) via vLLM-Omni's platform abstraction
 
@@ -224,7 +228,7 @@ vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --data-parallel-size 4
 
-# Without DLO AllGather (each rank streams standard-loader rank-local weights)
+# Without DLO AllGather (compatible TP1 models use node-shared mmap storage)
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --data-parallel-size 4 \
@@ -258,20 +262,26 @@ m = Omni(
 
 ### How model weights are loaded (mmap path)
 
-When DLO + AllGather is active, the offloader:
+When DLO and an mmap-compatible model are active, the offloader:
 
 1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`)
    from the normally-created transformer
 2. **Converts to meta device** via `to_empty(device="meta")`, releasing random
    initialization weights
 3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`,
-   which return views into the OS page cache (shared across all ranks, 0 RSS)
+   which return views into the OS page cache without eagerly copying the full
-4. **Calls `post_load_weights()`** to apply model-specific dtype conversions
+   tensor into anonymous process memory
-   (e.g. Cosmos3's `time_embedder` → FP32)
+4. **Applies declared per-parameter layout adapters while packing**, when the
+   checkpoint and TP1 runtime layouts differ (for example MiniMax-H3 grouped
+   QKV)
 5. **Restores non-persistent buffers** from saved copies
 
-This approach requires **zero model-specific code changes** — no pipeline or
+With AllGather, each rank copies only its persistent shard and then releases
-transformer modifications are needed.
+the source mappings. Without AllGather, mappings remain open and every worker
+owns two pinned staging slots sized to the largest streamed block. The latter
+shares file-backed pages, not process-local Python objects, and does not alter
+request routing or scheduling. Packing into the host slot is synchronous; the
+pinned H2D copy can overlap with computation.
 
 ### OffloadPlan (declarative topology metadata)
 
@@ -313,7 +323,7 @@ request while AllGather synchronizes only weight shards (request-independent).
   `--dlo-no-use-allgather` or disable online quantization.
 - Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
   weight-loader callbacks are bypassed. The no-AllGather path retains the
-  standard TP loader and is experimental.
+  standard TP loader instead of rank-local mmap and is experimental.
 - HSDP + AllGather is rejected because it would double-shard parameters. HSDP
   with no-AllGather is accepted at configuration level, but has limited
   end-to-end validation.
@@ -373,6 +383,7 @@ reasoner/generator components inside the model forward pass.
 | StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` |
 | Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
 | Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
+| MiniMaxH3Pipeline | `MiniMaxAI/MiniMax-H3` (`FL2VA`) | `MiniMaxH3DiTModel` | - | ✓ | ✓ | `"blocks"` |
 | SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` |
 | BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` |
 | Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` |
@@ -382,8 +393,9 @@ reasoner/generator components inside the model forward pass.
 - Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
 - Distributed Layerwise Offloading reuses the layerwise block topology, but the
   model and parallelism combination still needs validation. AllGather mode
-  additionally requires a compatible checkpoint/mmap layout; no-AllGather mode
+  additionally requires a compatible checkpoint/mmap layout. For explicitly
-  retains the standard loader's rank-local tensors. See the
+  compatible TP1 models, no-AllGather uses node-shared mmap sources; other
+  layouts retain the standard loader's rank-local tensors. See the
   [DLO design](../../design/feature/distributed_layerwise_offload.md) and
   [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
   for usage examples.

examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py

@@ -11,6 +11,13 @@ Examples:
     python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
         --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
 
+    # Rank-local mmap storage, with no DLO collective or DP wave scheduling.
+    VLLM_WORKER_MULTIPROC_METHOD=spawn \
+    VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
+    CUDA_VISIBLE_DEVICES=0,1 \
+    python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
+        --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2-no-allgather
+
     VLLM_WORKER_MULTIPROC_METHOD=spawn \
     VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
     CUDA_VISIBLE_DEVICES=0,1 \
@@ -45,9 +52,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
     parser.add_argument(
         "--mode",
-        choices=("dlo-dp2", "request"),
+        choices=("dlo-dp2", "dlo-dp2-no-allgather", "request"),
         required=True,
-        help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
+        help="DLO DP2 (with or without AllGather) or ordinary non-DLO TP2 request mode",
     )
     parser.add_argument("--steps", type=int, default=2)
     parser.add_argument("--duration", type=float, default=5.0)
@@ -74,14 +81,14 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
         "stage_init_timeout": args.init_timeout,
         "init_timeout": args.init_timeout,
     }
-    if args.mode == "dlo-dp2":
+    if args.mode in {"dlo-dp2", "dlo-dp2-no-allgather"}:
         common.update(
             tensor_parallel_size=1,
             data_parallel_size=2,
             text_encoder_tp_size=1,
             vae_patch_parallel_size=1,
             enable_distributed_layerwise_offload=True,
-            dlo_use_allgather=True,
+            dlo_use_allgather=args.mode == "dlo-dp2",
             dlo_resident_layers=0,
         )
     else:

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py

@@ -191,6 +191,30 @@ def test_model_load_weights_transforms_before_calling_vllm_loader():
     ]
 
 
+def test_unquantized_qkv_declares_equivalent_mmap_transform(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "Attention", _FakeAttention)
+
+    arch = h3.MiniMaxH3DiTArchConfig(
+        hidden_size=1,
+        num_attention_heads=2,
+        attention_head_dim=1,
+        rope_inv_freq_len=1,
+    )
+    attention = h3.MiniMaxH3Attention(
+        arch,
+        quant_config=None,
+        prefix="blocks.0.attn",
+    )
+    transform = attention.qkv_proj.weight.mmap_weight_transform
+    checkpoint_weight = torch.arange(6, dtype=torch.float32).reshape(6, 1)
+
+    assert transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
+
+
 def test_pipeline_resolves_transformer_component_quant_config():
     from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
         _resolve_component_quant_config,

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -30,7 +30,13 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
     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.diffusion.offloader.offload_plan import (
+    OffloadPlan,
+    get_offload_plan,
+    has_online_quantization,
+    should_select_dlo_mmap,
+    supports_mmap_loading,
+)
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -48,6 +54,9 @@ class DummyEvent:
     def record(self, _stream) -> None:
         return None
 
+    def synchronize(self) -> None:
+        return None
+
 
 @contextmanager
 def dummy_stream(_stream):
@@ -350,6 +359,79 @@ class TestDistributedLayerwiseOffloadHook:
 
         assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0]))
 
+    def test_rank_local_mmap_retains_source_and_uses_staging(self, dist_group, patched_offload_runtime):
+        current_block = nn.Linear(2, 2, bias=False)
+        next_block = nn.Linear(2, 2, bias=False)
+        next_block.weight.data.copy_(torch.arange(4, dtype=torch.float32).view(2, 2))
+        source_storage = next_block.weight.untyped_storage().data_ptr()
+        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=1,
+            rank=0,
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        hook.initialize_hook(current_block)
+
+        source = hook.cpu_sources[torch.float32][0]["tensor"]
+        assert source.untyped_storage().data_ptr() == source_storage
+        assert hook.cpu_shards == {}
+        assert next_block.weight.numel() == 0
+
+        hook.cpu_staging_buffers = [
+            {torch.float32: torch.empty(4)},
+            {torch.float32: torch.empty(4)},
+        ]
+        hook.prefetch_layer(slot=0, non_blocking=False)
+
+        assert torch.equal(
+            next_block.weight,
+            torch.tensor([[2.0, 3.0], [0.0, 1.0]]),
+        )
+        assert torch.equal(source, torch.arange(4, dtype=torch.float32).view(2, 2))
+
+    def test_rank_local_mmap_staging_is_bounded_by_largest_block(self, patched_offload_runtime):
+        hooks = []
+        for size in (4, 9):
+            current_block = nn.Linear(1, 1, bias=False)
+            next_block = nn.Module()
+            next_block.weight = nn.Parameter(torch.arange(size, dtype=torch.float32))
+            hook = DistributedLayerwiseOffloadHook(
+                next_block=next_block,
+                device=torch.device("cpu"),
+                dp_group=None,
+                dp_size=1,
+                rank=0,
+                pin_memory=False,
+                shared_buffers=[None, None],
+                rank_local_mmap=True,
+            )
+            hook.initialize_hook(current_block)
+            hooks.append(hook)
+
+        resident_block = nn.Module()
+        resident_block.weight = nn.Parameter(torch.arange(12, dtype=torch.float32))
+        resident_group = PinnedResidentLayerGroup(
+            [resident_block],
+            device=torch.device("cpu"),
+            copy_stream=DummyStream(),
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        staging = DistributedLayerwiseOffloadBackend._allocate_shared_cpu_staging_buffers(
+            hooks,
+            resident_group,
+        )
+
+        assert len(staging) == 2
+        assert staging[0][torch.float32].numel() == 12
+        assert staging[1][torch.float32].numel() == 12
+
 
 class TestPinnedResidentLayerGroup:
     def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -386,6 +468,42 @@ class TestPinnedResidentLayerGroup:
 
         group.offload()
 
+    def test_rank_local_mmap_resident_layers_retain_sources(self, patched_offload_runtime):
+        blocks = [nn.Linear(2, 2, bias=False) for _ in range(3)]
+        expected = []
+        source_storage = []
+        for index, block in enumerate(blocks):
+            block.weight.data.fill_(index + 1)
+            expected.append(block.weight.detach().clone())
+            source_storage.append(block.weight.untyped_storage().data_ptr())
+
+        group = PinnedResidentLayerGroup(
+            blocks,
+            device=torch.device("cpu"),
+            copy_stream=DummyStream(),
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+
+        for index, state in enumerate(group._states):
+            source = state["cpu_sources"][torch.float32][0]["tensor"]
+            assert source.untyped_storage().data_ptr() == source_storage[index]
+            assert state["cpu_shards"] == {}
+        assert len(group._cpu_staging_buffers) == 2
+
+        group.load()
+        for block, weight in zip(blocks, expected):
+            assert torch.equal(block.weight, weight)
+
+        blocks[0].weight.data.zero_()
+        group.offload()
+        group.load()
+        for block, weight in zip(blocks, expected):
+            assert torch.equal(block.weight, weight)
+
+        group.offload()
+        assert all(block.weight.numel() == 0 for block in blocks)
+
     def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
         module = nn.Linear(2, 2, bias=False)
 
@@ -779,6 +897,46 @@ class TestOffloadPlan:
         with pytest.raises(Exception):
             plan.block_attrs = {"x": ("y",)}  # type: ignore
 
+    def test_explicit_component_sources_can_opt_in_to_mmap(self):
+        class Pipeline(nn.Module):
+            _supports_mmap_loading = True
+            weights_sources = [object()]
+
+        assert supports_mmap_loading(Pipeline())
+
+    def test_component_sources_do_not_implicitly_enable_mmap(self):
+        class Pipeline(nn.Module):
+            weights_sources = [object()]
+
+        assert not supports_mmap_loading(Pipeline())
+
+    def test_online_quantization_detection(self):
+        model = nn.Sequential(nn.Linear(1, 1))
+        assert not has_online_quantization(model)
+
+        model[0].quant_method = SimpleNamespace(uses_meta_device=True)
+        assert has_online_quantization(model)
+
+    @pytest.mark.parametrize(
+        ("kwargs", "expected"),
+        [
+            ({"dlo_use_allgather": True, "tensor_parallel_size": 2}, True),
+            ({"dlo_use_allgather": False, "tensor_parallel_size": 1}, True),
+            ({"dlo_use_allgather": False, "tensor_parallel_size": 2}, False),
+            ({"dlo_use_allgather": False, "tensor_parallel_size": 1, "use_hsdp": True}, False),
+            ({"dlo_use_allgather": False, "tensor_parallel_size": 1, "has_online_quant": True}, False),
+        ],
+    )
+    def test_dlo_mmap_policy(self, kwargs, expected):
+        defaults = {
+            "mmap_supported": True,
+            "dlo_use_allgather": False,
+            "has_online_quant": False,
+            "tensor_parallel_size": 1,
+            "use_hsdp": False,
+        }
+        assert should_select_dlo_mmap(**(defaults | kwargs)) is expected
+
     def test_all_resident_dit_still_prepares_planned_submodules(self, patched_offload_runtime):
         class TokenRefiner(nn.Module):
             def __init__(self):
@@ -827,6 +985,93 @@ class TestOffloadPlan:
 class TestMmapValidation:
     """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
 
+    def test_component_source_prefix_builds_model_namespace(self, tmp_path):
+        source_root = tmp_path / "partition"
+        checkpoint_dir = source_root / "transformer"
+        checkpoint_dir.mkdir(parents=True)
+        checkpoint_file = checkpoint_dir / "model.safetensors"
+        save_file({"blocks.0.weight": torch.arange(4, dtype=torch.float32)}, str(checkpoint_file))
+
+        class Pipeline(nn.Module):
+            _supports_mmap_loading = True
+
+            def __init__(self):
+                super().__init__()
+                self.weights_sources = [
+                    SimpleNamespace(
+                        model_or_path=str(source_root),
+                        subfolder="transformer",
+                        revision=None,
+                        prefix="transformer.",
+                    )
+                ]
+
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+            ),
+            torch.device("cpu"),
+        )
+
+        model_map = backend._build_mmap_model_map(Pipeline())
+
+        assert model_map == {
+            "transformer.blocks.0.weight": (
+                "blocks.0.weight",
+                str(checkpoint_file),
+            )
+        }
+
+    def test_explicit_mmap_opt_in_rejects_missing_model_tensors(self, tmp_path, patched_offload_runtime, monkeypatch):
+        monkeypatch.setattr(
+            "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
+            lambda: 1,
+        )
+
+        source_root = tmp_path / "partition"
+        checkpoint_dir = source_root / "transformer"
+        checkpoint_dir.mkdir(parents=True)
+        save_file(
+            {"blocks.0.weight": torch.ones((2, 2), dtype=torch.float32)},
+            str(checkpoint_dir / "model.safetensors"),
+        )
+
+        class Transformer(nn.Module):
+            _layerwise_offload_blocks_attrs = ["blocks"]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False), nn.Linear(2, 2, bias=False)])
+
+        class Pipeline(nn.Module):
+            _supports_mmap_loading = True
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+                self.weights_sources = [
+                    SimpleNamespace(
+                        model_or_path=str(source_root),
+                        subfolder="transformer",
+                        revision=None,
+                        prefix="transformer.",
+                    )
+                ]
+
+        pipeline = Pipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+            ),
+            torch.device("cpu"),
+        )
+        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+
+        with pytest.raises(RuntimeError, match="Explicit mmap loading left 1"):
+            backend._load_weights_via_mmap(pipeline, modules)
+
     def test_tp_rejected_when_tp_world_size_gt_1(self, tmp_path, patched_offload_runtime, monkeypatch):
         """DLO+AllGather should reject when TP world size > 1."""
         import json

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -363,22 +363,39 @@ class DiffusersPipelineLoader:
             else:
                 model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
 
-                # Skip load_weights only for DLO+AllGather when the model
+                # Skip load_weights when DLO can keep checkpoint tensors as
-                # supports mmap loading and no online quantization is active.
+                # mmap-backed host masters.  Storage selection is independent
-                # This condition MUST match the gate in
+                # of whether DLO reconstructs weights with AllGather.  The
-                # DistributedLayerwiseOffloadBackend.enable() so that the
+                # no-AllGather path can use mmap only when the checkpoint is
-                # loader skips ⟺ enable() uses mmap.
+                # already the rank-local runtime layout (currently TP=1 and no
-                from vllm_omni.diffusion.offloader.offload_plan import supports_mmap_loading
+                # HSDP); TP/HSDP still require the regular loader callbacks.
+                # The loader and backend share should_select_dlo_mmap() so
+                # skipping load_weights always selects mmap in enable().
+                from vllm_omni.diffusion.offloader.offload_plan import (
+                    should_select_dlo_mmap,
+                    supports_mmap_loading,
+                )
 
                 _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
                 _use_ag = getattr(self.od_config, "dlo_use_allgather", True)
                 _has_online_quant = self._has_online_quant(model)
                 _supports_mmap = supports_mmap_loading(model)
+                _tp_size = int(getattr(self.parallel_config, "tensor_parallel_size", 1))
+                _use_hsdp = bool(getattr(self.parallel_config, "use_hsdp", False))
 
-                _skip_load = _dist_offload and _use_ag and _supports_mmap and not _has_online_quant
+                _skip_load = _dist_offload and should_select_dlo_mmap(
+                    mmap_supported=_supports_mmap,
+                    dlo_use_allgather=_use_ag,
+                    has_online_quant=_has_online_quant,
+                    tensor_parallel_size=_tp_size,
+                    use_hsdp=_use_hsdp,
+                )
 
                 if _skip_load:
-                    logger.info("DLO+AllGather active: skipping load_weights (will load via mmap in enable())")
+                    logger.info(
+                        "DLO mmap storage active (%s): skipping load_weights (will map checkpoint weights in enable())",
+                        "AllGather" if _use_ag else "rank-local",
+                    )
                 else:
                     if _dist_offload and _use_ag and _has_online_quant:
                         raise ValueError(
@@ -441,11 +458,11 @@ class DiffusersPipelineLoader:
         """Whether any layer uses an online-quant method that defers weight
         materialization onto the ``meta`` device (upstream vLLM
         ``uses_meta_device=True``, e.g. online FP8)."""
-        for module in model.modules():
+        from vllm_omni.diffusion.offloader.offload_plan import (
-            quant_method = getattr(module, "quant_method", None)
+            has_online_quantization,
-            if getattr(quant_method, "uses_meta_device", False):
+        )
-                return True
+
-        return False
+        return has_online_quantization(model)
 
     def _apply_skip_softmax_calibration(self, model: nn.Module) -> None:
         from vllm_omni.diffusion.attention.backends.trtllm_calibration import (

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -11,6 +11,7 @@ from __future__ import annotations
 import math
 from collections.abc import Iterable, Mapping
 from dataclasses import dataclass
+from functools import partial
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -350,6 +351,18 @@ class MiniMaxH3Attention(nn.Module):
             prefix=f"{prefix}.qkv_proj",
             return_bias=True,
         )
+        if quant_config is None:
+            # The checkpoint interleaves Q/K/V by query group, while vLLM's
+            # fused QKV linear consumes contiguous Q, then K, then V rows.  The
+            # regular loader performs this conversion before weight_loader().
+            # DLO mmap applies the same adapter while packing one bounded host
+            # staging block; the raw file-backed master remains node-shared.
+            self.qkv_proj.weight.mmap_weight_transform = partial(
+                _reorder_grouped_qkv_to_qkv,
+                num_query_groups=arch.num_attention_heads,
+                heads_per_group=1,
+                head_dim=arch.attention_head_dim,
+            )
         self.num_heads = self.qkv_proj.num_heads
         self.num_kv_heads = self.qkv_proj.num_kv_heads
         self.rot_dim = 6 * arch.rope_inv_freq_len

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -532,13 +532,12 @@ class MiniMaxH3Pipeline(
         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
+    # Explicit mmap opt-in: weights_sources below provide exact pipeline
-    # and fused MLP). Keep it on that path until mmap can run the same loader
+    # prefixes, and TP1 DLO mmap applies the grouped-QKV adapter while packing
-    # callbacks for every affected parameter.
+    # each block. The fused MLP checkpoint is already gate-then-up, matching
-    _supports_mmap_loading: ClassVar[bool] = False
+    # its TP1 runtime layout. TP>1 and quantized layouts stay on the regular
-    # TODO(offload): Re-enable after the generic rank-local mmap path can run
+    # loader. See the DLO design doc for the model opt-in contract.
-    # this model's grouped-QKV reorder and fused-MLP packing before TP sharding.
+    _supports_mmap_loading: ClassVar[bool] = True
-    # Do not bypass the regular loader until that equivalence is tested.
     _PROFILER_TARGETS: ClassVar[list[str]] = [
         "_prepare_reference_videos",
         "encode_prompt",
@@ -648,6 +647,10 @@ class MiniMaxH3Pipeline(
             od_config.quantization_config,
             "transformer",
         )
+        if transformer_quant_config is not None:
+            # Quantized parameters may have packed weights and scale tensors
+            # whose runtime layout is owned by quant-method loader callbacks.
+            self._supports_mmap_loading = False
         self.transformer = MiniMaxH3DiTModel(
             od_config,
             quant_config=transformer_quant_config,

vllm_omni/diffusion/offloader/base.py

@@ -32,6 +32,10 @@ class OffloadConfig:
     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
+    # Rank-local mmap bypasses TP-aware loader callbacks, so the backend needs
+    # the configured TP size when selecting mmap versus standard-loader host
+    # storage. Appended to preserve the positional order of existing fields.
+    tensor_parallel_size: int = 1
 
     @classmethod
     def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -58,6 +62,7 @@ class OffloadConfig:
 
         parallel_config = getattr(od_config, "parallel_config", None)
         use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False
+        tensor_parallel_size = int(getattr(parallel_config, "tensor_parallel_size", 1)) if parallel_config else 1
 
         # Derive dp_size from parallel_config — not user-configurable.
         # The offload adapts to whatever DP/SP is already configured.
@@ -116,7 +121,8 @@ class OffloadConfig:
             dp_size = 1
             logger.info(
                 "Distributed layerwise offload: dlo_use_allgather=False, "
-                "streaming standard-loader rank-local weights (no DP shard or AllGather)"
+                "streaming complete rank-local blocks (no DLO shard or AllGather); "
+                "the backend will select mmap or standard-loader host storage"
             )
 
         # HSDP already shards parameters into DTensors.  Running distributed
@@ -134,6 +140,7 @@ class OffloadConfig:
             strategy=strategy,
             pin_cpu_memory=pin_cpu_memory,
             use_hsdp=use_hsdp,
+            tensor_parallel_size=tensor_parallel_size,
             dp_size=dp_size,
             dlo_use_allgather=dlo_use_allgather,
             dlo_resident_layers=dlo_resident_layers,

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -1,15 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Distributed Layerwise Offload backend with H2D + AllGather overlap.
+"""Distributed Layerwise Offload backend with double-buffered H2D.
 
 This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
 
-* Shards model weights across DP ranks and stores only the local shard
+* Optionally shards model weights across DP ranks and reconstructs each block
-  (1/DP_size of full model) in host pinned memory per rank.
+  with AllGather, or streams a complete rank-local block without a collective.
+* Can retain compatible checkpoint tensors as node-shared mmap sources instead
+  of creating a persistent private host copy in every rank.
 * Uses a fixed double-buffer scheme that keeps only two layers' worth of
   weights on each device at any time.
-* Asynchronously pipelines both H2D transfers and AllGather communications
+* Pipelines H2D transfers and, when enabled, AllGather communications on
-  on dedicated streams, fully overlapping them with computation.
+  dedicated streams, overlapping them with computation.
 * Is hardware-agnostic, supporting both NVIDIA GPU (CUDA) and Ascend NPU
   (CANN) platforms via vLLM-Omni's platform abstraction layer.
 """
@@ -31,7 +33,13 @@ from vllm_omni.platforms import current_omni_platform
 from .base import OffloadBackend, OffloadConfig
 from .block_discovery import get_blocks_from_dit
 from .module_collector import ModuleDiscovery
-from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading
+from .offload_plan import (
+    OffloadPlan,
+    get_offload_plan,
+    has_online_quantization,
+    should_select_dlo_mmap,
+    supports_mmap_loading,
+)
 from .tensor_utils import (
     dtype_size as _dtype_size,
 )
@@ -74,6 +82,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         comm_stream: Any | None = None,
         pin_memory: bool = True,
         shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None,
+        rank_local_mmap: bool = False,
     ):
         assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`"
 
@@ -83,6 +92,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         self.dp_size = dp_size
         self.rank = rank
         self.pin_memory = pin_memory
+        self.rank_local_mmap = rank_local_mmap
 
         self.copy_stream = copy_stream or current_omni_platform.Stream()
         self.comm_stream = comm_stream or current_omni_platform.Stream()
@@ -98,8 +108,18 @@ class DistributedLayerwiseOffloadHook(ModelHook):
 
         # Sharded host weights for the next block, keyed by dtype
         self.cpu_shards: dict[torch.dtype, torch.Tensor] = {}
+        # File-backed source tensors for rank-local mmap.  Unlike cpu_shards,
+        # these remain immutable views of the checkpoint and are never pinned
+        # or flattened into a model-sized private allocation.
+        self.cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {}
         self.metadata: dict[torch.dtype, list[dict[str, Any]]] = {}
 
+        # Rank-local mmap uses two host staging slots shared by every hook in
+        # this worker.  They are assigned by the backend after all block sizes
+        # are known, mirroring the shared device-buffer allocation.
+        self.cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]
+        self.cpu_staging_events: list[Any | None] = [None, None]
+
         # Current slot index (0 or 1).  Updated dynamically by the previous
         # hook's prefetch_layer call via _prefetched_slot.  This ensures
         # correct slot tracking for ALL block counts (including odd N).
@@ -154,14 +174,20 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         self.next_block_parameters = dict(self.next_block.named_parameters())
         self.next_block_buffers = dict(self.next_block.named_buffers())
 
-        # Shard next block's weights and store local shard in pinned CPU memory
+        if self.rank_local_mmap:
-        self.cpu_shards, self.metadata = self._shard_and_pin(
+            self.cpu_sources, self.metadata = self._collect_mmap_sources(
-            self.next_block_parameters,
+                self.next_block_parameters,
-            self.next_block_buffers,
+                self.next_block_buffers,
-            self.dp_size,
+            )
-            self.rank,
+        else:
-            self.pin_memory,
+            # Shard next block's weights and store local shard in pinned CPU memory.
-        )
+            self.cpu_shards, self.metadata = self._shard_and_pin(
+                self.next_block_parameters,
+                self.next_block_buffers,
+                self.dp_size,
+                self.rank,
+                self.pin_memory,
+            )
 
         # Allocate device buffers only if not using shared buffers from backend
         if self._owns_buffers:
@@ -185,11 +211,65 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         self._ag_output_sizes: dict[torch.dtype, int] = {}
         for dtype, metas in self.metadata.items():
             total_numel = sum(m["numel"] for m in metas)
-            shard_numel = self.cpu_shards[dtype].numel()
+            if self.rank_local_mmap:
-            self._ag_output_sizes[dtype] = shard_numel * self.dp_size if self.dp_size > 1 else total_numel
+                self._ag_output_sizes[dtype] = total_numel
+            else:
+                shard_numel = self.cpu_shards[dtype].numel()
+                self._ag_output_sizes[dtype] = shard_numel * self.dp_size if self.dp_size > 1 else total_numel
 
         return module
 
+    @staticmethod
+    def _collect_mmap_sources(
+        params: dict[str, nn.Parameter],
+        bufs: dict[str, torch.Tensor],
+    ) -> tuple[dict[torch.dtype, list[dict[str, Any]]], dict[torch.dtype, list[dict[str, Any]]]]:
+        """Retain file-backed tensors and replace module storage with placeholders.
+
+        The returned sources preserve safetensors mmap storage.  Runtime-layout
+        adapters are applied only while packing a bounded staging slot, so no
+        full-model anonymous CPU copy is retained by a worker.
+        """
+        cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {}
+        metadata: dict[torch.dtype, list[dict[str, Any]]] = {}
+        offsets: dict[torch.dtype, int] = {}
+
+        for name, target in chain(params.items(), bufs.items()):
+            source = target.to_local() if hasattr(target, "to_local") else target
+            if source.device.type != "cpu":
+                raise ValueError(
+                    f"Rank-local mmap storage requires CPU checkpoint views, but {name!r} is on {source.device}."
+                )
+
+            dtype = source.dtype
+            offset = offsets.get(dtype, 0)
+            transform = getattr(target, "mmap_weight_transform", None)
+            if not getattr(target, "mmap_weight_transform_pending", False):
+                transform = None
+
+            cpu_sources.setdefault(dtype, []).append(
+                {
+                    "name": name,
+                    "tensor": source.detach(),
+                    "transform": transform,
+                }
+            )
+            metadata.setdefault(dtype, []).append(
+                {
+                    "name": name,
+                    "offset": offset,
+                    "numel": source.numel(),
+                    "shape": source.shape,
+                }
+            )
+            offsets[dtype] = offset + source.numel()
+
+            # The detached source above keeps the mmap storage alive while the
+            # module parameter/buffer is rebound to the rotating device slot.
+            set_tensor_storage(target, make_offload_placeholder(target))
+
+        return cpu_sources, metadata
+
     @staticmethod
     def _shard_and_pin(
         params: dict[str, nn.Parameter],
@@ -306,6 +386,52 @@ class DistributedLayerwiseOffloadHook(ModelHook):
     #  Prefetch: H2D + AllGather (overlapped on dedicated streams)      #
     # ------------------------------------------------------------------ #
 
+    @staticmethod
+    def _pack_mmap_sources(
+        cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
+        metadata: dict[torch.dtype, list[dict[str, Any]]],
+        slot_buffers: dict[torch.dtype, torch.Tensor],
+    ) -> dict[torch.dtype, torch.Tensor]:
+        staged: dict[torch.dtype, torch.Tensor] = {}
+        for dtype, metas in metadata.items():
+            total_numel = sum(meta["numel"] for meta in metas)
+            destination = slot_buffers[dtype][:total_numel]
+            sources = cpu_sources[dtype]
+            for source_info, meta in zip(sources, metas, strict=True):
+                source = source_info["tensor"]
+                transform = source_info["transform"]
+                if callable(transform):
+                    source = transform(source)
+                if source.dtype != dtype or source.numel() != meta["numel"]:
+                    raise ValueError(
+                        "mmap weight transform changed tensor layout for "
+                        f"{source_info['name']!r}: expected dtype={dtype}, "
+                        f"numel={meta['numel']}, got dtype={source.dtype}, "
+                        f"numel={source.numel()}"
+                    )
+                start = meta["offset"]
+                destination[start : start + meta["numel"]].copy_(source.reshape(-1))
+            staged[dtype] = destination
+        return staged
+
+    def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]:
+        """Pack this block's mmap views into one bounded host staging slot."""
+        previous_copy = self.cpu_staging_events[slot]
+        if previous_copy is not None:
+            synchronize = getattr(previous_copy, "synchronize", None)
+            if callable(synchronize):
+                synchronize()
+            else:
+                # Platform events normally expose synchronize().  Retain a
+                # correctness fallback for test and non-CUDA platform shims.
+                current_omni_platform.synchronize()
+            self.cpu_staging_events[slot] = None
+
+        slot_buffers = self.cpu_staging_buffers[slot]
+        if slot_buffers is None:
+            raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated")
+        return self._pack_mmap_sources(self.cpu_sources, self.metadata, slot_buffers)
+
     @torch.compiler.disable
     def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None:
         """Prepare next block's weights into the shared device buffer for *slot*.
@@ -321,11 +447,17 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"
 
         if self.dp_size <= 1 or self.dp_group is None:
+            cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
             with current_omni_platform.stream(self.copy_stream):
-                for dtype, cpu_shard in self.cpu_shards.items():
+                for dtype, cpu_shard in cpu_weights.items():
                     gw = gpu_weights[dtype]
-                    gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=non_blocking)
+                    async_copy = non_blocking and cpu_shard.is_pinned()
+                    gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
                 evt.record(self.copy_stream)
+            if self.rank_local_mmap:
+                # The CPU slot may be overwritten only after this H2D copy has
+                # finished.  The shared event protects reuse by another hook.
+                self.cpu_staging_events[slot] = evt
         else:
             gpu_shards: dict[torch.dtype, torch.Tensor] = {}
             shard_bufs = self.gpu_shard_buffers[slot]
@@ -468,6 +600,7 @@ def apply_distributed_block_hook(
     comm_stream: Any | None = None,
     pin_memory: bool = True,
     shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None,
+    rank_local_mmap: bool = False,
 ) -> DistributedLayerwiseOffloadHook:
     """Register a DistributedLayerwiseOffloadHook on *module*."""
     registry = HookRegistry.get_or_create(module)
@@ -481,6 +614,7 @@ def apply_distributed_block_hook(
         comm_stream=comm_stream,
         pin_memory=pin_memory,
         shared_buffers=shared_buffers,
+        rank_local_mmap=rank_local_mmap,
     )
     registry.register_hook(DistributedLayerwiseOffloadHook._HOOK_NAME, hook)
     return hook
@@ -495,15 +629,16 @@ def remove_distributed_block_hook(module: nn.Module) -> None:
 
 
 class PinnedResidentLayerGroup:
-    """Keep selected layers in pinned host memory between requests.
+    """Keep selected layers available for stage-scoped device residency.
 
 
     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
+    pinned CPU master copy (or mmap source plus bounded staging) and never
-    host.  Entering the denoise stage performs one asynchronous H2D pass;
+    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.
 
@@ -518,10 +653,14 @@ class PinnedResidentLayerGroup:
         device: torch.device,
         copy_stream: Any,
         pin_memory: bool,
+        rank_local_mmap: bool = False,
+        defer_staging: bool = False,
     ) -> None:
         self.device = device
         self.copy_stream = copy_stream
         self.loaded = False
+        self.rank_local_mmap = rank_local_mmap
+        self.pin_memory = pin_memory
         self._states: list[dict[str, Any]] = []
         self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
 
@@ -529,21 +668,47 @@ class PinnedResidentLayerGroup:
             params = dict(block.named_parameters())
             bufs = dict(block.named_buffers())
             targets: dict[str, torch.Tensor] = {**params, **bufs}
-            cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin(
+            if rank_local_mmap:
-                params,
+                cpu_sources, metadata = DistributedLayerwiseOffloadHook._collect_mmap_sources(
-                bufs,
+                    params,
-                dp_size=1,
+                    bufs,
-                rank=0,
+                )
-                pin_memory=pin_memory,
+                cpu_shards = {}
-            )
+            else:
+                cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin(
+                    params,
+                    bufs,
+                    dp_size=1,
+                    rank=0,
+                    pin_memory=pin_memory,
+                )
+                cpu_sources = {}
             self._states.append(
                 {
                     "targets": targets,
                     "cpu_shards": cpu_shards,
+                    "cpu_sources": cpu_sources,
                     "metadata": metadata,
                 }
             )
 
+        self._cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor]] = []
+        self._cpu_staging_events: list[Any | None] = [None, None]
+        if rank_local_mmap and not defer_staging:
+            max_sizes: dict[torch.dtype, int] = {}
+            for state in self._states:
+                for dtype, metas in state["metadata"].items():
+                    total = sum(meta["numel"] for meta in metas)
+                    max_sizes[dtype] = max(max_sizes.get(dtype, 0), total)
+            for _ in range(2):
+                buffers = {}
+                for dtype, total in max_sizes.items():
+                    buffer = torch.empty(total, dtype=dtype, device="cpu")
+                    if pin_memory:
+                        buffer = buffer.pin_memory()
+                    buffers[dtype] = buffer
+                self._cpu_staging_buffers.append(buffers)
+
     def load(self) -> None:
         if self.loaded:
             return
@@ -554,20 +719,41 @@ class PinnedResidentLayerGroup:
         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():
+            for dtype, metas in state["metadata"].items():
-                block_buffers[dtype] = torch.empty(
+                total = sum(meta["numel"] for meta in metas)
-                    cpu_shard.shape,
+                block_buffers[dtype] = torch.empty(total, dtype=dtype, device=self.device)
-                    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 index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
-                for dtype, cpu_shard in state["cpu_shards"].items():
+                if self.rank_local_mmap:
-                    block_buffers[dtype].copy_(cpu_shard, non_blocking=True)
+                    slot = index % 2
+                    previous_copy = self._cpu_staging_events[slot]
+                    if previous_copy is not None:
+                        synchronize = getattr(previous_copy, "synchronize", None)
+                        if callable(synchronize):
+                            synchronize()
+                        else:
+                            current_omni_platform.synchronize()
+                    cpu_weights = DistributedLayerwiseOffloadHook._pack_mmap_sources(
+                        state["cpu_sources"],
+                        state["metadata"],
+                        self._cpu_staging_buffers[slot],
+                    )
+                else:
+                    cpu_weights = state["cpu_shards"]
+
+                for dtype, cpu_weight in cpu_weights.items():
+                    block_buffers[dtype].copy_(
+                        cpu_weight,
+                        non_blocking=cpu_weight.is_pinned(),
+                    )
+                if self.rank_local_mmap:
+                    slot_ready = current_omni_platform.Event()
+                    slot_ready.record(self.copy_stream)
+                    self._cpu_staging_events[slot] = slot_ready
             ready.record(self.copy_stream)
 
         for state, block_buffers in zip(self._states, gpu_buffers):
@@ -630,6 +816,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = []
         self._resident_blocks: list[nn.Module] = []
         self._resident_layer_group: PinnedResidentLayerGroup | None = None
+        self._using_mmap = False
+        self._using_rank_local_mmap = False
 
     def load_resident_layers(self) -> None:
         """Load the model-declared leading blocks for the denoise stage."""
@@ -664,70 +852,146 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 if attr == "mmap_weight_transform":
                     replacement.mmap_weight_transform_pending = True
 
-    def _load_weights_via_mmap(self, pipeline: nn.Module, modules) -> None:
+    @staticmethod
-        """Load DiT weights from safetensors via mmap views (no RSS).
+    def _add_safetensors_source(
-
+        checkpoint_dir: str,
-        When the transformer is created on meta device, this method
+        prefix: str,
-        replaces meta params with mmap views of the checkpoint files.
+        remap_fn: Any | None,
-        The views point to OS page cache (shared across ranks), so
+        model_to_ckpt: dict[str, tuple[str, str]],
-        no private copy is created.  _shard_and_pin then copies only
+    ) -> int:
-        the 1/dp_size shard from the mmap view to a private buffer.
+        """Index one checkpoint directory into the pipeline parameter namespace."""
-
-        Non-DiT modules (VAE, encoders) are NOT affected — they were
-        created on CPU with real weights via from_pretrained.
-        """
         import glob
         import json
 
         from safetensors import safe_open
 
-        model_path = self.config.model_path
+        indexed_keys = 0
-        if not model_path:
+        index_files = glob.glob(os.path.join(checkpoint_dir, "*.safetensors.index.json"))
-            logger.warning("No model_path for mmap weight loading, skipping")
+        if index_files:
-            return
+            for index_file in index_files:
+                with open(index_file, encoding="utf-8") as handle:
+                    index = json.load(handle)
+                for checkpoint_key, filename in index.get("weight_map", {}).items():
+                    source_name = prefix + checkpoint_key
+                    model_name = remap_fn(source_name) if remap_fn is not None else source_name
+                    if model_name is not None:
+                        model_to_ckpt[model_name] = (
+                            checkpoint_key,
+                            os.path.join(checkpoint_dir, filename),
+                        )
+                    indexed_keys += 1
+            return indexed_keys
 
-        # Resolve HF repo ID to local snapshot path.
+        single_files = glob.glob(os.path.join(checkpoint_dir, "*.safetensors"))
-        # If model_path is a local directory, use it directly; otherwise
+        for filename in single_files:
-        # download the safetensors index + weight files via HuggingFace Hub.
+            with safe_open(filename, framework="pt", device="cpu") as handle:
-        if not os.path.isdir(model_path):
+                for checkpoint_key in handle.keys():
-            from vllm.model_executor.model_loader.weight_utils import (
+                    source_name = prefix + checkpoint_key
-                download_weights_from_hf,
+                    model_name = remap_fn(source_name) if remap_fn is not None else source_name
+                    if model_name is not None:
+                        model_to_ckpt[model_name] = (checkpoint_key, filename)
+                    indexed_keys += 1
+        return indexed_keys
+
+    def _build_mmap_model_map(self, pipeline: nn.Module) -> dict[str, tuple[str, str]]:
+        """Resolve safetensors sources to ``model name -> (key, file)``."""
+        remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
+        model_to_ckpt: dict[str, tuple[str, str]] = {}
+        indexed_keys = 0
+
+        # Prefer the same component sources consumed by DiffusersPipelineLoader.
+        # Besides avoiding unrelated checkpoint subdirectories, the source
+        # prefix gives models such as MiniMax-H3 an exact pipeline namespace.
+        sources = tuple(getattr(pipeline, "weights_sources", ()))
+        for source in sources:
+            source_root = source.model_or_path
+            if not os.path.isdir(source_root):
+                from vllm.model_executor.model_loader.weight_utils import (
+                    download_weights_from_hf,
+                )
+
+                source_root = download_weights_from_hf(
+                    model_name_or_path=source.model_or_path,
+                    cache_dir=None,
+                    allow_patterns=["*.safetensors", "*.safetensors.index.json"],
+                    revision=source.revision,
+                    subfolder=source.subfolder,
+                )
+            checkpoint_dir = (
+                os.path.join(source_root, source.subfolder) if source.subfolder is not None else source_root
+            )
+            indexed_keys += self._add_safetensors_source(
+                checkpoint_dir,
+                source.prefix,
+                remap_fn,
+                model_to_ckpt,
             )
 
-            logger.info("model_path %s is not local, downloading from HF", model_path)
+        if not sources:
-            model_path = download_weights_from_hf(
+            model_path = self.config.model_path
-                model_name_or_path=model_path,
+            if not model_path:
-                cache_dir=None,
+                raise RuntimeError("No model_path or weights_sources for mmap weight loading")
-                allow_patterns=["*.safetensors", "*.safetensors.index.json"],
+            if not os.path.isdir(model_path):
-            )
+                from vllm.model_executor.model_loader.weight_utils import (
+                    download_weights_from_hf,
+                )
 
-        # Build {checkpoint_key: file_path} from safetensors index
+                logger.info("model_path %s is not local, downloading from HF", model_path)
-        weight_map: dict[str, str] = {}
+                model_path = download_weights_from_hf(
-        for idx_file in glob.glob(os.path.join(model_path, "**", "*.safetensors.index.json"), recursive=True):
+                    model_name_or_path=model_path,
-            idx_dir = os.path.dirname(idx_file)
+                    cache_dir=None,
-            with open(idx_file) as f:
+                    allow_patterns=["*.safetensors", "*.safetensors.index.json"],
-                idx = json.load(f)
+                )
-            for key, filename in idx.get("weight_map", {}).items():
-                weight_map[key] = os.path.join(idx_dir, filename)
 
-        if not weight_map:
+            # Legacy mmap-enabled pipelines may keep the transformer one level
-            single_files = glob.glob(os.path.join(model_path, "**", "model.safetensors"), recursive=True)
+            # below model_path, so preserve recursive discovery in this fallback.
-            for sf in single_files:
+            import glob
-                from safetensors import safe_open as _safe_open
 
-                f = _safe_open(sf, framework="pt", device="cpu")
+            checkpoint_dirs = {
-                for key in f.keys():
+                os.path.dirname(path)
-                    weight_map[key] = sf
+                for pattern in ("*.safetensors.index.json", "*.safetensors")
-                logger.info("Found single-file safetensors: %s (%d keys)", sf, len(f.keys()))
+                for path in glob.glob(os.path.join(model_path, "**", pattern), recursive=True)
+            }
+            for checkpoint_dir in sorted(checkpoint_dirs):
+                indexed_keys += self._add_safetensors_source(
+                    checkpoint_dir,
+                    "",
+                    remap_fn,
+                    model_to_ckpt,
+                )
 
-        if not weight_map:
+        if not model_to_ckpt:
             raise RuntimeError(
-                "Distributed layerwise offload could not find safetensors "
+                "Distributed layerwise offload could not find compatible "
-                f"weights at {model_path}. Expected *.safetensors.index.json "
+                "safetensors weights in the configured model sources."
-                "or model.safetensors. Ensure the checkpoint path is correct."
             )
+        logger.info(
+            "Built mmap source map: %d model tensors from %d checkpoint keys",
+            len(model_to_ckpt),
+            indexed_keys,
+        )
+        return model_to_ckpt
 
-        logger.info("Loading DiT weights via mmap (meta → page cache, no RSS): %d keys", len(weight_map))
+    def _load_weights_via_mmap(self, pipeline: nn.Module, modules) -> None:
+        """Load DiT checkpoint tensors as file-backed safetensors views.
+
+        When the transformer is created on meta device, this method
+        replaces meta params with mmap views of the checkpoint files.
+        The views point to OS page cache shared across ranks. AllGather mode
+        copies only the rank's 1/dp_size shard to a persistent private buffer;
+        rank-local mode retains the views and packs one block at a time into
+        bounded staging storage.
+
+        Non-DiT modules (VAE, encoders) are NOT affected — they were
+        created on CPU with real weights via from_pretrained.
+        """
+        from safetensors import safe_open
+
+        model_to_ckpt = self._build_mmap_model_map(pipeline)
+        logger.info(
+            "Loading DiT weights via mmap (meta -> shared page cache): %d tensors",
+            len(model_to_ckpt),
+        )
 
         # --- Convert DiT modules to meta device ---
         # The transformer was created normally (with random weights and
@@ -812,25 +1076,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 len(saved_buffers.get(id(dit_module), {})),
             )
 
-        # Build reverse mapping {model_param_name: (ckpt_key, file_path)} using
-        # the pipeline's _remap_ckpt_key (if available).  This handles the
-        # model-specific checkpoint key remapping (e.g., Cosmos3's
-        # layers.0.mlp_moe_gen.gate_proj → transformer.gen_layers.0.mlp.gate_proj).
-        remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
-        model_to_ckpt: dict[str, tuple[str, str]] = {}
-        if remap_fn is not None:
-            for ckpt_key, file_path in weight_map.items():
-                model_name = remap_fn(ckpt_key)
-                if model_name is not None:
-                    model_to_ckpt[model_name] = (ckpt_key, file_path)
-            logger.info(
-                "Built reverse remap: %d model params from %d checkpoint keys", len(model_to_ckpt), len(weight_map)
-            )
-        else:
-            # No remap function — use checkpoint keys directly
-            for ckpt_key, file_path in weight_map.items():
-                model_to_ckpt[ckpt_key] = (ckpt_key, file_path)
-
         # Collect all meta params in DiT modules
         dit_ids = set()
         for dit_module in modules.dits:
@@ -850,8 +1095,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             for block in blocks:
                 block_module_ids.update(id(m) for m in block.modules())
 
-        for name, param in pipeline.named_parameters():
+        for name, target in chain(pipeline.named_parameters(), pipeline.named_buffers()):
-            if not (hasattr(param, "is_meta") and param.is_meta):
+            if not (hasattr(target, "is_meta") and target.is_meta):
                 continue
 
             # Check if this param is inside a block (will be handled by
@@ -895,13 +1140,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     parent = parent[int(part)]
                 else:
                     parent = getattr(parent, part)
-            replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad)
+            leaf_name = parts[-1]
-            self._attach_mmap_param_attrs(name, replacement, param)
+            if leaf_name in parent._parameters:
-            parent._parameters[parts[-1]] = replacement
+                replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad)
+                self._attach_mmap_param_attrs(name, replacement, target)
+                parent._parameters[leaf_name] = replacement
+            else:
+                parent._buffers[leaf_name] = tensor
             loaded += 1
             loaded_names.add(name)
 
-        logger.info("Mmap weight loading: %d non-block params loaded, %d skipped", loaded, skipped)
+        logger.info("Mmap weight loading: %d non-block tensors loaded, %d skipped", loaded, skipped)
 
         # Now load DiT block params from safetensors using mmap views.
         # Each block param is assigned an mmap view (no RSS).
@@ -927,8 +1176,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 # e.g. "transformer.language_model.layers.0" or "transformer.gen_layers.0"
                 block_full_prefix = f"{dit_name}.{blocks_attr}.{block_idx}"
 
-                for bname, bparam in block.named_parameters():
+                for bname, block_target in chain(block.named_parameters(), block.named_buffers()):
-                    if not (hasattr(bparam, "is_meta") and bparam.is_meta):
+                    if not (hasattr(block_target, "is_meta") and block_target.is_meta):
                         continue
 
                     # Use reverse remap to find checkpoint key
@@ -952,29 +1201,43 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                             parent = parent[int(part)]
                         else:
                             parent = getattr(parent, part)
-                    replacement = torch.nn.Parameter(tensor, requires_grad=bparam.requires_grad)
+                    leaf_name = bparts[-1]
-                    self._attach_mmap_param_attrs(full_param_name, replacement, bparam)
+                    if leaf_name in parent._parameters:
-                    parent._parameters[bparts[-1]] = replacement
+                        replacement = torch.nn.Parameter(tensor, requires_grad=block_target.requires_grad)
+                        self._attach_mmap_param_attrs(full_param_name, replacement, block_target)
+                        parent._parameters[leaf_name] = replacement
+                    else:
+                        parent._buffers[leaf_name] = tensor
                     block_loaded += 1
                     loaded_names.add(full_param_name)
 
-        logger.info("Mmap block loading: %d params loaded, %d skipped", block_loaded, skipped)
+        logger.info("Mmap block loading: %d tensors loaded, %d skipped", block_loaded, skipped)
 
         # Strict validation: check that all meta params were loaded.
         # This replaces the validation that AutoWeightsLoader.load_weights()
         # performs in the regular load path.
         remaining_meta = [
-            name for name, param in pipeline.named_parameters() if hasattr(param, "is_meta") and param.is_meta
+            name
+            for name, tensor in chain(pipeline.named_parameters(), pipeline.named_buffers())
+            if hasattr(tensor, "is_meta") and tensor.is_meta
         ]
         if remaining_meta:
             # Filter out params inside submodules that will be loaded later
             # by _load_module_weights_from_mmap (non-block submodules)
             dit_param_names = set()
-            for dit_module in modules.dits:
+            for dit_name, dit_module in zip(modules.dit_names, modules.dits):
-                for pname, _ in dit_module.named_parameters():
+                for pname, _ in chain(dit_module.named_parameters(), dit_module.named_buffers()):
-                    dit_param_names.add(pname)
+                    dit_param_names.add(f"{dit_name}.{pname}")
             truly_missing = [n for n in remaining_meta if n in dit_param_names]
             if truly_missing:
+                if getattr(pipeline, "_supports_mmap_loading", None) is True:
+                    raise RuntimeError(
+                        "Explicit mmap loading left "
+                        f"{len(truly_missing)} DiT tensors on the meta device "
+                        f"(first 5: {truly_missing[:5]}). Check that every "
+                        "runtime tensor has a matching weights_sources prefix "
+                        "and checkpoint key."
+                    )
                 logger.warning(
                     "Mmap loading: %d params still on meta device after "
                     "loading (first 5: %s). These may be loaded later by "
@@ -1021,10 +1284,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                         parent = parent[int(part)]
                     else:
                         parent = getattr(parent, part)
-                # Place on the same device as the module's parameters
+                # Rank-local hooks collect every block buffer as a CPU source.
-                # Place on the offloader's target device (not the module's
+                # AllGather preserves the existing behavior of restoring these
-                # current device, which may still be meta at this point).
+                # generated buffers on the target device before shard packing.
-                parent._buffers[parts[-1]] = buf.to(self.device)
+                restore_device = torch.device("cpu") if self._using_rank_local_mmap else self.device
+                parent._buffers[parts[-1]] = buf.to(restore_device)
 
     def _load_module_weights_from_mmap(self, module: nn.Module, dit_name: str, child_name: str) -> None:
         """Load a non-block DiT submodule's weights from safetensors via mmap."""
@@ -1037,8 +1301,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         model_to_ckpt = getattr(self, "_mmap_model_to_ckpt", {})
         loaded = 0
 
-        for pname, param in module.named_parameters():
+        for pname, target in chain(module.named_parameters(), module.named_buffers()):
-            if not (hasattr(param, "is_meta") and param.is_meta):
+            if not (hasattr(target, "is_meta") and target.is_meta):
                 continue
 
             # Build full model param name and look up in reverse remap
@@ -1065,9 +1329,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     parent = parent[int(part)]
                 else:
                     parent = getattr(parent, part)
-            replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad)
+            leaf_name = parts[-1]
-            self._attach_mmap_param_attrs(full_name, replacement, param)
+            if leaf_name in parent._parameters:
-            parent._parameters[parts[-1]] = replacement
+                replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad)
+                self._attach_mmap_param_attrs(full_name, replacement, target)
+                parent._parameters[leaf_name] = replacement
+            else:
+                parent._buffers[leaf_name] = tensor
             loaded += 1
 
         if loaded > 0:
@@ -1265,6 +1533,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             self.comm_stream,
             self.config.pin_cpu_memory,
             shared_buffers=[None, None],
+            rank_local_mmap=self._using_rank_local_mmap,
         )
         sub_hooks = [last_hook]
         for i, block in enumerate(blocks[:-1]):
@@ -1280,6 +1549,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 self.comm_stream,
                 self.config.pin_cpu_memory,
                 shared_buffers=[None, None],
+                rank_local_mmap=self._using_rank_local_mmap,
             )
             sub_hooks.append(hook)
 
@@ -1402,33 +1672,49 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 self.config.dlo_resident_layers,
             )
 
-        # Load weights via mmap for DLO+AllGather.
+        # Select checkpoint mmap independently from the DLO transfer protocol.
-        # TODO(offload): Add a rank-local mmap path for dlo_no_use_allgather.
+        # AllGather consumes rank shards copied from mmap.  Without AllGather,
-        # It must apply declarative checkpoint-to-runtime adapters before the
+        # TP1/no-HSDP workers retain the same file-backed tensors and pack only
-        # standard TP loader shards each block, and must cover staged encoders
+        # the next block into bounded, rank-private pinned staging buffers.
-        # and VAEs as well as DiT blocks. The current AllGather mmap path is
+        # TP/HSDP rank-local layouts still come from the standard loader.
-        # deliberately not reused because it changes the weight layout.
+        # The loader and backend share should_select_dlo_mmap(), so this
-        # Gate condition MUST match diffusers_loader.py:
+        # decision stays aligned with whether load_weights() was skipped.
-        #   supports_mmap_loading(pipeline) and not _has_online_quant
+        # When selection is False, the loader has already loaded weights via
-        # When the gate is False, the loader has already loaded weights via
         # regular load_weights() + _process_weights_after_loading().
-        if self.config.dlo_use_allgather and self.dp_size > 1:
+        _has_online_quant = has_online_quantization(pipeline)
-            _has_online_quant = any(
+        if self.config.dlo_use_allgather and _has_online_quant:
-                getattr(getattr(module, "quant_method", None), "uses_meta_device", False)
+            raise ValueError(
-                for module in pipeline.modules()
+                "Online quantization is incompatible with DLO+AllGather: "
+                "the sharding + AllGather mechanism flattens weights by "
+                "dtype, which breaks quantized weight/scale layouts. "
+                "Please use --dlo-no-use-allgather or disable online "
+                "quantization."
             )
-            if _has_online_quant:
+
-                raise ValueError(
+        mmap_supported = supports_mmap_loading(pipeline)
-                    "Online quantization is incompatible with DLO+AllGather: "
+        self._using_mmap = should_select_dlo_mmap(
-                    "the sharding + AllGather mechanism flattens weights by "
+            mmap_supported=mmap_supported,
-                    "dtype, which breaks quantized weight/scale layouts. "
+            dlo_use_allgather=self.config.dlo_use_allgather,
-                    "Please use --dlo-no-use-allgather or disable online "
+            has_online_quant=_has_online_quant,
-                    "quantization."
+            tensor_parallel_size=self.config.tensor_parallel_size,
+            use_hsdp=self.config.use_hsdp,
+        )
+        self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1
+        if self._using_mmap:
+            self._load_weights_via_mmap(pipeline, modules)
+            if self._using_rank_local_mmap:
+                logger.info(
+                    "DLO rank-local mmap storage enabled: checkpoint pages are "
+                    "node-shared; each worker owns only two bounded host staging slots"
                 )
-            if supports_mmap_loading(pipeline):
+        elif not mmap_supported:
-                self._load_weights_via_mmap(pipeline, modules)
+            logger.info("Weights loaded via regular loader - model does not support mmap")
-            else:
+        elif not self.config.dlo_use_allgather and (self.config.tensor_parallel_size != 1 or self.config.use_hsdp):
-                logger.info("Weights loaded via regular loader — skipping mmap (model does not support mmap)")
+            logger.info(
+                "Rank-local mmap disabled for TP=%d/HSDP=%s; using standard-loader tensors",
+                self.config.tensor_parallel_size,
+                self.config.use_hsdp,
+            )
 
         # Keep VAE/encoders on CPU; move to GPU on-demand via hooks.
         # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9)
@@ -1528,6 +1814,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 self.comm_stream,
                 self.config.pin_cpu_memory,
                 shared_buffers=[None, None],
+                rank_local_mmap=self._using_rank_local_mmap,
             )
 
             block_hooks: list[DistributedLayerwiseOffloadHook] = [last_hook]
@@ -1544,6 +1831,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     self.comm_stream,
                     self.config.pin_cpu_memory,
                     shared_buffers=[None, None],
+                    rank_local_mmap=self._using_rank_local_mmap,
                 )
                 block_hooks.append(hook)
 
@@ -1572,11 +1860,15 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 self.device,
                 self.copy_stream,
                 self.config.pin_cpu_memory,
+                rank_local_mmap=self._using_rank_local_mmap,
+                defer_staging=bool(self._all_hook_groups),
             )
             pipeline._dlo_residency_controller = self
 
         if not self._all_hook_groups:
             self.enabled = bool(self._resident_blocks)
+            if self._using_mmap and not self.enabled:
+                self._release_mmap_handles()
             return
 
         # Unified allocation: 2 shared output buffers + 2 shared shard buffers
@@ -1590,6 +1882,21 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         unified_shard_buffers = None
         if self.dp_size > 1:
             unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
+        unified_cpu_staging = None
+        cpu_staging_events = None
+        if self._using_rank_local_mmap:
+            unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
+                all_hooks,
+                self._resident_layer_group,
+            )
+            cpu_staging_events = [None, None]
+            if self._resident_layer_group is not None:
+                # Resident and streamed layers execute in the same stage and
+                # reuse the same host slots. Events serialize slot reuse.
+                self._resident_layer_group._cpu_staging_buffers = [
+                    buffers for buffers in unified_cpu_staging if buffers is not None
+                ]
+                self._resident_layer_group._cpu_staging_events = cpu_staging_events
 
         # Shared slot-group tracker: _shared_slot_group[slot] = group_id
         # that last wrote to that slot.  Group-first hooks use this to
@@ -1602,6 +1909,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 hook._owns_buffers = False
                 if unified_shard_buffers is not None:
                     hook.gpu_shard_buffers = unified_shard_buffers
+                if unified_cpu_staging is not None and cpu_staging_events is not None:
+                    hook.cpu_staging_buffers = unified_cpu_staging
+                    hook.cpu_staging_events = cpu_staging_events
                 hook._group_id = group_idx
                 hook._shared_slot_group = shared_slot_group
 
@@ -1628,9 +1938,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         self.enabled = True
 
-        # Release mmap file handles — _shard_and_pin has copied all shards,
+        if self._using_mmap and not self._using_rank_local_mmap:
-        # params now point to offload placeholders (not mmap views).
+            # AllGather mode copied each rank's persistent shard, so the source
-        self._release_mmap_handles()
+            # mappings are no longer needed. Rank-local mode retains them as
+            # the node-shared host master until disable().
+            self._release_mmap_handles()
 
         # Assign GPU buffers to sharded on-demand modules (VAE/encoders).
         # Each module gets a dedicated input (shard-sized) and output
@@ -1674,9 +1986,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             pass
 
     def disable(self) -> None:
-        if not self.enabled:
+        if not self.enabled and not hasattr(self, "_mmap_file_cache"):
             return
 
+        if self._using_rank_local_mmap:
+            current_omni_platform.synchronize()
+
         for blocks in self._blocks:
             for block in blocks:
                 remove_distributed_block_hook(block)
@@ -1691,6 +2006,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._all_hook_groups.clear()
         self._resident_blocks.clear()
         self._resident_layer_group = None
+        self._release_mmap_handles()
+        self._using_mmap = False
+        self._using_rank_local_mmap = False
         self.enabled = False
         logger.info("Distributed layer-wise offloading disabled")
 
@@ -1730,6 +2048,41 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         )
         return shared_buffers
 
+    @staticmethod
+    def _allocate_shared_cpu_staging_buffers(
+        hooks: list[DistributedLayerwiseOffloadHook],
+        resident_group: PinnedResidentLayerGroup | None = None,
+    ) -> list[dict[torch.dtype, torch.Tensor] | None]:
+        """Allocate two bounded host slots for rank-local mmap -> device copies."""
+        max_sizes: dict[torch.dtype, int] = {}
+        for hook in hooks:
+            for dtype, metas in hook.metadata.items():
+                total = sum(meta["numel"] for meta in metas)
+                max_sizes[dtype] = max(max_sizes.get(dtype, 0), total)
+        if resident_group is not None:
+            for state in resident_group._states:
+                for dtype, metas in state["metadata"].items():
+                    total = sum(meta["numel"] for meta in metas)
+                    max_sizes[dtype] = max(max_sizes.get(dtype, 0), total)
+
+        pin_memory = hooks[0].pin_memory if hooks else bool(resident_group and resident_group.pin_memory)
+        shared_staging: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]
+        for slot in range(2):
+            buffers: dict[torch.dtype, torch.Tensor] = {}
+            for dtype, total_numel in max_sizes.items():
+                buffer = torch.empty(total_numel, dtype=dtype, device="cpu")
+                if pin_memory:
+                    buffer = buffer.pin_memory()
+                buffers[dtype] = buffer
+            shared_staging[slot] = buffers
+
+        logger.info(
+            "Allocated 2 shared host staging buffers for rank-local mmap (max block size: %s, pinned=%s)",
+            {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_sizes.items()},
+            pin_memory,
+        )
+        return shared_staging
+
     @staticmethod
     def _allocate_shared_shard_buffers(
         hooks: list[DistributedLayerwiseOffloadHook],

vllm_omni/diffusion/offloader/offload_plan.py

@@ -60,15 +60,51 @@ def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None:
 def supports_mmap_loading(pipeline: nn.Module) -> bool:
     """Whether the pipeline supports mmap weight loading.
 
-    A pipeline supports mmap if any module defines ``_remap_ckpt_key``,
+    A pipeline supports mmap when either:
-    which maps checkpoint keys to model parameter names.  Models without
+
-    this method must use the regular ``load_weights()`` path.
+    * a module defines ``_remap_ckpt_key``, which maps checkpoint keys to
+      model parameter names; or
+    * the pipeline explicitly opts in with ``_supports_mmap_loading=True``
+      and declares ``weights_sources``.  In that case each source prefix is
+      used as the model namespace.
+
+    The explicit opt-in is important: many pipelines declare
+    ``weights_sources`` but still require arbitrary loader callbacks and must
+    stay on the regular ``load_weights()`` path.
 
     This check is shared between ``diffusers_loader.py`` (to decide whether
     to skip ``load_weights``) and ``DistributedLayerwiseOffloadBackend.enable()``
     (to decide whether to call ``_load_weights_via_mmap``), ensuring both
     paths use the same gate condition.
     """
-    return bool(getattr(pipeline, "_supports_mmap_loading", True)) and any(
+    declared_support = getattr(pipeline, "_supports_mmap_loading", None)
-        callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules()
+    if declared_support is False:
-    )
+        return False
+
+    has_remap = any(callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules())
+    if has_remap:
+        return True
+
+    return declared_support is True and bool(getattr(pipeline, "weights_sources", ()))
+
+
+def has_online_quantization(model: nn.Module) -> bool:
+    """Whether any module defers online-quantized weights on ``meta``."""
+    return any(getattr(getattr(module, "quant_method", None), "uses_meta_device", False) for module in model.modules())
+
+
+def should_select_dlo_mmap(
+    *,
+    mmap_supported: bool,
+    dlo_use_allgather: bool,
+    has_online_quant: bool,
+    tensor_parallel_size: int,
+    use_hsdp: bool,
+) -> bool:
+    """Whether loader and backend should enter DLO's mmap path.
+
+    AllGather preserves its existing fail-closed topology validation inside
+    the mmap loader. The extra TP/HSDP gate applies to rank-local mmap only.
+    """
+    rank_local_compatible = tensor_parallel_size == 1 and not use_hsdp
+    return mmap_supported and not has_online_quant and (dlo_use_allgather or rank_local_compatible)