[Diffusion] Add rank-local mmap storage for DLO

Decouple node-local checkpoint storage from the DLO AllGather transfer protocol. Compatible rank-local workers retain shared file-backed sources and stream complete blocks through bounded staging without changing request scheduling.

Refs #6195

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

docs/design/feature/distributed_layerwise_offload.md

@@ -5,6 +5,8 @@ models. DLO keeps only a small number of DiT blocks on the accelerator and
 streams the remaining blocks from host memory. The distributed backend can
 either shard those host-side weights across an existing parallel group or keep
 the standard loader's rank-local weights and avoid an additional collective.
+The rank-local mmap storage design is tracked in
+[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
 
 For user-facing commands, see the
 [distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
@@ -16,8 +18,10 @@ DLO is implemented for multi-device diffusion execution. The default
 AllGather path is the primary path for DP and SP deployments. The
 `--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
 for standard-loader sharding, workstation bring-up, and systems where an
-additional DLO collective is undesirable, but it does not reduce host weight
-storage across ranks.
+additional DLO collective is undesirable. For explicitly mmap-compatible
+models in a TP1, non-HSDP topology, this path shares file-backed DiT checkpoint
+pages across processes and keeps only bounded staging buffers private. TP,
+HSDP, and other models retain standard-loader rank-local host storage.
 
 The compatibility matrix below describes the current implementation. The
 unit-level guards are covered, but not every parallelism combination has a
@@ -68,9 +72,18 @@ requests must take the same execution path at every denoising step.
 ### Rank-local path without DLO AllGather
 
 With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
-one. The regular model loader remains responsible for preparing each rank's
-weights, including TP-local tensors or HSDP-managed parameters. DLO then
-streams those rank-local tensors block by block using H2D copies only.
+one and streams complete runtime blocks using H2D copies only. Host storage is
+selected independently from that transfer protocol:
+
+- For an explicitly mmap-compatible model with TP1 and no HSDP, each worker
+  retains immutable safetensors views backed by the node's shared OS page
+  cache. Two pinned host buffers, sized to the largest streamed block, stage
+  mmap pages for asynchronous H2D. Checkpoint-to-runtime adapters run while a
+  block is packed, so transformed weights do not become a persistent private
+  model copy.
+- With TP, HSDP, online quantization, or a model without mmap support, the
+  regular loader remains responsible for preparing rank-local tensors. This
+  preserves TP/HSDP/quantization loader callbacks.
 
 This mode means:
 
@@ -80,19 +93,26 @@ This mode means:
   not shard weights across SP ranks.
 - TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
   DLO's additional weight AllGather is disabled.
-- Pure DP deployments keep a full host-side model copy per rank, subject to
-  shared OS page-cache behavior.
+- Pure DP deployments with a compatible mmap checkpoint share the streamed
+  DiT's file-backed pages. Other components and fallback paths may still keep
+  rank-private host weights.
 - The scheduler does not require a synchronized DP request wave for DLO.
 
-This path is intentionally not implemented by reusing the AllGather mmap
-loader. It relies on the standard loader so model-specific TP/HSDP transforms
-remain intact.
+The mmap source remains immutable in both transfer modes. AllGather copies a
+persistent `1 / group_size` shard from it and then closes the mappings;
+rank-local mode keeps the mappings open and copies a complete block through
+bounded staging without adding a collective.
+
+The current implementation packs mmap views into a staging slot synchronously
+on the host; the subsequent pinned H2D transfer overlaps with block compute.
+This keeps host memory bounded but makes CPU packing bandwidth part of the
+rank-local path's performance profile.
 
 ## Parallelism compatibility
 
 | Parallelism | DLO + AllGather | DLO without AllGather |
 |---|---|---|
-| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
+| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight collective. Compatible TP1 models share file-backed DiT pages across replicas. |
 | **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
 | **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
 | **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
@@ -122,11 +142,12 @@ 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
-model has a compatible checkpoint layout. Online quantization is incompatible
-with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
-quantization. Models that do not meet the mmap requirements use the regular
-loader path.
+The mmap loader is used when the model explicitly declares a compatible
+checkpoint layout. It is independent of AllGather for TP1/non-HSDP execution.
+Models can declare a per-parameter checkpoint-to-runtime adapter; the adapter
+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.
 
 ## Validation coverage
 
@@ -136,6 +157,8 @@ Current source-level validation includes:
 - HSDP + DLO without AllGather acceptance at configuration level;
 - TP rejection in the DLO+AllGather mmap path;
 - resident-layer requests requiring no-AllGather;
+- rank-local mmap source retention, layout transforms, and bounded host
+  staging;
 - DP request-wave validation for denoising-step compatibility;
 - sharding, double-buffer, AllGather-size, and heterogeneous-block regression
   tests.
@@ -150,7 +173,8 @@ on the target CUDA/NCCL or CANN/HCCL hardware.
   scaling path.
 - Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
   not the goal.
-- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
-  the expectation of higher host-memory use and lower validation confidence.
+- Use **no-AllGather** when replicas must avoid DLO collectives. TP1 models with
+  mmap support can share DiT checkpoint pages; TP/HSDP and unsupported models
+  retain the higher-memory standard-loader path.
 - Prefer **HSDP alone** for production HSDP deployments until the combined
   HSDP + DLO no-AllGather path has broader end-to-end coverage.