From 1c3d59796f92b89555a4e7ecd9df73186651dffc Mon Sep 17 00:00:00 2001
From: Hongsheng Liu <liuhongsheng4@huawei.com>
Date: Sat, 15 Aug 2026 06:57:38 +0800
Subject: [PATCH] [Diffusion] Document DLO mmap model opt-in

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
---
 .../feature/distributed_layerwise_offload.md  | 23 +++++++++++++++++++
 .../models/minimax_h3/pipeline_minimax_h3.py  |  8 ++++---
 vllm_omni/diffusion/offloader/base.py         |  4 +++-
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/docs/design/feature/distributed_layerwise_offload.md b/docs/design/feature/distributed_layerwise_offload.md
index 0c0cca0f8..48d433914 100644
--- a/docs/design/feature/distributed_layerwise_offload.md
+++ b/docs/design/feature/distributed_layerwise_offload.md
@@ -153,6 +153,29 @@ must preserve dtype and element count so it can be applied while packing one
 staging block. MiniMax-H3 uses this mechanism for grouped QKV. Online
 quantization and TP/HSDP no-AllGather execution use the regular loader path.
 
+### Adding mmap support to another model
+
+A new pipeline should opt in only after its ordinary loader contract can be
+represented without materializing a private full-model copy:
+
+1. Populate `weights_sources` with safetensors component sources and prefixes
+   that place every DiT checkpoint key in the pipeline parameter namespace.
+2. If prefixed checkpoint keys already equal runtime parameter names, declare
+   `_supports_mmap_loading = True`. If names differ, provide
+   `_remap_ckpt_key` to map checkpoint names to runtime names (or return
+   `None` for intentionally ignored keys).
+3. If a parameter needs a TP1 layout conversion, attach a callable
+   `mmap_weight_transform` to that parameter during model construction. The
+   transform must preserve dtype and element count because it runs while one
+   bounded block is packed.
+4. Verify that mmap and the regular loader produce the same runtime tensors,
+   that every expected DiT tensor maps successfully, and that TP, HSDP, and
+   online-quantized configurations retain their standard-loader behavior.
+
+Setting `_supports_mmap_loading = True` alone is not sufficient: explicit
+opt-ins fail during initialization if any expected DiT tensor remains on the
+`meta` device after checkpoint mapping.
+
 ## Validation coverage
 
 Current source-level validation includes:
diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py
index fc11632d8..31d34b266 100644
--- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py
+++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py
@@ -532,9 +532,11 @@ class MiniMaxH3Pipeline(
         encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")},
         on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}),
     )
-    # TP1 DLO mmap applies the grouped-QKV adapter while packing each block.
-    # The fused MLP checkpoint is already gate-then-up, matching its TP1
-    # runtime layout. TP>1 and quantized layouts stay on the regular loader.
+    # Explicit mmap opt-in: weights_sources below provide exact pipeline
+    # prefixes, and TP1 DLO mmap applies the grouped-QKV adapter while packing
+    # each block. The fused MLP checkpoint is already gate-then-up, matching
+    # its TP1 runtime layout. TP>1 and quantized layouts stay on the regular
+    # loader. See the DLO design doc for the model opt-in contract.
     _supports_mmap_loading: ClassVar[bool] = True
     _PROFILER_TARGETS: ClassVar[list[str]] = [
         "_prepare_reference_videos",
diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py
index bed80256a..68860ee46 100644
--- a/vllm_omni/diffusion/offloader/base.py
+++ b/vllm_omni/diffusion/offloader/base.py
@@ -32,7 +32,9 @@ class OffloadConfig:
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0  # leading DiT layers kept on device
     model_path: str | None = None  # checkpoint path for mmap weight loading
-    # Appended to preserve the positional order of the pre-existing fields.
+    # Rank-local mmap bypasses TP-aware loader callbacks, so the backend needs
+    # the configured TP size when selecting mmap versus standard-loader host
+    # storage. Appended to preserve the positional order of existing fields.
     tensor_parallel_size: int = 1
 
     @classmethod