[Diffusion] Register runtime cache for direct DLO H2D

Avoid recurrent mmap-to-pinned packing in no-AllGather DLO by registering transform-complete runtime-cache mappings under an explicit per-worker budget. Copy mapped tensors directly into the rotating device buffers while preserving bounded-staging fallback and ordered shutdown cleanup.

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

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -7,7 +7,9 @@ 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. The optional
 runtime-weight cache extends node-local sharing to final ordinary-loader
-layouts, including matching TP coordinates across independent engines.
+layouts, including matching TP coordinates across independent engines. On
+CUDA, an explicit registration budget can make those shared mappings directly
+H2D-copyable and avoid recurrent mmap-to-pinned host packing.
 
 See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
 for the implementation contract and compatibility matrix.
@@ -48,7 +50,8 @@ vllm serve /path/to/model --omni \
   --tensor-parallel-size 2 \
   --dlo-no-use-allgather \
   --dlo-enable-runtime-cache \
-  --dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights
+  --dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights \
+  --dlo-runtime-cache-pin-limit-gib 80
 
 # Sequence parallel deployment
 vllm serve /path/to/model --omni \
@@ -65,6 +68,7 @@ omni = Omni(
     dlo_use_allgather=False,
     dlo_enable_runtime_cache=True,
     dlo_runtime_cache_dir="/var/cache/vllm-omni/dlo-runtime-weights",
+    dlo_runtime_cache_pin_limit_gib=80,
 )
 ```
 
@@ -79,6 +83,7 @@ omni = Omni(
 | `--dlo-enable-runtime-cache` | Normalize final ordinary-loader DiT weights into a node-local mmap cache; requires no-AllGather | `false` |
 | `--dlo-runtime-cache-dir PATH` | Shared local-disk cache root | `~/.cache/vllm-omni/dlo-runtime-weights` |
 | `--dlo-runtime-cache-lock-timeout SECONDS` | Maximum wait for another process publishing the same layout | `600` |
+| `--dlo-runtime-cache-pin-limit-gib GIB` | Per-worker limit for registering the complete read-only runtime-cache mapping for direct H2D; zero keeps bounded host staging | `0` |
 | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
 
 ## Host-weight loading
@@ -95,14 +100,21 @@ 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.
 
-When `--dlo-enable-runtime-cache` is set with no-AllGather, that fallback gains
-a second storage opportunity. Every rank still completes ordinary loading,
-loader callbacks, post-load processing, validation, and calibration. It then
-hashes the final CPU DiT parameters and persistent buffers. One process per
-equivalent runtime layout publishes sharded safetensors plus a checksummed
-manifest; peers validate and mmap the same files. DLO replaces only tensor
-storage, preserves the final `Parameter` objects, and releases the private
-allocator-backed copies before installing its streaming hooks.
+When `--dlo-enable-runtime-cache` is set with no-AllGather, ordinary loading
+gains a second storage opportunity. Every rank completes loader callbacks,
+post-load processing, validation, and calibration. It then hashes the final
+CPU DiT parameters and persistent buffers. One process per equivalent runtime
+layout publishes sharded safetensors plus a checksummed manifest; peers
+validate and mmap the same files. DLO replaces only tensor storage, preserves
+the final `Parameter` objects, and releases the private allocator-backed copies
+before installing its streaming hooks.
+
+A positive `--dlo-runtime-cache-pin-limit-gib` deliberately selects this final
+runtime layout even when a direct TP1 checkpoint plan is available. Raw
+checkpoint mappings may require deferred loader transforms such as grouped-QKV
+reordering; registering them would still leave recurrent host work. Publishing
+the fully transformed layout once makes every streamed source directly
+copyable.
 
 The runtime layout identity excludes DP rank and SP rank. It includes TP world
 size/rank and conservative SP implementation/world-size guards. Therefore two
@@ -110,13 +122,27 @@ independent `TP=2` engines normally create two entries: both TP0 workers map
 one entry and both TP1 workers map the other. Cache sharing does not create a
 process group or an inference-time collective.
 
-At inference time, cached tensors remain pageable mmap views. DLO packs each
-streamed block into one of two bounded pinned host slots before the asynchronous
-H2D copy. The private ordinary-loader path instead creates model-sized pinned
-storage once during initialization. The cache therefore exchanges recurrent
-host-to-host copies and DDR bandwidth for lower steady-state PSS; competing TP
-ranks or engines can turn staging skew into collective wait time. Treat this as
-a host-capacity mode and benchmark it on the target CPU/memory topology.
+With a zero registration budget, cached tensors remain pageable mmap views.
+DLO packs each streamed block into one of two bounded pinned host slots before
+the asynchronous H2D copy. This lowers steady-state PSS but adds recurrent
+host-to-host copies; competing TP ranks or engines can turn DDR contention and
+staging skew into collective wait time.
+
+With a positive budget on CUDA, DLO instead attempts to register the complete
+page-aligned cache mapping as read-only pinned host memory. Registration is
+all-or-nothing: the total mapped range must fit the per-worker budget, and any
+CUDA registration or verification failure rolls back partial registrations and
+uses the bounded staging path. If CUDA cannot roll back a partial registration,
+worker initialization aborts instead of unmapping memory that CUDA may still
+own. On success, each immutable tensor view is copied directly into the
+existing rotating device buffers on the H2D stream, so no private host staging
+slots are allocated.
+
+CUDA registration is process-local, but it does not clone file contents. The
+same physical page-cache pages remain shared among equivalent workers; each
+worker does, however, hold its own registration reference and must satisfy the
+host's page-locking limits. Shutdown synchronizes outstanding GPU work,
+unregisters every range, and only then closes the mmap handles.
 
 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
@@ -165,6 +191,13 @@ entry is published. Operators may remove entries or the complete cache root
 only after all workers using its mmap files have stopped. The feature is
 opt-in for this reason.
 
+For direct H2D, set the registration limit above the complete page-aligned
+layout size reported in the worker log. Partial registration is intentionally
+unsupported because mixing direct and staged tensors would complicate buffer
+lifetime and make performance topology-dependent. Registration is CUDA-only;
+an insufficient budget, unsupported runtime, OS/CUDA pinning limit, or
+registration error is reported and falls back to bounded staging.
+
 ## Declarative topology
 
 Models may declare an `OffloadPlan` instead of embedding offload logic:
@@ -196,7 +229,9 @@ must enter each collective.
 - Direct checkpoint mmap currently requires TP1. TP greater than one falls
   back to the ordinary TP-aware loader. With the optional runtime cache,
   equivalent processes at the same TP coordinate can then share that final
-  rank-local layout in no-AllGather mode.
+  rank-local layout in no-AllGather mode. A positive registration budget also
+  chooses the final runtime cache at TP1 so deferred checkpoint transforms are
+  not repeated during inference.
 - HSDP plus AllGather is rejected to avoid double sharding. HSDP without
   AllGather has limited end-to-end validation.
 - Per-tensor online FP8 linears use the ordinary loader and can run with either
@@ -212,16 +247,19 @@ must enter each collective.
 
 The runtime cache improves steady-state host PSS, not startup peak: every rank
 loads private weights before remapping, and validation adds full content-hash
-passes. Skip-load-on-cache-hit and cache eviction remain follow-up work in
+passes. Direct H2D removes recurrent staging but does not change that startup
+contract. Skip-load-on-cache-hit and cache eviction remain follow-up work in
 [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
 
-It is also not latency-neutral. In the four-L20X MiniMax-H3 smoke documented in
-the L20X runtime-cache matrix in the
+Zero-budget staging is also not latency-neutral. In the four-L20X MiniMax-H3
+smoke documented in the L20X runtime-cache matrix in the
 [feature design](../../../design/feature/offloader/distributed_layerwise_offload.md#l20x-runtime-cache-smoke-matrix),
 two independent TP2 engines reduced request-time host PSS by 36.6% with the
 cache, but recurrent mmap-to-pinned staging reduced combined throughput by
-69.3%. HBM and GPU compute were unchanged. Use the cache only when that
-measured memory/latency tradeoff is appropriate for the deployment.
+69.3%. The registered path recovered near-private per-engine latency while
+retaining the shared-memory benefit; a matched TP1 run improved mean wave
+latency from 33.87 seconds to 4.88 seconds with byte-identical outputs. See the
+feature design for the full smoke-test caveats and metrics.
 
 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.