[Diffusion] Add loader-owned host-weight plans for DLO (TP=1) (#6213)

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

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -4,7 +4,7 @@ This document describes distributed layerwise offload (DLO) for diffusion
 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.
+complete rank-local block sources and avoid an additional collective.
 
 For user-facing commands, see the
 [distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
@@ -14,10 +14,20 @@ and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/rec
 
 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
+`--dlo-no-use-allgather` path streams complete blocks independently and adds no
-for standard-loader sharding, workstation bring-up, and systems where an
+DLO weight collective.
-additional DLO collective is undesirable, but it does not reduce host weight
+
-storage across ranks.
+Host storage is selected separately from the transfer protocol. The loader can
+produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
+otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
+the same node can share immutable checkpoint pages when direct mmap is
+selected, while the ordinary-loader fallback still keeps a private runtime
+copy per process.
+
+The Phase A shared-mmap support boundary is TP1. TP greater than one is an
+ordinary-loader compatibility path: DLO can consume the resulting TP-local
+tensors, but those configurations do not use checkpoint mmap and must not be
+used to claim shared-mmap host-memory savings.
 
 The compatibility matrix below describes the current implementation. The
 unit-level guards are covered, but not every parallelism combination has a
@@ -41,6 +51,36 @@ TP is deliberately not used as DLO's AllGather group. HSDP has its own
 parameter-sharding lifecycle and is not allowed to be sharded a second time by
 DLO's AllGather path.
 
+### The loader owns host-weight planning
+
+Before it decides whether ordinary weight materialization can be skipped, the
+diffusion loader builds one `HostWeightPlan`. A direct-checkpoint mmap plan is
+accepted only when preflight proves all of the following:
+
+- every required DiT parameter and persistent buffer has exactly one source;
+- runtime names, checkpoint keys, shapes, and dtypes match;
+- the runtime topology is TP1 without HSDP or online quantization; and
+- every custom loader operation is represented by a loader-owned checkpoint
+  adapter.
+
+The exact plan object is handed to DLO. The backend does not rescan checkpoint
+files, repeat the capability decision, or reconstruct names from its block
+topology. If preflight fails, the loader materializes weights normally and DLO
+consumes those runtime tensors.
+
+The plan owns only dedicated DiT component sources. If a pipeline also exposes
+ordinary sources for a text encoder or another non-DiT component, the loader
+still consumes those sources and includes their loaded names in its strict
+coverage check. Only the source prefixes covered by the plan skip ordinary
+materialization. A source that mixes DiT and non-DiT weights fails closed to
+the complete ordinary-loader path because it cannot be skipped safely as a
+unit.
+
+This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline
+flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
+transforms. Model-specific direct-layout knowledge, when required, lives in a
+checkpoint adapter beside the ordinary loader.
+
 ### AllGather path
 
 With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -61,6 +101,15 @@ Buffers:    [current slot]       [prefetch slot]       [current slot]
 The backend uses two shared device buffers, so accelerator weight residency is
 bounded by the largest streamed blocks rather than the complete model.
 
+When direct checkpoint mmap is selected, the checkpoint mappings are only the
+source used to prepare each rank's persistent shard. They can be closed after
+shard preparation. Across the AllGather group, those private shards total
+approximately one runtime model copy.
+
+An effective DLO group size of one performs no collective, even when
+`dlo_use_allgather=True`; it follows the rank-local transfer path described
+below.
+
 When DP is greater than one, the engine can process one request per DP rank in
 the same denoising wave. Because AllGather is a collective, all participating
 requests must take the same execution path at every denoising step.
@@ -68,9 +117,20 @@ 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 blocks using H2D copies only. The host backing may be
-weights, including TP-local tensors or HSDP-managed parameters. DLO then
+either a loader-approved checkpoint mapping or ordinary runtime tensors.
-streams those rank-local tensors block by block using H2D copies only.
+
+For direct mmap, each process retains immutable safetensors views and uses two
+bounded pinned host staging slots. Processes on the same node that map the same
+files share physical checkpoint pages through the OS page cache. This removes
+the persistent private full-model copy per pure-DP process, but each process
+still packs and transfers every complete block. Sharing is node-local; each
+node has its own page cache.
+
+When direct mmap preflight fails, the regular model loader remains responsible
+for preparing each rank's weights, including TP-local tensors or HSDP-managed
+parameters. In that fallback, each pure-DP process keeps a private full runtime
+copy.
 
 This mode means:
 
@@ -80,21 +140,18 @@ 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 share one checkpoint-backed copy per node when direct
-  shared OS page-cache behavior.
+  mmap is selected; the ordinary-loader fallback keeps one private runtime
+  copy per rank.
 - The scheduler does not require a synchronized DP request wave for DLO.
 
-This path is intentionally not implemented by reusing the AllGather mmap
-loader. It relies on the standard loader so model-specific TP/HSDP transforms
-remain intact.
-
 ## 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. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
 | **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. |
+| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
 | **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. |
 
 ### Combined dimensions
@@ -122,11 +179,15 @@ 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
+Direct checkpoint mmap can back either transfer path. It is currently limited
-model has a compatible checkpoint layout. Online quantization is incompatible
+to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
-with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
+ordinary loader. Online quantization remains incompatible with DLO AllGather;
-quantization. Models that do not meet the mmap requirements use the regular
+use `--dlo-no-use-allgather` or disable online quantization.
-loader path.
+
+A normalized runtime mmap cache, built through the ordinary loader, is the
+proposed general mechanism for sharing transformed TP or quantized layouts.
+That cache and its publication/lifecycle protocol are intentionally outside
+this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
 
 ## Validation coverage
 
@@ -134,15 +195,75 @@ Current source-level validation includes:
 
 - HSDP + DLO + AllGather rejection;
 - HSDP + DLO without AllGather acceptance at configuration level;
-- TP rejection in the DLO+AllGather mmap path;
+- loader preflight fallback for TP, HSDP, online quantization, unknown custom
+  loaders, missing keys, and shape/dtype mismatches;
+- exact loader-to-backend plan transfer and ordinary-loader fallback;
+- rank-local mmap source retention, bounded two-slot staging, and adapter
+  transforms without parameter-side flags;
 - resident-layer requests requiring no-AllGather;
 - DP request-wave validation for denoising-step compatibility;
 - sharding, double-buffer, AllGather-size, and heterogeneous-block regression
   tests.
 
-The highest-value missing coverage is end-to-end numerical comparison against
+### B300 parallel-topology smoke matrix
-ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
+
-on the target CUDA/NCCL or CANN/HCCL hardware.
+A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed,
+CUDNN attention backend, 256x256 output, two denoising steps, and
+`dlo_resident_layers=0`. The TP2 rows used DiT DP2xTP2 with the text encoder and
+VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap
+or shared-mmap host-memory savings.
+
+| Configuration | Result | Warm E2E | Peak device memory | Host PSS |
+|---|---:|---:|---:|---:|
+| DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB |
+| DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB |
+| DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB |
+| DP2xTP2 no-AllGather | Passed, 1 request | 3.51 s | 11.88 GiB | 314.01 GiB |
+
+Within each topology, the AllGather and no-AllGather video and audio outputs
+were byte-identical. All four runs completed without an `ERROR` or traceback
+and released their device allocations. For DP4xTP1, no-AllGather direct mmap
+reduced total PSS by 24.22 GiB (11.4%) and `Private_Dirty` from 211.33 to
+125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the
+ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above
+AllGather, because DP replicas did not share checkpoint-backed runtime
+weights. This is a functional and memory smoke test, not a production-quality
+performance or output-quality benchmark.
+
+### Host-memory measurement
+
+A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
+ordinary-loader fallback with direct mmap. Both runs used
+DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a
+256x256 four-second request. The ordinary-loader workers were sampled after
+initialization. The mmap workers were sampled after one completed request, so
+the checkpoint working set had been faulted into the page cache; this is the
+more conservative point for mmap.
+
+The values below come from `/proc/<worker>/smaps_rollup` and include the whole
+worker, not only the DiT. The stable rank-to-rank difference comes from other
+pipeline components, so each worker should be compared with the same worker in
+the other storage mode.
+
+| Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction |
+|---|---:|---:|---:|---:|---:|
+| DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB |
+| DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB |
+| **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** |
+
+The direct-mmap workers each reported 62.45 GiB `Shared_Clean` but only
+31.20 GiB `Pss_File`, which is the proportional charge expected when the same
+resident checkpoint pages are mapped by two workers. `Private_Dirty` also fell
+from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
+1, a reduction of about 97–98 GiB per worker. RSS understates this benefit
+because it counts a shared physical page in every process that maps it; summed
+PSS is the appropriate node-memory comparison.
+
+The highest-value missing coverage is broader end-to-end numerical and
+lifecycle comparison against ordinary layerwise offload for DP+SP,
+HSDP+SP+no-AllGather, and TP greater than one across additional models and
+target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
+change the Phase A direct-mmap TP1 support boundary.
 
 ## Recommendations
 
@@ -150,7 +271,9 @@ 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 independent replica execution is required. TP1
-  the expectation of higher host-memory use and lower validation confidence.
+  direct-mmap deployments can share checkpoint pages per node; other layouts
+  retain the ordinary loader's private host memory behavior and are outside
+  the Phase A shared-mmap support scope.
 - 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/offloader/distributed_layerwise_offload.md

@@ -3,7 +3,9 @@
 Distributed layerwise offloading (DLO) extends block streaming to multi-device
 deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
 the host weights and reconstructs each layer at runtime. Without AllGather,
-each rank streams its standard-loader rank-local weights independently.
+each rank streams a complete block independently. Compatible TP1 deployments
+can share checkpoint-backed host pages among processes on the same node;
+otherwise DLO streams the ordinary loader's rank-local tensors.
 
 See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
 for the implementation contract and compatibility matrix.
@@ -59,21 +61,48 @@ omni = Omni(
 | `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
 | `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
 | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
-| `--dlo-no-use-allgather` | Keep standard-loader rank-local weights | `false` |
+| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
 | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
 
-## mmap weight loading
+## Host-weight loading
 
-The DLO plus AllGather path:
+The diffusion loader chooses host storage before DLO is enabled. It first
+attempts to build a complete, validated direct-checkpoint mmap plan. If names,
+coverage, shape, dtype, topology, or loader-callback compatibility cannot be
+proven, it runs the ordinary model loader instead. DLO consumes that result and
+does not make a second checkpoint-compatibility decision.
+
+The shared-mmap optimization in this phase is supported only with TP1. TP
+greater than one falls back before model mutation to ordinary TP-aware loading.
+DLO may still consume those TP-local tensors, but this is a compatibility path:
+it does not share checkpoint-backed runtime weights across DP replicas and
+provides no shared-mmap host-memory guarantee.
+
+The mmap plan skips only dedicated DiT weight sources. Other component sources,
+such as a text encoder loaded through the shared diffusion loader, continue to
+use their ordinary component loader. A checkpoint source that mixes DiT and
+non-DiT weights falls back completely rather than leaving an unplanned
+component uninitialized.
+
+With direct checkpoint mmap, the loader:
 
 1. saves non-persistent buffers such as RoPE frequencies;
 2. moves the normally created transformer to the meta device;
 3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
-4. calls model-specific `post_load_weights()` conversions; and
+4. applies any loader-owned bounded layout adapters while packing blocks;
-5. restores the saved non-persistent buffers.
+5. restores saved non-persistent buffers; and
+6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle
+   hooks.
 
-This avoids one full checkpoint RSS copy per rank and does not require
+For AllGather with a group larger than one, each process copies only its
-model-specific loading code.
+persistent shard and then releases the source mapping. For no-AllGather, each
+process keeps the mapping open and packs complete blocks through two bounded
+pinned staging slots. Processes mapping the same files on one node share the
+immutable pages; no-AllGather still performs a complete-block H2D copy in each
+process.
+
+When the effective DLO group size is one, `dlo_use_allgather=True` does not
+perform a collective and uses the same rank-local transfer behavior.
 
 ## Declarative topology
 
@@ -103,16 +132,22 @@ must enter each collective.
 
 ## Limitations
 
-- Online FP8 quantization is rejected with the DLO plus AllGather mmap path.
+- Direct checkpoint mmap currently requires TP1. TP greater than one is
-  Use `--dlo-no-use-allgather` or disable online quantization.
+  outside the Phase A shared-mmap support scope and falls back before model
-- Tensor parallel size greater than one is rejected in the mmap path because
+  mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
-  it bypasses TP-aware loader callbacks. The no-AllGather path retains the
+  but it provides no shared-mmap host-memory benefit or guarantee.
-  standard loader and remains experimental with TP.
 - HSDP plus AllGather is rejected to avoid double sharding. HSDP without
   AllGather has limited end-to-end validation.
+- Online quantization uses the ordinary loader with no-AllGather. It remains
+  incompatible with DLO AllGather.
 - Resident leading layers require `--dlo-no-use-allgather` and a model
   `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.
 
+Sharing transformed TP or quantized runtime layouts through a normalized mmap
+cache is a follow-up design in
+[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
+the direct-checkpoint path.
+
 See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
 for an end-to-end example.

examples/offline_inference/minimax_h3/dlo_lifecycle.py

@@ -1,21 +1,32 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
+"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
+
+The DLO modes accept any positive data-parallel size.  With AllGather enabled,
+the script submits one complete DP wave and verifies recovery after an invalid
+wave.  It also checks that every mode shuts down all worker processes cleanly.
 
 Examples:
 
     VLLM_WORKER_MULTIPROC_METHOD=spawn \
     VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
     VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
+    CUDA_VISIBLE_DEVICES=0,1,2,3 \
+    python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
+        --model /path/to/MiniMax-H3/FL2VA --mode dlo --dp-size 4
+
+    # 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 \
+    python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
-        --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
+        --model /path/to/MiniMax-H3/FL2VA --mode dlo-no-allgather --dp-size 2
 
     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 \
+    python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
-        --model /path/to/MiniMax-H3/FL2VA --mode request
+        --model /path/to/MiniMax-H3/FL2VA --mode request --tp-size 2
 """
 
 from __future__ import annotations
@@ -40,14 +51,33 @@ DEFAULT_PROMPTS = (
 )
 
 
+def positive_int(value: str) -> int:
+    parsed = int(value)
+    if parsed < 1:
+        raise argparse.ArgumentTypeError("must be a positive integer")
+    return parsed
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
     parser.add_argument(
         "--mode",
-        choices=("dlo-dp2", "request"),
+        choices=("dlo", "dlo-no-allgather", "request"),
         required=True,
-        help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
+        help="DLO (with or without AllGather) or ordinary non-DLO request mode",
+    )
+    parser.add_argument(
+        "--dp-size",
+        type=positive_int,
+        default=2,
+        help="DLO data-parallel size (for example 1, 2, 4, or 8; default: 2)",
+    )
+    parser.add_argument(
+        "--tp-size",
+        type=positive_int,
+        default=2,
+        help="Tensor-parallel size used only by request mode (default: 2)",
     )
     parser.add_argument("--steps", type=int, default=2)
     parser.add_argument("--duration", type=float, default=5.0)
@@ -60,10 +90,11 @@ def parse_args() -> argparse.Namespace:
 
 
 def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
+    is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
     common: dict[str, Any] = {
         "model": args.model,
         "trust_remote_code": True,
-        "num_gpus": 2,
+        "num_gpus": args.dp_size if is_dlo else args.tp_size,
         "usp": 1,
         "ring": 1,
         "vae_parallel_mode": "tile",
@@ -74,22 +105,22 @@ 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 is_dlo:
         common.update(
             tensor_parallel_size=1,
-            data_parallel_size=2,
+            data_parallel_size=args.dp_size,
             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",
             dlo_resident_layers=0,
         )
     else:
         common.update(
-            tensor_parallel_size=2,
+            tensor_parallel_size=args.tp_size,
             data_parallel_size=1,
-            text_encoder_tp_size=2,
+            text_encoder_tp_size=args.tp_size,
-            vae_patch_parallel_size=2,
+            vae_patch_parallel_size=args.tp_size,
             enable_distributed_layerwise_offload=False,
         )
     return common
@@ -172,7 +203,8 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
         "engine_kwargs": engine_kwargs(args),
     }
     try:
-        if args.mode == "dlo-dp2":
+        uses_dp_wave = args.mode == "dlo" and args.dp_size > 1
+        if uses_dp_wave:
             started = time.perf_counter()
             asymmetric = await asyncio.gather(
                 generate_one(
@@ -182,12 +214,15 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
                     prompt="",
                     seed=1001,
                 ),
-                generate_one(
+                *(
-                    engine,
+                    generate_one(
-                    args,
+                        engine,
-                    request_id="invalid-peer",
+                        args,
-                    prompt=DEFAULT_PROMPTS[0],
+                        request_id=f"invalid-peer-{index}",
-                    seed=1002,
+                        prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
+                        seed=1001 + index,
+                    )
+                    for index in range(1, args.dp_size)
                 ),
                 return_exceptions=True,
             )
@@ -195,18 +230,18 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
             summary["asymmetric_errors"] = [
                 f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
             ]
-            if len(summary["asymmetric_errors"]) != 2:
+            if len(summary["asymmetric_errors"]) != args.dp_size:
-                raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
+                raise RuntimeError("Every request in the asymmetric DP wave must fail before dispatch")
 
         started = time.perf_counter()
-        request_count = 2 if args.mode == "dlo-dp2" else 1
+        request_count = args.dp_size if uses_dp_wave else 1
         outputs = await asyncio.gather(
             *(
                 generate_one(
                     engine,
                     args,
                     request_id=f"recovery-{index}",
-                    prompt=DEFAULT_PROMPTS[index],
+                    prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
                     seed=2000 + index,
                 )
                 for index in range(request_count)

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -16,6 +16,11 @@ from vllm.config.load import LoadConfig
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+    HostWeightPlan,
+    HostWeightPlanResult,
+    TensorBinding,
+)
 from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.registry import initialize_model
 
@@ -258,6 +263,137 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     assert get_current_diffusion_config_or_none() is None
 
 
+def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    plan = HostWeightPlan(
+        backing_kind="checkpoint_mmap",
+        bindings={},
+    )
+    loaded_ordinary_weights = False
+
+    def load_weights(_model):
+        nonlocal loaded_ordinary_weights
+        loaded_ordinary_weights = True
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = load_weights  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(plan),
+    )
+
+    assert loader.load_model(load_device="cpu") is model
+    assert not loaded_ordinary_weights
+    assert loader.take_host_weight_plan() is plan
+    assert loader.take_host_weight_plan() is None
+
+
+def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
+    class MixedSourceModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.transformer = nn.Linear(2, 2, bias=False)
+            self.text_encoder = nn.Linear(2, 2, bias=False)
+            self.weights_sources = (
+                DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="transformer."),
+                DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="text_encoder."),
+            )
+            self.loaded_weight_names: list[str] = []
+
+        def load_weights(self, weights):
+            self.loaded_weight_names = [name for name, _ in weights]
+            return set(self.loaded_weight_names)
+
+    model = MixedSourceModel()
+    plan = HostWeightPlan(
+        backing_kind="checkpoint_mmap",
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="weight",
+                file_path="unused",
+            )
+        },
+        planned_source_prefixes=frozenset({"transformer."}),
+    )
+    requested_prefixes: list[str] = []
+
+    def get_weights(source, model=None):
+        del model
+        requested_prefixes.append(source.prefix)
+        yield source.prefix + "weight", torch.ones(2, 2)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader._get_weights_iterator = get_weights  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(plan),
+    )
+
+    assert loader.load_model(load_device="cpu") is model
+    assert requested_prefixes == ["text_encoder."]
+    assert model.loaded_weight_names == ["text_encoder.weight"]
+    assert loader.take_host_weight_plan() is plan
+
+
+def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    calls: list[str] = []
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: calls.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
+    )
+
+    assert loader.load_model(load_device="cpu") is model
+    assert calls == ["load", "process"]
+    assert loader.take_host_weight_plan() is None
+
+
 def test_hsdp_processes_quantized_weights_before_sharding(mocker):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
     from vllm_omni.diffusion.offloader.module_collector import PipelineModules

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py

@@ -191,6 +191,49 @@ def test_model_load_weights_transforms_before_calling_vllm_loader():
     ]
 
 
+def test_loader_adapter_declares_equivalent_direct_mmap_transform(monkeypatch):
+    from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
+        get_direct_mmap_adapter,
+    )
+    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",
+    )
+    transformer = object.__new__(h3.MiniMaxH3DiTModel)
+    nn.Module.__init__(transformer)
+    transformer.blocks = nn.ModuleList([nn.Module()])
+    transformer.blocks[0].attn = attention
+    pipeline = nn.Module()
+    pipeline.transformer = transformer
+
+    adapter = get_direct_mmap_adapter(pipeline)
+    assert adapter is not None
+    policy = adapter.policy_for(
+        "transformer.blocks.0.attn.qkv_proj.weight",
+        attention.qkv_proj.weight,
+    )
+    assert policy is not None
+    assert policy.allow_custom_loader
+    assert policy.transform is not None
+    checkpoint_weight = torch.arange(6, dtype=torch.float32).reshape(6, 1)
+
+    assert policy.transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
+    assert not hasattr(attention.qkv_proj.weight, "mmap_weight_transform")
+
+
 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

@@ -18,6 +18,11 @@ from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+    HostWeightPlan,
+    TensorBinding,
+    build_checkpoint_mmap_plan,
+)
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
 from vllm_omni.diffusion.offloader.block_discovery import (
     get_blocks_attr_names,
@@ -30,7 +35,10 @@ 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,
+)
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -48,6 +56,9 @@ class DummyEvent:
     def record(self, _stream) -> None:
         return None
 
+    def synchronize(self) -> None:
+        return None
+
 
 @contextmanager
 def dummy_stream(_stream):
@@ -351,8 +362,9 @@ class TestDistributedLayerwiseOffloadHook:
         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
+        def transform(tensor):
+            return tensor.flip(0)
 
         hook = DistributedLayerwiseOffloadHook(
             next_block=next_block,
@@ -361,17 +373,17 @@ class TestDistributedLayerwiseOffloadHook:
             dp_size=2,
             rank=0,
             pin_memory=False,
+            tensor_transforms={id(next_block.weight): transform},
         )
         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):
+    def test_transform_is_not_applied_without_a_loader_plan(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,
@@ -385,6 +397,82 @@ 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()
+
+        def transform(tensor):
+            return tensor.t()
+
+        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,
+            tensor_transforms={id(next_block.weight): transform},
+        )
+        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([[0.0, 2.0], [1.0, 3.0]]),
+        )
+        assert next_block.weight.stride() == (1, 2)
+        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):
@@ -439,23 +527,41 @@ class TestPinnedResidentLayerGroup:
         assert block.weight.stride() == expected_stride
         group.offload()
 
-    def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
+    def test_rank_local_mmap_resident_layers_retain_sources(self, patched_offload_runtime):
-        module = nn.Linear(2, 2, bias=False)
+        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())
 
-        def transform(tensor):
+        group = PinnedResidentLayerGroup(
-            return tensor.flip(0)
+            blocks,
+            device=torch.device("cpu"),
+            copy_stream=DummyStream(),
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
 
-        module.weight.mmap_weight_transform = transform
+        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
 
-        backend = object.__new__(DistributedLayerwiseOffloadBackend)
+        group.load()
-        backend._remember_mmap_param_attrs(module)
+        for block, weight in zip(blocks, expected):
-        module.to_empty(device="meta")
+            assert torch.equal(block.weight, weight)
 
-        replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float())
+        blocks[0].weight.data.zero_()
-        backend._attach_mmap_param_attrs("weight", replacement, module.weight)
+        group.offload()
+        group.load()
+        for block, weight in zip(blocks, expected):
+            assert torch.equal(block.weight, weight)
 
-        assert replacement.mmap_weight_transform is transform
+        group.offload()
-        assert replacement.mmap_weight_transform_pending is True
+        assert all(block.weight.numel() == 0 for block in blocks)
 
     def test_shared_allgather_output_is_narrowed_to_current_block(
         self, dist_group, patched_offload_runtime, monkeypatch
@@ -598,7 +704,6 @@ class TestMmapWeightLoading:
             OffloadConfig(
                 strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                 pin_cpu_memory=False,
-                model_path=str(tmp_path),
             ),
             torch.device("cpu"),
         )
@@ -606,13 +711,81 @@ class TestMmapWeightLoading:
             dits=[pipeline.transformer],
             dit_names=["transformer"],
         )
+        plan = HostWeightPlan(
+            backing_kind="checkpoint_mmap",
+            bindings={
+                name: TensorBinding(
+                    checkpoint_key=name,
+                    file_path=str(weight_file),
+                )
+                for name in weights
+            },
+        )
 
-        backend._load_weights_via_mmap(pipeline, modules)
+        backend._load_weights_via_mmap(pipeline, modules, plan)
 
         assert pipeline.transformer.post_load_calls == 1
         assert pipeline.transformer.time_embedder.weight.dtype == torch.float32
         assert pipeline.transformer.blocks[0].weight.dtype == torch.bfloat16
 
+    def test_allgather_flag_with_group_size_one_uses_rank_local_mmap(
+        self,
+        tmp_path,
+        patched_offload_runtime,
+    ):
+        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):
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+            @staticmethod
+            def _remap_ckpt_key(key):
+                return key
+
+        pipeline = Pipeline()
+        weights = {name: torch.ones_like(param) for name, param in pipeline.named_parameters()}
+        checkpoint_file = tmp_path / "model.safetensors"
+        save_file(weights, str(checkpoint_file))
+        result = build_checkpoint_mmap_plan(
+            pipeline,
+            dit_modules=(("transformer", pipeline.transformer),),
+            sources=(),
+            model_path=str(tmp_path),
+            tensor_parallel_size=1,
+            use_hsdp=False,
+            online_quantization=False,
+        )
+        assert result.plan is not None
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=1,
+                dlo_use_allgather=True,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=result.plan,
+        )
+
+        backend.enable(pipeline)
+
+        assert backend._using_rank_local_mmap
+        assert backend.dp_group is None
+        assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
+        backend.disable()
+
 
 class TestGetBlocksFromDit:
     def test_get_blocks_from_dit_single_block_attr(self):
@@ -878,23 +1051,123 @@ class TestOffloadPlan:
 
 
 class TestMmapValidation:
-    """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
+    """Tests for loader preflight and backend plan realization."""
 
-    def test_tp_rejected_when_tp_world_size_gt_1(self, tmp_path, patched_offload_runtime, monkeypatch):
+    def test_component_source_prefix_builds_model_namespace(self, tmp_path):
-        """DLO+AllGather should reject when TP world size > 1."""
+        source_root = tmp_path / "partition"
-        import json
+        checkpoint_dir = source_root / "transformer"
-        from types import SimpleNamespace
+        checkpoint_dir.mkdir(parents=True)
-
+        checkpoint_file = checkpoint_dir / "model.safetensors"
-        from safetensors.torch import save_file
+        save_file(
-
+            {"blocks.0.weight": torch.arange(4, dtype=torch.float32).reshape(1, 4)},
-        from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+            str(checkpoint_file),
-
-        # Mock TP world size = 2
-        monkeypatch.setattr(
-            "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
-            lambda: 2,
         )
 
+        class Pipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = nn.Module()
+                self.transformer.blocks = nn.ModuleList([nn.Linear(4, 1, bias=False)])
+                self.weights_sources = [
+                    SimpleNamespace(
+                        model_or_path=str(source_root),
+                        subfolder="transformer",
+                        revision=None,
+                        prefix="transformer.",
+                    )
+                ]
+
+        pipeline = Pipeline()
+        result = build_checkpoint_mmap_plan(
+            pipeline,
+            dit_modules=(("transformer", pipeline.transformer),),
+            sources=tuple(pipeline.weights_sources),
+            model_path=None,
+            tensor_parallel_size=1,
+            use_hsdp=False,
+            online_quantization=False,
+        )
+
+        assert result.fallback_reason is None
+        assert result.plan is not None
+        assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
+            checkpoint_key="blocks.0.weight",
+            file_path=str(checkpoint_file),
+        )
+        assert result.plan.planned_source_prefixes == frozenset({"transformer."})
+
+    def test_non_dedicated_component_source_falls_back_before_discovery(self):
+        pipeline = nn.Module()
+        pipeline.transformer = nn.Linear(2, 2, bias=False)
+        sources = (
+            SimpleNamespace(
+                model_or_path="unused",
+                subfolder=None,
+                revision=None,
+                prefix="",
+            ),
+        )
+
+        result = build_checkpoint_mmap_plan(
+            pipeline,
+            dit_modules=(("transformer", pipeline.transformer),),
+            sources=sources,
+            model_path=None,
+            tensor_parallel_size=1,
+            use_hsdp=False,
+            online_quantization=False,
+        )
+
+        assert result.plan is None
+        assert "dedicated component weight source" in result.fallback_reason
+
+    def test_preflight_falls_back_before_mutation_for_missing_model_tensors(self, tmp_path):
+        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):
+            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()
+        original_weight = pipeline.transformer.blocks[0].weight
+        result = build_checkpoint_mmap_plan(
+            pipeline,
+            dit_modules=(("transformer", pipeline.transformer),),
+            sources=tuple(pipeline.weights_sources),
+            model_path=None,
+            tensor_parallel_size=1,
+            use_hsdp=False,
+            online_quantization=False,
+        )
+
+        assert result.plan is None
+        assert "1 required DiT tensors" in result.fallback_reason
+        assert pipeline.transformer.blocks[0].weight is original_weight
+        assert not original_weight.is_meta
+
+    def test_tp_falls_back_before_checkpoint_discovery(self):
         class SimpleModel(nn.Module):
             _layerwise_offload_blocks_attrs = ["blocks"]
 
@@ -912,41 +1185,118 @@ class TestMmapValidation:
                 return key
 
         pipeline = SimplePipeline()
-        weights = {name: torch.ones(p.shape, dtype=p.dtype) for name, p in pipeline.named_parameters() if not p.is_meta}
+        result = build_checkpoint_mmap_plan(
-        save_file(weights, str(tmp_path / "model.safetensors"))
+            pipeline,
-        (tmp_path / "model.safetensors.index.json").write_text(
+            dit_modules=(("transformer", pipeline.transformer),),
-            json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
+            sources=(),
+            model_path="unused",
+            tensor_parallel_size=2,
+            use_hsdp=False,
+            online_quantization=False,
         )
 
-        backend = DistributedLayerwiseOffloadBackend(
+        assert result.plan is None
-            OffloadConfig(
+        assert result.fallback_reason == "TP=2 requires the ordinary loader"
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
-                model_path=str(tmp_path),
-            ),
-            torch.device("cpu"),
-        )
-        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
 
-        with pytest.raises(ValueError, match="Tensor Parallel"):
+    @pytest.mark.parametrize(
-            backend._load_weights_via_mmap(pipeline, modules)
+        ("use_hsdp", "online_quantization", "expected_reason"),
+        [
+            (True, False, "HSDP requires the ordinary loader"),
+            (False, True, "online quantization requires the ordinary loader"),
+        ],
+    )
+    def test_runtime_layouts_that_need_loader_callbacks_fall_back(
+        self,
+        use_hsdp,
+        online_quantization,
+        expected_reason,
+    ):
+        pipeline = nn.Module()
+        pipeline.transformer = nn.Linear(2, 2, bias=False)
 
-    def test_tp_allowed_when_tp_world_size_eq_1(self, tmp_path, patched_offload_runtime, monkeypatch):
+        result = build_checkpoint_mmap_plan(
-        """DLO+AllGather should work when TP world size = 1, even if params
+            pipeline,
-        have custom weight_loader attributes (e.g. QKVParallelLinear at TP=1)."""
+            dit_modules=(("transformer", pipeline.transformer),),
-        import json
+            sources=(),
-        from types import SimpleNamespace
+            model_path="unused",
-
+            tensor_parallel_size=1,
-        from safetensors.torch import save_file
+            use_hsdp=use_hsdp,
-
+            online_quantization=online_quantization,
-        from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
-
-        # Mock TP world size = 1 (default)
-        monkeypatch.setattr(
-            "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
-            lambda: 1,
         )
 
+        assert result.plan is None
+        assert result.fallback_reason == expected_reason
+
+    @pytest.mark.parametrize(
+        ("checkpoint_tensor", "expected_reason"),
+        [
+            (torch.ones((3, 2), dtype=torch.float32), "shape mismatch"),
+            (torch.ones((2, 2), dtype=torch.bfloat16), "dtype mismatch"),
+        ],
+    )
+    def test_source_metadata_mismatch_falls_back(self, tmp_path, checkpoint_tensor, expected_reason):
+        checkpoint_file = tmp_path / "model.safetensors"
+        save_file({"transformer.weight": checkpoint_tensor}, str(checkpoint_file))
+
+        class Pipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = nn.Linear(2, 2, bias=False)
+
+            @staticmethod
+            def _remap_ckpt_key(key):
+                return key
+
+        pipeline = Pipeline()
+
+        result = build_checkpoint_mmap_plan(
+            pipeline,
+            dit_modules=(("transformer", pipeline.transformer),),
+            sources=(),
+            model_path=str(tmp_path),
+            tensor_parallel_size=1,
+            use_hsdp=False,
+            online_quantization=False,
+        )
+
+        assert result.plan is None
+        assert expected_reason in result.fallback_reason
+
+    def test_non_persistent_buffers_do_not_need_checkpoint_bindings(self, tmp_path):
+        class Transformer(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.weight = nn.Parameter(torch.ones(2, 2))
+                self.register_buffer("derived", torch.ones(2), persistent=False)
+
+        checkpoint_file = tmp_path / "model.safetensors"
+        save_file({"transformer.weight": torch.ones(2, 2)}, str(checkpoint_file))
+
+        class Pipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+            @staticmethod
+            def _remap_ckpt_key(key):
+                return key
+
+        pipeline = Pipeline()
+
+        result = build_checkpoint_mmap_plan(
+            pipeline,
+            dit_modules=(("transformer", pipeline.transformer),),
+            sources=(),
+            model_path=str(tmp_path),
+            tensor_parallel_size=1,
+            use_hsdp=False,
+            online_quantization=False,
+        )
+
+        assert result.plan is not None
+        assert set(result.plan.bindings) == {"transformer.weight"}
+
+    def test_unadapted_custom_weight_loader_falls_back_at_tp1(self, tmp_path):
         def custom_weight_loader(param, weight):
             param.data.copy_(weight)
 
@@ -976,27 +1326,67 @@ class TestMmapValidation:
             json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
         )
 
-        backend = DistributedLayerwiseOffloadBackend(
+        result = build_checkpoint_mmap_plan(
-            OffloadConfig(
+            pipeline,
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            dit_modules=(("transformer", pipeline.transformer),),
-                pin_cpu_memory=False,
+            sources=(),
-                model_path=str(tmp_path),
+            model_path=str(tmp_path),
-            ),
+            tensor_parallel_size=1,
-            torch.device("cpu"),
+            use_hsdp=False,
+            online_quantization=False,
         )
-        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
 
-        # Should NOT raise — TP=1 is allowed even with custom weight_loader
+        assert result.plan is None
-        backend._load_weights_via_mmap(pipeline, modules)
+        assert "custom weight_loader" in result.fallback_reason
+
+    def test_cosmos3_adapter_preserves_existing_tp1_direct_mmap_contract(self, tmp_path):
+        from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+            Cosmos3OmniDiffusersPipeline,
+        )
+        from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+            Cosmos3VFMTransformer,
+        )
+
+        transformer = object.__new__(Cosmos3VFMTransformer)
+        nn.Module.__init__(transformer)
+        transformer.proj_in = nn.Linear(2, 2, bias=False)
+        transformer.proj_in.weight.weight_loader = lambda *_args: None
+        pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
+        nn.Module.__init__(pipeline)
+        pipeline.transformer = transformer
+
+        source_root = tmp_path / "partition"
+        checkpoint_dir = source_root / "transformer"
+        checkpoint_dir.mkdir(parents=True)
+        checkpoint_file = checkpoint_dir / "model.safetensors"
+        save_file({"proj_in.weight": torch.ones(2, 2)}, str(checkpoint_file))
+        sources = (
+            SimpleNamespace(
+                model_or_path=str(source_root),
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+            ),
+        )
+
+        result = build_checkpoint_mmap_plan(
+            pipeline,
+            dit_modules=(("transformer", transformer),),
+            sources=sources,
+            model_path=None,
+            tensor_parallel_size=1,
+            use_hsdp=False,
+            online_quantization=False,
+        )
+
+        assert result.plan is not None
+        assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
+            checkpoint_key="proj_in.weight",
+            file_path=str(checkpoint_file),
+        )
 
     def test_validate_loaded_weights_called(self, tmp_path, patched_offload_runtime):
         """validate_loaded_weights should be called after mmap loading."""
-        import json
-        from types import SimpleNamespace
-
-        from safetensors.torch import save_file
-
-        from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
 
         class ModelWithValidate(nn.Module):
             _layerwise_offload_blocks_attrs = ["blocks"]
@@ -1038,19 +1428,64 @@ class TestMmapValidation:
 
         backend = DistributedLayerwiseOffloadBackend(
             OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False, model_path=str(tmp_path)
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
             ),
             torch.device("cpu"),
         )
         modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+        result = build_checkpoint_mmap_plan(
+            pipeline,
+            dit_modules=(("transformer", pipeline.transformer),),
+            sources=(),
+            model_path=str(tmp_path),
+            tensor_parallel_size=1,
+            use_hsdp=False,
+            online_quantization=False,
+        )
 
-        backend._load_weights_via_mmap(pipeline, modules)
+        assert result.plan is not None
+        backend._load_weights_via_mmap(pipeline, modules, result.plan)
         assert pipeline.transformer.validate_called, "validate_loaded_weights should be called"
 
 
 class TestConfigValidation:
     """Tests for configuration validation in OffloadConfig / execute_request."""
 
+    def test_loader_plan_cannot_be_silently_dropped_on_unsupported_platform(self, monkeypatch):
+        import vllm_omni.diffusion.offloader as offloader_module
+
+        od_config = SimpleNamespace(
+            enable_cpu_offload=False,
+            enable_layerwise_offload=False,
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=False,
+            dlo_resident_layers=0,
+            pin_cpu_memory=False,
+            model="unused",
+            parallel_config=SimpleNamespace(
+                use_hsdp=False,
+                data_parallel_size=1,
+                sequence_parallel_size=1,
+            ),
+        )
+        plan = HostWeightPlan(
+            backing_kind="checkpoint_mmap",
+            bindings={},
+        )
+        monkeypatch.setattr(
+            offloader_module.current_omni_platform,
+            "supports_cpu_offload",
+            lambda: False,
+        )
+
+        with pytest.raises(RuntimeError, match="skipped ordinary weight materialization"):
+            offloader_module.get_offload_backend(
+                od_config,
+                device=torch.device("cpu"),
+                host_weight_plan=plan,
+            )
+
     def test_hsdp_with_allgather_rejected(self):
         """HSDP + DLO + AllGather should raise ValueError (double sharding)."""
         from vllm_omni.diffusion.offloader.base import OffloadConfig

tests/diffusion/test_diffusion_model_runner.py

@@ -740,6 +740,9 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
             del kwargs
             return SimpleNamespace(transformer=torch.nn.Identity())
 
+        def take_host_weight_plan(self):
+            return None
+
     class _DummyMemoryProfiler:
         consumed_memory = 0
 
@@ -779,7 +782,11 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
     monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
     monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
-    monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
+    monkeypatch.setattr(
+        model_runner_module,
+        "get_offload_backend",
+        lambda od_config, device, host_weight_plan: None,
+    )
     monkeypatch.setattr(
         model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
     )

vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py

@@ -9,6 +9,13 @@ from .modelopt import (
 )
 
 
+def get_direct_mmap_adapter(model: nn.Module):
+    """Resolve direct-mmap adapters lazily to avoid model-loader cycles."""
+    from .direct_mmap import get_direct_mmap_adapter as _get_direct_mmap_adapter
+
+    return _get_direct_mmap_adapter(model)
+
+
 def get_checkpoint_adapter(
     model: nn.Module,
     source: object,
@@ -29,4 +36,5 @@ __all__ = [
     "ModelOptMixedPrecisionCheckpointAdapter",
     "ModelOptNvFp4CheckpointAdapter",
     "get_checkpoint_adapter",
+    "get_direct_mmap_adapter",
 ]

vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py

@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Loader-side adapters for proven direct-checkpoint mmap layouts."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from functools import partial
+from typing import Protocol
+
+import torch
+from torch import nn
+
+
+@dataclass(frozen=True)
+class DirectMmapTensorPolicy:
+    """Loader-owned proof for one tensor's direct mmap behavior."""
+
+    allow_custom_loader: bool = False
+    transform: Callable[[torch.Tensor], torch.Tensor] | None = None
+
+
+class DirectMmapAdapter(Protocol):
+    def policy_for(
+        self,
+        runtime_name: str,
+        target: torch.Tensor,
+    ) -> DirectMmapTensorPolicy | None: ...
+
+
+class _MiniMaxH3DirectMmapAdapter:
+    """TP1 runtime-layout contract already enforced by MiniMax-H3's loader."""
+
+    def __init__(self, pipeline: nn.Module) -> None:
+        self.pipeline = pipeline
+
+    def policy_for(
+        self,
+        runtime_name: str,
+        target: torch.Tensor,
+    ) -> DirectMmapTensorPolicy:
+        del target
+        transform = None
+        suffix = ".qkv_proj.weight"
+        if runtime_name.endswith(suffix):
+            from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+                MiniMaxH3Attention,
+                _reorder_grouped_qkv_to_qkv,
+            )
+
+            attention_path = runtime_name[: -len(suffix)]
+            attention = self.pipeline.get_submodule(attention_path)
+            if not isinstance(attention, MiniMaxH3Attention):
+                raise ValueError(
+                    f"MiniMax-H3 direct-mmap adapter expected attention at {attention_path!r}, "
+                    f"got {type(attention).__name__}"
+                )
+            transform = partial(
+                _reorder_grouped_qkv_to_qkv,
+                num_query_groups=attention.total_num_heads,
+                heads_per_group=1,
+                head_dim=attention.head_dim,
+            )
+        return DirectMmapTensorPolicy(
+            allow_custom_loader=True,
+            transform=transform,
+        )
+
+
+class _Cosmos3DirectMmapAdapter:
+    """Preserve Cosmos3's existing TP1 direct-mmap load contract."""
+
+    def policy_for(
+        self,
+        runtime_name: str,
+        target: torch.Tensor,
+    ) -> DirectMmapTensorPolicy:
+        del runtime_name, target
+        # Cosmos3 has no packed module mapping or TP1 semantic transform after
+        # its checkpoint-key remap. Its custom vLLM loaders reduce to exact
+        # copies at TP1; topology and metadata are checked by plan preflight.
+        return DirectMmapTensorPolicy(allow_custom_loader=True)
+
+
+def get_direct_mmap_adapter(pipeline: nn.Module) -> DirectMmapAdapter | None:
+    """Return a loader-owned adapter without requiring a pipeline flag."""
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTModel,
+    )
+
+    if any(isinstance(module, MiniMaxH3DiTModel) for module in pipeline.modules()):
+        return _MiniMaxH3DirectMmapAdapter(pipeline)
+
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        Cosmos3VFMTransformer,
+    )
+
+    if any(isinstance(module, Cosmos3VFMTransformer) for module in pipeline.modules()):
+        return _Cosmos3DirectMmapAdapter()
+    return None
+
+
+__all__ = [
+    "DirectMmapAdapter",
+    "DirectMmapTensorPolicy",
+    "get_direct_mmap_adapter",
+]

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -6,7 +6,7 @@ import glob
 import os
 import re
 import time
-from collections.abc import Generator, Iterable
+from collections.abc import Generator, Iterable, Sequence
 from pathlib import Path
 from typing import cast
 
@@ -34,6 +34,11 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     get_checkpoint_adapter,
 )
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+    HostWeightPlan,
+    build_checkpoint_mmap_plan,
+    has_online_quantization,
+)
 from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.registry import initialize_model
@@ -135,6 +140,13 @@ class DiffusersPipelineLoader:
         self.od_config = od_config
         self.quant_config = od_config.quantization_config
         self.parallel_config = od_config.parallel_config
+        self.host_weight_plan: HostWeightPlan | None = None
+
+    def take_host_weight_plan(self) -> HostWeightPlan | None:
+        """Transfer the loader-produced plan to the offload backend."""
+        plan = self.host_weight_plan
+        self.host_weight_plan = None
+        return plan
 
     def _prepare_weights(
         self,
@@ -293,8 +305,10 @@ class DiffusersPipelineLoader:
     def get_all_weights(
         self,
         model: nn.Module,
+        sources: Sequence["ComponentSource"] | None = None,
     ) -> Generator[tuple[str, torch.Tensor], None, None]:
-        sources = self._get_weight_sources(model)
+        if sources is None:
+            sources = self._get_weight_sources(model)
         for source in sources:
             yield from self._get_weights_iterator(source, model=model)
 
@@ -386,6 +400,7 @@ class DiffusersPipelineLoader:
         device: torch.device | None = None,
     ) -> nn.Module:
         """Load a model with the given configurations."""
+        self.host_weight_plan = None
         if load_format is None:
             load_format = "default"
         # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8),
@@ -417,22 +432,54 @@ 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
-                # supports mmap loading and no online quantization is active.
-                # This condition MUST match the gate in
-                # DistributedLayerwiseOffloadBackend.enable() so that the
-                # loader skips ⟺ enable() uses mmap.
-                from vllm_omni.diffusion.offloader.offload_plan import 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))
+                _dp_size = int(getattr(self.parallel_config, "data_parallel_size", 1))
+                _sp_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1))
+                _dlo_group_size = _dp_size if _dp_size > 1 else _sp_size
 
-                _skip_load = _dist_offload and _use_ag and _supports_mmap and not _has_online_quant
+                plan_result = None
+                weight_sources = self._get_weight_sources(model)
+                if _dist_offload:
+                    modules = ModuleDiscovery.discover(model)
+                    plan_result = build_checkpoint_mmap_plan(
+                        model,
+                        dit_modules=tuple(zip(modules.dit_names, modules.dits)),
+                        sources=weight_sources,
+                        model_path=str(getattr(self.od_config, "model", "")) or None,
+                        tensor_parallel_size=_tp_size,
+                        use_hsdp=_use_hsdp,
+                        online_quantization=_has_online_quant,
+                    )
+                    self.host_weight_plan = plan_result.plan
+
+                _skip_load = self.host_weight_plan is not None
 
                 if _skip_load:
-                    logger.info("DLO+AllGather active: skipping load_weights (will load via mmap in enable())")
+                    logger.info(
+                        "DLO host-weight plan active (%s, %s): skipping ordinary materialization for %s",
+                        "AllGather" if _use_ag and _dlo_group_size > 1 else "rank-local",
+                        self.host_weight_plan.backing_kind,
+                        sorted(self.host_weight_plan.planned_source_prefixes) or "legacy DiT sources",
+                    )
+                    ordinary_sources = tuple(
+                        source
+                        for source in weight_sources
+                        if source.prefix not in self.host_weight_plan.planned_source_prefixes
+                    )
+                    if ordinary_sources:
+                        logger.info(
+                            "Loading %d component weight source(s) outside the DLO host-weight plan",
+                            len(ordinary_sources),
+                        )
+                        self.load_weights(
+                            model,
+                            sources=ordinary_sources,
+                            planned_weights=self.host_weight_plan.bindings,
+                        )
                 else:
                     if _dist_offload and _use_ag and _has_online_quant:
                         raise ValueError(
@@ -442,9 +489,10 @@ class DiffusersPipelineLoader:
                             "Please use --dlo-no-use-allgather or disable online "
                             "quantization."
                         )
-                    if _dist_offload and _use_ag and not _supports_mmap:
+                    if _dist_offload and plan_result is not None:
                         logger.info(
-                            "DLO+AllGather active but model does not support mmap: loading weights via regular loader."
+                            "DLO direct checkpoint mmap unavailable; using ordinary loader: %s",
+                            plan_result.fallback_reason,
                         )
                     logger.debug("Loading weights on %s ...", load_device)
                     if offload_after_quant:
@@ -498,11 +546,7 @@ 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():
+        return has_online_quantization(model)
-            quant_method = getattr(module, "quant_method", None)
-            if getattr(quant_method, "uses_meta_device", False):
-                return True
-        return False
 
     def _apply_skip_softmax_calibration(self, model: nn.Module) -> None:
         from vllm_omni.diffusion.attention.backends.trtllm_calibration import (
@@ -605,12 +649,16 @@ class DiffusersPipelineLoader:
         model: nn.Module,
         *,
         stream_online_quant_to_cpu: bool = False,
+        sources: Sequence["ComponentSource"] | None = None,
+        planned_weights: Iterable[str] = (),
     ) -> None:
         weights_to_load = self._get_expected_parameter_names(model)
-        weights = self.get_all_weights(model)
+        weights = self.get_all_weights(model) if sources is None else self.get_all_weights(model, sources=sources)
         if stream_online_quant_to_cpu:
             weights = self._stream_online_quant_weights_to_cpu(model, weights)
         loaded_weights = model.load_weights(weights)
+        if loaded_weights is not None:
+            loaded_weights = set(loaded_weights).union(planned_weights)
 
         self.counter_after_loading_weights = time.perf_counter()
         logger.info_once(

vllm_omni/diffusion/model_loader/host_weight_plan.py

@@ -0,0 +1,355 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Loader-owned host-weight plans shared with diffusion offload backends."""
+
+from __future__ import annotations
+
+import glob
+import json
+import os
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+
+import torch
+from safetensors import safe_open
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
+    get_direct_mmap_adapter,
+)
+
+logger = init_logger(__name__)
+
+TensorTransform = Callable[[torch.Tensor], torch.Tensor]
+
+
+@dataclass(frozen=True)
+class TensorBinding:
+    """One runtime tensor backed by one safetensors entry."""
+
+    checkpoint_key: str
+    file_path: str
+    transform: TensorTransform | None = None
+
+
+@dataclass(frozen=True)
+class HostWeightPlan:
+    """Complete, prevalidated host backing consumed by an offload backend."""
+
+    backing_kind: str
+    bindings: dict[str, TensorBinding]
+    planned_source_prefixes: frozenset[str] = frozenset()
+
+
+@dataclass(frozen=True)
+class HostWeightPlanResult:
+    """A complete plan or a fail-closed reason to use the ordinary loader."""
+
+    plan: HostWeightPlan | None
+    fallback_reason: str | None = None
+
+
+class _PlanIncompatibleError(RuntimeError):
+    pass
+
+
+_SAFETENSORS_DTYPES: dict[str, torch.dtype] = {
+    "BOOL": torch.bool,
+    "U8": torch.uint8,
+    "I8": torch.int8,
+    "I16": torch.int16,
+    "I32": torch.int32,
+    "I64": torch.int64,
+    "F16": torch.float16,
+    "BF16": torch.bfloat16,
+    "F32": torch.float32,
+    "F64": torch.float64,
+}
+for _safe_name, _torch_name in (
+    ("F8_E4M3", "float8_e4m3fn"),
+    ("F8_E5M2", "float8_e5m2"),
+):
+    if (_dtype := getattr(torch, _torch_name, None)) is not None:
+        _SAFETENSORS_DTYPES[_safe_name] = _dtype
+
+
+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 _resolve_source_root(source: object) -> str:
+    source_root = os.fspath(getattr(source, "model_or_path"))
+    subfolder = getattr(source, "subfolder", None)
+    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_root,
+            cache_dir=None,
+            allow_patterns=["*.safetensors", "*.safetensors.index.json"],
+            revision=getattr(source, "revision", None),
+            subfolder=subfolder,
+        )
+    return os.path.join(source_root, subfolder) if subfolder is not None else source_root
+
+
+def _record_binding(
+    model_to_ckpt: dict[str, tuple[str, str]],
+    model_name: str,
+    checkpoint_key: str,
+    file_path: str,
+) -> None:
+    existing = model_to_ckpt.get(model_name)
+    candidate = (checkpoint_key, file_path)
+    if existing is not None and existing != candidate:
+        raise _PlanIncompatibleError(
+            f"multiple checkpoint entries map to runtime tensor {model_name!r}: {existing!r} and {candidate!r}"
+        )
+    model_to_ckpt[model_name] = candidate
+
+
+def _add_safetensors_source(
+    checkpoint_dir: str,
+    prefix: str,
+    remap_fn: Callable[[str], str | None] | None,
+    model_to_ckpt: dict[str, tuple[str, str]],
+) -> int:
+    indexed_keys = 0
+    index_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors.index.json")))
+    if index_files:
+        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:
+                    _record_binding(
+                        model_to_ckpt,
+                        model_name,
+                        checkpoint_key,
+                        os.path.join(checkpoint_dir, filename),
+                    )
+                indexed_keys += 1
+        return indexed_keys
+
+    for filename in sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors"))):
+        with safe_open(filename, framework="pt", device="cpu") as handle:
+            for checkpoint_key in handle.keys():
+                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:
+                    _record_binding(model_to_ckpt, model_name, checkpoint_key, filename)
+                indexed_keys += 1
+    return indexed_keys
+
+
+def _build_source_map(
+    *,
+    sources: Sequence[object],
+    model_path: str | None,
+    remap_fn: Callable[[str], str | None] | None,
+) -> dict[str, tuple[str, str]]:
+    model_to_ckpt: dict[str, tuple[str, str]] = {}
+    indexed_keys = 0
+
+    for source in sources:
+        indexed_keys += _add_safetensors_source(
+            _resolve_source_root(source),
+            getattr(source, "prefix", ""),
+            remap_fn,
+            model_to_ckpt,
+        )
+
+    if not sources:
+        if remap_fn is None:
+            raise _PlanIncompatibleError(
+                "the loader has no component weight sources and the pipeline has no checkpoint-key remapper"
+            )
+        if not model_path:
+            raise _PlanIncompatibleError("the loader has no model path for legacy checkpoint remapping")
+
+        resolved_model_path = model_path
+        if not os.path.isdir(resolved_model_path):
+            from vllm.model_executor.model_loader.weight_utils import (
+                download_weights_from_hf,
+            )
+
+            resolved_model_path = download_weights_from_hf(
+                model_name_or_path=resolved_model_path,
+                cache_dir=None,
+                allow_patterns=["*.safetensors", "*.safetensors.index.json"],
+            )
+        checkpoint_dirs = {
+            os.path.dirname(path)
+            for pattern in ("*.safetensors.index.json", "*.safetensors")
+            for path in glob.glob(os.path.join(resolved_model_path, "**", pattern), recursive=True)
+        }
+        for checkpoint_dir in sorted(checkpoint_dirs):
+            indexed_keys += _add_safetensors_source(
+                checkpoint_dir,
+                "",
+                remap_fn,
+                model_to_ckpt,
+            )
+
+    if not model_to_ckpt:
+        raise _PlanIncompatibleError("no compatible safetensors entries were found in the loader's model sources")
+    logger.info(
+        "Indexed %d runtime tensor names from %d checkpoint keys for DLO preflight",
+        len(model_to_ckpt),
+        indexed_keys,
+    )
+    return model_to_ckpt
+
+
+def _is_persistent_buffer(module: nn.Module, local_name: str) -> bool:
+    parts = local_name.split(".")
+    owner = module.get_submodule(".".join(parts[:-1])) if len(parts) > 1 else module
+    return parts[-1] not in owner._non_persistent_buffers_set
+
+
+def _collect_required_dit_tensors(
+    dit_modules: Sequence[tuple[str, nn.Module]],
+) -> dict[str, torch.Tensor]:
+    required: dict[str, torch.Tensor] = {}
+    for dit_name, dit_module in dit_modules:
+        for local_name, tensor in dit_module.named_parameters():
+            required[f"{dit_name}.{local_name}"] = tensor
+        for local_name, tensor in dit_module.named_buffers():
+            if _is_persistent_buffer(dit_module, local_name):
+                required[f"{dit_name}.{local_name}"] = tensor
+    return required
+
+
+def _planned_source_prefixes(
+    sources: Sequence[object],
+    dit_modules: Sequence[tuple[str, nn.Module]],
+) -> frozenset[str]:
+    """Identify component sources that the DiT-only plan owns completely."""
+    if not sources:
+        return frozenset()
+
+    expected = frozenset(f"{dit_name}." for dit_name, _ in dit_modules)
+    available = {getattr(source, "prefix", "") for source in sources}
+    missing = sorted(expected - available)
+    if missing:
+        raise _PlanIncompatibleError(
+            f"direct mmap requires a dedicated component weight source for each DiT; missing prefixes: {missing}"
+        )
+    return expected
+
+
+def _validate_source_metadata(
+    required: dict[str, torch.Tensor],
+    bindings: dict[str, TensorBinding],
+) -> None:
+    by_file: dict[str, list[tuple[str, torch.Tensor, TensorBinding]]] = {}
+    for runtime_name, target in required.items():
+        by_file.setdefault(bindings[runtime_name].file_path, []).append((runtime_name, target, bindings[runtime_name]))
+
+    for file_path, entries in by_file.items():
+        if not os.path.isfile(file_path):
+            raise _PlanIncompatibleError(f"checkpoint file does not exist: {file_path}")
+        with safe_open(file_path, framework="pt", device="cpu") as handle:
+            available = set(handle.keys())
+            for runtime_name, target, binding in entries:
+                if binding.checkpoint_key not in available:
+                    raise _PlanIncompatibleError(
+                        f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}"
+                    )
+                source = handle.get_slice(binding.checkpoint_key)
+                source_shape = tuple(source.get_shape())
+                source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
+                if source_shape != tuple(target.shape):
+                    raise _PlanIncompatibleError(
+                        f"shape mismatch for {runtime_name!r}: checkpoint={source_shape}, runtime={tuple(target.shape)}"
+                    )
+                if source_dtype is None or source_dtype != target.dtype:
+                    raise _PlanIncompatibleError(
+                        f"dtype mismatch for {runtime_name!r}: checkpoint={source.get_dtype()}, runtime={target.dtype}"
+                    )
+
+
+def build_checkpoint_mmap_plan(
+    pipeline: nn.Module,
+    *,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+    sources: Sequence[object],
+    model_path: str | None,
+    tensor_parallel_size: int,
+    use_hsdp: bool,
+    online_quantization: bool,
+) -> HostWeightPlanResult:
+    """Build a complete direct-checkpoint plan or return a fallback reason."""
+    if tensor_parallel_size != 1:
+        return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader")
+    if use_hsdp:
+        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
+    if online_quantization:
+        return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
+
+    remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
+    if not callable(remap_fn):
+        remap_fn = None
+    adapter = get_direct_mmap_adapter(pipeline)
+
+    try:
+        planned_source_prefixes = _planned_source_prefixes(sources, dit_modules)
+        planned_sources = tuple(
+            source for source in sources if getattr(source, "prefix", "") in planned_source_prefixes
+        )
+        model_to_ckpt = _build_source_map(
+            sources=planned_sources,
+            model_path=model_path,
+            remap_fn=remap_fn,
+        )
+        required = _collect_required_dit_tensors(dit_modules)
+        if not required:
+            raise _PlanIncompatibleError("no DiT parameters or persistent buffers were discovered")
+
+        missing = sorted(set(required) - set(model_to_ckpt))
+        if missing:
+            raise _PlanIncompatibleError(
+                f"{len(missing)} required DiT tensors have no checkpoint binding (first 5: {missing[:5]})"
+            )
+
+        bindings: dict[str, TensorBinding] = {}
+        for runtime_name, target in required.items():
+            policy = adapter.policy_for(runtime_name, target) if adapter is not None else None
+            has_custom_loader = callable(getattr(target, "weight_loader", None))
+            if has_custom_loader and not (policy and policy.allow_custom_loader):
+                raise _PlanIncompatibleError(
+                    f"{runtime_name!r} requires a custom weight_loader with no direct-mmap adapter"
+                )
+            checkpoint_key, file_path = model_to_ckpt[runtime_name]
+            bindings[runtime_name] = TensorBinding(
+                checkpoint_key=checkpoint_key,
+                file_path=file_path,
+                transform=policy.transform if policy is not None else None,
+            )
+
+        _validate_source_metadata(required, bindings)
+    except (OSError, ValueError, _PlanIncompatibleError) as exc:
+        return HostWeightPlanResult(None, str(exc))
+
+    return HostWeightPlanResult(
+        HostWeightPlan(
+            backing_kind="checkpoint_mmap",
+            bindings=bindings,
+            planned_source_prefixes=planned_source_prefixes,
+        )
+    )
+
+
+__all__ = [
+    "HostWeightPlan",
+    "HostWeightPlanResult",
+    "TensorBinding",
+    "build_checkpoint_mmap_plan",
+    "has_online_quantization",
+]

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -540,13 +540,6 @@ 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
-    # 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",

vllm_omni/diffusion/offloader/__init__.py

@@ -5,6 +5,7 @@ import torch
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan
 from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig, OffloadStrategy
@@ -17,7 +18,7 @@ from .distributed_layerwise_backend import (
 )
 from .layerwise_backend import LayerWiseOffloadBackend
 from .module_residency import PinnedModuleStager
-from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading
+from .offload_plan import OffloadPlan, get_offload_plan
 from .sequential_backend import (
     ModelLevelOffloadBackend,
     apply_sequential_offload,
@@ -49,7 +50,6 @@ __all__ = [
     "remove_distributed_block_hook",
     "get_offload_backend",
     "get_offload_plan",
-    "supports_mmap_loading",
     "get_blocks_attr_names",
     "get_blocks_from_dit",
     "set_blocks_attr_names",
@@ -64,12 +64,15 @@ __all__ = [
 def get_offload_backend(
     od_config: OmniDiffusionConfig,
     device: torch.device | None = None,
+    host_weight_plan: HostWeightPlan | None = None,
 ) -> OffloadBackend | None:
     """Create appropriate offload backend based on configuration.
 
     Args:
         od_config: OmniDiffusionConfig with offload settings
         device: Target device (auto-detected if None)
+        host_weight_plan: Exact loader-produced backing plan, if ordinary
+            weight materialization was skipped.
 
     Returns:
         OffloadBackend instance or None if offloading disabled
@@ -82,12 +85,22 @@ def get_offload_backend(
     # Extract and validate configuration
     config = OffloadConfig.from_od_config(od_config)
 
+    if host_weight_plan is not None and config.strategy != OffloadStrategy.DISTRIBUTED_LAYER_WISE:
+        raise RuntimeError(
+            "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
+        )
+
     # Return None if no offloading requested
     if config.strategy == OffloadStrategy.NONE:
         return None
 
     # Validate platform (CUDA required for now)
     if not current_omni_platform.supports_cpu_offload() or current_omni_platform.get_device_count() < 1:
+        if host_weight_plan is not None:
+            raise RuntimeError(
+                "The loader skipped ordinary weight materialization for DLO, "
+                "but this platform cannot create the required offload backend"
+            )
         logger.warning(
             "Current device: %s does not support CPU offloading. Skipping offloading.",
             current_omni_platform.get_device_name(),
@@ -99,6 +112,11 @@ def get_offload_backend(
         try:
             device = current_omni_platform.get_torch_device()
         except (NotImplementedError, AttributeError) as exc:
+            if host_weight_plan is not None:
+                raise RuntimeError(
+                    "The loader skipped ordinary weight materialization for DLO, "
+                    "but the target device could not be resolved"
+                ) from exc
             logger.error("Failed to detect device: %s. Skipping offloading.", exc)
             return None
 
@@ -108,7 +126,11 @@ def get_offload_backend(
     elif config.strategy == OffloadStrategy.LAYER_WISE:
         return LayerWiseOffloadBackend(config, device)
     elif config.strategy == OffloadStrategy.DISTRIBUTED_LAYER_WISE:
-        return DistributedLayerwiseOffloadBackend(config, device)
+        return DistributedLayerwiseOffloadBackend(
+            config,
+            device,
+            host_weight_plan=host_weight_plan,
+        )
     else:
         logger.error("Unknown offload strategy: %s", config.strategy)
         return None

vllm_omni/diffusion/offloader/base.py

@@ -27,11 +27,10 @@ class OffloadConfig:
     pin_cpu_memory: bool = True
     use_hsdp: bool = False
     dp_size: int = 1  # derived from parallel_config, not user-configurable
-    # True: add DP sharding + AllGather. False: stream the standard loader's
+    # True: add DP sharding + AllGather. False: stream complete rank-local
-    # rank-local tensors (including TP-local shards) with H2D only.
+    # blocks from the loader-selected host backing 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
     def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -58,7 +57,6 @@ class OffloadConfig:
 
         parallel_config = getattr(od_config, "parallel_config", None)
         use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False
-
         # Derive dp_size from parallel_config — not user-configurable.
         # The offload adapts to whatever DP/SP is already configured.
         dp_size = 1
@@ -116,7 +114,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
@@ -137,7 +136,6 @@ class OffloadConfig:
             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/offload_plan.py

@@ -55,20 +55,3 @@ class OffloadPlan:
 def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None:
     """Retrieve the OffloadPlan declared by the pipeline, if any."""
     return getattr(pipeline, "_offload_plan", 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``,
-    which maps checkpoint keys to model parameter names.  Models without
-    this method must use 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(
-        callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules()
-    )

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -292,7 +292,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             )
 
         # Apply CPU offloading
-        self.offload_backend = get_offload_backend(self.od_config, device=self.device)
+        self.offload_backend = get_offload_backend(
+            self.od_config,
+            device=self.device,
+            host_weight_plan=model_loader.take_host_weight_plan(),
+        )
         if self.offload_backend is not None:
             logger.info(f" Enabling offloader backend: {self.offload_backend.__class__.__name__}")
             self.offload_backend.enable(self.pipeline)